[Mlir-commits] [mlir] [mlir][vector] Fix multi-reduction unrolling with a lower rank target shape (PR #216799)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 17 11:10:47 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir
Author: Hamza Qureshi (hamzaqureshi5)
<details>
<summary>Changes</summary>
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.
---
Full diff: https://github.com/llvm/llvm-project/pull/216799.diff
2 Files Affected:
- (modified) mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp (+13-6)
- (modified) mlir/test/Dialect/Vector/vector-unroll-options.mlir (+27)
``````````diff
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>
``````````
</details>
https://github.com/llvm/llvm-project/pull/216799
More information about the Mlir-commits
mailing list