[Mlir-commits] [mlir] [mlir][xegpu] Raise the reduction vector size to match maxReduceVectorSize (PR #225122)
Jianhui Li
llvmlistbot at llvm.org
Mon Sep 21 08:50:29 PDT 2026
https://github.com/Jianhui-Li created https://github.com/llvm/llvm-project/pull/225122
`maxReduceVectorSize` was pinned at 1, so a lane reduced one element at a time no
matter how much data it held. This raises it to 16 and makes the layout
computation actually spend it, so a lane's reduce vector matches the per-lane
budget.
```cpp
- int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
+ int64_t maxReduceVectorSize = 16; // could extend to spirv vector Size
```
>From df8ac1580e107fca33d595a880f1a0f9eb0c2313 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 19 Sep 2026 22:24:09 +0000
Subject: [PATCH 1/3] [mlir][xegpu] Raise maxReduceVectorSize and derive
reduction lane layouts
Bump maxReduceVectorSize from 1 to 16 so a multi_reduction's source layout can
pack several elements per lane instead of one, and derive that layout from a
single rule instead of two competing policies.
The previous code chose between a consumer-driven path that spent spare lanes on
the reduction dim and a fallback that kept the reduction inside a lane. Which
one ran depended on the op's shape at the time, so XeGPUPeepHoleOptimizer
splitting a 2D reduction changed the answer between propagation runs and left an
unlowerable convert_layout behind. computeReductionLaneLayoutAndData now applies
one rule:
* lane_layout: lanes live only in the innermost two dims. Surviving dims take
the consumer's lane_layout; leftover lanes go to the innermost dim when it
is reduced. Lanes are not claimed further out, because that is where a
producer spreads its own and a redistribution has no lowering. A consumer
already spending the whole subgroup therefore leaves the reduction dim
lane-less and the reduction lane-local, so [16, 1] falls out of the rule
rather than from a separate policy.
* lane_data: maxReduceVectorSize is a per-lane budget shared across dims. The
consumer's lane_data counts against it and only the remainder is placed, on
a lane-less unpacked dim, preferring the one just outside the innermost two
-- the scale computation shape, where those two come from a shape_cast that
split a single dim.
* A consumer sliced over exactly the reduction dims already describes every
source dim, so its grid and packing are reused directly.
inst_data is now lane_layout * lane_data on every dim, with no exception; the
previous fold of the reduce-vector size into inst_data on an arbitrary
lane-less dim is gone.
Also clamp lane_data in setupGenericLoadAnchorLayout, shared by load_gather and
load_matrix. Those read one contiguous span per lane, so a lane can only pack
along the innermost dim and only as far as the op's chunk size allows. A
consumer requesting lane_data > 1 on an outer, strided dim is not lowerable:
load_matrix becomes a single contiguous llvm.load, which would fetch adjacent
inner elements instead of the intended strided ones and silently drop rows in a
cross-subgroup reduction SLM round-trip.
Also initialize `intra` in XeGPUPeepHoleOptimizer's getReductionDimOrder. Only
`cross` was initialized, so when every reduction dim carries lanes the loop
never assigns `intra` and the caller's `== -1` guard cannot fire, indexing out
of bounds in accShape.erase.
Fixes seven imex integration failures (flash_attention_fwd and its two
array_len variants, four reduction cases), all of which failed to legalize
convert_layout and multi_reduction. All XeGPU and VectorToXeGPU tests pass.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 213 +++++++++++++-----
.../Transforms/XeGPUPeepHoleOptimizer.cpp | 2 +-
.../XeGPU/propagate-layout-inst-data.mlir | 53 +++--
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 26 +--
4 files changed, 202 insertions(+), 92 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 05b3f24d660b4..85a4fcdd69168 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1269,42 +1269,136 @@ compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
return {laneLayout, laneData, order};
}
-/// Computes the (lane_layout, lane_data) for a multi-reduction's source layout.
-/// Only the innermost two dims are distributed; leading dims are assumed unit.
-/// `subgroupSize` lanes go on one dim; up to `maxReduceVectorSize` elements are
-/// packed into lane_data on the other. To minimize cross-lane reduction, lanes
-/// are spread across a non-reduction dim when possible so the reduction happens
-/// within a lane. inst_data is the element-wise product lane_layout *
-/// lane_data.
+/// Computes the (lane_layout, lane_data, inst_data) for a multi-reduction's
+/// source layout, with inst_data = lane_layout * lane_data on every dim.
///
-/// e.g. with srcShape=[32, 128], subgroupSize=16, maxReduceVectorSize=2:
-/// - Switch: reductionDims=[1] and consumerReductionDims=[] -> lanes move
-/// to the non-reduction dim 0: lane_layout=[16, 1], lane_data=[1, 2].
-/// - Default: reductionDims=[0, 1] (both reduced) -> lanes stay on the
-/// innermost dim: lane_layout=[1, 16], lane_data=[2, 1].
-static std::pair<SmallVector<int64_t>, SmallVector<int64_t>>
-computeReductionLaneLayoutAndData(ArrayRef<int64_t> srcShape,
- ArrayRef<int64_t> reductionDims,
- int subgroupSize, int64_t maxReduceVectorSize,
- bool verticalLaneLayout = false) {
+/// Lanes live only in the innermost two dims. Surviving dims take the
+/// consumer's lane_layout; leftover lanes go to the innermost dim when it is
+/// reduced, so the reduction stays inside the subgroup. A consumer already
+/// spending the whole subgroup thus leaves the reduction dim lane-less and the
+/// reduction lane-local, which is where [16, 1] comes from.
+///
+/// maxReduceVectorSize is a per-lane budget shared across dims: the consumer's
+/// lane_data counts against it and only the remainder is placed here.
+///
+/// e.g. src=[16,32,32] reduce [2], consumer lane=[1,2] ->
+/// lane_layout=[1,2,8], lane_data=[16,1,1], inst_data=[16,2,8].
+/// e.g. src=[16,16] reduce [1], consumer lane=[16] ->
+/// lane_layout=[16,1], lane_data=[1,16], inst_data=[16,16].
+static std::tuple<SmallVector<int64_t>, SmallVector<int64_t>,
+ SmallVector<int64_t>>
+computeReductionLaneLayoutAndData(
+ ArrayRef<int64_t> srcShape, ArrayRef<int64_t> reductionDims,
+ int subgroupSize, int64_t maxReduceVectorSize,
+ xegpu::DistributeLayoutAttr consumerLayout = nullptr) {
int srcRank = srcShape.size();
SmallVector<int64_t> laneLayout(srcRank, 1), laneData(srcRank, 1);
+ auto isReduced = [&](int dim) {
+ return llvm::is_contained(reductionDims, dim);
+ };
- int innermost = srcRank - 1;
- int secondInnermost = srcRank - 2;
+ // A consumer sliced over exactly the reduction dims already describes every
+ // source dim, so reuse its grid and packing instead of re-deriving something
+ // that could disagree with the IR. The reduce-vector size is raised on top.
+ auto consumerSlice = dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+ if (consumerSlice &&
+ consumerSlice.getDims().asArrayRef().equals(reductionDims)) {
+ auto parent = consumerSlice.getParent();
+ SmallVector<int64_t> parentLaneLayout =
+ parent.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> parentLaneData = parent.getEffectiveLaneDataAsInt();
+ // Only trust a parent that spends the whole subgroup. One that does not
+ // distributes less than the producer already does.
+ int64_t parentLanes = 1;
+ for (int64_t l : parentLaneLayout)
+ parentLanes *= l;
+ if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
+ static_cast<int>(parentLaneData.size()) == srcRank &&
+ parentLanes == subgroupSize) {
+ // Raise, never lower, what the parent asks for.
+ for (int i = 0; i < srcRank; ++i)
+ if (parentLaneLayout[i] == 1 && srcShape[i] > 1) {
+ parentLaneData[i] = std::max(
+ parentLaneData[i], std::min(maxReduceVectorSize, srcShape[i]));
+ break;
+ }
+ SmallVector<int64_t> parentInstData(srcRank);
+ for (int i = 0; i < srcRank; ++i)
+ parentInstData[i] = parentLaneLayout[i] * parentLaneData[i];
+ return {parentLaneLayout, parentLaneData, parentInstData};
+ }
+ }
- if (verticalLaneLayout && secondInnermost >= 0) {
- std::swap(innermost, secondInnermost);
+ // Otherwise the consumer covers only the surviving dims, at reduced rank.
+ SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
+ if (consumerLayout) {
+ consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+ consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
}
- int laneDim = innermost;
- int vectorDim = secondInnermost; // negative for rank 1
- laneLayout[laneDim] =
- std::min(static_cast<int64_t>(subgroupSize), srcShape[laneDim]);
+ // Lanes are confined to the innermost two dims.
+ int firstLaneDim = std::max(0, srcRank - 2);
+
+ // Surviving dims inherit the consumer's grid where lanes may live, and its
+ // packing as the base. Lanes it wants on a leading dim cannot be honored;
+ // convert_layout reconciles that.
+ int64_t lanesUsed = 1;
+ for (int i = 0, c = 0; i < srcRank; ++i) {
+ if (isReduced(i))
+ continue;
+ int64_t consumerLanes = c < static_cast<int>(consumerLaneLayout.size())
+ ? consumerLaneLayout[c]
+ : 1;
+ if (c < static_cast<int>(consumerLaneData.size()))
+ laneData[i] = consumerLaneData[c];
+ ++c;
+ if (i < firstLaneDim)
+ continue;
+ laneLayout[i] = std::min<int64_t>(consumerLanes, srcShape[i]);
+ lanesUsed *= laneLayout[i];
+ }
+
+ // Leftover lanes reduce within the subgroup, and only on the innermost dim:
+ // that is where producers spread theirs, so claiming them further out needs a
+ // redistribution convert_layout cannot lower.
+ int64_t spareLanes = subgroupSize / lanesUsed;
+ int innermost = srcRank - 1;
+ if (spareLanes > 1 && isReduced(innermost) && srcShape[innermost] > 1)
+ laneLayout[innermost] = std::min(spareLanes, srcShape[innermost]);
+
+ // Spend what is left of the budget on a lane-less, unpacked dim, since a lane
+ // owns such a dim whole. Prefer the dim just outside the innermost two: in
+ // the scale computation those two come from a shape_cast that split one dim,
+ // so the dim above them is what should carry the per-lane vector.
+ int64_t packedByConsumer = 1;
+ for (int64_t d : laneData)
+ packedByConsumer *= d;
+ int64_t vectorBudget =
+ maxReduceVectorSize / std::min(packedByConsumer, maxReduceVectorSize);
+
+ int vectorDim = -1;
+ int outsideLaneDims = srcRank - 3;
+ auto usable = [&](int i) {
+ return laneLayout[i] == 1 && laneData[i] == 1 && srcShape[i] > 1;
+ };
+ if (vectorBudget > 1) {
+ if (outsideLaneDims >= 0 && usable(outsideLaneDims))
+ vectorDim = outsideLaneDims;
+ else
+ for (int i = srcRank - 1; i >= 0; --i)
+ if (usable(i)) {
+ vectorDim = i;
+ break;
+ }
+ }
if (vectorDim >= 0)
- laneData[vectorDim] = std::min(maxReduceVectorSize, srcShape[vectorDim]);
+ laneData[vectorDim] = std::min(vectorBudget, srcShape[vectorDim]);
- return {laneLayout, laneData};
+ SmallVector<int64_t> instData(srcRank);
+ for (int i = 0; i < srcRank; ++i)
+ instData[i] = laneLayout[i] * laneData[i];
+
+ return {laneLayout, laneData, instData};
}
//===----------------------------------------------------------------------===//
@@ -1767,7 +1861,8 @@ xegpu::setupPrefetchNdAnchorLayout(xegpu::LayoutKind layoutKind,
/// consumer layout (from its result's downstream uses) and validates it
/// against uArch constraints; if valid, the consumer's `inst_data` /
/// `sg_layout` are honored. Otherwise the helper falls back to defaults
-/// derived from uArch block parameters.
+/// derived from uArch block parameters. `lane_data` never comes from the
+/// consumer; see compute2DBlockIOLaneLayout.
xegpu::DistributeLayoutAttr
xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
VectorType resVecTy,
@@ -1925,6 +2020,20 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
laneLayout.assign(consumerLaneLayout.begin(), consumerLaneLayout.end());
laneData.assign(consumerLaneData.begin(), consumerLaneData.end());
+ // A matrix/gather load reads a single contiguous span of memory per lane, so
+ // a lane can only pack (lane_data > 1) elements along the innermost,
+ // contiguous dim. A consumer that requests lane_data > 1 on an outer, strided
+ // dim (e.g. a downstream reduction packing across mem_desc rows) is not
+ // lowerable: the contiguous load would fetch adjacent inner elements instead
+ // of the strided outer ones. Clamp the innermost dim to maxChunkSize and
+ // force outer dims to 1; a downstream convert_layout reconciles the
+ // difference with the consumer.
+ for (size_t i = 0; i + 1 < laneData.size(); ++i)
+ laneData[i] = 1;
+ if (!laneData.empty())
+ laneData.back() =
+ std::min(laneData.back(), static_cast<int64_t>(maxChunkSize));
+
if (layoutKind == xegpu::LayoutKind::InstData) {
SmallVector<int64_t> instData;
instData.resize(resShape.size());
@@ -2436,7 +2545,7 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
auto context = srcVecTy.getContext();
const int subgroupSize = uArch->getSubgroupSize();
- int64_t maxReduceVectorSize = 1; // could extend to spirv vector Size
+ int64_t maxReduceVectorSize = 16; // could extend to spirv vector Size
xegpu::DistributeLayoutAttr srcLayout;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
xegpu::SliceAttr consumerSliceLayout =
@@ -2506,39 +2615,14 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
/*laneData=*/{}, resOrderAttr);
}
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- xegpu::SliceAttr consumerSliceLayout =
- dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- auto consumerReductionDims =
- consumerSliceLayout
- ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
- : SmallVector<int64_t>({});
- // A[i] reduced from A[i, j] is stored out directly, use vertical Lane
- // layout like [16, 1]
- bool verticalLaneLayout = consumerReductionDims.empty() &&
- reductionDims.size() == 1 &&
- reductionDims[0] == (srcRank - 1);
- auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
- srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
- verticalLaneLayout);
- // inst_data is the per-instruction data, i.e. the element-wise product of
- // lane_layout and lane_data.
- SmallVector<int64_t> instData(srcRank);
- for (int i = 0; i < srcRank; i++)
- instData[i] = laneLayout[i] * laneData[i];
+ auto [laneLayout, laneData, instData] =
+ computeReductionLaneLayoutAndData(srcShape, reductionDims, subgroupSize,
+ maxReduceVectorSize, consumerLayout);
srcLayout =
buildInstDataLayoutWithLane(context, instData, laneLayout, laneData);
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- // Only the innermost two dimensions are distributed; all leading dimensions
- // are assumed to be unit dimensions.
- assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
- "Lane reduction layout assumes all leading (non-innermost-two) "
- "dimensions are unit dimensions");
xegpu::SliceAttr consumerSliceLayout =
dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
- auto consumerReductionDims =
- consumerSliceLayout
- ? SmallVector<int64_t>(consumerSliceLayout.getDims().asArrayRef())
- : SmallVector<int64_t>({});
if (consumerSliceLayout &&
consumerSliceLayout.getDims().asArrayRef().equals(reductionDims)) {
// at the lane level, the consumerSliceLayout can be directly reused
@@ -2546,12 +2630,17 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
// the layout is not consistent
srcLayout = consumerSliceLayout.getParent();
} else {
- bool verticalLaneLayout = consumerReductionDims.empty() &&
- reductionDims.size() == 1 &&
- reductionDims[0] == (srcRank - 1);
- auto [laneLayout, laneData] = computeReductionLaneLayoutAndData(
+ // The re-derivation path below only distributes the innermost two dims,
+ // so it requires the leading dims to be unit. The reused-consumer branch
+ // above has no such restriction (a consumer may legitimately keep the
+ // reduce-vector size as per-lane data on a non-unit leading dim).
+ assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
+ "Lane reduction layout assumes all leading (non-innermost-two) "
+ "dimensions are unit dimensions");
+ auto [laneLayout, laneData, instData] = computeReductionLaneLayoutAndData(
srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
- verticalLaneLayout);
+ consumerLayout);
+ (void)instData;
srcLayout = buildLaneLayout(context, laneLayout, laneData);
}
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
index c64defc471de8..5ce5491121400 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPeepHoleOptimizer.cpp
@@ -569,7 +569,7 @@ class MultiRed2dOpPattern
xegpu::DistributeLayoutAttr layout) const {
assert(layout.isForSubgroup() && "Must know the lane layout");
assert(reductionDims.size() == 2 && "Expected 2D reduction");
- int64_t intra, cross = -1;
+ int64_t intra = -1, cross = -1;
xegpu::LayoutAttr layoutAttr = dyn_cast<xegpu::LayoutAttr>(layout);
if (auto layoutSliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
layoutAttr =
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 841f06545eb19..d3bd493178c1f 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -162,9 +162,9 @@ func.func @load_matrix_with_coalesce_chunksize(%arg0: memref<2048x512xf32>, %arg
gpu.module @test {
// CHECK-LABEL: func.func @scatter_ops_coalesce_chunksize(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<512xf32>) {
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} dense<true> : vector<16x32xi1>
-// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} dense<12> : vector<16x32xindex>
-// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<true> : vector<16x32xi1>
+// CHECK: %{{.*}} = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>} dense<12> : vector<16x32xindex>
+// CHECK: %{{.*}} = xegpu.load %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}> :
// CHECK-SAME: memref<512xf32>, vector<16x32xindex>, vector<16x32xi1> -> vector<16x32xf32>
// CHECK: xegpu.store %0, %[[ARG0]][%{{.*}}], %{{.*}} <{layout = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>}> :
// CHECK-SAME: vector<16x32xf32>, memref<512xf32>, vector<16x32xindex>, vector<16x32xi1>
@@ -183,10 +183,10 @@ func.func @scatter_ops_coalesce_chunksize(%src: memref<512xf32>) {
gpu.module @test {
// CHECK-LABEL: func.func @load_gather_with_coalesce_chunksize(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>}
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 1], lane_layout = [16, 1], lane_data = [1, 1]>}
// CHECK-SAME: dense<0> : vector<16x16xindex>
-// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
-// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<inst_data = [16, 2], lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+// CHECK-NEXT: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16, 1], lane_layout = [16, 1], lane_data = [1, 1]>} dense<true> : vector<16x16xi1>
+// CHECK-NEXT: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<inst_data = [16, 1], lane_layout = [16, 1], lane_data = [1, 1]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
@@ -208,11 +208,11 @@ func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: m
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>}> : memref<1024xf16>, vector<1024xindex>, vector<1024xi1> -> vector<1024xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
-// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
-// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
-// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 16], lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
-// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>} : vector<1024xf16> to vector<8x8x16xf16>
+// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
+// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [2, 8, 16], lane_layout = [1, 1, 16], lane_data = [2, 8, 1]>, dims = [0]>} [0] : vector<8x8x16xf16> to vector<8x16xf16>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [8, 16], lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<1024xi1>
%0 = vector.step : vector<1024xindex>
@@ -232,7 +232,7 @@ func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<1024xf16>, %arg1
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size(
-// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 1], lane_layout = [1, 1, 1], lane_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x1xf16> to vector<1xf16>
+// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16, 1], lane_layout = [1, 1, 1], lane_data = [1, 16, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x1xf16> to vector<1xf16>
func.func @vector_2d_reduction_with_fractional_subgroup_size(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<16xi1>
%0 = vector.step : vector<16xindex>
@@ -250,7 +250,7 @@ func.func @vector_2d_reduction_with_fractional_subgroup_size(%arg0: memref<1024x
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(
-// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 1, 4], lane_layout = [1, 1, 4], lane_data = [1, 1, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x4xf16> to vector<1xf16>
+// CHECK: %[[ReduceVal:.*]] = vector.multi_reduction <add>, %[[Val:.*]], %[[CST:.*]] {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [1, 16, 4], lane_layout = [1, 1, 4], lane_data = [1, 16, 1]>, dims = [1, 2]>} [1, 2] : vector<1x16x4xf16> to vector<1xf16>
func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(%arg0: memref<1024xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<64xi1>
%0 = vector.step : vector<64xindex>
@@ -268,9 +268,9 @@ func.func @vector_2d_reduction_with_fractional_subgroup_size_1x4x1(%arg0: memref
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_and_merge(
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} dense<true> : vector<256xi1>
-// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} : vector<256xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>} dense<true> : vector<256xi1>
+// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>} : vector<256xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<inst_data = [16], lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
// CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<inst_data = [1, 1, 32], lane_layout = [1, 1, 16], lane_data = [1, 1, 2]>} : vector<256xf16> to vector<2x4x32xf16>
// CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<inst_data = [1, 32], lane_layout = [1, 16], lane_data = [1, 2]>} : vector<2x4x32xf16> to vector<1x256xf16>
// CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<inst_data = [32], lane_layout = [16], lane_data = [2]>} : vector<1x256xf16> to vector<256xf16>
@@ -700,3 +700,24 @@ gpu.module @test {
gpu.return
}
}
+
+// -----
+// A 32-bit block load cannot pack along the innermost dim, so the load_nd
+// anchor insists on unit lane_data even though its multi_reduction consumer
+// asks to pack 16 elements per lane into lane_data.
+gpu.module @test {
+// CHECK-LABEL: func.func @f32_reduction_from_load_nd(
+// CHECK: %[[TD:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32, #xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %[[TD]][0, 0] <{layout = #xegpu.layout<inst_data = [16, 8], lane_layout = [16, 1], lane_data = [1, 1], order = [0, 1]>}>
+// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[LOAD]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<inst_data = [16, 16], lane_layout = [16, 1], lane_data = [1, 16]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
+func.func @f32_reduction_from_load_nd(%arg0: memref<16x16xf32>, %arg1: memref<256xf32>) {
+ %cst = arith.constant dense<0.000000e+00> : vector<16xf32>
+ %mask = arith.constant dense<true> : vector<16xi1>
+ %offset = vector.step : vector<16xindex>
+ %t = xegpu.create_nd_tdesc %arg0 : memref<16x16xf32> -> !xegpu.tensor_desc<16x16xf32>
+ %v = xegpu.load_nd %t[0, 0] : !xegpu.tensor_desc<16x16xf32> -> vector<16x16xf32>
+ %0 = vector.multi_reduction <add>, %v, %cst [1] : vector<16x16xf32> to vector<16xf32>
+ xegpu.store %0, %arg1[%offset], %mask : vector<16xf32>, memref<256xf32>, vector<16xindex>, vector<16xi1>
+ return
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index d19c0ca17c0cf..b3fe76c4da17c 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -138,9 +138,9 @@ func.func @extf_truncf(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.tensor
gpu.module @test {
// CHECK-LABEL: func.func @load_gather_with_coalesce_chunksize(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: memref<8x16xf16>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf16>, %[[ARG2:[0-9a-zA-Z]+]]: memref<8x16xf32>) {
-// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} dense<0> : vector<16x16xindex>
-// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>} dense<true> : vector<16x16xi1>
-// CHECK: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 2]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
+// CHECK: %[[OFFSET:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} dense<0> : vector<16x16xindex>
+// CHECK: %[[MASK:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>} dense<true> : vector<16x16xi1>
+// CHECK: %{{.*}} = xegpu.load %arg1[%[[OFFSET]]], %[[MASK]] <{layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>}> : memref<256xf16>, vector<16x16xindex>, vector<16x16xi1> -> vector<16x16xf16>
func.func @load_gather_with_coalesce_chunksize(%arg0: memref<8x16xf16>, %arg1: memref<256xf16>, %arg2: memref<8x16xf32>) {
%c0 = arith.constant 0 : index
%0 = xegpu.create_nd_tdesc %arg0 : memref<8x16xf16> -> !xegpu.tensor_desc<8x16xf16>
@@ -517,7 +517,7 @@ func.func @if_multiple_uses(%arg0: !xegpu.tensor_desc<8x16xf16>, %arg1: !xegpu.t
gpu.module @test {
// CHECK-LABEL: func.func @vector_outer_reduction(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf32>) {
-// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<16x16xf32> to vector<16xf32>
+// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [16, 1]>, dims = [0]>} [0] : vector<16x16xf32> to vector<16xf32>
func.func @vector_outer_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf32>) {
%cst = arith.constant dense<0.000000e+00> : vector<16xf32>
%mask = arith.constant dense<true> : vector<16xi1>
@@ -531,7 +531,7 @@ func.func @vector_outer_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf3
gpu.module @test {
// CHECK-LABEL: func.func @vector_inner_reduction(
// CHECK-SAME: %[[ARG0:[0-9a-zA-Z]+]]: vector<16x16xf32>, %[[ARG1:[0-9a-zA-Z]+]]: memref<256xf32>) {
-// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
+// CHECK: %{{.*}} = vector.multi_reduction <add>, %[[ARG0]], %{{.*}} {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1], lane_data = [1, 16]>, dims = [1]>} [1] : vector<16x16xf32> to vector<16xf32>
func.func @vector_inner_reduction(%arg0: vector<16x16xf32>, %arg1: memref<256xf32>) {
%cst = arith.constant dense<0.000000e+00> : vector<16xf32>
%mask = arith.constant dense<true> : vector<16xi1>
@@ -708,11 +708,11 @@ gpu.module @test{
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_non_unit_dims(
// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP:.*]]], %[[CST:.*]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<128xf16>, vector<128xindex>, vector<128xi1> -> vector<128xf16>
-// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>} : vector<128xf16> to vector<1x8x16xf16>
-// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
-// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
-// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0]>} [0] : vector<1x8x16xf16> to vector<8x16xf16>
-// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>} : vector<128xf16> to vector<1x8x16xf16>
+// CHECK: %[[CST_0:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<8x16xf16>
+// CHECK: %[[CST_1:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} dense<0.000000e+00> : vector<16xf16>
+// CHECK: %[[REDUCE_0:.*]] = vector.multi_reduction <add>, %[[CAST]], %[[CST_0]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 8, 1]>, dims = [0]>} [0] : vector<1x8x16xf16> to vector<8x16xf16>
+// CHECK: %[[REDUCE_1:.*]] = vector.multi_reduction <add>, %[[REDUCE_0]], %[[CST_1]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, dims = [0]>} [0] : vector<8x16xf16> to vector<16xf16>
func.func @vector_shape_cast_expand_non_unit_dims(%arg0: memref<128xf16>, %arg1: memref<16xf16>) {
%cst = arith.constant dense<true> : vector<128xi1>
%0 = vector.step : vector<128xindex>
@@ -862,9 +862,9 @@ func.func @vector_reduction_broadcast_transpose(%arg0: memref<1024x64xf32>, %arg
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_shape_cast_expand_and_merge(
-// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} dense<true> : vector<256xi1>
-// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<256xindex>
-// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [2]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<256xi1>
+// CHECK: %[[STEP:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<256xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[STEP]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<256xf16>, vector<256xindex>, vector<256xi1> -> vector<256xf16>
// CHECK: %[[CAST_0:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 2]>} : vector<256xf16> to vector<2x4x32xf16>
// CHECK: %[[CAST_1:.*]] = vector.shape_cast %[[CAST_0]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>} : vector<2x4x32xf16> to vector<1x256xf16>
// CHECK: %[[CAST_2:.*]] = vector.shape_cast %[[CAST_1]] {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [2]>} : vector<1x256xf16> to vector<256xf16>
>From 63126b2ace43572386482823ed65e858c20c2c1a Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 20 Sep 2026 03:38:16 +0000
Subject: [PATCH 2/3] [mlir][xegpu] Bound load and reduction layouts by what
the shape can carry
Three places let a consumer's layout through unchecked, each producing a layout
the value it is attached to cannot distribute.
setupLoadNdAnchorLayout inferred the block variant from the consumer's
lane_data: a 32-bit consumer packing rows made it look for a VNNI transform
block, which exists only below 32 bits. getBlockWidthHeightCount then returned
nothing, the anchor gave up with a null layout, and the load kept the consumer's
lane_data unfiltered -- which requirePacked turned into pack_register on a
32-bit blockload2d, rejected by the XeVM verifier. Gate the transform variant on
the element being narrower than the general packed format size, so a wide load
takes the regular variant and derives unit lane_data as it should.
setupGenericLoadAnchorLayout copied the consumer's lane_layout verbatim. It
already receives resShape, so bound the lane count by it: a cross-subgroup
reduction reloading a 16x4 slab from SLM asked for the 16 lanes it uses on the
wider pre-SLM source, and the load_matrix verifier rejected 16 lanes over 4
columns.
computeReductionLaneLayoutAndData reused a slice consumer's parent layout whole.
A cross-subgroup reduction shares one result layout across both of its stages,
so that parent may describe the 16x32 source of the first stage rather than the
16x4 slab reduced by the second. Clamp the reused lane_layout to what this
source supplies; the parent's lane_data is still right, only the lane count was
borrowed, so clamping keeps agreement with the consumer where deriving afresh
would lose it.
Together these fix the remaining imex integration failures:
softmax_intra_sg_cross_lane and flash_attention_kernel_from_Tile_IR. All nine
previously failing cases now lower, and the full MLIR test suite is unchanged.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 28 ++++++++++++++++++-
1 file changed, 27 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 85a4fcdd69168..e3291e1fb756c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1315,6 +1315,15 @@ computeReductionLaneLayoutAndData(
if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
static_cast<int>(parentLaneData.size()) == srcRank &&
parentLanes == subgroupSize) {
+ // Lanes cannot exceed what this source supplies. A cross-subgroup
+ // reduction reuses one result layout for both of its stages, so the
+ // parent may describe the wider pre-SLM source (16x32) rather than the
+ // slab being reduced here (16x4); spreading 16 lanes over 4 columns is
+ // not something the shape can carry.
+ for (int i = 0; i < srcRank; ++i)
+ parentLaneLayout[i] =
+ std::min(parentLaneLayout[i],
+ std::max<int64_t>(1, srcShape[i] / parentLaneData[i]));
// Raise, never lower, what the parent asks for.
for (int i = 0; i < srcRank; ++i)
if (parentLaneLayout[i] == 1 && srcShape[i] > 1) {
@@ -1903,7 +1912,16 @@ xegpu::setupLoadNdAnchorLayout(xegpu::LayoutKind layoutKind,
// attr
bool hasTranspose =
consumerLaneLayout[rank - 2] > 1 && consumerLaneLayout[rank - 1] == 1;
- bool hasTransform = !hasTranspose && consumerLaneData[rank - 2] > 1 &&
+ // VNNI exists only below 32 bits, so a wider consumer packing rows into
+ // lane_data is not asking for a transform. Reading it as one sends us looking
+ // for a block variant the hardware does not have, and the whole layout is
+ // then abandoned -- leaving the consumer's lane_data on the load, which is
+ // what pack_register would wrongly be set from.
+ bool canTransform =
+ elemTy.isIntOrFloat() &&
+ elemTy.getIntOrFloatBitWidth() < uArch->getGeneralPackedFormatBitSize();
+ bool hasTransform = !hasTranspose && canTransform &&
+ consumerLaneData[rank - 2] > 1 &&
consumerLaneData[rank - 1] == 1;
assert((consumerLaneData[rank - 2] == 1 || consumerLaneData[rank - 1] == 1) &&
"Expected consumer lane data to have at most one non-unit dim");
@@ -2034,6 +2052,14 @@ static xegpu::DistributeLayoutAttr setupGenericLoadAnchorLayout(
laneData.back() =
std::min(laneData.back(), static_cast<int64_t>(maxChunkSize));
+ // Lanes are bounded by what this load's own tile can supply. A consumer may
+ // distribute a dim more finely than the tile is wide -- a cross-subgroup
+ // reduction reloading a 16x4 slab wants the 16 lanes it uses elsewhere -- and
+ // keeping that would leave the layout unable to distribute the result.
+ for (size_t i = 0; i < laneLayout.size() && i < resShape.size(); ++i)
+ laneLayout[i] = std::min(laneLayout[i],
+ std::max<int64_t>(1, resShape[i] / laneData[i]));
+
if (layoutKind == xegpu::LayoutKind::InstData) {
SmallVector<int64_t> instData;
instData.resize(resShape.size());
>From 76d9e765dcc3549c3381baf91f08ca52f6635afa Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Mon, 21 Sep 2026 15:35:19 +0000
Subject: [PATCH 3/3] [mlir][xegpu] Restructure reduction lane layout around
three rules
computeReductionLaneLayoutAndData had two paths that each re-derived the
whole layout, so the shape bound and the reduce-vector placement were
written twice with slightly different meanings. Restate it as three
sequential rules -- take the layout from the consumer, bound lane_layout
by the source shape, then spend the vector budget -- with the reuse and
derive cases differing only in rule 1.
This drops the restriction that lanes live only in the innermost two
dims. A consumer's lane_layout is now honored on whatever dim it lands
on, leading dims included. Reducing src [16, 1, 4] over dim 2 previously
discarded the consumer's 16 lanes on dim 0, reassigned 4 of them to the
reduced dim, and left the feeding gather on lane_layout [4] -- 4 lanes of
16, and a cross-lane combine where a lane-local one was available. It now
yields lane_layout [16, 1, 1]. The assert requiring leading dims to be
unit guarded exactly this case and goes away with it.
Rule 3 also multiplies the chosen dim's lane_data up to the budget
instead of overwriting it, so a dim the consumer already packs can absorb
the remaining budget rather than being skipped.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 191 +++++++-----------
mlir/test/Dialect/XeGPU/propagate-layout.mlir | 29 +++
2 files changed, 106 insertions(+), 114 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index e3291e1fb756c..657bcab4248b4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -435,17 +435,6 @@ void xegpu::removeTemporaryLayoutAttrs(Operation *op) {
});
}
-/// Returns true if every dimension of `shape` except the innermost
-/// `numInnerDims` is a unit (size-1) dimension.
-[[maybe_unused]] static bool leadingDimsAreUnit(ArrayRef<int64_t> shape,
- int numInnerDims) {
- int numLeading = static_cast<int>(shape.size()) - numInnerDims;
- if (numLeading <= 0)
- return true;
- return llvm::all_of(shape.take_front(numLeading),
- [](int64_t dim) { return dim == 1; });
-}
-
static xegpu::LayoutAttr buildInstDataLayoutWithLane(
mlir::MLIRContext *context, ArrayRef<int64_t> instData,
ArrayRef<int64_t> laneLayout, ArrayRef<int64_t> laneData,
@@ -1269,17 +1258,27 @@ compute2DBlockIOLaneLayout(ArrayRef<int64_t> instShape, int64_t subgroupSize,
return {laneLayout, laneData, order};
}
-/// Computes the (lane_layout, lane_data, inst_data) for a multi-reduction's
-/// source layout, with inst_data = lane_layout * lane_data on every dim.
+/// Computes a multi_reduction's source layout as (lane_layout, lane_data,
+/// inst_data), where inst_data = lane_layout * lane_data on every dim.
+///
+/// `maxReduceVectorSize` is the per-lane vector budget.
///
-/// Lanes live only in the innermost two dims. Surviving dims take the
-/// consumer's lane_layout; leftover lanes go to the innermost dim when it is
-/// reduced, so the reduction stays inside the subgroup. A consumer already
-/// spending the whole subgroup thus leaves the reduction dim lane-less and the
-/// reduction lane-local, which is where [16, 1] comes from.
+/// 1. Take lane_layout and lane_data from the consumer. A consumer sliced over
+/// exactly `reductionDims` covers every source dim, so its parent is reused
+/// as-is. Otherwise it covers only the surviving dims, one entry each, and
+/// each entry is honored on its own dim -- leading dims included, which is
+/// what keeps a reduction over a trailing dim lane-local. Lanes the consumer
+/// leaves unclaimed go to the innermost reduced dim rather than sit idle.
///
-/// maxReduceVectorSize is a per-lane budget shared across dims: the consumer's
-/// lane_data counts against it and only the remainder is placed here.
+/// 2. Limit lane_layout by `srcShape`: a dim can spread over no more lanes than
+/// it has elements to give.
+///
+/// 3. Raise lane_data until a lane packs `maxReduceVectorSize` elements, on one
+/// undistributed dim -- lane_layout 1, so every lane already spans its full
+/// extent and only inst_data widens. What lane_data packs elsewhere comes off
+/// the budget first, and the chosen dim is multiplied up, never overwritten.
+/// Prefer the dim just outside the innermost two, else the innermost
+/// undistributed dim.
///
/// e.g. src=[16,32,32] reduce [2], consumer lane=[1,2] ->
/// lane_layout=[1,2,8], lane_data=[16,1,1], inst_data=[16,2,8].
@@ -1297,111 +1296,85 @@ computeReductionLaneLayoutAndData(
return llvm::is_contained(reductionDims, dim);
};
- // A consumer sliced over exactly the reduction dims already describes every
- // source dim, so reuse its grid and packing instead of re-deriving something
- // that could disagree with the IR. The reduce-vector size is raised on top.
+ // Rule 1. Only a parent that covers the source rank and spends the whole
+ // subgroup is trusted; one spending less distributes less than the producer
+ // already does.
auto consumerSlice = dyn_cast_if_present<xegpu::SliceAttr>(consumerLayout);
+ bool reusedParent = false;
if (consumerSlice &&
consumerSlice.getDims().asArrayRef().equals(reductionDims)) {
auto parent = consumerSlice.getParent();
SmallVector<int64_t> parentLaneLayout =
parent.getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> parentLaneData = parent.getEffectiveLaneDataAsInt();
- // Only trust a parent that spends the whole subgroup. One that does not
- // distributes less than the producer already does.
- int64_t parentLanes = 1;
- for (int64_t l : parentLaneLayout)
- parentLanes *= l;
if (static_cast<int>(parentLaneLayout.size()) == srcRank &&
static_cast<int>(parentLaneData.size()) == srcRank &&
- parentLanes == subgroupSize) {
- // Lanes cannot exceed what this source supplies. A cross-subgroup
- // reduction reuses one result layout for both of its stages, so the
- // parent may describe the wider pre-SLM source (16x32) rather than the
- // slab being reduced here (16x4); spreading 16 lanes over 4 columns is
- // not something the shape can carry.
- for (int i = 0; i < srcRank; ++i)
- parentLaneLayout[i] =
- std::min(parentLaneLayout[i],
- std::max<int64_t>(1, srcShape[i] / parentLaneData[i]));
- // Raise, never lower, what the parent asks for.
- for (int i = 0; i < srcRank; ++i)
- if (parentLaneLayout[i] == 1 && srcShape[i] > 1) {
- parentLaneData[i] = std::max(
- parentLaneData[i], std::min(maxReduceVectorSize, srcShape[i]));
- break;
- }
- SmallVector<int64_t> parentInstData(srcRank);
- for (int i = 0; i < srcRank; ++i)
- parentInstData[i] = parentLaneLayout[i] * parentLaneData[i];
- return {parentLaneLayout, parentLaneData, parentInstData};
+ computeProduct(parentLaneLayout) == subgroupSize) {
+ laneLayout = std::move(parentLaneLayout);
+ laneData = std::move(parentLaneData);
+ reusedParent = true;
}
}
- // Otherwise the consumer covers only the surviving dims, at reduced rank.
- SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
- if (consumerLayout) {
- consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
- consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+ if (!reusedParent) {
+ SmallVector<int64_t> consumerLaneLayout, consumerLaneData;
+ if (consumerLayout) {
+ consumerLaneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+ consumerLaneData = consumerLayout.getEffectiveLaneDataAsInt();
+ }
+ int64_t lanesUsed = 1;
+ for (int i = 0, c = 0; i < srcRank; ++i) {
+ if (isReduced(i))
+ continue;
+ if (c < static_cast<int>(consumerLaneLayout.size())) {
+ laneLayout[i] = consumerLaneLayout[c];
+ lanesUsed *= laneLayout[i];
+ }
+ if (c < static_cast<int>(consumerLaneData.size()))
+ laneData[i] = consumerLaneData[c];
+ ++c;
+ }
+
+ // The innermost reduced dim, not an outer one: claiming lanes further out
+ // needs a redistribution convert_layout cannot lower.
+ int64_t spareLanes = std::max<int64_t>(1, subgroupSize / lanesUsed);
+ int innermost = srcRank - 1;
+ if (spareLanes > 1 && isReduced(innermost))
+ laneLayout[innermost] = spareLanes;
}
- // Lanes are confined to the innermost two dims.
- int firstLaneDim = std::max(0, srcRank - 2);
+ // Rule 2. This also bounds a reused parent, which may describe a wider source
+ // than the one reduced here: a cross-subgroup reduction shares one result
+ // layout across both stages, so the parent can ask for 16 lanes on a dim this
+ // slab has only 4 columns of.
+ for (int i = 0; i < srcRank; ++i)
+ laneLayout[i] = std::min(laneLayout[i],
+ std::max<int64_t>(1, srcShape[i] / laneData[i]));
- // Surviving dims inherit the consumer's grid where lanes may live, and its
- // packing as the base. Lanes it wants on a leading dim cannot be honored;
- // convert_layout reconciles that.
- int64_t lanesUsed = 1;
- for (int i = 0, c = 0; i < srcRank; ++i) {
- if (isReduced(i))
- continue;
- int64_t consumerLanes = c < static_cast<int>(consumerLaneLayout.size())
- ? consumerLaneLayout[c]
- : 1;
- if (c < static_cast<int>(consumerLaneData.size()))
- laneData[i] = consumerLaneData[c];
- ++c;
- if (i < firstLaneDim)
- continue;
- laneLayout[i] = std::min<int64_t>(consumerLanes, srcShape[i]);
- lanesUsed *= laneLayout[i];
- }
-
- // Leftover lanes reduce within the subgroup, and only on the innermost dim:
- // that is where producers spread theirs, so claiming them further out needs a
- // redistribution convert_layout cannot lower.
- int64_t spareLanes = subgroupSize / lanesUsed;
- int innermost = srcRank - 1;
- if (spareLanes > 1 && isReduced(innermost) && srcShape[innermost] > 1)
- laneLayout[innermost] = std::min(spareLanes, srcShape[innermost]);
-
- // Spend what is left of the budget on a lane-less, unpacked dim, since a lane
- // owns such a dim whole. Prefer the dim just outside the innermost two: in
- // the scale computation those two come from a shape_cast that split one dim,
- // so the dim above them is what should carry the per-lane vector.
- int64_t packedByConsumer = 1;
- for (int64_t d : laneData)
- packedByConsumer *= d;
+ // Rule 3. The dim just outside the innermost two is preferred because in the
+ // scale computation those two come from a shape_cast that split one dim, so
+ // the dim above them is what should carry the per-lane vector.
int64_t vectorBudget =
- maxReduceVectorSize / std::min(packedByConsumer, maxReduceVectorSize);
-
- int vectorDim = -1;
- int outsideLaneDims = srcRank - 3;
- auto usable = [&](int i) {
- return laneLayout[i] == 1 && laneData[i] == 1 && srcShape[i] > 1;
+ maxReduceVectorSize /
+ std::min(computeProduct(laneData), maxReduceVectorSize);
+ auto canCarryVector = [&](int i) {
+ return laneLayout[i] == 1 && srcShape[i] > 1;
};
+ int vectorDim = -1;
if (vectorBudget > 1) {
- if (outsideLaneDims >= 0 && usable(outsideLaneDims))
- vectorDim = outsideLaneDims;
+ int preferredDim = srcRank - 3;
+ if (preferredDim >= 0 && canCarryVector(preferredDim))
+ vectorDim = preferredDim;
else
for (int i = srcRank - 1; i >= 0; --i)
- if (usable(i)) {
+ if (canCarryVector(i)) {
vectorDim = i;
break;
}
}
if (vectorDim >= 0)
- laneData[vectorDim] = std::min(vectorBudget, srcShape[vectorDim]);
+ laneData[vectorDim] =
+ std::min(vectorBudget * laneData[vectorDim], srcShape[vectorDim]);
SmallVector<int64_t> instData(srcRank);
for (int i = 0; i < srcRank; ++i)
@@ -2466,13 +2439,10 @@ xegpu::completeDpasMxLaneLayoutFromInstData(
/// (`ResolveLayoutConflicts`), which inserts a `convert_layout` op - this
/// function never has to give up.
///
-/// For the InstData and Lane layout kinds only the innermost two dimensions
-/// are distributed; all leading dimensions are assumed to be unit dimensions.
-/// This assumption is checked via `leadingDimsAreUnit`. The lane_layout and
-/// lane_data are computed by `computeReductionLaneLayoutAndData`, which picks
-/// a layout that minimizes cross-lane reduction (reducing within a lane when
-/// only one of the innermost two dims is a reduction dim). The inst_data is
-/// simply the element-wise product lane_layout * lane_data.
+/// For the InstData and Lane layout kinds the lane_layout and lane_data are
+/// computed by `computeReductionLaneLayoutAndData`, which prefers a lane-local
+/// reduction and spends only otherwise-idle lanes on a cross-lane one. The
+/// inst_data is simply the element-wise product lane_layout * lane_data.
///
/// The function returns the *result* layout (the SliceAttr). The *source*
/// layout it decides on is the parent of that slice; both are listed below so
@@ -2656,13 +2626,6 @@ xegpu::SliceAttr xegpu::setupMultiReductionResultLayout(
// the layout is not consistent
srcLayout = consumerSliceLayout.getParent();
} else {
- // The re-derivation path below only distributes the innermost two dims,
- // so it requires the leading dims to be unit. The reused-consumer branch
- // above has no such restriction (a consumer may legitimately keep the
- // reduce-vector size as per-lane data on a non-unit leading dim).
- assert(leadingDimsAreUnit(srcShape, /*numInnerDims=*/2) &&
- "Lane reduction layout assumes all leading (non-innermost-two) "
- "dimensions are unit dimensions");
auto [laneLayout, laneData, instData] = computeReductionLaneLayoutAndData(
srcShape, reductionDims, subgroupSize, maxReduceVectorSize,
consumerLayout);
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index b3fe76c4da17c..0aba1faeb1cba 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -1201,3 +1201,32 @@ func.func @broadcast_slice_expanded_dim(%dst: !xegpu.mem_desc<16x4xbf16>) {
return
}
}
+
+// -----
+// A reduction whose only non-unit surviving dim is a *leading* dim: src
+// [16, 1, 4] reducing dim 2. The 16 lanes belong on dim 0, which leaves the
+// reduction lane-local and lets the gather feeding it use the whole subgroup.
+gpu.module @test {
+// CHECK-LABEL: func.func @reduce_innermost_lanes_on_leading(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} dense<true> : vector<64xi1>
+// CHECK: %[[IDX:.*]] = vector.step {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<64xindex>
+// CHECK: %[[LOAD:.*]] = xegpu.load %arg0[%[[IDX]]], %[[CST]] <{layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>}> : memref<64xf16>, vector<64xindex>, vector<64xi1> -> vector<64xf16>
+// CHECK: %[[SC:.*]] = vector.shape_cast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>} : vector<64xf16> to vector<16x1x4xf16>
+// CHECK: %[[ACC:.*]] = arith.constant {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>, dims = [2]>} dense<0.000000e+00> : vector<16x1xf16>
+// CHECK: %[[RED:.*]] = vector.multi_reduction <add>, %[[SC]], %[[ACC]] {layout_result_0 = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1], lane_data = [1, 1, 4]>, dims = [2]>} [2] : vector<16x1x4xf16> to vector<16x1xf16>
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[RED]] {layout_result_0 = #xegpu.layout<lane_layout = [16], lane_data = [1]>} : vector<16x1xf16> to vector<16xf16>
+// CHECK: xegpu.store %[[FLAT]]
+func.func @reduce_innermost_lanes_on_leading(%arg0: memref<64xf16>, %arg1: memref<16xf16>) {
+ %cst = arith.constant dense<true> : vector<64xi1>
+ %0 = vector.step : vector<64xindex>
+ %1 = xegpu.load %arg0[%0], %cst : memref<64xf16>, vector<64xindex>, vector<64xi1> -> vector<64xf16>
+ %2 = vector.shape_cast %1 : vector<64xf16> to vector<16x1x4xf16>
+ %cst_0 = arith.constant dense<0.000000e+00> : vector<16x1xf16>
+ %3 = vector.multi_reduction <add>, %2, %cst_0 [2] : vector<16x1x4xf16> to vector<16x1xf16>
+ %4 = vector.shape_cast %3 : vector<16x1xf16> to vector<16xf16>
+ %cst_2 = arith.constant dense<true> : vector<16xi1>
+ %cst_3 = arith.constant dense<1> : vector<16xindex>
+ xegpu.store %4, %arg1[%cst_3], %cst_2 : vector<16xf16>, memref<16xf16>, vector<16xindex>, vector<16xi1>
+ return
+ }
+}
More information about the Mlir-commits
mailing list