[Mlir-commits] [mlir] [mlir] Interface-ify updating starting positions on vector.transfer_* (PR #195186)
Krzysztof Drewniak
llvmlistbot at llvm.org
Mon May 4 08:32:50 PDT 2026
https://github.com/krzysz00 updated https://github.com/llvm/llvm-project/pull/195186
>From 00005e41f7ae3cc49f74aee5e2e91819a0b46d24 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 19 Jan 2026 19:09:37 +0000
Subject: [PATCH 1/2] [mlir] Interface-ify updating starting positions on
vector.transfer_*
This commit adds methods to VectorTransferOpInterface that allow
transfer operations to be queried for whether their base memref (or
tensor) and permutation map can be updated in some particular way and
then for performing this update. This is part of a series of changes
designed to make passes like fold-memref-alias-ops more generic,
allowing downstream operations, like IREE's transfer_gather, to
participate in them without needing to duplicate patterns.
In order to test this new method, migrate FoldMemrefAliasOps to use
these methods to fold memref.subview, memref.expand_shape,and
memref.collapse_shape into tranfer_read and transfer_write.
AI note: the tranfer_read / transfer_write patterns, which are taken
from a previous PR, were written with Claude 4.5.
---
.../mlir/Interfaces/VectorInterfaces.h | 1 +
.../mlir/Interfaces/VectorInterfaces.td | 68 ++++-
.../MemRef/Transforms/FoldMemRefAliasOps.cpp | 282 +++++++++++-------
.../Dialect/MemRef/fold-memref-alias-ops.mlir | 41 ++-
4 files changed, 288 insertions(+), 104 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/VectorInterfaces.h b/mlir/include/mlir/Interfaces/VectorInterfaces.h
index 7ae4ee35de337..551d2c1dbc445 100644
--- a/mlir/include/mlir/Interfaces/VectorInterfaces.h
+++ b/mlir/include/mlir/Interfaces/VectorInterfaces.h
@@ -16,6 +16,7 @@
#include "mlir/IR/AffineMap.h"
#include "mlir/IR/BuiltinTypes.h"
#include "mlir/IR/OpDefinition.h"
+#include "mlir/IR/PatternMatch.h"
/// Include the generated interface declarations.
#include "mlir/Interfaces/VectorInterfaces.h.inc"
diff --git a/mlir/include/mlir/Interfaces/VectorInterfaces.td b/mlir/include/mlir/Interfaces/VectorInterfaces.td
index 6838c16fdf0fe..d79588e6d42d7 100644
--- a/mlir/include/mlir/Interfaces/VectorInterfaces.td
+++ b/mlir/include/mlir/Interfaces/VectorInterfaces.td
@@ -135,7 +135,8 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
InterfaceMethod<
/*desc=*/[{
Return the indices that specify the starting offsets into the source
- operand. The starting offsets are guaranteed to be in-bounds.
+ operand. The starting offsets are guaranteed to be in-bounds except
+ when the transfer operation is masked.
}],
/*retTy=*/"::mlir::OperandRange",
/*methodName=*/"getIndices",
@@ -174,6 +175,71 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
/*retTy=*/"Value",
/*methodName=*/"getMask",
/*args=*/(ins)
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Returns whether replacing the base operand with a shaped value of
+ `newBaseType` and updating the permutation map according to
+ `newPermutationMap`, where the replacement would adhere to the
+ conditinos specified by `updateStartingPosition`, passes any addihtional
+ op-specific constraints.
+
+ (Note: this method has been added in case it ends up being needed
+ and, if some upstream or downstream use requires additional argument,
+ they should be added.)
+ }],
+ /*retTy=*/"::mlir::LogicalResult",
+ /*methodName=*/"mayUpdateStartingPosition",
+ /*args=*/(ins "::mlir::ShapedType":$newBaseType,
+ "::mlir::AffineMap":$newPermutationMap),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ return ::mlir::success();
+ }]
+ >,
+ InterfaceMethod<
+ /*desc=*/[{
+ Updates the base of this transfer operation to `newBase`, the indices into
+ that base into `newIndices`, and the permutation map to `newPermutationMap`.
+ The new base, indices, and map must:
+ - Keep the same element type as the base (this condition could be relaxed
+ in the future)
+ - Preserve the number of transefr dimensions and avoid inserting new
+ dimensions unaccounted for in the permutation map - that is, if the transfer operation is
+ reading or writing a 4x2 vector to/from an Nx4x2 tensor or memref, it is
+ not valid to update that base to an Nx8 one, but updating in an N0xN1x4x2
+ base - or even an Nx16x32 one - is acceptable. Similarly, if the transfer
+ rank is 2 and the base is MxN, replacing it with a base that is Mx1xN or
+ MxKxN is only permitted if the permutation map is updated to skip the new
+ dimension. This restriction means that the updater does not need to reshape
+ or otherwise adjust the vector type of the transfer or other related attributes
+ like the inbounds attribute array.
+ - Must not shrink the size of a transfer dimension such that an `inbounds`
+ annotation becomes incorrect.
+ - Must pass any op-specific conditions in `mayUpdateStartingPosition`
+
+ The update is performed in-place. Implementations of the interface
+ should use `RewriterBase::modifyOpInPlace` to ensure proper pattern
+ rewriter operation.
+
+ Note: The initial motivating factcor for adding this interface method is to
+ enable folding operations like `memref.subview` with transfer opes in a
+ generic way. If future usecases require extending this method, it should be
+ done.
+ }],
+ /*retTy=*/"void",
+ /*methodName=*/"updateStartingPosition",
+ /*args=*/(ins "::mlir::RewriterBase&":$rewriter, "::mlir::Value":$newBase,
+ "::mlir::ValueRange":$newIndices,
+ "::mlir::AffineMapAttr":$newPermutationMap),
+ /*methodBody=*/"",
+ /*defaultImplementation=*/[{
+ rewriter.modifyOpInPlace($_op, [&]() {
+ $_op.getBaseMutable().assign(newBase);
+ $_op.getIndicesMutable().assign(newIndices);
+ $_op.setPermutationMapAttr(newPermutationMap);
+ });
+ }]
>
];
diff --git a/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp b/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
index df42cfeefa1c1..e36ddfa063e11 100644
--- a/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
+++ b/mlir/lib/Dialect/MemRef/Transforms/FoldMemRefAliasOps.cpp
@@ -70,10 +70,6 @@ static Value getMemRefOperand(LoadOrStoreOpTy op) {
return op.getMemref();
}
-static Value getMemRefOperand(vector::TransferReadOp op) {
- return op.getBase();
-}
-
static Value getMemRefOperand(vector::LoadOp op) { return op.getBase(); }
static Value getMemRefOperand(vector::StoreOp op) { return op.getBase(); }
@@ -82,10 +78,6 @@ static Value getMemRefOperand(vector::MaskedLoadOp op) { return op.getBase(); }
static Value getMemRefOperand(vector::MaskedStoreOp op) { return op.getBase(); }
-static Value getMemRefOperand(vector::TransferWriteOp op) {
- return op.getBase();
-}
-
//===----------------------------------------------------------------------===//
// Patterns
//===----------------------------------------------------------------------===//
@@ -251,40 +243,53 @@ struct IndexedMemCopyOpOfCollapseShapeOpFolder final
LogicalResult matchAndRewrite(memref::IndexedMemCopyOpInterface op,
PatternRewriter &rewriter) const override;
};
-} // namespace
-template <typename XferOp>
-static LogicalResult
-preconditionsFoldSubViewOpImpl(RewriterBase &rewriter, XferOp xferOp,
- memref::SubViewOp subviewOp) {
- static_assert(
- !llvm::is_one_of<vector::TransferReadOp, vector::TransferWriteOp>::value,
- "must be a vector transfer op");
- if (xferOp.hasOutOfBoundsDim())
- return rewriter.notifyMatchFailure(xferOp, "out of bounds transfer dim");
- if (!subviewOp.hasUnitStride()) {
- return rewriter.notifyMatchFailure(
- xferOp, "non-1 stride subview, need to track strides in folded memref");
- }
- return success();
-}
+/// Merges memref.subview ops on the base argument to vector transfer operations
+/// into the base and indices of that transfer if:
+/// - The subview has unit strides on transfer dimensions
+/// - All the transfer dimensions are in-bounds
+/// This will correctly update said permutation map to account for dropped
+/// dimensions in rank-reducing subviews.
+struct TransferOpOfSubViewOpFolder final
+ : OpInterfaceRewritePattern<VectorTransferOpInterface> {
+ using Base::Base;
-static LogicalResult preconditionsFoldSubViewOp(RewriterBase &rewriter,
- Operation *op,
- memref::SubViewOp subviewOp) {
- return success();
-}
+ LogicalResult matchAndRewrite(VectorTransferOpInterface op,
+ PatternRewriter &rewriter) const override;
+};
-static LogicalResult preconditionsFoldSubViewOp(RewriterBase &rewriter,
- vector::TransferReadOp readOp,
- memref::SubViewOp subviewOp) {
- return preconditionsFoldSubViewOpImpl(rewriter, readOp, subviewOp);
-}
+/// Merges memref.expand_shape ops that create the base of a vector transfer
+/// operation into the base and indices of that transfer. Does not act when the
+/// a dimension is potentially out of bounds, if one of the transfer dimensions
+/// would need to be strided because of the collapse, or if it would merge two
+/// dimensions that are both transfer dimensions.
+/// TODO: become more sophisticated about length-1 dimensions that are the
+/// result of an expansion becoming broadcasts.
+struct TransferOpOfExpandShapeOpFolder final
+ : OpInterfaceRewritePattern<VectorTransferOpInterface> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(VectorTransferOpInterface op,
+ PatternRewriter &rewriter) const override;
+};
+
+/// Merges memref.collapse_shape ops that create the base of a vector transfer
+/// operation into the base and indices of that transfer. Does not act when the
+/// permutation map is not trivial, a dimension could be performing out of
+/// bounds reads, or if it would break apart a transfer dimension.
+struct TransferOpOfCollapseShapeOpFolder final
+ : OpInterfaceRewritePattern<VectorTransferOpInterface> {
+ using Base::Base;
+
+ LogicalResult matchAndRewrite(VectorTransferOpInterface op,
+ PatternRewriter &rewriter) const override;
+};
+} // namespace
static LogicalResult preconditionsFoldSubViewOp(RewriterBase &rewriter,
- vector::TransferWriteOp writeOp,
+ Operation *op,
memref::SubViewOp subviewOp) {
- return preconditionsFoldSubViewOpImpl(rewriter, writeOp, subviewOp);
+ return success();
}
template <typename OpTy>
@@ -321,14 +326,6 @@ LogicalResult LoadOpOfSubViewOpFolder<OpTy>::matchAndRewrite(
op, op.getType(), subViewOp.getSource(), sourceIndices,
op.getMask(), op.getPassThru());
})
- .Case([&](vector::TransferReadOp op) {
- rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
- op, op.getVectorType(), subViewOp.getSource(), sourceIndices,
- AffineMapAttr::get(expandDimsToRank(
- op.getPermutationMap(), subViewOp.getSourceType().getRank(),
- subViewOp.getDroppedDims())),
- op.getPadding(), op.getMask(), op.getInBoundsAttr());
- })
.DefaultUnreachable("unexpected operation");
return success();
}
@@ -342,40 +339,6 @@ LogicalResult LoadOpOfExpandShapeOpFolder<OpTy>::matchAndRewrite(
if (!expandShapeOp)
return failure();
- // For vector::TransferReadOp, validate preconditions before creating any IR.
- // resolveSourceIndicesExpandShape creates new ops, so all checks that can
- // fail must happen before that call to avoid "pattern returned failure but
- // IR did change" errors (caught by MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS).
- SmallVector<AffineExpr> transferReadNewResults;
- if (auto transferOp =
- dyn_cast<vector::TransferReadOp>(loadOp.getOperation())) {
- const int64_t vectorRank = transferOp.getVectorType().getRank();
- const int64_t sourceRank =
- cast<MemRefType>(expandShapeOp.getViewSource().getType()).getRank();
- if (sourceRank < vectorRank)
- return failure();
-
- // We can only fold if the permutation map uses only the least significant
- // dimension from each expanded reassociation group.
- for (AffineExpr result : transferOp.getPermutationMap().getResults()) {
- bool foundExpr = false;
- for (auto reassocationIndices :
- llvm::enumerate(expandShapeOp.getReassociationIndices())) {
- auto reassociation = reassocationIndices.value();
- AffineExpr dim = getAffineDimExpr(
- reassociation[reassociation.size() - 1], rewriter.getContext());
- if (dim == result) {
- transferReadNewResults.push_back(getAffineDimExpr(
- reassocationIndices.index(), rewriter.getContext()));
- foundExpr = true;
- break;
- }
- }
- if (!foundExpr)
- return failure();
- }
- }
-
SmallVector<Value> sourceIndices;
// memref.load guarantees that indexes start inbounds while the vector
// operations don't. This impacts if our linearization is `disjoint`
@@ -402,16 +365,6 @@ LogicalResult LoadOpOfExpandShapeOpFolder<OpTy>::matchAndRewrite(
op.getMask(), op.getPassThru());
return success();
})
- .Case([&](vector::TransferReadOp op) {
- const int64_t sourceRank = sourceIndices.size();
- auto newMap = AffineMap::get(sourceRank, 0, transferReadNewResults,
- op.getContext());
- rewriter.replaceOpWithNewOp<vector::TransferReadOp>(
- op, op.getVectorType(), expandShapeOp.getViewSource(),
- sourceIndices, newMap, op.getPadding(), op.getMask(),
- op.getInBounds());
- return success();
- })
.DefaultUnreachable("unexpected operation");
}
@@ -473,14 +426,6 @@ LogicalResult StoreOpOfSubViewOpFolder<OpTy>::matchAndRewrite(
op, op.getValue(), subViewOp.getSource(), sourceIndices,
op.getNontemporal());
})
- .Case([&](vector::TransferWriteOp op) {
- rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
- op, op.getValue(), subViewOp.getSource(), sourceIndices,
- AffineMapAttr::get(expandDimsToRank(
- op.getPermutationMap(), subViewOp.getSourceType().getRank(),
- subViewOp.getDroppedDims())),
- op.getMask(), op.getInBoundsAttr());
- })
.Case([&](vector::StoreOp op) {
rewriter.replaceOpWithNewOp<vector::StoreOp>(
op, op.getValueToStore(), subViewOp.getSource(), sourceIndices);
@@ -764,23 +709,160 @@ LogicalResult IndexedMemCopyOpOfCollapseShapeOpFolder::matchAndRewrite(
return success();
}
+LogicalResult
+TransferOpOfSubViewOpFolder::matchAndRewrite(VectorTransferOpInterface op,
+ PatternRewriter &rewriter) const {
+ auto subview = op.getBase().getDefiningOp<memref::SubViewOp>();
+ if (!subview)
+ return rewriter.notifyMatchFailure(op, "not accessing a subview");
+
+ AffineMap perm = op.getPermutationMap();
+ // Note: no identity permutation check here, since subview folding can handle
+ // complex permutations because it doesn't merge or split any individual
+ // dimension.
+ if (op.hasOutOfBoundsDim())
+ return rewriter.notifyMatchFailure(op, "out of bounds dimension");
+ VectorType vecTy = op.getVectorType();
+ // Because we know the permutation map is a minor identity, we know that the
+ // last N dimensions must have unit stride, where N is the vector rank.
+ if (!hasTrailingUnitStrides(subview, vecTy.getRank()))
+ return rewriter.notifyMatchFailure(subview, "non-unit stride within last " +
+ Twine(vecTy.getRank()) +
+ " dimensions");
+
+ AffineMap newPerm = expandDimsToRank(perm, subview.getSourceType().getRank(),
+ subview.getDroppedDims());
+
+ if (failed(op.mayUpdateStartingPosition(subview.getSourceType(), newPerm)))
+ return rewriter.notifyMatchFailure(subview,
+ "failed op-specific preconditions");
+
+ SmallVector<Value> newIndices;
+ affine::resolveIndicesIntoOpWithOffsetsAndStrides(
+ rewriter, op.getLoc(), subview.getMixedOffsets(),
+ subview.getMixedStrides(), subview.getDroppedDims(), op.getIndices(),
+ newIndices);
+ op.updateStartingPosition(rewriter, subview.getSource(), newIndices,
+ AffineMapAttr::get(newPerm));
+ return success();
+}
+
+LogicalResult TransferOpOfExpandShapeOpFolder::matchAndRewrite(
+ VectorTransferOpInterface op, PatternRewriter &rewriter) const {
+ auto expand = op.getBase().getDefiningOp<memref::ExpandShapeOp>();
+ if (!expand)
+ return rewriter.notifyMatchFailure(op, "not accessing an expand_shape");
+
+ if (op.hasOutOfBoundsDim())
+ return rewriter.notifyMatchFailure(op, "out of bounds dimension");
+
+ int64_t srcRank = expand.getSrc().getType().getRank();
+ int64_t vecRank = op.getVectorType().getRank();
+ if (srcRank < vecRank)
+ return rewriter.notifyMatchFailure(op,
+ "source rank is less than vector rank");
+
+ llvm::SmallDenseMap<int64_t, int64_t, 8> unstridedResDimToSrcDim;
+ for (auto [srcIdx, reassoc] :
+ llvm::enumerate(expand.getReassociationIndices())) {
+ unstridedResDimToSrcDim.insert({reassoc.back(), srcIdx});
+ }
+ // If every dimension of the expanded shape that appears in the permutation
+ // map is also present in the final entry of the expansions (meaning that
+ // collapsing in more values won't cause us to need to stride the index), we
+ // can fold in the expansion. (This doesn't currently account for expanding
+ // length X to X by 1, but it could in the future).
+ AffineMap permMap = op.getPermutationMap();
+ SmallVector<AffineExpr> newPermMapResults;
+ newPermMapResults.reserve(permMap.getNumResults());
+ for (AffineExpr permRes : permMap.getResults()) {
+ auto resDim = dyn_cast<AffineDimExpr>(permRes);
+ if (!resDim)
+ return rewriter.notifyMatchFailure(
+ op, "has non-dim entry in permutation map");
+ auto dimInSrc = unstridedResDimToSrcDim.find(resDim.getPosition());
+ if (dimInSrc == unstridedResDimToSrcDim.end())
+ return rewriter.notifyMatchFailure(op,
+ "permutation map result would be made "
+ "strided by expand_shape folding");
+ newPermMapResults.push_back(rewriter.getAffineDimExpr(dimInSrc->second));
+ }
+
+ auto newPerm = AffineMap::get(srcRank, 0, newPermMapResults, op.getContext());
+
+ if (failed(op.mayUpdateStartingPosition(expand.getSrc().getType(), newPerm)))
+ return rewriter.notifyMatchFailure(op, "failed op-specific preconditions");
+
+ SmallVector<Value> newIndices;
+ // We can use a disjoint linearization if we aren't masking, because then all
+ // indicators show that the start position will be in bounds.
+ memref::resolveSourceIndicesExpandShape(op.getLoc(), rewriter, expand,
+ op.getIndices(), newIndices,
+ /*startsInbounds=*/!op.getMask());
+
+ op.updateStartingPosition(rewriter, expand.getViewSource(), newIndices,
+ AffineMapAttr::get(newPerm));
+ return success();
+}
+
+LogicalResult TransferOpOfCollapseShapeOpFolder::matchAndRewrite(
+ VectorTransferOpInterface op, PatternRewriter &rewriter) const {
+ auto collapse = op.getBase().getDefiningOp<memref::CollapseShapeOp>();
+ if (!collapse)
+ return rewriter.notifyMatchFailure(op, "not accessing a collapse_shape");
+
+ if (!op.getPermutationMap().isMinorIdentity())
+ return rewriter.notifyMatchFailure(op,
+ "non-minor identity permutation map");
+
+ if (op.hasOutOfBoundsDim())
+ return rewriter.notifyMatchFailure(op, "out of bounds dimension");
+
+ int64_t srcRank = collapse.getSrc().getType().getRank();
+ int64_t vecRank = op.getVectorType().getRank();
+ if (srcRank < vecRank)
+ return rewriter.notifyMatchFailure(op,
+ "source rank is less than vector rank");
+
+ // Note: no - 1 on the rank here. While we could treat the collapse of [1, 1,
+ // N] into N as a special case, that is left as future work for those who need
+ // such a pattern.
+ SmallVector<ReassociationIndices> reassocs =
+ collapse.getReassociationIndices();
+ if (!hasTrivialReassociationSuffix(reassocs, vecRank))
+ return rewriter.notifyMatchFailure(
+ op, "collapse_shape folding would split a transfer dimension");
+
+ AffineMap newPerm =
+ AffineMap::getMinorIdentityMap(srcRank, vecRank, op.getContext());
+ if (failed(
+ op.mayUpdateStartingPosition(collapse.getSrc().getType(), newPerm)))
+ return rewriter.notifyMatchFailure(op, "failed op-specific preconditions");
+
+ SmallVector<Value> newIndices;
+ memref::resolveSourceIndicesCollapseShape(op.getLoc(), rewriter, collapse,
+ op.getIndices(), newIndices);
+
+ op.updateStartingPosition(rewriter, collapse.getViewSource(), newIndices,
+ AffineMapAttr::get(newPerm));
+ return success();
+}
+
void memref::populateFoldMemRefAliasOpPatterns(RewritePatternSet &patterns) {
patterns.add<
// Interface-based patterns to which we will be migrating.
AccessOpOfSubViewOpFolder, AccessOpOfExpandShapeOpFolder,
AccessOpOfCollapseShapeOpFolder, IndexedMemCopyOpOfSubViewOpFolder,
IndexedMemCopyOpOfExpandShapeOpFolder,
- IndexedMemCopyOpOfCollapseShapeOpFolder,
+ IndexedMemCopyOpOfCollapseShapeOpFolder, TransferOpOfSubViewOpFolder,
+ TransferOpOfExpandShapeOpFolder, TransferOpOfCollapseShapeOpFolder,
// The old way of doing things. Don't add more of these.
LoadOpOfSubViewOpFolder<vector::LoadOp>,
LoadOpOfSubViewOpFolder<vector::MaskedLoadOp>,
- LoadOpOfSubViewOpFolder<vector::TransferReadOp>,
- StoreOpOfSubViewOpFolder<vector::TransferWriteOp>,
StoreOpOfSubViewOpFolder<vector::StoreOp>,
StoreOpOfSubViewOpFolder<vector::MaskedStoreOp>,
LoadOpOfExpandShapeOpFolder<vector::LoadOp>,
LoadOpOfExpandShapeOpFolder<vector::MaskedLoadOp>,
- LoadOpOfExpandShapeOpFolder<vector::TransferReadOp>,
StoreOpOfExpandShapeOpFolder<vector::StoreOp>,
StoreOpOfExpandShapeOpFolder<vector::MaskedStoreOp>,
LoadOpOfCollapseShapeOpFolder<vector::LoadOp>,
diff --git a/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir b/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
index 2084dbc0e35a4..50c7ebaff1e6a 100644
--- a/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
+++ b/mlir/test/Dialect/MemRef/fold-memref-alias-ops.mlir
@@ -792,7 +792,7 @@ func.func @fold_vector_transfer_read_expand_shape(
// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
// CHECK: %[[C0:.*]] = arith.constant 0
// CHECK: %[[PAD:.*]] = ub.poison : f32
-// CHECK: %[[IDX:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (4, 8)
+// CHECK: %[[IDX:.*]] = affine.linearize_index disjoint [%[[ARG1]], %[[C0]]] by (4, 8)
// CHECK: vector.transfer_read %[[ARG0]][%[[IDX]]], %[[PAD]] {in_bounds = [true]}
// -----
@@ -812,8 +812,8 @@ func.func @fold_vector_transfer_read_expand_shape_non_identity(
// CHECK-SAME: %[[ARG2:[a-zA-Z0-9_]+]]: index
// CHECK: %[[C0:.*]] = arith.constant 0
// CHECK: %[[PAD:.*]] = ub.poison : f32
-// CHECK: %[[IDX1:.*]] = affine.linearize_index [%[[ARG1]], %[[C0]]] by (4, 8)
-// CHECK: %[[IDX2:.*]] = affine.linearize_index [%[[ARG2]], %[[C0]]] by (4, 8)
+// CHECK: %[[IDX1:.*]] = affine.linearize_index disjoint [%[[ARG1]], %[[C0]]] by (4, 8)
+// CHECK: %[[IDX2:.*]] = affine.linearize_index disjoint [%[[ARG2]], %[[C0]]] by (4, 8)
// CHECK: vector.transfer_read %[[ARG0]][%[[IDX1]], %[[IDX2]]], %[[PAD]] {in_bounds = [true, true]}
// -----
@@ -927,6 +927,41 @@ func.func @fold_vector_maskedstore_collapse_shape(
// -----
+func.func @no_fold_collapse_shape_transfer_read(
+ %arg0 : memref<4x4x8xf32>, %arg1 : index) -> vector<4x8xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = ub.poison : f32
+ %0 = memref.collapse_shape %arg0 [[0, 1], [2]] : memref<4x4x8xf32> into memref<16x8xf32>
+ %1 = vector.transfer_read %0[%arg1, %c0], %pad {in_bounds = [true, true]} : memref<16x8xf32>, vector<4x8xf32>
+ return %1 : vector<4x8xf32>
+}
+
+// CHECK-LABEL: func @no_fold_collapse_shape_transfer_read
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x4x8xf32>
+// CHECK: memref.collapse_shape %[[ARG0]]
+// CHECK: vector.transfer_read
+
+// -----
+
+func.func @fold_collapse_shape_transfer_read(
+ %arg0 : memref<4x4x8xf32>, %arg1 : index) -> vector<8xf32> {
+ %c0 = arith.constant 0 : index
+ %pad = ub.poison : f32
+ %0 = memref.collapse_shape %arg0 [[0, 1], [2]] : memref<4x4x8xf32> into memref<16x8xf32>
+ %1 = vector.transfer_read %0[%arg1, %c0], %pad {in_bounds = [true]} : memref<16x8xf32>, vector<8xf32>
+ return %1 : vector<8xf32>
+}
+
+// CHECK-LABEL: func @fold_collapse_shape_transfer_read
+// CHECK-SAME: %[[ARG0:[a-zA-Z0-9_]+]]: memref<4x4x8xf32>
+// CHECK-SAME: %[[ARG1:[a-zA-Z0-9_]+]]: index
+// CHECK: %[[C0:.*]] = arith.constant 0
+// CHECK: %[[PAD:.*]] = ub.poison : f32
+// CHECK: %[[IDXS:.*]]:2 = affine.delinearize_index %[[ARG1]] into (4, 4)
+// CHECK: vector.transfer_read %[[ARG0]][%[[IDXS]]#0, %[[IDXS]]#1, %[[C0]]], %[[PAD]] {in_bounds = [true]}
+
+// -----
+
func.func @fold_dma_start_subview_src(
%src : memref<128x64xf32>, %dst : memref<32xf32, 1>, %tag : memref<1xi32>,
%off0 : index, %off1 : index) {
>From 1e9db29931adc5afefd6fc6170dbabc450946ba6 Mon Sep 17 00:00:00 2001
From: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
Date: Mon, 4 May 2026 15:32:30 +0000
Subject: [PATCH 2/2] Review feedback
---
.../include/mlir/Interfaces/VectorInterfaces.td | 17 +++++++----------
1 file changed, 7 insertions(+), 10 deletions(-)
diff --git a/mlir/include/mlir/Interfaces/VectorInterfaces.td b/mlir/include/mlir/Interfaces/VectorInterfaces.td
index d79588e6d42d7..3a6dbf0d44fcb 100644
--- a/mlir/include/mlir/Interfaces/VectorInterfaces.td
+++ b/mlir/include/mlir/Interfaces/VectorInterfaces.td
@@ -136,7 +136,9 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
/*desc=*/[{
Return the indices that specify the starting offsets into the source
operand. The starting offsets are guaranteed to be in-bounds except
- when the transfer operation is masked.
+ when the transfer operation is masked or the `in_bounds` array doesn't
+ specify that that dimension is not necessarily in bounds. It is undefinied
+ behavior if these guarantees are violated.
}],
/*retTy=*/"::mlir::OperandRange",
/*methodName=*/"getIndices",
@@ -183,10 +185,6 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
`newPermutationMap`, where the replacement would adhere to the
conditinos specified by `updateStartingPosition`, passes any addihtional
op-specific constraints.
-
- (Note: this method has been added in case it ends up being needed
- and, if some upstream or downstream use requires additional argument,
- they should be added.)
}],
/*retTy=*/"::mlir::LogicalResult",
/*methodName=*/"mayUpdateStartingPosition",
@@ -197,6 +195,10 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
return ::mlir::success();
}]
>,
+ // Note: The initial motivating factcor for adding this interface method is to
+ // enable folding operations like `memref.subview` with transfer opes in a
+ // generic way. If future usecases require extending this method, it should be
+ // done.
InterfaceMethod<
/*desc=*/[{
Updates the base of this transfer operation to `newBase`, the indices into
@@ -221,11 +223,6 @@ def VectorTransferOpInterface : OpInterface<"VectorTransferOpInterface"> {
The update is performed in-place. Implementations of the interface
should use `RewriterBase::modifyOpInPlace` to ensure proper pattern
rewriter operation.
-
- Note: The initial motivating factcor for adding this interface method is to
- enable folding operations like `memref.subview` with transfer opes in a
- generic way. If future usecases require extending this method, it should be
- done.
}],
/*retTy=*/"void",
/*methodName=*/"updateStartingPosition",
More information about the Mlir-commits
mailing list