[Mlir-commits] [mlir] [MLIR][XeGPU] Improve multi-reduction layout inference (PR #224618)
Artem Kroviakov
llvmlistbot at llvm.org
Fri Sep 18 05:27:03 PDT 2026
https://github.com/akroviakov created https://github.com/llvm/llvm-project/pull/224618
This PR improves the order construction in layout inference and fixes layout's dim collapse.
1. Layout inference now puts the reduction dim order row-major adjacent to all dims except for the outermost.
2. Dim collapse now checks order adjacency in both directions.
>From 2861c7940bfad688298f126ea39fbf2792c00bba Mon Sep 17 00:00:00 2001
From: Artem Kroviakov <artem.kroviakov at intel.com>
Date: Fri, 18 Sep 2026 12:23:50 +0000
Subject: [PATCH] [MLIR][XeGPU] Improve multi-reduction layout inference
---
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 30 +++++--------
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 44 ++++++++++++++-----
.../XeGPU/propagate-layout-subgroup.mlir | 33 +++++++++++---
3 files changed, 73 insertions(+), 34 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 0951350e2fe826..4c32f5a6635e7b 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -605,24 +605,18 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
SmallVector<int64_t> sortedDimGroup = dimGroup;
llvm::sort(sortedDimGroup);
int64_t dimBeforeCurrent = -1;
- for (auto dimIdx : sortedDimGroup) {
- // when order attr is present, adjacency dims are values like [3, 2, 1, 0]
- // in decreasing order; otherwise based on dim indices like [0, 1, 2, 3]
- // in increasing order
- if (dimBeforeCurrent >= 0) {
- if (getOrder() && !getOrder().empty()) {
- int64_t orderBefore = origOrder[dimBeforeCurrent];
- int64_t orderCurrent = origOrder[dimIdx];
- if (orderBefore != (orderCurrent - 1))
- llvm::report_fatal_error(
- "dimensions being collapsed must be adjacent in order");
- } else {
- if (dimIdx != (dimBeforeCurrent + 1))
- llvm::report_fatal_error(
- "dimensions being collapsed must be adjacent");
- }
- }
- dimBeforeCurrent = dimIdx;
+ // Dims must be adjacent in the walk order
+ bool hasExplicitWalkOrder = getOrder() && !getOrder().empty();
+ for (size_t dimIdx = 1; dimIdx < sortedDimGroup.size(); ++dimIdx) {
+ int64_t prev = sortedDimGroup[dimIdx - 1];
+ int64_t curr = sortedDimGroup[dimIdx];
+
+ if (hasExplicitWalkOrder) {
+ if (std::abs(origOrder[prev] - origOrder[curr]) != 1)
+ llvm::report_fatal_error(
+ "dimensions being collapsed must be adjacent in order");
+ } else if (curr - prev != 1)
+ llvm::report_fatal_error("dimensions being collapsed must be adjacent");
}
int firstDim = sortedDimGroup.front();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index d46269338888f0..5d5b62946b1ae4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2437,21 +2437,19 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
consumerLayout ? consumerLayout.getNumSubgroups() : numSg;
int consumerIdx = 0;
- // First pass: Match consumer's layout on non-reduction dimensions
+ // First pass: match the consumer's layout on non-reduction dims.
for (int i = 0; i < srcRank; i++) {
if (!llvm::is_contained(reductionDims, i) &&
consumerIdx < static_cast<int>(consumerSgLayout.size())) {
sgLayout[i] = consumerSgLayout[consumerIdx];
sgData[i] = consumerSgData[consumerIdx];
remainingSgCount /= sgLayout[i];
- order[i] = consumerOrder[consumerIdx];
consumerIdx++;
}
}
- // Second pass: Distribute remaining subgroups across reduction dimensions
- // the reduction to scalar case is handled only by this loop
- int64_t remainOrder = consumerSgLayout.size();
+ // Second pass: distribute remaining subgroups across the reduction dims.
+ // The reduction-to-scalar case is handled only by this loop.
for (int i = 0; i < srcRank; i++) {
if (llvm::is_contained(reductionDims, i)) {
sgLayout[i] =
@@ -2460,13 +2458,39 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
"source shape not divisible by sg_layout");
sgData[i] = srcShape[i] / sgLayout[i];
remainingSgCount /= sgLayout[i];
- order[i] = remainOrder++;
}
}
- DenseI32ArrayAttr resOrderAttr = DenseI32ArrayAttr::get(
- context, SmallVector<int32_t>(order.begin(), order.end()));
- if (!orderAttr || orderAttr.empty())
- resOrderAttr = nullptr;
+
+ DenseI32ArrayAttr resOrderAttr = nullptr;
+ int numRetainedDims = srcRank - static_cast<int>(reductionDims.size());
+ if (orderAttr && !orderAttr.empty() &&
+ static_cast<int>(consumerOrder.size()) == numRetainedDims) {
+ // Fill known order
+ int retainedRank = 0;
+ for (int dim = 0; dim < srcRank; dim++)
+ if (!llvm::is_contained(reductionDims, dim))
+ order[dim] = consumerOrder[retainedRank++];
+ // Each new order is placed as close as possible to its neighbors,
+ // row-major. Re-index the order of higher walk orders.
+ for (int reducedDim = 0; reducedDim < srcRank; reducedDim++) {
+ if (!llvm::is_contained(reductionDims, reducedDim))
+ continue;
+ int64_t insertRank = 0;
+ // Add as an inner dim to the outer neighbor
+ if (reducedDim)
+ insertRank = order[reducedDim - 1];
+ // Add as an outer dim to the inner neighbor if no left neighbor
+ // exists.
+ else if (reducedDim + 1 < srcRank)
+ insertRank = order[reducedDim + 1] + 1;
+ for (int otherDim = 0; otherDim < srcRank; otherDim++)
+ if (order[otherDim] >= insertRank)
+ order[otherDim]++;
+ order[reducedDim] = insertRank;
+ }
+ resOrderAttr = DenseI32ArrayAttr::get(
+ context, SmallVector<int32_t>(order.begin(), order.end()));
+ }
assert(remainingSgCount == 1 && "not all subgroups distributed");
srcLayout = buildLayout(context, sgLayout, sgData,
/*instData=*/{}, /*laneLayout=*/{},
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index a48d34072016f8..3dd784f9719e54 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -447,19 +447,19 @@ gpu.module @test {
// CHECK-LABEL: sum_reduction_broadcast_transpose
// CHECK-SAME: %[[ARG0:.*]]: memref<64x256xf16>
// CHECK-SAME: %[[ARG1:.*]]: memref<64x32xf16>
- // CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>}
+ // CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [1, 0]>, dims = [1]>}
// CHECK-SAME: dense<0.000000e+00> : vector<64xf16>
// CHECK: %[[TDESC0:.*]] = xegpu.create_nd_tdesc %[[ARG0]] : memref<64x256xf16> -> !xegpu.tensor_desc<64x32xf16,
- // CHECK-SAME: #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>>
+ // CHECK-SAME: #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [1, 0]>>
// CHECK: %[[FOR:.*]] = scf.for %[[ARG2:.*]] = %{{.*}} iter_args(%[[ARG3:.*]] = %[[CST]]) -> (vector<64xf16>) {
// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TDESC0]][0, %[[ARG2]]]
- // CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>}> :
- // CHECK-SAME: !xegpu.tensor_desc<64x32xf16, #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>> -> vector<64x32xf16>
+ // CHECK-SAME: <{layout = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [1, 0]>}> :
+ // CHECK-SAME: !xegpu.tensor_desc<64x32xf16, #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [1, 0]>> -> vector<64x32xf16>
// CHECK: %[[REDUCE:.*]] = vector.multi_reduction <add>, %[[LOAD]], %[[ARG3]]
- // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>} [1]
+ // CHECK-SAME: {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [1, 0]>, dims = [1]>} [1]
// CHECK-SAME: : vector<64x32xf16> to vector<64xf16>
// CHECK: scf.yield %[[REDUCE]] : vector<64xf16>
- // CHECK: } {layout_operand_3 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>, layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [0, 1]>, dims = [1]>}
+ // CHECK: } {layout_operand_3 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [1, 0]>, dims = [1]>, layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16], order = [1, 0]>, dims = [1]>}
// CHECK: %[[BCAST:.*]] = vector.broadcast %[[FOR]] {layout_result_0 = #xegpu.layout<sg_layout = [2, 4], sg_data = [16, 16], order = [0, 1]>}
// CHECK-SAME: : vector<64xf16> to vector<32x64xf16>
// CHECK: %[[TRANS:.*]] = vector.transpose %[[BCAST]], [1, 0] {layout_result_0 = #xegpu.layout<sg_layout = [4, 2], sg_data = [16, 16]>}
@@ -487,6 +487,27 @@ gpu.module @test {
}
}
+// -----
+gpu.module @test {
+ // CHECK-LABEL: reduce_split_transpose
+ // CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 1, 1], sg_data = [8, 8, 8], order = [0, 2, 1]>, dims = [2]>}
+ // CHECK: xegpu.create_nd_tdesc {{.*}} -> !xegpu.tensor_desc<64x64xf16, #xegpu.layout<sg_layout = [8, 1], sg_data = [8, 64], order = [0, 1]>>
+ // CHECK: %[[SPLIT:.*]] = vector.shape_cast %{{.*}} {layout_result_0 = #xegpu.layout<sg_layout = [8, 1, 1], sg_data = [8, 8, 8], order = [0, 2, 1]>} : vector<64x64xf16> to vector<64x8x8xf16>
+ // CHECK: vector.multi_reduction <maximumf>, %[[SPLIT]], %[[CST]] {layout_result_0 = #xegpu.slice<#xegpu.layout<sg_layout = [8, 1, 1], sg_data = [8, 8, 8], order = [0, 2, 1]>, dims = [2]>} [2] : vector<64x8x8xf16> to vector<64x8xf16>
+ // CHECK: vector.transpose %{{.*}}, [1, 0] {layout_result_0 = #xegpu.layout<sg_layout = [1, 8], sg_data = [8, 8]>} : vector<64x8xf16> to vector<8x64xf16>
+ gpu.func @reduce_split_transpose(%input: memref<64x64xf16>, %output: memref<8x64xf16>) kernel attributes {known_block_size = array<i32: 1, 32, 16>} {
+ %cst = arith.constant dense<0.000000e+00> : vector<64x8xf16>
+ %tdesc_input = xegpu.create_nd_tdesc %input : memref<64x64xf16> -> !xegpu.tensor_desc<64x64xf16>
+ %loaded = xegpu.load_nd %tdesc_input[0, 0] : !xegpu.tensor_desc<64x64xf16> -> vector<64x64xf16>
+ %split = vector.shape_cast %loaded : vector<64x64xf16> to vector<64x8x8xf16>
+ %reduced = vector.multi_reduction <maximumf>, %split, %cst [2] : vector<64x8x8xf16> to vector<64x8xf16>
+ %transposed = vector.transpose %reduced, [1, 0] : vector<64x8xf16> to vector<8x64xf16>
+ %tdesc_output = xegpu.create_nd_tdesc %output : memref<8x64xf16> -> !xegpu.tensor_desc<8x64xf16>
+ xegpu.store_nd %transposed, %tdesc_output[0, 0] <{layout = #xegpu.layout<sg_layout = [1, 8], sg_data = [8, 8]>}> : vector<8x64xf16>, !xegpu.tensor_desc<8x64xf16>
+ gpu.return
+ }
+}
+
// -----
// shape_cast collapse: all src dims fold into a single innermost dst dim.
// srcShape=[8, 16, 32], resShape=[4096], consumer sg_layout=[8],
More information about the Mlir-commits
mailing list