PaddlePaddle
diff --git a/‎paddle/ir/core/operation.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/ir/core/operation.cc
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/ir/core/operation.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions b/‎paddle/ir/core/operation.h
Copy file name to clipboard
+8Lines changed: 8 additions & 0 deletions
diff --git a/‎paddle/ir/core/region.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions b/‎paddle/ir/core/region.cc
Copy file name to clipboard
+6Lines changed: 6 additions & 0 deletions
diff --git a/‎paddle/ir/core/region.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions b/‎paddle/ir/core/region.h
Copy file name to clipboard
+3Lines changed: 3 additions & 0 deletions
diff --git a/‎paddle/ir/core/value.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions b/‎paddle/ir/core/value.cc
Copy file name to clipboard
+2Lines changed: 2 additions & 0 deletions
diff --git a/‎paddle/ir/core/value.h
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions b/‎paddle/ir/core/value.h
Copy file name to clipboard
+4-2Lines changed: 4 additions & 2 deletions
diff --git a/‎paddle/ir/core/value_impl.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions b/‎paddle/ir/core/value_impl.h
Copy file name to clipboard
+4Lines changed: 4 additions & 0 deletions
diff --git a/‎paddle/ir/pattern_rewrite/frozen_rewrite_pattern_set.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion b/‎paddle/ir/pattern_rewrite/frozen_rewrite_pattern_set.h
Copy file name to clipboard
+2-1Lines changed: 2 additions & 1 deletion
diff --git a/‎paddle/ir/pattern_rewrite/pattern_match.cc
Copy file name to clipboard
+50-39Lines changed: 50 additions & 39 deletions b/‎paddle/ir/pattern_rewrite/pattern_match.cc
Copy file name to clipboard
+50-39Lines changed: 50 additions & 39 deletions
diff --git a/‎paddle/ir/pattern_rewrite/pattern_match.h
Copy file name to clipboard
+29-23Lines changed: 29 additions & 23 deletions b/‎paddle/ir/pattern_rewrite/pattern_match.h
Copy file name to clipboard
+29-23Lines changed: 29 additions & 23 deletions
@@ -224,4 +224,12 @@ void Operation::SetParent(Block *parent, const Block::iterator &position) {
   position_ = position;
 }
 
+void Operation::ReplaceAllUsesWith(const std::vector<Value> &values) {
+  IR_ENFORCE(num_results_ == values.size(),
+             "the num of result should be the same.");
+  for (uint32_t i = 0; i < num_results_; ++i) {
+    result(i).ReplaceAllUsesWith(values[i]);
+  }
+}
+
 }  // namespace ir
@@ -15,6 +15,7 @@
 #pragma once
 
 #include <ostream>
+#include <vector>
 #include "paddle/ir/core/block.h"
 #include "paddle/ir/core/op_info.h"
 #include "paddle/ir/core/operation_utils.h"
