[Mlir-commits] [mlir] Revert "[MLIR][Transforms] Fix Mem2Reg removal order to respect dominance (#68687)" (PR #68732)

[Balaji V. Iyer.]

https://github.com/bviyer created https://github.com/llvm/llvm-project/pull/68732

This commit causes the following issue with sanitizers:

`include/c++/v1/__debug_utils/strict_weak_ordering_check.h:52: assertion !__comp(*(__first + __b), *(__first + __a)) failed: Your comparator is not a valid strict-weak ordering`

probably due to an invalid sort().

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

This reverts commit be81f42b551c8b3c520132c3d60bc19cfc1c72fb.

>From 21a6756c5c9f8a221f4b34acc162efa6a9e15a15 Mon Sep 17 00:00:00 2001
From: "Balaji V. Iyer" <bviyer at gmail.com>
Date: Tue, 10 Oct 2023 18:26:29 +0000
Subject: [PATCH] This commit causes the following issue with sanitizers:

`include/c++/v1/__debug_utils/strict_weak_ordering_check.h:52: assertion !__comp(*(__first + __b), *(__first + __a)) failed: Your comparator is not a valid strict-weak ordering`

I think it is due to an invalid sort().

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

This reverts commit be81f42b551c8b3c520132c3d60bc19cfc1c72fb.
---
 mlir/lib/Transforms/Mem2Reg.cpp       | 28 ++++++++++++---------------
 mlir/test/Dialect/LLVMIR/mem2reg.mlir | 13 -------------
 2 files changed, 12 insertions(+), 29 deletions(-)

diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 62246a87b291ac8..65de25dd2f32663 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/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,16 +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.
-  llvm::sort(usersToRemoveUses, [&](Operation *lhs, Operation *rhs) {
-    return dominance.properlyDominates(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
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg.mlir b/mlir/test/Dialect/LLVMIR/mem2reg.mlir
index 32e3fed7e5485df..30ba459d07a49f3 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg.mlir
+++ b/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
-}