[Mlir-commits] [mlir] 3363a0e - [MLIR][Shard] Fix NormalizeSharding and FoldDuplicateShardOp direct mutations (#188981)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-27T16:22:12Z
New Revision: 3363a0ead9cf3235d45fcf508e3cd9e33558d597

URL: https://github.com/llvm/llvm-project/commit/3363a0ead9cf3235d45fcf508e3cd9e33558d597
DIFF: https://github.com/llvm/llvm-project/commit/3363a0ead9cf3235d45fcf508e3cd9e33558d597.diff

LOG: [MLIR][Shard] Fix NormalizeSharding and FoldDuplicateShardOp direct mutations (#188981)

Calling attribute setters and MutableOperandRange::assign() without
going through the PatternRewriter, bypassing the rewriter's
change-tracking triggered "operation finger print changed" after the
pattern returned success under
MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS.

Assisted-by: Claude Code

Added: 
    

Modified: 
    mlir/lib/Dialect/Shard/IR/ShardOps.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
index a173da3db1d18..ff790a0bf961d 100644
--- a/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
+++ b/mlir/lib/Dialect/Shard/IR/ShardOps.cpp
@@ -631,8 +631,12 @@ class NormalizeSharding final : public OpRewritePattern<ShardingOp> {
     bool modified = succeeded(foldDynamicIndexList(mixedHalos, true)) ||
                     succeeded(foldDynamicIndexList(mixedOffs, true));
 
-    auto [staticHalos, dynamicHalos] = decomposeMixedValues(mixedHalos);
-    auto [staticOffs, dynamicOffs] = decomposeMixedValues(mixedOffs);
+    auto decomposedHalos = decomposeMixedValues(mixedHalos);
+    auto staticHalos = decomposedHalos.first;
+    auto dynamicHalos = decomposedHalos.second;
+    auto decomposedOffs = decomposeMixedValues(mixedOffs);
+    auto staticOffs = decomposedOffs.first;
+    auto dynamicOffs = decomposedOffs.second;
 
     if (dynamicHalos.empty() && !staticHalos.empty()) {
       if (staticHalos[0] == 0 && llvm::all_equal(staticHalos)) {
@@ -666,11 +670,12 @@ class NormalizeSharding final : public OpRewritePattern<ShardingOp> {
       return failure();
     }
 
-    op.setStaticHaloSizes(staticHalos);
-    op.getDynamicHaloSizesMutable().assign(dynamicHalos);
-    op.setStaticShardedDimsOffsets(staticOffs);
-    op.getDynamicShardedDimsOffsetsMutable().assign(dynamicOffs);
-
+    b.modifyOpInPlace(op, [&]() {
+      op.setStaticHaloSizes(staticHalos);
+      op.getDynamicHaloSizesMutable().assign(dynamicHalos);
+      op.setStaticShardedDimsOffsets(staticOffs);
+      op.getDynamicShardedDimsOffsetsMutable().assign(dynamicOffs);
+    });
     return success();
   }
 };
@@ -877,7 +882,8 @@ class FoldDuplicateShardOp final : public OpRewritePattern<ShardOp> {
           b.eraseOp(op.getOperation());
         } else {
           // use the other sharding as input for op
-          op.getSrcMutable().assign(otherOp.getResult());
+          b.modifyOpInPlace(
+              op, [&]() { op.getSrcMutable().assign(otherOp.getResult()); });
         }
         return success();
       }