[Mlir-commits] [mlir] [mlir][XeGPU] Fix 2D reduction dim ordering for nested slice layouts (PR #230850)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Mon Oct 19 00:34:15 PDT 2026
https://github.com/edisongz updated https://github.com/llvm/llvm-project/pull/230850
>From c5a4d4239105c0f52bc6fa8dab215cc66ba98ab9 Mon Sep 17 00:00:00 2001
From: Yijie Jiang <edisongz123 at gmail.com>
Date: Sun, 11 Oct 2026 03:41:19 +0800
Subject: [PATCH 1/2] [mlir][XeGPU] Fix 2D reduction dim ordering for nested
slice layouts
The reduction dims were indexed against the flattened root layout instead
of the source layout, so nested-slice sources misclassified intra/cross
lanes and reduced in the wrong order.
Fixes #228644
---
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 10 +--
.../test/Dialect/XeGPU/peephole-optimize.mlir | 76 +++++++++++++++++++
2 files changed, 81 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index b430a5430c3c7..97f57e32433dd 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -565,12 +565,12 @@ class MultiRed2dOpPattern
assert(layout.isForSubgroup() && "Must know the lane layout");
assert(reductionDims.size() == 2 && "Expected 2D reduction");
int64_t intra = -1, cross = -1;
- xegpu::LayoutAttr layoutAttr = dyn_cast<xegpu::LayoutAttr>(layout);
+ // reductionDims index the source layout (the slice's parent), not the
+ // flattened root.
+ xegpu::DistributeLayoutAttr srcLayout = layout;
if (auto layoutSliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
- layoutAttr =
- dyn_cast<xegpu::LayoutAttr>(layoutSliceAttr.flatten().getParent());
- assert(layoutAttr);
- SmallVector<int64_t> laneLayout = layoutAttr.getEffectiveLaneLayoutAsInt();
+ srcLayout = layoutSliceAttr.getParent();
+ SmallVector<int64_t> laneLayout = srcLayout.getEffectiveLaneLayoutAsInt();
assert(laneLayout.size() && "Expected a non-empty layout");
// try to pick a dim that does not communicate
diff --git a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index 11c58d28377d2..29a7604c5d986 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -512,6 +512,82 @@ gpu.module @xevm_test {
}
}
+// -----
+// Nested slice layout: source effective lane layout is [16, 1], so dim 1
+// (intra-lane) reduces first, then the cross-lane dim 0.
+// CHECK-LABEL: gpu.func @nested_slice_reduce_order(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<4x16x8xi32>) -> i32 {
+// CHECK: %[[ACC_VEC:.*]] = arith.constant dense<0> : vector<16xi32>
+// CHECK: %[[ACC_SCALAR:.*]] = arith.constant 0 : i32
+// CHECK: %[[EXTRACTED:.*]] = vector.extract %[[ARG0]][0] : vector<16x8xi32> from vector<4x16x8xi32>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[EXTRACTED]], %[[ACC_VEC]] [1] : vector<16x8xi32> to vector<16xi32>
+// CHECK: %[[REDUCE_2:.*]] = vector.multi_reduction <add>, %[[REDUCE_1]], %[[ACC_SCALAR]] [0] : vector<16xi32> to i32
+// CHECK: xegpu.convert_layout %[[REDUCE_2]]
+#root = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [1, 1, 1]>
+#src = #xegpu.slice<#root, dims = [0]>
+#res = #xegpu.slice<#src, dims = [0, 1]>
+gpu.module @xevm_module {
+ gpu.func @nested_slice_reduce_order(%arg: vector<4x16x8xi32>) -> i32 {
+ %v = vector.extract %arg[0]
+ {layout_operand_0 = #root, layout_result_0 = #src}
+ : vector<16x8xi32> from vector<4x16x8xi32>
+ %zero = arith.constant 0 : i32
+ %r = vector.multi_reduction <add>, %v, %zero
+ {layout_result_0 = #res}
+ [0, 1] : vector<16x8xi32> to i32
+ gpu.return %r : i32
+ }
+}
+
+// -----
+// Reducing non-leading dims [1, 2]: source effective lane layout is [1, 16, 1],
+// so dim 2 (intra-lane) reduces first, then dim 1.
+// CHECK-LABEL: gpu.func @nested_slice_reduce_dims_12(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<2x2x16x8xi32>) -> vector<2xi32> {
+// CHECK: %[[ACC_2D:.*]] = arith.constant dense<0> : vector<2x16xi32>
+// CHECK: %[[ACC_1D:.*]] = arith.constant dense<0> : vector<2xi32>
+// CHECK: %[[EXTRACTED:.*]] = vector.extract %[[ARG0]][0] : vector<2x16x8xi32> from vector<2x2x16x8xi32>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[EXTRACTED]], %[[ACC_2D]] [2] : vector<2x16x8xi32> to vector<2x16xi32>
+// CHECK: %[[REDUCE_2:.*]] = vector.multi_reduction <add>, %[[REDUCE_1]], %[[ACC_1D]] [1] : vector<2x16xi32> to vector<2xi32>
+// CHECK: xegpu.convert_layout %[[REDUCE_2]]
+#root3 = #xegpu.layout<lane_layout = [1, 1, 16, 1], lane_data = [1, 1, 1, 1]>
+#src3 = #xegpu.slice<#root3, dims = [0]>
+#res3 = #xegpu.slice<#src3, dims = [1, 2]>
+gpu.module @xevm_module {
+ gpu.func @nested_slice_reduce_dims_12(%arg: vector<2x2x16x8xi32>) -> vector<2xi32> {
+ %v = vector.extract %arg[0]
+ {layout_operand_0 = #root3, layout_result_0 = #src3}
+ : vector<2x16x8xi32> from vector<2x2x16x8xi32>
+ %zero = arith.constant dense<0> : vector<2xi32>
+ %r = vector.multi_reduction <add>, %v, %zero
+ {layout_result_0 = #res3}
+ [1, 2] : vector<2x16x8xi32> to vector<2xi32>
+ gpu.return %r : vector<2xi32>
+ }
+}
+
+// -----
+// Both dims are cross-lane ([2, 8]), so order is irrelevant; fall back to
+// reducing dims[0] then dims[1] without reading an uninitialized intra dim.
+// CHECK-LABEL: gpu.func @reduce_both_cross(
+// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<4x8xi32>) -> i32 {
+// CHECK-DAG: %[[ACC_VEC:.*]] = arith.constant dense<0> : vector<8xi32>
+// CHECK-DAG: %[[ACC_SCALAR:.*]] = arith.constant 0 : i32
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[ARG0]], %[[ACC_VEC]] [0] : vector<4x8xi32> to vector<8xi32>
+// CHECK: %[[REDUCE_2:.*]] = vector.multi_reduction <add>, %[[REDUCE_1]], %[[ACC_SCALAR]] [0] : vector<8xi32> to i32
+// CHECK: xegpu.convert_layout %[[REDUCE_2]]
+#root4 = #xegpu.layout<lane_layout = [2, 8], lane_data = [1, 1]>
+#res4 = #xegpu.slice<#root4, dims = [0, 1]>
+gpu.module @xevm_module {
+ gpu.func @reduce_both_cross(%arg: vector<4x8xi32>) -> i32 {
+ %zero = arith.constant 0 : i32
+ %r = vector.multi_reduction <add>, %arg, %zero
+ {layout_result_0 = #res4}
+ [0, 1] : vector<4x8xi32> to i32
+ gpu.return %r : i32
+ }
+}
+
// -----
// Transpose optimization on a >2D descriptor with unit leading dims.
// CHECK-LABEL: gpu.func @transpose_4d(
>From da23502ddfdff84a1221cd04e91cced94236c75b Mon Sep 17 00:00:00 2001
From: Yijie Jiang <edisongz123 at gmail.com>
Date: Mon, 19 Oct 2026 15:32:56 +0800
Subject: [PATCH 2/2] [mlir][XeGPU] Use reduction chains in nested slice
reduction tests
---
.../test/Dialect/XeGPU/peephole-optimize.mlir | 52 +++++++++++--------
1 file changed, 29 insertions(+), 23 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
index 29a7604c5d986..876461f7dbef4 100644
--- a/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
+++ b/mlir/test/Dialect/XeGPU/peephole-optimize.mlir
@@ -513,14 +513,16 @@ gpu.module @xevm_test {
}
// -----
-// Nested slice layout: source effective lane layout is [16, 1], so dim 1
+// Nested slice layout: the 2D reduction's source is a 1D reduction result
+// (slice<root, [0]>), so its effective lane layout is [16, 1]; dim 1
// (intra-lane) reduces first, then the cross-lane dim 0.
// CHECK-LABEL: gpu.func @nested_slice_reduce_order(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<4x16x8xi32>) -> i32 {
-// CHECK: %[[ACC_VEC:.*]] = arith.constant dense<0> : vector<16xi32>
-// CHECK: %[[ACC_SCALAR:.*]] = arith.constant 0 : i32
-// CHECK: %[[EXTRACTED:.*]] = vector.extract %[[ARG0]][0] : vector<16x8xi32> from vector<4x16x8xi32>
-// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[EXTRACTED]], %[[ACC_VEC]] [1] : vector<16x8xi32> to vector<16xi32>
+// CHECK-DAG: %[[ACC_VEC:.*]] = arith.constant dense<0> : vector<16xi32>
+// CHECK-DAG: %[[ACC_2D:.*]] = arith.constant dense<0> : vector<16x8xi32>
+// CHECK-DAG: %[[ACC_SCALAR:.*]] = arith.constant 0 : i32
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[ARG0]], %[[ACC_2D]] [0] : vector<4x16x8xi32> to vector<16x8xi32>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[ACC_VEC]] [1] : vector<16x8xi32> to vector<16xi32>
// CHECK: %[[REDUCE_2:.*]] = vector.multi_reduction <add>, %[[REDUCE_1]], %[[ACC_SCALAR]] [0] : vector<16xi32> to i32
// CHECK: xegpu.convert_layout %[[REDUCE_2]]
#root = #xegpu.layout<lane_layout = [1, 16, 1], lane_data = [1, 1, 1]>
@@ -528,26 +530,29 @@ gpu.module @xevm_test {
#res = #xegpu.slice<#src, dims = [0, 1]>
gpu.module @xevm_module {
gpu.func @nested_slice_reduce_order(%arg: vector<4x16x8xi32>) -> i32 {
- %v = vector.extract %arg[0]
- {layout_operand_0 = #root, layout_result_0 = #src}
- : vector<16x8xi32> from vector<4x16x8xi32>
- %zero = arith.constant 0 : i32
- %r = vector.multi_reduction <add>, %v, %zero
+ %c1 = arith.constant dense<0> : vector<16x8xi32>
+ %c2 = arith.constant 0 : i32
+ %r1 = vector.multi_reduction <add>, %arg, %c1
+ {layout_result_0 = #src}
+ [0] : vector<4x16x8xi32> to vector<16x8xi32>
+ %r2 = vector.multi_reduction <add>, %r1, %c2
{layout_result_0 = #res}
[0, 1] : vector<16x8xi32> to i32
- gpu.return %r : i32
+ gpu.return %r2 : i32
}
}
// -----
-// Reducing non-leading dims [1, 2]: source effective lane layout is [1, 16, 1],
-// so dim 2 (intra-lane) reduces first, then dim 1.
+// Reducing non-leading dims [1, 2] on a nested slice source (slice<root, [0]>):
+// effective lane layout is [1, 16, 1], so dim 2 (intra-lane) reduces first,
+// then the cross-lane dim 1.
// CHECK-LABEL: gpu.func @nested_slice_reduce_dims_12(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<2x2x16x8xi32>) -> vector<2xi32> {
-// CHECK: %[[ACC_2D:.*]] = arith.constant dense<0> : vector<2x16xi32>
-// CHECK: %[[ACC_1D:.*]] = arith.constant dense<0> : vector<2xi32>
-// CHECK: %[[EXTRACTED:.*]] = vector.extract %[[ARG0]][0] : vector<2x16x8xi32> from vector<2x2x16x8xi32>
-// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[EXTRACTED]], %[[ACC_2D]] [2] : vector<2x16x8xi32> to vector<2x16xi32>
+// CHECK-DAG: %[[ACC_2D:.*]] = arith.constant dense<0> : vector<2x16xi32>
+// CHECK-DAG: %[[ACC_3D:.*]] = arith.constant dense<0> : vector<2x16x8xi32>
+// CHECK-DAG: %[[ACC_1D:.*]] = arith.constant dense<0> : vector<2xi32>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[ARG0]], %[[ACC_3D]] [0] : vector<2x2x16x8xi32> to vector<2x16x8xi32>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[ACC_2D]] [2] : vector<2x16x8xi32> to vector<2x16xi32>
// CHECK: %[[REDUCE_2:.*]] = vector.multi_reduction <add>, %[[REDUCE_1]], %[[ACC_1D]] [1] : vector<2x16xi32> to vector<2xi32>
// CHECK: xegpu.convert_layout %[[REDUCE_2]]
#root3 = #xegpu.layout<lane_layout = [1, 1, 16, 1], lane_data = [1, 1, 1, 1]>
@@ -555,14 +560,15 @@ gpu.module @xevm_module {
#res3 = #xegpu.slice<#src3, dims = [1, 2]>
gpu.module @xevm_module {
gpu.func @nested_slice_reduce_dims_12(%arg: vector<2x2x16x8xi32>) -> vector<2xi32> {
- %v = vector.extract %arg[0]
- {layout_operand_0 = #root3, layout_result_0 = #src3}
- : vector<2x16x8xi32> from vector<2x2x16x8xi32>
- %zero = arith.constant dense<0> : vector<2xi32>
- %r = vector.multi_reduction <add>, %v, %zero
+ %c1 = arith.constant dense<0> : vector<2x16x8xi32>
+ %c2 = arith.constant dense<0> : vector<2xi32>
+ %r1 = vector.multi_reduction <add>, %arg, %c1
+ {layout_result_0 = #src3}
+ [0] : vector<2x2x16x8xi32> to vector<2x16x8xi32>
+ %r2 = vector.multi_reduction <add>, %r1, %c2
{layout_result_0 = #res3}
[1, 2] : vector<2x16x8xi32> to vector<2xi32>
- gpu.return %r : vector<2xi32>
+ gpu.return %r2 : vector<2xi32>
}
}
More information about the Mlir-commits
mailing list