@@ -102,6 +103,13 @@ class IR_API alignas(8) Operation final {
 
   operator Block::const_iterator() const { return position_; }
 
+  /// Replace all uses of results of this operation with the provided 'values'.
+  void ReplaceAllUsesWith(const std::vector<Value> &values);
+
+  inline void ReplaceAllUsesWith(Value value) {
+    ReplaceAllUsesWith(std::vector<Value>{value});
+  }
+
  private:
   Operation(const AttributeMap &attribute,
             ir::OpInfo op_info,
 
@@ -15,6 +15,7 @@
 #include "paddle/ir/core/region.h"
 #include "paddle/ir/core/block.h"
 #include "paddle/ir/core/enforce.h"
+#include "paddle/ir/core/operation.h"
 
 namespace ir {
 Region::~Region() { clear(); }
@@ -50,4 +51,9 @@ void Region::clear() {
     blocks_.pop_back();
   }
 }
+
+IrContext *Region::ir_context() const {
+  IR_ENFORCE(parent_, "Region is not attached to a container.");
+  return parent_->ir_context();
+}
 }  // namespace ir
@@ -23,6 +23,7 @@ namespace ir {
 
 class Block;
 class Operation;
+class IrContext;
 
 class IR_API Region {
  public:
@@ -55,6 +56,8 @@ class IR_API Region {
 
   Operation *GetParent() const { return parent_; }
 
+  IrContext *ir_context() const;
+
  private:
   Region(Region &) = delete;
   Region &operator=(const Region &) = delete;
 
@@ -85,6 +85,8 @@ OpOperand Value::first_use() const { return impl()->first_use(); }
 
 bool Value::use_empty() const { return !first_use(); }
 
+bool Value::HasOneUse() const { return impl()->HasOneUse(); }
+
 void Value::ReplaceUsesWithIf(
     Value new_value,
     const std::function<bool(OpOperand)> &should_replace) const {
 
@@ -158,10 +158,12 @@ class IR_API Value {
 
   OpOperand first_use() const;
 
-  friend struct std::hash<Value>;
-
   bool use_empty() const;
 
+  bool HasOneUse() const;
+
+  friend struct std::hash<Value>;
+
   void ReplaceUsesWithIf(
       Value new_value,
       const std::function<bool(OpOperand)> &should_replace) const;
 
@@ -98,6 +98,10 @@ class alignas(8) ValueImpl {
 
   bool use_empty() const { return first_use() == nullptr; }
 
+  bool HasOneUse() const {
+    return (first_use() != nullptr) && (first_use()->next_use() == nullptr);
+  }
+
   std::string PrintUdChain();
 
  protected:
 
@@ -21,12 +21,13 @@
 #include <unordered_map>
 #include <vector>
 
+#include "paddle/ir/core/dll_decl.h"
 #include "paddle/ir/core/op_info.h"
 #include "paddle/ir/pattern_rewrite/pattern_match.h"
 
 namespace ir {
 
-class FrozenRewritePatternSet {
+class IR_API FrozenRewritePatternSet {
   using NativePatternListT = std::vector<std::unique_ptr<RewritePattern>>;
 
  public:
 
@@ -15,9 +15,9 @@
 #include "paddle/ir/pattern_rewrite/pattern_match.h"
 
 #include <algorithm>
-#include <cassert>
 #include <cstdint>
 
+#include "paddle/ir/core/enforce.h"
 #include "paddle/ir/core/operation.h"
 
 namespace ir {
@@ -90,62 +90,73 @@ RewritePattern::~RewritePattern() = default;
 //===----------------------------------------------------------------------===//
 RewriterBase::~RewriterBase() = default;
 
-// TODO(wilber): value support replace method.
-// void RewriterBase::ReplaceOpWithIf(Operation* op,
-//                                    ValueRange new_values,
-//                                    bool* all_uses_replaced,
-//                                    std::function<bool(OpOperand&)> functor) {
-//   // assert(op->num_results() == new_values.size() && "incorrect number of
-//   values to replace operation"); NotifyRootReplaced(op, new_values); bool
-//   replace_all_uses = true; for (uint32_t i = 0; i < op->num_results(); ++i) {
-//     // op->result(0)
-//   }
-// }
-// void RewriterBase::ReplaceOpWithIf(Operation* op,
-//                        ValueRange new_values,
-//                        std::function<bool(OpOperand&)> functor) {
-//   ReplaceOpWithIf(op, new_values, nullptr, functor);
-// }
-
-// TODO(wilber): support erase.
-// void ReplaceOp(Operation* op, ValueRange new_values) {
-//   NotifyRootReplaced(op, new_values);
-//   assert(op->num_results() == new_values.size() && "incorrect # of
-//   replacement values"); op->ReplaceAllUsesWith(new_values);
-//   NotifyOperationRemoved(op);
-//   op->erase();
-// }
+void RewriterBase::ReplaceOpWithIf(
+    Operation* op,
+    const std::vector<Value>& new_values,
+    bool* all_uses_replaced,
+    const std::function<bool(OpOperand)>& functor) {
+  IR_ENFORCE(op->num_results() == new_values.size(),
+             "incorrect number of values to replace operation");
+  NotifyRootReplaced(op, new_values);
+
+  // Replace each use of the results when the functor is true.
+  bool replace_all_uses = true;
+  for (uint32_t i = 0; i < op->num_results(); ++i) {
+    auto src_res = op->result(i);
+    src_res.ReplaceUsesWithIf(new_values[i], functor);
+    replace_all_uses &= src_res.use_empty();
+  }
+  if (replace_all_uses) {
+    *all_uses_replaced = replace_all_uses;
+  }
+}
+
+void RewriterBase::ReplaceOpWithIf(
+    Operation* op,
+    const std::vector<Value>& new_values,
+    const std::function<bool(OpOperand)>& functor) {
+  ReplaceOpWithIf(op, new_values, nullptr, functor);
+}
+
+void RewriterBase::ReplaceOp(Operation* op,
+                             const std::vector<Value>& new_values) {
+  NotifyRootReplaced(op, new_values);
+  IR_ENFORCE(op->num_results() == new_values.size(),
+             "incorrect # of replacement values");
+  op->ReplaceAllUsesWith(new_values);
+  NotifyOperationRemoved(op);
+  op->GetParent()->erase(*op);
+}
+
 void RewriterBase::EraseOp(Operation* op) {
-  //   assert(op->use_empty() && "expected 'op' to have no uses");
-  //   NotifyOperationRemoved(op);
-  //   op->erase();
+  // TODO(wilber): Operation support use_empty.
+  // IR_ENFORCE(op->use_empty(), "expected 'op' to have no uses");
+  NotifyOperationRemoved(op);
+  op->GetParent()->erase(*op);
 }
 
+/// Find uses of `from` and replace it with `to`
 void RewriterBase::ReplaceAllUsesWith(Value from, Value to) {
-  // from.
-  // for (OpOperand& operand : llvm::make_early_inc_range(from.getUses()))
-  // {
-  //   Operation* op = operand.getOwner();
-  //   UpdateRootInPlace(op, [&]() { operand.set(to); });
-  // }
+  // TODO(wilber): Substitue a low level impl.
+  from.ReplaceAllUsesWith(to);
 }
 
 // TODO(wilber): iterator maybe should support modify inplace.
 void RewriterBase::ReplaceUseIf(Value from,
                                 Value to,
                                 std::function<bool(OpOperand&)> functor) {
   // for (auto it = from.begin(); it != from.end(); ++it) {
-  //   // TODO: need a lvalue.
-  //   if (functor(it.get())) {
+  // //   // TODO: need a lvalue.
+  //   if (functor(*it)) {
   //     UpdateRootInplace(it.owner(), [&](){it.get().set(to)});
   //   }
   // }
 }
 
 void RewriterBase::ReplaceOpWithResultsOfAnotherOp(Operation* op,
                                                    Operation* new_op) {
-  assert(op->num_results() == new_op->num_results() &&
-         "replacement op doesn't match results of original op");
+  IR_ENFORCE(op->num_results() == new_op->num_results(),
+             "replacement op doesn't match results of original op");
   // TODO(wilber): Op support results method.
   // if (op->num_results() == 1) return ReplaceOp(op,
   // new_op->result(0)); return ReplaceOp(op, new_op->GetResults());
 
@@ -25,6 +25,7 @@
 #include <vector>
 
 #include "paddle/ir/core/builder.h"
+#include "paddle/ir/core/dll_decl.h"
 #include "paddle/ir/core/ir_context.h"
 #include "paddle/ir/core/op_info.h"
 #include "paddle/ir/core/operation.h"
@@ -36,7 +37,7 @@ namespace ir {
 
 // This class reprensents the benefit of a pattern. The most common
 // unit to use is the `numver of operations` in the pattern.
-class PatternBenefit {
+class IR_API PatternBenefit {
  public:
   PatternBenefit() = default;
   PatternBenefit(uint32_t val) : val_(val) {}  // NOLINT
@@ -257,30 +258,21 @@ class RewriterBase : public Builder {
  public:
   // TODO(wilber): Supplementary methods of block and region.
 
-  // TODO(wilber): Support ValueRange.
-  // virtual void ReplaceOpWithIf(Operation* op,
-  //                              ValueRange new_values,
-  //                              bool* all_uses_replaced,
-  //                              std::function<bool(OpOperand&)> functor);
-  // void ReplaceOpWithIf(Operation* op,
-  //                      ValueRange new_values,
-  //                      std::function<bool(OpOperand&)> functor);
-  // virtual void ReplaceOp(Operation* op, ValueRange new_values);
+  virtual void ReplaceOpWithIf(Operation* op,
+                               const std::vector<Value>& new_values,
+                               bool* all_uses_replaced,
+                               const std::function<bool(OpOperand)>& functor);
 
-  // virtual void ReplaceOpWithNewOp()
+  void ReplaceOpWithIf(Operation* op,
+                       const std::vector<Value>& new_values,
+                       const std::function<bool(OpOperand)>& functor);
 
-  virtual void EraseOp(Operation* op);
+  virtual void ReplaceOp(Operation* op, const std::vector<Value>& new_values);
 
-  virtual void StartRootUpdate(Operation* op) {}
-  virtual void FinalizeRootUpdate(Operation* op) {}
-  virtual void CancleRootUpdate(Operation* op) {}
+  // template <typename OpTy, typename... Args>
+  // OpTy ReplaceOpWithNewOp(Operation *op, Args &&...args);
 
-  template <typename CallableT>
-  void UpdateRootInplace(Operation* root, CallableT&& callable) {
-    StartRootUpdate(root);
-    callable();
-    FinalizeRootUpdate(root);
-  }
+  virtual void EraseOp(Operation* op);
 
   void ReplaceAllUsesWith(Value from, Value to);
 
@@ -293,11 +285,25 @@ class RewriterBase : public Builder {
 
   virtual ~RewriterBase();
 
-  // virtual void NotifyRootReplaced(Operation* op, ValueRange replacement) {}
+  virtual void NotifyRootReplaced(Operation* op,
+                                  const std::vector<Value>& replacement) {}
 
   virtual void NotifyOperationRemoved(Operation* op) {}
 
-  // virtual bool NotifyMatchFailure()
+  virtual void NotifyOperationInserted(Operation* op) {}
+
+  virtual void StartRootUpdate(Operation* op) {}
+
+  virtual void FinalizeRootUpdate(Operation* op) {}
+
+  virtual void CancleRootUpdate(Operation* op) {}
+
+  template <typename CallableT>
+  void UpdateRootInplace(Operation* root, CallableT&& callable) {
+    StartRootUpdate(root);
+    callable();
+    FinalizeRootUpdate(root);
+  }
 
  private:
   void operator=(const RewriterBase&) = delete;