[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();
----------------
fschlimb wrote:

Nice.

This changes behavior for every operation that uses `defaultGetShardingOption`. Please call this out in the PR description and add a targeted regression test covering the change.

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