[Mlir-commits] [mlir] [mlir][XeGPU] Add an SLM round-trip fallback for sg-to-lane convert_layout (PR #227884)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed Sep 30 14:49:06 PDT 2026
llvmorg-github-actions[bot] wrote:
<!--LLVM PR SUMMARY COMMENT-->
@llvm/pr-subscribers-mlir-gpu
Author: Jianhui Li (Jianhui-Li)
<details>
<summary>Changes</summary>
This PR adds a general fallback lowering for xegpu.convert_layout in the sg-to-lane distribution pass, which round-trips the value through shared local memory.
The existing lane-level lowerings are all special cases: the layouts fold into each other, or they differ in a way a xegpu.lane_shuffle or a gpu.shuffle can express. Anything else failed to legalize and stopped the pipeline — for instance a conversion that only moves the lane_data of a dimension distributed over part of the subgroup:
```mlir
xegpu.convert_layout %src
<{input_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 4]>,
target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>}>
: vector<1x2x32xbf16>
```
Here dim 2 is distributed over 8 of the 16 lanes and dim 1 over the other factor of 2, so the subgroup splits into two groups that each hold a run of their own. A single xegpu.lane_shuffle round-robins with a stride of the whole subgroup size and cannot express that.
The new pattern works for any pair of layouts which can both distribute the value: each lane writes its fragment to SLM at the coordinates input_layout assigns it, then reads it back from the ones target_layout assigns. It mirrors what XeGPUWgToSgDistribute already does for a workgroup-level convert_layout whose sg_layout changes, and is registered at a lower pattern benefit than every specialized lowering, so it only runs once they have all declined.
assisted-by-claude
---
Full diff: https://github.com/llvm/llvm-project/pull/227884.diff
5 Files Affected:
- (modified) mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h (+7)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp (+4-12)
- (modified) mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp (+202)
- (modified) mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp (+18)
- (modified) mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir (+51)
``````````diff
diff --git a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
index d54d159fe2120..4db69bda113dc 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Utils/XeGPUUtils.h
@@ -104,6 +104,13 @@ Value createVectorWithShapeFromValues(OpBuilder &builder, Location loc,
/// if no GPU module parent or XeVM target attribute exists.
std::optional<std::string> getChipStr(Operation *op);
+/// Returns the number of subgroups the kernel enclosing `op` runs, derived from
+/// the `known_block_size` of its parent `gpu.func`. Fails when there is no such
+/// parent, when the attribute is absent, when a block dimension is not a
+/// positive power of two, or when the block does not cover a whole subgroup.
+FailureOr<int64_t> getNumSubgroupsFromBlockSize(Operation *op,
+ int64_t subgroupSize);
+
/// Generates element-wise addition ops of two arrays with same length.
SmallVector<OpFoldResult> addElementwise(OpBuilder &builder, Location loc,
ArrayRef<OpFoldResult> lhs,
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index f05193504116b..08512f67a87e1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -486,18 +486,10 @@ FailureOr<int64_t> LayoutInfoPropagation::getNumSgOrFail(
return llvm::product_of(sgLayout);
}
// Otherwise fall back to the kernel's known_block_size.
- if (auto gpuFunc = op->getParentOfType<gpu::GPUFuncOp>()) {
- std::optional<ArrayRef<int32_t>> knownBlockSize =
- gpuFunc.getKnownBlockSize();
- if (knownBlockSize) {
- bool isPowerOf2Block = llvm::all_of(*knownBlockSize, [](int32_t dim) {
- return dim > 0 && llvm::isPowerOf2_32(dim);
- });
- int64_t numSg = llvm::product_of(*knownBlockSize) / sgSize;
- if (isPowerOf2Block && numSg > 0)
- return numSg;
- }
- }
+ if (FailureOr<int64_t> numSg =
+ xegpu::getNumSubgroupsFromBlockSize(op, sgSize);
+ succeeded(numSg))
+ return *numSg;
// Only subgroup mode needs the count; elsewhere a missing one is benign.
if (layoutKind == xegpu::LayoutKind::Subgroup) {
markFailure(op, "Unable to determine the number of subgroups for the "
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 06f769cefb4fd..539175be768ac 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1939,6 +1939,203 @@ struct SgToLaneConvertLayout
}
};
+/// The offsets and shape of one `lane_data`-sized piece of a lane's fragment.
+/// `computeDistributedCoords` returns one coordinate per distribution unit,
+/// enumerated row major over `shape / (lane_layout * lane_data)`, and unit `u`
+/// owns the piece of the fragment at `delinearize(u) * lane_data`.
+struct LaneFragmentPiece {
+ SmallVector<int64_t> offsets;
+ SmallVector<int64_t> sizes;
+};
+
+static LaneFragmentPiece getLaneFragmentPiece(int64_t unit,
+ ArrayRef<int64_t> shape,
+ ArrayRef<int64_t> laneLayout,
+ ArrayRef<int64_t> laneData) {
+ int64_t rank = shape.size();
+ SmallVector<int64_t> units(rank);
+ for (int64_t d = 0; d < rank; ++d)
+ units[d] = shape[d] / std::min(shape[d], laneLayout[d] * laneData[d]);
+ SmallVector<int64_t> unitCoords = delinearize(unit, computeStrides(units));
+ LaneFragmentPiece piece;
+ for (int64_t d = 0; d < rank; ++d) {
+ piece.offsets.push_back(unitCoords[d] * laneData[d]);
+ piece.sizes.push_back(laneData[d]);
+ }
+ return piece;
+}
+
+/// Last-resort lowering for a `convert_layout` that none of the patterns above
+/// handle: round-trip the value through shared local memory, writing each
+/// lane's fragment to the coordinates `input_layout` gives it and reading it
+/// back from the ones `target_layout` gives it. Any pair of layouts that can
+/// both distribute the value works, at the cost of an SLM write and read.
+///
+/// The op is subgroup level, but every subgroup of the workgroup executes it on
+/// its own data, so the scratch holds one tile per subgroup and each subgroup
+/// addresses its own through a `gpu.subgroup_id` offset on the outermost
+/// dimension. The subgroup count comes from the kernel's `known_block_size`.
+///
+/// A lane's fragment is one `lane_data`-sized piece per distribution unit, so
+/// the transfer is one `store_matrix`/`load_matrix` per unit; layouts with a
+/// single unit, which is the common case, give a single pair.
+struct SgToLaneConvertLayoutViaSLM
+ : public OpConversionPattern<xegpu::ConvertLayoutOp> {
+ using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(xegpu::ConvertLayoutOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
+ xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
+ if (!inputLayout || !targetLayout || !inputLayout.isForSubgroup() ||
+ !targetLayout.isForSubgroup())
+ return rewriter.notifyMatchFailure(op, "both layouts must be lane level");
+
+ auto valueTy = dyn_cast<VectorType>(op.getResult().getType());
+ if (!valueTy)
+ return rewriter.notifyMatchFailure(op, "value type must be a vector");
+ Type elemTy = valueTy.getElementType();
+ // The scratch is sized in bytes, so sub-byte elements would not get a
+ // well-defined address.
+ if (!elemTy.isIntOrFloat() || elemTy.getIntOrFloatBitWidth() % 8 != 0)
+ return rewriter.notifyMatchFailure(
+ op, "element type must be a whole number of bytes");
+
+ ArrayRef<int64_t> sgShape = valueTy.getShape();
+ int64_t rank = sgShape.size();
+ if (rank != inputLayout.getRank() || rank != targetLayout.getRank())
+ return rewriter.notifyMatchFailure(
+ op, "both layouts must have the rank of the value");
+
+ FailureOr<VectorType> distInputTy =
+ xegpu::getDistVecTypeBasedOnLaneLayout(inputLayout, valueTy);
+ FailureOr<VectorType> distTargetTy =
+ xegpu::getDistVecTypeBasedOnLaneLayout(targetLayout, valueTy);
+ if (failed(distInputTy) || failed(distTargetTy))
+ return rewriter.notifyMatchFailure(
+ op, "value type must be distributable by both layouts");
+
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+ if (!uArch)
+ return rewriter.notifyMatchFailure(
+ op, "target attribute is required to determine the subgroup size");
+ FailureOr<int64_t> numSubgroups =
+ xegpu::getNumSubgroupsFromBlockSize(op, uArch->getSubgroupSize());
+ if (failed(numSubgroups))
+ return rewriter.notifyMatchFailure(
+ op, "the scratch holds one tile per subgroup, so @known_block_size "
+ "must be attached to the kernel, with power-of-two dimensions "
+ "covering at least one subgroup");
+
+ Location loc = op.getLoc();
+
+ // One tile per subgroup, stacked along the outermost dimension.
+ SmallVector<int64_t> slmShape(sgShape);
+ slmShape[0] *= *numSubgroups;
+ int64_t slmBytes =
+ computeProduct(slmShape) * elemTy.getIntOrFloatBitWidth() / 8;
+ auto slmTy = MemRefType::get({slmBytes}, rewriter.getI8Type(), {}, 3);
+ Value slm = memref::AllocaOp::create(rewriter, loc, slmTy);
+ Value memDesc = xegpu::CreateMemDescOp::create(
+ rewriter, loc,
+ xegpu::MemDescType::get(rewriter.getContext(), slmShape, elemTy,
+ nullptr),
+ slm);
+
+ // This subgroup's tile starts at `subgroup_id` tiles into the scratch.
+ Value sgId = gpu::SubgroupIdOp::create(rewriter, loc,
+ rewriter.getIndexType(), nullptr);
+ SmallVector<Value> base(rank);
+ base[0] = arith::MulIOp::create(
+ rewriter, loc, sgId,
+ arith::ConstantIndexOp::create(rewriter, loc, sgShape[0]));
+ for (int64_t d = 1; d < rank; ++d)
+ base[d] = arith::ConstantIndexOp::create(rewriter, loc, 0);
+
+ Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+ /*upperBound=*/mlir::IntegerAttr());
+ auto dynamicOffsets = rewriter.getDenseI64ArrayAttr(
+ SmallVector<int64_t>(rank, ShapedType::kDynamic));
+
+ // Adds this subgroup's base to one distribution unit's coordinates.
+ auto withBase = [&](ArrayRef<Value> coords) {
+ SmallVector<Value> offsets;
+ for (auto [coord, b] : llvm::zip_equal(coords, base))
+ offsets.push_back(arith::AddIOp::create(rewriter, loc, coord, b));
+ return offsets;
+ };
+
+ // Write phase: every lane stores the pieces `input_layout` gives it.
+ auto storeCoords =
+ inputLayout.computeDistributedCoords(rewriter, loc, laneId, sgShape);
+ if (failed(storeCoords))
+ return rewriter.notifyMatchFailure(
+ op, "failed to compute the input_layout coordinates");
+ SmallVector<int64_t> inputLaneLayout =
+ inputLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> inputLaneData =
+ inputLayout.getEffectiveLaneDataAsInt();
+ Value fragment =
+ castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getSource()),
+ *distInputTy);
+ for (auto [unit, coords] : llvm::enumerate(*storeCoords)) {
+ LaneFragmentPiece piece =
+ getLaneFragmentPiece(unit, sgShape, inputLaneLayout, inputLaneData);
+ Value data = fragment;
+ if (piece.sizes != SmallVector<int64_t>(distInputTy->getShape()))
+ data = vector::ExtractStridedSliceOp::create(
+ rewriter, loc, fragment, piece.offsets, piece.sizes,
+ SmallVector<int64_t>(rank, 1));
+ xegpu::StoreMatrixOp::create(rewriter, loc, TypeRange{}, data, memDesc,
+ ValueRange(withBase(coords)), dynamicOffsets,
+ nullptr, xegpu::DistributeLayoutAttr{});
+ }
+
+ // The scratch is only exchanged between the lanes of one subgroup, which is
+ // a single hardware thread, so ordering the write before the read is enough
+ // and no workgroup barrier is needed.
+ xegpu::FenceOp::create(rewriter, loc, xegpu::MemorySpace::SLM,
+ xegpu::FenceScope::Workgroup);
+
+ // Read phase: every lane loads the pieces `target_layout` gives it.
+ auto loadCoords =
+ targetLayout.computeDistributedCoords(rewriter, loc, laneId, sgShape);
+ if (failed(loadCoords))
+ return rewriter.notifyMatchFailure(
+ op, "failed to compute the target_layout coordinates");
+ SmallVector<int64_t> targetLaneLayout =
+ targetLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> targetLaneData =
+ targetLayout.getEffectiveLaneDataAsInt();
+ Value result;
+ if (loadCoords->size() > 1)
+ result = arith::ConstantOp::create(rewriter, loc, *distTargetTy,
+ rewriter.getZeroAttr(*distTargetTy));
+ for (auto [unit, coords] : llvm::enumerate(*loadCoords)) {
+ LaneFragmentPiece piece =
+ getLaneFragmentPiece(unit, sgShape, targetLaneLayout, targetLaneData);
+ auto pieceTy = VectorType::get(piece.sizes, elemTy);
+ Value loaded = xegpu::LoadMatrixOp::create(
+ rewriter, loc, pieceTy, memDesc, ValueRange(withBase(coords)),
+ dynamicOffsets, nullptr, xegpu::DistributeLayoutAttr{});
+ if (!result) {
+ result = loaded;
+ break;
+ }
+ result = vector::InsertStridedSliceOp::create(
+ rewriter, loc, loaded, result, piece.offsets,
+ SmallVector<int64_t>(rank, 1));
+ }
+
+ rewriter.replaceOp(op, castValueTo(rewriter,
+ cast<TypedValue<VectorType>>(result),
+ *distTargetTy));
+ return success();
+ }
+};
+
/// `getEffectiveLaneDataAsInt` is empty when `lane_data` is unset, so the unit
/// check also rejects layouts that are not lane level.
static bool hasDefaultOrderAndUnitLaneData(xegpu::DistributeLayoutAttr layout) {
@@ -2687,4 +2884,9 @@ void xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
SgToLaneConvertLayoutPartialBroadcastExtractShuffle,
SgToLaneConvertLayoutDeinterleaveSelect>(typeConverter,
patterns.getContext());
+ // The SLM round-trip handles any pair of layouts but pays for a memory
+ // round-trip, so it only runs when every pattern above has declined.
+ patterns.add<SgToLaneConvertLayoutViaSLM>(typeConverter,
+ patterns.getContext(),
+ /*benefit=*/PatternBenefit(0));
}
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index bf50fd3b3af7e..d1d4f044c8ed0 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -490,6 +490,24 @@ std::optional<std::string> xegpu::getChipStr(Operation *op) {
return std::nullopt;
}
+FailureOr<int64_t> xegpu::getNumSubgroupsFromBlockSize(Operation *op,
+ int64_t subgroupSize) {
+ auto gpuFunc = op->getParentOfType<gpu::GPUFuncOp>();
+ if (!gpuFunc)
+ return failure();
+ std::optional<ArrayRef<int32_t>> blockSize = gpuFunc.getKnownBlockSize();
+ if (!blockSize)
+ return failure();
+ if (!llvm::all_of(*blockSize, [](int32_t dim) {
+ return dim > 0 && llvm::isPowerOf2_32(dim);
+ }))
+ return failure();
+ int64_t numSubgroups = llvm::product_of(*blockSize) / subgroupSize;
+ if (numSubgroups < 1)
+ return failure();
+ return numSubgroups;
+}
+
/// Generates element-wise addition ops of two arrays with same length.
SmallVector<OpFoldResult> xegpu::addElementwise(OpBuilder &builder,
Location loc,
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..cb5872a0973fb 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1283,6 +1283,57 @@ gpu.func @convert_layout_repack_innermost_lane_data_3d() {
}
}
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_via_slm_lane_layout
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<16x16xf32> to vector<16x1xf32>
+// CHECK: %[[SLM:.*]] = memref.alloca() : memref<4096xi8, 3>
+// CHECK: %[[MD:.*]] = xegpu.create_mem_desc %[[SLM]] : memref<4096xi8, 3> -> !xegpu.mem_desc<64x16xf32>
+// CHECK: %[[SGID:.*]] = gpu.subgroup_id : index
+// CHECK: %[[ROWS:.*]] = arith.constant 16 : index
+// CHECK: %[[BASE:.*]] = arith.muli %[[SGID]], %[[ROWS]] : index
+// CHECK: gpu.lane_id
+// CHECK-COUNT-16: xegpu.store_matrix %{{.*}}, %[[MD]]{{\[}}%{{.*}}, %{{.*}}] : vector<1x1xf32>, !xegpu.mem_desc<64x16xf32>, index, index
+// CHECK: xegpu.fence memory_kind = slm, fence_scope = workgroup
+// CHECK-COUNT-16: xegpu.load_matrix %[[MD]]{{\[}}%{{.*}}, %{{.*}}] : !xegpu.mem_desc<64x16xf32>, index, index -> vector<1x1xf32>
+// CHECK: builtin.unrealized_conversion_cast %{{.*}} : vector<1x16xf32> to vector<16x16xf32>
+gpu.func @convert_layout_via_slm_lane_layout() kernel attributes {known_block_size = array<i32: 64, 1, 1>} {
+ %src = "test.some_op"() : () -> vector<16x16xf32>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
+ }> : vector<16x16xf32>
+ "test.some_use"(%cvt) : (vector<16x16xf32>) -> ()
+ gpu.return
+}
+}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_via_slm_lane_data
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<1x2x32xbf16> to vector<1x1x4xbf16>
+// CHECK: %[[SLM:.*]] = memref.alloca() : memref<512xi8, 3>
+// CHECK: %[[MD:.*]] = xegpu.create_mem_desc %[[SLM]] : memref<512xi8, 3> -> !xegpu.mem_desc<4x2x32xbf16>
+// CHECK: gpu.subgroup_id : index
+// CHECK: xegpu.store_matrix %[[SRC]], %[[MD]]{{\[}}%{{.*}}, %{{.*}}, %{{.*}}] : vector<1x1x4xbf16>, !xegpu.mem_desc<4x2x32xbf16>, index, index, index
+// CHECK: xegpu.fence memory_kind = slm, fence_scope = workgroup
+// CHECK-COUNT-4: xegpu.load_matrix %[[MD]]{{\[}}%{{.*}}, %{{.*}}, %{{.*}}] : !xegpu.mem_desc<4x2x32xbf16>, index, index, index -> vector<1x1x1xbf16>
+// CHECK: builtin.unrealized_conversion_cast %{{.*}} : vector<1x1x4xbf16> to vector<1x2x32xbf16>
+gpu.func @convert_layout_via_slm_lane_data() kernel attributes {known_block_size = array<i32: 64, 1, 1>} {
+ %src = "test.some_op"() : () -> vector<1x2x32xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 4]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 2, 8], lane_data = [1, 1, 1]>
+ }> : vector<1x2x32xbf16>
+ "test.some_use"(%cvt) : (vector<1x2x32xbf16>) -> ()
+ gpu.return
+}
+}
+
// -----
// load_matrix and store_matrix with coordinate computation (offsets [0,0])
gpu.module @xevm_module {
``````````
</details>
https://github.com/llvm/llvm-project/pull/227884
More information about the Mlir-commits
mailing list