[llvm-branch-commits] [mlir] [mlir][Transforms] Support `replaceAllUsesWith` in dialect conversion (PR #84725)

[Mehdi Amini via llvm-branch-commits]

================
@@ -751,6 +731,44 @@ class UnresolvedMaterializationRewrite : public OperationRewrite {
   /// The original output type. This is only used for argument conversions.
   Type origOutputType;
 };
+
+/// A value rewrite.
+class ValueRewrite : public IRRewrite {
+public:
+  /// Return the operation that this rewrite operates on.
+  Value getValue() const { return value; }
+
+  static bool classof(const IRRewrite *rewrite) {
+    return rewrite->getKind() >= Kind::ReplaceAllUses &&
+           rewrite->getKind() <= Kind::ReplaceAllUses;
+  }
+
+protected:
+  ValueRewrite(Kind kind, ConversionPatternRewriterImpl &rewriterImpl,
+               Value value)
+      : IRRewrite(kind, rewriterImpl), value(value) {}
+
+  // The value that this rewrite operates on.
+  Value value;
+};
+
+/// Replacing a value. This rewrite is not immediately reflected in the IR. An
+/// internal IR mapping is updated, but the actual replacement is delayed until
+/// the rewrite is committed.
+class ReplaceAllUsesRewrite : public ValueRewrite {
+public:
+  ReplaceAllUsesRewrite(ConversionPatternRewriterImpl &rewriterImpl,
+                        Value value)
+      : ValueRewrite(Kind::ReplaceAllUses, rewriterImpl, value) {}
+
+  static bool classof(const IRRewrite *rewrite) {
+    return rewrite->getKind() == Kind::ReplaceAllUses;
+  }
+
+  void commit(RewriterBase &rewriter) override;
+
+  void rollback() override;
+};
----------------
joker-eph wrote:

The whole flow of Dialect conversion tracking of these changes is too complicated for me to know whether the commit/rollback logic is safe and complete here :(
It's likely that only testing will tell, but that's unfortunate!

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