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

[Artem Kroviakov]

================
@@ -202,16 +243,18 @@ XeGPUBlockingPass::getTileShape(Operation *op) const {
     std::optional<SmallVector<int64_t>> aScaleTile =
         getTileShape(op->getOpOperand(scaleAOperandIdx));
 
-    if (!aScaleTile || aScaleTile->size() != 2)
+    if (!aScaleTile || aScaleTile->size() < 2)
       return std::nullopt;
 
-    // Validate scale_a tile: [M_tile, K_scale]
-    // M dimension must match A's M dimension
-    if ((*aScaleTile)[0] != aTile[0])
+    // Validate scale_a tile: [batch..., M_tile, K_scale]
+    // M dimension (second-to-last) must match A's M dimension
+    int64_t scaleRank = aScaleTile->size();
+    int64_t aBatchRank = aTile.size() - 2;
+    if ((*aScaleTile)[scaleRank - 2] != aTile[aBatchRank])
       return std::nullopt;
 
-    // Return the K scale factor
-    return (*aScaleTile)[1];
+    // Return the K scale factor (last dim)
+    return (*aScaleTile)[scaleRank - 1];
----------------
akroviakov wrote:

```suggestion
    return aScaleTile->back();
```


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