[Mlir-commits] [mlir] [mlir][Vector] Promote buffers accessed via dynamic subviews and memref.copy (PR #217793)
Andrzej WarzyĆski
llvmlistbot at llvm.org
Sun Oct 11 09:44:17 PDT 2026
================
@@ -84,43 +95,273 @@ getPromotableSubViewOffsets(memref::SubViewOp subView) {
llvm::zip_equal(subView.getMixedSizes(), resType.getShape())) {
std::optional<int64_t> s = getConstantIntValue(size);
if (!s || *s != dim)
- return std::nullopt;
+ return false;
+ }
+ return true;
+}
+
+/// Returns the static offsets of `subView` within the parent vector for
+/// `vector.extract_strided_slice` and `vector.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;
}
+
+/// Returns whether `subView` can be promoted as an alias of the whole parent
+/// vector. Static shapes are handled by `isAliasableStaticShapeSubView`
+/// instead.
+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 false;
+
+ // At least one dynamic result dimension; a fully static sub-slice is handled
+ // by `isAliasableStaticShapeSubView` instead.
+ if (resType.hasStaticShape())
+ return false;
+
+ // Same rank: the read/write vector matches the parent rank.
+ if (srcType.getRank() != resType.getRank())
+ 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).
+ 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 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());
+}
+
+/// If `slotPtr` is a dynamic subview, returns it.
+static memref::SubViewOp getDynamicSubView(Value slotPtr) {
+ auto subView = slotPtr.getDefiningOp<memref::SubViewOp>();
+ if (subView && isAliasableDynamicShapeSubView(subView))
+ return subView;
+ return {};
+}
+
+/// Builds the mask of `subView`'s valid region: a `vector.create_mask` of the
+/// subview's sizes, in the shape of the whole parent that its alias holds.
+static Value buildSubViewMask(OpBuilder &builder, Location loc,
+ memref::SubViewOp subView) {
+ VectorType parentVecType = getWholeParentVectorType(subView);
+ SmallVector<Value> bounds =
+ getValueOrCreateConstantIndexOp(builder, loc, subView.getMixedSizes());
+ return vector::CreateMaskOp::create(
+ builder, loc,
+ VectorType::get(parentVecType.getShape(), builder.getI1Type()), bounds);
+}
+
+/// 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 value does not affect the copy: both operands must have the same
+/// runtime shape. For a whole-buffer slot, the source covers the whole vector.
+/// For a dynamic-subview slot, the aliaser's masked select discards lanes
+/// outside the subview's extent.
+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));
+}
+
+//===----------------------------------------------------------------------===//
+// memref.copy
+//===----------------------------------------------------------------------===//
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.
+/// Mem2Reg model for `memref.copy`.
+///
+/// A copy into the slot becomes a `vector.transfer_read` of the source in
+/// `getStored`. A copy out of the slot becomes a `vector.transfer_write` of the
+/// reaching definition in `removeBlockingUses`.
+///
+/// A copy between two slots is rewritten one slot at a time, in either order.
+/// For dynamic subviews, the aliaser masks writes into the slot, and this model
+/// masks copies out of it.
+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;
+ }
+
+ // 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 {
+ auto copyOp = cast<memref::CopyOp>(op);
+ return readMemRefAsVector(builder, op->getLoc(), copyOp.getSource(),
+ cast<VectorType>(slot.valueType));
+ }
+
+ 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.valueType);
+ 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;
+ // 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);
+ if (copyOp.getSource() == slot.ptr) {
+ Location loc = op->getLoc();
+ Value mask;
+ if (memref::SubViewOp subView = getDynamicSubView(slot.ptr))
+ mask = buildSubViewMask(builder, loc, subView);
+ writeVectorToMemRef(builder, loc, reachingDefinition, copyOp.getTarget(),
+ mask);
+ }
+ return DeletionKind::Delete;
+ }
+};
+
+//===----------------------------------------------------------------------===//
+// memref.subview aliaser
----------------
banach-space wrote:
Here and in other places, instead of using terms without clear definition (what is "aliaser"), could you refer to interfaces from `MemorySlotInterfaces.td`?
In particular, I would these block comments to tell me which interfaces from "memory slot interfaces" is being implemented. That sort of organisation would make more sense to me and would also provide clear hint as to which interface to look forl
https://github.com/llvm/llvm-project/pull/217793
More information about the Mlir-commits
mailing list