[Mlir-commits] [mlir] [mlir][vector] Fix mask rank extension for TransferWriteNonPermutationLowering (PR #229852)
Dhairyashil R G
llvmlistbot at llvm.org
Thu Oct 8 09:57:00 PDT 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/2] [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 adaa3a1995d14..7162b0668eef2 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 cbf4d22808a69..81ea721ff01bc 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/2] [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 7162b0668eef2..1c166e8dc79e1 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 81ea721ff01bc..a809c06d6196d 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>,
More information about the Mlir-commits
mailing list