[Mlir-commits] [mlir] [mlir][mesh] add support in spmdization for incomplete sharding annotations (PR #82442)

[Boian Petkantchin]

================
@@ -615,34 +614,59 @@ static SmallVector<MeshShardingAttr> getResultShardings(Operation &op) {
                     assert(result.hasOneUse());
                     Operation *userOp = *result.getUsers().begin();
                     ShardOp shardOp = llvm::cast<ShardOp>(userOp);
-                    assert(!shardOp.getAnnotateForUsers());
                     return shardOp.getShard();
                   });
   return res;
 }
 
+static ShardOp getSourceShardOpOrNull(ShardOp targetShardOp) {
+  Operation *srcOp = targetShardOp.getOperand().getDefiningOp();
+  if (!srcOp) {
+    return ShardOp();
+  }
+  ShardOp srcShardOp =
+      llvm::dyn_cast<ShardOp>(targetShardOp.getOperand().getDefiningOp());
----------------
sogartar wrote:

Done.

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