[Mlir-commits] [mlir] [mlir][Mem2Reg] Support promotion of vector.mask-wrapped transfers (PR #233705)
Jianhui Li
llvmlistbot at llvm.org
Mon Oct 19 22:54:02 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/233705
This PR enables Mem2Reg promotion of vector.mask wrapping transfer reads and writes.
Add ownsEmbeddedBlockingUses to let an operation handle a single wrapped memory access. Promotion uses the wrapper’s block and resolves aliases through the wrapped access’s operands. Masked reads preserve passthru or padding behavior; masked writes retain previous values in disabled lanes.
assisted-by-claude
>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] 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 b2b5418e1bac1..30891edbfa970 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 3997c84963207..32128c8269b34 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 33abbf2644f0f..acccbbeb1fd2a 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 57c00aa031d37..d3dcc9a830faf 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 770c5705a7080..c6027a3ae397f 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 da39bc356a109..0f5a8011a6f7e 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 e650f7e3fb09a..c098c9aa0a88c 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 62a35c90d324e..806c54a6f7523 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 ad08c5c6501b7..8c9d5014250ef 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 f95d4419a83ad..f8d8250787385 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
//===----------------------------------------------------------------------===//
More information about the Mlir-commits
mailing list