Reapply "[MLIR] [Mem2Reg] Fix unused block argument removal logic (#188484)" (#188571) (#188599)

This reverts commit d9402d087ab90610d3ff8a78a50eb66d3be4cffd. This re-applies commit e5adddc5be63b8bb8c36572f68ac64c8042cb282 along with 62eafb5cd1 Co-authored-by: Yi Zhang <cathyzhyi@google.com> Co-authored-by: Yi Zhang <cathyzhyi@google.com>
2026-03-25 14:08:50 -07:00 · 2026-03-25 14:08:50 -07:00 · 443e4cb2df
commit 443e4cb2df
parent dd57b45522
3 changed files with 288 additions and 33 deletions
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@ -19,6 +19,7 @@
 #include "mlir/Interfaces/MemorySlotInterfaces.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/STLExtras.h"
+#include "llvm/ADT/SetVector.h"
 #include "llvm/Support/DebugLog.h"
 #include "llvm/Support/GenericIteratedDominanceFrontier.h"

@ -264,9 +265,9 @@ private:
  /// to a different region, the new region will be processed instead.
  void removeBlockingUses(Region *region);

-  /// Links merge point block arguments to the terminators targeting the merge
-  /// point or remove the argument if it is not used.
-  void linkMergePoints();
+  /// Removes operations and merge point block arguments that ended up not being
+  /// necessary.
+  void removeUnusedItems();

  /// Lazily-constructed default value representing the content of the slot when
  /// no store has been executed. This function may mutate IR.
@ -294,7 +295,7 @@ private:
  /// the promotion.
  llvm::SmallVector<PromotableOpInterface> toVisitReplacedValues;
  /// Operations to be erased at the end of the promotion.
-  llvm::SmallVector<Operation *> toErase;
+  llvm::SmallSetVector<Operation *, 8> toErase;

  DominanceInfo &dominance;
  const DataLayout &dataLayout;
@ -664,6 +665,18 @@ void MemorySlotPromoter::promoteInRegion(Region *region, Value reachingDef) {

    job.reachingDef = promoteInBlock(block, job.reachingDef);

+    if (auto terminator = dyn_cast<BranchOpInterface>(block->getTerminator())) {
+      for (BlockOperand &blockOperand : terminator->getBlockOperands()) {
+        if (info.mergePoints.contains(blockOperand.get())) {
+          if (!job.reachingDef)
+            job.reachingDef = getOrCreateDefaultValue();
+
+          terminator.getSuccessorOperands(blockOperand.getOperandNumber())
+              .append(job.reachingDef);
+        }
+      }
+    }
+
    for (auto *child : job.block->children())
      dfsStack.emplace_back<DfsJob>({child, job.reachingDef});
  }
@ -753,7 +766,7 @@ void MemorySlotPromoter::removeBlockingUses(Region *region) {
      if (toPromoteMemOp.removeBlockingUses(slot, blockingUsesMap[toPromote],
                                            builder, reachingDef,
                                            dataLayout) == DeletionKind::Delete)
-        toErase.push_back(toPromote);
+        toErase.insert(toPromote);
      if (toPromoteMemOp.storesTo(slot))
        if (Value replacedValue = replacedValuesMap[toPromoteMemOp])
          replacedValues.push_back({toPromoteMemOp, replacedValue});
@ -764,46 +777,99 @@ void MemorySlotPromoter::removeBlockingUses(Region *region) {
    builder.setInsertionPointAfter(toPromote);
    if (toPromoteBasic.removeBlockingUses(blockingUsesMap[toPromote],
                                          builder) == DeletionKind::Delete)
-      toErase.push_back(toPromote);
+      toErase.insert(toPromote);
    if (toPromoteBasic.requiresReplacedValues())
      toVisitReplacedValues.push_back(toPromoteBasic);
  }
 }

-void MemorySlotPromoter::linkMergePoints() {
-  // We want to eliminate unused block arguments. In case connecting a block
-  // argument to its predecessor would trigger the use of the predecessor's
-  // unused block argument, we need to process merge points in an expanding
-  // worklist, `mergePointArgsToProcess`.
+void MemorySlotPromoter::removeUnusedItems() {
+  // We want to eliminate unused block arguments. Because block arguments can be
+  // used to populate other block arguments, there might be cycles of arguments
+  // that are only used to populate each-other. We therefore need a small
+  // dataflow analysis to identify which block arguments are truly used.

  SmallPtrSet<BlockArgument, 8> mergePointArgsUnused;
-  SmallVector<BlockArgument> mergePointArgsToProcess;
+  SmallVector<BlockArgument> usedMergePointArgsToProcess;
+
+  // First, separate the block arguments that are not used or only used for the
+  // purpose of populating a merge point block argument from the others. These
+  // block arguments are potentially unused. Meanwhile, arguments that are
+  // definitely used will be the starting point of the propagation of the
+  // analysis.
+  auto isDefinitelyUsed = [&](BlockArgument arg) {
+    for (auto &use : arg.getUses()) {
+      if (llvm::is_contained(toErase, use.getOwner()))
+        continue;
+
+      // We now want to detect whether the use is to populate a merge point
+      // block argument. If it is not, the argument is definitely used.
+
+      auto branchOp = dyn_cast<BranchOpInterface>(use.getOwner());
+      if (!branchOp)
+        return true;
+
+      std::optional<BlockArgument> successorArgument =
+          branchOp.getSuccessorBlockArgument(use.getOperandNumber());
+      if (!successorArgument)
+        return true;
+
+      if (!info.mergePoints.contains(successorArgument->getOwner()))
+        return true;
+
+      // The last block argument of a merge point is its reaching definition
+      // argument. If the argument being populated is not the last one, it is a
+      // genuine use of the value.
+      bool isLastBlockArgument =
+          successorArgument->getArgNumber() ==
+          successorArgument->getOwner()->getNumArguments() - 1;
+      if (!isLastBlockArgument)
+        return true;
+    }
+    return false;
+  };
+
  for (Block *mergePoint : info.mergePoints) {
    BlockArgument arg = mergePoint->getArguments().back();
-    if (arg.use_empty())
-      mergePointArgsUnused.insert(arg);
+    if (isDefinitelyUsed(arg))
+      usedMergePointArgsToProcess.push_back(arg);
    else
-      mergePointArgsToProcess.push_back(arg);
+      mergePointArgsUnused.insert(arg);
  }

-  while (!mergePointArgsToProcess.empty()) {
-    BlockArgument arg = mergePointArgsToProcess.pop_back_val();
+  // We now refine mergePointArgsUnused from the information of which block
+  // arguments are definitely used.
+  while (!usedMergePointArgsToProcess.empty()) {
+    BlockArgument arg = usedMergePointArgsToProcess.pop_back_val();
    Block *mergePoint = arg.getOwner();

+    assert(arg.getArgNumber() == mergePoint->getNumArguments() - 1 &&
+           "merge point argument must be the last argument of the merge point");
+
    for (BlockOperand &use : mergePoint->getUses()) {
-      Value reachingDef = reachingAtBlockEnd[use.getOwner()->getBlock()];
-      if (!reachingDef)
-        reachingDef = getOrCreateDefaultValue();
+      // If a value used to populate this used merge point argument is another
+      // merge point block argument that is currently considered unused, it must
+      // now be considered used and processed as such later.

-      // If the reaching definition is a block argument of an unused merge
-      // point, mark it as used and process it as such later.
-      auto reachingDefArgument = dyn_cast<BlockArgument>(reachingDef);
-      if (reachingDefArgument &&
-          mergePointArgsUnused.erase(reachingDefArgument))
-        mergePointArgsToProcess.push_back(reachingDefArgument);
+      auto branch = cast<BranchOpInterface>(use.getOwner());
+      SuccessorOperands succOperands =
+          branch.getSuccessorOperands(use.getOperandNumber());

-      BranchOpInterface user = cast<BranchOpInterface>(use.getOwner());
-      user.getSuccessorOperands(use.getOperandNumber()).append(reachingDef);
+      // The successor operand is either the last one or is not present if the
+      // user block is dead.
+      assert(succOperands.size() == mergePoint->getNumArguments() ||
+             succOperands.size() + 1 == mergePoint->getNumArguments());
+
+      // If the user block is dead, the default value acts as a placeholder
+      // dummy value.
+      if (succOperands.size() + 1 == mergePoint->getNumArguments())
+        succOperands.append(getOrCreateDefaultValue());
+
+      Value populatedValue = succOperands[arg.getArgNumber()];
+      auto populatedValueAsArg = dyn_cast<BlockArgument>(populatedValue);
+      if (populatedValueAsArg &&
+          mergePointArgsUnused.erase(populatedValueAsArg))
+        usedMergePointArgsToProcess.push_back(populatedValueAsArg);
    }

    builder.setInsertionPointToStart(mergePoint);
@ -812,6 +878,25 @@ void MemorySlotPromoter::linkMergePoints() {
      (*statistics.newBlockArgumentAmount)++;
  }

+  for (Operation *toEraseOp : toErase)
+    toEraseOp->erase();
+
+  // First, erase all successor operands that feed into unused merge point
+  // block arguments. This must be done before erasing the block arguments
+  // themselves because an unused merge point argument may be used to
+  // populate another unused merge point argument via a branch operation.
+  for (BlockArgument arg : mergePointArgsUnused) {
+    Block *mergePoint = arg.getOwner();
+    for (BlockOperand &use : mergePoint->getUses()) {
+      auto branch = cast<BranchOpInterface>(use.getOwner());
+      SuccessorOperands succOperands =
+          branch.getSuccessorOperands(use.getOperandNumber());
+      succOperands.erase(arg.getArgNumber());
+    }
+  }
+
+  // Now that all successor operands feeding unused args have been removed,
+  // erase the block arguments themselves.
  for (BlockArgument arg : mergePointArgsUnused) {
    Block *mergePoint = arg.getOwner();
    mergePoint->eraseArgument(mergePoint->getNumArguments() - 1);
@ -840,11 +925,8 @@ MemorySlotPromoter::promoteSlot() {
    op.visitReplacedValues(replacedValues, builder);
  }

-  // Finally, connect merge points to their predecessor's reaching definitions.
-  linkMergePoints();
-
-  for (Operation *toEraseOp : toErase)
-    toEraseOp->erase();
+  // Finally, remove unused operations and merge point block arguments.
+  removeUnusedItems();

  assert(slot.ptr.use_empty() &&
         "after promotion, the slot pointer should not be used anymore");
--- a/mlir/test/Dialect/MemRef/mem2reg.mlir
+++ b/mlir/test/Dialect/MemRef/mem2reg.mlir
@ -171,8 +171,141 @@ func.func @unused_alloca_store_loop() {
  // CHECK-NOT: memref.alloca
  %cst = arith.constant 1 : i32
  %alloca = memref.alloca() : memref<i32>
+  // CHECK: cf.br ^[[BB1:.*]]
  cf.br ^bb1
+
+// CHECK: ^[[BB1]]:
 ^bb1:
+  // CHECK-NOT: memref.store
  memref.store %cst, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[BB1]]
  cf.br ^bb1
 }
+
+// -----
+
+// CHECK-LABEL: func.func @store_back_to_alloca
+// CHECK-SAME: (%[[COND:.*]]: i1)
+func.func @store_back_to_alloca(%cond: i1) -> i32 {
+  // CHECK-NOT: memref.alloca
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
+  %c0 = arith.constant 0 : i32
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+  %c1 = arith.constant 1 : i32
+  // CHECK-NOT: memref.alloca
+  %alloca = memref.alloca() : memref<i32>
+  memref.store %c0, %alloca[] : memref<i32>
+  %loaded = memref.load %alloca[] : memref<i32>
+  // CHECK: cf.cond_br %[[COND]], ^[[STORE_BACK:.*]], ^[[SKIP:.*]]
+  cf.cond_br %cond, ^store_back, ^skip
+
+// CHECK: ^[[STORE_BACK]]:
+^store_back:
+  memref.store %loaded, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[MERGE:.*]](%[[C0]] : i32)
+  cf.br ^merge
+
+// CHECK: ^[[SKIP]]:
+^skip:
+  memref.store %c1, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[MERGE]](%[[C1]] : i32)
+  cf.br ^merge
+
+// CHECK: ^[[MERGE]](%[[RESULT:.*]]: i32):
+^merge:
+  %result = memref.load %alloca[] : memref<i32>
+  // CHECK: return %[[RESULT]] : i32
+  return %result : i32
+}
+
+// -----
+
+// Ensure that a merge point used by an erased operation is not considered used.
+
+// CHECK-LABEL: func.func @merge_point_used_by_erased_op
+// CHECK-SAME: (%[[COND:.*]]: i1)
+func.func @merge_point_used_by_erased_op(%cond: i1) -> i32 {
+  // CHECK-NOT: memref.alloca
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
+  %c0 = arith.constant 0 : i32
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+  %c1 = arith.constant 1 : i32
+  // CHECK-NOT: memref.alloca
+  %alloca = memref.alloca() : memref<i32>
+  // CHECK: cf.cond_br %[[COND]], ^[[PRED1:.*]], ^[[PRED2:.*]]
+  cf.cond_br %cond, ^pred1, ^pred2
+
+// CHECK: ^[[PRED1]]:
+^pred1:
+  memref.store %c0, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[MERGE:.*]]{{$}}
+  cf.br ^merge
+
+// CHECK: ^[[PRED2]]:
+^pred2:
+  memref.store %c1, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[MERGE]]{{$}}
+  cf.br ^merge
+
+// CHECK: ^[[MERGE]]:
+^merge:
+  %result = memref.load %alloca[] : memref<i32>
+  memref.store %result, %alloca[] : memref<i32>
+  // CHECK: return %[[C0]] : i32
+  return %c0 : i32
+}
+
+// -----
+
+// Two consecutive merge points: pred1 and pred2 merge at merge1, then merge1
+// and pred3 merge at merge2.
+
+// CHECK-LABEL: func.func @two_consecutive_merge_points
+// CHECK-SAME: (%[[COND1:.*]]: i1, %[[COND2:.*]]: i1)
+func.func @two_consecutive_merge_points(%cond1: i1, %cond2: i1) -> i32 {
+  // CHECK-NOT: memref.alloca
+  // CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
+  // CHECK-DAG: %[[C1:.*]] = arith.constant 1 : i32
+  // CHECK-DAG: %[[C2:.*]] = arith.constant 2 : i32
+  // CHECK-NOT: memref.alloca
+  %c0 = arith.constant 0 : i32
+  %c1 = arith.constant 1 : i32
+  %c2 = arith.constant 2 : i32
+  %alloca = memref.alloca() : memref<i32>
+  // CHECK: cf.cond_br %[[COND1]], ^[[PRED1:.*]], ^[[MID:.*]]
+  cf.cond_br %cond1, ^pred1, ^mid
+
+// CHECK: ^[[MID]]:
+^mid:
+  // CHECK: cf.cond_br %[[COND2]], ^[[PRED2:.*]], ^[[PRED3:.*]]
+  cf.cond_br %cond2, ^pred2, ^pred3
+
+// CHECK: ^[[PRED1]]:
+^pred1:
+  memref.store %c0, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[MERGE1:.*]](%[[C0]] : i32)
+  cf.br ^merge1
+
+// CHECK: ^[[PRED2]]:
+^pred2:
+  memref.store %c1, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[MERGE1]](%[[C1]] : i32)
+  cf.br ^merge1
+
+// CHECK: ^[[MERGE1]](%[[MARG:.*]]: i32):
+^merge1:
+  // CHECK: cf.br ^[[MERGE2:.*]](%[[MARG]] : i32)
+  cf.br ^merge2
+
+// CHECK: ^[[PRED3]]:
+^pred3:
+  memref.store %c2, %alloca[] : memref<i32>
+  // CHECK: cf.br ^[[MERGE2]](%[[C2]] : i32)
+  cf.br ^merge2
+
+// CHECK: ^[[MERGE2]](%[[RESULT:.*]]: i32):
+^merge2:
+  %result = memref.load %alloca[] : memref<i32>
+  // CHECK: return %[[RESULT]] : i32
+  return %result : i32
+}
--- a/mlir/test/Transforms/mem2reg.mlir
+++ b/mlir/test/Transforms/mem2reg.mlir
@ -113,3 +113,43 @@ func.func @unknown_region_op_load() {
  }) : () -> ()
  return
 }
