[IR] Type system stage3: add class Dialect

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,38 @@
+// 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.
+  ///
+  static const std::string name() { return "builtin"; }
+
+ private:
+  void initialize();
+};
+
+}  // namespace ir

paddle/ir/builtin_type.h

@@ -17,6 +17,11 @@
 #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);

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_(name), context_(context), id_(id) {}
+
+void Dialect::RegisterType(ir::TypeId type_id,
+                           ir::AbstractType &&abstract_type) {
+  ir::AbstractType *new_abstract_type =
+      new ir::AbstractType(std::move(abstract_type));
+  this->ir_context()->RegisterAbstractType(type_id, new_abstract_type);
+}
+
+}  // namespace ir

paddle/ir/dialect.h

@@ -0,0 +1,75 @@
+// 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);
+
+  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() {
+    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());
+  }
+
+  ///
+  /// \brief Register type_id and 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::TypeId type_id, ir::AbstractType &&abstract_type);
+
+ private:
+  std::string name_;
+
+  ir::IrContext *context_;  // not owned
+
+  ir::TypeId id_;
+};
+}  // namespace ir

paddle/ir/ir_context.cc

@@ -14,7 +14,9 @@
 
 #include <unordered_map>
 
+#include "paddle/ir/builtin_dialect.h"
 #include "paddle/ir/builtin_type.h"
+#include "paddle/ir/dialect.h"
 #include "paddle/ir/ir_context.h"
 #include "paddle/ir/spin_lock.h"
 #include "paddle/ir/type_base.h"
@@ -26,11 +28,16 @@ class IrContextImpl {
   IrContextImpl() {}
 
   ~IrContextImpl() {
-    std::lock_guard<ir::SpinLock> guard(registed_abstract_types_lock_);
+    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) {
@@ -41,12 +48,13 @@ class IrContextImpl {
     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) << "].";
+      LOG(WARNING)
+          << "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="
@@ -56,6 +64,27 @@ class IrContextImpl {
     }
   }
 
+  void RegisterDialect(std::string name, Dialect *dialect) {
+    std::lock_guard<ir::SpinLock> guard(registed_dialect_lock_);
+    VLOG(4) << "IrContext 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()) {
+      LOG(WARNING) << "IrContext not fonund cached dialect of: [name=" << name
+                   << "].";
+      return nullptr;
+    } else {
+      VLOG(4) << "IrContext fonund a cached dialect of: [name=" << name
+              << ", dialect_ptr=" << iter->second << "].";
+      return iter->second;
+    }
+  }
+
   ir::SpinLock registed_abstract_types_lock_;
 
   // Cached AbstractType instances.
@@ -64,9 +93,16 @@ class IrContextImpl {
   // TypeStorage uniquer and cache instances.
   StorageManager registed_storage_manager_;
 
+  ir::SpinLock registed_dialect_lock_;
+
+  // The dialcet registered in the context.
+  std::unordered_map<std::string, Dialect *> registed_dialect_;
+
   // Some built-in type.
   Float32Type fp32_type;
   Int32Type int32_type;
+
+  ir::SpinLock destructor_lock_;
 };
 
 IrContext *IrContext::Instance() {
@@ -75,13 +111,11 @@ 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>();
+
   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 +132,43 @@ std::unordered_map<TypeId, AbstractType *>
   return impl().registed_abstract_types_;
 }
 
+Dialect *IrContext::GetOrRegisterDialect(
+    std::string dialect_name, std::function<Dialect *()> constructor) {
+  VLOG(4) << "IrContext get or register a dialect of: [name=" << dialect_name
+          << "].";
+  Dialect *dialect = impl().GetDialect(dialect_name);
+  if (dialect == nullptr) {
+    VLOG(4) << "Not fonund cached dialect, create and register a new dialect.";
+    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(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) {
   VLOG(4) << "Lookup abstract type [TypeId_hash="
           << std::hash<ir::TypeId>()(type_id) << "] from IrContext [ptr=" << 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 {