[IR] Type system stage3: add class Dialect (#50959)

* add dialect

* add some interface for dialect

* add some dialect interfaces for class Type

* set WITH_NEWIR=OFF

* refine code by comment

* polish code

* refine include style

* refine log for debug

Author: zhangbo9674

paddle/ir/builtin_dialect.cc

@@ -0,0 +1,29 @@
+// 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 "paddle/ir/builtin_dialect.h"
+#include "paddle/ir/builtin_type.h"
+
+namespace ir {
+BuiltinDialect::BuiltinDialect(ir::IrContext *context)
+    : ir::Dialect(name(), context, ir::TypeId::get<BuiltinDialect>()) {
+  initialize();
+}
+
+void BuiltinDialect::initialize() {
+  // Register all built-in types defined in builtin_type.h.
+  RegisterTypes<GET_BUILT_IN_TYPE_LIST>();
+}
+
+}  // namespace ir

paddle/ir/builtin_dialect.h

@@ -0,0 +1,40 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ir/dialect.h"
+
+namespace ir {
+///
+/// \brief Built-in Dialect: automatically registered into global IrContext,
+/// all built-in types defined in builtin_type.h will be registered in this
+/// Dialect.
+///
+class BuiltinDialect : public ir::Dialect {
+ public:
+  explicit BuiltinDialect(ir::IrContext *context);
+  ///
+  /// \brief Each Dialect needs to provide a name function to return the name of
+  /// the Dialect.
+  ///
+  /// \return The name of this Dialect.
+  ///
+  static const char *name() { return "builtin"; }
+
+ private:
+  void initialize();
+};
+
+}  // namespace ir

paddle/ir/builtin_type.h

@@ -17,9 +17,18 @@
 #include "paddle/ir/type.h"
 
 namespace ir {
+///
+/// \brief This macro is used to get a list of all built-in types in this file.
+///
+#define GET_BUILT_IN_TYPE_LIST ir::Float32Type, ir::Int32Type
+
 ///
 /// \brief Definitions of built-in type classes. The built-in type object get
-/// method is as follows: Type fp32 = Float32Type::get(ctx);
+/// method is as follows:
+/// \code{cpp}
+///   ir::IrContext *ctx = ir::IrContext::Instance();
+///   Type fp32 = Float32Type::get(ctx);
+/// \endcode
 ///
 class Float32Type : public ir::Type {
  public:

paddle/ir/dialect.cc

@@ -0,0 +1,28 @@
+// 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 "paddle/ir/dialect.h"
+
+namespace ir {
+Dialect::Dialect(std::string name, ir::IrContext *context, ir::TypeId id)
+    : name_(std::move(name)), context_(context), id_(id) {}
+
+void Dialect::RegisterType(ir::AbstractType &&abstract_type) {
+  ir::AbstractType *new_abstract_type =
+      new ir::AbstractType(std::move(abstract_type));
+  this->ir_context()->RegisterAbstractType(new_abstract_type->type_id(),
+                                           new_abstract_type);
+}
+
+}  // namespace ir

paddle/ir/dialect.h

@@ -0,0 +1,77 @@
+// 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.
+
+#pragma once
+
+#include "paddle/ir/ir_context.h"
+#include "paddle/ir/type_base.h"
+
+namespace ir {
+///
+/// \brief Dialect can basically be understood as a namespace. In Dialect, we
+/// can define a series of types, operations, etc. An instance of the dialect
+/// object will be loaded into the global IrContext. Specific compilers only
+/// need to combine existing dialects and add their own extensions or
+/// customizations.
+///
+class Dialect {
+ public:
+  Dialect(std::string name, ir::IrContext *context, ir::TypeId id);
+
+  const std::string &name() const { return name_; }
+
+  ir::IrContext *ir_context() const { return context_; }
+
+  ir::TypeId id() const { return id_; }
+
+  ///
+  /// \brief Register all types contained in the template parameter Args.
+  /// To register only one Type, you can use the RegisterType template function.
+  ///
+  template <typename... Args>
+  void RegisterTypes() {
+    (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());
+    VLOG(4) << "----------------------------------";
+  }
+
+  ///
+  /// \brief Register abstract_type into context.
+  /// NOTE: It's not recommended to use this interface directly. This interface
+  /// only registers abstract_type. To register TypeStorage into context, you
+  /// need to call ir::TypeManager::RegisterType<T>() additionally,
+  /// RegisterType<T>() is recommended to use.
+  ///
+  void RegisterType(ir::AbstractType &&abstract_type);
+
+ private:
+  std::string name_;
+
+  ir::IrContext *context_;  // not owned
+
+  ir::TypeId id_;
+};
+}  // namespace ir

paddle/ir/ir_context.cc

@@ -12,61 +12,95 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "paddle/ir/ir_context.h"
+
 #include <unordered_map>
 
+#include "paddle/ir/builtin_dialect.h"
 #include "paddle/ir/builtin_type.h"
-#include "paddle/ir/ir_context.h"
+#include "paddle/ir/dialect.h"
 #include "paddle/ir/spin_lock.h"
 #include "paddle/ir/type_base.h"
 
 namespace ir {
-// The implementation class of the IrContext class
+// The implementation class of the IrContext class, cache registered
+// AbstractType, TypeStorage, Dialect.
 class IrContextImpl {
  public:
   IrContextImpl() {}
 
   ~IrContextImpl() {
-    std::lock_guard<ir::SpinLock> guard(registed_abstract_types_lock_);
-    for (auto abstract_type_map : registed_abstract_types_) {
+    std::lock_guard<ir::SpinLock> guard(destructor_lock_);
+    for (auto &abstract_type_map : registed_abstract_types_) {
       delete abstract_type_map.second;
     }
     registed_abstract_types_.clear();
+
+    for (auto &dialect_map : registed_dialect_) {
+      delete dialect_map.second;
+    }
+    registed_dialect_.clear();
   }
 
   void RegisterAbstractType(ir::TypeId type_id, AbstractType *abstract_type) {
     std::lock_guard<ir::SpinLock> guard(registed_abstract_types_lock_);
-    VLOG(4) << "IrContext register an abstract_type of: [TypeId_hash="
+    VLOG(4) << "Register an abstract_type of: [TypeId_hash="
             << std::hash<ir::TypeId>()(type_id)
             << ", AbstractType_ptr=" << abstract_type << "].";
     registed_abstract_types_.emplace(type_id, abstract_type);
   }
 
-  AbstractType *lookup(ir::TypeId type_id) {
+  AbstractType *GetAbstractType(ir::TypeId type_id) {
     std::lock_guard<ir::SpinLock> guard(registed_abstract_types_lock_);
     auto iter = registed_abstract_types_.find(type_id);
-    if (iter == registed_abstract_types_.end()) {
-      VLOG(4) << "IrContext not fonund cached abstract_type of: [TypeId_hash="
-              << std::hash<ir::TypeId>()(type_id) << "].";
-      return nullptr;
-    } else {
-      VLOG(4) << "IrContext fonund a cached abstract_type of: [TypeId_hash="
+    if (iter != registed_abstract_types_.end()) {
+      VLOG(4) << "Fonund a cached abstract_type of: [TypeId_hash="
               << std::hash<ir::TypeId>()(type_id)
               << ", AbstractType_ptr=" << iter->second << "].";
       return iter->second;
     }
+    LOG(WARNING) << "No cache found abstract_type of: [TypeId_hash="
+                 << std::hash<ir::TypeId>()(type_id) << "].";
+    return nullptr;
   }
 
-  ir::SpinLock registed_abstract_types_lock_;
+  void RegisterDialect(std::string name, Dialect *dialect) {
+    std::lock_guard<ir::SpinLock> guard(registed_dialect_lock_);
+    VLOG(4) << "Register a dialect of: [name=" << name
+            << ", dialect_ptr=" << dialect << "].";
+    registed_dialect_.emplace(name, dialect);
+  }
+
+  Dialect *GetDialect(std::string name) {
+    std::lock_guard<ir::SpinLock> guard(registed_dialect_lock_);
+    auto iter = registed_dialect_.find(name);
+    if (iter != registed_dialect_.end()) {
+      VLOG(4) << "Fonund a cached dialect of: [name=" << name
+              << ", dialect_ptr=" << iter->second << "].";
+      return iter->second;
+    }
+    LOG(WARNING) << "No cache fonund dialect of: [name=" << name << "].";
+    return nullptr;
+  }
 
   // Cached AbstractType instances.
   std::unordered_map<TypeId, AbstractType *> registed_abstract_types_;
 
+  ir::SpinLock registed_abstract_types_lock_;
+
   // TypeStorage uniquer and cache instances.
   StorageManager registed_storage_manager_;
 
-  // Some built-in type.
+  // The dialcet registered in the context.
+  std::unordered_map<std::string, Dialect *> registed_dialect_;
+
+  ir::SpinLock registed_dialect_lock_;
+
+  // Some built-in types.
   Float32Type fp32_type;
   Int32Type int32_type;
+
+  ir::SpinLock destructor_lock_;
 };
 
 IrContext *IrContext::Instance() {
@@ -75,13 +109,12 @@ IrContext *IrContext::Instance() {
 }
 
 IrContext::IrContext() : impl_(new IrContextImpl()) {
-  VLOG(4) << "IrContext register built-in type...";
-  REGISTER_TYPE_2_IRCONTEXT(Float32Type, this);
+  VLOG(4) << "BuiltinDialect registered into IrContext. ===>";
+  GetOrRegisterDialect<BuiltinDialect>();
+  VLOG(4) << "==============================================";
+
   impl_->fp32_type = TypeManager::get<Float32Type>(this);
-  VLOG(4) << "Float32Type registration complete";
-  REGISTER_TYPE_2_IRCONTEXT(Int32Type, this);
   impl_->int32_type = TypeManager::get<Int32Type>(this);
-  VLOG(4) << "Int32Type registration complete";
 }
 
 void IrContext::RegisterAbstractType(ir::TypeId type_id,
@@ -98,12 +131,41 @@ std::unordered_map<TypeId, AbstractType *>
   return impl().registed_abstract_types_;
 }
 
-const AbstractType &AbstractType::lookup(TypeId type_id, IrContext *ctx) {
-  VLOG(4) << "Lookup abstract type [TypeId_hash="
-          << std::hash<ir::TypeId>()(type_id) << "] from IrContext [ptr=" << ctx
+Dialect *IrContext::GetOrRegisterDialect(
+    std::string dialect_name, std::function<Dialect *()> constructor) {
+  VLOG(4) << "Try to get or register a Dialect of: [name=" << dialect_name
           << "].";
+  Dialect *dialect = impl().GetDialect(dialect_name);
+  if (dialect == nullptr) {
+    VLOG(4) << "Create and register a new Dialect of: [name=" << dialect_name
+            << "].";
+    dialect = constructor();
+    impl().RegisterDialect(dialect_name, dialect);
+  }
+  return dialect;
+}
+
+std::vector<Dialect *> IrContext::GetRegisteredDialects() {
+  std::vector<Dialect *> result;
+  for (auto dialect_map : impl().registed_dialect_) {
+    result.push_back(dialect_map.second);
+  }
+  return result;
+}
+
+Dialect *IrContext::GetRegisteredDialect(const std::string &dialect_name) {
+  for (auto dialect_map : impl().registed_dialect_) {
+    if (dialect_map.first == dialect_name) {
+      return dialect_map.second;
+    }
+  }
+  LOG(WARNING) << "No dialect registered for " << dialect_name;
+  return nullptr;
+}
+
+const AbstractType &AbstractType::lookup(TypeId type_id, IrContext *ctx) {
   auto &impl = ctx->impl();
-  AbstractType *abstract_type = impl.lookup(type_id);
+  AbstractType *abstract_type = impl.GetAbstractType(type_id);
   if (abstract_type) {
     return *abstract_type;
   } else {