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

[Frank Schlimbach]

================
@@ -304,25 +307,41 @@ shard::detail::defaultGetShardingOption(Operation *op,
     if (!shardAttr)
       continue;
 
-    anyShardingInResultsOrOperands = !shardAttr.getSplitAxes().empty();
+    // Record the grid before processing mapped dimensions. Constant-indexed
+    // dimensions represent replication and do not call fillShardingOption.
+    FlatSymbolRefAttr grid = shardAttr.getGridAttr();
+    if (shardingOption.grid && grid && shardingOption.grid != grid)
+      return failure();
+    if (grid)
+      shardingOption.grid = grid;
+    anyShardingInResultsOrOperands |= !shardAttr.getSplitAxes().empty();
     AffineMap map = maps[shardingIt.index()];
     unsigned numDims = map.getNumDims();
 
     // Handle the split axes.
     //
     // TODO: Change to process the operands with single loop index first and
     // then the operands with multiple loop indices.
-    for (auto it : llvm::zip(map.getResults(), shardAttr.getSplitAxes())) {
-      AffineExpr expr = std::get<0>(it);
-      ArrayRef<GridAxis> axes = std::get<1>(it).asArrayRef();
+    for (auto [dim, exprAndAxes] : llvm::enumerate(
+             llvm::zip(map.getResults(), shardAttr.getSplitAxes()))) {
+      AffineExpr expr = std::get<0>(exprAndAxes);
+      ArrayRef<GridAxis> axes = std::get<1>(exprAndAxes).asArrayRef();
       FailureOr<llvm::SmallSet<unsigned, 2>> loopIndices =
           checkOperandAffineExpr(expr, numDims);
       if (failed(loopIndices))
         return op->emitOpError()
                << "operand's affine expression is restricted to const_i * "
                   "dim_i + const_j + dim_j + ...";
-      if (loopIndices->empty())
+      // A constant expression selects the same element for every loop
+      // iteration. It can represent replication, but cannot carry split axes.
+      if (loopIndices->empty()) {
+        if (!axes.empty())
+          return op->emitOpError()
+                 << "cannot shard operand " << shardingIt.index()
+                 << " dimension " << dim + 1
+                 << ": its indexing expression is constant";
----------------
fschlimb wrote:

This should not return an error, it only rejects a candidate. Printing an error (possibly twice) is confusing if the propagation actually succeeds. This is not really an error.

Better issue a `LLVM_DEBUG` message and return plain failure().

While at it, you could also similarly fix the 2 other error cases lines 331 and 363.

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