[Mlir-commits] [mlir] [mlir][Mem2Reg] Support promotion of vector.mask-wrapped transfers (PR #233705)
Jianhui Li
llvmlistbot at llvm.org
Tue Oct 20 23:19:51 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/233705
>From 503057f2452428ca7540eb8b00b3ef44793d621c Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 20 Oct 2026 05:52:09 +0000
Subject: [PATCH 1/2] adding masked transfer support and enhance mem2reg to
support embeded blocking use interface
---
.../mlir/Interfaces/MemorySlotInterfaces.h | 4 +-
.../mlir/Interfaces/MemorySlotInterfaces.td | 19 +++
.../Transforms/MemorySlotInterfacesImpl.cpp | 85 +++++++++-
mlir/lib/Interfaces/MemorySlotInterfaces.cpp | 18 +-
mlir/lib/Transforms/Mem2Reg.cpp | 50 ++++--
mlir/test/Transforms/Mem2Reg/mem2reg.mlir | 155 ++++++++++++++++++
.../Mem2Reg/vector-dyn-subview.mlir | 66 ++++++++
mlir/test/Transforms/Mem2Reg/vector.mlir | 49 ++++--
mlir/test/lib/Dialect/Test/TestOpDefs.cpp | 64 ++++++++
mlir/test/lib/Dialect/Test/TestOps.td | 22 +++
10 files changed, 498 insertions(+), 34 deletions(-)
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/Dialect/Vector/Transforms/MemorySlotInterfacesImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotInterfacesImpl.cpp
index 33abbf2644f0fc..acccbbeb1fd2ac 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotInterfacesImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotInterfacesImpl.cpp
@@ -7,8 +7,8 @@
//===----------------------------------------------------------------------===//
//
// This file implements memory slot interface models for
-// `vector.transfer_read` and `vector.transfer_write`. These models implement
-// `PromotableMemOpInterface` so Mem2Reg can replace vector transfers with uses
+// `vector.transfer_read`, `vector.transfer_write`, and `vector.mask`. They
+// implement `PromotableMemOpInterface` to replace vector transfers with uses
// and definitions of a vector SSA value. The models for `memref.subview`
// live in the MemRef dialect.
//
@@ -21,9 +21,10 @@
// of those lanes. `arith.select` implements this behavior and also handles
// lanes disabled by a transfer's mask operand.
//
-// Transfers wrapped in `vector.mask` prevent promotion because that op does
-// not implement `PromotableRegionOpInterface`, which Mem2Reg requires to
-// promote accesses inside a region.
+// `vector.mask` owns its wrapped transfer's blocking use and is promoted as a
+// whole. Its region cannot accept extra operations, so `arith.select` is built
+// outside it. Disabled read lanes use the passthru, or padding if none is
+// given; disabled write lanes retain their previous values.
//
//===----------------------------------------------------------------------===//
@@ -189,6 +190,79 @@ struct TransferWriteOpMemOpModel
return DeletionKind::Delete;
}
};
+
+/// Promotes `vector.mask` wrapping a transfer read or write by owning the
+/// transfer's slot use and applying the wrapper's mask.
+struct MaskOpMemOpModel
+ : public PromotableMemOpInterface::ExternalModel<MaskOpMemOpModel,
+ vector::MaskOp> {
+ bool ownsEmbeddedBlockingUses(Operation *op) const {
+ return isa_and_nonnull<vector::TransferReadOp, vector::TransferWriteOp>(
+ cast<vector::MaskOp>(op).getMaskableOp());
+ }
+
+ bool loadsFrom(Operation *op, const MemorySlot &slot) const {
+ auto readOp = dyn_cast_or_null<vector::TransferReadOp>(
+ cast<vector::MaskOp>(op).getMaskableOp());
+ return readOp && readOp.getBase() == slot.ptr;
+ }
+
+ bool storesTo(Operation *op, const MemorySlot &slot) const {
+ auto writeOp = dyn_cast_or_null<vector::TransferWriteOp>(
+ cast<vector::MaskOp>(op).getMaskableOp());
+ return writeOp && writeOp.getBase() == slot.ptr;
+ }
+
+ // Disabled lanes retain the reaching definition.
+ Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
+ Value reachingDef, const DataLayout &dataLayout) const {
+ auto maskOp = cast<vector::MaskOp>(op);
+ auto writeOp = cast<vector::TransferWriteOp>(maskOp.getMaskableOp());
+ return arith::SelectOp::create(builder, op->getLoc(), maskOp.getMask(),
+ writeOp.getValueToStore(), reachingDef);
+ }
+
+ bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) const {
+ auto xferOp = dyn_cast_or_null<VectorTransferOpInterface>(
+ cast<vector::MaskOp>(op).getMaskableOp());
+ // Transfers with their own mask operand are unsupported.
+ if (!xferOp || xferOp.getMask())
+ return false;
+ return isPromotableTransfer(xferOp, slot, blockingUses);
+ }
+
+ // Replace a read with the masked reaching definition, inserting IR after
+ // `vector.mask`. Writes need no further work after `getStored`.
+ DeletionKind
+ removeBlockingUses(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder, Value reachingDefinition,
+ const DataLayout &dataLayout) const {
+ auto maskOp = cast<vector::MaskOp>(op);
+ auto readOp = dyn_cast<vector::TransferReadOp>(maskOp.getMaskableOp());
+ if (!readOp)
+ return DeletionKind::Delete;
+
+ Location loc = op->getLoc();
+ Value padSplat = vector::BroadcastOp::create(
+ builder, loc, readOp.getVectorType(), readOp.getPadding());
+ Value result = reachingDefinition;
+ // Lanes outside a dynamic view take the padding.
+ if (Value viewMask =
+ memref::buildDynamicSubViewMask(builder, loc, slot.ptr))
+ result =
+ arith::SelectOp::create(builder, loc, viewMask, result, padSplat);
+ // Disabled lanes use the passthru, or padding if none is given.
+ Value disabled = maskOp.hasPassthru() ? maskOp.getPassthru() : padSplat;
+ result = arith::SelectOp::create(builder, loc, maskOp.getMask(), result,
+ disabled);
+ maskOp->getResult(0).replaceAllUsesWith(result);
+ return DeletionKind::Delete;
+ }
+};
} // namespace
//===----------------------------------------------------------------------===//
@@ -199,5 +273,6 @@ void mlir::vector::registerMemorySlotExternalModels(DialectRegistry ®istry) {
registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
TransferReadOp::attachInterface<TransferReadOpMemOpModel>(*ctx);
TransferWriteOp::attachInterface<TransferWriteOpMemOpModel>(*ctx);
+ MaskOp::attachInterface<MaskOpMemOpModel>(*ctx);
});
}
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..c6027a3ae397f0 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -355,6 +355,16 @@ Value MemorySlotPromoter::getOrCreateDefaultValue() {
return defaultValue = allocator.getDefaultValue(slot, builder);
}
+/// Returns `op`'s parent if it owns embedded blocking uses, or `op` otherwise.
+/// Uses directly in `slotRegion` stay with `op`, since its parent is outside
+/// the slot's region.
+static Operation *getBlockingUseOwner(Operation *op, Region *slotRegion) {
+ if (op->getParentRegion() == slotRegion)
+ return 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 +383,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(), slotPtrRegion);
+ 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 +404,15 @@ 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, slotPtrRegion);
+ // Visit an owner only at its wrapped access, even if the slice also
+ // includes the owner or its terminator.
+ if (auto memOp = dyn_cast<PromotableMemOpInterface>(user);
+ memOp && memOp.ownsEmbeddedBlockingUses() &&
+ op != &user->getRegion(0).front().front())
+ continue;
// If the next operation has no blocking uses, everything is fine.
auto *blockingUsesMapIt = userToBlockingUses.find(user->getParentRegion());
if (blockingUsesMapIt == userToBlockingUses.end())
@@ -445,14 +464,17 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
return failure();
}
+ // Reject new blocking uses from owners: users of their results are absent
+ // from the slice and would escape analysis.
+ auto memOp = dyn_cast<PromotableMemOpInterface>(user);
+ if (memOp && memOp.ownsEmbeddedBlockingUses() && !newBlockingUses.empty())
+ return failure();
+
// Then, register any new blocking uses for coming operations.
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 +575,9 @@ 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(), slot.ptr.getParentRegion());
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..0f5a8011a6f7e4 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 @owning_access_declines
+// CHECK: test.multi_slot_alloca
+// CHECK: test.owning_access decline
+// CHECK: memref.load
+func.func @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 @owning_access_disowned
+// CHECK: test.multi_slot_alloca
+// CHECK: test.owning_access disown
+// CHECK: memref.load
+func.func @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 @owning_access_multiple_accesses
+// CHECK: test.multi_slot_alloca
+// CHECK: test.owning_access
+// CHECK: memref.store
+// CHECK: memref.load
+func.func @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/Transforms/Mem2Reg/vector-dyn-subview.mlir b/mlir/test/Transforms/Mem2Reg/vector-dyn-subview.mlir
index e650f7e3fb09a7..c098c9aa0a88c7 100644
--- a/mlir/test/Transforms/Mem2Reg/vector-dyn-subview.mlir
+++ b/mlir/test/Transforms/Mem2Reg/vector-dyn-subview.mlir
@@ -72,6 +72,72 @@ func.func @read_with_mask_dyn_subview(%v: vector<8x16xf32>, %n: index, %m: vecto
%r = vector.transfer_read %sv[%c0, %c0], %pad, %m {in_bounds = [true, false]} : memref<8x?xf32, strided<[16, 1]>>, vector<8x16xf32>
return %r : vector<8x16xf32>
}
+// -----
+
+// Read lanes outside the dynamic subview or disabled by `vector.mask` use
+// padding.
+// CHECK-LABEL: func.func @masked_read_dyn_subview(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[M:.*]]: vector<8x16xi1>, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.mask
+// CHECK: %[[PS:.*]] = vector.broadcast %[[PAD]] : f32 to vector<8x16xf32>
+// CHECK: %[[CM:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK: %[[IN:.*]] = arith.select %[[CM]], %[[V]], %[[PS]]
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[IN]], %[[PS]]
+// CHECK: return %[[SEL]]
+func.func @masked_read_dyn_subview(%v: vector<8x16xf32>, %n: index, %m: vector<8x16xi1>, %pad: f32) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8x16xf32>
+ vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<8x16xf32>, memref<8x16xf32>
+ %sv = memref.subview %a[0, 0] [8, %n] [1, 1] : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+ %r = vector.mask %m { vector.transfer_read %sv[%c0, %c0], %pad {in_bounds = [true, false]} : memref<8x?xf32, strided<[16, 1]>>, vector<8x16xf32> } : vector<8x16xi1> -> vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Disabled read lanes use the passthru; enabled lanes outside the dynamic
+// subview use padding.
+// CHECK-LABEL: func.func @masked_read_passthru_dyn_subview(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[M:.*]]: vector<8x16xi1>, %[[PT:.*]]: vector<8x16xf32>, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.mask
+// CHECK: %[[PS:.*]] = vector.broadcast %[[PAD]] : f32 to vector<8x16xf32>
+// CHECK: %[[CM:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK: %[[IN:.*]] = arith.select %[[CM]], %[[V]], %[[PS]]
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[IN]], %[[PT]]
+// CHECK: return %[[SEL]]
+func.func @masked_read_passthru_dyn_subview(%v: vector<8x16xf32>, %n: index, %m: vector<8x16xi1>, %pt: vector<8x16xf32>, %pad: f32) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8x16xf32>
+ vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<8x16xf32>, memref<8x16xf32>
+ %sv = memref.subview %a[0, 0] [8, %n] [1, 1] : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+ %r = vector.mask %m, %pt { vector.transfer_read %sv[%c0, %c0], %pad {in_bounds = [true, false]} : memref<8x?xf32, strided<[16, 1]>>, vector<8x16xf32> } : vector<8x16xi1> -> vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Update only write lanes enabled by the mask and inside the dynamic subview.
+// Preserve the parent's previous values elsewhere.
+// CHECK-LABEL: func.func @masked_write_dyn_subview(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[W:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[M:.*]]: vector<8x16xi1>
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.mask
+// CHECK: %[[STORED:.*]] = arith.select %[[M]], %[[W]], %[[V]]
+// CHECK: %[[CM:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[STORED]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @masked_write_dyn_subview(%v: vector<8x16xf32>, %w: vector<8x16xf32>, %n: index, %m: vector<8x16xi1>, %pad: f32) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8x16xf32>
+ vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<8x16xf32>, memref<8x16xf32>
+ %sv = memref.subview %a[0, 0] [8, %n] [1, 1] : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+ vector.mask %m { vector.transfer_write %w, %sv[%c0, %c0] {in_bounds = [true, false]} : vector<8x16xf32>, memref<8x?xf32, strided<[16, 1]>> } : vector<8x16xi1>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
// -----
diff --git a/mlir/test/Transforms/Mem2Reg/vector.mlir b/mlir/test/Transforms/Mem2Reg/vector.mlir
index 62a35c90d324e3..806c54a6f75233 100644
--- a/mlir/test/Transforms/Mem2Reg/vector.mlir
+++ b/mlir/test/Transforms/Mem2Reg/vector.mlir
@@ -659,15 +659,17 @@ func.func @write_with_mask_static_subview(%v: vector<4xf32>, %init: vector<8xf32
// -----
-// NEGATIVE: the masked form (`vector.mask` wrapping the transfer) is not
-// promoted. The transfer's use of the buffer is inside a region, and Mem2Reg
-// only promotes through regions whose parent op implements
-// `PromotableRegionOpInterface`, which `vector.mask` does not.
+// Promote `vector.mask` and its wrapped read together. Disabled lanes use the
+// read's padding.
-// CHECK-LABEL: func.func @negative_masked_read(
-// CHECK: memref.alloca
-// CHECK: vector.mask
-func.func @negative_masked_read(%v: vector<8xf32>, %m: vector<8xi1>, %pad: f32) -> vector<8xf32> {
+// CHECK-LABEL: func.func @masked_read(
+// CHECK-SAME: %[[V:.*]]: vector<8xf32>, %[[M:.*]]: vector<8xi1>, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.mask
+// CHECK: %[[PS:.*]] = vector.broadcast %[[PAD]] : f32 to vector<8xf32>
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[V]], %[[PS]]
+// CHECK: return %[[SEL]]
+func.func @masked_read(%v: vector<8xf32>, %m: vector<8xi1>, %pad: f32) -> vector<8xf32> {
%c0 = arith.constant 0 : index
%a = memref.alloca() : memref<8xf32>
vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
@@ -677,12 +679,33 @@ func.func @negative_masked_read(%v: vector<8xf32>, %m: vector<8xi1>, %pad: f32)
// -----
-// NEGATIVE: a write wrapped in vector.mask also prevents buffer promotion.
+// Disabled read lanes use the passthru when provided.
-// CHECK-LABEL: func.func @negative_masked_write(
-// CHECK: memref.alloca
-// CHECK: vector.mask
-func.func @negative_masked_write(%v: vector<8xf32>, %w: vector<8xf32>, %m: vector<8xi1>, %pad: f32) -> vector<8xf32> {
+// CHECK-LABEL: func.func @masked_read_passthru(
+// CHECK-SAME: %[[V:.*]]: vector<8xf32>, %[[M:.*]]: vector<8xi1>, %[[PT:.*]]: vector<8xf32>, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.mask
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[V]], %[[PT]]
+// CHECK: return %[[SEL]]
+func.func @masked_read_passthru(%v: vector<8xf32>, %m: vector<8xi1>, %pt: vector<8xf32>, %pad: f32) -> vector<8xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
+ %r = vector.mask %m, %pt { vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<8xf32>, vector<8xf32> } : vector<8xi1> -> vector<8xf32>
+ return %r : vector<8xf32>
+}
+
+// -----
+
+// Disabled write lanes retain the buffer's previous value.
+
+// CHECK-LABEL: func.func @masked_write(
+// CHECK-SAME: %[[V:.*]]: vector<8xf32>, %[[W:.*]]: vector<8xf32>, %[[M:.*]]: vector<8xi1>, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: vector.mask
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[W]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @masked_write(%v: vector<8xf32>, %w: vector<8xf32>, %m: vector<8xi1>, %pad: f32) -> vector<8xf32> {
%c0 = arith.constant 0 : index
%a = memref.alloca() : memref<8xf32>
vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
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
//===----------------------------------------------------------------------===//
>From 78d11c2eeceadb276d0b8b82557c9e939c427ef8 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 21 Oct 2026 06:18:56 +0000
Subject: [PATCH 2/2] [mlir][mem2reg] Simplify the embedded blocking-use
handling
Drop three pieces of the driver that guard against owners no operation
implements:
* `getBlockingUseOwner` no longer keeps uses directly in the slot's region
with their own operation. That only matters when an owner contains the
slot's allocation, which a single-access wrapper such as `vector.mask`
cannot. Its `slotRegion` parameter goes with it.
* The forward-slice walk no longer skips an owner except at its wrapped
access. The slice reaches the owner again through its terminator, and
through its own operands when one of them is in the slice, but asking it
again gets the same answers and records the same uses.
* Owners requesting new blocking uses are no longer rejected. No memory
operation requests any, and the interface already forbids owners from
doing so.
Prefix the owning-access tests that keep the slot in memory with
`negative_`.
Co-Authored-By: Claude Opus 5.5 (1M context) <noreply at anthropic.com>
---
mlir/lib/Transforms/Mem2Reg.cpp | 25 ++++-------------------
mlir/test/Transforms/Mem2Reg/mem2reg.mlir | 12 +++++------
2 files changed, 10 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Transforms/Mem2Reg.cpp b/mlir/lib/Transforms/Mem2Reg.cpp
index c6027a3ae397f0..7abfe029fca3da 100644
--- a/mlir/lib/Transforms/Mem2Reg.cpp
+++ b/mlir/lib/Transforms/Mem2Reg.cpp
@@ -356,11 +356,7 @@ Value MemorySlotPromoter::getOrCreateDefaultValue() {
}
/// Returns `op`'s parent if it owns embedded blocking uses, or `op` otherwise.
-/// Uses directly in `slotRegion` stay with `op`, since its parent is outside
-/// the slot's region.
-static Operation *getBlockingUseOwner(Operation *op, Region *slotRegion) {
- if (op->getParentRegion() == slotRegion)
- return op;
+static Operation *getBlockingUseOwner(Operation *op) {
auto memOp = dyn_cast_or_null<PromotableMemOpInterface>(op->getParentOp());
return memOp && memOp.ownsEmbeddedBlockingUses() ? memOp : op;
}
@@ -384,7 +380,7 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
return failure();
auto addBlockingUse = [&](OpOperand *use) {
- Operation *owner = getBlockingUseOwner(use->getOwner(), slotPtrRegion);
+ Operation *owner = getBlockingUseOwner(use->getOwner());
userToBlockingUses[owner->getParentRegion()][owner].insert(use);
};
@@ -406,13 +402,7 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
mlir::getForwardSlice(slot.ptr, &forwardSlice);
for (Operation *op : forwardSlice) {
// Each owner wraps a single access, so replacing it preserves slice order.
- Operation *user = getBlockingUseOwner(op, slotPtrRegion);
- // Visit an owner only at its wrapped access, even if the slice also
- // includes the owner or its terminator.
- if (auto memOp = dyn_cast<PromotableMemOpInterface>(user);
- memOp && memOp.ownsEmbeddedBlockingUses() &&
- op != &user->getRegion(0).front().front())
- continue;
+ 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())
@@ -464,12 +454,6 @@ LogicalResult MemorySlotPromotionAnalyzer::computeBlockingUses(
return failure();
}
- // Reject new blocking uses from owners: users of their results are absent
- // from the slice and would escape analysis.
- auto memOp = dyn_cast<PromotableMemOpInterface>(user);
- if (memOp && memOp.ownsEmbeddedBlockingUses() && !newBlockingUses.empty())
- return failure();
-
// Then, register any new blocking uses for coming operations.
for (OpOperand *blockingUse : newBlockingUses) {
assert(llvm::is_contained(user->getResults(), blockingUse->get()));
@@ -576,8 +560,7 @@ MemorySlotPromotionAnalyzer::computeInfo() {
auto collectStoringBlocks = [&](Value ptr, const MemorySlot &ptrSlot) {
for (OpOperand &use : ptr.getUses()) {
// Attribute an embedded store to its blocking-use owner's block.
- Operation *user =
- getBlockingUseOwner(use.getOwner(), slot.ptr.getParentRegion());
+ 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 0f5a8011a6f7e4..40cbedec958f4f 100644
--- a/mlir/test/Transforms/Mem2Reg/mem2reg.mlir
+++ b/mlir/test/Transforms/Mem2Reg/mem2reg.mlir
@@ -656,11 +656,11 @@ func.func @owning_access_through_alias(%a: i32) -> f32 {
// Keep the slot in memory when the wrapper rejects blocking-use removal.
-// CHECK-LABEL: func.func @owning_access_declines
+// CHECK-LABEL: func.func @negative_owning_access_declines
// CHECK: test.multi_slot_alloca
// CHECK: test.owning_access decline
// CHECK: memref.load
-func.func @owning_access_declines(%a: i32) -> i32 {
+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 {
@@ -675,11 +675,11 @@ func.func @owning_access_declines(%a: i32) -> i32 {
// Keep the slot in memory when the wrapper disowns its embedded uses, leaving
// the load in an unpromotable region.
-// CHECK-LABEL: func.func @owning_access_disowned
+// CHECK-LABEL: func.func @negative_owning_access_disowned
// CHECK: test.multi_slot_alloca
// CHECK: test.owning_access disown
// CHECK: memref.load
-func.func @owning_access_disowned(%a: i32) -> i32 {
+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 {
@@ -693,12 +693,12 @@ func.func @owning_access_disowned(%a: i32) -> i32 {
// A wrapper with multiple accesses cannot own their blocking uses.
-// CHECK-LABEL: func.func @owning_access_multiple_accesses
+// 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 @owning_access_multiple_accesses(%a: i32) -> i32 {
+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>
More information about the Mlir-commits
mailing list