[Mlir-commits] [mlir] 833166e - [mlir][xegpu] Distribute three more convert_layout cases to lanes (#222693)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Thu Sep 24 08:04:42 PDT 2026
Author: Sang Ik Lee
Date: 2026-09-24T08:04:35-07:00
New Revision: 833166e22fe20c3513a77a02e34385c16128ea31
URL: https://github.com/llvm/llvm-project/commit/833166e22fe20c3513a77a02e34385c16128ea31
DIFF: https://github.com/llvm/llvm-project/commit/833166e22fe20c3513a77a02e34385c16128ea31.diff
LOG: [mlir][xegpu] Distribute three more convert_layout cases to lanes (#222693)
`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 two `gpu.shuffle`
- fully broadcast source to one column per lane group:
`vector.deinterleave` plus a select on the lane group
Added:
mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
Modified:
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm_quantizeA_F4.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
index 8bc8a56bc3991..06f769cefb4fd 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToLaneDistribute.cpp
@@ -1939,6 +1939,420 @@ struct SgToLaneConvertLayout
}
};
+/// `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;
+ 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
diff er, 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 `hasDefaultOrderAndUnitLaneData`
+/// and must be able to distribute the value, the two can only
diff er 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 (!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]");
+
+ 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.
+///
+/// The input layout has effective `lane_layout` [1, 1], so the distributed
+/// 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.
+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.
+/// 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`
+/// 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.
+///
+/// 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;
+
+ 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 `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");
+ // 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 divide 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});
+ // 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{column0, column1});
+ 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.
+///
+/// The input layout has effective `lane_layout` [1, 1], so the distributed
+/// 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;
+
+ 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> {
@@ -2257,17 +2671,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..21ff2bac0d7d7
--- /dev/null
+++ b/mlir/test/Dialect/XeGPU/sg-to-lane-distribute-convert-layout-slice.mlir
@@ -0,0 +1,85 @@
+// RUN: mlir-opt --xevm-attach-target='module=xevm_* chip=cri' \
+// RUN: --test-xegpu-sg-to-lane-distribute --split-input-file %s | FileCheck %s
+
+// -----
+// 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"()
+// 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
+}
+}
+
+// -----
+// 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.
+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>
+// 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
+ <{
+ 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
+}
+}
+
+// -----
+// 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"()
+// 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
+}
+}
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]>
More information about the Mlir-commits
mailing list