[Mlir-commits] [mlir] [mlir][xegpu] Distribute extract_strided_slice over multiple dims (PR #227478)
Jianhui Li
llvmlistbot at llvm.org
Wed Sep 30 19:05:51 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/227478
>From 0013be041e1b2dece6b12d78f05545e41fa97322 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 29 Sep 2026 21:04:46 +0000
Subject: [PATCH 1/3] [mlir][xegpu] Distribute extract_strided_slice over
multiple dims
SgToLaneVectorExtractStridedSlice only handled a single distributed
dimension, and it rescaled that dimension by the subgroup size rather
than by the dimension's own lane count -- deriving an effective subgroup
size from the uArch and then clamping it to lane_layout[distDim]. That
clamp is exactly the lane count, so the uArch lookup was only ever a
detour, and it made the multi-dim case look harder than it is.
Scale each distributed dim by its own lane_layout entry instead. Both
the size and the offset along a dim shrink by the number of lanes that
split it, independently of the other dims, so the single-dim case is
unchanged and the general case falls out. A dim the lanes do not split
keeps its offset untouched, and may still carry non-unit lane_data as
before. The uArch/target-attribute lookup goes away with it, so the
pattern no longer fails on modules without a chip attribute.
The divisibility and unit-lane_data checks are now per distributed dim.
A packed distributed dim still declines, since its offset would have to
be scaled by lane_data as well.
Found on an mx-fp quantize-A kernel, where a reduction derives
lane_layout [1, 2, 8] and the extract_strided_slice feeding it
distributes two dims. Unit test added for that shape; the XeGPU lit
suite passes (27/27), and the test fails without this change with
"failed to legalize operation 'vector.extract_strided_slice'".
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 70 ++++++++-----------
.../XeGPU/sg-to-lane-distribute-unit.mlir | 24 +++++++
2 files changed, 53 insertions(+), 41 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 06f769cefb4fd..e2b61e281dfcf 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1359,52 +1359,40 @@ struct SgToLaneVectorExtractStridedSlice
updatedStrides.push_back(rewriter.getI64IntegerAttr(1));
}
- // If the result is distributed, adjust offsets and sizes in the
- // distributed dimension.
+ // Each distributed dim shrinks by its own lane count, so its size and
+ // offset are rescaled by that count.
if (!distributedDims.empty()) {
- if (distributedDims.size() != 1)
- return rewriter.notifyMatchFailure(
- op, "only single dimension distribution is supported");
- int64_t distDim = distributedDims[0];
- const auto *uArch =
- xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
- if (!uArch)
- return rewriter.notifyMatchFailure(
- op, "target attribute required to determine subgroup size");
- int subgroupSize = uArch->getSubgroupSize();
auto sourceLayout = xegpu::getTemporaryLayout(op->getOpOperand(0));
if (!sourceLayout || sourceLayout.getEffectiveLaneLayoutAsInt().empty())
return rewriter.notifyMatchFailure(
op, "source of extract_strided_slice lacks distribution layout");
- int sourceDistrDimSize = op.getSourceVectorType().getShape()[distDim];
- auto laneLayout = sourceLayout.getEffectiveLaneLayoutAsInt();
- // Effective subgroup size needs to be adjusted if laneLayout along
- // the distributed dimension is smaller than subgroup size.
- if (laneLayout[distDim] < subgroupSize &&
- subgroupSize % laneLayout[distDim] == 0)
- subgroupSize = laneLayout[distDim];
- if (sourceDistrDimSize % subgroupSize != 0)
- return rewriter.notifyMatchFailure(
- op, "source size along distributed dim is not a multiple of "
- "subgroup size");
- auto sourceLaneData = sourceLayout.getEffectiveLaneDataAsInt();
- // Only check lane_data for the distributed dimension. Non-distributed
- // dimensions may have non-unit lane_data (e.g., packed layouts).
- if (distDim < static_cast<int64_t>(sourceLaneData.size()) &&
- sourceLaneData[distDim] != 1)
- return rewriter.notifyMatchFailure(
- op, "expecting unit lane data along the distributed dimension");
- int64_t distrDimOffset =
- cast<IntegerAttr>(updatedOffsets[distDim]).getInt();
- if (distrDimOffset % subgroupSize != 0)
- return rewriter.notifyMatchFailure(
- op, "offset along distributed dim is not a multiple of "
- "subgroup size");
- // Adjust sizes and offsets for the distributed dimension.
- updatedSizes[distDim] =
- rewriter.getI64IntegerAttr(distResultTy.getDimSize(distDim));
- updatedOffsets[distDim] =
- rewriter.getI64IntegerAttr(distrDimOffset / subgroupSize);
+ SmallVector<int64_t> laneLayout =
+ sourceLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> sourceLaneData =
+ sourceLayout.getEffectiveLaneDataAsInt();
+ ArrayRef<int64_t> sourceShape = op.getSourceVectorType().getShape();
+ for (int64_t distDim : distributedDims) {
+ int64_t lanes = laneLayout[distDim];
+ // A packed dim would need the offset scaled by lane_data as well.
+ if (distDim < static_cast<int64_t>(sourceLaneData.size()) &&
+ sourceLaneData[distDim] != 1)
+ return rewriter.notifyMatchFailure(
+ op, "expecting unit lane data along a distributed dimension");
+ if (lanes == 0 || sourceShape[distDim] % lanes != 0)
+ return rewriter.notifyMatchFailure(
+ op, "source size along a distributed dim is not a multiple of "
+ "its lane count");
+ int64_t distrDimOffset =
+ cast<IntegerAttr>(updatedOffsets[distDim]).getInt();
+ if (distrDimOffset % lanes != 0)
+ return rewriter.notifyMatchFailure(
+ op, "offset along a distributed dim is not a multiple of its "
+ "lane count");
+ updatedSizes[distDim] =
+ rewriter.getI64IntegerAttr(distResultTy.getDimSize(distDim));
+ updatedOffsets[distDim] =
+ rewriter.getI64IntegerAttr(distrDimOffset / lanes);
+ }
}
auto newOp = vector::ExtractStridedSliceOp::create(
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index e69fabbce0dc3..ff781b06a5217 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1736,3 +1736,27 @@ gpu.func @convert_layout_partial_subgroup() {
gpu.return
}
}
+
+// -----
+// Two distributed dims: lane_layout [1, 2, 8] splits dims 1 and 2, so each
+// carries its own lane count and both the size and the offset along it are
+// rescaled by that count. Dim 0 is not split and keeps its offset.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @extract_strided_slice_two_distributed_dims
+// CHECK: %[[SRC:.*]] = "test.some_op"()
+// CHECK: %[[DIST:.*]] = builtin.unrealized_conversion_cast %[[SRC]] : vector<16x32x32xbf16> to vector<16x16x4xbf16>
+// CHECK: vector.extract_strided_slice %[[DIST]] offsets = [3, 1, 1], sizes = [1, 1, 1], strides = [1, 1, 1] : vector<16x16x4xbf16> to vector<1x1x1xbf16>
+gpu.func @extract_strided_slice_two_distributed_dims() {
+ %src = "test.some_op"() : () -> vector<16x32x32xbf16>
+ %0 = vector.extract_strided_slice %src offsets = [3, 2, 8], sizes = [1, 2, 8], strides = [1, 1, 1]
+ : vector<16x32x32xbf16> to vector<1x2x8xbf16>
+ // Anchors the layout on %0; recoverTemporaryLayouts derives the operand's
+ // layout from it.
+ %1 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>
+ }> : vector<1x2x8xbf16>
+ gpu.return
+}
+}
>From dad920f8514e943d2aea710b5d828a91867a15f6 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 29 Sep 2026 22:03:28 +0000
Subject: [PATCH 2/3] [mlir][xegpu] Address review on extract_strided_slice
distribution
Two test changes in response to review:
* Add @vector_extract_strided_slice_lane_count_below_subgroup_size,
pinning the case the old code handled by clamping the subgroup size
down to lane_layout[distDim]. Every pre-existing extract_strided_slice
test uses lane_layout 16 on the distributed dim, so the clamp never
fired and nothing covered it. Verified to produce identical IR before
and after the parent commit.
* Rework @extract_strided_slice_two_distributed_dims so the extracted
shape is a multiple of the lane layout instead of equal to it
(sizes [1, 4, 16] rather than [1, 2, 8]). Previously every distributed
dim divided down to 1, which hid the per-dim division; now each lane
keeps a 2x2 slab, so a wrong divisor on either dim changes the output.
This subsumes the exact-match case, so the separate test for it is
dropped.
No functional change.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../XeGPU/sg-to-lane-distribute-unit.mlir | 24 +++++++++++++++----
1 file changed, 20 insertions(+), 4 deletions(-)
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index ff781b06a5217..5b3654a266afc 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -998,6 +998,22 @@ gpu.func @vector_extract_strided_slice_partial_offsets() {
gpu.return
}
+// Only 8 of the subgroup's lanes split dim 1, so that dim's size and offset are
+// rescaled by 8 rather than by the subgroup size.
+// CHECK-LABEL: gpu.func @vector_extract_strided_slice_lane_count_below_subgroup_size
+// CHECK: %[[ESS:.*]] = vector.extract_strided_slice %{{.*}} offsets = [2, 2], sizes = [4, 2], strides = [1, 1] : vector<8x4xf32> to vector<4x2xf32>
+gpu.func @vector_extract_strided_slice_lane_count_below_subgroup_size() {
+ %0 = "test.some_op"()
+ : () -> vector<8x32xf32>
+ %1 = vector.extract_strided_slice %0 offsets = [2, 16], sizes = [4, 16], strides = [1, 1]
+ : vector<8x32xf32> to vector<4x16xf32>
+ %cl1 = xegpu.convert_layout %1
+ <{
+ target_layout = #xegpu.layout<lane_layout = [1, 8], lane_data = [1, 1]>
+ }> : vector<4x16xf32>
+ gpu.return
+}
+
// A convert_layout that only repacks lane_data along the non-distributed outer
// dimension (from [4, 1] to [1, 1] with order = [1, 0]), keeping lane_layout
// unchanged, folds to its source when consumed by exactly 4 (== outer
@@ -1745,18 +1761,18 @@ gpu.module @xevm_module {
// CHECK-LABEL: gpu.func @extract_strided_slice_two_distributed_dims
// CHECK: %[[SRC:.*]] = "test.some_op"()
// CHECK: %[[DIST:.*]] = builtin.unrealized_conversion_cast %[[SRC]] : vector<16x32x32xbf16> to vector<16x16x4xbf16>
-// CHECK: vector.extract_strided_slice %[[DIST]] offsets = [3, 1, 1], sizes = [1, 1, 1], strides = [1, 1, 1] : vector<16x16x4xbf16> to vector<1x1x1xbf16>
+// CHECK: vector.extract_strided_slice %[[DIST]] offsets = [3, 2, 2], sizes = [1, 2, 2], strides = [1, 1, 1] : vector<16x16x4xbf16> to vector<1x2x2xbf16>
gpu.func @extract_strided_slice_two_distributed_dims() {
%src = "test.some_op"() : () -> vector<16x32x32xbf16>
- %0 = vector.extract_strided_slice %src offsets = [3, 2, 8], sizes = [1, 2, 8], strides = [1, 1, 1]
- : vector<16x32x32xbf16> to vector<1x2x8xbf16>
+ %0 = vector.extract_strided_slice %src offsets = [3, 4, 16], sizes = [1, 4, 16], strides = [1, 1, 1]
+ : vector<16x32x32xbf16> to vector<1x4x16xbf16>
// Anchors the layout on %0; recoverTemporaryLayouts derives the operand's
// layout from it.
%1 = xegpu.convert_layout %0
<{
input_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>,
target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>
- }> : vector<1x2x8xbf16>
+ }> : vector<1x4x16xbf16>
gpu.return
}
}
>From 4accf41a0801a81bd20ab2b7fb1124fe494acab0 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 1 Oct 2026 01:34:36 +0000
Subject: [PATCH 3/3] [mlir][xegpu] Distribute a packed distributed dim in
extract_strided_slice
The per-dim loop refused any distributed dim whose lane_data is not 1, on
the grounds that a packed dim would need the offset scaled by lane_data as
well. It does not: a lane owns lane_data contiguous elements along the dim,
so packing changes where a slice may start, not how the offset divides --
offset / (lanes * laneData) * laneData == offset / lanes.
Replace the refusal with the alignment check it stands in for: the offset
must be a multiple of lanes * lane_data, so the slice begins on a lane-tile
boundary rather than cutting one lane's run. The size needs no extra
condition, since computing the distributed type already required the shape
to be a multiple of lane_layout * lane_data.
This lets a slice over a dim that is both distributed and packed through,
which is what a shape_cast result layout carrying non-unit lane_data on a
split dim produces.
Co-Authored-By: Claude Opus 5 (1M context) <noreply at anthropic.com>
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 16 +++++++-------
.../XeGPU/sg-to-lane-distribute-unit.mlir | 21 +++++++++++++++++++
2 files changed, 30 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index e2b61e281dfcf..da1c97ffe0da1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1373,21 +1373,23 @@ struct SgToLaneVectorExtractStridedSlice
ArrayRef<int64_t> sourceShape = op.getSourceVectorType().getShape();
for (int64_t distDim : distributedDims) {
int64_t lanes = laneLayout[distDim];
- // A packed dim would need the offset scaled by lane_data as well.
- if (distDim < static_cast<int64_t>(sourceLaneData.size()) &&
- sourceLaneData[distDim] != 1)
- return rewriter.notifyMatchFailure(
- op, "expecting unit lane data along a distributed dimension");
if (lanes == 0 || sourceShape[distDim] % lanes != 0)
return rewriter.notifyMatchFailure(
op, "source size along a distributed dim is not a multiple of "
"its lane count");
+ // A lane owns lane_data contiguous elements along this dim, so the slice
+ // has to start on a lane-tile boundary or it would cut one of those
+ // runs. Packing changes only that alignment requirement, not the
+ // divisor: offset / (lanes * laneData) * laneData == offset / lanes.
+ int64_t laneData = distDim < static_cast<int64_t>(sourceLaneData.size())
+ ? sourceLaneData[distDim]
+ : 1;
int64_t distrDimOffset =
cast<IntegerAttr>(updatedOffsets[distDim]).getInt();
- if (distrDimOffset % lanes != 0)
+ if (laneData == 0 || distrDimOffset % (lanes * laneData) != 0)
return rewriter.notifyMatchFailure(
op, "offset along a distributed dim is not a multiple of its "
- "lane count");
+ "lane tile");
updatedSizes[distDim] =
rewriter.getI64IntegerAttr(distResultTy.getDimSize(distDim));
updatedOffsets[distDim] =
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
index 5b3654a266afc..2431ac9492126 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1776,3 +1776,24 @@ gpu.func @extract_strided_slice_two_distributed_dims() {
gpu.return
}
}
+
+// -----
+// A distributed dim that is also packed: lane_data 4 on dim 1, so a lane owns 4
+// contiguous columns. The slice covers the whole dim, so only its size rescales.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @extract_strided_slice_packed_distributed_dim
+// CHECK: %[[SRC:.*]] = "test.some_op"()
+// CHECK: %[[DIST:.*]] = builtin.unrealized_conversion_cast %[[SRC]] : vector<8x64xbf16> to vector<8x4xbf16>
+// CHECK: vector.extract_strided_slice %[[DIST]] offsets = [0, 0], sizes = [1, 4], strides = [1, 1] : vector<8x4xbf16> to vector<1x4xbf16>
+gpu.func @extract_strided_slice_packed_distributed_dim() {
+ %src = "test.some_op"() : () -> vector<8x64xbf16>
+ %0 = vector.extract_strided_slice %src offsets = [0, 0], sizes = [1, 64], strides = [1, 1]
+ : vector<8x64xbf16> to vector<1x64xbf16>
+ %1 = xegpu.convert_layout %0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>
+ }> : vector<1x64xbf16>
+ gpu.return
+}
+}
More information about the Mlir-commits
mailing list