[Mlir-commits] [mlir] [MLIR][XeGPU] Redistribute broadcasted data into non-fastest lane dims (PR #217104)
Sang Ik Lee
llvmlistbot at llvm.org
Tue Aug 18 14:48:11 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/217104
>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 1/2] [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
+}
}
>From de050c8a9253cdf4de298f810b0336f89c7be0e1 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 12 Aug 2026 20:22:15 +0000
Subject: [PATCH 2/2] [MLIR][XeGPU] Redistribute broadcasted data into
non-fastest lane dims
The A-scale operand of a quantized-A mxfp GEMM is produced broadcast over the
subgroup and consumed distributed. sg-to-lane distribution rejected that
`convert_layout` with "lowering incompatible convert_layout not yet supported":
%cvt = xegpu.convert_layout %scale <{
input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16],
lane_data = [1, 1, 1]>, dims = [2]>,
target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2],
lane_data = [4, 1, 1], order = [0, 2, 1]>, dims = [0]>
}> : vector<8x2xbf16>
Each lane extracts the element it owns out of the copy it holds, at an index
computed from its lane id. The lowering only supported indices that advance by a
constant per lane, so it could not handle a target whose broadcast dimension
varies fastest: above, lane `i` owns column `i / 8`, so groups of eight lanes
share an index.
Generalize the index so that a group of consecutive lanes can share one, and can
wrap. Layouts supported before generate identical code, and layouts fitting
neither form are still rejected rather than lowered incorrectly.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 95 ++++++++++++++++---
.../XeGPU/sg-to-lane-distribute-unit.mlir | 39 ++++++++
2 files changed, 120 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 68931ea16b1cb..27ae3a38e6ac5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1790,12 +1790,22 @@ computePerLaneElementCoords(xegpu::DistributeLayoutAttr layout,
/// 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`.
+/// target lane layout, the lane extracts element
+/// `stride * ((t / divisor) % modulus) + 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`.
+///
+/// `divisor` and `modulus` single out the one coordinate of the target lane
+/// layout that the element index varies with: `divisor` is the number of
+/// consecutive lanes sharing that coordinate and `modulus` the number of
+/// distinct values it takes. A `divisor` of 1 paired with a `modulus`
+/// spanning the whole target lane layout recovers the plain `stride * t +
+/// offset` of a layout whose distributed dimension varies fastest.
struct BroadcastRedistribution {
int64_t stride;
int64_t offset;
+ int64_t divisor;
+ int64_t modulus;
int64_t laneOffset;
bool needsShuffle;
};
@@ -1834,13 +1844,42 @@ static FailureOr<BroadcastRedistribution> deriveBroadcastRedistribution(
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.
+ // to be a strided function of a single coordinate of the lane's position
+ // in the target lane layout. Lanes sharing that coordinate are adjacent,
+ // so the index is constant over blocks of `divisor` consecutive lanes.
int64_t offset = elements[0];
- int64_t stride = numTargetLanes > 1 ? elements[1] - offset : 0;
- if (stride < 0)
+ int64_t divisor = 1;
+ while (divisor < numTargetLanes && elements[divisor] == offset)
+ divisor++;
+ int64_t stride = divisor < numTargetLanes ? elements[divisor] - offset : 0;
+ if (stride < 0 || numTargetLanes % divisor != 0)
continue;
+
+ // How many distinct values the coordinate takes before repeating. Derived
+ // from the largest block index actually observed, so that a coordinate
+ // that wraps around within the target lane layout is picked up as well.
+ int64_t modulus = 1;
+ if (stride == 0) {
+ // A lane-invariant index: the divisor is immaterial, normalize it away
+ // so equal indices share a single extract.
+ divisor = 1;
+ } else {
+ bool strided = true;
+ for (int64_t t = 0; t < numTargetLanes && strided; t++) {
+ int64_t diff = elements[t] - offset;
+ strided = diff >= 0 && diff % stride == 0;
+ if (strided)
+ modulus = std::max(modulus, diff / stride + 1);
+ }
+ if (!strided)
+ continue;
+ }
+
+ auto elementIndex = [=](int64_t t) {
+ return stride * ((t / divisor) % modulus) + offset;
+ };
if (!llvm::all_of(llvm::seq<int64_t>(0, numTargetLanes), [&](int64_t t) {
- return elements[t] == stride * t + offset;
+ return elements[t] == elementIndex(t);
}))
continue;
@@ -1850,10 +1889,11 @@ static FailureOr<BroadcastRedistribution> deriveBroadcastRedistribution(
bool needsShuffle =
laneOffset != 0 ||
!llvm::all_of(llvm::seq<int64_t>(0, numLanes), [&](int64_t lane) {
- return srcCoords[lane][stride * (lane % numTargetLanes) + offset] ==
+ return srcCoords[lane][elementIndex(lane % numTargetLanes)] ==
resCoords[lane][pos];
});
- return BroadcastRedistribution{stride, offset, laneOffset, needsShuffle};
+ return BroadcastRedistribution{stride, offset, divisor,
+ modulus, laneOffset, needsShuffle};
}
return failure();
}
@@ -1944,17 +1984,44 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
rewriter.getZeroAttr(flatResTy));
Value width;
Value laneIdxI32;
- llvm::DenseMap<std::pair<int64_t, int64_t>, Value> extracted;
+
+ // The lane coordinate the extract index is derived from,
+ // `(laneIdx / divisor) % modulus`, shared by every element that reads the
+ // same coordinate. Both steps are skipped when they are no-ops over the
+ // range of `laneIdx`, so a layout distributed along its fastest varying
+ // dimension keeps indexing by the lane index directly.
+ llvm::DenseMap<std::pair<int64_t, int64_t>, Value> laneTerms;
+ auto getLaneTerm = [&](int64_t divisor, int64_t modulus) -> Value {
+ Value &term = laneTerms[{divisor, modulus}];
+ if (term)
+ return term;
+ term = laneIdx;
+ if (divisor != 1)
+ term = arith::DivUIOp::create(
+ rewriter, loc, term,
+ arith::ConstantIndexOp::create(rewriter, loc, divisor));
+ if (modulus < numTargetLanes / divisor)
+ term = arith::RemUIOp::create(
+ rewriter, loc, term,
+ arith::ConstantIndexOp::create(rewriter, loc, modulus));
+ return term;
+ };
+
+ llvm::DenseMap<std::tuple<int64_t, int64_t, 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}];
+ // Extract element `stride * ((laneIdx / divisor) % modulus) + 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,
+ redistribution.divisor, redistribution.modulus}];
if (!value) {
OpFoldResult element;
if (redistribution.stride == 0) {
element = rewriter.getIndexAttr(redistribution.offset);
} else {
- Value index = laneIdx;
+ Value index =
+ getLaneTerm(redistribution.divisor, redistribution.modulus);
if (redistribution.stride != 1)
index =
arith::MulIOp::create(rewriter, loc, index,
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 2a24cdc805802..d13020ca48598 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1649,4 +1649,43 @@ gpu.func @convert_layout_broadcast_all_lanes() {
"some_use"(%cvt) : (vector<8x2xf8E8M0FNU>) -> ()
gpu.return
}
+
+// Redistribution into a target layout whose distributed dimension is not the
+// fastest varying one: the parent `order` puts the sliced (broadcast)
+// dimension first, so lane `i` owns column `i / 8` rather than column `i % 2`.
+// The extract index therefore divides the lane id instead of taking a
+// remainder of it. The source is broadcast to the whole subgroup, so every
+// lane already holds its column and nothing is moved across lanes.
+// CHECK-LABEL: gpu.func @convert_layout_broadcast_to_sliced_target
+// CHECK: %[[SRC:.*]] = arith.constant dense<1.000000e+00> : vector<8x2xbf16>
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<8x2xbf16> to vector<16xbf16>
+// CHECK: %[[BITS:.*]] = vector.bitcast %[[FLAT]] : vector<16xbf16> to vector<16xi16>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<8xi16>
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[COL:.*]] = arith.divui %[[LANE]], %[[C8]] : index
+// CHECK: %[[ELEM0:.*]] = vector.extract %[[BITS]][%[[COL]]] : i16 from vector<16xi16>
+// CHECK: %[[INS0:.*]] = vector.insert %[[ELEM0]], %[[ZERO]] [0] : i16 into vector<8xi16>
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[IDX1:.*]] = arith.addi %[[COL]], %[[C2]] : index
+// CHECK: %[[ELEM1:.*]] = vector.extract %[[BITS]][%[[IDX1]]] : i16 from vector<16xi16>
+// CHECK: %[[INS1:.*]] = vector.insert %[[ELEM1]], %[[INS0]] [1] : i16 into vector<8xi16>
+// CHECK: %[[C14:.*]] = arith.constant 14 : index
+// CHECK: %[[IDX7:.*]] = arith.addi %[[COL]], %[[C14]] : index
+// CHECK: %[[ELEM7:.*]] = vector.extract %[[BITS]][%[[IDX7]]] : i16 from vector<16xi16>
+// CHECK: %[[INS7:.*]] = vector.insert %[[ELEM7]], %{{.*}} [7] : i16 into vector<8xi16>
+// CHECK: %[[BACK:.*]] = vector.bitcast %[[INS7]] : vector<8xi16> to vector<8xbf16>
+// CHECK: vector.shape_cast %[[BACK]] : vector<8xbf16> to vector<8x1xbf16>
+// CHECK: gpu.return
+// CHECK-NOT: gpu.shuffle
+gpu.func @convert_layout_broadcast_to_sliced_target() {
+ %a_scale_src = arith.constant dense<1.0> : vector<8x2xbf16>
+ %cvt = xegpu.convert_layout %a_scale_src
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [2]>,
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2], lane_data = [4, 1, 1], order = [0, 2, 1]>, dims = [0]>
+ }> : vector<8x2xbf16>
+ "some_use"(%cvt) : (vector<8x2xbf16>) -> ()
+ gpu.return
+}
}
More information about the Mlir-commits
mailing list