[Mlir-commits] [mlir] 6c8782b - [MLIR][Vector] Fix direct operand.set() bypassing rewriter in WarpOpScfIfOp/ForOp (#188948)

[llvmlistbot at llvm.org]

Author: Mehdi Amini
Date: 2026-03-30T12:28:07+02:00
New Revision: 6c8782b347b56f8499d4d914052e7f819f8e9fab

URL: https://github.com/llvm/llvm-project/commit/6c8782b347b56f8499d4d914052e7f819f8e9fab
DIFF: https://github.com/llvm/llvm-project/commit/6c8782b347b56f8499d4d914052e7f819f8e9fab.diff

LOG: [MLIR][Vector] Fix direct operand.set() bypassing rewriter in WarpOpScfIfOp/ForOp (#188948)

In WarpOpScfIfOp and WarpOpScfForOp, the walk that updates users of
escaping values (after moving them to the inner WarpOp) was calling
operand.set() directly, bypassing the rewriter API. This causes the
MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS fingerprint check to fail.

Fix by wrapping the operand updates with rewriter.modifyOpInPlace().

Assisted-by: Claude Code
Fix a failure present with MLIR_ENABLE_EXPENSIVE_PATTERN_API_CHECKS=ON.

Added: 
    

Modified: 
    mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp

Removed: 
    


################################################################################
diff  --git a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
index b4d500212c770..2e0e650f2bb9c 100644
--- a/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
+++ b/mlir/lib/Dialect/Vector/Transforms/VectorDistribute.cpp
@@ -1978,11 +1978,20 @@ struct WarpOpScfIfOp : public WarpDistributionPattern {
           // Update any users of escaping values that were forwarded to the
           // inner `WarpOp`. These values are arguments of the inner `WarpOp`.
           innerWarp.walk([&](Operation *op) {
+            SmallVector<std::pair<unsigned, Value>> replacements;
             for (OpOperand &operand : op->getOpOperands()) {
               auto it = escapeValToBlockArgIndex.find(operand.get());
               if (it == escapeValToBlockArgIndex.end())
                 continue;
-              operand.set(innerWarp.getBodyRegion().getArgument(it->second));
+              replacements.emplace_back(
+                  operand.getOperandNumber(),
+                  innerWarp.getBodyRegion().getArgument(it->second));
+            }
+            if (!replacements.empty()) {
+              rewriter.modifyOpInPlace(op, [&]() {
+                for (auto [idx, newVal] : replacements)
+                  op->setOperand(idx, newVal);
+              });
             }
           });
           mlir::vector::moveScalarUniformCode(innerWarp);
@@ -2218,11 +2227,20 @@ struct WarpOpScfForOp : public WarpDistributionPattern {
     // Update any users of escaping values that were forwarded to the
     // inner `WarpOp`. These values are now arguments of the inner `WarpOp`.
     newForOp.walk([&](Operation *op) {
+      SmallVector<std::pair<unsigned, Value>> replacements;
       for (OpOperand &operand : op->getOpOperands()) {
         auto it = argIndexMapping.find(operand.get());
         if (it == argIndexMapping.end())
           continue;
-        operand.set(innerWarp.getBodyRegion().getArgument(it->second));
+        replacements.emplace_back(
+            operand.getOperandNumber(),
+            innerWarp.getBodyRegion().getArgument(it->second));
+      }
+      if (!replacements.empty()) {
+        rewriter.modifyOpInPlace(op, [&]() {
+          for (auto [idx, newVal] : replacements)
+            op->setOperand(idx, newVal);
+        });
       }
     });