Revert "Reland: [MLIR][Transforms] Fix Mem2Reg removal order to respect dominance (#68767)"

This reverts commit 59fec73.

This reland did not properly fix the partial order.

Author: Dinistro

mlir/lib/Transforms/Mem2Reg.cpp

@@ -96,9 +96,6 @@ using namespace mlir;
 
 namespace {
 
-using BlockingUsesMap =
-    llvm::MapVector<Operation *, SmallPtrSet<OpOperand *, 4>>;
-
 /// Information computed during promotion analysis used to perform actual
 /// promotion.
 struct MemorySlotPromotionInfo {
@@ -109,7 +106,7 @@ struct MemorySlotPromotionInfo {
   /// its uses, it is because the defining ops of the blocking uses requested
   /// it. The defining ops therefore must also have blocking uses or be the
   /// starting point of the bloccking uses.
-  BlockingUsesMap userToBlockingUses;
+  DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> userToBlockingUses;
 };
 
 /// Computes information for basic slot promotion. This will check that direct
@@ -132,7 +129,8 @@ class MemorySlotPromotionAnalyzer {
   /// uses (typically, removing its users because it will delete itself to
   /// resolve its own blocking uses). This will fail if one of the transitive
   /// users cannot remove a requested use, and should prevent promotion.
-  LogicalResult computeBlockingUses(BlockingUsesMap &userToBlockingUses);
+  LogicalResult computeBlockingUses(
+      DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> &userToBlockingUses);
 
   /// Computes in which blocks the value stored in the slot is actually used,
   /// meaning blocks leading to a load. This method uses `definingBlocks`, the
@@ -235,7 +233,7 @@ Value MemorySlotPromoter::getLazyDefaultValue() {
 }
 
 LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
-    BlockingUsesMap &userToBlockingUses) {
+    DenseMap<Operation *, SmallPtrSet<OpOperand *, 4>> &userToBlockingUses) {
   // The promotion of an operation may require the promotion of further
   // operations (typically, removing operations that use an operation that must
   // delete itself). We thus need to start from the use of the slot pointer and
@@ -245,7 +243,7 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
   // use it.
   for (OpOperand &use : slot.ptr.getUses()) {
     SmallPtrSet<OpOperand *, 4> &blockingUses =
-        userToBlockingUses[use.getOwner()];
+        userToBlockingUses.getOrInsertDefault(use.getOwner());
     blockingUses.insert(&use);
   }
 
@@ -283,7 +281,7 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
       assert(llvm::is_contained(user->getResults(), blockingUse->get()));
 
       SmallPtrSetImpl<OpOperand *> &newUserBlockingUseSet =
-          userToBlockingUses[blockingUse->getOwner()];
+          userToBlockingUses.getOrInsertDefault(blockingUse->getOwner());
       newUserBlockingUseSet.insert(blockingUse);
     }
   }
@@ -518,21 +516,14 @@ void MemorySlotPromoter::computeReachingDefInRegion(Region *region,
 }
 
 void MemorySlotPromoter::removeBlockingUses() {
-  llvm::SmallVector<Operation *> usersToRemoveUses(
-      llvm::make_first_range(info.userToBlockingUses));
-
-  // The uses need to be traversed in *reverse dominance* order to ensure that
-  // transitive replacements are performed correctly.
-  // NOTE: The order can be non-deterministic, due to a pointer comparision, but
-  // this has no effect on the result of the pattern. This is necessary to get a
-  // strict weak order relation.
-  llvm::sort(usersToRemoveUses, [&](Operation *lhs, Operation *rhs) {
-    return dominance.properlyDominates(rhs, lhs) ||
-           (!dominance.properlyDominates(lhs, rhs) && rhs < lhs);
-  });
+  llvm::SetVector<Operation *> usersToRemoveUses;
+  for (auto &user : llvm::make_first_range(info.userToBlockingUses))
+    usersToRemoveUses.insert(user);
+  SetVector<Operation *> sortedUsersToRemoveUses =
+      mlir::topologicalSort(usersToRemoveUses);
 
   llvm::SmallVector<Operation *> toErase;
-  for (Operation *toPromote : usersToRemoveUses) {
+  for (Operation *toPromote : llvm::reverse(sortedUsersToRemoveUses)) {
     if (auto toPromoteMemOp = dyn_cast<PromotableMemOpInterface>(toPromote)) {
       Value reachingDef = reachingDefs.lookup(toPromoteMemOp);
       // If no reaching definition is known, this use is outside the reach of

mlir/test/Dialect/LLVMIR/mem2reg.mlir

@@ -683,16 +683,3 @@ llvm.func @no_inner_alloca_promotion(%arg: i64) -> i64 {
   // CHECK: llvm.return %[[RES]] : i64
   llvm.return %2 : i64
 }
-
-// -----
-
-// CHECK-LABEL: @transitive_reaching_def
-llvm.func @transitive_reaching_def() -> !llvm.ptr {
-  %0 = llvm.mlir.constant(1 : i32) : i32
-  // CHECK-NOT: alloca
-  %1 = llvm.alloca %0 x !llvm.ptr {alignment = 8 : i64} : (i32) -> !llvm.ptr
-  %2 = llvm.load %1 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
-  llvm.store %2, %1 {alignment = 8 : i64} : !llvm.ptr, !llvm.ptr
-  %3 = llvm.load %1 {alignment = 8 : i64} : !llvm.ptr -> !llvm.ptr
-  llvm.return %3 : !llvm.ptr
-}