[IR&PASS] part 3-2: add PatternApplicator and FrozenRewritePatternSet, refine PatternMatch code, add some api for Builder

paddle/ir/core/block.h

@@ -23,9 +23,10 @@ class Operation;
 
 class Block {
  public:
-  using iterator = std::list<Operation *>::iterator;
-  using reverse_iterator = std::list<Operation *>::reverse_iterator;
-  using const_iterator = std::list<Operation *>::const_iterator;
+  using OpListType = std::list<Operation *>;
+  using iterator = OpListType::iterator;
+  using reverse_iterator = OpListType::reverse_iterator;
+  using const_iterator = OpListType::const_iterator;
 
   Block() = default;
   ~Block();
@@ -56,7 +57,7 @@ class Block {
   void SetParent(Region *parent) { parent_ = parent; }
 
  private:
-  Region *parent_;              // not owned
-  std::list<Operation *> ops_;  // owned
+  Region *parent_;  // not owned
+  OpListType ops_;  // owned
 };
 }  // namespace ir

paddle/ir/core/builder.h

@@ -20,17 +20,24 @@
 #include "paddle/ir/core/operation.h"
 
 namespace ir {
+
 ///
 /// \brief Unified interface of the Attribute class. Derivation of all Attribute
 /// classes only derives interfaces, not members.
 ///
 class Builder {
  public:
   Builder(IrContext *context, Block *block, Block::iterator insert_point)
-      : context_(context), block_(block), insert_point_(insert_point) {}
+      : context_(context) {
+    SetInsertionPoint(block, insert_point);
+  }
+
   Builder(IrContext *context, Block *block)
       : Builder(context, block, block->end()) {}
 
+  explicit Builder(IrContext *context)
+      : Builder(context, nullptr, Block::iterator{}) {}
+
   static Builder AtBlockBegin(IrContext *context, Block *block) {
     return Builder(context, block, block->begin());
   }
@@ -39,7 +46,37 @@ class Builder {
     return Builder(context, block, block->end());
   }
 
-  IrContext *context() const { return context_; }
+  /// Set the insertion point to the specified location.
+  void SetInsertionPoint(Block *block, Block::iterator insert_point) {
+    // TODO(liuyuanle): check that insertPoint is in this rather than some other
+    // block.
+    this->block_ = block;
+    this->insert_point_ = insert_point;
+  }
+
+  /// Set the insertion point to the specified operation, which will cause
+  /// subsequent insertions to go right before it.
+  void SetInsertionPoint(Operation *op) {
+    SetInsertionPoint(op->GetParent(), Block::iterator{*op});
+  }
+
+  /// Set the insertion point to the node after the specified operation, which
+  /// will cause subsequent insertions to go right after it.
+  void SetInsertionPointAfter(Operation *op) {
+    SetInsertionPoint(op->GetParent(), std::next(Block::iterator{*op}));
+  }
+
+  /// Set the insertion point to the start of the specified block.
+  void SetInsertionPointToStart(Block *block) {
+    SetInsertionPoint(block, block->begin());
+  }
+
+  /// Set the insertion point to the end of the specified block.
+  void SetInsertionPointToEnd(Block *block) {
+    SetInsertionPoint(block, block->end());
+  }
+
+  IrContext *ir_context() const { return context_; }
 
   Block *block() const { return block_; }
 
@@ -65,8 +102,9 @@ class Builder {
   Operation *Insert(Operation *op);
 
   IrContext *context_;
-  Block *block_ = nullptr;
+  Block *block_;
   // The insertion point within the list that this builder is inserting before.
   Block::iterator insert_point_;
 };
+
 }  // namespace ir

paddle/ir/core/op_base.h

@@ -142,16 +142,16 @@ class ConstructInterfacesOrTraits {
   static void PlacementConstrctInterface(
       InterfaceValue *&p_interface) {  // NOLINT
     p_interface->swap(InterfaceValue::get<ConcreteOp, T>());
-    VLOG(4) << "New a interface: id[" << (p_interface->type_id()).storage()
-            << "].";
+    VLOG(4) << "New a interface: id["
+            << (p_interface->type_id()).AsOpaquePointer() << "].";
     ++p_interface;
   }
 
   /// Placement new trait.
   template <typename T>
   static void PlacementConstrctTrait(ir::TypeId *&p_trait) {  // NOLINT
     *p_trait = TypeId::get<T>();
-    VLOG(4) << "New a trait: id[" << p_trait->storage() << "].";
+    VLOG(4) << "New a trait: id[" << p_trait->AsOpaquePointer() << "].";
     ++p_trait;
   }
 };
@@ -206,4 +206,5 @@ class Op : public OpBase {
     return trait_set;
   }
 };
