[Mlir-commits] [mlir] [MLIR][XeGPU] Fix expandDim sg_data cap for replicated layouts (PR #203299)
Sang Ik Lee
llvmlistbot at llvm.org
Thu Jun 11 14:46:05 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/203299
>From 1944c5a4ae790866f04af0abdb89b0209451c4cc Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Thu, 11 Jun 2026 02:40:14 +0000
Subject: [PATCH 1/2] [MLIR][XeGPU] Fix expandDim sg_data cap for replicated
layouts
LayoutAttr::expandDim capped sg_data per dim by the per-subgroup share
targetShape[i] / sgLayout[i], which assumes an evenly-tiled layout. For a
replicated/broadcast dim, where a dimension is shared across subgroups so
sg_layout[dim] * sg_data[dim] > extent, sg_data can be as large as the full
extent and does not fit within that share. The inner-first spread then failed
the 'remaining == 1' assertion (XeGPUDialect.cpp:729).
This is exactly the layout shape of WG-level mxfp GEMM operands: the per-block
scale spans the full K extent shared across the N-parallel subgroups. The
shape_cast collapse inference (inferShapeCastSourceLayout) calls expandDim on
such a layout and crashed.
Cap sg_data by the full extent targetShape[i] instead. This is the inverse of
collapseDims (which simply multiplies sg_data with no cap). It is
behavior-preserving for evenly-distributed layouts: sg_layout fills
outer-to-inner, so dims inner of the boundary have sg_layout == 1 (full cap ==
share cap), and when the inner-first sg_data fill reaches the boundary dim the
remaining value equals the per-sg share, so min(remaining, fullExtent) still
selects the share. For replicated layouts it lets a subgroup hold up to the
whole dim, and collapseDims(expandDim(...)) round-trips back to the original.
Add a regression test (shape_cast_collapse_replicated) covering the replicated
sg_data collapse.
---
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 19 +++++++++----
.../XeGPU/propagate-layout-subgroup.mlir | 27 +++++++++++++++++++
2 files changed, 41 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index b780c66594eb0..0703d05729497 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -660,8 +660,14 @@ DistributeLayoutAttr LayoutAttr::collapseDims(SmallVector<int64_t> dimGroup) {
// - sg_layout / lane_layout: spread outer-to-inner; each dim takes
// min(remaining, targetShape[i]); leftover spills into the next inner
// dim.
-// - sg_data: fill innermost-first, capped per dim by
-// targetShape[i] / sgLayout[i] (the per-sg share of the extent).
+// - sg_data: fill innermost-first, capped per dim by targetShape[i] (the
+// full extent of the expanded dim). Capping by the full extent (rather
+// than the per-sg share targetShape[i] / sgLayout[i]) keeps the inverse of
+// collapseDims well-defined for replicated/broadcast layouts, where a dim
+// is shared across subgroups so sg_layout[dim] * sg_data[dim] > extent and
+// sg_data can be as large as the full extent. For evenly-distributed
+// layouts the inner-first fill still lands on the per-sg share, so the
+// result is unchanged.
// - lane_data: fill innermost-first, capped per dim by
// (targetShape[i] / sgLayout[i]) / laneLayout[i] (the per-lane share of
// the per-sg extent).
@@ -752,10 +758,13 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
splice(sgLayout, expSgLayout);
}
if (hasSgData) {
+ // Cap by the full extent targetShape[i], not the per-sg share
+ // targetShape[i] / sgLayout[i]: a replicated dim (sg_layout[dim] *
+ // sg_data[dim] > extent) has sg_data up to the full extent, which would not
+ // fit within the per-sg share. For evenly-distributed layouts the
+ // inner-first fill lands on the per-sg share regardless, so this is
+ // equivalent.
SmallVector<int64_t> dimSizeCap(targetShape.begin(), targetShape.end());
- if (hasSgLayout)
- for (int64_t i = 0; i < expCount; ++i)
- dimSizeCap[i] /= expSgLayout[i];
SmallVector<int64_t> expSgData =
spread(origSgDataDim, dimSizeCap, /*outerToInner=*/false);
splice(sgData, expSgData);
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index 5021c8a746045..f2f970ccbbf80 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -575,3 +575,30 @@ gpu.module @test {
gpu.return
}
}
+
+// -----
+// shape_cast collapse where the collapsed dst dim is replicated across
+// subgroups: sg_layout[1] * sg_data[1] = 2 * 64 = 128 > extent 64, i.e. the
+// two subgroups along dim1 hold the same data (broadcast). This is the layout
+// shape produced by WG-level mxfp GEMM operands, where the per-block scale
+// spans the full K extent shared across the N-parallel subgroups.
+// srcShape=[16, 8, 8], resShape=[16, 64], consumer sg_layout=[2, 2],
+// sg_data=[8, 64]
+// - dst[1]=64 collapses src[1, 2]: sg_layout outer-to-inner: dim1 take=
+// min(2, 8)=2 (rem=1) -> [_, 2, 1]; sg_data innermost-first capped by the
+// full extent (not the per-sg share): dim2 take=min(64, 8)=8 (rem=8);
+// dim1 take=min(8, 8)=8 (rem=1) -> [_, 8, 8]. The replicated sg_data 64
+// would not fit the per-sg share cap 64/2=32 on dim1, so the full extent
+// is used.
+gpu.module @test {
+// CHECK-LABEL: gpu.func @shape_cast_collapse_replicated(
+// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [2, 2, 1], sg_data = [8, 8, 8]>} dense<0.000000e+00> : vector<16x8x8xf16>
+// CHECK: %[[CAST:.*]] = vector.shape_cast %[[CST]] {layout_result_0 = #xegpu.layout<sg_layout = [2, 2], sg_data = [8, 64]>} : vector<16x8x8xf16> to vector<16x64xf16>
+ gpu.func @shape_cast_collapse_replicated(%dst: memref<16x64xf16>) kernel {
+ %cst = arith.constant dense<0.000000e+00> : vector<16x8x8xf16>
+ %0 = vector.shape_cast %cst : vector<16x8x8xf16> to vector<16x64xf16>
+ %tdesc = xegpu.create_nd_tdesc %dst : memref<16x64xf16> -> !xegpu.tensor_desc<16x64xf16>
+ xegpu.store_nd %0, %tdesc[0, 0] <{layout = #xegpu.layout<sg_layout = [2, 2], sg_data = [8, 64]>}> : vector<16x64xf16>, !xegpu.tensor_desc<16x64xf16>
+ gpu.return
+ }
+}
>From 7e145dd28619d353b47d46c76958f672ff98dd48 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Thu, 11 Jun 2026 21:43:11 +0000
Subject: [PATCH 2/2] Address reviewer comments.
---
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 13 ++++++-------
.../Dialect/XeGPU/propagate-layout-subgroup.mlir | 13 -------------
2 files changed, 6 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 0703d05729497..89eb254831e5f 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -757,14 +757,13 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
expSgLayout = spread(origSgLayoutDim, targetShape, /*outerToInner=*/true);
splice(sgLayout, expSgLayout);
}
+ bool sgDataReplicated =
+ hasSgData && origSgDataDim == computeProduct(targetShape);
if (hasSgData) {
- // Cap by the full extent targetShape[i], not the per-sg share
- // targetShape[i] / sgLayout[i]: a replicated dim (sg_layout[dim] *
- // sg_data[dim] > extent) has sg_data up to the full extent, which would not
- // fit within the per-sg share. For evenly-distributed layouts the
- // inner-first fill lands on the per-sg share regardless, so this is
- // equivalent.
SmallVector<int64_t> dimSizeCap(targetShape.begin(), targetShape.end());
+ if (hasSgLayout && !sgDataReplicated)
+ for (int64_t i = 0; i < expCount; ++i)
+ dimSizeCap[i] /= expSgLayout[i];
SmallVector<int64_t> expSgData =
spread(origSgDataDim, dimSizeCap, /*outerToInner=*/false);
splice(sgData, expSgData);
@@ -774,7 +773,7 @@ DistributeLayoutAttr LayoutAttr::expandDim(int64_t dim,
// targetShape[i] / sg_layout[i] when sg_layout is present, else
// targetShape itself.
SmallVector<int64_t> perSgShape(targetShape.begin(), targetShape.end());
- if (hasSgLayout)
+ if (hasSgLayout && !sgDataReplicated)
for (int64_t i = 0; i < expCount; ++i)
perSgShape[i] /= expSgLayout[i];
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index f2f970ccbbf80..d44497d0bba34 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -577,19 +577,6 @@ gpu.module @test {
}
// -----
-// shape_cast collapse where the collapsed dst dim is replicated across
-// subgroups: sg_layout[1] * sg_data[1] = 2 * 64 = 128 > extent 64, i.e. the
-// two subgroups along dim1 hold the same data (broadcast). This is the layout
-// shape produced by WG-level mxfp GEMM operands, where the per-block scale
-// spans the full K extent shared across the N-parallel subgroups.
-// srcShape=[16, 8, 8], resShape=[16, 64], consumer sg_layout=[2, 2],
-// sg_data=[8, 64]
-// - dst[1]=64 collapses src[1, 2]: sg_layout outer-to-inner: dim1 take=
-// min(2, 8)=2 (rem=1) -> [_, 2, 1]; sg_data innermost-first capped by the
-// full extent (not the per-sg share): dim2 take=min(64, 8)=8 (rem=8);
-// dim1 take=min(8, 8)=8 (rem=1) -> [_, 8, 8]. The replicated sg_data 64
-// would not fit the per-sg share cap 64/2=32 on dim1, so the full extent
-// is used.
gpu.module @test {
// CHECK-LABEL: gpu.func @shape_cast_collapse_replicated(
// CHECK: %[[CST:.*]] = arith.constant {layout_result_0 = #xegpu.layout<sg_layout = [2, 2, 1], sg_data = [8, 8, 8]>} dense<0.000000e+00> : vector<16x8x8xf16>
More information about the Mlir-commits
mailing list