[Mlir-commits] [mlir] 12470b3 - [MLIR][XeGPU] Improve deinterleave/interleave/dpas_mx ops handling (#197223)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue May 19 08:57:29 PDT 2026
Author: Artem Kroviakov
Date: 2026-05-19T17:57:23+02:00
New Revision: 12470b3109bfa6e5fb3afda6bb22001add4d788b
URL: https://github.com/llvm/llvm-project/commit/12470b3109bfa6e5fb3afda6bb22001add4d788b
DIFF: https://github.com/llvm/llvm-project/commit/12470b3109bfa6e5fb3afda6bb22001add4d788b.diff
LOG: [MLIR][XeGPU] Improve deinterleave/interleave/dpas_mx ops handling (#197223)
Added:
Modified:
mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
mlir/test/Conversion/XeGPUToXeVM/dpas_mx.mlir
mlir/test/Dialect/XeGPU/invalid.mlir
mlir/test/Dialect/XeGPU/ops.mlir
mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
mlir/test/Dialect/XeGPU/propagate-layout.mlir
mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
Removed:
################################################################################
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index ae29c4f3726e2..f0c380752306d 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1138,9 +1138,9 @@ def XeGPU_FenceOp: XeGPU_Op<"fence", []> {
let extraClassDeclaration = extraBaseClassDeclaration;
}
-def XeGPU_ConvertLayoutOp: XeGPU_Op<"convert_layout", [AllTypesMatch<["source", "result"]>,
- MemoryEffects<[MemRead, MemWrite]>,
- AnchorLayoutInterface]> {
+def XeGPU_ConvertLayoutOp: XeGPU_Op<"convert_layout", [AllTypesMatch<["source", "result"]>,
+ MemoryEffects<[MemRead, MemWrite]>,
+ AnchorLayoutInterface]> {
let summary = "Convert the layout of the input operand";
let description = [{
`convert_layout` redistribute data across subgroups and/or lanes from the `input_layout` to
@@ -1492,10 +1492,7 @@ def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
$a `,` $b (`,` $acc^)?
(`scale_a` `=` $scale_a^)?
(`scale_b` `=` $scale_b^)?
- attr-dict `:` type($a)`,` type($b)
- (`,` type($acc)^)?
- (`,` type($scale_a)^)?
- (`,` type($scale_b)^)? `->` type($result)
+ attr-dict `:` functional-type(operands, results)
}];
let hasVerifier = 1;
}
diff --git a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
index 7600ec39fb3f5..6ee13747af0d1 100644
--- a/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
+++ b/mlir/lib/Dialect/GPU/Pipelines/GPUToXeVMPipeline.cpp
@@ -96,6 +96,8 @@ void buildGPUPassPipeline(OpPassManager &pm,
pm.addNestedPass<gpu::GPUModuleOp>(createLoopInvariantCodeMotionPass());
pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
pm.addNestedPass<gpu::GPUModuleOp>(xegpu::createXeGPUVectorLinearize());
+ pm.addNestedPass<gpu::GPUModuleOp>(createCanonicalizerPass());
+ pm.addNestedPass<gpu::GPUModuleOp>(createCSEPass());
}
pm.addNestedPass<gpu::GPUModuleOp>(createConvertMathToXeVM());
ConvertXeGPUToXeVMPassOptions xegpuToXeVMOptions;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 813e9608dbf8e..80dec0f33ee0e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1009,43 +1009,47 @@ xegpu::DistributeLayoutAttr xegpu::setupBitCastResultLayout(
SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
+
assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
"laneData must be available for all dimensions");
- size_t dim = srcShape.size() - 1;
+ size_t innerMostDim = srcShape.size() - 1;
int64_t sgDataValue = -1;
int64_t instDataValue = -1;
int64_t laneDataValue = -1;
- const int subgroupSize = uArch->getSubgroupSize();
if (srcElemTyBitWidth > resElemTyBitWidth) {
// When casting to a smaller bitwidth, multiply the result layout
// accordingly to ensure it can be divided by the ratio back to the
// source layout.
int bitWidthRatio = srcElemTyBitWidth / resElemTyBitWidth;
- int innermostDimLaneLayout = subgroupSize;
if (layoutKind == xegpu::LayoutKind::Subgroup) {
- sgDataValue = sgData[dim];
- while ((sgDataValue <= resShape[dim]) &&
+ sgDataValue = sgData[innerMostDim];
+ while ((sgDataValue <= resShape[innerMostDim]) &&
(sgDataValue % bitWidthRatio) != 0)
sgDataValue *= 2;
} else if (layoutKind == xegpu::LayoutKind::InstData) {
- instDataValue = instData[dim];
+ instDataValue = instData[innerMostDim];
+ const int innermostDimLaneLayout = laneLayout.empty()
+ ? uArch->getSubgroupSize()
+ : laneLayout[innerMostDim];
// Adjust instDataValue so it still fits within an instruction after
// dividing by bitWidthRatio
- while ((instDataValue <= resShape[dim]) &&
+ while ((instDataValue <= resShape[innerMostDim]) &&
(instDataValue % (innermostDimLaneLayout * bitWidthRatio) != 0))
instDataValue *= 2;
- assert((resShape[dim] % instDataValue) == 0 &&
+ assert((resShape[innerMostDim] % instDataValue) == 0 &&
"resShape, instData, and lanelayout for innermost must be 2^n !");
} else if (layoutKind == xegpu::LayoutKind::Lane) {
- laneDataValue = laneData[dim];
- while ((laneDataValue <= resShape[dim]) &&
+ laneDataValue = laneData[innerMostDim];
+ while ((laneDataValue <= resShape[innerMostDim]) &&
(laneDataValue % bitWidthRatio != 0))
laneDataValue *= 2;
}
// Now set only instData and laneData, preserving sgData
xegpu::DistributeLayoutAttr resLayout;
- resLayout = consumerLayout.setDimData(dim, sgDataValue, instDataValue,
- laneDataValue);
+ resLayout = consumerLayout.setDimData(innerMostDim, sgDataValue,
+ instDataValue, laneDataValue);
return resLayout;
}
return consumerLayout;
@@ -1074,6 +1078,8 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
SmallVector<int64_t> sgData = consumerLayout.getEffectiveSgDataAsInt();
SmallVector<int64_t> instData = consumerLayout.getEffectiveInstDataAsInt();
SmallVector<int64_t> laneData = consumerLayout.getEffectiveLaneDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ consumerLayout.getEffectiveLaneLayoutAsInt();
assert(consumerLayout.getRank() == static_cast<int64_t>(srcShape.size()) &&
"consumer layout rank must match source shape rank");
@@ -1084,7 +1090,6 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
// Interleave doubles the innermost dimension (ratio = 2)
constexpr int ratio = 2;
- int innermostDimLaneLayout = uArch->getSubgroupSize();
if (layoutKind == xegpu::LayoutKind::Subgroup) {
sgDataValue = sgData[innerMostDim];
@@ -1094,6 +1099,9 @@ xegpu::DistributeLayoutAttr xegpu::setupInterleaveResultLayout(
sgDataValue *= ratio;
} else if (layoutKind == xegpu::LayoutKind::InstData) {
instDataValue = instData[innerMostDim];
+ const int innermostDimLaneLayout = laneLayout.empty()
+ ? uArch->getSubgroupSize()
+ : laneLayout[innerMostDim];
// Adjust instDataValue so it can be divided by (innermostDimLaneLayout *
// ratio) when inferring the source layout
while ((instDataValue <= srcShape[innerMostDim]) &&
@@ -1456,6 +1464,8 @@ getDpasInstDataVectors(VectorType aTy, VectorType bTy, VectorType cdTy,
int kDimSize = subgroupSize;
if (isDpasMx) {
auto supportedKLen = uArchInstruction->getSupportedK(aTy.getElementType());
+ if (supportedKLen.empty())
+ return std::nullopt;
kDimSize = supportedKLen[0];
}
@@ -1910,7 +1920,7 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// For non-anchor ops, derive the operand layout from the op's result
// layout via op-specific semantics.
xegpu::DistributeLayoutAttr resLayout;
- if (op->getNumResults() == 1)
+ if (op->getNumResults() == 1 || isa<vector::DeinterleaveOp>(op))
resLayout = xegpu::getDistributeLayoutAttr(op->getResult(0));
return inferSourceLayoutFromResultForNonAnchorOp(operand, resLayout);
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
index b70c8bd245853..2af5429395526 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
@@ -1535,6 +1535,153 @@ struct SgToWiConvertLayout
}
};
+// Trivially distribute `vector.interleave`
+struct SgToWiVectorInterleave
+ : public OpConversionPattern<vector::InterleaveOp> {
+ using OpConversionPattern<vector::InterleaveOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::InterleaveOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+
+ auto newOp = vector::InterleaveOp::create(
+ rewriter, op.getLoc(), adaptor.getLhs(), adaptor.getRhs());
+ rewriter.replaceOp(op, newOp.getResult());
+ return success();
+ }
+};
+
+// Trivially distribute `vector.deinterleave`
+struct SgToWiVectorDeinterleave
+ : public OpConversionPattern<vector::DeinterleaveOp> {
+ using OpConversionPattern<vector::DeinterleaveOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(vector::DeinterleaveOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+
+ auto newOp = vector::DeinterleaveOp::create(rewriter, op.getLoc(),
+ adaptor.getSource());
+ rewriter.replaceOp(op, newOp.getResults());
+ return success();
+ }
+};
+
+struct SgToWiDpasMx : public OpConversionPattern<xegpu::DpasMxOp> {
+ using OpConversionPattern<xegpu::DpasMxOp>::OpConversionPattern;
+
+ LogicalResult
+ matchAndRewrite(xegpu::DpasMxOp op, OpAdaptor adaptor,
+ ConversionPatternRewriter &rewriter) const override {
+ const uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ if (!uArch)
+ return failure();
+ if (!uArch->isSupportedInstruction(
+ xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc))
+ return rewriter.notifyMatchFailure(
+ op, "target uArch does not support scaled subgroup mma");
+ // Check if the op has A, B and CD layouts attached.
+ auto layoutA = cast<xegpu::LayoutAttr>(op.getLayoutAAttr());
+ auto layoutB = cast<xegpu::LayoutAttr>(op.getLayoutBAttr());
+ auto layoutCd = cast<xegpu::LayoutAttr>(op.getLayoutCdAttr());
+ if (!layoutA || !layoutB || !layoutCd)
+ return rewriter.notifyMatchFailure(
+ op, "missing required layout attributes for DpasMxOp distribution");
+
+ // Retrieve expected types, according to anchor layouts.
+ auto expected1DTypeResult =
+ xegpu::getDistributedVectorType(op.getType(), layoutCd);
+ auto expected1DTypeA =
+ xegpu::getDistributedVectorType(op.getA().getType(), layoutA);
+ auto expected1DTypeB =
+ xegpu::getDistributedVectorType(op.getB().getType(), layoutB);
+
+ VectorType expected1DTypeScaleA, expected1DTypeScaleB;
+ if (op.getScaleA()) {
+ auto layoutScaleA = cast<xegpu::LayoutAttr>(op.getLayoutAScaleAttr());
+ auto expected1DTypeScaleAOrFailure = xegpu::getDistributedVectorType(
+ cast<VectorType>(op.getScaleA().getType()), layoutScaleA);
+ if (failed(expected1DTypeScaleAOrFailure))
+ return rewriter.notifyMatchFailure(
+ op, "failed to calculate expected 1D vector type for scale A");
+ expected1DTypeScaleA = expected1DTypeScaleAOrFailure.value();
+ }
+ if (op.getScaleB()) {
+ auto layoutScaleB = cast<xegpu::LayoutAttr>(op.getLayoutBScaleAttr());
+ auto expected1DTypeScaleBOrFailure = xegpu::getDistributedVectorType(
+ cast<VectorType>(op.getScaleB().getType()), layoutScaleB);
+ if (failed(expected1DTypeScaleBOrFailure))
+ return rewriter.notifyMatchFailure(
+ op, "failed to calculate expected 1D vector type for scale B");
+ expected1DTypeScaleB = expected1DTypeScaleBOrFailure.value();
+ }
+
+ auto expectedNDTypeResult =
+ xegpu::getDistVecTypeBasedOnLaneLayout(layoutCd, op.getType());
+ if (failed(expected1DTypeResult) || failed(expected1DTypeA) ||
+ failed(expected1DTypeB))
+ return rewriter.notifyMatchFailure(
+ op,
+ "failed to calculate supported workitem 1D vector types for DpasOp "
+ "from layouts");
+ if (failed(expectedNDTypeResult))
+ return rewriter.notifyMatchFailure(
+ op, "unable to compute expected workitem vector type for DpasOp from "
+ "lane layout");
+
+ // Validate bit widths match uArch packed format requirements
+ const auto *uArchInstruction = dyn_cast<
+ xegpu::uArch::SubgroupScaledMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupScaledMatrixMultiplyAcc));
+ assert(uArchInstruction);
+ auto wiAType = expected1DTypeA.value();
+ auto wiBType = expected1DTypeB.value();
+ // Calculate total packed bit width = element bit width * vector size
+ unsigned aPackedBitWidth =
+ wiAType.getElementTypeBitWidth() * wiAType.getNumElements();
+ unsigned bPackedBitWidth =
+ wiBType.getElementTypeBitWidth() * wiBType.getNumElements();
+ if (aPackedBitWidth % uArchInstruction->getPackedFormatBitSizeA())
+ return rewriter.notifyMatchFailure(
+ op, "A operand packed bit width must be a multiple of uArch packed "
+ "format requirement");
+ if (bPackedBitWidth % uArchInstruction->getPackedFormatBitSizeB())
+ return rewriter.notifyMatchFailure(
+ op, "B operand packed bit width must be a multiple of uArch packed "
+ "format requirement");
+
+ auto newOp = xegpu::DpasMxOp::create(
+ rewriter, op->getLoc(), expected1DTypeResult.value(),
+ castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getA()),
+ expected1DTypeA.value()),
+ castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getB()),
+ expected1DTypeB.value()),
+ op.getAcc()
+ ? castValueTo(rewriter,
+ cast<TypedValue<VectorType>>(adaptor.getAcc()),
+ expected1DTypeResult.value())
+ : nullptr,
+
+ op.getScaleA()
+ ? castValueTo(rewriter,
+ cast<TypedValue<VectorType>>(adaptor.getScaleA()),
+ expected1DTypeScaleA)
+ : nullptr,
+ op.getScaleB()
+ ? castValueTo(rewriter,
+ cast<TypedValue<VectorType>>(adaptor.getScaleB()),
+ expected1DTypeScaleB)
+ : nullptr,
+ /** layoutA**/ nullptr,
+ /** layoutB**/ nullptr, /** layoutCd**/ nullptr,
+ /** layoutAScale**/ nullptr, /** layoutBScale**/ nullptr);
+ // Explicitly set the new types to enable correct type materializations.
+ rewriter.replaceOp(op, castValueTo(rewriter, newOp.getResult(),
+ expectedNDTypeResult.value()));
+ return success();
+ }
+};
+
struct XeGPUSgToWiDistributeExperimentalPass
: public xegpu::impl::XeGPUSgToWiDistributeExperimentalBase<
XeGPUSgToWiDistributeExperimentalPass> {
@@ -1686,6 +1833,8 @@ void xegpu::populateXeGPUSgToWiDistributeTypeConversionAndLegality(
[&](xegpu::CreateNdDescOp op) { return !op.getType().getLayoutAttr(); });
// Any anchor XeGPU op is legal only if it has no anchor layout.
target.addDynamicallyLegalDialect<xegpu::XeGPUDialect>([](Operation *op) {
+ if (isa<xegpu::ConvertLayoutOp>(op))
+ return false;
auto anchorOp = dyn_cast<AnchorLayoutInterface>(op);
if (!anchorOp)
return true;
@@ -1760,6 +1909,10 @@ void xegpu::populateXeGPUSgToWiDistributeTypeConversionAndLegality(
[=](vector::InsertStridedSliceOp op) -> bool {
return !xegpu::getTemporaryLayout(op->getOpResult(0));
});
+ target.addDynamicallyLegalOp<vector::InterleaveOp, vector::DeinterleaveOp>(
+ [=](Operation *op) -> bool {
+ return !xegpu::getTemporaryLayout(op->getOpResult(0));
+ });
target.markUnknownOpDynamicallyLegal([](Operation *op) { return true; });
patterns.add<SgToWiCreateNdDesc, SgToWiLoadNd, SgToWiStoreNd, SgToWiDpas,
SgToWiElementWise, SgToWiArithConstant, SgToWiPrefetchNd,
@@ -1770,6 +1923,7 @@ void xegpu::populateXeGPUSgToWiDistributeTypeConversionAndLegality(
SgToWiVectorTranspose, SgToWiVectorBitcast, SgToWiVectorStep,
SgToWiVectorShapeCast, SgToWiBroadcast,
SgToWiCreateMask<vector::CreateMaskOp>,
- SgToWiCreateMask<vector::ConstantMaskOp>>(typeConverter,
- patterns.getContext());
+ SgToWiCreateMask<vector::ConstantMaskOp>,
+ SgToWiVectorDeinterleave, SgToWiVectorInterleave, SgToWiDpasMx>(
+ typeConverter, patterns.getContext());
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
index 119ec59daf765..1aa03ebc0f376 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUWgToSgDistribute.cpp
@@ -1344,9 +1344,8 @@ struct WgToSgVectorTransposeOp
xegpu::getTemporaryLayout(dyn_cast<OpResult>(op.getResult()));
if (!layout || !layout.isForWorkgroup())
return failure();
- // TODO-LayoutRefactor: handle the case using getTemporaryLayout
xegpu::DistributeLayoutAttr sourceLayout =
- xegpu::getDistributeLayoutAttr(op.getVector());
+ xegpu::getTemporaryLayout(op->getOpOperand(0));
if (!sourceLayout || !sourceLayout.isForWorkgroup())
return failure();
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 71bcf3d53f9b7..9c2d8e6aa5247 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -975,6 +975,12 @@ bool xegpu::matchSplitDimExpansion(
currentDstDims.push_back(dstIdx);
if (accumulatedSize == src[srcIdx]) {
+ // Also collect trailing unit dims in destination, if any.
+ // Leading unit dims were implicitly collected.
+ if (srcIdx == src.size() - 1) {
+ while (++dstIdx < dst.size() && dst[dstIdx] == 1)
+ currentDstDims.push_back(dstIdx);
+ }
// Record the mapping: srcIdx -> currentDstDims
splitDimGroups.push_back(currentDstDims);
// move to next src dim
diff --git a/mlir/test/Conversion/XeGPUToXeVM/dpas_mx.mlir b/mlir/test/Conversion/XeGPUToXeVM/dpas_mx.mlir
index 1c40c300f5937..e3dd34f5c9e38 100644
--- a/mlir/test/Conversion/XeGPUToXeVM/dpas_mx.mlir
+++ b/mlir/test/Conversion/XeGPUToXeVM/dpas_mx.mlir
@@ -17,7 +17,7 @@ gpu.module @dpas_mx_bf8 [#xevm.target<chip = "cri">] {
// CHECK-SAME: {shape = <m = 8, n = 16, k = 32>, types = <d = f32, a = bf8, b = bf8, c = f32>}
// CHECK-SAME: : (vector<16xi8>, vector<32xi8>, i8, i8, vector<8xf32>) -> vector<8xf32>
%res = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a scale_b = %scale_b :
- vector<16xf8E5M2>, vector<32xf8E5M2>, vector<8xf32>, vector<1xf8E8M0FNU>, vector<1xf8E8M0FNU> -> vector<8xf32>
+ (vector<16xf8E5M2>, vector<32xf8E5M2>, vector<8xf32>, vector<1xf8E8M0FNU>, vector<1xf8E8M0FNU>) -> vector<8xf32>
gpu.return
}
}
@@ -41,7 +41,7 @@ gpu.module @dpas_mx_f8 [#xevm.target<chip = "cri">] {
// CHECK-SAME: {shape = <m = 8, n = 16, k = 32>, types = <d = f32, a = f8, b = f8, c = f32>}
// CHECK-SAME: : (vector<16xi8>, vector<32xi8>, i8, i8, vector<8xf32>) -> vector<8xf32>
%res = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a scale_b = %scale_b :
- vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<8xf32>, vector<1xf8E8M0FNU>, vector<1xf8E8M0FNU> -> vector<8xf32>
+ (vector<16xf8E4M3FN>, vector<32xf8E4M3FN>, vector<8xf32>, vector<1xf8E8M0FNU>, vector<1xf8E8M0FNU>) -> vector<8xf32>
gpu.return
}
}
@@ -65,7 +65,7 @@ gpu.module @dpas_mx_e2m1 [#xevm.target<chip = "cri">] {
// CHECK-SAME: {shape = <m = 8, n = 16, k = 64>, types = <d = f32, a = e2m1, b = e2m1, c = f32>}
// CHECK-SAME: : (vector<16xi8>, vector<32xi8>, vector<2xi8>, vector<2xi8>, vector<8xf32>) -> vector<8xf32>
%res = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a scale_b = %scale_b :
- vector<32xf4E2M1FN>, vector<64xf4E2M1FN>, vector<8xf32>, vector<2xf8E8M0FNU>, vector<2xf8E8M0FNU> -> vector<8xf32>
+ (vector<32xf4E2M1FN>, vector<64xf4E2M1FN>, vector<8xf32>, vector<2xf8E8M0FNU>, vector<2xf8E8M0FNU>) -> vector<8xf32>
gpu.return
}
}
@@ -79,7 +79,7 @@ gpu.module @dpas_mx_no_acc [#xevm.target<chip = "cri">] {
// CHECK: %[[ACC:.*]] = arith.constant dense<0.000000e+00> : vector<8xf32>
// CHECK: %{{.*}} = xevm.mma_mx %{{.*}}, %{{.*}}, %{{.*}}, %{{.*}}, %[[ACC]]
%res = xegpu.dpas_mx %a, %b scale_a = %scale_a scale_b = %scale_b :
- vector<32xf4E2M1FN>, vector<64xf4E2M1FN>, vector<8xf32>, vector<2xf8E8M0FNU>, vector<2xf8E8M0FNU> -> vector<8xf32>
+ (vector<32xf4E2M1FN>, vector<64xf4E2M1FN>, vector<2xf8E8M0FNU>, vector<2xf8E8M0FNU>) -> vector<8xf32>
gpu.return
}
}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index d0b5e2e07f02f..c469a40310607 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -696,49 +696,49 @@ func.func @truncf_invalid_result_size(%a: vector<8x16xf16>) {
// -----
func.func @dpas_mx_acc_result_type_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xbf16>) {
// expected-error at +1 {{Accumulator type must match result type.}}
- %1 = xegpu.dpas_mx %a, %b, %acc : vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xbf16> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xbf16>) -> vector<8x16xf32>
return
}
// -----
func.func @dpas_mx_a_not_2d(%a : vector<128xf8E5M2>, %b: vector<16x16xf8E5M2>) {
// expected-error at +1 {{A operand must be a 2D vector.}}
- %1 = xegpu.dpas_mx %a, %b : vector<128xf8E5M2>, vector<16x16xf8E5M2> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b : (vector<128xf8E5M2>, vector<16x16xf8E5M2>) -> vector<8x16xf32>
return
}
// -----
func.func @dpas_mx_b_not_2d(%a : vector<8x16xf8E5M2>, %b: vector<256xf8E5M2>) {
// expected-error at +1 {{B operand must be a 2D or 3D vector.}}
- %1 = xegpu.dpas_mx %a, %b : vector<8x16xf8E5M2>, vector<256xf8E5M2> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b : (vector<8x16xf8E5M2>, vector<256xf8E5M2>) -> vector<8x16xf32>
return
}
// -----
func.func @dpas_mx_result_not_2d(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>) {
// expected-error at +1 {{Result must be a 2D vector.}}
- %1 = xegpu.dpas_mx %a, %b : vector<8x16xf8E5M2>, vector<16x16xf8E5M2> -> vector<128xf32>
+ %1 = xegpu.dpas_mx %a, %b : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>) -> vector<128xf32>
return
}
// -----
func.func @dpas_mx_k_dimension_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<8x16xf8E5M2>) {
// expected-error at +1 {{K-dimension mismatch: A has K=16 but B has K=8.}}
- %1 = xegpu.dpas_mx %a, %b : vector<8x16xf8E5M2>, vector<8x16xf8E5M2> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b : (vector<8x16xf8E5M2>, vector<8x16xf8E5M2>) -> vector<8x16xf32>
return
}
// -----
func.func @dpas_mx_m_dimension_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>) {
// expected-error at +1 {{M-dimension mismatch: A has M=8 but result has M=16.}}
- %1 = xegpu.dpas_mx %a, %b : vector<8x16xf8E5M2>, vector<16x16xf8E5M2> -> vector<16x16xf32>
+ %1 = xegpu.dpas_mx %a, %b : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>) -> vector<16x16xf32>
return
}
// -----
func.func @dpas_mx_n_dimension_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>) {
// expected-error at +1 {{N-dimension mismatch: B has N=16 but result has N=8.}}
- %1 = xegpu.dpas_mx %a, %b : vector<8x16xf8E5M2>, vector<16x16xf8E5M2> -> vector<8x8xf32>
+ %1 = xegpu.dpas_mx %a, %b : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>) -> vector<8x8xf32>
return
}
@@ -746,21 +746,21 @@ func.func @dpas_mx_n_dimension_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x
// -----
func.func @dpas_mx_scale_a_m_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xf32>, %scale_a: vector<4x2xf8E8M0FNU>) {
// expected-error at +1 {{Scale A M dimension [4] must match A M dimension [8].}}
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a : vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<4x2xf8E8M0FNU> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<4x2xf8E8M0FNU>) -> vector<8x16xf32>
return
}
// -----
func.func @dpas_mx_scale_b_n_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xf32>, %scale_a: vector<8x2xf8E8M0FNU>, %scale_b: vector<2x8xf8E8M0FNU>) {
// expected-error at +1 {{Scale B N dimension [8] must match B N dimension [16].}}
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a scale_b = %scale_b : vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<2x8xf8E8M0FNU> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a scale_b = %scale_b : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<2x8xf8E8M0FNU>) -> vector<8x16xf32>
return
}
// -----
func.func @dpas_mx_scale_k_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xf32>, %scale_a_val: vector<8x2xf8E8M0FNU>, %scale_b_val: vector<4x16xf8E8M0FNU>) {
// expected-error at +1 {{Scale K dimension mismatch: scale_a has K=2 but scale_b has K=4.}}
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val scale_b = %scale_b_val : vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<4x16xf8E8M0FNU> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val scale_b = %scale_b_val : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<4x16xf8E8M0FNU>) -> vector<8x16xf32>
return
}
@@ -770,7 +770,7 @@ func.func @dpas_mx_scale_k_mismatch(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf
#layout_cd = #xegpu.layout<sg_layout = [1, 1], sg_data = [8, 16]>
func.func @dpas_mx_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>) {
// expected-error at +1 {{A shape is not distributable with the layout}}
- %1 = xegpu.dpas_mx %a, %b {layout_a = #layout_a, layout_b = #layout_b, layout_cd = #layout_cd} : vector<8x16xf8E5M2>, vector<16x16xf8E5M2> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b {layout_a = #layout_a, layout_b = #layout_b, layout_cd = #layout_cd} : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>) -> vector<8x16xf32>
return
}
@@ -778,7 +778,7 @@ func.func @dpas_mx_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector
#layout_a_scale_invalid = #xegpu.layout<sg_layout = [1, 1], sg_data = [5, 3]>
func.func @dpas_mx_scale_a_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xf32>, %scale_a_val: vector<8x2xf8E8M0FNU>) {
// expected-error at +1 {{ScaleA shape is not distributable with the layout}}
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val {layout_a_scale = #layout_a_scale_invalid} : vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val {layout_a_scale = #layout_a_scale_invalid} : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>) -> vector<8x16xf32>
return
}
@@ -786,6 +786,6 @@ func.func @dpas_mx_scale_a_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b
#layout_b_scale_invalid = #xegpu.layout<sg_layout = [1, 1], sg_data = [3, 11]>
func.func @dpas_mx_scale_b_layout_not_distributable(%a : vector<8x16xf8E5M2>, %b: vector<16x16xf8E5M2>, %acc: vector<8x16xf32>, %scale_a_val: vector<8x2xf8E8M0FNU>, %scale_b_val: vector<2x16xf8E8M0FNU>) {
// expected-error at +1 {{ScaleB shape is not distributable with the layout}}
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val scale_b = %scale_b_val {layout_b_scale = #layout_b_scale_invalid} : vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU> -> vector<8x16xf32>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %scale_a_val scale_b = %scale_b_val {layout_b_scale = #layout_b_scale_invalid} : (vector<8x16xf8E5M2>, vector<16x16xf8E5M2>, vector<8x16xf32>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
return
}
diff --git a/mlir/test/Dialect/XeGPU/ops.mlir b/mlir/test/Dialect/XeGPU/ops.mlir
index edba8103a4075..f28fb39fd278b 100644
--- a/mlir/test/Dialect/XeGPU/ops.mlir
+++ b/mlir/test/Dialect/XeGPU/ops.mlir
@@ -609,10 +609,9 @@ gpu.func @truncf(%a: vector<8x16xf16>) {
// CHECK-LABEL: gpu.func @dpas_mx
gpu.func @dpas_mx(%a : vector<8x32xf8E5M2>, %b: vector<32x16xf8E5M2>, %acc: vector<8x16xbf16>, %a_scale: vector<8x1xf8E8M0FNU>, %b_scale: vector<1x16xf8E8M0FNU>) {
- // CHECK: %{{.+}} = xegpu.dpas_mx %{{.+}}, %{{.+}}, %{{.+}} scale_a = %{{.+}} scale_b = %{{.+}} : vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU> -> vector<8x16xbf16>
- %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %a_scale scale_b = %b_scale : vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU> -> vector<8x16xbf16>
+ // CHECK: %{{.+}} = xegpu.dpas_mx %{{.+}}, %{{.+}}, %{{.+}} scale_a = %{{.+}} scale_b = %{{.+}} : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
+ %1 = xegpu.dpas_mx %a, %b, %acc scale_a = %a_scale scale_b = %b_scale : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
gpu.return
}
}
-
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
index ab398ba524bf7..6f587959b697d 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-inst-data.mlir
@@ -466,7 +466,7 @@ func.func @bitcast_ui16_to_f4(%arg0: memref<256x16xui16>) {
// CHECK-SAME: !xegpu.tensor_desc<2x32xf8E8M0FNU, #xegpu.layout<inst_data = [1, 16]>> -> vector<2x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 32]>, layout_a_scale = #xegpu.layout<inst_data = [8, 1]>, layout_b = #xegpu.layout<inst_data = [32, 16]>, layout_b_scale = #xegpu.layout<inst_data = [1, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
-// CHECK-SAME: vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU> -> vector<16x32xbf16>
+// CHECK-SAME: (vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
gpu.module @test {
@@ -482,7 +482,7 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x64xf8E5M2>, %arg1: memref<64x32xf8E5M
%5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x2xf8E8M0FNU> -> vector<16x2xf8E8M0FNU>
%6 = xegpu.create_nd_tdesc %arg4 : memref<2x32xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU>
%7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<2x32xf8E8M0FNU> -> vector<2x32xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU> -> vector<16x32xbf16>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x64xf8E5M2>, vector<64x32xf8E5M2>, vector<16x32xbf16>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16>
return
@@ -508,7 +508,7 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<16x64xf8E5M2>, %arg1: memref<64x32xf8E5M
// CHECK-SAME: !xegpu.tensor_desc<4x32xf8E8M0FNU, #xegpu.layout<inst_data = [2, 16]>> -> vector<4x32xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: {layout_a = #xegpu.layout<inst_data = [8, 64]>, layout_a_scale = #xegpu.layout<inst_data = [8, 2]>, layout_b = #xegpu.layout<inst_data = [64, 16]>, layout_b_scale = #xegpu.layout<inst_data = [2, 16]>, layout_cd = #xegpu.layout<inst_data = [8, 16]>} :
-// CHECK-SAME: vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU> -> vector<16x32xbf16>
+// CHECK-SAME: (vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU>) -> vector<16x32xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<inst_data = [8, 16]>}> : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16, #xegpu.layout<inst_data = [8, 16]>>
gpu.module @test {
@@ -524,7 +524,7 @@ func.func @dpas_mx_f4e2m1(%arg0: memref<16x128xf4E2M1FN>, %arg1: memref<128x32xf
%5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<16x4xf8E8M0FNU> -> vector<16x4xf8E8M0FNU>
%6 = xegpu.create_nd_tdesc %arg4 : memref<4x32xf8E8M0FNU> -> !xegpu.tensor_desc<4x32xf8E8M0FNU>
%7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<4x32xf8E8M0FNU> -> vector<4x32xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU> -> vector<16x32xbf16>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<16x128xf4E2M1FN>, vector<128x32xf4E2M1FN>, vector<16x32xbf16>, vector<16x4xf8E8M0FNU>, vector<4x32xf8E8M0FNU>) -> vector<16x32xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<16x32xbf16> -> !xegpu.tensor_desc<16x32xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<16x32xbf16>, !xegpu.tensor_desc<16x32xbf16>
return
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index b4a34fb01456f..0917d325d4271 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -399,7 +399,7 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<16x256xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>> -> vector<16x256xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T1]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 32]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>} :
- // CHECK-SAME: vector<128x512xf8E5M2>, vector<512x256xf8E5M2>, vector<128x256xbf16>, vector<128x16xf8E8M0FNU>, vector<16x256xf8E8M0FNU> -> vector<128x256xbf16>
+ // CHECK-SAME: (vector<128x512xf8E5M2>, vector<512x256xf8E5M2>, vector<128x256xbf16>, vector<128x16xf8E8M0FNU>, vector<16x256xf8E8M0FNU>) -> vector<128x256xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x256xbf16> -> !xegpu.tensor_desc<128x256xbf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> : vector<128x256xbf16>, !xegpu.tensor_desc<128x256xbf16, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>>
gpu.func @dpas_mx(%arg0: memref<128x512xf8E5M2>, %arg1: memref<512x256xf8E5M2>, %arg2: memref<128x256xbf16>,
@@ -414,7 +414,7 @@ gpu.module @test {
%load_a_scale = xegpu.load_nd %tdesc_a_scale[0, 0] : !xegpu.tensor_desc<128x16xf8E8M0FNU> -> vector<128x16xf8E8M0FNU>
%tdesc_b_scale = xegpu.create_nd_tdesc %arg4 : memref<16x256xf8E8M0FNU> -> !xegpu.tensor_desc<16x256xf8E8M0FNU>
%load_b_scale = xegpu.load_nd %tdesc_b_scale[0, 0] : !xegpu.tensor_desc<16x256xf8E8M0FNU> -> vector<16x256xf8E8M0FNU>
- %dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale : vector<128x512xf8E5M2>, vector<512x256xf8E5M2>, vector<128x256xbf16>, vector<128x16xf8E8M0FNU>, vector<16x256xf8E8M0FNU> -> vector<128x256xbf16>
+ %dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale : (vector<128x512xf8E5M2>, vector<512x256xf8E5M2>, vector<128x256xbf16>, vector<128x16xf8E8M0FNU>, vector<16x256xf8E8M0FNU>) -> vector<128x256xbf16>
%tdesc_cd = xegpu.create_nd_tdesc %arg2 : memref<128x256xbf16> -> !xegpu.tensor_desc<128x256xbf16>
xegpu.store_nd %dpas_mx, %tdesc_cd[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 32]>}> : vector<128x256xbf16>, !xegpu.tensor_desc<128x256xbf16>
gpu.return
@@ -440,7 +440,7 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<16x128xf8E8M0FNU, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>> -> vector<16x128xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T1]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: {layout_a = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 512]>, layout_a_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16]>, layout_b_scale = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>, layout_cd = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>} :
- // CHECK-SAME: vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>, vector<128x128xf32>, vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU> -> vector<128x128xf32>
+ // CHECK-SAME: (vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>, vector<128x128xf32>, vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>) -> vector<128x128xf32>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %{{.*}} : memref<128x128xf32> -> !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32, #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>>
gpu.func @dpas_mx_fp4(%arg0: memref<128x512xf4E2M1FN>, %arg1: memref<512x128xf4E2M1FN>, %arg2: memref<128x128xf32>,
@@ -455,7 +455,7 @@ gpu.module @test {
%load_a_scale = xegpu.load_nd %tdesc_a_scale[0, 0] : !xegpu.tensor_desc<128x16xf8E8M0FNU> -> vector<128x16xf8E8M0FNU>
%tdesc_b_scale = xegpu.create_nd_tdesc %arg4 : memref<16x128xf8E8M0FNU> -> !xegpu.tensor_desc<16x128xf8E8M0FNU>
%load_b_scale = xegpu.load_nd %tdesc_b_scale[0, 0] : !xegpu.tensor_desc<16x128xf8E8M0FNU> -> vector<16x128xf8E8M0FNU>
- %dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale : vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>, vector<128x128xf32>, vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU> -> vector<128x128xf32>
+ %dpas_mx = xegpu.dpas_mx %load_a, %load_b, %cst scale_a = %load_a_scale scale_b = %load_b_scale : (vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>, vector<128x128xf32>, vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>) -> vector<128x128xf32>
%tdesc_cd = xegpu.create_nd_tdesc %arg2 : memref<128x128xf32> -> !xegpu.tensor_desc<128x128xf32>
xegpu.store_nd %dpas_mx, %tdesc_cd[0, 0] <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [16, 16]>}> : vector<128x128xf32>, !xegpu.tensor_desc<128x128xf32>
gpu.return
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout.mlir b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
index 3ceddde54dede..f8cbfed14d743 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout.mlir
@@ -315,7 +315,7 @@ func.func @vector_bitcast_i16_to_i32(%arg0: memref<8x32xi16>, %arg1: memref<8x16
// -----
gpu.module @test {
// CHECK-LABEL: func.func @vector_bitcast_require_cross_lane_shuffle(
-// CHECK: %[[LOAD:.*]] = xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
+// CHECK: %[[LOAD:.*]] = xegpu.load_nd %{{.*}} <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}>
// CHECK-SAME: !xegpu.tensor_desc<8x16xi32, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: %{{.*}} = vector.bitcast %[[LOAD]] {layout_result_0 = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>}
// CHECK-SAME: vector<8x16xi32> to vector<8x32xi16>
@@ -1001,7 +1001,7 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<1x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>> -> vector<1x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 2]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [4, 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]>} :
-// CHECK-SAME: vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU> -> vector<8x16xbf16>
+// CHECK-SAME: (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
func.func @dpas_mx_f8e5m2(%arg0: memref<8x32xf8E5M2>, %arg1: memref<32x16xf8E5M2>, %arg2: memref<8x16xbf16>,
@@ -1016,7 +1016,7 @@ func.func @dpas_mx_f8e5m2(%arg0: memref<8x32xf8E5M2>, %arg1: memref<32x16xf8E5M2
%5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<8x1xf8E8M0FNU> -> vector<8x1xf8E8M0FNU>
%6 = xegpu.create_nd_tdesc %arg4 : memref<1x16xf8E8M0FNU> -> !xegpu.tensor_desc<1x16xf8E8M0FNU>
%7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<1x16xf8E8M0FNU> -> vector<1x16xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU> -> vector<8x16xbf16>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<8x32xf8E5M2>, vector<32x16xf8E5M2>, vector<8x16xbf16>, vector<8x1xf8E8M0FNU>, vector<1x16xf8E8M0FNU>) -> vector<8x16xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16>
return
@@ -1043,7 +1043,7 @@ gpu.module @test {
// CHECK-SAME: !xegpu.tensor_desc<2x16xf8E8M0FNU, #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>> -> vector<2x16xf8E8M0FNU>
// CHECK: %[[T8:.*]] = xegpu.dpas_mx %[[T2]], %[[T3]], %[[CST]] scale_a = %[[T5]] scale_b = %[[T7]]
// CHECK-SAME: {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 4]>, layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 2]>, layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [8, 1]>, layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [2, 1]>, layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>} :
-// CHECK-SAME: vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU> -> vector<8x16xbf16>
+// CHECK-SAME: (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
// CHECK: %[[T9:.*]] = xegpu.create_nd_tdesc %[[ARG2]] : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
// CHECK: xegpu.store_nd %[[T8]], %[[T9]][0, 0] <{layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}> : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16, #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>>
func.func @dpas_mx_fp4(%arg0: memref<8x64xf4E2M1FN>, %arg1: memref<64x16xf4E2M1FN>, %arg2: memref<8x16xbf16>,
@@ -1058,7 +1058,7 @@ func.func @dpas_mx_fp4(%arg0: memref<8x64xf4E2M1FN>, %arg1: memref<64x16xf4E2M1F
%5 = xegpu.load_nd %4[0, 0] : !xegpu.tensor_desc<8x2xf8E8M0FNU> -> vector<8x2xf8E8M0FNU>
%6 = xegpu.create_nd_tdesc %arg4 : memref<2x16xf8E8M0FNU> -> !xegpu.tensor_desc<2x16xf8E8M0FNU>
%7 = xegpu.load_nd %6[0, 0] : !xegpu.tensor_desc<2x16xf8E8M0FNU> -> vector<2x16xf8E8M0FNU>
- %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU> -> vector<8x16xbf16>
+ %8 = xegpu.dpas_mx %2, %3, %cst scale_a = %5 scale_b = %7 : (vector<8x64xf4E2M1FN>, vector<64x16xf4E2M1FN>, vector<8x16xbf16>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xbf16>
%9 = xegpu.create_nd_tdesc %arg2 : memref<8x16xbf16> -> !xegpu.tensor_desc<8x16xbf16>
xegpu.store_nd %8, %9[0, 0] : vector<8x16xbf16>, !xegpu.tensor_desc<8x16xbf16>
return
diff --git a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
index 0e890ab625b0a..f1e56f4493ec7 100644
--- a/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
+++ b/mlir/test/Dialect/XeGPU/sg-to-wi-experimental-unit.mlir
@@ -1,7 +1,8 @@
-// RUN: mlir-opt --xevm-attach-target='module=xevm_* chip=pvc' --allow-unregistered-dialect \
+// RUN: mlir-opt --xevm-attach-target='module=xevm_* chip=cri' --allow-unregistered-dialect \
// RUN: --test-xegpu-sg-to-wi-distribute-experimental --split-input-file %s | FileCheck %s
+
gpu.module @xevm_module {
// CHECK-LABEL: gpu.func @create_nd_tdesc
// CHECK: %[[C0:.*]] = arith.constant 0 : index
@@ -895,12 +896,12 @@ gpu.func @convert_layout_removed_when_compatible() {
target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>}>
: vector<16xf32>
%3 = xegpu.convert_layout %2
- <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,
- target_layout = #xegpu.layout<lane_layout = [1], lane_data = [1]>}>
+ <{input_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>, dims = [0]>,
+ target_layout = #xegpu.layout<lane_layout = [1], lane_data = [1]>}>
: vector<1xf32>
%4 = xegpu.convert_layout %3
<{input_layout = #xegpu.layout<lane_layout = [1], lane_data = [1]>,
- target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0, 1]>}>
+ target_layout = #xegpu.slice<#xegpu.layout<lane_layout = [1, 1, 16], lane_data = [1, 1, 1]>, dims = [0, 1]>}>
: vector<1xf32>
gpu.return
}
@@ -1244,3 +1245,78 @@ gpu.func @vector_multi_reduction_1d_to_scalar() {
gpu.return
}
}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @vector_interleave
+// CHECK: %[[LHS:.*]] = arith.constant dense<2.000000e+00> : vector<1x1xf32>
+// CHECK: %[[RHS:.*]] = arith.constant dense<3.000000e+00> : vector<1x1xf32>
+// CHECK: %[[INTERLEAVED:.*]] = vector.interleave %[[LHS]], %[[RHS]] : vector<1x1xf32> -> vector<1x2xf32>
+gpu.func @vector_interleave() {
+ %0 = arith.constant dense<2.> : vector<8x2xf32>
+ %1 = arith.constant dense<3.> : vector<8x2xf32>
+ %2 = vector.interleave %0, %1 : vector<8x2xf32> -> vector<8x4xf32>
+
+ %cl1 = xegpu.convert_layout %2
+ <{
+ input_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 2]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 2]>
+ }> : vector<8x4xf32>
+
+ gpu.return
+}
+}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @vector_deinterleave
+// CHECK: %[[SRC:.*]] = arith.constant dense<3.000000e+00> : vector<1x2xf32>
+// CHECK: %[[DEINTERLEAVED:.*]] = vector.deinterleave %[[SRC]] : vector<1x2xf32> -> vector<1x1xf32>
+gpu.func @vector_deinterleave() {
+ %src = arith.constant dense<3.> : vector<8x4xf32>
+ %even, %odd = vector.deinterleave %src : vector<8x4xf32> -> vector<8x2xf32>
+
+ %cl1 = xegpu.convert_layout %even
+ <{
+ input_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [8, 2], lane_data = [1, 1]>
+ }> : vector<8x2xf32>
+
+ gpu.return
+}
+}
+
+// -----
+gpu.module @xevm_module {
+// CHECK-LABEL: gpu.func @xegpu_dpas_mx
+// CHECK-DAG: %[[CST:.*]] = arith.constant dense<3.000000e+00> : vector<8x1xf8E4M3FN>
+// CHECK-DAG: %[[CST_0:.*]] = arith.constant dense<3.000000e+00> : vector<16x1xf8E4M3FN>
+// CHECK-DAG: %[[CST_1:.*]] = arith.constant dense<1.000000e+00> : vector<1x2xf8E8M0FNU>
+// CHECK-DAG: %[[CST_2:.*]] = arith.constant dense<5.000000e-01> : vector<2x1xf8E8M0FNU>
+// CHECK-DAG: %[[SCALE_B:.*]] = vector.shape_cast %[[CST_2]] : vector<2x1xf8E8M0FNU> to vector<2xf8E8M0FNU>
+// CHECK-DAG: %[[SCALE_A:.*]] = vector.shape_cast %[[CST_1]] : vector<1x2xf8E8M0FNU> to vector<2xf8E8M0FNU>
+// CHECK-DAG: %[[B:.*]] = vector.shape_cast %[[CST_0]] : vector<16x1xf8E4M3FN> to vector<16xf8E4M3FN>
+// CHECK-DAG: %[[A:.*]] = vector.shape_cast %[[CST]] : vector<8x1xf8E4M3FN> to vector<8xf8E4M3FN>
+// CHECK: %[[RESULT:.*]] = xegpu.dpas_mx %[[A]], %[[B]] scale_a = %[[SCALE_A]] scale_b = %[[SCALE_B]] : (vector<8xf8E4M3FN>, vector<16xf8E4M3FN>, vector<2xf8E8M0FNU>, vector<2xf8E8M0FNU>) -> vector<8xf32>
+gpu.func @xegpu_dpas_mx(%arg0: !xegpu.mem_desc<8x8xf16>, %arg1: !xegpu.mem_desc<8x8xf16>) {
+ %A = arith.constant dense<3.> : vector<8x16xf8E4M3FN>
+ %B = arith.constant dense<3.> : vector<16x16xf8E4M3FN>
+ %scale_A = arith.constant dense<1.> : vector<8x2xf8E8M0FNU>
+ %scale_B = arith.constant dense<0.5> : vector<2x16xf8E8M0FNU>
+
+ %4 = xegpu.dpas_mx %A, %B scale_a = %scale_A scale_b = %scale_B
+ {layout_a = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ layout_b = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ layout_cd = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ layout_a_scale = #xegpu.layout<lane_layout = [8, 1], lane_data = [1, 1]>,
+ layout_b_scale = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>}
+ : (vector<8x16xf8E4M3FN>, vector<16x16xf8E4M3FN>, vector<8x2xf8E8M0FNU>, vector<2x16xf8E8M0FNU>) -> vector<8x16xf32>
+
+ %anchor = xegpu.convert_layout %4
+ <{
+ input_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>,
+ target_layout = #xegpu.layout<lane_layout = [1, 16], lane_data = [1, 1]>
+ }> : vector<8x16xf32>
+ gpu.return
+}
+}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index e5dd5000fc3a1..56a00508eda15 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -736,7 +736,7 @@ gpu.module @test_kernel {
//CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<1x16xf8E8M0FNU>
%sb = xegpu.load_nd %scale_b_tdesc[%c0, %c0] {layout = #l2_scale}: !xegpu.tensor_desc<2x32xf8E8M0FNU, #l2_scale> -> vector<2x32xf8E8M0FNU>
//CHECK-COUNT-8: xegpu.dpas_mx {{.*}}
- %c = xegpu.dpas_mx %a, %b, %arg2 scale_a = %sa scale_b = %sb {layout_a=#l1, layout_b = #l2, layout_cd = #l3, layout_a_scale = #l1_scale, layout_b_scale = #l2_scale, layout_result_0 = #l3}: vector<16x64xf4E2M1FN>, vector<64x32xf4E2M1FN>, vector<16x32xf32>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU> -> vector<16x32xf32>
+ %c = xegpu.dpas_mx %a, %b, %arg2 scale_a = %sa scale_b = %sb {layout_a=#l1, layout_b = #l2, layout_cd = #l3, layout_a_scale = #l1_scale, layout_b_scale = #l2_scale, layout_result_0 = #l3}: (vector<16x64xf4E2M1FN>, vector<64x32xf4E2M1FN>, vector<16x32xf32>, vector<16x2xf8E8M0FNU>, vector<2x32xf8E8M0FNU>) -> vector<16x32xf32>
scf.yield %c : vector<16x32xf32>
} {layout_result_0 = #l3}
//CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
@@ -744,4 +744,3 @@ gpu.module @test_kernel {
gpu.return
}
}
-
diff --git a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
index 97499a2b10a0c..3efe3871bb236 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-wg-to-sg.mlir
@@ -96,7 +96,7 @@ gpu.module @test_distribution {
// 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>
+ // 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]>}
@@ -124,7 +124,7 @@ gpu.module @test_distribution {
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>
+ : (vector<128x128xf8E5M2>, vector<128x128xf8E5M2>, vector<128x128xbf16>, vector<128x4xf8E8M0FNU>, vector<4x128xf8E8M0FNU>) -> vector<128x128xbf16>
gpu.return
}
@@ -419,7 +419,7 @@ gpu.module @test_distribution {
%reduce = vector.multi_reduction <add>, %load, %cst [0]
: vector<4x128xf32> to vector<128xf32>
%anchor = xegpu.convert_layout %reduce
- <{input_layout = #xegpu.slice<#xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>, dims = [0]>,
+ <{input_layout = #xegpu.slice<#xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>, dims = [0]>,
target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [1, 32], sg_data = [4, 4]>, dims = [0]>}>
: vector<128xf32>
gpu.return
@@ -437,7 +437,7 @@ gpu.module @test_distribution {
%reduce = vector.multi_reduction <add>, %load, %cst [1]
: vector<256x64xf32> to vector<256xf32>
%anchor = xegpu.convert_layout %reduce
- <{input_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>,
+ <{input_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>,
target_layout = #xegpu.slice<#xegpu.layout<sg_layout = [16, 1], sg_data = [16, 64]>, dims = [1]>}>
: vector<256xf32>
gpu.return
@@ -1272,7 +1272,7 @@ gpu.module @test_distribution {
<{
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>
+ }> : vector<256x128xf32>
gpu.return
}
diff --git a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
index e59599ee4855f..e2ca280096edf 100644
--- a/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
+++ b/mlir/test/Integration/Dialect/XeGPU/WG/simple_mxfp_gemm.mlir
@@ -61,9 +61,9 @@ gpu.module @test {
layout_cd = #c,
layout_a_scale = #a_scale,
layout_b_scale = #b_scale}
- : vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>,
+ : (vector<128x512xf4E2M1FN>, vector<512x128xf4E2M1FN>,
vector<128x128xf32>,
- vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>
+ vector<128x16xf8E8M0FNU>, vector<16x128xf8E8M0FNU>)
-> vector<128x128xf32>
// store_nd with offset
More information about the Mlir-commits
mailing list