[Mlir-commits] [mlir] [mlir][Vector] Fix out-of-bounds crash unrolling create_mask/constant_mask with rank-mismatched native shape (PR #217514)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir-vector

Author: Aman Singh (amanyagami)

<details>
<summary>Changes</summary>

`UnrollCreateMaskPattern` and `UnrollConstantMaskPattern` assume the
native/target unroll shape returned by `getTargetShape()` always has
the same rank as the `vector.create_mask`/`vector.constant_mask` op
being unrolled. That is not guaranteed: `computeShapeRatio` permits a
shorter target shape, whose entries are matched against the trailing
dimensions of the op's shape. When that happens, both patterns index
the (shorter) offsets range with the op's full operand/dimension
count, going out of bounds and hitting the `SmallVector` assertion
`idx < size()`.

Fix by bailing out with a match failure when the target shape's rank
does not match the op's rank, rather than indexing out of bounds.

Verified: reverting this fix reproduces the reported crash; with the
fix, `mlir-opt -test-vector-unrolling-patterns` on the reported
reproducer succeeds, and
`mlir/test/Dialect/Vector/vector-unroll-options.mlir` passes.

Fixes #<!-- -->217175

🤖 Generated with [Claude Code](https://claude.com/claude-code)

---
Full diff: https://github.com/llvm/llvm-project/pull/217514.diff


2 Files Affected:

- (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+14) 
- (modified) mlir/test/Dialect/Vector/vector-unroll-options.mlir (+28) 


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 62869111496d1..9febf5e562985 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -1064,6 +1064,13 @@ struct UnrollCreateMaskPattern : public OpRewritePattern<vector::CreateMaskOp> {
 
     VectorType resultType = createMaskOp.getVectorType();
     SmallVector<int64_t> originalSize = *createMaskOp.getShapeForUnroll();
+    // The native/target shape coming from `getTargetShape` may have fewer
+    // dimensions than the op being unrolled (it is only required to be a
+    // divisor of the trailing dimensions, see `computeShapeRatio`). The
+    // logic below assumes a 1-1 correspondence between `targetShape` and the
+    // op's mask operands/dimensions, so bail out if that does not hold.
+    if (targetShape->size() != originalSize.size())
+      return failure();
     Location loc = createMaskOp.getLoc();
 
     Value result = arith::ConstantOp::create(rewriter, loc, resultType,
@@ -1156,6 +1163,13 @@ struct UnrollConstantMaskPattern
 
     VectorType resultType = constantMaskOp.getVectorType();
     SmallVector<int64_t> originalSize = *constantMaskOp.getShapeForUnroll();
+    // The native/target shape coming from `getTargetShape` may have fewer
+    // dimensions than the op being unrolled (it is only required to be a
+    // divisor of the trailing dimensions, see `computeShapeRatio`). The
+    // logic below assumes a 1-1 correspondence between `targetShape` and the
+    // op's mask dimensions, so bail out if that does not hold.
+    if (targetShape->size() != originalSize.size())
+      return failure();
     Location loc = constantMaskOp.getLoc();
 
     Value result = arith::ConstantOp::create(rewriter, loc, resultType,
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index da058d5d2410d..542f01e8a91bd 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -602,6 +602,34 @@ func.func @vector_constant_mask() -> vector<16x16xi1> {
 
 // -----
 
+// 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 than the
+// target shape's rank (as below, rank 3 vs rank 2), the unroll patterns must
+// not fire, since they assume a 1-1 correspondence between the target shape
+// and the op's mask operands/dimensions. See #217175.
+
+// CHECK-LABEL: func @vector_create_mask_rank_mismatch
+//  CHECK-NEXT:   vector.create_mask %{{.*}}, %{{.*}}, %{{.*}} : vector<8x8x8xi1>
+//  CHECK-NEXT:   return
+func.func @vector_create_mask_rank_mismatch(%size1: index) -> vector<8x8x8xi1> {
+  %c8 = arith.constant 8 : index
+  %0 = vector.create_mask %c8, %size1, %c8 : vector<8x8x8xi1>
+  return %0 : vector<8x8x8xi1>
+}
+
+// -----
+
+// CHECK-LABEL: func @vector_constant_mask_rank_mismatch
+//  CHECK-NEXT:   vector.constant_mask [4, 8, 8] : vector<8x8x8xi1>
+//  CHECK-NEXT:   return
+func.func @vector_constant_mask_rank_mismatch() -> vector<8x8x8xi1> {
+  %0 = vector.constant_mask [4, 8, 8] : vector<8x8x8xi1>
+  return %0 : vector<8x8x8xi1>
+}
+
+// -----
+
 func.func @shape_cast_1D(%v: vector<16xf32>) -> vector<2x2x4xf32> {
   %0 = vector.shape_cast %v : vector<16xf32> to vector<2x2x4xf32>
   return %0 : vector<2x2x4xf32>

``````````

</details>


https://github.com/llvm/llvm-project/pull/217514