[Mlir-commits] [mlir] [MLIR][XeGPU] Add unrolling/blocking support for 3D+ batched operations (PR #201725)

[Charitha Saumya]

================
@@ -170,11 +191,24 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
     std::optional<SmallVector<int64_t>> bTile =
         getTileShape(op->getOpOperand(1));
 
-    if (!aTile || aTile->size() != 2 || !bTile || bTile->size() != 2)
+    if (!aTile || aTile->size() < 2 || !bTile || bTile->size() < 2)
+      return std::nullopt;
+
+    // Both must have the same number of batch dimensions.
+    int64_t aBatchRank = aTile->size() - 2;
+    int64_t bBatchRank = bTile->size() - 2;
+    if (aBatchRank != bBatchRank)
       return std::nullopt;
 
-    // semantic check for A and B
-    if ((*aTile)[1] != (*bTile)[0])
+    // Batch dimensions must match.
+    for (int64_t i = 0; i < aBatchRank; ++i) {
+      if ((*aTile)[i] != (*bTile)[i])
+        return std::nullopt;
+    }
+
+    // semantic check for A and B: K dimension must match
----------------
charithaintc wrote:

nit: comments must start with UC letters. 

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