[Mlir-commits] [mlir] [mlir][xegpu] Distribute three more convert_layout cases to lanes (PR #222693)
Sang Ik Lee
llvmlistbot at llvm.org
Fri Sep 11 14:51:45 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/3] [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/3] [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/3] [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)
More information about the Mlir-commits
mailing list