[Mlir-commits] [mlir] 8da9d92 - [mlir][mem2reg] fix 197158 by moving visitReplacedValues call (#198552)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 10 01:10:36 PDT 2026
Author: jeanPerier
Date: 2026-06-10T10:10:31+02:00
New Revision: 8da9d925ac73becc1733c6733ffabc2020131912
URL: https://github.com/llvm/llvm-project/commit/8da9d925ac73becc1733c6733ffabc2020131912
DIFF: https://github.com/llvm/llvm-project/commit/8da9d925ac73becc1733c6733ffabc2020131912.diff
LOG: [mlir][mem2reg] fix 197158 by moving visitReplacedValues call (#198552)
Fix #197158 and #200844 by moving the `visitReplacedValues` calls
between `promoteInRegion` and `removeBlockingUses` , as well as setting
the insertion point before the replaced store operation before calling
the `PromotableMemOpInterface::getStored` API (instead of setting the
insertion point after).
The action order change is done at the top level. The `promoteInRegion`
are done for all regions in post order, then the `visitReplacedValues`
are done for all regions, and then only the `removeBlockingUses` are
done for all regions in post order. This ensures that any load results
that would happen to be used in a later stored is not deleted by
`removeBlockingUses` before it is used by `visitReplacedValues`.
The insertion point change ensures that the stored values passed to
`visitReplacedValues` dominate the related store operations. Otherwise,
typical `visitReplacedValues` that set insertion points at the store
operation and use the stored values generated invalid IR when
`getStored` generates new IR (like bitcasts for the LLVM dialect
implementation).
Added:
Modified:
mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
mlir/lib/Transforms/Mem2Reg.cpp
mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
mlir/test/Transforms/mem2reg.mlir
mlir/test/lib/Dialect/Test/TestOpDefs.cpp
mlir/test/lib/Dialect/Test/TestOps.td
Removed:
################################################################################
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index cfb914c798729..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.
}],
@@ -235,29 +237,30 @@ 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 once for this operation
+ after all reaching definitions have been computed but before any
+ blocking uses are removed.
- 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 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 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 +307,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..12f50d777e79c 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
@@ -161,6 +164,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 +195,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 +248,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 +258,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.
@@ -264,6 +270,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` so that replacedValues are still
+ /// valid (`removeBlockingUses` can later replace them).
+ 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
/// to a
diff erent region, the new region will be processed instead.
@@ -291,13 +303,10 @@ 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;
+ /// 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;
@@ -349,7 +358,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
@@ -408,6 +417,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
@@ -527,7 +538,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
@@ -583,7 +595,11 @@ Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
MemorySlot aliasSlot =
getOpAliasSlot(memOp, slot, info.aliasMap).value_or(slot);
if (memOp.storesTo(aliasSlot)) {
- builder.setInsertionPointAfter(memOp);
+ // Ensures helper IR introduced by `getStored` is created before the
+ // storing op. This ensures the provided stored 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.
@@ -671,13 +687,6 @@ Value MemorySlotPromoter::promoteInBlock(Block *block, Value reachingDef) {
builder.setInsertionPointAfter(op);
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.
- for (auto &[region, reachingDef] : regionsToProcess)
- removeBlockingUses(region);
}
}
}
@@ -689,6 +698,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;
}
@@ -731,6 +741,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
@@ -829,9 +844,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 +852,24 @@ 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, ArrayRef<std::pair<Operation *, Value>> replacedValues) {
+ auto *blockingUsesMapIt = info.userToBlockingUses.find(region);
+ if (blockingUsesMapIt == info.userToBlockingUses.end())
+ return;
+ BlockingUsesMap &blockingUsesMap = blockingUsesMapIt->second;
+ if (blockingUsesMap.empty())
+ return;
+
+ for (auto &[op, _] : blockingUsesMap) {
+ auto toVisit = dyn_cast<PromotableOpInterface>(op);
+ if (!toVisit || !toVisit.requiresReplacedValues())
+ continue;
+ builder.setInsertionPointAfter(op);
+ toVisit.visitReplacedValues(replacedValues, builder);
}
}
@@ -967,26 +995,29 @@ 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
+ // Step 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.
- // 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.
+ // Blocking uses are not removed yet.
promoteInRegion(slot.ptr.getParentRegion(), nullptr);
- // 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);
+ // Step 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);
}
+ // Step 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);
+
// 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..cf979ee786861 100644
--- a/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
+++ b/mlir/test/Dialect/LLVMIR/mem2reg-dbginfo.mlir
@@ -126,3 +126,66 @@ 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.
+
+// 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
+}
+
+// Regression test for https://github.com/llvm/llvm-project/issues/200844.
+
+// 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
+}
+
+// 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
+}
diff --git a/mlir/test/Transforms/mem2reg.mlir b/mlir/test/Transforms/mem2reg.mlir
index 551f913b49313..da39bc356a109 100644
--- a/mlir/test/Transforms/mem2reg.mlir
+++ b/mlir/test/Transforms/mem2reg.mlir
@@ -401,3 +401,154 @@ func.func @promotable_through_partial_alias(%x: f32) -> f32 {
%v = memref.load %alias[] : memref<f32>
return %v : f32
}
+
+// -----
+
+// `test.slot_tracker` tests `PromotableOpInterface::visitReplacedValues`.
+// It emits a `test.tracked_value` after each store to the slot.
+
+// Basic case: two unconditional stores followed by a load.
+
+// 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. A value loaded from
+// the slot is stored back to it.
+
+// 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
diff erent blocks merged at a load.
+
+// 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 as `@tracker_basic`, but the tracker is placed last.
+
+// 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.
+
+// 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
+}
+
+// -----
+
+// 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/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index 41f5a6f69fa21..c48b16a2ad211 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());
+ }
+}
+
::mlir::LogicalResult test::TestDummyTensorOp::bufferize(
::mlir::RewriterBase &rewriter,
const ::mlir::bufferization::BufferizationOptions &options,
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index d008343e43a39..31002d1f17a75 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -4024,6 +4024,26 @@ 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`.
+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.
+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
diff erent signednesses (signed and
// unsigned 32-bit integers, both aliasing the signless i32 storage).
More information about the Mlir-commits
mailing list