[Mlir-commits] [mlir] [mlir][vector] Update `CastAwayTransfer{Read|Write}LeadingOneDim` (PR #219499)
Andrzej Warzyński
llvmlistbot at llvm.org
Fri Aug 28 08:13:04 PDT 2026
https://github.com/banach-space created https://github.com/llvm/llvm-project/pull/219499
Updates `CastAwayTransfer{Read|Write}LeadingOneDim` to use
`vector.shape_cast`, rather than `vector.extract`, as the canonical form
for stripping unit dimensions.
This change was originally implemented by @krzysz00 in #196206, but was
subsequently reverted in #199546. This PR intentionally restores only a
subset of #196206, making it easier to identify and triage any potential
regressions.
Co-authored-by: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
>From e89b7fb721a01208d6a6805c5f4d3cefd41a2afe Mon Sep 17 00:00:00 2001
From: Andrzej Warzynski <andrzej.warzynski at arm.com>
Date: Fri, 28 Aug 2026 15:11:14 +0000
Subject: [PATCH] [mlir][vector] Update
`CastAwayTransfer{Read|Write}LeadingOneDim`
Updates `CastAwayTransfer{Read|Write}LeadingOneDim` to use
`vector.shape_cast`, rather than `vector.extract`, as the canonical form
for stripping unit dimensions.
This change was originally implemented by @krzysz00 in #196206, but was
subsequently reverted in #199546. This PR intentionally restores only a
subset of #196206, making it easier to identify and triage any potential
regressions.
Co-authored-by: Krzysztof Drewniak <Krzysztof.Drewniak at amd.com>
---
.../Transforms/VectorDropLeadUnitDim.cpp | 62 ++++++++-----------
.../vector-dropleadunitdim-transforms.mlir | 12 ++--
2 files changed, 31 insertions(+), 43 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
index a5e1c838de60a..3a10305446f16 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDropLeadUnitDim.cpp
@@ -24,7 +24,7 @@ using namespace mlir::vector;
// Trims leading one dimensions from `oldType` and returns the result type.
// Returns `vector<1xT>` if `oldType` only has one element.
-static VectorType trimLeadingOneDims(VectorType oldType) {
+static VectorType trimLeadingUnitDims(VectorType oldType) {
ArrayRef<int64_t> oldShape = oldType.getShape();
ArrayRef<int64_t> newShape = oldShape;
@@ -63,7 +63,7 @@ struct CastAwayExtractStridedSliceLeadingOneDim
// the same rank. Here we drop leading one dimensions from the input vector
// type to make sure we don't cause mismatch.
VectorType oldSrcType = extractOp.getSourceVectorType();
- VectorType newSrcType = trimLeadingOneDims(oldSrcType);
+ VectorType newSrcType = trimLeadingUnitDims(oldSrcType);
if (newSrcType.getRank() == oldSrcType.getRank())
return failure();
@@ -110,9 +110,9 @@ struct CastAwayInsertStridedSliceLeadingOneDim
LogicalResult matchAndRewrite(vector::InsertStridedSliceOp insertOp,
PatternRewriter &rewriter) const override {
VectorType oldSrcType = insertOp.getSourceVectorType();
- VectorType newSrcType = trimLeadingOneDims(oldSrcType);
+ VectorType newSrcType = trimLeadingUnitDims(oldSrcType);
VectorType oldDstType = insertOp.getDestVectorType();
- VectorType newDstType = trimLeadingOneDims(oldDstType);
+ VectorType newDstType = trimLeadingUnitDims(oldDstType);
int64_t srcDropCount = oldSrcType.getRank() - newSrcType.getRank();
int64_t dstDropCount = oldDstType.getRank() - newDstType.getRank();
@@ -154,13 +154,13 @@ struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
Type newSrcType = oldSrcType;
int64_t oldSrcRank = 0, newSrcRank = 0;
if (auto type = dyn_cast<VectorType>(oldSrcType)) {
- newSrcType = trimLeadingOneDims(type);
+ newSrcType = trimLeadingUnitDims(type);
oldSrcRank = type.getRank();
newSrcRank = cast<VectorType>(newSrcType).getRank();
}
VectorType oldDstType = insertOp.getDestVectorType();
- VectorType newDstType = trimLeadingOneDims(oldDstType);
+ VectorType newDstType = trimLeadingUnitDims(oldDstType);
int64_t srcDropCount = oldSrcRank - newSrcRank;
int64_t dstDropCount = oldDstType.getRank() - newDstType.getRank();
@@ -201,19 +201,9 @@ struct CastAwayInsertLeadingOneDim : public OpRewritePattern<vector::InsertOp> {
};
static Value dropUnitDimsFromMask(OpBuilder &b, Location loc, Value mask,
- VectorType newType, AffineMap newMap,
- VectorType oldMaskType) {
- // Infer the type of the new mask from the new map.
+ VectorType newType, AffineMap newMap) {
VectorType newMaskType = inferTransferOpMaskType(newType, newMap);
- // If the new mask is broadcastable to the old result type, we can safely
- // use a `vector.extract` to get the new mask. Otherwise the best we can
- // do is shape cast.
- if (vector::isBroadcastableTo(newMaskType, oldMaskType) ==
- BroadcastableToResult::Success) {
- int64_t dropDim = oldMaskType.getRank() - newMaskType.getRank();
- return vector::ExtractOp::create(b, loc, mask, splatZero(dropDim));
- }
return vector::ShapeCastOp::create(b, loc, newMaskType, mask);
}
@@ -229,16 +219,17 @@ struct CastAwayTransferReadLeadingOneDim
// TODO(#78787): Not supported masked op yet.
if (cast<MaskableOpInterface>(read.getOperation()).isMasked())
return failure();
- // TODO: support 0-d corner case.
+
if (read.getTransferRank() == 0)
- return failure();
+ return rewriter.notifyMatchFailure(
+ read, "Nothing to trim - the transfer itself has rank zero");
auto shapedType = cast<ShapedType>(read.getBase().getType());
if (shapedType.getElementType() != read.getVectorType().getElementType())
return failure();
VectorType oldType = read.getVectorType();
- VectorType newType = trimLeadingOneDims(oldType);
+ VectorType newType = trimLeadingUnitDims(oldType);
if (newType == oldType)
return failure();
@@ -256,11 +247,9 @@ struct CastAwayTransferReadLeadingOneDim
read.getInBoundsAttr().getValue().take_back(newType.getRank()));
Value mask = Value();
- if (read.getMask()) {
- VectorType maskType = read.getMaskType();
+ if (read.getMask())
mask = dropUnitDimsFromMask(rewriter, read.getLoc(), read.getMask(),
- newType, newMap, maskType);
- }
+ newType, newMap);
auto newRead = vector::TransferReadOp::create(
rewriter, read.getLoc(), newType, read.getBase(), read.getIndices(),
@@ -283,19 +272,19 @@ struct CastAwayTransferWriteLeadingOneDim
// TODO(#78787): Not supported masked op yet.
if (cast<MaskableOpInterface>(write.getOperation()).isMasked())
return failure();
- // TODO: support 0-d corner case.
+
if (write.getTransferRank() == 0)
- return failure();
+ return rewriter.notifyMatchFailure(
+ write, "Nothing to trim - the transfer itself has rank zero");
auto shapedType = dyn_cast<ShapedType>(write.getBase().getType());
if (shapedType.getElementType() != write.getVectorType().getElementType())
return failure();
VectorType oldType = write.getVectorType();
- VectorType newType = trimLeadingOneDims(oldType);
+ VectorType newType = trimLeadingUnitDims(oldType);
if (newType == oldType)
return failure();
- int64_t dropDim = oldType.getRank() - newType.getRank();
AffineMap oldMap = write.getPermutationMap();
ArrayRef<AffineExpr> newResults =
@@ -309,13 +298,12 @@ struct CastAwayTransferWriteLeadingOneDim
inBoundsAttr = rewriter.getArrayAttr(
write.getInBoundsAttr().getValue().take_back(newType.getRank()));
- auto newVector = vector::ExtractOp::create(
- rewriter, write.getLoc(), write.getVector(), splatZero(dropDim));
+ auto newVector = rewriter.createOrFold<vector::ShapeCastOp>(
+ write.getLoc(), newType, write.getVector());
if (write.getMask()) {
- VectorType maskType = write.getMaskType();
- Value newMask = dropUnitDimsFromMask(
- rewriter, write.getLoc(), write.getMask(), newType, newMap, maskType);
+ Value newMask = dropUnitDimsFromMask(rewriter, write.getLoc(),
+ write.getMask(), newType, newMap);
rewriter.replaceOpWithNewOp<vector::TransferWriteOp>(
write, newVector, write.getBase(), write.getIndices(),
AffineMapAttr::get(newMap), newMask, inBoundsAttr);
@@ -516,7 +504,7 @@ class CastAwayElementwiseLeadingOneDim : public RewritePattern {
auto vecType = dyn_cast<VectorType>(op->getResultTypes()[0]);
if (!vecType)
return failure();
- VectorType newVecType = trimLeadingOneDims(vecType);
+ VectorType newVecType = trimLeadingUnitDims(vecType);
if (newVecType == vecType)
return failure();
int64_t dropDim = vecType.getRank() - newVecType.getRank();
@@ -572,7 +560,7 @@ struct CastAwayLoadLikeLeadingOneDim : public OpRewritePattern<OpTy> {
LogicalResult matchAndRewrite(OpTy op,
PatternRewriter &rewriter) const override {
VectorType oldResultType = op.getVectorType();
- VectorType newResultType = trimLeadingOneDims(oldResultType);
+ VectorType newResultType = trimLeadingUnitDims(oldResultType);
if (newResultType == oldResultType)
return failure();
int64_t nDropped = oldResultType.getRank() - newResultType.getRank();
@@ -608,7 +596,7 @@ struct CastAwayStoreLikeLeadingOneDim : public OpRewritePattern<OpTy> {
LogicalResult matchAndRewrite(OpTy op,
PatternRewriter &rewriter) const override {
VectorType oldVecType = op.getVectorType();
- VectorType newVecType = trimLeadingOneDims(oldVecType);
+ VectorType newVecType = trimLeadingUnitDims(oldVecType);
if (newVecType == oldVecType)
return failure();
int64_t nDropped = oldVecType.getRank() - newVecType.getRank();
@@ -642,7 +630,7 @@ struct CastAwayConstantMaskLeadingOneDim
LogicalResult matchAndRewrite(vector::ConstantMaskOp mask,
PatternRewriter &rewriter) const override {
VectorType oldType = mask.getType();
- VectorType newType = trimLeadingOneDims(oldType);
+ VectorType newType = trimLeadingUnitDims(oldType);
if (newType == oldType)
return failure();
diff --git a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
index 749c400d21d12..6cf05b71dc570 100644
--- a/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
+++ b/mlir/test/Dialect/Vector/vector-dropleadunitdim-transforms.mlir
@@ -350,7 +350,7 @@ func.func @cast_away_masked_transfer_read_leading_one_dims(%arg0: memref<1x4x8x1
%c0 = arith.constant 0 : index
// CHECK: %[[F0:.+]] = arith.constant 0.000000e+00 : f16
%f0 = arith.constant 0. : f16
- // CHECK: %[[MASK_CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
+ // CHECK: %[[MASK_CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
// CHECK: %[[READ:.+]] = vector.transfer_read %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]], %[[F0]], %[[MASK_CAST]] {in_bounds = [true]} : memref<1x4x8x16xf16>, vector<4xf16>
// CHECK: %[[CAST:.+]] = vector.broadcast %[[READ]] : vector<4xf16> to vector<1x4xf16>
%0 = vector.transfer_read %arg0[%c0, %c0, %c0, %c0], %f0, %arg1 {in_bounds = [true, true]} : memref<1x4x8x16xf16>, vector<1x4xf16>
@@ -410,7 +410,7 @@ func.func @not_insert_cast_fo4_transfer_read_under_mask(%arg0: memref<1x1x4xf16>
func.func @cast_away_transfer_write_leading_one_dims(%arg0: memref<1x4x8x16xf16>, %arg1: vector<1x4xf16>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- // CHECK: %[[CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xf16> from vector<1x4xf16>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf16> to vector<4xf16>
// CHECK: vector.transfer_write %[[CAST]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]] {in_bounds = [true]} : vector<4xf16>, memref<1x4x8x16xf16>
vector.transfer_write %arg1, %arg0[%c0, %c0, %c0, %c0] {in_bounds = [true, true]} : vector<1x4xf16>, memref<1x4x8x16xf16>
@@ -421,8 +421,8 @@ func.func @cast_away_transfer_write_leading_one_dims(%arg0: memref<1x4x8x16xf16>
func.func @cast_away_masked_transfer_write_leading_one_dims(%arg0: memref<1x4x8x16xf16>, %arg1: vector<1x4xf16>, %arg2: vector<1x4xi1>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- // CHECK: %[[CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xf16> from vector<1x4xf16>
- // CHECK: %[[MASK_CAST:.+]] = vector.extract %{{.*}}[0] : vector<4xi1> from vector<1x4xi1>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xf16> to vector<4xf16>
+ // CHECK: %[[MASK_CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4xi1> to vector<4xi1>
// CHECK: vector.transfer_write %[[CAST]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]], %[[C0]]], %[[MASK_CAST]] {in_bounds = [true]} : vector<4xf16>, memref<1x4x8x16xf16>
vector.transfer_write %arg1, %arg0[%c0, %c0, %c0, %c0], %arg2 {in_bounds = [true, true]} : vector<1x4xf16>, memref<1x4x8x16xf16>
@@ -432,7 +432,7 @@ func.func @cast_away_masked_transfer_write_leading_one_dims(%arg0: memref<1x4x8x
// CHECK-LABEL: func @cast_away_transfer_write_leading_one_dims_one_element
func.func @cast_away_transfer_write_leading_one_dims_one_element(%arg0: memref<1x1x1x1xf16>, %arg1: vector<1x1xf16>) {
%c0 = arith.constant 0 : index
- // CHECK: vector.extract %{{.+}}[0] : vector<1xf16> from vector<1x1xf16>
+ // CHECK: vector.shape_cast %{{.+}} : vector<1x1xf16> to vector<1xf16>
vector.transfer_write %arg1, %arg0[%c0, %c0, %c0, %c0] {in_bounds = [true, true]} : vector<1x1xf16>, memref<1x1x1x1xf16>
return
}
@@ -461,7 +461,7 @@ func.func @not_insert_cast_for_transfer_write_under_mask(%arg0: memref<1x1x4xf16
func.func @cast_away_nontrivial_map_masked_transfer_write(%arg0: memref<1x4x8xf16>, %arg1: vector<1x1x4xf16>, %arg2: vector<1x4x1xi1>) {
// CHECK: %[[C0:.+]] = arith.constant 0 : index
%c0 = arith.constant 0 : index
- // CHECK: %[[CAST:.+]] = vector.extract %{{.*}}[0, 0] : vector<4xf16> from vector<1x1x4xf16>
+ // CHECK: %[[CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x1x4xf16> to vector<4xf16>
// CHECK: %[[MASK_CAST:.+]] = vector.shape_cast %{{.*}} : vector<1x4x1xi1> to vector<4xi1>
// CHECK: vector.transfer_write %[[CAST]], %{{.*}}[%[[C0]], %[[C0]], %[[C0]]], %[[MASK_CAST]] {in_bounds = [true]
// CHECK-SAME: permutation_map = #[[$MAP]]} : vector<4xf16>, memref<1x4x8xf16>
More information about the Mlir-commits
mailing list