[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 23:14:54 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 1/4] [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>
+}
>From 32386aec45bfaf3ae14499d7ff35d3c291b1ad18 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 16 Sep 2026 23:08:19 +0000
Subject: [PATCH 2/4] [mlir][mem2reg] Support reassociated views of promotable
buffers
Extend the vector Mem2Reg models to `memref.expand_shape` and
`memref.collapse_shape`, so a buffer accessed through a reshaped view
promotes as well. A contiguous reshape keeps element order, so both alias
projections are a `vector.shape_cast`. The alias shape is reassociated from
the parent slot's vector shape rather than from the view's memref type,
which makes a reshape of a dynamically-shaped subview typeable; its mask is
reshaped along with the value, so it stays exact for any reassociation.
An expansion of a static subview is accepted when the subview is a
contiguous slab. That is decided from its strides, not from its layout being
the identity: a subview always carries a strided layout, so the previous
`isIdentity()` test rejected every one of them.
For a dynamic view the alias over-approximates the view, being shaped from
the parent's extents while the view's extent is only known at runtime. This
is sound only where the disagreeing extent cannot shift another dimension's
stride, that is, where every dimension ahead of the dynamic one is unit and
so has index zero. `neg_nonunit_dim_before_dynamic` pins the case where
violating it would silently read the wrong elements.
Also split each predicate from its data getter, so that no helper doubles as
a boolean test and a computation, and collect them in the utilities section.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 798 ++++++++++++------
mlir/test/Dialect/Vector/mem2reg-copy.mlir | 53 ++
.../Vector/mem2reg-dynamic-subview.mlir | 162 ++++
mlir/test/Dialect/Vector/mem2reg.mlir | 150 +++-
4 files changed, 903 insertions(+), 260 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index 54c053fe48cff..1e9a0b3a543f5 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -7,50 +7,62 @@
//===----------------------------------------------------------------------===//
//
// This file implements Mem2Reg-related interfaces that let a statically-shaped
-// 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:
+// memref be promoted into a single vector SSA value: `PromotableMemOpInterface`
+// models for the ops that access such a memref and `PromotableAliaserInterface`
+// models for the ops that view it. Mem2Reg calls the memory it promotes a
+// *slot*: a pointer paired with the type of the value it becomes, here a memref
+// and its vector type. Promoting a slot replaces it with that value, used as
+// the slot's reaching definition: the value the slot holds at a given program
+// point, merged with a block argument where control flow joins.
+//
+// A slot is promoted when each of its uses is an access these models can
+// rewrite: a `vector.transfer_read` or `vector.transfer_write` meeting the
+// criteria in `isPromotableTransfer`, or a `memref.copy` that has the slot as
+// its source or its target. A view of the slot (`memref.subview`,
+// `memref.expand_shape`, `memref.collapse_shape`) is allowed as well: the view
+// becomes an alias slot of its own, promoted with its parent, as long as every
+// use of the view is in turn such an access, or another such view.
+//
+// The accesses are rewritten as follows:
//
// * `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.
+// `vector.transfer_write` becomes a new definition of it. A masked transfer
+// covers only its active lanes, so it is composed with an `arith.select`.
+//
+// * `memref.copy` is modeled as a transfer of the slot value: if the slot is
+// the copy's target, the value stored into it is a `vector.transfer_read`
+// of the copy's source; if the slot is the copy's source, its value is
+// written to the target with a `vector.transfer_write`. Either way the copy
+// itself is removed. A copy between two slots is resolved one slot at a
+// time, each promotion rewriting its own side (see `CopyOpMemOpModel`).
//
-// * a static, same-rank `memref.subview` is exposed as a promotable sub-slice
-// alias of the buffer's slot (via `PromotableAliaserInterface`): a read of
-// the subview projects out of the vector value with
-// `vector.extract_strided_slice`, and a write into it composes back into
-// the value with `vector.insert_strided_slice`. This lets a buffer that is
-// only ever accessed through static subviews promote as well, with partial
-// and overlapping sub-writes composing in program order.
+// * a same-rank, unit-stride `memref.subview` is exposed as an alias slot
+// (via `PromotableAliaserInterface`), so a slot accessed through subviews
+// promotes as well:
//
-// * 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).
+// - a static sub-slice becomes an alias of that sub-vector: a read projects
+// it out of the slot value with `vector.extract_strided_slice` and a
+// write composes back into the value with `vector.insert_strided_slice`;
//
-// 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.
+// - a dynamic sub-slice has no vector type, since vector shapes must be
+// static. Its alias therefore holds the WHOLE parent vector, and accesses
+// through it are out-of-bounds transfers of that shape. Promotion masks
+// the value down to the extent carried by the subview's size operands,
+// using `vector.create_mask` and `arith.select`. Such a subview must
+// start at the buffer origin.
//
-// 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.
+// * a `memref.expand_shape` / `memref.collapse_shape` is exposed as an alias
+// of the reshaped value: a contiguous reshape keeps the element order, so
+// both projections are a `vector.shape_cast`. Reshapes compose with the
+// subview aliases above, so a reshaped dynamic view (e.g. `collapse_shape`
+// of a dynamic subview) promotes as well: its alias shape comes from the
+// parent's extents and its mask is reshaped along with the value.
+//
+// The promoted memref must be statically shaped, so that its slot has a
+// fixed-shape vector type. Accesses that do not meet the criteria above --
+// dynamic offsets, rank-reducing or non-unit-stride subviews, non-contiguous
+// reshapes, non-zero transfer indices -- are left untouched, so the memref is
+// not promoted.
//
//===----------------------------------------------------------------------===//
@@ -59,6 +71,7 @@
#include "mlir/Dialect/Arith/IR/Arith.h"
#include "mlir/Dialect/Arith/Utils/Utils.h"
#include "mlir/Dialect/MemRef/IR/MemRef.h"
+#include "mlir/Dialect/MemRef/Utils/MemRefUtils.h"
#include "mlir/Dialect/Utils/StaticValueUtils.h"
#include "mlir/Dialect/Vector/IR/VectorOps.h"
#include "mlir/IR/BuiltinTypes.h"
@@ -71,63 +84,347 @@ using namespace mlir::vector;
// Utilities
//===----------------------------------------------------------------------===//
-/// 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.
+// Defined below, after the subview helpers it builds on.
+static memref::SubViewOp getUnderlyingDynamicSubView(Value slotPtr);
+
+/// 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.
///
-/// 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).
+/// 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, possibly
+/// reached through a reassociative reshape.
+/// Their active lanes take the reaching value; the inactive lanes take the
+/// transfer's padding (read) or keep the reaching value (write).
+static bool
+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;
+ Value blockingUse = (*blockingUses.begin())->get();
+ if (blockingUse != slot.ptr || xferOp.getBase() != slot.ptr)
+ return false;
+
+ // Reject the tensor form (already implied, since slot pointers are memrefs).
+ if (!isa<MemRefType>(xferOp.getBase().getType()))
+ return false;
+
+ // Exact type match pins rank/extents/element type/scalable dims.
+ if (xferOp.getVectorType() != slot.elemType)
+ return false;
+
+ // Access must start at the buffer origin in every dimension.
+ for (Value index : xferOp.getIndices()) {
+ std::optional<int64_t> constIndex = getConstantIntValue(index);
+ if (!constIndex || *constIndex != 0)
+ return false;
+ }
+
+ // Identity map: no broadcast or transpose.
+ if (!xferOp.getPermutationMap().isIdentity())
+ return false;
+
+ // Out-of-bounds is allowed only for a dynamic view.
+ if (xferOp.hasOutOfBoundsDim() && !getUnderlyingDynamicSubView(slot.ptr))
+ return false;
+
+ return true;
+}
+
+/// Whether `subView` is a statically-shaped slice that can be aliased as the
+/// sub-vector it covers, the alias its promotion uses. Promotion projects the
+/// parent buffer's vector value through `vector.extract_strided_slice` /
+/// `insert_strided_slice`, which require:
+/// * fully static offsets and sizes,
+/// * unit strides,
+/// * no rank reduction (result rank == source rank),
+/// so a dropped or dynamic dimension disqualifies the subview. A dynamically-
+/// shaped slice aliases the whole parent buffer instead
+/// (`isAliasableDynamicShapeSubView`).
+static bool isAliasableStaticShapeSubView(memref::SubViewOp subView) {
+ auto srcType = dyn_cast<MemRefType>(subView.getSource().getType());
+ auto resType = dyn_cast<MemRefType>(subView.getResult().getType());
+ if (!srcType || !resType || !srcType.hasStaticShape() ||
+ !resType.hasStaticShape())
+ return false;
+
+ // No rank reduction: extract/insert_strided_slice operate at a single rank.
+ if (srcType.getRank() != resType.getRank())
+ return false;
+
+ // Unit strides only.
+ for (OpFoldResult stride : subView.getMixedStrides()) {
+ std::optional<int64_t> s = getConstantIntValue(stride);
+ if (!s || *s != 1)
+ return false;
+ }
+
+ // Static offsets.
+ for (OpFoldResult offset : subView.getMixedOffsets()) {
+ if (!getConstantIntValue(offset))
+ return false;
+ }
+
+ // Static sizes (already implied by the result's static shape, but the sizes
+ // must match the result shape so the slice covers exactly the subview).
+ for (auto [size, dim] :
+ llvm::zip_equal(subView.getMixedSizes(), resType.getShape())) {
+ std::optional<int64_t> s = getConstantIntValue(size);
+ if (!s || *s != dim)
+ return false;
+ }
+ return true;
+}
+
+/// The offsets at which the alias sub-vector of a statically-shaped `subView`
+/// sits within the parent slot's value, for `vector.extract_strided_slice` /
+/// `insert_strided_slice`.
+static SmallVector<int64_t> getStaticSubViewOffsets(memref::SubViewOp subView) {
+ SmallVector<int64_t> offsets;
+ for (OpFoldResult offset : subView.getMixedOffsets()) {
+ std::optional<int64_t> value = getConstantIntValue(offset);
+ assert(value && "expected a static offset");
+ offsets.push_back(*value);
+ }
+ return offsets;
+}
+
+/// Whether `subView` is a dynamically-shaped slice that can be aliased as the
+/// whole parent buffer, the alias its promotion uses. A statically-shaped slice
+/// aliases just its sub-vector instead (`isAliasableStaticShapeSubView`).
+static bool isAliasableDynamicShapeSubView(memref::SubViewOp subView) {
+ // The parent must be statically shaped so the slot has a fixed-shape vector
+ // type; only the subview's sizes may be dynamic.
auto srcType = dyn_cast<MemRefType>(subView.getSource().getType());
auto resType = dyn_cast<MemRefType>(subView.getResult().getType());
if (!srcType || !resType || !srcType.hasStaticShape())
- return std::nullopt;
+ return false;
- // Only dynamic sub-slices; static ones go to `getPromotableSubViewOffsets`.
+ // At least one dynamic result dimension; a fully static sub-slice is handled
+ // by `isAliasableStaticShapeSubView` instead.
if (resType.hasStaticShape())
- return std::nullopt;
+ return false;
// Same rank: the read/write vector matches the parent rank.
if (srcType.getRank() != resType.getRank())
- return std::nullopt;
+ return false;
// 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());
+ if (!subView.hasUnitStride() || !subView.hasZeroOffset())
+ return false;
+
+ return VectorType::isValidElementType(srcType.getElementType());
+}
+
+/// The vector type covering `subView`'s whole parent buffer, which is the type
+/// the alias slot of a dynamically-shaped subview takes.
+static VectorType getWholeParentBufferVectorType(memref::SubViewOp subView) {
+ auto parentType = cast<MemRefType>(subView.getSource().getType());
+ assert(parentType.hasStaticShape() && "expected a statically-shaped parent");
+ return VectorType::get(parentType.getShape(), parentType.getElementType());
+}
+
+/// Returns the dynamic subview `slotPtr` is a view of, looking through
+/// reassociative reshapes (`expand_shape` / `collapse_shape`), or null if there
+/// is none. A slot carries only a pointer and a vector type, so the dynamic
+/// extent is not part of it: it is recovered from this subview's size operands
+/// whenever a mask is needed. A reshaped dynamic subview is still a dynamic
+/// view of the parent: the reshape preserves element order, so the valid region
+/// is the same set of elements, just indexed differently (see
+/// `buildDynamicViewMask`).
+static memref::SubViewOp getUnderlyingDynamicSubView(Value slotPtr) {
+ Operation *def = slotPtr.getDefiningOp();
+ while (def) {
+ if (auto subView = dyn_cast<memref::SubViewOp>(def)) {
+ if (isAliasableDynamicShapeSubView(subView))
+ return subView;
+ return {};
+ }
+ if (auto expand = dyn_cast<memref::ExpandShapeOp>(def)) {
+ def = expand.getSrc().getDefiningOp();
+ continue;
+ }
+ if (auto collapse = dyn_cast<memref::CollapseShapeOp>(def)) {
+ def = collapse.getSrc().getDefiningOp();
+ continue;
+ }
+ return {};
+ }
+ return {};
}
-/// 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) {
+/// Returns whether the view lays its elements out in the same (row-major) order
+/// as `parentShape`, the parent slot's vector shape, which is what makes a
+/// `vector.shape_cast` model it. These are the conditions under which
+/// `getReassociatedShape` below is meaningful, and it may be called only when
+/// they hold.
+///
+/// Collapsing is order preserving exactly when each reassociation group is
+/// contiguous in the source; that also admits a dynamic but contiguous source,
+/// such as a dynamic subview whose trailing dimensions are full.
+static bool isOrderPreservingReshape(memref::CollapseShapeOp op,
+ ArrayRef<int64_t> parentShape) {
+ SmallVector<ReassociationIndices, 4> groups = op.getReassociationIndices();
+ // The groups must partition the parent's dimensions, which they index into.
+ int64_t numGroupedDims = 0;
+ for (ReassociationIndices group : groups) {
+ for (int64_t dim : group) {
+ if (dim < 0 || dim >= static_cast<int64_t>(parentShape.size()))
+ return false;
+ }
+ numGroupedDims += group.size();
+ }
+ if (numGroupedDims != static_cast<int64_t>(parentShape.size()))
+ return false;
+
+ return memref::CollapseShapeOp::isGuaranteedCollapsible(
+ op.getSrcType(), op.getReassociationIndices());
+}
+
+/// An expansion only subdivides dimensions, so element order carries over once
+/// the source's own elements are row-major.
+///
+/// A STATIC source must have row-major strides -- read off the strides, since a
+/// subview always carries a strided layout: `[4, 16]` of `memref<8x16xf32>`
+/// qualifies (strides `[16, 1]`), `[4, 8]` does not (strides stay `[16, 1]`
+/// where `[8, 1]` is needed). Nothing more: a static view's alias takes the
+/// view's own shape.
+///
+/// A DYNAMIC source needs no strides. It qualifies as a supported dynamic view
+/// (`getUnderlyingDynamicSubView`), whose index space coincides with its static
+/// parent's.
+///
+/// Its alias, though, is shaped from the parent's extents, as if the view were
+/// as large as the parent -- and at runtime the view can be smaller. Take a
+/// `vector<8xf32>` parent and `[2, %m]`: the alias is always a `shape_cast` to
+/// `2 x 4`, while a subview extent of 4 makes the view `2 x 2`, so `%m` is 2,
+/// not 4. Row 1 then begins at element 4 in the alias and at element 2 in the
+/// view -- the same lane reads different data, which no mask can repair.
+///
+/// So the group rules below require every dimension ahead of the dynamic one to
+/// be unit: its index is then always 0, so a wrong extent behind it shifts
+/// nothing. `[%m, 2]` passes for that reason -- nothing precedes the dynamic
+/// dimension, and the `2` after it is the same in both shapes.
+static bool isOrderPreservingReshape(memref::ExpandShapeOp op,
+ ArrayRef<int64_t> parentShape) {
+ if (!memref::isStaticShapeAndContiguousRowMajor(op.getSrcType()) &&
+ !getUnderlyingDynamicSubView(op.getSrc()))
+ return false;
+
+ ArrayRef<int64_t> resShape = cast<MemRefType>(op.getType()).getShape();
+ SmallVector<ReassociationIndices, 4> groups = op.getReassociationIndices();
+ if (groups.size() != parentShape.size())
+ return false;
+
+ for (auto [srcDim, group] : llvm::enumerate(groups)) {
+ int64_t staticProduct = 1;
+ int64_t dynamicDim = -1;
+ for (int64_t dim : group) {
+ if (dim < 0 || dim >= static_cast<int64_t>(resShape.size()))
+ return false;
+ if (!ShapedType::isDynamic(resShape[dim])) {
+ staticProduct *= resShape[dim];
+ continue;
+ }
+ if (dynamicDim >= 0)
+ return false; // More than one dynamic dimension: not determined.
+ dynamicDim = dim;
+ }
+
+ int64_t parentDim = parentShape[srcDim];
+ // Fully static group: its extents must account for the parent's exactly.
+ if (dynamicDim < 0) {
+ if (staticProduct != parentDim)
+ return false;
+ continue;
+ }
+ // Only unit dimensions may precede the dynamic one within the group.
+ for (int64_t dim : group) {
+ if (dim == dynamicDim)
+ break;
+ if (resShape[dim] != 1)
+ return false;
+ }
+ // The dynamic extent must come out of the parent's as a whole number.
+ if (staticProduct == 0 || parentDim % staticProduct != 0)
+ return false;
+ }
+ return true;
+}
+
+/// Returns the shape `parentShape` takes under the view's reassociation, i.e.
+/// the alias value's shape. It is derived from the parent slot's shape rather
+/// than the result memref so that a reshape of a dynamically-shaped view is
+/// still typeable. Requires `isOrderPreservingReshape`.
+static SmallVector<int64_t>
+getReassociatedShape(memref::CollapseShapeOp op,
+ ArrayRef<int64_t> parentShape) {
+ SmallVector<int64_t> shape;
+ for (ReassociationIndices group : op.getReassociationIndices()) {
+ int64_t size = 1;
+ for (int64_t dim : group)
+ size *= parentShape[dim];
+ shape.push_back(size);
+ }
+ return shape;
+}
+
+/// An expansion keeps its result's static extents; each dynamic one is the
+/// parent extent its group splits, divided by the group's static extents.
+static SmallVector<int64_t>
+getReassociatedShape(memref::ExpandShapeOp op, ArrayRef<int64_t> parentShape) {
+ ArrayRef<int64_t> resShape = cast<MemRefType>(op.getType()).getShape();
+ SmallVector<int64_t> shape(resShape);
+ for (auto [srcDim, group] : llvm::enumerate(op.getReassociationIndices())) {
+ int64_t staticProduct = 1;
+ int64_t dynamicDim = -1;
+ for (int64_t dim : group) {
+ if (ShapedType::isDynamic(resShape[dim]))
+ dynamicDim = dim;
+ else
+ staticProduct *= resShape[dim];
+ }
+ if (dynamicDim >= 0)
+ shape[dynamicDim] = parentShape[srcDim] / staticProduct;
+ }
+ return shape;
+}
+
+/// Builds the mask of `subView`'s valid region, expressed in `vecType`'s shape:
+/// a `vector.create_mask` of the subview's sizes in the parent's shape,
+/// reshaped with `vector.shape_cast` when the slot is a reassociated view of
+/// the parent (both describe the same elements in row-major order).
+static Value buildDynamicViewMask(OpBuilder &builder, Location loc,
+ memref::SubViewOp subView,
+ VectorType vecType) {
+ VectorType parentVecType = getWholeParentBufferVectorType(subView);
SmallVector<Value> bounds =
getValueOrCreateConstantIndexOp(builder, loc, subView.getMixedSizes());
- auto maskType = VectorType::get(vecType.getShape(), builder.getI1Type());
- return vector::CreateMaskOp::create(builder, loc, maskType, bounds);
+ Value mask = vector::CreateMaskOp::create(
+ builder, loc,
+ VectorType::get(parentVecType.getShape(), builder.getI1Type()), bounds);
+ if (parentVecType.getShape() != vecType.getShape())
+ mask = vector::ShapeCastOp::create(
+ builder, loc, VectorType::get(vecType.getShape(), builder.getI1Type()),
+ mask);
+ return mask;
}
-/// Reads `mem` at the origin into `vecType` (identity map), marking a dimension
+/// Produces the value that a `memref.copy` stores into a slot: reads `mem` at
+/// the origin into `vecType` with an identity map, marking a dimension
/// in-bounds only when the memref extent is statically at least the vector
/// extent.
+///
+/// The padding is a don't-care: `memref.copy` requires both sides to have the
+/// same shape, so the source covers the whole vector at runtime and a dynamic
+/// extent only forces a conservative out-of-bounds marking. For a
+/// dynamic-subview slot, lanes past its extent are in any case discarded by the
+/// aliaser's masked select.
static Value readMemRefAsVector(OpBuilder &builder, Location loc, Value mem,
VectorType vecType) {
auto memType = cast<MemRefType>(mem.getType());
@@ -146,11 +443,12 @@ static Value readMemRefAsVector(OpBuilder &builder, Location loc, Value mem,
/*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.
+/// Stores a slot's value to the target of a `memref.copy`: writes `vec` into
+/// `mem` at the origin, with the identity map and in-bounds rule of
+/// `readMemRefAsVector`. `mask`, when set, restricts the write to the lanes it
+/// selects.
static void writeVectorToMemRef(OpBuilder &builder, Location loc, Value vec,
- Value mem) {
+ Value mem, Value mask = {}) {
auto memType = cast<MemRefType>(mem.getType());
auto vecType = cast<VectorType>(vec.getType());
int64_t rank = vecType.getRank();
@@ -162,59 +460,8 @@ static void writeVectorToMemRef(OpBuilder &builder, Location loc, Value vec,
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
-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;
- Value blockingUse = (*blockingUses.begin())->get();
- if (blockingUse != slot.ptr || xferOp.getBase() != slot.ptr)
- return false;
-
- // Reject the tensor form (already implied, since slot pointers are memrefs).
- if (!isa<MemRefType>(xferOp.getBase().getType()))
- return false;
-
- // Exact type match pins rank/extents/element type/scalable dims.
- if (xferOp.getVectorType() != slot.elemType)
- return false;
-
- // Access must start at the buffer origin in every dimension.
- for (Value index : xferOp.getIndices()) {
- std::optional<int64_t> constIndex = getConstantIntValue(index);
- if (!constIndex || *constIndex != 0)
- return false;
- }
-
- // Identity map: no broadcast or transpose.
- if (!xferOp.getPermutationMap().isIdentity())
- 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;
+ AffineMapAttr::get(builder.getMultiDimIdentityMap(rank)), mask,
+ builder.getBoolArrayAttr(inBounds));
}
//===----------------------------------------------------------------------===//
@@ -245,9 +492,10 @@ struct TransferReadOpMemOpModel
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).
+ // Promotion replaces the slot with a vector value, so the read becomes a use
+ // of that value: the value itself, or a select bringing in the transfer's
+ // padding on the lanes it does not read -- those outside the dynamic-view
+ // extent or its own mask.
DeletionKind
removeBlockingUses(Operation *op, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses,
@@ -256,10 +504,9 @@ struct TransferReadOpMemOpModel
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 (memref::SubViewOp subView = getUnderlyingDynamicSubView(slot.ptr))
+ mask =
+ buildDynamicViewMask(builder, loc, subView, readOp.getVectorType());
if (Value opMask = readOp.getMask())
mask = mask
? arith::AndIOp::create(builder, loc, mask, opMask).getResult()
@@ -286,13 +533,15 @@ struct TransferWriteOpMemOpModel
return cast<vector::TransferWriteOp>(op).getBase() == slot.ptr;
}
+ // Promotion replaces the slot with a vector value, so the write becomes the
+ // code producing that value: the transfer's vector, or a select of it over
+ // the previous value when the write is masked.
Value getStored(Operation *op, const MemorySlot &slot, OpBuilder &builder,
Value reachingDef, const DataLayout &dataLayout) const {
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.
+ // A dynamic-view extent is composed separately, by the aliaser's
+ // `projectAliasValueToSlotValue`.
if (Value mask = writeOp.getMask())
stored = arith::SelectOp::create(builder, op->getLoc(), mask, stored,
reachingDef);
@@ -307,6 +556,9 @@ struct TransferWriteOpMemOpModel
blockingUses);
}
+ // `getStored` already produced the value and the driver threaded it to the
+ // slot's later uses, so nothing is emitted here: the write disappears along
+ // with the buffer.
DeletionKind
removeBlockingUses(Operation *op, const MemorySlot &slot,
const SmallPtrSetImpl<OpOperand *> &blockingUses,
@@ -329,9 +581,9 @@ struct TransferWriteOpMemOpModel
/// * 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.
+/// * copy OUT of the slot (source == slot): on removal the slot's value is
+/// written to the target with a `vector.transfer_write` -- the copy becomes
+/// that write, masked to the valid region if the slot is a dynamic view.
/// 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.
@@ -346,10 +598,10 @@ struct CopyOpMemOpModel
return cast<memref::CopyOp>(op).getTarget() == slot.ptr;
}
+ // Promotion replaces the slot with a vector value, so a copy into the slot
+ // becomes the code producing that value: a read of the copy's source.
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));
@@ -369,28 +621,28 @@ struct CopyOpMemOpModel
// 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();
+ // No further check: `memref.copy` verifies that both operands have the same
+ // shape and element type.
+ return true;
}
+ // A copy out of the slot is a use of the slot's vector value: it becomes a
+ // `vector.transfer_write` of that value into the copy's target.
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());
+ if (copyOp.getSource() == slot.ptr) {
+ Location loc = op->getLoc();
+ Value mask;
+ if (memref::SubViewOp subView = getUnderlyingDynamicSubView(slot.ptr))
+ mask = buildDynamicViewMask(builder, loc, subView,
+ cast<VectorType>(slot.elemType));
+ writeVectorToMemRef(builder, loc, reachingDefinition, copyOp.getTarget(),
+ mask);
+ }
return DeletionKind::Delete;
}
};
@@ -398,57 +650,36 @@ struct CopyOpMemOpModel
} // namespace
//===----------------------------------------------------------------------===//
-// memref.subview aliaser
+// memref view aliasers
//===----------------------------------------------------------------------===//
-/// 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 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,
-/// * unit strides,
-/// * no rank reduction (result rank == source rank),
-/// so a dropped or dynamic dimension disqualifies the subview.
-static std::optional<SmallVector<int64_t>>
-getPromotableSubViewOffsets(memref::SubViewOp subView) {
- auto srcType = dyn_cast<MemRefType>(subView.getSource().getType());
- auto resType = dyn_cast<MemRefType>(subView.getResult().getType());
- if (!srcType || !resType || !srcType.hasStaticShape() ||
- !resType.hasStaticShape())
- return std::nullopt;
-
- // No rank reduction: extract/insert_strided_slice operate at a single rank.
- if (srcType.getRank() != resType.getRank())
- return std::nullopt;
-
- // Unit strides only.
- for (OpFoldResult stride : subView.getMixedStrides()) {
- std::optional<int64_t> s = getConstantIntValue(stride);
- if (!s || *s != 1)
- return std::nullopt;
- }
+namespace {
- // Static offsets.
- SmallVector<int64_t> offsets;
- for (OpFoldResult offset : subView.getMixedOffsets()) {
- std::optional<int64_t> o = getConstantIntValue(offset);
- if (!o)
- return std::nullopt;
- offsets.push_back(*o);
+/// Companion `PromotableOpInterface` model for the view ops aliased below
+/// (`memref.subview`, `memref.expand_shape`, `memref.collapse_shape`): once the
+/// slot is promoted, the view has no remaining memory uses and is erased.
+template <typename OpTy>
+struct ViewOpPromotableModel
+ : public PromotableOpInterface::ExternalModel<ViewOpPromotableModel<OpTy>,
+ OpTy> {
+ bool canUsesBeRemoved(Operation *op,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ SmallVectorImpl<OpOperand *> &newBlockingUses,
+ const DataLayout &dataLayout) const {
+ // The view result is itself a blocking use of the parent slot; its own
+ // users (the transfers) are resolved through the alias projections.
+ for (OpOperand &use : op->getResult(0).getUses())
+ newBlockingUses.push_back(&use);
+ return true;
}
- // Static sizes (already implied by the result's static shape, but the sizes
- // must match the result shape so the slice covers exactly the subview).
- for (auto [size, dim] :
- llvm::zip_equal(subView.getMixedSizes(), resType.getShape())) {
- std::optional<int64_t> s = getConstantIntValue(size);
- if (!s || *s != dim)
- return std::nullopt;
+ DeletionKind
+ removeBlockingUses(Operation *op,
+ const SmallPtrSetImpl<OpOperand *> &blockingUses,
+ OpBuilder &builder) const {
+ return DeletionKind::Delete;
}
- return offsets;
-}
-
-namespace {
+};
/// 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
@@ -462,11 +693,10 @@ namespace {
/// `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`.
+/// * static sub-slice: `extract_strided_slice` down, `insert_strided_slice`
+/// up;
+/// * dynamic sub-slice: identity down, `select(create_mask(sizes), value,
+/// reachingDef)` up -- see `isAliasableDynamicShapeSubView`.
struct SubViewOpAliasModel
: public PromotableAliaserInterface::ExternalModel<SubViewOpAliasModel,
memref::SubViewOp> {
@@ -475,6 +705,8 @@ struct SubViewOpAliasModel
const MemorySlot &parentSlot,
SmallVectorImpl<MemorySlot> &newSlots) const {
auto subView = cast<memref::SubViewOp>(op);
+ // Called once per operand holding the slot pointer; only the viewed source
+ // exposes an alias.
if (aliasedSlotPointerOperand.get() != subView.getSource())
return;
@@ -484,23 +716,21 @@ struct SubViewOpAliasModel
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)) {
+ // A dynamic slice aliases the whole parent value.
+ if (isAliasableDynamicShapeSubView(subView)) {
newSlots.push_back(MemorySlot{subView.getResult(), parentVecType});
return;
}
- if (!getPromotableSubViewOffsets(subView))
- return;
-
- // The alias's value type is the sub-vector matching the subview's shape.
- auto resType = cast<MemRefType>(subView.getResult().getType());
- if (!VectorType::isValidElementType(resType.getElementType()))
- return;
- VectorType aliasVecType =
- VectorType::get(resType.getShape(), resType.getElementType());
- newSlots.push_back(MemorySlot{subView.getResult(), aliasVecType});
+ // A static slice aliases the sub-vector matching the subview's shape.
+ if (isAliasableStaticShapeSubView(subView)) {
+ auto resType = cast<MemRefType>(subView.getResult().getType());
+ if (!VectorType::isValidElementType(resType.getElementType()))
+ return;
+ VectorType aliasVecType =
+ VectorType::get(resType.getShape(), resType.getElementType());
+ newSlots.push_back(MemorySlot{subView.getResult(), aliasVecType});
+ }
}
Value projectSlotValueToAliasValue(Operation *op,
@@ -510,15 +740,13 @@ 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))
+ // Identity: the alias holds the whole parent value. Masking is left to
+ // `TransferReadOpMemOpModel::removeBlockingUses`, which unlike this hook
+ // can see the consuming read's padding value.
+ if (isAliasableDynamicShapeSubView(subView))
return slotValue;
- SmallVector<int64_t> offsets = *getPromotableSubViewOffsets(subView);
+ SmallVector<int64_t> offsets = getStaticSubViewOffsets(subView);
auto aliasVecType = cast<VectorType>(aliasSlot.elemType);
SmallVector<int64_t> strides(offsets.size(), 1);
return vector::ExtractStridedSliceOp::create(
@@ -534,18 +762,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)) {
+ // Compose the stored value over the subview's extent only, leaving the rest
+ // of the parent value untouched.
+ if (isAliasableDynamicShapeSubView(subView)) {
Location loc = op->getLoc();
- Value mask = buildSubViewMask(builder, loc, subView, *vecType);
+ VectorType parentVecType = getWholeParentBufferVectorType(subView);
+ Value mask = buildDynamicViewMask(builder, loc, subView, parentVecType);
return arith::SelectOp::create(builder, loc, mask, aliasValue,
reachingDef);
}
- SmallVector<int64_t> offsets = *getPromotableSubViewOffsets(subView);
+ SmallVector<int64_t> offsets = getStaticSubViewOffsets(subView);
SmallVector<int64_t> strides(offsets.size(), 1);
return vector::InsertStridedSliceOp::create(
builder, op->getLoc(), aliasValue, reachingDef, offsets, strides)
@@ -553,27 +780,73 @@ struct SubViewOpAliasModel
}
};
-/// Companion `PromotableOpInterface` model: once the slot is promoted, the
-/// subview has no remaining memory uses and is erased.
-struct SubViewOpPromotableModel
- : public PromotableOpInterface::ExternalModel<SubViewOpPromotableModel,
- memref::SubViewOp> {
- bool canUsesBeRemoved(Operation *op,
- const SmallPtrSetImpl<OpOperand *> &blockingUses,
- SmallVectorImpl<OpOperand *> &newBlockingUses,
- const DataLayout &dataLayout) const {
- // The subview result is itself a blocking use of the parent slot; its own
- // users (the transfers) are resolved through the alias projections.
- for (OpOperand &use : op->getResult(0).getUses())
- newBlockingUses.push_back(&use);
- return true;
+/// Exposes a `memref.expand_shape` / `memref.collapse_shape` as an alias of a
+/// vector slot: the view holds the same elements in the same order, only with a
+/// different rank, so both projections are a `vector.shape_cast` (see
+/// `isOrderPreservingReshape` for when that holds).
+///
+/// The alias shape comes from the parent slot's vector shape rather than the
+/// result memref (see `getReassociatedShape`), so a reshape of a dynamic view
+/// is typeable as well; the dynamic extent is masked where the underlying
+/// subview is composed, not here. The view covers the alias value in full, so
+/// the up projection does not need `reachingDef`.
+template <typename OpTy>
+struct ReassociativeViewOpAliasModel
+ : public PromotableAliaserInterface::ExternalModel<
+ ReassociativeViewOpAliasModel<OpTy>, OpTy> {
+ void getPromotableSlotAliases(Operation *op,
+ OpOperand &aliasedSlotPointerOperand,
+ const MemorySlot &parentSlot,
+ SmallVectorImpl<MemorySlot> &newSlots) const {
+ auto viewOp = cast<OpTy>(op);
+ // Called once per operand holding the slot pointer; only the reshaped
+ // source exposes an alias.
+ if (aliasedSlotPointerOperand.get() != viewOp.getSrc())
+ return;
+
+ // The parent slot must promote to a vector; a scalar slot has no shape to
+ // reassociate.
+ auto parentVecType = dyn_cast<VectorType>(parentSlot.elemType);
+ if (!parentVecType || parentVecType.isScalable())
+ return;
+
+ auto resType = cast<MemRefType>(viewOp.getResult().getType());
+ if (!VectorType::isValidElementType(resType.getElementType()) ||
+ resType.getElementType() != parentVecType.getElementType())
+ return;
+ if (!isOrderPreservingReshape(viewOp, parentVecType.getShape()))
+ return;
+ SmallVector<int64_t> aliasShape =
+ getReassociatedShape(viewOp, parentVecType.getShape());
+ auto aliasVecType = VectorType::get(aliasShape, resType.getElementType());
+ // Guard the `vector.shape_cast` contract; implied by the view's semantics.
+ if (aliasVecType.getNumElements() != parentVecType.getNumElements())
+ return;
+ newSlots.push_back(MemorySlot{viewOp.getResult(), aliasVecType});
}
- DeletionKind
- removeBlockingUses(Operation *op,
- const SmallPtrSetImpl<OpOperand *> &blockingUses,
- OpBuilder &builder) const {
- return DeletionKind::Delete;
+ Value projectSlotValueToAliasValue(Operation *op,
+ OpOperand & /*aliasedSlotPointerOperand*/,
+ const MemorySlot & /*parentSlot*/,
+ const MemorySlot &aliasSlot,
+ Value slotValue,
+ OpBuilder &builder) const {
+ return vector::ShapeCastOp::create(builder, op->getLoc(),
+ cast<VectorType>(aliasSlot.elemType),
+ slotValue)
+ .getResult();
+ }
+
+ Value projectAliasValueToSlotValue(Operation *op,
+ OpOperand & /*aliasedSlotPointerOperand*/,
+ const MemorySlot &parentSlot,
+ const MemorySlot & /*aliasSlot*/,
+ Value aliasValue, Value /*reachingDef*/,
+ OpBuilder &builder) const {
+ return vector::ShapeCastOp::create(builder, op->getLoc(),
+ cast<VectorType>(parentSlot.elemType),
+ aliasValue)
+ .getResult();
}
};
@@ -593,7 +866,16 @@ void mlir::vector::registerMemorySlotOpInterfaceExternalModels(
// projections build Vector ops; Vector already depends on MemRef.
registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
memref::SubViewOp::attachInterface<SubViewOpAliasModel>(*ctx);
- memref::SubViewOp::attachInterface<SubViewOpPromotableModel>(*ctx);
+ memref::SubViewOp::attachInterface<
+ ViewOpPromotableModel<memref::SubViewOp>>(*ctx);
+ memref::ExpandShapeOp::attachInterface<
+ ReassociativeViewOpAliasModel<memref::ExpandShapeOp>>(*ctx);
+ memref::ExpandShapeOp::attachInterface<
+ ViewOpPromotableModel<memref::ExpandShapeOp>>(*ctx);
+ memref::CollapseShapeOp::attachInterface<
+ ReassociativeViewOpAliasModel<memref::CollapseShapeOp>>(*ctx);
+ memref::CollapseShapeOp::attachInterface<
+ ViewOpPromotableModel<memref::CollapseShapeOp>>(*ctx);
memref::CopyOp::attachInterface<CopyOpMemOpModel>(*ctx);
});
}
diff --git a/mlir/test/Dialect/Vector/mem2reg-copy.mlir b/mlir/test/Dialect/Vector/mem2reg-copy.mlir
index 3fef3130590ba..e68d5f8ed38ea 100644
--- a/mlir/test/Dialect/Vector/mem2reg-copy.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg-copy.mlir
@@ -116,3 +116,56 @@ func.func @copy_dynsub_to_dynsub(%v: vector<8x16xf32>, %src: 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 global into a dynamic subview of a static local buffer (an "SLM"
+// staging buffer): the global is not promotable, so the copy becomes a read of
+// the global's slice, composed onto the buffer's value over the subview extent.
+// The read is not in bounds on the dynamic dimension and its padding is a
+// don't-care: those lanes keep the buffer's previous value through the select.
+// The second function repeats it in a non-default memory space.
+memref.global "private" constant @tile : memref<128x64xf16> = dense<1.000000e+00>
+
+// CHECK-LABEL: func.func @copy_global_to_dyn_subview(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f16
+// CHECK: %[[G:.*]] = memref.get_global @tile : memref<128x64xf16>
+// CHECK: %[[GSV:.*]] = memref.subview %[[G]][0, 0] [%[[N]], 64] [1, 1]
+// CHECK: %[[RD:.*]] = vector.transfer_read %[[GSV]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [false, true]} : memref<?x64xf16, strided<[64, 1]>>, vector<128x64xf16>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<128x64xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[RD]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @copy_global_to_dyn_subview(%v: vector<128x64xf16>, %n: index, %pad: f16) -> vector<128x64xf16> {
+ %c0 = arith.constant 0 : index
+ %g = memref.get_global @tile : memref<128x64xf16>
+ %gsv = memref.subview %g[0, 0] [%n, 64] [1, 1] : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
+ %slm = memref.alloca() : memref<128x64xf16>
+ vector.transfer_write %v, %slm[%c0, %c0] {in_bounds = [true, true]} : vector<128x64xf16>, memref<128x64xf16>
+ %sv = memref.subview %slm[0, 0] [%n, 64] [1, 1] : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
+ memref.copy %gsv, %sv : memref<?x64xf16, strided<[64, 1]>> to memref<?x64xf16, strided<[64, 1]>>
+ %r = vector.transfer_read %slm[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16>, vector<128x64xf16>
+ return %r : vector<128x64xf16>
+}
+
+// CHECK-LABEL: func.func @copy_global_to_dyn_subview_memspace(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK: %[[RD:.*]] = vector.transfer_read %{{.*}} {in_bounds = [false, true]} : memref<?x64xf16, strided<[64, 1]>>, vector<128x64xf16>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<128x64xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[RD]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @copy_global_to_dyn_subview_memspace(%v: vector<128x64xf16>, %n: index, %pad: f16) -> vector<128x64xf16> {
+ %c0 = arith.constant 0 : index
+ %g = memref.get_global @tile : memref<128x64xf16>
+ %gsv = memref.subview %g[0, 0] [%n, 64] [1, 1] : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
+ %slm = memref.alloca() : memref<128x64xf16, 3>
+ vector.transfer_write %v, %slm[%c0, %c0] {in_bounds = [true, true]} : vector<128x64xf16>, memref<128x64xf16, 3>
+ %sv = memref.subview %slm[0, 0] [%n, 64] [1, 1] : memref<128x64xf16, 3> to memref<?x64xf16, strided<[64, 1]>, 3>
+ memref.copy %gsv, %sv : memref<?x64xf16, strided<[64, 1]>> to memref<?x64xf16, strided<[64, 1]>, 3>
+ %r = vector.transfer_read %slm[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16, 3>, vector<128x64xf16>
+ return %r : vector<128x64xf16>
+}
diff --git a/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
index e340e99a573b0..38ff7b20e8b2b 100644
--- a/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
@@ -95,3 +95,165 @@ func.func @masked_read_dyn_subview(%v: vector<8x16xf32>, %n: index, %m: vector<8
%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>
}
+
+// -----
+
+// A dynamic subview reshaped by `memref.collapse_shape`: the reshape keeps the
+// whole parent value, so the read is the parent value shape_cast to the collapsed
+// shape, masked by the subview extent -- the mask is built in the parent's shape
+// and shape_cast along with the value (both describe the same elements).
+// CHECK-LABEL: func.func @collapse_shape_dyn_subview(
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.collapse_shape
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[V]] : vector<8x16xf32> to vector<128xf32>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<8x16xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[CM]] : vector<8x16xi1> to vector<128xi1>
+// CHECK: %[[PS:.*]] = vector.broadcast %[[PAD]] : f32 to vector<128xf32>
+// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[SC]], %[[PS]]
+// CHECK: return %[[SEL]]
+func.func @collapse_shape_dyn_subview(%v: vector<8x16xf32>, %n: index, %pad: f32) -> vector<128xf32> {
+ %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] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
+ %flat = memref.collapse_shape %sv [[0, 1]] : memref<?x16xf32, strided<[16, 1]>> into memref<?xf32, strided<[1]>>
+ %r = vector.transfer_read %flat[%c0], %pad {in_bounds = [false]} : memref<?xf32, strided<[1]>>, vector<128xf32>
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// `memref.expand_shape` of a dynamic view: the dynamic result dimension is
+// recovered from the parent extent it splits (16 = ? x 2 gives 8), so the alias
+// is `vector<8x2xf32>` and the mask is reshaped alongside the value.
+// CHECK-LABEL: func.func @expand_shape_dyn_subview(
+// CHECK-SAME: %[[V:.*]]: vector<16xf32>, %[[M:.*]]: index, %[[PAD:.*]]: f32
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[V]] : vector<16xf32> to vector<8x2xf32>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[M]] : vector<16xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[CM]] : vector<16xi1> to vector<8x2xi1>
+// CHECK: %[[PS:.*]] = vector.broadcast %[[PAD]] : f32 to vector<8x2xf32>
+// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[SC]], %[[PS]]
+// CHECK: return %[[SEL]]
+func.func @expand_shape_dyn_subview(%v: vector<16xf32>, %m: index, %pad: f32) -> vector<8x2xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<16xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<16xf32>, memref<16xf32>
+ %sv = memref.subview %a[0] [%m] [1] : memref<16xf32> to memref<?xf32, strided<[1]>>
+ %exp = memref.expand_shape %sv [[0, 1]] output_shape [%m, 2] : memref<?xf32, strided<[1]>> into memref<?x2xf32>
+ %r = vector.transfer_read %exp[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x2xf32>, vector<8x2xf32>
+ return %r : vector<8x2xf32>
+}
+
+// -----
+
+// Negative: two dynamic dimensions in one reassociation group -- the alias shape
+// is not determined by the parent extent, so the buffer is not promoted.
+// CHECK-LABEL: func.func @neg_two_dynamic_dims_in_group(
+// CHECK: memref.alloca
+// CHECK: memref.expand_shape
+// CHECK: vector.transfer_read
+func.func @neg_two_dynamic_dims_in_group(%v: vector<16xf32>, %m: index, %d0: index,
+ %d1: index, %pad: f32) -> vector<8x2xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<16xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<16xf32>, memref<16xf32>
+ %sv = memref.subview %a[0] [%m] [1] : memref<16xf32> to memref<?xf32, strided<[1]>>
+ %exp = memref.expand_shape %sv [[0, 1]] output_shape [%d0, %d1] : memref<?xf32, strided<[1]>> into memref<?x?xf32>
+ %r = vector.transfer_read %exp[%c0, %c0], %pad : memref<?x?xf32>, vector<8x2xf32>
+ return %r : vector<8x2xf32>
+}
+
+// -----
+
+// Negative: a non-unit dimension precedes the dynamic one in its group, so the
+// view and the alias disagree on that dimension's stride -- the view strides by
+// the runtime %m, while an alias shaped from the parent strides by 8 / 2 = 4.
+// Promoting would move data rather than merely mask it: for a buffer holding
+// 1..8 with %n = 4 (so %m = 2), the view reads
+// ((1, 2, pad, pad), (3, 4, pad, pad)) whereas a masked shape_cast of the parent
+// value yields ((1, 2, 3, 4), (pad, pad, pad, pad)).
+// CHECK-LABEL: func.func @neg_nonunit_dim_before_dynamic(
+// CHECK: memref.alloca
+// CHECK: memref.subview
+// CHECK: memref.expand_shape
+// CHECK: vector.transfer_read
+func.func @neg_nonunit_dim_before_dynamic(%v: vector<8xf32>, %n: index, %m: index,
+ %pad: f32) -> vector<2x4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
+ %sv = memref.subview %a[0] [%n] [1] : memref<8xf32> to memref<?xf32, strided<[1]>>
+ %exp = memref.expand_shape %sv [[0, 1]] output_shape [2, %m] : memref<?xf32, strided<[1]>> into memref<2x?xf32>
+ %r = vector.transfer_read %exp[%c0, %c0], %pad {in_bounds = [true, false]} : memref<2x?xf32>, vector<2x4xf32>
+ return %r : vector<2x4xf32>
+}
+
+// -----
+
+// The reported bufferization pattern: a padded static buffer, a dynamic subview
+// of its real region, expanded to a higher rank (with a dynamic dim), copied out
+// to another buffer. The whole chain promotes: the parent value is shape_cast to
+// the expanded shape and written masked to the copy's valid region.
+// CHECK-LABEL: func.func @expand_shape_dyn_subview_copy_out(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.expand_shape
+// CHECK-NOT: memref.copy
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[V]] : vector<128x64xf16> to vector<1x1x128x64xf16>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<128x64xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[CM]] : vector<128x64xi1> to vector<1x1x128x64xi1>
+// CHECK: vector.transfer_write %[[SC]], %{{.*}}[%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}], %[[MASK]] {in_bounds = [true, true, false, true]} : vector<1x1x128x64xf16>, memref<1x1x?x64xf16, {{.*}}>
+func.func @expand_shape_dyn_subview_copy_out(%v: vector<128x64xf16>, %n: index,
+ %dst: memref<?x16x?x64xf16>, %i: index, %j: index, %k: index) {
+ %c0 = arith.constant 0 : index
+ %buf = memref.alloca() : memref<128x64xf16>
+ vector.transfer_write %v, %buf[%c0, %c0] {in_bounds = [true, true]} : vector<128x64xf16>, memref<128x64xf16>
+ %sv = memref.subview %buf[0, 0] [%n, 64] [1, 1]
+ : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
+ %ex = memref.expand_shape %sv [[0, 1, 2], [3]] output_shape [1, 1, %n, 64]
+ : memref<?x64xf16, strided<[64, 1]>> into memref<1x1x?x64xf16, strided<[?, ?, 64, 1]>>
+ %dsv = memref.subview %dst[%i, %j, %k, 0] [1, 1, %n, 64] [1, 1, 1, 1]
+ : memref<?x16x?x64xf16> to memref<1x1x?x64xf16, strided<[?, ?, 64, 1], offset: ?>>
+ memref.copy %ex, %dsv : memref<1x1x?x64xf16, strided<[?, ?, 64, 1]>> to memref<1x1x?x64xf16, strided<[?, ?, 64, 1], offset: ?>>
+ return
+}
+
+// -----
+
+// The same chain with the copy reversed -- loading a tile from the other buffer
+// into the staging buffer through the expanded dynamic view. The copy becomes a
+// read of the source's slice, reshaped back and composed onto the buffer's value
+// over the subview extent. The read carries a padding operand and is NOT in
+// bounds on the dynamic dimension, which is what makes reading past the source's
+// extent defined; those padded lanes are discarded by the select, so the lanes
+// outside the extent keep the buffer's previous value.
+// CHECK-LABEL: func.func @expand_shape_dyn_subview_copy_in(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.expand_shape
+// CHECK-NOT: memref.copy
+// CHECK: %[[SSV:.*]] = memref.subview %{{.*}} : memref<?x16x?x64xf16> to memref<1x1x?x64xf16, {{.*}}>
+// CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f16
+// CHECK: %[[RD:.*]] = vector.transfer_read %[[SSV]][%{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [true, true, false, true]} : memref<1x1x?x64xf16, {{.*}}>, vector<1x1x128x64xf16>
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[RD]] : vector<1x1x128x64xf16> to vector<128x64xf16>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<128x64xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[SC]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @expand_shape_dyn_subview_copy_in(%v: vector<128x64xf16>, %n: index,
+ %src: memref<?x16x?x64xf16>, %i: index, %j: index, %k: index, %pad: f16) -> vector<128x64xf16> {
+ %c0 = arith.constant 0 : index
+ %buf = memref.alloca() : memref<128x64xf16>
+ vector.transfer_write %v, %buf[%c0, %c0] {in_bounds = [true, true]} : vector<128x64xf16>, memref<128x64xf16>
+ %sv = memref.subview %buf[0, 0] [%n, 64] [1, 1]
+ : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
+ %ex = memref.expand_shape %sv [[0, 1, 2], [3]] output_shape [1, 1, %n, 64]
+ : memref<?x64xf16, strided<[64, 1]>> into memref<1x1x?x64xf16, strided<[?, ?, 64, 1]>>
+ %ssv = memref.subview %src[%i, %j, %k, 0] [1, 1, %n, 64] [1, 1, 1, 1]
+ : memref<?x16x?x64xf16> to memref<1x1x?x64xf16, strided<[?, ?, 64, 1], offset: ?>>
+ memref.copy %ssv, %ex : memref<1x1x?x64xf16, strided<[?, ?, 64, 1], offset: ?>> to memref<1x1x?x64xf16, strided<[?, ?, 64, 1]>>
+ %r = vector.transfer_read %buf[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16>, vector<128x64xf16>
+ return %r : vector<128x64xf16>
+}
diff --git a/mlir/test/Dialect/Vector/mem2reg.mlir b/mlir/test/Dialect/Vector/mem2reg.mlir
index 70be1e0742038..efd5dd298f495 100644
--- a/mlir/test/Dialect/Vector/mem2reg.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg.mlir
@@ -615,9 +615,9 @@ func.func @masked_write(%v: vector<8xf32>, %w: vector<8xf32>, %m: vector<8xi1>,
// 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: %[[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: %[[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
@@ -628,3 +628,149 @@ func.func @masked_write_static_subview(%v: vector<4xf32>, %init: vector<8xf32>,
%r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<8xf32>, vector<8xf32>
return %r : vector<8xf32>
}
+
+// -----
+
+// A `memref.collapse_shape` view is an alias of the reshaped value: it covers the
+// same elements in the same row-major order, so a read through it is the slot
+// value `vector.shape_cast` to the collapsed shape.
+
+// CHECK-LABEL: func.func @collapse_shape_read
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.collapse_shape
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[V]] : vector<8x16xf32> to vector<128xf32>
+// CHECK: return %[[SC]]
+func.func @collapse_shape_read(%v: vector<8x16xf32>, %pad: f32) -> vector<128xf32> {
+ %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>
+ %flat = memref.collapse_shape %a [[0, 1]] : memref<8x16xf32> into memref<128xf32>
+ %r = vector.transfer_read %flat[%c0], %pad {in_bounds = [true]} : memref<128xf32>, vector<128xf32>
+ return %r : vector<128xf32>
+}
+
+// -----
+
+// `memref.expand_shape` is handled the same way, in the other direction.
+
+// CHECK-LABEL: func.func @expand_shape_read
+// CHECK-SAME: %[[V:.*]]: vector<128xf32>
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[V]] : vector<128xf32> to vector<8x16xf32>
+// CHECK: return %[[SC]]
+func.func @expand_shape_read(%v: vector<128xf32>, %pad: f32) -> vector<8x16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<128xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<128xf32>, memref<128xf32>
+ %exp = memref.expand_shape %a [[0, 1]] output_shape [8, 16] : memref<128xf32> into memref<8x16xf32>
+ %r = vector.transfer_read %exp[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// A write through the view covers the whole buffer, so it redefines the slot
+// value: the stored value is reshaped back to the buffer's shape.
+
+// CHECK-LABEL: func.func @collapse_shape_write
+// CHECK-SAME: %{{.*}}: vector<8x16xf32>, %[[W:.*]]: vector<128xf32>
+// CHECK-NOT: memref.alloca
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[W]] : vector<128xf32> to vector<8x16xf32>
+// CHECK: return %[[SC]]
+func.func @collapse_shape_write(%v: vector<8x16xf32>, %w: vector<128xf32>, %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>
+ %flat = memref.collapse_shape %a [[0, 1]] : memref<8x16xf32> into memref<128xf32>
+ vector.transfer_write %w, %flat[%c0] {in_bounds = [true]} : vector<128xf32>, memref<128xf32>
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Views compose: a subview aliases the sub-slice, and collapsing it reshapes that
+// sub-slice value (the collapse is contiguous within the subview).
+
+// CHECK-LABEL: func.func @subview_then_collapse_shape
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>
+// CHECK-NOT: memref.alloca
+// CHECK: %[[EX:.*]] = vector.extract_strided_slice %[[V]] offsets = [0, 0], sizes = [4, 16]
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[EX]] : vector<4x16xf32> to vector<64xf32>
+// CHECK: return %[[SC]]
+func.func @subview_then_collapse_shape(%v: vector<8x16xf32>, %pad: f32) -> vector<64xf32> {
+ %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] [4, 16] [1, 1] : memref<8x16xf32> to memref<4x16xf32, strided<[16, 1]>>
+ %flat = memref.collapse_shape %sv [[0, 1]] : memref<4x16xf32, strided<[16, 1]>> into memref<64xf32, strided<[1]>>
+ %r = vector.transfer_read %flat[%c0], %pad {in_bounds = [true]} : memref<64xf32, strided<[1]>>, vector<64xf32>
+ return %r : vector<64xf32>
+}
+
+// -----
+
+// An expansion composes with a subview the same way. The subview's layout is
+// strided rather than identity, so what matters is that it is a contiguous slab:
+// its elements keep their row-major order under the expansion.
+
+// CHECK-LABEL: func.func @subview_then_expand_shape
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[EX:.*]] = vector.extract_strided_slice %[[V]] offsets = [0, 0], sizes = [4, 16]
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[EX]] : vector<4x16xf32> to vector<1x4x16xf32>
+// CHECK: return %[[SC]]
+func.func @subview_then_expand_shape(%v: vector<8x16xf32>, %pad: f32) -> vector<1x4x16xf32> {
+ %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] [4, 16] [1, 1] : memref<8x16xf32> to memref<4x16xf32, strided<[16, 1]>>
+ %exp = memref.expand_shape %sv [[0, 1], [2]] output_shape [1, 4, 16] : memref<4x16xf32, strided<[16, 1]>> into memref<1x4x16xf32, strided<[64, 16, 1]>>
+ %r = vector.transfer_read %exp[%c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x4x16xf32, strided<[64, 16, 1]>>, vector<1x4x16xf32>
+ return %r : vector<1x4x16xf32>
+}
+
+// -----
+
+// The slab need not start at the origin: a non-zero constant offset only shifts
+// which sub-vector the alias extracts, it does not affect element order.
+
+// CHECK-LABEL: func.func @subview_offset_then_expand_shape
+// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>
+// CHECK-NOT: memref.alloca
+// CHECK: %[[EX:.*]] = vector.extract_strided_slice %[[V]] offsets = [2, 0], sizes = [4, 16]
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[EX]] : vector<4x16xf32> to vector<1x4x16xf32>
+// CHECK: return %[[SC]]
+func.func @subview_offset_then_expand_shape(%v: vector<8x16xf32>, %pad: f32) -> vector<1x4x16xf32> {
+ %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[2, 0] [4, 16] [1, 1] : memref<8x16xf32> to memref<4x16xf32, strided<[16, 1], offset: 32>>
+ %exp = memref.expand_shape %sv [[0, 1], [2]] output_shape [1, 4, 16] : memref<4x16xf32, strided<[16, 1], offset: 32>> into memref<1x4x16xf32, strided<[64, 16, 1], offset: 32>>
+ %r = vector.transfer_read %exp[%c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x4x16xf32, strided<[64, 16, 1], offset: 32>>, vector<1x4x16xf32>
+ return %r : vector<1x4x16xf32>
+}
+
+// -----
+
+// Negative: the subview keeps only part of each row, so its rows are strided in
+// the parent (stride 16 over an 8-element row). Expanding it would not preserve
+// element order, and the buffer stays in memory.
+
+// CHECK-LABEL: func.func @neg_expand_shape_strided_subview
+// CHECK: memref.alloca
+// CHECK: memref.subview
+// CHECK: memref.expand_shape
+// CHECK: vector.transfer_read
+func.func @neg_expand_shape_strided_subview(%v: vector<8x16xf32>, %pad: f32) -> vector<1x4x8xf32> {
+ %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] [4, 8] [1, 1] : memref<8x16xf32> to memref<4x8xf32, strided<[16, 1]>>
+ %exp = memref.expand_shape %sv [[0, 1], [2]] output_shape [1, 4, 8] : memref<4x8xf32, strided<[16, 1]>> into memref<1x4x8xf32, strided<[64, 16, 1]>>
+ %r = vector.transfer_read %exp[%c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x4x8xf32, strided<[64, 16, 1]>>, vector<1x4x8xf32>
+ return %r : vector<1x4x8xf32>
+}
>From eacf88939c9ad7682eac704747ed2359eb97f6d4 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 17 Sep 2026 01:56:33 +0000
Subject: [PATCH 3/4] [mlir][mem2reg] Harden and test the vector promotion
models
Findings from a review pass over the vector Mem2Reg models. No change to
which memrefs promote, except that the file now says what it does.
A reassociation group may hold a zero static extent, which the parent extent
cannot be divided by: `memref<?x0x16xf32>` is valid IR, and without the
`staticProduct == 0` guard the expansion predicate faults. Zero extents are
screened at the allocation, but only on the buffer, whereas the divisor comes
from the view's own type. Covered by @neg_zero_extent_in_group.
Correct the header's claim that a promoted memref is always statically
shaped: a 1-D extent that is a `vector.vscale` multiple takes a scalable
slot. Such a slot is promoted through whole-buffer transfers alone, and the
static requirement belongs to the views and copies that need static extents.
Add tests for the compositions the suite did not reach: a dynamic view merged
across control flow by a block argument, a dynamic view nested on a static
subview, and two reshapes composed over a dynamic view. Add negatives
isolating the zero-offset requirement, a collapse whose group's dynamic
extent is not its outermost, and a scalable slot reaching `memref.copy`.
Rename `getWholeParentBufferVectorType` to `getWholeParentVectorType`, since
that parent is often another view rather than the buffer.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 40 +++--
mlir/test/Dialect/Vector/mem2reg-copy.mlir | 21 +++
.../Vector/mem2reg-dynamic-subview.mlir | 153 ++++++++++++++++++
3 files changed, 197 insertions(+), 17 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index 1e9a0b3a543f5..d3e835e3db050 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -6,14 +6,13 @@
//
//===----------------------------------------------------------------------===//
//
-// This file implements Mem2Reg-related interfaces that let a statically-shaped
-// memref be promoted into a single vector SSA value: `PromotableMemOpInterface`
-// models for the ops that access such a memref and `PromotableAliaserInterface`
-// models for the ops that view it. Mem2Reg calls the memory it promotes a
-// *slot*: a pointer paired with the type of the value it becomes, here a memref
-// and its vector type. Promoting a slot replaces it with that value, used as
-// the slot's reaching definition: the value the slot holds at a given program
-// point, merged with a block argument where control flow joins.
+// This file implements Mem2Reg-related interfaces that let a memref be promoted
+// into a single vector SSA value: `PromotableMemOpInterface` models for the ops
+// that access such a memref and `PromotableAliaserInterface` models for the ops
+// that view it. Mem2Reg calls the memory it promotes a *slot*: a pointer paired
+// with the type of the value it becomes, here a memref and its vector type.
+// Promoting a slot replaces it with that value, used as the slot's reaching
+// definition.
//
// A slot is promoted when each of its uses is an access these models can
// rewrite: a `vector.transfer_read` or `vector.transfer_write` meeting the
@@ -58,11 +57,14 @@
// of a dynamic subview) promotes as well: its alias shape comes from the
// parent's extents and its mask is reshaped along with the value.
//
-// The promoted memref must be statically shaped, so that its slot has a
-// fixed-shape vector type. Accesses that do not meet the criteria above --
-// dynamic offsets, rank-reducing or non-unit-stride subviews, non-contiguous
-// reshapes, non-zero transfer indices -- are left untouched, so the memref is
-// not promoted.
+// A memref accessed through the subviews, reshapes or copies above must be
+// statically shaped, so that its slot has a fixed-shape vector type. A scalable
+// slot may instead be dynamically shaped, but it is 1-D and promotes through
+// whole-buffer transfers alone: it cannot be subviewed, reshaped or copied.
+//
+// Accesses that do not meet the criteria above -- dynamic offsets,
+// rank-reducing or non-unit-stride subviews, non-contiguous reshapes, non-zero
+// transfer indices -- are left untouched, so the memref is not promoted.
//
//===----------------------------------------------------------------------===//
@@ -223,7 +225,7 @@ static bool isAliasableDynamicShapeSubView(memref::SubViewOp subView) {
/// The vector type covering `subView`'s whole parent buffer, which is the type
/// the alias slot of a dynamically-shaped subview takes.
-static VectorType getWholeParentBufferVectorType(memref::SubViewOp subView) {
+static VectorType getWholeParentVectorType(memref::SubViewOp subView) {
auto parentType = cast<MemRefType>(subView.getSource().getType());
assert(parentType.hasStaticShape() && "expected a statically-shaped parent");
return VectorType::get(parentType.getShape(), parentType.getElementType());
@@ -398,11 +400,13 @@ getReassociatedShape(memref::ExpandShapeOp op, ArrayRef<int64_t> parentShape) {
/// Builds the mask of `subView`'s valid region, expressed in `vecType`'s shape:
/// a `vector.create_mask` of the subview's sizes in the parent's shape,
/// reshaped with `vector.shape_cast` when the slot is a reassociated view of
-/// the parent (both describe the same elements in row-major order).
+/// the parent (both describe the same elements in row-major order). One cast
+/// suffices however many reshapes the slot is behind, since each preserves that
+/// order and so they compose.
static Value buildDynamicViewMask(OpBuilder &builder, Location loc,
memref::SubViewOp subView,
VectorType vecType) {
- VectorType parentVecType = getWholeParentBufferVectorType(subView);
+ VectorType parentVecType = getWholeParentVectorType(subView);
SmallVector<Value> bounds =
getValueOrCreateConstantIndexOp(builder, loc, subView.getMixedSizes());
Value mask = vector::CreateMaskOp::create(
@@ -427,6 +431,7 @@ static Value buildDynamicViewMask(OpBuilder &builder, Location loc,
/// aliaser's masked select.
static Value readMemRefAsVector(OpBuilder &builder, Location loc, Value mem,
VectorType vecType) {
+ assert(!vecType.isScalable() && "expected a fixed-size vector");
auto memType = cast<MemRefType>(mem.getType());
int64_t rank = vecType.getRank();
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
@@ -451,6 +456,7 @@ static void writeVectorToMemRef(OpBuilder &builder, Location loc, Value vec,
Value mem, Value mask = {}) {
auto memType = cast<MemRefType>(mem.getType());
auto vecType = cast<VectorType>(vec.getType());
+ assert(!vecType.isScalable() && "expected a fixed-size vector");
int64_t rank = vecType.getRank();
Value zero = arith::ConstantIndexOp::create(builder, loc, 0);
SmallVector<Value> indices(rank, zero);
@@ -766,7 +772,7 @@ struct SubViewOpAliasModel
// of the parent value untouched.
if (isAliasableDynamicShapeSubView(subView)) {
Location loc = op->getLoc();
- VectorType parentVecType = getWholeParentBufferVectorType(subView);
+ VectorType parentVecType = getWholeParentVectorType(subView);
Value mask = buildDynamicViewMask(builder, loc, subView, parentVecType);
return arith::SelectOp::create(builder, loc, mask, aliasValue,
reachingDef);
diff --git a/mlir/test/Dialect/Vector/mem2reg-copy.mlir b/mlir/test/Dialect/Vector/mem2reg-copy.mlir
index e68d5f8ed38ea..f884fc924d351 100644
--- a/mlir/test/Dialect/Vector/mem2reg-copy.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg-copy.mlir
@@ -169,3 +169,24 @@ func.func @copy_global_to_dyn_subview_memspace(%v: vector<128x64xf16>, %n: index
%r = vector.transfer_read %slm[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16, 3>, vector<128x64xf16>
return %r : vector<128x64xf16>
}
+
+// -----
+
+// NEGATIVE: the copy models keep to fixed-size slots, so a scalable slot is not
+// promoted. The dynamic buffer size is not what blocks it: being a vscale
+// multiple is exactly what gives the slot a scalable vector type, and such a
+// buffer promotes through transfers (see @scalable_whole_buffer_in_loop).
+// CHECK-LABEL: func.func @negative_scalable_copy(
+// CHECK: memref.alloca
+// CHECK: memref.copy
+// CHECK: vector.transfer_read
+func.func @negative_scalable_copy(%src: memref<8xf32>, %pad: f32) -> vector<[4]xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %vs = vector.vscale
+ %sz = arith.muli %vs, %c4 : index
+ %a = memref.alloca(%sz) : memref<?xf32>
+ memref.copy %src, %a : memref<8xf32> to memref<?xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<?xf32>, vector<[4]xf32>
+ return %r : vector<[4]xf32>
+}
diff --git a/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
index 38ff7b20e8b2b..7ad8f265ddf6f 100644
--- a/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
@@ -257,3 +257,156 @@ func.func @expand_shape_dyn_subview_copy_in(%v: vector<128x64xf16>, %n: index,
%r = vector.transfer_read %buf[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16>, vector<128x64xf16>
return %r : vector<128x64xf16>
}
+
+// -----
+
+// A dynamic view nested on a static subview: the two aliases compose, and the
+// mask is built in the extent the subview's sizes are relative to -- the static
+// subview's shape (4x16), not the buffer's (8x16).
+
+// CHECK-LABEL: func.func @dyn_subview_of_static_subview(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.subview
+// CHECK: %[[EX:.*]] = vector.extract_strided_slice %[[V]] offsets = [2, 0], sizes = [4, 16]
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<4x16xi1>
+// CHECK: %[[PS:.*]] = vector.broadcast
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[EX]], %[[PS]]
+// CHECK: return %[[SEL]]
+func.func @dyn_subview_of_static_subview(%v: vector<8x16xf32>, %n: index,
+ %pad: f32) -> vector<4x16xf32> {
+ %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>
+ %sv1 = memref.subview %a[2, 0] [4, 16] [1, 1] : memref<8x16xf32> to memref<4x16xf32, strided<[16, 1], offset: 32>>
+ %sv2 = memref.subview %sv1[0, 0] [%n, 16] [1, 1] : memref<4x16xf32, strided<[16, 1], offset: 32>> to memref<?x16xf32, strided<[16, 1], offset: 32>>
+ %r = vector.transfer_read %sv2[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x16xf32, strided<[16, 1], offset: 32>>, vector<4x16xf32>
+ return %r : vector<4x16xf32>
+}
+
+// -----
+
+// Two reshapes compose over a dynamic view. The value is projected one alias at
+// a time, so it picks up a shape_cast per reshape (8x16 -> 128 -> 64x2). The
+// mask is not projected: it is built in the parent's shape where it is needed
+// and cast once to the final shape, which is equivalent because shape_casts
+// preserve row-major order and therefore compose.
+
+// CHECK-LABEL: func.func @collapse_then_expand_dyn_subview(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.collapse_shape
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[F:.*]] = vector.shape_cast %[[V]] : vector<8x16xf32> to vector<128xf32>
+// CHECK: %[[E:.*]] = vector.shape_cast %[[F]] : vector<128xf32> to vector<64x2xf32>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<8x16xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[CM]] : vector<8x16xi1> to vector<64x2xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[E]], %{{.*}}
+// CHECK: return %[[SEL]]
+func.func @collapse_then_expand_dyn_subview(%v: vector<8x16xf32>, %n: index,
+ %d: index, %pad: f32) -> vector<64x2xf32> {
+ %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] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
+ %f = memref.collapse_shape %sv [[0, 1]] : memref<?x16xf32, strided<[16, 1]>> into memref<?xf32, strided<[1]>>
+ %e = memref.expand_shape %f [[0, 1]] output_shape [%d, 2] : memref<?xf32, strided<[1]>> into memref<?x2xf32>
+ %r = vector.transfer_read %e[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x2xf32>, vector<64x2xf32>
+ return %r : vector<64x2xf32>
+}
+
+// -----
+
+// The promoted value flows across control flow: the masked store on one path is
+// merged with the untouched value on the other by a block argument. The mask is
+// built where the view is, so it is only computed on the path that stores.
+
+// CHECK-LABEL: func.func @dyn_subview_across_cfg(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[W:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK: cf.cond_br %{{.*}}, ^[[BB1:.*]], ^[[BB2:.*]]
+// CHECK: ^[[BB1]]:
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<8x16xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[W]], %[[V]]
+// CHECK: cf.br ^[[BB3:.*]](%[[SEL]] : vector<8x16xf32>)
+// CHECK: ^[[BB2]]:
+// CHECK: cf.br ^[[BB3]](%[[V]] : vector<8x16xf32>)
+// CHECK: ^[[BB3]](%[[PHI:.*]]: vector<8x16xf32>):
+// CHECK: return %[[PHI]]
+func.func @dyn_subview_across_cfg(%v: vector<8x16xf32>, %w: vector<8x16xf32>,
+ %n: index, %pad: f32, %c: i1) -> 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>
+ cf.cond_br %c, ^bb1, ^bb2
+^bb1:
+ %sv = memref.subview %a[0, 0] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
+ vector.transfer_write %w, %sv[%c0, %c0] {in_bounds = [false, true]} : vector<8x16xf32>, memref<?x16xf32, strided<[16, 1]>>
+ cf.br ^bb3
+^bb2:
+ cf.br ^bb3
+^bb3:
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Negative: a dynamic size with a static but non-zero offset. The valid region
+// no longer starts at the parent's origin, which `vector.create_mask` cannot
+// express, so the buffer is not promoted.
+// CHECK-LABEL: func.func @neg_static_offset_dyn_size(
+// CHECK: memref.alloca
+// CHECK: memref.subview
+// CHECK: vector.transfer_read
+func.func @neg_static_offset_dyn_size(%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[2, 0] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1], offset: 32>>
+ %r = vector.transfer_read %sv[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x16xf32, strided<[16, 1], offset: 32>>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// Negative: the collapsed group's dynamic extent is not its outermost, so the
+// group's stride would have to be that runtime extent. This is the collapse
+// counterpart of @neg_nonunit_dim_before_dynamic, rejected one step earlier:
+// such a collapse is not guaranteed contiguous, so its result stride is dynamic.
+// CHECK-LABEL: func.func @neg_collapse_dyn_inner_dim(
+// CHECK: memref.alloca
+// CHECK: memref.collapse_shape
+// CHECK: vector.transfer_read
+func.func @neg_collapse_dyn_inner_dim(%v: vector<4x4xf32>, %n: index,
+ %pad: f32) -> vector<16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<4x4xf32>
+ vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<4x4xf32>, memref<4x4xf32>
+ %sv = memref.subview %a[0, 0] [4, %n] [1, 1] : memref<4x4xf32> to memref<4x?xf32, strided<[4, 1]>>
+ %f = memref.collapse_shape %sv [[0, 1]] : memref<4x?xf32, strided<[4, 1]>> into memref<?xf32, strided<[?]>>
+ %r = vector.transfer_read %f[%c0], %pad {in_bounds = [false]} : memref<?xf32, strided<[?]>>, vector<16xf32>
+ return %r : vector<16xf32>
+}
+
+// -----
+
+// Negative, and a crash regression: a reassociation group may hold a zero static
+// extent, which the parent extent cannot be divided by. The buffer is not
+// promoted; without the guard the division faults.
+// CHECK-LABEL: func.func @neg_zero_extent_in_group(
+// CHECK: memref.alloca
+// CHECK: memref.subview
+// CHECK: memref.expand_shape
+// CHECK: vector.transfer_read
+func.func @neg_zero_extent_in_group(%v: vector<8x16xf32>, %n: index, %d: 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] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
+ %e = memref.expand_shape %sv [[0, 1], [2]] output_shape [%d, 0, 16] : memref<?x16xf32, strided<[16, 1]>> into memref<?x0x16xf32, strided<[0, 16, 1]>>
+ %r = vector.transfer_read %e[%c0, %c0, %c0], %pad : memref<?x0x16xf32, strided<[0, 16, 1]>>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
>From 0e666dfa6107fd9be9a3de76078e555b25742947 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 17 Sep 2026 06:12:16 +0000
Subject: [PATCH 4/4] [mlir][mem2reg][NFC] Reorganise the vector promotion
tests and comments
Fold mem2reg-copy.mlir into the two files organised by the axis the suite
already uses. `memref.copy` is a different axis from static-versus-dynamic
view, so a file named after it crossed the other two in both directions: it
held dynamic-subview tests while the dynamic-subview file held copy tests.
Whole-buffer copies move to mem2reg.mlir, joining the other whole-buffer and
scalable cases; copies through a dynamic view move to
mem2reg-dynamic-subview.mlir, joining the two expanded-view copies already
there. Both files now run the same plain -mem2reg pipeline: only two CHECKs
depended on -canonicalize, and a mem2reg test should not assert on folder
output anyway.
Order mem2reg-dynamic-subview.mlir as positives grouped by feature -- plain
dynamic subview, reshapes of one, copies through one -- followed by the
negatives grouped by the rule each violates.
Reword the test comments to say what each test does rather than how the
models do it, and drop cross-references between tests. Merge the two
anonymous namespaces in the models file, move the reshape helpers into the
utilities section, and give the transfer models the same one-line
introduction as the copy model.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Transforms/MemorySlotOpInterfaceImpl.cpp | 164 ++++---
mlir/test/Dialect/Vector/mem2reg-copy.mlir | 192 --------
.../Vector/mem2reg-dynamic-subview.mlir | 422 ++++++++++--------
mlir/test/Dialect/Vector/mem2reg.mlir | 94 ++++
4 files changed, 417 insertions(+), 455 deletions(-)
delete mode 100644 mlir/test/Dialect/Vector/mem2reg-copy.mlir
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index d3e835e3db050..751f0abc2364e 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -260,6 +260,79 @@ static memref::SubViewOp getUnderlyingDynamicSubView(Value slotPtr) {
return {};
}
+/// Builds the mask of `subView`'s valid region, expressed in `vecType`'s shape:
+/// a `vector.create_mask` of the subview's sizes in the parent's shape,
+/// reshaped with `vector.shape_cast` when the slot is a reassociated view of
+/// the parent (both describe the same elements in row-major order). One cast
+/// suffices however many reshapes the slot is behind, since each preserves that
+/// order and so they compose.
+static Value buildDynamicViewMask(OpBuilder &builder, Location loc,
+ memref::SubViewOp subView,
+ VectorType vecType) {
+ VectorType parentVecType = getWholeParentVectorType(subView);
+ SmallVector<Value> bounds =
+ getValueOrCreateConstantIndexOp(builder, loc, subView.getMixedSizes());
+ Value mask = vector::CreateMaskOp::create(
+ builder, loc,
+ VectorType::get(parentVecType.getShape(), builder.getI1Type()), bounds);
+ if (parentVecType.getShape() != vecType.getShape())
+ mask = vector::ShapeCastOp::create(
+ builder, loc, VectorType::get(vecType.getShape(), builder.getI1Type()),
+ mask);
+ return mask;
+}
+
+/// Produces the value that a `memref.copy` stores into a slot: reads `mem` at
+/// the origin into `vecType` with an identity map, marking a dimension
+/// in-bounds only when the memref extent is statically at least the vector
+/// extent.
+///
+/// The padding is a don't-care: `memref.copy` requires both sides to have the
+/// same shape, so the source covers the whole vector at runtime and a dynamic
+/// extent only forces a conservative out-of-bounds marking. For a
+/// dynamic-subview slot, lanes past its extent are in any case discarded by the
+/// aliaser's masked select.
+static Value readMemRefAsVector(OpBuilder &builder, Location loc, Value mem,
+ VectorType vecType) {
+ assert(!vecType.isScalable() && "expected a fixed-size vector");
+ 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));
+}
+
+/// Stores a slot's value to the target of a `memref.copy`: writes `vec` into
+/// `mem` at the origin, with the identity map and in-bounds rule of
+/// `readMemRefAsVector`. `mask`, when set, restricts the write to the lanes it
+/// selects.
+static void writeVectorToMemRef(OpBuilder &builder, Location loc, Value vec,
+ Value mem, Value mask = {}) {
+ auto memType = cast<MemRefType>(mem.getType());
+ auto vecType = cast<VectorType>(vec.getType());
+ assert(!vecType.isScalable() && "expected a fixed-size vector");
+ 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,
+ builder.getBoolArrayAttr(inBounds));
+}
+
/// Returns whether the view lays its elements out in the same (row-major) order
/// as `parentShape`, the parent slot's vector shape, which is what makes a
/// `vector.shape_cast` model it. These are the conditions under which
@@ -397,85 +470,13 @@ getReassociatedShape(memref::ExpandShapeOp op, ArrayRef<int64_t> parentShape) {
return shape;
}
-/// Builds the mask of `subView`'s valid region, expressed in `vecType`'s shape:
-/// a `vector.create_mask` of the subview's sizes in the parent's shape,
-/// reshaped with `vector.shape_cast` when the slot is a reassociated view of
-/// the parent (both describe the same elements in row-major order). One cast
-/// suffices however many reshapes the slot is behind, since each preserves that
-/// order and so they compose.
-static Value buildDynamicViewMask(OpBuilder &builder, Location loc,
- memref::SubViewOp subView,
- VectorType vecType) {
- VectorType parentVecType = getWholeParentVectorType(subView);
- SmallVector<Value> bounds =
- getValueOrCreateConstantIndexOp(builder, loc, subView.getMixedSizes());
- Value mask = vector::CreateMaskOp::create(
- builder, loc,
- VectorType::get(parentVecType.getShape(), builder.getI1Type()), bounds);
- if (parentVecType.getShape() != vecType.getShape())
- mask = vector::ShapeCastOp::create(
- builder, loc, VectorType::get(vecType.getShape(), builder.getI1Type()),
- mask);
- return mask;
-}
-
-/// Produces the value that a `memref.copy` stores into a slot: reads `mem` at
-/// the origin into `vecType` with an identity map, marking a dimension
-/// in-bounds only when the memref extent is statically at least the vector
-/// extent.
-///
-/// The padding is a don't-care: `memref.copy` requires both sides to have the
-/// same shape, so the source covers the whole vector at runtime and a dynamic
-/// extent only forces a conservative out-of-bounds marking. For a
-/// dynamic-subview slot, lanes past its extent are in any case discarded by the
-/// aliaser's masked select.
-static Value readMemRefAsVector(OpBuilder &builder, Location loc, Value mem,
- VectorType vecType) {
- assert(!vecType.isScalable() && "expected a fixed-size vector");
- 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));
-}
-
-/// Stores a slot's value to the target of a `memref.copy`: writes `vec` into
-/// `mem` at the origin, with the identity map and in-bounds rule of
-/// `readMemRefAsVector`. `mask`, when set, restricts the write to the lanes it
-/// selects.
-static void writeVectorToMemRef(OpBuilder &builder, Location loc, Value vec,
- Value mem, Value mask = {}) {
- auto memType = cast<MemRefType>(mem.getType());
- auto vecType = cast<VectorType>(vec.getType());
- assert(!vecType.isScalable() && "expected a fixed-size vector");
- 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,
- builder.getBoolArrayAttr(inBounds));
-}
-
//===----------------------------------------------------------------------===//
// Interface models
//===----------------------------------------------------------------------===//
namespace {
+/// Mem2Reg model for a `vector.transfer_read` of the slot.
struct TransferReadOpMemOpModel
: public PromotableMemOpInterface::ExternalModel<TransferReadOpMemOpModel,
vector::TransferReadOp> {
@@ -530,6 +531,8 @@ struct TransferReadOpMemOpModel
}
};
+/// Mem2Reg model for a `vector.transfer_write` to the slot, which qualifies
+/// under the same `isPromotableTransfer` criteria.
struct TransferWriteOpMemOpModel
: public PromotableMemOpInterface::ExternalModel<TransferWriteOpMemOpModel,
vector::TransferWriteOp> {
@@ -653,14 +656,10 @@ struct CopyOpMemOpModel
}
};
-} // namespace
-
//===----------------------------------------------------------------------===//
// memref view aliasers
//===----------------------------------------------------------------------===//
-namespace {
-
/// Companion `PromotableOpInterface` model for the view ops aliased below
/// (`memref.subview`, `memref.expand_shape`, `memref.collapse_shape`): once the
/// slot is promoted, the view has no remaining memory uses and is erased.
@@ -722,12 +721,6 @@ struct SubViewOpAliasModel
if (!parentVecType)
return;
- // A dynamic slice aliases the whole parent value.
- if (isAliasableDynamicShapeSubView(subView)) {
- newSlots.push_back(MemorySlot{subView.getResult(), parentVecType});
- return;
- }
-
// A static slice aliases the sub-vector matching the subview's shape.
if (isAliasableStaticShapeSubView(subView)) {
auto resType = cast<MemRefType>(subView.getResult().getType());
@@ -737,6 +730,11 @@ struct SubViewOpAliasModel
VectorType::get(resType.getShape(), resType.getElementType());
newSlots.push_back(MemorySlot{subView.getResult(), aliasVecType});
}
+
+ // A dynamic slice aliases the whole parent value.
+ if (isAliasableDynamicShapeSubView(subView)) {
+ newSlots.push_back(MemorySlot{subView.getResult(), parentVecType});
+ }
}
Value projectSlotValueToAliasValue(Operation *op,
diff --git a/mlir/test/Dialect/Vector/mem2reg-copy.mlir b/mlir/test/Dialect/Vector/mem2reg-copy.mlir
deleted file mode 100644
index f884fc924d351..0000000000000
--- a/mlir/test/Dialect/Vector/mem2reg-copy.mlir
+++ /dev/null
@@ -1,192 +0,0 @@
-// 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>
-}
-
-// -----
-
-// Copy from a global into a dynamic subview of a static local buffer (an "SLM"
-// staging buffer): the global is not promotable, so the copy becomes a read of
-// the global's slice, composed onto the buffer's value over the subview extent.
-// The read is not in bounds on the dynamic dimension and its padding is a
-// don't-care: those lanes keep the buffer's previous value through the select.
-// The second function repeats it in a non-default memory space.
-memref.global "private" constant @tile : memref<128x64xf16> = dense<1.000000e+00>
-
-// CHECK-LABEL: func.func @copy_global_to_dyn_subview(
-// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
-// CHECK-NOT: memref.alloca
-// CHECK-NOT: memref.copy
-// CHECK-DAG: %[[PAD:.*]] = arith.constant 0.000000e+00 : f16
-// CHECK: %[[G:.*]] = memref.get_global @tile : memref<128x64xf16>
-// CHECK: %[[GSV:.*]] = memref.subview %[[G]][0, 0] [%[[N]], 64] [1, 1]
-// CHECK: %[[RD:.*]] = vector.transfer_read %[[GSV]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [false, true]} : memref<?x64xf16, strided<[64, 1]>>, vector<128x64xf16>
-// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<128x64xi1>
-// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[RD]], %[[V]]
-// CHECK: return %[[SEL]]
-func.func @copy_global_to_dyn_subview(%v: vector<128x64xf16>, %n: index, %pad: f16) -> vector<128x64xf16> {
- %c0 = arith.constant 0 : index
- %g = memref.get_global @tile : memref<128x64xf16>
- %gsv = memref.subview %g[0, 0] [%n, 64] [1, 1] : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
- %slm = memref.alloca() : memref<128x64xf16>
- vector.transfer_write %v, %slm[%c0, %c0] {in_bounds = [true, true]} : vector<128x64xf16>, memref<128x64xf16>
- %sv = memref.subview %slm[0, 0] [%n, 64] [1, 1] : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
- memref.copy %gsv, %sv : memref<?x64xf16, strided<[64, 1]>> to memref<?x64xf16, strided<[64, 1]>>
- %r = vector.transfer_read %slm[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16>, vector<128x64xf16>
- return %r : vector<128x64xf16>
-}
-
-// CHECK-LABEL: func.func @copy_global_to_dyn_subview_memspace(
-// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
-// CHECK-NOT: memref.alloca
-// CHECK-NOT: memref.copy
-// CHECK: %[[RD:.*]] = vector.transfer_read %{{.*}} {in_bounds = [false, true]} : memref<?x64xf16, strided<[64, 1]>>, vector<128x64xf16>
-// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<128x64xi1>
-// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[RD]], %[[V]]
-// CHECK: return %[[SEL]]
-func.func @copy_global_to_dyn_subview_memspace(%v: vector<128x64xf16>, %n: index, %pad: f16) -> vector<128x64xf16> {
- %c0 = arith.constant 0 : index
- %g = memref.get_global @tile : memref<128x64xf16>
- %gsv = memref.subview %g[0, 0] [%n, 64] [1, 1] : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
- %slm = memref.alloca() : memref<128x64xf16, 3>
- vector.transfer_write %v, %slm[%c0, %c0] {in_bounds = [true, true]} : vector<128x64xf16>, memref<128x64xf16, 3>
- %sv = memref.subview %slm[0, 0] [%n, 64] [1, 1] : memref<128x64xf16, 3> to memref<?x64xf16, strided<[64, 1]>, 3>
- memref.copy %gsv, %sv : memref<?x64xf16, strided<[64, 1]>> to memref<?x64xf16, strided<[64, 1]>, 3>
- %r = vector.transfer_read %slm[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16, 3>, vector<128x64xf16>
- return %r : vector<128x64xf16>
-}
-
-// -----
-
-// NEGATIVE: the copy models keep to fixed-size slots, so a scalable slot is not
-// promoted. The dynamic buffer size is not what blocks it: being a vscale
-// multiple is exactly what gives the slot a scalable vector type, and such a
-// buffer promotes through transfers (see @scalable_whole_buffer_in_loop).
-// CHECK-LABEL: func.func @negative_scalable_copy(
-// CHECK: memref.alloca
-// CHECK: memref.copy
-// CHECK: vector.transfer_read
-func.func @negative_scalable_copy(%src: memref<8xf32>, %pad: f32) -> vector<[4]xf32> {
- %c0 = arith.constant 0 : index
- %c4 = arith.constant 4 : index
- %vs = vector.vscale
- %sz = arith.muli %vs, %c4 : index
- %a = memref.alloca(%sz) : memref<?xf32>
- memref.copy %src, %a : memref<8xf32> to memref<?xf32>
- %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<?xf32>, vector<[4]xf32>
- return %r : vector<[4]xf32>
-}
diff --git a/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
index 7ad8f265ddf6f..ca6abe8a76fe3 100644
--- a/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir
@@ -1,8 +1,8 @@
// 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.
+// Reading a static buffer through a DYNAMIC subview: the read returns the value
+// that was written, and takes the transfer's padding wherever it reads past the
+// subview's extent.
// CHECK-LABEL: func.func @read_dyn_subview(
// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[PAD:.*]]: f32
@@ -29,9 +29,9 @@ func.func @read_dyn_subview(%v: vector<8x16xf32>, %n: index, %pad: f32)
// -----
-// 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.
+// Write through a dynamic subview, then read the whole buffer. The write updates
+// only the elements within the dynamic extent, so the read returns the written
+// value for those and the buffer's earlier value for the rest.
// CHECK-LABEL: func.func @write_then_read_dyn_subview(
// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[W:.*]]: vector<8x16xf32>, %[[N:.*]]: index
@@ -56,28 +56,6 @@ func.func @write_then_read_dyn_subview(%v: vector<8x16xf32>, %w: 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(
@@ -98,10 +76,68 @@ func.func @masked_read_dyn_subview(%v: vector<8x16xf32>, %n: index, %m: vector<8
// -----
-// A dynamic subview reshaped by `memref.collapse_shape`: the reshape keeps the
-// whole parent value, so the read is the parent value shape_cast to the collapsed
-// shape, masked by the subview extent -- the mask is built in the parent's shape
-// and shape_cast along with the value (both describe the same elements).
+// A dynamic subview of a static subview: both views apply, and the mask uses
+// the shape the subview's sizes are relative to (4x16, not the buffer's 8x16).
+
+// CHECK-LABEL: func.func @dyn_subview_of_static_subview(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.subview
+// CHECK: %[[EX:.*]] = vector.extract_strided_slice %[[V]] offsets = [2, 0], sizes = [4, 16]
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<4x16xi1>
+// CHECK: %[[PS:.*]] = vector.broadcast
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[EX]], %[[PS]]
+// CHECK: return %[[SEL]]
+func.func @dyn_subview_of_static_subview(%v: vector<8x16xf32>, %n: index,
+ %pad: f32) -> vector<4x16xf32> {
+ %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>
+ %sv1 = memref.subview %a[2, 0] [4, 16] [1, 1] : memref<8x16xf32> to memref<4x16xf32, strided<[16, 1], offset: 32>>
+ %sv2 = memref.subview %sv1[0, 0] [%n, 16] [1, 1] : memref<4x16xf32, strided<[16, 1], offset: 32>> to memref<?x16xf32, strided<[16, 1], offset: 32>>
+ %r = vector.transfer_read %sv2[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x16xf32, strided<[16, 1], offset: 32>>, vector<4x16xf32>
+ return %r : vector<4x16xf32>
+}
+
+// -----
+
+// The promoted value flows across control flow: a block argument merges the
+// masked store on one path with the untouched value on the other.
+
+// CHECK-LABEL: func.func @dyn_subview_across_cfg(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[W:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK: cf.cond_br %{{.*}}, ^[[BB1:.*]], ^[[BB2:.*]]
+// CHECK: ^[[BB1]]:
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<8x16xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[W]], %[[V]]
+// CHECK: cf.br ^[[BB3:.*]](%[[SEL]] : vector<8x16xf32>)
+// CHECK: ^[[BB2]]:
+// CHECK: cf.br ^[[BB3]](%[[V]] : vector<8x16xf32>)
+// CHECK: ^[[BB3]](%[[PHI:.*]]: vector<8x16xf32>):
+// CHECK: return %[[PHI]]
+func.func @dyn_subview_across_cfg(%v: vector<8x16xf32>, %w: vector<8x16xf32>,
+ %n: index, %pad: f32, %c: i1) -> 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>
+ cf.cond_br %c, ^bb1, ^bb2
+^bb1:
+ %sv = memref.subview %a[0, 0] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
+ vector.transfer_write %w, %sv[%c0, %c0] {in_bounds = [false, true]} : vector<8x16xf32>, memref<?x16xf32, strided<[16, 1]>>
+ cf.br ^bb3
+^bb2:
+ cf.br ^bb3
+^bb3:
+ %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<8x16xf32>
+ return %r : vector<8x16xf32>
+}
+
+// -----
+
+// A dynamic subview reshaped by `memref.collapse_shape`: the read is the
+// buffer's value shape_cast to the collapsed shape, and the mask is shape_cast
+// the same way so that it still marks the same elements.
// CHECK-LABEL: func.func @collapse_shape_dyn_subview(
// CHECK-SAME: %[[V:.*]]: vector<8x16xf32>, %[[N:.*]]: index, %[[PAD:.*]]: f32
// CHECK-NOT: memref.alloca
@@ -125,8 +161,8 @@ func.func @collapse_shape_dyn_subview(%v: vector<8x16xf32>, %n: index, %pad: f32
// -----
// `memref.expand_shape` of a dynamic view: the dynamic result dimension is
-// recovered from the parent extent it splits (16 = ? x 2 gives 8), so the alias
-// is `vector<8x2xf32>` and the mask is reshaped alongside the value.
+// known from the buffer extent it splits (16 = ? x 2 gives 8), so the read is a
+// `vector<8x2xf32>`, and the mask is reshaped the same way.
// CHECK-LABEL: func.func @expand_shape_dyn_subview(
// CHECK-SAME: %[[V:.*]]: vector<16xf32>, %[[M:.*]]: index, %[[PAD:.*]]: f32
// CHECK-NOT: memref.alloca
@@ -149,54 +185,117 @@ func.func @expand_shape_dyn_subview(%v: vector<16xf32>, %m: index, %pad: f32) ->
// -----
-// Negative: two dynamic dimensions in one reassociation group -- the alias shape
-// is not determined by the parent extent, so the buffer is not promoted.
-// CHECK-LABEL: func.func @neg_two_dynamic_dims_in_group(
-// CHECK: memref.alloca
-// CHECK: memref.expand_shape
-// CHECK: vector.transfer_read
-func.func @neg_two_dynamic_dims_in_group(%v: vector<16xf32>, %m: index, %d0: index,
- %d1: index, %pad: f32) -> vector<8x2xf32> {
+// Two reshapes in a row: the value goes through one shape_cast per reshape
+// (8x16 -> 128 -> 64x2), while the mask is built once in the buffer's shape and
+// cast straight to the final one.
+
+// CHECK-LABEL: func.func @collapse_then_expand_dyn_subview(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.collapse_shape
+// CHECK-NOT: memref.expand_shape
+// CHECK: %[[F:.*]] = vector.shape_cast %[[V]] : vector<8x16xf32> to vector<128xf32>
+// CHECK: %[[E:.*]] = vector.shape_cast %[[F]] : vector<128xf32> to vector<64x2xf32>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<8x16xi1>
+// CHECK: %[[MASK:.*]] = vector.shape_cast %[[CM]] : vector<8x16xi1> to vector<64x2xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[E]], %{{.*}}
+// CHECK: return %[[SEL]]
+func.func @collapse_then_expand_dyn_subview(%v: vector<8x16xf32>, %n: index,
+ %d: index, %pad: f32) -> vector<64x2xf32> {
%c0 = arith.constant 0 : index
- %a = memref.alloca() : memref<16xf32>
- vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<16xf32>, memref<16xf32>
- %sv = memref.subview %a[0] [%m] [1] : memref<16xf32> to memref<?xf32, strided<[1]>>
- %exp = memref.expand_shape %sv [[0, 1]] output_shape [%d0, %d1] : memref<?xf32, strided<[1]>> into memref<?x?xf32>
- %r = vector.transfer_read %exp[%c0, %c0], %pad : memref<?x?xf32>, vector<8x2xf32>
- return %r : vector<8x2xf32>
+ %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] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
+ %f = memref.collapse_shape %sv [[0, 1]] : memref<?x16xf32, strided<[16, 1]>> into memref<?xf32, strided<[1]>>
+ %e = memref.expand_shape %f [[0, 1]] output_shape [%d, 2] : memref<?xf32, strided<[1]>> into memref<?x2xf32>
+ %r = vector.transfer_read %e[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x2xf32>, vector<64x2xf32>
+ return %r : vector<64x2xf32>
}
// -----
-// Negative: a non-unit dimension precedes the dynamic one in its group, so the
-// view and the alias disagree on that dimension's stride -- the view strides by
-// the runtime %m, while an alias shaped from the parent strides by 8 / 2 = 4.
-// Promoting would move data rather than merely mask it: for a buffer holding
-// 1..8 with %n = 4 (so %m = 2), the view reads
-// ((1, 2, pad, pad), (3, 4, pad, pad)) whereas a masked shape_cast of the parent
-// value yields ((1, 2, 3, 4), (pad, pad, pad, pad)).
-// CHECK-LABEL: func.func @neg_nonunit_dim_before_dynamic(
-// CHECK: memref.alloca
-// CHECK: memref.subview
-// CHECK: memref.expand_shape
-// CHECK: vector.transfer_read
-func.func @neg_nonunit_dim_before_dynamic(%v: vector<8xf32>, %n: index, %m: index,
- %pad: f32) -> vector<2x4xf32> {
+// A `memref.copy` into a dynamic subview, read back afterwards: the copy
+// becomes a read of the source, kept only where it falls inside the extent.
+
+// CHECK-LABEL: func.func @copy_into_dyn_subview(
+// CHECK-SAME: %[[SRC:.*]]: memref<?x?xf32>, %[[N:.*]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK: %[[Z:.*]] = vector.broadcast %{{.*}} : f32 to 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_into_dyn_subview(%src: memref<?x?xf32>, %n: index, %pad: f32) -> vector<8x16xf32> {
%c0 = arith.constant 0 : index
- %a = memref.alloca() : memref<8xf32>
- vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
- %sv = memref.subview %a[0] [%n] [1] : memref<8xf32> to memref<?xf32, strided<[1]>>
- %exp = memref.expand_shape %sv [[0, 1]] output_shape [2, %m] : memref<?xf32, strided<[1]>> into memref<2x?xf32>
- %r = vector.transfer_read %exp[%c0, %c0], %pad {in_bounds = [true, false]} : memref<2x?xf32>, vector<2x4xf32>
- return %r : vector<2x4xf32>
+ %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>
}
// -----
-// The reported bufferization pattern: a padded static buffer, a dynamic subview
-// of its real region, expanded to a higher rank (with a dynamic dim), copied out
-// to another buffer. The whole chain promotes: the parent value is shape_cast to
-// the expanded shape and written masked to the copy's valid region.
+// Copy from one dynamic subview into another: the source subview is read, and
+// only the elements inside the extent reach the buffer's value.
+
+// 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>
+}
+
+// -----
+
+// Copy from a global into a dynamic subview of a static staging buffer, in a
+// non-default memory space.
+memref.global "private" constant @tile : memref<128x64xf16> = dense<1.000000e+00>
+
+// CHECK-LABEL: func.func @copy_global_to_dyn_subview(
+// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
+// CHECK-NOT: memref.alloca
+// CHECK-NOT: memref.copy
+// CHECK: %[[G:.*]] = memref.get_global @tile : memref<128x64xf16>
+// CHECK: %[[GSV:.*]] = memref.subview %[[G]][0, 0] [%[[N]], 64] [1, 1]
+// CHECK: %[[PAD:.*]] = arith.constant 0.000000e+00 : f16
+// CHECK: %[[RD:.*]] = vector.transfer_read %[[GSV]][%{{.*}}, %{{.*}}], %[[PAD]] {in_bounds = [false, true]} : memref<?x64xf16, strided<[64, 1]>>, vector<128x64xf16>
+// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<128x64xi1>
+// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[RD]], %[[V]]
+// CHECK: return %[[SEL]]
+func.func @copy_global_to_dyn_subview(%v: vector<128x64xf16>, %n: index, %pad: f16) -> vector<128x64xf16> {
+ %c0 = arith.constant 0 : index
+ %g = memref.get_global @tile : memref<128x64xf16>
+ %gsv = memref.subview %g[0, 0] [%n, 64] [1, 1] : memref<128x64xf16> to memref<?x64xf16, strided<[64, 1]>>
+ %slm = memref.alloca() : memref<128x64xf16, 3>
+ vector.transfer_write %v, %slm[%c0, %c0] {in_bounds = [true, true]} : vector<128x64xf16>, memref<128x64xf16, 3>
+ %sv = memref.subview %slm[0, 0] [%n, 64] [1, 1] : memref<128x64xf16, 3> to memref<?x64xf16, strided<[64, 1]>, 3>
+ memref.copy %gsv, %sv : memref<?x64xf16, strided<[64, 1]>> to memref<?x64xf16, strided<[64, 1]>, 3>
+ %r = vector.transfer_read %slm[%c0, %c0], %pad {in_bounds = [true, true]} : memref<128x64xf16, 3>, vector<128x64xf16>
+ return %r : vector<128x64xf16>
+}
+
+// -----
+
+// A padded static buffer copied OUT to a 1x1x?x64 tile of a larger buffer,
+// through a dynamic subview of its real region expanded to that rank.
// CHECK-LABEL: func.func @expand_shape_dyn_subview_copy_out(
// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
// CHECK-NOT: memref.alloca
@@ -223,13 +322,8 @@ func.func @expand_shape_dyn_subview_copy_out(%v: vector<128x64xf16>, %n: index,
// -----
-// The same chain with the copy reversed -- loading a tile from the other buffer
-// into the staging buffer through the expanded dynamic view. The copy becomes a
-// read of the source's slice, reshaped back and composed onto the buffer's value
-// over the subview extent. The read carries a padding operand and is NOT in
-// bounds on the dynamic dimension, which is what makes reading past the source's
-// extent defined; those padded lanes are discarded by the select, so the lanes
-// outside the extent keep the buffer's previous value.
+// A 1x1x?x64 tile of a larger buffer copied IN to a padded static buffer,
+// through a dynamic subview of it expanded to that rank.
// CHECK-LABEL: func.func @expand_shape_dyn_subview_copy_in(
// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<128x64xf16>, %[[N:[a-z0-9_]+]]: index
// CHECK-NOT: memref.alloca
@@ -260,93 +354,23 @@ func.func @expand_shape_dyn_subview_copy_in(%v: vector<128x64xf16>, %n: index,
// -----
-// A dynamic view nested on a static subview: the two aliases compose, and the
-// mask is built in the extent the subview's sizes are relative to -- the static
-// subview's shape (4x16), not the buffer's (8x16).
-
-// CHECK-LABEL: func.func @dyn_subview_of_static_subview(
-// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
-// CHECK-NOT: memref.alloca
-// CHECK-NOT: memref.subview
-// CHECK: %[[EX:.*]] = vector.extract_strided_slice %[[V]] offsets = [2, 0], sizes = [4, 16]
-// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<4x16xi1>
-// CHECK: %[[PS:.*]] = vector.broadcast
-// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[EX]], %[[PS]]
-// CHECK: return %[[SEL]]
-func.func @dyn_subview_of_static_subview(%v: vector<8x16xf32>, %n: index,
- %pad: f32) -> vector<4x16xf32> {
- %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>
- %sv1 = memref.subview %a[2, 0] [4, 16] [1, 1] : memref<8x16xf32> to memref<4x16xf32, strided<[16, 1], offset: 32>>
- %sv2 = memref.subview %sv1[0, 0] [%n, 16] [1, 1] : memref<4x16xf32, strided<[16, 1], offset: 32>> to memref<?x16xf32, strided<[16, 1], offset: 32>>
- %r = vector.transfer_read %sv2[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x16xf32, strided<[16, 1], offset: 32>>, vector<4x16xf32>
- return %r : vector<4x16xf32>
-}
-
-// -----
-
-// Two reshapes compose over a dynamic view. The value is projected one alias at
-// a time, so it picks up a shape_cast per reshape (8x16 -> 128 -> 64x2). The
-// mask is not projected: it is built in the parent's shape where it is needed
-// and cast once to the final shape, which is equivalent because shape_casts
-// preserve row-major order and therefore compose.
-
-// CHECK-LABEL: func.func @collapse_then_expand_dyn_subview(
-// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
-// CHECK-NOT: memref.alloca
-// CHECK-NOT: memref.collapse_shape
-// CHECK-NOT: memref.expand_shape
-// CHECK: %[[F:.*]] = vector.shape_cast %[[V]] : vector<8x16xf32> to vector<128xf32>
-// CHECK: %[[E:.*]] = vector.shape_cast %[[F]] : vector<128xf32> to vector<64x2xf32>
-// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<8x16xi1>
-// CHECK: %[[MASK:.*]] = vector.shape_cast %[[CM]] : vector<8x16xi1> to vector<64x2xi1>
-// CHECK: %[[SEL:.*]] = arith.select %[[MASK]], %[[E]], %{{.*}}
-// CHECK: return %[[SEL]]
-func.func @collapse_then_expand_dyn_subview(%v: vector<8x16xf32>, %n: index,
- %d: index, %pad: f32) -> vector<64x2xf32> {
- %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] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
- %f = memref.collapse_shape %sv [[0, 1]] : memref<?x16xf32, strided<[16, 1]>> into memref<?xf32, strided<[1]>>
- %e = memref.expand_shape %f [[0, 1]] output_shape [%d, 2] : memref<?xf32, strided<[1]>> into memref<?x2xf32>
- %r = vector.transfer_read %e[%c0, %c0], %pad {in_bounds = [false, true]} : memref<?x2xf32>, vector<64x2xf32>
- return %r : vector<64x2xf32>
-}
-
-// -----
-
-// The promoted value flows across control flow: the masked store on one path is
-// merged with the untouched value on the other by a block argument. The mask is
-// built where the view is, so it is only computed on the path that stores.
+// Negative: a dynamic OFFSET (not just size) subview is not promotable; the
+// buffer is left alone.
-// CHECK-LABEL: func.func @dyn_subview_across_cfg(
-// CHECK-SAME: %[[V:[a-z0-9_]+]]: vector<8x16xf32>, %[[W:[a-z0-9_]+]]: vector<8x16xf32>, %[[N:[a-z0-9_]+]]: index
-// CHECK-NOT: memref.alloca
-// CHECK: cf.cond_br %{{.*}}, ^[[BB1:.*]], ^[[BB2:.*]]
-// CHECK: ^[[BB1]]:
-// CHECK: %[[CM:.*]] = vector.create_mask %[[N]], %{{.*}} : vector<8x16xi1>
-// CHECK: %[[SEL:.*]] = arith.select %[[CM]], %[[W]], %[[V]]
-// CHECK: cf.br ^[[BB3:.*]](%[[SEL]] : vector<8x16xf32>)
-// CHECK: ^[[BB2]]:
-// CHECK: cf.br ^[[BB3]](%[[V]] : vector<8x16xf32>)
-// CHECK: ^[[BB3]](%[[PHI:.*]]: vector<8x16xf32>):
-// CHECK: return %[[PHI]]
-func.func @dyn_subview_across_cfg(%v: vector<8x16xf32>, %w: vector<8x16xf32>,
- %n: index, %pad: f32, %c: i1) -> vector<8x16xf32> {
+// 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>
- cf.cond_br %c, ^bb1, ^bb2
-^bb1:
- %sv = memref.subview %a[0, 0] [%n, 16] [1, 1] : memref<8x16xf32> to memref<?x16xf32, strided<[16, 1]>>
- vector.transfer_write %w, %sv[%c0, %c0] {in_bounds = [false, true]} : vector<8x16xf32>, memref<?x16xf32, strided<[16, 1]>>
- cf.br ^bb3
-^bb2:
- cf.br ^bb3
-^bb3:
- %r = vector.transfer_read %a[%c0, %c0], %pad {in_bounds = [true, true]} : memref<8x16xf32>, vector<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>
}
@@ -371,30 +395,48 @@ func.func @neg_static_offset_dyn_size(%v: vector<8x16xf32>, %n: index,
// -----
-// Negative: the collapsed group's dynamic extent is not its outermost, so the
-// group's stride would have to be that runtime extent. This is the collapse
-// counterpart of @neg_nonunit_dim_before_dynamic, rejected one step earlier:
-// such a collapse is not guaranteed contiguous, so its result stride is dynamic.
-// CHECK-LABEL: func.func @neg_collapse_dyn_inner_dim(
+// Negative: two dynamic dimensions in one reassociation group -- the alias
+// shape is not determined by the parent extent, so the buffer is not promoted.
+// CHECK-LABEL: func.func @neg_two_dynamic_dims_in_group(
// CHECK: memref.alloca
-// CHECK: memref.collapse_shape
+// CHECK: memref.expand_shape
// CHECK: vector.transfer_read
-func.func @neg_collapse_dyn_inner_dim(%v: vector<4x4xf32>, %n: index,
- %pad: f32) -> vector<16xf32> {
+func.func @neg_two_dynamic_dims_in_group(%v: vector<16xf32>, %m: index, %d0: index,
+ %d1: index, %pad: f32) -> vector<8x2xf32> {
%c0 = arith.constant 0 : index
- %a = memref.alloca() : memref<4x4xf32>
- vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<4x4xf32>, memref<4x4xf32>
- %sv = memref.subview %a[0, 0] [4, %n] [1, 1] : memref<4x4xf32> to memref<4x?xf32, strided<[4, 1]>>
- %f = memref.collapse_shape %sv [[0, 1]] : memref<4x?xf32, strided<[4, 1]>> into memref<?xf32, strided<[?]>>
- %r = vector.transfer_read %f[%c0], %pad {in_bounds = [false]} : memref<?xf32, strided<[?]>>, vector<16xf32>
- return %r : vector<16xf32>
+ %a = memref.alloca() : memref<16xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<16xf32>, memref<16xf32>
+ %sv = memref.subview %a[0] [%m] [1] : memref<16xf32> to memref<?xf32, strided<[1]>>
+ %exp = memref.expand_shape %sv [[0, 1]] output_shape [%d0, %d1] : memref<?xf32, strided<[1]>> into memref<?x?xf32>
+ %r = vector.transfer_read %exp[%c0, %c0], %pad : memref<?x?xf32>, vector<8x2xf32>
+ return %r : vector<8x2xf32>
+}
+
+// -----
+
+// Negative: a non-unit dimension precedes the dynamic one, so the view strides
+// by the runtime %m where an alias shaped from the parent strides by 8 / 2 = 4.
+// Promoting would move data, which no mask can undo.
+// CHECK-LABEL: func.func @neg_nonunit_dim_before_dynamic(
+// CHECK: memref.alloca
+// CHECK: memref.subview
+// CHECK: memref.expand_shape
+// CHECK: vector.transfer_read
+func.func @neg_nonunit_dim_before_dynamic(%v: vector<8xf32>, %n: index, %m: index,
+ %pad: f32) -> vector<2x4xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<8xf32>
+ vector.transfer_write %v, %a[%c0] {in_bounds = [true]} : vector<8xf32>, memref<8xf32>
+ %sv = memref.subview %a[0] [%n] [1] : memref<8xf32> to memref<?xf32, strided<[1]>>
+ %exp = memref.expand_shape %sv [[0, 1]] output_shape [2, %m] : memref<?xf32, strided<[1]>> into memref<2x?xf32>
+ %r = vector.transfer_read %exp[%c0, %c0], %pad {in_bounds = [true, false]} : memref<2x?xf32>, vector<2x4xf32>
+ return %r : vector<2x4xf32>
}
// -----
-// Negative, and a crash regression: a reassociation group may hold a zero static
-// extent, which the parent extent cannot be divided by. The buffer is not
-// promoted; without the guard the division faults.
+// Negative, and a crash regression: a group may hold a zero static extent,
+// which the parent extent cannot be divided by.
// CHECK-LABEL: func.func @neg_zero_extent_in_group(
// CHECK: memref.alloca
// CHECK: memref.subview
@@ -410,3 +452,23 @@ func.func @neg_zero_extent_in_group(%v: vector<8x16xf32>, %n: index, %d: index,
%r = vector.transfer_read %e[%c0, %c0, %c0], %pad : memref<?x0x16xf32, strided<[0, 16, 1]>>, vector<8x16xf32>
return %r : vector<8x16xf32>
}
+
+// -----
+
+// Negative: the collapsed group's dynamic extent is not its outermost, so the
+// collapse is not guaranteed contiguous. The collapse counterpart of
+// @neg_nonunit_dim_before_dynamic, rejected one step earlier.
+// CHECK-LABEL: func.func @neg_collapse_dyn_inner_dim(
+// CHECK: memref.alloca
+// CHECK: memref.collapse_shape
+// CHECK: vector.transfer_read
+func.func @neg_collapse_dyn_inner_dim(%v: vector<4x4xf32>, %n: index,
+ %pad: f32) -> vector<16xf32> {
+ %c0 = arith.constant 0 : index
+ %a = memref.alloca() : memref<4x4xf32>
+ vector.transfer_write %v, %a[%c0, %c0] {in_bounds = [true, true]} : vector<4x4xf32>, memref<4x4xf32>
+ %sv = memref.subview %a[0, 0] [4, %n] [1, 1] : memref<4x4xf32> to memref<4x?xf32, strided<[4, 1]>>
+ %f = memref.collapse_shape %sv [[0, 1]] : memref<4x?xf32, strided<[4, 1]>> into memref<?xf32, strided<[?]>>
+ %r = vector.transfer_read %f[%c0], %pad {in_bounds = [false]} : memref<?xf32, strided<[?]>>, vector<16xf32>
+ return %r : vector<16xf32>
+}
diff --git a/mlir/test/Dialect/Vector/mem2reg.mlir b/mlir/test/Dialect/Vector/mem2reg.mlir
index efd5dd298f495..1182a0bf09412 100644
--- a/mlir/test/Dialect/Vector/mem2reg.mlir
+++ b/mlir/test/Dialect/Vector/mem2reg.mlir
@@ -774,3 +774,97 @@ func.func @neg_expand_shape_strided_subview(%v: vector<8x16xf32>, %pad: f32) ->
%r = vector.transfer_read %exp[%c0, %c0, %c0], %pad {in_bounds = [true, true, true]} : memref<1x4x8xf32, strided<[64, 16, 1]>>, vector<1x4x8xf32>
return %r : vector<1x4x8xf32>
}
+
+// -----
+
+// `memref.copy` participates in promotion as a transfer of the buffer's value.
+// 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>
+}
+
+// -----
+
+// 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.
+
+// CHECK-LABEL: func.func @negative_self_copy(
+// CHECK: memref.alloca
+// CHECK: memref.copy
+// 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>
+}
+
+// -----
+
+// NEGATIVE: the copy models keep to fixed-size slots. The dynamic size is not
+// what blocks it -- being a vscale multiple is what makes the slot scalable, and
+// such a buffer does promote through transfers (@scalable_whole_buffer_in_loop).
+
+// CHECK-LABEL: func.func @negative_scalable_copy(
+// CHECK: memref.alloca
+// CHECK: memref.copy
+// CHECK: vector.transfer_read
+func.func @negative_scalable_copy(%src: memref<8xf32>, %pad: f32) -> vector<[4]xf32> {
+ %c0 = arith.constant 0 : index
+ %c4 = arith.constant 4 : index
+ %vs = vector.vscale
+ %sz = arith.muli %vs, %c4 : index
+ %a = memref.alloca(%sz) : memref<?xf32>
+ memref.copy %src, %a : memref<8xf32> to memref<?xf32>
+ %r = vector.transfer_read %a[%c0], %pad {in_bounds = [true]} : memref<?xf32>, vector<[4]xf32>
+ return %r : vector<[4]xf32>
+}
More information about the Mlir-commits
mailing list