[Mlir-commits] [mlir] [mlir][xegpu] Lower 1D vector transfers to scattered load/store (PR #213469)
Jianhui Li
llvmlistbot at llvm.org
Sat Aug 1 15:52:06 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/213469
>From 44905d20e6827b15257865295f528ac129e0ec83 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Sat, 1 Aug 2026 22:44:37 +0000
Subject: [PATCH] [mlir][xegpu] Lower lane_data repack convert_layout to
lane_shuffle
Extend the SgToLaneConvertLayout pattern to lower a convert_layout that
only repacks lane_data between round-robin and contiguous form (keeping
lane_layout and order unchanged) into xegpu.lane_shuffle. Each lane keeps
the same elements along the repacked dimension, but their assignment to
lanes changes, so the data is moved across lanes with a type-preserving
lane_shuffle (pack for round-robin -> contiguous, unpack for the reverse).
Because lane_shuffle preserves the element type, no bitcast to a storage
integer type is needed. The repacked dimension must be the distributed
one of the two innermost dimensions. When the lane fragment is a single
run it is shape_cast to 1D and shuffled once; when the run is strided the
k-long slice is extracted, shuffled, and inserted back per run.
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply at anthropic.com>
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 142 ++++++++++++++++++
.../XeGPU/sg-to-lane-distribute-unit.mlir | 81 ++++++++++
2 files changed, 223 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index a5a998f9b4427..13b7fe52a3e6c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1665,6 +1665,101 @@ shuffleDataAsLaneLayoutChange(ConversionPatternRewriter &rewriter, Location loc,
return res;
}
+/// Repacks `src`'s `lane_data` along `repackDim` between round-robin and
+/// contiguous form with an `xegpu.lane_shuffle`, which moves each lane's run of
+/// `k` elements across lanes while preserving the element type.
+///
+/// `inputData`/`targetData` are the `repackDim` `lane_data` of the input and
+/// target layouts; exactly one must be 1 (round-robin) and the other `k`
+/// (contiguous). Returns failure if that does not hold.
+static FailureOr<Value>
+repackLaneData(ConversionPatternRewriter &rewriter, Location loc, Value src,
+ int64_t repackDim, int64_t inputData, int64_t targetData) {
+ auto srcTy = dyn_cast<VectorType>(src.getType());
+ if (!srcTy)
+ return failure();
+ int64_t rank = srcTy.getRank();
+ Type elemTy = srcTy.getElementType();
+ int64_t k = srcTy.getShape()[repackDim];
+
+ bool roundRobinToContig = inputData == 1 && targetData == k;
+ bool contigToRoundRobin = inputData == k && targetData == 1;
+ if (!roundRobinToContig && !contigToRoundRobin)
+ return failure();
+
+ // Round-robin -> contiguous gathers a lane's strided elements into
+ // consecutive positions (pack); the reverse scatters them back (unpack).
+ xegpu::LaneShuffleMode mode = roundRobinToContig
+ ? xegpu::LaneShuffleMode::Pack
+ : xegpu::LaneShuffleMode::Unpack;
+ VectorType runTy = VectorType::get({k}, elemTy);
+
+ // Common case: the lane fragment is a single run (every dimension other than
+ // `repackDim` is unit), so collapse it to 1D, shuffle once, and restore it.
+ if (srcTy.getNumElements() == k) {
+ if (rank == 1)
+ return Value(xegpu::LaneShuffleOp::create(rewriter, loc, runTy, src,
+ mode));
+ Value flat = vector::ShapeCastOp::create(rewriter, loc, runTy, src);
+ Value shuffled =
+ xegpu::LaneShuffleOp::create(rewriter, loc, runTy, flat, mode);
+ return Value(vector::ShapeCastOp::create(rewriter, loc, srcTy, shuffled));
+ }
+
+ // When `repackDim` is innermost each run is a contiguous sub-vector, so it is
+ // extracted and re-inserted as a whole.
+ if (repackDim == rank - 1) {
+ SmallVector<int64_t> outerShape(srcTy.getShape().drop_back());
+ int64_t numRuns = computeProduct(outerShape);
+ SmallVector<int64_t> outerStrides = computeStrides(outerShape);
+ Value result = arith::ConstantOp::create(rewriter, loc, srcTy,
+ rewriter.getZeroAttr(srcTy));
+ for (int64_t i = 0; i < numRuns; ++i) {
+ SmallVector<int64_t> pos = delinearize(i, outerStrides);
+ Value run = vector::ExtractOp::create(rewriter, loc, src, pos);
+ Value shuffled =
+ xegpu::LaneShuffleOp::create(rewriter, loc, runTy, run, mode);
+ result = vector::InsertOp::create(rewriter, loc, shuffled, result, pos);
+ }
+ return result;
+ }
+
+ // Otherwise each run is strided along `repackDim`: extract the `k`-long slice
+ // (a sub-vector that is unit along every other dim), flatten it to 1D,
+ // shuffle, and insert it back.
+ SmallVector<int64_t> keptShape;
+ SmallVector<int64_t> keptDims;
+ for (int64_t d = 0; d < rank; ++d)
+ if (d != repackDim) {
+ keptShape.push_back(srcTy.getShape()[d]);
+ keptDims.push_back(d);
+ }
+ int64_t numRuns = computeProduct(keptShape);
+ SmallVector<int64_t> keptStrides = computeStrides(keptShape);
+ SmallVector<int64_t> sliceSizes(rank, 1);
+ sliceSizes[repackDim] = k;
+ SmallVector<int64_t> sliceStrides(rank, 1);
+ VectorType sliceTy = VectorType::get(sliceSizes, elemTy);
+ Value result = arith::ConstantOp::create(rewriter, loc, srcTy,
+ rewriter.getZeroAttr(srcTy));
+ for (int64_t i = 0; i < numRuns; ++i) {
+ SmallVector<int64_t> keptPos = delinearize(i, keptStrides);
+ SmallVector<int64_t> offsets(rank, 0);
+ for (auto [dim, coord] : llvm::zip_equal(keptDims, keptPos))
+ offsets[dim] = coord;
+ Value slice = vector::ExtractStridedSliceOp::create(
+ rewriter, loc, src, offsets, sliceSizes, sliceStrides);
+ Value run = vector::ShapeCastOp::create(rewriter, loc, runTy, slice);
+ Value repacked =
+ xegpu::LaneShuffleOp::create(rewriter, loc, runTy, run, mode);
+ Value repackedSlice =
+ vector::ShapeCastOp::create(rewriter, loc, sliceTy, repacked);
+ result = vector::InsertStridedSliceOp::create(
+ rewriter, loc, repackedSlice, result, offsets, sliceStrides);
+ }
+ return result;
+}
+
/// Folds a subgroup-level ConvertLayout op with compatible lane layouts.
struct SgToLaneConvertLayout
: public OpConversionPattern<xegpu::ConvertLayoutOp> {
@@ -1721,6 +1816,53 @@ struct SgToLaneConvertLayout
}
}
+ // Handle a pure `lane_data` repack: `lane_layout` and `order` are unchanged
+ // and exactly one dimension's `lane_data` switches between round-robin
+ // (lane_data 1) and contiguous (lane_data == run length). The elements per
+ // lane are unchanged, but their assignment to lanes is not, so the data is
+ // moved across lanes with `xegpu.lane_shuffle`. The changed dimension must
+ // be one of the two innermost ones, since sg-to-lane distribution is 2D.
+ if (inputLayout.getEffectiveOrderAsInt() ==
+ targetLayout.getEffectiveOrderAsInt() &&
+ inputLayout.getEffectiveLaneLayoutAsInt() ==
+ targetLayout.getEffectiveLaneLayoutAsInt()) {
+ auto laneLayout = inputLayout.getEffectiveLaneLayoutAsInt();
+ auto laneData = inputLayout.getEffectiveLaneDataAsInt();
+ auto targetLaneData = targetLayout.getEffectiveLaneDataAsInt();
+ // Find the single dimension whose lane_data changed; bail out if more
+ // than one differs.
+ int64_t rank = laneData.size();
+ int64_t repackDim = -1;
+ bool multipleChanged = false;
+ for (int64_t d = 0; d < rank; ++d)
+ if (laneData[d] != targetLaneData[d]) {
+ if (repackDim != -1)
+ multipleChanged = true;
+ repackDim = d;
+ }
+
+ // `repackDim` must be the distributed dim (lane_layout != 1) and the
+ // other innermost dim non-distributed (lane_layout == 1).
+ int64_t otherDim = repackDim == rank - 1 ? rank - 2 : rank - 1;
+ bool laneLayoutOk =
+ repackDim != -1 && laneLayout[repackDim] != 1 &&
+ (rank < 2 || laneLayout[otherDim] == 1);
+
+ // Exactly one dimension must change, and it must be one of the two
+ // innermost (>= rank - 2).
+ if (repackDim != -1 && repackDim >= rank - 2 && !multipleChanged &&
+ laneLayoutOk) {
+ FailureOr<Value> res =
+ repackLaneData(rewriter, op.getLoc(), adaptor.getSource(),
+ repackDim, laneData[repackDim],
+ targetLaneData[repackDim]);
+ if (succeeded(res)) {
+ rewriter.replaceOp(op, *res);
+ 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 a7fe2d9534351..8607d45d5828c 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1045,6 +1045,87 @@ gpu.func @convert_layout_scalar() {
}
}
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_repack_innermost_lane_data
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<1x64xbf16> to vector<1x4xbf16>
+// CHECK: %[[F:.*]] = vector.shape_cast %[[SRC]] : vector<1x4xbf16> to vector<4xbf16>
+// CHECK: %[[S0:.*]] = xegpu.lane_shuffle %[[F]] pack : vector<4xbf16>
+// CHECK: vector.shape_cast %[[S0]] : vector<4xbf16> to vector<1x4xbf16>
+gpu.func @convert_layout_repack_innermost_lane_data() {
+ %src = "some_op"() : () -> vector<1x64xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>
+ }> : vector<1x64xbf16>
+ "some_use"(%cvt) : (vector<1x64xbf16>) -> ()
+ gpu.return
+}
+}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_repack_innermost_lane_data_1d
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<64xbf16> to vector<4xbf16>
+// CHECK: xegpu.lane_shuffle %[[SRC]] pack : vector<4xbf16>
+gpu.func @convert_layout_repack_innermost_lane_data_1d() {
+ %src = "some_op"() : () -> vector<64xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16], lane_data = [1]>,
+ target_layout = #xegpu.layout<lane_layout = [16], lane_data = [4]>
+ }> : vector<64xbf16>
+ "some_use"(%cvt) : (vector<64xbf16>) -> ()
+ gpu.return
+}
+}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_repack_second_innermost_lane_data
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<64x4xbf16> to vector<4x4xbf16>
+// CHECK: %[[SL0:.*]] = vector.extract_strided_slice %[[SRC]] {offsets = [0, 0], sizes = [4, 1], strides = [1, 1]} : vector<4x4xbf16> to vector<4x1xbf16>
+// CHECK: %[[F0:.*]] = vector.shape_cast %[[SL0]] : vector<4x1xbf16> to vector<4xbf16>
+// CHECK: %[[S0:.*]] = xegpu.lane_shuffle %[[F0]] unpack : vector<4xbf16>
+// CHECK: %[[C0:.*]] = vector.shape_cast %[[S0]] : vector<4xbf16> to vector<4x1xbf16>
+// CHECK: vector.insert_strided_slice %[[C0]], %{{.*}} {offsets = [0, 0], strides = [1, 1]} : vector<4x1xbf16> into vector<4x4xbf16>
+// CHECK-COUNT-3: xegpu.lane_shuffle %{{.*}} unpack : vector<4xbf16>
+gpu.func @convert_layout_repack_second_innermost_lane_data() {
+ %src = "some_op"() : () -> vector<64x4xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [4, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>
+ }> : vector<64x4xbf16>
+ "some_use"(%cvt) : (vector<64x4xbf16>) -> ()
+ gpu.return
+}
+}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @convert_layout_repack_innermost_lane_data_3d
+// CHECK-NOT: xegpu.convert_layout
+// CHECK: %[[SRC:.*]] = builtin.unrealized_conversion_cast %{{.*}} : vector<1x1x64xbf16> to vector<1x1x4xbf16>
+// CHECK: %[[F:.*]] = vector.shape_cast %[[SRC]] : vector<1x1x4xbf16> to vector<4xbf16>
+// CHECK: %[[S0:.*]] = xegpu.lane_shuffle %[[F]] pack : vector<4xbf16>
+// CHECK: vector.shape_cast %[[S0]] : vector<4xbf16> to vector<1x1x4xbf16>
+gpu.func @convert_layout_repack_innermost_lane_data_3d() {
+ %src = "some_op"() : () -> vector<1x1x64xbf16>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 4]>
+ }> : vector<1x1x64xbf16>
+ "some_use"(%cvt) : (vector<1x1x64xbf16>) -> ()
+ gpu.return
+}
+}
+
// -----
// load_matrix and store_matrix with coordinate computation (offsets [0,0])
gpu.module @xevm_module {
More information about the Mlir-commits
mailing list