[Mlir-commits] [mlir] [mlir][vector] Pad the target shape when unrolling masks (PR #234448)
Hamza Qureshi
llvmlistbot at llvm.org
Wed Oct 21 18:59:30 PDT 2026
https://github.com/hamzaqureshi5 created https://github.com/llvm/llvm-project/pull/234448
`getTargetShape` only needs to divide the trailing dimensions, so it can have
a smaller rank than the op. Both mask unroll patterns index the offsets by the
op's rank, so a rank 3 mask with a rank 2 target shape reads past the end and
asserts. Pad with leading unit dims, the same way `UnrollElementwisePattern`
does.
With that, `constant_mask` hits a second issue: a tile past the mask ends up
with a zero in one dimension and its full extent in the others, which the op
rejects since it only accepts all zeros for an empty mask. Normalize those
tiles. This one also happens without a rank mismatch, e.g.
`vector.constant_mask [8, 8] : vector<16x16xi1>`.
Fixes #217175
>From dbbd42eb07a22b6eb533e8806d9fb53502c8a702 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] [mlir][vector] Pad the target shape when unrolling masks
getTargetShape is only required to divide the trailing dimensions of the op
being unrolled, so it may have a smaller rank. UnrollCreateMaskPattern and
UnrollConstantMaskPattern assumed a 1-1 correspondence with the op's mask
operands and dimensions, indexing offsets by the op's rank while the offsets
carried the rank of the target shape. A rank 3 mask against the rank 2 target
shape used by the test pass therefore read past the end and tripped
`idx < size()`.
Pad the target shape with leading unit dimensions, as UnrollElementwisePattern
already does, so the leading dimensions are tiled one slice at a time.
Unrolling constant_mask then exposes a second problem. A tile that starts past
the mask in one dimension gets a zero there while the others keep their extent,
and vector.constant_mask is the conjunction of its dimensions, so such a tile
is empty and the op accepts only the all zeros spelling. Normalize it.
That is reachable without any rank mismatch: `vector.constant_mask [8, 8] :
vector<16x16xi1>` fails the same way. The existing @vector_constant_mask test
uses [12, 10], where both cutoffs exceed the tile size and so no tile can come
out partly zero, which is why it did not catch this.
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 19abe44081eb43..14bcc27f1fe96e 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 61a74744e7a72e..4d029f84c1501e 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>
More information about the Mlir-commits
mailing list