[IR] OpTrait & OpInterface & OpInfo (PaddlePaddle#52846)

* add OpTrait OpInterface ValueIterator TypeList

* refine code

* refine code

* refine code

* add opinfo

* add typeid copy constructor

* add trait interface construct method for opinfo_impl

* add trait interface construct method for opinfo_impl

* add trait interface construct method for opinfo_impl

* add trait interface construct method for opinfo_impl

* add trait interface construct method for opinfo_impl

* add create

* add member func for opinfo

* fix compile bug

* add op interface in ircontext

* fix compile bug

* fix compile bug

* refine code

* fix compile bug

* add ut

* refine ut

* refine code of opinfo_impl

* delete unused code

* add dyncast for operation

* refine comment

* refine opinfo_impl

* delete unused code

* refine code by comment

* refine code

* refine code

* refine code for registerOp

* refine opfin create

* refine code of search method of ircontext

* refine op attribute

* change opinfo_map key from type_id to string

Author: zhangbo9674

paddle/ir/builtin_attribute_storage.cc

@@ -59,19 +59,16 @@ DictionaryAttributeStorage::ParamKey DictionaryAttributeStorage::GetAsKey()
 }
 
 Attribute DictionaryAttributeStorage::GetValue(const StrAttribute &name) const {
-  if (size_ > 0) {
-    size_t left = 0;
-    size_t right = size_ - 1;
-    size_t mid = 0;
-    while (left <= right) {
-      mid = (left + right) / 2;
-      if (data_[mid].name() == name) {
-        return data_[mid].value();
-      } else if (data_[mid].name() < name) {
-        left = mid + 1;
-      } else {
-        right = mid - 1;
-      }
+  size_t left = 0;
+  size_t right = size_;
+  while (left < right) {
+    size_t mid = left + (right - left) / 2;
+    if (data_[mid].name() == name) {
+      return data_[mid].value();
+    } else if (data_[mid].name() < name) {
+      left = mid + 1;
+    } else {
+      right = mid;
     }
   }
   return nullptr;

paddle/ir/dialect.cc

@@ -31,4 +31,8 @@ void Dialect::RegisterAttribute(ir::AbstractAttribute &&abstract_attribute) {
   this->ir_context()->RegisterAbstractAttribute(
       new_abstract_attribute->type_id(), new_abstract_attribute);
 }
+
+void Dialect::RegisterOp(const std::string &name, OpInfoImpl *op_info) {
+  this->ir_context()->RegisterOpInfo(name, op_info);
+}
 }  // namespace ir

paddle/ir/dialect.h

@@ -16,6 +16,7 @@
 
 #include "paddle/ir/attribute_base.h"
 #include "paddle/ir/ir_context.h"
+#include "paddle/ir/op_info_impl.h"
 #include "paddle/ir/type_base.h"
 
 namespace ir {
@@ -45,17 +46,19 @@ class Dialect {
     (void)std::initializer_list<int>{0, (RegisterType<Args>(), 0)...};
   }
 
-  ///
-  /// \brief Register type of class T.
-  ///
   template <typename T>
   void RegisterType() {
     VLOG(4) << "Type registered into Dialect. --->";
-    ir::AbstractType *abstract_type =
-        new ir::AbstractType(std::move(ir::AbstractType::get<T>(*this)));
-    this->ir_context()->RegisterAbstractType(ir::TypeId::get<T>(),
-                                             abstract_type);
-    ir::TypeManager::RegisterType<T>(this->ir_context());
+    // if (this->ir_context()->registed_abstract_type().count(
+    //         ir::TypeId::get<T>()) == 0) {
+    if (this->ir_context()->GetRegisteredAbstractType(ir::TypeId::get<T>()) ==
+        nullptr) {
+      ir::AbstractType *abstract_type =
+          new ir::AbstractType(std::move(ir::AbstractType::get<T>(*this)));
+      this->ir_context()->RegisterAbstractType(ir::TypeId::get<T>(),
+                                               abstract_type);
+      ir::TypeManager::RegisterType<T>(this->ir_context());
+    }
     VLOG(4) << "----------------------------------";
   }
 
@@ -78,24 +81,42 @@ class Dialect {
     (void)std::initializer_list<int>{0, (RegisterAttribute<Args>(), 0)...};
   }
 
-  ///
-  /// \brief Register attribute of class T.
-  ///
   template <typename T>
   void RegisterAttribute() {
     VLOG(4) << "Attribute registered into Dialect. --->";
-    ir::AbstractAttribute *abstract_attribute = new ir::AbstractAttribute(
-        std::move(ir::AbstractAttribute::get<T>(*this)));
-    this->ir_context()->RegisterAbstractAttribute(ir::TypeId::get<T>(),
-                                                  abstract_attribute);
-    ir::AttributeManager::RegisterAttribute<T>(this->ir_context());
+    if (this->ir_context()->GetRegisteredAbstractAttribute(
+            ir::TypeId::get<T>()) == nullptr) {
+      ir::AbstractAttribute *abstract_attribute = new ir::AbstractAttribute(
+          std::move(ir::AbstractAttribute::get<T>(*this)));
+      this->ir_context()->RegisterAbstractAttribute(ir::TypeId::get<T>(),
+                                                    abstract_attribute);
+      ir::AttributeManager::RegisterAttribute<T>(this->ir_context());
+    }
     VLOG(4) << "----------------------------------";
   }
 
+  void RegisterAttribute(ir::AbstractAttribute &&abstract_attribute);
+
   ///
-  /// \brief Register abstract_attribute into context.
+  /// \brief Register Operation methods.
   ///
-  void RegisterAttribute(ir::AbstractAttribute &&abstract_attribute);
+  template <typename... Args>
+  void RegisterOps() {
+    (void)std::initializer_list<int>{0, (RegisterOp<Args>(), 0)...};
+  }
+
+  template <typename ConcertOp>
+  void RegisterOp() {
+    std::string name = this->name() + "." + std::string(ConcertOp::name());
+    VLOG(4) << "Op " << name << " registered into Dialect. --->";
+    if (this->ir_context()->GetRegisteredOpInfo(name) == nullptr) {
+      ir::OpInfoImpl *op_info = ir::OpInfoImpl::create<ConcertOp>(this);
+      this->ir_context()->RegisterOpInfo(name, op_info);
+    }
+    VLOG(4) << "----------------------------------";
+  }
+
+  void RegisterOp(const std::string &name, OpInfoImpl *op_info);
 
  private:
   std::string name_;

paddle/ir/ir_context.cc

@@ -20,6 +20,7 @@
 #include "paddle/ir/builtin_dialect.h"
 #include "paddle/ir/builtin_type.h"
 #include "paddle/ir/dialect.h"
+#include "paddle/ir/op_info_impl.h"
 #include "paddle/ir/spin_lock.h"
 #include "paddle/ir/type_base.h"
 
@@ -46,6 +47,11 @@ class IrContextImpl {
       delete dialect_map.second;
     }
     registed_dialect_.clear();
+
+    for (auto &op_map : registed_op_infos_) {
+      op_map.second->destroy();
+    }
+    registed_op_infos_.clear();
   }
 
   void RegisterAbstractType(ir::TypeId type_id, AbstractType *abstract_type) {
@@ -93,6 +99,25 @@ class IrContextImpl {
     return nullptr;
   }
 
+  void RegisterOpInfo(const std::string &name, OpInfoImpl *opinfo) {
+    std::lock_guard<ir::SpinLock> guard(registed_op_infos_lock_);
+    VLOG(4) << "Register an operation of: [Name=" << name
+            << ", OpInfoImpl ptr=" << opinfo << "].";
+    registed_op_infos_.emplace(name, opinfo);
+  }
+
+  OpInfoImpl *GetOpInfo(const std::string &name) {
+    std::lock_guard<ir::SpinLock> guard(registed_op_infos_lock_);
+    auto iter = registed_op_infos_.find(name);
+    if (iter != registed_op_infos_.end()) {
+      VLOG(4) << "Fonund a cached operation of: [name=" << name
+              << ", OpInfoImpl ptr=" << iter->second << "].";
+      return iter->second;
+    }
+    LOG(WARNING) << "No cache found operation of: [Name=" << name << "].";
+    return nullptr;
+  }
+
   void RegisterDialect(std::string name, Dialect *dialect) {
     std::lock_guard<ir::SpinLock> guard(registed_dialect_lock_);
     VLOG(4) << "Register a dialect of: [name=" << name
@@ -135,6 +160,10 @@ class IrContextImpl {
   std::unordered_map<std::string, Dialect *> registed_dialect_;
   ir::SpinLock registed_dialect_lock_;
 
+  // The Op registered in the context.
+  std::unordered_map<std::string, OpInfoImpl *> registed_op_infos_;
+  ir::SpinLock registed_op_infos_lock_;
+
   ir::SpinLock destructor_lock_;
 };
 
@@ -165,9 +194,12 @@ StorageManager &IrContext::type_storage_manager() {
   return impl().registed_type_storage_manager_;
 }
 
-std::unordered_map<TypeId, AbstractType *>
-    &IrContext::registed_abstracted_type() {
-  return impl().registed_abstract_types_;
+AbstractType *IrContext::GetRegisteredAbstractType(TypeId id) {
+  auto search = impl().registed_abstract_types_.find(id);
+  if (search != impl().registed_abstract_types_.end()) {
+    return search->second;
+  }
+  return nullptr;
 }
 
 void IrContext::RegisterAbstractAttribute(
@@ -179,9 +211,12 @@ StorageManager &IrContext::attribute_storage_manager() {
   return impl().registed_attribute_storage_manager_;
 }
 
-std::unordered_map<TypeId, AbstractAttribute *>
-    &IrContext::registed_abstracted_attribute() {
-  return impl().registed_abstract_attributes_;
+AbstractAttribute *IrContext::GetRegisteredAbstractAttribute(TypeId id) {
+  auto search = impl().registed_abstract_attributes_.find(id);
+  if (search != impl().registed_abstract_attributes_.end()) {
+    return search->second;
+  }
+  return nullptr;
 }
 
 Dialect *IrContext::GetOrRegisterDialect(
@@ -216,6 +251,17 @@ Dialect *IrContext::GetRegisteredDialect(const std::string &dialect_name) {
   return nullptr;
 }
 
+OpInfoImpl *IrContext::GetRegisteredOpInfo(const std::string &name) {
+  OpInfoImpl *rtn = impl().GetOpInfo(name);
+  return rtn ? rtn : nullptr;
+}
+
+void IrContext::RegisterOpInfo(const std::string &name, OpInfoImpl *opinfo) {
+  if (impl().GetOpInfo(name) == nullptr) {
+    impl().RegisterOpInfo(name, opinfo);
+  }
+}
+
 const AbstractType &AbstractType::lookup(TypeId type_id, IrContext *ctx) {
   auto &impl = ctx->impl();
   AbstractType *abstract_type = impl.GetAbstractType(type_id);

paddle/ir/ir_context.h

@@ -26,6 +26,7 @@ class AbstractType;
 class AbstractAttribute;
 class TypeId;
 class Dialect;
+class OpInfoImpl;
 
 ///
 /// \brief IrContext is a global parameterless class used to store and manage
@@ -47,7 +48,7 @@ class IrContext {
   IrContextImpl &impl() { return *impl_; }
 
   ///
-  /// \brief Register an AbstractType to IrContext
+  /// \brief Register an AbstractType to IrContext.
   ///
   /// \param type_id The type id of the AbstractType.
   /// \param abstract_type AbstractType* provided by user.
@@ -64,13 +65,9 @@ class IrContext {
   StorageManager &type_storage_manager();
 
   ///
-  /// \brief Returns the storage uniquer used for constructing TypeStorage
-  /// instances.
-  ///
-  /// \return The storage uniquer used for constructing TypeStorage
-  /// instances.
+  /// \brief Get registered AbstractType from IrContext.
   ///
-  std::unordered_map<TypeId, AbstractType *> &registed_abstracted_type();
+  AbstractType *GetRegisteredAbstractType(TypeId id);
 
   ///
   /// \brief Register an AbstractAttribute to IrContext
@@ -91,14 +88,16 @@ class IrContext {
   StorageManager &attribute_storage_manager();
 
   ///
-  /// \brief Returns the storage uniquer used for constructing AttributeStorage
-  /// instances.
+  /// \brief Get registered AbstractAttribute from IrContext.
   ///
-  /// \return The storage uniquer used for constructing AttributeStorage
-  /// instances.
+  AbstractAttribute *GetRegisteredAbstractAttribute(TypeId id);
+
+  ///
+  /// \brief Get or register operaiton.
   ///
-  std::unordered_map<TypeId, AbstractAttribute *>
-      &registed_abstracted_attribute();
+  void RegisterOpInfo(const std::string &name, OpInfoImpl *opinfo);
+
+  OpInfoImpl *GetRegisteredOpInfo(const std::string &name);
 
   ///
   /// \brief Get the dialect of the DialectT class in the context, ff not found,

paddle/ir/op_base.h

@@ -15,23 +15,59 @@
 #pragma once
 
 #include "paddle/ir/operation.h"
+#include "paddle/ir/utils.h"
 
 namespace ir {
 class OpBase {
  public:
-  Operation *operation() { return operation_; }
+  explicit OpBase(const Operation *operation) : operation_(operation) {}
 
-  explicit operator bool() { return operation() != nullptr; }
+  const Operation *operation() const { return operation_; }
 
-  operator Operation *() const { return operation_; }
+  explicit operator bool() const { return operation() != nullptr; }
 
-  Operation *operator->() const { return operation_; }
+  operator const Operation *() const { return operation_; }
 
- protected:
-  explicit OpBase(Operation *operation) : operation_(operation) {}
+  const Operation *operator->() const { return operation_; }
 
  private:
-  Operation *operation_;
+  const Operation *operation_;  // Not owned
+};
+
+///
+/// \brief OpTrait
+///
+template <class ConcreteTrait>
+class OpTraitBase : public OpBase {
+ public:
+  explicit OpTraitBase(const Operation *op) : OpBase(op) {}
+
+  static TypeId GetTraitId() { return TypeId::get<ConcreteTrait>(); }
+};
+
+///
+/// \brief OpInterface
+///
+template <typename ConcreteInterface>
+class OpInterfaceBase : public OpBase {
+ public:
+  // explicit OpInterfaceBase(Operation *op) : OpBase(op) {}
+
+  explicit OpInterfaceBase(const Operation *op) : OpBase(op) {}
+
+  static TypeId GetInterfaceId() { return TypeId::get<ConcreteInterface>(); }
+};
+
+template <typename ConcreteOp, class... TraitOrInterface>
+class Op : public OpBase {
+ public:
+  using OpBase::OpBase;
+
+  using TraitList =
+      typename Filter<OpTraitBase, std::tuple<TraitOrInterface...>>::Type;
+
+  using InterfaceList =
+      typename Filter<OpInterfaceBase, std::tuple<TraitOrInterface...>>::Type;
 };
 
 }  // namespace ir