[Mlir-commits] [mlir] [MLIR][shard] checking for correct&full sharding annotations (PR #176000)

[Jakub Kuderski]

================
@@ -663,13 +663,87 @@ partitionOperation(ShardOp shardOp, IRMapping &partitionMap,
   return success();
 }
 
+// Check if the block args are correctly annotated with sharding information:
+//   - non-tensor and 0d-tensor args are ignored
+//   - each tensor arg must have exactly one use, which must be a shard.shard
+//   operation
+static LogicalResult checkFullyAnnotated(Block &block) {
+  for (auto arg : block.getArguments()) {
+    auto rankedTensorArg = dyn_cast<TypedValue<RankedTensorType>>(arg);
+    if (!rankedTensorArg || rankedTensorArg.getType().getRank() == 0)
+      continue;
+
+    if (rankedTensorArg.getNumUses() > 1)
+      return emitError(block.getParent()->getLoc())
+             << "Cannot partition: expected a single use for block argument "
+             << arg.getArgNumber() << " in block "
+             << block.computeBlockNumber();
+    Operation *useOp = *rankedTensorArg.getUsers().begin();
+    ShardOp shardOp = llvm::dyn_cast<ShardOp>(useOp);
----------------
kuhar wrote:

```suggestion
    ShardOp shardOp = dyn_cast<ShardOp>(useOp);
```

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