+
 }  // namespace ir

paddle/ir/core/op_info.h

@@ -68,8 +68,8 @@ class OpInfo {
   typename Interface::Concept *GetInterfaceImpl() const;
 
   void *AsOpaquePointer() const { return impl_; }
-  static OpInfo RecoverFromOpaquePointer(void *impl) {
-    return static_cast<OpInfoImpl *>(impl);
+  static OpInfo RecoverFromOpaquePointer(void *pointer) {
+    return OpInfo(reinterpret_cast<OpInfoImpl *>(pointer));
   }
 
   friend class OpInfoImpl;

paddle/ir/core/type_id.h

@@ -51,7 +51,10 @@ class TypeId {
 
   TypeId &operator=(const TypeId &other) = default;
 
-  const Storage *storage() const { return storage_; }
+  void *AsOpaquePointer() const { return storage_; }
+  static TypeId RecoverFromOpaquePointer(void *pointer) {
+    return TypeId(reinterpret_cast<Storage *>(pointer));
+  }
 
   ///
   /// \brief Comparison operations.
@@ -77,10 +80,11 @@ class TypeId {
   ///
   /// \param storage The storage of this TypeId.
   ///
-  explicit TypeId(const Storage *storage) : storage_(storage) {}
+  TypeId(Storage *storage) : storage_(storage) {}  // NOLINT
 
-  const Storage *storage_{nullptr};
+  Storage *storage_{nullptr};
 };
+
 }  // namespace ir
 
 namespace std {

paddle/ir/pattern_rewrite/frozen_rewrite_pattern_set.cc

@@ -0,0 +1,99 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+#include <memory>
+#include <optional>
+#include <set>
+#include <string>
+
+#include "paddle/ir/core/op_info.h"
+#include "paddle/ir/pattern_rewrite/frozen_rewrite_pattern_set.h"
+#include "paddle/utils/optional.h"
+
+namespace ir {
+
+FrozenRewritePatternSet::FrozenRewritePatternSet()
+    : impl_(std::make_shared<Impl>()) {}
+
+FrozenRewritePatternSet::FrozenRewritePatternSet(
+    RewritePatternSet&& patterns,
+    const std::vector<std::string>& disabled_pattern_labels,
+    const std::vector<std::string>& enabled_pattern_labels)
+    : impl_(std::make_shared<Impl>()) {
+  std::set<std::string> disabled_patterns, enabled_patterns;
+  disabled_patterns.insert(disabled_pattern_labels.begin(),
+                           disabled_pattern_labels.end());
+  enabled_patterns.insert(enabled_pattern_labels.begin(),
+                          enabled_pattern_labels.end());
+
+  ir::OpInfoMap op_info_map;
+  auto AddToOpsWhen = [&](std::unique_ptr<RewritePattern>& pattern,
+                          std::function<bool(OpInfo)> callback) {
+    if (op_info_map.empty())
+      op_info_map = pattern->ir_context()->registered_op_info_map();
+    for (auto& info_map : op_info_map) {
+      if (callback(info_map.second))
+        impl_->op_specific_native_pattern_map_[info_map.second].push_back(
+            pattern.get());
+      impl_->op_specific_native_patterns_.push_back(std::move(pattern));
+    }
+  };
+
+  for (std::unique_ptr<RewritePattern>& pat : patterns.native_patterns()) {
+    // Don't add patterns that haven't been enabled by the user.
+    if (!enabled_patterns.empty()) {
+      auto IsEnableFn = [&](const std::string& label) {
+        return enabled_patterns.count(label);
+      };
+      if (!IsEnableFn(pat->debug_name()) &&
+          std::none_of(pat->debug_labels().begin(),
+                       pat->debug_labels().end(),
+                       IsEnableFn))
+        continue;
+    }
+
+    // Don't add patterns that have been disabled by the user.
+    if (!disabled_patterns.empty()) {
+      auto IsDisabledFn = [&](const std::string& label) {
+        return disabled_patterns.count(label);
+      };
+      if (IsDisabledFn(pat->debug_name()) ||
+          std::any_of(pat->debug_labels().begin(),
+                      pat->debug_labels().end(),
+                      IsDisabledFn))
+        continue;
+    }
+
+    if (paddle::optional<OpInfo> root_name = pat->root_kind()) {
+      impl_->op_specific_native_pattern_map_[*root_name].push_back(pat.get());
+      impl_->op_specific_native_patterns_.push_back(std::move(pat));
+      continue;
+    }
+
+    if (paddle::optional<TypeId> interface_id = pat->GetRootInterfaceID()) {
+      AddToOpsWhen(
+          pat, [&](OpInfo info) { return info.HasInterface(*interface_id); });
+      continue;
+    }
+
+    if (paddle::optional<TypeId> trait_id = pat->GetRootTraitID()) {
+      AddToOpsWhen(pat, [&](OpInfo info) { return info.HasTrait(*trait_id); });
+      continue;
+    }
+
+    impl_->match_any_op_native_patterns_.push_back(std::move(pat));
+  }
+}
+
+}  // namespace ir

paddle/ir/pattern_rewrite/frozen_rewrite_pattern_set.h

@@ -0,0 +1,73 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+/// The design and code is mainly from MLIR, thanks to the greate project.
+
+#pragma once
+
+#include <memory>
+#include <string>
+#include <unordered_map>
+#include <vector>
+
+#include "paddle/ir/core/op_info.h"
+#include "paddle/ir/pattern_rewrite/pattern_match.h"
+
+namespace ir {
+
+class FrozenRewritePatternSet {
+  using NativePatternListT = std::vector<std::unique_ptr<RewritePattern>>;
+
+ public:
+  using OpSpecificNativePatternListT =
+      std::unordered_map<OpInfo, std::vector<RewritePattern*>>;
+
+  FrozenRewritePatternSet();
+  FrozenRewritePatternSet(FrozenRewritePatternSet&& patterns) = default;
+  FrozenRewritePatternSet(const FrozenRewritePatternSet& patterns) = default;
+  FrozenRewritePatternSet& operator=(FrozenRewritePatternSet&& patterns) =
+      default;
+  FrozenRewritePatternSet& operator=(const FrozenRewritePatternSet& patterns) =
+      default;
+  ~FrozenRewritePatternSet() = default;
+
+  /// Freeze the patterns held in `patterns`, and take ownership.
+  FrozenRewritePatternSet(
+      RewritePatternSet&& patterns,
+      const std::vector<std::string>& disabled_pattern_labels = {},
+      const std::vector<std::string>& enabled_pattern_labels = {});
+
+  /// Return the op specific native patterns held by this list.
+  const OpSpecificNativePatternListT& op_specific_native_patterns() const {
+    return impl_->op_specific_native_pattern_map_;
+  }
+
+  /// Return the "match any" native patterns held by this list.
+  const NativePatternListT& match_any_op_native_patterns() const {
+    return impl_->match_any_op_native_patterns_;
+  }
+
+ private:
+  struct Impl {
+    OpSpecificNativePatternListT op_specific_native_pattern_map_;
+
+    NativePatternListT op_specific_native_patterns_;
+
+    NativePatternListT match_any_op_native_patterns_;
+  };
+
+  std::shared_ptr<Impl> impl_;
+};
+
+}  // namespace ir