[Mlir-commits] [mlir] [MLIR][XeGPU] Sg to Lane distribute convert_layout redistributing broadcasted data (PR #215645)
Sang Ik Lee
llvmlistbot at llvm.org
Wed Sep 2 14:53:12 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/215645
>From 2f585e200576d3d57545a0d9fddbad9848153796 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Fri, 21 Aug 2026 22:48:27 +0000
Subject: [PATCH 01/46] [MLIR][XeGPU] Sg to Lane distribute convert_layout
redistributing broadcasted data
A convert_layout whose input layout replicates a value over lanes, while the
target hands each lane a different part of it, failed to legalize. Both the
group of lanes sharing a copy and the target dimension a lane's share varies
with can differ:
#xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2], order = [0, 2, 1]>, dims = [0]>
-> #xegpu.layout<lane_layout = [8, 1]>
#xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16]>, dims = [2]>
-> #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2], order = [0, 2, 1]>, dims = [0]>
Scale operands of scaled matrix multiplication need this, as they are produced
broadcasted and consumed distributed.
Distribute it by having each lane extract the element it needs from the copy it
already holds, using gpu.shuffle only for the elements that are not in it.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 288 ++++++++++++++++++
.../XeGPU/sg-to-lane-distribute-unit.mlir | 104 +++++++
2 files changed, 392 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index bc85140f9f121..56e2e4017a3b8 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1761,6 +1761,274 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
return result;
}
+/// Number of lanes the subgroup is split into by `layout`. Lanes covered by
+/// sliced dimensions count too: they hold a broadcasted copy, not nothing.
+static int64_t getLaneLayoutSize(xegpu::DistributeLayoutAttr layout) {
+ if (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
+ layout = sliceAttr.flatten().getParent();
+ return computeProduct(layout.getEffectiveLaneLayoutAsInt());
+}
+
+/// For every lane, the coordinates of the elements of `shape` it owns under
+/// `layout`. Fails 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;
+}
+
+/// How one element of the distributed result is obtained. With `t` the lane's
+/// index within the target lane layout, the lane extracts element
+/// `stride * ((t / divisor) % modulus) + offset` of its copy, and `gpu.shuffle
+/// idx` pulls it from lane `t + laneOffset` unless it already owns it.
+///
+/// `divisor` and `modulus` single out the one target lane layout coordinate the
+/// element index varies with: `divisor` lanes share a coordinate value, which
+/// takes `modulus` distinct values.
+struct BroadcastRedistribution {
+ int64_t stride;
+ int64_t offset;
+ int64_t divisor;
+ int64_t modulus;
+ int64_t laneOffset;
+ bool needsShuffle;
+};
+
+/// Derives the redistribution of element `pos` from the per-lane element
+/// coordinates of both layouts. Fails when the movement is not 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) {
+ 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 serves the target lane sharing its position in the target lane
+ // layout, so the shuffled-from lane is a whole number of layouts away.
+ for (int64_t laneOffset = 0; laneOffset < numLanes;
+ laneOffset += numTargetLanes) {
+ 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 index is computed from the lane id at runtime, so it has to be
+ // strided in one coordinate. Lanes sharing that coordinate are adjacent.
+ int64_t offset = elements[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;
+
+ int64_t modulus = 1;
+ if (stride == 0) {
+ 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] == elementIndex(t);
+ }))
+ continue;
+
+ // Lanes outside the target lane layout hold a replicated copy nothing
+ // reads, so only the lanes carrying a result need the value they own. That
+ // lets an already-local element skip the shuffle.
+ bool needsShuffle =
+ laneOffset != 0 ||
+ !llvm::all_of(llvm::seq<int64_t>(0, numTargetLanes), [&](int64_t lane) {
+ return srcCoords[lane][elementIndex(lane % numTargetLanes)] ==
+ resCoords[lane][pos];
+ });
+ return BroadcastRedistribution{stride, offset, divisor,
+ modulus, laneOffset, needsShuffle};
+ }
+ return failure();
+}
+
+/// Redistributes `src` for a `convert_layout` whose input layout replicates the
+/// value over groups of lanes while the target hands each lane a different part
+/// of it. Scale operands of scaled matrix multiplication need this: they are
+/// produced broadcasted and consumed distributed.
+///
+/// Any lane can read any element of its own copy with a dynamic extract, but a
+/// lane can only contribute one value to a `gpu.shuffle`, so both sides have to
+/// agree on what is exchanged: each lane extracts the element its counterpart
+/// in the target lane layout needs, and one shuffle per result element moves it
+/// to the owner.
+///
+/// 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 only moves the integer widths a lane can hold.
+ 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();
+ // Require the layouts to distribute what the vector types account for, and
+ // the target to 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);
+ }
+
+ // Shuffles move same-width integers, which any lane data type bitcasts to.
+ 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);
+
+ 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;
+
+ // `(laneIdx / divisor) % modulus`, shared by elements reading the same
+ // coordinate. Either step is skipped when it is a no-op over `laneIdx`.
+ 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)) {
+ // Result elements reading 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 =
+ getLaneTerm(redistribution.divisor, redistribution.modulus);
+ 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 +2130,26 @@ struct SgToLaneConvertLayout
}
}
+ // Input layout broadcasts over groups of lanes while the target
+ // distributes, so data has to move 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 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..c36a70635d400 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,108 @@ gpu.func @convert_layout_partial_subgroup() {
} : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
gpu.return
}
+
+// Input layout broadcasts over two groups of eight lanes, each holding one
+// column; the target gives each of the first eight lanes a row. Element 0 is
+// already local, element 1 has to cross lanes.
+// 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: %[[INS0:.*]] = vector.insert %[[ELEM]], %[[ZERO]] [0] : i8 into vector<2xi8>
+
+// CHECK: %[[WIDTH:.*]] = arith.constant 16 : i32
+// CHECK: %[[ROW_I32:.*]] = arith.index_cast %[[ROW]] : index to i32
+// 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
+}
+
+// Broadcast to all lanes: every lane already holds its row, so no shuffles.
+// 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
+}
+
+// Target whose distributed dimension is not the fastest varying one: the parent
+// `order` puts the broadcast dimension first, so lane `i` owns column `i / 8`
+// and the extract divides the lane id rather than taking a remainder.
+// 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
+}
}
>From b7736734422abab01725c44076b75acf829a6d18 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 31 Aug 2026 18:29:07 +0000
Subject: [PATCH 02/46] Describe the broadcast redistribution algorithm and
name the code after it
Review feedback: the derivation and the emission of the redistribution were
hard to follow. Neither said what model it was implementing, and the index
formula they agree on appeared twice, 150 lines apart, spelled differently.
Add a section header that introduces the concepts the lowering rests on --
fragment, broadcast group, slot, donor and element source -- then states the
algorithm as pseudo code in those terms, with a worked example. Two things that
were left implicit are now stated: replication makes the donor ambiguous, so the
smallest delta is taken because it is the only one whose shuffle can be dropped;
and only the lanes of a slot are required to hold what the target layout assigns
them, which is what allows dropping it.
Keep fragment and broadcast group apart, since a lane holding a fragment says
nothing about sharing on its own. A broadcast group is formed by a sliced
lane_layout dimension: nothing is distributed over its lanes, so they are all
given the same data. A donor group is a different thing again -- a partition by
the target layout -- so it is named separately.
Rename the code to match, and to the words the dialect already uses:
getLaneLayoutSize -> getNumParticipatingLanes, computePerLaneElementCoords ->
computeOwnedCoords returning a named OwnedCoords table, srcCoords/resCoords ->
inputOwned/targetOwned after the input_layout and target_layout they come from,
numLanes -> subgroupSize as in LayoutAttr::isCompatibleWith, numTargetLanes ->
numSlots, laneOffset -> donorDelta, laneIdx -> slot, flatSrc -> fragment. Keep
coordinate for the element coordinates of the ownership tables, as the rest of
the file does, and name the lane_layout side after what it is: the opaque
divisor and modulus become slotStride and extent, the index along the one
dimension the fragment index varies with.
Split the six flat fields of BroadcastRedistribution into a FragmentIndex, the
index a slot reads, and an ElementSource that pairs it with the donor to read
from -- which is all the pass needs to know per element of the target fragment.
FragmentIndex::at() is now the only compile-time statement of the index formula,
used by both the fit and the needsShuffle decision, and getDimIndex/
getFragmentIndex the only one in the emitter. Split the derivation accordingly
into fitFragmentIndex, the integer fit, and deriveElementSource, the donor
search.
Also say donor lane in shuffleDataAsLaneLayoutChange, which called the same
thing a partner lane, so the file has one vocabulary.
No functional change: the IR emitted for both tests is identical, op for op.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 509 +++++++++++-------
1 file changed, 322 insertions(+), 187 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 56e2e4017a3b8..4fb35ac292d01 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1630,7 +1630,7 @@ shuffleDataAsLaneLayoutChange(ConversionPatternRewriter &rewriter, Location loc,
// -- create a temp 1D vector of i32 initialized to zero
// -- for each i32 element:
// ---- extract it from the source bundle
- // ---- gpu.shuffle to gather the value from the partner lane
+ // ---- gpu.shuffle to gather the value from the donor lane
// ---- insert it into the temp bundle
// -- cast the temp back to the source vector type
// -- vector.shuffle the source and temp to concatenate along the outer dim
@@ -1761,147 +1761,274 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
return result;
}
-/// Number of lanes the subgroup is split into by `layout`. Lanes covered by
-/// sliced dimensions count too: they hold a broadcasted copy, not nothing.
-static int64_t getLaneLayoutSize(xegpu::DistributeLayoutAttr layout) {
+//===----------------------------------------------------------------------===//
+// Broadcast redistribution
+//===----------------------------------------------------------------------===//
+//
+// Lowers a `convert_layout` that only relocates data between lanes: the input
+// layout replicates the data over lanes, while the target layout hands each
+// lane a different part of it. Every element of the result already exists in
+// some lane, so nothing is recomputed. The scale operands of scaled matrix
+// multiplication need this, being produced broadcasted and consumed
+// distributed.
+//
+// Concepts
+// --------
+//
+// fragment The elements a lane holds under a layout, as a flat list --
+// the lane's distributed vector. Both layouts give a lane a
+// fragment, and the conversion turns the input one into the
+// target one.
+//
+// broadcast group Lanes that hold the same fragment. A sliced lane_layout
+// dimension forms one: nothing is distributed over its lanes,
+// so they are all given the same data. The group spans that
+// dimension -- all 16 lanes if it covers the subgroup, 8 if it
+// covers half of it.
+//
+// slot A lane's linear id within the target lane layout,
+// `lane % numSlots`. The target layout gives the same fragment
+// to every lane of a slot, so there are `numSlots` distinct
+// target fragments and the lanes past the first `numSlots` are
+// replicas.
+//
+// donor The lane a value is read from. Only `gpu.shuffle idx` moves
+// a value between lanes, and it exchanges exactly one value
+// per lane: each lane offers one value and receives one lane's
+// offering. So all lanes have to run the same extract, and a
+// lane can only be served by a lane that arrives at the same
+// extract index -- one of its own slot, i.e. `slot + delta`
+// for a `delta` that is a multiple of `numSlots`. Lanes
+// `delta` apart form a donor group, which partitions the
+// subgroup by the target layout and is not a broadcast group.
+//
+// element source Where one element of the target fragment comes from: the
+// donor to take it from, and the index to extract from that
+// donor's fragment. A shuffle is needed exactly when that
+// donor cannot be the lane itself. The index is a function of
+// the slot,
+//
+// stride * ((slot / slotStride) % extent) + offset
+//
+// which says "one dimension of the target lane layout selects
+// one element of the fragment": `(slot / slotStride) % extent`
+// is the lane's index along that dimension, `slotStride` being
+// how many slots apart two lanes have to be for it to differ
+// by one and `extent` how many values it takes. It costs a
+// couple of `arith` ops on the lane id, where an arbitrary
+// per-slot index would need a materialized table and a gather.
+//
+// Algorithm
+// ---------
+//
+// tabulate inputOwned[lane] and targetOwned[slot] // element coordinates
+// for each element `pos` of the target fragment:
+// for delta in 0, numSlots, 2 * numSlots, ...: // smallest first
+// needed[slot] = index of targetOwned[slot][pos]
+// in inputOwned[slot + delta]
+// if a slot's donor does not hold its element: next delta
+// index = fit `needed` to stride/offset/slotStride/extent,
+// else next delta
+// accept (delta, index) and stop
+// if no delta fits: fail, the layout change is not of this form
+// for each element source:
+// %idx = index(slot) // arith on gpu.lane_id
+// %val = vector.extract %fragment[%idx]
+// %val = gpu.shuffle idx %val, slot + delta // unless donor == lane
+// %res = vector.insert %val, %res[pos]
+//
+// Replication leaves the donor ambiguous: several donor groups may hold the
+// element and they all yield the same value, so the smallest `delta` is taken.
+// Only `delta = 0` can make the donor the lane itself, so it is the only one
+// that may need no shuffle, and taking it is what leaves a fully broadcast
+// input shuffle-free.
+//
+// For example, redistributing a `vector<8x2>` whose input fragments hold column
+// 0 in lanes [0, 8) and column 1 in lanes [8, 16) to a target that gives lane
+// `i` row `i` (`numSlots` = 8, so lanes [8, 16) are replicas) yields
+//
+// pos 0 -> delta = 0, idx = slot // column 0 is in the lane's own group
+// pos 1 -> delta = 8, idx = slot // column 1 is one donor group away
+//
+// An index is derived from the tables and then verified against them, so a
+// layout change that does not have this form is reported as a match failure
+// rather than lowered incorrectly.
+//
+// Assumption: only the first `numSlots` lanes have to end up with the elements
+// the target layout assigns them, the rest being replicas. Whether the donor
+// can be the lane itself is therefore decided over the slots alone, and for all
+// of them at once since the index is uniform over the subgroup: one slot whose
+// local extract misses the element it owns forces the shuffle on every lane.
+//
+//===----------------------------------------------------------------------===//
+
+/// Number of lanes of the subgroup `layout` distributes over. Lanes covered by
+/// sliced dimensions count too: they hold a replicated fragment, not nothing.
+static int64_t getNumParticipatingLanes(xegpu::DistributeLayoutAttr layout) {
if (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
layout = sliceAttr.flatten().getParent();
return computeProduct(layout.getEffectiveLaneLayoutAsInt());
}
-/// For every lane, the coordinates of the elements of `shape` it owns under
+/// Ownership table: `owned[lane][i]` is the coordinate, in the undistributed
+/// shape, of element `i` of the fragment `lane` holds.
+using OwnedCoords = SmallVector<SmallVector<SmallVector<int64_t>>>;
+
+/// Tabulates the fragment every lane of the subgroup owns of `shape` under
/// `layout`. Fails 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) {
+static FailureOr<OwnedCoords>
+computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
+ int64_t subgroupSize) {
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;
+ OwnedCoords owned;
+ for (int64_t lane = 0; lane < subgroupSize; lane++)
+ owned.push_back(layout.computeStaticDistributedCoords(lane, shape));
+ return owned;
}
-/// How one element of the distributed result is obtained. With `t` the lane's
-/// index within the target lane layout, the lane extracts element
-/// `stride * ((t / divisor) % modulus) + offset` of its copy, and `gpu.shuffle
-/// idx` pulls it from lane `t + laneOffset` unless it already owns it.
-///
-/// `divisor` and `modulus` single out the one target lane layout coordinate the
-/// element index varies with: `divisor` lanes share a coordinate value, which
-/// takes `modulus` distinct values.
-struct BroadcastRedistribution {
+/// Index into the lane's fragment as a function of its slot:
+/// `stride * ((slot / slotStride) % extent) + offset`.
+struct FragmentIndex {
int64_t stride;
int64_t offset;
- int64_t divisor;
- int64_t modulus;
- int64_t laneOffset;
+ /// Slots apart two lanes have to be for their index along the selected
+ /// lane_layout dimension to differ by one.
+ int64_t slotStride;
+ /// Number of values that index takes.
+ int64_t extent;
+
+ /// The index `slot` reads. Compile-time twin of the emitted arithmetic.
+ int64_t at(int64_t slot) const {
+ return stride * ((slot / slotStride) % extent) + offset;
+ }
+
+ /// Key identifying the value this extracts, for reuse across elements.
+ std::tuple<int64_t, int64_t, int64_t, int64_t> asTuple() const {
+ return {stride, offset, slotStride, extent};
+ }
+};
+
+/// Where one element of the target fragment comes from: `index` picks it out of
+/// the fragment of the donor `donorDelta` lanes away, and `gpu.shuffle idx`
+/// moves it from there to the lane that owns it, unless the donor is that lane
+/// itself.
+struct ElementSource {
+ FragmentIndex index;
+ int64_t donorDelta;
bool needsShuffle;
};
-/// Derives the redistribution of element `pos` from the per-lane element
-/// coordinates of both layouts. Fails when the movement is not 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) {
- auto findElement = [&](int64_t lane, ArrayRef<int64_t> coord) -> int64_t {
- for (auto [idx, candidate] : llvm::enumerate(srcCoords[lane]))
+/// Fits `needed`, the fragment index each of the `numSlots` slots has to read,
+/// to a `FragmentIndex`. Fails when no index reproduces `needed`.
+static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
+ int64_t numSlots) {
+ // Slots sharing an index along the selected lane_layout dimension are
+ // adjacent, so the first change in `needed` tells how many of them there are
+ // (`slotStride`) and how far apart the elements of consecutive indices are
+ // (`stride`).
+ int64_t offset = needed[0];
+ int64_t slotStride = 1;
+ while (slotStride < numSlots && needed[slotStride] == offset)
+ slotStride++;
+ int64_t stride = slotStride < numSlots ? needed[slotStride] - offset : 0;
+ if (stride < 0 || numSlots % slotStride != 0)
+ return failure();
+
+ // The index wraps around once it has taken all its values, which is how far
+ // `needed` gets from `offset`.
+ int64_t extent = 1;
+ if (stride == 0) {
+ // Every slot reads `offset`; no dimension is involved.
+ slotStride = 1;
+ } else {
+ for (int64_t index : needed) {
+ int64_t diff = index - offset;
+ if (diff < 0 || diff % stride != 0)
+ return failure();
+ extent = std::max(extent, diff / stride + 1);
+ }
+ }
+
+ // Derived from the table, so verify against the table before relying on it.
+ FragmentIndex index{stride, offset, slotStride, extent};
+ if (!llvm::all_of(llvm::seq<int64_t>(0, numSlots), [&](int64_t slot) {
+ return needed[slot] == index.at(slot);
+ }))
+ return failure();
+ return index;
+}
+
+/// Works out where element `pos` of the target fragment comes from, using the
+/// two ownership tables. Fails when no donor group provides that element
+/// through a single extract, i.e. when the data movement is not of the form
+/// described above.
+static FailureOr<ElementSource>
+deriveElementSource(const OwnedCoords &inputOwned,
+ const OwnedCoords &targetOwned, int64_t pos,
+ int64_t subgroupSize, int64_t numSlots) {
+ // Index of the element with coordinate `coord` in the fragment `lane` holds,
+ // or -1 if it holds none.
+ auto findInFragment = [&](int64_t lane, ArrayRef<int64_t> coord) -> int64_t {
+ for (auto [idx, candidate] : llvm::enumerate(inputOwned[lane]))
if (ArrayRef<int64_t>(candidate) == coord)
return idx;
return -1;
};
- // A lane only serves the target lane sharing its position in the target lane
- // layout, so the shuffled-from lane is a whole number of layouts away.
- for (int64_t laneOffset = 0; laneOffset < numLanes;
- laneOffset += numTargetLanes) {
- SmallVector<int64_t> elements;
- for (int64_t t = 0; t < numTargetLanes; t++) {
- int64_t element = findElement(t + laneOffset, resCoords[t][pos]);
- if (element < 0)
+ // A donor shares the slot of the lane it serves, so it is a whole number of
+ // target lane layouts away. Replication can put the element in more than one
+ // donor group; the smallest `delta` wins, as it is the only one that can read
+ // the lane's own fragment and so skip the shuffle.
+ for (int64_t donorDelta = 0; donorDelta < subgroupSize;
+ donorDelta += numSlots) {
+ // Where in its fragment each slot's donor keeps the element it needs.
+ SmallVector<int64_t> needed;
+ for (int64_t slot = 0; slot < numSlots; slot++) {
+ int64_t index = findInFragment(slot + donorDelta, targetOwned[slot][pos]);
+ if (index < 0)
break;
- elements.push_back(element);
+ needed.push_back(index);
}
- if (static_cast<int64_t>(elements.size()) != numTargetLanes)
+ if (static_cast<int64_t>(needed.size()) != numSlots)
continue;
- // The index is computed from the lane id at runtime, so it has to be
- // strided in one coordinate. Lanes sharing that coordinate are adjacent.
- int64_t offset = elements[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)
+ FailureOr<FragmentIndex> index = fitFragmentIndex(needed, numSlots);
+ if (failed(index))
continue;
- int64_t modulus = 1;
- if (stride == 0) {
- 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] == elementIndex(t);
- }))
- continue;
-
- // Lanes outside the target lane layout hold a replicated copy nothing
- // reads, so only the lanes carrying a result need the value they own. That
- // lets an already-local element skip the shuffle.
+ // A shuffle is needed unless the donor is the lane itself, that is unless
+ // the local extract already lands on the element the lane owns. Only the
+ // lanes of a slot decide that, the rest being replicas.
bool needsShuffle =
- laneOffset != 0 ||
- !llvm::all_of(llvm::seq<int64_t>(0, numTargetLanes), [&](int64_t lane) {
- return srcCoords[lane][elementIndex(lane % numTargetLanes)] ==
- resCoords[lane][pos];
+ donorDelta != 0 ||
+ !llvm::all_of(llvm::seq<int64_t>(0, numSlots), [&](int64_t slot) {
+ return inputOwned[slot][index->at(slot)] == targetOwned[slot][pos];
});
- return BroadcastRedistribution{stride, offset, divisor,
- modulus, laneOffset, needsShuffle};
+ return ElementSource{*index, donorDelta, needsShuffle};
}
return failure();
}
-/// Redistributes `src` for a `convert_layout` whose input layout replicates the
-/// value over groups of lanes while the target hands each lane a different part
-/// of it. Scale operands of scaled matrix multiplication need this: they are
-/// produced broadcasted and consumed distributed.
-///
-/// Any lane can read any element of its own copy with a dynamic extract, but a
-/// lane can only contribute one value to a `gpu.shuffle`, so both sides have to
-/// agree on what is exchanged: each lane extracts the element its counterpart
-/// in the target lane layout needs, and one shuffle per result element moves it
-/// to the owner.
-///
-/// Returns failure if the redistribution is not expressible in that form.
+/// Redistributes `src`, the fragment the lane holds under `inputLayout`, into
+/// the `resTy` fragment of the `shape`-shaped value it owns under
+/// `targetLayout`, over a subgroup of `subgroupSize` lanes. See the algorithm
+/// above; returns failure if the layout change is not of 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) {
+ ArrayRef<int64_t> shape, int64_t subgroupSize) {
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)
+ int64_t numSlots = getNumParticipatingLanes(targetLayout);
+ if (numSlots < 1 || subgroupSize % numSlots != 0)
return failure();
// gpu.shuffle only moves the integer widths a lane can hold.
int64_t elemBitWidth = srcTy.getElementTypeBitWidth();
@@ -1909,118 +2036,126 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
elemBitWidth > 64)
return failure();
- auto srcCoords = computePerLaneElementCoords(inputLayout, shape, numLanes);
- auto resCoords = computePerLaneElementCoords(targetLayout, shape, numLanes);
- if (failed(srcCoords) || failed(resCoords))
+ // Tabulate ownership: what each lane brings and what each slot needs.
+ auto inputOwned = computeOwnedCoords(inputLayout, shape, subgroupSize);
+ auto targetOwned = computeOwnedCoords(targetLayout, shape, subgroupSize);
+ if (failed(inputOwned) || failed(targetOwned))
return failure();
// Require the layouts to distribute what the vector types account for, and
- // the target to 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])
+ // the lanes past the first `numSlots` to be replicas.
+ for (int64_t lane = 0; lane < subgroupSize; lane++) {
+ if (static_cast<int64_t>((*inputOwned)[lane].size()) != srcNumElems ||
+ static_cast<int64_t>((*targetOwned)[lane].size()) != resNumElems ||
+ (*targetOwned)[lane] != (*targetOwned)[lane % numSlots])
return failure();
}
- SmallVector<BroadcastRedistribution> redistributions;
+ // Work out where every element of the result comes from before emitting
+ // anything, so that an unsupported layout change leaves no ops behind.
+ SmallVector<ElementSource> sources;
for (int64_t pos = 0; pos < resNumElems; pos++) {
- auto redistribution = deriveBroadcastRedistribution(
- *srcCoords, *resCoords, pos, numLanes, numTargetLanes);
- if (failed(redistribution))
+ FailureOr<ElementSource> source = deriveElementSource(
+ *inputOwned, *targetOwned, pos, subgroupSize, numSlots);
+ if (failed(source))
return failure();
- redistributions.push_back(*redistribution);
+ sources.push_back(*source);
}
- // Shuffles move same-width integers, which any lane data type bitcasts to.
+ // The lane's fragment, flattened and bitcast to integers: shuffles move
+ // same-width integers, which any lane data type bitcasts to.
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);
+ Value fragment = src;
+ auto flatFragmentTy = VectorType::get({srcNumElems}, elemTy);
+ if (srcTy != flatFragmentTy)
+ fragment =
+ vector::ShapeCastOp::create(rewriter, loc, flatFragmentTy, fragment);
if (elemTy != shuffleTy)
- flatSrc = vector::BitCastOp::create(
- rewriter, loc, VectorType::get({srcNumElems}, shuffleTy), flatSrc);
+ fragment = vector::BitCastOp::create(
+ rewriter, loc, VectorType::get({srcNumElems}, shuffleTy), fragment);
- 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));
+ // The lane's slot, which every fragment index is a function of.
+ Value slot = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+ /*upperBound=*/mlir::IntegerAttr());
+ if (numSlots != subgroupSize)
+ slot = arith::RemUIOp::create(
+ rewriter, loc, slot,
+ arith::ConstantIndexOp::create(rewriter, loc, numSlots));
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;
-
- // `(laneIdx / divisor) % modulus`, shared by elements reading the same
- // coordinate. Either step is skipped when it is a no-op over `laneIdx`.
- 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;
+ Value slotI32;
+
+ // `(slot / slotStride) % extent`, the lane's index along the lane_layout
+ // dimension the fragment index varies with. Indices over the same dimension
+ // share it, and either step is skipped when it is a no-op over `slot`.
+ llvm::DenseMap<std::pair<int64_t, int64_t>, Value> dimIndices;
+ auto getDimIndex = [&](const FragmentIndex &index) -> Value {
+ Value &dimIndex = dimIndices[{index.slotStride, index.extent}];
+ if (dimIndex)
+ return dimIndex;
+ dimIndex = slot;
+ if (index.slotStride != 1)
+ dimIndex = arith::DivUIOp::create(
+ rewriter, loc, dimIndex,
+ arith::ConstantIndexOp::create(rewriter, loc, index.slotStride));
+ if (index.extent < numSlots / index.slotStride)
+ dimIndex = arith::RemUIOp::create(
+ rewriter, loc, dimIndex,
+ arith::ConstantIndexOp::create(rewriter, loc, index.extent));
+ return dimIndex;
};
+ // `stride * dimIndex + offset`, the index the lane extracts.
+ auto getFragmentIndex = [&](const FragmentIndex &index) -> OpFoldResult {
+ // Every slot reads the same element; no dimension is involved.
+ if (index.stride == 0)
+ return rewriter.getIndexAttr(index.offset);
+ Value value = getDimIndex(index);
+ if (index.stride != 1)
+ value = arith::MulIOp::create(
+ rewriter, loc, value,
+ arith::ConstantIndexOp::create(rewriter, loc, index.stride));
+ if (index.offset != 0)
+ value = arith::AddIOp::create(
+ rewriter, loc, value,
+ arith::ConstantIndexOp::create(rewriter, loc, index.offset));
+ return value;
+ };
+
+ // Extract, shuffle and insert one element of the result at a time.
llvm::DenseMap<std::tuple<int64_t, int64_t, int64_t, int64_t>, Value>
extracted;
- for (auto [pos, redistribution] : llvm::enumerate(redistributions)) {
- // Result elements reading 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 =
- getLaneTerm(redistribution.divisor, redistribution.modulus);
- 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) {
+ for (auto [pos, source] : llvm::enumerate(sources)) {
+ // Elements reading the same index of the fragment share the extract.
+ Value &value = extracted[source.index.asTuple()];
+ if (!value)
+ value = vector::ExtractOp::create(rewriter, loc, fragment,
+ getFragmentIndex(source.index));
+
+ Value owned = value;
+ if (source.needsShuffle) {
if (!width) {
- width = arith::ConstantIntOp::create(rewriter, loc, i32Ty, numLanes);
- laneIdxI32 = arith::IndexCastOp::create(rewriter, loc, i32Ty, laneIdx);
+ width =
+ arith::ConstantIntOp::create(rewriter, loc, i32Ty, subgroupSize);
+ slotI32 = arith::IndexCastOp::create(rewriter, loc, i32Ty, slot);
}
- 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);
+ Value donor = slotI32;
+ if (source.donorDelta != 0)
+ donor =
+ arith::AddIOp::create(rewriter, loc, donor,
+ arith::ConstantIntOp::create(
+ rewriter, loc, i32Ty, source.donorDelta));
+ owned = gpu::ShuffleOp::create(rewriter, loc, owned, donor, width,
+ gpu::ShuffleMode::IDX)
+ .getResult(0);
}
- res = vector::InsertOp::create(rewriter, loc, result, res, pos);
+ res = vector::InsertOp::create(rewriter, loc, owned, res, pos);
}
+
if (elemTy != shuffleTy)
res = vector::BitCastOp::create(
rewriter, loc, VectorType::get({resNumElems}, elemTy), res);
@@ -2130,19 +2265,19 @@ struct SgToLaneConvertLayout
}
}
- // Input layout broadcasts over groups of lanes while the target
+ // Input layout replicates over broadcast groups of lanes while the target
// distributes, so data has to move 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();
+ int64_t subgroupSize = uArch->getSubgroupSize();
// The target may leave the remaining lanes with a replicated value.
- if (getLaneLayoutSize(inputLayout) == numLanes) {
+ if (getNumParticipatingLanes(inputLayout) == subgroupSize) {
FailureOr<Value> res = redistributeBroadcastedValue(
rewriter, op.getLoc(), adaptor.getSource(), *resDistTy, inputLayout,
- targetLayout, resShapeVec, numLanes);
+ targetLayout, resShapeVec, subgroupSize);
if (succeeded(res)) {
rewriter.replaceOp(op, *res);
return success();
>From dbd5014c3f6f0acaf0c4bbc24c0c1792cd564a98 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 31 Aug 2026 20:54:51 +0000
Subject: [PATCH 03/46] State the precondition of the broadcast redistribution
as a predicate
Review feedback on two points.
First, the guard read `getLaneLayoutSize(inputLayout) == numLanes`, which said
nothing about what the pattern requires of the two layouts. Give it a name and
the layout vocabulary: isBroadcastRedistribution checks that both layouts hand
the lanes the same distribution unit -- equal effective lane_data -- so that only
the assignment of elements to lanes changes, that the input lays out the whole
subgroup, and that a target laying out only part of it does so evenly. The
unit-lane_data restriction that computeOwnedCoords imposes is now a TODO next to
the lane_data check rather than a silent failure further down.
Not checked is that the lane_layouts differ. The layouts are already known to be
incompatible at that point, since isCompatibleWith returns earlier, and requiring
a lane_layout difference would reject a change that is only in `order`, which the
ownership tables handle like any other.
Second, getNumParticipatingLanes unwrapped a slice with flatten() without saying
why that is enough for a nested one. flatten() coalesces the whole chain, so its
parent is always a plain LayoutAttr -- the same idiom SliceAttr::isCompatibleWith
uses -- which the cast now makes explicit. Cover it with a test whose input
layout nests two slices: it lowers exactly like the single-slice case, one
extract and one shuffle.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 81 ++++++++++---------
.../XeGPU/sg-to-lane-distribute-unit.mlir | 34 ++++++++
2 files changed, 79 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 4fb35ac292d01..5ce31697f8923 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1865,11 +1865,35 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
/// Number of lanes of the subgroup `layout` distributes over. Lanes covered by
/// sliced dimensions count too: they hold a replicated fragment, not nothing.
static int64_t getNumParticipatingLanes(xegpu::DistributeLayoutAttr layout) {
+ // flatten() coalesces nested slices, so the parent it leaves is always a
+ // plain LayoutAttr carrying the lane_layout of the whole subgroup.
if (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
- layout = sliceAttr.flatten().getParent();
+ layout = cast<xegpu::LayoutAttr>(sliceAttr.flatten().getParent());
return computeProduct(layout.getEffectiveLaneLayoutAsInt());
}
+/// Matches the layout change `redistributeBroadcastedValue` implements: both
+/// layouts hand the lanes the same distribution unit (equal effective
+/// lane_data), so only the assignment of elements to lanes changes, and the
+/// input lays out the whole subgroup. A target laying out only part of it
+/// replicates its fragments over the lanes it leaves out, which
+/// `redistributeBroadcastedValue` verifies element-wise.
+///
+/// That the layouts differ is not checked: they are already known to be
+/// incompatible here, and the difference may equally be in `order`.
+static bool isBroadcastRedistribution(xegpu::DistributeLayoutAttr inputLayout,
+ xegpu::DistributeLayoutAttr targetLayout,
+ int64_t subgroupSize) {
+ // TODO: computeOwnedCoords only tabulates unit lane_data, so equal but
+ // non-unit lane_data is rejected further down rather than here.
+ if (inputLayout.getEffectiveLaneDataAsInt() !=
+ targetLayout.getEffectiveLaneDataAsInt())
+ return false;
+ int64_t numSlots = getNumParticipatingLanes(targetLayout);
+ return getNumParticipatingLanes(inputLayout) == subgroupSize &&
+ numSlots > 0 && subgroupSize % numSlots == 0;
+}
+
/// Ownership table: `owned[lane][i]` is the coordinate, in the undistributed
/// shape, of element `i` of the fragment `lane` holds.
using OwnedCoords = SmallVector<SmallVector<SmallVector<int64_t>>>;
@@ -1889,18 +1913,16 @@ computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
return owned;
}
-/// Index into the lane's fragment as a function of its slot:
-/// `stride * ((slot / slotStride) % extent) + offset`.
+/// Index into the lane's fragment as a function of its slot.
struct FragmentIndex {
int64_t stride;
int64_t offset;
- /// Slots apart two lanes have to be for their index along the selected
- /// lane_layout dimension to differ by one.
+ /// Slots apart for the index along the selected dimension to advance by one.
int64_t slotStride;
/// Number of values that index takes.
int64_t extent;
- /// The index `slot` reads. Compile-time twin of the emitted arithmetic.
+ /// Compile-time twin of the arithmetic the emitter builds.
int64_t at(int64_t slot) const {
return stride * ((slot / slotStride) % extent) + offset;
}
@@ -1911,10 +1933,8 @@ struct FragmentIndex {
}
};
-/// Where one element of the target fragment comes from: `index` picks it out of
-/// the fragment of the donor `donorDelta` lanes away, and `gpu.shuffle idx`
-/// moves it from there to the lane that owns it, unless the donor is that lane
-/// itself.
+/// Where one element of the target fragment comes from: which donor, and where
+/// in that donor's fragment.
struct ElementSource {
FragmentIndex index;
int64_t donorDelta;
@@ -1925,10 +1945,8 @@ struct ElementSource {
/// to a `FragmentIndex`. Fails when no index reproduces `needed`.
static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
int64_t numSlots) {
- // Slots sharing an index along the selected lane_layout dimension are
- // adjacent, so the first change in `needed` tells how many of them there are
- // (`slotStride`) and how far apart the elements of consecutive indices are
- // (`stride`).
+ // Slots sharing an index are adjacent, so the first change in `needed` gives
+ // both how many of them there are and how far apart their elements sit.
int64_t offset = needed[0];
int64_t slotStride = 1;
while (slotStride < numSlots && needed[slotStride] == offset)
@@ -1937,8 +1955,8 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
if (stride < 0 || numSlots % slotStride != 0)
return failure();
- // The index wraps around once it has taken all its values, which is how far
- // `needed` gets from `offset`.
+ // The index wraps once it has taken all its values, so how far `needed` gets
+ // from `offset` gives the extent.
int64_t extent = 1;
if (stride == 0) {
// Every slot reads `offset`; no dimension is involved.
@@ -1961,10 +1979,8 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
return index;
}
-/// Works out where element `pos` of the target fragment comes from, using the
-/// two ownership tables. Fails when no donor group provides that element
-/// through a single extract, i.e. when the data movement is not of the form
-/// described above.
+/// Works out where element `pos` of the target fragment comes from. Fails when
+/// no donor group provides it through a single extract.
static FailureOr<ElementSource>
deriveElementSource(const OwnedCoords &inputOwned,
const OwnedCoords &targetOwned, int64_t pos,
@@ -1978,10 +1994,7 @@ deriveElementSource(const OwnedCoords &inputOwned,
return -1;
};
- // A donor shares the slot of the lane it serves, so it is a whole number of
- // target lane layouts away. Replication can put the element in more than one
- // donor group; the smallest `delta` wins, as it is the only one that can read
- // the lane's own fragment and so skip the shuffle.
+ // Donor groups, smallest `delta` first.
for (int64_t donorDelta = 0; donorDelta < subgroupSize;
donorDelta += numSlots) {
// Where in its fragment each slot's donor keeps the element it needs.
@@ -1999,9 +2012,8 @@ deriveElementSource(const OwnedCoords &inputOwned,
if (failed(index))
continue;
- // A shuffle is needed unless the donor is the lane itself, that is unless
- // the local extract already lands on the element the lane owns. Only the
- // lanes of a slot decide that, the rest being replicas.
+ // The donor is the lane itself only if the local extract lands on the
+ // element the lane owns, which just the slots get to decide.
bool needsShuffle =
donorDelta != 0 ||
!llvm::all_of(llvm::seq<int64_t>(0, numSlots), [&](int64_t slot) {
@@ -2013,9 +2025,8 @@ deriveElementSource(const OwnedCoords &inputOwned,
}
/// Redistributes `src`, the fragment the lane holds under `inputLayout`, into
-/// the `resTy` fragment of the `shape`-shaped value it owns under
-/// `targetLayout`, over a subgroup of `subgroupSize` lanes. See the algorithm
-/// above; returns failure if the layout change is not of that form.
+/// the `resTy` fragment it owns under `targetLayout`. Fails if the layout
+/// change is not of the form described above.
static FailureOr<Value>
redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
Value src, VectorType resTy,
@@ -2089,9 +2100,8 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
Value width;
Value slotI32;
- // `(slot / slotStride) % extent`, the lane's index along the lane_layout
- // dimension the fragment index varies with. Indices over the same dimension
- // share it, and either step is skipped when it is a no-op over `slot`.
+ // The lane's index along the selected dimension, shared by sources over the
+ // same one. Either step is skipped when it is a no-op over `slot`.
llvm::DenseMap<std::pair<int64_t, int64_t>, Value> dimIndices;
auto getDimIndex = [&](const FragmentIndex &index) -> Value {
Value &dimIndex = dimIndices[{index.slotStride, index.extent}];
@@ -2109,7 +2119,7 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
return dimIndex;
};
- // `stride * dimIndex + offset`, the index the lane extracts.
+ // The index the lane extracts.
auto getFragmentIndex = [&](const FragmentIndex &index) -> OpFoldResult {
// Every slot reads the same element; no dimension is involved.
if (index.stride == 0)
@@ -2273,8 +2283,7 @@ struct SgToLaneConvertLayout
targetLayout, cast<VectorType>(valType));
if (uArch && succeeded(resDistTy)) {
int64_t subgroupSize = uArch->getSubgroupSize();
- // The target may leave the remaining lanes with a replicated value.
- if (getNumParticipatingLanes(inputLayout) == subgroupSize) {
+ if (isBroadcastRedistribution(inputLayout, targetLayout, subgroupSize)) {
FailureOr<Value> res = redistributeBroadcastedValue(
rewriter, op.getLoc(), adaptor.getSource(), *resDistTy, inputLayout,
targetLayout, resShapeVec, subgroupSize);
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 c36a70635d400..b6bcc6521e6c0 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1682,4 +1682,38 @@ gpu.func @convert_layout_broadcast_to_sliced_target() {
"some_use"(%cvt) : (vector<8x2xbf16>) -> ()
gpu.return
}
+
+// The same redistribution as @convert_layout_broadcast_to_lane_distributed, but
+// reaching it through two nested slices. They are coalesced, so what counts is
+// the parent lane_layout covering all 16 lanes and the dims the slices drop
+// together: the 8 lanes of the dropped dim 0 hold the same fragment, while the
+// kept dim of extent 2 splits the columns between the two groups.
+// CHECK-LABEL: gpu.func @convert_layout_broadcast_nested_slice
+// 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: %[[SLOT:.*]] = arith.remui %[[LANE]], %[[C8]] : index
+// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<2xi8>
+// CHECK: %[[ELEM:.*]] = vector.extract %[[BITS]][%[[SLOT]]] : i8 from vector<8xi8>
+// CHECK: %[[INS0:.*]] = vector.insert %[[ELEM]], %[[ZERO]] [0] : i8 into vector<2xi8>
+// CHECK: %[[WIDTH:.*]] = arith.constant 16 : i32
+// CHECK: %[[SLOT_I32:.*]] = arith.index_cast %[[SLOT]] : index to i32
+// CHECK: %[[C8_I32:.*]] = arith.constant 8 : i32
+// CHECK: %[[DONOR:.*]] = arith.addi %[[SLOT_I32]], %[[C8_I32]] : i32
+// CHECK: %[[SHUF:.*]], %{{.*}} = gpu.shuffle idx %[[ELEM]], %[[DONOR]], %[[WIDTH]] : i8
+// CHECK: %[[INS1:.*]] = vector.insert %[[SHUF]], %[[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_nested_slice() {
+ %scale_a_src = arith.constant dense<1.0> : vector<8x2xf8E8M0FNU>
+ %cvt = xegpu.convert_layout %scale_a_src
+ <{
+ input_layout = #xegpu.slice<#xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 1, 2], lane_data = [1, 1, 1, 1], order = [0, 3, 2, 1]>, dims = [2]>, dims = [0]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>
+ }> : vector<8x2xf8E8M0FNU>
+ "some_use"(%cvt) : (vector<8x2xf8E8M0FNU>) -> ()
+ gpu.return
+}
}
>From f0d97a7245e638d5697d57d66b5e2cece8489ad9 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 1 Sep 2026 16:00:03 +0000
Subject: [PATCH 04/46] Lead the broadcast redistribution with what it does and
does not handle
Review feedback: the scope of the pattern was not clear, and the explanation was
prose where examples would do.
Open the section with the capability, as layout pairs rather than description.
Three handled cases, which are the three tests: broadcast over the subgroup to 8
lanes (2 extracts, no shuffle), broadcast over two groups of 8 to 8 lanes (1
extract, 1 shuffle, the running example), and broadcast over the subgroup to two
groups of 8 (8 extracts, no shuffle). Then what it declines and why: lane_data
differing, equal but non-unit lane_data, an input covering part of the subgroup,
a target that does not replicate over the lanes it leaves out, elements narrower
than a byte, and indices outside the restricted form. A reader can now tell in
one screen whether their conversion is covered.
The concepts and the index shapes now refer to those three cases instead of
carrying their own examples, which is where the added lines come back.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 249 ++++++++++++------
1 file changed, 171 insertions(+), 78 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 5ce31697f8923..e189ea7150b96 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1767,104 +1767,197 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// Lowers a `convert_layout` that only relocates data between lanes: the input
// layout replicates the data over lanes, while the target layout hands each
-// lane a different part of it. Every element of the result already exists in
-// some lane, so nothing is recomputed. The scale operands of scaled matrix
-// multiplication need this, being produced broadcasted and consumed
-// distributed.
+// lane a different part of it. Every element already exists in some lane, so
+// nothing is recomputed. The scale operands of scaled matrix multiplication
+// need this, being produced broadcasted and consumed distributed.
+//
+// What it handles
+// ---------------
+//
+// The three cases below, all on a `vector<8x2>` over 16 lanes, are the three
+// tests. `#xegpu.` prefixes are dropped for width.
+//
+// 1. Broadcast over the whole subgroup, distributed to 8 lanes.
+//
+// slice<layout<lane_layout = [1, 1, 16]>, dims = [2]>
+// -> layout<lane_layout = [8, 1]>
+//
+// Every lane arrives holding all 16 elements and leaves holding one row.
+// Emits 2 extracts and no shuffle: every lane already has what it needs.
+//
+// 2. Broadcast over two groups of 8 lanes, distributed to 8 lanes.
+//
+// slice<layout<lane_layout = [8, 1, 2], order = [0, 2, 1]>, dims = [0]>
+// -> layout<lane_layout = [8, 1]>
+//
+// Lanes 0-7 arrive holding all of column 0, lanes 8-15 all of column 1, and
+// lane `i` leaves holding row `i`. So lane 3 wants (3, 0) and (3, 1); it
+// holds (3, 0) already, while (3, 1) is in lane 11. Emits 1 extract and 1
+// shuffle. This is the running example used below.
+//
+// 3. Broadcast over the whole subgroup, distributed to two groups of 8 lanes.
+//
+// slice<layout<lane_layout = [1, 1, 16]>, dims = [2]>
+// -> slice<layout<lane_layout = [8, 1, 2], order = [0, 2, 1]>,
+// dims = [0]>
+//
+// Every lane arrives holding all 16 elements; lanes 0-7 leave holding column
+// 0 and lanes 8-15 column 1. Emits 8 extracts and no shuffle.
+//
+// Nested slices are coalesced before anything else looks at a layout, so only
+// the lane_layout of the underlying layout and the union of the sliced dims
+// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
+//
+// What it declines
+// ----------------
+//
+// Each of these reports a match failure rather than lowering. Same notation,
+// and `vector<8x2>` over 16 lanes unless said otherwise.
+//
+// * `lane_data` differing between the two layouts:
+//
+// layout<[8, 1], [1, 1]> -> layout<[1, 8], [1, 2]>
+//
+// A `lane_data`-only change is the repack pattern's job; changing it and
+// `lane_layout` together is handled by neither.
+//
+// * `lane_data` equal but not unit:
+//
+// slice<layout<[1, 1, 16], [1, 2, 1]>, dims = [2]>
+// -> layout<[8, 1], [1, 2]>
+//
+// A lane then holds two elements per distribution unit, which the ownership
+// tables do not enumerate.
+//
+// * an input layout that covers only part of the subgroup:
+//
+// layout<[8, 1], [1, 1]> -> layout<[1, 8], [1, 1]>
+//
+// Only the target may leave lanes out; the input has to name all 16.
+//
+// * elements narrower than a byte:
+//
+// slice<layout<[1, 1, 16]>, dims = [2]> -> layout<[8, 1]>
+// on vector<8x2xf4E2M1FN>
+//
+// There is no type for `gpu.shuffle` to carry a 4-bit element in.
+//
+// * a lane's element index not being one dimension of the target lane layout,
+// as when `order` transposes it:
+//
+// slice<layout<[1, 1, 16]>, dims = [2]>
+// -> layout<[4, 4], order = [0, 1]> on vector<4x4>
+//
+// Lane `i` would need element `(i % 4) * 4 + i / 4`, which the restricted
+// index form below cannot express.
//
// Concepts
// --------
//
-// fragment The elements a lane holds under a layout, as a flat list --
-// the lane's distributed vector. Both layouts give a lane a
-// fragment, and the conversion turns the input one into the
-// target one.
+// fragment
+// What a lane holds under a layout: 8x1 under case 2's input layout, 1x2
+// under its target.
+//
+// broadcast group
+// Lanes holding the same fragment. A sliced lane_layout dimension forms
+// one, since nothing is distributed over its lanes: in case 2 dim 0 is
+// sliced and spans 8 lanes, so lanes 0-7 are one group and 8-15 another.
+//
+// slot
+// A lane's index among the lanes the target layout distributes to, `lane %
+// numSlots`. In case 2 `numSlots` is 8, so the lanes pair up: lane 11 has
+// slot 3 like lane 3, and runs the same arithmetic on its own fragment. How
+// many elements a lane holds is a separate matter -- it can own several,
+// which show up as several positions in its fragment, not as several slots.
//
-// broadcast group Lanes that hold the same fragment. A sliced lane_layout
-// dimension forms one: nothing is distributed over its lanes,
-// so they are all given the same data. The group spans that
-// dimension -- all 16 lanes if it covers the subgroup, 8 if it
-// covers half of it.
+// donor
+// The lane a value is read from, `slot + delta`. `gpu.shuffle idx` is the
+// only way to move data between lanes and it exchanges one value per lane,
+// so every lane runs the same extract; a donor can therefore only serve
+// lanes of its own slot, which makes `delta` a multiple of `numSlots`. In
+// case 2, slot 3 reads (3, 1) from lane 3 + 8.
//
-// slot A lane's linear id within the target lane layout,
-// `lane % numSlots`. The target layout gives the same fragment
-// to every lane of a slot, so there are `numSlots` distinct
-// target fragments and the lanes past the first `numSlots` are
-// replicas.
+// The two indices come from opposite sides: `slot` from the target layout,
+// which decides how many lanes the value is distributed to, and `delta`
+// from the input layout, which decides how many copies of it there are.
//
-// donor The lane a value is read from. Only `gpu.shuffle idx` moves
-// a value between lanes, and it exchanges exactly one value
-// per lane: each lane offers one value and receives one lane's
-// offering. So all lanes have to run the same extract, and a
-// lane can only be served by a lane that arrives at the same
-// extract index -- one of its own slot, i.e. `slot + delta`
-// for a `delta` that is a multiple of `numSlots`. Lanes
-// `delta` apart form a donor group, which partitions the
-// subgroup by the target layout and is not a broadcast group.
+// element source
+// Where one element of the target fragment comes from: the donor, and the
+// index to extract from that donor's fragment. A shuffle is needed exactly
+// when the donor cannot be the lane itself. Case 2 has two, one per element
+// of its 1x2 fragment:
//
-// element source Where one element of the target fragment comes from: the
-// donor to take it from, and the index to extract from that
-// donor's fragment. A shuffle is needed exactly when that
-// donor cannot be the lane itself. The index is a function of
-// the slot,
+// element 0 -> delta = 0, index = slot // (3, 0), already local
+// element 1 -> delta = 8, index = slot // (3, 1), from lane 11
//
-// stride * ((slot / slotStride) % extent) + offset
+// The index has to be computable from the lane id, so it is restricted to
//
-// which says "one dimension of the target lane layout selects
-// one element of the fragment": `(slot / slotStride) % extent`
-// is the lane's index along that dimension, `slotStride` being
-// how many slots apart two lanes have to be for it to differ
-// by one and `extent` how many values it takes. It costs a
-// couple of `arith` ops on the lane id, where an arbitrary
-// per-slot index would need a materialized table and a gather.
+// stride * ((slot / slotStride) % extent) + offset
+//
+// which walks one dimension of the target lane layout: `(slot / slotStride) %
+// extent` is the lane's index along that dimension and `stride`/`offset` place
+// it in the fragment. The three cases give it three shapes:
+//
+// index = slot consecutive slots take consecutive elements (case 2:
+// slotStride = 1, extent = 8)
+// index = 2 * slot every other element, because the fragment holds whole
+// 2-element rows (case 1)
+// index = slot / 8 eight consecutive slots take the same element, the
+// column their half of the subgroup owns (case 3:
+// slotStride = 8, extent = 2)
+//
+// An arbitrary index per slot would need a materialized table and a gather;
+// this costs a couple of `arith` ops on the lane id.
//
// Algorithm
// ---------
//
-// tabulate inputOwned[lane] and targetOwned[slot] // element coordinates
+// tabulate inputOwned[lane] and targetOwned[lane] // element coordinates
// for each element `pos` of the target fragment:
-// for delta in 0, numSlots, 2 * numSlots, ...: // smallest first
-// needed[slot] = index of targetOwned[slot][pos]
-// in inputOwned[slot + delta]
-// if a slot's donor does not hold its element: next delta
-// index = fit `needed` to stride/offset/slotStride/extent,
-// else next delta
+// for delta = 0, numSlots, 2 * numSlots, ... // smallest first
+// for each slot:
+// needed[slot] = index of the element targetOwned[slot][pos] in the
+// fragment of its donor, lane `slot + delta`
+// if that donor does not hold the element: give up on this delta
+// look for one stride/offset/slotStride/extent whose index reproduces
+// every needed[slot]; if there is none: next delta
// accept (delta, index) and stop
// if no delta fits: fail, the layout change is not of this form
// for each element source:
-// %idx = index(slot) // arith on gpu.lane_id
-// %val = vector.extract %fragment[%idx]
-// %val = gpu.shuffle idx %val, slot + delta // unless donor == lane
-// %res = vector.insert %val, %res[pos]
-//
-// Replication leaves the donor ambiguous: several donor groups may hold the
-// element and they all yield the same value, so the smallest `delta` is taken.
-// Only `delta = 0` can make the donor the lane itself, so it is the only one
-// that may need no shuffle, and taking it is what leaves a fully broadcast
-// input shuffle-free.
+// %idx = index(slot) // arith on gpu.lane_id
+// %val = vector.extract %fragment[%idx]
+// %donor = slot + delta
+// %val = gpu.shuffle idx %val, %donor // unless %donor is %lane
+// %res = vector.insert %val, %res[pos]
//
-// For example, redistributing a `vector<8x2>` whose input fragments hold column
-// 0 in lanes [0, 8) and column 1 in lanes [8, 16) to a target that gives lane
-// `i` row `i` (`numSlots` = 8, so lanes [8, 16) are replicas) yields
+// `%donor` is the lane itself when `delta` is 0 and the extract already lands
+// on an element the lane owns, and that is the only case where the shuffle is
+// dropped. Replication makes it a choice: several donor groups may hold the
+// element and all give the same value, so the smallest `delta` is taken, which
+// is what leaves a fully broadcast input shuffle-free.
//
-// pos 0 -> delta = 0, idx = slot // column 0 is in the lane's own group
-// pos 1 -> delta = 8, idx = slot // column 1 is one donor group away
+// Only lanes 0 to `numSlots` - 1 are required to end up holding the elements
+// the target layout assigns their slot; each is the one lane of its slot the
+// layout distributes to. Lanes sharing a slot run the same arithmetic on their
+// own fragment, so when a shuffle is emitted they all receive the donor's value
+// and are equally correct, and when it is dropped they keep whatever that index
+// holds in their own fragment -- which is why dropping it relies on nobody
+// reading them.
//
// An index is derived from the tables and then verified against them, so a
-// layout change that does not have this form is reported as a match failure
-// rather than lowered incorrectly.
+// layout change that is not of this form is reported as a match failure rather
+// than lowered incorrectly.
//
-// Assumption: only the first `numSlots` lanes have to end up with the elements
-// the target layout assigns them, the rest being replicas. Whether the donor
-// can be the lane itself is therefore decided over the slots alone, and for all
-// of them at once since the index is uniform over the subgroup: one slot whose
-// local extract misses the element it owns forces the shuffle on every lane.
+// One caveat: after the rewrite only the first `numSlots` lanes are
+// guaranteed to hold what the target layout assigns them. The others are
+// replicas whose local extract can land on a different element, so dropping a
+// shuffle relies on nobody reading them.
//
//===----------------------------------------------------------------------===//
-/// Number of lanes of the subgroup `layout` distributes over. Lanes covered by
-/// sliced dimensions count too: they hold a replicated fragment, not nothing.
-static int64_t getNumParticipatingLanes(xegpu::DistributeLayoutAttr layout) {
+/// Number of lanes `layout` distributes to. Lanes covered by sliced dimensions
+/// count too: they hold a replicated fragment, not nothing.
+static int64_t getNumDistributedLanes(xegpu::DistributeLayoutAttr layout) {
// flatten() coalesces nested slices, so the parent it leaves is always a
// plain LayoutAttr carrying the lane_layout of the whole subgroup.
if (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
@@ -1889,9 +1982,9 @@ static bool isBroadcastRedistribution(xegpu::DistributeLayoutAttr inputLayout,
if (inputLayout.getEffectiveLaneDataAsInt() !=
targetLayout.getEffectiveLaneDataAsInt())
return false;
- int64_t numSlots = getNumParticipatingLanes(targetLayout);
- return getNumParticipatingLanes(inputLayout) == subgroupSize &&
- numSlots > 0 && subgroupSize % numSlots == 0;
+ int64_t numSlots = getNumDistributedLanes(targetLayout);
+ return getNumDistributedLanes(inputLayout) == subgroupSize && numSlots > 0 &&
+ subgroupSize % numSlots == 0;
}
/// Ownership table: `owned[lane][i]` is the coordinate, in the undistributed
@@ -1941,8 +2034,8 @@ struct ElementSource {
bool needsShuffle;
};
-/// Fits `needed`, the fragment index each of the `numSlots` slots has to read,
-/// to a `FragmentIndex`. Fails when no index reproduces `needed`.
+/// Fits `needed`, the fragment index each of the `numSlots` slots has to
+/// read, to a `FragmentIndex`. Fails when no index reproduces `needed`.
static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
int64_t numSlots) {
// Slots sharing an index are adjacent, so the first change in `needed` gives
@@ -2038,7 +2131,7 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
return failure();
int64_t srcNumElems = srcTy.getNumElements();
int64_t resNumElems = resTy.getNumElements();
- int64_t numSlots = getNumParticipatingLanes(targetLayout);
+ int64_t numSlots = getNumDistributedLanes(targetLayout);
if (numSlots < 1 || subgroupSize % numSlots != 0)
return failure();
// gpu.shuffle only moves the integer widths a lane can hold.
>From 52c127da71d6639dde57c3d636a61ac94afdee35 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 1 Sep 2026 21:36:45 +0000
Subject: [PATCH 05/46] Decide the shuffle by the donor group alone
The shuffle condition also tested whether each slot`s local extract lands on the
element it owns, which cannot fail: `needed[slot]` is by construction the index of
that element in the fragment of lane `slot + delta`, and the fit is verified
against `needed`, so at delta 0 the extract is local for every slot. It became
dead when the test narrowed from the whole subgroup to the slots, so drop it and
let the donor group decide.
Say why the narrow test is the right one: the target layout assigns no elements to
the lanes it does not distribute to, so what they hold is unspecified.
Note the missing packing while here: a shuffle moves a whole 32-bit lane word but
carries one element, so an f8 fragment costs four shuffles where one would do.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 19 +++++++++----------
1 file changed, 9 insertions(+), 10 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index e189ea7150b96..c716d627dfa26 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1941,8 +1941,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// layout distributes to. Lanes sharing a slot run the same arithmetic on their
// own fragment, so when a shuffle is emitted they all receive the donor's value
// and are equally correct, and when it is dropped they keep whatever that index
-// holds in their own fragment -- which is why dropping it relies on nobody
-// reading them.
+// holds in their own fragment. The target layout assigns those lanes no
+// elements, so what they hold is unspecified.
//
// An index is derived from the tables and then verified against them, so a
// layout change that is not of this form is reported as a match failure rather
@@ -2105,14 +2105,10 @@ deriveElementSource(const OwnedCoords &inputOwned,
if (failed(index))
continue;
- // The donor is the lane itself only if the local extract lands on the
- // element the lane owns, which just the slots get to decide.
- bool needsShuffle =
- donorDelta != 0 ||
- !llvm::all_of(llvm::seq<int64_t>(0, numSlots), [&](int64_t slot) {
- return inputOwned[slot][index->at(slot)] == targetOwned[slot][pos];
- });
- return ElementSource{*index, donorDelta, needsShuffle};
+ // The donor is the lane itself exactly when `delta` is 0: `needed[slot]`
+ // was located in the fragment of lane `slot + delta`, so at 0 every slot's
+ // extract is local by construction.
+ return ElementSource{*index, donorDelta, /*needsShuffle=*/donorDelta != 0};
}
return failure();
}
@@ -2230,6 +2226,9 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
};
// Extract, shuffle and insert one element of the result at a time.
+ // TODO: one shuffle per element, though `gpu.shuffle` moves a whole 32-bit
+ // lane word. Elements sharing a donor could travel in one, as
+ // shuffleDataAsLaneLayoutChange does.
llvm::DenseMap<std::tuple<int64_t, int64_t, int64_t, int64_t>, Value>
extracted;
for (auto [pos, source] : llvm::enumerate(sources)) {
>From 6308d2180c48401288db605d49b2a4fecf801cb0 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 1 Sep 2026 23:46:08 +0000
Subject: [PATCH 06/46] Lift the input-coverage limit and ground the vocabulary
Addresses review feedback that the input and target restrictions were
asymmetric without explanation, and that the index expression was hard to
follow.
The input no longer has to lay out the whole subgroup. A lane_layout covering
fewer lanes than the subgroup has is already a replicated layout:
computeStaticDistributedCoords delinearizes the lane id modulo each extent, so
the lanes past the product repeat the ones before it, and computeOwnedCoords
returns a defined fragment for every lane either way. Nothing in the algorithm
depended on the coverage, so the condition is simply dropped, with no
replacement guard: delta is always a multiple of the lane period, so a donor
reduces to the same slot as its reader and extracts the same index however many
lanes the input covers. That holds even when the period does not divide the
subgroup size.
Renames what the code got wrong about itself. getNumDistributedLanes returned
the parent lane_layout product, which is the modulus after which the layout
repeats, not the number of lanes holding distinct data -- a sliced dimension
contributes its lanes to it while replicating one fragment. It is getLanePeriod
now, and FragmentIndex::slotStride/extent, which describe the target
lane_layout dimension being walked, are dimStride/dimExtent.
The section header opened by asserting that the input replicates and the target
distributes. Both are wrong: either side may replicate, and every target in the
tests does. It now states what the two sides mean -- the input supplies the
values, the target states the obligation -- and describes the two spellings of
replication, a sliced dimension and a short lane_layout, as the same thing.
The index expression is introduced through one worked example with every term
given its value, followed by a table of the four cases with all four constants,
and a note on where the literals come from. On a vector<8x2> over 16 lanes the
lane period, case 2 delta and case 3 dimStride are all 8 and are three
different quantities, which the text now says.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 214 +++++++++++-------
.../XeGPU/sg-to-lane-distribute-unit.mlir | 41 ++++
2 files changed, 169 insertions(+), 86 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index c716d627dfa26..8ccc0956e4d5d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1765,17 +1765,31 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Broadcast redistribution
//===----------------------------------------------------------------------===//
//
-// Lowers a `convert_layout` that only relocates data between lanes: the input
-// layout replicates the data over lanes, while the target layout hands each
-// lane a different part of it. Every element already exists in some lane, so
-// nothing is recomputed. The scale operands of scaled matrix multiplication
-// need this, being produced broadcasted and consumed distributed.
+// Lowers a `convert_layout` that only relocates data between lanes: the two
+// layouts hand the lanes the same distribution unit, and differ only in which
+// lane holds which part of the value. Every element already exists in some
+// lane, so nothing is recomputed. The scale operands of scaled matrix
+// multiplication need this, being produced broadcasted and consumed
+// distributed.
+//
+// Either layout may replicate. A layout replicates over a set of lanes when a
+// `lane_layout` dimension is sliced away, and equally when its `lane_layout`
+// lays out fewer lanes than the subgroup has: the lane id delinearizes modulo
+// each extent, so the lanes past the product repeat the ones before it. The two
+// spellings are treated alike, and case 4 below is the second one.
+//
+// The two sides are not symmetric in what they mean, though:
+//
+// input supplies the values. Every lane is a potential donor, so what each
+// one holds has to be known -- which it is, for any lane_layout.
+// target states the obligation. Only the lanes it distributes to have to end
+// up holding what it assigns them; the rest are left unspecified.
//
// What it handles
// ---------------
//
-// The three cases below, all on a `vector<8x2>` over 16 lanes, are the three
-// tests. `#xegpu.` prefixes are dropped for width.
+// The cases below are the tests. Cases 1-3 are all on a `vector<8x2>` over 16
+// lanes. `#xegpu.` prefixes are dropped for width.
//
// 1. Broadcast over the whole subgroup, distributed to 8 lanes.
//
@@ -1804,6 +1818,15 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Every lane arrives holding all 16 elements; lanes 0-7 leave holding column
// 0 and lanes 8-15 column 1. Emits 8 extracts and no shuffle.
//
+// 4. Partial input layout, distributed to 2 lanes.
+//
+// layout<lane_layout = [8, 1]> -> layout<lane_layout = [2, 1]>
+//
+// The input names 8 of the 16 lanes, so lanes 8-15 repeat lanes 0-7 and lane
+// `i` arrives holding row `i % 8`. The target hands out 2 fragments of 4
+// rows each, so slot is `lane % 2` and the donors sit 2, 4 and 6 lanes away.
+// Emits 2 extracts and 6 shuffles.
+//
// Nested slices are coalesced before anything else looks at a layout, so only
// the lane_layout of the underlying layout and the union of the sliced dims
// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
@@ -1829,12 +1852,6 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// A lane then holds two elements per distribution unit, which the ownership
// tables do not enumerate.
//
-// * an input layout that covers only part of the subgroup:
-//
-// layout<[8, 1], [1, 1]> -> layout<[1, 8], [1, 1]>
-//
-// Only the target may leave lanes out; the input has to name all 16.
-//
// * elements narrower than a byte:
//
// slice<layout<[1, 1, 16]>, dims = [2]> -> layout<[8, 1]>
@@ -1864,62 +1881,76 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// sliced and spans 8 lanes, so lanes 0-7 are one group and 8-15 another.
//
// slot
-// A lane's index among the lanes the target layout distributes to, `lane %
-// numSlots`. In case 2 `numSlots` is 8, so the lanes pair up: lane 11 has
-// slot 3 like lane 3, and runs the same arithmetic on its own fragment. How
-// many elements a lane holds is a separate matter -- it can own several,
-// which show up as several positions in its fragment, not as several slots.
+// A lane's index within the target's lane period, `lane % lanePeriod`. In
+// case 2 `lanePeriod` is 8, so the lanes pair up: lane 11 has slot 3 like
+// lane 3, and runs the same arithmetic on its own fragment. How many
+// elements a lane holds is a separate matter -- it can own several, which
+// show up as several positions in its fragment, not as several slots.
//
// donor
// The lane a value is read from, `slot + delta`. `gpu.shuffle idx` is the
// only way to move data between lanes and it exchanges one value per lane,
// so every lane runs the same extract; a donor can therefore only serve
-// lanes of its own slot, which makes `delta` a multiple of `numSlots`. In
+// lanes of its own slot, which makes `delta` a multiple of `lanePeriod`. In
// case 2, slot 3 reads (3, 1) from lane 3 + 8.
//
-// The two indices come from opposite sides: `slot` from the target layout,
-// which decides how many lanes the value is distributed to, and `delta`
-// from the input layout, which decides how many copies of it there are.
+// That is also why the donor extracts the right thing without knowing who
+// asked: its own slot is `(slot + delta) % lanePeriod`, which is `slot`, so
+// it computes the same index.
//
// element source
// Where one element of the target fragment comes from: the donor, and the
// index to extract from that donor's fragment. A shuffle is needed exactly
// when the donor cannot be the lane itself. Case 2 has two, one per element
-// of its 1x2 fragment:
+// of its 1x2 fragment. Slot 3 wants (3, 0) and (3, 1):
//
-// element 0 -> delta = 0, index = slot // (3, 0), already local
-// element 1 -> delta = 8, index = slot // (3, 1), from lane 11
+// element 0 -> delta = 0, donor = 3, index = 3 already local
+// element 1 -> delta = 8, donor = 11, index = 3 shuffled
+//
+// The index expression
+// --------------------
//
// The index has to be computable from the lane id, so it is restricted to
//
-// stride * ((slot / slotStride) % extent) + offset
+// index(slot) = stride * ((slot / dimStride) % dimExtent) + offset
+//
+// `(slot / dimStride) % dimExtent` is the lane's coordinate along one dimension
+// of the target lane_layout: `dimExtent` is that dimension's size and
+// `dimStride` how many consecutive slots share a coordinate. `stride` and
+// `offset` place that coordinate in the donor's fragment. An arbitrary index
+// per slot would need a materialized table and a gather; this costs a couple of
+// `arith` ops on the lane id.
//
-// which walks one dimension of the target lane layout: `(slot / slotStride) %
-// extent` is the lane's index along that dimension and `stride`/`offset` place
-// it in the fragment. The three cases give it three shapes:
+// The cases, `p` being the position in the target fragment:
//
-// index = slot consecutive slots take consecutive elements (case 2:
-// slotStride = 1, extent = 8)
-// index = 2 * slot every other element, because the fragment holds whole
-// 2-element rows (case 1)
-// index = slot / 8 eight consecutive slots take the same element, the
-// column their half of the subgroup owns (case 3:
-// slotStride = 8, extent = 2)
+// case index(slot) stride dimStride dimExtent offset
+// 1 2*slot + p 2 1 8 p
+// 2 slot 1 1 8 0
+// 3 slot/8 + 2*p 1 8 2 2*p
+// 4 p % 2 0 1 1 p % 2
//
-// An arbitrary index per slot would need a materialized table and a gather;
-// this costs a couple of `arith` ops on the lane id.
+// Where the constants come from: the `2`s are the innermost extent of the input
+// fragment, which in cases 1 and 3 is the whole 8x2 value, so stepping one row
+// costs 2 elements. In case 2 the fragment is a single column, its row stride
+// is 1, and no `2` appears. Case 4's index does not depend on the slot at all:
+// each target row arrives whole from one donor, so `stride` is 0 and `offset`
+// alone selects the column, with `delta` stepping 0, 2, 4, 6 over the rows.
+//
+// Beware: on a `vector<8x2>` over 16 lanes the lane period (8), case 2's delta
+// (8) and case 3's `dimStride` (8) all have the same value, and they are three
+// different quantities -- a lane count, a lane offset, and an index divisor.
//
// Algorithm
// ---------
//
// tabulate inputOwned[lane] and targetOwned[lane] // element coordinates
// for each element `pos` of the target fragment:
-// for delta = 0, numSlots, 2 * numSlots, ... // smallest first
+// for delta = 0, lanePeriod, 2 * lanePeriod, ... // smallest first
// for each slot:
// needed[slot] = index of the element targetOwned[slot][pos] in the
// fragment of its donor, lane `slot + delta`
// if that donor does not hold the element: give up on this delta
-// look for one stride/offset/slotStride/extent whose index reproduces
+// look for one stride/offset/dimStride/dimExtent whose index reproduces
// every needed[slot]; if there is none: next delta
// accept (delta, index) and stop
// if no delta fits: fail, the layout change is not of this form
@@ -1936,7 +1967,7 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// element and all give the same value, so the smallest `delta` is taken, which
// is what leaves a fully broadcast input shuffle-free.
//
-// Only lanes 0 to `numSlots` - 1 are required to end up holding the elements
+// Only lanes 0 to `lanePeriod` - 1 are required to end up holding the elements
// the target layout assigns their slot; each is the one lane of its slot the
// layout distributes to. Lanes sharing a slot run the same arithmetic on their
// own fragment, so when a shuffle is emitted they all receive the donor's value
@@ -1948,16 +1979,24 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// layout change that is not of this form is reported as a match failure rather
// than lowered incorrectly.
//
-// One caveat: after the rewrite only the first `numSlots` lanes are
+// One caveat: after the rewrite only the first `lanePeriod` lanes are
// guaranteed to hold what the target layout assigns them. The others are
// replicas whose local extract can land on a different element, so dropping a
// shuffle relies on nobody reading them.
//
//===----------------------------------------------------------------------===//
-/// Number of lanes `layout` distributes to. Lanes covered by sliced dimensions
-/// count too: they hold a replicated fragment, not nothing.
-static int64_t getNumDistributedLanes(xegpu::DistributeLayoutAttr layout) {
+/// The lane period of `layout`: the modulus after which its assignment of
+/// elements to lanes repeats. Lane `l` and lane `l + getLanePeriod(layout)`
+/// always hold the same fragment.
+///
+/// This is not the number of lanes holding distinct data. A sliced dimension
+/// contributes its lanes to the period even though they replicate one fragment:
+/// `slice<layout<lane_layout = [8, 1, 2]>, dims = [0]>` has period 16 and hands
+/// out 2 distinct fragments. Nothing here needs the latter count -- the
+/// ownership tables are consulted per lane -- and the period is what the lane
+/// id has to be reduced by.
+static int64_t getLanePeriod(xegpu::DistributeLayoutAttr layout) {
// flatten() coalesces nested slices, so the parent it leaves is always a
// plain LayoutAttr carrying the lane_layout of the whole subgroup.
if (auto sliceAttr = dyn_cast<xegpu::SliceAttr>(layout))
@@ -1967,10 +2006,14 @@ static int64_t getNumDistributedLanes(xegpu::DistributeLayoutAttr layout) {
/// Matches the layout change `redistributeBroadcastedValue` implements: both
/// layouts hand the lanes the same distribution unit (equal effective
-/// lane_data), so only the assignment of elements to lanes changes, and the
-/// input lays out the whole subgroup. A target laying out only part of it
-/// replicates its fragments over the lanes it leaves out, which
-/// `redistributeBroadcastedValue` verifies element-wise.
+/// lane_data), so only the assignment of elements to lanes changes.
+///
+/// Neither layout has to lay out the whole subgroup. A layout whose lane_layout
+/// covers fewer lanes replicates its fragments over the rest, because
+/// `computeStaticDistributedCoords` delinearizes the lane id modulo each
+/// lane_layout extent; the fragment of every lane is therefore defined either
+/// way. `redistributeBroadcastedValue` verifies the assignment element-wise, so
+/// this only has to be permissive enough not to reject what it can lower.
///
/// That the layouts differ is not checked: they are already known to be
/// incompatible here, and the difference may equally be in `order`.
@@ -1982,9 +2025,8 @@ static bool isBroadcastRedistribution(xegpu::DistributeLayoutAttr inputLayout,
if (inputLayout.getEffectiveLaneDataAsInt() !=
targetLayout.getEffectiveLaneDataAsInt())
return false;
- int64_t numSlots = getNumDistributedLanes(targetLayout);
- return getNumDistributedLanes(inputLayout) == subgroupSize && numSlots > 0 &&
- subgroupSize % numSlots == 0;
+ int64_t lanePeriod = getLanePeriod(targetLayout);
+ return lanePeriod > 0 && subgroupSize % lanePeriod == 0;
}
/// Ownership table: `owned[lane][i]` is the coordinate, in the undistributed
@@ -2011,18 +2053,18 @@ struct FragmentIndex {
int64_t stride;
int64_t offset;
/// Slots apart for the index along the selected dimension to advance by one.
- int64_t slotStride;
+ int64_t dimStride;
/// Number of values that index takes.
- int64_t extent;
+ int64_t dimExtent;
/// Compile-time twin of the arithmetic the emitter builds.
int64_t at(int64_t slot) const {
- return stride * ((slot / slotStride) % extent) + offset;
+ return stride * ((slot / dimStride) % dimExtent) + offset;
}
/// Key identifying the value this extracts, for reuse across elements.
std::tuple<int64_t, int64_t, int64_t, int64_t> asTuple() const {
- return {stride, offset, slotStride, extent};
+ return {stride, offset, dimStride, dimExtent};
}
};
@@ -2034,38 +2076,38 @@ struct ElementSource {
bool needsShuffle;
};
-/// Fits `needed`, the fragment index each of the `numSlots` slots has to
+/// Fits `needed`, the fragment index each of the `lanePeriod` slots has to
/// read, to a `FragmentIndex`. Fails when no index reproduces `needed`.
static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
- int64_t numSlots) {
+ int64_t lanePeriod) {
// Slots sharing an index are adjacent, so the first change in `needed` gives
// both how many of them there are and how far apart their elements sit.
int64_t offset = needed[0];
- int64_t slotStride = 1;
- while (slotStride < numSlots && needed[slotStride] == offset)
- slotStride++;
- int64_t stride = slotStride < numSlots ? needed[slotStride] - offset : 0;
- if (stride < 0 || numSlots % slotStride != 0)
+ int64_t dimStride = 1;
+ while (dimStride < lanePeriod && needed[dimStride] == offset)
+ dimStride++;
+ int64_t stride = dimStride < lanePeriod ? needed[dimStride] - offset : 0;
+ if (stride < 0 || lanePeriod % dimStride != 0)
return failure();
// The index wraps once it has taken all its values, so how far `needed` gets
- // from `offset` gives the extent.
- int64_t extent = 1;
+ // from `offset` gives the dimExtent.
+ int64_t dimExtent = 1;
if (stride == 0) {
// Every slot reads `offset`; no dimension is involved.
- slotStride = 1;
+ dimStride = 1;
} else {
for (int64_t index : needed) {
int64_t diff = index - offset;
if (diff < 0 || diff % stride != 0)
return failure();
- extent = std::max(extent, diff / stride + 1);
+ dimExtent = std::max(dimExtent, diff / stride + 1);
}
}
// Derived from the table, so verify against the table before relying on it.
- FragmentIndex index{stride, offset, slotStride, extent};
- if (!llvm::all_of(llvm::seq<int64_t>(0, numSlots), [&](int64_t slot) {
+ FragmentIndex index{stride, offset, dimStride, dimExtent};
+ if (!llvm::all_of(llvm::seq<int64_t>(0, lanePeriod), [&](int64_t slot) {
return needed[slot] == index.at(slot);
}))
return failure();
@@ -2077,7 +2119,7 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
static FailureOr<ElementSource>
deriveElementSource(const OwnedCoords &inputOwned,
const OwnedCoords &targetOwned, int64_t pos,
- int64_t subgroupSize, int64_t numSlots) {
+ int64_t subgroupSize, int64_t lanePeriod) {
// Index of the element with coordinate `coord` in the fragment `lane` holds,
// or -1 if it holds none.
auto findInFragment = [&](int64_t lane, ArrayRef<int64_t> coord) -> int64_t {
@@ -2089,19 +2131,19 @@ deriveElementSource(const OwnedCoords &inputOwned,
// Donor groups, smallest `delta` first.
for (int64_t donorDelta = 0; donorDelta < subgroupSize;
- donorDelta += numSlots) {
+ donorDelta += lanePeriod) {
// Where in its fragment each slot's donor keeps the element it needs.
SmallVector<int64_t> needed;
- for (int64_t slot = 0; slot < numSlots; slot++) {
+ for (int64_t slot = 0; slot < lanePeriod; slot++) {
int64_t index = findInFragment(slot + donorDelta, targetOwned[slot][pos]);
if (index < 0)
break;
needed.push_back(index);
}
- if (static_cast<int64_t>(needed.size()) != numSlots)
+ if (static_cast<int64_t>(needed.size()) != lanePeriod)
continue;
- FailureOr<FragmentIndex> index = fitFragmentIndex(needed, numSlots);
+ FailureOr<FragmentIndex> index = fitFragmentIndex(needed, lanePeriod);
if (failed(index))
continue;
@@ -2127,8 +2169,8 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
return failure();
int64_t srcNumElems = srcTy.getNumElements();
int64_t resNumElems = resTy.getNumElements();
- int64_t numSlots = getNumDistributedLanes(targetLayout);
- if (numSlots < 1 || subgroupSize % numSlots != 0)
+ int64_t lanePeriod = getLanePeriod(targetLayout);
+ if (lanePeriod < 1 || subgroupSize % lanePeriod != 0)
return failure();
// gpu.shuffle only moves the integer widths a lane can hold.
int64_t elemBitWidth = srcTy.getElementTypeBitWidth();
@@ -2142,11 +2184,11 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
if (failed(inputOwned) || failed(targetOwned))
return failure();
// Require the layouts to distribute what the vector types account for, and
- // the lanes past the first `numSlots` to be replicas.
+ // the lanes past the first `lanePeriod` to be replicas.
for (int64_t lane = 0; lane < subgroupSize; lane++) {
if (static_cast<int64_t>((*inputOwned)[lane].size()) != srcNumElems ||
static_cast<int64_t>((*targetOwned)[lane].size()) != resNumElems ||
- (*targetOwned)[lane] != (*targetOwned)[lane % numSlots])
+ (*targetOwned)[lane] != (*targetOwned)[lane % lanePeriod])
return failure();
}
@@ -2155,7 +2197,7 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
SmallVector<ElementSource> sources;
for (int64_t pos = 0; pos < resNumElems; pos++) {
FailureOr<ElementSource> source = deriveElementSource(
- *inputOwned, *targetOwned, pos, subgroupSize, numSlots);
+ *inputOwned, *targetOwned, pos, subgroupSize, lanePeriod);
if (failed(source))
return failure();
sources.push_back(*source);
@@ -2177,10 +2219,10 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
// The lane's slot, which every fragment index is a function of.
Value slot = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
/*upperBound=*/mlir::IntegerAttr());
- if (numSlots != subgroupSize)
+ if (lanePeriod != subgroupSize)
slot = arith::RemUIOp::create(
rewriter, loc, slot,
- arith::ConstantIndexOp::create(rewriter, loc, numSlots));
+ arith::ConstantIndexOp::create(rewriter, loc, lanePeriod));
Type i32Ty = rewriter.getI32Type();
auto flatResTy = VectorType::get({resNumElems}, shuffleTy);
@@ -2193,18 +2235,18 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
// same one. Either step is skipped when it is a no-op over `slot`.
llvm::DenseMap<std::pair<int64_t, int64_t>, Value> dimIndices;
auto getDimIndex = [&](const FragmentIndex &index) -> Value {
- Value &dimIndex = dimIndices[{index.slotStride, index.extent}];
+ Value &dimIndex = dimIndices[{index.dimStride, index.dimExtent}];
if (dimIndex)
return dimIndex;
dimIndex = slot;
- if (index.slotStride != 1)
+ if (index.dimStride != 1)
dimIndex = arith::DivUIOp::create(
rewriter, loc, dimIndex,
- arith::ConstantIndexOp::create(rewriter, loc, index.slotStride));
- if (index.extent < numSlots / index.slotStride)
+ arith::ConstantIndexOp::create(rewriter, loc, index.dimStride));
+ if (index.dimExtent < lanePeriod / index.dimStride)
dimIndex = arith::RemUIOp::create(
rewriter, loc, dimIndex,
- arith::ConstantIndexOp::create(rewriter, loc, index.extent));
+ arith::ConstantIndexOp::create(rewriter, loc, index.dimExtent));
return dimIndex;
};
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 b6bcc6521e6c0..c7e6a6cfe5783 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1717,3 +1717,44 @@ gpu.func @convert_layout_broadcast_nested_slice() {
gpu.return
}
}
+
+// -----
+
+// An input layout that lays out only part of the subgroup. `lane_layout = [8, 1]`
+// names 8 of the 16 lanes, and the lane id delinearizes modulo that, so lanes
+// 8-15 replicate lanes 0-7: the input is a replicated layout spelled without a
+// slice. The target hands out 2 fragments, so `slot` is `lane % 2` and the
+// donors sit 2, 4 and 6 lanes away, one per row of the target fragment beyond
+// the local one.
+// CHECK-LABEL: gpu.func @convert_layout_broadcast_partial_input
+// CHECK: %[[SRC:.*]] = arith.constant dense<1.000000e+00> : vector<1x2xf8E8M0FNU>
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<1x2xf8E8M0FNU> to vector<2xf8E8M0FNU>
+// CHECK: %[[BITS:.*]] = vector.bitcast %[[FLAT]] : vector<2xf8E8M0FNU> to vector<2xi8>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[C2:.*]] = arith.constant 2 : index
+// CHECK: %[[SLOT:.*]] = arith.remui %[[LANE]], %[[C2]] : index
+// CHECK: %[[E0:.*]] = vector.extract %[[BITS]][0] : i8 from vector<2xi8>
+// CHECK: vector.insert %[[E0]], %{{.*}} [0] : i8 into vector<8xi8>
+// CHECK: %[[E1:.*]] = vector.extract %[[BITS]][1] : i8 from vector<2xi8>
+// CHECK: vector.insert %[[E1]], %{{.*}} [1] : i8 into vector<8xi8>
+// CHECK: %[[SLOT_I32:.*]] = arith.index_cast %[[SLOT]] : index to i32
+// CHECK: %[[D2:.*]] = arith.addi %[[SLOT_I32]], %{{.*}} : i32
+// CHECK: gpu.shuffle idx %[[E0]], %[[D2]], %{{.*}} : i8
+// CHECK: vector.insert %{{.*}} [2] : i8 into vector<8xi8>
+// CHECK: gpu.shuffle idx %[[E1]], %{{.*}}, %{{.*}} : i8
+// CHECK: vector.insert %{{.*}} [3] : i8 into vector<8xi8>
+// CHECK-COUNT-4: gpu.shuffle
+// CHECK: %[[BACK:.*]] = vector.bitcast %{{.*}} : vector<8xi8> to vector<8xf8E8M0FNU>
+// CHECK: vector.shape_cast %[[BACK]] : vector<8xf8E8M0FNU> to vector<4x2xf8E8M0FNU>
+gpu.module @xevm_module {
+gpu.func @convert_layout_broadcast_partial_input() {
+ %scale_src = arith.constant dense<1.0> : vector<8x2xf8E8M0FNU>
+ %cvt = xegpu.convert_layout %scale_src
+ <{
+ input_layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [2, 1], lane_data = [1, 1]>
+ }> : vector<8x2xf8E8M0FNU>
+ "some_use"(%cvt) : (vector<8x2xf8E8M0FNU>) -> ()
+ gpu.return
+}
+}
>From 04a4b596c7b2717fc9d1c99ad72e9854d3e8678f Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Tue, 1 Sep 2026 23:57:25 +0000
Subject: [PATCH 07/46] Name the limitations and split the gate by which side
they constrain
Follow-up on review feedback that the capability boundary has to be legible
and the code extensible: the limitations were spread over five functions as
bare failure returns, and the list in the section header was a separate prose
copy that had already drifted from them.
Every failure path now reports a RedistributionLimit, and the declined list is
grouped by what each value constrains -- relational, per layout, target only,
value only -- with each tagged implementation or fundamental. The list and the
enum are therefore the same list, and the match failure names the reason
instead of saying only that the layout change is unsupported.
The gate is split into haveSameDistributionUnit and isSupportedLanePeriod, so
the remaining asymmetry is a named predicate rather than a clause: the target
period has to divide the subgroup size because the donor walk steps by it
across the subgroup, while the input, only ever indexed, has no counterpart.
Two limitations were documented in the wrong place or twice. Non-unit lane_data
was described in the gate as a TODO but enforced in computeOwnedCoords, whose
comment now explains it -- the tables hold one coordinate per distribution unit
while the algorithm indexes elements -- and points at expandBlockCoords for a
lift, noting its enumeration order matches the fragment only when lane_data is
non-unit on the innermost dimension. Sub-byte elements were listed separately
although they are declined by lane_data first in practice.
The target periodicity check is now an assertion. Both computeStaticDistributedCoords
implementations delinearize the lane id modulo the lane_layout extents, whose
product is the lane period, so ownership is periodic by construction and the
check could not fail.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 223 +++++++++++++-----
1 file changed, 170 insertions(+), 53 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 8ccc0956e4d5d..9b2e1e12a0525 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1834,39 +1834,64 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// What it declines
// ----------------
//
-// Each of these reports a match failure rather than lowering. Same notation,
+// Each of these reports a match failure naming a `RedistributionLimit`, so this
+// list and the enum are the same list. Grouped by what the limit constrains,
+// which is where the two layouts are and are not treated alike. Same notation,
// and `vector<8x2>` over 16 lanes unless said otherwise.
//
-// * `lane_data` differing between the two layouts:
+// Relational -- properties of the pair, neither side alone:
+//
+// * LaneDataDiffers, `lane_data` differing between the layouts (fundamental
+// here):
//
// layout<[8, 1], [1, 1]> -> layout<[1, 8], [1, 2]>
//
// A `lane_data`-only change is the repack pattern's job; changing it and
// `lane_layout` together is handled by neither.
//
-// * `lane_data` equal but not unit:
+// * NoDonorDelta, no donor group holding an element through a single extract
+// (fundamental: `gpu.shuffle idx` reads one value per lane).
+//
+// * IndexNotLaneAffine, an index the restricted form cannot express
+// (implementation: an arbitrary one needs a materialized table and a
+// gather), as when `order` transposes it:
+//
+// slice<layout<[1, 1, 16]>, dims = [2]>
+// -> layout<[4, 4], order = [0, 1]> on vector<4x4>
+//
+// Lane `i` would need element `(i % 4) * 4 + i / 4`.
+//
+// Per layout -- the same rule applied to each side:
+//
+// * LaneDataNotUnit, `lane_data` equal but not all ones (implementation):
//
// slice<layout<[1, 1, 16], [1, 2, 1]>, dims = [2]>
// -> layout<[8, 1], [1, 2]>
//
-// A lane then holds two elements per distribution unit, which the ownership
-// tables do not enumerate.
+// The ownership tables hold one coordinate per distribution unit while the
+// algorithm indexes elements; the two coincide only when a unit is a single
+// element. See `computeOwnedCoords` for what lifting it needs. This also
+// covers sub-byte elements in practice, which come with `lane_data` 2 so
+// that a pair fills a byte.
//
-// * elements narrower than a byte:
+// * FragmentSizeMismatch, a layout not distributing what its vector type
+// accounts for (fundamental: the caller derived one from the other).
//
-// slice<layout<[1, 1, 16]>, dims = [2]> -> layout<[8, 1]>
-// on vector<8x2xf4E2M1FN>
+// Target only -- the one asymmetry:
//
-// There is no type for `gpu.shuffle` to carry a 4-bit element in.
+// * LanePeriodNotDivisor, the target's lane period not dividing the subgroup
+// size (fundamental to the donor walk, see `isSupportedLanePeriod`). The
+// input's period is unconstrained because it is only ever indexed.
//
-// * a lane's element index not being one dimension of the target lane layout,
-// as when `order` transposes it:
+// Value only -- neither layout:
//
-// slice<layout<[1, 1, 16]>, dims = [2]>
-// -> layout<[4, 4], order = [0, 1]> on vector<4x4>
+// * SubByteElement, an element `gpu.shuffle` has no type for:
+//
+// slice<layout<[1, 1, 16]>, dims = [2]> -> layout<[8, 1]>
+// on vector<8x2xf4E2M1FN>
//
-// Lane `i` would need element `(i % 4) * 4 + i / 4`, which the restricted
-// index form below cannot express.
+// Reachable only with unit `lane_data`; otherwise LaneDataNotUnit declines
+// first.
//
// Concepts
// --------
@@ -2004,9 +2029,62 @@ static int64_t getLanePeriod(xegpu::DistributeLayoutAttr layout) {
return computeProduct(layout.getEffectiveLaneLayoutAsInt());
}
-/// Matches the layout change `redistributeBroadcastedValue` implements: both
-/// layouts hand the lanes the same distribution unit (equal effective
-/// lane_data), so only the assignment of elements to lanes changes.
+/// Why a `convert_layout` is not the redistribution
+/// `redistributeBroadcastedValue` implements. Every failure path below reports
+/// one of these, so the limitations listed in the section header are exactly
+/// the values here.
+enum class RedistributionLimit {
+ LaneDataDiffers,
+ LaneDataNotUnit,
+ LanePeriodNotDivisor,
+ SubByteElement,
+ FragmentSizeMismatch,
+ NoDonorDelta,
+ IndexNotLaneAffine,
+};
+
+static StringRef describe(RedistributionLimit limit) {
+ switch (limit) {
+ case RedistributionLimit::LaneDataDiffers:
+ return "input and target lane_data differ";
+ case RedistributionLimit::LaneDataNotUnit:
+ return "lane_data is not all ones";
+ case RedistributionLimit::LanePeriodNotDivisor:
+ return "the target lane period does not divide the subgroup size";
+ case RedistributionLimit::SubByteElement:
+ return "no gpu.shuffle type carries the element";
+ case RedistributionLimit::FragmentSizeMismatch:
+ return "a layout does not distribute what its vector type accounts for";
+ case RedistributionLimit::NoDonorDelta:
+ return "no donor group holds an element through a single extract";
+ case RedistributionLimit::IndexNotLaneAffine:
+ return "the fragment index is not an affine function of the lane id";
+ }
+ llvm_unreachable("unhandled RedistributionLimit");
+}
+
+/// Relational rule: both layouts have to hand the lanes the same distribution
+/// unit, so that only the assignment of elements to lanes changes.
+static bool haveSameDistributionUnit(xegpu::DistributeLayoutAttr inputLayout,
+ xegpu::DistributeLayoutAttr targetLayout) {
+ return inputLayout.getEffectiveLaneDataAsInt() ==
+ targetLayout.getEffectiveLaneDataAsInt();
+}
+
+/// Target-side rule: `deriveElementSource` walks the donor groups by stepping
+/// `delta` in multiples of the lane period while `delta < subgroupSize`, and a
+/// slot is smaller than the period. The period therefore has to divide the
+/// subgroup size, or `slot + delta` would name a lane past the last one.
+///
+/// The input has no counterpart: it is only ever indexed, never used to derive
+/// a delta, so its own period is unconstrained.
+static bool isSupportedLanePeriod(xegpu::DistributeLayoutAttr targetLayout,
+ int64_t subgroupSize) {
+ int64_t lanePeriod = getLanePeriod(targetLayout);
+ return lanePeriod > 0 && subgroupSize % lanePeriod == 0;
+}
+
+/// Matches the layout change `redistributeBroadcastedValue` implements.
///
/// Neither layout has to lay out the whole subgroup. A layout whose lane_layout
/// covers fewer lanes replicates its fragments over the rest, because
@@ -2017,16 +2095,20 @@ static int64_t getLanePeriod(xegpu::DistributeLayoutAttr layout) {
///
/// That the layouts differ is not checked: they are already known to be
/// incompatible here, and the difference may equally be in `order`.
-static bool isBroadcastRedistribution(xegpu::DistributeLayoutAttr inputLayout,
- xegpu::DistributeLayoutAttr targetLayout,
- int64_t subgroupSize) {
- // TODO: computeOwnedCoords only tabulates unit lane_data, so equal but
- // non-unit lane_data is rejected further down rather than here.
- if (inputLayout.getEffectiveLaneDataAsInt() !=
- targetLayout.getEffectiveLaneDataAsInt())
+static bool
+isBroadcastRedistribution(xegpu::DistributeLayoutAttr inputLayout,
+ xegpu::DistributeLayoutAttr targetLayout,
+ int64_t subgroupSize,
+ std::optional<RedistributionLimit> &limit) {
+ if (!haveSameDistributionUnit(inputLayout, targetLayout)) {
+ limit = RedistributionLimit::LaneDataDiffers;
return false;
- int64_t lanePeriod = getLanePeriod(targetLayout);
- return lanePeriod > 0 && subgroupSize % lanePeriod == 0;
+ }
+ if (!isSupportedLanePeriod(targetLayout, subgroupSize)) {
+ limit = RedistributionLimit::LanePeriodNotDivisor;
+ return false;
+ }
+ return true;
}
/// Ownership table: `owned[lane][i]` is the coordinate, in the undistributed
@@ -2034,14 +2116,25 @@ static bool isBroadcastRedistribution(xegpu::DistributeLayoutAttr inputLayout,
using OwnedCoords = SmallVector<SmallVector<SmallVector<int64_t>>>;
/// Tabulates the fragment every lane of the subgroup owns of `shape` under
-/// `layout`. Fails for non-unit lane_data, where a lane's elements are not a
-/// plain enumeration of the distribution unit starts.
+/// `layout`: `owned[lane][i]` is the coordinate of element `i` of that
+/// fragment, numbered as the flattened fragment `vector.extract` indexes.
+///
+/// Fails for non-unit lane_data. `computeStaticDistributedCoords` yields one
+/// coordinate per distribution unit -- the block start -- and those coincide
+/// with the elements only when a unit is a single element. Lifting this needs
+/// the block starts expanded in fragment order; `expandBlockCoords` in
+/// XeGPUDialect.cpp does such an expansion for layout comparison, but emits a
+/// whole unit before moving to the next, which matches the fragment only when
+/// lane_data is non-unit on the innermost dimension alone.
static FailureOr<OwnedCoords>
computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
- int64_t subgroupSize) {
+ int64_t subgroupSize,
+ std::optional<RedistributionLimit> &limit) {
SmallVector<int64_t> laneData = layout.getEffectiveLaneDataAsInt();
- if (!llvm::all_of(laneData, [](int64_t d) { return d == 1; }))
+ if (!llvm::all_of(laneData, [](int64_t d) { return d == 1; })) {
+ limit = RedistributionLimit::LaneDataNotUnit;
return failure();
+ }
OwnedCoords owned;
for (int64_t lane = 0; lane < subgroupSize; lane++)
owned.push_back(layout.computeStaticDistributedCoords(lane, shape));
@@ -2115,11 +2208,13 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
}
/// Works out where element `pos` of the target fragment comes from. Fails when
-/// no donor group provides it through a single extract.
+/// no donor group provides it through a single extract whose index is affine in
+/// the lane id.
static FailureOr<ElementSource>
deriveElementSource(const OwnedCoords &inputOwned,
const OwnedCoords &targetOwned, int64_t pos,
- int64_t subgroupSize, int64_t lanePeriod) {
+ int64_t subgroupSize, int64_t lanePeriod,
+ std::optional<RedistributionLimit> &limit) {
// Index of the element with coordinate `coord` in the fragment `lane` holds,
// or -1 if it holds none.
auto findInFragment = [&](int64_t lane, ArrayRef<int64_t> coord) -> int64_t {
@@ -2129,7 +2224,10 @@ deriveElementSource(const OwnedCoords &inputOwned,
return -1;
};
- // Donor groups, smallest `delta` first.
+ // Donor groups, smallest `delta` first. `sawCompleteDonor` records whether
+ // any group held the element at all, which is what distinguishes the two ways
+ // this fails.
+ bool sawCompleteDonor = false;
for (int64_t donorDelta = 0; donorDelta < subgroupSize;
donorDelta += lanePeriod) {
// Where in its fragment each slot's donor keeps the element it needs.
@@ -2142,6 +2240,7 @@ deriveElementSource(const OwnedCoords &inputOwned,
}
if (static_cast<int64_t>(needed.size()) != lanePeriod)
continue;
+ sawCompleteDonor = true;
FailureOr<FragmentIndex> index = fitFragmentIndex(needed, lanePeriod);
if (failed(index))
@@ -2152,44 +2251,55 @@ deriveElementSource(const OwnedCoords &inputOwned,
// extract is local by construction.
return ElementSource{*index, donorDelta, /*needsShuffle=*/donorDelta != 0};
}
+ limit = sawCompleteDonor ? RedistributionLimit::IndexNotLaneAffine
+ : RedistributionLimit::NoDonorDelta;
return failure();
}
/// Redistributes `src`, the fragment the lane holds under `inputLayout`, into
/// the `resTy` fragment it owns under `targetLayout`. Fails if the layout
/// change is not of the form described above.
-static FailureOr<Value>
-redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
- Value src, VectorType resTy,
- xegpu::DistributeLayoutAttr inputLayout,
- xegpu::DistributeLayoutAttr targetLayout,
- ArrayRef<int64_t> shape, int64_t subgroupSize) {
+static FailureOr<Value> redistributeBroadcastedValue(
+ ConversionPatternRewriter &rewriter, Location loc, Value src,
+ VectorType resTy, xegpu::DistributeLayoutAttr inputLayout,
+ xegpu::DistributeLayoutAttr targetLayout, ArrayRef<int64_t> shape,
+ int64_t subgroupSize, std::optional<RedistributionLimit> &limit) {
auto srcTy = dyn_cast<VectorType>(src.getType());
if (!srcTy)
return failure();
int64_t srcNumElems = srcTy.getNumElements();
int64_t resNumElems = resTy.getNumElements();
int64_t lanePeriod = getLanePeriod(targetLayout);
- if (lanePeriod < 1 || subgroupSize % lanePeriod != 0)
+ if (!isSupportedLanePeriod(targetLayout, subgroupSize)) {
+ limit = RedistributionLimit::LanePeriodNotDivisor;
return failure();
+ }
// gpu.shuffle only moves the integer widths a lane can hold.
int64_t elemBitWidth = srcTy.getElementTypeBitWidth();
if (!llvm::isPowerOf2_64(elemBitWidth) || elemBitWidth < 8 ||
- elemBitWidth > 64)
+ elemBitWidth > 64) {
+ limit = RedistributionLimit::SubByteElement;
return failure();
+ }
// Tabulate ownership: what each lane brings and what each slot needs.
- auto inputOwned = computeOwnedCoords(inputLayout, shape, subgroupSize);
- auto targetOwned = computeOwnedCoords(targetLayout, shape, subgroupSize);
+ auto inputOwned = computeOwnedCoords(inputLayout, shape, subgroupSize, limit);
+ auto targetOwned =
+ computeOwnedCoords(targetLayout, shape, subgroupSize, limit);
if (failed(inputOwned) || failed(targetOwned))
return failure();
- // Require the layouts to distribute what the vector types account for, and
- // the lanes past the first `lanePeriod` to be replicas.
+ // Require the layouts to distribute what the vector types account for. The
+ // lanes past the first `lanePeriod` are replicas by construction: both
+ // computeStaticDistributedCoords implementations delinearize the lane id
+ // modulo the lane_layout extents, whose product is `lanePeriod`.
for (int64_t lane = 0; lane < subgroupSize; lane++) {
+ assert((*targetOwned)[lane] == (*targetOwned)[lane % lanePeriod] &&
+ "target ownership is not periodic in the lane period");
if (static_cast<int64_t>((*inputOwned)[lane].size()) != srcNumElems ||
- static_cast<int64_t>((*targetOwned)[lane].size()) != resNumElems ||
- (*targetOwned)[lane] != (*targetOwned)[lane % lanePeriod])
+ static_cast<int64_t>((*targetOwned)[lane].size()) != resNumElems) {
+ limit = RedistributionLimit::FragmentSizeMismatch;
return failure();
+ }
}
// Work out where every element of the result comes from before emitting
@@ -2197,7 +2307,7 @@ redistributeBroadcastedValue(ConversionPatternRewriter &rewriter, Location loc,
SmallVector<ElementSource> sources;
for (int64_t pos = 0; pos < resNumElems; pos++) {
FailureOr<ElementSource> source = deriveElementSource(
- *inputOwned, *targetOwned, pos, subgroupSize, lanePeriod);
+ *inputOwned, *targetOwned, pos, subgroupSize, lanePeriod, limit);
if (failed(source))
return failure();
sources.push_back(*source);
@@ -2409,18 +2519,20 @@ struct SgToLaneConvertLayout
}
}
- // Input layout replicates over broadcast groups of lanes while the target
- // distributes, so data has to move across lanes.
+ // The two layouts assign the same distribution unit to different lanes, so
+ // data has to move across lanes.
const auto *uArch =
xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
FailureOr<VectorType> resDistTy = xegpu::getDistVecTypeBasedOnLaneLayout(
targetLayout, cast<VectorType>(valType));
+ std::optional<RedistributionLimit> limit;
if (uArch && succeeded(resDistTy)) {
int64_t subgroupSize = uArch->getSubgroupSize();
- if (isBroadcastRedistribution(inputLayout, targetLayout, subgroupSize)) {
+ if (isBroadcastRedistribution(inputLayout, targetLayout, subgroupSize,
+ limit)) {
FailureOr<Value> res = redistributeBroadcastedValue(
rewriter, op.getLoc(), adaptor.getSource(), *resDistTy, inputLayout,
- targetLayout, resShapeVec, subgroupSize);
+ targetLayout, resShapeVec, subgroupSize, limit);
if (succeeded(res)) {
rewriter.replaceOp(op, *res);
return success();
@@ -2428,6 +2540,11 @@ struct SgToLaneConvertLayout
}
}
+ if (limit)
+ return rewriter.notifyMatchFailure(
+ op, Twine("convert_layout is not a redistribution this pattern "
+ "lowers: ") +
+ describe(*limit));
return rewriter.notifyMatchFailure(
op, "lowering incompatible convert_layout not yet supported");
}
>From 9ec51282d8b43cc686ee3ee4d3e8f5773aa119ab Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 00:08:07 +0000
Subject: [PATCH 08/46] Make each declaration readable without the section
header
Review feedback was that terms like delta were hard to follow, and the check
suggested was to read each comment in isolation, which is how the code gets
reviewed. Doing that, the header was fine and the declarations were not: the
whole vocabulary lived only in the header, three hundred lines from the
functions using it.
Each declaration now grounds the terms it uses, with a value from the running
example where that helps.
getLanePeriod says why the assignment repeats -- the lane id is
delinearized modulo each lane_layout extent, so nothing
can depend on it beyond their product -- rather than
asserting that it does. "Period" was a new term
implying something cyclical without saying what.
isSupportedLanePeriod restates the donor search locally instead of forward
referencing deriveElementSource, and defines donor.
FragmentIndex documents all four fields, stride and offset having had
none, and gives the expression a meaning and an
instance next to at().
ElementSource defines donor and both fields.
haveSameDistributionUnit says a distribution unit is the lane_data-sized
block.
fitFragmentIndex, deriveElementSource a clause each for slot and donor.
The header called it delta while the field is donorDelta, so grepping the term
found a definition under another name. It is donorDelta throughout now.
Adds a fifth case, a target handing out 4 fragments rather than 8. On the
layout<[8, 1]> targets the lane period, case 2 donorDelta and case 3 dimStride
are all 8 while being a lane count, a lane offset and an index divisor, so the
index expression could not be read off them. Case 5 separates the four
quantities and the table leads with it. Cases 1 to 4 stay, being what the
dpas_mx scale operands actually produce, and keep their tests.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 257 +++++++++++++-----
.../XeGPU/sg-to-lane-distribute-unit.mlir | 41 +++
2 files changed, 229 insertions(+), 69 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 9b2e1e12a0525..241a3a2355fc2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1774,9 +1774,10 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// Either layout may replicate. A layout replicates over a set of lanes when a
// `lane_layout` dimension is sliced away, and equally when its `lane_layout`
-// lays out fewer lanes than the subgroup has: the lane id delinearizes modulo
-// each extent, so the lanes past the product repeat the ones before it. The two
-// spellings are treated alike, and case 4 below is the second one.
+// lays out fewer lanes than the subgroup has: the lane id is delinearized
+// modulo each `lane_layout` extent, so the lanes past their product repeat the
+// ones before it. The two spellings are treated alike, and case 4 below is the
+// second one.
//
// The two sides are not symmetric in what they mean, though:
//
@@ -1827,10 +1828,67 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// rows each, so slot is `lane % 2` and the donors sit 2, 4 and 6 lanes away.
// Emits 2 extracts and 6 shuffles.
//
+// 5. Broadcast over the whole subgroup, distributed to 4 lanes.
+//
+// slice<layout<lane_layout = [1, 1, 16]>, dims = [2]>
+// -> layout<lane_layout = [4, 1]>
+//
+// As case 1, but the target hands out 4 fragments of 2 rows instead of 8 of
+// one. Kept because it is the one case where the quantities below take
+// distinct values, so it is the clearest to read the index expression off.
+// Emits 4 extracts and no shuffle.
+//
// Nested slices are coalesced before anything else looks at a layout, so only
// the lane_layout of the underlying layout and the union of the sliced dims
// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
//
+// What it costs
+// -------------
+//
+// There is one element source per element of the target fragment, so both
+// counts scale with how much the target gives one lane, not with the size of
+// the value. The two are not counted the same way, which is why they can differ
+// so widely:
+//
+// extracts one per *distinct* fragment index among the sources. Elements
+// reading the same index share one, which is why case 4 needs only
+// 2 for its 8 elements: its index does not depend on the slot, so
+// the two columns are the only distinct reads.
+//
+// shuffles one per element the reading lane does not already hold, that is
+// per source whose `donorDelta` is not 0. These are not shared,
+// even between elements taken from the same donor; see the TODO in
+// the emitter, `gpu.shuffle` moving a whole 32-bit lane word.
+//
+// So a fully broadcast input never shuffles: every lane holds everything, so
+// the smallest `donorDelta` that works is always 0, as in cases 1, 3 and 5.
+// Case 2 shuffles once, for the column its own group does not hold. Case 4
+// shuffles six times, for the three rows beyond its local one, both columns
+// each.
+//
+// How close that is to optimal
+// ----------------------------
+//
+// Extracts are minimal: they are register reads and equal indices are shared.
+//
+// Shuffles are not. `gpu.shuffle idx` reads one donor lane and moves a whole
+// 32-bit lane word, so within this scheme -- every lane reading one donor per
+// instruction -- the floor is one shuffle per distinct nonzero `donorDelta`,
+// carrying up to 32 bits of that group's elements at once. Emitting one per
+// remote element instead costs a factor of however many elements a lane takes
+// from the same group: nothing where that is one, as in case 2, and 2x in case
+// 4, whose 6 shuffles move 3 groups of two bytes. Closing that is the TODO in
+// the emitter; `shuffleDataAsLaneLayoutChange` in this file already packs
+// elements into lane words that way.
+//
+// A lower floor exists, but only outside this scheme. If lanes forwarded data
+// in several hops, a lane needing `k` remote elements of `w` bits could receive
+// them in `ceil(k * w / 32)` shuffles however many lanes they start on -- 2
+// rather than 3 for case 4. That trades instruction count for dependent latency
+// and a routing schedule. For the scale operands this lowers, a lane's remote
+// elements come from one or two donors, so the single-hop floor is at or near
+// the absolute one and the simpler scheme costs nothing.
+//
// What it declines
// ----------------
//
@@ -1849,7 +1907,7 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// A `lane_data`-only change is the repack pattern's job; changing it and
// `lane_layout` together is handled by neither.
//
-// * NoDonorDelta, no donor group holding an element through a single extract
+// * NoDonorDelta, no donor group holding the element for every slot
// (fundamental: `gpu.shuffle idx` reads one value per lane).
//
// * IndexNotLaneAffine, an index the restricted form cannot express
@@ -1913,15 +1971,21 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// show up as several positions in its fragment, not as several slots.
//
// donor
-// The lane a value is read from, `slot + delta`. `gpu.shuffle idx` is the
-// only way to move data between lanes and it exchanges one value per lane,
-// so every lane runs the same extract; a donor can therefore only serve
-// lanes of its own slot, which makes `delta` a multiple of `lanePeriod`. In
-// case 2, slot 3 reads (3, 1) from lane 3 + 8.
+// The lane a value is read from, `slot + donorDelta`. `gpu.shuffle idx` is
+// the only way to move data between lanes and it exchanges one value per
+// lane, so every lane runs the same extract; a donor can therefore only
+// serve lanes of its own slot, which makes `donorDelta` a multiple of
+// `lanePeriod`. In case 2, slot 3 reads (3, 1) from lane 3 + 8.
//
// That is also why the donor extracts the right thing without knowing who
-// asked: its own slot is `(slot + delta) % lanePeriod`, which is `slot`, so
-// it computes the same index.
+// asked: its own slot is `(slot + donorDelta) % lanePeriod`, which is
+// `slot`, so it computes the same index.
+//
+// donor group
+// The lanes one `donorDelta` names, one per slot: `donorDelta` through
+// `donorDelta + lanePeriod - 1`. Candidate donors are considered a group at
+// a time, since every slot has to be served by the same `donorDelta`. In
+// case 2 there are two, lanes 0-7 and lanes 8-15.
//
// element source
// Where one element of the target fragment comes from: the donor, and the
@@ -1929,8 +1993,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// when the donor cannot be the lane itself. Case 2 has two, one per element
// of its 1x2 fragment. Slot 3 wants (3, 0) and (3, 1):
//
-// element 0 -> delta = 0, donor = 3, index = 3 already local
-// element 1 -> delta = 8, donor = 11, index = 3 shuffled
+// element 0 -> donorDelta = 0, donor = 3, index = 3 already local
+// element 1 -> donorDelta = 8, donor = 11, index = 3 shuffled
//
// The index expression
// --------------------
@@ -1946,51 +2010,57 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// per slot would need a materialized table and a gather; this costs a couple of
// `arith` ops on the lane id.
//
-// The cases, `p` being the position in the target fragment:
-//
-// case index(slot) stride dimStride dimExtent offset
-// 1 2*slot + p 2 1 8 p
-// 2 slot 1 1 8 0
-// 3 slot/8 + 2*p 1 8 2 2*p
-// 4 p % 2 0 1 1 p % 2
-//
-// Where the constants come from: the `2`s are the innermost extent of the input
-// fragment, which in cases 1 and 3 is the whole 8x2 value, so stepping one row
-// costs 2 elements. In case 2 the fragment is a single column, its row stride
-// is 1, and no `2` appears. Case 4's index does not depend on the slot at all:
-// each target row arrives whole from one donor, so `stride` is 0 and `offset`
-// alone selects the column, with `delta` stepping 0, 2, 4, 6 over the rows.
-//
-// Beware: on a `vector<8x2>` over 16 lanes the lane period (8), case 2's delta
-// (8) and case 3's `dimStride` (8) all have the same value, and they are three
-// different quantities -- a lane count, a lane offset, and an index divisor.
+// The cases, `p` being the position in the target fragment. Case 5 first, being
+// the one whose numbers are all distinct:
+//
+// case lanePeriod index(slot) stride dimStride dimExtent offset
+// 5 4 2*slot + off 2 1 4 0,1,8,9
+// 1 8 2*slot + p 2 1 8 p
+// 2 8 slot 1 1 8 0
+// 3 16 slot/8 + 2*p 1 8 2 2*p
+// 4 2 p % 2 0 1 1 p % 2
+//
+// Where the constants come from: the `2`s multiplying `slot` are the size of
+// the input fragment's innermost dimension, which in cases 1, 3 and 5 is the
+// whole 8x2 value, so stepping one row costs 2 elements. In case 2 the fragment
+// is a single column, its row stride is 1, and no `2` appears. Case 4's index
+// does not depend on the slot at all: each target row arrives whole from one
+// donor, so `stride` is 0 and `offset` alone selects the column, with
+// `donorDelta` stepping 0, 2, 4, 6 over the rows.
+//
+// Read the expression off case 5, where the four quantities are 4, 2, 1 and 4.
+// On the `layout<[8, 1]>` targets of cases 1 and 2 they collide: the lane
+// period (8), case 2's `donorDelta` (8) and case 3's `dimStride` (8) all have
+// the same value while being three different things -- a lane count, a lane
+// offset and an index divisor. Those cases are kept because they are what the
+// scale operands of `dpas_mx` actually produce.
//
// Algorithm
// ---------
//
// tabulate inputOwned[lane] and targetOwned[lane] // element coordinates
// for each element `pos` of the target fragment:
-// for delta = 0, lanePeriod, 2 * lanePeriod, ... // smallest first
+// for donorDelta = 0, lanePeriod, 2 * lanePeriod, ... // smallest first
// for each slot:
// needed[slot] = index of the element targetOwned[slot][pos] in the
-// fragment of its donor, lane `slot + delta`
-// if that donor does not hold the element: give up on this delta
+// fragment of its donor, lane `slot + donorDelta`
+// if that donor does not hold the element: give up on this donorDelta
// look for one stride/offset/dimStride/dimExtent whose index reproduces
-// every needed[slot]; if there is none: next delta
-// accept (delta, index) and stop
-// if no delta fits: fail, the layout change is not of this form
+// every needed[slot]; if there is none: next donorDelta
+// accept (donorDelta, index) and stop
+// if no donorDelta fits: fail, the layout change is not of this form
// for each element source:
// %idx = index(slot) // arith on gpu.lane_id
// %val = vector.extract %fragment[%idx]
-// %donor = slot + delta
+// %donor = slot + donorDelta
// %val = gpu.shuffle idx %val, %donor // unless %donor is %lane
// %res = vector.insert %val, %res[pos]
//
-// `%donor` is the lane itself when `delta` is 0 and the extract already lands
-// on an element the lane owns, and that is the only case where the shuffle is
-// dropped. Replication makes it a choice: several donor groups may hold the
-// element and all give the same value, so the smallest `delta` is taken, which
-// is what leaves a fully broadcast input shuffle-free.
+// `%donor` is the lane itself when `donorDelta` is 0 and the extract already
+// lands on an element the lane owns, and that is the only case where the
+// shuffle is dropped. Replication makes it a choice: several donor groups may
+// hold the element and all give the same value, so the smallest `donorDelta` is
+// taken, which is what leaves a fully broadcast input shuffle-free.
//
// Only lanes 0 to `lanePeriod` - 1 are required to end up holding the elements
// the target layout assigns their slot; each is the one lane of its slot the
@@ -2012,8 +2082,14 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//===----------------------------------------------------------------------===//
/// The lane period of `layout`: the modulus after which its assignment of
-/// elements to lanes repeats. Lane `l` and lane `l + getLanePeriod(layout)`
-/// always hold the same fragment.
+/// elements to lanes repeats, so that lane `l` and lane `l + getLanePeriod()`
+/// hold the same *fragment* -- the set of elements one lane owns.
+///
+/// It repeats because `computeStaticDistributedCoords` delinearizes the lane id
+/// into digits, one per lane_layout dimension, taking each modulo that
+/// dimension's extent. Nothing a layout derives from the lane id can therefore
+/// depend on it beyond the product of those extents, which is what this
+/// returns. A subgroup larger than the product simply repeats.
///
/// This is not the number of lanes holding distinct data. A sliced dimension
/// contributes its lanes to the period even though they replicate one fragment:
@@ -2056,7 +2132,7 @@ static StringRef describe(RedistributionLimit limit) {
case RedistributionLimit::FragmentSizeMismatch:
return "a layout does not distribute what its vector type accounts for";
case RedistributionLimit::NoDonorDelta:
- return "no donor group holds an element through a single extract";
+ return "no donor lane holds the element for every slot";
case RedistributionLimit::IndexNotLaneAffine:
return "the fragment index is not an affine function of the lane id";
}
@@ -2064,20 +2140,28 @@ static StringRef describe(RedistributionLimit limit) {
}
/// Relational rule: both layouts have to hand the lanes the same distribution
-/// unit, so that only the assignment of elements to lanes changes.
+/// unit -- the `lane_data`-sized block of elements a lane owns as a whole -- so
+/// that only the assignment of those blocks to lanes changes.
static bool haveSameDistributionUnit(xegpu::DistributeLayoutAttr inputLayout,
xegpu::DistributeLayoutAttr targetLayout) {
return inputLayout.getEffectiveLaneDataAsInt() ==
targetLayout.getEffectiveLaneDataAsInt();
}
-/// Target-side rule: `deriveElementSource` walks the donor groups by stepping
-/// `delta` in multiples of the lane period while `delta < subgroupSize`, and a
-/// slot is smaller than the period. The period therefore has to divide the
-/// subgroup size, or `slot + delta` would name a lane past the last one.
+/// Target-side rule: the target's lane period has to divide the subgroup size.
+///
+/// A value is moved between lanes by having every lane extract one element and
+/// `gpu.shuffle idx` it from a chosen *donor* lane. All lanes run the same
+/// code, so a lane can only be served by a donor that agrees with it on which
+/// element to extract -- one whose lane id differs by a whole number of lane
+/// periods. The lowering therefore searches donors at `donorDelta = 0,
+/// lanePeriod, 2 * lanePeriod, ...` up to the subgroup size, reading lane `slot
+/// + donorDelta`, where `slot` is `lane % lanePeriod` and so below the period.
+/// If the period did not divide the subgroup size that sum could name a lane
+/// past the last one.
///
/// The input has no counterpart: it is only ever indexed, never used to derive
-/// a delta, so its own period is unconstrained.
+/// a donor, so its own period is unconstrained.
static bool isSupportedLanePeriod(xegpu::DistributeLayoutAttr targetLayout,
int64_t subgroupSize) {
int64_t lanePeriod = getLanePeriod(targetLayout);
@@ -2141,13 +2225,33 @@ computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
return owned;
}
-/// Index into the lane's fragment as a function of its slot.
+/// Which element of the donor's fragment a lane extracts -- the donor being the
+/// lane it reads from, `slot + donorDelta` -- as a function of its `slot`
+/// (`lane % lanePeriod`).
+///
+/// `at()` is restricted to one shape: the lane's coordinate along a single
+/// dimension of the target lane_layout, scaled and offset into the fragment.
+/// That keeps the index a couple of `arith` ops on the lane id; an arbitrary
+/// index per slot would need a materialized table and a gather. A `stride` of 0
+/// is the degenerate case where no dimension is walked at all and every slot
+/// reads `offset`.
+///
+/// For a `vector<8x2>` whose input layout gives every lane all 16 elements,
+/// row-major, and whose target gives lane `i` row `i`, element `p` of the
+/// target fragment sits at `2 * slot + p`: `stride` 2, `offset` p, `dimStride`
+/// 1, `dimExtent` 8.
struct FragmentIndex {
+ /// Distance in the fragment between consecutive coordinates along the
+ /// dimension being walked, or 0 when every slot reads the same element.
int64_t stride;
+ /// Fragment index that coordinate 0 maps to, and the whole index when
+ /// `stride` is 0.
int64_t offset;
- /// Slots apart for the index along the selected dimension to advance by one.
+ /// Slots apart for the coordinate to advance by one.
int64_t dimStride;
- /// Number of values that index takes.
+ /// How many values the coordinate takes before wrapping. This is the extent
+ /// of the target `lane_layout` dimension being walked, except when `stride`
+ /// is 0 and no dimension is: then it is 1, as is `dimStride`.
int64_t dimExtent;
/// Compile-time twin of the arithmetic the emitter builds.
@@ -2161,16 +2265,23 @@ struct FragmentIndex {
}
};
-/// Where one element of the target fragment comes from: which donor, and where
-/// in that donor's fragment.
+/// Where one element of the target fragment comes from.
struct ElementSource {
+ /// Which element of the donor's fragment to extract.
FragmentIndex index;
+ /// Lanes from the reader to the donor it reads through `gpu.shuffle idx`: the
+ /// donor is lane `slot + donorDelta`. Always a multiple of the lane period,
+ /// so that donor and reader agree on the index to extract.
int64_t donorDelta;
+ /// False when the donor is the reader itself, which is exactly `donorDelta`
+ /// being 0, and the extract is already local.
bool needsShuffle;
};
-/// Fits `needed`, the fragment index each of the `lanePeriod` slots has to
-/// read, to a `FragmentIndex`. Fails when no index reproduces `needed`.
+/// Fits `needed[slot]`, the fragment index each slot has to read, to the
+/// restricted form `FragmentIndex` can emit. A slot is a lane reduced modulo
+/// the lane period, so there are `lanePeriod` of them. Fails when no index in
+/// that form reproduces the whole table.
static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
int64_t lanePeriod) {
// Slots sharing an index are adjacent, so the first change in `needed` gives
@@ -2207,9 +2318,15 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
return index;
}
-/// Works out where element `pos` of the target fragment comes from. Fails when
-/// no donor group provides it through a single extract whose index is affine in
-/// the lane id.
+/// Works out where element `pos` of the target fragment comes from: which donor
+/// lane, at `slot + donorDelta` for some multiple of the lane period, and which
+/// element of that donor's fragment. Candidates are tried a donor group at a
+/// time -- the lanes one `donorDelta` names, one per slot -- because every slot
+/// has to be served by the same `donorDelta`. The smallest is preferred, so a
+/// lane that already holds the element needs no shuffle.
+///
+/// Fails when no donor group holds the element for every slot, or when the
+/// indices they would be read at are not of `FragmentIndex`'s form.
static FailureOr<ElementSource>
deriveElementSource(const OwnedCoords &inputOwned,
const OwnedCoords &targetOwned, int64_t pos,
@@ -2224,9 +2341,9 @@ deriveElementSource(const OwnedCoords &inputOwned,
return -1;
};
- // Donor groups, smallest `delta` first. `sawCompleteDonor` records whether
- // any group held the element at all, which is what distinguishes the two ways
- // this fails.
+ // Donor groups, smallest `donorDelta` first. `sawCompleteDonor` records
+ // whether any group held the element at all, which is what distinguishes the
+ // two ways this fails.
bool sawCompleteDonor = false;
for (int64_t donorDelta = 0; donorDelta < subgroupSize;
donorDelta += lanePeriod) {
@@ -2246,9 +2363,9 @@ deriveElementSource(const OwnedCoords &inputOwned,
if (failed(index))
continue;
- // The donor is the lane itself exactly when `delta` is 0: `needed[slot]`
- // was located in the fragment of lane `slot + delta`, so at 0 every slot's
- // extract is local by construction.
+ // The donor is the lane itself exactly when `donorDelta` is 0:
+ // `needed[slot]` was located in the fragment of lane `slot + donorDelta`,
+ // so at 0 every slot's extract is local by construction.
return ElementSource{*index, donorDelta, /*needsShuffle=*/donorDelta != 0};
}
limit = sawCompleteDonor ? RedistributionLimit::IndexNotLaneAffine
@@ -2377,7 +2494,9 @@ static FailureOr<Value> redistributeBroadcastedValue(
return value;
};
- // Extract, shuffle and insert one element of the result at a time.
+ // Extract, shuffle and insert one element of the result at a time. One
+ // extract per distinct fragment index, one shuffle per element the lane does
+ // not already hold; see "What it costs" above.
// TODO: one shuffle per element, though `gpu.shuffle` moves a whole 32-bit
// lane word. Elements sharing a donor could travel in one, as
// shuffleDataAsLaneLayoutChange does.
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 c7e6a6cfe5783..0163a8e1db926 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1758,3 +1758,44 @@ gpu.func @convert_layout_broadcast_partial_input() {
gpu.return
}
}
+
+// -----
+
+// A target handing out 4 fragments rather than 8, which separates the constants
+// the index expression is built from: the lane period is 4 while the input
+// fragment's row stride is 2. Every lane arrives with all 16 elements, so
+// element `p` of the 2x2 target fragment is a local extract at `2*slot + off`
+// with `off` running 0, 1, 8, 9 over the two rows the fragment covers.
+// CHECK-LABEL: gpu.func @convert_layout_broadcast_four_slots
+// 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: %[[C4:.*]] = arith.constant 4 : index
+// CHECK: %[[SLOT:.*]] = arith.remui %[[LANE]], %[[C4]] : index
+// CHECK: %[[M0:.*]] = arith.muli %[[SLOT]], %{{.*}} : index
+// CHECK: %[[E0:.*]] = vector.extract %[[BITS]][%[[M0]]] : i8 from vector<16xi8>
+// CHECK: vector.insert %[[E0]], %{{.*}} [0] : i8 into vector<4xi8>
+// CHECK: %[[C1:.*]] = arith.constant 1 : index
+// CHECK: %[[I1:.*]] = arith.addi %{{.*}}, %[[C1]] : index
+// CHECK: vector.extract %[[BITS]][%[[I1]]] : i8 from vector<16xi8>
+// CHECK: %[[C8:.*]] = arith.constant 8 : index
+// CHECK: %[[I2:.*]] = arith.addi %{{.*}}, %[[C8]] : index
+// CHECK: vector.extract %[[BITS]][%[[I2]]] : i8 from vector<16xi8>
+// CHECK: %[[C9:.*]] = arith.constant 9 : index
+// CHECK: %[[I3:.*]] = arith.addi %{{.*}}, %[[C9]] : index
+// CHECK: vector.extract %[[BITS]][%[[I3]]] : i8 from vector<16xi8>
+// CHECK: vector.shape_cast %{{.*}} : vector<4xf8E8M0FNU> to vector<2x2xf8E8M0FNU>
+// CHECK-NOT: gpu.shuffle
+gpu.module @xevm_module {
+gpu.func @convert_layout_broadcast_four_slots() {
+ %scale_src = arith.constant dense<1.0> : vector<8x2xf8E8M0FNU>
+ %cvt = xegpu.convert_layout %scale_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 = [4, 1], lane_data = [1, 1]>
+ }> : vector<8x2xf8E8M0FNU>
+ "some_use"(%cvt) : (vector<8x2xf8E8M0FNU>) -> ()
+ gpu.return
+}
+}
>From 48c8925f5070974ae6206f933718266d131841c3 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 00:50:24 +0000
Subject: [PATCH 09/46] Add worked instances where the mechanics are
non-obvious, drop restatement
A pass over the comments against two questions: does a reader meet any
non-trivial formula or procedure without a worked instance, and is anything
explained that the code already says.
Two places introduced mechanics with no instance:
isSupportedLanePeriod why the period must divide the subgroup size was a
chain of sentences. Now ends on the case that breaks:
a period of 3 on 16 lanes reaches donorDelta 15 and
slot 2, naming lane 17.
fitFragmentIndex the fit derives four constants from the first change
in the table, which is hard to picture. Now walks
needed = [0, 0, 2, 2] through to all four.
Two removals, both restatement rather than explanation:
The pseudo code transcribed the emitter line for line -- extract, shuffle,
insert -- which is what the function under it and the CHECK lines in the tests
already show. Reduced to one line, keeping the donor search, whose control flow
is the part that cannot be read off the code.
That only the first lanePeriod lanes end up correct was said twice, once as a
paragraph and again as a caveat, in different words. Merged.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 74 +++++++++----------
1 file changed, 35 insertions(+), 39 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 241a3a2355fc2..35c3dfda03953 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1926,11 +1926,16 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// slice<layout<[1, 1, 16], [1, 2, 1]>, dims = [2]>
// -> layout<[8, 1], [1, 2]>
//
-// The ownership tables hold one coordinate per distribution unit while the
-// algorithm indexes elements; the two coincide only when a unit is a single
-// element. See `computeOwnedCoords` for what lifting it needs. This also
-// covers sub-byte elements in practice, which come with `lane_data` 2 so
-// that a pair fills a byte.
+// `computeStaticDistributedCoords` yields one coordinate per distribution
+// unit -- the block start -- while the algorithm indexes elements, and the
+// two coincide only when a unit is a single element. Lifting it means
+// expanding the block starts in fragment order. `expandBlockCoords` in
+// XeGPUDialect.cpp expands them for layout comparison, but emits a whole
+// unit before moving to the next, which matches the fragment only when
+// lane_data is non-unit on the innermost dimension alone.
+//
+// This also covers sub-byte elements in practice, which come with
+// `lane_data` 2 so that a pair fills a byte.
//
// * FragmentSizeMismatch, a layout not distributing what its vector type
// accounts for (fundamental: the caller derived one from the other).
@@ -2046,39 +2051,30 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// fragment of its donor, lane `slot + donorDelta`
// if that donor does not hold the element: give up on this donorDelta
// look for one stride/offset/dimStride/dimExtent whose index reproduces
-// every needed[slot]; if there is none: next donorDelta
+// every needed[slot]; if there is none: next donorDelta
// accept (donorDelta, index) and stop
// if no donorDelta fits: fail, the layout change is not of this form
-// for each element source:
-// %idx = index(slot) // arith on gpu.lane_id
-// %val = vector.extract %fragment[%idx]
-// %donor = slot + donorDelta
-// %val = gpu.shuffle idx %val, %donor // unless %donor is %lane
-// %res = vector.insert %val, %res[pos]
-//
-// `%donor` is the lane itself when `donorDelta` is 0 and the extract already
-// lands on an element the lane owns, and that is the only case where the
-// shuffle is dropped. Replication makes it a choice: several donor groups may
-// hold the element and all give the same value, so the smallest `donorDelta` is
-// taken, which is what leaves a fully broadcast input shuffle-free.
-//
-// Only lanes 0 to `lanePeriod` - 1 are required to end up holding the elements
-// the target layout assigns their slot; each is the one lane of its slot the
-// layout distributes to. Lanes sharing a slot run the same arithmetic on their
-// own fragment, so when a shuffle is emitted they all receive the donor's value
-// and are equally correct, and when it is dropped they keep whatever that index
-// holds in their own fragment. The target layout assigns those lanes no
-// elements, so what they hold is unspecified.
+// emit, per element source: an extract at index(slot), a `gpu.shuffle idx`
+// from `slot + donorDelta` unless that is the lane itself, and an insert.
+//
+// The shuffle is dropped exactly when `donorDelta` is 0, since the extract then
+// lands on an element the lane owns. Replication makes that a choice: several
+// donor groups may hold the element and all give the same value, so the
+// smallest `donorDelta` is taken, which is what leaves a fully broadcast input
+// shuffle-free.
+//
+// Only the first `lanePeriod` lanes are required to end up correct -- each is
+// the one lane of its slot the target layout distributes to, and the layout
+// assigns the rest no elements. Lanes sharing a slot run the same arithmetic on
+// their own fragment: when a shuffle is emitted they all receive the donor's
+// value and are equally correct, and when it is dropped their local extract can
+// land on a different element. Dropping one therefore relies on nobody reading
+// them.
//
// An index is derived from the tables and then verified against them, so a
// layout change that is not of this form is reported as a match failure rather
// than lowered incorrectly.
//
-// One caveat: after the rewrite only the first `lanePeriod` lanes are
-// guaranteed to hold what the target layout assigns them. The others are
-// replicas whose local extract can land on a different element, so dropping a
-// shuffle relies on nobody reading them.
-//
//===----------------------------------------------------------------------===//
/// The lane period of `layout`: the modulus after which its assignment of
@@ -2158,7 +2154,8 @@ static bool haveSameDistributionUnit(xegpu::DistributeLayoutAttr inputLayout,
/// lanePeriod, 2 * lanePeriod, ...` up to the subgroup size, reading lane `slot
/// + donorDelta`, where `slot` is `lane % lanePeriod` and so below the period.
/// If the period did not divide the subgroup size that sum could name a lane
-/// past the last one.
+/// past the last one: with a period of 3 on 16 lanes, `donorDelta` reaches 15
+/// and a slot reaches 2, naming lane 17.
///
/// The input has no counterpart: it is only ever indexed, never used to derive
/// a donor, so its own period is unconstrained.
@@ -2203,13 +2200,7 @@ using OwnedCoords = SmallVector<SmallVector<SmallVector<int64_t>>>;
/// `layout`: `owned[lane][i]` is the coordinate of element `i` of that
/// fragment, numbered as the flattened fragment `vector.extract` indexes.
///
-/// Fails for non-unit lane_data. `computeStaticDistributedCoords` yields one
-/// coordinate per distribution unit -- the block start -- and those coincide
-/// with the elements only when a unit is a single element. Lifting this needs
-/// the block starts expanded in fragment order; `expandBlockCoords` in
-/// XeGPUDialect.cpp does such an expansion for layout comparison, but emits a
-/// whole unit before moving to the next, which matches the fragment only when
-/// lane_data is non-unit on the innermost dimension alone.
+/// Fails for non-unit lane_data, `LaneDataNotUnit` above.
static FailureOr<OwnedCoords>
computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
int64_t subgroupSize,
@@ -2282,6 +2273,11 @@ struct ElementSource {
/// restricted form `FragmentIndex` can emit. A slot is a lane reduced modulo
/// the lane period, so there are `lanePeriod` of them. Fails when no index in
/// that form reproduces the whole table.
+///
+/// For `needed = [0, 0, 2, 2]`: the first change is at slot 2, so pairs of
+/// slots share a coordinate (`dimStride` 2) and consecutive coordinates are 2
+/// elements apart (`stride` 2); the table reaches 2, one step, so the
+/// coordinate takes 2 values (`dimExtent` 2); and slot 0 reads 0 (`offset` 0).
static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
int64_t lanePeriod) {
// Slots sharing an index are adjacent, so the first change in `needed` gives
>From 85093f1714a4a05676678453068ccbb995c40cfb Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 02:18:08 +0000
Subject: [PATCH 10/46] State the capability as categories rather than a list
of cases
The "what it handles" section read as an enumeration of five concrete layout
pairs, which says nothing about what else would work. It is really two
properties, one per side, and the cases are one example each.
The input side is its replication: fully broadcast, so nothing has to travel, or
partially broadcast, so a lane has to be given what its group lacks -- the
latter in either spelling, a sliced dimension or a lane_layout naming fewer
lanes than the subgroup has.
The target side is its lane period, which decides where a lane may be given data
from: donors sit at multiples of it, so a period P over S lanes offers S/P donor
groups.
Putting the two together produces the boundary the list left implicit. A period
equal to the subgroup size leaves only donorDelta 0, so nothing can cross a
lane and a partially broadcast input cannot be completed. Verified: layout<[8,
1]> into slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims = [0]> is declined
with NoDonorDelta, a lane arriving with 2 of the 8 elements it needs. That row
is now in the table, and NoDonorDelta points at it.
It also shows cases 1 and 5 are one category, differing only in constants; 5 is
kept for having distinct ones, which the index expression is read off.
Net 52 insertions against 39 deletions, the per-case prose collapsing into the
table.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 91 +++++++++++--------
1 file changed, 52 insertions(+), 39 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 35c3dfda03953..618e50c545e72 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1789,54 +1789,67 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// What it handles
// ---------------
//
-// The cases below are the tests. Cases 1-3 are all on a `vector<8x2>` over 16
-// lanes. `#xegpu.` prefixes are dropped for width.
+// Two properties decide it, one per side. Examples are on a `vector<8x2>` over
+// 16 lanes, `#xegpu.` prefixes dropped for width; the case numbers are used
+// throughout the rest of this comment.
//
-// 1. Broadcast over the whole subgroup, distributed to 8 lanes.
+// The input's replication decides what a lane arrives with, and so what has to
+// travel:
//
-// slice<layout<lane_layout = [1, 1, 16]>, dims = [2]>
-// -> layout<lane_layout = [8, 1]>
+// fully broadcast every lane holds the whole value, so nothing travels
+// and
+// the result is extracts alone. Spelled by slicing away
+// every distributed dimension:
//
-// Every lane arrives holding all 16 elements and leaves holding one row.
-// Emits 2 extracts and no shuffle: every lane already has what it needs.
+// slice<layout<[1, 1, 16]>, dims = [2]>
//
-// 2. Broadcast over two groups of 8 lanes, distributed to 8 lanes.
+// partially broadcast groups of lanes each hold part of it, so a lane has to
+// be given whatever its group lacks. Two spellings,
+// treated alike -- a sliced dimension,
//
-// slice<layout<lane_layout = [8, 1, 2], order = [0, 2, 1]>, dims = [0]>
-// -> layout<lane_layout = [8, 1]>
+// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims =
+// [0]>
//
-// Lanes 0-7 arrive holding all of column 0, lanes 8-15 all of column 1, and
-// lane `i` leaves holding row `i`. So lane 3 wants (3, 0) and (3, 1); it
-// holds (3, 0) already, while (3, 1) is in lane 11. Emits 1 extract and 1
-// shuffle. This is the running example used below.
+// or a lane_layout naming fewer lanes than the subgroup
+// has, the rest repeating those:
//
-// 3. Broadcast over the whole subgroup, distributed to two groups of 8 lanes.
+// layout<[8, 1]>
//
-// slice<layout<lane_layout = [1, 1, 16]>, dims = [2]>
-// -> slice<layout<lane_layout = [8, 1, 2], order = [0, 2, 1]>,
-// dims = [0]>
+// The target's lane period decides where a lane may be given it from. Donors
+// sit at multiples of the period, so a period `P` over `S` lanes offers `S / P`
+// donor groups:
//
-// Every lane arrives holding all 16 elements; lanes 0-7 leave holding column
-// 0 and lanes 8-15 column 1. Emits 8 extracts and no shuffle.
+// period < S several groups, so elements can cross lanes.
+// layout<[8, 1]> offers 2, layout<[2, 1]> offers 8.
//
-// 4. Partial input layout, distributed to 2 lanes.
+// period == S one group, `donorDelta` 0, so nothing can cross a lane:
+// only an input already holding each lane's target
+// fragment works. A replicating target has this period,
+// slicing not reducing it:
//
-// layout<lane_layout = [8, 1]> -> layout<lane_layout = [2, 1]>
+// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims =
+// [0]>
//
-// The input names 8 of the 16 lanes, so lanes 8-15 repeat lanes 0-7 and lane
-// `i` arrives holding row `i % 8`. The target hands out 2 fragments of 4
-// rows each, so slot is `lane % 2` and the donors sit 2, 4 and 6 lanes away.
-// Emits 2 extracts and 6 shuffles.
+// The combinations, and what the tests show them emitting:
//
-// 5. Broadcast over the whole subgroup, distributed to 4 lanes.
+// case input target emits
+// 1 fully broadcast period 8 2 extracts
+// 5 fully broadcast period 4 4 extracts
+// 3 fully broadcast period 16 8 extracts
+// 2 partial, sliced period 8 1 extract, 1 shuffle
+// 4 partial, short period 2 2 extracts, 6 shuffles
+// -- partially broadcast period 16 declined, NoDonorDelta
//
-// slice<layout<lane_layout = [1, 1, 16]>, dims = [2]>
-// -> layout<lane_layout = [4, 1]>
+// Cases 1 and 5 are one category; 5 is kept because its constants are all
+// distinct, making it the clearest to read the index expression off. The last
+// row is the boundary the two properties imply: with only `donorDelta` 0
+// available, a lane arriving with less than its target fragment cannot be
+// completed. Verified for `layout<[8, 1]>` into the sliced target above, where
+// a lane arrives with 2 of the 8 elements it needs.
//
-// As case 1, but the target hands out 4 fragments of 2 rows instead of 8 of
-// one. Kept because it is the one case where the quantities below take
-// distinct values, so it is the clearest to read the index expression off.
-// Emits 4 extracts and no shuffle.
+// Case 2 is the running example below: lanes 0-7 arrive holding all of column
+// 0, lanes 8-15 all of column 1, and lane `i` leaves holding row `i`. Lane 3
+// wants (3, 0) and (3, 1); it holds (3, 0) already, while (3, 1) is in lane 11.
//
// Nested slices are coalesced before anything else looks at a layout, so only
// the lane_layout of the underlying layout and the union of the sliced dims
@@ -1860,11 +1873,10 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// even between elements taken from the same donor; see the TODO in
// the emitter, `gpu.shuffle` moving a whole 32-bit lane word.
//
-// So a fully broadcast input never shuffles: every lane holds everything, so
-// the smallest `donorDelta` that works is always 0, as in cases 1, 3 and 5.
-// Case 2 shuffles once, for the column its own group does not hold. Case 4
-// shuffles six times, for the three rows beyond its local one, both columns
-// each.
+// Against the categories above: cases 1, 3 and 5 shuffle nothing, a fully
+// broadcast input always finding `donorDelta` 0. Case 2 shuffles once, for the
+// column its own group does not hold, and case 4 six times, for the three rows
+// beyond its local one, both columns each.
//
// How close that is to optimal
// ----------------------------
@@ -1908,7 +1920,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// `lane_layout` together is handled by neither.
//
// * NoDonorDelta, no donor group holding the element for every slot
-// (fundamental: `gpu.shuffle idx` reads one value per lane).
+// (fundamental: `gpu.shuffle idx` reads one value per lane). The last row of
+// the table above is this limit reached through the target's period.
//
// * IndexNotLaneAffine, an index the restricted form cannot express
// (implementation: an arbitrary one needs a materialized table and a
>From c8f2919e671cbac8bae6f56b63274c1c6d07738c Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 02:26:44 +0000
Subject: [PATCH 11/46] Drop the cost and optimality analysis, and the
duplicate first case
Two trims to the section header.
The shuffle counting rules and the comparison against a theoretical floor were
analysis of the lowering rather than description of it, and nothing needs them
to read the code. Moved out; the emitter keeps its one-line TODO. The counts
stay in the capability table, where they are facts read off the tests.
Cases 1 and 5 were the same category, a fully broadcast input into a target of
period less than the subgroup size, differing only in the period and so in the
constants. Only one belongs in a section describing categories. Kept the one
whose four index quantities are distinct, renumbered to 1, and noted that the
suite carries the other as @convert_layout_broadcast_all_lanes. References
elsewhere in the comment were updated; case 2, the running example, keeps its
number.
Header goes from 301 to 255 lines.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 102 +++++-------------
1 file changed, 27 insertions(+), 75 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 618e50c545e72..faef07893e960 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1830,20 +1830,23 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims =
// [0]>
//
-// The combinations, and what the tests show them emitting:
+// One example per combination, and what the tests show it emitting:
//
// case input target emits
-// 1 fully broadcast period 8 2 extracts
-// 5 fully broadcast period 4 4 extracts
+// 1 fully broadcast period 4 4 extracts
// 3 fully broadcast period 16 8 extracts
// 2 partial, sliced period 8 1 extract, 1 shuffle
// 4 partial, short period 2 2 extracts, 6 shuffles
// -- partially broadcast period 16 declined, NoDonorDelta
//
-// Cases 1 and 5 are one category; 5 is kept because its constants are all
-// distinct, making it the clearest to read the index expression off. The last
-// row is the boundary the two properties imply: with only `donorDelta` 0
-// available, a lane arriving with less than its target fragment cannot be
+// The period is the only thing that varies within a category, so the test suite
+// carries more than one of some -- @convert_layout_broadcast_all_lanes is case
+// 1 with period 8 -- but they behave alike and only one is described here. Case
+// 1 is the period whose constants are all distinct, which the index expression
+// below is read off.
+//
+// The last row is the boundary the two properties imply: with only `donorDelta`
+// 0 available, a lane arriving with less than its target fragment cannot be
// completed. Verified for `layout<[8, 1]>` into the sliced target above, where
// a lane arrives with 2 of the 8 elements it needs.
//
@@ -1855,52 +1858,6 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// the lane_layout of the underlying layout and the union of the sliced dims
// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
//
-// What it costs
-// -------------
-//
-// There is one element source per element of the target fragment, so both
-// counts scale with how much the target gives one lane, not with the size of
-// the value. The two are not counted the same way, which is why they can differ
-// so widely:
-//
-// extracts one per *distinct* fragment index among the sources. Elements
-// reading the same index share one, which is why case 4 needs only
-// 2 for its 8 elements: its index does not depend on the slot, so
-// the two columns are the only distinct reads.
-//
-// shuffles one per element the reading lane does not already hold, that is
-// per source whose `donorDelta` is not 0. These are not shared,
-// even between elements taken from the same donor; see the TODO in
-// the emitter, `gpu.shuffle` moving a whole 32-bit lane word.
-//
-// Against the categories above: cases 1, 3 and 5 shuffle nothing, a fully
-// broadcast input always finding `donorDelta` 0. Case 2 shuffles once, for the
-// column its own group does not hold, and case 4 six times, for the three rows
-// beyond its local one, both columns each.
-//
-// How close that is to optimal
-// ----------------------------
-//
-// Extracts are minimal: they are register reads and equal indices are shared.
-//
-// Shuffles are not. `gpu.shuffle idx` reads one donor lane and moves a whole
-// 32-bit lane word, so within this scheme -- every lane reading one donor per
-// instruction -- the floor is one shuffle per distinct nonzero `donorDelta`,
-// carrying up to 32 bits of that group's elements at once. Emitting one per
-// remote element instead costs a factor of however many elements a lane takes
-// from the same group: nothing where that is one, as in case 2, and 2x in case
-// 4, whose 6 shuffles move 3 groups of two bytes. Closing that is the TODO in
-// the emitter; `shuffleDataAsLaneLayoutChange` in this file already packs
-// elements into lane words that way.
-//
-// A lower floor exists, but only outside this scheme. If lanes forwarded data
-// in several hops, a lane needing `k` remote elements of `w` bits could receive
-// them in `ceil(k * w / 32)` shuffles however many lanes they start on -- 2
-// rather than 3 for case 4. That trades instruction count for dependent latency
-// and a routing schedule. For the scale operands this lowers, a lane's remote
-// elements come from one or two donors, so the single-hop floor is at or near
-// the absolute one and the simpler scheme costs nothing.
-//
// What it declines
// ----------------
//
@@ -2028,30 +1985,26 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// per slot would need a materialized table and a gather; this costs a couple of
// `arith` ops on the lane id.
//
-// The cases, `p` being the position in the target fragment. Case 5 first, being
-// the one whose numbers are all distinct:
+// The cases, `p` being the position in the target fragment:
//
// case lanePeriod index(slot) stride dimStride dimExtent offset
-// 5 4 2*slot + off 2 1 4 0,1,8,9
-// 1 8 2*slot + p 2 1 8 p
+// 1 4 2*slot + off 2 1 4 0,1,8,9
// 2 8 slot 1 1 8 0
// 3 16 slot/8 + 2*p 1 8 2 2*p
// 4 2 p % 2 0 1 1 p % 2
//
// Where the constants come from: the `2`s multiplying `slot` are the size of
-// the input fragment's innermost dimension, which in cases 1, 3 and 5 is the
-// whole 8x2 value, so stepping one row costs 2 elements. In case 2 the fragment
-// is a single column, its row stride is 1, and no `2` appears. Case 4's index
-// does not depend on the slot at all: each target row arrives whole from one
-// donor, so `stride` is 0 and `offset` alone selects the column, with
-// `donorDelta` stepping 0, 2, 4, 6 over the rows.
-//
-// Read the expression off case 5, where the four quantities are 4, 2, 1 and 4.
-// On the `layout<[8, 1]>` targets of cases 1 and 2 they collide: the lane
-// period (8), case 2's `donorDelta` (8) and case 3's `dimStride` (8) all have
-// the same value while being three different things -- a lane count, a lane
-// offset and an index divisor. Those cases are kept because they are what the
-// scale operands of `dpas_mx` actually produce.
+// the input fragment's innermost dimension, which in cases 1 and 3 is the whole
+// 8x2 value, so stepping one row costs 2 elements. In case 2 the fragment is a
+// single column, its row stride is 1, and no `2` appears. Case 4's index does
+// not depend on the slot at all: each target row arrives whole from one donor,
+// so `stride` is 0 and `offset` alone selects the column.
+//
+// Read the expression off case 1, where the four quantities are 4, 2, 1 and 4.
+// On case 2's `layout<[8, 1]>` target they collide: the lane period (8), its
+// own `donorDelta` (8) and case 3's `dimStride` (8) all have the same value
+// while being three different things -- a lane count, a lane offset and an
+// index divisor.
//
// Algorithm
// ---------
@@ -2503,12 +2456,11 @@ static FailureOr<Value> redistributeBroadcastedValue(
return value;
};
- // Extract, shuffle and insert one element of the result at a time. One
+ // Extract, shuffle and insert one element of the result at a time: one
// extract per distinct fragment index, one shuffle per element the lane does
- // not already hold; see "What it costs" above.
- // TODO: one shuffle per element, though `gpu.shuffle` moves a whole 32-bit
- // lane word. Elements sharing a donor could travel in one, as
- // shuffleDataAsLaneLayoutChange does.
+ // not already hold.
+ // TODO: elements sharing a donor could travel in one shuffle, `gpu.shuffle`
+ // moving a whole 32-bit lane word.
llvm::DenseMap<std::tuple<int64_t, int64_t, int64_t, int64_t>, Value>
extracted;
for (auto [pos, source] : llvm::enumerate(sources)) {
>From 88965da746b297a4282ae0a99d2f7b68634da8d7 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 02:31:27 +0000
Subject: [PATCH 12/46] Reduce the declined list to the categories themselves
The list carried the cause of each declined pair and, for the two
implementation limits, what lifting them would take. None of that is needed to
read the code, and it made a capability statement three times longer than the
capability.
It is now the seven categories, grouped by what each constrains and tagged
implementation or inherent -- 26 lines against 68. computeOwnedCoords already
points at its category rather than restating it.
The causes and the lift guidance moved to an issue draft, which keeps the parts
worth not rediscovering: that lifting LaneDataNotUnit needs the block starts
expanded in fragment order and cannot just reuse expandBlockCoords, whose order
only agrees when lane_data is non-unit on the innermost dimension alone; that
decomposing LaneDataDiffers fails in both orderings, one step folding away and
the other being a transpose; and the verified instances behind NoDonorDelta and
LanePeriodNotDivisor.
Section header goes from 255 to 239 lines.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 70 ++++---------------
1 file changed, 14 insertions(+), 56 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index faef07893e960..0a42f715b701e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1861,70 +1861,28 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// What it declines
// ----------------
//
-// Each of these reports a match failure naming a `RedistributionLimit`, so this
-// list and the enum are the same list. Grouped by what the limit constrains,
-// which is where the two layouts are and are not treated alike. Same notation,
-// and `vector<8x2>` over 16 lanes unless said otherwise.
+// Every failure path reports a `RedistributionLimit`, so these are all of them.
+// Grouped by what the limit constrains; (i) is an implementation limit, (f) one
+// inherent to the scheme.
//
-// Relational -- properties of the pair, neither side alone:
+// The pair:
//
-// * LaneDataDiffers, `lane_data` differing between the layouts (fundamental
-// here):
+// LaneDataDiffers (f) the layouts disagree on the distribution unit
+// NoDonorDelta (f) the element is in no reachable donor group
+// IndexNotLaneAffine (i) the per-slot index is not affine in the lane id
//
-// layout<[8, 1], [1, 1]> -> layout<[1, 8], [1, 2]>
+// Either layout:
//
-// A `lane_data`-only change is the repack pattern's job; changing it and
-// `lane_layout` together is handled by neither.
+// LaneDataNotUnit (i) lane_data is not all ones
+// FragmentSizeMismatch (f) a layout disagrees with its vector type
//
-// * NoDonorDelta, no donor group holding the element for every slot
-// (fundamental: `gpu.shuffle idx` reads one value per lane). The last row of
-// the table above is this limit reached through the target's period.
+// The target alone:
//
-// * IndexNotLaneAffine, an index the restricted form cannot express
-// (implementation: an arbitrary one needs a materialized table and a
-// gather), as when `order` transposes it:
+// LanePeriodNotDivisor (f) its lane period does not divide the subgroup size
//
-// slice<layout<[1, 1, 16]>, dims = [2]>
-// -> layout<[4, 4], order = [0, 1]> on vector<4x4>
+// The value:
//
-// Lane `i` would need element `(i % 4) * 4 + i / 4`.
-//
-// Per layout -- the same rule applied to each side:
-//
-// * LaneDataNotUnit, `lane_data` equal but not all ones (implementation):
-//
-// slice<layout<[1, 1, 16], [1, 2, 1]>, dims = [2]>
-// -> layout<[8, 1], [1, 2]>
-//
-// `computeStaticDistributedCoords` yields one coordinate per distribution
-// unit -- the block start -- while the algorithm indexes elements, and the
-// two coincide only when a unit is a single element. Lifting it means
-// expanding the block starts in fragment order. `expandBlockCoords` in
-// XeGPUDialect.cpp expands them for layout comparison, but emits a whole
-// unit before moving to the next, which matches the fragment only when
-// lane_data is non-unit on the innermost dimension alone.
-//
-// This also covers sub-byte elements in practice, which come with
-// `lane_data` 2 so that a pair fills a byte.
-//
-// * FragmentSizeMismatch, a layout not distributing what its vector type
-// accounts for (fundamental: the caller derived one from the other).
-//
-// Target only -- the one asymmetry:
-//
-// * LanePeriodNotDivisor, the target's lane period not dividing the subgroup
-// size (fundamental to the donor walk, see `isSupportedLanePeriod`). The
-// input's period is unconstrained because it is only ever indexed.
-//
-// Value only -- neither layout:
-//
-// * SubByteElement, an element `gpu.shuffle` has no type for:
-//
-// slice<layout<[1, 1, 16]>, dims = [2]> -> layout<[8, 1]>
-// on vector<8x2xf4E2M1FN>
-//
-// Reachable only with unit `lane_data`; otherwise LaneDataNotUnit declines
-// first.
+// SubByteElement (f) no gpu.shuffle type carries the element
//
// Concepts
// --------
>From bc18bce8bf0f876a009490ebfbf4203a7674d4b4 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 02:34:41 +0000
Subject: [PATCH 13/46] Stop calling any decline fundamental
The declined list tagged each category implementation or inherent to the scheme.
That distinction was wrong. convert_layout is meaningful for any pair of
layouts, so every pair has to be lowerable by something; a decline that were
truly fundamental would be a hole in the dialect, not a property of this
pattern. Each of the six categories is reachable another way -- arbitrary
shuffle indices, several hops, a round trip through SLM, a materialized table,
or the repack pattern -- so tagging any of them inherent would tell the next
reader not to bother.
The tags are gone, replaced by one sentence saying the limits are the scheme s
and that another pattern or a more general scheme lowers any of them. That is
also the answer to whether each is implementation or fundamental: all
implementation, said once.
FragmentSizeMismatch is separated out. It fires on a layout disagreeing with the
vector type the caller derived from it, so it is an invariant check rather than
a layout pair, and listing it beside the others overstated it.
The issue draft carrying the causes had the same split and is corrected the same
way, its second group now named for needing a different data movement scheme
rather than for being inherent.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 23 +++++++++++--------
1 file changed, 14 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 0a42f715b701e..f810d2a319ca7 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1862,27 +1862,32 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// ----------------
//
// Every failure path reports a `RedistributionLimit`, so these are all of them.
-// Grouped by what the limit constrains; (i) is an implementation limit, (f) one
-// inherent to the scheme.
+// None is a limit of `convert_layout`, which is meaningful for any pair of
+// layouts and so always lowerable; each is a limit of the scheme here, one
+// shuffle per element from a donor a whole number of lane periods away. Another
+// pattern or a more general scheme lowers any of them. Grouped by what each
+// constrains.
//
// The pair:
//
-// LaneDataDiffers (f) the layouts disagree on the distribution unit
-// NoDonorDelta (f) the element is in no reachable donor group
-// IndexNotLaneAffine (i) the per-slot index is not affine in the lane id
+// LaneDataDiffers the layouts disagree on the distribution unit
+// NoDonorDelta the element is in no reachable donor group
+// IndexNotLaneAffine the per-slot index is not affine in the lane id
//
// Either layout:
//
-// LaneDataNotUnit (i) lane_data is not all ones
-// FragmentSizeMismatch (f) a layout disagrees with its vector type
+// LaneDataNotUnit lane_data is not all ones
//
// The target alone:
//
-// LanePeriodNotDivisor (f) its lane period does not divide the subgroup size
+// LanePeriodNotDivisor its lane period does not divide the subgroup size
//
// The value:
//
-// SubByteElement (f) no gpu.shuffle type carries the element
+// SubByteElement no gpu.shuffle type carries the element
+//
+// FragmentSizeMismatch is reported too, for a layout disagreeing with the vector
+// type the caller derived from it. That is an invariant check, not a layout pair.
//
// Concepts
// --------
>From 4be145e0a7d7031ac74d0ad7f39a755d33092f56 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 02:45:00 +0000
Subject: [PATCH 14/46] Give the declined side umbrella requirements rather
than fail points
The two capability sections were at different levels. What it handles states two
properties and reads four combinations off them. What it declines listed six
enum values, grouped by which side each constrains -- but an enum value is where
a check gave way, not a category of layout pair, so the section was an
enumeration wearing a grouping.
The six collapse into three requirements, two fail points each:
the change is a relocation of single elements
LaneDataDiffers, LaneDataNotUnit
the data can be reached
NoDonorDelta, LanePeriodNotDivisor
the read can be emitted
IndexNotLaneAffine, SubByteElement
Which is what the pattern needs, said as three statements, with the enum names
kept so a match failure still maps back. Three umbrellas against four covered
combinations also stops the declined side looking richer than the handled one.
The issue draft carrying the causes notes the same grouping, so the two agree.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 39 +++++++++----------
1 file changed, 19 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index f810d2a319ca7..ecf2cd59d2a6c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1861,33 +1861,32 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// What it declines
// ----------------
//
-// Every failure path reports a `RedistributionLimit`, so these are all of them.
-// None is a limit of `convert_layout`, which is meaningful for any pair of
-// layouts and so always lowerable; each is a limit of the scheme here, one
-// shuffle per element from a donor a whole number of lane periods away. Another
-// pattern or a more general scheme lowers any of them. Grouped by what each
-// constrains.
+// Three things have to hold. None is a limit of `convert_layout`, which is
+// meaningful for any pair of layouts and so always lowerable; each is a limit
+// of the scheme here, one shuffle per element from a donor a whole number of
+// lane periods away. Another pattern or a more general scheme lowers any of
+// them. The `RedistributionLimit` a match failure names is the point within one
+// of these that gave way.
//
-// The pair:
+// 1. The change is a relocation of single elements. Both layouts have to hand
+// the lanes the same unit, and that unit has to be one element.
//
-// LaneDataDiffers the layouts disagree on the distribution unit
-// NoDonorDelta the element is in no reachable donor group
-// IndexNotLaneAffine the per-slot index is not affine in the lane id
+// LaneDataDiffers, LaneDataNotUnit
//
-// Either layout:
+// 2. The data can be reached. Every element a lane needs has to sit in a lane
+// it
+// may read from, and the donor walk has to stay inside the subgroup.
//
-// LaneDataNotUnit lane_data is not all ones
+// NoDonorDelta, LanePeriodNotDivisor
//
-// The target alone:
+// 3. The read can be emitted. Its index has to follow from the lane id in a
+// couple of `arith` ops, and a `gpu.shuffle` type has to carry the element.
//
-// LanePeriodNotDivisor its lane period does not divide the subgroup size
+// IndexNotLaneAffine, SubByteElement
//
-// The value:
-//
-// SubByteElement no gpu.shuffle type carries the element
-//
-// FragmentSizeMismatch is reported too, for a layout disagreeing with the vector
-// type the caller derived from it. That is an invariant check, not a layout pair.
+// FragmentSizeMismatch is reported too, for a layout disagreeing with the
+// vector type the caller derived from it. That is an invariant check, not a
+// layout pair.
//
// Concepts
// --------
>From 51dbdb18ed4a6984a18eb5c571a8718e503e4bd8 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 02:50:29 +0000
Subject: [PATCH 15/46] Describe the declined side as layout combinations
The three requirements were indirect: a reader had to invert them to work out
which layout pairs are refused. Every covered category satisfies all three, so
the useful statement is which combinations do not, in the same terms the handled
side uses.
They divide by how they are decided rather than by what they constrain, which is
the honest split:
by inspecting the layouts lane_data differing, a common lane_data that is
not all ones, a target lane period that does not
divide the subgroup
by building the tables a partially broadcast input into a target of
period S, and a target whose distributed
dimension the index cannot walk because order
transposed it
The second group is shapes known to fail rather than a characterisation, and says
so -- other combinations fail the same way.
Two of the seven are not layout conditions at all and are named as such:
SubByteElement is about the element type, FragmentSizeMismatch about an
invariant.
The issue draft gains the closed form for the reachability half of NoDonorDelta,
gcd of the two lane periods dividing the distance between the group a slot needs
and the slot, checked against one lowering and one declined pair. That is what
would turn the search into an inspectable check.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 47 +++++++++++--------
1 file changed, 28 insertions(+), 19 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index ecf2cd59d2a6c..9e59550d72b73 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1861,32 +1861,41 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// What it declines
// ----------------
//
-// Three things have to hold. None is a limit of `convert_layout`, which is
-// meaningful for any pair of layouts and so always lowerable; each is a limit
-// of the scheme here, one shuffle per element from a donor a whole number of
-// lane periods away. Another pattern or a more general scheme lowers any of
-// them. The `RedistributionLimit` a match failure names is the point within one
-// of these that gave way.
+// The layout combinations, in the terms above, with the `RedistributionLimit` a
+// match failure names. None is a limit of `convert_layout`, which is meaningful
+// for any pair and so always lowerable; each is a limit of the scheme here, one
+// shuffle per element from a donor a whole number of lane periods away. Another
+// pattern or a more general scheme lowers any of them.
//
-// 1. The change is a relocation of single elements. Both layouts have to hand
-// the lanes the same unit, and that unit has to be one element.
+// Decided by inspecting the layouts:
//
-// LaneDataDiffers, LaneDataNotUnit
+// they differ in `lane_data` LaneDataDiffers
//
-// 2. The data can be reached. Every element a lane needs has to sit in a lane
-// it
-// may read from, and the donor walk has to stay inside the subgroup.
+// layout<[8, 1], [1, 1]> -> layout<[1, 8], [1, 2]>
//
-// NoDonorDelta, LanePeriodNotDivisor
+// A `lane_data`-only change is the repack pattern's; both changing is
+// neither's.
//
-// 3. The read can be emitted. Its index has to follow from the lane id in a
-// couple of `arith` ops, and a `gpu.shuffle` type has to carry the element.
+// their common `lane_data` is not all ones LaneDataNotUnit
+// the target's lane period does not divide the subgroup LanePeriodNotDivisor
//
-// IndexNotLaneAffine, SubByteElement
+// Decided by building the ownership tables, so these are shapes known to fail
+// rather than a characterisation -- other combinations fail the same way:
//
-// FragmentSizeMismatch is reported too, for a layout disagreeing with the
-// vector type the caller derived from it. That is an invariant check, not a
-// layout pair.
+// a partially broadcast input into a target of period S NoDonorDelta
+//
+// The last row of the table above: only `donorDelta` 0 is available, so a
+// lane arriving with less than its target fragment cannot be completed.
+//
+// a target whose distributed dimension the index cannot IndexNotLaneAffine
+// walk, `order` having transposed it
+//
+// slice<layout<[1, 1, 16]>, dims = [2]>
+// -> layout<[4, 4], order = [0, 1]> on vector<4x4>
+//
+// Not layout conditions: SubByteElement, `gpu.shuffle` having no type for an
+// element narrower than a byte, and FragmentSizeMismatch, an invariant check
+// for a layout disagreeing with the vector type the caller derived from it.
//
// Concepts
// --------
>From e138d41567ff67d7654a6b9640e1eccb2d6de585 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 02:53:03 +0000
Subject: [PATCH 16/46] Lead with the layout, not the period derived from it
Layout is what a reader sees in their own IR. Lane period is computed from it, so
a section that names only the period asks the reader to work backwards.
The target axis now says what the period is -- the product of the lane_layout,
which slicing does not reduce -- so it can be read off a layout instead of looked
up. The capability table gives the target as its lane_layout with the period
beside it rather than instead of it. LanePeriodNotDivisor is phrased as the
product of the target lane_layout not dividing the subgroup size, which is
something you can check by looking, matching the group it sits in.
Only getLanePeriod still names a period on its own, that being where it is
defined, with the layout it comes from alongside.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 34 ++++++++++---------
1 file changed, 18 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 9e59550d72b73..320dc32ea768d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1815,9 +1815,10 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// layout<[8, 1]>
//
-// The target's lane period decides where a lane may be given it from. Donors
-// sit at multiples of the period, so a period `P` over `S` lanes offers `S / P`
-// donor groups:
+// The target's lane period decides where a lane may be given it from. The period
+// is the product of its `lane_layout`, which slicing does not reduce, and donors
+// sit at multiples of it, so a period `P` over `S` lanes offers `S / P` donor
+// groups:
//
// period < S several groups, so elements can cross lanes.
// layout<[8, 1]> offers 2, layout<[2, 1]> offers 8.
@@ -1830,20 +1831,20 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims =
// [0]>
//
-// One example per combination, and what the tests show it emitting:
+// One example per combination, and what the tests show it emitting. The target is
+// given as its `lane_layout`, the period following from it:
//
-// case input target emits
-// 1 fully broadcast period 4 4 extracts
-// 3 fully broadcast period 16 8 extracts
-// 2 partial, sliced period 8 1 extract, 1 shuffle
-// 4 partial, short period 2 2 extracts, 6 shuffles
-// -- partially broadcast period 16 declined, NoDonorDelta
+// case input target lane_layout period emits
+// 1 fully broadcast [4, 1] 4 4 extracts
+// 3 fully broadcast [8, 1, 2] sliced 16 8 extracts
+// 2 partial, sliced [8, 1] 8 1 extract, 1 shuffle
+// 4 partial, short [2, 1] 2 2 extracts, 6 shuffles
+// -- partially broadcast [8, 1, 2] sliced 16 declined, NoDonorDelta
//
-// The period is the only thing that varies within a category, so the test suite
-// carries more than one of some -- @convert_layout_broadcast_all_lanes is case
-// 1 with period 8 -- but they behave alike and only one is described here. Case
-// 1 is the period whose constants are all distinct, which the index expression
-// below is read off.
+// Only the target varies within a category, so the test suite carries more than
+// one of some -- @convert_layout_broadcast_all_lanes is case 1 with `[8, 1]` --
+// but they behave alike and only one is described here. Case 1 is the one whose
+// constants are all distinct, which the index expression below is read off.
//
// The last row is the boundary the two properties imply: with only `donorDelta`
// 0 available, a lane arriving with less than its target fragment cannot be
@@ -1877,7 +1878,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// neither's.
//
// their common `lane_data` is not all ones LaneDataNotUnit
-// the target's lane period does not divide the subgroup LanePeriodNotDivisor
+// the product of the target's `lane_layout` does not LanePeriodNotDivisor
+// divide the subgroup size
//
// Decided by building the ownership tables, so these are shapes known to fail
// rather than a characterisation -- other combinations fail the same way:
>From 9d7732adc176dd5ace17cb960e8fd17f17f6c7e0 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 03:06:13 +0000
Subject: [PATCH 17/46] Introduce every derived term before it is used, from
the layouts
Reading the patch top-down, lanePeriod, fragment, slot and donor all
appeared before anything said what they were. Add a Terms block right
after the opening paragraph that derives each one from lane_layout,
lane_data and slicing -- the only things a user writes -- and place it
ahead of every use.
Give replication a worked instance of each of its two spellings, a
sliced dimension and a lane_layout shorter than the subgroup, showing
how fragment and lane period come out for both; the later sections now
refer back to it instead of repeating the spellings.
Also reflow the decline lists, whose enum column had been mangled, and
drop the donor rename in the repack pattern's comment, which used the
term outside the section that defines it.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 178 +++++++++---------
1 file changed, 91 insertions(+), 87 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 320dc32ea768d..45276c318f496 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1630,7 +1630,7 @@ shuffleDataAsLaneLayoutChange(ConversionPatternRewriter &rewriter, Location loc,
// -- create a temp 1D vector of i32 initialized to zero
// -- for each i32 element:
// ---- extract it from the source bundle
- // ---- gpu.shuffle to gather the value from the donor lane
+ // ---- gpu.shuffle to gather the value from the partner lane
// ---- insert it into the temp bundle
// -- cast the temp back to the source vector type
// -- vector.shuffle the source and temp to concatenate along the outer dim
@@ -1766,80 +1766,96 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//===----------------------------------------------------------------------===//
//
// Lowers a `convert_layout` that only relocates data between lanes: the two
-// layouts hand the lanes the same distribution unit, and differ only in which
-// lane holds which part of the value. Every element already exists in some
-// lane, so nothing is recomputed. The scale operands of scaled matrix
-// multiplication need this, being produced broadcasted and consumed
-// distributed.
+// layouts agree on `lane_data`, and differ only in which lane holds which part
+// of the value. Every element already exists in some lane, so nothing is
+// recomputed. The scale operands of scaled matrix multiplication need this,
+// being produced broadcasted and consumed distributed.
//
-// Either layout may replicate. A layout replicates over a set of lanes when a
-// `lane_layout` dimension is sliced away, and equally when its `lane_layout`
-// lays out fewer lanes than the subgroup has: the lane id is delinearized
-// modulo each `lane_layout` extent, so the lanes past their product repeat the
-// ones before it. The two spellings are treated alike, and case 4 below is the
-// second one.
+// Terms
+// -----
//
-// The two sides are not symmetric in what they mean, though:
+// All of these follow from the layouts, for a value of shape `Sh` on a subgroup
+// of `S` lanes. `#xegpu.` prefixes are dropped for width throughout.
//
-// input supplies the values. Every lane is a potential donor, so what each
-// one holds has to be known -- which it is, for any lane_layout.
-// target states the obligation. Only the lanes it distributes to have to end
-// up holding what it assigns them; the rest are left unspecified.
+// distribution unit what a lane owns as one block: `lane_data`
+// fragment everything one lane owns. Slicing drops dimensions from
+// `lane_layout`; over what is left a lane owns
+// `Sh[i] / lane_layout[i]` per dimension, flattened
+// row-major
+// lane period `product(lane_layout)` counting the dimensions slicing
+// dropped, so slicing does not reduce it. Lanes this far
+// apart hold the same fragment. `getLanePeriod`,
+// `lanePeriod` below
+// slot `lane % lanePeriod`, a lane's index within one period
+// donor the lane a value is read from, `slot + donorDelta`, with
+// `donorDelta` a multiple of `lanePeriod`
//
-// What it handles
-// ---------------
+// Either layout may replicate -- a set of lanes all holding the same fragment
+// -- and a layout has two ways to spell that. Both of these are on a
+// `vector<8x2>` over 16 lanes:
//
-// Two properties decide it, one per side. Examples are on a `vector<8x2>` over
-// 16 lanes, `#xegpu.` prefixes dropped for width; the case numbers are used
-// throughout the rest of this comment.
+// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims = [0]> ("sliced")
+// slicing leaves `lane_layout` [1, 2], so a fragment is 8x1, one whole
+// column, while the lane period still counts the sliced dimension and is
+// 8 * 1 * 2 = 16. Sixteen lanes hold two distinct columns between them,
+// eight lanes to a column.
//
-// The input's replication decides what a lane arrives with, and so what has to
-// travel:
+// layout<[8, 1]> ("short")
+// a `lane_layout` naming fewer lanes than the subgroup has. A fragment is
+// 1x2, one whole row, and the lane period is 8: the lane id is delinearized
+// modulo each extent, so a lane's coordinates depend on `lane % 8` alone
+// and lanes 8-15 repeat lanes 0-7.
//
-// fully broadcast every lane holds the whole value, so nothing travels
-// and
-// the result is extracts alone. Spelled by slicing away
-// every distributed dimension:
+// The two spellings are treated alike throughout.
//
-// slice<layout<[1, 1, 16]>, dims = [2]>
+// The two sides are not symmetric in what they mean:
//
-// partially broadcast groups of lanes each hold part of it, so a lane has to
-// be given whatever its group lacks. Two spellings,
-// treated alike -- a sliced dimension,
+// input supplies the values. Any lane may be a donor, so every lane's
+// fragment has to be known -- which it is, for any lane_layout.
+// target states the obligation. Only the lanes it distributes to have to end
+// up holding the fragment it assigns them; the rest are left
+// unspecified.
//
-// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims =
-// [0]>
+// What it handles
+// ---------------
+//
+// Two properties decide it, one per side. Examples stay on a `vector<8x2>` over
+// 16 lanes; the case numbers are used throughout the rest of this comment.
//
-// or a lane_layout naming fewer lanes than the subgroup
-// has, the rest repeating those:
+// The input's replication decides what a lane arrives with, and so what has to
+// travel:
//
-// layout<[8, 1]>
+// fully broadcast every lane already holds the whole value, so nothing
+// travels and the result is extracts alone. Spelled by
+// slicing away every distributed dimension:
+// slice<layout<[1, 1, 16]>, dims = [2]>
//
-// The target's lane period decides where a lane may be given it from. The period
-// is the product of its `lane_layout`, which slicing does not reduce, and donors
-// sit at multiples of it, so a period `P` over `S` lanes offers `S / P` donor
-// groups:
+// partially broadcast groups of lanes each hold part of it, so a lane has to
+// be given whatever its group lacks. Either spelling of
+// replication from Terms above.
//
-// period < S several groups, so elements can cross lanes.
-// layout<[8, 1]> offers 2, layout<[2, 1]> offers 8.
+// The target's lane period decides where a lane may be given that from. Donors
+// sit at multiples of the period, so a period `P` over `S` lanes offers `S / P`
+// donor groups:
//
-// period == S one group, `donorDelta` 0, so nothing can cross a lane:
-// only an input already holding each lane's target
-// fragment works. A replicating target has this period,
-// slicing not reducing it:
+// period < S several groups, so elements can cross lanes. `layout<[8, 1]>`
+// offers 2, `layout<[2, 1]>` offers 8.
//
-// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims =
-// [0]>
+// period == S one group, `donorDelta` 0, so nothing can cross a lane: only
+// an input already holding each lane's target fragment works.
+// A replicating target can have this period, slicing not
+// reducing it.
//
-// One example per combination, and what the tests show it emitting. The target is
-// given as its `lane_layout`, the period following from it:
+// One example per combination, and what the tests show it emitting. The target
+// is given as its `lane_layout`, the period following from it; the input as its
+// category, `sliced` and `short` being the two spellings named in Terms:
//
-// case input target lane_layout period emits
-// 1 fully broadcast [4, 1] 4 4 extracts
-// 3 fully broadcast [8, 1, 2] sliced 16 8 extracts
-// 2 partial, sliced [8, 1] 8 1 extract, 1 shuffle
-// 4 partial, short [2, 1] 2 2 extracts, 6 shuffles
-// -- partially broadcast [8, 1, 2] sliced 16 declined, NoDonorDelta
+// case input target period emits
+// 1 fully broadcast [4, 1] 4 4 extracts
+// 3 fully broadcast [8, 1, 2] sliced 16 8 extracts
+// 2 partial, sliced [8, 1] 8 1 extract, 1 shuffle
+// 4 partial, short [2, 1] 2 2 extracts, 6 shuffles
+// -- partial, short [8, 1, 2] sliced 16 declined
//
// Only the target varies within a category, so the test suite carries more than
// one of some -- @convert_layout_broadcast_all_lanes is case 1 with `[8, 1]` --
@@ -1848,8 +1864,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// The last row is the boundary the two properties imply: with only `donorDelta`
// 0 available, a lane arriving with less than its target fragment cannot be
-// completed. Verified for `layout<[8, 1]>` into the sliced target above, where
-// a lane arrives with 2 of the 8 elements it needs.
+// completed, and the match fails with NoDonorDelta. There a lane arrives with 2
+// of the 8 elements it needs.
//
// Case 2 is the running example below: lanes 0-7 arrive holding all of column
// 0, lanes 8-15 all of column 1, and lane `i` leaves holding row `i`. Lane 3
@@ -1870,27 +1886,26 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// Decided by inspecting the layouts:
//
-// they differ in `lane_data` LaneDataDiffers
+// they differ in `lane_data` LaneDataDiffers
//
// layout<[8, 1], [1, 1]> -> layout<[1, 8], [1, 2]>
//
// A `lane_data`-only change is the repack pattern's; both changing is
// neither's.
//
-// their common `lane_data` is not all ones LaneDataNotUnit
-// the product of the target's `lane_layout` does not LanePeriodNotDivisor
-// divide the subgroup size
+// their common `lane_data` is not all ones LaneDataNotUnit
//
-// Decided by building the ownership tables, so these are shapes known to fail
-// rather than a characterisation -- other combinations fail the same way:
+// the target's lane period does not divide `S` LanePeriodNotDivisor
//
-// a partially broadcast input into a target of period S NoDonorDelta
+// Decided by building the ownership tables, so these name layout combinations
+// known to fail rather than a full characterisation -- others fail the same
+// way:
//
-// The last row of the table above: only `donorDelta` 0 is available, so a
-// lane arriving with less than its target fragment cannot be completed.
+// a partially broadcast input into a target of NoDonorDelta
+// period `S`, the last row of the table above
//
-// a target whose distributed dimension the index cannot IndexNotLaneAffine
-// walk, `order` having transposed it
+// a target whose distributed dimension the index IndexNotLaneAffine
+// cannot walk, `order` having transposed it
//
// slice<layout<[1, 1, 16]>, dims = [2]>
// -> layout<[4, 4], order = [0, 1]> on vector<4x4>
@@ -1902,28 +1917,17 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Concepts
// --------
//
-// fragment
-// What a lane holds under a layout: 8x1 under case 2's input layout, 1x2
-// under its target.
-//
// broadcast group
// Lanes holding the same fragment. A sliced lane_layout dimension forms
// one, since nothing is distributed over its lanes: in case 2 dim 0 is
// sliced and spans 8 lanes, so lanes 0-7 are one group and 8-15 another.
//
-// slot
-// A lane's index within the target's lane period, `lane % lanePeriod`. In
-// case 2 `lanePeriod` is 8, so the lanes pair up: lane 11 has slot 3 like
-// lane 3, and runs the same arithmetic on its own fragment. How many
-// elements a lane holds is a separate matter -- it can own several, which
-// show up as several positions in its fragment, not as several slots.
-//
-// donor
-// The lane a value is read from, `slot + donorDelta`. `gpu.shuffle idx` is
-// the only way to move data between lanes and it exchanges one value per
-// lane, so every lane runs the same extract; a donor can therefore only
-// serve lanes of its own slot, which makes `donorDelta` a multiple of
-// `lanePeriod`. In case 2, slot 3 reads (3, 1) from lane 3 + 8.
+// donor, continued
+// `gpu.shuffle idx` is the only way to move data between lanes and it
+// exchanges one value per lane, so every lane runs the same extract; a
+// donor can therefore only serve lanes of its own slot, which is why
+// `donorDelta` is a multiple of `lanePeriod`. In case 2, slot 3 reads (3,
+// 1) from lane 3 + 8.
//
// That is also why the donor extracts the right thing without knowing who
// asked: its own slot is `(slot + donorDelta) % lanePeriod`, which is
>From 79019646df908176b377c7792a7d1af7d7836451 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 03:13:05 +0000
Subject: [PATCH 18/46] Define every coined term in Terms, not in a glossary
further down
donor group, element source and ownership table were coined names whose
meaning is neither plain English nor standard, and each was either used
before the glossary that defined it or defined so far from its uses that
the definition was easy to miss. Move all three into Terms alongside the
layout-derived ones, so nothing in the section is used before it is
introduced.
Drop broadcast group, a term nothing used -- replication already covers
what it named.
What remains of the glossary is the reason donorDelta is a multiple of
lanePeriod, which is a rationale rather than a definition. Give it its
own section and put it after the cases it draws its example from.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 70 +++++++++----------
1 file changed, 34 insertions(+), 36 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 45276c318f496..98539c2b3ffba 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1774,8 +1774,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Terms
// -----
//
-// All of these follow from the layouts, for a value of shape `Sh` on a subgroup
-// of `S` lanes. `#xegpu.` prefixes are dropped for width throughout.
+// These follow from the layouts, for a value of shape `Sh` on a subgroup of `S`
+// lanes. `#xegpu.` prefixes are dropped for width throughout.
//
// distribution unit what a lane owns as one block: `lane_data`
// fragment everything one lane owns. Slicing drops dimensions from
@@ -1787,8 +1787,23 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// apart hold the same fragment. `getLanePeriod`,
// `lanePeriod` below
// slot `lane % lanePeriod`, a lane's index within one period
+// ownership table every lane's fragment, tabulated as element coordinates,
+// one table per side. `computeOwnedCoords`, `inputOwned`
+// and `targetOwned` below
+//
+// These name the moving parts of the scheme:
+//
// donor the lane a value is read from, `slot + donorDelta`, with
// `donorDelta` a multiple of `lanePeriod`
+// donor group the lanes one `donorDelta` names, one per slot:
+// `donorDelta` through `donorDelta + lanePeriod - 1`.
+// Candidates are tried a group at a time, since every slot
+// has to be served by the same `donorDelta`
+// element source where one element of the target fragment comes from: a
+// donor, and the index to extract from that donor's
+// fragment. One per element, and it needs a shuffle
+// exactly when the donor is not the lane itself.
+// `ElementSource` below
//
// Either layout may replicate -- a set of lanes all holding the same fragment
// -- and a layout has two ways to spell that. Both of these are on a
@@ -1875,6 +1890,23 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// the lane_layout of the underlying layout and the union of the sliced dims
// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
//
+// Why donors sit a whole period apart
+// -----------------------------------
+//
+// `gpu.shuffle idx` is the only way to move data between lanes and it exchanges
+// one value per lane, so every lane runs the same extract. A donor can
+// therefore only serve lanes of its own slot, which is why `donorDelta` is a
+// multiple of `lanePeriod`. It also means the donor extracts the right thing
+// without knowing who asked: its own slot is
+// `(slot + donorDelta) % lanePeriod`, which is `slot`, so it computes the same
+// index.
+//
+// Case 2 has two element sources, one per element of its 1x2 fragment. Slot 3
+// wants (3, 0) and (3, 1):
+//
+// element 0 -> donorDelta = 0, donor = 3, index = 3 already local
+// element 1 -> donorDelta = 8, donor = 11, index = 3 shuffled
+//
// What it declines
// ----------------
//
@@ -1914,40 +1946,6 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// element narrower than a byte, and FragmentSizeMismatch, an invariant check
// for a layout disagreeing with the vector type the caller derived from it.
//
-// Concepts
-// --------
-//
-// broadcast group
-// Lanes holding the same fragment. A sliced lane_layout dimension forms
-// one, since nothing is distributed over its lanes: in case 2 dim 0 is
-// sliced and spans 8 lanes, so lanes 0-7 are one group and 8-15 another.
-//
-// donor, continued
-// `gpu.shuffle idx` is the only way to move data between lanes and it
-// exchanges one value per lane, so every lane runs the same extract; a
-// donor can therefore only serve lanes of its own slot, which is why
-// `donorDelta` is a multiple of `lanePeriod`. In case 2, slot 3 reads (3,
-// 1) from lane 3 + 8.
-//
-// That is also why the donor extracts the right thing without knowing who
-// asked: its own slot is `(slot + donorDelta) % lanePeriod`, which is
-// `slot`, so it computes the same index.
-//
-// donor group
-// The lanes one `donorDelta` names, one per slot: `donorDelta` through
-// `donorDelta + lanePeriod - 1`. Candidate donors are considered a group at
-// a time, since every slot has to be served by the same `donorDelta`. In
-// case 2 there are two, lanes 0-7 and lanes 8-15.
-//
-// element source
-// Where one element of the target fragment comes from: the donor, and the
-// index to extract from that donor's fragment. A shuffle is needed exactly
-// when the donor cannot be the lane itself. Case 2 has two, one per element
-// of its 1x2 fragment. Slot 3 wants (3, 0) and (3, 1):
-//
-// element 0 -> donorDelta = 0, donor = 3, index = 3 already local
-// element 1 -> donorDelta = 8, donor = 11, index = 3 shuffled
-//
// The index expression
// --------------------
//
>From bdbd86816e9e1ab07173a082469faa1845f9e784 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 03:55:38 +0000
Subject: [PATCH 19/46] Drop donor group, which named nothing beyond donorDelta
itself
The term was mine and it was a bad one. A donor group was the lanes one
donorDelta selects -- one per slot, so they hold different fragments,
which is nearly the opposite of a broadcast group. The only example given
for it was case 2, where the target period happens to equal the input's
broadcast group size and the two coincide, so the term invited exactly
the wrong reading.
Nothing needed a name here: a donor group was one value of donorDelta,
so say that instead. Use broadcast group freely, as a term that carries
its own meaning.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 50 +++++++++----------
1 file changed, 23 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 98539c2b3ffba..bb025244bc175 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1793,12 +1793,10 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// These name the moving parts of the scheme:
//
-// donor the lane a value is read from, `slot + donorDelta`, with
-// `donorDelta` a multiple of `lanePeriod`
-// donor group the lanes one `donorDelta` names, one per slot:
-// `donorDelta` through `donorDelta + lanePeriod - 1`.
-// Candidates are tried a group at a time, since every slot
-// has to be served by the same `donorDelta`
+// donor the lane a value is read from, `slot + donorDelta`.
+// `donorDelta` is one multiple of `lanePeriod`, shared by
+// every slot, so it is the single choice the scheme makes
+// per element
// element source where one element of the target fragment comes from: a
// donor, and the index to extract from that donor's
// fragment. One per element, and it needs a shuffle
@@ -1845,19 +1843,19 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// slicing away every distributed dimension:
// slice<layout<[1, 1, 16]>, dims = [2]>
//
-// partially broadcast groups of lanes each hold part of it, so a lane has to
-// be given whatever its group lacks. Either spelling of
-// replication from Terms above.
+// partially broadcast each broadcast group holds part of it, so a lane has
+// to be given whatever its group lacks. Either spelling
+// of replication from Terms above.
//
// The target's lane period decides where a lane may be given that from. Donors
-// sit at multiples of the period, so a period `P` over `S` lanes offers `S / P`
-// donor groups:
+// sit at multiples of the period, so a period `P` over `S` lanes leaves `S / P`
+// values of `donorDelta` to choose from:
//
-// period < S several groups, so elements can cross lanes. `layout<[8, 1]>`
-// offers 2, `layout<[2, 1]>` offers 8.
+// period < S more than one, so elements can cross lanes. `layout<[8, 1]>`
+// leaves 2, `layout<[2, 1]>` leaves 8.
//
-// period == S one group, `donorDelta` 0, so nothing can cross a lane: only
-// an input already holding each lane's target fragment works.
+// period == S only `donorDelta` 0, so nothing can cross a lane: only an
+// input already holding each lane's target fragment works.
// A replicating target can have this period, slicing not
// reducing it.
//
@@ -2000,9 +1998,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// The shuffle is dropped exactly when `donorDelta` is 0, since the extract then
// lands on an element the lane owns. Replication makes that a choice: several
-// donor groups may hold the element and all give the same value, so the
-// smallest `donorDelta` is taken, which is what leaves a fully broadcast input
-// shuffle-free.
+// values of `donorDelta` may work and all give the same value, so the smallest
+// is taken, which is what leaves a fully broadcast input shuffle-free.
//
// Only the first `lanePeriod` lanes are required to end up correct -- each is
// the one lane of its slot the target layout distributes to, and the layout
@@ -2257,13 +2254,12 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
/// Works out where element `pos` of the target fragment comes from: which donor
/// lane, at `slot + donorDelta` for some multiple of the lane period, and which
-/// element of that donor's fragment. Candidates are tried a donor group at a
-/// time -- the lanes one `donorDelta` names, one per slot -- because every slot
-/// has to be served by the same `donorDelta`. The smallest is preferred, so a
-/// lane that already holds the element needs no shuffle.
+/// element of that donor's fragment. Candidates are tried one `donorDelta` at a
+/// time because every slot has to be served by the same one; the smallest is
+/// preferred so that a lane already holding the element needs no shuffle.
///
-/// Fails when no donor group holds the element for every slot, or when the
-/// indices they would be read at are not of `FragmentIndex`'s form.
+/// Fails when no `donorDelta` puts the element within reach of every slot, or
+/// when the indices it would be read at are not of `FragmentIndex`'s form.
static FailureOr<ElementSource>
deriveElementSource(const OwnedCoords &inputOwned,
const OwnedCoords &targetOwned, int64_t pos,
@@ -2278,9 +2274,9 @@ deriveElementSource(const OwnedCoords &inputOwned,
return -1;
};
- // Donor groups, smallest `donorDelta` first. `sawCompleteDonor` records
- // whether any group held the element at all, which is what distinguishes the
- // two ways this fails.
+ // Each `donorDelta`, smallest first. `sawCompleteDonor` records whether any
+ // of them held the element at all, which is what distinguishes the two ways
+ // this fails.
bool sawCompleteDonor = false;
for (int64_t donorDelta = 0; donorDelta < subgroupSize;
donorDelta += lanePeriod) {
>From ab46a7433715b108a2ee167646c76899acbcce84 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 04:00:54 +0000
Subject: [PATCH 20/46] Derive broadcast group from the layout instead of
assuming it is obvious
A broadcast group looked like a term that carried its own meaning, but
its size is not something a reader can eyeball: it is product(lane_layout)
taken after slicing, counted against the subgroup size. Put that
derivation in Terms.
It pairs with the lane period, which is the same product with the sliced
dimensions still counted. Stating them together explains the one result
that otherwise looks contradictory -- a layout can have a lane period of
S and still replicate, because slicing raises the period without adding
fragments -- and gives the group count for both spellings of replication.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 30 ++++++++++++-------
1 file changed, 19 insertions(+), 11 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index bb025244bc175..46070a60df0f4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1782,8 +1782,13 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// `lane_layout`; over what is left a lane owns
// `Sh[i] / lane_layout[i]` per dimension, flattened
// row-major
-// lane period `product(lane_layout)` counting the dimensions slicing
-// dropped, so slicing does not reduce it. Lanes this far
+// broadcast group the lanes that hold one fragment. `product(lane_layout)`
+// taken *after* slicing counts the distinct fragments, so
+// each group is `S / that` lanes: all `S` of them when
+// every dimension is sliced away, and 1 -- no replication
+// -- when the product is already `S`
+// lane period the same product with the sliced dimensions still
+// counted, so slicing does not reduce it. Lanes this far
// apart hold the same fragment. `getLanePeriod`,
// `lanePeriod` below
// slot `lane % lanePeriod`, a lane's index within one period
@@ -1803,21 +1808,24 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// exactly when the donor is not the lane itself.
// `ElementSource` below
//
-// Either layout may replicate -- a set of lanes all holding the same fragment
-// -- and a layout has two ways to spell that. Both of these are on a
-// `vector<8x2>` over 16 lanes:
+// A broadcast group and the lane period are that one product taken either side
+// of slicing, so a layout can have a lane period of `S` and still replicate:
+// the sliced dimensions raise the period without adding fragments. That is one
+// of the two ways a layout spells replication. Both, on a `vector<8x2>` over 16
+// lanes:
//
// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims = [0]> ("sliced")
// slicing leaves `lane_layout` [1, 2], so a fragment is 8x1, one whole
-// column, while the lane period still counts the sliced dimension and is
-// 8 * 1 * 2 = 16. Sixteen lanes hold two distinct columns between them,
-// eight lanes to a column.
+// column, and there are 1 * 2 = 2 of them: two groups of 8, lanes 0-7 and
+// lanes 8-15. The lane period still counts the sliced dimension and is
+// 8 * 1 * 2 = 16, well above the group size.
//
// layout<[8, 1]> ("short")
// a `lane_layout` naming fewer lanes than the subgroup has. A fragment is
-// 1x2, one whole row, and the lane period is 8: the lane id is delinearized
-// modulo each extent, so a lane's coordinates depend on `lane % 8` alone
-// and lanes 8-15 repeat lanes 0-7.
+// 1x2, one whole row, and there are 8: eight groups of 2, lane `l` paired
+// with lane `l + 8`. Here the period is 8 as well, because the lane id is
+// delinearized modulo each extent, so a lane's coordinates depend on
+// `lane % 8` alone.
//
// The two spellings are treated alike throughout.
//
>From 93e8094ba812e32948b5fc844d7bcc40358d4eee Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 04:27:14 +0000
Subject: [PATCH 21/46] Name the two lane periods apart and define donorDelta
before donor
Five wording problems, all found by reading the section cold:
lanePeriod is always the target's, but Terms defined lane period for a
layout generically and every later use meant the target silently. Rename
the variable to targetLanePeriod and say in Terms that both layouts have
a period, that the target's is the one the scheme steps by, and that the
input's is never computed because computeOwnedCoords has already folded
it in. The input's period is what makes several donorDelta values reach
the element, so it was invisible while doing real work.
donorDelta was defined inside the donor entry, which inverted the
dependency: donorDelta is chosen from a candidate set fixed by the target
period, and the donor follows from it. Give it its own entry, state that
it is non-negative so data always moves from higher lane ids to lower,
and leave donor as the consequence.
Add reader for the lane a fragment is assembled for, so target always
means the target layout, and use it to disambiguate
donorSlot = (slot + donorDelta) % targetLanePeriod = slot
which previously wrote slot for both the donor's and the reader's slot
and so read as a definition rather than the check it is. Say outright
that gpu.shuffle is not a remote read -- the donor extracts at its own
index before any shuffle and never learns who asked -- since that is what
makes the check necessary.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 142 +++++++++++-------
1 file changed, 84 insertions(+), 58 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 46070a60df0f4..f9d15c09c4239 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1789,23 +1789,38 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// -- when the product is already `S`
// lane period the same product with the sliced dimensions still
// counted, so slicing does not reduce it. Lanes this far
-// apart hold the same fragment. `getLanePeriod`,
-// `lanePeriod` below
-// slot `lane % lanePeriod`, a lane's index within one period
+// apart hold the same fragment. `getLanePeriod`. Both
+// layouts have one and they are not interchangeable: the
+// target's, `targetLanePeriod` below, is what the scheme
+// steps and groups lanes by, while the input's is never
+// computed -- it is already folded into the ownership
+// table
+// slot `lane % targetLanePeriod`, a lane's index within one
+// target period. A target-side notion: lanes sharing a
+// slot are assigned the same target fragment
// ownership table every lane's fragment, tabulated as element coordinates,
// one table per side. `computeOwnedCoords`, `inputOwned`
// and `targetOwned` below
//
// These name the moving parts of the scheme:
//
-// donor the lane a value is read from, `slot + donorDelta`.
-// `donorDelta` is one multiple of `lanePeriod`, shared by
-// every slot, so it is the single choice the scheme makes
-// per element
+// reader the lane a fragment is being assembled for. Everything
+// below depends on it only through its slot
+// donorDelta the one choice the scheme makes per element of the
+// target fragment: a lane offset from the reader's slot to
+// the lane supplying that element. Non-negative and a
+// multiple of `targetLanePeriod`, so there are
+// `S / targetLanePeriod` candidates and data always moves
+// from a higher lane id to a lower one
+// donor where that offset lands, `slot + donorDelta`. Derived
+// from `donorDelta`, not the other way round. More than
+// one offset may reach a lane holding the element, since
+// the input's own lane period replicates it; the smallest
+// is taken
// element source where one element of the target fragment comes from: a
// donor, and the index to extract from that donor's
// fragment. One per element, and it needs a shuffle
-// exactly when the donor is not the lane itself.
+// exactly when the donor is not the reader itself.
// `ElementSource` below
//
// A broadcast group and the lane period are that one product taken either side
@@ -1899,20 +1914,29 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Why donors sit a whole period apart
// -----------------------------------
//
-// `gpu.shuffle idx` is the only way to move data between lanes and it exchanges
-// one value per lane, so every lane runs the same extract. A donor can
-// therefore only serve lanes of its own slot, which is why `donorDelta` is a
-// multiple of `lanePeriod`. It also means the donor extracts the right thing
-// without knowing who asked: its own slot is
-// `(slot + donorDelta) % lanePeriod`, which is `slot`, so it computes the same
-// index.
+// `gpu.shuffle idx` is the only way to move data between lanes, and it is not a
+// remote read: every lane extracts one element from its own fragment and
+// contributes it, then each lane receives whatever the lane it names
+// contributed. A donor has therefore already extracted, at its own index,
+// before any shuffle happens, and never learns who asked.
//
-// Case 2 has two element sources, one per element of its 1x2 fragment. Slot 3
-// wants (3, 0) and (3, 1):
+// That works only if the index the donor used is the one the reader wanted. The
+// index is a function of the slot, so the two have to share one:
+//
+// donorSlot = (slot + donorDelta) % targetLanePeriod = slot
+//
+// which holds exactly because `donorDelta` is a multiple of
+// `targetLanePeriod`. That is where the constraint comes from.
+//
+// Case 2 has two element sources, one per element of its 1x2 fragment. Reader
+// slot 3 wants (3, 0) and (3, 1):
//
// element 0 -> donorDelta = 0, donor = 3, index = 3 already local
// element 1 -> donorDelta = 8, donor = 11, index = 3 shuffled
//
+// Both lanes extract at index 3, from different fragments: the input gives lane
+// 3 column 0 and lane 11 column 1, so one index yields the two elements wanted.
+//
// What it declines
// ----------------
//
@@ -1968,11 +1992,11 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// The cases, `p` being the position in the target fragment:
//
-// case lanePeriod index(slot) stride dimStride dimExtent offset
-// 1 4 2*slot + off 2 1 4 0,1,8,9
-// 2 8 slot 1 1 8 0
-// 3 16 slot/8 + 2*p 1 8 2 2*p
-// 4 2 p % 2 0 1 1 p % 2
+// case period index(slot) stride dimStride dimExtent offset
+// 1 4 2*slot + off 2 1 4 0,1,8,9
+// 2 8 slot 1 1 8 0
+// 3 16 slot/8 + 2*p 1 8 2 2*p
+// 4 2 p % 2 0 1 1 p % 2
//
// Where the constants come from: the `2`s multiplying `slot` are the size of
// the input fragment's innermost dimension, which in cases 1 and 3 is the whole
@@ -1992,7 +2016,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// tabulate inputOwned[lane] and targetOwned[lane] // element coordinates
// for each element `pos` of the target fragment:
-// for donorDelta = 0, lanePeriod, 2 * lanePeriod, ... // smallest first
+// for donorDelta = 0, targetLanePeriod, 2 * targetLanePeriod, ...
+// // smallest first
// for each slot:
// needed[slot] = index of the element targetOwned[slot][pos] in the
// fragment of its donor, lane `slot + donorDelta`
@@ -2009,13 +2034,13 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// values of `donorDelta` may work and all give the same value, so the smallest
// is taken, which is what leaves a fully broadcast input shuffle-free.
//
-// Only the first `lanePeriod` lanes are required to end up correct -- each is
-// the one lane of its slot the target layout distributes to, and the layout
-// assigns the rest no elements. Lanes sharing a slot run the same arithmetic on
-// their own fragment: when a shuffle is emitted they all receive the donor's
-// value and are equally correct, and when it is dropped their local extract can
-// land on a different element. Dropping one therefore relies on nobody reading
-// them.
+// Only the first `targetLanePeriod` lanes are required to end up correct --
+// each is the one lane of its slot the target layout distributes to, and the
+// layout assigns the rest no elements. Lanes sharing a slot run the same
+// arithmetic on their own fragment: when a shuffle is emitted they all receive
+// the donor's value and are equally correct, and when it is dropped their local
+// extract can land on a different element. Dropping one therefore relies on
+// nobody reading them.
//
// An index is derived from the tables and then verified against them, so a
// layout change that is not of this form is reported as a match failure rather
@@ -2097,18 +2122,18 @@ static bool haveSameDistributionUnit(xegpu::DistributeLayoutAttr inputLayout,
/// code, so a lane can only be served by a donor that agrees with it on which
/// element to extract -- one whose lane id differs by a whole number of lane
/// periods. The lowering therefore searches donors at `donorDelta = 0,
-/// lanePeriod, 2 * lanePeriod, ...` up to the subgroup size, reading lane `slot
-/// + donorDelta`, where `slot` is `lane % lanePeriod` and so below the period.
-/// If the period did not divide the subgroup size that sum could name a lane
-/// past the last one: with a period of 3 on 16 lanes, `donorDelta` reaches 15
-/// and a slot reaches 2, naming lane 17.
+/// targetLanePeriod, 2 * targetLanePeriod, ...` up to the subgroup size,
+/// reading lane `slot + donorDelta`, where `slot` is `lane % targetLanePeriod`
+/// and so below the period. If the period did not divide the subgroup size that
+/// sum could name a lane past the last one: with a period of 3 on 16 lanes,
+/// `donorDelta` reaches 15 and a slot reaches 2, naming lane 17.
///
/// The input has no counterpart: it is only ever indexed, never used to derive
/// a donor, so its own period is unconstrained.
static bool isSupportedLanePeriod(xegpu::DistributeLayoutAttr targetLayout,
int64_t subgroupSize) {
- int64_t lanePeriod = getLanePeriod(targetLayout);
- return lanePeriod > 0 && subgroupSize % lanePeriod == 0;
+ int64_t targetLanePeriod = getLanePeriod(targetLayout);
+ return targetLanePeriod > 0 && subgroupSize % targetLanePeriod == 0;
}
/// Matches the layout change `redistributeBroadcastedValue` implements.
@@ -2164,7 +2189,7 @@ computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
/// Which element of the donor's fragment a lane extracts -- the donor being the
/// lane it reads from, `slot + donorDelta` -- as a function of its `slot`
-/// (`lane % lanePeriod`).
+/// (`lane % targetLanePeriod`).
///
/// `at()` is restricted to one shape: the lane's coordinate along a single
/// dimension of the target lane_layout, scaled and offset into the fragment.
@@ -2217,23 +2242,24 @@ struct ElementSource {
/// Fits `needed[slot]`, the fragment index each slot has to read, to the
/// restricted form `FragmentIndex` can emit. A slot is a lane reduced modulo
-/// the lane period, so there are `lanePeriod` of them. Fails when no index in
-/// that form reproduces the whole table.
+/// the target's lane period, so there are `targetLanePeriod` of them. Fails
+/// when no index in that form reproduces the whole table.
///
/// For `needed = [0, 0, 2, 2]`: the first change is at slot 2, so pairs of
/// slots share a coordinate (`dimStride` 2) and consecutive coordinates are 2
/// elements apart (`stride` 2); the table reaches 2, one step, so the
/// coordinate takes 2 values (`dimExtent` 2); and slot 0 reads 0 (`offset` 0).
static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
- int64_t lanePeriod) {
+ int64_t targetLanePeriod) {
// Slots sharing an index are adjacent, so the first change in `needed` gives
// both how many of them there are and how far apart their elements sit.
int64_t offset = needed[0];
int64_t dimStride = 1;
- while (dimStride < lanePeriod && needed[dimStride] == offset)
+ while (dimStride < targetLanePeriod && needed[dimStride] == offset)
dimStride++;
- int64_t stride = dimStride < lanePeriod ? needed[dimStride] - offset : 0;
- if (stride < 0 || lanePeriod % dimStride != 0)
+ int64_t stride =
+ dimStride < targetLanePeriod ? needed[dimStride] - offset : 0;
+ if (stride < 0 || targetLanePeriod % dimStride != 0)
return failure();
// The index wraps once it has taken all its values, so how far `needed` gets
@@ -2253,7 +2279,7 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
// Derived from the table, so verify against the table before relying on it.
FragmentIndex index{stride, offset, dimStride, dimExtent};
- if (!llvm::all_of(llvm::seq<int64_t>(0, lanePeriod), [&](int64_t slot) {
+ if (!llvm::all_of(llvm::seq<int64_t>(0, targetLanePeriod), [&](int64_t slot) {
return needed[slot] == index.at(slot);
}))
return failure();
@@ -2271,7 +2297,7 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
static FailureOr<ElementSource>
deriveElementSource(const OwnedCoords &inputOwned,
const OwnedCoords &targetOwned, int64_t pos,
- int64_t subgroupSize, int64_t lanePeriod,
+ int64_t subgroupSize, int64_t targetLanePeriod,
std::optional<RedistributionLimit> &limit) {
// Index of the element with coordinate `coord` in the fragment `lane` holds,
// or -1 if it holds none.
@@ -2287,20 +2313,20 @@ deriveElementSource(const OwnedCoords &inputOwned,
// this fails.
bool sawCompleteDonor = false;
for (int64_t donorDelta = 0; donorDelta < subgroupSize;
- donorDelta += lanePeriod) {
+ donorDelta += targetLanePeriod) {
// Where in its fragment each slot's donor keeps the element it needs.
SmallVector<int64_t> needed;
- for (int64_t slot = 0; slot < lanePeriod; slot++) {
+ for (int64_t slot = 0; slot < targetLanePeriod; slot++) {
int64_t index = findInFragment(slot + donorDelta, targetOwned[slot][pos]);
if (index < 0)
break;
needed.push_back(index);
}
- if (static_cast<int64_t>(needed.size()) != lanePeriod)
+ if (static_cast<int64_t>(needed.size()) != targetLanePeriod)
continue;
sawCompleteDonor = true;
- FailureOr<FragmentIndex> index = fitFragmentIndex(needed, lanePeriod);
+ FailureOr<FragmentIndex> index = fitFragmentIndex(needed, targetLanePeriod);
if (failed(index))
continue;
@@ -2327,7 +2353,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
return failure();
int64_t srcNumElems = srcTy.getNumElements();
int64_t resNumElems = resTy.getNumElements();
- int64_t lanePeriod = getLanePeriod(targetLayout);
+ int64_t targetLanePeriod = getLanePeriod(targetLayout);
if (!isSupportedLanePeriod(targetLayout, subgroupSize)) {
limit = RedistributionLimit::LanePeriodNotDivisor;
return failure();
@@ -2347,11 +2373,11 @@ static FailureOr<Value> redistributeBroadcastedValue(
if (failed(inputOwned) || failed(targetOwned))
return failure();
// Require the layouts to distribute what the vector types account for. The
- // lanes past the first `lanePeriod` are replicas by construction: both
+ // lanes past the first `targetLanePeriod` are replicas by construction: both
// computeStaticDistributedCoords implementations delinearize the lane id
- // modulo the lane_layout extents, whose product is `lanePeriod`.
+ // modulo the lane_layout extents, whose product is `targetLanePeriod`.
for (int64_t lane = 0; lane < subgroupSize; lane++) {
- assert((*targetOwned)[lane] == (*targetOwned)[lane % lanePeriod] &&
+ assert((*targetOwned)[lane] == (*targetOwned)[lane % targetLanePeriod] &&
"target ownership is not periodic in the lane period");
if (static_cast<int64_t>((*inputOwned)[lane].size()) != srcNumElems ||
static_cast<int64_t>((*targetOwned)[lane].size()) != resNumElems) {
@@ -2365,7 +2391,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
SmallVector<ElementSource> sources;
for (int64_t pos = 0; pos < resNumElems; pos++) {
FailureOr<ElementSource> source = deriveElementSource(
- *inputOwned, *targetOwned, pos, subgroupSize, lanePeriod, limit);
+ *inputOwned, *targetOwned, pos, subgroupSize, targetLanePeriod, limit);
if (failed(source))
return failure();
sources.push_back(*source);
@@ -2387,10 +2413,10 @@ static FailureOr<Value> redistributeBroadcastedValue(
// The lane's slot, which every fragment index is a function of.
Value slot = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
/*upperBound=*/mlir::IntegerAttr());
- if (lanePeriod != subgroupSize)
+ if (targetLanePeriod != subgroupSize)
slot = arith::RemUIOp::create(
rewriter, loc, slot,
- arith::ConstantIndexOp::create(rewriter, loc, lanePeriod));
+ arith::ConstantIndexOp::create(rewriter, loc, targetLanePeriod));
Type i32Ty = rewriter.getI32Type();
auto flatResTy = VectorType::get({resNumElems}, shuffleTy);
@@ -2411,7 +2437,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
dimIndex = arith::DivUIOp::create(
rewriter, loc, dimIndex,
arith::ConstantIndexOp::create(rewriter, loc, index.dimStride));
- if (index.dimExtent < lanePeriod / index.dimStride)
+ if (index.dimExtent < targetLanePeriod / index.dimStride)
dimIndex = arith::RemUIOp::create(
rewriter, loc, dimIndex,
arith::ConstantIndexOp::create(rewriter, loc, index.dimExtent));
>From 368b88b83862e85e4508e758574f5707802a07a0 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 05:05:40 +0000
Subject: [PATCH 22/46] Make the Algorithm section a map into the code
The section stated the algorithm correctly but gave a reader nothing to
descend into: no step said which function carried it out, so after
reading it the implementation was still a group of functions that each
looked fine alone.
Tag every step with its function, and say up front that
redistributeBroadcastedValue is the enclosing one and that
SgToLaneConvertLayout only reaches it once the two narrower convert_layout
paths decline.
No functional change.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 16 ++++++++++++++++
1 file changed, 16 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index f9d15c09c4239..918940e666c83 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2014,7 +2014,18 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Algorithm
// ---------
//
+// Also the map into the code: a `[name]` tag gives the function carrying out
+// the step under it. `redistributeBroadcastedValue` is the enclosing one -- it
+// owns both tables and emits the result -- and `SgToLaneConvertLayout` reaches
+// it only after the two narrower `convert_layout` paths decline.
+//
+// [isBroadcastRedistribution]
+// gate the layout pair: `lane_data` equal, target period divides `S`
+//
+// [computeOwnedCoords] once per side
// tabulate inputOwned[lane] and targetOwned[lane] // element coordinates
+//
+// [deriveElementSource] once per element of the target fragment
// for each element `pos` of the target fragment:
// for donorDelta = 0, targetLanePeriod, 2 * targetLanePeriod, ...
// // smallest first
@@ -2022,10 +2033,15 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// needed[slot] = index of the element targetOwned[slot][pos] in the
// fragment of its donor, lane `slot + donorDelta`
// if that donor does not hold the element: give up on this donorDelta
+//
+// [fitFragmentIndex]
// look for one stride/offset/dimStride/dimExtent whose index reproduces
// every needed[slot]; if there is none: next donorDelta
+//
// accept (donorDelta, index) and stop
// if no donorDelta fits: fail, the layout change is not of this form
+//
+// [redistributeBroadcastedValue, after the loop]
// emit, per element source: an extract at index(slot), a `gpu.shuffle idx`
// from `slot + donorDelta` unless that is the lane itself, and an insert.
//
>From 1963e47ba786b85b91436c4106ca974d32610ca8 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 15:36:33 +0000
Subject: [PATCH 23/46] Show what the emitted sequence computes at runtime
The Algorithm section describes compile-time steps, and left a reader to
work out which of them survive into the emitted code. Add the runtime
dataflow after the pseudocode: lane_id is the only input, and exactly two
values derived from it reach an op that needs one -- the extract position
and the shuffle source lane. The insert needs neither.
Names the shuffle operand donor at the point of use, matching the local
variable in redistributeBroadcastedValue rather than respelling the
arithmetic.
No functional change.
---
.../Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 8 ++++++++
1 file changed, 8 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 918940e666c83..8c3a0cb310664 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2045,6 +2045,14 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// emit, per element source: an extract at index(slot), a `gpu.shuffle idx`
// from `slot + donorDelta` unless that is the lane itself, and an insert.
//
+// At runtime `lane_id` is the only input; everything the search settles is a
+// constant by then. Two values reach an op that needs one:
+//
+// lane_id --> slot --+--> index(slot) --> extract position
+// +--> donor = slot + donorDelta --> shuffle source lane
+//
+// The insert needs neither, its position being the compile-time `pos`.
+//
// The shuffle is dropped exactly when `donorDelta` is 0, since the extract then
// lands on an element the lane owns. Replication makes that a choice: several
// values of `donorDelta` may work and all give the same value, so the smallest
>From c1049bf075409821b23e2edf84be4ad55ff87744 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 15:44:42 +0000
Subject: [PATCH 24/46] Describe the index as the reverse map it is, and name
it after the donor
The section led with the emitted vector.extract -- "the index the lane
extracts" -- which is the bottom of the abstraction, and called the
quantity index(slot). In an MLIR file that reads as a function returning
the builtin index type, and it never said which fragment the index is
into.
Recast the section around what the search actually produces, a map from
where an element is wanted back to where it already lives:
target (pos, slot) --> input (donor, donorIndex)
That also explains why the search runs once per position and why
ElementSource holds exactly two things.
Rename to match: DonorIndex, fitDonorIndex, getDonorIndex, and the field
ElementSource::donorIndex. donorIndex pairs with donor -- one says which
lane, the other where in that lane's fragment -- and no longer collides
with the index type.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 71 ++++++++++++-------
1 file changed, 44 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 8c3a0cb310664..bf2a87bbcdf5f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1976,12 +1976,26 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// element narrower than a byte, and FragmentSizeMismatch, an invariant check
// for a layout disagreeing with the vector type the caller derived from it.
//
-// The index expression
-// --------------------
+// The reverse map
+// ---------------
+//
+// For one position `pos` of the target fragment, the search produces a map from
+// where an element is wanted back to where it already lives:
+//
+// target (pos, slot) --> input (donor, donorIndex)
+//
+// donor = slot + donorDelta which lane holds it
+// donorIndex(slot) where in that lane's fragment
+//
+// Both are functions of `slot` alone, since every lane runs the same code.
+// `donorDelta` and the constants of `donorIndex` are fixed per `pos`, which is
+// why the search runs once per position and `ElementSource` holds exactly these
+// two things.
//
-// The index has to be computable from the lane id, so it is restricted to
+// `donor` costs one add. `donorIndex` is the half that needs a form, and it is
+// restricted to
//
-// index(slot) = stride * ((slot / dimStride) % dimExtent) + offset
+// donorIndex(slot) = stride * ((slot / dimStride) % dimExtent) + offset
//
// `(slot / dimStride) % dimExtent` is the lane's coordinate along one dimension
// of the target lane_layout: `dimExtent` is that dimension's size and
@@ -1992,11 +2006,11 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// The cases, `p` being the position in the target fragment:
//
-// case period index(slot) stride dimStride dimExtent offset
-// 1 4 2*slot + off 2 1 4 0,1,8,9
-// 2 8 slot 1 1 8 0
-// 3 16 slot/8 + 2*p 1 8 2 2*p
-// 4 2 p % 2 0 1 1 p % 2
+// case period donorIndex(slot) stride dimStride dimExtent offset
+// 1 4 2*slot + off 2 1 4 0,1,8,9
+// 2 8 slot 1 1 8 0
+// 3 16 slot/8 + 2*p 1 8 2 2*p
+// 4 2 p % 2 0 1 1 p % 2
//
// Where the constants come from: the `2`s multiplying `slot` are the size of
// the input fragment's innermost dimension, which in cases 1 and 3 is the whole
@@ -2034,21 +2048,22 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// fragment of its donor, lane `slot + donorDelta`
// if that donor does not hold the element: give up on this donorDelta
//
-// [fitFragmentIndex]
+// [fitDonorIndex]
// look for one stride/offset/dimStride/dimExtent whose index reproduces
// every needed[slot]; if there is none: next donorDelta
//
-// accept (donorDelta, index) and stop
+// accept (donorDelta, donorIndex) and stop
// if no donorDelta fits: fail, the layout change is not of this form
//
// [redistributeBroadcastedValue, after the loop]
-// emit, per element source: an extract at index(slot), a `gpu.shuffle idx`
-// from `slot + donorDelta` unless that is the lane itself, and an insert.
+// emit, per element source: an extract at donorIndex(slot), a
+// `gpu.shuffle idx` from the donor unless that is the reader itself, and an
+// insert
//
// At runtime `lane_id` is the only input; everything the search settles is a
// constant by then. Two values reach an op that needs one:
//
-// lane_id --> slot --+--> index(slot) --> extract position
+// lane_id --> slot --+--> donorIndex(slot) --> extract position
// +--> donor = slot + donorDelta --> shuffle source lane
//
// The insert needs neither, its position being the compile-time `pos`.
@@ -2226,7 +2241,7 @@ computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
/// row-major, and whose target gives lane `i` row `i`, element `p` of the
/// target fragment sits at `2 * slot + p`: `stride` 2, `offset` p, `dimStride`
/// 1, `dimExtent` 8.
-struct FragmentIndex {
+struct DonorIndex {
/// Distance in the fragment between consecutive coordinates along the
/// dimension being walked, or 0 when every slot reads the same element.
int64_t stride;
@@ -2251,10 +2266,12 @@ struct FragmentIndex {
}
};
-/// Where one element of the target fragment comes from.
+/// Where one element of the target fragment comes from: the two halves of the
+/// reverse map, `target (pos, slot) -> input (donor, donorIndex)`, for one
+/// `pos`.
struct ElementSource {
/// Which element of the donor's fragment to extract.
- FragmentIndex index;
+ DonorIndex donorIndex;
/// Lanes from the reader to the donor it reads through `gpu.shuffle idx`: the
/// donor is lane `slot + donorDelta`. Always a multiple of the lane period,
/// so that donor and reader agree on the index to extract.
@@ -2265,7 +2282,7 @@ struct ElementSource {
};
/// Fits `needed[slot]`, the fragment index each slot has to read, to the
-/// restricted form `FragmentIndex` can emit. A slot is a lane reduced modulo
+/// restricted form `DonorIndex` can emit. A slot is a lane reduced modulo
/// the target's lane period, so there are `targetLanePeriod` of them. Fails
/// when no index in that form reproduces the whole table.
///
@@ -2273,8 +2290,8 @@ struct ElementSource {
/// slots share a coordinate (`dimStride` 2) and consecutive coordinates are 2
/// elements apart (`stride` 2); the table reaches 2, one step, so the
/// coordinate takes 2 values (`dimExtent` 2); and slot 0 reads 0 (`offset` 0).
-static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
- int64_t targetLanePeriod) {
+static FailureOr<DonorIndex> fitDonorIndex(ArrayRef<int64_t> needed,
+ int64_t targetLanePeriod) {
// Slots sharing an index are adjacent, so the first change in `needed` gives
// both how many of them there are and how far apart their elements sit.
int64_t offset = needed[0];
@@ -2302,7 +2319,7 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
}
// Derived from the table, so verify against the table before relying on it.
- FragmentIndex index{stride, offset, dimStride, dimExtent};
+ DonorIndex index{stride, offset, dimStride, dimExtent};
if (!llvm::all_of(llvm::seq<int64_t>(0, targetLanePeriod), [&](int64_t slot) {
return needed[slot] == index.at(slot);
}))
@@ -2317,7 +2334,7 @@ static FailureOr<FragmentIndex> fitFragmentIndex(ArrayRef<int64_t> needed,
/// preferred so that a lane already holding the element needs no shuffle.
///
/// Fails when no `donorDelta` puts the element within reach of every slot, or
-/// when the indices it would be read at are not of `FragmentIndex`'s form.
+/// when the indices it would be read at are not of `DonorIndex`'s form.
static FailureOr<ElementSource>
deriveElementSource(const OwnedCoords &inputOwned,
const OwnedCoords &targetOwned, int64_t pos,
@@ -2350,7 +2367,7 @@ deriveElementSource(const OwnedCoords &inputOwned,
continue;
sawCompleteDonor = true;
- FailureOr<FragmentIndex> index = fitFragmentIndex(needed, targetLanePeriod);
+ FailureOr<DonorIndex> index = fitDonorIndex(needed, targetLanePeriod);
if (failed(index))
continue;
@@ -2452,7 +2469,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
// The lane's index along the selected dimension, shared by sources over the
// same one. Either step is skipped when it is a no-op over `slot`.
llvm::DenseMap<std::pair<int64_t, int64_t>, Value> dimIndices;
- auto getDimIndex = [&](const FragmentIndex &index) -> Value {
+ auto getDimIndex = [&](const DonorIndex &index) -> Value {
Value &dimIndex = dimIndices[{index.dimStride, index.dimExtent}];
if (dimIndex)
return dimIndex;
@@ -2469,7 +2486,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
};
// The index the lane extracts.
- auto getFragmentIndex = [&](const FragmentIndex &index) -> OpFoldResult {
+ auto getDonorIndex = [&](const DonorIndex &index) -> OpFoldResult {
// Every slot reads the same element; no dimension is involved.
if (index.stride == 0)
return rewriter.getIndexAttr(index.offset);
@@ -2494,10 +2511,10 @@ static FailureOr<Value> redistributeBroadcastedValue(
extracted;
for (auto [pos, source] : llvm::enumerate(sources)) {
// Elements reading the same index of the fragment share the extract.
- Value &value = extracted[source.index.asTuple()];
+ Value &value = extracted[source.donorIndex.asTuple()];
if (!value)
value = vector::ExtractOp::create(rewriter, loc, fragment,
- getFragmentIndex(source.index));
+ getDonorIndex(source.donorIndex));
Value owned = value;
if (source.needsShuffle) {
>From 1fa00a63ee9ac61f493476c960d614e5fb996208 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 15:47:52 +0000
Subject: [PATCH 25/46] Use donorIndex for the locals and parameters holding
one
The type, the field and the notation were renamed, but every local and
parameter of that type was still called index, so the term stopped at the
declarations. Rename them, and the two loop variables holding a wanted
index to wanted, which is what they are -- an entry of the needed table,
not the fitted expression.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 58 ++++++++++---------
1 file changed, 30 insertions(+), 28 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index bf2a87bbcdf5f..d5fdf7a90518d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2310,8 +2310,8 @@ static FailureOr<DonorIndex> fitDonorIndex(ArrayRef<int64_t> needed,
// Every slot reads `offset`; no dimension is involved.
dimStride = 1;
} else {
- for (int64_t index : needed) {
- int64_t diff = index - offset;
+ for (int64_t wanted : needed) {
+ int64_t diff = wanted - offset;
if (diff < 0 || diff % stride != 0)
return failure();
dimExtent = std::max(dimExtent, diff / stride + 1);
@@ -2319,12 +2319,12 @@ static FailureOr<DonorIndex> fitDonorIndex(ArrayRef<int64_t> needed,
}
// Derived from the table, so verify against the table before relying on it.
- DonorIndex index{stride, offset, dimStride, dimExtent};
+ DonorIndex donorIndex{stride, offset, dimStride, dimExtent};
if (!llvm::all_of(llvm::seq<int64_t>(0, targetLanePeriod), [&](int64_t slot) {
- return needed[slot] == index.at(slot);
+ return needed[slot] == donorIndex.at(slot);
}))
return failure();
- return index;
+ return donorIndex;
}
/// Works out where element `pos` of the target fragment comes from: which donor
@@ -2358,23 +2358,25 @@ deriveElementSource(const OwnedCoords &inputOwned,
// Where in its fragment each slot's donor keeps the element it needs.
SmallVector<int64_t> needed;
for (int64_t slot = 0; slot < targetLanePeriod; slot++) {
- int64_t index = findInFragment(slot + donorDelta, targetOwned[slot][pos]);
- if (index < 0)
+ int64_t wanted =
+ findInFragment(slot + donorDelta, targetOwned[slot][pos]);
+ if (wanted < 0)
break;
- needed.push_back(index);
+ needed.push_back(wanted);
}
if (static_cast<int64_t>(needed.size()) != targetLanePeriod)
continue;
sawCompleteDonor = true;
- FailureOr<DonorIndex> index = fitDonorIndex(needed, targetLanePeriod);
- if (failed(index))
+ FailureOr<DonorIndex> donorIndex = fitDonorIndex(needed, targetLanePeriod);
+ if (failed(donorIndex))
continue;
// The donor is the lane itself exactly when `donorDelta` is 0:
// `needed[slot]` was located in the fragment of lane `slot + donorDelta`,
// so at 0 every slot's extract is local by construction.
- return ElementSource{*index, donorDelta, /*needsShuffle=*/donorDelta != 0};
+ return ElementSource{*donorIndex, donorDelta,
+ /*needsShuffle=*/donorDelta != 0};
}
limit = sawCompleteDonor ? RedistributionLimit::IndexNotLaneAffine
: RedistributionLimit::NoDonorDelta;
@@ -2466,39 +2468,39 @@ static FailureOr<Value> redistributeBroadcastedValue(
Value width;
Value slotI32;
- // The lane's index along the selected dimension, shared by sources over the
- // same one. Either step is skipped when it is a no-op over `slot`.
+ // The lane's coordinate along the selected dimension, shared by sources over
+ // the same one. Either step is skipped when it is a no-op over `slot`.
llvm::DenseMap<std::pair<int64_t, int64_t>, Value> dimIndices;
- auto getDimIndex = [&](const DonorIndex &index) -> Value {
- Value &dimIndex = dimIndices[{index.dimStride, index.dimExtent}];
+ auto getDimIndex = [&](const DonorIndex &donorIndex) -> Value {
+ Value &dimIndex = dimIndices[{donorIndex.dimStride, donorIndex.dimExtent}];
if (dimIndex)
return dimIndex;
dimIndex = slot;
- if (index.dimStride != 1)
+ if (donorIndex.dimStride != 1)
dimIndex = arith::DivUIOp::create(
rewriter, loc, dimIndex,
- arith::ConstantIndexOp::create(rewriter, loc, index.dimStride));
- if (index.dimExtent < targetLanePeriod / index.dimStride)
+ arith::ConstantIndexOp::create(rewriter, loc, donorIndex.dimStride));
+ if (donorIndex.dimExtent < targetLanePeriod / donorIndex.dimStride)
dimIndex = arith::RemUIOp::create(
rewriter, loc, dimIndex,
- arith::ConstantIndexOp::create(rewriter, loc, index.dimExtent));
+ arith::ConstantIndexOp::create(rewriter, loc, donorIndex.dimExtent));
return dimIndex;
};
- // The index the lane extracts.
- auto getDonorIndex = [&](const DonorIndex &index) -> OpFoldResult {
+ // `donorIndex(slot)`: which element of its own fragment the lane extracts.
+ auto getDonorIndex = [&](const DonorIndex &donorIndex) -> OpFoldResult {
// Every slot reads the same element; no dimension is involved.
- if (index.stride == 0)
- return rewriter.getIndexAttr(index.offset);
- Value value = getDimIndex(index);
- if (index.stride != 1)
+ if (donorIndex.stride == 0)
+ return rewriter.getIndexAttr(donorIndex.offset);
+ Value value = getDimIndex(donorIndex);
+ if (donorIndex.stride != 1)
value = arith::MulIOp::create(
rewriter, loc, value,
- arith::ConstantIndexOp::create(rewriter, loc, index.stride));
- if (index.offset != 0)
+ arith::ConstantIndexOp::create(rewriter, loc, donorIndex.stride));
+ if (donorIndex.offset != 0)
value = arith::AddIOp::create(
rewriter, loc, value,
- arith::ConstantIndexOp::create(rewriter, loc, index.offset));
+ arith::ConstantIndexOp::create(rewriter, loc, donorIndex.offset));
return value;
};
>From 234812270ef5b25caec474cf5053c12be2ee322d Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 15:51:20 +0000
Subject: [PATCH 26/46] Open the Algorithm section with the algorithm, not with
the tags
The section began by explaining its own notation, so the first thing a
reader met was a note about [name] tags rather than what the thing does.
Lead with the whole algorithm in four clauses -- agree the block, find an
equal block, make it travel by extract/shuffle/insert, and note that only
the search costs anything because the match must hold for every slot at
once -- then explain the tags.
Label the two search loops with what they are testing, so the candidate
being rejected is obvious at each level, and drop the dangling
"// smallest first" continuation that had been carrying that job.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 17 ++++++++++++-----
1 file changed, 12 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index d5fdf7a90518d..2f298deede418 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2028,10 +2028,16 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Algorithm
// ---------
//
-// Also the map into the code: a `[name]` tag gives the function carrying out
-// the step under it. `redistributeBroadcastedValue` is the enclosing one -- it
-// owns both tables and emits the result -- and `SgToLaneConvertLayout` reaches
-// it only after the two narrower `convert_layout` paths decline.
+// Check the two layouts agree on the block, then for each block the target
+// wants, find an equal block in some lane and have that lane's copy travel:
+// extract, shuffle, insert. Only the search costs anything, because the match
+// has to hold for every slot at once -- through one lane offset and one index
+// expression.
+//
+// A `[name]` tag below gives the function carrying out the step under it.
+// `redistributeBroadcastedValue` is the enclosing one -- it owns both tables
+// and emits the result -- and `SgToLaneConvertLayout` reaches it only after the
+// two narrower `convert_layout` paths decline.
//
// [isBroadcastRedistribution]
// gate the layout pair: `lane_data` equal, target period divides `S`
@@ -2041,8 +2047,9 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// [deriveElementSource] once per element of the target fragment
// for each element `pos` of the target fragment:
+// // candidate lane offsets, smallest first
// for donorDelta = 0, targetLanePeriod, 2 * targetLanePeriod, ...
-// // smallest first
+// // one offset, and below one expression, has to serve every slot
// for each slot:
// needed[slot] = index of the element targetOwned[slot][pos] in the
// fragment of its donor, lane `slot + donorDelta`
>From 29ead37a60c5200b07abf18a6b188f23f03e36e0 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 15:58:13 +0000
Subject: [PATCH 27/46] Trace the remaining magic numbers back to the layouts
Three sets of numbers in the walkthroughs were stated rather than
derived, so a reader could not check them:
Case 1's offsets 0,1,8,9. The paragraph explained the 2s multiplying slot
and said nothing about these, in the very case the text tells you to read
the expression off. They are where the 2x2 target fragment sits in the
whole 8x2 value, and the jump from 1 to 8 is four rows at 2 elements, the
4 being the target lane_layout extent.
The emits column. Provenance was "what the tests show", which is honest
but not checkable. State the three rules instead -- fragment size, one
shuffle per travelling element, one extract per distinct position read --
and work case 2 and case 4, the two whose counts are not simply the
fragment size.
The boundary paragraph's 2 and 8, now attributed to the two fragments
they come from.
No functional change.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 15 ++++++++++++++-
1 file changed, 14 insertions(+), 1 deletion(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 2f298deede418..ca84eb9cf4e06 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1893,6 +1893,13 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// 4 partial, short [2, 1] 2 2 extracts, 6 shuffles
// -- partial, short [8, 1, 2] sliced 16 declined
//
+// The counts follow from the target fragment. One shuffle per element that has
+// to travel, and one extract per distinct position read in the donor's
+// fragment, elements reading the same position sharing one. So case 2's 1x2
+// fragment is two elements read at the same position: one extract, and one of
+// the two travels. Case 4's 4x2 fragment is eight elements over two positions:
+// two extracts, six of the eight travelling.
+//
// Only the target varies within a category, so the test suite carries more than
// one of some -- @convert_layout_broadcast_all_lanes is case 1 with `[8, 1]` --
// but they behave alike and only one is described here. Case 1 is the one whose
@@ -1901,7 +1908,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// The last row is the boundary the two properties imply: with only `donorDelta`
// 0 available, a lane arriving with less than its target fragment cannot be
// completed, and the match fails with NoDonorDelta. There a lane arrives with 2
-// of the 8 elements it needs.
+// elements, `layout<[8, 1]>` giving it one row, and needs the sliced target's
+// 8, one whole column.
//
// Case 2 is the running example below: lanes 0-7 arrive holding all of column
// 0, lanes 8-15 all of column 1, and lane `i` leaves holding row `i`. Lane 3
@@ -2019,6 +2027,11 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// not depend on the slot at all: each target row arrives whole from one donor,
// so `stride` is 0 and `offset` alone selects the column.
//
+// Case 1's four offsets are where its 2x2 target fragment sits inside that
+// whole 8x2 value: lane `s` owns rows `s` and `s + 4`, columns 0 and 1, so
+// `2s+0, 2s+1, 2s+8, 2s+9`. The jump from 1 to 8 is four rows down at 2
+// elements each, the 4 being the extent of the target `lane_layout` [4, 1].
+//
// Read the expression off case 1, where the four quantities are 4, 2, 1 and 4.
// On case 2's `layout<[8, 1]>` target they collide: the lane period (8), its
// own `donorDelta` (8) and case 3's `dimStride` (8) all have the same value
>From 31ad37ae6ca1685309c126006ba6a52dc973058f Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 16:10:18 +0000
Subject: [PATCH 28/46] Cut the doc comments that the section header now
duplicates
getLanePeriod re-derived the delinearization and re-argued that the period
is not the count of lanes holding distinct data, with the same sliced
example; Terms now carries both, and states the period and the broadcast
group as one product taken either side of slicing, which is clearer than
what was here. Keep the local statement of what is returned and a pointer
for the distinction.
isSupportedLanePeriod re-derived the whole donor argument, which is now
its own section. Keep the divisibility consequence and the lane 17
counterexample, which are what this predicate is about.
DonorIndex re-explained the restricted form and why an arbitrary index
would need a table and a gather, both now in The reverse map, and repeated
the stride 0 case that its own field already documents. Keep the
definition and the worked instance, which is a concrete case the header's
table does not cover.
31 doc lines for 12 lines of code become 19. No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 45 ++++++-------------
1 file changed, 13 insertions(+), 32 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index ca84eb9cf4e06..2903f9e79534a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2108,21 +2108,14 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//===----------------------------------------------------------------------===//
/// The lane period of `layout`: the modulus after which its assignment of
-/// elements to lanes repeats, so that lane `l` and lane `l + getLanePeriod()`
-/// hold the same *fragment* -- the set of elements one lane owns.
-///
-/// It repeats because `computeStaticDistributedCoords` delinearizes the lane id
-/// into digits, one per lane_layout dimension, taking each modulo that
-/// dimension's extent. Nothing a layout derives from the lane id can therefore
-/// depend on it beyond the product of those extents, which is what this
-/// returns. A subgroup larger than the product simply repeats.
+/// elements to lanes repeats, so lane `l` and lane `l + getLanePeriod()` hold
+/// the same fragment. It is `product(lane_layout)` because
+/// `computeStaticDistributedCoords` delinearizes the lane id modulo each
+/// extent, so nothing a layout derives from the lane id depends on it beyond
+/// that product.
///
-/// This is not the number of lanes holding distinct data. A sliced dimension
-/// contributes its lanes to the period even though they replicate one fragment:
-/// `slice<layout<lane_layout = [8, 1, 2]>, dims = [0]>` has period 16 and hands
-/// out 2 distinct fragments. Nothing here needs the latter count -- the
-/// ownership tables are consulted per lane -- and the period is what the lane
-/// id has to be reduced by.
+/// Sliced dimensions are counted, so this is not the number of lanes holding
+/// distinct data -- see `broadcast group` in the section header.
static int64_t getLanePeriod(xegpu::DistributeLayoutAttr layout) {
// flatten() coalesces nested slices, so the parent it leaves is always a
// plain LayoutAttr carrying the lane_layout of the whole subgroup.
@@ -2176,16 +2169,10 @@ static bool haveSameDistributionUnit(xegpu::DistributeLayoutAttr inputLayout,
/// Target-side rule: the target's lane period has to divide the subgroup size.
///
-/// A value is moved between lanes by having every lane extract one element and
-/// `gpu.shuffle idx` it from a chosen *donor* lane. All lanes run the same
-/// code, so a lane can only be served by a donor that agrees with it on which
-/// element to extract -- one whose lane id differs by a whole number of lane
-/// periods. The lowering therefore searches donors at `donorDelta = 0,
-/// targetLanePeriod, 2 * targetLanePeriod, ...` up to the subgroup size,
-/// reading lane `slot + donorDelta`, where `slot` is `lane % targetLanePeriod`
-/// and so below the period. If the period did not divide the subgroup size that
-/// sum could name a lane past the last one: with a period of 3 on 16 lanes,
-/// `donorDelta` reaches 15 and a slot reaches 2, naming lane 17.
+/// `donorDelta` steps by the period up to the subgroup size and a slot is below
+/// the period, so otherwise `slot + donorDelta` could name a lane past the last
+/// one: with a period of 3 on 16 lanes, `donorDelta` reaches 15 and a slot
+/// reaches 2, naming lane 17.
///
/// The input has no counterpart: it is only ever indexed, never used to derive
/// a donor, so its own period is unconstrained.
@@ -2248,14 +2235,8 @@ computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
/// Which element of the donor's fragment a lane extracts -- the donor being the
/// lane it reads from, `slot + donorDelta` -- as a function of its `slot`
-/// (`lane % targetLanePeriod`).
-///
-/// `at()` is restricted to one shape: the lane's coordinate along a single
-/// dimension of the target lane_layout, scaled and offset into the fragment.
-/// That keeps the index a couple of `arith` ops on the lane id; an arbitrary
-/// index per slot would need a materialized table and a gather. A `stride` of 0
-/// is the degenerate case where no dimension is walked at all and every slot
-/// reads `offset`.
+/// (`lane % targetLanePeriod`). `at()` is the restricted form the section
+/// header derives; the fields below are its four constants.
///
/// For a `vector<8x2>` whose input layout gives every lane all 16 elements,
/// row-major, and whose target gives lane `i` row `i`, element `p` of the
>From 67d7a91d9bb810f820adbd7a4018ffbc84309468 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 16:12:29 +0000
Subject: [PATCH 29/46] Stop explaining the lane-id delinearization three times
Terms derives it once, on the short-lane_layout spelling of replication.
getLanePeriod re-derived it to justify returning a product, and
isBroadcastRedistribution re-derived it again to justify accepting a short
layout. Both are leaf predicates whose bodies are two to four obvious
lines, so the mechanism was carrying no weight there -- state the
consequence and leave the derivation to Terms.
Also drop haveSameDistributionUnit's restatement of what a distribution
unit is, which Terms defines.
The one remaining copy outside Terms justifies the periodicity assertion
in redistributeBroadcastedValue, where a reader meeting the assert wants
it in place.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 22 +++++++------------
1 file changed, 8 insertions(+), 14 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 2903f9e79534a..6f7ad7cee3157 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2107,12 +2107,9 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
//===----------------------------------------------------------------------===//
-/// The lane period of `layout`: the modulus after which its assignment of
-/// elements to lanes repeats, so lane `l` and lane `l + getLanePeriod()` hold
-/// the same fragment. It is `product(lane_layout)` because
-/// `computeStaticDistributedCoords` delinearizes the lane id modulo each
-/// extent, so nothing a layout derives from the lane id depends on it beyond
-/// that product.
+/// The lane period of `layout`: `product(lane_layout)`, the modulus after which
+/// its assignment of elements to lanes repeats, so lane `l` and lane
+/// `l + getLanePeriod()` hold the same fragment.
///
/// Sliced dimensions are counted, so this is not the number of lanes holding
/// distinct data -- see `broadcast group` in the section header.
@@ -2159,8 +2156,7 @@ static StringRef describe(RedistributionLimit limit) {
}
/// Relational rule: both layouts have to hand the lanes the same distribution
-/// unit -- the `lane_data`-sized block of elements a lane owns as a whole -- so
-/// that only the assignment of those blocks to lanes changes.
+/// unit, so that only the assignment of those blocks to lanes changes.
static bool haveSameDistributionUnit(xegpu::DistributeLayoutAttr inputLayout,
xegpu::DistributeLayoutAttr targetLayout) {
return inputLayout.getEffectiveLaneDataAsInt() ==
@@ -2184,12 +2180,10 @@ static bool isSupportedLanePeriod(xegpu::DistributeLayoutAttr targetLayout,
/// Matches the layout change `redistributeBroadcastedValue` implements.
///
-/// Neither layout has to lay out the whole subgroup. A layout whose lane_layout
-/// covers fewer lanes replicates its fragments over the rest, because
-/// `computeStaticDistributedCoords` delinearizes the lane id modulo each
-/// lane_layout extent; the fragment of every lane is therefore defined either
-/// way. `redistributeBroadcastedValue` verifies the assignment element-wise, so
-/// this only has to be permissive enough not to reject what it can lower.
+/// Neither layout has to lay out the whole subgroup: a short `lane_layout` is
+/// already a replicating one, so every lane's fragment is defined either way.
+/// `redistributeBroadcastedValue` verifies the assignment element-wise, so this
+/// only has to be permissive enough not to reject what it can lower.
///
/// That the layouts differ is not checked: they are already known to be
/// incompatible here, and the difference may equally be in `order`.
>From 26015317defeb7294b5fa6bcb42194a3d6c5cb37 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 19:03:02 +0000
Subject: [PATCH 30/46] Checkpoint: top-down review pass over the section
header
Joint editing pass, roughly the first third of the header. Not final.
Terms is now one list of introductions rather than two groups of mixed
definitions and rationale. element is defined and distribution unit is
gone -- the latter is established XeGPU vocabulary for lane_layout x
lane_data, so using it for the lane_data block alone contradicted the
dialect and the same file 1300 lines earlier. With lane_data forced to
ones an element is the unit, so no coined term is needed at all.
reader lane and donor lane replace reader and donor, since both are lanes
and reading them as data was the reported failure mode. reader lane now
precedes slot, the raw lane before the many-to-one collapse. element
source is a tuple of three matching ElementSource's fields, derived at
compile time.
Dropped twenty lines that taught how sliced and short lane_layouts work.
Reviewers and maintainers know how layouts work; only the pattern's own
consequences need stating. The input/target asymmetry paragraph collapsed
to one sentence at the head of the declined list, where the question
arises, and bounded to the up-front gate -- the input can still cause
NoDonorDelta once the tables exist.
Vocabulary sweeps: dimExtent to dimSize and extent out of prose, since
rank already names the count so size is unambiguous; travel to cross-lane
shuffle, which also disambiguates against the within-lane vector.shuffle
in the neighbouring patterns; the arrive/arrives metaphor to holds.
No functional change.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 206 +++++++-----------
1 file changed, 79 insertions(+), 127 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 6f7ad7cee3157..fecc80e61d568 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1765,110 +1765,59 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Broadcast redistribution
//===----------------------------------------------------------------------===//
//
-// Lowers a `convert_layout` that only relocates data between lanes: the two
-// layouts agree on `lane_data`, and differ only in which lane holds which part
-// of the value. Every element already exists in some lane, so nothing is
-// recomputed. The scale operands of scaled matrix multiplication need this,
-// being produced broadcasted and consumed distributed.
+// Lowers a `convert_layout` that relocates data between lanes: the two layouts
+// must agree on `lane_data`, and differ only in which lane holds which part
+// of the value.
//
// Terms
// -----
//
-// These follow from the layouts, for a value of shape `Sh` on a subgroup of `S`
-// lanes. `#xegpu.` prefixes are dropped for width throughout.
-//
-// distribution unit what a lane owns as one block: `lane_data`
-// fragment everything one lane owns. Slicing drops dimensions from
-// `lane_layout`; over what is left a lane owns
-// `Sh[i] / lane_layout[i]` per dimension, flattened
-// row-major
-// broadcast group the lanes that hold one fragment. `product(lane_layout)`
-// taken *after* slicing counts the distinct fragments, so
-// each group is `S / that` lanes: all `S` of them when
-// every dimension is sliced away, and 1 -- no replication
-// -- when the product is already `S`
-// lane period the same product with the sliced dimensions still
-// counted, so slicing does not reduce it. Lanes this far
-// apart hold the same fragment. `getLanePeriod`. Both
-// layouts have one and they are not interchangeable: the
-// target's, `targetLanePeriod` below, is what the scheme
-// steps and groups lanes by, while the input's is never
-// computed -- it is already folded into the ownership
-// table
-// slot `lane % targetLanePeriod`, a lane's index within one
-// target period. A target-side notion: lanes sharing a
-// slot are assigned the same target fragment
-// ownership table every lane's fragment, tabulated as element coordinates,
-// one table per side. `computeOwnedCoords`, `inputOwned`
-// and `targetOwned` below
-//
-// These name the moving parts of the scheme:
-//
-// reader the lane a fragment is being assembled for. Everything
+// element scalar of some element type
+// fragment everything one lane owns. Consists of elements.
+// broadcast group the lanes that hold the same fragment.
+// lane period how many lanes the layout lays out,
+// `product(lane_layout)`. A larger subgroup repeats: lane
+// `l` and lane `l + period` get the same fragment
+// reader lane the lane a fragment is being assembled for. Everything
// below depends on it only through its slot
-// donorDelta the one choice the scheme makes per element of the
-// target fragment: a lane offset from the reader's slot to
-// the lane supplying that element. Non-negative and a
-// multiple of `targetLanePeriod`, so there are
-// `S / targetLanePeriod` candidates and data always moves
-// from a higher lane id to a lower one
-// donor where that offset lands, `slot + donorDelta`. Derived
-// from `donorDelta`, not the other way round. More than
-// one offset may reach a lane holding the element, since
-// the input's own lane period replicates it; the smallest
-// is taken
-// element source where one element of the target fragment comes from: a
-// donor, and the index to extract from that donor's
-// fragment. One per element, and it needs a shuffle
-// exactly when the donor is not the reader itself.
-// `ElementSource` below
-//
-// A broadcast group and the lane period are that one product taken either side
-// of slicing, so a layout can have a lane period of `S` and still replicate:
-// the sliced dimensions raise the period without adding fragments. That is one
-// of the two ways a layout spells replication. Both, on a `vector<8x2>` over 16
-// lanes:
-//
-// slice<layout<[8, 1, 2], order = [0, 2, 1]>, dims = [0]> ("sliced")
-// slicing leaves `lane_layout` [1, 2], so a fragment is 8x1, one whole
-// column, and there are 1 * 2 = 2 of them: two groups of 8, lanes 0-7 and
-// lanes 8-15. The lane period still counts the sliced dimension and is
-// 8 * 1 * 2 = 16, well above the group size.
-//
-// layout<[8, 1]> ("short")
-// a `lane_layout` naming fewer lanes than the subgroup has. A fragment is
-// 1x2, one whole row, and there are 8: eight groups of 2, lane `l` paired
-// with lane `l + 8`. Here the period is 8 as well, because the lane id is
-// delinearized modulo each extent, so a lane's coordinates depend on
-// `lane % 8` alone.
-//
-// The two spellings are treated alike throughout.
-//
-// The two sides are not symmetric in what they mean:
-//
-// input supplies the values. Any lane may be a donor, so every lane's
-// fragment has to be known -- which it is, for any lane_layout.
-// target states the obligation. Only the lanes it distributes to have to end
-// up holding the fragment it assigns them; the rest are left
-// unspecified.
+// slot `reader lane % target lane period`, a lane's index
+// within one target period. A target-side notion: lanes
+// sharing a slot are assigned the same target fragment
+// ownership table tabulated "element of a fragment" to coordinate mapping
+// with respect to undistributed value. one table per input
+// and target
+// donorDelta a lane offset from the reader lane to
+// the lane supplying an element. Can have multiple
+// candidates.
+// donor lane supplying lane, computed as `slot + donorDelta`. First
+// fit is taken by enumerating through the donorDeltas.
+// element source for one element of the target fragment, a tuple: which
+// donor lane supplies it, where in that lane's fragment
+// it sits, and whether a shuffle is needed. One is
+// derived per element position, at compile time.
//
// What it handles
// ---------------
//
-// Two properties decide it, one per side. Examples stay on a `vector<8x2>` over
-// 16 lanes; the case numbers are used throughout the rest of this comment.
+// Two properties shape the outcome: the input's replication sets how many
+// elements need a cross-lane shuffle -- `gpu.shuffle idx`, as opposed to the
+// within-lane `vector.shuffle` the neighbouring patterns use -- and the
+// target's lane period sets where they may be shuffled from. Only the second
+// can decline the conversion. Examples stay on a `vector<8x2>` over 16 lanes;
+// the case numbers are used throughout the rest of this comment.
//
-// The input's replication decides what a lane arrives with, and so what has to
-// travel:
+// The input's replication decides what each lane already holds, and so how many
+// shuffles are needed:
//
-// fully broadcast every lane already holds the whole value, so nothing
-// travels and the result is extracts alone. Spelled by
-// slicing away every distributed dimension:
+// fully broadcast every lane already holds the whole value, so the
+// result is extracts alone. Spelled by slicing away
+// every distributed dimension:
// slice<layout<[1, 1, 16]>, dims = [2]>
//
// partially broadcast each broadcast group holds part of it, so a lane has
-// to be given whatever its group lacks. Either spelling
-// of replication from Terms above.
+// to be given whatever its group lacks. Spelled either
+// by slicing a dimension away or by a `lane_layout`
+// shorter than the subgroup.
//
// The target's lane period decides where a lane may be given that from. Donors
// sit at multiples of the period, so a period `P` over `S` lanes leaves `S / P`
@@ -1884,7 +1833,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// One example per combination, and what the tests show it emitting. The target
// is given as its `lane_layout`, the period following from it; the input as its
-// category, `sliced` and `short` being the two spellings named in Terms:
+// category, `sliced` meaning a sliced dimension and `short` a `lane_layout`
+// shorter than the subgroup:
//
// case input target period emits
// 1 fully broadcast [4, 1] 4 4 extracts
@@ -1893,12 +1843,12 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// 4 partial, short [2, 1] 2 2 extracts, 6 shuffles
// -- partial, short [8, 1, 2] sliced 16 declined
//
-// The counts follow from the target fragment. One shuffle per element that has
-// to travel, and one extract per distinct position read in the donor's
-// fragment, elements reading the same position sharing one. So case 2's 1x2
-// fragment is two elements read at the same position: one extract, and one of
-// the two travels. Case 4's 4x2 fragment is eight elements over two positions:
-// two extracts, six of the eight travelling.
+// The counts follow from the target fragment. One shuffle per element the
+// reader lane does not already hold, and one extract per distinct position read
+// in the donor lane's fragment, elements reading the same position sharing one.
+// So case 2's 1x2 fragment is two elements read at the same position: one
+// extract, and one of the two shuffled. Case 4's 4x2 fragment is eight elements
+// over two positions: two extracts, six of the eight shuffled.
//
// Only the target varies within a category, so the test suite carries more than
// one of some -- @convert_layout_broadcast_all_lanes is case 1 with `[8, 1]` --
@@ -1906,12 +1856,12 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// constants are all distinct, which the index expression below is read off.
//
// The last row is the boundary the two properties imply: with only `donorDelta`
-// 0 available, a lane arriving with less than its target fragment cannot be
-// completed, and the match fails with NoDonorDelta. There a lane arrives with 2
+// 0 available, a lane holding less than its target fragment cannot be
+// completed, and the match fails with NoDonorDelta. There a lane holds 2
// elements, `layout<[8, 1]>` giving it one row, and needs the sliced target's
// 8, one whole column.
//
-// Case 2 is the running example below: lanes 0-7 arrive holding all of column
+// Case 2 is the running example below: lanes 0-7 hold all of column
// 0, lanes 8-15 all of column 1, and lane `i` leaves holding row `i`. Lane 3
// wants (3, 0) and (3, 1); it holds (3, 0) already, while (3, 1) is in lane 11.
//
@@ -1954,7 +1904,9 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// shuffle per element from a donor a whole number of lane periods away. Another
// pattern or a more general scheme lowers any of them.
//
-// Decided by inspecting the layouts:
+// Decided by inspecting the layouts. None of these rejects the input on its
+// own -- the input is only read, so any `lane_layout` is admissible until the
+// tables show what it actually holds:
//
// they differ in `lane_data` LaneDataDiffers
//
@@ -2003,10 +1955,10 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// `donor` costs one add. `donorIndex` is the half that needs a form, and it is
// restricted to
//
-// donorIndex(slot) = stride * ((slot / dimStride) % dimExtent) + offset
+// donorIndex(slot) = stride * ((slot / dimStride) % dimSize) + offset
//
-// `(slot / dimStride) % dimExtent` is the lane's coordinate along one dimension
-// of the target lane_layout: `dimExtent` is that dimension's size and
+// `(slot / dimStride) % dimSize` is the lane's coordinate along one dimension
+// of the target lane_layout: `dimSize` is that dimension's size and
// `dimStride` how many consecutive slots share a coordinate. `stride` and
// `offset` place that coordinate in the donor's fragment. An arbitrary index
// per slot would need a materialized table and a gather; this costs a couple of
@@ -2014,7 +1966,7 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
//
// The cases, `p` being the position in the target fragment:
//
-// case period donorIndex(slot) stride dimStride dimExtent offset
+// case period donorIndex(slot) stride dimStride dimSize offset
// 1 4 2*slot + off 2 1 4 0,1,8,9
// 2 8 slot 1 1 8 0
// 3 16 slot/8 + 2*p 1 8 2 2*p
@@ -2024,13 +1976,13 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// the input fragment's innermost dimension, which in cases 1 and 3 is the whole
// 8x2 value, so stepping one row costs 2 elements. In case 2 the fragment is a
// single column, its row stride is 1, and no `2` appears. Case 4's index does
-// not depend on the slot at all: each target row arrives whole from one donor,
+// not depend on the slot at all: each target row comes whole from one donor,
// so `stride` is 0 and `offset` alone selects the column.
//
// Case 1's four offsets are where its 2x2 target fragment sits inside that
// whole 8x2 value: lane `s` owns rows `s` and `s + 4`, columns 0 and 1, so
// `2s+0, 2s+1, 2s+8, 2s+9`. The jump from 1 to 8 is four rows down at 2
-// elements each, the 4 being the extent of the target `lane_layout` [4, 1].
+// elements each, the 4 coming from the target `lane_layout` [4, 1].
//
// Read the expression off case 1, where the four quantities are 4, 2, 1 and 4.
// On case 2's `layout<[8, 1]>` target they collide: the lane period (8), its
@@ -2042,7 +1994,7 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// ---------
//
// Check the two layouts agree on the block, then for each block the target
-// wants, find an equal block in some lane and have that lane's copy travel:
+// wants, find an equal block in some lane and move that lane's copy to it:
// extract, shuffle, insert. Only the search costs anything, because the match
// has to hold for every slot at once -- through one lane offset and one index
// expression.
@@ -2069,7 +2021,7 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// if that donor does not hold the element: give up on this donorDelta
//
// [fitDonorIndex]
-// look for one stride/offset/dimStride/dimExtent whose index reproduces
+// look for one stride/offset/dimStride/dimSize whose index reproduces
// every needed[slot]; if there is none: next donorDelta
//
// accept (donorDelta, donorIndex) and stop
@@ -2235,7 +2187,7 @@ computeOwnedCoords(xegpu::DistributeLayoutAttr layout, ArrayRef<int64_t> shape,
/// For a `vector<8x2>` whose input layout gives every lane all 16 elements,
/// row-major, and whose target gives lane `i` row `i`, element `p` of the
/// target fragment sits at `2 * slot + p`: `stride` 2, `offset` p, `dimStride`
-/// 1, `dimExtent` 8.
+/// 1, `dimSize` 8.
struct DonorIndex {
/// Distance in the fragment between consecutive coordinates along the
/// dimension being walked, or 0 when every slot reads the same element.
@@ -2243,21 +2195,21 @@ struct DonorIndex {
/// Fragment index that coordinate 0 maps to, and the whole index when
/// `stride` is 0.
int64_t offset;
- /// Slots apart for the coordinate to advance by one.
+ /// Slots apart for the coordinate to advance by one, or 1 when no dimension
+ /// is walked.
int64_t dimStride;
- /// How many values the coordinate takes before wrapping. This is the extent
- /// of the target `lane_layout` dimension being walked, except when `stride`
- /// is 0 and no dimension is: then it is 1, as is `dimStride`.
- int64_t dimExtent;
+ /// How many values the coordinate takes before wrapping: the size of the
+ /// target `lane_layout` dimension being walked, or 1 when none is.
+ int64_t dimSize;
/// Compile-time twin of the arithmetic the emitter builds.
int64_t at(int64_t slot) const {
- return stride * ((slot / dimStride) % dimExtent) + offset;
+ return stride * ((slot / dimStride) % dimSize) + offset;
}
/// Key identifying the value this extracts, for reuse across elements.
std::tuple<int64_t, int64_t, int64_t, int64_t> asTuple() const {
- return {stride, offset, dimStride, dimExtent};
+ return {stride, offset, dimStride, dimSize};
}
};
@@ -2284,7 +2236,7 @@ struct ElementSource {
/// For `needed = [0, 0, 2, 2]`: the first change is at slot 2, so pairs of
/// slots share a coordinate (`dimStride` 2) and consecutive coordinates are 2
/// elements apart (`stride` 2); the table reaches 2, one step, so the
-/// coordinate takes 2 values (`dimExtent` 2); and slot 0 reads 0 (`offset` 0).
+/// coordinate takes 2 values (`dimSize` 2); and slot 0 reads 0 (`offset` 0).
static FailureOr<DonorIndex> fitDonorIndex(ArrayRef<int64_t> needed,
int64_t targetLanePeriod) {
// Slots sharing an index are adjacent, so the first change in `needed` gives
@@ -2299,8 +2251,8 @@ static FailureOr<DonorIndex> fitDonorIndex(ArrayRef<int64_t> needed,
return failure();
// The index wraps once it has taken all its values, so how far `needed` gets
- // from `offset` gives the dimExtent.
- int64_t dimExtent = 1;
+ // from `offset` gives the dimSize.
+ int64_t dimSize = 1;
if (stride == 0) {
// Every slot reads `offset`; no dimension is involved.
dimStride = 1;
@@ -2309,12 +2261,12 @@ static FailureOr<DonorIndex> fitDonorIndex(ArrayRef<int64_t> needed,
int64_t diff = wanted - offset;
if (diff < 0 || diff % stride != 0)
return failure();
- dimExtent = std::max(dimExtent, diff / stride + 1);
+ dimSize = std::max(dimSize, diff / stride + 1);
}
}
// Derived from the table, so verify against the table before relying on it.
- DonorIndex donorIndex{stride, offset, dimStride, dimExtent};
+ DonorIndex donorIndex{stride, offset, dimStride, dimSize};
if (!llvm::all_of(llvm::seq<int64_t>(0, targetLanePeriod), [&](int64_t slot) {
return needed[slot] == donorIndex.at(slot);
}))
@@ -2413,7 +2365,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
// Require the layouts to distribute what the vector types account for. The
// lanes past the first `targetLanePeriod` are replicas by construction: both
// computeStaticDistributedCoords implementations delinearize the lane id
- // modulo the lane_layout extents, whose product is `targetLanePeriod`.
+ // modulo the `lane_layout` dimensions, whose product is `targetLanePeriod`.
for (int64_t lane = 0; lane < subgroupSize; lane++) {
assert((*targetOwned)[lane] == (*targetOwned)[lane % targetLanePeriod] &&
"target ownership is not periodic in the lane period");
@@ -2467,7 +2419,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
// the same one. Either step is skipped when it is a no-op over `slot`.
llvm::DenseMap<std::pair<int64_t, int64_t>, Value> dimIndices;
auto getDimIndex = [&](const DonorIndex &donorIndex) -> Value {
- Value &dimIndex = dimIndices[{donorIndex.dimStride, donorIndex.dimExtent}];
+ Value &dimIndex = dimIndices[{donorIndex.dimStride, donorIndex.dimSize}];
if (dimIndex)
return dimIndex;
dimIndex = slot;
@@ -2475,10 +2427,10 @@ static FailureOr<Value> redistributeBroadcastedValue(
dimIndex = arith::DivUIOp::create(
rewriter, loc, dimIndex,
arith::ConstantIndexOp::create(rewriter, loc, donorIndex.dimStride));
- if (donorIndex.dimExtent < targetLanePeriod / donorIndex.dimStride)
+ if (donorIndex.dimSize < targetLanePeriod / donorIndex.dimStride)
dimIndex = arith::RemUIOp::create(
rewriter, loc, dimIndex,
- arith::ConstantIndexOp::create(rewriter, loc, donorIndex.dimExtent));
+ arith::ConstantIndexOp::create(rewriter, loc, donorIndex.dimSize));
return dimIndex;
};
@@ -2502,7 +2454,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
// Extract, shuffle and insert one element of the result at a time: one
// extract per distinct fragment index, one shuffle per element the lane does
// not already hold.
- // TODO: elements sharing a donor could travel in one shuffle, `gpu.shuffle`
+ // TODO: elements sharing a donor could move in one shuffle, `gpu.shuffle`
// moving a whole 32-bit lane word.
llvm::DenseMap<std::tuple<int64_t, int64_t, int64_t, int64_t>, Value>
extracted;
>From 364d3ab93dbd609f602e46952a6d1ce44d3ea0f8 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 20:12:11 +0000
Subject: [PATCH 31/46] Define donor index, and catch the comments up with the
rename
The struct, its field and the fitting function are DonorIndex, donorIndex and
fitDonorIndex, but five comments still called the thing a fragment index: the
describe() string that reaches a diagnostic, the offset field, fitDonorIndex,
and two in the emitter.
It was also the one component of an element source without a Terms entry, while
donorDelta and donor lane both have one, so the name appeared in a diagnostic
without being defined anywhere. Terms now defines it, and the element source
entry names it where it lists the components rather than only describing it as
where in that lane fragment the element sits.
The entry says the index is a function of the slot, which is the phrasing the
Why donors sit a whole period apart section already uses for the same fact, and
it settles what the name leaves open: DonorIndex holds four constants and an
at(slot), so it is a rule every lane evaluates with its own slot, not a value
per lane. Without that, describe() calling it a function of the lane id and
fitDonorIndex calling it the index each slot has to read read as two different
things.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 17 ++++++++++-------
1 file changed, 10 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index fecc80e61d568..fa7656aa08fe4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1791,9 +1791,12 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// candidates.
// donor lane supplying lane, computed as `slot + donorDelta`. First
// fit is taken by enumerating through the donorDeltas.
+// donor index where in the donor lane's fragment the element sits. A
+// function of the slot, so every lane computes it with its
+// own
// element source for one element of the target fragment, a tuple: which
-// donor lane supplies it, where in that lane's fragment
-// it sits, and whether a shuffle is needed. One is
+// donor lane supplies it, the donor index within that
+// lane's fragment, and whether a shuffle is needed. One is
// derived per element position, at compile time.
//
// What it handles
@@ -2102,7 +2105,7 @@ static StringRef describe(RedistributionLimit limit) {
case RedistributionLimit::NoDonorDelta:
return "no donor lane holds the element for every slot";
case RedistributionLimit::IndexNotLaneAffine:
- return "the fragment index is not an affine function of the lane id";
+ return "the donor index is not an affine function of the lane id";
}
llvm_unreachable("unhandled RedistributionLimit");
}
@@ -2192,7 +2195,7 @@ struct DonorIndex {
/// Distance in the fragment between consecutive coordinates along the
/// dimension being walked, or 0 when every slot reads the same element.
int64_t stride;
- /// Fragment index that coordinate 0 maps to, and the whole index when
+ /// Donor index that coordinate 0 maps to, and the whole index when
/// `stride` is 0.
int64_t offset;
/// Slots apart for the coordinate to advance by one, or 1 when no dimension
@@ -2228,7 +2231,7 @@ struct ElementSource {
bool needsShuffle;
};
-/// Fits `needed[slot]`, the fragment index each slot has to read, to the
+/// Fits `needed[slot]`, the donor index each slot has to read, to the
/// restricted form `DonorIndex` can emit. A slot is a lane reduced modulo
/// the target's lane period, so there are `targetLanePeriod` of them. Fails
/// when no index in that form reproduces the whole table.
@@ -2400,7 +2403,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
fragment = vector::BitCastOp::create(
rewriter, loc, VectorType::get({srcNumElems}, shuffleTy), fragment);
- // The lane's slot, which every fragment index is a function of.
+ // The lane's slot, which every donor index is a function of.
Value slot = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
/*upperBound=*/mlir::IntegerAttr());
if (targetLanePeriod != subgroupSize)
@@ -2452,7 +2455,7 @@ static FailureOr<Value> redistributeBroadcastedValue(
};
// Extract, shuffle and insert one element of the result at a time: one
- // extract per distinct fragment index, one shuffle per element the lane does
+ // extract per distinct donor index, one shuffle per element the lane does
// not already hold.
// TODO: elements sharing a donor could move in one shuffle, `gpu.shuffle`
// moving a whole 32-bit lane word.
>From c5cbe9dbeb438ee1712c5760de7f0564d0f44fa6 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 20:43:50 +0000
Subject: [PATCH 32/46] [MLIR][XeGPU] Keep declined cases out of the "What it
handles" section
The table listing what the broadcast redistribution handles carried a
final row for a combination it declines, and a paragraph below it existed
only to explain that row. Both belong in "What it declines".
Move the row's content there, folded into the NoDonorDelta entry so that
entry reads on its own, and drop the paragraph and the sentence
cross-referencing the row.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 69 ++++++-------------
1 file changed, 20 insertions(+), 49 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index fa7656aa08fe4..f4316ce1197a1 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1803,67 +1803,34 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// ---------------
//
// Two properties shape the outcome: the input's replication sets how many
-// elements need a cross-lane shuffle -- `gpu.shuffle idx`, as opposed to the
-// within-lane `vector.shuffle` the neighbouring patterns use -- and the
-// target's lane period sets where they may be shuffled from. Only the second
-// can decline the conversion. Examples stay on a `vector<8x2>` over 16 lanes;
-// the case numbers are used throughout the rest of this comment.
+// elements need a cross-lane shuffle, and the target's lane period sets where
+// they may be shuffled from. The period alone can decline the conversion; the
+// input's replication only does so together with it. Examples stay on a
+// `vector<8x2>` over 16 lanes; the case numbers are used throughout the rest of
+// this comment.
//
// The input's replication decides what each lane already holds, and so how many
-// shuffles are needed:
+// shuffles are needed.
//
-// fully broadcast every lane already holds the whole value, so the
-// result is extracts alone. Spelled by slicing away
-// every distributed dimension:
-// slice<layout<[1, 1, 16]>, dims = [2]>
-//
-// partially broadcast each broadcast group holds part of it, so a lane has
-// to be given whatever its group lacks. Spelled either
-// by slicing a dimension away or by a `lane_layout`
-// shorter than the subgroup.
-//
-// The target's lane period decides where a lane may be given that from. Donors
-// sit at multiples of the period, so a period `P` over `S` lanes leaves `S / P`
-// values of `donorDelta` to choose from:
-//
-// period < S more than one, so elements can cross lanes. `layout<[8, 1]>`
-// leaves 2, `layout<[2, 1]>` leaves 8.
-//
-// period == S only `donorDelta` 0, so nothing can cross a lane: only an
-// input already holding each lane's target fragment works.
-// A replicating target can have this period, slicing not
-// reducing it.
-//
-// One example per combination, and what the tests show it emitting. The target
-// is given as its `lane_layout`, the period following from it; the input as its
-// category, `sliced` meaning a sliced dimension and `short` a `lane_layout`
-// shorter than the subgroup:
+// Below are the categories of input and target combination, with one example
+// each and what the tests emit for it. The target is given as its `lane_layout`
+// and the period following from it; the input as its category, where `sliced`
+// means a sliced dimension and `short` a `lane_layout` shorter than the
+// subgroup:
//
// case input target period emits
// 1 fully broadcast [4, 1] 4 4 extracts
// 3 fully broadcast [8, 1, 2] sliced 16 8 extracts
// 2 partial, sliced [8, 1] 8 1 extract, 1 shuffle
// 4 partial, short [2, 1] 2 2 extracts, 6 shuffles
-// -- partial, short [8, 1, 2] sliced 16 declined
//
-// The counts follow from the target fragment. One shuffle per element the
-// reader lane does not already hold, and one extract per distinct position read
-// in the donor lane's fragment, elements reading the same position sharing one.
-// So case 2's 1x2 fragment is two elements read at the same position: one
+// The `emits` column follows from the target fragment. One shuffle per element
+// the reader lane does not already hold, and one extract per distinct position
+// read in the donor lane's fragment, elements reading the same position sharing
+// one. So case 2's 1x2 fragment is two elements read at the same position: one
// extract, and one of the two shuffled. Case 4's 4x2 fragment is eight elements
// over two positions: two extracts, six of the eight shuffled.
//
-// Only the target varies within a category, so the test suite carries more than
-// one of some -- @convert_layout_broadcast_all_lanes is case 1 with `[8, 1]` --
-// but they behave alike and only one is described here. Case 1 is the one whose
-// constants are all distinct, which the index expression below is read off.
-//
-// The last row is the boundary the two properties imply: with only `donorDelta`
-// 0 available, a lane holding less than its target fragment cannot be
-// completed, and the match fails with NoDonorDelta. There a lane holds 2
-// elements, `layout<[8, 1]>` giving it one row, and needs the sliced target's
-// 8, one whole column.
-//
// Case 2 is the running example below: lanes 0-7 hold all of column
// 0, lanes 8-15 all of column 1, and lane `i` leaves holding row `i`. Lane 3
// wants (3, 0) and (3, 1); it holds (3, 0) already, while (3, 1) is in lane 11.
@@ -1927,7 +1894,11 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// way:
//
// a partially broadcast input into a target of NoDonorDelta
-// period `S`, the last row of the table above
+// period `S`, where only `donorDelta` 0 is available
+// and a lane holding less than its target fragment
+// cannot be completed:
+//
+// layout<[8, 1]> -> slice<layout<[8, 1, 2]>, dims = [0]>
//
// a target whose distributed dimension the index IndexNotLaneAffine
// cannot walk, `order` having transposed it
>From 67638633bb7d9d283995a9d391040584964ff973 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 20:45:54 +0000
Subject: [PATCH 33/46] [MLIR][XeGPU] Trim the donor-period rationale to what
it has to say
The section explaining why a donor sits a whole lane period from its
reader restated its own conclusion, re-counted case 2's element sources,
and repeated which lane holds which element from the running example
above. Keep the shuffle semantics, the equation they imply, and a single
instance of the shared index.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 26 +++++++------------
1 file changed, 9 insertions(+), 17 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index f4316ce1197a1..3cae129984a96 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1842,28 +1842,20 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Why donors sit a whole period apart
// -----------------------------------
//
-// `gpu.shuffle idx` is the only way to move data between lanes, and it is not a
-// remote read: every lane extracts one element from its own fragment and
-// contributes it, then each lane receives whatever the lane it names
+// `gpu.shuffle idx` is not a remote read: every lane extracts one element from
+// its own fragment and contributes it, then receives whatever the lane it names
// contributed. A donor has therefore already extracted, at its own index,
-// before any shuffle happens, and never learns who asked.
-//
-// That works only if the index the donor used is the one the reader wanted. The
-// index is a function of the slot, so the two have to share one:
+// before any shuffle happens, and never learns who asked. So the index the
+// donor used has to be the one the reader wanted, and since the index is a
+// function of the slot, the two have to share one:
//
// donorSlot = (slot + donorDelta) % targetLanePeriod = slot
//
-// which holds exactly because `donorDelta` is a multiple of
-// `targetLanePeriod`. That is where the constraint comes from.
-//
-// Case 2 has two element sources, one per element of its 1x2 fragment. Reader
-// slot 3 wants (3, 0) and (3, 1):
-//
-// element 0 -> donorDelta = 0, donor = 3, index = 3 already local
-// element 1 -> donorDelta = 8, donor = 11, index = 3 shuffled
+// which holds exactly when `donorDelta` is a multiple of `targetLanePeriod`.
//
-// Both lanes extract at index 3, from different fragments: the input gives lane
-// 3 column 0 and lane 11 column 1, so one index yields the two elements wanted.
+// In case 2, reader slot 3 gets (3, 0) from donorDelta 0 and (3, 1) from
+// donorDelta 8, lane 11 -- and both lanes extract at index 3, the input giving
+// lane 3 column 0 and lane 11 column 1.
//
// What it declines
// ----------------
>From e02df353e29e116ad2e795ffa0330be6f5e06a87 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 20:50:10 +0000
Subject: [PATCH 34/46] [MLIR][XeGPU] Put the capability boundary before the
mechanism
"What it declines" now follows "What it handles" directly, so the two
halves of the boundary read together, ahead of any explanation of how the
scheme works.
The running example and the nested-slice normalisation are assumptions of
that explanation rather than statements about what is handled, so they
move down to the head of the first section using them. The example loses
the lane 3 walkthrough, which the donorDelta example right below it now
gives.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 51 +++++++++----------
1 file changed, 25 insertions(+), 26 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 3cae129984a96..e8305fea21c21 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1831,32 +1831,6 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// extract, and one of the two shuffled. Case 4's 4x2 fragment is eight elements
// over two positions: two extracts, six of the eight shuffled.
//
-// Case 2 is the running example below: lanes 0-7 hold all of column
-// 0, lanes 8-15 all of column 1, and lane `i` leaves holding row `i`. Lane 3
-// wants (3, 0) and (3, 1); it holds (3, 0) already, while (3, 1) is in lane 11.
-//
-// Nested slices are coalesced before anything else looks at a layout, so only
-// the lane_layout of the underlying layout and the union of the sliced dims
-// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
-//
-// Why donors sit a whole period apart
-// -----------------------------------
-//
-// `gpu.shuffle idx` is not a remote read: every lane extracts one element from
-// its own fragment and contributes it, then receives whatever the lane it names
-// contributed. A donor has therefore already extracted, at its own index,
-// before any shuffle happens, and never learns who asked. So the index the
-// donor used has to be the one the reader wanted, and since the index is a
-// function of the slot, the two have to share one:
-//
-// donorSlot = (slot + donorDelta) % targetLanePeriod = slot
-//
-// which holds exactly when `donorDelta` is a multiple of `targetLanePeriod`.
-//
-// In case 2, reader slot 3 gets (3, 0) from donorDelta 0 and (3, 1) from
-// donorDelta 8, lane 11 -- and both lanes extract at index 3, the input giving
-// lane 3 column 0 and lane 11 column 1.
-//
// What it declines
// ----------------
//
@@ -1902,6 +1876,31 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// element narrower than a byte, and FragmentSizeMismatch, an invariant check
// for a layout disagreeing with the vector type the caller derived from it.
//
+// Why donors sit a whole period apart
+// -----------------------------------
+//
+// Case 2 is the running example from here on: lanes 0-7 hold all of column 0,
+// lanes 8-15 all of column 1, and lane `i` leaves holding row `i`.
+//
+// Nested slices are coalesced before anything else looks at a layout, so only
+// the lane_layout of the underlying layout and the union of the sliced dims
+// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
+//
+// `gpu.shuffle idx` is not a remote read: every lane extracts one element from
+// its own fragment and contributes it, then receives whatever the lane it names
+// contributed. A donor has therefore already extracted, at its own index,
+// before any shuffle happens, and never learns who asked. So the index the
+// donor used has to be the one the reader wanted, and since the index is a
+// function of the slot, the two have to share one:
+//
+// donorSlot = (slot + donorDelta) % targetLanePeriod = slot
+//
+// which holds exactly when `donorDelta` is a multiple of `targetLanePeriod`.
+//
+// In case 2, reader slot 3 gets (3, 0) from donorDelta 0 and (3, 1) from
+// donorDelta 8, lane 11 -- and both lanes extract at index 3, the input giving
+// lane 3 column 0 and lane 11 column 1.
+//
// The reverse map
// ---------------
//
>From 8441e8d6206ab947172bf256cf7ef51f96184bde Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 20:51:51 +0000
Subject: [PATCH 35/46] [MLIR][XeGPU] Write the declines lead-in as sentences
The paragraph opened with a verbless noun phrase and a relative clause
missing its pronoun. Say the same thing in full sentences.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 10 +++++-----
1 file changed, 5 insertions(+), 5 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index e8305fea21c21..1bc60c5584dec 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1834,11 +1834,11 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// What it declines
// ----------------
//
-// The layout combinations, in the terms above, with the `RedistributionLimit` a
-// match failure names. None is a limit of `convert_layout`, which is meaningful
-// for any pair and so always lowerable; each is a limit of the scheme here, one
-// shuffle per element from a donor a whole number of lane periods away. Another
-// pattern or a more general scheme lowers any of them.
+// Each combination below is declined, with the `RedistributionLimit` that the
+// match failure reports. None means a limit of `convert_layout` itself, which
+// is meaningful for any pair of layouts and always lowerable by some pattern.
+// They limit only the scheme used here -- one shuffle per element, from a donor
+// a whole number of lane periods away -- and a more general scheme lifts them.
//
// Decided by inspecting the layouts. None of these rejects the input on its
// own -- the input is only read, so any `lane_layout` is admissible until the
>From c8747dbb22df4651b3a80ff6e30d2f0d239c6fd3 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 20:56:07 +0000
Subject: [PATCH 36/46] [MLIR][XeGPU] Drop the nested-slice note from the
section header
Nested slices are coalesced by the one flatten() call in the file, whose
call site in getLanePeriod already says so. The header paragraph repeated
that, added a SliceAttr fact this pass never uses, and named a test.
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 4 ----
1 file changed, 4 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 1bc60c5584dec..0f405eabc1844 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1882,10 +1882,6 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Case 2 is the running example from here on: lanes 0-7 hold all of column 0,
// lanes 8-15 all of column 1, and lane `i` leaves holding row `i`.
//
-// Nested slices are coalesced before anything else looks at a layout, so only
-// the lane_layout of the underlying layout and the union of the sliced dims
-// matter; @convert_layout_broadcast_nested_slice covers a two-level input.
-//
// `gpu.shuffle idx` is not a remote read: every lane extracts one element from
// its own fragment and contributes it, then receives whatever the lane it names
// contributed. A donor has therefore already extracted, at its own index,
>From dbf6636da2f2db22d03dd26765ab0e892beee1f7 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:03:48 +0000
Subject: [PATCH 37/46] [MLIR][XeGPU] Say that the lanes past the lane period
are undefined
The paragraph claimed the target layout assigns those lanes no elements,
and then argued from whether anything reads them. Neither is the point:
the layout distributes to the first lane period of lanes and asks nothing
of the rest, so what they hold is undefined and any code depending on it
is ill-defined. The wrapped fragment computeStaticDistributedCoords
returns for them follows from delinearising the lane id, and is not a
fragment the layout requires.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 10 +++-------
1 file changed, 3 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 0f405eabc1844..596ca88670fb8 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2006,13 +2006,9 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// values of `donorDelta` may work and all give the same value, so the smallest
// is taken, which is what leaves a fully broadcast input shuffle-free.
//
-// Only the first `targetLanePeriod` lanes are required to end up correct --
-// each is the one lane of its slot the target layout distributes to, and the
-// layout assigns the rest no elements. Lanes sharing a slot run the same
-// arithmetic on their own fragment: when a shuffle is emitted they all receive
-// the donor's value and are equally correct, and when it is dropped their local
-// extract can land on a different element. Dropping one therefore relies on
-// nobody reading them.
+// The target layout distributes to the first `targetLanePeriod` lanes and asks
+// nothing of the rest, so only those lanes are made correct. What the others
+// hold is undefined, and code depending on it is ill-defined.
//
// An index is derived from the tables and then verified against them, so a
// layout change that is not of this form is reported as a match failure rather
>From 415fc29269ab58f64c54003225ec8065fba532f6 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:18:05 +0000
Subject: [PATCH 38/46] [MLIR][XeGPU] Name replication as a choice, not the
layout's meaning
The lanes past the first lane period are undefined by the layout;
replication is what computeStaticDistributedCoords picks for them, and the
assert below checks that choice. Say it that way round, matching the
section header.
---
.../Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 7 ++++---
1 file changed, 4 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 596ca88670fb8..db8829d4486cb 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2320,9 +2320,10 @@ static FailureOr<Value> redistributeBroadcastedValue(
if (failed(inputOwned) || failed(targetOwned))
return failure();
// Require the layouts to distribute what the vector types account for. The
- // lanes past the first `targetLanePeriod` are replicas by construction: both
- // computeStaticDistributedCoords implementations delinearize the lane id
- // modulo the `lane_layout` dimensions, whose product is `targetLanePeriod`.
+ // layout leaves the lanes past the first `targetLanePeriod` undefined, and
+ // replication is what computeStaticDistributedCoords picks for them: both
+ // implementations delinearize the lane id modulo the `lane_layout`
+ // dimensions, whose product is `targetLanePeriod`.
for (int64_t lane = 0; lane < subgroupSize; lane++) {
assert((*targetOwned)[lane] == (*targetOwned)[lane % targetLanePeriod] &&
"target ownership is not periodic in the lane period");
>From 1ca4f8f8726e9d660f34a8de316e0871579d2a0f Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:32:24 +0000
Subject: [PATCH 39/46] [MLIR][XeGPU] Fix three inconsistencies in the section
header
donorDelta was defined as an offset from the reader lane, while every
other mention of it, and the code, adds it to the slot. The two coincide
only for the lanes below the lane period, which are the only ones the
layout defines.
The donor index entry trailed off mid-phrase, and the table lead-in
credited the tests with the ops the pattern emits.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 17 ++++++++---------
1 file changed, 8 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index db8829d4486cb..8bebf4bb3f769 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1786,14 +1786,13 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// ownership table tabulated "element of a fragment" to coordinate mapping
// with respect to undistributed value. one table per input
// and target
-// donorDelta a lane offset from the reader lane to
-// the lane supplying an element. Can have multiple
-// candidates.
+// donorDelta a lane offset from the slot to the lane supplying an
+// element. Can have multiple candidates.
// donor lane supplying lane, computed as `slot + donorDelta`. First
// fit is taken by enumerating through the donorDeltas.
// donor index where in the donor lane's fragment the element sits. A
-// function of the slot, so every lane computes it with its
-// own
+// function of the slot, so every lane computes it from its
+// own slot
// element source for one element of the target fragment, a tuple: which
// donor lane supplies it, the donor index within that
// lane's fragment, and whether a shuffle is needed. One is
@@ -1813,10 +1812,10 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// shuffles are needed.
//
// Below are the categories of input and target combination, with one example
-// each and what the tests emit for it. The target is given as its `lane_layout`
-// and the period following from it; the input as its category, where `sliced`
-// means a sliced dimension and `short` a `lane_layout` shorter than the
-// subgroup:
+// each and what the pattern emits for it. The target is given as its
+// `lane_layout` and the period following from it; the input as its category,
+// where `sliced` means a sliced dimension and `short` a `lane_layout` shorter
+// than the subgroup:
//
// case input target period emits
// 1 fully broadcast [4, 1] 4 4 extracts
>From 00fe3af38f9a239ef2ae4c63df47d563bade9ec8 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:34:02 +0000
Subject: [PATCH 40/46] [MLIR][XeGPU] Fold the orphaned replication lead-in
into the intro
The sentence introduced two table rows that no longer exist, leaving it
restating the opening paragraph. Keep the step it contributed -- that
replication decides what a lane already holds, hence the shuffle count --
in that paragraph, and drop the sentence.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 15 ++++++---------
1 file changed, 6 insertions(+), 9 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 8bebf4bb3f769..fcd66ff4a7db9 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1801,15 +1801,12 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// What it handles
// ---------------
//
-// Two properties shape the outcome: the input's replication sets how many
-// elements need a cross-lane shuffle, and the target's lane period sets where
-// they may be shuffled from. The period alone can decline the conversion; the
-// input's replication only does so together with it. Examples stay on a
-// `vector<8x2>` over 16 lanes; the case numbers are used throughout the rest of
-// this comment.
-//
-// The input's replication decides what each lane already holds, and so how many
-// shuffles are needed.
+// Two properties shape the outcome: the input's replication decides what each
+// lane already holds and so how many elements need a cross-lane shuffle, and
+// the target's lane period sets where they may be shuffled from. The period
+// alone can decline the conversion; the input's replication only does so
+// together with it. Examples stay on a `vector<8x2>` over 16 lanes; the case
+// numbers are used throughout the rest of this comment.
//
// Below are the categories of input and target combination, with one example
// each and what the pattern emits for it. The target is given as its
>From e830aa4aafd8361337a72b6f63d001a1a68a0807 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:36:32 +0000
Subject: [PATCH 41/46] [MLIR][XeGPU] Retire "block" for the vocabulary already
in use
The Algorithm section opened by calling two different things a block: the
distribution unit the two layouts must agree on, and the single element
moved per shuffle, which is what a unit is once lane_data has to be all
ones. Name them `lane_data` and element, as the rest of the comment and
the tag list do.
haveSameDistributionUnit switched from "distribution unit" to "blocks"
inside one sentence; keep the one word.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 11 +++++------
1 file changed, 5 insertions(+), 6 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index fcd66ff4a7db9..f8b431db32254 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1950,11 +1950,10 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Algorithm
// ---------
//
-// Check the two layouts agree on the block, then for each block the target
-// wants, find an equal block in some lane and move that lane's copy to it:
-// extract, shuffle, insert. Only the search costs anything, because the match
-// has to hold for every slot at once -- through one lane offset and one index
-// expression.
+// Check the two layouts agree on `lane_data`, then for each element the target
+// wants, find a lane already holding it and move its copy over: extract,
+// shuffle, insert. The search is the only hard part, because one lane offset
+// and one index expression have to serve every slot at once.
//
// A `[name]` tag below gives the function carrying out the step under it.
// `redistributeBroadcastedValue` is the enclosing one -- it owns both tables
@@ -2061,7 +2060,7 @@ static StringRef describe(RedistributionLimit limit) {
}
/// Relational rule: both layouts have to hand the lanes the same distribution
-/// unit, so that only the assignment of those blocks to lanes changes.
+/// unit, so that only the assignment of those units to lanes changes.
static bool haveSameDistributionUnit(xegpu::DistributeLayoutAttr inputLayout,
xegpu::DistributeLayoutAttr targetLayout) {
return inputLayout.getEffectiveLaneDataAsInt() ==
>From 071af44bdaa5a2f6262120dfbc0608049fe7a7a5 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:40:36 +0000
Subject: [PATCH 42/46] [MLIR][XeGPU] Drop the unlabelled foursome and the
third selection rule
"Read the expression off case 1, where the four quantities are 4, 2, 1
and 4" named neither the quantities nor the order, and the four it lists
are not the four the table's columns lead with. It also recommended the
one row where two of them collide: case 1's period and dimSize are both 4,
the target [4, 1] walking a dimension whose size is the period. Keep the
sentence after it, which makes the point on its own, with the roles named.
The donor lane entry in Terms also carried the donorDelta selection rule,
which the Algorithm section states twice already, in the words the code
uses.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 11 ++++-------
1 file changed, 4 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index f8b431db32254..70ea8644ad376 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1788,8 +1788,7 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// and target
// donorDelta a lane offset from the slot to the lane supplying an
// element. Can have multiple candidates.
-// donor lane supplying lane, computed as `slot + donorDelta`. First
-// fit is taken by enumerating through the donorDeltas.
+// donor lane supplying lane, computed as `slot + donorDelta`.
// donor index where in the donor lane's fragment the element sits. A
// function of the slot, so every lane computes it from its
// own slot
@@ -1941,11 +1940,9 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// `2s+0, 2s+1, 2s+8, 2s+9`. The jump from 1 to 8 is four rows down at 2
// elements each, the 4 coming from the target `lane_layout` [4, 1].
//
-// Read the expression off case 1, where the four quantities are 4, 2, 1 and 4.
-// On case 2's `layout<[8, 1]>` target they collide: the lane period (8), its
-// own `donorDelta` (8) and case 3's `dimStride` (8) all have the same value
-// while being three different things -- a lane count, a lane offset and an
-// index divisor.
+// The same value can stand for different things: on case 2's `layout<[8, 1]>`
+// target, 8 is the lane period, its own `donorDelta`, and case 3's `dimStride`
+// -- a lane count, a lane offset and an index divisor.
//
// Algorithm
// ---------
>From 577a4fb5cc2f8c168e14ce96584df3e2f1640b4e Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:43:19 +0000
Subject: [PATCH 43/46] [MLIR][XeGPU] Bring the broadcast tests onto the same
vocabulary
Two of them captured `lane % period` as ROW while the others call it SLOT,
which is the term the section header defines. One comment also still said
"extent", the name dimSize replaced.
---
.../XeGPU/sg-to-lane-distribute-unit.mlir | 16 ++++++++--------
1 file changed, 8 insertions(+), 8 deletions(-)
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 20bed3064cadd..7b38b53b7e280 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-unit.mlir
@@ -1600,16 +1600,16 @@ gpu.func @convert_layout_partial_subgroup() {
// 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: %[[SLOT:.*]] = arith.remui %[[LANE]], %[[C8]] : index
// CHECK: %[[ZERO:.*]] = arith.constant dense<0> : vector<2xi8>
-// CHECK: %[[ELEM:.*]] = vector.extract %[[BITS]][%[[ROW]]] : i8 from vector<8xi8>
+// CHECK: %[[ELEM:.*]] = vector.extract %[[BITS]][%[[SLOT]]] : i8 from vector<8xi8>
// CHECK: %[[INS0:.*]] = vector.insert %[[ELEM]], %[[ZERO]] [0] : i8 into vector<2xi8>
// CHECK: %[[WIDTH:.*]] = arith.constant 16 : i32
-// CHECK: %[[ROW_I32:.*]] = arith.index_cast %[[ROW]] : index to i32
+// CHECK: %[[SLOT_I32:.*]] = arith.index_cast %[[SLOT]] : index to i32
// CHECK: %[[C8_I32:.*]] = arith.constant 8 : i32
-// CHECK: %[[LANE1:.*]] = arith.addi %[[ROW_I32]], %[[C8_I32]] : i32
+// CHECK: %[[LANE1:.*]] = arith.addi %[[SLOT_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>
@@ -1632,14 +1632,14 @@ gpu.func @convert_layout_broadcast_to_lane_distributed() {
// 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: %[[SLOT:.*]] = 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: %[[IDX0:.*]] = arith.muli %[[SLOT]], %[[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: %[[MUL1:.*]] = arith.muli %[[SLOT]], %[[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>
@@ -1699,7 +1699,7 @@ gpu.func @convert_layout_broadcast_to_sliced_target() {
// reaching it through two nested slices. They are coalesced, so what counts is
// the parent lane_layout covering all 16 lanes and the dims the slices drop
// together: the 8 lanes of the dropped dim 0 hold the same fragment, while the
-// kept dim of extent 2 splits the columns between the two groups.
+// kept dim of size 2 splits the columns between the two groups.
// CHECK-LABEL: gpu.func @convert_layout_broadcast_nested_slice
// CHECK: %[[SRC:.*]] = arith.constant dense<1.000000e+00> : vector<8x1xf8E8M0FNU>
// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<8x1xf8E8M0FNU> to vector<8xf8E8M0FNU>
>From 0281b1c35ac8f8abed243a3def04284b80eacf0d Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:46:45 +0000
Subject: [PATCH 44/46] [MLIR][XeGPU] Mark Terms as local definitions
The dialect documents `lane_data` as defining the minimal distribution
unit, calling that a fragment, while Terms uses fragment for everything a
lane owns. Say that Terms redefines words for this comment and that the
dialect's own keep their meanings, which is also why distribution unit has
no entry.
---
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 3 +++
1 file changed, 3 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 70ea8644ad376..4225ccea9735e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1772,6 +1772,9 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// Terms
// -----
//
+// Defined here for this comment; `lane_layout`, `lane_data` and distribution
+// unit keep the meanings the dialect gives them.
+//
// element scalar of some element type
// fragment everything one lane owns. Consists of elements.
// broadcast group the lanes that hold the same fragment.
>From 6f8063d2ed3c5eda391bbfc9d4df97b7a2be546e Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:49:12 +0000
Subject: [PATCH 45/46] [MLIR][XeGPU] Punctuate the Terms entries consistently
Five of the ten ended with a full stop and five did not. Drop the trailing
stops so all ten read alike, keeping the ones that separate sentences
within an entry.
The ownership table entry also started a sentence lowercase and was
missing an article; make it one sentence.
---
.../XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 14 +++++++-------
1 file changed, 7 insertions(+), 7 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 4225ccea9735e..1817494b6436b 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1776,8 +1776,8 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// unit keep the meanings the dialect gives them.
//
// element scalar of some element type
-// fragment everything one lane owns. Consists of elements.
-// broadcast group the lanes that hold the same fragment.
+// fragment everything one lane owns. Consists of elements
+// broadcast group the lanes that hold the same fragment
// lane period how many lanes the layout lays out,
// `product(lane_layout)`. A larger subgroup repeats: lane
// `l` and lane `l + period` get the same fragment
@@ -1787,18 +1787,18 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// within one target period. A target-side notion: lanes
// sharing a slot are assigned the same target fragment
// ownership table tabulated "element of a fragment" to coordinate mapping
-// with respect to undistributed value. one table per input
-// and target
+// with respect to the undistributed value, one table per
+// input and target
// donorDelta a lane offset from the slot to the lane supplying an
-// element. Can have multiple candidates.
-// donor lane supplying lane, computed as `slot + donorDelta`.
+// element. Can have multiple candidates
+// donor lane supplying lane, computed as `slot + donorDelta`
// donor index where in the donor lane's fragment the element sits. A
// function of the slot, so every lane computes it from its
// own slot
// element source for one element of the target fragment, a tuple: which
// donor lane supplies it, the donor index within that
// lane's fragment, and whether a shuffle is needed. One is
-// derived per element position, at compile time.
+// derived per element position, at compile time
//
// What it handles
// ---------------
>From 2108a4569f393c4768641f3149faf194423fa5c3 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 2 Sep 2026 21:50:54 +0000
Subject: [PATCH 46/46] [MLIR][XeGPU] Retire donorSlot for the reason it
existed
The equation compares the reader's slot with the donor's, and the left
side had been given a name rather than an explanation -- applying a
target-side word to a lane being viewed as a donor. The donor has a slot
because every lane runs the same code and so reads for itself; say that,
and the equation needs no new term.
---
.../Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp | 7 ++++---
1 file changed, 4 insertions(+), 3 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 1817494b6436b..f86b6a7c1959d 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1884,10 +1884,11 @@ static FailureOr<Value> repackLaneData(ConversionPatternRewriter &rewriter,
// its own fragment and contributes it, then receives whatever the lane it names
// contributed. A donor has therefore already extracted, at its own index,
// before any shuffle happens, and never learns who asked. So the index the
-// donor used has to be the one the reader wanted, and since the index is a
-// function of the slot, the two have to share one:
+// donor used has to be the one the reader wanted. Every lane runs the same
+// code, so the donor computed its index from its own slot, and the two slots
+// have to agree:
//
-// donorSlot = (slot + donorDelta) % targetLanePeriod = slot
+// (slot + donorDelta) % targetLanePeriod = slot
//
// which holds exactly when `donorDelta` is a multiple of `targetLanePeriod`.
//
More information about the Mlir-commits
mailing list