MLIR: Add ReplaceOpAction into DialectConversion and PatternMatch

Author: mgehre-amd

mlir/include/mlir/IR/PatternMatch.h

@@ -13,6 +13,7 @@
 #include "mlir/IR/BuiltinOps.h"
 #include "llvm/ADT/FunctionExtras.h"
 #include "llvm/Support/TypeName.h"
+#include <mlir/IR/PatternMatchAction.h>
 #include <optional>
 
 using llvm::SmallPtrSetImpl;
@@ -652,6 +653,8 @@ class RewriterBase : public OpBuilder {
   /// Find uses of `from` and replace them with `to`. Also notify the listener
   /// about every in-place op modification (for every use that was replaced).
   void replaceAllUsesWith(Value from, Value to) {
+    if(auto* fromOp = from.getDefiningOp())
+      getContext()->executeAction<ReplaceOpAction>([]() {}, ArrayRef<IRUnit>{fromOp}, to);
     for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
       Operation *op = operand.getOwner();
       modifyOpInPlace(op, [&]() { operand.set(to); });

mlir/include/mlir/IR/PatternMatchAction.h

@@ -0,0 +1,20 @@
+#ifndef MLIR_IR_PATTERNMATCHACTION_H
+#define MLIR_IR_PATTERNMATCHACTION_H
+
+#include "mlir/IR/Action.h"
+
+namespace mlir {
+struct ReplaceOpAction : public tracing::ActionImpl<ReplaceOpAction> {
+  using Base = tracing::ActionImpl<ReplaceOpAction>;
+  ReplaceOpAction(ArrayRef<IRUnit> irUnits, ValueRange replacement);
+  static constexpr StringLiteral tag = "op-replacement";
+  void print(raw_ostream &os) const override;
+
+  Operation *getOp() const;
+
+public:
+  ValueRange replacement;
+};
+}
+
+#endif

mlir/lib/IR/PatternMatch.cpp

@@ -127,6 +127,35 @@ void RewriterBase::replaceAllOpUsesWith(Operation *from, Operation *to) {
   replaceAllUsesWith(from->getResults(), to->getResults());
 }
 
+ReplaceOpAction::ReplaceOpAction(ArrayRef<IRUnit> irUnits,
+                                 ValueRange replacement)
+    : Base(irUnits), replacement(replacement) {
+  assert(irUnits.size() == 1);
+  assert(irUnits[0]);
+  assert(isa<Operation *>(irUnits[0]));
+}
+
+void ReplaceOpAction::print(raw_ostream &os) const {
+  OpPrintingFlags flags;
+  flags.elideLargeElementsAttrs(10);
+  os << "`" << tag << "` replacing operation `";
+  getOp()->print(os, flags);
+  os << "` by ";
+  bool first = true;
+  for (auto r : replacement) {
+    if (!first)
+      os << ", ";
+    os << "`";
+    r.print(os, flags);
+    os << "`";
+    first = false;
+  }
+}
+
+Operation *ReplaceOpAction::getOp() const {
+  return cast<Operation *>(getContextIRUnits()[0]);
+}
+
 /// This method replaces the results of the operation with the specified list of
 /// values. The number of provided values must match the number of results of
 /// the operation. The replaced op is erased.
@@ -265,8 +294,12 @@ void RewriterBase::replaceUsesWithIf(Value from, Value to,
   bool allReplaced = true;
   for (OpOperand &operand : llvm::make_early_inc_range(from.getUses())) {
     bool replace = functor(operand);
-    if (replace)
+    if (replace) {
+      if (auto *fromOp = from.getDefiningOp())
+        getContext()->executeAction<ReplaceOpAction>(
+            []() {}, ArrayRef<IRUnit>{fromOp}, to);
       modifyOpInPlace(operand.getOwner(), [&]() { operand.set(to); });
+    }
     allReplaced &= replace;
   }
   if (allUsesReplaced)