[Mlir-commits] [mlir] [mlir][mem2reg] Promote buffers accessed via dynamic subviews, memref.copy, and masks (PR #217793)
Jianhui Li
llvmlistbot at llvm.org
Thu Aug 20 17:20:35 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/217793
**Context**
#211880 taught Mem2Reg to promote a statically-shaped memref into a single vector SSA value, but only for
whole-buffer, in-bounds, unmasked vector.transfer_read/transfer_write (plus fully-static memref.subview
sub-slices). This PR extends the vector Mem2Reg models so buffers accessed through dynamic subviews,
memref.copy, and masked transfers also promote, reconstructing any partial access with an arith.select during
promotion instead of bailing.
**Motivation**
Padding a dynamically-shaped op to a static shape and vectorizing it is a common recipe — e.g. a matmul on
dynamic shapes (dynamic M/N, contracted over a static K) whose padded result feeds a max/min reduction over the
padded dimension, where the padded tail must carry the reduction's neutral (-inf/+inf). After
one-shot-bufferize, the padded result lives in a static buffer, but everything around it is expressed in terms
of the real (dynamic) region — and that's what blocked promotion. A representative bufferized fragment:
```mlir
%buf = memref.alloc() : memref<128x64xf32> // padded (static) buffer
%realC = memref.subview %buf[0, 0] [%m, %n] [1, 1] // real region (dynamic)
: memref<128x64xf32> to memref<?x?xf32, strided<[64, 1]>>
memref.copy %C, %realC // (a) copy real input in
...
vector.transfer_write %mm, %buf[%c0, %c0] {in_bounds=[true,true]} // matmul result, whole buffer
: memref<128x?xf32, strided<[64, 1]>>, vector<128x64xf32>
```
Neither (a) nor (b) promoted before, so %buf (and the transfers around it) stayed in memory. This PR handles
both: (a) memref.copy into a dynamic subview, and (b) the out-of-bounds read/write of the padded buffer through
a dynamic subview.
**What it adds**
All in mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp:
1. Dynamic subviews. A same-rank, unit-stride, zero-offset memref.subview of a static parent with a dynamic
result is exposed as an alias whose value type is the whole parent vector (the sub-slice isn't statically
typeable). An out-of-bounds transfer of the parent extent through it expresses the dynamic valid region:
- read → select(create_mask(subview sizes), reachingDef, broadcast(padding)),
- write → select(create_mask(subview sizes), stored, reachingDef) (the aliaser's up-projection).
isPromotableTransfer allows the out-of-bounds transfer only for such an alias.
2. memref.copy. Modeled as a vector transfer of the buffer: copy into the slot is a store of
transfer_read(source); copy out is a transfer_write of the slot value into the target; a dynamic-subview
target is masked by the aliaser.
3. Masked transfers (completeness). Since the partial-access reconstruction is just a select, a masked
transfer is handled the same way: read → select(mask, reachingDef, padding), write → select(mask, stored,
reachingDef); combined with a dynamic subview the two masks are and-ed.
Renames isWholeBufferTransfer → isPromotableTransfer (it now accepts partial accesses that are composed, not
just whole ones).
>From 94e43ba4e1d5bd910d1d8e795eba1f736357300c Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 20 Aug 2026 16:44:13 +0000
Subject: [PATCH] [mlir][mem2reg] Promote buffers accessed via dynamic
subviews, memref.copy, and masks
Extends the vector Mem2Reg support added in #211880 so a statically-shaped buffer
still promotes to a vector SSA value under more access forms, reconstructing any
partial access with an arith.select during promotion:
* Dynamic subview: a same-rank, unit-stride, zero-offset memref.subview of a
static parent with a dynamic result is exposed as an alias whose element type
is the whole parent vector. An out-of-bounds transfer through it expresses the
dynamic valid region; a read becomes select(create_mask(sizes), reachingDef,
padding) and a write composes symmetrically.
* memref.copy: modeled as a vector transfer of the buffer -- copy into the slot
is a store of transfer_read(source); copy out is a load plus a transfer_write
into the target; a dynamic-subview target is masked by the aliaser. A copy
overwritten before any read becomes a dead def and is eliminated; one that is
read is preserved -- no dead-store heuristic. (Constant fills are handled by
lowering them to transfer_write via structured vectorization, which the
existing model already promotes, so they need no dedicated model.)
* Masked transfers: a masked read becomes select(mask, reachingDef, padding) and
a masked write becomes select(mask, stored, reachingDef); combined with a
dynamic subview the two masks are and-ed together.
Renames isWholeBufferTransfer to isPromotableTransfer (it now accepts partial
accesses that are composed rather than only whole-buffer ones).
Tests: mem2reg-dynamic-subview.mlir, mem2reg-copy.mlir, mem2reg-masked.mlir, and
updated masked cases in mem2reg.mlir.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 354 ++++++++++++++++--
mlir/test/Dialect/Vector/mem2reg-copy.mlir | 118 ++++++
.../Vector/mem2reg-dynamic-subview.mlir | 97 +++++
mlir/test/Dialect/Vector/mem2reg.mlir | 93 +++--
4 files changed, 585 insertions(+), 77 deletions(-)
create mode 100644 mlir/test/Dialect/Vector/mem2reg-copy.mlir
create mode 100644 mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index 1acebcb215439..54c053fe48cff 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -7,14 +7,20 @@
//===----------------------------------------------------------------------===//
//
// This file implements Mem2Reg-related interfaces that let a statically-shaped
-// memref be promoted into a single vector SSA value, provided every access to
-// the buffer is a whole-buffer read or write (or a whole-sub-region access of
-// such a buffer via a subview). With these models, Mem2Reg replaces the memory
-// slot with a vector value, threading it as the reaching definition:
+// memref be promoted into a single vector SSA value. A memref is promoted when
+// every access to it is a `vector` transfer meeting the criteria in
+// `isPromotableTransfer` (or a promotable `memref.subview` / `memref.copy`
+// built on such transfers). An access need not cover the whole buffer: a masked
+// or dynamic-subview transfer that touches only part of it is reconstructed
+// with an `arith.select` during promotion. Mem2Reg replaces the memory slot
+// with a vector value, used as its reaching definition:
//
-// * `vector.transfer_read` of the whole buffer becomes a use of the current
-// vector value; `vector.transfer_write` of the whole buffer becomes a new
-// definition of it (see the `PromotableMemOpInterface` models below).
+// * `vector.transfer_read` becomes a use of the current vector value;
+// `vector.transfer_write` becomes a new definition of it (see the
+// `PromotableMemOpInterface` models below). A masked or out-of-bounds
+// transfer reads/writes only part of the slot and is composed with an
+// `arith.select`: on the inactive lanes a read yields the transfer's
+// padding and a write keeps the reaching value.
//
// * a static, same-rank `memref.subview` is exposed as a promotable sub-slice
// alias of the buffer's slot (via `PromotableAliaserInterface`): a read of
@@ -24,14 +30,34 @@
// only ever accessed through static subviews promote as well, with partial
// and overlapping sub-writes composing in program order.
//
-// Accesses that are not whole-(sub-)buffer -- dynamic offsets, rank-reducing or
-// non-unit-stride subviews, masked or partial transfers, non-zero transfer
-// indices -- are left untouched, so the buffer is not promoted.
+// * a DYNAMIC, same-rank, zero-offset, unit-stride `memref.subview` of a
+// statically-shaped parent is exposed as an alias whose value type is the
+// WHOLE parent vector (the sub-slice is not statically typeable). An
+// out-of-bounds `vector.transfer_read`/`transfer_write` of the parent
+// extent through it is how the dynamic valid region is expressed: the read
+// masks the tail beyond the dynamic size with its padding value via
+// `vector.create_mask` + `arith.select`, and a write composes the stored
+// value onto the parent within the dynamic extent the same way. This is
+// what typically appears after bufferizing a padded, dynamically-shaped
+// op (the padded buffer is static, the real region is a dynamic subview).
+//
+// Supported subview forms (the parent buffer must be statically shaped so its
+// slot has a fixed-shape vector type):
+// - fully static sub-slice -> extract/insert_strided_slice;
+// - dynamic sub-slice (>=1 dynamic result dim), same-rank, zero-offset,
+// unit-stride -> create_mask + arith.select.
+// A dynamically-shaped parent is never promoted.
+//
+// Accesses that do not meet these criteria -- dynamic offsets, rank-reducing or
+// non-unit-stride subviews, non-zero transfer indices -- are left untouched, so
+// the memref is not promoted.
//
//===----------------------------------------------------------------------===//
#include "mlir/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.h"
+#include "mlir/Dialect/Arith/IR/Arith.h"
+#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
@@ -45,11 +71,115 @@ using namespace mlir::vector;
// Utilities
//===----------------------------------------------------------------------===//
-/// Returns whether `xferOp` accesses exactly the whole contents of `slot`, so
-/// it can act as a plain whole-buffer load/store during Mem2Reg.
+/// If `subView` is a same-rank, unit-stride, zero-offset slice of a
+/// statically-shaped parent that has at least one dynamic dimension in its
+/// result, returns the parent's vector type; otherwise returns nullopt.
+///
+/// Such a subview is exposed to Mem2Reg as an alias whose element type is the
+/// WHOLE parent vector (not the sub-slice, which is not statically typeable).
+/// An out-of-bounds `vector.transfer_read`/`transfer_write` of the parent
+/// extent through it reads/writes the whole parent value and masks the tail
+/// beyond the (dynamic) subview size with `vector.create_mask` + `arith.select`
+/// (see the transfer models and the aliaser projections below). This is the
+/// dynamic-shape counterpart of `getPromotableSubViewOffsets`, which only
+/// handles fully static sub-slices via extract/insert_strided_slice.
+static std::optional<VectorType>
+getDynamicWholeParentSubView(memref::SubViewOp subView) {
+ // The parent (slot) must be statically shaped so the slot has a fixed-shape
+ // vector type; the subview may be dynamic (the out-of-bounds transfer through
+ // it expresses the valid region).
+ auto srcType = dyn_cast<MemRefType>(subView.getSource().getType());
+ auto resType = dyn_cast<MemRefType>(subView.getResult().getType());
+ if (!srcType || !resType || !srcType.hasStaticShape())
+ return std::nullopt;
+
+ // Only dynamic sub-slices; static ones go to `getPromotableSubViewOffsets`.
+ if (resType.hasStaticShape())
+ return std::nullopt;
+
+ // Same rank: the read/write vector matches the parent rank.
+ if (srcType.getRank() != resType.getRank())
+ return std::nullopt;
+
+ // Unit strides and zero offsets: a leading, contiguous sub-region starting at
+ // the parent's origin, so the valid region on each dim is exactly [0, size).
+ for (OpFoldResult stride : subView.getMixedStrides())
+ if (getConstantIntValue(stride) != std::optional<int64_t>(1))
+ return std::nullopt;
+ for (OpFoldResult offset : subView.getMixedOffsets())
+ if (getConstantIntValue(offset) != std::optional<int64_t>(0))
+ return std::nullopt;
+
+ if (!VectorType::isValidElementType(srcType.getElementType()))
+ return std::nullopt;
+ return VectorType::get(srcType.getShape(), srcType.getElementType());
+}
+
+/// Builds a `vector.create_mask` of `vecType`'s shape marking `subView`'s sizes
+/// per dimension -- the in-bounds region of the zero-offset slice.
+static Value buildSubViewMask(OpBuilder &builder, Location loc,
+ memref::SubViewOp subView, VectorType vecType) {
+ SmallVector<Value> bounds =
+ getValueOrCreateConstantIndexOp(builder, loc, subView.getMixedSizes());
+ auto maskType = VectorType::get(vecType.getShape(), builder.getI1Type());
+ return vector::CreateMaskOp::create(builder, loc, maskType, bounds);
+}
+
+/// Reads `mem` at the origin into `vecType` (identity map), marking a dimension
+/// in-bounds only when the memref extent is statically at least the vector
+/// extent.
+static Value readMemRefAsVector(OpBuilder &builder, Location loc, Value mem,
+ VectorType vecType) {
+ auto memType = cast<MemRefType>(mem.getType());
+ int64_t rank = vecType.getRank();
+ Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
+ SmallVector<Value> indices(rank, zero);
+ Value padding = arith::ConstantOp::create(
+ builder, loc, builder.getZeroAttr(vecType.getElementType()));
+ SmallVector<bool> inBounds(rank);
+ for (int64_t d = 0; d < rank; ++d)
+ inBounds[d] = !memType.isDynamicDim(d) &&
+ memType.getDimSize(d) >= vecType.getDimSize(d);
+ return vector::TransferReadOp::create(
+ builder, loc, vecType, mem, indices,
+ AffineMapAttr::get(builder.getMultiDimIdentityMap(rank)), padding,
+ /*mask=*/Value(), builder.getBoolArrayAttr(inBounds));
+}
+
+/// Writes `vec` into `mem` at the origin (identity map), marking a dimension
+/// in-bounds only when the memref extent is statically at least the vector
+/// extent.
+static void writeVectorToMemRef(OpBuilder &builder, Location loc, Value vec,
+ Value mem) {
+ auto memType = cast<MemRefType>(mem.getType());
+ auto vecType = cast<VectorType>(vec.getType());
+ int64_t rank = vecType.getRank();
+ Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
+ SmallVector<Value> indices(rank, zero);
+ SmallVector<bool> inBounds(rank);
+ for (int64_t d = 0; d < rank; ++d)
+ inBounds[d] = !memType.isDynamicDim(d) &&
+ memType.getDimSize(d) >= vecType.getDimSize(d);
+ vector::TransferWriteOp::create(
+ builder, loc, vec, mem, indices,
+ AffineMapAttr::get(builder.getMultiDimIdentityMap(rank)),
+ /*mask=*/Value(), builder.getBoolArrayAttr(inBounds));
+}
+
+/// Returns whether `xferOp` can be promoted to a load/store of `slot`'s vector
+/// value. This requires that the transfer's sole use of the slot is as its
+/// base, the transferred vector type equals `slot.elemType`, the indices are
+/// all zero (origin), and the permutation map is the identity.
+///
+/// Two forms of partial access are accepted (rather than rejected) and
+/// reconstructed with a `select` during promotion (see the transfer models):
+/// - a masked transfer, and
+/// - an out-of-bounds transfer through a dynamic-subview alias.
+/// Their active lanes take the reaching value; the inactive lanes take the
+/// transfer's padding (read) or keep the reaching value (write).
static bool
-isWholeBufferTransfer(VectorTransferOpInterface xferOp, const MemorySlot &slot,
- const SmallPtrSetImpl<OpOperand *> &blockingUses) {
+isPromotableTransfer(VectorTransferOpInterface xferOp, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses) {
// The sole blocking use must be the slot pointer as the transfer's base.
if (blockingUses.size() != 1)
return false;
@@ -76,16 +206,13 @@ isWholeBufferTransfer(VectorTransferOpInterface xferOp, const MemorySlot &slot,
if (!xferOp.getPermutationMap().isIdentity())
return false;
- // All dimensions must be in bounds. An out-of-bounds dimension means the
- // transfer reaches past the buffer, so a read would materialize padding
- // rather than buffer contents and a write would only cover part of the
- // buffer: in neither case does the transfer stand in for the whole slot.
- if (xferOp.hasOutOfBoundsDim())
- return false;
-
- // A mask could make the access partial.
- if (xferOp.getMask())
- return false;
+ // Out-of-bounds is allowed only for a dynamic-subview alias, whose tail is
+ // masked in during promotion; otherwise it would cover only part of the slot.
+ if (xferOp.hasOutOfBoundsDim()) {
+ auto subView = slot.ptr.getDefiningOp<memref::SubViewOp>();
+ if (!subView || !getDynamicWholeParentSubView(subView))
+ return false;
+ }
return true;
}
@@ -114,19 +241,38 @@ struct TransferReadOpMemOpModel
const SmallPtrSetImpl<OpOperand *> &blockingUses,
SmallVectorImpl<OpOperand *> &newBlockingUses,
const DataLayout &dataLayout) const {
- return isWholeBufferTransfer(cast<VectorTransferOpInterface>(op), slot,
- blockingUses);
+ return isPromotableTransfer(cast<VectorTransferOpInterface>(op), slot,
+ blockingUses);
}
+ // Replaces the read with the reaching value, masking in the transfer's
+ // padding on lanes it does not read -- bounded by the dynamic-subview extent
+ // and/or the transfer's own mask (neither: the read covers the whole slot).
DeletionKind
removeBlockingUses(Operation *op, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses,
OpBuilder &builder, Value reachingDefinition,
const DataLayout &dataLayout) const {
- // Whole-buffer read: replace the loaded vector with the reaching
- // definition.
- cast<vector::TransferReadOp>(op).getVector().replaceAllUsesWith(
- reachingDefinition);
+ auto readOp = cast<vector::TransferReadOp>(op);
+ Location loc = op->getLoc();
+ Value mask;
+ if (auto subView = slot.ptr.getDefiningOp<memref::SubViewOp>())
+ if (std::optional<VectorType> vecType =
+ getDynamicWholeParentSubView(subView))
+ mask = buildSubViewMask(builder, loc, subView, *vecType);
+ if (Value opMask = readOp.getMask())
+ mask = mask
+ ? arith::AndIOp::create(builder, loc, mask, opMask).getResult()
+ : opMask;
+
+ Value result = reachingDefinition;
+ if (mask) {
+ Value padSplat = vector::BroadcastOp::create(
+ builder, loc, readOp.getVectorType(), readOp.getPadding());
+ result = arith::SelectOp::create(builder, loc, mask, reachingDefinition,
+ padSplat);
+ }
+ readOp.getVector().replaceAllUsesWith(result);
return DeletionKind::Delete;
}
};
@@ -142,16 +288,23 @@ struct TransferWriteOpMemOpModel
Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
Value reachingDef, const DataLayout &dataLayout) const {
- return cast<vector::TransferWriteOp>(op).getValueToStore();
+ auto writeOp = cast<vector::TransferWriteOp>(op);
+ Value stored = writeOp.getValueToStore();
+ // Compose the transfer's own mask here (a masked write updates only its
+ // active lanes). The dynamic-subview extent is composed separately, by the
+ // aliaser's projectAliasValueToSlotValue.
+ if (Value mask = writeOp.getMask())
+ stored = arith::SelectOp::create(builder, op->getLoc(), mask, stored,
+ reachingDef);
+ return stored;
}
bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses,
SmallVectorImpl<OpOperand *> &newBlockingUses,
const DataLayout &dataLayout) const {
- // No self-store guard needed: a vector value can never equal a memref slot.
- return isWholeBufferTransfer(cast<VectorTransferOpInterface>(op), slot,
- blockingUses);
+ return isPromotableTransfer(cast<VectorTransferOpInterface>(op), slot,
+ blockingUses);
}
DeletionKind
@@ -163,6 +316,85 @@ struct TransferWriteOpMemOpModel
}
};
+/// Mem2Reg model for `memref.copy`.
+///
+/// Mem2Reg turns a memref slot into one vector SSA value and tracks which value
+/// the buffer holds at each point. Three hooks drive this: (1)
+/// `canUsesBeRemoved` checks every access is one we can handle (otherwise the
+/// buffer stays in memory); (2) `getStored`, called at each op that writes the
+/// buffer, returns the vector value it stores -- this becomes the buffer's
+/// value from then on; (3) `removeBlockingUses` rewrites each access to use
+/// that vector value instead of the memref. A `memref.copy` fits this as a
+/// vector transfer of the value:
+/// * copy INTO the slot (target == slot): `getStored` reads the source into a
+/// vector (`vector.transfer_read`); that becomes the slot's value, and the
+/// copy is deleted.
+/// * copy OUT of the slot (source == slot): the slot's value is written to
+/// the
+/// target with a `vector.transfer_write` -- the copy becomes that write.
+/// A copy between two slots promotes one slot at a time, in either order (each
+/// promotion rewrites its own side); a dynamic-subview target is masked by the
+/// aliaser's projections.
+struct CopyOpMemOpModel
+ : public PromotableMemOpInterface::ExternalModel<CopyOpMemOpModel,
+ memref::CopyOp> {
+ bool loadsFrom(Operation *op, const MemorySlot &slot) const {
+ return cast<memref::CopyOp>(op).getSource() == slot.ptr;
+ }
+
+ bool storesTo(Operation *op, const MemorySlot &slot) const {
+ return cast<memref::CopyOp>(op).getTarget() == slot.ptr;
+ }
+
+ Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
+ Value reachingDef, const DataLayout &dataLayout) const {
+ // Only reached when storing into the slot (target == slot.ptr): the value
+ // is the whole source read as a vector matching the slot's element type.
+ auto copyOp = cast<memref::CopyOp>(op);
+ return readMemRefAsVector(builder, op->getLoc(), copyOp.getSource(),
+ cast<VectorType>(slot.elemType));
+ }
+
+ bool canUsesBeRemoved(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) const {
+ auto copyOp = cast<memref::CopyOp>(op);
+ auto vecType = dyn_cast<VectorType>(slot.elemType);
+ if (!vecType || vecType.isScalable())
+ return false;
+ bool srcIsSlot = copyOp.getSource() == slot.ptr;
+ bool dstIsSlot = copyOp.getTarget() == slot.ptr;
+ // Exactly one side must be this slot. A self-copy (both sides the slot) is
+ // not modeled here.
+ if (srcIsSlot == dstIsSlot)
+ return false;
+ // memref.copy requires both operands to have the same shape and element
+ // type, so the other side already matches the slot's extent (a plain slot
+ // is fully covered; a dynamic-subview alias covers its sub-region and is
+ // masked by the aliaser). We only need it to be a rank-matching memref; if
+ // it is itself a promotable slot, the transfer emitted here becomes a use
+ // of that slot and is resolved when it is promoted.
+ Value other = srcIsSlot ? copyOp.getTarget() : copyOp.getSource();
+ auto otherType = dyn_cast<MemRefType>(other.getType());
+ return otherType && otherType.getRank() == vecType.getRank();
+ }
+
+ DeletionKind
+ removeBlockingUses(Operation *op, const MemorySlot &slot,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder, Value reachingDefinition,
+ const DataLayout &dataLayout) const {
+ auto copyOp = cast<memref::CopyOp>(op);
+ // Copy out of the slot: materialize the slot value into the target memref.
+ // (Copy into the slot is captured through getStored; nothing to emit.)
+ if (copyOp.getSource() == slot.ptr)
+ writeVectorToMemRef(builder, op->getLoc(), reachingDefinition,
+ copyOp.getTarget());
+ return DeletionKind::Delete;
+ }
+};
+
} // namespace
//===----------------------------------------------------------------------===//
@@ -170,7 +402,7 @@ struct TransferWriteOpMemOpModel
//===----------------------------------------------------------------------===//
/// Returns the offsets of `subView` as a static, contiguous, same-rank slice of
-/// its source, or nullopt if the subview is not promotable as a whole-buffer
+/// its source, or nullopt if the subview is not promotable as a static
/// sub-slice. Promotion projects the parent buffer's vector value through
/// `vector.extract_strided_slice` / `insert_strided_slice`, which require:
/// * fully static offsets and sizes,
@@ -218,10 +450,23 @@ getPromotableSubViewOffsets(memref::SubViewOp subView) {
namespace {
-/// Exposes a static, same-rank `memref.subview` as a sub-slice alias of a
-/// whole-buffer vector slot. Reads of the subview become
-/// `vector.extract_strided_slice` of the parent value; writes become
-/// `vector.insert_strided_slice` into the current reaching definition.
+/// Exposes a same-rank `memref.subview` as a sub-slice alias of a vector slot,
+/// so a buffer accessed through subviews still promotes. When an access (a
+/// transfer or a copy) goes through the subview, Mem2Reg converts between the
+/// parent value and the alias value with two hooks, run around that access's
+/// own mem-op hooks:
+/// * a load reads the parent value projected DOWN to the alias
+/// (`projectSlotValueToAliasValue`);
+/// * a store runs down-project -> `getStored` -> up-project: the parent value
+/// is projected down to feed `getStored`'s `reachingDef`, then
+/// `getStored`'s result is projected UP to the parent
+/// (`projectAliasValueToSlotValue`).
+/// The projections depend on the subview's shape:
+/// * static sub-slice: `extract_strided_slice` (down) /
+/// `insert_strided_slice`
+/// (up);
+/// * dynamic sub-slice: identity (down) / `select(create_mask(sizes), value,
+/// reachingDef)` (up) -- see `getDynamicWholeParentSubView`.
struct SubViewOpAliasModel
: public PromotableAliaserInterface::ExternalModel<SubViewOpAliasModel,
memref::SubViewOp> {
@@ -233,12 +478,19 @@ struct SubViewOpAliasModel
if (aliasedSlotPointerOperand.get() != subView.getSource())
return;
- // The parent slot must promote to a vector (whole-buffer promotion). A
- // scalar (single-element) parent slot cannot be sliced.
+ // The parent slot must promote to a vector. A scalar (single-element)
+ // parent slot cannot be sliced.
auto parentVecType = dyn_cast<VectorType>(parentSlot.elemType);
if (!parentVecType)
return;
+ // Dynamic sub-region: the alias exposes the WHOLE parent vector; readers /
+ // writers mask the tail beyond the dynamic size with arith.select.
+ if (getDynamicWholeParentSubView(subView)) {
+ newSlots.push_back(MemorySlot{subView.getResult(), parentVecType});
+ return;
+ }
+
if (!getPromotableSubViewOffsets(subView))
return;
@@ -258,6 +510,14 @@ struct SubViewOpAliasModel
Value slotValue,
OpBuilder &builder) const {
auto subView = cast<memref::SubViewOp>(op);
+ // Dynamic sub-region alias holds the whole parent value, so the down
+ // projection is the identity. Masking is not applied here: it needs the
+ // consuming read's padding value, which this hook cannot see (one alias
+ // feeds reads with different paddings), so it is applied per-read in
+ // `TransferReadOpMemOpModel::removeBlockingUses` (create_mask + select).
+ if (getDynamicWholeParentSubView(subView))
+ return slotValue;
+
SmallVector<int64_t> offsets = *getPromotableSubViewOffsets(subView);
auto aliasVecType = cast<VectorType>(aliasSlot.elemType);
SmallVector<int64_t> strides(offsets.size(), 1);
@@ -274,6 +534,17 @@ struct SubViewOpAliasModel
Value aliasValue, Value reachingDef,
OpBuilder &builder) const {
auto subView = cast<memref::SubViewOp>(op);
+ // Dynamic sub-region write: compose the stored value within the (dynamic)
+ // subview extent back onto the parent, keeping the reaching value
+ // elsewhere: select(create_mask(subview sizes), stored, reachingDef).
+ if (std::optional<VectorType> vecType =
+ getDynamicWholeParentSubView(subView)) {
+ Location loc = op->getLoc();
+ Value mask = buildSubViewMask(builder, loc, subView, *vecType);
+ return arith::SelectOp::create(builder, loc, mask, aliasValue,
+ reachingDef);
+ }
+
SmallVector<int64_t> offsets = *getPromotableSubViewOffsets(subView);
SmallVector<int64_t> strides(offsets.size(), 1);
return vector::InsertStridedSliceOp::create(
@@ -323,5 +594,6 @@ void mlir::vector::registerMemorySlotOpInterfaceExternalModels(
registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
memref::SubViewOp::attachInterface<SubViewOpAliasModel>(*ctx);
memref::SubViewOp::attachInterface<SubViewOpPromotableModel>(*ctx);
+ memref::CopyOp::attachInterface<CopyOpMemOpModel>(*ctx);
});
}
diff --git a/mlir/test/Dialect/Vector/mem2reg-copy.mlir b/mlir/test/Dialect/Vector/mem2reg-copy.mlir
new file mode 100644
index 0000000000000..3fef3130590ba
--- /dev/null
+++ b/mlir/test/Dialect/Vector/mem2reg-copy.mlir
@@ -0,0 +1,118 @@
+// RUN: mlir-opt %s -mem2reg -canonicalize -split-input-file | FileCheck %s
+
+// memref.copy participates in Mem2Reg as a vector transfer of the buffer.
+
+// Copy INTO the whole slot then read -> read of the source; buffer eliminated.
+// CHECK-LABEL: func.func @copy_in_whole(
+// CHECK-SAME: %[[SRC:.*]]: memref<8x16xf32>, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK: %[[R:.*]] = vector.transfer_read %[[SRC]]{{.*}} : memref<8x16xf32>, vector<8x16xf32>
+// CHECK: return %[[R]]
+func.func @copy_in_whole(%src: memref<8x16xf32>, %pad: f32) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8x16xf32>
+ memref.copy %src, %a : memref<8x16xf32> to memref<8x16xf32>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// A copy into a dynamic subview whose value IS read is preserved as a masked
+// select composing the source over the prior buffer value.
+// CHECK-LABEL: func.func @copy_live(
+// CHECK-SAME: %[[SRC:.*]]: memref<?x?xf32>, %[[N:.*]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK-DAG: %[[Z:.*]] = arith.constant dense<0.000000e+00> : vector<8x16xf32>
+// CHECK: %[[RD:.*]] = vector.transfer_read %[[SRC]]{{.*}} : memref<?x?xf32>, vector<8x16xf32>
+// CHECK: %[[M:.*]] = vector.create_mask %[[N]], %[[N]] : vector<8x16xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[RD]], %[[Z]]
+// CHECK: return %[[SEL]]
+func.func @copy_live(%src: memref<?x?xf32>, %n: index, %pad: f32) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %z = arith.constant 0.000000e+00 : f32
+ %zv = vector.broadcast %z : f32 to vector<8x16xf32>
+ %a = memref.alloca() : memref<8x16xf32>
+ vector.transfer_write %zv, %a[%c0, %c0] {in_bounds = [true, true]} : vector<8x16xf32>, memref<8x16xf32>
+ %sv = memref.subview %a[0, 0] [%n, %n] [1, 1] : memref<8x16xf32> to memref<?x?xf32, strided<[16, 1]>>
+ memref.copy %src, %sv : memref<?x?xf32> to memref<?x?xf32, strided<[16, 1]>>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Copy OUT of the slot -> a transfer_write of the slot value into the target.
+// CHECK-LABEL: func.func @copy_out(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[DST:.*]]: memref<8x16xf32>
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK: vector.transfer_write %[[V]], %[[DST]]
+func.func @copy_out(%v: vector<8x16xf32>, %dst: memref<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>
+ memref.copy %a, %dst : memref<8x16xf32> to memref<8x16xf32>
+ return
+}
+
+// -----
+
+// Copy between two promotable slots: both promote, the value threads across.
+// CHECK-LABEL: func.func @copy_between_slots(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK: return %[[V]]
+func.func @copy_between_slots(%v: vector<8x16xf32>, %pad: f32) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8x16xf32>
+ %b = memref.alloca() : memref<8x16xf32>
+ vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<8x16xf32>, memref<8x16xf32>
+ memref.copy %a, %b : memref<8x16xf32> to memref<8x16xf32>
+ %r = vector.transfer_read %b[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// NEGATIVE: a self-copy references the slot on both sides and is not modeled,
+// so the buffer is not promoted (mem2reg leaves the alloca and its transfers;
+// canonicalize then folds the self-copy away, but promotion has already bailed).
+// CHECK-LABEL: func.func @negative_self_copy(
+// CHECK: memref.alloca
+// CHECK: vector.transfer_read
+func.func @negative_self_copy(%v: 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>
+ memref.copy %a, %a : memref<8x16xf32> to memref<8x16xf32>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Copy from a dynamic subview into a dynamic subview of the slot: the source
+// subview is read and composed into the slot with the subview mask.
+// CHECK-LABEL: func.func @copy_dynsub_to_dynsub(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[SRC:.*]]: memref<8x16xf32>, %[[N:.*]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK: %[[SS:.*]] = memref.subview %[[SRC]][0, 0] [8, %[[N]]] [1, 1]
+// CHECK: %[[RD:.*]] = vector.transfer_read %[[SS]]{{.*}} {in_bounds = [true, false]}
+// CHECK: %[[M:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[RD]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @copy_dynsub_to_dynsub(%v: vector<8x16xf32>, %src: memref<8x16xf32>, %n: index, %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>
+ %ssrc = memref.subview %src[0, 0] [8, %n] [1, 1] : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+ %sdst = memref.subview %a[0, 0] [8, %n] [1, 1] : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1]>>
+ memref.copy %ssrc, %sdst : memref<8x?xf32, strided<[16, 1]>> to memref<8x?xf32, strided<[16, 1]>>
+ %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/Dialect/Vector/mem2reg-dynamic-subview.mlir b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
new file mode 100644
index 0000000000000..e340e99a573b0
--- /dev/null
+++ b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
@@ -0,0 +1,97 @@
+// RUN: mlir-opt %s -mem2reg -split-input-file | FileCheck %s
+
+// A static buffer read through a DYNAMIC subview with an out-of-bounds transfer
+// is promoted directly by mem2reg: the written vector is threaded into the read
+// and the sliced-away tail is masked in with the transfer's padding value.
+
+// CHECK-LABEL: func.func @read_dyn_subview(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.subview
+// CHECK-NOT: vector.transfer_read
+// CHECK-NOT: vector.transfer_write
+// CHECK: %[[MASK:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK: %[[PS:.*]] = vector.broadcast %[[PAD]] : f32 to vector<8x16xf32>
+// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[V]], %[[PS]]
+// CHECK: return %[[SEL]]
+func.func @read_dyn_subview(%v: vector<8x16xf32>, %n: index, %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.transfer_read %sv[%c0, %c0], %pad {in_bounds = [true, false]}
+ : memref<8x?xf32, strided<[16, 1]>>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Write through a dynamic subview then read back: the store composes onto the
+// parent within the dynamic extent (select), so a subsequent whole-buffer read
+// sees the masked value.
+
+// CHECK-LABEL: func.func @write_then_read_dyn_subview(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[W:.*]]: vector<8x16xf32>, %[[N:.*]]: index
+// CHECK-NOT: memref.alloca
+// CHECK: %[[MASK:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[W]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @write_then_read_dyn_subview(%v: vector<8x16xf32>, %w: vector<8x16xf32>,
+ %n: index, %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.transfer_write %w, %sv[%c0, %c0] {in_bounds = [true, false]}
+ : vector<8x16xf32>, memref<8x?xf32, strided<[16, 1]>>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]}
+ : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Negative: a dynamic OFFSET (not just size) subview is not promotable; the
+// buffer is left alone.
+
+// CHECK-LABEL: func.func @neg_dynamic_offset(
+// CHECK: memref.alloca
+// CHECK: memref.subview
+// CHECK: vector.transfer_read
+func.func @neg_dynamic_offset(%v: vector<8x16xf32>, %off: index, %n: index,
+ %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, %off] [8, %n] [1, 1]
+ : memref<8x16xf32> to memref<8x?xf32, strided<[16, 1], offset: ?>>
+ %r = vector.transfer_read %sv[%c0, %c0], %pad {in_bounds = [true, false]}
+ : memref<8x?xf32, strided<[16, 1], offset: ?>>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// A masked read through a dynamic subview combines both masks: the subview
+// extent (create_mask) AND the transfer's own mask.
+// 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: %[[CM:.*]] = vector.create_mask %{{.*}}, %[[N]] : vector<8x16xi1>
+// CHECK: %[[AND:.*]] = arith.andi %[[CM]], %[[M]]
+// CHECK: %[[SEL:.*]] = arith.select %[[AND]], %[[V]], %{{.*}}
+// 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.transfer_read %sv[%c0, %c0], %pad, %m {in_bounds = [true, false]} : memref<8x?xf32, strided<[16, 1]>>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
diff --git a/mlir/test/Dialect/Vector/mem2reg.mlir b/mlir/test/Dialect/Vector/mem2reg.mlir
index e80effcc49003..0623d4538ffa3 100644
--- a/mlir/test/Dialect/Vector/mem2reg.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg.mlir
@@ -81,24 +81,6 @@ func.func @negative_nonzero_index(%pad: f32) -> vector<4xf32> {
return %r : vector<4xf32>
}
-// -----
-
-// A masked transfer only touches part of the buffer: must NOT be promoted.
-
-// CHECK-LABEL: func.func @negative_masked
-// CHECK: memref.alloca
-// CHECK: vector.transfer_write
-// CHECK: vector.transfer_read
-func.func @negative_masked(%pad: f32, %m: vector<4xi1>) -> vector<4xf32> {
- %c0 = arith.constant 0 : index
- %cst = arith.constant dense<1.0> : vector<4xf32>
- %a = memref.alloca() : memref<4xf32>
- vector.transfer_write %cst, %a[%c0], %m {in_bounds = [true]} : vector<4xf32>, memref<4xf32>
- %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<4xf32>, vector<4xf32>
- return %r : vector<4xf32>
-}
-
-// -----
// A partial (out-of-bounds) transfer must NOT be promoted.
@@ -460,24 +442,6 @@ func.func @negative_subview_rank_reducing(%v: vector<4xf32>, %init: vector<2x4xf
return %r : vector<2x4xf32>
}
-// -----
-
-// A masked write into the subview is not a whole-sub-region access: not promoted.
-
-// CHECK-LABEL: func.func @negative_subview_masked
-// CHECK: memref.alloca
-// CHECK: memref.subview
-func.func @negative_subview_masked(%v: vector<4xf32>, %init: vector<8xf32>, %pad: f32, %m: vector<4xi1>) -> vector<8xf32> {
- %c0 = arith.constant 0 : index
- %a = memref.alloca() : memref<8xf32>
- vector.transfer_write %init, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
- %sv = memref.subview %a[2] [4] [1] : memref<8xf32> to memref<4xf32, strided<[1], offset: 2>>
- vector.transfer_write %v, %sv[%c0], %m {in_bounds = [true]} : vector<4xf32>, memref<4xf32, strided<[1], offset: 2>>
- %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<8xf32>, vector<8xf32>
- return %r : vector<8xf32>
-}
-
-// -----
// A prefetching software-pipelined loop. The stage buffer double-buffers the
// prefetched tile: the prefetch of the next tile is skipped on the last iteration.
@@ -607,3 +571,60 @@ func.func @gemm_k_early_exit(%A: memref<4x16xf32>, %B: memref<16x4xf32>,
vector.transfer_write %sum, %C[%c0, %c0] {in_bounds = [true, true]} : vector<4x4xf32>, memref<4x4xf32>
return
}
+
+// -----
+
+// A masked whole-buffer read promotes to select(mask, reachingDef, padding).
+// CHECK-LABEL: func.func @masked_read(
+// CHECK-SAME: %[[V:.*]]: vector<8xf32>, %[[M:.*]]: vector<8xi1>, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// 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>
+ %r = vector.transfer_read %a[%c0], %pad, %m {in_bounds = [true]} : memref<8xf32>, vector<8xf32>
+ return %r : vector<8xf32>
+}
+
+// -----
+
+// A masked whole-buffer write composes select(mask, stored, reachingDef); a
+// later read observes that composed value.
+// CHECK-LABEL: func.func @masked_write(
+// CHECK-SAME: %[[V:.*]]: vector<8xf32>, %[[W:.*]]: vector<8xf32>, %[[M:.*]]: vector<8xi1>
+// CHECK-NOT: memref.alloca
+// 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>
+ vector.transfer_write %w, %a[%c0], %m {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<8xf32>, vector<8xf32>
+ return %r : vector<8xf32>
+}
+
+// -----
+
+// A masked write into a static subview composes through the subview aliaser:
+// the region is projected out (extract_strided_slice), the mask selects the
+// stored value, and the result is inserted back (insert_strided_slice).
+// CHECK-LABEL: func.func @masked_write_static_subview(
+// CHECK-SAME: %[[V:.*]]: vector<4xf32>, %[[INIT:.*]]: vector<8xf32>, %{{.*}}: f32, %[[M:.*]]: vector<4xi1>
+// CHECK-NOT: memref.alloca
+// CHECK: %[[EX:.*]] = vector.extract_strided_slice %[[INIT]] {offsets = [2], sizes = [4]
+// CHECK: %[[SEL:.*]] = arith.select %[[M]], %[[V]], %[[EX]]
+// CHECK: %[[INS:.*]] = vector.insert_strided_slice %[[SEL]], %[[INIT]] {offsets = [2]
+// CHECK: return %[[INS]]
+func.func @masked_write_static_subview(%v: vector<4xf32>, %init: vector<8xf32>, %pad: f32, %m: vector<4xi1>) -> vector<8xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8xf32>
+ vector.transfer_write %init, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
+ %sv = memref.subview %a[2] [4] [1] : memref<8xf32> to memref<4xf32, strided<[1], offset: 2>>
+ vector.transfer_write %v, %sv[%c0], %m {in_bounds = [true]} : vector<4xf32>, memref<4xf32, strided<[1], offset: 2>>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<8xf32>, vector<8xf32>
+ return %r : vector<8xf32>
+}
More information about the Mlir-commits
mailing list