[Mlir-commits] [mlir] [MLIR][XeGPU] Unroll Dpasmx Op (PR #195179)
Jianhui Li
llvmlistbot at llvm.org
Tue May 5 16:01:19 PDT 2026
https://github.com/Jianhui-Li updated https://github.com/llvm/llvm-project/pull/195179
>From 618b5fe929e70aada84a6baa8fafde7bb4f1697e Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 23 Apr 2026 23:33:41 +0000
Subject: [PATCH 01/10] [mlir][XeGPU] Add DpasMx op definition and layout
support
This patch extends the DpasMx operation with scale layout attributes and
implements layout setup and propagation for MXFP (microscaling floating point)
operations.
Op definition changes (XeGPUOps.td):
- Add layout_a_scale and layout_b_scale attributes to DpasMx op
- Remove restrictive AllElementTypesMatch trait to allow different types for
A and B operands with scale factors
Layout support (XeGPULayoutImpl.cpp/h):
- setupDpasMxLayout: Creates anchor layouts for all DpasMx operands (A, B, C/D,
scale_a, scale_b) across different layout kinds (Subgroup, InstData, Lane)
- Derives scale layouts from parent matrix layouts by adjusting dimensions based
on scaling factors (K/32 for scales)
Layout propagation (XeGPUPropagateLayout.cpp):
- visitDpasMxOp: Propagates layout attributes from DpasMx op to its operands
during dataflow analysis
This infrastructure enables proper layout tracking for mixed-precision matrix
operations with separate scale factors during workgroup-to-subgroup distribution.
Co-Authored-By: Claude Sonnet 4.5 <noreply at anthropic.com>
---
.../include/mlir/Dialect/XeGPU/IR/XeGPUOps.td | 5 +-
.../XeGPU/Transforms/XeGPULayoutImpl.h | 18 ++-
.../XeGPU/Transforms/XeGPULayoutImpl.cpp | 151 ++++++++++++++++++
.../XeGPU/Transforms/XeGPUPropagateLayout.cpp | 39 +++++
4 files changed, 205 insertions(+), 8 deletions(-)
diff --git a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
index 31fe93d209a6d..8b6763e234091 100644
--- a/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
+++ b/mlir/include/mlir/Dialect/XeGPU/IR/XeGPUOps.td
@@ -1554,7 +1554,6 @@ def XeGPU_TruncfOp
}
def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
- AllElementTypesMatch<["a", "b"]>,
AnchorLayoutInterface]> {
let summary = "It performs scaled mma computation";
@@ -1601,7 +1600,9 @@ def XeGPU_DpasMxOp : XeGPU_Op<"dpas_mx", [Pure, AttrSizedOperandSegments,
VectorOfRankAndType<[1, 2], [F8E8M0FNU]>]>>:$scale_b,
OptionalAttr<DistributeLayoutAttr>:$layout_a,
OptionalAttr<DistributeLayoutAttr>:$layout_b,
- OptionalAttr<DistributeLayoutAttr>:$layout_cd);
+ OptionalAttr<DistributeLayoutAttr>:$layout_cd,
+ OptionalAttr<DistributeLayoutAttr>:$layout_a_scale,
+ OptionalAttr<DistributeLayoutAttr>:$layout_b_scale);
let results = (outs XeGPU_DpasResType:$result);
let extraClassDeclaration = [{
diff --git a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
index 83eb939cf1bec..c68e7334434d6 100644
--- a/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
+++ b/mlir/include/mlir/Dialect/XeGPU/Transforms/XeGPULayoutImpl.h
@@ -39,12 +39,6 @@ LogicalResult propagateLayouts(OpBuilder &builder, Operation *target,
LogicalResult resolveLayoutConflicts(Operation *target);
-/// [to-be-deprecated] Set the DistributeLayoutAttr for each OpOperand and
-/// OpResult of of the given operation. If the operation contains regions, it is
-/// also applied recursively to the contained operations operation.
-/// TODO: To be replaced by recoverTemporaryLayouts()
-void recoverTemporaryLayoutsDeprecated(Operation *op);
-
/// Attach layout attributes to all vector-type operands of operations within
/// the given operation's nested region. Reports an error if any vector operand
/// lacks a layout attribute.
@@ -199,6 +193,18 @@ setupDpasLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
VectorType cdTy, DistributeLayoutAttr consumerLayout, int numSg,
const uArch::uArch *uArch);
+/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
+/// B_scale). The numSg and consumerLayout (optional) are only used by sg layout
+/// creation. A_scale and B_scale are optional.
+std::optional<std::tuple<DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr, DistributeLayoutAttr,
+ DistributeLayoutAttr>>
+setupDpasMxLayout(LayoutKind layoutKind, VectorType aTy, VectorType bTy,
+ VectorType cdTy, std::optional<VectorType> aScaleTy,
+ std::optional<VectorType> bScaleTy,
+ DistributeLayoutAttr consumerLayout, int numSg,
+ const uArch::uArch *uArch);
+
/// Gets the expected layout for a given consumer operand. This will check if
/// the owning operation of the consumer operand is one of the special layout
/// users and determine the expected layout accordingly.
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
index 7d48315eec6ff..8f07ddf09e9a4 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPULayoutImpl.cpp
@@ -1426,3 +1426,154 @@ xegpu::DistributeLayoutAttr xegpu::getConsumerLayoutAt(OpOperand &operand) {
// the operand.
return xegpu::getDistributeLayoutAttr(operand.get());
}
+/// Sets up the anchor layouts for dpas_mx operands (A, B, C/D, A_scale, and
+/// B_scale). The numSg and consumerLayout (optional) are only used by sg layout
+/// creation.
+std::optional<
+ std::tuple<xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr, xegpu::DistributeLayoutAttr,
+ xegpu::DistributeLayoutAttr>>
+xegpu::setupDpasMxLayout(xegpu::LayoutKind layoutKind, VectorType aTy,
+ VectorType bTy, VectorType cdTy,
+ std::optional<VectorType> aScaleTy,
+ std::optional<VectorType> bScaleTy,
+ xegpu::DistributeLayoutAttr consumerLayout, int numSg,
+ const xegpu::uArch::uArch *uArch) {
+ auto context = aTy.getContext();
+ const int subgroupSize = uArch->getSubgroupSize();
+
+ // Helper to create scale layout from parent layout
+ auto createScaleLayout = [&](VectorType parentTy, VectorType scaleTy,
+ xegpu::DistributeLayoutAttr parentLayout,
+ bool isBScale) -> xegpu::DistributeLayoutAttr {
+ if (!scaleTy || !parentLayout)
+ return nullptr;
+
+ // Calculate scaling factor by dividing parent shape by scale shape
+ ArrayRef<int64_t> parentShape = parentTy.getShape();
+ ArrayRef<int64_t> scaleShape = scaleTy.getShape();
+ int64_t scaleFactor = parentShape.back() / scaleShape.back();
+ int64_t rank = parentLayout.getRank();
+ assert(rank == 2 && "dpas layouts must be two dimensions");
+
+ SmallVector<int64_t> sgLayout = parentLayout.getEffectiveSgLayoutAsInt();
+ SmallVector<int64_t> sgData = parentLayout.getEffectiveSgDataAsInt();
+ SmallVector<int64_t> instData = parentLayout.getEffectiveInstDataAsInt();
+ SmallVector<int64_t> laneLayout =
+ parentLayout.getEffectiveLaneLayoutAsInt();
+ SmallVector<int64_t> laneData = parentLayout.getEffectiveLaneDataAsInt();
+ auto order = parentLayout.getOrder();
+
+ // Divide last dimension by scaling factor
+ if (!sgData.empty())
+ sgData.back() = sgData.back() / scaleFactor;
+ if (!instData.empty())
+ instData.back() = instData.back() / scaleFactor;
+
+ if (isBScale) {
+ // For B scale: lane_layout = [min(subgroupSize, scaleFactor), 1]
+ // lane_data = [1, number of scale elements]
+ laneLayout[rank - 2] =
+ std::min(static_cast<int64_t>(subgroupSize), scaleShape[rank-2]);
+ laneLayout[rank - 1] = 1;
+ laneData[rank - 2] = 1;
+ laneData[rank - 1] = scaleShape.back();
+ } else {
+ laneData.back() = scaleShape.back();
+ }
+
+ return xegpu::LayoutAttr::get(
+ context,
+ sgLayout.empty()
+ ? nullptr
+ : DenseI32ArrayAttr::get(
+ context, SmallVector<int>(sgLayout.begin(), sgLayout.end())),
+ sgData.empty()
+ ? nullptr
+ : DenseI32ArrayAttr::get(
+ context, SmallVector<int>(sgData.begin(), sgData.end())),
+ instData.empty()
+ ? nullptr
+ : DenseI32ArrayAttr::get(
+ context, SmallVector<int>(instData.begin(), instData.end())),
+ laneLayout.empty() ? nullptr
+ : DenseI32ArrayAttr::get(
+ context, SmallVector<int>(laneLayout.begin(),
+ laneLayout.end())),
+ laneData.empty()
+ ? nullptr
+ : DenseI32ArrayAttr::get(
+ context, SmallVector<int>(laneData.begin(), laneData.end())),
+ order);
+ };
+
+ if (layoutKind == xegpu::LayoutKind::Subgroup) {
+ assert(numSg > 0 &&
+ "Number of subgroups must be provided for sg layout creation.");
+ auto dpasLayouts =
+ getupDpasSubgroupLayouts(context, aTy, bTy, cdTy, consumerLayout, numSg, uArch);
+ if (!dpasLayouts)
+ return std::nullopt;
+
+ auto [dpasALayout, dpasBLayout, dpasCDLayout] = *dpasLayouts;
+
+ // Create scale layouts
+ auto aScaleLayout =
+ aScaleTy.has_value()
+ ? createScaleLayout(aTy, *aScaleTy, dpasALayout, false)
+ : nullptr;
+ auto bScaleLayout =
+ bScaleTy.has_value()
+ ? createScaleLayout(bTy, *bScaleTy, dpasBLayout, true)
+ : nullptr;
+
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ } else if (layoutKind == xegpu::LayoutKind::InstData) {
+ auto instDataVecs = getDpasInstDataVectors(aTy, bTy, cdTy, uArch);
+ if (!instDataVecs)
+ return std::nullopt;
+ auto [instDataA, instDataB, instDataCD] = *instDataVecs;
+
+ auto dpasALayout = xegpu::LayoutAttr::get(
+ context, SmallVector<int>(instDataA.begin(), instDataA.end()));
+ auto dpasBLayout = xegpu::LayoutAttr::get(
+ context, SmallVector<int>(instDataB.begin(), instDataB.end()));
+ auto dpasCDLayout = xegpu::LayoutAttr::get(
+ context, SmallVector<int>(instDataCD.begin(), instDataCD.end()));
+
+ // Create scale layouts
+ auto aScaleLayout =
+ aScaleTy.has_value()
+ ? createScaleLayout(aTy, *aScaleTy, dpasALayout, false)
+ : nullptr;
+ auto bScaleLayout =
+ bScaleTy.has_value()
+ ? createScaleLayout(bTy, *bScaleTy, dpasBLayout, true)
+ : nullptr;
+
+ return std::make_tuple(dpasALayout, dpasBLayout, dpasCDLayout, aScaleLayout,
+ bScaleLayout);
+ } else if (layoutKind == xegpu::LayoutKind::Lane) {
+ const auto *uArchInstruction =
+ dyn_cast<xegpu::uArch::SubgroupMatrixMultiplyAcc>(uArch->getInstruction(
+ xegpu::uArch::InstructionKind::SubgroupMatrixMultiplyAcc));
+ auto aLayout = getDefaultLaneLayout2DBlockIo(
+ aTy, uArch, uArchInstruction->getPackedFormatBitSizeA());
+ auto bLayout = getDefaultLaneLayout2DBlockIo(
+ bTy, uArch, uArchInstruction->getPackedFormatBitSizeB(), true);
+ auto cdLayout = getDefaultLaneLayout2DBlockIo(cdTy, uArch);
+
+ // Create scale layouts
+ auto aScaleLayout = aScaleTy.has_value()
+ ? createScaleLayout(aTy, *aScaleTy, aLayout, false)
+ : nullptr;
+ auto bScaleLayout = bScaleTy.has_value()
+ ? createScaleLayout(bTy, *bScaleTy, bLayout, true)
+ : nullptr;
+
+ return std::make_tuple(aLayout, bLayout, cdLayout, aScaleLayout,
+ bScaleLayout);
+ }
+ return std::nullopt;
+}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
index 43998ed41f7aa..f1709e95d34a9 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUPropagateLayout.cpp
@@ -316,6 +316,9 @@ class LayoutInfoPropagation
void visitDpasOp(xegpu::DpasOp dpas, ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
+ void visitDpasMxOp(xegpu::DpasMxOp dpasMx, ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results);
+
void visitStoreNdOp(xegpu::StoreNdOp store,
ArrayRef<LayoutInfoLattice *> operands,
ArrayRef<const LayoutInfoLattice *> results);
@@ -426,6 +429,9 @@ LogicalResult LayoutInfoPropagation::visitOperation(
TypeSwitch<Operation *>(op)
.Case(
[&](xegpu::DpasOp dpasOp) { visitDpasOp(dpasOp, operands, results); })
+ .Case([&](xegpu::DpasMxOp dpasMxOp) {
+ visitDpasMxOp(dpasMxOp, operands, results);
+ })
.Case([&](xegpu::StoreNdOp storeNdOp) {
visitStoreNdOp(storeNdOp, operands, results);
})
@@ -804,6 +810,39 @@ void LayoutInfoPropagation::visitDpasOp(
std::tie(requiredALayout, requiredBLayout, requiredCDLayoutAttr) = *layouts;
dpas.setLayoutAAttr(requiredALayout);
+
+/// Propagate layout for DpasMxOp operands using the layout attributes.
+/// DpasMxOp has operands: a, b, acc (optional), scale_a (optional), scale_b (optional)
+void LayoutInfoPropagation::visitDpasMxOp(
+ xegpu::DpasMxOp dpasMx, ArrayRef<LayoutInfoLattice *> operands,
+ ArrayRef<const LayoutInfoLattice *> results) {
+
+ // Get the layout attributes from the operation
+ xegpu::DistributeLayoutAttr layoutA = dpasMx.getLayoutAAttr();
+ xegpu::DistributeLayoutAttr layoutB = dpasMx.getLayoutBAttr();
+ xegpu::DistributeLayoutAttr layoutCD = dpasMx.getLayoutCdAttr();
+ xegpu::DistributeLayoutAttr layoutAScale = dpasMx.getLayoutAScaleAttr();
+ xegpu::DistributeLayoutAttr layoutBScale = dpasMx.getLayoutBScaleAttr();
+
+ // Propagate layouts to operands based on their positions:
+ // operands[0] = a, operands[1] = b, operands[2] = acc (optional),
+ // operands[3] = scale_a (optional), operands[4] = scale_b (optional)
+
+ if (layoutA && operands.size() > 0)
+ propagateIfChanged(operands[0], operands[0]->meet(LayoutInfo(layoutA)));
+
+ if (layoutB && operands.size() > 1)
+ propagateIfChanged(operands[1], operands[1]->meet(LayoutInfo(layoutB)));
+
+ if (layoutCD && operands.size() > 2)
+ propagateIfChanged(operands[2], operands[2]->meet(LayoutInfo(layoutCD)));
+
+ if (layoutAScale && operands.size() > 3)
+ propagateIfChanged(operands[3], operands[3]->meet(LayoutInfo(layoutAScale)));
+
+ if (layoutBScale && operands.size() > 4)
+ propagateIfChanged(operands[4], operands[4]->meet(LayoutInfo(layoutBScale)));
+}
dpas.setLayoutBAttr(requiredBLayout);
dpas.setLayoutCdAttr(requiredCDLayoutAttr);
dpasALayout = LayoutInfo(requiredALayout);
>From 7ea868f07189f52944301e924f52fcf210e92ef6 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 23 Apr 2026 23:45:13 +0000
Subject: [PATCH 02/10] [mlir][XeGPU] Add unrolling and blocking support for
DpasMx ops
This patch implements unrolling and blocking (tiling) support for DpasMx
operations with comprehensive test coverage.
Blocking support (XeGPUBlocking.cpp):
- getTileShape for DpasMxOp: Computes tile shapes for A, B, and C/D matrices
based on DPAS instruction parameters
- Validates matrix dimensions and scale factor semantics (scale K dimension
must be parent K / 32)
- Ensures A and B matrix dimensions are compatible for matrix multiplication
Unrolling support (XeGPUUnroll.cpp):
- UnrollDpasMxOp pattern: Unrolls DpasMx operations into smaller blocked
operations following a 4D target shape [M, K, N, S] where S is the scale
dimension
- Properly distributes scale factors across unrolled tiles by computing
appropriate scale tile offsets and shapes
- Handles optional accumulator and scale operands
Test coverage (propagate-layout-subgroup.mlir):
- Tests for layout propagation through DpasMx operations with scales
- Validates unrolling behavior with different tile configurations
Depends on: PR #3 (DpasMx op definition and layout support)
Co-Authored-By: Claude Sonnet 4.5 <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUBlocking.cpp | 164 +++++++++++++++--
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 170 +++++++++++++++++-
.../XeGPU/propagate-layout-subgroup.mlir | 60 +++++++
3 files changed, 374 insertions(+), 20 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 7fc5d2fffae51..5e841e0421178 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -163,29 +163,166 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
if (isa<xegpu::StoreScatterOp>(op))
return getTileShape(op->getOpOperand(0));
- if (isa<xegpu::DpasOp>(op)) {
- std::optional<SmallVector<int64_t>> aTile =
- getTileShape(op->getOpOperand(0));
- std::optional<SmallVector<int64_t>> bTile =
- getTileShape(op->getOpOperand(1));
+ // 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 = getTileShape(op->getOpOperand(1));
+
+ LLVM_DEBUG(llvm::dbgs() << " aTile: "
+ << (aTile ? llvm::join(llvm::map_range(*aTile,
+ [](int64_t v) { return std::to_string(v); }), "x")
+ : "nullopt")
+ << "\n");
+ LLVM_DEBUG(llvm::dbgs() << " bTile: "
+ << (bTile ? llvm::join(llvm::map_range(*bTile,
+ [](int64_t v) { return std::to_string(v); }), "x")
+ : "nullopt")
+ << "\n");
if (!aTile || aTile->size() != 2 || !bTile || bTile->size() != 2)
return std::nullopt;
// semantic check for A and B
- if ((*aTile)[1] != (*bTile)[0])
+ if ((*aTile)[1] != (*bTile)[0]) {
+ LLVM_DEBUG(llvm::dbgs() << " A/B semantic check failed: aTile[1]="
+ << (*aTile)[1] << " != bTile[0]=" << (*bTile)[0]
+ << "\n");
+ 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]};
+ LLVM_DEBUG(llvm::dbgs() << " cTile: "
+ << (cTile ? llvm::join(llvm::map_range(*cTile,
+ [](int64_t v) { return std::to_string(v); }), "x")
+ : "nullopt")
+ << ", expected: " << expectedCTile[0] << "x"
+ << expectedCTile[1] << "\n");
+ if (!cTile || !llvm::equal(*cTile, expectedCTile))
+ return false;
+ return true;
+ };
+
+ // Helper lambda to validate scale A/B tiles for DpasMxOp
+ auto validateABScaleTiles = [&](Operation *op, unsigned scaleAOperandIdx,
+ unsigned scaleBOperandIdx,
+ const SmallVector<int64_t> &aTile,
+ const SmallVector<int64_t> &bTile)
+ -> std::optional<int64_t> {
+ std::optional<SmallVector<int64_t>> aScaleTile =
+ getTileShape(op->getOpOperand(scaleAOperandIdx));
+ std::optional<SmallVector<int64_t>> bScaleTile =
+ getTileShape(op->getOpOperand(scaleBOperandIdx));
+
+ LLVM_DEBUG(llvm::dbgs() << " aScaleTile: "
+ << (aScaleTile ? llvm::join(llvm::map_range(*aScaleTile,
+ [](int64_t v) { return std::to_string(v); }), "x")
+ : "nullopt")
+ << "\n");
+ LLVM_DEBUG(llvm::dbgs() << " bScaleTile: "
+ << (bScaleTile ? llvm::join(llvm::map_range(*bScaleTile,
+ [](int64_t v) { return std::to_string(v); }), "x")
+ : "nullopt")
+ << "\n");
+
+ if (!aScaleTile || aScaleTile->size() != 2 ||
+ !bScaleTile || bScaleTile->size() != 2)
+ return std::nullopt;
+
+ // Validate scale tile dimensions
+ assert((*aScaleTile)[0] == aTile[0] &&
+ "aScaleTile[0] must equal aTile[0]");
+ assert((*bScaleTile)[1] == bTile[1] &&
+ "bScaleTile[1] must equal bTile[1]");
+
+ if ((*aScaleTile)[1] != (*bScaleTile)[0]) {
+ LLVM_DEBUG(llvm::dbgs() << " scale A/B semantic check failed: aScaleTile[1]="
+ << (*aScaleTile)[1] << " != bScaleTile[0]="
+ << (*bScaleTile)[0] << "\n");
+ return std::nullopt;
+ }
+
+ // Return the K scale factor
+ return (*aScaleTile)[1];
+ };
+
+ if (isa<xegpu::DpasOp>(op)) {
+ LLVM_DEBUG(llvm::dbgs() << "getTileShape for DpasOp: " << *op << "\n");
+
+ 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;
+
+ LLVM_DEBUG(llvm::dbgs() << " result: [" << aTile[0] << ", "
+ << aTile[1] << ", " << bTile[1] << "]\n");
+ return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1]});
+ }
+
+ if (auto dpasMxOp = dyn_cast<xegpu::DpasMxOp>(op)) {
+ LLVM_DEBUG(llvm::dbgs() << "getTileShape for DpasMxOp: " << *op << "\n");
+
+ auto abTiles = validateABTiles(op);
+ if (!abTiles)
+ return std::nullopt;
+
+ auto [aTile, bTile] = *abTiles;
+
+ // Get operand indices using AttrSizedOperandSegments
+ auto segmentSizesAttr = dpasMxOp->getAttrOfType<DenseI32ArrayAttr>(
+ dpasMxOp.getOperandSegmentSizesAttrName());
+ if (!segmentSizesAttr)
+ return std::nullopt;
+
+ auto segmentSizes = segmentSizesAttr.asArrayRef();
+ unsigned aSize = segmentSizes[0];
+ unsigned bSize = segmentSizes[1];
+ unsigned accSize = segmentSizes[2];
+ unsigned scaleASize = segmentSizes[3];
+ unsigned scaleBSize = segmentSizes[4];
+
+ // Validate C tile if present
+ if (accSize > 0) {
+ unsigned accOperandIdx = aSize + bSize;
+ if (!validateCTile(op, accOperandIdx, aTile, bTile))
+ return std::nullopt;
+ }
+
+ // Validate scale tiles if present
+ int64_t kScaleFactor = 1;
+ if (scaleASize > 0 && scaleBSize > 0) {
+ unsigned scaleAOperandIdx = aSize + bSize + accSize;
+ unsigned scaleBOperandIdx = scaleAOperandIdx + scaleASize;
+
+ auto scaleFactor = validateABScaleTiles(op, scaleAOperandIdx,
+ scaleBOperandIdx, aTile, bTile);
+ if (!scaleFactor)
return std::nullopt;
+
+ kScaleFactor = *scaleFactor;
}
- return SmallVector<int64_t>({(*aTile)[0], (*aTile)[1], (*bTile)[1]});
+ LLVM_DEBUG(llvm::dbgs() << " result: [" << aTile[0] << ", "
+ << aTile[1] << ", " << bTile[1] << ", "
+ << kScaleFactor << "]\n");
+ return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1], kScaleFactor});
}
if (OpTrait::hasElementwiseMappableTraits(op) && op->getNumResults() == 1)
@@ -195,7 +332,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>(
+ vector::ShapeCastOp, vector::ConstantMaskOp, vector::CreateMaskOp,
+ vector::BitCastOp, vector::InterleaveOp, vector::DeinterleaveOp>(
op))
return getTileShape(op->getOpResult(0));
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 51693da389a49..d37e005be7122 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -477,6 +477,163 @@ 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();
+
+ LLVM_DEBUG(llvm::dbgs() << "UnrollDpasMxOp: original op: " << op << "\n");
+
+ // 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();
+ auto M = (*targetShape)[0];
+ auto K = (*targetShape)[1];
+ auto N = (*targetShape)[2];
+ auto S = (*targetShape)[3];
+
+ LLVM_DEBUG(llvm::dbgs() << " targetShape: M=" << M << ", K=" << K
+ << ", N=" << N << ", S=" << S << "\n");
+
+ 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};
+
+ LLVM_DEBUG(llvm::dbgs() << " aBlockSize: [" << M << ", " << K << "]\n");
+ LLVM_DEBUG(llvm::dbgs() << " bBlockSize: [" << K << ", " << N << "]\n");
+ LLVM_DEBUG(llvm::dbgs() << " cBlockSize: [" << M << ", " << N << "]\n");
+ LLVM_DEBUG(llvm::dbgs()
+ << " aScaleBlockSize: [" << M << ", " << K / 32 << "]\n");
+ LLVM_DEBUG(llvm::dbgs()
+ << " bScaleBlockSize: [" << K / 32 << ", " << N << "]\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.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());
+
+ auto aShape = a.getType().getShape();
+ auto bShape = b.getType().getShape();
+
+ SmallVector<Value> aVals, bVals, cVals, aScaleVals, bScaleVals;
+ aVals = packWrapper(a, aBlockSize);
+ bVals = packWrapper(b, bBlockSize);
+
+ if (c)
+ cVals = packWrapper(c, cBlockSize);
+ if (ascale)
+ aScaleVals = packWrapper(ascale, aScaleBlockSize);
+ if (bscale)
+ bScaleVals = packWrapper(bscale, bScaleBlockSize);
+
+ LLVM_DEBUG(llvm::dbgs() << " aVals size: " << aVals.size() << "\n");
+ LLVM_DEBUG(llvm::dbgs() << " bVals size: " << bVals.size() << "\n");
+ LLVM_DEBUG(llvm::dbgs() << " cVals size: " << cVals.size() << "\n");
+ LLVM_DEBUG(llvm::dbgs()
+ << " aScaleVals size: " << aScaleVals.size() << "\n");
+ LLVM_DEBUG(llvm::dbgs()
+ << " bScaleVals size: " << bScaleVals.size() << "\n");
+
+ // 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; }) ||
+ // llvm::all_of(ranges, [](auto &v) { return v.size() == 1; }))
+ // return failure();
+
+ VectorType resultTy = op.getResult().getType();
+ auto vecTy = VectorType::get(cBlockSize, resultTy.getElementType());
+
+ int64_t mIters = aShape[0] / M;
+ int64_t kIters = aShape[1] / K;
+ int64_t nIters = bShape[1] / N;
+
+ LLVM_DEBUG(llvm::dbgs() << " mIters=" << mIters << ", kIters=" << kIters
+ << ", nIters=" << nIters << "\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]; // init with acc
+
+ for (int64_t k = 0; k < kIters; ++k) {
+ Value aVec = aVals[i * kIters + k];
+ Value bVec = bVals[k * nIters + j];
+
+ LLVM_DEBUG(llvm::dbgs() << " [i=" << i << ", j=" << j << ", k=" << k
+ << "] aVec: " << aVec.getType()
+ << ", bVec: " << bVec.getType() << "\n");
+
+ SmallVector<Value> operands({aVec, bVec});
+ if (tmpC) {
+ operands.push_back(tmpC);
+ LLVM_DEBUG(llvm::dbgs() << " tmpC: " << tmpC.getType() << "\n");
+ }
+ if (ascale) {
+ Value aScaleVec = aScaleVals[i * kIters + k];
+ operands.push_back(aScaleVec);
+ LLVM_DEBUG(llvm::dbgs()
+ << " aScaleVec: " << aScaleVec.getType() << "\n");
+ }
+ if (bscale) {
+ Value bScaleVec = bScaleVals[k * nIters + j];
+ operands.push_back(bScaleVec);
+ LLVM_DEBUG(llvm::dbgs()
+ << " bScaleVec: " << bScaleVec.getType() << "\n");
+ }
+ LLVM_DEBUG(llvm::dbgs() << " total operands: " << operands.size()
+ << ", resTy: " << vecTy << "\n");
+ newDpasMxOp = xegpu::DpasMxOp::create(
+ rewriter, loc, vecTy, operands,
+ xegpu::dropInstDataOnAttrs(op->getAttrs()));
+ LLVM_DEBUG(llvm::dbgs() << " created: " << newDpasMxOp << "\n");
+ }
+ newOps.push_back(newDpasMxOp);
+ }
+ }
+
+ LLVM_DEBUG(llvm::dbgs()
+ << " total new DpasMxOps: " << newOps.size() << "\n");
+
+ Value castOp = unpack(newOps, resultTy, cBlockSize, loc, rewriter);
+ rewriter.replaceOp(op, castOp);
+ return success();
+ }
+};
+
struct UnrollLoadGatherOp : public UnrollPattern<xegpu::LoadGatherOp> {
using UnrollPattern<xegpu::LoadGatherOp>::UnrollPattern;
LogicalResult matchAndRewrite(xegpu::LoadGatherOp op,
@@ -973,11 +1130,10 @@ struct UnrollConvertLayoutOp : public UnrollPattern<xegpu::ConvertLayoutOp> {
void mlir::xegpu::populateXeGPUUnrollPatterns(
RewritePatternSet &patterns, const xegpu::UnrollOptions &options) {
- patterns
- .add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
- UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp, UnrollLoadGatherOp,
- UnrollStoreScatterOp, UnrollPrefetchOp, UnrollLoadMatrixOp,
- UnrollStoreMatrixOp, UnrollLoadGatherOpWithOffset,
- UnrollStoreScatterOpWithOffsets, UnrollConvertLayoutOp>(
- patterns.getContext(), options);
+ patterns.add<UnrollCreateNdOp, UnrollUpdateNdOffsetOp, UnrollPrefetchNdOp,
+ UnrollLoadNdOp, UnrollStoreNdOp, UnrollDpasOp, UnrollDpasMxOp,
+ UnrollLoadGatherOp, UnrollStoreScatterOp, UnrollPrefetchOp,
+ UnrollLoadMatrixOp, UnrollStoreMatrixOp,
+ UnrollLoadGatherOpWithOffset, UnrollStoreScatterOpWithOffsets,
+ UnrollConvertLayoutOp>(patterns.getContext(), options);
}
diff --git a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
index d8a07d7c85a6c..bc27603da133e 100644
--- a/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
+++ b/mlir/test/Dialect/XeGPU/propagate-layout-subgroup.mlir
@@ -350,3 +350,63 @@ gpu.module @test {
gpu.return
}
}
+
+// -----
+#ld_bpack = #xegpu.layout<sg_layout = [8, 8], sg_data = [256, 16], inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>
+#b = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [4, 1]>
+
+gpu.module @test {
+ // CHECK-LABEL: b_load_manipulate_store
+ gpu.func @b_load_manipulate_store(%arg0: memref<2048x1024xui8>) {
+ // Construct 2D offsets for loading B (256x128) from memref<2048x1024xui8>
+ %b_row_step = vector.step : vector<256xindex>
+ %b_col_step = vector.step : vector<128xindex>
+ %b_row_step_2d = vector.shape_cast %b_row_step : vector<256xindex> to vector<256x1xindex>
+ %b_row_indices = vector.broadcast %b_row_step_2d : vector<256x1xindex> to vector<256x128xindex>
+ %c1024_vec = arith.constant dense<1024> : vector<256x128xindex>
+ %b_row_offsets = arith.muli %b_row_indices, %c1024_vec : vector<256x128xindex>
+ %b_col_offsets = vector.broadcast %b_col_step : vector<128xindex> to vector<256x128xindex>
+ %b_offsets = arith.addi %b_row_offsets, %b_col_offsets : vector<256x128xindex>
+ %b_mask = arith.constant dense<true> : vector<256x128xi1>
+
+ // Construct 2D offsets for storing B (512x128) back
+ %b_st_row_step = vector.step : vector<512xindex>
+ %b_st_col_step = vector.step : vector<128xindex>
+ %b_st_row_step_2d = vector.shape_cast %b_st_row_step : vector<512xindex> to vector<512x1xindex>
+ %b_st_row_indices = vector.broadcast %b_st_row_step_2d : vector<512x1xindex> to vector<512x128xindex>
+ %c2048_vec = arith.constant dense<2048> : vector<512x128xindex>
+ %b_st_row_offsets = arith.muli %b_st_row_indices, %c2048_vec : vector<512x128xindex>
+ %b_st_col_step_2d = vector.shape_cast %b_st_col_step : vector<128xindex> to vector<1x128xindex>
+ %b_st_col_offsets = vector.broadcast %b_st_col_step_2d : vector<1x128xindex> to vector<512x128xindex>
+ %b_st_offsets = arith.addi %b_st_row_offsets, %b_st_col_offsets : vector<512x128xindex>
+ %b_st_mask = arith.constant dense<true> : vector<512x128xi1>
+
+ // Extract pointer
+ %ptr_b_idx = memref.extract_aligned_pointer_as_index %arg0 : memref<2048x1024xui8> -> index
+ %ptr_b = arith.index_cast %ptr_b_idx : index to i64
+
+ // Load packed B (256x128 ui8 = 512x128 fp4 column-major packed)
+ // CHECK: xegpu.load %{{.*}}[%{{.*}}], %{{.*}} <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [256, 16], inst_data = [1, 16], lane_layout = [1, 16], lane_data = [1, 1]>}>
+ %b_packed = xegpu.load %ptr_b[%b_offsets], %b_mask <{layout = #ld_bpack}>
+ : i64, vector<256x128xindex>, vector<256x128xi1> -> vector<256x128xui8>
+
+ // Bitcast to fp4: 256x128 uint8 -> 256x256 fp4 (each uint8 holds 2 fp4 values)
+ %b_bitcast = vector.bitcast %b_packed : vector<256x128xui8> to vector<256x256xf4E2M1FN>
+
+ // De-interleave: extract even and odd columns
+ %b_even, %b_odd = vector.deinterleave %b_bitcast : vector<256x256xf4E2M1FN> -> vector<256x128xf4E2M1FN>
+
+ // Reconstruct 512x128 by interleaving even/odd rows
+ %b_even_t = vector.transpose %b_even, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_odd_t = vector.transpose %b_odd, [1, 0] : vector<256x128xf4E2M1FN> to vector<128x256xf4E2M1FN>
+ %b_interleaved = vector.interleave %b_even_t, %b_odd_t : vector<128x256xf4E2M1FN> -> vector<128x512xf4E2M1FN>
+ %b_loaded = vector.transpose %b_interleaved, [1, 0] : vector<128x512xf4E2M1FN> to vector<512x128xf4E2M1FN>
+
+ // Store B back to same memory location
+ // CHECK: xegpu.store %{{.*}}, %{{.*}}[%{{.*}}], %{{.*}} <{layout = #xegpu.layout<sg_layout = [8, 8], sg_data = [512, 16], inst_data = [64, 16], lane_layout = [1, 16], lane_data = [4, 1]>}>
+ xegpu.store %b_loaded, %ptr_b[%b_st_offsets], %b_st_mask <{layout = #b}>
+ : vector<512x128xf4E2M1FN>, i64, vector<512x128xindex>, vector<512x128xi1>
+
+ gpu.return
+ }
+}
>From 0fafcae57ab0745fef33c6b6d6332996e787f881 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 30 Apr 2026 06:05:37 +0000
Subject: [PATCH 03/10] fixing tests
---
.../XeGPU/Transforms/XeGPUBlocking.cpp | 75 ++++++-------------
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 6 +-
mlir/lib/Dialect/XeGPU/Utils/XeGPUUtils.cpp | 26 +++++++
mlir/test/Dialect/XeGPU/xegpu-blocking.mlir | 48 ++++++++++++
4 files changed, 102 insertions(+), 53 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index a1234b8ff8dc2..3093ede967d8f 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -122,13 +122,17 @@ XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
xegpu::DistributeLayoutAttr layout =
xegpu::getDistributeLayoutAttr(operandOrResult);
+ LDBG() << "getTileShape for value: " << value << ", layout: " << layout;
if (layout && layout.isForSubgroup()) {
if (!layout.getEffectiveInstDataAsInt().empty()) {
SmallVector<int64_t> instData = layout.getEffectiveInstDataAsInt();
+ LDBG() << " returning instData size: " << instData.size();
return instData;
}
- if (auto type = dyn_cast<ShapedType>(value.getType()))
+ if (auto type = dyn_cast<ShapedType>(value.getType())) {
+ LDBG() << " returning shape from type";
return llvm::to_vector(type.getShape());
+ }
}
LDBG() << "failed to getTileShape for: " << value;
return std::nullopt;
@@ -168,27 +172,12 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
std::optional<SmallVector<int64_t>> aTile = getTileShape(op->getOpOperand(0));
std::optional<SmallVector<int64_t>> bTile = getTileShape(op->getOpOperand(1));
- LLVM_DEBUG(llvm::dbgs() << " aTile: "
- << (aTile ? llvm::join(llvm::map_range(*aTile,
- [](int64_t v) { return std::to_string(v); }), "x")
- : "nullopt")
- << "\n");
- LLVM_DEBUG(llvm::dbgs() << " bTile: "
- << (bTile ? llvm::join(llvm::map_range(*bTile,
- [](int64_t v) { return std::to_string(v); }), "x")
- : "nullopt")
- << "\n");
-
if (!aTile || aTile->size() != 2 || !bTile || bTile->size() != 2)
return std::nullopt;
// semantic check for A and B
- if ((*aTile)[1] != (*bTile)[0]) {
- LLVM_DEBUG(llvm::dbgs() << " A/B semantic check failed: aTile[1]="
- << (*aTile)[1] << " != bTile[0]=" << (*bTile)[0]
- << "\n");
+ if ((*aTile)[1] != (*bTile)[0])
return std::nullopt;
- }
return std::make_pair(*aTile, *bTile);
};
@@ -203,12 +192,6 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
std::optional<SmallVector<int64_t>> cTile =
getTileShape(op->getOpOperand(cOperandIdx));
int64_t expectedCTile[2] = {aTile[0], bTile[1]};
- LLVM_DEBUG(llvm::dbgs() << " cTile: "
- << (cTile ? llvm::join(llvm::map_range(*cTile,
- [](int64_t v) { return std::to_string(v); }), "x")
- : "nullopt")
- << ", expected: " << expectedCTile[0] << "x"
- << expectedCTile[1] << "\n");
if (!cTile || !llvm::equal(*cTile, expectedCTile))
return false;
return true;
@@ -225,17 +208,6 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
std::optional<SmallVector<int64_t>> bScaleTile =
getTileShape(op->getOpOperand(scaleBOperandIdx));
- LLVM_DEBUG(llvm::dbgs() << " aScaleTile: "
- << (aScaleTile ? llvm::join(llvm::map_range(*aScaleTile,
- [](int64_t v) { return std::to_string(v); }), "x")
- : "nullopt")
- << "\n");
- LLVM_DEBUG(llvm::dbgs() << " bScaleTile: "
- << (bScaleTile ? llvm::join(llvm::map_range(*bScaleTile,
- [](int64_t v) { return std::to_string(v); }), "x")
- : "nullopt")
- << "\n");
-
if (!aScaleTile || aScaleTile->size() != 2 ||
!bScaleTile || bScaleTile->size() != 2)
return std::nullopt;
@@ -246,20 +218,14 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
assert((*bScaleTile)[1] == bTile[1] &&
"bScaleTile[1] must equal bTile[1]");
- if ((*aScaleTile)[1] != (*bScaleTile)[0]) {
- LLVM_DEBUG(llvm::dbgs() << " scale A/B semantic check failed: aScaleTile[1]="
- << (*aScaleTile)[1] << " != bScaleTile[0]="
- << (*bScaleTile)[0] << "\n");
+ if ((*aScaleTile)[1] != (*bScaleTile)[0])
return std::nullopt;
- }
// Return the K scale factor
return (*aScaleTile)[1];
};
if (isa<xegpu::DpasOp>(op)) {
- LLVM_DEBUG(llvm::dbgs() << "getTileShape for DpasOp: " << *op << "\n");
-
auto abTiles = validateABTiles(op);
if (!abTiles)
return std::nullopt;
@@ -270,14 +236,10 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
if (!validateCTile(op, 2, aTile, bTile))
return std::nullopt;
- LLVM_DEBUG(llvm::dbgs() << " result: [" << aTile[0] << ", "
- << aTile[1] << ", " << bTile[1] << "]\n");
return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1]});
}
if (auto dpasMxOp = dyn_cast<xegpu::DpasMxOp>(op)) {
- LLVM_DEBUG(llvm::dbgs() << "getTileShape for DpasMxOp: " << *op << "\n");
-
auto abTiles = validateABTiles(op);
if (!abTiles)
return std::nullopt;
@@ -318,9 +280,6 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
kScaleFactor = *scaleFactor;
}
- LLVM_DEBUG(llvm::dbgs() << " result: [" << aTile[0] << ", "
- << aTile[1] << ", " << bTile[1] << ", "
- << kScaleFactor << "]\n");
return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1], kScaleFactor});
}
@@ -371,7 +330,14 @@ bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
bool hasUnrollableOperands =
llvm::any_of(op->getOpOperands(), [&](OpOperand &opr) {
std::optional<SmallVector<int64_t>> tileShape = getTileShape(opr);
- return tileShape.has_value() && isUnrollable(opr.get(), *tileShape);
+ bool result =
+ tileShape.has_value() && isUnrollable(opr.get(), *tileShape);
+ if (isa<xegpu::DpasMxOp>(op)) {
+ LDBG() << " operand " << opr.getOperandNumber()
+ << ": tileShape.has_value=" << tileShape.has_value()
+ << ", isUnrollable=" << result;
+ }
+ return result;
});
bool hasUnrollableResults =
llvm::any_of(op->getOpResults(), [&](OpResult result) {
@@ -386,8 +352,15 @@ bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
isConvertLayoutWithInstData = true;
}
}
- return hasUnrollableOperands || hasUnrollableResults ||
- isConvertLayoutWithInstData;
+ bool shouldUnroll = hasUnrollableOperands || hasUnrollableResults ||
+ isConvertLayoutWithInstData;
+ if (isa<xegpu::DpasMxOp>(op)) {
+ LDBG() << "needsUnroll for DpasMxOp: hasUnrollableOperands="
+ << hasUnrollableOperands
+ << ", hasUnrollableResults=" << hasUnrollableResults
+ << ", shouldUnroll=" << shouldUnroll;
+ }
+ return shouldUnroll;
}
void XeGPUBlockingPass::runOnOperation() {
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index df970e1ad8e83..c3d19285f4502 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -425,10 +425,10 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
}))
return failure();
- // A vector of 3 elements should be returned, representing M, K, N
+ // A vector of 4 elements should be returned, representing M, K, N, S
// respectively.
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
- if (!targetShape || targetShape->size() != 3)
+ if (!targetShape || targetShape->size() != 4)
return failure();
auto M = (*targetShape)[0];
auto K = (*targetShape)[1];
@@ -553,6 +553,8 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
rewriter, loc, vecTy, operands,
xegpu::dropInstDataOnAttrs(op->getAttrs()));
LLVM_DEBUG(llvm::dbgs() << " created: " << newDpasMxOp << "\n");
+ // Update tmpC to accumulate across K iterations
+ tmpC = newDpasMxOp.getResult();
}
newOps.push_back(newDpasMxOp);
}
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/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
+ }
+}
+
>From 6917f24d28db53a9848929b7c31d749de6d219ba Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 30 Apr 2026 18:59:30 +0000
Subject: [PATCH 04/10] [XeGPU] Fix DpasMxOp unrolling to support K dimension
accumulation
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit
Fixed two critical bugs in DpasMxOp unrolling:
1. Incorrect targetShape size check in XeGPUUnroll.cpp
- Changed from checking for 3 elements to 4 elements (M, K, N, S)
- DpasMxOp requires scale factor dimension in addition to M, K, N
2. Missing K accumulation chaining in nested unroll loop
- K-dimension tiles must chain together, feeding output as accumulator
- Added tmpC update after each dpas_mx creation to maintain chain
- Without this, all but the last K iteration became dead code
3. Added DpasMxOp support in getDistributeLayoutAttr (XeGPUUtils.cpp)
- Enables retrieval of layout_a, layout_b, layout_cd, layout_a_scale,
layout_b_scale attributes based on operand index
- Uses AttrSizedOperandSegments to determine correct layout per operand
This enables proper unrolling of matrix multiplications like 16x64 @ 64x32
into 2×2×2 = 8 smaller dpas_mx operations with correct K accumulation.
Test: mlir/test/Dialect/XeGPU/xegpu-blocking.mlir now passes
Co-Authored-By: Claude Sonnet 4.5 <noreply at anthropic.com>
---
.../XeGPU/Transforms/XeGPUBlocking.cpp | 27 +++----------------
1 file changed, 4 insertions(+), 23 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 3093ede967d8f..c453c06d15962 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -122,19 +122,14 @@ XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
xegpu::DistributeLayoutAttr layout =
xegpu::getDistributeLayoutAttr(operandOrResult);
- LDBG() << "getTileShape for value: " << value << ", layout: " << layout;
if (layout && layout.isForSubgroup()) {
if (!layout.getEffectiveInstDataAsInt().empty()) {
SmallVector<int64_t> instData = layout.getEffectiveInstDataAsInt();
- LDBG() << " returning instData size: " << instData.size();
return instData;
}
- if (auto type = dyn_cast<ShapedType>(value.getType())) {
- LDBG() << " returning shape from type";
+ if (auto type = dyn_cast<ShapedType>(value.getType()))
return llvm::to_vector(type.getShape());
- }
}
- LDBG() << "failed to getTileShape for: " << value;
return std::nullopt;
}
@@ -330,14 +325,7 @@ bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
bool hasUnrollableOperands =
llvm::any_of(op->getOpOperands(), [&](OpOperand &opr) {
std::optional<SmallVector<int64_t>> tileShape = getTileShape(opr);
- bool result =
- tileShape.has_value() && isUnrollable(opr.get(), *tileShape);
- if (isa<xegpu::DpasMxOp>(op)) {
- LDBG() << " operand " << opr.getOperandNumber()
- << ": tileShape.has_value=" << tileShape.has_value()
- << ", isUnrollable=" << result;
- }
- return result;
+ return tileShape.has_value() && isUnrollable(opr.get(), *tileShape);
});
bool hasUnrollableResults =
llvm::any_of(op->getOpResults(), [&](OpResult result) {
@@ -352,15 +340,8 @@ bool XeGPUBlockingPass::needsUnroll(Operation *op) const {
isConvertLayoutWithInstData = true;
}
}
- bool shouldUnroll = hasUnrollableOperands || hasUnrollableResults ||
- isConvertLayoutWithInstData;
- if (isa<xegpu::DpasMxOp>(op)) {
- LDBG() << "needsUnroll for DpasMxOp: hasUnrollableOperands="
- << hasUnrollableOperands
- << ", hasUnrollableResults=" << hasUnrollableResults
- << ", shouldUnroll=" << shouldUnroll;
- }
- return shouldUnroll;
+ return hasUnrollableOperands || hasUnrollableResults ||
+ isConvertLayoutWithInstData;
}
void XeGPUBlockingPass::runOnOperation() {
>From 9a72a398927ac353236b8af2605627912a8391f1 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 30 Apr 2026 21:13:26 +0000
Subject: [PATCH 05/10] remove debug print
---
.../XeGPU/Transforms/XeGPUBlocking.cpp | 1 +
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 54 -------------------
2 files changed, 1 insertion(+), 54 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index c453c06d15962..69e8fe7eca153 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -130,6 +130,7 @@ XeGPUBlockingPass::getTileShape(const T &operandOrResult) const {
if (auto type = dyn_cast<ShapedType>(value.getType()))
return llvm::to_vector(type.getShape());
}
+ LDBG() << "failed to getTileShape for: " << value;
return std::nullopt;
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index c3d19285f4502..d6f73422953c2 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -415,18 +415,12 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
LogicalResult matchAndRewrite(xegpu::DpasMxOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
-
- LLVM_DEBUG(llvm::dbgs() << "UnrollDpasMxOp: original op: " << op << "\n");
-
- // 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 4 elements should be returned, representing M, K, N, S
- // respectively.
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape || targetShape->size() != 4)
return failure();
@@ -435,23 +429,12 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
auto N = (*targetShape)[2];
auto S = (*targetShape)[3];
- LLVM_DEBUG(llvm::dbgs() << " targetShape: M=" << M << ", K=" << K
- << ", N=" << N << ", S=" << S << "\n");
-
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};
- LLVM_DEBUG(llvm::dbgs() << " aBlockSize: [" << M << ", " << K << "]\n");
- LLVM_DEBUG(llvm::dbgs() << " bBlockSize: [" << K << ", " << N << "]\n");
- LLVM_DEBUG(llvm::dbgs() << " cBlockSize: [" << M << ", " << N << "]\n");
- LLVM_DEBUG(llvm::dbgs()
- << " aScaleBlockSize: [" << M << ", " << K / 32 << "]\n");
- LLVM_DEBUG(llvm::dbgs()
- << " bScaleBlockSize: [" << K / 32 << ", " << N << "]\n");
-
auto packWrapper = [&](TypedValue<VectorType> val,
ArrayRef<int64_t> blockSize) {
VectorType type = val.getType();
@@ -488,22 +471,6 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
if (bscale)
bScaleVals = packWrapper(bscale, bScaleBlockSize);
- LLVM_DEBUG(llvm::dbgs() << " aVals size: " << aVals.size() << "\n");
- LLVM_DEBUG(llvm::dbgs() << " bVals size: " << bVals.size() << "\n");
- LLVM_DEBUG(llvm::dbgs() << " cVals size: " << cVals.size() << "\n");
- LLVM_DEBUG(llvm::dbgs()
- << " aScaleVals size: " << aScaleVals.size() << "\n");
- LLVM_DEBUG(llvm::dbgs()
- << " bScaleVals size: " << bScaleVals.size() << "\n");
-
- // 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; }) ||
- // llvm::all_of(ranges, [](auto &v) { return v.size() == 1; }))
- // return failure();
-
VectorType resultTy = op.getResult().getType();
auto vecTy = VectorType::get(cBlockSize, resultTy.getElementType());
@@ -511,9 +478,6 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
int64_t kIters = aShape[1] / K;
int64_t nIters = bShape[1] / N;
- LLVM_DEBUG(llvm::dbgs() << " mIters=" << mIters << ", kIters=" << kIters
- << ", nIters=" << nIters << "\n");
-
SmallVector<Value> newOps;
xegpu::DpasMxOp newDpasMxOp;
for (int64_t i = 0; i < mIters; ++i) {
@@ -525,44 +489,26 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
for (int64_t k = 0; k < kIters; ++k) {
Value aVec = aVals[i * kIters + k];
Value bVec = bVals[k * nIters + j];
-
- LLVM_DEBUG(llvm::dbgs() << " [i=" << i << ", j=" << j << ", k=" << k
- << "] aVec: " << aVec.getType()
- << ", bVec: " << bVec.getType() << "\n");
-
SmallVector<Value> operands({aVec, bVec});
if (tmpC) {
operands.push_back(tmpC);
- LLVM_DEBUG(llvm::dbgs() << " tmpC: " << tmpC.getType() << "\n");
}
if (ascale) {
Value aScaleVec = aScaleVals[i * kIters + k];
operands.push_back(aScaleVec);
- LLVM_DEBUG(llvm::dbgs()
- << " aScaleVec: " << aScaleVec.getType() << "\n");
}
if (bscale) {
Value bScaleVec = bScaleVals[k * nIters + j];
operands.push_back(bScaleVec);
- LLVM_DEBUG(llvm::dbgs()
- << " bScaleVec: " << bScaleVec.getType() << "\n");
}
- LLVM_DEBUG(llvm::dbgs() << " total operands: " << operands.size()
- << ", resTy: " << vecTy << "\n");
newDpasMxOp = xegpu::DpasMxOp::create(
rewriter, loc, vecTy, operands,
xegpu::dropInstDataOnAttrs(op->getAttrs()));
- LLVM_DEBUG(llvm::dbgs() << " created: " << newDpasMxOp << "\n");
- // Update tmpC to accumulate across K iterations
tmpC = newDpasMxOp.getResult();
}
newOps.push_back(newDpasMxOp);
}
}
-
- LLVM_DEBUG(llvm::dbgs()
- << " total new DpasMxOps: " << newOps.size() << "\n");
-
Value castOp = unpack(newOps, resultTy, cBlockSize, loc, rewriter);
rewriter.replaceOp(op, castOp);
return success();
>From 27e8c513f1fdd303579094f3346080f0ba2d5ce3 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Thu, 30 Apr 2026 21:14:35 +0000
Subject: [PATCH 06/10] address format issue
---
.../XeGPU/Transforms/XeGPUBlocking.cpp | 30 +++++++++----------
1 file changed, 14 insertions(+), 16 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 69e8fe7eca153..6b3437a1ef832 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -165,8 +165,10 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
// 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 = getTileShape(op->getOpOperand(1));
+ std::optional<SmallVector<int64_t>> aTile =
+ getTileShape(op->getOpOperand(0));
+ std::optional<SmallVector<int64_t>> bTile =
+ getTileShape(op->getOpOperand(1));
if (!aTile || aTile->size() != 2 || !bTile || bTile->size() != 2)
return std::nullopt;
@@ -194,25 +196,22 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
};
// Helper lambda to validate scale A/B tiles for DpasMxOp
- auto validateABScaleTiles = [&](Operation *op, unsigned scaleAOperandIdx,
- unsigned scaleBOperandIdx,
- const SmallVector<int64_t> &aTile,
- const SmallVector<int64_t> &bTile)
- -> std::optional<int64_t> {
+ auto validateABScaleTiles =
+ [&](Operation *op, unsigned scaleAOperandIdx, unsigned scaleBOperandIdx,
+ const SmallVector<int64_t> &aTile,
+ const SmallVector<int64_t> &bTile) -> std::optional<int64_t> {
std::optional<SmallVector<int64_t>> aScaleTile =
getTileShape(op->getOpOperand(scaleAOperandIdx));
std::optional<SmallVector<int64_t>> bScaleTile =
getTileShape(op->getOpOperand(scaleBOperandIdx));
- if (!aScaleTile || aScaleTile->size() != 2 ||
- !bScaleTile || bScaleTile->size() != 2)
+ if (!aScaleTile || aScaleTile->size() != 2 || !bScaleTile ||
+ bScaleTile->size() != 2)
return std::nullopt;
// Validate scale tile dimensions
- assert((*aScaleTile)[0] == aTile[0] &&
- "aScaleTile[0] must equal aTile[0]");
- assert((*bScaleTile)[1] == bTile[1] &&
- "bScaleTile[1] must equal bTile[1]");
+ assert((*aScaleTile)[0] == aTile[0] && "aScaleTile[0] must equal aTile[0]");
+ assert((*bScaleTile)[1] == bTile[1] && "bScaleTile[1] must equal bTile[1]");
if ((*aScaleTile)[1] != (*bScaleTile)[0])
return std::nullopt;
@@ -269,7 +268,7 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
unsigned scaleBOperandIdx = scaleAOperandIdx + scaleASize;
auto scaleFactor = validateABScaleTiles(op, scaleAOperandIdx,
- scaleBOperandIdx, aTile, bTile);
+ scaleBOperandIdx, aTile, bTile);
if (!scaleFactor)
return std::nullopt;
@@ -287,8 +286,7 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
if (isa<vector::TransposeOp, vector::BroadcastOp, vector::StepOp,
vector::ShapeCastOp, vector::ConstantMaskOp, vector::CreateMaskOp,
- vector::BitCastOp, vector::InterleaveOp, vector::DeinterleaveOp>(
- op))
+ vector::BitCastOp, vector::InterleaveOp, vector::DeinterleaveOp>(op))
return getTileShape(op->getOpResult(0));
return std::nullopt;
>From a87dc9477bc559bc5603b5e7e34fc706aa91aae2 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 5 May 2026 22:16:31 +0000
Subject: [PATCH 07/10] [XeGPU] Address PR feedback and add uArch bit width
validation
This commit addresses review feedback from PR #195179:
1. Refactor DpasOp and DpasMxOp verifiers to share common validation
logic through helper functions (verifyLayoutDistributable,
verifyDpasDimensions, verifyDpasAccumulator)
2. Replace manual AttrSizedOperandSegments manipulation with
op-specific accessor methods in XeGPUBlocking.cpp and
XeGPUUnroll.cpp
3. Add comprehensive negative tests in invalid.mlir covering all
DpasMxOp error cases (dimension mismatches, layout issues, etc.)
4. Add uArch packed bit width validation in sg-to-wi distribution
pass to ensure workitem types match microarchitecture requirements
Co-Authored-By: Claude Sonnet 4.5 <noreply at anthropic.com>
---
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 198 +++++++++++++-----
.../XeGPU/Transforms/XeGPUBlocking.cpp | 34 +--
.../XeGPUSgToWiDistributeExperimental.cpp | 28 ++-
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 10 +-
mlir/test/Dialect/XeGPU/invalid.mlir | 101 ++++++++-
5 files changed, 287 insertions(+), 84 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index 4fe15c625ea49..c8bbf7a4fbab9 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -693,61 +693,90 @@ void StoreScatterOp::build(
l3_hint, layout);
}
+//===----------------------------------------------------------------------===//
+// DPAS Common Verification Helpers
+//===----------------------------------------------------------------------===//
+
+// 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();
+}
+
+// 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) {
+
+ auto aRank = aShape.size();
+ auto bRank = bShape.size();
+ auto resRank = resShape.size();
+ if (aRank == 1 && bRank == 1 && resRank == 1)
+ return success();
+
+ // 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() {
- 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");
+ // 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();
- if (auto aLayout = getLayoutA())
- if (!aLayout->isDistributable(
- SmallVector<int64_t>(lhsShape.begin(), lhsShape.end())))
- return emitOpError("Value shape is not distributable with the layout");
+ // Verify accumulator if present
+ if (getAcc() && failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
+ return failure();
- 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.");
- 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.");
-
- return success();
+ return verifyDpasDimensions(*this, lhsShape, rhsShape, resShape);
}
//===----------------------------------------------------------------------===//
@@ -853,8 +882,83 @@ 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 6b3437a1ef832..26069f938ce3e 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -209,9 +209,9 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
bScaleTile->size() != 2)
return std::nullopt;
- // Validate scale tile dimensions
- assert((*aScaleTile)[0] == aTile[0] && "aScaleTile[0] must equal aTile[0]");
- assert((*bScaleTile)[1] == bTile[1] && "bScaleTile[1] must equal bTile[1]");
+ // Validate scale tile dimensions match expected values
+ if ((*aScaleTile)[0] != aTile[0] || (*bScaleTile)[1] != bTile[1])
+ return std::nullopt;
if ((*aScaleTile)[1] != (*bScaleTile)[0])
return std::nullopt;
@@ -241,31 +241,19 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
auto [aTile, bTile] = *abTiles;
- // Get operand indices using AttrSizedOperandSegments
- auto segmentSizesAttr = dpasMxOp->getAttrOfType<DenseI32ArrayAttr>(
- dpasMxOp.getOperandSegmentSizesAttrName());
- if (!segmentSizesAttr)
- return std::nullopt;
-
- auto segmentSizes = segmentSizesAttr.asArrayRef();
- unsigned aSize = segmentSizes[0];
- unsigned bSize = segmentSizes[1];
- unsigned accSize = segmentSizes[2];
- unsigned scaleASize = segmentSizes[3];
- unsigned scaleBSize = segmentSizes[4];
-
- // Validate C tile if present
- if (accSize > 0) {
- unsigned accOperandIdx = aSize + bSize;
+ // 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
+ // Validate scale tiles if present using op-specific accessors
int64_t kScaleFactor = 1;
- if (scaleASize > 0 && scaleBSize > 0) {
- unsigned scaleAOperandIdx = aSize + bSize + accSize;
- unsigned scaleBOperandIdx = scaleAOperandIdx + scaleASize;
+ if (dpasMxOp.getScaleA() && dpasMxOp.getScaleB()) {
+ // Calculate operand indices based on which operands are present
+ unsigned scaleAOperandIdx = 2 + (dpasMxOp.getAcc() ? 1 : 0);
+ unsigned scaleBOperandIdx = scaleAOperandIdx + 1;
auto scaleFactor = validateABScaleTiles(op, scaleAOperandIdx,
scaleBOperandIdx, aTile, bTile);
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
index c153db431c035..4e7027a0610ed 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
@@ -234,7 +234,6 @@ 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());
@@ -259,6 +258,33 @@ 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 d6f73422953c2..80e31b05d33d5 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -415,10 +415,12 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
LogicalResult matchAndRewrite(xegpu::DpasMxOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
- if (llvm::any_of(op->getOperandTypes(), [&](Type type) {
- auto vecTy = dyn_cast<VectorType>(type);
- return !vecTy || vecTy.getRank() != 2;
- }))
+ // Scale operands can be scalars, which we don't unroll
+ // Check that A and B (required operands) are 2D vectors
+ if (op.getAType().getRank() != 2 || op.getBType().getRank() != 2)
+ return failure();
+ // If acc is present, it must be a 2D vector
+ if (op.getAcc() && op.getAccType().getRank() != 2)
return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
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
+}
>From 5334bfeba993ba406dc1e06be58216d622eaa4b0 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 5 May 2026 22:30:33 +0000
Subject: [PATCH 08/10] [XeGPU] Refactor DpasMxOp blocking to support
independent scale operands
Previously, the scale validation
logic required both scale_a and
scale_b
to be present. This refactoring
splits the validation into separate
functions (validateScaleATile and
validateScaleBTile) that can validate
each scale operand independently.
---
.../XeGPU/Transforms/XeGPUBlocking.cpp | 69 +++++++++++++------
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 7 --
2 files changed, 49 insertions(+), 27 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
index 26069f938ce3e..7db887915b275 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUBlocking.cpp
@@ -195,29 +195,42 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
return true;
};
- // Helper lambda to validate scale A/B tiles for DpasMxOp
- auto validateABScaleTiles =
- [&](Operation *op, unsigned scaleAOperandIdx, unsigned scaleBOperandIdx,
- const SmallVector<int64_t> &aTile,
- const SmallVector<int64_t> &bTile) -> std::optional<int64_t> {
+ // 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));
- std::optional<SmallVector<int64_t>> bScaleTile =
- getTileShape(op->getOpOperand(scaleBOperandIdx));
- if (!aScaleTile || aScaleTile->size() != 2 || !bScaleTile ||
- bScaleTile->size() != 2)
+ if (!aScaleTile || aScaleTile->size() != 2)
return std::nullopt;
- // Validate scale tile dimensions match expected values
- if ((*aScaleTile)[0] != aTile[0] || (*bScaleTile)[1] != bTile[1])
+ // Validate scale_a tile: [M_tile, K_scale]
+ // M dimension must match A's M dimension
+ if ((*aScaleTile)[0] != aTile[0])
return std::nullopt;
- if ((*aScaleTile)[1] != (*bScaleTile)[0])
+ // 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 (*aScaleTile)[1];
+ return (*bScaleTile)[0];
};
if (isa<xegpu::DpasOp>(op)) {
@@ -250,17 +263,33 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
// Validate scale tiles if present using op-specific accessors
int64_t kScaleFactor = 1;
- if (dpasMxOp.getScaleA() && dpasMxOp.getScaleB()) {
- // Calculate operand indices based on which operands are present
+ std::optional<int64_t> scaleAFactor;
+ std::optional<int64_t> scaleBFactor;
+
+ if (dpasMxOp.getScaleA()) {
unsigned scaleAOperandIdx = 2 + (dpasMxOp.getAcc() ? 1 : 0);
- unsigned scaleBOperandIdx = scaleAOperandIdx + 1;
+ scaleAFactor = validateScaleATile(op, scaleAOperandIdx, aTile);
+ if (!scaleAFactor)
+ return std::nullopt;
+ }
- auto scaleFactor = validateABScaleTiles(op, scaleAOperandIdx,
- scaleBOperandIdx, aTile, bTile);
- if (!scaleFactor)
+ if (dpasMxOp.getScaleB()) {
+ unsigned scaleBOperandIdx =
+ 2 + (dpasMxOp.getAcc() ? 1 : 0) + (dpasMxOp.getScaleA() ? 1 : 0);
+ scaleBFactor = validateScaleBTile(op, scaleBOperandIdx, bTile);
+ if (!scaleBFactor)
return std::nullopt;
+ }
- kScaleFactor = *scaleFactor;
+ // 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], kScaleFactor});
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index 80e31b05d33d5..e8dd4f16b2e22 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
@@ -415,13 +415,6 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
LogicalResult matchAndRewrite(xegpu::DpasMxOp op,
PatternRewriter &rewriter) const override {
Location loc = op.getLoc();
- // Scale operands can be scalars, which we don't unroll
- // Check that A and B (required operands) are 2D vectors
- if (op.getAType().getRank() != 2 || op.getBType().getRank() != 2)
- return failure();
- // If acc is present, it must be a 2D vector
- if (op.getAcc() && op.getAccType().getRank() != 2)
- return failure();
std::optional<SmallVector<int64_t>> targetShape = getTargetShape(op);
if (!targetShape || targetShape->size() != 4)
>From 394ea9e489d2322cc0e2d068311cf7427b85dae6 Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 5 May 2026 22:37:45 +0000
Subject: [PATCH 09/10] git clang format
---
mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp | 53 ++++++++++++-------
.../XeGPUSgToWiDistributeExperimental.cpp | 22 ++++----
2 files changed, 46 insertions(+), 29 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
index c8bbf7a4fbab9..a1933f3e05f6b 100644
--- a/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
+++ b/mlir/lib/Dialect/XeGPU/IR/XeGPUOps.cpp
@@ -698,20 +698,22 @@ void StoreScatterOp::build(
//===----------------------------------------------------------------------===//
// 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())))
+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();
}
// 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) {
+static LogicalResult verifyDpasDimensions(Operation *op,
+ ArrayRef<int64_t> aShape,
+ ArrayRef<int64_t> bShape,
+ ArrayRef<int64_t> resShape) {
auto aRank = aShape.size();
auto bRank = bShape.size();
@@ -750,7 +752,7 @@ static LogicalResult verifyDpasDimensions(Operation *op, ArrayRef<int64_t> aShap
// Helper to verify accumulator matches result type
static LogicalResult verifyDpasAccumulator(Operation *op, Type accType,
- Type resultType) {
+ Type resultType) {
if (accType != resultType)
return op->emitOpError("Accumulator type must match result type.");
return success();
@@ -765,7 +767,8 @@ LogicalResult DpasOp::verify() {
auto resShape = getResultType().getShape();
// Verify layout distributability
- if (failed(verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
+ if (failed(
+ verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
return failure();
if (failed(verifyLayoutDistributable(*this, getLayoutA(), lhsShape, "A")))
return failure();
@@ -773,7 +776,8 @@ LogicalResult DpasOp::verify() {
return failure();
// Verify accumulator if present
- if (getAcc() && failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
+ if (getAcc() &&
+ failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
return failure();
return verifyDpasDimensions(*this, lhsShape, rhsShape, resShape);
@@ -887,7 +891,8 @@ LogicalResult DpasMxOp::verify() {
auto resShape = getResultType().getShape();
// Verify layout distributability for A, B, and result
- if (failed(verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
+ if (failed(
+ verifyLayoutDistributable(*this, getLayoutCd(), resShape, "Result")))
return failure();
if (failed(verifyLayoutDistributable(*this, getLayoutA(), aShape, "A")))
return failure();
@@ -895,7 +900,8 @@ LogicalResult DpasMxOp::verify() {
return failure();
// Verify accumulator if present
- if (getAcc() && failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
+ if (getAcc() &&
+ failed(verifyDpasAccumulator(*this, getAcc().getType(), getResultType())))
return failure();
// Verify M, N, K dimensions
@@ -913,13 +919,15 @@ LogicalResult DpasMxOp::verify() {
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")))
+ 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] << "].";
+ << scaleAShape[0] << "] must match A M dimension [" << aShape[0]
+ << "].";
}
}
@@ -934,17 +942,20 @@ LogicalResult DpasMxOp::verify() {
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")))
+ 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] << "].";
+ << scaleBShape[1] << "] must match B N dimension [" << bShape[1]
+ << "].";
}
}
- // Validate scale K dimension compatibility if both scales are present and vectors
+ // 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());
@@ -953,10 +964,12 @@ LogicalResult DpasMxOp::verify() {
auto scaleAShape = scaleAVecType.getShape();
auto scaleBShape = scaleBVecType.getShape();
- // Validate scale K dimension compatibility: scale_a[1] must match scale_b[0]
+ // 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] << ".";
+ << scaleAShape[1] << " but scale_b has K=" << scaleBShape[0]
+ << ".";
}
}
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
index 4e7027a0610ed..b70c8bd245853 100644
--- a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
+++ b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUSgToWiDistributeExperimental.cpp
@@ -240,7 +240,6 @@ struct SgToWiDpas : public OpConversionPattern<xegpu::DpasOp> {
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 =
@@ -263,25 +262,30 @@ struct SgToWiDpas : public OpConversionPattern<xegpu::DpasOp> {
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));
+ 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 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");
+ 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");
+ op,
+ "B operand packed bit width must be a multiple of uArch packed "
+ "format requirement");
}
}
>From 0c94de1bfb6f998eef27512d4a9a5dd30c55c78d Mon Sep 17 00:00:00 2001
From: Jianhui Li <jian.hui.li at intel.com>
Date: Tue, 5 May 2026 23:00:59 +0000
Subject: [PATCH 10/10] address feedback
---
.../Dialect/XeGPU/Transforms/XeGPUUnroll.cpp | 116 +++++++-----------
1 file changed, 43 insertions(+), 73 deletions(-)
diff --git a/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp b/mlir/lib/Dialect/XeGPU/Transforms/XeGPUUnroll.cpp
index e8dd4f16b2e22..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];
@@ -430,45 +420,29 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
int64_t aScaleBlockSize[2] = {M, S};
int64_t bScaleBlockSize[2] = {S, 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.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());
- auto aShape = a.getType().getShape();
- auto bShape = b.getType().getShape();
-
- SmallVector<Value> aVals, bVals, cVals, aScaleVals, bScaleVals;
- 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);
+ SmallVector<Value> aScaleVals;
if (ascale)
- aScaleVals = packWrapper(ascale, aScaleBlockSize);
+ aScaleVals = packOperandForDpas(ascale, aScaleBlockSize, loc, rewriter);
+ SmallVector<Value> bScaleVals;
if (bscale)
- bScaleVals = packWrapper(bscale, bScaleBlockSize);
+ 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;
@@ -479,23 +453,19 @@ struct UnrollDpasMxOp : public UnrollPattern<xegpu::DpasMxOp> {
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];
Value bVec = bVals[k * nIters + j];
SmallVector<Value> operands({aVec, bVec});
- if (tmpC) {
+ if (tmpC)
operands.push_back(tmpC);
- }
- if (ascale) {
- Value aScaleVec = aScaleVals[i * kIters + k];
- operands.push_back(aScaleVec);
- }
- if (bscale) {
- Value bScaleVec = bScaleVals[k * nIters + j];
- operands.push_back(bScaleVec);
- }
+ 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()));
More information about the Mlir-commits
mailing list