+
+// -----
+
+// A cycle of merge points where both merge point block arguments are unused.
+// merge1 branches to merge2, and merge2 branches back to merge1, so each
+// merge point's reaching definition arg is used as a successor operand
+// feeding the other. During removeUnusedItems, the successor operand erasure
+// and block argument erasure must be performed in separate phases. Otherwise,
+// regardless of iteration order, erasing either arg first will crash because
+// the other's successor operand still uses it.
+
+// CHECK-LABEL: func.func @cyclic_unused_merge_points
+// CHECK-SAME: (%[[COND:.*]]: i1)
+// CHECK-NOT: memref.alloca
+// CHECK-DAG: %[[C0:.*]] = arith.constant 0 : i32
+// CHECK: cf.br ^[[MERGE1:.*]]{{$}}
+// CHECK: ^[[MERGE1]]:
+// CHECK:   cf.cond_br %[[COND]], ^[[MERGE2:.*]], ^[[STORE:.*]]
+// CHECK: ^[[STORE]]:
+// CHECK:   cf.br ^[[MERGE2]]{{$}}
+// CHECK: ^[[MERGE2]]:
+// CHECK:   cf.cond_br %[[COND]], ^[[MERGE1]], ^[[EXIT:.*]]
+// CHECK: ^[[EXIT]]:
+// CHECK:   return %[[C0]] : i32
+func.func @cyclic_unused_merge_points(%cond: i1) -> i32 {
+  %c0 = arith.constant 0 : i32
+  %c1 = arith.constant 1 : i32
+  %alloca = memref.alloca() : memref<i32>
+  memref.store %c0, %alloca[] : memref<i32>
+  cf.br ^merge1
+^merge1:
+  cf.cond_br %cond, ^merge2, ^store
+^store:
+  memref.store %c1, %alloca[] : memref<i32>
+  cf.br ^merge2
+^merge2:
+  cf.cond_br %cond, ^merge1, ^exit
+^exit:
+  return %c0 : i32
+}