[Mlir-commits] [mlir] [mlir][vector] Fix multi-reduction unrolling with a lower rank target shape (PR #216799)
Hamza Qureshi
llvmlistbot at llvm.org
Mon Aug 17 11:10:05 PDT 2026
https://github.com/hamzaqureshi5 created https://github.com/llvm/llvm-project/pull/216799
Fixes #216640
`mlir-opt -test-vector-unrolling-patterns` crashes on a rank 3 `vector.multi_reduction`:
```
Assertion `offsets.size() == sizes.size() && offsets.size() == strides.size()' failed.
```
The unroll target shape is allowed to have fewer dimensions than the vector being unrolled, in which case it applies to the trailing dimensions. `computeShapeRatio` accepts that, and `UnrollElementwisePattern` handles it by padding the shape with leading unit dimensions.
`UnrollMultiReductionPattern` used the target shape as it is, both for the size of the slice taken from the source and to index the dimensions of the reduction. The offsets have the rank of the source, so with a shorter target shape the slice got fewer sizes than offsets, which is the assertion above. The test pass uses a native shape of `{2, 2}`, so any reduction on a rank 3 vector hits it.
The fix pads the target shape with leading unit dimensions before using it, the same way the elementwise pattern does.
## Testing
Added a rank 3 reduction to `mlir/test/Dialect/Vector/vector-unroll-options.mlir`. The existing RUN lines all pass `unroll-based-on-type`, so the plain form of the pass, the one with the `{2, 2}` native shape, was not covered at all; this adds a RUN line for it.
Checked on an assertions build that the reproducer aborts before the change and succeeds after, that the new test fails without it, that the unrolled IR round trips, and that a reduction whose rank already matches the native shape is still unrolled the same way. `mlir/test/Dialect/Vector` and `mlir/test/Conversion` are 513/513.
>From 03b6659a79a0a422f5350e802aa47e0c5eec3344 Mon Sep 17 00:00:00 2001
From: hamzaqureshi5 <hamza7771.861 at gmail.com>
Date: Mon, 17 Aug 2026 23:06:58 +0500
Subject: [PATCH] [mlir][vector] Fix multi-reduction unrolling with a lower
rank target shape
The unroll target shape may have fewer dimensions than the vector being
unrolled, in which case it applies to the trailing dimensions;
computeShapeRatio accepts that and UnrollElementwisePattern handles it by
padding the shape with leading unit dimensions.
UnrollMultiReductionPattern used the target shape directly, both as the sizes
of the slice taken from the source and to index the dimensions of the
reduction. With a target shape of a lower rank than the source, the slice
sizes had fewer entries than its offsets, which tripped an assertion in
inferStridedSliceOpResultType.
Pad the target shape the same way before using it.
Fixes #216640
---
.../Vector/Transforms/VectorUnroll.cpp | 19 ++++++++-----
.../Dialect/Vector/vector-unroll-options.mlir | 27 +++++++++++++++++++
2 files changed, 40 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 62869111496d1..565373b1866e5 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -383,17 +383,24 @@ struct UnrollMultiReductionPattern
Location loc = reductionOp.getLoc();
auto resultType = reductionOp->getResult(0).getType();
+ // A target shape with fewer dimensions than the source vector applies to
+ // its trailing dimensions. Add leading unit dimensions so that it can be
+ // used to slice the source and to index its dimensions, the same way
+ // UnrollElementwisePattern does.
+ SmallVector<int64_t> adjustedTargetShape(originalSize.size(), 1);
+ llvm::copy(*targetShape, adjustedTargetShape.end() - targetShape->size());
+
// Handle scalar result case: all dimensions are reduced.
// Each source tile is reduced to a scalar, and partial results are
// chained through the accumulator operand.
if (resultType.isIntOrFloat()) {
Value accumulator = reductionOp.getAcc();
for (SmallVector<int64_t> offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<int64_t> operandStrides(offsets.size(), 1);
Value slicedOperand =
rewriter.createOrFold<vector::ExtractStridedSliceOp>(
- loc, reductionOp.getSource(), offsets, *targetShape,
+ loc, reductionOp.getSource(), offsets, adjustedTargetShape,
operandStrides);
Operation *newOp = cloneOpWithOperandsAndTypes(
rewriter, loc, reductionOp, {slicedOperand, accumulator},
@@ -413,20 +420,20 @@ struct UnrollMultiReductionPattern
// Stride of the ratios, this gives us the offsets of sliceCount in a basis
// of multiples of the targetShape.
for (SmallVector<int64_t> offsets :
- StaticTileOffsetRange(originalSize, *targetShape)) {
+ StaticTileOffsetRange(originalSize, adjustedTargetShape)) {
SmallVector<Value> operands;
SmallVector<int64_t> operandStrides(offsets.size(), 1);
Value slicedOperand =
rewriter.createOrFold<vector::ExtractStridedSliceOp>(
- loc, reductionOp.getSource(), offsets, *targetShape,
+ loc, reductionOp.getSource(), offsets, adjustedTargetShape,
operandStrides);
operands.push_back(slicedOperand);
SmallVector<int64_t> dstShape;
SmallVector<int64_t> destOffset;
- for (size_t i : llvm::seq(size_t(0), targetShape->size())) {
+ for (size_t i : llvm::seq(size_t(0), adjustedTargetShape.size())) {
if (!reductionOp.isReducedDim(i)) {
destOffset.push_back(offsets[i]);
- dstShape.push_back((*targetShape)[i]);
+ dstShape.push_back(adjustedTargetShape[i]);
}
}
Value acc;
diff --git a/mlir/test/Dialect/Vector/vector-unroll-options.mlir b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
index da058d5d2410d..a2a7569bc016f 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -1,5 +1,6 @@
// RUN: mlir-opt %s -test-vector-unrolling-patterns=unroll-based-on-type -split-input-file | FileCheck %s
// RUN: mlir-opt %s -test-vector-unrolling-patterns="unroll-based-on-type unroll-order=2,0,1" -split-input-file | FileCheck %s --check-prefix=ORDER
+// RUN: mlir-opt %s -test-vector-unrolling-patterns -split-input-file | FileCheck %s --check-prefix=NATIVE
// RUN: mlir-opt %s -test-vector-unrolling-patterns="unroll-based-on-type unroll-order=0,3,1,2" -split-input-file | FileCheck %s --check-prefix=BATCHED
func.func @vector_contract_f32(%lhs : vector<8x4xf32>, %rhs : vector<8x4xf32>,
@@ -864,3 +865,29 @@ func.func @deinterleave_2d(%v: vector<4x8xi32>) -> (vector<4x4xi32>, vector<4x4x
// CHECK: {{.*}} = vector.insert_strided_slice {{.*}}, {{.*}} {offsets = [2, 0], strides = [1, 1]} : vector<2x4xi32> into vector<4x4xi32>
// CHECK: {{.*}} = vector.insert_strided_slice {{.*}}, {{.*}} {offsets = [2, 0], strides = [1, 1]} : vector<2x4xi32> into vector<4x4xi32>
// CHECK: return {{.*}}, {{.*}} : vector<4x4xi32>, vector<4x4xi32>
+
+// -----
+
+// The native unroll shape {2, 2} has fewer dimensions than the source vector.
+// It applies to the trailing dimensions, so the leading ones are sliced one
+// element at a time.
+
+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>
+}
+// NATIVE-LABEL: func @vector_multi_reduction_rank_mismatch
+// NATIVE: %[[V0:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
+// NATIVE: %[[E0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// NATIVE: %[[ACC0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [1, 2], strides = [1, 1]} : vector<2x2xf32> to vector<1x2xf32>
+// NATIVE: %[[R0:.*]] = vector.multi_reduction <add>, %[[E0]], %[[ACC0]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// NATIVE: %[[E1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// NATIVE: %[[R1:.*]] = vector.multi_reduction <add>, %[[E1]], %[[R0]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// NATIVE: %[[E2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// NATIVE: %[[ACC1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0], sizes = [1, 2], strides = [1, 1]} : vector<2x2xf32> to vector<1x2xf32>
+// NATIVE: %[[R2:.*]] = vector.multi_reduction <add>, %[[E2]], %[[ACC1]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// NATIVE: %[[E3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// NATIVE: %[[R3:.*]] = vector.multi_reduction <add>, %[[E3]], %[[R2]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// NATIVE: %[[V1:.*]] = vector.insert_strided_slice %[[R1]], %[[V0]] {offsets = [0, 0], strides = [1, 1]} : vector<1x2xf32> into vector<2x2xf32>
+// NATIVE: %[[V2:.*]] = vector.insert_strided_slice %[[R3]], %[[V1]] {offsets = [1, 0], strides = [1, 1]} : vector<1x2xf32> into vector<2x2xf32>
+// NATIVE: return %[[V2]] : vector<2x2xf32>
More information about the Mlir-commits
mailing list