[Mlir-commits] [mlir] [mlir][Vector] Fix crash unrolling multi_reduction with rank-mismatched native shape (PR #217513)

[llvmlistbot at llvm.org]

llvmorg-github-actions[bot] wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Aman Singh (amanyagami)

<details>
<summary>Changes</summary>

`UnrollMultiReductionPattern` assumes the native/target unroll shape
returned by `getTargetShape()` always has the same rank as the
`vector.multi_reduction` op's source vector. 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, the offsets computed from the op's full rank no longer match
the (shorter) sizes/strides built from the target shape, and building
the resulting `vector.extract_strided_slice` op hits the assertion
`offsets.size() == sizes.size() && offsets.size() == strides.size()`
in `inferStridedSliceOpResultType`.

Fix by bailing out with a match failure when the target shape's rank
does not match the op's rank, rather than asserting later while
building the strided-slice op.

Verified: reverting this fix reproduces the reported crash (via
`-test-vector-unrolling-patterns` on the reported reproducer, which
exercises this pattern through a permuted `vector.transfer_read`
feeding a rank-3 `vector.multi_reduction`); with the fix it succeeds,
and `mlir/test/Dialect/Vector/vector-unroll-options.mlir` passes.

Fixes #<!-- -->216640

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

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


2 Files Affected:

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


``````````diff
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 62869111496d1..8bde314f13ebc 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -380,6 +380,19 @@ struct UnrollMultiReductionPattern
     if (!targetShape)
       return failure();
     SmallVector<int64_t> originalSize = *reductionOp.getShapeForUnroll();
+    // This pattern unconditionally indexes into `targetShape` and the
+    // reduction's `isReducedDim` using the same rank as `originalSize`
+    // (e.g. to compute per-dimension offsets and the result shape). The
+    // native/target shape returned above may have a smaller rank than the
+    // op being unrolled (its trailing dimensions are then implicitly
+    // aligned with the leading dimensions of `originalSize`, see
+    // StaticTileOffsetRange), which this pattern does not handle. Bail out
+    // rather than building ops with mismatched offsets/sizes/strides.
+    if (targetShape->size() != originalSize.size())
+      return rewriter.notifyMatchFailure(
+          reductionOp, "unrolling multi_reduction with a native/target "
+                       "shape of different rank than the source vector is "
+                       "not supported");
     Location loc = reductionOp.getLoc();
     auto resultType = reductionOp->getResult(0).getType();
 
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index da058d5d2410d..18a49b378be2b 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -267,6 +267,24 @@ func.func @vector_multi_reduction_scalar(%v: vector<4x2xf32>, %acc: f32) -> f32
 
 // -----
 
+// This is a negative test case to ensure that multi_reduction ops whose
+// source rank (3) does not match the rank of the configured native/target
+// shape (2x2) are not unrolled, rather than crashing while building an
+// extract_strided_slice op with mismatched offsets/sizes/strides.
+// See https://github.com/llvm/llvm-project/issues/216640.
+func.func @negative_multi_reduction_rank_mismatch(
+    %v : vector<8x12x4xf32>, %acc: vector<8x12xf32>) -> vector<8x12xf32> {
+  %0 = vector.multi_reduction #vector.kind<add>, %v, %acc [2]
+      : vector<8x12x4xf32> to vector<8x12xf32>
+  return %0 : vector<8x12xf32>
+}
+// CHECK-LABEL: func @negative_multi_reduction_rank_mismatch
+//  CHECK-SAME:   %[[V:.*]]: vector<8x12x4xf32>, %[[ACC:.*]]: vector<8x12xf32>
+//       CHECK:   %[[R:.*]] = vector.multi_reduction <add>, %[[V]], %[[ACC]] [2] : vector<8x12x4xf32> to vector<8x12xf32>
+//       CHECK:   return %[[R]] : vector<8x12xf32>
+
+// -----
+
 func.func @vector_reduction(%v : vector<8xf32>) -> f32 {
   %0 = vector.reduction <add>, %v : vector<8xf32> into f32
   return %0 : f32

``````````

</details>


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