[Mlir-commits] [mlir] [MLIR][mesh] Mesh fixes (PR #124724)

[Frank Schlimbach]

================
@@ -796,6 +822,53 @@ void ShardOp::getAsmResultNames(
   setNameFn(getResult(), "sharding_annotated");
 }
 
+namespace {
+// Determine if the given ShardOp is a duplicate of another ShardOp
+// on the same value. This can happen if constant values are sharded.
+class FoldDuplicateShardOp final : public OpRewritePattern<ShardOp> {
+public:
+  using OpRewritePattern<ShardOp>::OpRewritePattern;
+
+  LogicalResult matchAndRewrite(ShardOp op, PatternRewriter &b) const override {
+    // Get the use-list of the value being sharded and check if it has more than
+    // one use.
+    Value value = op.getSrc();
+    if (value.hasOneUse() || value.getDefiningOp<ShardOp>()) {
+      return failure();
+    }
+
+    // Iterate through the uses of the value to find a duplicate ShardOp.
+    for (auto &use : value.getUses()) {
+      if (use.getOwner() != op.getOperation()) {
+        auto otherOp = dyn_cast<ShardOp>(use.getOwner());
+        if (!otherOp || !otherOp->isBeforeInBlock(op)) {
----------------
fschlimb wrote:

Your example does not cover the pattern which is canonicalized because it does not shard the same value mutliple times.

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