[Mlir-commits] [mlir] [MLIR][XeGPU] Fold lane_data repack convert_layout (PR #204016)
Sang Ik Lee
llvmlistbot at llvm.org
Wed Jun 17 09:53:23 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/204016
>From ca44d2391ded2cf1ccba7bd22e051e4c6716041c 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] Treat lane_data repacks as compatible layouts
A subgroup-level convert_layout that only repacks lane_data - e.g. from
[N, 1] to [1, 1] with order = [1, 0] while keeping lane_layout unchanged -
is a no-op after lane distribution: each lane owns the same elements in the
same order, and getDistributedVectorType ignores lane_data, so both layouts
yield the same per-lane vector.
Previously isCompatibleWith compared the per-distribution-unit block-start
coordinates from computeStaticDistributedCoords, which encode the lane_data
blocking via the sub-shape. Layouts that differ only by a lane_data repack
therefore looked incompatible even though the convert is redundant.
Fix this at the source: at the Lane level, expand each lane's block-start
coordinates into the full ordered element coordinates (using each layout's
lane_data) before comparing. Layouts that merely repack lane_data now compare
equal, so the existing isCompatibleWith check in SgToLaneConvertLayout folds
the convert to its source. This replaces the previous special case in the
pass, which matched a narrow lane_data pattern and relied on counting
vector.extract_strided_slice consumers.
Add a shared expandBlockCoords helper and apply the expanded comparison in
both LayoutAttr::isCompatibleWith and SliceAttr::isCompatibleWith. The
Subgroup and InstData paths are unchanged.
The fold is still covered by the lit test in sg-to-lane-distribute-unit.mlir.
---
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp | 52 ++++++++++++++++++
.../XeGPU/sg-to-lane-distribute-unit.mlir | 55 +++++++++++++++++++
2 files changed, 107 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 075442d443658..0f3a6bab6ad0e 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -123,6 +123,30 @@ static SmallVector<SmallVector<int64_t>> genStaticCoordinates(
return coordinates;
}
+/// Expands a list of per-distribution-unit block-start coordinates into the
+/// full list of element coordinates owned within each block. Each block start
+/// covers a `subShape`-sized region; this enumerates every element of that
+/// region (in row-major order) offset by the block start. Comparing these
+/// expanded coordinates - rather than the bare block starts - lets two layouts
+/// that differ only in how a lane's elements are blocked (i.e. `lane_data`),
+/// but that ultimately own the same elements in the same order, be recognized
+/// as equivalent.
+static SmallVector<SmallVector<int64_t>>
+expandBlockCoords(ArrayRef<SmallVector<int64_t>> blockStarts,
+ ArrayRef<int64_t> subShape) {
+ SmallVector<int64_t> unitTile(subShape.size(), 1);
+ SmallVector<SmallVector<int64_t>> expanded;
+ for (const SmallVector<int64_t> &start : blockStarts) {
+ for (SmallVector<int64_t> off : StaticTileOffsetRange(subShape, unitTile)) {
+ SmallVector<int64_t> coord(start.size());
+ for (size_t i = 0; i < start.size(); ++i)
+ coord[i] = start[i] + off[i];
+ expanded.push_back(std::move(coord));
+ }
+ }
+ return expanded;
+}
+
// Checks if the given memref type represents shared local memory (SLM).
bool XeGPUDialect::isSharedMemory(const MemRefType &memrefTy) {
Attribute attr = memrefTy.getMemorySpace();
@@ -951,9 +975,23 @@ bool LayoutAttr::isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
}
auto compareCoordsForAllIds = [&](int64_t size) {
+ // At the Lane level, `lane_data` only changes how a lane's elements are
+ // blocked across distribution units, not which elements it owns nor their
+ // order. Expand the per-unit block starts into the full element coordinate
+ // sequence (using each layout's `lane_data`) before comparing, so layouts
+ // that merely repack `lane_data` are treated as compatible.
+ SmallVector<int64_t> selfSubShape, otherSubShape;
+ if (level == xegpu::LayoutKind::Lane) {
+ selfSubShape = getEffectiveLaneDataAsInt();
+ otherSubShape = other.getEffectiveLaneDataAsInt();
+ }
for (int64_t id : llvm::seq<int64_t>(0, size)) {
auto coords = computeStaticDistributedCoords(id, shape);
auto otherCoords = other.computeStaticDistributedCoords(id, shape);
+ if (!selfSubShape.empty())
+ coords = expandBlockCoords(coords, selfSubShape);
+ if (!otherSubShape.empty())
+ otherCoords = expandBlockCoords(otherCoords, otherSubShape);
if (coords != otherCoords)
return false;
}
@@ -1172,9 +1210,23 @@ bool SliceAttr::isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
}
auto compareCoordsForAllIds = [&](int64_t size) {
+ // At the Lane level, `lane_data` only changes how a lane's elements are
+ // blocked across distribution units, not which elements it owns nor their
+ // order. Expand the per-unit block starts into the full element coordinate
+ // sequence (using each layout's `lane_data`) before comparing, so layouts
+ // that merely repack `lane_data` are treated as compatible.
+ SmallVector<int64_t> selfSubShape, otherSubShape;
+ if (level == xegpu::LayoutKind::Lane) {
+ selfSubShape = getEffectiveLaneDataAsInt();
+ otherSubShape = other.getEffectiveLaneDataAsInt();
+ }
for (int64_t id : llvm::seq<int64_t>(0, size)) {
auto coords = computeStaticDistributedCoords(id, shape);
auto otherCoords = other.computeStaticDistributedCoords(id, shape);
+ if (!selfSubShape.empty())
+ coords = expandBlockCoords(coords, selfSubShape);
+ if (!otherSubShape.empty())
+ otherCoords = expandBlockCoords(otherCoords, otherSubShape);
if (coords != otherCoords)
return false;
}
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