[Mlir-commits] [mlir] [mlir][mem2reg] Allow memory operations to own embedded blocking uses (PR #233705)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Oct 21 08:48:43 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-core
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
Allow PromotableMemOpInterface operations to own blocking uses of a wrapped memory access, enabling Mem2Reg to promote loads and stores through the wrapper. Add Test dialect coverage for supported and rejected cases.
assisted-by-claude
---
Full diff: https://github.com/llvm/llvm-project/pull/233705.diff
7 Files Affected:
- (modified) mlir/include/mlir/Interfaces/MemorySlotInterfaces.h (+3-1)
- (modified) mlir/include/mlir/Interfaces/MemorySlotInterfaces.td (+19)
- (modified) mlir/lib/Interfaces/MemorySlotInterfaces.cpp (+16-2)
- (modified) mlir/lib/Transforms/Mem2Reg.cpp (+20-13)
- (modified) mlir/test/Transforms/Mem2Reg/mem2reg.mlir (+155)
- (modified) mlir/test/lib/Dialect/Test/TestOpDefs.cpp (+64)
- (modified) mlir/test/lib/Dialect/Test/TestOps.td (+22)
``````````diff
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.h b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.h
index b2b5418e1bac10..30891edbfa9706 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.h
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.h
@@ -75,6 +75,8 @@ void populatePromotableAliasMap(PromotableAliaserInterface aliaser,
/// `rootSlot`. If `op` reaches `rootSlot` through multiple distinct aliases
/// (e.g., a memcpy between two aliases of the same root), the result is one
/// of them; use `referencesAtMostOneAliasOfSlot` to rule this out.
+/// If `op` owns its embedded blocking uses, also consider the wrapped access's
+/// operands.
std::optional<MemorySlot> getOpAliasSlot(Operation *op,
const MemorySlot &rootSlot,
const PromotableAliasMap &aliasMap);
@@ -83,7 +85,7 @@ std::optional<MemorySlot> getOpAliasSlot(Operation *op,
/// distinct alias pointer (the root itself or a single `aliasMap` entry).
/// Multiple operands referencing the same alias are allowed. This is used to
/// guard `PromotableMemOpInterface` calls, which assume a single slot per
-/// operation.
+/// operation. Includes embedded operands as described in `getOpAliasSlot`.
bool referencesAtMostOneAliasOfSlot(Operation *op, const MemorySlot &rootSlot,
const PromotableAliasMap &aliasMap);
diff --git a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
index 3997c84963207e..32128c8269b34c 100644
--- a/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
+++ b/mlir/include/mlir/Interfaces/MemorySlotInterfaces.td
@@ -182,6 +182,25 @@ def PromotableMemOpInterface : OpInterface<"PromotableMemOpInterface"> {
"::mlir::Value":$reachingDefinition,
"const ::mlir::DataLayout &":$dataLayout)
>,
+ InterfaceMethod<[{
+ Returns whether this operation owns the blocking uses of a single
+ wrapped memory access. If true, it must have one single-block region
+ containing only that access and a terminator. All blocking uses must
+ belong to the wrapped access, not this operation or its terminator.
+
+ Promotion calls this operation's interface methods with the wrapped
+ access's operands in `blockingUses`, without consulting
+ `PromotableRegionOpInterface` for its region. The builder inserts IR
+ outside the region, before or after this operation.
+
+ An operation that owns embedded blocking uses must not add uses to
+ `newBlockingUses`.
+
+ No IR mutation is allowed in this method.
+ }],
+ "bool", "ownsEmbeddedBlockingUses", (ins), [{}],
+ /*defaultImplementation=*/[{ return false; }]
+ >,
];
}
diff --git a/mlir/lib/Interfaces/MemorySlotInterfaces.cpp b/mlir/lib/Interfaces/MemorySlotInterfaces.cpp
index 57c00aa031d372..d3dcc9a830faf6 100644
--- a/mlir/lib/Interfaces/MemorySlotInterfaces.cpp
+++ b/mlir/lib/Interfaces/MemorySlotInterfaces.cpp
@@ -43,10 +43,24 @@ void mlir::populatePromotableAliasMap(PromotableAliaserInterface aliaser,
}
}
+/// Returns `op`'s operands, including the wrapped access's operands if `op`
+/// owns its embedded blocking uses.
+static SmallVector<Value> getPotentialSlotOperands(Operation *op) {
+ SmallVector<Value> operands(op->getOperands());
+ auto memOp = dyn_cast<PromotableMemOpInterface>(op);
+ if (!memOp || !memOp.ownsEmbeddedBlockingUses())
+ return operands;
+
+ Block &body = op->getRegion(0).front();
+ Operation &access = body.front();
+ llvm::append_range(operands, access.getOperands());
+ return operands;
+}
+
std::optional<MemorySlot>
mlir::getOpAliasSlot(Operation *op, const MemorySlot &rootSlot,
const PromotableAliasMap &aliasMap) {
- for (Value operand : op->getOperands())
+ for (Value operand : getPotentialSlotOperands(op))
if (std::optional<MemorySlot> slot =
getParentSlot(operand, rootSlot, aliasMap))
return slot;
@@ -57,7 +71,7 @@ bool mlir::referencesAtMostOneAliasOfSlot(Operation *op,
const MemorySlot &rootSlot,
const PromotableAliasMap &aliasMap) {
Value uniqueAliasPtr;
- for (Value operand : op->getOperands()) {
+ for (Value operand : getPotentialSlotOperands(op)) {
std::optional<MemorySlot> slot = getParentSlot(operand, rootSlot, aliasMap);
if (!slot)
continue;
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index 770c5705a7080b..7abfe029fca3da 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -355,6 +355,12 @@ Value MemorySlotPromoter::getOrCreateDefaultValue() {
return defaultValue = allocator.getDefaultValue(slot, builder);
}
+/// Returns `op`'s parent if it owns embedded blocking uses, or `op` otherwise.
+static Operation *getBlockingUseOwner(Operation *op) {
+ auto memOp = dyn_cast_or_null<PromotableMemOpInterface>(op->getParentOp());
+ return memOp && memOp.ownsEmbeddedBlockingUses() ? memOp : op;
+}
+
LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
RegionBlockingUsesMap &userToBlockingUses,
DenseMap<Region *, RegionPromotionInfo> ®ionsToPromote,
@@ -373,13 +379,14 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
RegionKind::Graph)
return failure();
- // First insert that all immediate users of the slot pointer must no longer
- // use it.
- for (OpOperand &use : slot.ptr.getUses()) {
- SmallPtrSet<OpOperand *, 4> &blockingUses =
- userToBlockingUses[use.getOwner()->getParentRegion()][use.getOwner()];
- blockingUses.insert(&use);
- }
+ auto addBlockingUse = [&](OpOperand *use) {
+ Operation *owner = getBlockingUseOwner(use->getOwner());
+ userToBlockingUses[owner->getParentRegion()][owner].insert(use);
+ };
+
+ // Register all direct uses of the slot pointer as blocking.
+ for (OpOperand &use : slot.ptr.getUses())
+ addBlockingUse(&use);
// Regions that immediately contain a slot memory use that is not a store.
RegionSet regionsWithDirectUse;
@@ -393,7 +400,9 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
// pointer.
SetVector<Operation *> forwardSlice;
mlir::getForwardSlice(slot.ptr, &forwardSlice);
- for (Operation *user : forwardSlice) {
+ for (Operation *op : forwardSlice) {
+ // Each owner wraps a single access, so replacing it preserves slice order.
+ Operation *user = getBlockingUseOwner(op);
// If the next operation has no blocking uses, everything is fine.
auto *blockingUsesMapIt = userToBlockingUses.find(user->getParentRegion());
if (blockingUsesMapIt == userToBlockingUses.end())
@@ -449,10 +458,7 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
for (OpOperand *blockingUse : newBlockingUses) {
assert(llvm::is_contained(user->getResults(), blockingUse->get()));
- Operation *useOwner = blockingUse->getOwner();
- SmallPtrSetImpl<OpOperand *> &newUserBlockingUseSet =
- userToBlockingUses[useOwner->getParentRegion()][useOwner];
- newUserBlockingUseSet.insert(blockingUse);
+ addBlockingUse(blockingUse);
}
}
@@ -553,7 +559,8 @@ MemorySlotPromotionAnalyzer::computeInfo() {
DenseMap<Region *, SmallPtrSet<Block *, 16>> definingBlocks;
auto collectStoringBlocks = [&](Value ptr, const MemorySlot &ptrSlot) {
for (OpOperand &use : ptr.getUses()) {
- Operation *user = use.getOwner();
+ // Attribute an embedded store to its blocking-use owner's block.
+ Operation *user = getBlockingUseOwner(use.getOwner());
if (auto storeOp = dyn_cast<PromotableMemOpInterface>(user))
if (storeOp.storesTo(ptrSlot))
definingBlocks[user->getParentRegion()].insert(user->getBlock());
diff --git a/mlir/test/Transforms/Mem2Reg/mem2reg.mlir b/mlir/test/Transforms/Mem2Reg/mem2reg.mlir
index da39bc356a109f..40cbedec958f4f 100644
--- a/mlir/test/Transforms/Mem2Reg/mem2reg.mlir
+++ b/mlir/test/Transforms/Mem2Reg/mem2reg.mlir
@@ -552,3 +552,158 @@ func.func @tracker_inside_nested_region(%a: i32, %b: i32, %c: i32, %cond: i1) {
memref.store %c, %slot[] : memref<i32>
return
}
+
+// -----
+
+// Promote the wrapper as a whole, replacing its load result with the reaching
+// definition.
+
+// CHECK-LABEL: func.func @owning_access_load
+// CHECK-SAME: (%[[A:.*]]: i32) -> i32
+// CHECK-NOT: test.multi_slot_alloca
+// CHECK-NOT: test.owning_access
+// CHECK: return %[[A]] : i32
+func.func @owning_access_load(%a: i32) -> i32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ %v = test.owning_access {
+ %l = memref.load %slot[] : memref<i32>
+ test.owning_access_yield %l : i32
+ } : i32
+ return %v : i32
+}
+
+// -----
+
+// Attribute the wrapped store to the wrapper's block so ^bb2 gets a block
+// argument that preserves the conditional update.
+
+// CHECK-LABEL: func.func @owning_access_store_defines_block
+// CHECK-SAME: (%[[A:.*]]: i32, %[[COND:.*]]: i1) -> i32
+// CHECK-NOT: test.owning_access
+// CHECK: %[[C42:.*]] = arith.constant 42 : i32
+// CHECK: cf.cond_br %[[COND]], ^[[BB1:.*]], ^[[BB2:.*]](%[[C42]] : i32)
+// CHECK: ^[[BB1]]:
+// CHECK: cf.br ^[[BB2]](%[[A]] : i32)
+// CHECK: ^[[BB2]](%[[MERGE:.*]]: i32):
+// CHECK: return %[[MERGE]] : i32
+func.func @owning_access_store_defines_block(%a: i32, %cond: i1) -> i32 {
+ %c42 = arith.constant 42 : i32
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %c42, %slot[] : memref<i32>
+ cf.cond_br %cond, ^bb1, ^bb2
+^bb1:
+ test.owning_access {
+ memref.store %a, %slot[] : memref<i32>
+ test.owning_access_yield
+ }
+ cf.br ^bb2
+^bb2:
+ %v = memref.load %slot[] : memref<i32>
+ return %v : i32
+}
+
+// -----
+
+// Promote a wrapper inside `scf.if`, yielding its stored value.
+
+// CHECK-LABEL: func.func @owning_access_in_scf_if
+// CHECK-SAME: (%[[A:.*]]: i32, %[[COND:.*]]: i1) -> i32
+// CHECK-NOT: test.owning_access
+// CHECK: %[[C42:.*]] = arith.constant 42 : i32
+// CHECK: %[[IF:.*]] = scf.if %[[COND]] -> (i32) {
+// CHECK: scf.yield %[[A]] : i32
+// CHECK: } else {
+// CHECK: scf.yield %[[C42]] : i32
+// CHECK: }
+// CHECK: return %[[IF]] : i32
+func.func @owning_access_in_scf_if(%a: i32, %cond: i1) -> i32 {
+ %c42 = arith.constant 42 : i32
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %c42, %slot[] : memref<i32>
+ scf.if %cond {
+ test.owning_access {
+ memref.store %a, %slot[] : memref<i32>
+ test.owning_access_yield
+ }
+ }
+ %v = memref.load %slot[] : memref<i32>
+ return %v : i32
+}
+
+// -----
+
+// Resolve the alias through the wrapped load's operands. The wrapper has no
+// pointer operand; using the root slot would give the f32 load an i32 slot.
+
+// CHECK-LABEL: func.func @owning_access_through_alias
+// CHECK-SAME: (%[[A:.*]]: i32) -> f32
+// CHECK-NOT: test.owning_access
+// CHECK: %[[CAST:.*]] = builtin.unrealized_conversion_cast %[[A]] : i32 to f32
+// CHECK: return %[[CAST]] : f32
+func.func @owning_access_through_alias(%a: i32) -> f32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ %alias = test.transparent_cast_alias %slot : (memref<i32>) -> memref<f32>
+ %v = test.owning_access {
+ %l = memref.load %alias[] : memref<f32>
+ test.owning_access_yield %l : f32
+ } : f32
+ return %v : f32
+}
+
+// -----
+
+// Keep the slot in memory when the wrapper rejects blocking-use removal.
+
+// CHECK-LABEL: func.func @negative_owning_access_declines
+// CHECK: test.multi_slot_alloca
+// CHECK: test.owning_access decline
+// CHECK: memref.load
+func.func @negative_owning_access_declines(%a: i32) -> i32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ %v = test.owning_access decline {
+ %l = memref.load %slot[] : memref<i32>
+ test.owning_access_yield %l : i32
+ } : i32
+ return %v : i32
+}
+
+// -----
+
+// Keep the slot in memory when the wrapper disowns its embedded uses, leaving
+// the load in an unpromotable region.
+
+// CHECK-LABEL: func.func @negative_owning_access_disowned
+// CHECK: test.multi_slot_alloca
+// CHECK: test.owning_access disown
+// CHECK: memref.load
+func.func @negative_owning_access_disowned(%a: i32) -> i32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ memref.store %a, %slot[] : memref<i32>
+ %v = test.owning_access disown {
+ %l = memref.load %slot[] : memref<i32>
+ test.owning_access_yield %l : i32
+ } : i32
+ return %v : i32
+}
+
+// -----
+
+// A wrapper with multiple accesses cannot own their blocking uses.
+
+// CHECK-LABEL: func.func @negative_owning_access_multiple_accesses
+// CHECK: test.multi_slot_alloca
+// CHECK: test.owning_access
+// CHECK: memref.store
+// CHECK: memref.load
+func.func @negative_owning_access_multiple_accesses(%a: i32) -> i32 {
+ %slot = test.multi_slot_alloca : () -> memref<i32>
+ %v = test.owning_access {
+ memref.store %a, %slot[] : memref<i32>
+ %l = memref.load %slot[] : memref<i32>
+ test.owning_access_yield %l : i32
+ } : i32
+ return %v : i32
+}
diff --git a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
index ad08c5c6501b7f..8c9d5014250efd 100644
--- a/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
+++ b/mlir/test/lib/Dialect/Test/TestOpDefs.cpp
@@ -2114,6 +2114,70 @@ Value TestPartialAlias::projectAliasValueToSlotValue(
.getResult(0);
}
+//===----------------------------------------------------------------------===//
+// TestOwningAccess
+//===----------------------------------------------------------------------===//
+
+/// Returns the memory operation wrapped by `op`, if it is the only operation
+/// before the terminator.
+static PromotableMemOpInterface getOwnedAccess(TestOwningAccess op) {
+ Block &body = op.getBody().front();
+ if (body.getOperations().size() != 2)
+ return {};
+ return dyn_cast<PromotableMemOpInterface>(body.front());
+}
+
+bool TestOwningAccess::ownsEmbeddedBlockingUses() {
+ return !getDisown() && getOwnedAccess(*this);
+}
+
+bool TestOwningAccess::loadsFrom(const MemorySlot &slot) {
+ PromotableMemOpInterface access = getOwnedAccess(*this);
+ return access && access.loadsFrom(slot);
+}
+
+bool TestOwningAccess::storesTo(const MemorySlot &slot) {
+ PromotableMemOpInterface access = getOwnedAccess(*this);
+ return access && access.storesTo(slot);
+}
+
+Value TestOwningAccess::getStored(const MemorySlot &slot, OpBuilder &builder,
+ Value reachingDef,
+ const DataLayout &dataLayout) {
+ return getOwnedAccess(*this).getStored(slot, builder, reachingDef,
+ dataLayout);
+}
+
+bool TestOwningAccess::canUsesBeRemoved(
+ const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) {
+ if (getDecline())
+ return false;
+ PromotableMemOpInterface access = getOwnedAccess(*this);
+ if (!access)
+ return false;
+ // The wrapped access can only remove its own uses.
+ for (OpOperand *use : blockingUses)
+ if (use->getOwner() != access)
+ return false;
+ return access.canUsesBeRemoved(slot, blockingUses, newBlockingUses,
+ dataLayout);
+}
+
+DeletionKind TestOwningAccess::removeBlockingUses(
+ const MemorySlot &slot, const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder, Value reachingDefinition,
+ const DataLayout &dataLayout) {
+ getOwnedAccess(*this).removeBlockingUses(slot, blockingUses, builder,
+ reachingDefinition, dataLayout);
+ // The wrapped load's result has been replaced with the reaching definition.
+ // Replace the wrapper's results with the updated yield operands.
+ getOperation()->replaceAllUsesWith(
+ getBody().front().getTerminator()->getOperands());
+ return DeletionKind::Delete;
+}
+
//===----------------------------------------------------------------------===//
// TestSlotTracker
//===----------------------------------------------------------------------===//
diff --git a/mlir/test/lib/Dialect/Test/TestOps.td b/mlir/test/lib/Dialect/Test/TestOps.td
index f95d4419a83ad1..f8d82507873853 100644
--- a/mlir/test/lib/Dialect/Test/TestOps.td
+++ b/mlir/test/lib/Dialect/Test/TestOps.td
@@ -4609,6 +4609,28 @@ def TestPartialAlias : TEST_Op<"partial_alias",
let assemblyFormat = "$source attr-dict `:` functional-type($source, $result)";
}
+// Tests `ownsEmbeddedBlockingUses` by wrapping a single memory access and
+// forwarding promotion methods to it. `disown` disables ownership; `decline`
+// rejects removal of blocking uses.
+def TestOwningAccess : TEST_Op<"owning_access",
+ [DeclareOpInterfaceMethods<PromotableMemOpInterface,
+ ["ownsEmbeddedBlockingUses"]>,
+ SingleBlock]> {
+ let arguments = (ins UnitAttr:$disown, UnitAttr:$decline);
+ let results = (outs Variadic<AnyType>:$results);
+ let regions = (region SizedRegion<1>:$body);
+ let assemblyFormat = [{
+ (`disown` $disown^)? (`decline` $decline^)? attr-dict-with-keyword $body
+ (`:` type($results)^)?
+ }];
+}
+
+def TestOwningAccessYield : TEST_Op<"owning_access_yield",
+ [Pure, Terminator, HasParent<"TestOwningAccess">]> {
+ let arguments = (ins Variadic<AnyType>:$values);
+ let assemblyFormat = "attr-dict ($values^ `:` type($values))?";
+}
+
//===----------------------------------------------------------------------===//
// Test allocation Ops
//===----------------------------------------------------------------------===//
``````````
</details>
https://github.com/llvm/llvm-project/pull/233705
More information about the Mlir-commits
mailing list