[Mlir-commits] [mlir] [mlir][memref][mem2reg] Support expand_shape and collapse_shape aliases (PR #232535)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Oct 15 15:53:03 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
Extend MLIR Mem2Reg to promote buffers accessed through memref.expand_shape and memref.collapse_shape (deferred from PR#<!-- -->217793). Reshaped views become alias slots whose values are converted with vector.shape_cast.
Support composition with static and dynamic subviews, reshaping dynamic subview masks to preserve read padding and elements outside write bounds. Reject reshapes that cannot preserve the element mapping or count.
Validation: test landed under vector Mem2Reg test files, but will refactor as we land PR#<!-- -->232078.
assisted-by-claude
---
Patch is 39.78 KiB, truncated to 20.00 KiB below, full version: https://github.com/llvm/llvm-project/pull/232535.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.h (+11-7)
- (modified) mlir/lib/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.cpp (+236-35)
- (modified) mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp (+6-5)
- (modified) mlir/test/Dialect/Vector/mem2reg-dynamic-subview.mlir (+184)
- (modified) mlir/test/Dialect/Vector/mem2reg.mlir (+152)
``````````diff
diff --git a/mlir/include/mlir/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.h b/mlir/include/mlir/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.h
index cb94ccf1a3541b..5ecd6572e55844 100644
--- a/mlir/include/mlir/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.h
+++ b/mlir/include/mlir/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.h
@@ -9,7 +9,7 @@
#ifndef MLIR_DIALECT_MEMREF_TRANSFORMS_MEMORYSLOTOPINTERFACEIMPL_H
#define MLIR_DIALECT_MEMREF_TRANSFORMS_MEMORYSLOTOPINTERFACEIMPL_H
-#include "mlir/IR/Value.h"
+#include "mlir/Interfaces/MemorySlotInterfaces.h"
namespace mlir {
class DialectRegistry;
@@ -19,13 +19,17 @@ class OpBuilder;
namespace memref {
void registerMemorySlotOpInterfaceExternalModels(DialectRegistry ®istry);
-/// Returns whether `slotPtr` is a dynamic `memref.subview` that can be promoted
-/// as an alias of the whole parent vector.
-bool isDynamicSubViewSlot(Value slotPtr);
+/// Returns whether `slot` represents a supported dynamic `memref.subview`,
+/// optionally followed by `memref.expand_shape` or `memref.collapse_shape`.
+/// Such a slot retains the whole parent vector, so accesses must mask out
+/// elements outside the subview's runtime bounds.
+bool isPartialViewSlot(const MemorySlot &slot);
-/// Builds a mask for the subview's valid region in the parent vector, or
-/// returns null if `isDynamicSubViewSlot(slotPtr)` is false.
-Value buildDynamicSubViewMask(OpBuilder &builder, Location loc, Value slotPtr);
+/// Builds a mask of the subview's runtime bounds in the shape of
+/// `slot.valueType`, reshaping the mask if needed. Returns a null value for
+/// slots that do not require a subview mask. Expects a promotable vector slot.
+Value buildPartialViewMask(OpBuilder &builder, Location loc,
+ const MemorySlot &slot);
} // namespace memref
} // namespace mlir
diff --git a/mlir/lib/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.cpp
index c562c67433f470..04c4cca765e3d7 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -8,13 +8,13 @@
//
// This file implements memory slot interface models used by Mem2Reg to promote
// a memref to a single vector SSA value. `PromotableAliaserInterface` exposes
-// `memref.subview` aliases, and `PromotableOpInterface` removes the subviews
-// after promotion. Vector transfer ops are handled by the models in the
-// Vector dialect.
+// `memref.subview`, `memref.expand_shape` and `memref.collapse_shape` aliases,
+// and `PromotableOpInterface` removes those views after promotion. Vector
+// transfer ops are handled by the models in the Vector dialect.
//
// A memory slot pairs a memref with the vector type used to represent its
// contents. Promotion replaces accesses to the slot with uses and definitions
-// of that vector value. All uses, including those through subviews, must be
+// of that vector value. All uses, including those through views, must be
// supported by the models.
//
// The accesses are rewritten as follows:
@@ -31,12 +31,18 @@
// accesses to the subview's extent. The subview must start at the
// parent's origin.
//
-// Subviews require a statically shaped parent buffer. A dynamically
-// shaped buffer can instead promote to a 1-D scalable vector through
-// whole-buffer transfers, but subviews are not supported.
+// * A supported `memref.expand_shape` or `memref.collapse_shape` reshapes
+// the parent slot's vector with `vector.shape_cast` in both directions.
+// For a dynamic subview, the alias retains all elements of the parent
+// vector; the subview mask is cast to the alias shape when reading.
+//
+// Views require a statically shaped parent buffer. A dynamically shaped buffer
+// can instead promote to a 1-D scalable vector through whole-buffer transfers,
+// but views are not supported.
//
// Unsupported accesses, such as dynamic subview offsets, rank reduction,
-// non-unit subview strides, or non-zero transfer indices, prevent promotion.
+// non-unit subview strides, reshapes that cannot be represented by
+// `vector.shape_cast`, or non-zero transfer indices, prevent promotion.
//
//===----------------------------------------------------------------------===//
@@ -130,54 +136,175 @@ static VectorType getWholeParentVectorType(memref::SubViewOp subView) {
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 {};
+/// Follows the sources of `memref.expand_shape` and `memref.collapse_shape` to
+/// find the underlying value. Used to locate a dynamic subview whose bounds
+/// still apply after reshaping.
+static Value skipReassociativeReshapes(Value viewPtr) {
+ while (Operation *def = viewPtr.getDefiningOp()) {
+ if (auto expand = dyn_cast<memref::ExpandShapeOp>(def))
+ viewPtr = expand.getSrc();
+ else if (auto collapse = dyn_cast<memref::CollapseShapeOp>(def))
+ viewPtr = collapse.getSrc();
+ else
+ break;
+ }
+ return viewPtr;
}
-/// 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.
+/// Builds a mask from the subview's sizes in the parent vector's shape, then
+/// casts it to `vecType`'s shape if needed. Supported reshapes preserve element
+/// order, so a single cast also handles a chain of reshapes.
static Value buildSubViewMask(OpBuilder &builder, Location loc,
- memref::SubViewOp subView) {
+ memref::SubViewOp subView, VectorType vecType) {
VectorType parentVecType = getWholeParentVectorType(subView);
SmallVector<Value> bounds =
getValueOrCreateConstantIndexOp(builder, loc, subView.getMixedSizes());
- return vector::CreateMaskOp::create(
+ 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;
}
-bool mlir::memref::isDynamicSubViewSlot(Value slotPtr) {
- return static_cast<bool>(getDynamicSubView(slotPtr));
+/// Checks whether a reassociation group preserves the mapping of valid view
+/// elements to vector lanes. `viewExtents` are the memref dimensions;
+/// `aliasExtents` are the corresponding static vector dimensions.
+///
+/// Leading unit dimensions contribute no offset. The first non-unit dimension
+/// may differ because its extent does not affect its own row-major stride.
+/// All subsequent extents must match to keep that stride unchanged. Thus a
+/// supported group with a dynamic extent has the form
+/// `[1, ..., 1, dynamic, static, ..., static]`.
+static bool shapeCastModelsGroup(ArrayRef<int64_t> viewExtents,
+ ArrayRef<int64_t> aliasExtents) {
+ assert(viewExtents.size() == aliasExtents.size() && "group size mismatch");
+ // Leading unit dimensions always have index zero.
+ size_t pos = 0;
+ while (pos < viewExtents.size() && viewExtents[pos] == 1)
+ ++pos;
+ // Allow the first non-unit extent to differ, but require its stride factors
+ // (all following extents) to match.
+ if (pos < viewExtents.size())
+ ++pos;
+ for (; pos < viewExtents.size(); ++pos)
+ if (viewExtents[pos] != aliasExtents[pos])
+ return false;
+ return true;
+}
+
+/// Computes the collapsed alias shape by multiplying the extents from
+/// `parentSlotShape` in each reassociation group. These are the parent slot's
+/// static vector extents, even when the source memref has dynamic dimensions.
+/// Fails if a group cannot be represented by `vector.shape_cast`.
+static FailureOr<SmallVector<int64_t>>
+getShapeCastableAliasShape(memref::CollapseShapeOp op,
+ ArrayRef<int64_t> parentSlotShape) {
+ ArrayRef<int64_t> viewShape = op.getSrcType().getShape();
+ // Reassociation indices must address both the memref and vector shapes.
+ if (viewShape.size() != parentSlotShape.size())
+ return failure();
+
+ SmallVector<int64_t> aliasShape;
+ for (ReassociationIndices group : op.getReassociationIndices()) {
+ SmallVector<int64_t> viewExtents, aliasExtents;
+ for (int64_t dim : group) {
+ viewExtents.push_back(viewShape[dim]);
+ aliasExtents.push_back(parentSlotShape[dim]);
+ }
+ if (!shapeCastModelsGroup(viewExtents, aliasExtents))
+ return failure();
+ aliasShape.push_back(ShapedType::getNumElements(aliasExtents));
+ }
+ if (ShapedType::getNumElements(aliasShape) !=
+ ShapedType::getNumElements(parentSlotShape))
+ return failure();
+ return aliasShape;
}
-Value mlir::memref::buildDynamicSubViewMask(OpBuilder &builder, Location loc,
- Value slotPtr) {
- memref::SubViewOp subView = getDynamicSubView(slotPtr);
- if (!subView)
+/// Computes a static expanded alias shape from `parentSlotShape`, the parent
+/// slot's static vector extents. For each reassociation group, divides the
+/// corresponding extent from `parentSlotShape` by the static result extents,
+/// starting at the innermost dimension. The first dynamic dimension receives
+/// the remaining extent; dimensions before it remain 1.
+///
+/// Fails if the parent extent cannot be evenly divided among the static result
+/// dimensions, or if a group fails `shapeCastModelsGroup`. A group with
+/// dynamic dimensions must have the form
+/// `[1, ..., 1, dynamic, static, ..., static]`.
+static FailureOr<SmallVector<int64_t>>
+getShapeCastableAliasShape(memref::ExpandShapeOp op,
+ ArrayRef<int64_t> parentSlotShape) {
+ ArrayRef<int64_t> viewShape = op.getType().getShape();
+ SmallVector<ReassociationIndices, 4> groups = op.getReassociationIndices();
+ if (groups.size() != parentSlotShape.size())
+ return failure();
+
+ SmallVector<int64_t> aliasShape(viewShape.size(), 1);
+ for (auto [srcDim, group] : llvm::enumerate(groups)) {
+ int64_t remaining = parentSlotShape[srcDim];
+ for (int64_t dim : llvm::reverse(group)) {
+ if (ShapedType::isDynamic(viewShape[dim])) {
+ aliasShape[dim] = remaining;
+ break;
+ }
+ if (viewShape[dim] == 0 || remaining % viewShape[dim] != 0)
+ return failure();
+ aliasShape[dim] = viewShape[dim];
+ remaining /= viewShape[dim];
+ }
+
+ SmallVector<int64_t> viewExtents, aliasExtents;
+ for (int64_t dim : group) {
+ viewExtents.push_back(viewShape[dim]);
+ aliasExtents.push_back(aliasShape[dim]);
+ }
+ if (!shapeCastModelsGroup(viewExtents, aliasExtents))
+ return failure();
+ }
+ if (ShapedType::getNumElements(aliasShape) !=
+ ShapedType::getNumElements(parentSlotShape))
+ return failure();
+ return aliasShape;
+}
+
+bool mlir::memref::isPartialViewSlot(const MemorySlot &slot) {
+ auto subView =
+ skipReassociativeReshapes(slot.ptr).getDefiningOp<memref::SubViewOp>();
+ return subView && isAliasableDynamicShapeSubView(subView);
+}
+
+Value mlir::memref::buildPartialViewMask(OpBuilder &builder, Location loc,
+ const MemorySlot &slot) {
+ auto subView =
+ skipReassociativeReshapes(slot.ptr).getDefiningOp<memref::SubViewOp>();
+ // Static subview slots contain only the selected slice and need no mask.
+ // Dynamic subviews reaching here have already passed alias eligibility
+ // checks when their slots were created.
+ if (!subView || subView.getType().hasStaticShape())
return {};
- return buildSubViewMask(builder, loc, subView);
+ return buildSubViewMask(builder, loc, subView,
+ cast<VectorType>(slot.valueType));
}
namespace {
//===----------------------------------------------------------------------===//
-// PromotableOpInterface model for memref.subview
+// PromotableOpInterface model for the memref view ops
//===----------------------------------------------------------------------===//
-/// `PromotableOpInterface` model for `memref.subview`. Works with the
-/// `PromotableAliaserInterface` model below to remove the subview once its
-/// result's uses have been promoted.
-struct SubViewOpPromotableModel
- : public PromotableOpInterface::ExternalModel<SubViewOpPromotableModel,
- memref::SubViewOp> {
+/// Removes `memref.subview`, `memref.expand_shape`, and `memref.collapse_shape`
+/// once all uses of their results have been promoted.
+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 {
- // Removing the subview requires removing all uses of its memref result.
+ // Removing the view requires removing all uses of its memref result.
for (OpOperand &use : op->getResult(0).getUses())
newBlockingUses.push_back(&use);
return true;
@@ -272,7 +399,8 @@ struct SubViewOpAliasModel
// of the parent value untouched.
if (isAliasableDynamicShapeSubView(subView)) {
Location loc = op->getLoc();
- Value mask = buildSubViewMask(builder, loc, subView);
+ VectorType parentVecType = getWholeParentVectorType(subView);
+ Value mask = buildSubViewMask(builder, loc, subView, parentVecType);
return arith::SelectOp::create(builder, loc, mask, aliasValue,
reachingDef);
}
@@ -284,6 +412,70 @@ struct SubViewOpAliasModel
.getResult();
}
};
+
+//===----------------------------------------------------------------------===//
+// PromotableAliaserInterface model for the memref reshape ops
+//===----------------------------------------------------------------------===//
+
+/// Exposes supported `memref.expand_shape` and `memref.collapse_shape` ops as
+/// aliases of vector slots. Both projections use `vector.shape_cast`, which
+/// preserves element order. `getShapeCastableAliasShape` checks eligibility
+/// and derives the alias shape from the parent vector's static extents.
+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;
+
+ // Only fixed-size vector slots support these shape casts.
+ auto parentVecType = dyn_cast<VectorType>(parentSlot.valueType);
+ if (!parentVecType || parentVecType.isScalable())
+ return;
+
+ auto resType = cast<MemRefType>(viewOp.getResult().getType());
+ if (!VectorType::isValidElementType(resType.getElementType()) ||
+ resType.getElementType() != parentVecType.getElementType())
+ return;
+ FailureOr<SmallVector<int64_t>> aliasShape =
+ getShapeCastableAliasShape(viewOp, parentVecType.getShape());
+ if (failed(aliasShape))
+ return;
+ auto aliasVecType = VectorType::get(*aliasShape, resType.getElementType());
+ newSlots.push_back(MemorySlot{viewOp.getResult(), aliasVecType});
+ }
+
+ 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.valueType),
+ 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.valueType),
+ aliasValue)
+ .getResult();
+ }
+};
} // namespace
//===----------------------------------------------------------------------===//
@@ -294,6 +486,15 @@ void mlir::memref::registerMemorySlotOpInterfaceExternalModels(
DialectRegistry ®istry) {
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);
});
}
diff --git a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
index c8c17f2039b025..71256e514058d4 100644
--- a/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/MemorySlotOpInterfaceImpl.cpp
@@ -9,7 +9,7 @@
// This file implements memory slot interface models for
// `vector.transfer_read` and `vector.transfer_write`. These models implement
// `PromotableMemOpInterface` so Mem2Reg can replace vector transfers with uses
-// and definitions of a vector SSA value. The models for `memref.subview`
+// and definitions of a vector SSA value. The models for the `memref` view ops
// live in the MemRef dialect.
//
// A read uses the slot's current value; a write defines its next value.
@@ -19,7 +19,8 @@
// whole static parent, and the subview's slot holds that whole vector. Reads
// pad lanes outside the runtime extent, and writes preserve the previous value
// of those lanes. `arith.select` implements this behavior and also handles
-// lanes disabled by a transfer's mask operand.
+// lanes disabled by a transfer's mask operand. For reads through reshaped
+// dynamic subviews, the subview mask is cast to the slot's vector shape.
//
// Transfers wrapped in `vector.mask` prevent promotion because that op does
// not implement `PromotableRegionOpInterface`, which Mem2Reg requires to
@@ -48,7 +49,7 @@ using namespace mlir::vector;
/// value. The transfer's memref base must be the slot pointer, its vector type
/// must equal `slot.valueType`, its indices must be zero, and its permutation
/// map must be the identity. Out-of-bounds dimensions are allowed only for a
-/// supported dynamic subview.
+/// supported dynamic subview, possibly reached through a reassociative reshape.
///
/// More general transfers may be supportable, but their value depends on
///...
[truncated]
``````````
</details>
https://github.com/llvm/llvm-project/pull/232535
More information about the Mlir-commits
mailing list