https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/223077
>From f37ae881697622c44f0926d2ceb58c3b938f339e Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 11 Sep 2026 20:41:14 +0000
Subject: [PATCH 1/2] [mlir][xegpu] Fix lane-local classification for
packed-lane-data reductions
isReductionLaneLocal decided whether a subgroup multi_reduction is
lane-local by comparing the result vector type to its distributed type.
That heuristic misclassifies degenerate layouts where the reduction
dimension is packed entirely into lane_data (lane_layout == 1 along
every dim): the result type equals its distributed type, so the check
returned false and routed a genuinely lane-local reduction through the
cross-lane shuffle path. This double-counted the already-local data
(e.g. a 16-element reduction produced a 16x overcount) in cross-subgroup
SLM reductions.
Key the decision off the source layout's lane_layout along the reduction
dimension instead: the reduction is lane-local exactly when that count
is 1. Fall back to the previous result-type comparison when no usable
source lane_layout is available.
Adds a unit test for a subgroup multi_reduction whose reduction dim is
packed entirely into lane_data: it must lower to a lane-local
vector.reduction with no cross-lane gpu.shuffle.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 21 ++++++++++-------
.../XeGPU/sg-to-lane-distribute-unit.mlir | 23 +++++++++++++++++++
2 files changed, 36 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index bc85140f9f121..bd3a4269bdd03 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -85,18 +85,23 @@ static bool isValidSubgroupMultiReductionOp(vector::MultiDimReductionOp op) {
return op.getReductionDims().size() == 1;
}
-/// A vector::MultiDimReductionOp is doing lane-local reduction if each lane
-/// is doing its own local reduction. In this case the result layout ensures
-/// that result vector is distributed to lanes, i.e. the result vector type is
-/// different from the distributed result vector type.
+/// A vector::MultiDimReductionOp reduces lane-locally when no data is combined
+/// across lanes, which holds exactly when the source's lane_layout along the
+/// reduction dimension is 1: every reduced element then lives in one lane.
static bool isReductionLaneLocal(vector::MultiDimReductionOp op) {
// Must be valid MultiDimReductionOp.
assert(isValidSubgroupMultiReductionOp(op) && "Expecting a valid subgroup "
"MultiDimReductionOp");
- auto resLayout = xegpu::getTemporaryLayout(op->getOpResult(0));
- VectorType resTy = dyn_cast<VectorType>(op.getType());
- auto resDistTypeOrFailure = getDistVecTypeBasedOnLaneLayout(resLayout, resTy);
- return resTy != resDistTypeOrFailure.value();
+ auto srcLayout = xegpu::getTemporaryLayout(op->getOpOperand(0));
+ ArrayRef<int64_t> reductionDims = op.getReductionDims();
+ assert(reductionDims.size() == 1 &&
+ "Expecting single reduction dimension for subgroup multi "
+ "reduction op");
+ int64_t reductionDim = reductionDims[0];
+ SmallVector<int64_t> srcLaneLayout = srcLayout.getEffectiveLaneLayoutAsInt();
+ assert(reductionDim < static_cast<int64_t>(srcLaneLayout.size()) &&
+ "Expecting a source lane_layout covering the reduction dimension");
+ return srcLaneLayout[reductionDim] == 1;
}
/// Given a vector type and its distributed vector type, return the list of
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index c193d83361e04..cc3f0355eeccb 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -740,6 +740,29 @@ gpu.func @vector_multi_reduction_3d_leading_unit_dim_cross_lane() {
gpu.return
}
+// lane_data packs all 16 reduced elements of dim1 into a single lane and
+// lane_layout is 1 everywhere, so the reduction is lane-local and must lower to
+// a plain vector.reduction. Classifying it by result type would route it
+// through butterfly shuffles and count the data 16 times over.
+// CHECK-LABEL: gpu.func @vector_multi_reduction_3d_packed_lane_data_lane_local
+// CHECK: %[[F0:.*]] = vector.shape_cast %{{.*}} : vector<1x16x1xf32> to vector<16xf32>
+// CHECK: %[[A0:.*]] = vector.extract %{{.*}}[0, 0] : f32 from vector<1x1xf32>
+// CHECK: %[[R0:.*]] = vector.reduction <add>, %[[F0]], %[[A0]] : vector<16xf32> into f32
+// CHECK: vector.insert %[[R0]], %{{.*}} [0, 0] : f32 into vector<1x1xf32>
+// CHECK-NOT: gpu.shuffle
+// CHECK: gpu.return
+gpu.func @vector_multi_reduction_3d_packed_lane_data_lane_local() {
+ %src = arith.constant dense<0.0> : vector<1x16x1xf32>
+ %acc = arith.constant dense<0.0> : vector<1x1xf32>
+ %1 = vector.multi_reduction <add>, %src, %acc
+ [1] : vector<1x16x1xf32> to vector<1x1xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 1], lane_data = [1, 16, 1]>, dims = [1]>
+ }> : vector<1x1xf32>
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @vector_extract_from_2d
// CHECK: %[[EXT:.*]] = vector.extract %{{.*}}[0] : vector<1xf32> from vector<4x1xf32>
gpu.func @vector_extract_from_2d() {
>From bb1cbd6d830874415a2a8a507a941617ede42251 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 18 Sep 2026 01:47:07 +0000
Subject: [PATCH 2/2] fix test comment
---
mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir | 3 +--
1 file changed, 1 insertion(+), 2 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 2ae33188df55c..1407acf9ee3dc 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -706,8 +706,7 @@ gpu.func @vector_multi_reduction_3d_leading_unit_dim_cross_lane() {
// lane_data packs all 16 reduced elements of dim1 into a single lane and
// lane_layout is 1 everywhere, so the reduction is lane-local and must lower to
-// a plain vector.reduction. Classifying it by result type would route it
-// through butterfly shuffles and count the data 16 times over.
+// a plain vector.reduction.
// CHECK-LABEL: gpu.func @vector_multi_reduction_3d_packed_lane_data_lane_local
// CHECK: %[[F0:.*]] = vector.shape_cast %{{.*}} : vector<1x16x1xf32> to vector<16xf32>
// CHECK: %[[A0:.*]] = vector.extract %{{.*}}[0, 0] : f32 from vector<1x1xf32>