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

[Aman Singh]

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

`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)

>From a8dd32470603c36a7b8bbbb4f88a0aa19115b320 Mon Sep 17 00:00:00 2001
From: amanyagami <2amansingh2 at gmail.com>
Date: Wed, 19 Aug 2026 19:45:40 -0700
Subject: [PATCH] [mlir][Vector] Fix crash unrolling multi_reduction with
 rank-mismatched native shape

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.

Fixes https://github.com/llvm/llvm-project/issues/216640
---
 .../Dialect/Vector/Transforms/VectorUnroll.cpp | 13 +++++++++++++
 .../Dialect/Vector/vector-unroll-options.mlir  | 18 ++++++++++++++++++
 2 files changed, 31 insertions(+)

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