[Mlir-commits] [mlir] [mlir][Vector] Fix out-of-bounds crash unrolling create_mask/constant_mask with rank-mismatched native shape (PR #217514)
Aman Singh
llvmlistbot at llvm.org
Sat Sep 5 02:42:29 PDT 2026
https://github.com/amanyagami updated https://github.com/llvm/llvm-project/pull/217514
>From d2ce61c0bcd9a07ac134895c09eaf222844c42c3 Mon Sep 17 00:00:00 2001
From: amanyagami <2amansingh2 at gmail.com>
Date: Wed, 19 Aug 2026 19:45:53 -0700
Subject: [PATCH 1/3] [mlir][Vector] Pad the native shape when unrolling
create_mask/constant_mask
`getTargetShape` only guarantees that the native shape divides the trailing
dimensions of the op being unrolled, so it can have a smaller rank. The mask
unroll patterns indexed `targetShape` with the op's own rank, reading out of
bounds and asserting in SmallVector::operator[].
Pad the target shape with leading unit dimensions so it is rank-aligned with
the op, the way `UnrollElementwisePattern` already does, rather than refusing
to unroll. Tiling with the padded shape is equivalent: the leading dimensions
are unrolled one element at a time.
Also normalize empty tiles for `constant_mask`. A tile past the mask in one
dimension computes a zero there while other dimensions stay non-zero, which
`constant_mask` rejects ("expected all mask dim sizes to be zeros, as a result
of conjunction with zero mask dim"). This is reachable without any rank
mismatch -- `constant_mask [4, 10] : vector<16x16xi1>` unrolled to {8, 8}
builds `constant_mask [0, 8]` and fails to verify -- so it is covered by its
own test.
Fixes #217175.
Assisted-by: Claude Code
---
.../Vector/Transforms/VectorUnroll.cpp | 55 +++++++++++++---
.../Dialect/Vector/vector-unroll-options.mlir | 62 +++++++++++++++++++
2 files changed, 108 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index a95bfba0814bd..0b5991761681c 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -127,6 +127,26 @@ getTargetShape(const vector::UnrollVectorOptions &options, Operation *op) {
return targetShape;
}
+/// Returns `targetShape` left-padded with unit dimensions so that it has the
+/// same rank as `originalShape`.
+///
+/// The native/target shape returned by `getTargetShape` is only required to
+/// divide the *trailing* dimensions of the op being unrolled (see
+/// `computeShapeRatio`), so it may have a smaller rank than the op itself.
+/// Tiling with the padded shape is equivalent - the leading dimensions are
+/// simply unrolled one element at a time - and keeps the offsets, sizes and
+/// strides rank-aligned with the op's vector type.
+static SmallVector<int64_t>
+padTargetShapeToRank(ArrayRef<int64_t> targetShape,
+ ArrayRef<int64_t> originalShape) {
+ assert(targetShape.size() <= originalShape.size() &&
+ "expected the target shape to have at most the rank of the op");
+ SmallVector<int64_t> paddedShape(originalShape.size() - targetShape.size(),
+ 1);
+ llvm::append_range(paddedShape, targetShape);
+ return paddedShape;
+}
+
static SmallVector<int64_t>
getUnrollOrder(unsigned numLoops, Operation *op,
const vector::UnrollVectorOptions &options) {
@@ -1070,18 +1090,24 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
VectorType resultType = createMaskOp.getVectorType();
SmallVector<int64_t> originalSize = *createMaskOp.getShapeForUnroll();
+ // The logic below assumes a 1-1 correspondence between the unroll shape and
+ // the op's mask operands, so pad the target shape with leading unit
+ // dimensions when it has a smaller rank than the op (as is done in
+ // `UnrollElementwisePattern`).
+ SmallVector<int64_t> unrollShape =
+ padTargetShapeToRank(*targetShape, originalSize);
Location loc = createMaskOp.getLoc();
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(unrollShape, rewriter.getI1Type());
+ SmallVector<int64_t> strides(unrollShape.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, unrollShape)) {
SmallVector<Value> unrolledOperands;
for (auto [i, originalMaskOperand] :
@@ -1092,7 +1118,7 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
loc, originalMaskOperand, offsetVal);
Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
Value unrolledDimSize =
- arith::ConstantIndexOp::create(rewriter, loc, (*targetShape)[i]);
+ arith::ConstantIndexOp::create(rewriter, loc, unrollShape[i]);
Value nonNegative =
rewriter.createOrFold<arith::MaxSIOp>(loc, adjustedMaskSize, zero);
Value unrolledOperand = rewriter.createOrFold<arith::MinSIOp>(
@@ -1162,18 +1188,24 @@ struct UnrollConstantMaskPattern
VectorType resultType = constantMaskOp.getVectorType();
SmallVector<int64_t> originalSize = *constantMaskOp.getShapeForUnroll();
+ // The logic below assumes a 1-1 correspondence between the unroll shape and
+ // the op's mask dimensions, so pad the target shape with leading unit
+ // dimensions when it has a smaller rank than the op (as is done in
+ // `UnrollElementwisePattern`).
+ SmallVector<int64_t> unrollShape =
+ padTargetShapeToRank(*targetShape, originalSize);
Location loc = constantMaskOp.getLoc();
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(unrollShape, rewriter.getI1Type());
+ SmallVector<int64_t> strides(unrollShape.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, unrollShape)) {
SmallVector<int64_t> unrolledMaskDims;
for (auto [i, originalMaskDim] :
@@ -1182,11 +1214,16 @@ struct UnrollConstantMaskPattern
// for this particular slice
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]));
+ int64_t unrolledMaskDim = std::min(adjustedMaskSize, unrollShape[i]);
unrolledMaskDims.push_back(unrolledMaskDim);
}
+ // A tile that lies entirely outside the mask in one dimension is empty.
+ // `vector.constant_mask` requires such a mask to be all zeros, since the
+ // mask region is the conjunction of the per-dimension intervals.
+ 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 38bebdc24757c..315c4f85d150b 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -602,6 +602,68 @@ func.func @vector_constant_mask() -> vector<16x16xi1> {
// -----
+// A tile that falls entirely outside the mask in one dimension is empty, but
+// the other dimensions still compute a non-zero size, which `constant_mask`
+// rejects ("expected all mask dim sizes to be zeros, as a result of
+// conjunction with zero mask dim"). Such a tile must be emitted as an
+// all-zeros mask. This is independent of any rank mismatch - here the target
+// shape ({8, 8}) already matches the op's rank.
+
+func.func @vector_constant_mask_empty_tile() -> vector<16x16xi1> {
+ %0 = vector.constant_mask [4, 10] : vector<16x16xi1>
+ return %0 : vector<16x16xi1>
+}
+
+// The [8, 0] and [8, 8] tiles lie past the mask in dim 0. Before this was
+// handled they were emitted as `constant_mask [0, 8]` / `[0, 2]`, which fails
+// the op's verifier outright, so this whole function used to fail to unroll.
+
+// CHECK-LABEL: func @vector_constant_mask_empty_tile
+// CHECK-DAG: arith.constant dense<false> : vector<16x16xi1>
+// CHECK-DAG: arith.constant dense<false> : vector<8x8xi1>
+// CHECK: vector.constant_mask [4, 8] : vector<8x8xi1>
+// CHECK: vector.constant_mask [4, 2] : vector<8x8xi1>
+// CHECK-NOT: vector.constant_mask [0,
+
+// -----
+
+// The native/target unroll shape for vector.create_mask/vector.constant_mask
+// used by this test pass is fixed at rank 2 ({8, 8}), independent of the rank
+// of the op being unrolled. When the op's own rank is higher (rank 3 below),
+// the target shape is padded with leading unit dimensions so that it is
+// rank-aligned with the op, as `UnrollElementwisePattern` does. See #217175.
+
+func.func @vector_create_mask_rank_mismatch(%size1: index) -> vector<2x8x8xi1> {
+ %c8 = arith.constant 8 : index
+ %0 = vector.create_mask %c8, %size1, %c8 : vector<2x8x8xi1>
+ return %0 : vector<2x8x8xi1>
+}
+
+// CHECK-LABEL: func @vector_create_mask_rank_mismatch
+// CHECK: vector.create_mask {{.*}} : vector<1x8x8xi1>
+// CHECK: vector.insert_strided_slice {{.*}} offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: vector.create_mask {{.*}} : vector<1x8x8xi1>
+// CHECK: vector.insert_strided_slice {{.*}} offsets = [1, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+
+// -----
+
+func.func @vector_constant_mask_rank_mismatch() -> vector<2x8x8xi1> {
+ %0 = vector.constant_mask [1, 4, 8] : vector<2x8x8xi1>
+ return %0 : vector<2x8x8xi1>
+}
+
+// The [1, 0, 0] tile is past the mask in dim 0, so it is an all-zeros mask.
+
+// CHECK-LABEL: func @vector_constant_mask_rank_mismatch
+// CHECK-DAG: arith.constant dense<false> : vector<2x8x8xi1>
+// CHECK-DAG: arith.constant dense<false> : vector<1x8x8xi1>
+// CHECK: vector.constant_mask [1, 4, 8] : vector<1x8x8xi1>
+// CHECK: vector.insert_strided_slice {{.*}} offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK: vector.insert_strided_slice {{.*}} offsets = [1, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// CHECK-NOT: vector.constant_mask [0,
+
+// -----
+
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 c521ec4b666a3d40b8777d975a67952b586556b8 Mon Sep 17 00:00:00 2001
From: amanyagami <aman at noahlabs.ai>
Date: Sat, 5 Sep 2026 02:04:54 -0700
Subject: [PATCH 2/3] Address review: tighten comments, mark negative checks
Take the three suggested comment rewrites in VectorUnroll.cpp, cut the prose in
vector-unroll-options.mlir down to one line per test, and label the CHECK-NOT
assertions as negative checks.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
.../Vector/Transforms/VectorUnroll.cpp | 17 +++++---------
.../Dialect/Vector/vector-unroll-options.mlir | 23 +++++--------------
2 files changed, 12 insertions(+), 28 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 0b5991761681c..e5a903509829b 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -131,11 +131,10 @@ getTargetShape(const vector::UnrollVectorOptions &options, Operation *op) {
/// same rank as `originalShape`.
///
/// The native/target shape returned by `getTargetShape` is only required to
-/// divide the *trailing* dimensions of the op being unrolled (see
-/// `computeShapeRatio`), so it may have a smaller rank than the op itself.
-/// Tiling with the padded shape is equivalent - the leading dimensions are
-/// simply unrolled one element at a time - and keeps the offsets, sizes and
-/// strides rank-aligned with the op's vector type.
+/// divide the *trailing* dimensions of the op being unrolled, so it may have a
+/// smaller rank than the op itself.
+/// Tiling with the padded shape is equivalent and keeps the offsets, sizes
+/// and strides rank-aligned with the op's vector type.
static SmallVector<int64_t>
padTargetShapeToRank(ArrayRef<int64_t> targetShape,
ArrayRef<int64_t> originalShape) {
@@ -1091,9 +1090,7 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
VectorType resultType = createMaskOp.getVectorType();
SmallVector<int64_t> originalSize = *createMaskOp.getShapeForUnroll();
// The logic below assumes a 1-1 correspondence between the unroll shape and
- // the op's mask operands, so pad the target shape with leading unit
- // dimensions when it has a smaller rank than the op (as is done in
- // `UnrollElementwisePattern`).
+ // the op's mask operands, so pad the target shape with leading unit dims.
SmallVector<int64_t> unrollShape =
padTargetShapeToRank(*targetShape, originalSize);
Location loc = createMaskOp.getLoc();
@@ -1189,9 +1186,7 @@ struct UnrollConstantMaskPattern
VectorType resultType = constantMaskOp.getVectorType();
SmallVector<int64_t> originalSize = *constantMaskOp.getShapeForUnroll();
// The logic below assumes a 1-1 correspondence between the unroll shape and
- // the op's mask dimensions, so pad the target shape with leading unit
- // dimensions when it has a smaller rank than the op (as is done in
- // `UnrollElementwisePattern`).
+ // the op's mask dimensions, so pad the target shape with leading unit dims.
SmallVector<int64_t> unrollShape =
padTargetShapeToRank(*targetShape, originalSize);
Location loc = constantMaskOp.getLoc();
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 315c4f85d150b..9adb854ec4d6c 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -602,36 +602,26 @@ func.func @vector_constant_mask() -> vector<16x16xi1> {
// -----
-// A tile that falls entirely outside the mask in one dimension is empty, but
-// the other dimensions still compute a non-zero size, which `constant_mask`
-// rejects ("expected all mask dim sizes to be zeros, as a result of
-// conjunction with zero mask dim"). Such a tile must be emitted as an
-// all-zeros mask. This is independent of any rank mismatch - here the target
-// shape ({8, 8}) already matches the op's rank.
+// A tile lying entirely outside the mask in one dim must be emitted as an
+// all-zeros mask: `constant_mask` rejects a mix of zero and non-zero sizes.
func.func @vector_constant_mask_empty_tile() -> vector<16x16xi1> {
%0 = vector.constant_mask [4, 10] : vector<16x16xi1>
return %0 : vector<16x16xi1>
}
-// The [8, 0] and [8, 8] tiles lie past the mask in dim 0. Before this was
-// handled they were emitted as `constant_mask [0, 8]` / `[0, 2]`, which fails
-// the op's verifier outright, so this whole function used to fail to unroll.
-
// CHECK-LABEL: func @vector_constant_mask_empty_tile
// CHECK-DAG: arith.constant dense<false> : vector<16x16xi1>
// CHECK-DAG: arith.constant dense<false> : vector<8x8xi1>
// CHECK: vector.constant_mask [4, 8] : vector<8x8xi1>
// CHECK: vector.constant_mask [4, 2] : vector<8x8xi1>
+// Negative check: no partially-zero tile is emitted.
// CHECK-NOT: vector.constant_mask [0,
// -----
-// The native/target unroll shape for vector.create_mask/vector.constant_mask
-// used by this test pass is fixed at rank 2 ({8, 8}), independent of the rank
-// of the op being unrolled. When the op's own rank is higher (rank 3 below),
-// the target shape is padded with leading unit dimensions so that it is
-// rank-aligned with the op, as `UnrollElementwisePattern` does. See #217175.
+// The target shape used by this test pass is fixed at rank 2, so a rank-3 op
+// exercises the leading-unit-dim padding. See #217175.
func.func @vector_create_mask_rank_mismatch(%size1: index) -> vector<2x8x8xi1> {
%c8 = arith.constant 8 : index
@@ -652,14 +642,13 @@ func.func @vector_constant_mask_rank_mismatch() -> vector<2x8x8xi1> {
return %0 : vector<2x8x8xi1>
}
-// The [1, 0, 0] tile is past the mask in dim 0, so it is an all-zeros mask.
-
// CHECK-LABEL: func @vector_constant_mask_rank_mismatch
// CHECK-DAG: arith.constant dense<false> : vector<2x8x8xi1>
// CHECK-DAG: arith.constant dense<false> : vector<1x8x8xi1>
// CHECK: vector.constant_mask [1, 4, 8] : vector<1x8x8xi1>
// CHECK: vector.insert_strided_slice {{.*}} offsets = [0, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
// CHECK: vector.insert_strided_slice {{.*}} offsets = [1, 0, 0], strides = [1, 1, 1] : vector<1x8x8xi1> into vector<2x8x8xi1>
+// Negative check: no partially-zero tile is emitted.
// CHECK-NOT: vector.constant_mask [0,
// -----
>From 5fa18685d3edf621a1a9083d8480425e85fca783 Mon Sep 17 00:00:00 2001
From: amanyagami <2amansingh2 at gmail.com>
Date: Sat, 5 Sep 2026 02:42:08 -0700
Subject: [PATCH 3/3] [mlir][Vector] Drop remaining explanatory comments from
the tests
Address review: remove the block comments above the new tests. The
negative-check markers requested in review are kept.
Co-Authored-By: Claude Opus 5 <noreply at anthropic.com>
---
mlir/test/Dialect/Vector/vector-unroll-options.mlir | 8 --------
1 file changed, 8 deletions(-)
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index 9adb854ec4d6c..bd0bbdc59aa8b 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -602,9 +602,6 @@ func.func @vector_constant_mask() -> vector<16x16xi1> {
// -----
-// A tile lying entirely outside the mask in one dim must be emitted as an
-// all-zeros mask: `constant_mask` rejects a mix of zero and non-zero sizes.
-
func.func @vector_constant_mask_empty_tile() -> vector<16x16xi1> {
%0 = vector.constant_mask [4, 10] : vector<16x16xi1>
return %0 : vector<16x16xi1>
@@ -620,9 +617,6 @@ func.func @vector_constant_mask_empty_tile() -> vector<16x16xi1> {
// -----
-// The target shape used by this test pass is fixed at rank 2, so a rank-3 op
-// exercises the leading-unit-dim padding. See #217175.
-
func.func @vector_create_mask_rank_mismatch(%size1: index) -> vector<2x8x8xi1> {
%c8 = arith.constant 8 : index
%0 = vector.create_mask %c8, %size1, %c8 : vector<2x8x8xi1>
@@ -918,8 +912,6 @@ func.func @deinterleave_2d(%v: vector<4x8xi32>) -> (vector<4x4xi32>, vector<4x4x
// -----
-// TargetShape [2, 2] has a lower rank than the reduced source <2x2x4>, so it
-// applies to the trailing dimensions and is padded to [1, 2, 2].
func.func @vector_multi_reduction_rank_mismatch(%v : vector<2x2x4xf32>, %acc: vector<2x2xf32>) -> vector<2x2xf32> {
%0 = vector.multi_reduction #vector.kind<add>, %v, %acc [2] : vector<2x2x4xf32> to vector<2x2xf32>
return %0 : vector<2x2xf32>
More information about the Mlir-commits
mailing list