[Mlir-commits] [mlir] [mlir][mem2reg] fix 197158 by moving visitReplacedValues call (PR #198552)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Jun 8 06:09:48 PDT 2026
https://github.com/jeanPerier updated https://github.com/llvm/llvm-project/pull/198552
>From 97777a20836a82e35016640308011d532be6fbcf Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Tue, 19 May 2026 08:28:30 -0700
Subject: [PATCH 1/7] [mlir][mem2reg] fix 197158 by moving visitReplacedValues
call
---
.../mlir/Interfaces/MemorySlotInterfaces.td | 32 +++++-----
mlir/lib/Transforms/Mem2Reg.cpp | 64 +++++++++++++------
mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir | 25 ++++++++
3 files changed, 86 insertions(+), 35 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index cfb914c798729..467dfcad49ced 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -235,29 +235,28 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
"::mlir::OpBuilder &":$builder)
>,
InterfaceMethod<[{
- This method allows the promoted operation to visit the SSA values used
- in place of the memory slot once the promotion process of the memory
- slot is complete.
+ Indicates whether the promoted operation needs to visit the SSA values
+ that replace the memory slot.
- If this method returns true, the `visitReplacedValues` method on this
- operation will be called after the main mutation stage finishes
- (i.e., after all ops have been processed with `removeBlockingUses`).
+ If true, `visitReplacedValues` will be called after computing the
+ slot's reaching definitions but before removing its blocking uses.
- Operations should only visit the replaced values if the intended
- transformation applies to all the replaced values. Furthermore, replaced
- values must not be deleted.
+ This should only return true if the intended transformation applies
+ to all replaced values. Replaced values must not be deleted.
}], "bool", "requiresReplacedValues", (ins), [{}],
[{ return false; }]
>,
InterfaceMethod<[{
- Transforms the IR using the SSA values that replaced the memory slot.
+ Transforms the IR using the SSA values replacing the memory slot.
- This method will only be called after all blocking uses have been
- scheduled for removal and if `requiresReplacedValues` returned
- true.
+ Called after computing reaching definitions but before removing
+ blocking uses, provided `requiresReplacedValues` returns true.
+ The `mutatedDefs` reference the original IR. Any new uses of a load
+ result created here will be automatically redirected when the
+ framework replaces that load with its reaching definition.
- The builder is located after the promotable operation on call. During
- the transformation, *no operation should be deleted*.
+ The builder is positioned after the promotable operation.
+ No operations may be deleted during this transformation.
}],
"void", "visitReplacedValues",
(ins "::llvm::ArrayRef<std::pair<::mlir::Operation*, ::mlir::Value>>":$mutatedDefs,
@@ -304,7 +303,8 @@ def PromotableAliaserInterface : OpInterface<"PromotableAliaserInterface"> {
Exposing an alias requires implementing the two projection methods
below to bridge values between the parent and new slot element types.
- If these projections cannot be performed, leave `newMemorySlots` empty.
+ If these projections cannot be performed, `newMemorySlots` must be left
+ empty.
No IR mutation is allowed in this method.
}],
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 277457f574f55..83222dc42e5e2 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -264,6 +264,12 @@ class MemorySlotPromoter {
/// This member function must only be called at most once per region.
void promoteInRegion(Region *region, Value reachingDef);
+ /// Calls `visitReplacedValues` on ops in `region` that requested it.
+ /// Must run before `removeBlockingUses(region)` so values in
+ /// `replacedValuesMap` remain live. Any new IR uses of load results
+ /// will be redirected when `removeBlockingUses` replaces those loads.
+ void visitReplacedValuesForRegion(Region *region);
+
/// Removes the blocking uses of the slot within the given region, in
/// reverse topological order. If the content of the region was moved out
/// to a different region, the new region will be processed instead.
@@ -291,13 +297,6 @@ class MemorySlotPromoter {
/// Contains the reaching definition at the end of the blocks visited so far.
DenseMap<Block *, Value> reachingAtBlockEnd;
- /// Lists all the values that have been set by a memory operation as a
- /// reaching definition at one point during the promotion. The accompanying
- /// operation is the memory operation that originally stored the value.
- llvm::SmallVector<std::pair<Operation *, Value>> replacedValues;
- /// Operations to visit with the `visitReplacedValues` method at the end of
- /// the promotion.
- llvm::SmallVector<PromotableOpInterface> toVisitReplacedValues;
/// Operations to be erased at the end of the promotion.
llvm::SmallSetVector<Operation *, 8> toErase;
@@ -676,6 +675,13 @@ Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
// Even though `finalizePromotion` may have moved regions to a new
// operation, `removeBlockingUses` handles this case and will redirect
// processing to the correct region.
+ // `visitReplacedValuesForRegion` must run before `removeBlockingUses`
+ // so that load results in `replacedValuesMap` remain live. The
+ // subsequent `replaceAllUsesWith` in `removeBlockingUses` will
+ // redirect any new uses created by `visitReplacedValues` to their
+ // post-promotion definitions.
+ for (auto &[region, reachingDef] : regionsToProcess)
+ visitReplacedValuesForRegion(region);
for (auto &[region, reachingDef] : regionsToProcess)
removeBlockingUses(region);
}
@@ -829,9 +835,6 @@ void MemorySlotPromoter::removeBlockingUses(Region *region) {
aliasSlot, blockingUsesMap[toPromote], builder,
reachingDefAtBlockingUse, dataLayout) == DeletionKind::Delete)
toErase.insert(toPromote);
- if (toPromoteMemOp.storesTo(aliasSlot))
- if (Value replacedValue = replacedValuesMap[toPromoteMemOp])
- replacedValues.push_back({toPromoteMemOp, replacedValue});
continue;
}
@@ -840,8 +843,32 @@ void MemorySlotPromoter::removeBlockingUses(Region *region) {
if (toPromoteBasic.removeBlockingUses(blockingUsesMap[toPromote],
builder) == DeletionKind::Delete)
toErase.insert(toPromote);
- if (toPromoteBasic.requiresReplacedValues())
- toVisitReplacedValues.push_back(toPromoteBasic);
+ }
+}
+
+void MemorySlotPromoter::visitReplacedValuesForRegion(Region *region) {
+ auto *blockingUsesMapIt = info.userToBlockingUses.find(region);
+ if (blockingUsesMapIt == info.userToBlockingUses.end())
+ return;
+ BlockingUsesMap &blockingUsesMap = blockingUsesMapIt->second;
+ if (blockingUsesMap.empty())
+ return;
+
+ // Build the replaced-values list from per-store snapshots. These values
+ // are still live. Since this runs before `removeBlockingUses` replaces
+ // loads, any new uses of load results created by `visitReplacedValues`
+ // will be correctly redirected to post-promotion definitions later.
+ SmallVector<std::pair<Operation *, Value>> replacedValues;
+ for (const auto &[memOp, value] : replacedValuesMap)
+ if (value)
+ replacedValues.push_back({memOp, value});
+
+ for (auto &[op, _] : blockingUsesMap) {
+ auto toVisit = dyn_cast<PromotableOpInterface>(op);
+ if (!toVisit || !toVisit.requiresReplacedValues())
+ continue;
+ builder.setInsertionPointAfter(op);
+ toVisit.visitReplacedValues(replacedValues, builder);
}
}
@@ -977,16 +1004,15 @@ MemorySlotPromoter::promoteSlot() {
// before removeBlockingUses.
promoteInRegion(slot.ptr.getParentRegion(), nullptr);
+ // Notify operations of reaching definitions. This runs before
+ // `removeBlockingUses` so values passed to `visitReplacedValues` remain
+ // live. Any new uses of load results will be redirected when those loads
+ // are replaced.
+ visitReplacedValuesForRegion(slot.ptr.getParentRegion());
+
// Blocking uses can then be removed for the outermost region.
removeBlockingUses(slot.ptr.getParentRegion());
- // Notify operations that requested it of the reaching definitions set by
- // storing memory operations.
- for (PromotableOpInterface op : toVisitReplacedValues) {
- builder.setInsertionPointAfter(op);
- op.visitReplacedValues(replacedValues, builder);
- }
-
// Finally, remove unused operations and merge point block arguments.
removeUnusedItems();
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
index b7cbd787f06e4..fae71e751e66d 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
@@ -126,3 +126,28 @@ llvm.func @keep_dbg_if_not_promoted() {
llvm.call @use_ptr(%1) : (!llvm.ptr) -> ()
llvm.return
}
+
+// Regression test for https://github.com/llvm/llvm-project/issues/197158:
+// When a loaded value is stored back, `replacedValuesMap` records the load's
+// result. Running `visitReplacedValues` before `replaceAllUsesWith` ensures
+// the new `dbg.value` references the live load result, which is then
+// correctly redirected to the reaching definition.
+
+// CHECK-LABEL: llvm.func @store_load_store_back
+// CHECK-NOT: = llvm.alloca
+// CHECK-NOT: llvm.intr.dbg.declare
+// CHECK-NOT: llvm.store
+// CHECK-NOT: llvm.load
+// CHECK: %[[CST:.*]] = llvm.mlir.constant({{.*}}) : i64
+// CHECK: llvm.intr.dbg.value #[[$VAR]] = %[[CST]] : i64
+// CHECK: llvm.return
+llvm.func @store_load_store_back() {
+ %one = llvm.mlir.constant(1 : i32) : i32
+ %cst = llvm.mlir.constant(42 : i64) : i64
+ %p = llvm.alloca %one x i64 : (i32) -> !llvm.ptr
+ llvm.intr.dbg.declare #di_local_variable = %p : !llvm.ptr
+ llvm.store %cst, %p : i64, !llvm.ptr
+ %v = llvm.load %p : !llvm.ptr -> i64
+ llvm.store %v, %p : i64, !llvm.ptr
+ llvm.return
+}
>From 02f0b029bc37cba0742e4efbbb577f7d59a73f12 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Wed, 20 May 2026 03:21:04 -0700
Subject: [PATCH 2/7] add more tests
---
mlir/test/Transforms/mem2reg.mlir | 140 ++++++++++++++++++++++
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 28 +++++
mlir/test/lib/Dialect/Test/TestOps.td | 23 ++++
3 files changed, 191 insertions(+)
diff --git a/mlir/test/Transforms/mem2reg.mlir b/mlir/test/Transforms/mem2reg.mlir
index 551f913b49313..0b4f4346c1e40 100644
--- a/mlir/test/Transforms/mem2reg.mlir
+++ b/mlir/test/Transforms/mem2reg.mlir
@@ -401,3 +401,143 @@ func.func @promotable_through_partial_alias(%x: f32) -> f32 {
%v = memref.load %alias[] : memref<f32>
return %v : f32
}
+
+// -----
+
+// `test.slot_tracker` tests `PromotableOpInterface::visitReplacedValues`.
+// The framework calls it with (store, stored-value) pairs computed during
+// promotion. The implementation emits a `test.tracked_value` after each store.
+
+// Basic case: two unconditional stores followed by a load. Each store emits
+// a `tracked_value` at its position. The load is replaced by the last
+// reaching definition.
+
+// CHECK-LABEL: func.func @tracker_basic
+// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32) -> i32
+// CHECK-NOT: test.multi_slot_alloca
+// CHECK-NOT: test.slot_tracker
+// CHECK-NOT: memref.store
+// CHECK-NOT: memref.load
+// CHECK-DAG: test.tracked_value %[[A]], "x" : i32
+// CHECK-DAG: test.tracked_value %[[B]], "x" : i32
+// CHECK: return %[[B]] : i32
+func.func @tracker_basic(%a: i32, %b: i32) -> i32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ test.slot_tracker %slot, "x" : memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ memref.store %b, %slot[] : memref<i32>
+ %v = memref.load %slot[] : memref<i32>
+ return %v : i32
+}
+
+// -----
+
+// Regression test for use-of-deleted-SSA-value bug (similar to the
+// `dbg.declare` reproducer in `mem2reg-dbginfo.mlir`). A value loaded from
+// the slot is stored back to it. `visitReplacedValues` must run before
+// `replaceAllUsesWith` so the new `tracked_value` references the live load
+// result, which is then redirected to the original store's value.
+
+// CHECK-LABEL: func.func @tracker_load_stored_back
+// CHECK-SAME: (%[[A:.*]]: i32)
+// CHECK-NOT: test.multi_slot_alloca
+// CHECK-NOT: test.slot_tracker
+// CHECK-NOT: memref.store
+// CHECK-NOT: memref.load
+// CHECK: test.tracked_value %[[A]], "y" : i32
+// CHECK: test.tracked_value %[[A]], "y" : i32
+// CHECK: return
+func.func @tracker_load_stored_back(%a: i32) {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ test.slot_tracker %slot, "y" : memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ %v = memref.load %slot[] : memref<i32>
+ memref.store %v, %slot[] : memref<i32>
+ return
+}
+
+// -----
+
+// CFG case: stores in different blocks merged at a load. Each store emits
+// a `tracked_value` in its block. The load is replaced by the merge-point
+// block argument.
+
+// CHECK-LABEL: func.func @tracker_blocks
+// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32, %[[COND:.*]]: i1) -> i32
+// CHECK-NOT: test.multi_slot_alloca
+// CHECK-NOT: test.slot_tracker
+// CHECK-NOT: memref.store
+// CHECK-NOT: memref.load
+// CHECK: cf.cond_br %[[COND]], ^[[BB1:.*]], ^[[BB2:.*]]
+// CHECK: ^[[BB1]]:
+// CHECK: test.tracked_value %[[A]], "z" : i32
+// CHECK: cf.br ^[[BB3:.*]](%[[A]] : i32)
+// CHECK: ^[[BB2]]:
+// CHECK: test.tracked_value %[[B]], "z" : i32
+// CHECK: cf.br ^[[BB3]](%[[B]] : i32)
+// CHECK: ^[[BB3]](%[[MERGE:.*]]: i32):
+// CHECK: return %[[MERGE]] : i32
+func.func @tracker_blocks(%a: i32, %b: i32, %cond: i1) -> i32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ test.slot_tracker %slot, "z" : memref<i32>
+ cf.cond_br %cond, ^bb1, ^bb2
+^bb1:
+ memref.store %a, %slot[] : memref<i32>
+ cf.br ^bb3
+^bb2:
+ memref.store %b, %slot[] : memref<i32>
+ cf.br ^bb3
+^bb3:
+ %v = memref.load %slot[] : memref<i32>
+ return %v : i32
+}
+
+// -----
+
+// Same shape as `@tracker_basic`, but the tracker is placed last (i.e. all
+// stores and the load dominate it). The framework should still find it
+// among the slot's blocking uses, hand it the same (store, value) pairs,
+// and emit `tracked_value` ops at the original store positions, regardless
+// of where the tracker itself sits.
+
+// CHECK-LABEL: func.func @tracker_at_end
+// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32) -> i32
+// CHECK-NOT: test.multi_slot_alloca
+// CHECK-NOT: test.slot_tracker
+// CHECK-NOT: memref.store
+// CHECK-NOT: memref.load
+// CHECK-DAG: test.tracked_value %[[A]], "k" : i32
+// CHECK-DAG: test.tracked_value %[[B]], "k" : i32
+// CHECK: return %[[B]] : i32
+func.func @tracker_at_end(%a: i32, %b: i32) -> i32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ memref.store %b, %slot[] : memref<i32>
+ %v = memref.load %slot[] : memref<i32>
+ test.slot_tracker %slot, "k" : memref<i32>
+ return %v : i32
+}
+
+// -----
+
+// Nested-region case: an outer store and a conditional inner store. The
+// tracker in the outer region receives both stores, emitting a
+// `tracked_value` after each store's position.
+
+// CHECK-LABEL: func.func @tracker_nested_region
+// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32, %[[COND:.*]]: i1)
+// CHECK-NOT: test.multi_slot_alloca
+// CHECK-NOT: test.slot_tracker
+// CHECK-NOT: memref.store
+// CHECK: test.tracked_value %[[A]], "w" : i32
+// CHECK: scf.if %[[COND]]
+// CHECK: test.tracked_value %[[B]], "w" : i32
+func.func @tracker_nested_region(%a: i32, %b: i32, %cond: i1) {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ test.slot_tracker %slot, "w" : memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ scf.if %cond {
+ memref.store %b, %slot[] : memref<i32>
+ }
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 8315bd7cef783..0224a1ff69ee8 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -1950,6 +1950,34 @@ Value TestPartialAlias::projectAliasValueToSlotValue(
.getResult(0);
}
+//===----------------------------------------------------------------------===//
+// TestSlotTracker
+//===----------------------------------------------------------------------===//
+
+bool TestSlotTracker::canUsesBeRemoved(
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
+ if (blockingUses.size() != 1)
+ return false;
+ return (*blockingUses.begin())->get() == getSource();
+}
+
+DeletionKind TestSlotTracker::removeBlockingUses(
+ const SmallPtrSetImpl<OpOperand *> &blockingUses, OpBuilder &builder) {
+ return DeletionKind::Delete;
+}
+
+bool TestSlotTracker::requiresReplacedValues() { return true; }
+
+void TestSlotTracker::visitReplacedValues(
+ ArrayRef<std::pair<Operation *, Value>> definitions, OpBuilder &builder) {
+ for (auto [op, value] : definitions) {
+ builder.setInsertionPointAfter(op);
+ TestTrackedValue::create(builder, getLoc(), value, getNameAttr());
+ }
+}
+
namespace {
/// Returns test dialect's memref layout for test dialect's tensor encoding when
/// applicable.
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index a1529e3020c82..8bf8314b2c500 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3969,6 +3969,29 @@ def TestTransparentCastAlias : TEST_Op<"transparent_cast_alias",
let assemblyFormat = "$source attr-dict `:` functional-type($source, $result)";
}
+// Records the value last stored to a memref slot to test
+// `PromotableOpInterface::visitReplacedValues`. When promoted, it emits a
+// `test.tracked_value` after every `memref.store` of the reaching
+// definitions (similar to `llvm.intr.dbg.declare` and `dbg.value`).
+def TestSlotTracker : TEST_Op<"slot_tracker",
+ [DeclareOpInterfaceMethods<PromotableOpInterface,
+ ["canUsesBeRemoved",
+ "removeBlockingUses",
+ "requiresReplacedValues",
+ "visitReplacedValues"]>]> {
+ let arguments = (ins AnyMemRef:$source, StrAttr:$name);
+ let assemblyFormat =
+ "$source `,` $name attr-dict `:` type($source)";
+}
+
+// Marker emitted by `test.slot_tracker` at each reaching-definition store.
+// Carries the stored value and the tracker's name.
+def TestTrackedValue : TEST_Op<"tracked_value"> {
+ let arguments = (ins AnyType:$value, StrAttr:$name);
+ let assemblyFormat =
+ "$value `,` $name attr-dict `:` type($value)";
+}
+
// Transparent alias of a memref slot exposing two simultaneously-usable
// aliases of the same bit-width at different signednesses (signed and
// unsigned 32-bit integers, both aliasing the signless i32 storage).
>From 387994ca9b2f3810b6beab0f5f10fa424d419bdc Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Mon, 1 Jun 2026 06:32:29 -0700
Subject: [PATCH 3/7] visit all regions
---
.../mlir/Interfaces/MemorySlotInterfaces.td | 16 +--
mlir/lib/Transforms/Mem2Reg.cpp | 111 ++++++++++--------
mlir/test/Transforms/mem2reg.mlir | 53 +++++----
mlir/test/lib/Dialect/Test/TestOps.td | 5 +-
4 files changed, 102 insertions(+), 83 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index 467dfcad49ced..a200306fbe50f 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -238,8 +238,9 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
Indicates whether the promoted operation needs to visit the SSA values
that replace the memory slot.
- If true, `visitReplacedValues` will be called after computing the
- slot's reaching definitions but before removing its blocking uses.
+ If true, `visitReplacedValues` will be called once for this operation
+ after all reaching definitions have been computed but before any
+ blocking uses are removed.
This should only return true if the intended transformation applies
to all replaced values. Replaced values must not be deleted.
@@ -249,11 +250,12 @@ def PromotableOpInterface : OpInterface<"PromotableOpInterface"> {
InterfaceMethod<[{
Transforms the IR using the SSA values replacing the memory slot.
- Called after computing reaching definitions but before removing
- blocking uses, provided `requiresReplacedValues` returns true.
- The `mutatedDefs` reference the original IR. Any new uses of a load
- result created here will be automatically redirected when the
- framework replaces that load with its reaching definition.
+ Called after all reaching definitions for the slot have been computed
+ but before any blocking uses are removed, provided
+ `requiresReplacedValues` returns true. `mutatedDefs` contains the
+ replacing values. Any new uses of a load result created here will be
+ automatically redirected to its reaching definition when the load is
+ replaced.
The builder is positioned after the promotable operation.
No operations may be deleted during this transformation.
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 83222dc42e5e2..3e452f77912fc 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -161,6 +161,8 @@ struct MemorySlotPromotionInfo {
/// Transitive aliases of `slot.ptr` via `PromotableAliaserInterface`,
/// mapping alias values to their exposed slot and aliased operand.
PromotableAliasMap aliasMap;
+ /// True if at least one blocking use requires visiting the replaced values.
+ bool needsAnyReplacedValuesVisit = false;
};
/// Computes information for basic slot promotion. This will check that direct
@@ -190,7 +192,8 @@ class MemorySlotPromotionAnalyzer {
LogicalResult
computeBlockingUses(RegionBlockingUsesMap &userToBlockingUses,
DenseMap<Region *, RegionPromotionInfo> ®ionsToPromote,
- PromotableAliasMap &aliasMap);
+ PromotableAliasMap &aliasMap,
+ bool &needsAnyReplacedValuesVisit);
/// Computes the points in the provided region where multiple re-definitions
/// of the slot's value (stores) may conflict.
@@ -242,9 +245,8 @@ class MemorySlotPromoter {
/// promotion, including within nested regions needing promotion.
/// `reachingDef` is the value the slot contains at the beginning of the
/// block. This member function returns the reached definition at the end of
- /// the block. If the block contains a region that needs promotion, the
- /// blocking uses of that region will have been removed. This member function
- /// will not remove the blocking uses contained directly in the block.
+ /// the block. Blocking uses are not removed and `visitReplacedValues` is
+ /// not called here; both are deferred to `promoteSlot`.
///
/// The `reachingDef` may be a null value. In that case, a lazily-created
/// default value will be used.
@@ -253,10 +255,11 @@ class MemorySlotPromoter {
Value promoteInBlock(Block *block, Value reachingDef);
/// Computes the reaching definition for all the operations that require
- /// promotion, including within nested regions needing promotion, and removes
- /// the blocking uses of the slot within the region.
+ /// promotion, including within nested regions needing promotion. Records
+ /// `region` in `regionsInPostOrder` after the DFS completes.
/// `reachingDef` is the value the slot contains at the beginning of the
- /// region.
+ /// region. Blocking uses are not removed here; that is deferred to
+ /// `promoteSlot`.
///
/// The `reachingDef` may be a null value. In that case, a lazily-created
/// default value will be used.
@@ -265,10 +268,9 @@ class MemorySlotPromoter {
void promoteInRegion(Region *region, Value reachingDef);
/// Calls `visitReplacedValues` on ops in `region` that requested it.
- /// Must run before `removeBlockingUses(region)` so values in
- /// `replacedValuesMap` remain live. Any new IR uses of load results
- /// will be redirected when `removeBlockingUses` replaces those loads.
- void visitReplacedValuesForRegion(Region *region);
+ /// Must run before `removeBlockingUses` so load results remain live.
+ void visitReplacedValuesForRegion(
+ Region *region, ArrayRef<std::pair<Operation *, Value>> replacedValues);
/// Removes the blocking uses of the slot within the given region, in
/// reverse topological order. If the content of the region was moved out
@@ -297,6 +299,10 @@ class MemorySlotPromoter {
/// Contains the reaching definition at the end of the blocks visited so far.
DenseMap<Block *, Value> reachingAtBlockEnd;
+ /// Regions visited by `promoteInRegion`, recorded in DFS post-order.
+ /// Inner regions appear before their enclosing regions.
+ llvm::SmallVector<Region *> regionsInPostOrder;
+
/// Operations to be erased at the end of the promotion.
llvm::SmallSetVector<Operation *, 8> toErase;
@@ -348,7 +354,7 @@ Value MemorySlotPromoter::getOrCreateDefaultValue() {
LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
RegionBlockingUsesMap &userToBlockingUses,
DenseMap<Region *, RegionPromotionInfo> ®ionsToPromote,
- PromotableAliasMap &aliasMap) {
+ PromotableAliasMap &aliasMap, bool &needsAnyReplacedValuesVisit) {
// 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
@@ -407,6 +413,8 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
dataLayout))
return failure();
regionsWithDirectUse.insert(user->getParentRegion());
+ if (promotable.requiresReplacedValues())
+ needsAnyReplacedValuesVisit = true;
} else if (auto promotable = dyn_cast<PromotableMemOpInterface>(user)) {
// If the memop reaches the root slot through multiple distinct alias
// operands, promotion fails. `PromotableMemOpInterface` expects a
@@ -526,7 +534,8 @@ MemorySlotPromotionAnalyzer::computeInfo() {
// We also compute at this stage the regions that will be analyzed for
// reaching definition information.
if (failed(computeBlockingUses(info.userToBlockingUses, info.regionsToPromote,
- info.aliasMap)))
+ info.aliasMap,
+ info.needsAnyReplacedValuesVisit)))
return {};
// Compute the blocks containing a store for each region, either directly or
@@ -671,19 +680,8 @@ Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
reachingDef = promotableRegionOp.finalizePromotion(
slot, reachingDef, hasValueStores, reachingAtBlockEnd, builder);
- // Blocking uses can then be removed for the regions that were promoted.
- // Even though `finalizePromotion` may have moved regions to a new
- // operation, `removeBlockingUses` handles this case and will redirect
- // processing to the correct region.
- // `visitReplacedValuesForRegion` must run before `removeBlockingUses`
- // so that load results in `replacedValuesMap` remain live. The
- // subsequent `replaceAllUsesWith` in `removeBlockingUses` will
- // redirect any new uses created by `visitReplacedValues` to their
- // post-promotion definitions.
- for (auto &[region, reachingDef] : regionsToProcess)
- visitReplacedValuesForRegion(region);
- for (auto &[region, reachingDef] : regionsToProcess)
- removeBlockingUses(region);
+ // `visitReplacedValuesForRegion` and `removeBlockingUses` for these
+ // inner regions are deferred until after the entire DFS finishes.
}
}
}
@@ -695,6 +693,7 @@ Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
void MemorySlotPromoter::promoteInRegion(Region *region, Value reachingDef) {
if (region->hasOneBlock()) {
promoteInBlock(®ion->front(), reachingDef);
+ regionsInPostOrder.push_back(region);
return;
}
@@ -737,6 +736,11 @@ void MemorySlotPromoter::promoteInRegion(Region *region, Value reachingDef) {
for (auto *child : job.block->children())
dfsStack.emplace_back<DfsJob>({child, job.reachingDef});
}
+
+ // Record this region after its (and any nested) promotion has completed,
+ // so that `regionsInPostOrder` lists inner regions before their enclosing
+ // regions.
+ regionsInPostOrder.push_back(region);
}
/// Gets or creates a block index mapping for the region of which the entry
@@ -846,7 +850,8 @@ void MemorySlotPromoter::removeBlockingUses(Region *region) {
}
}
-void MemorySlotPromoter::visitReplacedValuesForRegion(Region *region) {
+void MemorySlotPromoter::visitReplacedValuesForRegion(
+ Region *region, ArrayRef<std::pair<Operation *, Value>> replacedValues) {
auto *blockingUsesMapIt = info.userToBlockingUses.find(region);
if (blockingUsesMapIt == info.userToBlockingUses.end())
return;
@@ -854,15 +859,7 @@ void MemorySlotPromoter::visitReplacedValuesForRegion(Region *region) {
if (blockingUsesMap.empty())
return;
- // Build the replaced-values list from per-store snapshots. These values
- // are still live. Since this runs before `removeBlockingUses` replaces
- // loads, any new uses of load results created by `visitReplacedValues`
- // will be correctly redirected to post-promotion definitions later.
- SmallVector<std::pair<Operation *, Value>> replacedValues;
- for (const auto &[memOp, value] : replacedValuesMap)
- if (value)
- replacedValues.push_back({memOp, value});
-
+ // `replacedValues` is built once from per-store snapshots.
for (auto &[op, _] : blockingUsesMap) {
auto toVisit = dyn_cast<PromotableOpInterface>(op);
if (!toVisit || !toVisit.requiresReplacedValues())
@@ -994,24 +991,36 @@ void MemorySlotPromoter::removeUnusedItems() {
std::optional<PromotableAllocationOpInterface>
MemorySlotPromoter::promoteSlot() {
- // Perform the promotion recursively through nested regions. The reaching
- // definition starts with a null value that will be replaced by a
- // lazily-created default value if the value must be passed to a promotion
- // interface while no store has been encountered yet.
- // Innermost regions will see their blocking uses be removed, but not the
- // outermost region which we have to remove manually afterwards. This is
- // because PromotableRegionOpInterface::finalizePromotion must be called
- // before removeBlockingUses.
+ // Pass 1: compute reaching definitions and run `finalizePromotion` for
+ // nested `PromotableRegionOpInterface` ops.
promoteInRegion(slot.ptr.getParentRegion(), nullptr);
- // Notify operations of reaching definitions. This runs before
- // `removeBlockingUses` so values passed to `visitReplacedValues` remain
- // live. Any new uses of load results will be redirected when those loads
- // are replaced.
- visitReplacedValuesForRegion(slot.ptr.getParentRegion());
+#ifndef NDEBUG
+ // Every region tracked in `info.userToBlockingUses` must have been visited
+ // by `promoteInRegion` (either directly for the slot's parent region or
+ // recursively through `PromotableRegionOpInterface`).
+ llvm::SmallPtrSet<Region *, 4> visitedRegions(llvm::from_range,
+ regionsInPostOrder);
+ for (auto &[region, _] : info.userToBlockingUses)
+ assert(visitedRegions.contains(region) &&
+ "every region with blocking uses must be visited during promotion");
+#endif
+
+ // Pass 2: call `visitReplacedValues` on operations that requested it.
+ if (info.needsAnyReplacedValuesVisit) {
+ SmallVector<std::pair<Operation *, Value>> replacedValues;
+ replacedValues.reserve(replacedValuesMap.size());
+ for (const auto &[memOp, value] : replacedValuesMap)
+ if (value)
+ replacedValues.push_back({memOp, value});
+ for (Region *region : regionsInPostOrder)
+ visitReplacedValuesForRegion(region, replacedValues);
+ }
- // Blocking uses can then be removed for the outermost region.
- removeBlockingUses(slot.ptr.getParentRegion());
+ // Pass 3: remove the slot's blocking uses across all regions. Iterating in
+ // DFS post-order (innermost regions first) is required.
+ for (Region *region : regionsInPostOrder)
+ removeBlockingUses(region);
// Finally, remove unused operations and merge point block arguments.
removeUnusedItems();
diff --git a/mlir/test/Transforms/mem2reg.mlir b/mlir/test/Transforms/mem2reg.mlir
index 0b4f4346c1e40..da39bc356a109 100644
--- a/mlir/test/Transforms/mem2reg.mlir
+++ b/mlir/test/Transforms/mem2reg.mlir
@@ -405,12 +405,9 @@ func.func @promotable_through_partial_alias(%x: f32) -> f32 {
// -----
// `test.slot_tracker` tests `PromotableOpInterface::visitReplacedValues`.
-// The framework calls it with (store, stored-value) pairs computed during
-// promotion. The implementation emits a `test.tracked_value` after each store.
+// It emits a `test.tracked_value` after each store to the slot.
-// Basic case: two unconditional stores followed by a load. Each store emits
-// a `tracked_value` at its position. The load is replaced by the last
-// reaching definition.
+// Basic case: two unconditional stores followed by a load.
// CHECK-LABEL: func.func @tracker_basic
// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32) -> i32
@@ -432,11 +429,8 @@ func.func @tracker_basic(%a: i32, %b: i32) -> i32 {
// -----
-// Regression test for use-of-deleted-SSA-value bug (similar to the
-// `dbg.declare` reproducer in `mem2reg-dbginfo.mlir`). A value loaded from
-// the slot is stored back to it. `visitReplacedValues` must run before
-// `replaceAllUsesWith` so the new `tracked_value` references the live load
-// result, which is then redirected to the original store's value.
+// Regression test for use-of-deleted-SSA-value bug. A value loaded from
+// the slot is stored back to it.
// CHECK-LABEL: func.func @tracker_load_stored_back
// CHECK-SAME: (%[[A:.*]]: i32)
@@ -458,9 +452,7 @@ func.func @tracker_load_stored_back(%a: i32) {
// -----
-// CFG case: stores in different blocks merged at a load. Each store emits
-// a `tracked_value` in its block. The load is replaced by the merge-point
-// block argument.
+// CFG case: stores in different blocks merged at a load.
// CHECK-LABEL: func.func @tracker_blocks
// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32, %[[COND:.*]]: i1) -> i32
@@ -494,11 +486,7 @@ func.func @tracker_blocks(%a: i32, %b: i32, %cond: i1) -> i32 {
// -----
-// Same shape as `@tracker_basic`, but the tracker is placed last (i.e. all
-// stores and the load dominate it). The framework should still find it
-// among the slot's blocking uses, hand it the same (store, value) pairs,
-// and emit `tracked_value` ops at the original store positions, regardless
-// of where the tracker itself sits.
+// Same as `@tracker_basic`, but the tracker is placed last.
// CHECK-LABEL: func.func @tracker_at_end
// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32) -> i32
@@ -520,9 +508,7 @@ func.func @tracker_at_end(%a: i32, %b: i32) -> i32 {
// -----
-// Nested-region case: an outer store and a conditional inner store. The
-// tracker in the outer region receives both stores, emitting a
-// `tracked_value` after each store's position.
+// Nested-region case: an outer store and a conditional inner store.
// CHECK-LABEL: func.func @tracker_nested_region
// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32, %[[COND:.*]]: i1)
@@ -541,3 +527,28 @@ func.func @tracker_nested_region(%a: i32, %b: i32, %cond: i1) {
}
return
}
+
+// -----
+
+// Inverse of `@tracker_nested_region`: the tracker lives *inside* the
+// nested region.
+
+// CHECK-LABEL: func.func @tracker_inside_nested_region
+// CHECK-SAME: (%[[A:.*]]: i32, %[[B:.*]]: i32, %[[C:.*]]: i32, %[[COND:.*]]: i1)
+// CHECK-NOT: test.multi_slot_alloca
+// CHECK-NOT: test.slot_tracker
+// CHECK-NOT: memref.store
+// CHECK: test.tracked_value %[[A]], "n" : i32
+// CHECK: scf.if %[[COND]]
+// CHECK: test.tracked_value %[[B]], "n" : i32
+// CHECK: test.tracked_value %[[C]], "n" : i32
+func.func @tracker_inside_nested_region(%a: i32, %b: i32, %c: i32, %cond: i1) {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ scf.if %cond {
+ test.slot_tracker %slot, "n" : memref<i32>
+ memref.store %b, %slot[] : memref<i32>
+ }
+ memref.store %c, %slot[] : memref<i32>
+ return
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index 8bf8314b2c500..5814d9c7b3e72 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -3970,9 +3970,7 @@ def TestTransparentCastAlias : TEST_Op<"transparent_cast_alias",
}
// Records the value last stored to a memref slot to test
-// `PromotableOpInterface::visitReplacedValues`. When promoted, it emits a
-// `test.tracked_value` after every `memref.store` of the reaching
-// definitions (similar to `llvm.intr.dbg.declare` and `dbg.value`).
+// `PromotableOpInterface::visitReplacedValues`.
def TestSlotTracker : TEST_Op<"slot_tracker",
[DeclareOpInterfaceMethods<PromotableOpInterface,
["canUsesBeRemoved",
@@ -3985,7 +3983,6 @@ def TestSlotTracker : TEST_Op<"slot_tracker",
}
// Marker emitted by `test.slot_tracker` at each reaching-definition store.
-// Carries the stored value and the tracker's name.
def TestTrackedValue : TEST_Op<"tracked_value"> {
let arguments = (ins AnyType:$value, StrAttr:$name);
let assemblyFormat =
>From cf797d68d50d389462ad831fb98d4c70fa3071df Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Mon, 1 Jun 2026 08:27:06 -0700
Subject: [PATCH 4/7] set insertion point before memOp in promoteInBlock
---
.../mlir/Interfaces/MemorySlotInterfaces.td | 10 ++++----
mlir/lib/Transforms/Mem2Reg.cpp | 5 +++-
mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir | 23 +++++++++++++++++++
3 files changed, 33 insertions(+), 5 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index a200306fbe50f..3997c84963207 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -121,11 +121,13 @@ def PromotableMemOpInterface : OpInterface<"PromotableMemOpInterface"> {
storing a value to a slot must always be able to provide the value it
stores. This method is only called once per slot promotion, and only
on operations that store to the slot according to the `storesTo` method.
- The returned value must dominate all operations dominated by the storing
- operation.
- The builder is located immediately after the memory operation on call.
- No IR deletion is allowed in this method. IR mutations must not
+ The returned value must dominate the memory operation. This ensures
+ that new uses inserted by `visitReplacedValues` after the memory
+ operation are properly dominated.
+
+ The builder is positioned immediately before the memory operation on
+ call. No IR deletion is allowed in this method. IR mutations must not
introduce new uses of the memory slot. Existing control flow must not
be modified.
}],
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 3e452f77912fc..58011d0f6bf57 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -591,7 +591,10 @@ Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
MemorySlot aliasSlot =
getOpAliasSlot(memOp, slot, info.aliasMap).value_or(slot);
if (memOp.storesTo(aliasSlot)) {
- builder.setInsertionPointAfter(memOp);
+ // Insert helper IR introduced by `getStored` before the storing op.
+ // This ensures the returned value dominates the store, allowing
+ // `visitReplacedValues` to safely create new uses after the store.
+ builder.setInsertionPoint(memOp);
// To not expose default value creation to the interfaces, if we have
// no reaching definition by now, we set it to the default value.
// This is slightly too eager as `getStored` may not need it.
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
index fae71e751e66d..10c41d55a221a 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
@@ -151,3 +151,26 @@ llvm.func @store_load_store_back() {
llvm.store %v, %p : i64, !llvm.ptr
llvm.return
}
+
+// Regression test for https://github.com/llvm/llvm-project/issues/200844.
+// `getStored` synthesizes casts at the builder's insertion point. Setting the
+// insertion point before the store ensures the cast dominates the `dbg.value`
+// created by `visitReplacedValues` after the store.
+
+// CHECK-LABEL: llvm.func @dbg_declare_with_store_type_conversion
+// CHECK-SAME: (%[[VAL:.*]]: f32)
+// CHECK-NOT: = llvm.alloca
+// CHECK-NOT: llvm.intr.dbg.declare
+// CHECK-NOT: llvm.store
+// CHECK-NOT: llvm.load
+// CHECK: %[[BITCAST:.*]] = llvm.bitcast %[[VAL]] : f32 to i32
+// CHECK: llvm.intr.dbg.value #[[$VAR]] = %[[BITCAST]] : i32
+// CHECK: llvm.return %[[BITCAST]] : i32
+llvm.func @dbg_declare_with_store_type_conversion(%val : f32) -> i32 {
+ %0 = llvm.mlir.constant(1 : i32) : i32
+ %1 = llvm.alloca %0 x i32 {alignment = 4 : i64} : (i32) -> !llvm.ptr
+ llvm.intr.dbg.declare #di_local_variable = %1 : !llvm.ptr
+ llvm.store %val, %1 {alignment = 4 : i64} : f32, !llvm.ptr
+ %2 = llvm.load %1 {alignment = 4 : i64} : !llvm.ptr -> i32
+ llvm.return %2 : i32
+}
>From b76f08237f3398d2156e633d2591623d9be1f53d Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Mon, 1 Jun 2026 09:07:32 -0700
Subject: [PATCH 5/7] trim comments
---
mlir/lib/Transforms/Mem2Reg.cpp | 26 +++++--------------
mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir | 9 +------
2 files changed, 8 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 58011d0f6bf57..c645209afbb8f 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -682,9 +682,6 @@ Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
builder.setInsertionPointAfter(op);
reachingDef = promotableRegionOp.finalizePromotion(
slot, reachingDef, hasValueStores, reachingAtBlockEnd, builder);
-
- // `visitReplacedValuesForRegion` and `removeBlockingUses` for these
- // inner regions are deferred until after the entire DFS finishes.
}
}
}
@@ -862,7 +859,6 @@ void MemorySlotPromoter::visitReplacedValuesForRegion(
if (blockingUsesMap.empty())
return;
- // `replacedValues` is built once from per-store snapshots.
for (auto &[op, _] : blockingUsesMap) {
auto toVisit = dyn_cast<PromotableOpInterface>(op);
if (!toVisit || !toVisit.requiresReplacedValues())
@@ -994,21 +990,13 @@ void MemorySlotPromoter::removeUnusedItems() {
std::optional<PromotableAllocationOpInterface>
MemorySlotPromoter::promoteSlot() {
- // Pass 1: compute reaching definitions and run `finalizePromotion` for
- // nested `PromotableRegionOpInterface` ops.
+ // Pass 1: perform the promotion recursively through nested regions. The
+ // reaching definition starts with a null value that will be replaced by a
+ // lazily-created default value if the value must be passed to a promotion
+ // interface while no store has been encountered yet.
+ // Blocking uses are not removed yet.
promoteInRegion(slot.ptr.getParentRegion(), nullptr);
-#ifndef NDEBUG
- // Every region tracked in `info.userToBlockingUses` must have been visited
- // by `promoteInRegion` (either directly for the slot's parent region or
- // recursively through `PromotableRegionOpInterface`).
- llvm::SmallPtrSet<Region *, 4> visitedRegions(llvm::from_range,
- regionsInPostOrder);
- for (auto &[region, _] : info.userToBlockingUses)
- assert(visitedRegions.contains(region) &&
- "every region with blocking uses must be visited during promotion");
-#endif
-
// Pass 2: call `visitReplacedValues` on operations that requested it.
if (info.needsAnyReplacedValuesVisit) {
SmallVector<std::pair<Operation *, Value>> replacedValues;
@@ -1020,8 +1008,8 @@ MemorySlotPromoter::promoteSlot() {
visitReplacedValuesForRegion(region, replacedValues);
}
- // Pass 3: remove the slot's blocking uses across all regions. Iterating in
- // DFS post-order (innermost regions first) is required.
+ // Pass 3: remove the slot's blocking uses across all regions. Iterating
+ // in through the regions in same order as promoteInRegion.
for (Region *region : regionsInPostOrder)
removeBlockingUses(region);
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
index 10c41d55a221a..ef4ebede0d3ce 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
@@ -127,11 +127,7 @@ llvm.func @keep_dbg_if_not_promoted() {
llvm.return
}
-// Regression test for https://github.com/llvm/llvm-project/issues/197158:
-// When a loaded value is stored back, `replacedValuesMap` records the load's
-// result. Running `visitReplacedValues` before `replaceAllUsesWith` ensures
-// the new `dbg.value` references the live load result, which is then
-// correctly redirected to the reaching definition.
+// Regression test for https://github.com/llvm/llvm-project/issues/197158.
// CHECK-LABEL: llvm.func @store_load_store_back
// CHECK-NOT: = llvm.alloca
@@ -153,9 +149,6 @@ llvm.func @store_load_store_back() {
}
// Regression test for https://github.com/llvm/llvm-project/issues/200844.
-// `getStored` synthesizes casts at the builder's insertion point. Setting the
-// insertion point before the store ensures the cast dominates the `dbg.value`
-// created by `visitReplacedValues` after the store.
// CHECK-LABEL: llvm.func @dbg_declare_with_store_type_conversion
// CHECK-SAME: (%[[VAL:.*]]: f32)
>From 81342f411634662af1f1c2620057fcb9966ddbf2 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Mon, 1 Jun 2026 09:12:27 -0700
Subject: [PATCH 6/7] replace by in commentt
---
mlir/lib/Transforms/Mem2Reg.cpp | 3 ++-
1 file changed, 2 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index c645209afbb8f..419fdd9f14bac 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -268,7 +268,8 @@ class MemorySlotPromoter {
void promoteInRegion(Region *region, Value reachingDef);
/// Calls `visitReplacedValues` on ops in `region` that requested it.
- /// Must run before `removeBlockingUses` so load results remain live.
+ /// Must run before `removeBlockingUses` so that replacedValues are still
+ /// valid (`removeBlockingUses` can later replace them).
void visitReplacedValuesForRegion(
Region *region, ArrayRef<std::pair<Operation *, Value>> replacedValues);
>From 32d9562252fb082abf6f62345c70b5cdc3dc3125 Mon Sep 17 00:00:00 2001
From: Jean Perier <jperier at nvidia.com>
Date: Mon, 8 Jun 2026 06:04:08 -0700
Subject: [PATCH 7/7] update pass documentation and add test
---
mlir/lib/Transforms/Mem2Reg.cpp | 83 ++++++++++---------
mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir | 22 +++++
2 files changed, 65 insertions(+), 40 deletions(-)
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 419fdd9f14bac..9406384937bb5 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -37,13 +37,14 @@ using namespace mlir;
/// This pass turns unnecessary uses of automatically allocated memory slots
/// into direct Value-based operations. For example, it will simplify storing a
/// constant in a memory slot to immediately load it to a direct use of that
-/// constant. In other words, given a memory slot addressed by a non-aliased
-/// "pointer" Value, mem2reg removes all the uses of that pointer.
+/// constant. In other words, given a memory slot addressed by a "pointer" Value
+/// (which may be exposed through aliases), mem2reg removes all the uses of
+/// that pointer and its aliases.
///
/// Within a block, this is done by following the chain of stores and loads of
-/// the slot and replacing the results of loads with the values previously
-/// stored. If a load happens before any other store, a poison value is used
-/// instead.
+/// the slot (or its aliases) and replacing the results of loads with the values
+/// previously stored. If a load happens before any other store, a poison value
+/// is used instead.
///
/// Control flow can create situations where a load could be replaced by
/// multiple possible stores depending on the control flow path taken. As a
@@ -71,14 +72,14 @@ using namespace mlir;
/// - A first step computes the list of operations that transitively use the
/// memory slot we would like to promote. The purpose of this phase is to
/// identify which uses must be removed to promote the slot, either by rewiring
-/// the user or deleting it. Naturally, direct uses of the slot must be removed.
-/// Sometimes additional uses must also be removed: this is notably the case
-/// when a direct user of the slot cannot rewire its use and must delete itself,
-/// and thus must make its users no longer use it. If the allocation is used in
-/// nested regions, it is also ensured the region operations provide the right
-/// interface to analyze the values of the allocation at the edges of its
-/// regions. If any of those constraints cannot be satisfied, promotion cannot
-/// continue: this is decided at this step.
+/// the user or deleting it. Naturally, direct uses of the slot and its aliases
+/// must be removed. Sometimes additional uses must also be removed: this is
+/// notably the case when a direct user of the slot cannot rewire its use and
+/// must delete itself, and thus must make its users no longer use it. If the
+/// allocation is used in nested regions, it is also ensured the region
+/// operations provide the right interface to analyze the values of the
+/// allocation at the edges of its regions. If any of those constraints cannot
+/// be satisfied, promotion cannot continue: this is decided at this step.
/// - A second step computes the list of blocks where a block argument will be
/// needed ("merge points") without mutating the IR. These blocks are the blocks
/// leading to a definition clash between two predecessors. Such blocks happen
@@ -89,35 +90,37 @@ using namespace mlir;
/// do so aborts promotion at this step).
///
/// At this point, promotion is guaranteed to happen, and the transformation
-/// phase can begin. For each region of the program, a two step process is
-/// carried out.
-/// - The first step of the per-region process computes the reaching definition
-/// of the memory slot at each blocking user. This is the core of the mem2reg
-/// algorithm, also known as load-store forwarding. This analyses loads and
-/// stores and propagates which value must be stored in the slot at each
-/// blocking user. This is achieved by doing a depth-first walk of the dominator
-/// tree of the function. This is sufficient because the reaching definition at
-/// the beginning of a block is either its new block argument if it is a merge
-/// block, or the definition reaching the end of its immediate dominator (parent
-/// in the dominator tree). We can therefore propagate this information down the
+/// phase can begin. The transformation is a three-step process.
+/// - The first step computes the reaching definition of the memory slot at
+/// each blocking user. This is the core of the mem2reg algorithm, also known
+/// as load-store forwarding. This analyses loads and stores and propagates
+/// which value must be stored in the slot at each blocking user. This is
+/// achieved by doing a depth-first walk of the dominator tree of the function.
+/// This is sufficient because the reaching definition at the beginning of a
+/// block is either its new block argument if it is a merge block, or the
+/// definition reaching the end of its immediate dominator (parent in the
+/// dominator tree). We can therefore propagate this information down the
/// dominator tree to proceed with renaming within blocks. If at any point a
/// region operation that contains a use of the allocation is encountered, the
-/// transformation process is triggered on the child regions of the encountered
-/// operation, to obtain the reaching definition at its end and carry on with
-/// the value forwarding.
-/// - The second step of the per-region process uses the reaching definition to
-/// remove blocking uses in topological order. Some reaching definitions may
-/// be values that will be removed or modified during the blocking use removal
-/// step (typically, in the case of a store that stores the result of a load).
-/// To properly handle such values, this step traverses the operations to modify
-/// in reverse topological order. This way, if a value that will disappear is
-/// used in place of reaching definition, the logic to make it disappear will be
-/// executed after the value has been used to replace an operation. For regions
-/// within a PromotableRegionOpInterface, in order to correctly handle cases
-/// where the finalization logic would use a reaching definition that will be
-/// replaced, the finalization logic must be called before the blocking use
-/// removal step, so that any use of a value that will be removed gets properly
-/// replaced.
+/// reaching definition computation is recursively triggered on the child
+/// regions of the encountered operation, to obtain the reaching definition at
+/// its end and carry on with the value forwarding.
+/// - The second step visits the values that will replace the memory slot for
+/// operations that requested it. This must happen before the removal of
+/// blocking uses, so that operations can safely inspect reaching definitions
+/// before they may be removed or modified.
+/// - The third step uses the reaching definition to remove blocking uses. Some
+/// reaching definitions may be values that will be removed or modified during
+/// the blocking use removal step (typically, in the case of a store that stores
+/// the result of a load). To properly handle such values, this step traverses
+/// the regions in post-order, and the operations to modify in reverse
+/// topological order. This way, if a value that will disappear is used in place
+/// of a reaching definition, the logic to make it disappear will be executed
+/// after the value has been used to replace an operation. For regions within a
+/// PromotableRegionOpInterface, this ensures that the finalization logic (run
+/// during the first step) and visitReplacedValues (run during the second step)
+/// happen before the blocking use removal step, so that any use of a value that
+/// will be removed gets properly replaced.
///
/// For further reading, chapter three of SSA-based Compiler Design [1]
/// showcases SSA construction for control-flow graphs, where mem2reg is an
diff --git a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
index ef4ebede0d3ce..cf979ee786861 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
@@ -167,3 +167,25 @@ llvm.func @dbg_declare_with_store_type_conversion(%val : f32) -> i32 {
%2 = llvm.load %1 {alignment = 4 : i64} : !llvm.ptr -> i32
llvm.return %2 : i32
}
+
+// CHECK-LABEL: llvm.func @nested_store_load_store_back
+// CHECK-NOT: = llvm.alloca
+// CHECK-NOT: llvm.intr.dbg.declare
+// CHECK-NOT: llvm.store
+// CHECK-NOT: llvm.load
+// CHECK: %[[CST:.*]] = llvm.mlir.constant({{.*}}) : f64
+// CHECK: scf.if
+// CHECK: llvm.intr.dbg.value #[[$VAR]] = %[[CST]] : f64
+// CHECK: llvm.return
+llvm.func @nested_store_load_store_back(%cdt1 : i1, %cdt2 : i1) {
+ %one = llvm.mlir.constant(1 : i32) : i32
+ %cst = llvm.mlir.constant(4.000000e+00 : f64) : f64
+ %p = llvm.alloca %one x f64 : (i32) -> !llvm.ptr
+ llvm.intr.dbg.declare #di_local_variable = %p : !llvm.ptr
+ scf.if %cdt1 {
+ llvm.store %cst, %p : f64, !llvm.ptr
+ %v = llvm.load %p : !llvm.ptr -> f64
+ llvm.store %v, %p : f64, !llvm.ptr
+ }
+ llvm.return
+}
More information about the Mlir-commits
mailing list