[Mlir-commits] [mlir] [mlir][vector] Fix mask rank extension for TransferWriteNonPermutationLowering (PR #229852)
Dhairyashil R G
llvmlistbot at llvm.org
Sun Nov 1 03:14:41 PST 2026
https://github.com/dhairyashilRG updated https://github.com/llvm/llvm-project/pull/229852
>From 8757c42f1ddde538c7f284a0b1d11dbf8cbf0381 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Wed, 7 Oct 2026 20:52:38 +0530
Subject: [PATCH 1/3] [mlir][vector] Fix mask rank extension for
TransferWriteNonPermutationLowering
When a transfer_write's permutation map skips a dim between the dims that are
written to, e.g. `(d0, d1, d2, d3) -> (d1, d3)`, the pattern adds a unit dim
for it. The vector got that unit dim in the right place, but the mask always
got it at the end, so the new write did not verify.
Put the mask's unit dims at the positions of the missing dims.
Based on #144893 by Kunwar Grover.
Co-authored-by: Kunwar Grover <groverkss at gmail.com>
Assisted-by: Claude
---
.../Vector/Transforms/LowerVectorTransfer.cpp | 51 +++++++++++--------
.../vector-transfer-permutation-lowering.mlir | 26 ++++++++++
2 files changed, 55 insertions(+), 22 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
index adaa3a1995d145..7162b0668eef24 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
@@ -32,8 +32,8 @@ inverseTransposeInBoundsAttr(OpBuilder &builder, ArrayAttr attr,
/// Extend the rank of a vector Value by `addedRanks` by adding outer unit
/// dimensions.
-static Value extendVectorRank(OpBuilder &builder, Location loc, Value vec,
- int64_t addedRank) {
+static TypedValue<VectorType> extendVectorRank(OpBuilder &builder, Location loc,
+ Value vec, int64_t addedRank) {
auto originalVecType = cast<VectorType>(vec.getType());
SmallVector<int64_t> newShape(addedRank, 1);
newShape.append(originalVecType.getShape().begin(),
@@ -47,19 +47,25 @@ static Value extendVectorRank(OpBuilder &builder, Location loc, Value vec,
return vector::BroadcastOp::create(builder, loc, newVecType, vec);
}
-/// Extend the rank of a vector Value by `addedRanks` by adding inner unit
-/// dimensions.
+/// Extend the rank of a mask Value by adding unit dimensions at the positions
+/// `missingInnerDims` of the result.
static Value extendMaskRank(OpBuilder &builder, Location loc, Value vec,
- int64_t addedRank) {
- Value broadcasted = extendVectorRank(builder, loc, vec, addedRank);
- SmallVector<int64_t> permutation;
- for (int64_t i = addedRank,
- e = cast<VectorType>(broadcasted.getType()).getRank();
- i < e; ++i)
- permutation.push_back(i);
- for (int64_t i = 0; i < addedRank; ++i)
- permutation.push_back(i);
- return vector::TransposeOp::create(builder, loc, broadcasted, permutation);
+ ArrayRef<int64_t> missingInnerDims) {
+ TypedValue<VectorType> broadcasted =
+ extendVectorRank(builder, loc, vec, missingInnerDims.size());
+ // The broadcast puts the unit dims first. `inversePerm[i]` is the position
+ // that dim `i` of the broadcast moves to.
+ SmallVector<bool> missing(broadcasted.getType().getRank(), false);
+ SmallVector<int64_t> inversePerm;
+ for (int64_t i : missingInnerDims) {
+ inversePerm.push_back(i);
+ missing[i] = true;
+ }
+ for (auto [i, isMissing] : llvm::enumerate(missing))
+ if (!isMissing)
+ inversePerm.push_back(i);
+ return vector::TransposeOp::create(builder, loc, broadcasted,
+ invertPermutationVector(inversePerm));
}
//===----------------------------------------------------------------------===//
@@ -279,28 +285,29 @@ struct TransferWriteNonPermutationLowering
for (AffineExpr exp : map.getResults())
foundDim[cast<AffineDimExpr>(exp).getPosition()] = true;
SmallVector<AffineExpr> exprs;
- bool foundFirstDim = false;
+ std::optional<size_t> firstDim;
SmallVector<int64_t> missingInnerDim;
for (size_t i = 0; i < foundDim.size(); i++) {
if (foundDim[i]) {
- foundFirstDim = true;
+ if (!firstDim)
+ firstDim = i;
continue;
}
- if (!foundFirstDim)
+ if (!firstDim)
continue;
// Once we found one outer dimension existing in the map keep track of all
- // the missing dimensions after that.
- missingInnerDim.push_back(i);
+ // the missing dimensions after that, relative to that dimension.
+ missingInnerDim.push_back(i - *firstDim);
exprs.push_back(rewriter.getAffineDimExpr(i));
}
// Vector: add unit dims at the beginning of the shape.
Value newVec = extendVectorRank(rewriter, op.getLoc(), op.getVector(),
missingInnerDim.size());
- // Mask: add unit dims at the end of the shape.
+ // Mask: add unit dims at the positions of the missing dimensions.
Value newMask;
if (op.getMask())
- newMask = extendMaskRank(rewriter, op.getLoc(), op.getMask(),
- missingInnerDim.size());
+ newMask =
+ extendMaskRank(rewriter, op.getLoc(), op.getMask(), missingInnerDim);
exprs.append(map.getResults().begin(), map.getResults().end());
AffineMap newMap =
AffineMap::get(map.getNumDims(), 0, exprs, op.getContext());
diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
index cbf4d22808a69a..81ea721ff01bca 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
@@ -162,6 +162,32 @@ func.func @xfer_write_non_minor_identity_with_mask_out_of_bounds(
return
}
+// The missing dim (d2) is between the dims that are written to, so the mask
+// gets its unit dim in the middle as well.
+
+// CHECK-LABEL: func.func @xfer_write_non_minor_identity_with_mask_missing_inner_dim(
+// CHECK-SAME: %[[MEM:.*]]: memref<?x?x?x?xf32>,
+// CHECK-SAME: %[[VEC:.*]]: vector<7x8xf32>,
+// CHECK-SAME: %[[MASK:.*]]: vector<7x8xi1>,
+// CHECK-SAME: %[[IDX:.*]]: index) {
+// CHECK: %[[BC_VEC:.*]] = vector.broadcast %[[VEC]] : vector<7x8xf32> to vector<1x7x8xf32>
+// CHECK: %[[BC_MASK:.*]] = vector.broadcast %[[MASK]] : vector<7x8xi1> to vector<1x7x8xi1>
+// CHECK: %[[TR_MASK:.*]] = vector.transpose %[[BC_MASK]], [1, 0, 2] : vector<1x7x8xi1> to vector<7x1x8xi1>
+// CHECK: %[[TR_VEC:.*]] = vector.transpose %[[BC_VEC]], [1, 0, 2] : vector<1x7x8xf32> to vector<7x1x8xf32>
+// CHECK: vector.transfer_write %[[TR_VEC]], %[[MEM]]{{\[}}%[[IDX]], %[[IDX]], %[[IDX]], %[[IDX]]], %[[TR_MASK]] {in_bounds = [false, true, false]} : vector<7x1x8xf32>, memref<?x?x?x?xf32>
+func.func @xfer_write_non_minor_identity_with_mask_missing_inner_dim(
+ %mem : memref<?x?x?x?xf32>,
+ %vec : vector<7x8xf32>,
+ %mask : vector<7x8xi1>,
+ %idx : index) {
+
+ vector.transfer_write %vec, %mem[%idx, %idx, %idx, %idx], %mask {
+ permutation_map = affine_map<(d0, d1, d2, d3) -> (d1, d3)>
+ } : vector<7x8xf32>, memref<?x?x?x?xf32>
+
+ return
+}
+
// CHECK-LABEL: func.func @xfer_write_non_minor_identity_with_mask_scalable(
// CHECK-SAME: %[[VEC:.*]]: vector<4x[8]xi16>,
// CHECK-SAME: %[[MEM:.*]]: memref<1x4x?x1xi16>,
>From a76a2ff1dca7bac4b3b7048ba9d0016fa92d9eb0 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Thu, 8 Oct 2026 21:35:36 +0530
Subject: [PATCH 2/3] [mlir][vector] Address review: test two missing dims,
clarify comment
Assisted-by: Claude
---
.../Vector/Transforms/LowerVectorTransfer.cpp | 3 ++-
.../vector-transfer-permutation-lowering.mlir | 25 +++++++++++++++++++
2 files changed, 27 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
index 7162b0668eef24..1c166e8dc79e18 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
@@ -48,7 +48,8 @@ static TypedValue<VectorType> extendVectorRank(OpBuilder &builder, Location loc,
}
/// Extend the rank of a mask Value by adding unit dimensions at the positions
-/// `missingInnerDims` of the result.
+/// `missingInnerDims` of the result. The positions are relative to the first
+/// dim that the transfer accesses, which is dim 0 of the mask.
static Value extendMaskRank(OpBuilder &builder, Location loc, Value vec,
ArrayRef<int64_t> missingInnerDims) {
TypedValue<VectorType> broadcasted =
diff --git a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
index 81ea721ff01bca..a809c06d6196d4 100644
--- a/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
+++ b/mlir/test/Dialect/Vector/vector-transfer-permutation-lowering.mlir
@@ -188,6 +188,31 @@ func.func @xfer_write_non_minor_identity_with_mask_missing_inner_dim(
return
}
+// Two missing dims (d1 and d3) that are not next to each other.
+
+// CHECK-LABEL: func.func @xfer_write_non_minor_identity_with_mask_two_missing_inner_dims(
+// CHECK-SAME: %[[MEM:.*]]: memref<?x?x?x?x?xf32>,
+// CHECK-SAME: %[[VEC:.*]]: vector<6x7x8xf32>,
+// CHECK-SAME: %[[MASK:.*]]: vector<6x7x8xi1>,
+// CHECK-SAME: %[[IDX:.*]]: index) {
+// CHECK: %[[BC_VEC:.*]] = vector.broadcast %[[VEC]] : vector<6x7x8xf32> to vector<1x1x6x7x8xf32>
+// CHECK: %[[BC_MASK:.*]] = vector.broadcast %[[MASK]] : vector<6x7x8xi1> to vector<1x1x6x7x8xi1>
+// CHECK: %[[TR_MASK:.*]] = vector.transpose %[[BC_MASK]], [2, 0, 3, 1, 4] : vector<1x1x6x7x8xi1> to vector<6x1x7x1x8xi1>
+// CHECK: %[[TR_VEC:.*]] = vector.transpose %[[BC_VEC]], [2, 0, 3, 1, 4] : vector<1x1x6x7x8xf32> to vector<6x1x7x1x8xf32>
+// CHECK: vector.transfer_write %[[TR_VEC]], %[[MEM]]{{\[}}%[[IDX]], %[[IDX]], %[[IDX]], %[[IDX]], %[[IDX]]], %[[TR_MASK]] {in_bounds = [false, true, false, true, false]} : vector<6x1x7x1x8xf32>, memref<?x?x?x?x?xf32>
+func.func @xfer_write_non_minor_identity_with_mask_two_missing_inner_dims(
+ %mem : memref<?x?x?x?x?xf32>,
+ %vec : vector<6x7x8xf32>,
+ %mask : vector<6x7x8xi1>,
+ %idx : index) {
+
+ vector.transfer_write %vec, %mem[%idx, %idx, %idx, %idx, %idx], %mask {
+ permutation_map = affine_map<(d0, d1, d2, d3, d4) -> (d0, d2, d4)>
+ } : vector<6x7x8xf32>, memref<?x?x?x?x?xf32>
+
+ return
+}
+
// CHECK-LABEL: func.func @xfer_write_non_minor_identity_with_mask_scalable(
// CHECK-SAME: %[[VEC:.*]]: vector<4x[8]xi16>,
// CHECK-SAME: %[[MEM:.*]]: memref<1x4x?x1xi16>,
>From 3515e7e982cd66b5767f3af7a8ee0bfc41637fa5 Mon Sep 17 00:00:00 2001
From: Dhairyashil R G <dhairyashil25 at gmail.com>
Date: Tue, 20 Oct 2026 22:12:02 +0530
Subject: [PATCH 3/3] [mlir][vector] Use one helper for the unit dims in
TransferWriteNonPermutationLowering
extendVectorRank added outer unit dims and extendMaskRank inner ones, so
the names did not say what differed. Replace both with insertUnitDims,
which takes the positions of the unit dims.
Assisted-by: Claude
---
.../Vector/Transforms/LowerVectorTransfer.cpp | 65 ++++++++-----------
1 file changed, 28 insertions(+), 37 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
index 1c166e8dc79e18..303f1cfb074b87 100644
--- a/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/LowerVectorTransfer.cpp
@@ -30,42 +30,31 @@ inverseTransposeInBoundsAttr(OpBuilder &builder, ArrayAttr attr,
return builder.getBoolArrayAttr(newInBoundsValues);
}
-/// Extend the rank of a vector Value by `addedRanks` by adding outer unit
-/// dimensions.
-static TypedValue<VectorType> extendVectorRank(OpBuilder &builder, Location loc,
- Value vec, int64_t addedRank) {
- auto originalVecType = cast<VectorType>(vec.getType());
- SmallVector<int64_t> newShape(addedRank, 1);
- newShape.append(originalVecType.getShape().begin(),
- originalVecType.getShape().end());
-
- SmallVector<bool> newScalableDims(addedRank, false);
- newScalableDims.append(originalVecType.getScalableDims().begin(),
- originalVecType.getScalableDims().end());
- VectorType newVecType = VectorType::get(
- newShape, originalVecType.getElementType(), newScalableDims);
- return vector::BroadcastOp::create(builder, loc, newVecType, vec);
-}
-
-/// Extend the rank of a mask Value by adding unit dimensions at the positions
-/// `missingInnerDims` of the result. The positions are relative to the first
-/// dim that the transfer accesses, which is dim 0 of the mask.
-static Value extendMaskRank(OpBuilder &builder, Location loc, Value vec,
- ArrayRef<int64_t> missingInnerDims) {
- TypedValue<VectorType> broadcasted =
- extendVectorRank(builder, loc, vec, missingInnerDims.size());
+/// Insert unit dims into `vec` so that they are at `unitDimPositions` in the
+/// result.
+static Value insertUnitDims(OpBuilder &builder, Location loc, Value vec,
+ ArrayRef<int64_t> unitDimPositions) {
+ auto vecType = cast<VectorType>(vec.getType());
+ int64_t numUnitDims = unitDimPositions.size();
+ SmallVector<int64_t> shape(numUnitDims, 1);
+ llvm::append_range(shape, vecType.getShape());
+ SmallVector<bool> scalableDims(numUnitDims, false);
+ llvm::append_range(scalableDims, vecType.getScalableDims());
// The broadcast puts the unit dims first. `inversePerm[i]` is the position
// that dim `i` of the broadcast moves to.
- SmallVector<bool> missing(broadcasted.getType().getRank(), false);
- SmallVector<int64_t> inversePerm;
- for (int64_t i : missingInnerDims) {
- inversePerm.push_back(i);
- missing[i] = true;
- }
- for (auto [i, isMissing] : llvm::enumerate(missing))
- if (!isMissing)
+ Value broadcast = vector::BroadcastOp::create(
+ builder, loc,
+ VectorType::get(shape, vecType.getElementType(), scalableDims), vec);
+ SmallVector<int64_t> inversePerm(unitDimPositions);
+ SmallVector<bool> isUnitDim(shape.size(), false);
+ for (int64_t pos : unitDimPositions)
+ isUnitDim[pos] = true;
+ for (auto [i, isUnit] : llvm::enumerate(isUnitDim))
+ if (!isUnit)
inversePerm.push_back(i);
- return vector::TransposeOp::create(builder, loc, broadcasted,
+ if (isIdentityPermutation(inversePerm))
+ return broadcast;
+ return vector::TransposeOp::create(builder, loc, broadcast,
invertPermutationVector(inversePerm));
}
@@ -302,13 +291,15 @@ struct TransferWriteNonPermutationLowering
exprs.push_back(rewriter.getAffineDimExpr(i));
}
// Vector: add unit dims at the beginning of the shape.
- Value newVec = extendVectorRank(rewriter, op.getLoc(), op.getVector(),
- missingInnerDim.size());
- // Mask: add unit dims at the positions of the missing dimensions.
+ Value newVec = insertUnitDims(
+ rewriter, op.getLoc(), op.getVector(),
+ llvm::to_vector(llvm::seq<int64_t>(missingInnerDim.size())));
+ // Mask: add unit dims at the positions of the missing dimensions. Dim 0 of
+ // the mask is the first written dim.
Value newMask;
if (op.getMask())
newMask =
- extendMaskRank(rewriter, op.getLoc(), op.getMask(), missingInnerDim);
+ insertUnitDims(rewriter, op.getLoc(), op.getMask(), missingInnerDim);
exprs.append(map.getResults().begin(), map.getResults().end());
AffineMap newMap =
AffineMap::get(map.getNumDims(), 0, exprs, op.getContext());
More information about the Mlir-commits
mailing list