[Mlir-commits] [mlir] b5ab963 - [mlir][xegpu] Fill split-group inner dims on shape_cast result layout (#227560)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Sat Oct 3 19:28:12 PDT 2026
Author: Jianhui Li
Date: 2026-10-04T02:28:05Z
New Revision: b5ab963ec7498a5c2ac81a8955ba98549dc5d6d0
URL: https://github.com/llvm/llvm-project/commit/b5ab963ec7498a5c2ac81a8955ba98549dc5d6d0
DIFF: https://github.com/llvm/llvm-project/commit/b5ab963ec7498a5c2ac81a8955ba98549dc5d6d0.diff
LOG: [mlir][xegpu] Fill split-group inner dims on shape_cast result layout (#227560)
This PR set the result layout for shape_cast when it can't infer the
layout from incoming consumer layout. For example,
```mlir
vector.shape_cast %0 : vector<16x1024xbf16> to vector<16x32x32xbf16>
// consumer layout: inst_data = [1, 2, 8], lane_layout = [1, 2, 8], lane_data = [1, 1, 1]
```
Taking the consumer layout as is doesn't work for layout propagation.
The lanes cover 2 along dim 1 and only 8 of dim 2's 32 — a 2x8 box.
Collapsing dims 1 and 2 gives inst_data = [1, 16] on the 16x1024 source,
which is 16 contiguous elements. Those are not the same elements: the
2x8 box is two runs of 8, 32 apart.
This PR fixes the issue by stretching the lane_data to cover the
innermost dims of source shape, to set up result layout so that each
lane gets the same data before and after the shapecast, which is the
condition that shapecast op can be successfully distributed.
```
before: inst_data = [1, 2, 8], lane_layout = [1, 2, 8],lane_data = [1, 1, 1]
after: inst_data = [1, 2, 32], lane_layout = [1, 2, 8],lane_data = [1, 1, 4]
^^ ^
collapses to: inst_data = [1, 64], lane_layout = [1, 16],lane_data = [1, 4]
```
lane_data grows from 1 to 4 so that 8 lanes x 4 elements = 32, and
inst_data follows.
---------
Co-authored-by: Claude Opus 5 (1M context) <noreply at anthropic.com>
Added:
mlir/test/Dialect/XeGPU/propagate-layout-inst-data-invalid.mlir
Modified:
mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index e23bce46b16ee..a79de4983bc2e 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -267,6 +267,46 @@ 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].
+///
+/// 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);
+
/// 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..4dda10c6a7a26 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -2752,6 +2752,101 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
resShape[innerMostDim], uArch);
}
+/// 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;
+
+ 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();
+
+ 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();
+
+ 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;
+}
+
/// 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..9a4b35eba245c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -672,8 +672,23 @@ 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);
+ // 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(
+ 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-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 841f06545eb19..05c5ec0e4b036 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -700,3 +700,39 @@ 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
+}
+}
+
+// -----
+// 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_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, 8, 1], lane_layout = [1, 2, 8, 1], lane_data = [1, 1, 1, 1]>}>
+ : vector<16x2x8x16xbf16>
+ return
+}
+}
More information about the Mlir-commits
mailing list