[Mlir-commits] [mlir] [MLIR][XeGPU] Unroll Dpasmx Op (PR #195179)

[Charitha Saumya]

================
@@ -175,16 +177,105 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
     if ((*aTile)[1] != (*bTile)[0])
       return std::nullopt;
 
+    return std::make_pair(*aTile, *bTile);
+  };
+
+  // Helper lambda to validate C tile
+  auto validateCTile = [&](Operation *op, unsigned cOperandIdx,
+                           const SmallVector<int64_t> &aTile,
+                           const SmallVector<int64_t> &bTile) -> bool {
+    if (op->getNumOperands() <= cOperandIdx)
+      return true;
+
+    std::optional<SmallVector<int64_t>> cTile =
+        getTileShape(op->getOpOperand(cOperandIdx));
+    int64_t expectedCTile[2] = {aTile[0], bTile[1]};
+    if (!cTile || !llvm::equal(*cTile, expectedCTile))
+      return false;
+    return true;
+  };
+
+  // Helper lambda to validate scale A/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));
+
+    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])
+      return std::nullopt;
+
+    // Return the K scale factor
+    return (*aScaleTile)[1];
+  };
+
+  if (isa<xegpu::DpasOp>(op)) {
+    auto abTiles = validateABTiles(op);
+    if (!abTiles)
+      return std::nullopt;
+
+    auto [aTile, bTile] = *abTiles;
+
     // semantic check for C
-    if (op->getNumOperands() == 3) {
-      std::optional<SmallVector<int64_t>> cTile =
-          getTileShape(op->getOpOperand(2));
-      int64_t expectedCTile[2] = {(*aTile)[0], (*bTile)[1]};
-      if (!cTile || !llvm::equal(*cTile, expectedCTile))
+    if (!validateCTile(op, 2, aTile, bTile))
+      return std::nullopt;
+
+    return SmallVector<int64_t>({aTile[0], aTile[1], bTile[1]});
+  }
+
+  if (auto dpasMxOp = dyn_cast<xegpu::DpasMxOp>(op)) {
+    auto abTiles = validateABTiles(op);
+    if (!abTiles)
+      return std::nullopt;
+
+    auto [aTile, bTile] = *abTiles;
+
+    // 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];
----------------
charithaintc wrote:

I would encourage using the op specifc methods for these by casting the op to `DpasMxOp`. In the future if you make any change to operands these segmentSizesAttr based access will fail, but tablegen accessors always work. 

https://github.com/llvm/llvm-project/pull/195179