[Mlir-commits] [mlir] 2384593 - [MLIR][XeGPU] Unroll Dpasmx Op (#195179)
llvmlistbot at llvm.org
llvmlistbot at llvm.org
Tue May 5 17:59:25 PDT 2026
Author: Jianhui Li
Date: 2026-05-05T17:59:21-07:00
New Revision: 2384593d63a2b2d436e80a8775a26b980088e7ca
URL: https://github.com/llvm/llvm-project/commit/2384593d63a2b2d436e80a8775a26b980088e7ca
DIFF: https://github.com/llvm/llvm-project/commit/2384593d63a2b2d436e80a8775a26b980088e7ca.diff
LOG: [MLIR][XeGPU] Unroll Dpasmx Op (#195179)
This PR adds support to unroll Dpasmx.
Assisted by Claude
---------
Co-authored-by: Claude Sonnet 4.5 <noreply at anthropic.com>
Added:
Modified:
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
mlir/test/Dialect/XeGPU/invalid.mlir
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
Removed:
################################################################################
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index b87f41d31e02c..bab2dbc9505dc 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -703,62 +703,95 @@ void StoreScatterOp::build(
}
//===----------------------------------------------------------------------===//
-// XeGPU_DpasOp
+// DPAS Common Verification Helpers
//===----------------------------------------------------------------------===//
-LogicalResult DpasOp::verify() {
- int64_t lhsRank = getLhsType().getRank();
- int64_t rhsRank = getRhsType().getRank();
- int64_t resRank = getResultType().getRank();
- auto lhsShape = getLhsType().getShape();
- auto rhsShape = getRhsType().getShape();
- auto resShape = getResultType().getShape();
- if (auto cdLayout = getLayoutCd())
- if (!cdLayout->isDistributable(
- SmallVector<int64_t>(resShape.begin(), resShape.end())))
- return emitOpError("Value shape is not distributable with the layout");
+// Helper to verify layout distributability for a value
+static LogicalResult
+verifyLayoutDistributable(Operation *op,
+ std::optional<DistributeLayoutAttr> layout,
+ ArrayRef<int64_t> shape, StringRef operandName) {
+ if (layout && !layout->isDistributable(
+ SmallVector<int64_t>(shape.begin(), shape.end())))
+ return op->emitOpError(operandName)
+ << " shape is not distributable with the layout";
+ return success();
+}
- if (auto aLayout = getLayoutA())
- if (!aLayout->isDistributable(
- SmallVector<int64_t>(lhsShape.begin(), lhsShape.end())))
- return emitOpError("Value shape is not distributable with the layout");
+// Helper to verify M, N, K dimensions match between A, B, and result matrices
+static LogicalResult verifyDpasDimensions(Operation *op,
+ ArrayRef<int64_t> aShape,
+ ArrayRef<int64_t> bShape,
+ ArrayRef<int64_t> resShape) {
- if (auto bLayout = getLayoutB())
- if (!bLayout->isDistributable(
- SmallVector<int64_t>(rhsShape.begin(), rhsShape.end())))
- return emitOpError("Value shape is not distributable with the layout");
-
- if (getAcc() && getAcc().getType() != getResultType())
- return emitOpError("Expecting the acc type to be the same as result.");
-
- // SIMT code: the size of the B operand has to be a multiple of 32 bits.
- // It skips the semantic check since lack of architecture information.
- // Users need to ensure the correctness.
- if (lhsRank == 1 && rhsRank == 1 && resRank == 1) {
- auto numElems = getRhsType().getNumElements();
- auto elemTy = getRhsType().getElementType();
- auto factor = 32 / elemTy.getIntOrFloatBitWidth();
- if (numElems % factor != 0)
- return emitOpError("Expecting B operand to be a multiple of 32 bits.");
+ auto aRank = aShape.size();
+ auto bRank = bShape.size();
+ auto resRank = resShape.size();
+ if (aRank == 1 && bRank == 1 && resRank == 1)
return success();
- }
- // SIMD code
- if (lhsRank != 2 || (rhsRank != 2 && rhsRank != 3) || resRank != 2)
- return emitOpError(
- "expecting lhs and result to be a 2D vector, and rhs to be either "
- "2D or 3D (packed) vector.");
- auto bK = rhsRank == 3 ? rhsShape[0] * rhsShape[2] : rhsShape[0];
- if (bK != lhsShape[1])
- return emitOpError("K-dimension mismatch.");
- if (lhsShape[0] != resShape[0])
- return emitOpError("M-dimension mismatch.");
- if (rhsShape[1] != resShape[1])
- return emitOpError("N-dimension mismatch.");
+ // Validate A and B are 2D
+ if (aRank != 2)
+ return op->emitOpError("A operand must be a 2D vector.");
+ if (bRank < 2 || bRank > 3)
+ return op->emitOpError("B operand must be a 2D or 3D vector.");
+ if (resRank != 2)
+ return op->emitOpError("Result must be a 2D vector.");
+
+ // Calculate effective K dimension for B (handle 3D packed case)
+ int64_t bK = bRank == 3 ? bShape[0] * bShape[2] : bShape[0];
+
+ // Verify K dimension match between A and B
+ if (bK != aShape[1])
+ return op->emitOpError("K-dimension mismatch: A has K=")
+ << aShape[1] << " but B has K=" << bK << ".";
+
+ // Verify M dimension match between A and result
+ if (aShape[0] != resShape[0])
+ return op->emitOpError("M-dimension mismatch: A has M=")
+ << aShape[0] << " but result has M=" << resShape[0] << ".";
+
+ // Verify N dimension match between B and result
+ if (bShape[1] != resShape[1])
+ return op->emitOpError("N-dimension mismatch: B has N=")
+ << bShape[1] << " but result has N=" << resShape[1] << ".";
+
+ return success();
+}
+// Helper to verify accumulator matches result type
+static LogicalResult verifyDpasAccumulator(Operation *op, Type accType,
+ Type resultType) {
+ if (accType != resultType)
+ return op->emitOpError("Accumulator type must match result type.");
return success();
}
+//===----------------------------------------------------------------------===//
+// XeGPU_DpasOp
+//===----------------------------------------------------------------------===//
+LogicalResult DpasOp::verify() {
+ auto lhsShape = getLhsType().getShape();
+ auto rhsShape = getRhsType().getShape();
+ auto resShape = getResultType().getShape();
+
+ // Verify layout distributability
+ if (failed(
+ verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
+ return failure();
+ if (failed(verifyLayoutDistributable(*this, getLayoutA(), lhsShape, "A")))
+ return failure();
+ if (failed(verifyLayoutDistributable(*this, getLayoutB(), rhsShape, "B")))
+ return failure();
+
+ // Verify accumulator if present
+ if (getAcc() &&
+ failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
+ return failure();
+
+ return verifyDpasDimensions(*this, lhsShape, rhsShape, resShape);
+}
+
//===----------------------------------------------------------------------===//
// XeGPU_ConvertLayoutOp
//===----------------------------------------------------------------------===//
@@ -862,8 +895,92 @@ LogicalResult TruncfOp::verify() {
//===----------------------------------------------------------------------===//
LogicalResult DpasMxOp::verify() {
- if (getAcc() && getAcc().getType() != getResultType())
- return emitOpError("Expecting the acc type to be the same as result.");
+ auto aShape = getAType().getShape();
+ auto bShape = getBType().getShape();
+ auto resShape = getResultType().getShape();
+
+ // Verify layout distributability for A, B, and result
+ if (failed(
+ verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
+ return failure();
+ if (failed(verifyLayoutDistributable(*this, getLayoutA(), aShape, "A")))
+ return failure();
+ if (failed(verifyLayoutDistributable(*this, getLayoutB(), bShape, "B")))
+ return failure();
+
+ // Verify accumulator if present
+ if (getAcc() &&
+ failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
+ return failure();
+
+ // Verify M, N, K dimensions
+ if (failed(verifyDpasDimensions(*this, aShape, bShape, resShape)))
+ return failure();
+
+ // Validate scale_a if present
+ if (getScaleA()) {
+ auto scaleAVecType = dyn_cast<VectorType>(getScaleAType());
+ // Only validate if scale is a vector (scalars are always valid)
+ if (scaleAVecType) {
+ auto scaleAShape = scaleAVecType.getShape();
+
+ if (scaleAVecType.getRank() != 2)
+ return emitOpError("Scale A must be a 2D vector when not a scalar.");
+
+ // Verify layout distributability for scale_a
+ if (failed(verifyLayoutDistributable(*this, getLayoutAScale(),
+ scaleAShape, "ScaleA")))
+ return failure();
+
+ // Validate M dimension: scale_a[0] must match a[0]
+ if (scaleAShape[0] != aShape[0])
+ return emitOpError("Scale A M dimension [")
+ << scaleAShape[0] << "] must match A M dimension [" << aShape[0]
+ << "].";
+ }
+ }
+
+ // Validate scale_b if present
+ if (getScaleB()) {
+ auto scaleBVecType = dyn_cast<VectorType>(getScaleBType());
+ // Only validate if scale is a vector (scalars are always valid)
+ if (scaleBVecType) {
+ auto scaleBShape = scaleBVecType.getShape();
+
+ if (scaleBVecType.getRank() != 2)
+ return emitOpError("Scale B must be a 2D vector when not a scalar.");
+
+ // Verify layout distributability for scale_b
+ if (failed(verifyLayoutDistributable(*this, getLayoutBScale(),
+ scaleBShape, "ScaleB")))
+ return failure();
+
+ // Validate N dimension: scale_b[1] must match b[1]
+ if (scaleBShape[1] != bShape[1])
+ return emitOpError("Scale B N dimension [")
+ << scaleBShape[1] << "] must match B N dimension [" << bShape[1]
+ << "].";
+ }
+ }
+
+ // Validate scale K dimension compatibility if both scales are present and
+ // vectors
+ if (getScaleA() && getScaleB()) {
+ auto scaleAVecType = dyn_cast<VectorType>(getScaleAType());
+ auto scaleBVecType = dyn_cast<VectorType>(getScaleBType());
+
+ if (scaleAVecType && scaleBVecType) {
+ auto scaleAShape = scaleAVecType.getShape();
+ auto scaleBShape = scaleBVecType.getShape();
+
+ // Validate scale K dimension compatibility: scale_a[1] must match
+ // scale_b[0]
+ if (scaleAShape[1] != scaleBShape[0])
+ return emitOpError("Scale K dimension mismatch: scale_a has K=")
+ << scaleAShape[1] << " but scale_b has K=" << scaleBShape[0]
+ << ".";
+ }
+ }
return success();
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 98c9dc3f5e53a..7db887915b275 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -162,7 +162,9 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
if (isa<xegpu::StoreScatterOp>(op))
return getTileShape(op->getOpOperand(0));
- if (isa<xegpu::DpasOp>(op)) {
+ // Helper lambda to validate and get A/B tiles
+ auto validateABTiles = [&](Operation *op)
+ -> std::optional<std::pair<SmallVector<int64_t>, SmallVector<int64_t>>> {
std::optional<SmallVector<int64_t>> aTile =
getTileShape(op->getOpOperand(0));
std::optional<SmallVector<int64_t>> bTile =
@@ -175,16 +177,122 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
if ((*aTile)[1] != (*bTile)[0])
return std::nullopt;
+ return std::make_pair(*aTile, *bTile);
+ };
+
+ // Helper lambda to validate C tile
+ auto validateCTile = [&](Operation *op, unsigned cOperandIdx,
+ const SmallVector<int64_t> &aTile,
+ const SmallVector<int64_t> &bTile) -> bool {
+ if (op->getNumOperands() <= cOperandIdx)
+ return true;
+
+ std::optional<SmallVector<int64_t>> cTile =
+ getTileShape(op->getOpOperand(cOperandIdx));
+ int64_t expectedCTile[2] = {aTile[0], bTile[1]};
+ if (!cTile || !llvm::equal(*cTile, expectedCTile))
+ return false;
+ return true;
+ };
+
+ // Helper lambda to validate scale A tile for DpasMxOp
+ auto validateScaleATile =
+ [&](Operation *op, unsigned scaleAOperandIdx,
+ const SmallVector<int64_t> &aTile) -> std::optional<int64_t> {
+ std::optional<SmallVector<int64_t>> aScaleTile =
+ getTileShape(op->getOpOperand(scaleAOperandIdx));
+
+ if (!aScaleTile || aScaleTile->size() != 2)
+ return std::nullopt;
+
+ // Validate scale_a tile: [M_tile, K_scale]
+ // M dimension must match A's M dimension
+ if ((*aScaleTile)[0] != aTile[0])
+ return std::nullopt;
+
+ // Return the K scale factor
+ return (*aScaleTile)[1];
+ };
+
+ // Helper lambda to validate scale B tile for DpasMxOp
+ auto validateScaleBTile =
+ [&](Operation *op, unsigned scaleBOperandIdx,
+ const SmallVector<int64_t> &bTile) -> std::optional<int64_t> {
+ std::optional<SmallVector<int64_t>> bScaleTile =
+ getTileShape(op->getOpOperand(scaleBOperandIdx));
+
+ if (!bScaleTile || bScaleTile->size() != 2)
+ return std::nullopt;
+
+ // Validate scale_b tile: [K_scale, N_tile]
+ // N dimension must match B's N dimension
+ if ((*bScaleTile)[1] != bTile[1])
+ return std::nullopt;
+
+ // Return the K scale factor
+ return (*bScaleTile)[0];
+ };
+
+ if (isa<xegpu::DpasOp>(op)) {
+ auto abTiles = validateABTiles(op);
+ if (!abTiles)
+ return std::nullopt;
+
+ auto [aTile, bTile] = *abTiles;
+
// semantic check for C
- if (op->getNumOperands() == 3) {
- std::optional<SmallVector<int64_t>> cTile =
- getTileShape(op->getOpOperand(2));
- int64_t expectedCTile[2] = {(*aTile)[0], (*bTile)[1]};
- if (!cTile || !llvm::equal(*cTile, expectedCTile))
+ if (!validateCTile(op, 2, aTile, bTile))
+ return std::nullopt;
+
+ return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1]});
+ }
+
+ if (auto dpasMxOp = dyn_cast<xegpu::DpasMxOp>(op)) {
+ auto abTiles = validateABTiles(op);
+ if (!abTiles)
+ return std::nullopt;
+
+ auto [aTile, bTile] = *abTiles;
+
+ // Validate C tile if present using op-specific accessor
+ if (dpasMxOp.getAcc()) {
+ unsigned accOperandIdx = 2; // acc is the 3rd operand
+ if (!validateCTile(op, accOperandIdx, aTile, bTile))
+ return std::nullopt;
+ }
+
+ // Validate scale tiles if present using op-specific accessors
+ int64_t kScaleFactor = 1;
+ std::optional<int64_t> scaleAFactor;
+ std::optional<int64_t> scaleBFactor;
+
+ if (dpasMxOp.getScaleA()) {
+ unsigned scaleAOperandIdx = 2 + (dpasMxOp.getAcc() ? 1 : 0);
+ scaleAFactor = validateScaleATile(op, scaleAOperandIdx, aTile);
+ if (!scaleAFactor)
+ return std::nullopt;
+ }
+
+ if (dpasMxOp.getScaleB()) {
+ unsigned scaleBOperandIdx =
+ 2 + (dpasMxOp.getAcc() ? 1 : 0) + (dpasMxOp.getScaleA() ? 1 : 0);
+ scaleBFactor = validateScaleBTile(op, scaleBOperandIdx, bTile);
+ if (!scaleBFactor)
+ return std::nullopt;
+ }
+
+ // If both scales are present, their K dimensions must match
+ if (scaleAFactor && scaleBFactor) {
+ if (*scaleAFactor != *scaleBFactor)
return std::nullopt;
+ kScaleFactor = *scaleAFactor;
+ } else if (scaleAFactor) {
+ kScaleFactor = *scaleAFactor;
+ } else if (scaleBFactor) {
+ kScaleFactor = *scaleBFactor;
}
- return SmallVector<int64_t>({(*aTile)[0], (*aTile)[1], (*bTile)[1]});
+ return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1], kScaleFactor});
}
if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)
@@ -194,8 +302,8 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
return getTileShape(op->getOpOperand(0));
if (isa<vector::TransposeOp, vector::BroadcastOp, vector::StepOp,
- vector::ShapeCastOp, vector::ConstantMaskOp, vector::CreateMaskOp>(
- op))
+ vector::ShapeCastOp, vector::ConstantMaskOp, vector::CreateMaskOp,
+ vector::BitCastOp, vector::InterleaveOp, vector::DeinterleaveOp>(op))
return getTileShape(op->getOpResult(0));
return std::nullopt;
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
index c153db431c035..b70c8bd245853 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
@@ -234,14 +234,12 @@ struct SgToWiDpas : public OpConversionPattern<xegpu::DpasOp> {
LogicalResult
matchAndRewrite(xegpu::DpasOp op, OpAdaptor adaptor,
ConversionPatternRewriter &rewriter) const override {
- // llvm::errs() << "DpasOpPattern matchAndRewrite called\n";
// 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 failure();
- // llvm::errs() << "tryning to calculate wi types for dpas op\n";
auto wiResultTyOrFailure =
xegpu::getDistributedVectorType(op.getType(), layoutCd);
auto wiATypeOrFailure =
@@ -259,6 +257,38 @@ struct SgToWiDpas : public OpConversionPattern<xegpu::DpasOp> {
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 uArch *uArch = getUArch(xegpu::getChipStr(op).value_or(""));
+ if (uArch) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(
+ uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ if (uArchInstruction) {
+ auto wiAType = wiATypeOrFailure.value();
+ auto wiBType = wiBTypeOrFailure.value();
+ // Calculate total packed bit width = element bit width * vector size
+ unsigned aPackedBitWidth =
+ wiAType.getElementTypeBitWidth() * wiAType.getNumElements();
+ unsigned bPackedBitWidth =
+ wiBType.getElementTypeBitWidth() * wiBType.getNumElements();
+ unsigned expectedABitSize = uArchInstruction->getPackedFormatBitSizeA();
+ unsigned expectedBBitSize = uArchInstruction->getPackedFormatBitSizeB();
+
+ if (aPackedBitWidth % expectedABitSize != 0)
+ return rewriter.notifyMatchFailure(
+ op,
+ "A operand packed bit width must be a multiple of uArch packed "
+ "format requirement");
+ if (bPackedBitWidth % expectedBBitSize != 0)
+ return rewriter.notifyMatchFailure(
+ op,
+ "B operand packed bit width must be a multiple of uArch packed "
+ "format requirement");
+ }
+ }
+
auto newOp = xegpu::DpasOp::create(
rewriter, op->getLoc(), wiResultTyOrFailure.value(),
castValueTo(rewriter, cast<TypedValue<VectorType>>(adaptor.getLhs()),
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index d11ce207cc064..a6fd4ce09ebf4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -111,6 +111,25 @@ struct UnrollPattern : public OpRewritePattern<SourceOp> {
return SmallVector<Value>();
}
+ /// Helper to pack operands for DPAS-like operations with early return if
+ /// no unrolling is needed.
+ SmallVector<Value> packOperandForDpas(Value operand,
+ ArrayRef<int64_t> blockSize,
+ Location loc,
+ PatternRewriter &rewriter) const {
+ auto vecType = cast<VectorType>(operand.getType());
+ std::optional<SmallVector<int64_t>> grids =
+ computeShapeRatio(vecType.getShape(), blockSize);
+ assert(grids && "Expecting grids to be computed.");
+ auto numNewOps = computeProduct(*grids);
+ if (numNewOps == 1)
+ return SmallVector<Value>({operand});
+ VectorType newVecTy =
+ vecType.cloneWith(blockSize, vecType.getElementType());
+ SmallVector<Type> convertedTypes(numNewOps, newVecTy);
+ return pack(operand, convertedTypes, blockSize, loc, rewriter);
+ }
+
private:
const char *const packAttrName = "__xegpu_blocking_pack__";
const char *const unpackAttrName = "__xegpu_blocking_unpack__";
@@ -318,15 +337,6 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
- // expecting every operands is a 2D Vector
- if (llvm::any_of(op->getOperandTypes(), [&](Type type) {
- auto vecTy = dyn_cast<VectorType>(type);
- return !vecTy || vecTy.getRank() != 2;
- }))
- return failure();
-
- // A vector of 3 elements should be returned, representing M, K, N
- // respectively.
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape || targetShape->size() != 3)
return failure();
@@ -338,38 +348,16 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
int64_t bBlockSize[2] = {K, N};
int64_t cBlockSize[2] = {M, N};
- auto packWrapper = [&](TypedValue<VectorType> val,
- ArrayRef<int64_t> blockSize) {
- VectorType type = val.getType();
- std::optional<SmallVector<int64_t>> grids =
- computeShapeRatio(type.getShape(), blockSize);
- assert(grids && "Expecting grids to be computed.");
- auto numNewOps = computeProduct(*grids);
- if (numNewOps == 1)
- return SmallVector<Value>({val});
- VectorType newVecTy = type.cloneWith(blockSize, type.getElementType());
- SmallVector<Type> convertedTypes(numNewOps, newVecTy);
- SmallVector<Value> values =
- pack(val, convertedTypes, blockSize, loc, rewriter);
- return values;
- };
-
auto a = op.getLhs();
auto b = op.getRhs();
auto c = op.getAcc();
- auto aShape = a.getType().getShape();
- auto bShape = b.getType().getShape();
-
- SmallVector<Value> aVals, bVals, cVals;
- aVals = packWrapper(a, aBlockSize);
- bVals = packWrapper(b, bBlockSize);
-
+ SmallVector<Value> aVals = packOperandForDpas(a, aBlockSize, loc, rewriter);
+ SmallVector<Value> bVals = packOperandForDpas(b, bBlockSize, loc, rewriter);
+ SmallVector<Value> cVals;
if (c)
- cVals = packWrapper(c, cBlockSize);
+ cVals = packOperandForDpas(c, cBlockSize, loc, rewriter);
- // Skip the operation if every operand has an invalid blocking size (empty)
- // or if the original shape matches the blocking size (size == 1).
auto ranges = c ? SmallVector<ValueRange>({aVals, bVals, cVals})
: SmallVector<ValueRange>({aVals, bVals});
if (llvm::any_of(ranges, [](auto &v) { return v.size() == 0; }) ||
@@ -379,6 +367,8 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
VectorType resultTy = op.getResult().getType();
auto vecTy = VectorType::get(cBlockSize, resultTy.getElementType());
+ auto aShape = a.getType().getShape();
+ auto bShape = b.getType().getShape();
int64_t mIters = aShape[0] / M;
int64_t kIters = aShape[1] / K;
int64_t nIters = bShape[1] / N;
@@ -388,7 +378,7 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
for (int64_t j = 0; j < nIters; ++j) {
Value tmpC;
if (c)
- tmpC = cVals[i * nIters + j]; // init with acc
+ tmpC = cVals[i * nIters + j];
for (int64_t k = 0; k < kIters; ++k) {
Value aVec = aVals[i * kIters + k];
@@ -410,6 +400,86 @@ struct UnrollDpasOp : public UnrollPattern<xegpu::DpasOp> {
}
};
+struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
+ using UnrollPattern<xegpu::DpasMxOp>::UnrollPattern;
+ LogicalResult matchAndRewrite(xegpu::DpasMxOp op,
+ PatternRewriter &rewriter) const override {
+ Location loc = op.getLoc();
+
+ std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
+ if (!targetShape || targetShape->size() != 4)
+ return failure();
+ auto M = (*targetShape)[0];
+ auto K = (*targetShape)[1];
+ auto N = (*targetShape)[2];
+ auto S = (*targetShape)[3];
+
+ int64_t aBlockSize[2] = {M, K};
+ int64_t bBlockSize[2] = {K, N};
+ int64_t cBlockSize[2] = {M, N};
+ int64_t aScaleBlockSize[2] = {M, S};
+ int64_t bScaleBlockSize[2] = {S, N};
+
+ auto a = op.getA();
+ auto b = op.getB();
+ auto c = op.getAcc();
+ auto ascale = dyn_cast<TypedValue<VectorType>>(op.getScaleA());
+ auto bscale = dyn_cast<TypedValue<VectorType>>(op.getScaleB());
+
+ SmallVector<Value> aVals = packOperandForDpas(a, aBlockSize, loc, rewriter);
+ SmallVector<Value> bVals = packOperandForDpas(b, bBlockSize, loc, rewriter);
+ SmallVector<Value> cVals;
+ if (c)
+ cVals = packOperandForDpas(c, cBlockSize, loc, rewriter);
+ SmallVector<Value> aScaleVals;
+ if (ascale)
+ aScaleVals = packOperandForDpas(ascale, aScaleBlockSize, loc, rewriter);
+ SmallVector<Value> bScaleVals;
+ if (bscale)
+ bScaleVals = packOperandForDpas(bscale, bScaleBlockSize, loc, rewriter);
+
+ VectorType resultTy = op.getResult().getType();
+ auto vecTy = VectorType::get(cBlockSize, resultTy.getElementType());
+
+ auto aShape = a.getType().getShape();
+ auto bShape = b.getType().getShape();
+ int64_t mIters = aShape[0] / M;
+ int64_t kIters = aShape[1] / K;
+ int64_t nIters = bShape[1] / N;
+
+ SmallVector<Value> newOps;
+ xegpu::DpasMxOp newDpasMxOp;
+ for (int64_t i = 0; i < mIters; ++i) {
+ for (int64_t j = 0; j < nIters; ++j) {
+ Value tmpC;
+ if (c)
+ tmpC = cVals[i * nIters + j];
+
+ for (int64_t k = 0; k < kIters; ++k) {
+ Value aVec = aVals[i * kIters + k];
+ Value bVec = bVals[k * nIters + j];
+ SmallVector<Value> operands({aVec, bVec});
+ if (tmpC)
+ operands.push_back(tmpC);
+ if (ascale)
+ operands.push_back(aScaleVals[i * kIters + k]);
+ if (bscale)
+ operands.push_back(bScaleVals[k * nIters + j]);
+
+ newDpasMxOp = xegpu::DpasMxOp::create(
+ rewriter, loc, vecTy, operands,
+ xegpu::dropInstDataOnAttrs(op->getAttrs()));
+ tmpC = newDpasMxOp.getResult();
+ }
+ newOps.push_back(newDpasMxOp);
+ }
+ }
+ Value castOp = unpack(newOps, resultTy, cBlockSize, loc, rewriter);
+ rewriter.replaceOp(op, castOp);
+ return success();
+ }
+};
+
/// This pattern handles the unrolling of LoadGatherOp with offsets (gathered
/// load).
/// It unrolls the offsets and mask operands accordingly, and creates multiple
@@ -739,7 +809,8 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
void mlir::xegpu::populateXeGPUUnrollPatterns(
RewritePatternSet &patterns, const xegpu::UnrollOptions &options) {
patterns.add<UnrollCreateNdOp, UnrollPrefetchNdOp, UnrollLoadNdOp,
- UnrollStoreNdOp, UnrollDpasOp, UnrollLoadMatrixOp,
- UnrollStoreMatrixOp, UnrollLoadGatherOp, UnrollStoreScatterOp,
- UnrollConvertLayoutOp>(patterns.getContext(), options);
+ UnrollStoreNdOp, UnrollDpasOp, UnrollDpasMxOp,
+ UnrollLoadMatrixOp, UnrollStoreMatrixOp, UnrollLoadGatherOp,
+ UnrollStoreScatterOp, UnrollConvertLayoutOp>(
+ patterns.getContext(), options);
}
diff --git a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
index 2d1ce6eea17aa..41c4b2173eb38 100644
--- a/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
+++ b/mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp
@@ -183,6 +183,32 @@ xegpu::getDistributeLayoutAttr(const OpOperand &opr) {
return dpasOp.getLayoutCdAttr();
}
}
+ 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) {
+ return dpasMxOp.getLayoutAAttr();
+ } else if (idx < aSize + bSize) {
+ 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 (auto convertOp = dyn_cast<xegpu::ConvertLayoutOp>(op)) {
return convertOp.getInputLayoutAttr();
}
diff --git a/mlir/test/Dialect/XeGPU/invalid.mlir b/mlir/test/Dialect/XeGPU/invalid.mlir
index 7e6fb35cc6974..d0b5e2e07f02f 100644
--- a/mlir/test/Dialect/XeGPU/invalid.mlir
+++ b/mlir/test/Dialect/XeGPU/invalid.mlir
@@ -394,7 +394,7 @@ func.func @dpas_vc_1(%a : vector<8x8xf16>, %b: vector<8x16x2xf16>) {
// -----
func.func @dpas_vc_2(%a : vector<8x8x2xf16>, %b: vector<8x16x2xf16>) {
- // expected-error at +1 {{expecting lhs and result to be a 2D vector, and rhs to be either 2D or 3D (packed) vector}}
+ // expected-error at +1 {{op A operand must be a 2D vector}}
%1 = xegpu.dpas %a, %b : vector<8x8x2xf16>, vector<8x16x2xf16> -> vector<8x16xf32>
return
}
@@ -420,13 +420,6 @@ func.func @dpas_5(%a : vector<8x16xf16>, %b: vector<8x8x2xf16>) {
return
}
-// -----
-func.func @dpas_simt_1(%a : vector<8xf16>, %b: vector<15xf16>) {
- // expected-error at +1 {{Expecting B operand to be a multiple of 32 bits}}
- %1 = xegpu.dpas %a, %b : vector<8xf16>, vector<15xf16> -> vector<8xf32>
- return
-}
-
// -----
func.func @tensor_desc_invalid_rank_1(%src: memref<24x32xf32>) {
%0 = xegpu.create_nd_tdesc %src : memref<24x32xf32> ->
@@ -702,7 +695,97 @@ 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 {{Expecting the acc type to be the same as result.}}
+ // 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>
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>
+ 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>
+ 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>
+ 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>
+ 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>
+ 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>
+ return
+}
+
+
+// -----
+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>
+ 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>
+ 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>
+ return
+}
+
+// -----
+#layout_a = #xegpu.layout<sg_layout = [1, 1], sg_data = [8, 32]>
+#layout_b = #xegpu.layout<sg_layout = [1, 1], sg_data = [32, 16]>
+#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>
+ return
+}
+
+// -----
+#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>
+ return
+}
+
+// -----
+#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>
+ return
+}
diff --git a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
index c2aac8fa6cf0b..64e79e98f2d71 100644
--- a/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
+++ b/mlir/test/Dialect/XeGPU/xegpu-blocking.mlir
@@ -696,3 +696,51 @@ gpu.module @test_kernel {
gpu.return
}
}
+
+// -----
+#l1 = #xegpu.layout<inst_data = [8, 32]>
+#l2 = #xegpu.layout<inst_data = [32, 16]>
+#l3 = #xegpu.layout<inst_data = [8, 16]>
+#l1_scale = #xegpu.layout<inst_data = [8, 1]>
+#l2_scale = #xegpu.layout<inst_data = [1, 16]>
+gpu.module @test_kernel {
+ gpu.func @dpas_mx(%A: memref<1024x1024xf4E2M1FN>, %B: memref<1024x1024xf4E2M1FN>, %C: memref<1024x1024xf32>, %scale_a: memref<1024x64xf8E8M0FNU>, %scale_b: memref<64x1024xf8E8M0FNU>) {
+ %c0 = arith.constant 0 : index
+ %c16 = arith.constant 16 : index
+ %c32 = arith.constant 32 : index
+ %c64 = arith.constant 64 : index
+ %c1024 = arith.constant 1024 : index
+ %block_id_x = gpu.block_id x
+ %block_id_y = gpu.block_id y
+ %m = arith.muli %block_id_x, %c16 : index
+ %n = arith.muli %block_id_y, %c32 : index
+
+ %c_tdesc = xegpu.create_nd_tdesc %C : memref<1024x1024xf32> -> !xegpu.tensor_desc<16x32xf32, #l3>
+ %c_init = xegpu.load_nd %c_tdesc[0, 0] {layout = #l3}: !xegpu.tensor_desc<16x32xf32, #l3> -> vector<16x32xf32>
+
+ %a_tdesc = xegpu.create_nd_tdesc %A : memref<1024x1024xf4E2M1FN> -> !xegpu.tensor_desc<16x64xf4E2M1FN, #l1>
+ %b_tdesc = xegpu.create_nd_tdesc %B : memref<1024x1024xf4E2M1FN> -> !xegpu.tensor_desc<64x32xf4E2M1FN, #l2>
+ %scale_a_tdesc = xegpu.create_nd_tdesc %scale_a : memref<1024x64xf8E8M0FNU> -> !xegpu.tensor_desc<16x2xf8E8M0FNU, #l1_scale>
+ %scale_b_tdesc = xegpu.create_nd_tdesc %scale_b : memref<64x1024xf8E8M0FNU> -> !xegpu.tensor_desc<2x32xf8E8M0FNU, #l2_scale>
+
+ %out = scf.for %k = %c0 to %c1024 step %c64
+ iter_args(%arg2 = %c_init)
+ -> (vector<16x32xf32>) {
+ //CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x32xf4E2M1FN>
+ %a = xegpu.load_nd %a_tdesc[%c0, %k] {layout = #l1}: !xegpu.tensor_desc<16x64xf4E2M1FN, #l1> -> vector<16x64xf4E2M1FN>
+ //CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<32x16xf4E2M1FN>
+ %b = xegpu.load_nd %b_tdesc[%k, %c0] {layout = #l2}: !xegpu.tensor_desc<64x32xf4E2M1FN, #l2> -> vector<64x32xf4E2M1FN>
+ //CHECK-COUNT-4: xegpu.load_nd {{.*}} -> vector<8x1xf8E8M0FNU>
+ %sa = xegpu.load_nd %scale_a_tdesc[%c0, %c0] {layout = #l1_scale}: !xegpu.tensor_desc<16x2xf8E8M0FNU, #l1_scale> -> vector<16x2xf8E8M0FNU>
+ //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>
+ scf.yield %c : vector<16x32xf32>
+ } {layout_result_0 = #l3}
+ //CHECK-COUNT-4: xegpu.store_nd {{.*}} : vector<8x16xf32>, !xegpu.tensor_desc<8x16xf32>
+ xegpu.store_nd %out, %c_tdesc[0, 0] {layout = #l3}: vector<16x32xf32>, !xegpu.tensor_desc<16x32xf32, #l3>
+ gpu.return
+ }
+}
+
More information about the Mlir-commits
mailing list