[Mlir-commits] [mlir] [mlir][IR] Tweak `RewriterBase::replaceUsesWithIf` call (PR #172883)

[llvmlistbot at llvm.org]

llvmbot wrote:


<!--LLVM PR SUMMARY COMMENT-->

@llvm/pr-subscribers-mlir

Author: Babak (BabakkGraphcore)

<details>
<summary>Changes</summary>

`ConversionPatternRewriter::replaceUsesWithIf` does not support the `allUsesReplaced` flag and asserts that it's not set. However the `ValueRange` overload of  `RewriterBase::replaceUsesWithIf` always passes this flag when calling the virtual overload of `replaceUsesWithIf`. This means calls made to `RewriterBase::replaceUsesWithIf` from a `ConversionPattern`  crash, whether the `allUsesReplaced` flag is set or not.

This change tweaks `RewriterBase::replaceUsesWithIf` to only pass that flag if the callee has set it.

---
Full diff: https://github.com/llvm/llvm-project/pull/172883.diff


1 Files Affected:

- (modified) mlir/lib/IR/PatternMatch.cpp (+9-8) 


``````````diff
diff --git a/mlir/lib/IR/PatternMatch.cpp b/mlir/lib/IR/PatternMatch.cpp
index 8bc0fcd4517d8..07fd52e0812c6 100644
--- a/mlir/lib/IR/PatternMatch.cpp
+++ b/mlir/lib/IR/PatternMatch.cpp
@@ -276,15 +276,16 @@ void RewriterBase::replaceUsesWithIf(ValueRange from, ValueRange to,
                                      function_ref<bool(OpOperand &)> functor,
                                      bool *allUsesReplaced) {
   assert(from.size() == to.size() && "incorrect number of replacements");
-  bool allReplaced = true;
-  for (auto it : llvm::zip_equal(from, to)) {
-    bool r;
-    replaceUsesWithIf(std::get<0>(it), std::get<1>(it), functor,
-                      /*allUsesReplaced=*/&r);
-    allReplaced &= r;
+  for (auto [fromVal, toVal] : llvm::zip_equal(from, to)) {
+    if (allUsesReplaced) {
+      bool r;
+      replaceUsesWithIf(fromVal, toVal, functor,
+                        /*allUsesReplaced=*/&r);
+      *allUsesReplaced &= r;
+    } else {
+      replaceUsesWithIf(fromVal, toVal, functor);
+    }
   }
-  if (allUsesReplaced)
-    *allUsesReplaced = allReplaced;
 }
 
 void RewriterBase::inlineBlockBefore(Block *source, Block *dest,

``````````

</details>


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