[Mlir-commits] [mlir] 64bbf27 - [mlir][vector] Verify multi_reduction reduction dimensions (#216854)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Aug 24 00:54:45 PDT 2026
Author: Alessandro Potenza
Date: 2026-08-24T09:54:40+02:00
New Revision: 64bbf27bc7ab180bb01915f29065981df1e9a98b
URL: https://github.com/llvm/llvm-project/commit/64bbf27bc7ab180bb01915f29065981df1e9a98b
DIFF: https://github.com/llvm/llvm-project/commit/64bbf27bc7ab180bb01915f29065981df1e9a98b.diff
LOG: [mlir][vector] Verify multi_reduction reduction dimensions (#216854)
`vector.multi_reduction` does not bounds-check `reduction_dims`.
Out-of-range
entries verify cleanly and then crash the canonicalizer:
```mlir
%0 = vector.multi_reduction <add>, %a, %b [1, 2] : vector<1x4xf16> to vector<1xf16>
```
`getReductionMask()` builds `SmallVector<bool> res(sourceRank)` and
executes
`res[2] = true`. The `[1, -1]` variant indexes `res[(size_t)-1]`.
Reject out-of-range and duplicate dims in the verifier, following
`TransposeOp::verify`. The mask is reused by the shape-inference loop
below.
One existing test is itself the repro: `propagate-layout.mlir` uses `[1,
2]` on a
rank-2 vector. I changed the source shape to `vector<1x1x4xf16>` rather
than the
dims to `[1]`, because `[1]` would flip `verticalLaneLayout` and
collapse the
lane layout the test checks. XeGPU owners: please confirm that is what
was meant.
Assisted-by: Claude (Anthropic)
AI-assisted, disclosed per the LLVM AI Tool Use Policy.
Added:
Modified:
mlir/lib/Dialect/Vector/IR/VectorOps.cpp
mlir/test/Dialect/Vector/invalid.mlir
mlir/test/Dialect/XeGPU/propagate-layout.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
index 67bf1c68674a3..f8f3deb2e4789 100644
--- a/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
+++ b/mlir/lib/Dialect/Vector/IR/VectorOps.cpp
@@ -603,16 +603,24 @@ MultiDimReductionOp::getShapeForUnroll() {
}
LogicalResult MultiDimReductionOp::verify() {
+ // Verify the reduction dimensions.
+ 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 d3a5499eb3e24..489f5489030fd 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 ee420a78117b5..b4e69173c1144 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