[Mlir-commits] [mlir] 715492e - [MLIR][XeGPU] Add wg-to-sg distirbution for dpasmx, bitcast, interleave, and deinterleave (#194985)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Wed May 6 10:34:50 PDT 2026
Author: Jianhui Li
Date: 2026-05-06T10:34:45-07:00
New Revision: 715492eb036a16d68ffbfa3040da07b68fbe7ee8
URL: https://github.com/llvm/llvm-project/commit/715492eb036a16d68ffbfa3040da07b68fbe7ee8
DIFF: https://github.com/llvm/llvm-project/commit/715492eb036a16d68ffbfa3040da07b68fbe7ee8.diff
LOG: [MLIR][XeGPU] Add wg-to-sg distirbution for dpasmx, bitcast, interleave, and deinterleave (#194985)
As title.
Assisted by Claude
Added:
Modified:
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 820ec11ee4e95..4cab1e24bf9e6 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -127,9 +127,10 @@ static xegpu::DistributeLayoutAttr getLayoutFromUsePoints(Value result) {
// For regular operations: First the result layouts are propagated from uses.
// Then the result layouts are propagated to uses (operands).
static void propagateResultsToRegularOperands(Operation *op) {
- if (op->getNumResults() == 0 || op->getNumResults() > 1)
+ if (op->getNumResults() == 0)
+ return;
+ if (op->getNumResults() > 1 && !isa<vector::DeinterleaveOp>(op))
return;
-
OpResult result = op->getResult(0);
xegpu::DistributeLayoutAttr resLayout = getLayoutFromUsePoints(result);
Type resultType = result.getType();
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 8aa0758943cd1..119ec59daf765 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -325,9 +325,9 @@ struct WgToSgDpasOp : public OpConversionPattern<xegpu::DpasOp> {
}
ArrayRef<int64_t> aVecShape =
- llvm::cast<VectorType>(aVec.getType()).getShape();
+ cast<VectorType>(aVec.getType()).getShape();
ArrayRef<int64_t> bVecShape =
- llvm::cast<VectorType>(bVec.getType()).getShape();
+ cast<VectorType>(bVec.getType()).getShape();
VectorType resTy = VectorType::get({aVecShape[0], bVecShape[1]},
resultTy.getElementType());
auto newDpasOp = xegpu::DpasOp::create(rewriter, loc, resTy, operands);
@@ -343,6 +343,58 @@ struct WgToSgDpasOp : public OpConversionPattern<xegpu::DpasOp> {
}
};
+/// This pattern transforms the DpasMxOp to work at subgroup level.
+struct WgToSgDpasMxOp : public OpConversionPattern<xegpu::DpasMxOp> {
+ using OpConversionPattern<xegpu::DpasMxOp>::OpConversionPattern;
+ LogicalResult
+ matchAndRewrite(xegpu::DpasMxOp op, OneToNOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+
+ Location loc = op.getLoc();
+ VectorType resultTy = op.getResult().getType();
+
+ if (resultTy.getRank() != 2)
+ return failure();
+
+ auto layoutCd = op.getLayoutCdAttr();
+ auto layoutA = op.getLayoutAAttr();
+ auto layoutB = op.getLayoutBAttr();
+ auto layoutAScale = op.getLayoutAScaleAttr();
+ auto layoutBScale = op.getLayoutBScaleAttr();
+
+ if (!layoutCd || !layoutA || !layoutB || !layoutAScale || !layoutBScale)
+ return failure();
+
+ size_t index_c = 0;
+ SmallVector<Value> newDpasMxOps;
+ for (auto [index_a, aVec] : llvm::enumerate(adaptor.getA())) {
+ for (auto [index_b, bVec] : llvm::enumerate(adaptor.getB())) {
+ Value accVal = (op.getAcc()) ? adaptor.getAcc()[index_c++] : Value();
+ Value scaleAVal =
+ (op.getScaleA()) ? adaptor.getScaleA()[index_a] : Value();
+ Value scaleBVal =
+ (op.getScaleB()) ? adaptor.getScaleB()[index_b] : Value();
+
+ ArrayRef<int64_t> aVecShape =
+ cast<VectorType>(aVec.getType()).getShape();
+ ArrayRef<int64_t> bVecShape =
+ cast<VectorType>(bVec.getType()).getShape();
+ VectorType resTy = VectorType::get({aVecShape[0], bVecShape[1]},
+ resultTy.getElementType());
+ auto newDpasMxOp = xegpu::DpasMxOp::create(
+ rewriter, loc, resTy, aVec, bVec, accVal, scaleAVal, scaleBVal,
+ layoutA.dropSgLayoutAndData(), layoutB.dropSgLayoutAndData(),
+ layoutCd.dropSgLayoutAndData(), layoutAScale.dropSgLayoutAndData(),
+ layoutBScale.dropSgLayoutAndData());
+
+ newDpasMxOps.push_back(newDpasMxOp);
+ }
+ }
+ rewriter.replaceOpWithMultiple(op, {newDpasMxOps});
+ return success();
+ }
+};
+
/// This pattern transforms vector.broadcast ops to work at subgroup level.
struct WgToSgVectorBroadcastOp
: public OpConversionPattern<vector::BroadcastOp> {
@@ -1403,19 +1455,111 @@ struct WgToSgVectorMaskOp : public OpConversionPattern<MaskOpType> {
using WgToSgVectorConstantMaskOp = WgToSgVectorMaskOp<vector::ConstantMaskOp>;
using WgToSgVectorCreateMaskOp = WgToSgVectorMaskOp<vector::CreateMaskOp>;
+
+// This pattern transforms vector.bitcast ops to work at subgroup level.
+struct WgToSgVectorBitCastOp : public OpConversionPattern<vector::BitCastOp> {
+ using OpConversionPattern<vector::BitCastOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::BitCastOp op, OneToNOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ VectorType resultType = op.getResultVectorType();
+
+ ArrayRef<int64_t> wgShape = resultType.getShape();
+ xegpu::DistributeLayoutAttr layout =
+ xegpu::getTemporaryLayout(dyn_cast<OpResult>(op.getResult()));
+ if (!layout || !layout.isForWorkgroup())
+ return failure();
+
+ SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;
+ VectorType newResultType =
+ VectorType::get(sgShape, resultType.getElementType());
+
+ SmallVector<Value> newBitCastOps;
+ for (auto src : adaptor.getSource()) {
+ auto newBitCast =
+ vector::BitCastOp::create(rewriter, op.getLoc(), newResultType, src);
+ newBitCastOps.push_back(newBitCast.getResult());
+ }
+
+ rewriter.replaceOpWithMultiple(op, {newBitCastOps});
+ return success();
+ }
+};
+
+// This pattern transforms vector.interleave ops to work at subgroup level.
+struct WgToSgVectorInterleaveOp
+ : public OpConversionPattern<vector::InterleaveOp> {
+ using OpConversionPattern<vector::InterleaveOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::InterleaveOp op, OneToNOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ VectorType resultType = op.getResultVectorType();
+
+ ArrayRef<int64_t> wgShape = resultType.getShape();
+ xegpu::DistributeLayoutAttr layout =
+ xegpu::getTemporaryLayout(dyn_cast<OpResult>(op.getResult()));
+ if (!layout || !layout.isForWorkgroup())
+ return failure();
+
+ SmallVector<int64_t> sgShape = getSgShapeAndCount(wgShape, layout).first;
+ VectorType newResultType =
+ VectorType::get(sgShape, resultType.getElementType());
+
+ SmallVector<Value> newInterleaveOps;
+ // Interleave operates pairwise: each lhs value is interleaved with
+ // corresponding rhs value
+ for (auto [lhs, rhs] : llvm::zip(adaptor.getLhs(), adaptor.getRhs())) {
+ auto newInterleave = vector::InterleaveOp::create(
+ rewriter, op.getLoc(), newResultType, lhs, rhs);
+ newInterleaveOps.push_back(newInterleave.getResult());
+ }
+
+ rewriter.replaceOpWithMultiple(op, {newInterleaveOps});
+ return success();
+ }
+};
+
+// This pattern transforms vector.deinterleave ops to work at subgroup level.
+struct WgToSgVectorDeinterleaveOp
+ : public OpConversionPattern<vector::DeinterleaveOp> {
+ using OpConversionPattern<vector::DeinterleaveOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::DeinterleaveOp op, OneToNOpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ SmallVector<Value> newRes1Ops;
+ SmallVector<Value> newRes2Ops;
+
+ for (auto src : adaptor.getSource()) {
+ auto newDeinterleave =
+ vector::DeinterleaveOp::create(rewriter, op.getLoc(), src);
+ newRes1Ops.push_back(newDeinterleave.getRes1());
+ newRes2Ops.push_back(newDeinterleave.getRes2());
+ }
+
+ SmallVector<SmallVector<Value>> results = {newRes1Ops, newRes2Ops};
+ rewriter.replaceOpWithMultiple(op, results);
+ return success();
+ }
+};
+
} // namespace
namespace mlir {
namespace xegpu {
void populateXeGPUWgToSgDistributePatterns(RewritePatternSet &patterns) {
patterns.add<WgToSgCreateNdOp, WgToSgLoadNdOp, WgToSgStoreNdOp, WgToSgDpasOp,
- WgToSgPrefetchNdOp, UnrealizedConversionCastOpPattern,
- WgToSgElementwiseOp, WgToSgVectorBroadcastOp,
- WgToSgConvertLayoutOp, WgToSgArithConstantOp, WgToSgLoadGatherOp,
- WgToSgStoreScatterOp, WgToSgLoadMatrixOp, WgToSgStoreMatrixOp,
- WgToSgVectorStepOp, WgToSgVectorShapeCastOp,
- WgToSgMultiDimReductionOp, WgToSgVectorTransposeOp,
- WgToSgVectorConstantMaskOp, WgToSgVectorCreateMaskOp>(
+ WgToSgDpasMxOp, WgToSgPrefetchNdOp,
+ UnrealizedConversionCastOpPattern, WgToSgElementwiseOp,
+ WgToSgVectorBroadcastOp, WgToSgConvertLayoutOp,
+ WgToSgArithConstantOp, WgToSgLoadGatherOp, WgToSgStoreScatterOp,
+ WgToSgLoadMatrixOp, WgToSgStoreMatrixOp, WgToSgVectorStepOp,
+ WgToSgVectorShapeCastOp, WgToSgMultiDimReductionOp,
+ WgToSgVectorTransposeOp, WgToSgVectorConstantMaskOp,
+ WgToSgVectorCreateMaskOp, WgToSgVectorBitCastOp,
+ WgToSgVectorInterleaveOp, WgToSgVectorDeinterleaveOp>(
patterns.getContext());
}
} // namespace xegpu
@@ -1539,6 +1683,12 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
return isLegal(layout);
});
+ target.addDynamicallyLegalOp<xegpu::DpasMxOp>(
+ [=](xegpu::DpasMxOp op) -> bool {
+ auto layout = op.getLayoutCdAttr();
+ return isLegal(layout);
+ });
+
target.addDynamicallyLegalOp<xegpu::LoadMatrixOp>(
[=](xegpu::LoadMatrixOp op) -> bool {
return isLegal(op.getLayoutAttr());
@@ -1560,16 +1710,16 @@ void XeGPUWgToSgDistributePass::runOnOperation() {
return isLegal(layout);
});
- target.addDynamicallyLegalOp<vector::ShapeCastOp, vector::StepOp,
- vector::TransposeOp, vector::BroadcastOp,
- vector::MultiDimReductionOp,
- vector::ConstantMaskOp, vector::CreateMaskOp>(
- [=](Operation *op) -> bool {
- // Check for either a SliceAttr or LayoutAttr on the result.
- auto layout =
- xegpu::getTemporaryLayout(dyn_cast<OpResult>(op->getResult(0)));
- return isLegal(layout);
- });
+ target.addDynamicallyLegalOp<
+ vector::ShapeCastOp, vector::StepOp, vector::TransposeOp,
+ vector::BroadcastOp, vector::MultiDimReductionOp, vector::ConstantMaskOp,
+ vector::CreateMaskOp, vector::BitCastOp, vector::InterleaveOp,
+ vector::DeinterleaveOp>([=](Operation *op) -> bool {
+ // Check for either a SliceAttr or LayoutAttr on the result.
+ auto layout =
+ xegpu::getTemporaryLayout(dyn_cast<OpResult>(op->getResult(0)));
+ return isLegal(layout);
+ });
target.addDynamicallyLegalOp<xegpu::LoadGatherOp>(
[=](xegpu::LoadGatherOp op) -> bool {
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 1b7172a27dc8a..d7cbec3b3f85d 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -185,29 +185,28 @@ xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
}
if (auto dpasMxOp = dyn_cast<xegpu::DpasMxOp>(op)) {
// DpasMxOp has operands: a, b, optional acc, optional scale_a, optional
- // scale_b Use AttrSizedOperandSegments to determine which operand this is
- auto segmentSizesAttr = dpasMxOp->getAttrOfType<DenseI32ArrayAttr>(
- dpasMxOp.getOperandSegmentSizesAttrName());
- if (!segmentSizesAttr)
- return nullptr;
-
- auto segmentSizes = segmentSizesAttr.asArrayRef();
- unsigned aSize = segmentSizes[0];
- unsigned bSize = segmentSizes[1];
- unsigned accSize = segmentSizes[2];
- unsigned scaleASize = segmentSizes[3];
-
- if (idx < aSize) {
+ // scale_b
+ unsigned currentIdx = 0;
+
+ if (idx == currentIdx++)
return dpasMxOp.getLayoutAAttr();
- } else if (idx < aSize + bSize) {
+
+ if (idx == currentIdx++)
return dpasMxOp.getLayoutBAttr();
- } else if (idx < aSize + bSize + accSize) {
- return dpasMxOp.getLayoutCdAttr();
- } else if (idx < aSize + bSize + accSize + scaleASize) {
- return dpasMxOp.getLayoutAScaleAttr();
- } else {
- return dpasMxOp.getLayoutBScaleAttr();
- }
+
+ if (dpasMxOp.getAcc())
+ if (idx == currentIdx++)
+ return dpasMxOp.getLayoutCdAttr();
+
+ if (dpasMxOp.getScaleA())
+ if (idx == currentIdx++)
+ return dpasMxOp.getLayoutAScaleAttr();
+
+ if (dpasMxOp.getScaleB())
+ if (idx == currentIdx++)
+ return dpasMxOp.getLayoutBScaleAttr();
+
+ return nullptr;
}
if (auto convertOp = dyn_cast<xegpu::ConvertLayoutOp>(op)) {
return convertOp.getInputLayoutAttr();
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index ff4e0db629083..97499a2b10a0c 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -94,6 +94,40 @@ gpu.module @test_distribution {
gpu.return
}
+ // CHECK-LABEL: dpas_mx
+ gpu.func @dpas_mx(%a: memref<128x128xf8E5M2>, %b: memref<128x128xf8E5M2>, %a_scale: memref<128x4xf8E8M0FNU>, %b_scale: memref<4x128xf8E8M0FNU>) {
+ // CHECK: %[[DPAS_MX:.*]] = xegpu.dpas_mx %{{.*}}, %{{.*}}, %{{.*}} scale_a = %{{.*}} scale_b = %{{.*}} {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<lane_layout = [16, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} : vector<16x128xf8E5M2>, vector<128x16xf8E5M2>, vector<16x16xbf16>, vector<16x4xf8E8M0FNU>, vector<4x16xf8E8M0FNU> -> vector<16x16xbf16>
+ %tdesc_a = xegpu.create_nd_tdesc %a : memref<128x128xf8E5M2>
+ -> !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>>
+ %load_a = xegpu.load_nd %tdesc_a[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>}
+ : !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>>
+ -> vector<128x128xf8E5M2>
+ %tdesc_b = xegpu.create_nd_tdesc %b : memref<128x128xf8E5M2>
+ -> !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
+ %load_b = xegpu.load_nd %tdesc_b[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>}
+ : !xegpu.tensor_desc<128x128xf8E5M2, #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>>
+ -> vector<128x128xf8E5M2>
+ %tdesc_a_scale = xegpu.create_nd_tdesc %a_scale : memref<128x4xf8E8M0FNU>
+ -> !xegpu.tensor_desc<128x4xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>>
+ %load_a_scale = xegpu.load_nd %tdesc_a_scale[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>}
+ : !xegpu.tensor_desc<128x4xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>>
+ -> vector<128x4xf8E8M0FNU>
+ %tdesc_b_scale = xegpu.create_nd_tdesc %b_scale : memref<4x128xf8E8M0FNU>
+ -> !xegpu.tensor_desc<4x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+ %load_b_scale = xegpu.load_nd %tdesc_b_scale[0, 0] {layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ : !xegpu.tensor_desc<4x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>>
+ -> vector<4x128xf8E8M0FNU>
+ %cst = arith.constant dense<0.0> : vector<128x128xbf16>
+ %dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale
+ {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 128], lane_layout = [1, 16], lane_data = [1, 2]>,
+ layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [128, 16], lane_layout = [1, 16], lane_data = [2, 1]>,
+ layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16], lane_layout = [1, 16], lane_data = [1, 1]>,
+ layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 4], lane_layout = [16, 1], lane_data = [1, 1]>,
+ layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [4, 16], lane_layout = [1, 16], lane_data = [1, 1]>}
+ : vector<128x128xf8E5M2>, vector<128x128xf8E5M2>, vector<128x128xbf16>, vector<128x4xf8E8M0FNU>, vector<4x128xf8E8M0FNU> -> vector<128x128xbf16>
+ gpu.return
+ }
+
// CHECK-LABEL: dpas_no_sg_data
gpu.func @dpas_no_sg_data(%a: memref<128x128xf16>, %b: memref<128x128xf16>) {
// CHECK: %[[DPAS:.*]] = xegpu.dpas %{{.*}}, %{{.*}} {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1], order = [1, 0]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1], order = [1, 0]>} : vector<16x128xf16>, vector<128x16xf16> -> vector<16x16xf32>
@@ -1266,6 +1300,61 @@ gpu.module @test_distribution {
gpu.return
}
+ // CHECK-LABEL: @bitcast_distribution
+ gpu.func @bitcast_distribution(%src: memref<256x128xf32>) {
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
+ -> !xegpu.tensor_desc<256x128xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ : !xegpu.tensor_desc<256x128xf32>
+ -> vector<256x128xf32>
+ // CHECK: vector.bitcast {{.*}} : vector<32x32xf32> to vector<32x64xi16>
+ %bitcast = vector.bitcast %load : vector<256x128xf32> to vector<256x256xi16>
+ %add = arith.addi %bitcast, %bitcast : vector<256x256xi16>
+ // CHECK: vector.bitcast {{.*}} : vector<32x64xi16> to vector<32x32xi32>
+ %bitcast2 = vector.bitcast %add : vector<256x256xi16> to vector<256x128xi32>
+ %anchor = xegpu.convert_layout %bitcast2
+ <{
+ input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>,
+ target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>
+ }> : vector<256x128xi32>
+ gpu.return
+ }
+
+ // CHECK-LABEL: @interleave_distribution
+ gpu.func @interleave_distribution(%src: memref<256x128xf32>) {
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x128xf32>
+ -> !xegpu.tensor_desc<256x128xf32>
+ %load1 = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ : !xegpu.tensor_desc<256x128xf32>
+ -> vector<256x128xf32>
+ %load2 = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>}
+ : !xegpu.tensor_desc<256x128xf32>
+ -> vector<256x128xf32>
+ // CHECK: vector.interleave {{.*}}, {{.*}} : vector<32x32xf32> -> vector<32x64xf32>
+ %interleave = vector.interleave %load1, %load2
+ : vector<256x128xf32> -> vector<256x256xf32>
+ %anchor = xegpu.convert_layout %interleave
+ <{
+ input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64]>,
+ target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64]>
+ }> : vector<256x256xf32>
+ gpu.return
+ }
+
+ // CHECK-LABEL: @deinterleave_distribution
+ gpu.func @deinterleave_distribution(%src: memref<256x256xf32>) {
+ %tdesc = xegpu.create_nd_tdesc %src : memref<256x256xf32> -> !xegpu.tensor_desc<256x256xf32>
+ %load = xegpu.load_nd %tdesc[0, 0] {layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 64]>} : !xegpu.tensor_desc<256x256xf32> -> vector<256x256xf32>
+ // CHECK: {{.*}} = vector.deinterleave {{.*}} : vector<32x64xf32> -> vector<32x32xf32>
+ %deinterleave:2 = vector.deinterleave %load : vector<256x256xf32> -> vector<256x128xf32>
+ %anchor = xegpu.convert_layout %deinterleave#0
+ <{
+ input_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>,
+ target_layout = #xegpu.layout<sg_layout = [8, 4], sg_data = [32, 32]>
+ }> : vector<256x128xf32>
+ gpu.return
+ }
+
}
// -----
More information about the Mlir-commits
mailing list