[Mlir-commits] [mlir] [mlir][vector] Pad the target shape when unrolling masks (PR #234448)
Hamza Qureshi
llvmlistbot at llvm.org
Sat Oct 31 04:19:57 PDT 2026
https://github.com/hamzaqureshi5 updated https://github.com/llvm/llvm-project/pull/234448
>From fb3caa5cc2911a15d29a09983b80795d7c232699 Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Thu, 22 Oct 2026 06:25:52 +0500
Subject: [PATCH 1/2] [mlir][vector] Pad the target shape when unrolling masks
`getTargetShape` only has to divide the trailing dimensions of the op, so it
may have a smaller rank. Both mask unroll patterns assumed a 1-1
correspondence with the op's dimensions and indexed the offsets by the op's
rank, so a rank 3 mask against the rank 2 target shape read past the end and
tripped `idx < size()`.
Pad the target shape with leading unit dimensions, as `UnrollElementwisePattern`
already does.
That lets `constant_mask` unroll, which exposes a second problem: a tile
starting past the mask in one dimension gets a zero there while the others
keep their extent. `constant_mask` is the conjunction of its dimensions, so
such a tile is empty and only the all zeros spelling is accepted. Normalize it.
This happens without any rank mismatch too, e.g.
`vector.constant_mask [8, 8] : vector<16x16xi1>`. The existing
`@vector_constant_mask` test uses `[12, 10]`, where both cutoffs exceed the
tile size, so no tile comes out partly zero.
Fixes #217175
---
.../Vector/Transforms/VectorUnroll.cpp | 39 ++++++++++++++----
.../Dialect/Vector/vector-unroll-options.mlir | 41 +++++++++++++++++++
2 files changed, 71 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 19abe44081eb4..14bcc27f1fe96 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -1091,16 +1091,23 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
*createMaskOp.getIterationShapeToUnroll();
Location loc = createMaskOp.getLoc();
+ // `getTargetShape` is only required to divide the trailing dimensions, so
+ // it may have a smaller rank than the op. The logic below assumes a 1-1
+ // correspondence with the op's mask operands, so pad with leading unit
+ // dims.
+ SmallVector<int64_t> adjustedTargetShape(originalSize.size(), 1);
+ llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
+
Value result = arith::ConstantOp::create(rewriter, loc, resultType,
rewriter.getZeroAttr(resultType));
VectorType targetVectorType =
- VectorType::get(*targetShape, rewriter.getI1Type());
- SmallVector<int64_t> strides(targetShape->size(), 1);
+ VectorType::get(adjustedTargetShape, rewriter.getI1Type());
+ SmallVector<int64_t> strides(adjustedTargetShape.size(), 1);
// In each dimension (d), each unrolled vector computes its mask size as:
// min(max(originalMaskOperands[d] - offset[d], 0), unrolledDimSize[d]).
for (SmallVector<int64_t> offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<Value> unrolledOperands;
for (auto [i, originalMaskOperand] :
@@ -1110,8 +1117,8 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
Value adjustedMaskSize = rewriter.createOrFold<arith::SubIOp>(
loc, originalMaskOperand, offsetVal);
Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
- Value unrolledDimSize =
- arith::ConstantIndexOp::create(rewriter, loc, (*targetShape)[i]);
+ Value unrolledDimSize = arith::ConstantIndexOp::create(
+ rewriter, loc, adjustedTargetShape[i]);
Value nonNegative =
rewriter.createOrFold<arith::MaxSIOp>(loc, adjustedMaskSize, zero);
Value unrolledOperand = rewriter.createOrFold<arith::MinSIOp>(
@@ -1184,16 +1191,23 @@ struct UnrollConstantMaskPattern
*constantMaskOp.getIterationShapeToUnroll();
Location loc = constantMaskOp.getLoc();
+ // `getTargetShape` is only required to divide the trailing dimensions, so
+ // it may have a smaller rank than the op. The logic below assumes a 1-1
+ // correspondence with the op's mask dimensions, so pad with leading unit
+ // dims.
+ SmallVector<int64_t> adjustedTargetShape(originalSize.size(), 1);
+ llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
+
Value result = arith::ConstantOp::create(rewriter, loc, resultType,
rewriter.getZeroAttr(resultType));
VectorType targetVectorType =
- VectorType::get(*targetShape, rewriter.getI1Type());
- SmallVector<int64_t> strides(targetShape->size(), 1);
+ VectorType::get(adjustedTargetShape, rewriter.getI1Type());
+ SmallVector<int64_t> strides(adjustedTargetShape.size(), 1);
// In each dimension (d), each unrolled vector computes its mask size as:
// min(max(originalMaskDim[d] - offset[d], 0), unrolledDimSize[d]).
for (const SmallVector<int64_t> &offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<int64_t> unrolledMaskDims;
for (auto [i, originalMaskDim] :
@@ -1203,10 +1217,17 @@ struct UnrollConstantMaskPattern
int64_t adjustedMaskSize =
std::max(originalMaskDim - offsets[i], static_cast<int64_t>(0));
int64_t unrolledMaskDim =
- std::min(adjustedMaskSize, static_cast<int64_t>((*targetShape)[i]));
+ std::min(adjustedMaskSize, adjustedTargetShape[i]);
unrolledMaskDims.push_back(unrolledMaskDim);
}
+ // A tile that starts past the mask in some dimension gets a zero there
+ // while the remaining dimensions keep their extent. `constant_mask` is
+ // the conjunction of its dimensions, so such a tile is empty, and the op
+ // accepts only the all zeros spelling for an empty mask.
+ if (llvm::is_contained(unrolledMaskDims, 0))
+ unrolledMaskDims.assign(unrolledMaskDims.size(), 0);
+
auto unrolledMask = rewriter.createOrFold<vector::ConstantMaskOp>(
loc, targetVectorType, unrolledMaskDims);
result = rewriter.createOrFold<vector::InsertStridedSliceOp>(
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 61a74744e7a72..4d029f84c1501 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -621,6 +621,47 @@ func.func @vector_constant_mask() -> vector<16x16xi1> {
// -----
+func.func @vector_constant_mask_rank_mismatch() -> vector<2x8x8xi1> {
+ %0 = vector.constant_mask [1, 8, 8] : vector<2x8x8xi1>
+ return %0 : vector<2x8x8xi1>
+}
+
+// CHECK-LABEL: func @vector_constant_mask_rank_mismatch
+// CHECK-SAME: () -> vector<2x8x8xi1>
+// CHECK: %[[EMPTY:.*]] = arith.constant dense<false> : vector<1x8x8xi1>
+// CHECK: %[[ACC:.*]] = arith.constant dense<false> : vector<2x8x8xi1>
+// CHECK: %[[FULL:.*]] = arith.constant dense<true> : vector<1x8x8xi1>
+// CHECK: %[[INS0:.*]] = vector.insert_strided_slice %[[FULL]], %[[ACC]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: %[[INS1:.*]] = vector.insert_strided_slice %[[EMPTY]], %[[INS0]] offsets = [1, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: return %[[INS1]] : vector<2x8x8xi1>
+
+// -----
+
+func.func @vector_create_mask_rank_mismatch(%size: index) -> vector<2x8x8xi1> {
+ %c8 = arith.constant 8 : index
+ %0 = vector.create_mask %size, %c8, %c8 : vector<2x8x8xi1>
+ return %0 : vector<2x8x8xi1>
+}
+
+// CHECK-LABEL: func @vector_create_mask_rank_mismatch
+// CHECK-SAME: (%[[SIZE:.*]]: index) -> vector<2x8x8xi1>
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[ACC:.*]] = arith.constant dense<false> : vector<2x8x8xi1>
+// CHECK: %[[C0:.*]] = arith.constant 0 : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[MAX0:.*]] = arith.maxsi %[[SIZE]], %[[C0]] : index
+// CHECK: %[[MIN0:.*]] = arith.minsi %[[MAX0]], %[[C1]] : index
+// CHECK: %[[MASK0:.*]] = vector.create_mask %[[MIN0]], %[[C8]], %[[C8]] : vector<1x8x8xi1>
+// CHECK: %[[INS0:.*]] = vector.insert_strided_slice %[[MASK0]], %[[ACC]] offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: %[[SUB1:.*]] = arith.subi %[[SIZE]], %[[C1]] : index
+// CHECK: %[[MAX1:.*]] = arith.maxsi %[[SUB1]], %[[C0]] : index
+// CHECK: %[[MIN1:.*]] = arith.minsi %[[MAX1]], %[[C1]] : index
+// CHECK: %[[MASK1:.*]] = vector.create_mask %[[MIN1]], %[[C8]], %[[C8]] : vector<1x8x8xi1>
+// CHECK: %[[INS1:.*]] = vector.insert_strided_slice %[[MASK1]], %[[INS0]] offsets = [1, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: return %[[INS1]] : vector<2x8x8xi1>
+
+// -----
+
func.func @shape_cast_1D(%v: vector<16xf32>) -> vector<2x2x4xf32> {
%0 = vector.shape_cast %v : vector<16xf32> to vector<2x2x4xf32>
return %0 : vector<2x2x4xf32>
>From 610c7580eb6a7750a7edcf2d8689396a033a195d Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Sat, 31 Oct 2026 15:06:17 +0500
Subject: [PATCH 2/2] [mlir][vector] Add a test without a rank mismatch
vector.constant_mask [8, 8] on vector<16x16xi1> puts the cutoff on a
tile boundary, so three of the four tiles come out empty. This reaches
the zero normalization in UnrollConstantMaskPattern on its own, with no
rank mismatch involved, which the existing tests only cover together.
---
.../Dialect/Vector/vector-unroll-options.mlir | 18 ++++++++++++++++++
1 file changed, 18 insertions(+)
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 4d029f84c1501..ac369ed77543f 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -621,6 +621,24 @@ func.func @vector_constant_mask() -> vector<16x16xi1> {
// -----
+func.func @vector_constant_mask_no_rank_mismatch() -> vector<16x16xi1> {
+ %0 = vector.constant_mask [8, 8] : vector<16x16xi1>
+ return %0 : vector<16x16xi1>
+}
+
+// CHECK-LABEL: func @vector_constant_mask_no_rank_mismatch
+// CHECK-SAME: () -> vector<16x16xi1>
+// CHECK: %[[EMPTY:.*]] = arith.constant dense<false> : vector<8x8xi1>
+// CHECK: %[[ACC:.*]] = arith.constant dense<false> : vector<16x16xi1>
+// CHECK: %[[FULL:.*]] = arith.constant dense<true> : vector<8x8xi1>
+// CHECK: %[[INS00:.*]] = vector.insert_strided_slice %[[FULL]], %[[ACC]] offsets = [0, 0], strides = [1, 1] : vector<8x8xi1> into vector<16x16xi1>
+// CHECK: %[[INS01:.*]] = vector.insert_strided_slice %[[EMPTY]], %[[INS00]] offsets = [0, 8], strides = [1, 1] : vector<8x8xi1> into vector<16x16xi1>
+// CHECK: %[[INS10:.*]] = vector.insert_strided_slice %[[EMPTY]], %[[INS01]] offsets = [8, 0], strides = [1, 1] : vector<8x8xi1> into vector<16x16xi1>
+// CHECK: %[[INS11:.*]] = vector.insert_strided_slice %[[EMPTY]], %[[INS10]] offsets = [8, 8], strides = [1, 1] : vector<8x8xi1> into vector<16x16xi1>
+// CHECK: return %[[INS11]] : vector<16x16xi1>
+
+// -----
+
func.func @vector_constant_mask_rank_mismatch() -> vector<2x8x8xi1> {
%0 = vector.constant_mask [1, 8, 8] : vector<2x8x8xi1>
return %0 : vector<2x8x8xi1>
More information about the Mlir-commits
mailing list