[Mlir-commits] [mlir] [mlir][mem2reg] Promote buffers accessed via dynamic subviews and memref.copy (PR #217793)
Jianhui Li
llvmlistbot at llvm.org
Wed Sep 16 07:49:07 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/217793
>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