[Mlir-commits] [mlir] [mlir][tosa] Support variadic matmul sharding (PR #226180)

[Frank Schlimbach]

================
@@ -47,20 +47,95 @@ struct MatMulOpSharding
   }
 
   SmallVector<AffineMap> getIndexingMaps(Operation *op) const {
-    auto tensorType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
-    if (!tensorType)
+    auto aType = dyn_cast<RankedTensorType>(op->getOperand(0).getType());
+    auto bType = dyn_cast<RankedTensorType>(op->getOperand(1).getType());
+    auto outputType = dyn_cast<RankedTensorType>(op->getResult(0).getType());
+    if (!aType || !bType || !outputType || aType.getRank() < 2 ||
+        bType.getRank() < 2 || outputType.getRank() < aType.getRank() ||
+        outputType.getRank() < bType.getRank())
       return {};
+
     MLIRContext *ctx = op->getContext();
+    const int64_t outputRank = outputType.getRank();
+    const int64_t loopRank = outputRank + 1;
+    const int64_t reductionLoop = outputRank;
+
+    auto getAlignedBatchDim = [&](RankedTensorType operandType,
+                                  int64_t outputDim) {
+      const int64_t operandDim =
+          outputDim - (outputRank - operandType.getRank());
+      return operandDim < 0 ? 1 : operandType.getDimSize(operandDim);
+    };
+    for (int64_t i = 0; i < outputRank - 2; ++i) {
+      const int64_t aDim = getAlignedBatchDim(aType, i);
+      const int64_t bDim = getAlignedBatchDim(bType, i);
+      // A dynamic dimension may be either one (broadcast) or the corresponding
+      // output extent. These cases require different indexing maps unless the
+      // other operand is known to be a singleton.
+      if ((ShapedType::isDynamic(aDim) && bDim != 1) ||
+          (ShapedType::isDynamic(bDim) && aDim != 1))
+        return {};
+    }
----------------
fschlimb wrote:

This looks overly restrictive. It bails out as soon as any batch dim is ambiguous, even if nothing shards that dim. For example, `?x4x8 @ ?x8x16` with the result split on dim 1 no longer propagates at all. The ambiguity only matters if that batch loop actually gets split.

I suggest always building the maps (dynamic dims already get `d_i`) and moving this check into getShardingOption. There, run defaultGetShardingOption first, then decline only if the resulting option splits an ambiguous batch loop.

A test case would be good as well.

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