[Mlir-commits] [mlir] [MLIR][XeGPU] Fold lane_data repack convert_layout (PR #204016)
Sang Ik Lee
llvmlistbot at llvm.org
Mon Jun 15 15:47:56 PDT 2026
https://github.com/silee2 created https://github.com/llvm/llvm-project/pull/204016
into extract_strided_slice consumers
Fold a subgroup-level convert_layout that merely repacks lane_data along the non-distributed outer dimension (from [N, 1] to [1, 1] with order = [1, 0]), keeping lane_layout unchanged, when its result is consumed by exactly N vector.extract_strided_slice ops. After lane distribution both layouts yield the same per-lane vector (getDistributedVectorType ignores lane_data) and the N slices recover the original blocking, so the convert is redundant and folds to its source.
Add a lit test covering the fold in sg-to-lane-distribute-unit.mlir.
>From f6fdc33fd065d95a223b537b29939185e88be4ff Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 15 Jun 2026 22:36:29 +0000
Subject: [PATCH] [MLIR][XeGPU] Fold lane_data repack convert_layout into
extract_strided_slice consumers
Fold a subgroup-level convert_layout that merely repacks lane_data along the
non-distributed outer dimension (from [N, 1] to [1, 1] with order = [1, 0]),
keeping lane_layout unchanged, when its result is consumed by exactly N
vector.extract_strided_slice ops. After lane distribution both layouts yield
the same per-lane vector (getDistributedVectorType ignores lane_data) and the N
slices recover the original blocking, so the convert is redundant and folds to
its source.
Add a lit test covering the fold in sg-to-lane-distribute-unit.mlir.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 39 +++++++++++++
.../XeGPU/sg-to-lane-distribute-unit.mlir | 55 +++++++++++++++++++
2 files changed, 94 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 75a87f84b3da8..12a9b71220987 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1638,6 +1638,45 @@ struct SgToLaneConvertLayout
}
}
+ // Special case: fold a convert_layout that merely repacks `lane_data`
+ // along the non-distributed outer dimension - from `[N, 1]` to `[1, 1]`
+ // with `order = [1, 0]`, keeping `lane_layout` unchanged - when its result
+ // is consumed by exactly `N` `vector.extract_strided_slice` ops. After lane
+ // distribution both layouts yield the same per-lane vector
+ // (`getDistributedVectorType` ignores `lane_data`), and the `N` slices
+ // recover the original blocking, so the convert is redundant and folds to
+ // its source.
+ if (inputLayout.getRank() == 2 && targetLayout.getRank() == 2) {
+ auto laneLayout = inputLayout.getEffectiveLaneLayoutAsInt();
+ auto targetLaneLayout = targetLayout.getEffectiveLaneLayoutAsInt();
+ auto laneData = inputLayout.getEffectiveLaneDataAsInt();
+ auto targetLaneData = targetLayout.getEffectiveLaneDataAsInt();
+ auto targetOrder = targetLayout.getEffectiveOrderAsInt();
+ if (laneLayout.size() == 2 && targetLaneLayout.size() == 2 &&
+ laneData.size() == 2 && targetLaneData.size() == 2 &&
+ laneLayout == targetLaneLayout && laneLayout[0] == 1 &&
+ laneData[0] > 1 && laneData[1] == 1 &&
+ targetLaneData == SmallVector<int64_t>({1, 1}) &&
+ targetOrder == SmallVector<int64_t>({1, 0})) {
+ // The block factor `N` is the outer (non-distributed) lane_data, and
+ // it must match the number of extract_strided_slice consumers.
+ int64_t blockFactor = laneData[0];
+ int64_t numSlices = 0;
+ bool allSlices = true;
+ for (Operation *user : op.getResult().getUsers()) {
+ ++numSlices;
+ if (!isa<vector::ExtractStridedSliceOp>(user)) {
+ allSlices = false;
+ break;
+ }
+ }
+ if (allSlices && numSlices == blockFactor) {
+ rewriter.replaceOp(op, adaptor.getSource());
+ return success();
+ }
+ }
+ }
+
return rewriter.notifyMatchFailure(
op, "lowering incompatible convert_layout not yet supported");
}
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 56db482f033d3..fcc2da3d5005b 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -794,6 +794,61 @@ gpu.func @vector_extract_strided_slice_partial_offsets() {
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
+// lane_data) extract_strided_slice ops. After distribution both layouts yield
+// the same per-lane vector, so the convert is redundant and the slices operate
+// directly on the distributed source.
+// CHECK-LABEL: gpu.func @convert_layout_repack_lane_data
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<32x16xi8> to vector<32x1xi8>
+// CHECK: vector.extract_strided_slice %[[SRC]] {offsets = [0, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+// CHECK: vector.extract_strided_slice %[[SRC]] {offsets = [8, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+// CHECK: vector.extract_strided_slice %[[SRC]] {offsets = [16, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+// CHECK: vector.extract_strided_slice %[[SRC]] {offsets = [24, 0], sizes = [8, 1], strides = [1, 1]} : vector<32x1xi8> to vector<8x1xi8>
+gpu.func @convert_layout_repack_lane_data() {
+ %src = "some_op"() : () -> vector<32x16xi8>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+ }> : vector<32x16xi8>
+ %s0 = vector.extract_strided_slice %cvt
+ {offsets = [0, 0], sizes = [8, 16], strides = [1, 1]}
+ : vector<32x16xi8> to vector<8x16xi8>
+ %s1 = vector.extract_strided_slice %cvt
+ {offsets = [8, 0], sizes = [8, 16], strides = [1, 1]}
+ : vector<32x16xi8> to vector<8x16xi8>
+ %s2 = vector.extract_strided_slice %cvt
+ {offsets = [16, 0], sizes = [8, 16], strides = [1, 1]}
+ : vector<32x16xi8> to vector<8x16xi8>
+ %s3 = vector.extract_strided_slice %cvt
+ {offsets = [24, 0], sizes = [8, 16], strides = [1, 1]}
+ : vector<32x16xi8> to vector<8x16xi8>
+ %a0 = xegpu.convert_layout %s0
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+ }> : vector<8x16xi8>
+ %a1 = xegpu.convert_layout %s1
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+ }> : vector<8x16xi8>
+ %a2 = xegpu.convert_layout %s2
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+ }> : vector<8x16xi8>
+ %a3 = xegpu.convert_layout %s3
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>
+ }> : vector<8x16xi8>
+ gpu.return
+}
+
// CHECK-LABEL: gpu.func @vector_insert_strided_slice_distributed_dim_fully_inserted
// CHECK: %[[ISS:.*]] = vector.insert_strided_slice %{{.*}}, %{{.*}} {offsets = [24, 0], strides = [1, 1]} : vector<16x1xf32> into vector<64x1xf32>
gpu.func @vector_insert_strided_slice_distributed_dim_fully_inserted() {
More information about the Mlir-commits
mailing list