[Mlir-commits] [mlir] [mlir][vector] Verify multi_reduction reduction dimensions (PR #216854)
Alessandro Potenza
llvmlistbot at llvm.org
Mon Aug 17 14:51:12 PDT 2026
https://github.com/alepot55 created https://github.com/llvm/llvm-project/pull/216854
`vector.multi_reduction`'s verifier never bounds-checks `reduction_dims`.
Out-of-range entries are accepted and then crash the canonicalizer:
```mlir
func.func @oob(%a: vector<1x4xf16>, %b: vector<1xf16>) -> vector<1xf16> {
%0 = vector.multi_reduction <add>, %a, %b [1, 2] : vector<1x4xf16> to vector<1xf16>
return %0 : vector<1xf16>
}
```
`mlir-opt` alone returns 0 -- the verifier accepts dimension `2` on a rank-2
vector. `mlir-opt --canonicalize` aborts: `ElideUnitDimsInMultiDimReduction`
calls `isReducedDim(0)`, which evaluates `getReductionMask()`, which builds
`SmallVector<bool> res(sourceRank, false)` and then executes `res[2] = true`.
The negative variant `[1, -1]` indexes `res[(size_t)-1]`.
Reject out-of-range and duplicate reduction dimensions in
`MultiDimReductionOp::verify()`, following the in-tree precedent of
`TransposeOp::verify` (same shape of check, same diagnostic wording). The mask
built for validation is then reused by the existing shape-inference loop,
replacing a quadratic `llvm::any_of` lambda.
`getReductionMask()` / `isReducedDim()` have callers across the Vector and XeGPU
transforms, so the out-of-bounds access is reachable from more than the one
pattern above.
### One existing test changes, and it needs an owner's opinion
`@vector_2d_reduction_with_fractional_subgroup_size_1x4` in
`mlir/test/Dialect/XeGPU/propagate-layout.mlir` is itself the repro: `[1, 2]` on
a rank-2 `vector<1x4xf16>`. It is a copy of the rank-3 test above it whose
`shape_cast` target was changed to rank 2 while the dims were left at `[1, 2]`.
Verified necessary, not cosmetic: with this patch and the unmodified file,
`mlir-opt` reports `'vector.multi_reduction' op reduction dimension out of
range: 2` at line 766.
**I changed the source shape to `vector<1x1x4xf16>` and kept `[1, 2]`, rather
than changing the dims to `[1]`, and I would like XeGPU owners to confirm that
is what was meant.** The reason it matters: `computeReductionLaneLayoutAndData`
branches on `verticalLaneLayout`, which is
`consumerReductionDims.empty() && reductionDims.size() == 1 && reductionDims[0] == srcRank - 1`.
The consumer here is an `xegpu.store` with a plain `LayoutAttr`, so
`consumerReductionDims` is empty; changing the dims to `[1]` would flip that
branch true and yield `lane_layout = [1, 1]`, one lane covering a 1x4 vector,
which destroys exactly the 4-lane distribution the test is named for. Keeping two
reduction dims preserves it as `[1, 1, 4]`. The sibling
`@vector_reduction_broadcast_transpose` does use rank-2 `[1]` and is not vertical,
because its consumer carries a `SliceAttr` -- consistent with that reading.
The honest version: the pass appears unable to produce a 4-lane layout for a
genuine rank-2 `[1]` reduction, and this change sidesteps that rather than
fixing it. If you would rather have `[1]` with regenerated CHECK lines, say so
and I will switch.
### Verified by execution
- Both the accepted-then-crashing behaviour and the `[1, -1]` variant reproduce on `d4e78d7f5`.
- The three new `invalid.mlir` cases fail on unpatched `mlir-opt` ("expected error ... was not produced")
and pass with the patch.
- No regression: `Dialect/Vector` 101/101, `Dialect/XeGPU` 25/25, plus MemRef, SCF, Affine,
Transforms (368 total).
- I audited every `multi_reduction` in the tree (230 ops across 58 files); the XeGPU test above is
the only one affected.
- `git clang-format` reports no changes.
---
Assisted-by: Claude (Anthropic)
This patch was written with AI assistance, disclosed per the LLVM AI Tool Use Policy.
Everything reported above as verified was verified by building and running.
>From 73e8b6f44595142ac761e52b359790e287410c7d Mon Sep 17 00:00:00 2001
From: Alessandro Potenza <ap.alessandro.potenza at gmail.com>
Date: Mon, 17 Aug 2026 23:00:46 +0200
Subject: [PATCH] [mlir][vector] Bounds-check multi_reduction reduction_dims
`MultiDimReductionOp::verify()` never checked that the entries of
`reduction_dims` are valid indices into the source vector, so out-of-range
(and duplicate) dims were accepted as valid IR:
func.func @oob(%a: vector<1x4xf16>, %b: vector<1xf16>) -> vector<1xf16> {
%0 = vector.multi_reduction <add>, %a, %b [1, 2]
: vector<1x4xf16> to vector<1xf16>
return %0 : vector<1xf16>
}
`mlir-opt` accepts this, and `mlir-opt --canonicalize` then aborts:
`ElideUnitDimsInMultiDimReduction` calls `isReducedDim(0)`, which evaluates
`getReductionMask()`. That builds `SmallVector<bool> res(sourceRank, false)`
and writes `res[dim] = true`, i.e. `res[2]` on a rank-2 source. A negative
dim such as `[1, -1]` indexes `res[(size_t)-1]`.
`getReductionMask()`/`isReducedDim()` are used by seven call sites across
the Vector and XeGPU transforms, all of which assumed dims were validated.
Verify the dims up front, rejecting out-of-range and repeated entries. This
follows `TransposeOp::verify()`, which already bounds-checks its permutation
list the same way. The existing shape-inference loop now reuses the computed
mask instead of a quadratic `llvm::any_of` scan.
---
mlir/lib/Dialect/Vector/IR/VectorOps.cpp | 18 ++++++++++----
mlir/test/Dialect/Vector/invalid.mlir | 24 +++++++++++++++++++
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 14 +++++------
3 files changed, 45 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 7dbe1847f077d..10c0cad4b5730 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -603,16 +603,26 @@ MultiDimReductionOp::getShapeForUnroll() {
}
LogicalResult MultiDimReductionOp::verify() {
+ // Verify the reduction dimensions before using them to index into the source
+ // shape below. Out-of-range or repeated entries would otherwise silently
+ // corrupt every consumer of getReductionMask()/isReducedDim().
+ int64_t sourceRank = getSourceVectorType().getRank();
+ SmallVector<bool> isReduced(sourceRank, false);
+ for (int64_t dim : getReductionDims()) {
+ if (dim < 0 || dim >= sourceRank)
+ return emitOpError("reduction dimension out of range: ") << dim;
+ if (isReduced[dim])
+ return emitOpError("duplicate reduction dimension: ") << dim;
+ isReduced[dim] = true;
+ }
+
SmallVector<int64_t> targetShape;
SmallVector<bool> scalableDims;
Type inferredReturnType;
auto sourceScalableDims = getSourceVectorType().getScalableDims();
for (auto [dimIdx, dimSize] :
llvm::enumerate(getSourceVectorType().getShape()))
- if (!llvm::any_of(getReductionDims(),
- [dimIdx = dimIdx](int64_t reductionDimIdx) {
- return reductionDimIdx == static_cast<int64_t>(dimIdx);
- })) {
+ if (!isReduced[dimIdx]) {
targetShape.push_back(dimSize);
scalableDims.push_back(sourceScalableDims[dimIdx]);
}
diff --git a/mlir/test/Dialect/Vector/invalid.mlir b/mlir/test/Dialect/Vector/invalid.mlir
index 0097db39f6ed9..db2ea26f9bef7 100644
--- a/mlir/test/Dialect/Vector/invalid.mlir
+++ b/mlir/test/Dialect/Vector/invalid.mlir
@@ -1273,6 +1273,30 @@ func.func @multi_reduce_invalid_type(%arg0: vector<4x16xf32>, %acc: vector<16xf3
// -----
+func.func @multi_reduce_dim_out_of_range(%arg0: vector<4x16xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+ // expected-error at +1 {{'vector.multi_reduction' op reduction dimension out of range: 2}}
+ %0 = vector.multi_reduction <add>, %arg0, %acc [1, 2] : vector<4x16xf32> to vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// -----
+
+func.func @multi_reduce_negative_dim(%arg0: vector<4x16xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+ // expected-error at +1 {{'vector.multi_reduction' op reduction dimension out of range: -1}}
+ %0 = vector.multi_reduction <add>, %arg0, %acc [1, -1] : vector<4x16xf32> to vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// -----
+
+func.func @multi_reduce_duplicate_dim(%arg0: vector<4x16xf32>, %acc: vector<4xf32>) -> vector<4xf32> {
+ // expected-error at +1 {{'vector.multi_reduction' op duplicate reduction dimension: 1}}
+ %0 = vector.multi_reduction <add>, %arg0, %acc [1, 1] : vector<4x16xf32> to vector<4xf32>
+ return %0 : vector<4xf32>
+}
+
+// -----
+
func.func @transpose_rank_mismatch_0d(%arg0: vector<f32>) {
// expected-error at +1 {{'vector.transpose' op vector result rank mismatch: 1}}
%0 = vector.transpose %arg0, [] : vector<f32> to vector<100xf32>
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 5d45ec79a5a8a..a83a519f0315d 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -747,23 +747,23 @@ func.func @vector_reduction_scalar(%arg0: memref<1024xf16>, %arg1: memref<16xf16
// -----
gpu.module @test {
-// CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4(
+// CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size_1x1x4(
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [4], lane_data = [1]>} dense<true> : vector<4xi1>
// CHECK: %[[IDX:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [4], lane_data = [1]>} : vector<4xindex>
// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[IDX]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [4], lane_data = [1]>}> : memref<1024xf16>, vector<4xindex>, vector<4xi1> -> vector<4xf16>
-// CHECK: %[[SC:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 4], lane_data = [1, 1]>} : vector<4xf16> to vector<1x4xf16>
-// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 4], lane_data = [1, 1]>, dims = [1, 2]>} dense<0.000000e+00> : vector<1xf16>
-// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[SC]], %[[ACC]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 4], lane_data = [1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x4xf16> to vector<1xf16>
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 4], lane_data = [1, 1, 1]>} : vector<4xf16> to vector<1x1x4xf16>
+// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 4], lane_data = [1, 1, 1]>, dims = [1, 2]>} dense<0.000000e+00> : vector<1xf16>
+// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[SC]], %[[ACC]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 4], lane_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x1x4xf16> to vector<1xf16>
// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1], lane_data = [1]>} dense<true> : vector<1xi1>
// CHECK: %[[OFF:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [1], lane_data = [1]>} dense<1> : vector<1xindex>
// CHECK: xegpu.store %[[RED]], %arg1[%[[OFF]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [1], lane_data = [1]>}> : vector<1xf16>, memref<16xf16>, vector<1xindex>, vector<1xi1>
-func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
+func.func @vector_2d_reduction_with_fractional_subgroup_size_1x1x4(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<4xi1>
%0 = vector.step : vector<4xindex>
%1 = xegpu.load %arg0[%0], %cst : memref<1024xf16>, vector<4xindex>, vector<4xi1> -> vector<4xf16>
- %2 = vector.shape_cast %1 : vector<4xf16> to vector<1x4xf16>
+ %2 = vector.shape_cast %1 : vector<4xf16> to vector<1x1x4xf16>
%cst_0 = arith.constant dense<0.000000e+00> : vector<1xf16>
- %4 = vector.multi_reduction <add>, %2, %cst_0 [1, 2] : vector<1x4xf16> to vector<1xf16>
+ %4 = vector.multi_reduction <add>, %2, %cst_0 [1, 2] : vector<1x1x4xf16> to vector<1xf16>
%cst_2 = arith.constant dense<true> : vector<1xi1>
%cst_3 = arith.constant dense<1> : vector<1xindex>
xegpu.store %4, %arg1[%cst_3], %cst_2 : vector<1xf16>, memref<16xf16>, vector<1xindex>, vector<1xi1>
More information about the Mlir-commits
mailing list