[Mlir-commits] [mlir] 5e1bcd6 - [mlir][vector] Fix multi-reduction unrolling with a lower rank target shape (#216799)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 24 11:46:36 PDT 2026
Author: Hamza Qureshi
Date: 2026-08-24T19:46:31+01:00
New Revision: 5e1bcd683d37d9ea024be18fd8419e497bd614fc
URL: https://github.com/llvm/llvm-project/commit/5e1bcd683d37d9ea024be18fd8419e497bd614fc
DIFF: https://github.com/llvm/llvm-project/commit/5e1bcd683d37d9ea024be18fd8419e497bd614fc.diff
LOG: [mlir][vector] Fix multi-reduction unrolling with a lower rank target shape (#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.
---------
Co-authored-by: Federico Bruzzone <federico.bruzzone.i at gmail.com>
Added:
Modified:
mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
mlir/test/Dialect/Vector/vector-unroll-options.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
index 62869111496d1..a95bfba0814bd 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorUnroll.cpp
@@ -383,17 +383,23 @@ 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.
+ 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 +419,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 ae31530837c1a..58d35ef35c0bd 100644
--- a/mlir/test/Dialect/Vector/vector-unroll-options.mlir
+++ b/mlir/test/Dialect/Vector/vector-unroll-options.mlir
@@ -864,3 +864,27 @@ 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>
+
+// -----
+
+// TargetShape [2, 2] has a lower rank than the reduced source <2x2x4>, so it
+// applies to the trailing dimensions and is padded to [1, 2, 2].
+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>
+}
+// CHECK-LABEL: func @vector_multi_reduction_rank_mismatch
+// CHECK: %[[V0:.*]] = arith.constant dense<0.000000e+00> : vector<2x2xf32>
+// CHECK: %[[E0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[ACC0:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0], sizes = [1, 2], strides = [1, 1]} : vector<2x2xf32> to vector<1x2xf32>
+// CHECK: %[[R0:.*]] = vector.multi_reduction <add>, %[[E0]], %[[ACC0]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[E1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [0, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[R1:.*]] = vector.multi_reduction <add>, %[[E1]], %[[R0]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[E2:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 0], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[ACC1:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0], sizes = [1, 2], strides = [1, 1]} : vector<2x2xf32> to vector<1x2xf32>
+// CHECK: %[[R2:.*]] = vector.multi_reduction <add>, %[[E2]], %[[ACC1]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[E3:.*]] = vector.extract_strided_slice %{{.*}} {offsets = [1, 0, 2], sizes = [1, 2, 2], strides = [1, 1, 1]} : vector<2x2x4xf32> to vector<1x2x2xf32>
+// CHECK: %[[R3:.*]] = vector.multi_reduction <add>, %[[E3]], %[[R2]] [2] : vector<1x2x2xf32> to vector<1x2xf32>
+// CHECK: %[[V1:.*]] = vector.insert_strided_slice %[[R1]], %[[V0]] {offsets = [0, 0], strides = [1, 1]} : vector<1x2xf32> into vector<2x2xf32>
+// CHECK: %[[V2:.*]] = vector.insert_strided_slice %[[R3]], %[[V1]] {offsets = [1, 0], strides = [1, 1]} : vector<1x2xf32> into vector<2x2xf32>
+// CHECK: return %[[V2]] : vector<2x2xf32>
More information about the Mlir-commits
mailing list