[Mlir-commits] [mlir] [mlir][xegpu] Distribute three more convert_layout cases to lanes (PR #222693)
Sang Ik Lee
llvmlistbot at llvm.org
Thu Sep 24 07:43:14 PDT 2026
https://github.com/silee2 updated https://github.com/llvm/llvm-project/pull/222693
>From 09decb0413f5bde5e2e1159b7662194b46f186e8 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 9 Sep 2026 23:01:46 +0000
Subject: [PATCH 1/8] [mlir][xegpu] Distribute three more convert_layout cases
to lanes
`xegpu-sg-to-lane-distribute` could not lower an `xegpu.convert_layout` whose
effective `lane_layout` changes in a way that moves elements between lanes. Add
one pattern per case found in the workgroup-level mx-fp GEMM quantize-A kernels,
where `#xegpu.slice` layouts of a rank-3 layout meet the rank-2 layouts of
`load_nd` and `dpas_mx`:
- fully broadcast source to one row per lane: `vector.extract`
- source broadcast over two lane groups to one row per lane: `vector.extract`
plus `gpu.shuffle xor` for the element held by the other group
- fully broadcast source to one column per lane group: `vector.deinterleave`
plus a select on the lane group
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 463 +++++++++++++++++-
...-lane-distribute-convert-layout-slice.mlir | 101 ++++
2 files changed, 551 insertions(+), 13 deletions(-)
create mode 100644 mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 812233047fb88..1b1ba85158c4e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1869,6 +1869,440 @@ struct SgToLaneConvertLayout
}
};
+/// Check if layout distributes a rank-2 value with unit lane_data
+/// in row major or default order
+static bool assignsElementsToLanes(xegpu::DistributeLayoutAttr layout) {
+ if (layout.getRank() != 2)
+ return false;
+ return layout.getEffectiveLaneDataAsInt() == SmallVector<int64_t>{1, 1} &&
+ layout.getEffectiveOrderAsInt() == SmallVector<int64_t>{1, 0};
+}
+
+/// The quantities every element-to-lane redistribution needs, see
+/// `matchElementLaneRedistribution`.
+struct ElementLaneRedistribution {
+ /// Subgroup-level type of the converted value.
+ VectorType valueType;
+ /// `valueType` as distributed by `input_layout` and by `target_layout`.
+ VectorType distributedInput;
+ VectorType distributedTarget;
+ /// Effective `lane_layout` of `input_layout` and of `target_layout`. These
+ /// always differ, otherwise the conversion would already have folded.
+ SmallVector<int64_t> inputLaneLayout;
+ SmallVector<int64_t> targetLaneLayout;
+ int64_t subgroupSize;
+};
+
+/// Recognizes an `xegpu.convert_layout` that redistributes individual elements
+/// between the lanes of a subgroup, which is the common condition the three
+/// slice-attributed lowerings below have.
+/// Both input and target layouts must satisfy `assignsElementsToLanes` and must
+/// be able to distribute the value, the two can only differ in which lane
+/// owns which element.
+static FailureOr<ElementLaneRedistribution>
+matchElementLaneRedistribution(xegpu::ConvertLayoutOp op,
+ ConversionPatternRewriter &rewriter) {
+ xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
+ xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
+ if (!assignsElementsToLanes(inputLayout) ||
+ !assignsElementsToLanes(targetLayout))
+ return rewriter.notifyMatchFailure(
+ op, "both layouts must be rank 2 with effective lane_data [1, 1] and "
+ "effective order [1, 0]");
+
+ auto valueType = dyn_cast<VectorType>(op.getResult().getType());
+ if (!valueType)
+ return rewriter.notifyMatchFailure(op, "value type must be a vector");
+
+ FailureOr<VectorType> distributedInput =
+ xegpu::getDistVecTypeBasedOnLaneLayout(inputLayout, valueType);
+ FailureOr<VectorType> distributedTarget =
+ xegpu::getDistVecTypeBasedOnLaneLayout(targetLayout, valueType);
+ if (failed(distributedInput) || failed(distributedTarget))
+ return rewriter.notifyMatchFailure(
+ op, "value type must be distributable by both layouts");
+
+ const auto *uArch =
+ xegpu::uArch::getUArch(xegpu::getChipStr(op).value_or(""));
+ if (!uArch)
+ return rewriter.notifyMatchFailure(
+ op, "target attribute is required to determine the subgroup size");
+
+ ElementLaneRedistribution redistribution;
+ redistribution.valueType = valueType;
+ redistribution.distributedInput = *distributedInput;
+ redistribution.distributedTarget = *distributedTarget;
+ redistribution.inputLaneLayout = inputLayout.getEffectiveLaneLayoutAsInt();
+ redistribution.targetLaneLayout = targetLayout.getEffectiveLaneLayoutAsInt();
+ redistribution.subgroupSize = uArch->getSubgroupSize();
+ return redistribution;
+}
+
+/// Returns the lane stride of the single distributed dimension of `slice`, i.e.
+/// the distance in lane ids between two adjacent coordinates along that
+/// dimension. `slice` must have exactly one non-sliced dimension whose parent
+/// `lane_layout` extent is greater than one; the stride is the product of the
+/// parent `lane_layout` extents that precede it in the parent `order`.
+///
+/// Example: for
+/// #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2], lane_data = [4, 1, 1],
+/// order = [0, 2, 1]>, dims = [0]>
+/// the only such dimension is parent dim 2, of extent 2. `order` makes dim 0
+/// the fastest varying, with extent 8, so lanes 0..7 hold coordinate 0 and
+/// lanes 8..15 hold coordinate 1: the stride is 8.
+static FailureOr<int64_t> getDistributedDimLaneStride(xegpu::SliceAttr slice) {
+ xegpu::SliceAttr flattened = slice.flatten();
+ auto parent = dyn_cast<xegpu::LayoutAttr>(flattened.getParent());
+ if (!parent)
+ return failure();
+
+ SmallVector<int64_t> parentLaneLayout = parent.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> parentOrder = parent.getEffectiveOrderAsInt();
+ if (parentLaneLayout.size() != parentOrder.size())
+ return failure();
+
+ ArrayRef<int64_t> slicedDims = flattened.getDims().asArrayRef();
+ std::optional<int64_t> distributedDim;
+ for (int64_t dim = 0, rank = static_cast<int64_t>(parentLaneLayout.size());
+ dim < rank; ++dim) {
+ if (llvm::is_contained(slicedDims, dim) || parentLaneLayout[dim] == 1)
+ continue;
+ if (distributedDim)
+ return failure();
+ distributedDim = dim;
+ }
+ if (!distributedDim)
+ return failure();
+
+ int64_t stride = 1;
+ for (int64_t dim : parentOrder) {
+ if (dim == *distributedDim)
+ return stride;
+ stride *= parentLaneLayout[dim];
+ }
+ return failure();
+}
+
+/// Distributes the slice-attributed `xegpu.convert_layout` whose source is
+/// fully broadcast, i.e. every lane holds the whole value, onto the rows of a
+/// subset of the lanes. Each lane keeps a single element, so no data crosses
+/// lanes and one `vector.extract` suffices.
+///
+/// xegpu.convert_layout %src
+/// <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1],
+/// lane_data = [2, 1, 1],
+/// order = [0, 2, 1]>,
+/// dims = [0]>,
+/// target_layout = #xegpu.layout<lane_layout = [8, 1],
+/// lane_data = [1, 1]>}>
+/// : vector<8x1xf8E8M0FNU>
+///
+/// The input layout has effective `lane_layout` [1, 1], so the distributed
+/// source is the full `vector<8x1xf8E8M0FNU>`. The target spreads the 8 rows
+/// over 8 lanes, one row each, so the distributed result is
+/// `vector<1x1xf8E8M0FNU>` holding row `lane_id % 8`:
+///
+/// %flat = vector.shape_cast %src : vector<8x1xf8E8M0FNU> to
+/// vector<8xf8E8M0FNU>
+/// %lane = gpu.lane_id
+/// %row = arith.remui %lane, %c8 : index
+/// %elem = vector.extract %flat[%row] : f8E8M0FNU from vector<8xf8E8M0FNU>
+/// %res = vector.from_elements %elem : vector<1x1xf8E8M0FNU>
+///
+/// The source is flattened first because `xegpu-vector-linearize` cannot
+/// linearize a `vector.extract` with a dynamic position out of a rank-2 value.
+struct SgToLaneConvertLayoutBroadcastExtract
+ : public OpConversionPattern<xegpu::ConvertLayoutOp> {
+ using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(xegpu::ConvertLayoutOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
+ xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
+
+ if (!isa<xegpu::SliceAttr>(inputLayout))
+ return rewriter.notifyMatchFailure(op,
+ "input_layout must be #xegpu.slice");
+ if (!isa<xegpu::LayoutAttr>(targetLayout))
+ return rewriter.notifyMatchFailure(op,
+ "target_layout must be #xegpu.layout");
+
+ FailureOr<ElementLaneRedistribution> redistribution =
+ matchElementLaneRedistribution(op, rewriter);
+ if (failed(redistribution))
+ return failure();
+
+ if (redistribution->inputLaneLayout != SmallVector<int64_t>{1, 1})
+ return rewriter.notifyMatchFailure(
+ op, "input_layout effective lane_layout must be [1, 1]");
+ SmallVector<int64_t> targetLaneLayout = redistribution->targetLaneLayout;
+ if (targetLaneLayout[0] <= 1 || targetLaneLayout[1] != 1)
+ return rewriter.notifyMatchFailure(
+ op, "target_layout effective lane_layout must be [n, 1] with n > 1");
+ if (redistribution->distributedInput != redistribution->valueType)
+ return rewriter.notifyMatchFailure(
+ op, "distributed input_layout type must equal the value type");
+ if (redistribution->distributedTarget.getShape() != ArrayRef<int64_t>{1, 1})
+ return rewriter.notifyMatchFailure(
+ op, "distributed target_layout type must be vector<1x1>");
+ if (targetLaneLayout[0] > redistribution->subgroupSize)
+ return rewriter.notifyMatchFailure(
+ op, "target_layout effective lane_layout[0] must not exceed the "
+ "subgroup size");
+
+ VectorType valueType = redistribution->valueType;
+ Location loc = op.getLoc();
+ Value src =
+ castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getSource()),
+ redistribution->distributedInput);
+ auto flatType = VectorType::get({valueType.getNumElements()},
+ valueType.getElementType());
+ Value flat = vector::ShapeCastOp::create(rewriter, loc, flatType, src);
+ Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+ /*upperBound=*/mlir::IntegerAttr());
+ Value rowCount =
+ arith::ConstantIndexOp::create(rewriter, loc, targetLaneLayout[0]);
+ Value row = arith::RemUIOp::create(rewriter, loc, laneId, rowCount);
+ Value element = vector::ExtractOp::create(rewriter, loc, flat,
+ ArrayRef<OpFoldResult>{row});
+ rewriter.replaceOpWithNewOp<vector::FromElementsOp>(
+ op, redistribution->distributedTarget, element);
+ return success();
+ }
+};
+
+/// Distributes the slice-attributed `xegpu.convert_layout` whose source is
+/// broadcast over two lane groups onto the rows of a subset of the lanes. Each
+/// lane needs one element from its own group and one from the other group, so
+/// one `vector.extract` is paired with one `gpu.shuffle xor`.
+///
+/// xegpu.convert_layout %src
+/// <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2],
+/// lane_data = [4, 1, 1],
+/// order = [0, 2, 1]>,
+/// dims = [0]>,
+/// target_layout = #xegpu.layout<lane_layout = [8, 1],
+/// lane_data = [1, 1]>}>
+/// : vector<8x2xf8E8M0FNU>
+///
+/// The input layout has effective `lane_layout` [1, 2], so the distributed
+/// source is `vector<8x1xf8E8M0FNU>`: lane `l` holds all 8 rows of column
+/// `l / 8`. The target puts both columns of one row in a single lane, so the
+/// distributed result is `vector<1x2xf8E8M0FNU>`. Lane `l` owns `[l % 8, 0]`
+/// of its own column and gets the other column from lane `l ^ 8`, which holds
+/// the same row:
+///
+/// %flat = vector.shape_cast %src : vector<8x1xf8E8M0FNU> to
+/// vector<8xf8E8M0FNU>
+/// %lane = gpu.lane_id
+/// %row = arith.remui %lane, %c8 : index
+/// %own = vector.extract %flat[%row] : f8E8M0FNU from
+/// vector<8xf8E8M0FNU>
+/// %part, %v = gpu.shuffle xor %own, %c8_i32, %c16_i32 : f8E8M0FNU
+/// %res = vector.from_elements %own, %part : vector<1x2xf8E8M0FNU>
+///
+/// The source is flattened first because `xegpu-vector-linearize` cannot
+/// linearize a `vector.extract` with a dynamic position out of a rank-2 value.
+///
+/// Lanes 8..15 receive the two columns in the opposite order, which is unused
+/// because the target `lane_layout` [8, 1] only occupies 8 lanes.
+struct SgToLaneConvertLayoutPartialBroadcastExtractShuffle
+ : public OpConversionPattern<xegpu::ConvertLayoutOp> {
+ using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(xegpu::ConvertLayoutOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
+ xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
+
+ auto inputSlice = dyn_cast<xegpu::SliceAttr>(inputLayout);
+ if (!inputSlice)
+ return rewriter.notifyMatchFailure(op,
+ "input_layout must be #xegpu.slice");
+ if (!isa<xegpu::LayoutAttr>(targetLayout))
+ return rewriter.notifyMatchFailure(op,
+ "target_layout must be #xegpu.layout");
+
+ FailureOr<ElementLaneRedistribution> redistribution =
+ matchElementLaneRedistribution(op, rewriter);
+ if (failed(redistribution))
+ return failure();
+
+ VectorType valueType = redistribution->valueType;
+ Type elementType = valueType.getElementType();
+ if (!elementType.isIntOrFloat() ||
+ !llvm::is_contained({8u, 16u, 32u, 64u},
+ elementType.getIntOrFloatBitWidth()))
+ return rewriter.notifyMatchFailure(
+ op, "element type must be an int or float of bit width 8, 16, 32 or "
+ "64 to be carried by gpu.shuffle");
+
+ if (redistribution->inputLaneLayout != SmallVector<int64_t>{1, 2})
+ return rewriter.notifyMatchFailure(
+ op, "input_layout effective lane_layout must be [1, 2]");
+ SmallVector<int64_t> targetLaneLayout = redistribution->targetLaneLayout;
+ if (targetLaneLayout[0] <= 1 || targetLaneLayout[1] != 1)
+ return rewriter.notifyMatchFailure(
+ op, "target_layout effective lane_layout must be [n, 1] with n > 1");
+ if (redistribution->distributedInput.getShape() !=
+ ArrayRef<int64_t>{valueType.getDimSize(0), 1})
+ return rewriter.notifyMatchFailure(
+ op, "distributed input_layout type must be vector<shape[0]x1>");
+ if (redistribution->distributedTarget.getShape() != ArrayRef<int64_t>{1, 2})
+ return rewriter.notifyMatchFailure(
+ op, "distributed target_layout type must be vector<1x2>");
+
+ FailureOr<int64_t> stride = getDistributedDimLaneStride(inputSlice);
+ if (failed(stride))
+ return rewriter.notifyMatchFailure(
+ op, "input_layout parent must have exactly one non-sliced dimension "
+ "with lane_layout extent greater than one");
+ // The two broadcast groups must together cover the whole subgroup, so that
+ // `lane_id ^ stride` is the lane holding the other column of the same row.
+ if (*stride * redistribution->inputLaneLayout[1] !=
+ redistribution->subgroupSize)
+ return rewriter.notifyMatchFailure(
+ op, "input_layout distributed dimension lane stride times its extent "
+ "must equal the subgroup size");
+ if (targetLaneLayout[0] > *stride)
+ return rewriter.notifyMatchFailure(
+ op, "target_layout effective lane_layout[0] must not exceed the "
+ "input_layout distributed dimension lane stride");
+
+ Location loc = op.getLoc();
+ Value src =
+ castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getSource()),
+ redistribution->distributedInput);
+ auto flatType =
+ VectorType::get({redistribution->distributedInput.getNumElements()},
+ valueType.getElementType());
+ Value flat = vector::ShapeCastOp::create(rewriter, loc, flatType, src);
+ Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+ /*upperBound=*/mlir::IntegerAttr());
+ Value rowCount =
+ arith::ConstantIndexOp::create(rewriter, loc, targetLaneLayout[0]);
+ Value row = arith::RemUIOp::create(rewriter, loc, laneId, rowCount);
+ Value own = vector::ExtractOp::create(rewriter, loc, flat,
+ ArrayRef<OpFoldResult>{row});
+ Value partner =
+ gpu::ShuffleOp::create(
+ rewriter, loc, own,
+ /*offset=*/static_cast<int32_t>(*stride),
+ /*width=*/static_cast<int32_t>(redistribution->subgroupSize),
+ /*mode=*/gpu::ShuffleMode::XOR)
+ .getShuffleResult();
+ rewriter.replaceOpWithNewOp<vector::FromElementsOp>(
+ op, redistribution->distributedTarget, ValueRange{own, partner});
+ return success();
+ }
+};
+
+/// Distributes the slice-attributed `xegpu.convert_layout` whose source is
+/// fully broadcast and whose target splits the two columns of the value over
+/// two lane groups. Each lane keeps one whole column, which is a stride-2
+/// subset of the row-major source, so the two columns are separated with one
+/// `vector.deinterleave` and the lane's group selects between them.
+///
+/// xegpu.convert_layout %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>
+///
+/// The input layout has effective `lane_layout` [1, 1], so the distributed
+/// source is the full `vector<8x2xbf16>`. The target has effective
+/// `lane_layout` [1, 2], so the distributed result is `vector<8x1xbf16>`:
+/// lanes 0..7 keep column 0 and lanes 8..15 keep column 1.
+///
+/// %flat = vector.shape_cast %src : vector<8x2xbf16> to vector<16xbf16>
+/// %even, %odd = vector.deinterleave %flat : vector<16xbf16> ->
+/// vector<8xbf16>
+/// %lane = gpu.lane_id
+/// %group = arith.divui %lane, %c8 : index
+/// %isFirst = arith.cmpi eq, %group, %c0 : index
+/// %sel = arith.select %isFirst, %even, %odd : vector<8xbf16>
+/// %res = vector.shape_cast %sel : vector<8xbf16> to vector<8x1xbf16>
+struct SgToLaneConvertLayoutDeinterleaveSelect
+ : public OpConversionPattern<xegpu::ConvertLayoutOp> {
+ using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(xegpu::ConvertLayoutOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
+ xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
+
+ if (!isa<xegpu::SliceAttr>(inputLayout))
+ return rewriter.notifyMatchFailure(op,
+ "input_layout must be #xegpu.slice");
+ auto targetSlice = dyn_cast<xegpu::SliceAttr>(targetLayout);
+ if (!targetSlice)
+ return rewriter.notifyMatchFailure(op,
+ "target_layout must be #xegpu.slice");
+
+ FailureOr<ElementLaneRedistribution> redistribution =
+ matchElementLaneRedistribution(op, rewriter);
+ if (failed(redistribution))
+ return failure();
+
+ VectorType valueType = redistribution->valueType;
+ if (redistribution->inputLaneLayout != SmallVector<int64_t>{1, 1})
+ return rewriter.notifyMatchFailure(
+ op, "input_layout effective lane_layout must be [1, 1]");
+ SmallVector<int64_t> targetLaneLayout = redistribution->targetLaneLayout;
+ if (targetLaneLayout != SmallVector<int64_t>{1, 2})
+ return rewriter.notifyMatchFailure(
+ op, "target_layout effective lane_layout must be [1, 2]");
+ if (redistribution->distributedInput != valueType)
+ return rewriter.notifyMatchFailure(
+ op, "distributed input_layout type must equal the value type");
+ if (redistribution->distributedTarget.getShape() !=
+ ArrayRef<int64_t>{valueType.getDimSize(0), 1})
+ return rewriter.notifyMatchFailure(
+ op, "distributed target_layout type must be vector<shape[0]x1>");
+
+ FailureOr<int64_t> stride = getDistributedDimLaneStride(targetSlice);
+ if (failed(stride))
+ return rewriter.notifyMatchFailure(
+ op, "target_layout parent must have exactly one non-sliced dimension "
+ "with lane_layout extent greater than one");
+ // The two lane groups must together cover the whole subgroup, so that
+ // `lane_id / stride` is the index of the column the lane keeps.
+ if (*stride * targetLaneLayout[1] != redistribution->subgroupSize)
+ return rewriter.notifyMatchFailure(
+ op, "target_layout distributed dimension lane stride times its "
+ "extent must equal the subgroup size");
+
+ Location loc = op.getLoc();
+ Value src =
+ castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getSource()),
+ redistribution->distributedInput);
+ auto flatType = VectorType::get({valueType.getNumElements()},
+ valueType.getElementType());
+ Value flat = vector::ShapeCastOp::create(rewriter, loc, flatType, src);
+ auto deinterleaved = vector::DeinterleaveOp::create(rewriter, loc, flat);
+ Value laneId = gpu::LaneIdOp::create(rewriter, loc, rewriter.getIndexType(),
+ /*upperBound=*/mlir::IntegerAttr());
+ Value strideVal = arith::ConstantIndexOp::create(rewriter, loc, *stride);
+ Value zero = arith::ConstantIndexOp::create(rewriter, loc, 0);
+ Value group = arith::DivUIOp::create(rewriter, loc, laneId, strideVal);
+ Value isFirstGroup = arith::CmpIOp::create(
+ rewriter, loc, arith::CmpIPredicate::eq, group, zero);
+ Value selected = arith::SelectOp::create(rewriter, loc, isFirstGroup,
+ deinterleaved.getRes1(),
+ deinterleaved.getRes2());
+ rewriter.replaceOpWithNewOp<vector::ShapeCastOp>(
+ op, redistribution->distributedTarget, selected);
+ return success();
+ }
+};
+
// Trivially distribute `vector.interleave`
struct SgToLaneVectorInterleave
: public OpConversionPattern<vector::InterleaveOp> {
@@ -2187,17 +2621,20 @@ void xegpu::populateXeGPUSgToLaneDistributeTypeConversionAndLegality(
return !xegpu::getTemporaryLayout(op->getOpResult(0));
});
target.markUnknownOpDynamicallyLegal([](Operation *op) { return true; });
- patterns.add<
- SgToLaneCreateNdDesc, SgToLaneLoadNd, SgToLaneStoreNd, SgToLaneDpas,
- SgToLaneElementWise, SgToLaneArithConstant, SgToLanePrefetchNd,
- SgToLaneLoadGather, SgToLaneStoreScatter, SgToLaneVectorReduction,
- SgToLaneMultiDimReduction, SgToLaneVectorExtract, SgToLaneVectorInsert,
- SgToLaneVectorExtractStridedSlice, SgToLaneVectorInsertStridedSlice,
- SgToLaneLoadMatrix, SgToLaneStoreMatrix, SgToLaneConvertLayout,
- SgToLaneVectorTranspose, SgToLaneVectorBitcast, SgToLaneVectorStep,
- SgToLaneVectorShapeCast, SgToLaneBroadcast,
- SgToLaneCreateMask<vector::CreateMaskOp>,
- SgToLaneCreateMask<vector::ConstantMaskOp>, SgToLaneVectorDeinterleave,
- SgToLaneVectorInterleave, SgToLaneDpasMx>(typeConverter,
- patterns.getContext());
+ patterns
+ .add<SgToLaneCreateNdDesc, SgToLaneLoadNd, SgToLaneStoreNd, SgToLaneDpas,
+ SgToLaneElementWise, SgToLaneArithConstant, SgToLanePrefetchNd,
+ SgToLaneLoadGather, SgToLaneStoreScatter, SgToLaneVectorReduction,
+ SgToLaneMultiDimReduction, SgToLaneVectorExtract,
+ SgToLaneVectorInsert, SgToLaneVectorExtractStridedSlice,
+ SgToLaneVectorInsertStridedSlice, SgToLaneLoadMatrix,
+ SgToLaneStoreMatrix, SgToLaneConvertLayout, SgToLaneVectorTranspose,
+ SgToLaneVectorBitcast, SgToLaneVectorStep, SgToLaneVectorShapeCast,
+ SgToLaneBroadcast, SgToLaneCreateMask<vector::CreateMaskOp>,
+ SgToLaneCreateMask<vector::ConstantMaskOp>,
+ SgToLaneVectorDeinterleave, SgToLaneVectorInterleave, SgToLaneDpasMx,
+ SgToLaneConvertLayoutBroadcastExtract,
+ SgToLaneConvertLayoutPartialBroadcastExtractShuffle,
+ SgToLaneConvertLayoutDeinterleaveSelect>(typeConverter,
+ patterns.getContext());
}
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
new file mode 100644
index 0000000000000..80c1fd946e563
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
@@ -0,0 +1,101 @@
+// RUN: mlir-opt --xevm-attach-target='module=xevm_* chip=cri' \
+// RUN: --test-xegpu-sg-to-lane-distribute --split-input-file %s | FileCheck %s
+
+// -----
+// The input layout leaves effective lane_layout [1, 1], so the value is fully
+// broadcast: every lane holds all 8 rows. The target spreads the 8 rows over 8
+// lanes, one row each, which is the layout `xegpu.dpas_mx` wants for its scale
+// operands. Lane `l` keeps row `l % 8` and no data crosses lanes, so a single
+// extract suffices.
+//
+// The source is flattened to rank 1 first because `xegpu-vector-linearize`
+// cannot linearize a `vector.extract` with a dynamic position out of a rank-2
+// value.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @broadcast_to_row_per_lane
+// CHECK: %[[SRC:.*]] = "test.some_op"()
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<8x1xf8E8M0FNU> to vector<8xf8E8M0FNU>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[ROWS:.*]] = arith.constant 8 : index
+// CHECK: %[[ROW:.*]] = arith.remui %[[LANE]], %[[ROWS]] : index
+// CHECK: %[[ELEM:.*]] = vector.extract %[[FLAT]][%[[ROW]]] : f8E8M0FNU from vector<8xf8E8M0FNU>
+// CHECK: vector.from_elements %[[ELEM]] : vector<1x1xf8E8M0FNU>
+gpu.func @broadcast_to_row_per_lane() {
+ %src = "test.some_op"() : () -> vector<8x1xf8E8M0FNU>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1], lane_data = [2, 1, 1], order = [0, 2, 1]>, dims = [0]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>
+ }> : vector<8x1xf8E8M0FNU>
+ gpu.return
+}
+}
+
+// -----
+// The input layout leaves effective lane_layout [1, 2], with the distributed
+// dimension at lane stride 8, so the value is broadcast over two lane groups:
+// lane `l` holds all 8 rows of column `l / 8`. The target wants both columns of
+// row `l % 8` in lane `l`, so each lane extracts its own column and shuffles
+// the other one in from lane `l ^ 8`, which holds the same row of the other
+// column.
+//
+// Lanes 8..15 receive the two columns in the opposite order, which is unused
+// because the target lane_layout [8, 1] only occupies 8 lanes.
+//
+// The source `test.some_op` is left undistributed by the pass, hence the cast to
+// the distributed input type; a real, distributed producer needs no cast.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @lane_group_columns_to_row_per_lane
+// CHECK: %[[SRC:.*]] = "test.some_op"()
+// CHECK: %[[DIST:.*]] = builtin.unrealized_conversion_cast %[[SRC]] : vector<8x2xf8E8M0FNU> to vector<8x1xf8E8M0FNU>
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[DIST]] : vector<8x1xf8E8M0FNU> to vector<8xf8E8M0FNU>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[ROWS:.*]] = arith.constant 8 : index
+// CHECK: %[[ROW:.*]] = arith.remui %[[LANE]], %[[ROWS]] : index
+// CHECK: %[[OWN:.*]] = vector.extract %[[FLAT]][%[[ROW]]] : f8E8M0FNU from vector<8xf8E8M0FNU>
+// The offset and width constants are created as arguments of the same call in
+// `gpu::ShuffleOp::build`, so their relative order is up to the host compiler.
+// CHECK-DAG: %[[OFFSET:.*]] = arith.constant 8 : i32
+// CHECK-DAG: %[[WIDTH:.*]] = arith.constant 16 : i32
+// CHECK: %[[PARTNER:.*]], %{{.*}} = gpu.shuffle xor %[[OWN]], %[[OFFSET]], %[[WIDTH]] : f8E8M0FNU
+// CHECK: vector.from_elements %[[OWN]], %[[PARTNER]] : vector<1x2xf8E8M0FNU>
+gpu.func @lane_group_columns_to_row_per_lane() {
+ %src = "test.some_op"() : () -> vector<8x2xf8E8M0FNU>
+ %cvt = xegpu.convert_layout %src
+ <{
+ input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2], lane_data = [4, 1, 1], order = [0, 2, 1]>, dims = [0]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>
+ }> : vector<8x2xf8E8M0FNU>
+ gpu.return
+}
+}
+
+// -----
+// The input layout leaves effective lane_layout [1, 1], so the value is fully
+// broadcast. The target leaves effective lane_layout [1, 2], with the
+// distributed dimension at lane stride 8: lanes 0..7 keep column 0 and lanes
+// 8..15 keep column 1. A column is a stride-2 subset of the row-major value, so
+// `vector.deinterleave` separates the two and the lane group selects which one
+// the lane keeps.
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @broadcast_to_column_per_lane_group
+// CHECK: %[[SRC:.*]] = "test.some_op"()
+// CHECK: %[[FLAT:.*]] = vector.shape_cast %[[SRC]] : vector<8x2xbf16> to vector<16xbf16>
+// CHECK: %[[EVEN:.*]], %[[ODD:.*]] = vector.deinterleave %[[FLAT]] : vector<16xbf16> -> vector<8xbf16>
+// CHECK: %[[LANE:.*]] = gpu.lane_id
+// CHECK: %[[STRIDE:.*]] = arith.constant 8 : index
+// CHECK: %[[ZERO:.*]] = arith.constant 0 : index
+// CHECK: %[[GROUP:.*]] = arith.divui %[[LANE]], %[[STRIDE]] : index
+// CHECK: %[[FIRST:.*]] = arith.cmpi eq, %[[GROUP]], %[[ZERO]] : index
+// CHECK: %[[SEL:.*]] = arith.select %[[FIRST]], %[[EVEN]], %[[ODD]] : vector<8xbf16>
+// CHECK: vector.shape_cast %[[SEL]] : vector<8xbf16> to vector<8x1xbf16>
+gpu.func @broadcast_to_column_per_lane_group() {
+ %src = "test.some_op"() : () -> vector<8x2xbf16>
+ %cvt = xegpu.convert_layout %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>
+ gpu.return
+}
+}
>From 96b326e028d43ed059c8208b9351432fffe62a52 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Fri, 11 Sep 2026 21:28:35 +0000
Subject: [PATCH 2/8] [mlir][xegpu] Rename assignsElementsToLanes, NFC
The name did not say what the predicate checks. It requires rank 2, the default (row major) effective order [1, 0] and effective lane_data [1, 1], and says nothing about lane_layout, which is what the convert_layout conversions redistribute. Spell that out in the name and the comment, including that an unset lane_data is rejected because getEffectiveLaneDataAsInt returns an empty vector for it.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 21 ++++++++++++-------
1 file changed, 13 insertions(+), 8 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 1b1ba85158c4e..30d5eb270be5c 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1869,9 +1869,14 @@ struct SgToLaneConvertLayout
}
};
-/// Check if layout distributes a rank-2 value with unit lane_data
-/// in row major or default order
-static bool assignsElementsToLanes(xegpu::DistributeLayoutAttr layout) {
+/// Returns true if `layout` has rank 2, the default (row major) effective order
+/// [1, 0] and effective lane_data [1, 1], i.e. it maps the elements of the value
+/// to lanes one by one, without packing several of them into a lane. An unset
+/// lane_data is not accepted: `getEffectiveLaneDataAsInt` returns an empty
+/// vector for it, so layouts that are not lane level are rejected here. The
+/// effective lane_layout is deliberately left unconstrained: it is what the
+/// conversions below redistribute.
+static bool hasDefaultOrderAndUnitLaneData(xegpu::DistributeLayoutAttr layout) {
if (layout.getRank() != 2)
return false;
return layout.getEffectiveLaneDataAsInt() == SmallVector<int64_t>{1, 1} &&
@@ -1896,16 +1901,16 @@ struct ElementLaneRedistribution {
/// Recognizes an `xegpu.convert_layout` that redistributes individual elements
/// between the lanes of a subgroup, which is the common condition the three
/// slice-attributed lowerings below have.
-/// Both input and target layouts must satisfy `assignsElementsToLanes` and must
-/// be able to distribute the value, the two can only differ in which lane
-/// owns which element.
+/// Both input and target layouts must satisfy `hasDefaultOrderAndUnitLaneData`
+/// and must be able to distribute the value, the two can only differ in which
+/// lane owns which element.
static FailureOr<ElementLaneRedistribution>
matchElementLaneRedistribution(xegpu::ConvertLayoutOp op,
ConversionPatternRewriter &rewriter) {
xegpu::DistributeLayoutAttr inputLayout = op.getEffectiveInputLayout();
xegpu::DistributeLayoutAttr targetLayout = op.getTargetLayoutAttr();
- if (!assignsElementsToLanes(inputLayout) ||
- !assignsElementsToLanes(targetLayout))
+ if (!hasDefaultOrderAndUnitLaneData(inputLayout) ||
+ !hasDefaultOrderAndUnitLaneData(targetLayout))
return rewriter.notifyMatchFailure(
op, "both layouts must be rank 2 with effective lane_data [1, 1] and "
"effective order [1, 0]");
>From ebf7f552a863beab8079ab48dd7c8197a8448335 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Fri, 11 Sep 2026 21:51:01 +0000
Subject: [PATCH 3/8] [mlir][xegpu] Fix clang-format of
hasDefaultOrderAndUnitLaneData comment, NFC
Reflow the doc comment lines that exceeded 80 columns, as reported by the clang-format CI check.
---
.../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 30d5eb270be5c..aeab3217bc18b 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1870,11 +1870,11 @@ struct SgToLaneConvertLayout
};
/// Returns true if `layout` has rank 2, the default (row major) effective order
-/// [1, 0] and effective lane_data [1, 1], i.e. it maps the elements of the value
-/// to lanes one by one, without packing several of them into a lane. An unset
-/// lane_data is not accepted: `getEffectiveLaneDataAsInt` returns an empty
-/// vector for it, so layouts that are not lane level are rejected here. The
-/// effective lane_layout is deliberately left unconstrained: it is what the
+/// [1, 0] and effective lane_data [1, 1], i.e. it maps the elements of the
+/// value to lanes one by one, without packing several of them into a lane. An
+/// unset lane_data is not accepted: `getEffectiveLaneDataAsInt` returns an
+/// empty vector for it, so layouts that are not lane level are rejected here.
+/// The effective lane_layout is deliberately left unconstrained: it is what the
/// conversions below redistribute.
static bool hasDefaultOrderAndUnitLaneData(xegpu::DistributeLayoutAttr layout) {
if (layout.getRank() != 2)
>From 98abe673288ac5bde280d0a98de708441459a485 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 21 Sep 2026 19:13:48 +0000
Subject: [PATCH 4/8] [mlir][xegpu] Replicate the partial lane layout in
convert_layout gather
SgToLaneConvertLayoutPartialBroadcastExtractShuffle paired the lane-local element with a gpu.shuffle xor partner, so the lanes of the second broadcast group received the two columns in the opposite order. A partial lane layout is not a mask: lane ids are delinearized modulo the lane_layout extents, so those lanes address the same elements as the covered ones and must hold replicas of their values.
Gather column c of row r from lane r + c * stride, which owns it, so every lane runs the same gather and holds the same pair. This also requires the row count to divide the lane stride, which the xor version silently assumed; check it instead.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 57 ++++++++++---------
...-lane-distribute-convert-layout-slice.mlir | 21 +++----
2 files changed, 38 insertions(+), 40 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index aeab3217bc18b..8305fe3f5cea4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2078,9 +2078,10 @@ struct SgToLaneConvertLayoutBroadcastExtract
};
/// Distributes the slice-attributed `xegpu.convert_layout` whose source is
-/// broadcast over two lane groups onto the rows of a subset of the lanes. Each
-/// lane needs one element from its own group and one from the other group, so
-/// one `vector.extract` is paired with one `gpu.shuffle xor`.
+/// broadcast over two lane groups onto the rows of a subset of the lanes.
+/// Column `c` of row `r` is owned by lane `r + c * stride`, so each lane
+/// extracts its own element and gathers the columns of its row with one
+/// `gpu.shuffle idx` per column.
///
/// xegpu.convert_layout %src
/// <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2],
@@ -2094,24 +2095,14 @@ struct SgToLaneConvertLayoutBroadcastExtract
/// The input layout has effective `lane_layout` [1, 2], so the distributed
/// source is `vector<8x1xf8E8M0FNU>`: lane `l` holds all 8 rows of column
/// `l / 8`. The target puts both columns of one row in a single lane, so the
-/// distributed result is `vector<1x2xf8E8M0FNU>`. Lane `l` owns `[l % 8, 0]`
-/// of its own column and gets the other column from lane `l ^ 8`, which holds
-/// the same row:
-///
-/// %flat = vector.shape_cast %src : vector<8x1xf8E8M0FNU> to
-/// vector<8xf8E8M0FNU>
-/// %lane = gpu.lane_id
-/// %row = arith.remui %lane, %c8 : index
-/// %own = vector.extract %flat[%row] : f8E8M0FNU from
-/// vector<8xf8E8M0FNU>
-/// %part, %v = gpu.shuffle xor %own, %c8_i32, %c16_i32 : f8E8M0FNU
-/// %res = vector.from_elements %own, %part : vector<1x2xf8E8M0FNU>
+/// distributed result is `vector<1x2xf8E8M0FNU>` holding row `l % 8`.
///
/// The source is flattened first because `xegpu-vector-linearize` cannot
/// linearize a `vector.extract` with a dynamic position out of a rank-2 value.
///
-/// Lanes 8..15 receive the two columns in the opposite order, which is unused
-/// because the target `lane_layout` [8, 1] only occupies 8 lanes.
+/// All lanes gather every column, including the one they already hold, so that
+/// lanes outside the target `lane_layout` hold replicas as partial lane layouts
+/// require.
struct SgToLaneConvertLayoutPartialBroadcastExtractShuffle
: public OpConversionPattern<xegpu::ConvertLayoutOp> {
using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
@@ -2165,15 +2156,18 @@ struct SgToLaneConvertLayoutPartialBroadcastExtractShuffle
op, "input_layout parent must have exactly one non-sliced dimension "
"with lane_layout extent greater than one");
// The two broadcast groups must together cover the whole subgroup, so that
- // `lane_id ^ stride` is the lane holding the other column of the same row.
+ // lane `r + c * stride` is the lane holding column `c` of row `r`.
if (*stride * redistribution->inputLaneLayout[1] !=
redistribution->subgroupSize)
return rewriter.notifyMatchFailure(
op, "input_layout distributed dimension lane stride times its extent "
"must equal the subgroup size");
- if (targetLaneLayout[0] > *stride)
+ // Lane `r + c * stride` must hold row `r`, i.e. (r + stride) %
+ // lane_layout[0] must be r, which requires the row count to divide the
+ // stride.
+ if (*stride % targetLaneLayout[0] != 0)
return rewriter.notifyMatchFailure(
- op, "target_layout effective lane_layout[0] must not exceed the "
+ op, "target_layout effective lane_layout[0] must divide the "
"input_layout distributed dimension lane stride");
Location loc = op.getLoc();
@@ -2191,15 +2185,22 @@ struct SgToLaneConvertLayoutPartialBroadcastExtractShuffle
Value row = arith::RemUIOp::create(rewriter, loc, laneId, rowCount);
Value own = vector::ExtractOp::create(rewriter, loc, flat,
ArrayRef<OpFoldResult>{row});
- Value partner =
- gpu::ShuffleOp::create(
- rewriter, loc, own,
- /*offset=*/static_cast<int32_t>(*stride),
- /*width=*/static_cast<int32_t>(redistribution->subgroupSize),
- /*mode=*/gpu::ShuffleMode::XOR)
- .getShuffleResult();
+ // Gather column `c` of `row` from lane `row + c * stride`, which owns it.
+ Type i32Type = rewriter.getI32Type();
+ Value width = arith::ConstantIntOp::create(rewriter, loc, i32Type,
+ redistribution->subgroupSize);
+ Value rowI32 = arith::IndexCastOp::create(rewriter, loc, i32Type, row);
+ Value strideI32 =
+ arith::ConstantIntOp::create(rewriter, loc, i32Type, *stride);
+ Value otherOwner = arith::AddIOp::create(rewriter, loc, rowI32, strideI32);
+ Value column0 = gpu::ShuffleOp::create(rewriter, loc, own, rowI32, width,
+ gpu::ShuffleMode::IDX)
+ .getShuffleResult();
+ Value column1 = gpu::ShuffleOp::create(rewriter, loc, own, otherOwner,
+ width, gpu::ShuffleMode::IDX)
+ .getShuffleResult();
rewriter.replaceOpWithNewOp<vector::FromElementsOp>(
- op, redistribution->distributedTarget, ValueRange{own, partner});
+ op, redistribution->distributedTarget, ValueRange{column0, column1});
return success();
}
};
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
index 80c1fd946e563..a6af2b8a689b5 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
@@ -35,12 +35,8 @@ gpu.func @broadcast_to_row_per_lane() {
// The input layout leaves effective lane_layout [1, 2], with the distributed
// dimension at lane stride 8, so the value is broadcast over two lane groups:
// lane `l` holds all 8 rows of column `l / 8`. The target wants both columns of
-// row `l % 8` in lane `l`, so each lane extracts its own column and shuffles
-// the other one in from lane `l ^ 8`, which holds the same row of the other
-// column.
-//
-// Lanes 8..15 receive the two columns in the opposite order, which is unused
-// because the target lane_layout [8, 1] only occupies 8 lanes.
+// row `l % 8` in lane `l`, and column `c` of a row `r` is owned by lane
+// `r + c * 8`, so every lane gathers both columns from there.
//
// The source `test.some_op` is left undistributed by the pass, hence the cast to
// the distributed input type; a real, distributed producer needs no cast.
@@ -53,12 +49,13 @@ gpu.module @xevm_module {
// CHECK: %[[ROWS:.*]] = arith.constant 8 : index
// CHECK: %[[ROW:.*]] = arith.remui %[[LANE]], %[[ROWS]] : index
// CHECK: %[[OWN:.*]] = vector.extract %[[FLAT]][%[[ROW]]] : f8E8M0FNU from vector<8xf8E8M0FNU>
-// The offset and width constants are created as arguments of the same call in
-// `gpu::ShuffleOp::build`, so their relative order is up to the host compiler.
-// CHECK-DAG: %[[OFFSET:.*]] = arith.constant 8 : i32
-// CHECK-DAG: %[[WIDTH:.*]] = arith.constant 16 : i32
-// CHECK: %[[PARTNER:.*]], %{{.*}} = gpu.shuffle xor %[[OWN]], %[[OFFSET]], %[[WIDTH]] : f8E8M0FNU
-// CHECK: vector.from_elements %[[OWN]], %[[PARTNER]] : vector<1x2xf8E8M0FNU>
+// CHECK: %[[WIDTH:.*]] = arith.constant 16 : i32
+// CHECK: %[[ROWI:.*]] = arith.index_cast %[[ROW]] : index to i32
+// CHECK: %[[STRIDE:.*]] = arith.constant 8 : i32
+// CHECK: %[[OTHER:.*]] = arith.addi %[[ROWI]], %[[STRIDE]] : i32
+// CHECK: %[[COL0:.*]], %{{.*}} = gpu.shuffle idx %[[OWN]], %[[ROWI]], %[[WIDTH]] : f8E8M0FNU
+// CHECK: %[[COL1:.*]], %{{.*}} = gpu.shuffle idx %[[OWN]], %[[OTHER]], %[[WIDTH]] : f8E8M0FNU
+// CHECK: vector.from_elements %[[COL0]], %[[COL1]] : vector<1x2xf8E8M0FNU>
gpu.func @lane_group_columns_to_row_per_lane() {
%src = "test.some_op"() : () -> vector<8x2xf8E8M0FNU>
%cvt = xegpu.convert_layout %src
>From 44e3f5622c182643a623045feb5714a85b188d09 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 21 Sep 2026 19:15:23 +0000
Subject: [PATCH 5/8] [mlir][xegpu] Trim convert_layout distribution comments,
NFC
The three slice-attributed convert_layout patterns documented both their input IR, which the test file already contains and verifies, and the IR they generate, which restates the rewrite right below it. Drop both and keep what the code does not say: what each conversion does, how the effective lane_layouts map elements to lanes, and the few whys.
In the test, replace the prose that repeated the implementation with a one line description of the case and the pattern that handles it.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 61 +++----------------
...-lane-distribute-convert-layout-slice.mlir | 27 +++-----
2 files changed, 15 insertions(+), 73 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 8305fe3f5cea4..2143e5451514a 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1993,26 +1993,9 @@ static FailureOr<int64_t> getDistributedDimLaneStride(xegpu::SliceAttr slice) {
/// subset of the lanes. Each lane keeps a single element, so no data crosses
/// lanes and one `vector.extract` suffices.
///
-/// xegpu.convert_layout %src
-/// <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [16, 1, 1],
-/// lane_data = [2, 1, 1],
-/// order = [0, 2, 1]>,
-/// dims = [0]>,
-/// target_layout = #xegpu.layout<lane_layout = [8, 1],
-/// lane_data = [1, 1]>}>
-/// : vector<8x1xf8E8M0FNU>
-///
/// The input layout has effective `lane_layout` [1, 1], so the distributed
-/// source is the full `vector<8x1xf8E8M0FNU>`. The target spreads the 8 rows
-/// over 8 lanes, one row each, so the distributed result is
-/// `vector<1x1xf8E8M0FNU>` holding row `lane_id % 8`:
-///
-/// %flat = vector.shape_cast %src : vector<8x1xf8E8M0FNU> to
-/// vector<8xf8E8M0FNU>
-/// %lane = gpu.lane_id
-/// %row = arith.remui %lane, %c8 : index
-/// %elem = vector.extract %flat[%row] : f8E8M0FNU from vector<8xf8E8M0FNU>
-/// %res = vector.from_elements %elem : vector<1x1xf8E8M0FNU>
+/// source is the whole value; the target has [n, 1], so lane `l` keeps row
+/// `l % n` and the distributed result is `vector<1x1>`.
///
/// The source is flattened first because `xegpu-vector-linearize` cannot
/// linearize a `vector.extract` with a dynamic position out of a rank-2 value.
@@ -2083,19 +2066,10 @@ struct SgToLaneConvertLayoutBroadcastExtract
/// extracts its own element and gathers the columns of its row with one
/// `gpu.shuffle idx` per column.
///
-/// xegpu.convert_layout %src
-/// <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2],
-/// lane_data = [4, 1, 1],
-/// order = [0, 2, 1]>,
-/// dims = [0]>,
-/// target_layout = #xegpu.layout<lane_layout = [8, 1],
-/// lane_data = [1, 1]>}>
-/// : vector<8x2xf8E8M0FNU>
-///
/// The input layout has effective `lane_layout` [1, 2], so the distributed
-/// source is `vector<8x1xf8E8M0FNU>`: lane `l` holds all 8 rows of column
-/// `l / 8`. The target puts both columns of one row in a single lane, so the
-/// distributed result is `vector<1x2xf8E8M0FNU>` holding row `l % 8`.
+/// source holds one column per lane group; the target has [n, 1], so lane `l`
+/// keeps row `l % n` of both columns and the distributed result is
+/// `vector<1x2>`.
///
/// The source is flattened first because `xegpu-vector-linearize` cannot
/// linearize a `vector.extract` with a dynamic position out of a rank-2 value.
@@ -2211,29 +2185,10 @@ struct SgToLaneConvertLayoutPartialBroadcastExtractShuffle
/// subset of the row-major source, so the two columns are separated with one
/// `vector.deinterleave` and the lane's group selects between them.
///
-/// xegpu.convert_layout %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>
-///
/// The input layout has effective `lane_layout` [1, 1], so the distributed
-/// source is the full `vector<8x2xbf16>`. The target has effective
-/// `lane_layout` [1, 2], so the distributed result is `vector<8x1xbf16>`:
-/// lanes 0..7 keep column 0 and lanes 8..15 keep column 1.
-///
-/// %flat = vector.shape_cast %src : vector<8x2xbf16> to vector<16xbf16>
-/// %even, %odd = vector.deinterleave %flat : vector<16xbf16> ->
-/// vector<8xbf16>
-/// %lane = gpu.lane_id
-/// %group = arith.divui %lane, %c8 : index
-/// %isFirst = arith.cmpi eq, %group, %c0 : index
-/// %sel = arith.select %isFirst, %even, %odd : vector<8xbf16>
-/// %res = vector.shape_cast %sel : vector<8xbf16> to vector<8x1xbf16>
+/// source is the whole value; the target has [1, 2], so the lanes below the
+/// group stride keep column 0, the rest column 1, and the distributed result
+/// has one column per lane.
struct SgToLaneConvertLayoutDeinterleaveSelect
: public OpConversionPattern<xegpu::ConvertLayoutOp> {
using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
diff --git a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
index a6af2b8a689b5..21ff2bac0d7d7 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
@@ -2,15 +2,9 @@
// RUN: --test-xegpu-sg-to-lane-distribute --split-input-file %s | FileCheck %s
// -----
-// The input layout leaves effective lane_layout [1, 1], so the value is fully
-// broadcast: every lane holds all 8 rows. The target spreads the 8 rows over 8
-// lanes, one row each, which is the layout `xegpu.dpas_mx` wants for its scale
-// operands. Lane `l` keeps row `l % 8` and no data crosses lanes, so a single
-// extract suffices.
-//
-// The source is flattened to rank 1 first because `xegpu-vector-linearize`
-// cannot linearize a `vector.extract` with a dynamic position out of a rank-2
-// value.
+// Fully broadcast source to one row per lane, the layout `xegpu.dpas_mx` wants
+// for its scale operands.
+// Handled by SgToLaneConvertLayoutBroadcastExtract.
gpu.module @xevm_module {
// CHECK-LABEL: gpu.func @broadcast_to_row_per_lane
// CHECK: %[[SRC:.*]] = "test.some_op"()
@@ -32,11 +26,8 @@ gpu.func @broadcast_to_row_per_lane() {
}
// -----
-// The input layout leaves effective lane_layout [1, 2], with the distributed
-// dimension at lane stride 8, so the value is broadcast over two lane groups:
-// lane `l` holds all 8 rows of column `l / 8`. The target wants both columns of
-// row `l % 8` in lane `l`, and column `c` of a row `r` is owned by lane
-// `r + c * 8`, so every lane gathers both columns from there.
+// Source broadcast over two lane groups to one row per lane.
+// Handled by SgToLaneConvertLayoutPartialBroadcastExtractShuffle.
//
// The source `test.some_op` is left undistributed by the pass, hence the cast to
// the distributed input type; a real, distributed producer needs no cast.
@@ -68,12 +59,8 @@ gpu.func @lane_group_columns_to_row_per_lane() {
}
// -----
-// The input layout leaves effective lane_layout [1, 1], so the value is fully
-// broadcast. The target leaves effective lane_layout [1, 2], with the
-// distributed dimension at lane stride 8: lanes 0..7 keep column 0 and lanes
-// 8..15 keep column 1. A column is a stride-2 subset of the row-major value, so
-// `vector.deinterleave` separates the two and the lane group selects which one
-// the lane keeps.
+// Fully broadcast source to one column per lane group.
+// Handled by SgToLaneConvertLayoutDeinterleaveSelect.
gpu.module @xevm_module {
// CHECK-LABEL: gpu.func @broadcast_to_column_per_lane_group
// CHECK: %[[SRC:.*]] = "test.some_op"()
>From b4c8b868057613ea478173e27ff23e8d8c67c23e Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 21 Sep 2026 19:28:13 +0000
Subject: [PATCH 6/8] [mlir][xegpu] Un-XFAIL the WG mx-fp quantize-A
integration test
The convert_layout distributions added here are what the test was missing: the workgroup level mx-fp GEMM with in-kernel A quantization now lowers through xegpu-sg-to-lane-distribute, so the compile-only RUN line passes and the test XPASSes without this change.
---
.../Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir | 2 --
1 file changed, 2 deletions(-)
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
index 2766214459f05..7bf6d5aa71fdd 100644
--- a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
@@ -7,8 +7,6 @@
// RUN-DISABLED: --entry-point-result=void \
// RUN-DISABLED: | FileCheck --check-prefix=MISMATCH %s
-// XFAIL: *
-
// Note: layouts used by dpas_mx need to match HW constaint. Otherwise dpas_mx is not unrolled.
#a = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 64], lane_layout = [1, 16], lane_data = [1, 4]>
#a_ld = #xegpu.layout<sg_layout = [2, 2], sg_data = [16, 1024], inst_data = [8, 16], lane_layout = [1, 16], lane_data = [1, 1]>
>From 14e2d2d81ea02212a72dd7cb4240b3033fe15142 Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Mon, 21 Sep 2026 22:54:00 +0000
Subject: [PATCH 7/8] [mlir][xegpu] Clarify lane stride wording in
convert_layout comments, NFC
Anchor the term stride where it is first used, and drop the hasDefaultOrderAndUnitLaneData comment down to what its name does not already say.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 22 ++++++++-----------
1 file changed, 9 insertions(+), 13 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 5858b079891db..81a2e02a97fb5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1857,13 +1857,8 @@ struct SgToLaneConvertLayout
}
};
-/// Returns true if `layout` has rank 2, the default (row major) effective order
-/// [1, 0] and effective lane_data [1, 1], i.e. it maps the elements of the
-/// value to lanes one by one, without packing several of them into a lane. An
-/// unset lane_data is not accepted: `getEffectiveLaneDataAsInt` returns an
-/// empty vector for it, so layouts that are not lane level are rejected here.
-/// The effective lane_layout is deliberately left unconstrained: it is what the
-/// conversions below redistribute.
+/// `getEffectiveLaneDataAsInt` is empty when `lane_data` is unset, so the unit
+/// check also rejects layouts that are not lane level.
static bool hasDefaultOrderAndUnitLaneData(xegpu::DistributeLayoutAttr layout) {
if (layout.getRank() != 2)
return false;
@@ -2050,9 +2045,10 @@ struct SgToLaneConvertLayoutBroadcastExtract
/// Distributes the slice-attributed `xegpu.convert_layout` whose source is
/// broadcast over two lane groups onto the rows of a subset of the lanes.
-/// Column `c` of row `r` is owned by lane `r + c * stride`, so each lane
-/// extracts its own element and gathers the columns of its row with one
-/// `gpu.shuffle idx` per column.
+/// Column `c` of row `r` is owned by lane `r + c * stride`, where `stride` is
+/// the lane stride of the input's distributed dimension (see
+/// `getDistributedDimLaneStride`), so each lane extracts its own element and
+/// gathers the columns of its row with one `gpu.shuffle idx` per column.
///
/// The input layout has effective `lane_layout` [1, 2], so the distributed
/// source holds one column per lane group; the target has [n, 1], so lane `l`
@@ -2174,9 +2170,9 @@ struct SgToLaneConvertLayoutPartialBroadcastExtractShuffle
/// `vector.deinterleave` and the lane's group selects between them.
///
/// The input layout has effective `lane_layout` [1, 1], so the distributed
-/// source is the whole value; the target has [1, 2], so the lanes below the
-/// group stride keep column 0, the rest column 1, and the distributed result
-/// has one column per lane.
+/// source is the whole value; the target has [1, 2] and its two groups tile the
+/// subgroup, so the first half of the lanes keeps column 0 and the second half
+/// column 1, one column per lane.
struct SgToLaneConvertLayoutDeinterleaveSelect
: public OpConversionPattern<xegpu::ConvertLayoutOp> {
using OpConversionPattern<xegpu::ConvertLayoutOp>::OpConversionPattern;
>From fc424943e44150e6492420aca6360fe7df14298b Mon Sep 17 00:00:00 2001
From: "Lee, Sang Ik" <sang.ik.lee at intel.com>
Date: Wed, 23 Sep 2026 20:50:17 +0000
Subject: [PATCH 8/8] [mlir][xegpu] Restore the convert_layout gather example,
NFC
Put back the worked example for the partial broadcast gather, updated to the two shuffle sequence the pattern now emits.
---
.../Transforms/XeGPUSgToLaneDistribute.cpp | 23 +++++++++++++++++++
1 file changed, 23 insertions(+)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 81a2e02a97fb5..c3d99bf1d4bee 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -2055,6 +2055,29 @@ struct SgToLaneConvertLayoutBroadcastExtract
/// keeps row `l % n` of both columns and the distributed result is
/// `vector<1x2>`.
///
+/// xegpu.convert_layout %src
+/// <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [8, 1, 2],
+/// lane_data = [4, 1, 1],
+/// order = [0, 2, 1]>,
+/// dims = [0]>,
+/// target_layout = #xegpu.layout<lane_layout = [8, 1],
+/// lane_data = [1, 1]>}>
+/// : vector<8x2xf8E8M0FNU>
+///
+/// becomes, with lane `l` holding all 8 rows of column `l / 8`:
+///
+/// %flat = vector.shape_cast %src : vector<8x1xf8E8M0FNU> to
+/// vector<8xf8E8M0FNU>
+/// %lane = gpu.lane_id
+/// %row = arith.remui %lane, %c8 : index
+/// %own = vector.extract %flat[%row] : f8E8M0FNU from
+/// vector<8xf8E8M0FNU>
+/// %rowI32 = arith.index_cast %row : index to i32
+/// %owner1 = arith.addi %rowI32, %c8_i32 : i32
+/// %col0, %v0 = gpu.shuffle idx %own, %rowI32, %c16_i32 : f8E8M0FNU
+/// %col1, %v1 = gpu.shuffle idx %own, %owner1, %c16_i32 : f8E8M0FNU
+/// %res = vector.from_elements %col0, %col1 : vector<1x2xf8E8M0FNU>
+///
/// The source is flattened first because `xegpu-vector-linearize` cannot
/// linearize a `vector.extract` with a dynamic position out of a rank-2 value.
///
More information about the Mlir-commits
mailing list