[Mlir-commits] [mlir] [MLIR][XeGPU] Sg to Lane distribute convert_layout redistributing broadcasted data (PR #215645)
Sang Ik Lee
llvmlistbot at llvm.org
Tue Aug 18 14:48:08 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/215645
>From 2b9238fda86aedea5841b493f7be1ce1fe7d7859 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 11 Aug 2026 18:56:08 +0000
Subject: [PATCH] [MLIR][XeGPU] Distribute convert_layout redistributing
broadcasted data
The scale operands of scaled matrix multiplication are produced with a layout that
replicates them over groups of lanes and consumed with one that gives each lane a
different part, expressed as an `xegpu.convert_layout` that sg-to-lane distribution
could not lower:
%cvt = xegpu.convert_layout %src <{
input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2],
lane_data = [4, 1, 1], order = [0, 2, 1]>, dims = [0]>,
target_layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>
}> : vector<8x2xf8E8M0FNU>
Every lane holds a copy of the data its group is responsible for, so it can read
any element of that copy locally with an extract driven by its lane id. When the
element a lane needs is not in its own copy, one `gpu.shuffle idx` per result
element moves it; because a lane can only contribute a single value to a shuffle,
each lane extracts the element its counterpart in the target layout is about to ask
for. When the source is replicated over the whole subgroup, every lane already
holds what it owns and no shuffle is emitted at all.
Indices and source lanes are derived at compile time from the coordinates both
layouts assign to each lane, so the generated code is only arithmetic on the lane
id, with no lookup tables. Layout changes that do not fit this form are reported as
a match failure rather than lowered incorrectly.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 257 ++++++++++++++++++
.../XeGPU/sg-to-lane-distribute-unit.mlir | 71 +++++
2 files changed, 328 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index bc85140f9f121..68931ea16b1cb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1761,6 +1761,241 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
return result;
}
+/// Returns the number of lanes the subgroup is split into by `layout`, i.e.
+/// the product of the lane_layout of the underlying (unsliced) layout. Lanes
+/// covered by sliced dimensions are counted as well, since they hold a
+/// broadcasted copy of the data rather than no data at all.
+static int64_t getLaneLayoutSize(xegpu::DistributeLayoutAttr layout) {
+ if (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
+ layout = sliceAttr.flatten().getParent();
+ return computeProduct(layout.getEffectiveLaneLayoutAsInt());
+}
+
+/// Computes, for every lane of the subgroup, the coordinates of the elements
+/// of `shape` the lane owns under `layout`. Entry `i` of a lane's list is the
+/// coordinate of element `i` of that lane's distributed vector. Returns
+/// failure for non-unit lane_data, where a lane's elements are not a plain
+/// enumeration of the distribution unit starts.
+static FailureOr<SmallVector<SmallVector<SmallVector<int64_t>>>>
+computePerLaneElementCoords(xegpu::DistributeLayoutAttr layout,
+ ArrayRef<int64_t> shape, int64_t numLanes) {
+ SmallVector<int64_t> laneData = layout.getEffectiveLaneDataAsInt();
+ if (!llvm::all_of(laneData, [](int64_t d) { return d == 1; }))
+ return failure();
+ SmallVector<SmallVector<SmallVector<int64_t>>> coords;
+ for (int64_t lane = 0; lane < numLanes; lane++)
+ coords.push_back(layout.computeStaticDistributedCoords(lane, shape));
+ return coords;
+}
+
+/// Describes how one element of the distributed result is obtained from the
+/// broadcasted source. With `t` the index of the current lane within the
+/// target lane layout, the lane extracts element `stride * t + offset` of the
+/// copy it holds and, unless that value is already the one it owns,
+/// `gpu.shuffle idx` pulls it from lane `t + laneOffset`.
+struct BroadcastRedistribution {
+ int64_t stride;
+ int64_t offset;
+ int64_t laneOffset;
+ bool needsShuffle;
+};
+
+/// Derives the redistribution of element `pos` of the distributed result from
+/// the per-lane element coordinates of both layouts. Returns failure when the
+/// data movement cannot be expressed as a lane-invariant extract followed by
+/// an optional shuffle.
+static FailureOr<BroadcastRedistribution> deriveBroadcastRedistribution(
+ ArrayRef<SmallVector<SmallVector<int64_t>>> srcCoords,
+ ArrayRef<SmallVector<SmallVector<int64_t>>> resCoords, int64_t pos,
+ int64_t numLanes, int64_t numTargetLanes) {
+ // Position of the element with coordinates `coord` within the copy held by
+ // `lane`, or -1 if that lane does not hold it.
+ auto findElement = [&](int64_t lane, ArrayRef<int64_t> coord) -> int64_t {
+ for (auto [idx, candidate] : llvm::enumerate(srcCoords[lane]))
+ if (ArrayRef<int64_t>(candidate) == coord)
+ return idx;
+ return -1;
+ };
+
+ // A lane only extracts elements for the target lane it shares its position
+ // in the target lane layout with, so the shuffled-from lane can only be a
+ // whole number of target lane layouts away.
+ for (int64_t laneOffset = 0; laneOffset < numLanes;
+ laneOffset += numTargetLanes) {
+ // Collect where in its own copy each lane finds the element it provides.
+ SmallVector<int64_t> elements;
+ for (int64_t t = 0; t < numTargetLanes; t++) {
+ int64_t element = findElement(t + laneOffset, resCoords[t][pos]);
+ if (element < 0)
+ break;
+ elements.push_back(element);
+ }
+ if (static_cast<int64_t>(elements.size()) != numTargetLanes)
+ continue;
+
+ // The extracted element is computed from the lane id at runtime, so it has
+ // to be an affine function of the lane's position in the target layout.
+ int64_t offset = elements[0];
+ int64_t stride = numTargetLanes > 1 ? elements[1] - offset : 0;
+ if (stride < 0)
+ continue;
+ if (!llvm::all_of(llvm::seq<int64_t>(0, numTargetLanes), [&](int64_t t) {
+ return elements[t] == stride * t + offset;
+ }))
+ continue;
+
+ // Lanes beyond the target lane layout hold a replicated result. They can
+ // keep the locally extracted value only if it happens to be the one they
+ // own, otherwise the value has to be shuffled in.
+ bool needsShuffle =
+ laneOffset != 0 ||
+ !llvm::all_of(llvm::seq<int64_t>(0, numLanes), [&](int64_t lane) {
+ return srcCoords[lane][stride * (lane % numTargetLanes) + offset] ==
+ resCoords[lane][pos];
+ });
+ return BroadcastRedistribution{stride, offset, laneOffset, needsShuffle};
+ }
+ return failure();
+}
+
+/// Redistributes `src` for a `convert_layout` whose input layout replicates
+/// (broadcasts) the value over groups of lanes while the target layout hands
+/// each lane a different part of it. This is the layout change required by the
+/// scale operands of scaled matrix multiplication, where the scales are
+/// produced broadcasted but consumed distributed.
+///
+/// Every lane holds a full copy of the data its group is responsible for, so
+/// any lane can read any element of that copy locally with a dynamic extract
+/// driven by the lane id. A lane can however only contribute a single value to
+/// a `gpu.shuffle`, so when the element a lane needs is not in the copy it
+/// holds, both sides have to agree on what is exchanged: every lane extracts
+/// the element its counterpart in the target lane layout needs, and one
+/// `gpu.shuffle idx` per element of the distributed result moves it to the
+/// lane that owns it. See `deriveBroadcastRedistribution` for the exact form.
+///
+/// Returns failure if the redistribution is not expressible in that form.
+static FailureOr<Value>
+redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
+ Value src, VectorType resTy,
+ xegpu::DistributeLayoutAttr inputLayout,
+ xegpu::DistributeLayoutAttr targetLayout,
+ ArrayRef<int64_t> shape, int64_t numLanes) {
+ auto srcTy = dyn_cast<VectorType>(src.getType());
+ if (!srcTy)
+ return failure();
+ int64_t srcNumElems = srcTy.getNumElements();
+ int64_t resNumElems = resTy.getNumElements();
+ int64_t numTargetLanes = getLaneLayoutSize(targetLayout);
+ if (numTargetLanes < 1 || numLanes % numTargetLanes != 0)
+ return failure();
+ // gpu.shuffle is only defined for the integer widths a lane can move.
+ int64_t elemBitWidth = srcTy.getElementTypeBitWidth();
+ if (!llvm::isPowerOf2_64(elemBitWidth) || elemBitWidth < 8 ||
+ elemBitWidth > 64)
+ return failure();
+
+ auto srcCoords = computePerLaneElementCoords(inputLayout, shape, numLanes);
+ auto resCoords = computePerLaneElementCoords(targetLayout, shape, numLanes);
+ if (failed(srcCoords) || failed(resCoords))
+ return failure();
+ // Bail out if the layouts do not distribute the elements the distributed
+ // vector types account for, or if the target does not simply replicate the
+ // result over the lanes it leaves out.
+ for (int64_t lane = 0; lane < numLanes; lane++) {
+ if (static_cast<int64_t>((*srcCoords)[lane].size()) != srcNumElems ||
+ static_cast<int64_t>((*resCoords)[lane].size()) != resNumElems ||
+ (*resCoords)[lane] != (*resCoords)[lane % numTargetLanes])
+ return failure();
+ }
+
+ SmallVector<BroadcastRedistribution> redistributions;
+ for (int64_t pos = 0; pos < resNumElems; pos++) {
+ auto redistribution = deriveBroadcastRedistribution(
+ *srcCoords, *resCoords, pos, numLanes, numTargetLanes);
+ if (failed(redistribution))
+ return failure();
+ redistributions.push_back(*redistribution);
+ }
+
+ // Values are shuffled as same-width integers, which any lane data type can
+ // be bitcast to, and are bitcast back to the original element type at the
+ // end.
+ Type elemTy = srcTy.getElementType();
+ Type shuffleTy = rewriter.getIntegerType(elemBitWidth);
+ Value flatSrc = src;
+ auto flatSrcTy = VectorType::get({srcNumElems}, elemTy);
+ if (srcTy != flatSrcTy)
+ flatSrc = vector::ShapeCastOp::create(rewriter, loc, flatSrcTy, flatSrc);
+ if (elemTy != shuffleTy)
+ flatSrc = vector::BitCastOp::create(
+ rewriter, loc, VectorType::get({srcNumElems}, shuffleTy), flatSrc);
+
+ // Index of the lane within the target lane layout.
+ Value laneIdx = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+ /*upperBound=*/mlir::IntegerAttr());
+ if (numTargetLanes != numLanes)
+ laneIdx = arith::RemUIOp::create(
+ rewriter, loc, laneIdx,
+ arith::ConstantIndexOp::create(rewriter, loc, numTargetLanes));
+
+ Type i32Ty = rewriter.getI32Type();
+ auto flatResTy = VectorType::get({resNumElems}, shuffleTy);
+ Value res = arith::ConstantOp::create(rewriter, loc, flatResTy,
+ rewriter.getZeroAttr(flatResTy));
+ Value width;
+ Value laneIdxI32;
+ llvm::DenseMap<std::pair<int64_t, int64_t>, Value> extracted;
+ for (auto [pos, redistribution] : llvm::enumerate(redistributions)) {
+ // Extract element `stride * laneIdx + offset` of the local copy. Result
+ // elements coming from the same element of the copy share the extract.
+ Value &value = extracted[{redistribution.stride, redistribution.offset}];
+ if (!value) {
+ OpFoldResult element;
+ if (redistribution.stride == 0) {
+ element = rewriter.getIndexAttr(redistribution.offset);
+ } else {
+ Value index = laneIdx;
+ if (redistribution.stride != 1)
+ index =
+ arith::MulIOp::create(rewriter, loc, index,
+ arith::ConstantIndexOp::create(
+ rewriter, loc, redistribution.stride));
+ if (redistribution.offset != 0)
+ index =
+ arith::AddIOp::create(rewriter, loc, index,
+ arith::ConstantIndexOp::create(
+ rewriter, loc, redistribution.offset));
+ element = index;
+ }
+ value = vector::ExtractOp::create(rewriter, loc, flatSrc, element);
+ }
+
+ Value result = value;
+ if (redistribution.needsShuffle) {
+ if (!width) {
+ width = arith::ConstantIntOp::create(rewriter, loc, i32Ty, numLanes);
+ laneIdxI32 = arith::IndexCastOp::create(rewriter, loc, i32Ty, laneIdx);
+ }
+ Value srcLane = laneIdxI32;
+ if (redistribution.laneOffset != 0)
+ srcLane = arith::AddIOp::create(
+ rewriter, loc, srcLane,
+ arith::ConstantIntOp::create(rewriter, loc, i32Ty,
+ redistribution.laneOffset));
+ result = gpu::ShuffleOp::create(rewriter, loc, result, srcLane, width,
+ gpu::ShuffleMode::IDX)
+ .getResult(0);
+ }
+ res = vector::InsertOp::create(rewriter, loc, result, res, pos);
+ }
+ if (elemTy != shuffleTy)
+ res = vector::BitCastOp::create(
+ rewriter, loc, VectorType::get({resNumElems}, elemTy), res);
+ if (res.getType() != resTy)
+ res = vector::ShapeCastOp::create(rewriter, loc, resTy, res);
+ return res;
+}
+
/// Folds a subgroup-level ConvertLayout op with compatible lane layouts.
struct SgToLaneConvertLayout
: public OpConversionPattern<xegpu::ConvertLayoutOp> {
@@ -1862,6 +2097,28 @@ struct SgToLaneConvertLayout
}
}
+ // Handle the case where the input layout broadcasts the value over groups
+ // of lanes and the target layout distributes it, which requires moving
+ // data across lanes.
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+ FailureOr<VectorType> resDistTy = xegpu::getDistVecTypeBasedOnLaneLayout(
+ targetLayout, cast<VectorType>(valType));
+ if (uArch && succeeded(resDistTy)) {
+ int64_t numLanes = uArch->getSubgroupSize();
+ // The input has to be distributed over the whole subgroup, while the
+ // target may leave the remaining lanes with a replicated value.
+ if (getLaneLayoutSize(inputLayout) == numLanes) {
+ FailureOr<Value> res = redistributeBroadcastedValue(
+ rewriter, op.getLoc(), adaptor.getSource(), *resDistTy, inputLayout,
+ targetLayout, resShapeVec, numLanes);
+ 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 8607d45d5828c..2a24cdc805802 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1578,4 +1578,75 @@ gpu.func @convert_layout_partial_subgroup() {
} : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
gpu.return
}
+
+// A convert_layout from a layout that broadcasts the value over two groups of
+// eight lanes to a layout that gives each of the first eight lanes a row of
+// it. Lane `i` extracts row `i % 8` of the copy it holds, so lanes [0, 8) end
+// up with the first column and lanes [8, 16) with the second one, and the
+// shuffles hand lane `i` the two values of its row.
+// CHECK-LABEL: gpu.func @convert_layout_broadcast_to_lane_distributed
+// CHECK: %[[SRC:.*]] = arith.constant dense<1.000000e+00> : vector<8x1xf8E8M0FNU>
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<8x1xf8E8M0FNU> to vector<8xf8E8M0FNU>
+// CHECK: %[[BITS:.*]] = vector.bitcast %[[FLAT]] : vector<8xf8E8M0FNU> to vector<8xi8>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[ROW:.*]] = arith.remui %[[LANE]], %[[C8]] : index
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<2xi8>
+// CHECK: %[[ELEM:.*]] = vector.extract %[[BITS]][%[[ROW]]] : i8 from vector<8xi8>
+// CHECK: %[[WIDTH:.*]] = arith.constant 16 : i32
+// CHECK: %[[ROW_I32:.*]] = arith.index_cast %[[ROW]] : index to i32
+// CHECK: %[[SHUF0:.*]], %{{.*}} = gpu.shuffle idx %[[ELEM]], %[[ROW_I32]], %[[WIDTH]] : i8
+// CHECK: %[[INS0:.*]] = vector.insert %[[SHUF0]], %[[ZERO]] [0] : i8 into vector<2xi8>
+// CHECK: %[[C8_I32:.*]] = arith.constant 8 : i32
+// CHECK: %[[LANE1:.*]] = arith.addi %[[ROW_I32]], %[[C8_I32]] : i32
+// CHECK: %[[SHUF1:.*]], %{{.*}} = gpu.shuffle idx %[[ELEM]], %[[LANE1]], %[[WIDTH]] : i8
+// CHECK: %[[INS1:.*]] = vector.insert %[[SHUF1]], %[[INS0]] [1] : i8 into vector<2xi8>
+// CHECK: %[[BACK:.*]] = vector.bitcast %[[INS1]] : vector<2xi8> to vector<2xf8E8M0FNU>
+// CHECK: vector.shape_cast %[[BACK]] : vector<2xf8E8M0FNU> to vector<1x2xf8E8M0FNU>
+gpu.func @convert_layout_broadcast_to_lane_distributed() {
+ %scale_a_src = arith.constant dense<1.0> : vector<8x2xf8E8M0FNU>
+ %cvt = xegpu.convert_layout %scale_a_src
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2], lane_data = [4, 1, 1], order = [0, 2, 1]>, dims = [0]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>
+ }> : vector<8x2xf8E8M0FNU>
+ "some_use"(%cvt) : (vector<8x2xf8E8M0FNU>) -> ()
+ gpu.return
+}
+
+// Same redistribution, but from a layout that broadcasts the value to all
+// lanes. Every lane already holds the row it owns, so the values are only
+// extracted using the lane id and no data has to be moved across lanes.
+// CHECK-LABEL: gpu.func @convert_layout_broadcast_all_lanes
+// CHECK: %[[SRC:.*]] = arith.constant dense<1.000000e+00> : vector<8x2xf8E8M0FNU>
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<8x2xf8E8M0FNU> to vector<16xf8E8M0FNU>
+// CHECK: %[[BITS:.*]] = vector.bitcast %[[FLAT]] : vector<16xf8E8M0FNU> to vector<16xi8>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[ROW:.*]] = arith.remui %[[LANE]], %[[C8]] : index
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<2xi8>
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[IDX0:.*]] = arith.muli %[[ROW]], %[[C2]] : index
+// CHECK: %[[ELEM0:.*]] = vector.extract %[[BITS]][%[[IDX0]]] : i8 from vector<16xi8>
+// CHECK: %[[INS0:.*]] = vector.insert %[[ELEM0]], %[[ZERO]] [0] : i8 into vector<2xi8>
+// CHECK: %[[C2_1:.*]] = arith.constant 2 : index
+// CHECK: %[[MUL1:.*]] = arith.muli %[[ROW]], %[[C2_1]] : index
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[IDX1:.*]] = arith.addi %[[MUL1]], %[[C1]] : index
+// CHECK: %[[ELEM1:.*]] = vector.extract %[[BITS]][%[[IDX1]]] : i8 from vector<16xi8>
+// CHECK: %[[INS1:.*]] = vector.insert %[[ELEM1]], %[[INS0]] [1] : i8 into vector<2xi8>
+// CHECK: %[[BACK:.*]] = vector.bitcast %[[INS1]] : vector<2xi8> to vector<2xf8E8M0FNU>
+// CHECK: vector.shape_cast %[[BACK]] : vector<2xf8E8M0FNU> to vector<1x2xf8E8M0FNU>
+// CHECK: gpu.return
+// CHECK-NOT: gpu.shuffle
+gpu.func @convert_layout_broadcast_all_lanes() {
+ %scale_a_src = arith.constant dense<1.0> : vector<8x2xf8E8M0FNU>
+ %cvt = xegpu.convert_layout %scale_a_src
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [2]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>
+ }> : vector<8x2xf8E8M0FNU>
+ "some_use"(%cvt) : (vector<8x2xf8E8M0FNU>) -> ()
+ gpu.return
+}
}
More information about the Mlir-commits
mailing list