[Mlir-commits] [mlir] [mlir][xegpu] Fill split-group inner dims on shape_cast result layout (PR #227560)
Jianhui Li
llvmlistbot at llvm.org
Sat Oct 3 19:11:26 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/227560
>From 1a724ede41f7b532c5b02cc2921b3657562b9c66 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Wed, 30 Sep 2026 01:25:46 +0000
Subject: [PATCH 1/4] [mlir][xegpu] Fill split-group inner dims on shape_cast
result layout
A dim-split shape_cast (e.g. 1024 -> 32x32) derives its source layout in
inferShapeCastSourceLayout by collapsing the split group with
DistributeLayoutAttr::collapseDims, which multiplies the group's per-dim
sizes together. That product only describes the same elements as the result
layout when the dims the lanes split are covered in full: a partially
covered dim is strided in the collapsed dim, not contiguous.
On an mx-fp gemm the result layout of
vector.shape_cast %absf : vector<16x1024xbf16> to vector<16x32x32xbf16>
is inst_data [1, 2, 8], lane_layout [1, 2, 8], lane_data [1, 1, 1],
demanded by a multi_reduction over dim 2. The split group is dims 1 and 2,
and the lanes cover only 8 of dim 2's 32. Collapsing yields inst_data
[1, 16] on the 16x1024 source, which names 16 contiguous elements rather
than the 2x8 strided box the result layout asked for.
Add setupShapeCastResultLayout and call it from visitShapeCastOp before
inferring the source layout -- the same consumer-driven setup step bitcast
and interleave already use. Scanning a group from dim 0 to its last dim, a
dim the lanes do not split adds no stride, so the leading such dims are
skipped, and so is the first dim the lanes do split: that one is free to
stay partial. For every dim after it, lane_data is stretched so that
lane_layout * lane_data covers the dim, and inst_data follows as
lane_layout * lane_data. For the case above that demands inst_data
[1, 2, 32], lane_layout [1, 2, 8], lane_data [1, 1, 4], which collapses to
inst_data [1, 64], lane_layout [1, 16], lane_data [1, 4].
Skipping the dims the lanes do not split is what keeps a layout like
lane_layout [1, 1, 16] on an 8x8x16 result legal -- it is 16 contiguous
elements -- and the existing vector_shape_cast_expand_non_unit_dims test
unchanged. inst_data is deliberately not used to decide which dims are
strided: an inst_data larger than lane_layout * lane_data means several
iterations of the lane box, not one strided box.
The layout's `order` is not consulted either, since shape_cast splits a dim
row-major while `order` maps sg / lane ids.
When the lanes cannot tile a dim with whole lane_data values the consumer's
layout is returned untouched and the mismatch is left to
ResolveLayoutConflicts. Only the inst_data and lane phases are handled; a
subgroup-level layout is returned unchanged, with a TODO.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 24 ++++++++++++
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 38 +++++++++++++++++++
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 9 ++++-
.../XeGPU/propagate-layout-inst-data.mlir | 30 +++++++++++++++
4 files changed, 99 insertions(+), 2 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index e23bce46b16ee..f2cc963bcb885 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -267,6 +267,30 @@ DistributeLayoutAttr setupInterleaveResultLayout(
LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy,
DistributeLayoutAttr consumerLayout, const uArch::uArch *uArch);
+/// Sets up the result layout for a shape cast that splits one source dim into
+/// several consecutive result dims, so that the source layout can be safely
+/// derived by collapsing each split group.
+///
+/// Within a group, the leading dims the lanes do not split and the first dim
+/// they do split may stay partial. Every dim after that is stretched to full:
+/// lane_data becomes dim_size / lane_layout, and inst_data the dim size.
+///
+/// Example:
+/// shape_cast: vector<16x1024xf32> -> vector<16x32x32xf32>
+/// Consumer layout: inst_data = [1, 2, 8], lane_layout = [1, 2, 8],
+/// lane_data = [1, 1, 1]
+/// Adjusted: inst_data = [1, 2, 32], lane_layout = [1, 2, 8],
+/// lane_data = [1, 1, 4]
+/// The adjusted layout collapses to inst_data = [1, 64],
+/// lane_layout = [1, 16], lane_data = [1, 4] on the source, whereas the
+/// unadjusted one would have collapsed to a strided inst_data = [1, 16].
+///
+/// Only the inst_data and lane phases are handled; a subgroup-level layout is
+/// returned unchanged.
+DistributeLayoutAttr setupShapeCastResultLayout(
+ LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy,
+ DistributeLayoutAttr consumerLayout);
+
/// Sets up the result layout for an insert strided slice operation.
/// Creates a result layout based on the specified layout kind (InstData or
/// Lane).
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index e7c945f535044..fcf82edba0a0e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2752,6 +2752,44 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
resShape[innerMostDim], uArch);
}
+xegpu::DistributeLayoutAttr xegpu::setupShapeCastResultLayout(
+ xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
+ DistributeLayoutAttr consumerLayout) {
+ // TODO: work out the subgroup-level rule; leave such layouts alone for now.
+ if (!consumerLayout || layoutKind == xegpu::LayoutKind::Subgroup)
+ return consumerLayout;
+
+ ArrayRef<int64_t> resShape = resVecTy.getShape();
+ SmallVector<SmallVector<int64_t>> splitDimGroups;
+ if (!xegpu::matchSplitDimExpansion(srcVecTy.getShape(), resShape,
+ splitDimGroups))
+ return consumerLayout;
+
+ SmallVector<int64_t> laneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
+ if (laneLayout.empty())
+ return consumerLayout;
+
+ xegpu::DistributeLayoutAttr resLayout = consumerLayout;
+ for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
+ size_t i = 0;
+ while (i < dimGroup.size() && laneLayout[dimGroup[i]] == 1)
+ ++i;
+
+ for (++i; i < dimGroup.size(); ++i) {
+ int64_t dim = dimGroup[i];
+ int64_t dimSize = resShape[dim];
+ int64_t lanes = laneLayout[dim];
+ if (lanes <= 0 || dimSize % lanes != 0)
+ return consumerLayout;
+ resLayout = resLayout.setDimData(dim, /*sgData=*/-1,
+ instData.empty() ? -1 : dimSize,
+ dimSize / lanes);
+ }
+ }
+ return resLayout;
+}
+
/// Sets up the result layout for an insert strided slice operation.
/// Creates a result layout based on the specified layout kind (InstData or
/// Lane).
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index f05193504116b..15cf2ef599021 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -672,8 +672,13 @@ void LayoutInfoPropagation::visitShapeCastOp(
auto resultLayoutAttr =
dyn_cast<xegpu::DistributeLayoutAttr>(resLayoutInfo.get());
- xegpu::DistributeLayoutAttr srcLayoutAttr =
- xegpu::inferShapeCastSourceLayout(resultLayoutAttr, resShape, srcShape);
+ auto requiredResLayoutAttr = xegpu::setupShapeCastResultLayout(
+ layoutKind, shapeCast.getSourceVectorType(),
+ shapeCast.getResultVectorType(), resultLayoutAttr);
+ xegpu::setTemporaryLayout(shapeCast->getResult(0), requiredResLayoutAttr);
+
+ xegpu::DistributeLayoutAttr srcLayoutAttr = xegpu::inferShapeCastSourceLayout(
+ requiredResLayoutAttr, resShape, srcShape);
// shape_cast is not an anchor op: another consumer of the source value may
// still supply a valid layout, so warn instead of stopping the propagation.
if (!srcLayoutAttr) {
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index 841f06545eb19..eeeb6cc84f308 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -700,3 +700,33 @@ gpu.module @test {
gpu.return
}
}
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @shape_cast_split_fills_strided_inner_dim(
+// CHECK: %[[CAST:.*]] = vector.shape_cast %{{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 2, 32], lane_layout = [1, 2, 8], lane_data = [1, 1, 4]>} : vector<16x1024xbf16> to vector<16x32x32xbf16>
+func.func @shape_cast_split_fills_strided_inner_dim(%arg0: memref<16x1024xbf16>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x1024xbf16> -> !xegpu.tensor_desc<16x1024xbf16>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x1024xbf16> -> vector<16x1024xbf16>
+ %2 = vector.shape_cast %1 : vector<16x1024xbf16> to vector<16x32x32xbf16>
+ %3 = xegpu.convert_layout %2
+ <{target_layout = #xegpu.layout<inst_data = [1, 2, 8], lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>}>
+ : vector<16x32x32xbf16>
+ return
+}
+}
+
+// -----
+gpu.module @test {
+// CHECK-LABEL: func.func @shape_cast_split_fills_middle_dim(
+// CHECK: %[[CAST:.*]] = vector.shape_cast %{{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 2, 4, 16], lane_layout = [1, 2, 2, 4], lane_data = [1, 1, 2, 4]>} : vector<16x128xbf16> to vector<16x2x4x16xbf16>
+func.func @shape_cast_split_fills_middle_dim(%arg0: memref<16x128xbf16>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x128xbf16> -> !xegpu.tensor_desc<16x128xbf16>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x128xbf16> -> vector<16x128xbf16>
+ %2 = vector.shape_cast %1 : vector<16x128xbf16> to vector<16x2x4x16xbf16>
+ %3 = xegpu.convert_layout %2
+ <{target_layout = #xegpu.layout<inst_data = [1, 2, 2, 4], lane_layout = [1, 2, 2, 4], lane_data = [1, 1, 1, 1]>}>
+ : vector<16x2x4x16xbf16>
+ return
+}
+}
>From 2fc0fee8e00deb01fef77637edf25ce1cb22193d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Fri, 2 Oct 2026 17:50:51 +0000
Subject: [PATCH 2/4] [mlir][xegpu] Reject shape_cast layouts that cannot be
collapsed
Stretching the inner dims of a split group does not guarantee the result
layout can be collapsed onto the source dim: collapsing also requires each
lane to own a single contiguous run of that dim, which the consumer's
lane_layout can rule out. For
shape_cast: vector<16x128xbf16> -> vector<16x2x4x16xbf16>
consumer: lane_layout = [1, 2, 2, 4], lane_data = [1, 1, 1, 1]
stretching gives lane_data = [1, 1, 2, 4], where lane 0 owns the result
offsets {0..3} and {16..19} - two runs, 16 apart - while the collapsed
lane_data = 8 claims the contiguous source columns [0, 8). collapseDims
just multiplies the fields, so nothing downstream catches it and the lane
ends up with different elements before and after the cast.
Check the three layout fields of the stretched layout with one stride walk
(isContiguousTiling) and return nullptr when a group does not collapse;
shape_cast is not an anchor op, so the propagation warns and leaves the
value un-laid-out instead.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPULayoutImpl.h | 22 ++++-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 99 +++++++++++++++----
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 10 ++
.../propagate-layout-inst-data-invalid.mlir | 64 ++++++++++++
.../XeGPU/propagate-layout-inst-data.mlir | 22 +++--
5 files changed, 187 insertions(+), 30 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/propagate-layout-inst-data-invalid.mlir
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index f2cc963bcb885..a79de4983bc2e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -285,11 +285,27 @@ DistributeLayoutAttr setupInterleaveResultLayout(
/// lane_layout = [1, 16], lane_data = [1, 4] on the source, whereas the
/// unadjusted one would have collapsed to a strided inst_data = [1, 16].
///
+/// Stretching does not guarantee the result layout is collapsible: collapsing
+/// also requires each lane to own one contiguous run of the source dim, which
+/// the consumer's lane_layout can rule out. This function checks that
+/// restriction on the stretched layout and returns nullptr when it does not
+/// hold, rejecting the consumer layout.
+///
+/// Example of a rejected layout:
+/// shape_cast: vector<16x128xf32> -> vector<16x2x4x16xf32>
+/// Consumer layout: inst_data = [1, 2, 2, 4], lane_layout = [1, 2, 2, 4],
+/// lane_data = [1, 1, 1, 1]
+/// Adjusted: inst_data = [1, 2, 4, 16], lane_layout = [1, 2, 2, 4],
+/// lane_data = [1, 1, 2, 4]
+/// Each lane owns two runs of 4 elements, 16 apart, instead of the 8
+/// contiguous source elements the collapsed lane_data = [1, 8] claims.
+///
/// Only the inst_data and lane phases are handled; a subgroup-level layout is
/// returned unchanged.
-DistributeLayoutAttr setupShapeCastResultLayout(
- LayoutKind layoutKind, VectorType srcVectorTy, VectorType resVectorTy,
- DistributeLayoutAttr consumerLayout);
+DistributeLayoutAttr
+setupShapeCastResultLayout(LayoutKind layoutKind, VectorType srcVectorTy,
+ VectorType resVectorTy,
+ DistributeLayoutAttr consumerLayout);
/// Sets up the result layout for an insert strided slice operation.
/// Creates a result layout based on the specified layout kind (InstData or
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index fcf82edba0a0e..ab14b031df406 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2752,9 +2752,38 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
resShape[innerMostDim], uArch);
}
-xegpu::DistributeLayoutAttr xegpu::setupShapeCastResultLayout(
- xegpu::LayoutKind layoutKind, VectorType srcVecTy, VectorType resVecTy,
- DistributeLayoutAttr consumerLayout) {
+/// Checks that the tiles a split group is cut into walk the collapsed source
+/// dim with one constant stride, which is what `collapseDims` assumes when it
+/// multiplies a layout field across the group.
+///
+/// Along each dim the group is cut into `counts[dim]` tiles of `tileShape[dim]`
+/// elements (empty for single-element tiles) taken from `shape[dim]`, which
+/// gives the row-major strides. Walking from the innermost dim out, a dim cut
+/// into several tiles steps by `tileShape[dim] * stride(dim)` and has to pick
+/// up where the inner dims left off, `expectedStride`; a dim cut into one tile
+/// does not step at all, so it only contributes its size.
+static bool isContiguousTiling(ArrayRef<int64_t> dimGroup,
+ ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> counts,
+ ArrayRef<int64_t> tileShape,
+ int64_t expectedStride) {
+ int64_t stride = 1;
+ for (int64_t dim : llvm::reverse(dimGroup)) {
+ int64_t tileSize = tileShape.empty() ? 1 : tileShape[dim];
+ if (counts[dim] != 1) {
+ if (stride * tileSize != expectedStride)
+ return false;
+ expectedStride = stride * tileSize * counts[dim];
+ }
+ stride *= shape[dim];
+ }
+ return true;
+}
+
+xegpu::DistributeLayoutAttr
+xegpu::setupShapeCastResultLayout(xegpu::LayoutKind layoutKind,
+ VectorType srcVecTy, VectorType resVecTy,
+ DistributeLayoutAttr consumerLayout) {
// TODO: work out the subgroup-level rule; leave such layouts alone for now.
if (!consumerLayout || layoutKind == xegpu::LayoutKind::Subgroup)
return consumerLayout;
@@ -2765,27 +2794,59 @@ xegpu::DistributeLayoutAttr xegpu::setupShapeCastResultLayout(
splitDimGroups))
return consumerLayout;
- SmallVector<int64_t> laneLayout = consumerLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
if (laneLayout.empty())
return consumerLayout;
- xegpu::DistributeLayoutAttr resLayout = consumerLayout;
- for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
- size_t i = 0;
- while (i < dimGroup.size() && laneLayout[dimGroup[i]] == 1)
- ++i;
-
- for (++i; i < dimGroup.size(); ++i) {
- int64_t dim = dimGroup[i];
- int64_t dimSize = resShape[dim];
- int64_t lanes = laneLayout[dim];
- if (lanes <= 0 || dimSize % lanes != 0)
- return consumerLayout;
- resLayout = resLayout.setDimData(dim, /*sgData=*/-1,
- instData.empty() ? -1 : dimSize,
- dimSize / lanes);
+ xegpu::DistributeLayoutAttr resLayout = [&]() -> DistributeLayoutAttr {
+ DistributeLayoutAttr layout = consumerLayout;
+ for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
+ size_t i = 0;
+ while (i < dimGroup.size() && laneLayout[dimGroup[i]] == 1)
+ ++i;
+
+ for (++i; i < dimGroup.size(); ++i) {
+ int64_t dim = dimGroup[i];
+ int64_t dimSize = resShape[dim];
+ int64_t lanes = laneLayout[dim];
+ if (lanes <= 0 || dimSize % lanes != 0)
+ return consumerLayout;
+ layout =
+ layout.setDimData(dim, /*sgData=*/-1,
+ instData.empty() ? -1 : dimSize, dimSize / lanes);
+ }
}
+ return layout;
+ }();
+
+ // Stretching does not make every layout collapsible; reject the rest.
+ SmallVector<int64_t> resInstData = resLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> resLaneData = resLayout.getEffectiveLaneDataAsInt();
+ if (resLaneLayout.empty() || resLaneData.empty())
+ return resLayout;
+
+ for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
+ // One lane's elements must be a single contiguous run of the collapsed dim.
+ if (!isContiguousTiling(dimGroup, resShape, resLaneData, /*tileShape=*/{},
+ /*expectedStride=*/1))
+ return nullptr;
+
+ // Consecutive lanes must be exactly one such run apart.
+ int64_t collapsedLaneData = 1;
+ for (int64_t dim : dimGroup)
+ collapsedLaneData *= resLaneData[dim];
+ if (!isContiguousTiling(dimGroup, resShape, resLaneLayout, resLaneData,
+ /*expectedStride=*/collapsedLaneData))
+ return nullptr;
+
+ // The inst tile must walk the result shape the same way.
+ if (!resInstData.empty() &&
+ !isContiguousTiling(dimGroup, resShape, resInstData, /*tileShape=*/{},
+ /*expectedStride=*/1))
+ return nullptr;
}
return resLayout;
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 15cf2ef599021..9a4b35eba245c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -675,6 +675,16 @@ void LayoutInfoPropagation::visitShapeCastOp(
auto requiredResLayoutAttr = xegpu::setupShapeCastResultLayout(
layoutKind, shapeCast.getSourceVectorType(),
shapeCast.getResultVectorType(), resultLayoutAttr);
+ // The consumer layout cannot be expressed on the source: no lane_data makes
+ // each lane's data a contiguous run of the collapsed source dim. shape_cast
+ // is not an anchor op, so warn and leave the value un-laid-out instead of
+ // propagating a layout that would move data between lanes.
+ if (!requiredResLayoutAttr) {
+ shapeCast.emitWarning("Failed to infer source layout for shape_cast; the "
+ "result layout required by its consumers cannot be "
+ "collapsed onto the source shape.");
+ return;
+ }
xegpu::setTemporaryLayout(shapeCast->getResult(0), requiredResLayoutAttr);
xegpu::DistributeLayoutAttr srcLayoutAttr = xegpu::inferShapeCastSourceLayout(
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data-invalid.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data-invalid.mlir
new file mode 100644
index 0000000000000..7dc239847444a
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data-invalid.mlir
@@ -0,0 +1,64 @@
+// RUN: mlir-opt -xevm-attach-target='chip=cri' -test-xegpu-propagate-layouts="layout-kind=inst" -split-input-file -verify-diagnostics %s
+
+// A shape cast that splits one source dim into several result dims can only
+// take a consumer layout whose lanes own data the collapsed source dim can
+// describe, i.e. one contiguous run per lane.
+
+// Here the lanes split dim 2 (2 of 4) as well as dim 3 (4 of 16). Stretching the
+// inner dims gives lane_data = [1, 1, 2, 4], so lane 0 owns dim 2 in [0, 2) and
+// dim 3 in [0, 4): the result offsets {0, 1, 2, 3} and {16, 17, 18, 19} - two
+// runs, 16 apart - whereas the collapsed lane_data = 8 claims the contiguous
+// source columns [0, 8). No lane_data fixes this, so the consumer layout is
+// rejected and no layout reaches the source of the cast.
+gpu.module @test {
+ func.func @shape_cast_split_lanes_split_two_dims(%arg0: memref<16x128xbf16>) {
+ // expected-warning at +1 {{op has users but no layout assigned for its result}}
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x128xbf16> -> !xegpu.tensor_desc<16x128xbf16>
+ // expected-warning at +1 {{op has users but no layout assigned for its result}}
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x128xbf16> -> vector<16x128xbf16>
+ // expected-warning at +1 {{Failed to infer source layout for shape_cast; the result layout required by its consumers cannot be collapsed onto the source shape.}}
+ %2 = vector.shape_cast %1 : vector<16x128xbf16> to vector<16x2x4x16xbf16>
+ %3 = xegpu.convert_layout %2
+ <{target_layout = #xegpu.layout<inst_data = [1, 2, 2, 4], lane_layout = [1, 2, 2, 4], lane_data = [1, 1, 1, 1]>}>
+ : vector<16x2x4x16xbf16>
+ return
+ }
+}
+
+// -----
+// Same requirement one dim out: the consumer asks each lane for 2 elements of
+// dim 1, which are 64 elements apart in the result, so the lane's data is not a
+// single run of the collapsed source dim either.
+gpu.module @test {
+ func.func @shape_cast_split_outer_dim_lane_data(%arg0: memref<16x128xbf16>) {
+ // expected-warning at +1 {{op has users but no layout assigned for its result}}
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x128xbf16> -> !xegpu.tensor_desc<16x128xbf16>
+ // expected-warning at +1 {{op has users but no layout assigned for its result}}
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x128xbf16> -> vector<16x128xbf16>
+ // expected-warning at +1 {{Failed to infer source layout for shape_cast; the result layout required by its consumers cannot be collapsed onto the source shape.}}
+ %2 = vector.shape_cast %1 : vector<16x128xbf16> to vector<16x2x4x16xbf16>
+ %3 = xegpu.convert_layout %2
+ <{target_layout = #xegpu.layout<inst_data = [1, 2, 2, 8], lane_layout = [1, 1, 2, 8], lane_data = [1, 2, 1, 1]>}>
+ : vector<16x2x4x16xbf16>
+ return
+ }
+}
+
+// -----
+// The inst tile has to collapse too: inst_data = [1, 2, 2, 8] takes only 2 of
+// the 4 elements of dim 2, so the tile is two 16-element chunks of the source
+// row rather than the 32 contiguous elements the collapsed inst_data claims.
+gpu.module @test {
+ func.func @shape_cast_split_partial_inst_data(%arg0: memref<16x128xbf16>) {
+ // expected-warning at +1 {{op has users but no layout assigned for its result}}
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x128xbf16> -> !xegpu.tensor_desc<16x128xbf16>
+ // expected-warning at +1 {{op has users but no layout assigned for its result}}
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x128xbf16> -> vector<16x128xbf16>
+ // expected-warning at +1 {{Failed to infer source layout for shape_cast; the result layout required by its consumers cannot be collapsed onto the source shape.}}
+ %2 = vector.shape_cast %1 : vector<16x128xbf16> to vector<16x2x4x16xbf16>
+ %3 = xegpu.convert_layout %2
+ <{target_layout = #xegpu.layout<inst_data = [1, 2, 2, 8], lane_layout = [1, 1, 2, 8], lane_data = [1, 1, 1, 1]>}>
+ : vector<16x2x4x16xbf16>
+ return
+ }
+}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index eeeb6cc84f308..05c5ec0e4b036 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -717,16 +717,22 @@ func.func @shape_cast_split_fills_strided_inner_dim(%arg0: memref<16x1024xbf16>)
}
// -----
+// A group of three dims: the lanes cover dim 1 and dim 2 exactly once (2 and 8
+// lanes) and are not distributed over dim 3, so the fill hands dim 3 to each
+// lane whole - lane_data = [1, 1, 1, 16], one contiguous run of 16 per lane,
+// which collapses to inst_data = [1, 256], lane_layout = [1, 16],
+// lane_data = [1, 16] on the 16x256 source. Dim 2 is already covered exactly
+// once, so the fill leaves it as it is.
gpu.module @test {
-// CHECK-LABEL: func.func @shape_cast_split_fills_middle_dim(
-// CHECK: %[[CAST:.*]] = vector.shape_cast %{{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 2, 4, 16], lane_layout = [1, 2, 2, 4], lane_data = [1, 1, 2, 4]>} : vector<16x128xbf16> to vector<16x2x4x16xbf16>
-func.func @shape_cast_split_fills_middle_dim(%arg0: memref<16x128xbf16>) {
- %0 = xegpu.create_nd_tdesc %arg0 : memref<16x128xbf16> -> !xegpu.tensor_desc<16x128xbf16>
- %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x128xbf16> -> vector<16x128xbf16>
- %2 = vector.shape_cast %1 : vector<16x128xbf16> to vector<16x2x4x16xbf16>
+// CHECK-LABEL: func.func @shape_cast_split_fills_undistributed_inner_dim(
+// CHECK: %[[CAST:.*]] = vector.shape_cast %{{.*}} {layout_result_0 = #xegpu.layout<inst_data = [1, 2, 8, 16], lane_layout = [1, 2, 8, 1], lane_data = [1, 1, 1, 16]>} : vector<16x256xbf16> to vector<16x2x8x16xbf16>
+func.func @shape_cast_split_fills_undistributed_inner_dim(%arg0: memref<16x256xbf16>) {
+ %0 = xegpu.create_nd_tdesc %arg0 : memref<16x256xbf16> -> !xegpu.tensor_desc<16x256xbf16>
+ %1 = xegpu.load_nd %0[0, 0] : !xegpu.tensor_desc<16x256xbf16> -> vector<16x256xbf16>
+ %2 = vector.shape_cast %1 : vector<16x256xbf16> to vector<16x2x8x16xbf16>
%3 = xegpu.convert_layout %2
- <{target_layout = #xegpu.layout<inst_data = [1, 2, 2, 4], lane_layout = [1, 2, 2, 4], lane_data = [1, 1, 1, 1]>}>
- : vector<16x2x4x16xbf16>
+ <{target_layout = #xegpu.layout<inst_data = [1, 2, 8, 1], lane_layout = [1, 2, 8, 1], lane_data = [1, 1, 1, 1]>}>
+ : vector<16x2x8x16xbf16>
return
}
}
>From aee9a6c829a3d817c3317add7ae800f5482d8d26 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 3 Oct 2026 02:41:49 +0000
Subject: [PATCH 3/4] [mlir][xegpu] Drop redundant lane layout checks in
shape_cast setup
Layout propagation only carries layouts that have lane_layout and lane_data,
an invariant setupLoadNdAnchorLayout already asserts, so the two emptiness
checks in setupShapeCastResultLayout were dead code.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 4 ----
1 file changed, 4 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index ab14b031df406..9d38c9c9d5eee 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2797,8 +2797,6 @@ xegpu::setupShapeCastResultLayout(xegpu::LayoutKind layoutKind,
SmallVector<int64_t> laneLayout =
consumerLayout.getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
- if (laneLayout.empty())
- return consumerLayout;
xegpu::DistributeLayoutAttr resLayout = [&]() -> DistributeLayoutAttr {
DistributeLayoutAttr layout = consumerLayout;
@@ -2825,8 +2823,6 @@ xegpu::setupShapeCastResultLayout(xegpu::LayoutKind layoutKind,
SmallVector<int64_t> resInstData = resLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> resLaneLayout = resLayout.getEffectiveLaneLayoutAsInt();
SmallVector<int64_t> resLaneData = resLayout.getEffectiveLaneDataAsInt();
- if (resLaneLayout.empty() || resLaneData.empty())
- return resLayout;
for (const SmallVector<int64_t> &dimGroup : splitDimGroups) {
// One lane's elements must be a single contiguous run of the collapsed dim.
>From f7a2b2b045c0f1e8c585ef198b24e22e7c465f8d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sun, 4 Oct 2026 02:10:57 +0000
Subject: [PATCH 4/4] add space to trigger the auto-merge process
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 9d38c9c9d5eee..4dda10c6a7a26 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2784,7 +2784,7 @@ xegpu::DistributeLayoutAttr
xegpu::setupShapeCastResultLayout(xegpu::LayoutKind layoutKind,
VectorType srcVecTy, VectorType resVecTy,
DistributeLayoutAttr consumerLayout) {
- // TODO: work out the subgroup-level rule; leave such layouts alone for now.
+ // TODO: work out the subgroup level rule; leave such layouts alone for now.
if (!consumerLayout || layoutKind == xegpu::LayoutKind::Subgroup)
return consumerLayout;
More information about the Mlir-commits
mailing list