[Mlir-commits] [mlir] 8b0361b - [MLIR][XeGPU] Treat lane_data repacks as compatible layouts (#204016)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Jun 17 16:48:27 PDT 2026
Author: Sang Ik Lee
Date: 2026-06-17T16:48:22-07:00
New Revision: 8b0361bba242d0f69560a717b52f26b289a7ebc2
URL: https://github.com/llvm/llvm-project/commit/8b0361bba242d0f69560a717b52f26b289a7ebc2
DIFF: https://github.com/llvm/llvm-project/commit/8b0361bba242d0f69560a717b52f26b289a7ebc2.diff
LOG: [MLIR][XeGPU] Treat lane_data repacks as compatible layouts (#204016)
A subgroup-level convert_layout that only repacks lane_data while keeping
lane_layout unchanged (e.g. [N, 1] to [1, 1] with order = [1, 0]) is a no-op
after lane distribution: each lane owns the same elements in the same order.
Previously isCompatibleWith compared per-distribution-unit block starts, which
encode the lane_data blocking, so such layouts looked incompatible.
Handle this at the Lane level in isCompatibleWith by expanding the block
starts into per-element coordinates before comparing. The expansion only runs
when lane_data differ; otherwise the cheaper block-start comparison is exact.
The shared logic lives in a compareDistributedCoords helper used by both
LayoutAttr and SliceAttr. The Subgroup level is left for a follow-up (TODO).
Add a lit test covering the fold in sg-to-lane-distribute-unit.mlir.
Added:
Modified:
mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
index 075442d443658..4aa1c0d666a94 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUDialect.cpp
@@ -123,6 +123,63 @@ static SmallVector<SmallVector<int64_t>> genStaticCoordinates(
return coordinates;
}
+/// Expands per-distribution-unit block-start coordinates into the full list of
+/// element coordinates each block covers: every element of the `subShape`-sized
+/// region (row-major) offset by the block start. Comparing these instead of the
+/// bare block starts lets layouts that
diff er only in `lane_data` blocking, but
+/// 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;
+}
+
+/// Returns true if `self` and `other` distribute `shape` identically at
+/// `level`: every id in `[0, size)` owns the same coordinates under both.
+///
+/// At the Lane level, layouts that pack `lane_data`
diff erently can still own
+/// the same per-lane elements in the same order; their block starts
diff er but
+/// the expanded per-element coordinates match. So block starts are expanded
+/// (via `expandBlockCoords`) before comparing, but only when it can change the
+/// result (Lane level with
diff ering `lane_data`) - otherwise comparing the
+/// cheaper block starts is already exact.
+///
+/// TODO: Extend the same handling to the Subgroup level (sg_data repacks).
+static bool compareDistributedCoords(xegpu::DistributeLayoutAttr self,
+ const xegpu::DistributeLayoutAttr &other,
+ ArrayRef<int64_t> shape,
+ xegpu::LayoutKind level, int64_t size) {
+ bool expandCoords =
+ level == xegpu::LayoutKind::Lane &&
+ self.getEffectiveLaneDataAsInt() != other.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> selfSubShape, otherSubShape;
+ if (expandCoords) {
+ selfSubShape = self.getEffectiveLaneDataAsInt();
+ otherSubShape = other.getEffectiveLaneDataAsInt();
+ }
+ for (int64_t id : llvm::seq<int64_t>(0, size)) {
+ auto coords = self.computeStaticDistributedCoords(id, shape);
+ auto otherCoords = other.computeStaticDistributedCoords(id, shape);
+ if (expandCoords) {
+ coords = expandBlockCoords(coords, selfSubShape);
+ otherCoords = expandBlockCoords(otherCoords, otherSubShape);
+ }
+ if (coords != otherCoords)
+ return false;
+ }
+ return true;
+}
+
// Checks if the given memref type represents shared local memory (SLM).
bool XeGPUDialect::isSharedMemory(const MemRefType &memrefTy) {
Attribute attr = memrefTy.getMemorySpace();
@@ -951,13 +1008,7 @@ bool LayoutAttr::isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
}
auto compareCoordsForAllIds = [&](int64_t size) {
- for (int64_t id : llvm::seq<int64_t>(0, size)) {
- auto coords = computeStaticDistributedCoords(id, shape);
- auto otherCoords = other.computeStaticDistributedCoords(id, shape);
- if (coords != otherCoords)
- return false;
- }
- return true;
+ return compareDistributedCoords(*this, other, shape, level, size);
};
if (level == xegpu::LayoutKind::Subgroup) {
@@ -1172,13 +1223,7 @@ bool SliceAttr::isCompatibleWith(const xegpu::DistributeLayoutAttr &other,
}
auto compareCoordsForAllIds = [&](int64_t size) {
- for (int64_t id : llvm::seq<int64_t>(0, size)) {
- auto coords = computeStaticDistributedCoords(id, shape);
- auto otherCoords = other.computeStaticDistributedCoords(id, shape);
- if (coords != otherCoords)
- return false;
- }
- return true;
+ return compareDistributedCoords(*this, other, shape, level, size);
};
auto flattenedThis = flatten();
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