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

[Mehdi Amini]

https://github.com/joker-eph created https://github.com/llvm/llvm-project/pull/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.

>From e3f73f3cb88a67b86eccefccf167f16b2311100d Mon Sep 17 00:00:00 2001
From: Mehdi Amini <joker.eph at gmail.com>
Date: Thu, 26 Mar 2026 15:57:51 -0700
Subject: [PATCH] [MLIR][Vector] Fix direct operand.set() bypassing rewriter in
 WarpOpScfIfOp/ForOp

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.
---
 .../Vector/Transforms/VectorDistribute.cpp    | 22 +++++++++++++++++--
 1 file changed, 20 insertions(+), 2 deletions(-)

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);
+        });
       }
     });