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

[Lei Zhang]

================
@@ -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());
----------------
antiagainst wrote:

This can be merged with the above using `dyn_cast_or_null`?

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