From d8fcd158d30c95da6e2f15817aae9356450e430b Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Mon, 27 Feb 2023 05:36:13 +0000
Subject: [PATCH 1/8] add dialect

---
 CMakeLists.txt               |  2 +-
 paddle/ir/builtin_dialect.cc | 29 +++++++++++++++
 paddle/ir/builtin_dialect.h  | 35 ++++++++++++++++++
 paddle/ir/builtin_type.h     |  5 +++
 paddle/ir/dialect.cc         | 28 ++++++++++++++
 paddle/ir/dialect.h          | 71 ++++++++++++++++++++++++++++++++++++
 paddle/ir/ir_context.cc      | 61 +++++++++++++++++++++++++++----
 paddle/ir/ir_context.h       | 32 ++++++++++++++++
 paddle/ir/storage_manager.cc | 12 +++---
 paddle/ir/storage_manager.h  |  4 +-
 paddle/ir/tests/type_test.cc | 33 ++++++++++++-----
 paddle/ir/type_base.h        | 44 +++++++++++++++++-----
 12 files changed, 319 insertions(+), 37 deletions(-)
 create mode 100644 paddle/ir/builtin_dialect.cc
 create mode 100644 paddle/ir/builtin_dialect.h
 create mode 100644 paddle/ir/dialect.cc
 create mode 100644 paddle/ir/dialect.h

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 2b8caf7addd16..8893723c9ef57 100755
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -315,7 +315,7 @@ option(WITH_CUDNN_FRONTEND
        "Compile with CUDNN Frontend API support (experimental)" OFF)
 option(WITH_CUDNN_DSO "Compile PaddlePaddle with cuDNN dynamic-link libraries"
        OFF)
-option(WITH_NEWIR "Compile PaddlePaddle with NEWIR" OFF)
+option(WITH_NEWIR "Compile PaddlePaddle with NEWIR" ON)
 
 if(WITH_RECORD_BUILDTIME)
   set_property(
diff --git a/paddle/ir/builtin_dialect.cc b/paddle/ir/builtin_dialect.cc
new file mode 100644
index 0000000000000..5c797c4214c35
--- /dev/null
+++ b/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
diff --git a/paddle/ir/builtin_dialect.h b/paddle/ir/builtin_dialect.h
new file mode 100644
index 0000000000000..ca2c6b160ed67
--- /dev/null
+++ b/paddle/ir/builtin_dialect.h
@@ -0,0 +1,35 @@
+// 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);
+
+  static const std::string name() { return "builtin"; }
+
+ private:
+  void initialize();
+};
+
+}  // namespace ir
diff --git a/paddle/ir/builtin_type.h b/paddle/ir/builtin_type.h
index 77159794bf11e..0636b5234a548 100644
--- a/paddle/ir/builtin_type.h
+++ b/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);
diff --git a/paddle/ir/dialect.cc b/paddle/ir/dialect.cc
new file mode 100644
index 0000000000000..6a2055c411817
--- /dev/null
+++ b/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
diff --git a/paddle/ir/dialect.h b/paddle/ir/dialect.h
new file mode 100644
index 0000000000000..241117a6c853c
--- /dev/null
+++ b/paddle/ir/dialect.h
@@ -0,0 +1,71 @@
+// 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 is a group of types, an instance of the dialect object will
+/// be loaded into the global IrContext.
+///
+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.
+  ///
+  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
diff --git a/paddle/ir/ir_context.cc b/paddle/ir/ir_context.cc
index 6aed600903438..bb520dd35370a 100644
--- a/paddle/ir/ir_context.cc
+++ b/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,7 +48,7 @@ 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()) {
@@ -56,6 +63,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()) {
+      VLOG(4) << "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 +92,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 +110,10 @@ 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 +130,25 @@ 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;
+}
+
 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 {
diff --git a/paddle/ir/ir_context.h b/paddle/ir/ir_context.h
index 146497e6c6c94..943be2df122eb 100644
--- a/paddle/ir/ir_context.h
+++ b/paddle/ir/ir_context.h
@@ -15,6 +15,7 @@
 #pragma once
 
 #include <glog/logging.h>
+#include <functional>
 #include <memory>
 #include <unordered_map>
 
@@ -23,6 +24,7 @@ class IrContextImpl;
 class StorageManager;
 class AbstractType;
 class TypeId;
+class Dialect;
 
 ///
 /// \brief IrContext is a global parameterless class used to store and manage
@@ -69,6 +71,36 @@ class IrContext {
   ///
   std::unordered_map<TypeId, AbstractType *> &registed_abstracted_type();
 
+  ///
+  /// \brief Get the dialect of the DialectT class in the context, ff not found,
+  /// create and register to context.
+  ///
+  /// \param DialectT The Dialect class that needs to be found or register.
+  ///
+  /// \return The dialect of the DialectT class in the context.
+  ///
+  template <typename DialectT>
+  DialectT *GetOrRegisterDialect() {
+    return static_cast<DialectT *>(
+        GetOrRegisterDialect(DialectT::name(), [this]() {
+          DialectT *dialect = new DialectT(this);
+          return dialect;
+        }));
+  }
+
+  ///
+  /// \brief Get the dialect of the DialectT class in the context, ff not found,
+  /// create and register to context.
+  ///
+  /// \param dialect_name The dialect name.
+  /// \param dialect_id The TypeId of the dialect.
+  /// \param constructor The dialect constructor.
+  ///
+  /// \return The dialect named "dialect_name" in the context.
+  ///
+  Dialect *GetOrRegisterDialect(std::string dialect_name,
+                                std::function<Dialect *()> constructor);
+
   IrContext(const IrContext &) = delete;
 
   void operator=(const IrContext &) = delete;
diff --git a/paddle/ir/storage_manager.cc b/paddle/ir/storage_manager.cc
index 991077e8777c4..91bc1ba3d360c 100644
--- a/paddle/ir/storage_manager.cc
+++ b/paddle/ir/storage_manager.cc
@@ -83,12 +83,12 @@ StorageManager::StorageBase *StorageManager::GetParametricStorageTypeImpl(
 
 StorageManager::StorageBase *StorageManager::GetParameterlessStorageTypeImpl(
     TypeId type_id) {
-  std::lock_guard<ir::SpinLock> guard(parameterless_instances_lock_);
+  std::lock_guard<ir::SpinLock> guard(parameterless_instance_lock_);
   VLOG(4) << "StorageManager get parameterless storage of: [TypeId_hash="
           << std::hash<ir::TypeId>()(type_id) << "].";
-  if (parameterless_instances_.find(type_id) == parameterless_instances_.end())
+  if (parameterless_instance_.find(type_id) == parameterless_instance_.end())
     throw("TypeId not found in IrContext.");
-  StorageBase *parameterless_instance = parameterless_instances_[type_id];
+  StorageBase *parameterless_instance = parameterless_instance_[type_id];
   return parameterless_instance;
 }
 
@@ -102,12 +102,12 @@ void StorageManager::RegisterParametricStorageTypeImpl(TypeId type_id) {
 
 void StorageManager::RegisterParameterlessStorageTypeImpl(
     TypeId type_id, std::function<StorageBase *()> constructor) {
-  std::lock_guard<ir::SpinLock> guard(parameterless_instances_lock_);
+  std::lock_guard<ir::SpinLock> guard(parameterless_instance_lock_);
   VLOG(4) << "StorageManager register parameterless storage of: [TypeId_hash="
           << std::hash<ir::TypeId>()(type_id) << "].";
-  if (parameterless_instances_.find(type_id) != parameterless_instances_.end())
+  if (parameterless_instance_.find(type_id) != parameterless_instance_.end())
     throw("storage class already registered");
-  parameterless_instances_.emplace(type_id, constructor());
+  parameterless_instance_.emplace(type_id, constructor());
 }
 
 }  // namespace ir
diff --git a/paddle/ir/storage_manager.h b/paddle/ir/storage_manager.h
index f94174586bc97..8b6c1a330e42b 100644
--- a/paddle/ir/storage_manager.h
+++ b/paddle/ir/storage_manager.h
@@ -141,9 +141,9 @@ class StorageManager {
   ir::SpinLock parametric_instance_lock_;
 
   // This map is a mapping between type id and parameterless type storage.
-  std::unordered_map<TypeId, StorageBase *> parameterless_instances_;
+  std::unordered_map<TypeId, StorageBase *> parameterless_instance_;
 
-  ir::SpinLock parameterless_instances_lock_;
+  ir::SpinLock parameterless_instance_lock_;
 };
 
 }  // namespace ir
diff --git a/paddle/ir/tests/type_test.cc b/paddle/ir/tests/type_test.cc
index 85deb51b694d5..0679953b57efb 100644
--- a/paddle/ir/tests/type_test.cc
+++ b/paddle/ir/tests/type_test.cc
@@ -16,6 +16,7 @@
 #include <unordered_map>
 
 #include "paddle/ir/builtin_type.h"
+#include "paddle/ir/dialect.h"
 #include "paddle/ir/ir_context.h"
 #include "paddle/ir/type_base.h"
 
@@ -39,20 +40,22 @@ TEST(type_test, type_id) {
   }
 }
 
-TEST(type_test, abstract_type) {
+TEST(type_test, type_base) {
   class TypeA {};
 
-  ir::TypeId a_id = ir::TypeId::get<TypeA>();
-  ir::AbstractType abstract_type_a = ir::AbstractType::get(a_id);
-
-  EXPECT_EQ(abstract_type_a.type_id(), a_id);
-}
+  struct FakeDialect : ir::Dialect {
+    explicit FakeDialect(ir::IrContext *context)
+        : ir::Dialect(name(), context, ir::TypeId::get<FakeDialect>()) {}
+    static const std::string name() { return "fake"; }
+  };
 
-TEST(type_test, type_storage) {
-  class TypeA {};
+  ir::IrContext *ctx = ir::IrContext::Instance();
+  ir::Dialect *fake_dialect = ctx->GetOrRegisterDialect<FakeDialect>();
 
   ir::TypeId a_id = ir::TypeId::get<TypeA>();
-  ir::AbstractType abstract_type_a = ir::AbstractType::get(a_id);
+  ir::AbstractType abstract_type_a = ir::AbstractType::get(a_id, *fake_dialect);
+
+  EXPECT_EQ(abstract_type_a.type_id(), a_id);
 
   ir::TypeStorage storage_a(&abstract_type_a);
 
@@ -119,9 +122,19 @@ class IntegerType : public ir::Type {
   DECLARE_TYPE_UTILITY_FUNCTOR(IntegerType, IntegerTypeStorage);
 };
 
+struct IntegerDialect : ir::Dialect {
+  explicit IntegerDialect(ir::IrContext *context)
+      : ir::Dialect(name(), context, ir::TypeId::get<IntegerDialect>()) {
+    RegisterType<IntegerType>();
+  }
+  static const std::string name() { return "integer"; }
+};
+
 TEST(type_test, parameteric_type) {
   ir::IrContext *ctx = ir::IrContext::Instance();
-  REGISTER_TYPE_2_IRCONTEXT(IntegerType, ctx);
+
+  ctx->GetOrRegisterDialect<IntegerDialect>();
+
   ir::Type int1_1 = IntegerType::get(ctx, 1, 0);
   ir::Type int1_2 = IntegerType::get(ctx, 1, 0);
   EXPECT_EQ(int1_1 == int1_2, 1);
diff --git a/paddle/ir/type_base.h b/paddle/ir/type_base.h
index aa800498f6ead..582c37050920b 100644
--- a/paddle/ir/type_base.h
+++ b/paddle/ir/type_base.h
@@ -19,6 +19,8 @@
 #include "paddle/ir/type_id.h"
 
 namespace ir {
+class Dialect;
+
 ///
 /// \brief Abstract the properties and behaviors common to all Type classes into
 /// an AbstractType class. There are two types in Type system:
@@ -32,8 +34,21 @@ class AbstractType {
   /// \brief Construct an AbstractType by TypeId directly.
   ///
   /// \param type_id The type id of the AbstractType.
+  /// \param dialect The Dialect which the type registered to.
   ///
-  static AbstractType get(TypeId type_id) { return AbstractType(type_id); }
+  static AbstractType get(TypeId type_id, const Dialect &dialect) {
+    return AbstractType(type_id, dialect);
+  }
+
+  ///
+  /// \brief Construct an AbstractType by TypeId directly.
+  ///
+  /// \param dialect The Dialect which the type registered to.
+  ///
+  template <typename T>
+  static AbstractType get(const Dialect &dialect) {
+    return AbstractType(TypeId::get<T>(), dialect);
+  }
 
   ///
   /// \brief Returns the type id of the AbstractType.
@@ -42,6 +57,11 @@ class AbstractType {
   ///
   TypeId type_id() const { return type_id_; }
 
+  ///
+  /// \brief Get the dialect this type was registered to.
+  ///
+  Dialect &dialect() const { return const_cast<Dialect &>(dialect_); }
+
   ///
   /// \brief Find the AbstractType instance whose TypeId is type_id from
   /// IrContext.
@@ -58,10 +78,14 @@ class AbstractType {
   /// get method to obtain and manage the AstractType.
   ///
   /// \param type_id The type id of the AbstractType.
+  /// \param dialect The Dialect which the type registered to.
   ///
-  explicit AbstractType(TypeId type_id) : type_id_(type_id) {}
+  explicit AbstractType(TypeId type_id, const Dialect &dialect)
+      : type_id_(type_id), dialect_(dialect) {}
 
   TypeId type_id_;
+
+  const Dialect &dialect_;
 };
 
 struct TypeManager;
@@ -239,13 +263,13 @@ struct TypeManager {
 ///
 /// \brief This macro definition is used to register custom Type class.
 ///
-#define REGISTER_TYPE_2_IRCONTEXT(concrete_type, ir_context)               \
-  ir::AbstractType *abstract_type_##concrete_type = new ir::AbstractType(  \
-      std::move(ir::AbstractType::get(ir::TypeId::get<concrete_type>()))); \
-                                                                           \
-  ir_context->RegisterAbstractType(ir::TypeId::get<concrete_type>(),       \
-                                   abstract_type_##concrete_type);         \
-                                                                           \
-  ir::TypeManager::RegisterType<concrete_type>(ir_context);
+#define REGISTER_TYPE_2_IRCONTEXT(concrete_type, dialect)                 \
+  ir::AbstractType *abstract_type_##concrete_type = new ir::AbstractType( \
+      std::move(ir::AbstractType::get<concrete_type>(*dialect)));         \
+                                                                          \
+  dialect->ir_context()->RegisterAbstractType(                            \
+      ir::TypeId::get<concrete_type>(), abstract_type_##concrete_type);   \
+                                                                          \
+  ir::TypeManager::RegisterType<concrete_type>(dialect->ir_context());
 
 }  // namespace ir

From 376ad12cb366fce29aab926e4c1a91f1d3ef8bfc Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Mon, 27 Feb 2023 07:55:42 +0000
Subject: [PATCH 2/8] add some interface for dialect

---
 paddle/ir/builtin_dialect.h  |  5 ++++-
 paddle/ir/dialect.h          |  8 ++++++--
 paddle/ir/ir_context.cc      | 28 ++++++++++++++++++++++++----
 paddle/ir/ir_context.h       | 27 +++++++++++++++++++++++++++
 paddle/ir/tests/type_test.cc | 17 ++++++++++++++++-
 paddle/ir/type_base.h        |  2 ++
 6 files changed, 79 insertions(+), 8 deletions(-)

diff --git a/paddle/ir/builtin_dialect.h b/paddle/ir/builtin_dialect.h
index ca2c6b160ed67..25c22110c1d5a 100644
--- a/paddle/ir/builtin_dialect.h
+++ b/paddle/ir/builtin_dialect.h
@@ -25,7 +25,10 @@ namespace ir {
 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:
diff --git a/paddle/ir/dialect.h b/paddle/ir/dialect.h
index 241117a6c853c..b50db790a99bb 100644
--- a/paddle/ir/dialect.h
+++ b/paddle/ir/dialect.h
@@ -19,8 +19,11 @@
 
 namespace ir {
 ///
-/// \brief Dialect is a group of types, an instance of the dialect object will
-/// be loaded into the global IrContext.
+/// \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:
@@ -34,6 +37,7 @@ class Dialect {
 
   ///
   /// \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() {
diff --git a/paddle/ir/ir_context.cc b/paddle/ir/ir_context.cc
index bb520dd35370a..1a96758c11918 100644
--- a/paddle/ir/ir_context.cc
+++ b/paddle/ir/ir_context.cc
@@ -52,8 +52,9 @@ class IrContextImpl {
     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="
@@ -74,8 +75,8 @@ class IrContextImpl {
     std::lock_guard<ir::SpinLock> guard(registed_dialect_lock_);
     auto iter = registed_dialect_.find(name);
     if (iter == registed_dialect_.end()) {
-      VLOG(4) << "IrContext not fonund cached dialect of: [name=" << name
-              << "].";
+      LOG(WARNING) << "IrContext not fonund cached dialect of: [name=" << name
+                   << "].";
       return nullptr;
     } else {
       VLOG(4) << "IrContext fonund a cached dialect of: [name=" << name
@@ -112,6 +113,7 @@ IrContext *IrContext::Instance() {
 IrContext::IrContext() : impl_(new IrContextImpl()) {
   VLOG(4) << "BuiltinDialect registered into IrContext.";
   GetOrRegisterDialect<BuiltinDialect>();
+
   impl_->fp32_type = TypeManager::get<Float32Type>(this);
   impl_->int32_type = TypeManager::get<Int32Type>(this);
 }
@@ -143,6 +145,24 @@ Dialect *IrContext::GetOrRegisterDialect(
   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
diff --git a/paddle/ir/ir_context.h b/paddle/ir/ir_context.h
index 943be2df122eb..a7be7f6948fc0 100644
--- a/paddle/ir/ir_context.h
+++ b/paddle/ir/ir_context.h
@@ -101,6 +101,33 @@ class IrContext {
   Dialect *GetOrRegisterDialect(std::string dialect_name,
                                 std::function<Dialect *()> constructor);
 
+  ///
+  /// \brief Get the dialect list registered to the context.
+  ///
+  /// \return The dialect list registered to the context.
+  ///
+  std::vector<Dialect *> GetRegisteredDialects();
+
+  ///
+  /// \brief Get the dialect named "name" from the context.
+  ///
+  /// \param name The name of the dialect to be obtained.
+  ///
+  /// \return The dialect named "name" from the context.
+  ///
+  Dialect *GetRegisteredDialect(std::string dialect_name);
+
+  ///
+  /// \brief Get a registered dialect for the given dialect type T. The
+  /// Dialect must provide a static 'name' method.
+  ///
+  /// \return The registered dialect for the given dialect type T.
+  ///
+  template <typename T>
+  T *GetRegisteredDialect() {
+    return static_cast<T *>(GetRegisteredDialect(T::name()));
+  }
+
   IrContext(const IrContext &) = delete;
 
   void operator=(const IrContext &) = delete;
diff --git a/paddle/ir/tests/type_test.cc b/paddle/ir/tests/type_test.cc
index 0679953b57efb..acead54ac50ea 100644
--- a/paddle/ir/tests/type_test.cc
+++ b/paddle/ir/tests/type_test.cc
@@ -15,6 +15,7 @@
 #include <gtest/gtest.h>
 #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"
@@ -130,9 +131,10 @@ struct IntegerDialect : ir::Dialect {
   static const std::string name() { return "integer"; }
 };
 
-TEST(type_test, parameteric_type) {
+TEST(type_test, custom_type_dialect) {
   ir::IrContext *ctx = ir::IrContext::Instance();
 
+  // Test Custom Type and Dialect
   ctx->GetOrRegisterDialect<IntegerDialect>();
 
   ir::Type int1_1 = IntegerType::get(ctx, 1, 0);
@@ -141,4 +143,17 @@ TEST(type_test, parameteric_type) {
 
   ir::Type int8 = IntegerType::get(ctx, 8, 0);
   EXPECT_EQ(int8 == int1_2, 0);
+
+  // Test Dialect interface
+  std::vector<ir::Dialect *> dialect_list = ctx->GetRegisteredDialects();
+  EXPECT_EQ(dialect_list.size() == 3, 1);  // integer, builtin, fake
+
+  ir::Dialect *dialect_builtin1 = ctx->GetRegisteredDialect("builtin");
+  ir::Dialect *dialect_builtin2 =
+      ctx->GetRegisteredDialect<ir::BuiltinDialect>();
+  EXPECT_EQ(dialect_builtin1 == dialect_builtin2, 1);
+
+  ir::Dialect *dialect_integer1 = ctx->GetRegisteredDialect("integer");
+  ir::Dialect *dialect_integer2 = ctx->GetRegisteredDialect<IntegerDialect>();
+  EXPECT_EQ(dialect_integer1 == dialect_integer2, 1);
 }
diff --git a/paddle/ir/type_base.h b/paddle/ir/type_base.h
index 582c37050920b..c3c6a9a7cea7b 100644
--- a/paddle/ir/type_base.h
+++ b/paddle/ir/type_base.h
@@ -60,6 +60,8 @@ class AbstractType {
   ///
   /// \brief Get the dialect this type was registered to.
   ///
+  /// \return The dialect this type was registered to.
+  ///
   Dialect &dialect() const { return const_cast<Dialect &>(dialect_); }
 
   ///

From ca0f82712a189d43d10d1c5260d0b71412fc5f84 Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Mon, 27 Feb 2023 08:56:50 +0000
Subject: [PATCH 3/8] add some dialect interfaces for class Type

---
 paddle/ir/tests/type_test.cc |  5 +++++
 paddle/ir/type.cc            | 21 +++++++++++++++++++++
 paddle/ir/type.h             |  4 ++++
 3 files changed, 30 insertions(+)
 create mode 100644 paddle/ir/type.cc

diff --git a/paddle/ir/tests/type_test.cc b/paddle/ir/tests/type_test.cc
index acead54ac50ea..412014d19f096 100644
--- a/paddle/ir/tests/type_test.cc
+++ b/paddle/ir/tests/type_test.cc
@@ -19,6 +19,7 @@
 #include "paddle/ir/builtin_type.h"
 #include "paddle/ir/dialect.h"
 #include "paddle/ir/ir_context.h"
+#include "paddle/ir/type.h"
 #include "paddle/ir/type_base.h"
 
 TEST(type_test, type_id) {
@@ -145,6 +146,10 @@ TEST(type_test, custom_type_dialect) {
   EXPECT_EQ(int8 == int1_2, 0);
 
   // Test Dialect interface
+  EXPECT_EQ(ctx == int8.ir_context(), 1);
+
+  EXPECT_EQ(int8.dialect().id() == ir::TypeId::get<IntegerDialect>(), 1);
+
   std::vector<ir::Dialect *> dialect_list = ctx->GetRegisteredDialects();
   EXPECT_EQ(dialect_list.size() == 3, 1);  // integer, builtin, fake
 
diff --git a/paddle/ir/type.cc b/paddle/ir/type.cc
new file mode 100644
index 0000000000000..bde3194f8fd41
--- /dev/null
+++ b/paddle/ir/type.cc
@@ -0,0 +1,21 @@
+// 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/type.h"
+#include "paddle/ir/dialect.h"
+
+namespace ir {
+IrContext *Type::ir_context() const { return dialect().ir_context(); }
+
+}  // namespace ir
diff --git a/paddle/ir/type.h b/paddle/ir/type.h
index c4c5663dda9db..3df2834d90eca 100644
--- a/paddle/ir/type.h
+++ b/paddle/ir/type.h
@@ -52,6 +52,10 @@ class Type {
 
   StorageType *storage() const { return storage_; }
 
+  Dialect &dialect() const { return storage_->abstract_type().dialect(); }
+
+  IrContext *ir_context() const;
+
   ///
   /// \brief Enable hashing Type.
   ///

From 5dad6afab0fb8c72e989d3a53e769a679da04579 Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Tue, 28 Feb 2023 01:35:14 +0000
Subject: [PATCH 4/8] set WITH_NEWIR=OFF

---
 CMakeLists.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/CMakeLists.txt b/CMakeLists.txt
index 8893723c9ef57..2b8caf7addd16 100755
--- a/CMakeLists.txt
+++ b/CMakeLists.txt
@@ -315,7 +315,7 @@ option(WITH_CUDNN_FRONTEND
        "Compile with CUDNN Frontend API support (experimental)" OFF)
 option(WITH_CUDNN_DSO "Compile PaddlePaddle with cuDNN dynamic-link libraries"
        OFF)
-option(WITH_NEWIR "Compile PaddlePaddle with NEWIR" ON)
+option(WITH_NEWIR "Compile PaddlePaddle with NEWIR" OFF)
 
 if(WITH_RECORD_BUILDTIME)
   set_property(

From 3e8691b41919b34ac23f8285171462e59558982b Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Tue, 28 Feb 2023 03:56:01 +0000
Subject: [PATCH 5/8] refine code by comment

---
 paddle/ir/builtin_dialect.h |  2 +-
 paddle/ir/dialect.cc        |  8 ++++----
 paddle/ir/dialect.h         |  6 +++---
 paddle/ir/ir_context.cc     | 26 ++++++++++++--------------
 paddle/ir/ir_context.h      |  2 +-
 paddle/ir/type.h            |  2 +-
 paddle/ir/type_base.h       |  2 +-
 7 files changed, 23 insertions(+), 25 deletions(-)

diff --git a/paddle/ir/builtin_dialect.h b/paddle/ir/builtin_dialect.h
index 25c22110c1d5a..93e2195253c92 100644
--- a/paddle/ir/builtin_dialect.h
+++ b/paddle/ir/builtin_dialect.h
@@ -29,7 +29,7 @@ class BuiltinDialect : public ir::Dialect {
   /// \brief Each Dialect needs to provide a name function to return the name of
   /// the Dialect.
   ///
-  static const std::string name() { return "builtin"; }
+  static const char *name() { return "builtin"; }
 
  private:
   void initialize();
diff --git a/paddle/ir/dialect.cc b/paddle/ir/dialect.cc
index 6a2055c411817..5a913fdf4bb9f 100644
--- a/paddle/ir/dialect.cc
+++ b/paddle/ir/dialect.cc
@@ -16,13 +16,13 @@
 
 namespace ir {
 Dialect::Dialect(std::string name, ir::IrContext *context, ir::TypeId id)
-    : name_(name), context_(context), id_(id) {}
+    : name_(std::move(name)), context_(context), id_(id) {}
 
-void Dialect::RegisterType(ir::TypeId type_id,
-                           ir::AbstractType &&abstract_type) {
+void Dialect::RegisterType(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);
+  this->ir_context()->RegisterAbstractType(new_abstract_type->type_id(),
+                                           new_abstract_type);
 }
 
 }  // namespace ir
diff --git a/paddle/ir/dialect.h b/paddle/ir/dialect.h
index b50db790a99bb..f9051ab89267e 100644
--- a/paddle/ir/dialect.h
+++ b/paddle/ir/dialect.h
@@ -29,7 +29,7 @@ class Dialect {
  public:
   Dialect(std::string name, ir::IrContext *context, ir::TypeId id);
 
-  std::string name() const { return name_; }
+  const std::string &name() const { return name_; }
 
   ir::IrContext *ir_context() const { return context_; }
 
@@ -57,13 +57,13 @@ class Dialect {
   }
 
   ///
-  /// \brief Register type_id and abstract_type into context.
+  /// \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::TypeId type_id, ir::AbstractType &&abstract_type);
+  void RegisterType(ir::AbstractType &&abstract_type);
 
  private:
   std::string name_;
diff --git a/paddle/ir/ir_context.cc b/paddle/ir/ir_context.cc
index 1a96758c11918..e2c0c9bfec866 100644
--- a/paddle/ir/ir_context.cc
+++ b/paddle/ir/ir_context.cc
@@ -29,12 +29,12 @@ class IrContextImpl {
 
   ~IrContextImpl() {
     std::lock_guard<ir::SpinLock> guard(destructor_lock_);
-    for (auto abstract_type_map : registed_abstract_types_) {
+    for (auto &abstract_type_map : registed_abstract_types_) {
       delete abstract_type_map.second;
     }
     registed_abstract_types_.clear();
 
-    for (auto dialect_map : registed_dialect_) {
+    for (auto &dialect_map : registed_dialect_) {
       delete dialect_map.second;
     }
     registed_dialect_.clear();
@@ -51,17 +51,16 @@ class IrContextImpl {
   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()) {
-      LOG(WARNING)
-          << "IrContext not fonund cached abstract_type of: [TypeId_hash="
-          << std::hash<ir::TypeId>()(type_id) << "].";
-      return nullptr;
-    } else {
+    if (iter != registed_abstract_types_.end()) {
       VLOG(4) << "IrContext fonund a cached abstract_type of: [TypeId_hash="
               << std::hash<ir::TypeId>()(type_id)
               << ", AbstractType_ptr=" << iter->second << "].";
       return iter->second;
     }
+    LOG(WARNING)
+        << "IrContext not fonund cached abstract_type of: [TypeId_hash="
+        << std::hash<ir::TypeId>()(type_id) << "].";
+    return nullptr;
   }
 
   void RegisterDialect(std::string name, Dialect *dialect) {
@@ -74,15 +73,14 @@ class IrContextImpl {
   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 {
+    if (iter != registed_dialect_.end()) {
       VLOG(4) << "IrContext fonund a cached dialect of: [name=" << name
               << ", dialect_ptr=" << iter->second << "].";
       return iter->second;
     }
+    LOG(WARNING) << "IrContext not fonund cached dialect of: [name=" << name
+                 << "].";
+    return nullptr;
   }
 
   ir::SpinLock registed_abstract_types_lock_;
@@ -153,7 +151,7 @@ std::vector<Dialect *> IrContext::GetRegisteredDialects() {
   return result;
 }
 
-Dialect *IrContext::GetRegisteredDialect(std::string dialect_name) {
+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;
diff --git a/paddle/ir/ir_context.h b/paddle/ir/ir_context.h
index a7be7f6948fc0..7343c57487445 100644
--- a/paddle/ir/ir_context.h
+++ b/paddle/ir/ir_context.h
@@ -115,7 +115,7 @@ class IrContext {
   ///
   /// \return The dialect named "name" from the context.
   ///
-  Dialect *GetRegisteredDialect(std::string dialect_name);
+  Dialect *GetRegisteredDialect(const std::string &dialect_name);
 
   ///
   /// \brief Get a registered dialect for the given dialect type T. The
diff --git a/paddle/ir/type.h b/paddle/ir/type.h
index 3df2834d90eca..64e65e99688a0 100644
--- a/paddle/ir/type.h
+++ b/paddle/ir/type.h
@@ -52,7 +52,7 @@ class Type {
 
   StorageType *storage() const { return storage_; }
 
-  Dialect &dialect() const { return storage_->abstract_type().dialect(); }
+  const Dialect &dialect() const { return storage_->abstract_type().dialect(); }
 
   IrContext *ir_context() const;
 
diff --git a/paddle/ir/type_base.h b/paddle/ir/type_base.h
index c3c6a9a7cea7b..803ac7eb721b6 100644
--- a/paddle/ir/type_base.h
+++ b/paddle/ir/type_base.h
@@ -62,7 +62,7 @@ class AbstractType {
   ///
   /// \return The dialect this type was registered to.
   ///
-  Dialect &dialect() const { return const_cast<Dialect &>(dialect_); }
+  const Dialect &dialect() const { return dialect_; }
 
   ///
   /// \brief Find the AbstractType instance whose TypeId is type_id from

From 5c9f16250dca9f2a04164083c6027ddd46ca0018 Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Tue, 28 Feb 2023 06:44:13 +0000
Subject: [PATCH 6/8] polish code

---
 paddle/ir/builtin_dialect.h  |  2 ++
 paddle/ir/builtin_type.h     |  6 +++++-
 paddle/ir/ir_context.cc      | 13 +++++++------
 paddle/ir/tests/type_test.cc | 30 ++++++++++++++++++++----------
 4 files changed, 34 insertions(+), 17 deletions(-)

diff --git a/paddle/ir/builtin_dialect.h b/paddle/ir/builtin_dialect.h
index 93e2195253c92..5016c1abea074 100644
--- a/paddle/ir/builtin_dialect.h
+++ b/paddle/ir/builtin_dialect.h
@@ -29,6 +29,8 @@ class BuiltinDialect : public ir::Dialect {
   /// \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:
diff --git a/paddle/ir/builtin_type.h b/paddle/ir/builtin_type.h
index 0636b5234a548..8b15ae6eed0db 100644
--- a/paddle/ir/builtin_type.h
+++ b/paddle/ir/builtin_type.h
@@ -24,7 +24,11 @@ namespace ir {
 
 ///
 /// \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:
diff --git a/paddle/ir/ir_context.cc b/paddle/ir/ir_context.cc
index e2c0c9bfec866..243d0655af72d 100644
--- a/paddle/ir/ir_context.cc
+++ b/paddle/ir/ir_context.cc
@@ -22,7 +22,8 @@
 #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() {}
@@ -83,20 +84,20 @@ class IrContextImpl {
     return nullptr;
   }
 
-  ir::SpinLock registed_abstract_types_lock_;
-
   // 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_;
 
-  ir::SpinLock registed_dialect_lock_;
-
   // The dialcet registered in the context.
   std::unordered_map<std::string, Dialect *> registed_dialect_;
 
-  // Some built-in type.
+  ir::SpinLock registed_dialect_lock_;
+
+  // Some built-in types.
   Float32Type fp32_type;
   Int32Type int32_type;
 
diff --git a/paddle/ir/tests/type_test.cc b/paddle/ir/tests/type_test.cc
index 412014d19f096..e8901be6c3539 100644
--- a/paddle/ir/tests/type_test.cc
+++ b/paddle/ir/tests/type_test.cc
@@ -23,17 +23,19 @@
 #include "paddle/ir/type_base.h"
 
 TEST(type_test, type_id) {
+  // Define two empty classes, just for testing.
   class TypeA {};
   class TypeB {};
 
-  // (1) Test construct TypeId by TypeId::Get()
+  // Test 1: Test construct TypeId by TypeId::get<T>() and overloaded operator==
+  // method.
   ir::TypeId a_id = ir::TypeId::get<TypeA>();
   ir::TypeId a_other_id = ir::TypeId::get<TypeA>();
   ir::TypeId b_id = ir::TypeId::get<TypeB>();
   EXPECT_EQ(a_id, a_other_id);
   EXPECT_NE(a_id, b_id);
 
-  // (2) Test TypeId hash
+  // Test 2: Test the hash function of TypeId.
   std::unordered_map<ir::TypeId, ir::TypeId *> type_id_register;
   type_id_register.emplace(a_id, &a_id);
   type_id_register.emplace(b_id, &b_id);
@@ -43,33 +45,37 @@ TEST(type_test, type_id) {
 }
 
 TEST(type_test, type_base) {
+  // Define two empty classes, just for testing.
   class TypeA {};
 
+  // Define a FakeDialect without registering any types.
   struct FakeDialect : ir::Dialect {
     explicit FakeDialect(ir::IrContext *context)
         : ir::Dialect(name(), context, ir::TypeId::get<FakeDialect>()) {}
-    static const std::string name() { return "fake"; }
+    static const char *name() { return "fake"; }
   };
 
+  // Test 1: Test the function of IrContext to register Dialect.
   ir::IrContext *ctx = ir::IrContext::Instance();
   ir::Dialect *fake_dialect = ctx->GetOrRegisterDialect<FakeDialect>();
 
+  // Test 2: Test the get method of AbstractType.
   ir::TypeId a_id = ir::TypeId::get<TypeA>();
   ir::AbstractType abstract_type_a = ir::AbstractType::get(a_id, *fake_dialect);
-
   EXPECT_EQ(abstract_type_a.type_id(), a_id);
 
+  // Test 3: Test the constructor of TypeStorage.
   ir::TypeStorage storage_a(&abstract_type_a);
-
   EXPECT_EQ(storage_a.abstract_type().type_id(), abstract_type_a.type_id());
 }
 
 TEST(type_test, built_in_type) {
-  // Test creation of built-in parameterless type.
+  // Test 1: Test the built-in type of IrContext.
   ir::IrContext *ctx = ir::IrContext::Instance();
   ir::Type fp32_1 = ir::Float32Type::get(ctx);
 
-  // Test interfaces of class Type
+  // Test 2: Test the interfaces of class Type: judgment, type_id,
+  // abstract_type, classof.
   ir::Type fp32_2 = ir::Float32Type::get(ctx);
   EXPECT_EQ(fp32_1 == fp32_2, 1);
   EXPECT_EQ(fp32_1 != fp32_2, 0);
@@ -89,6 +95,7 @@ TEST(type_test, built_in_type) {
   EXPECT_EQ(ir::Int32Type::classof(int32_1), 1);
 }
 
+// Customize a parameterized TypeStorage IntegerTypeStorage.
 struct IntegerTypeStorage : public ir::TypeStorage {
   IntegerTypeStorage(unsigned width, unsigned signedness)
       : width_(width), signedness_(signedness) {}
@@ -118,24 +125,27 @@ struct IntegerTypeStorage : public ir::TypeStorage {
   }
 };
 
+// Customize a parameterized type: IntegerType, storage type is
+// IntegerTypeStorage.
 class IntegerType : public ir::Type {
  public:
   using Type::Type;
   DECLARE_TYPE_UTILITY_FUNCTOR(IntegerType, IntegerTypeStorage);
 };
 
+// Customize a Dialect IntegerDialect, registration type of IntegerType.
 struct IntegerDialect : ir::Dialect {
   explicit IntegerDialect(ir::IrContext *context)
       : ir::Dialect(name(), context, ir::TypeId::get<IntegerDialect>()) {
     RegisterType<IntegerType>();
   }
-  static const std::string name() { return "integer"; }
+  static const char *name() { return "integer"; }
 };
 
 TEST(type_test, custom_type_dialect) {
   ir::IrContext *ctx = ir::IrContext::Instance();
 
-  // Test Custom Type and Dialect
+  // Test 1: Test the function of IrContext to register Dialect.
   ctx->GetOrRegisterDialect<IntegerDialect>();
 
   ir::Type int1_1 = IntegerType::get(ctx, 1, 0);
@@ -145,7 +155,7 @@ TEST(type_test, custom_type_dialect) {
   ir::Type int8 = IntegerType::get(ctx, 8, 0);
   EXPECT_EQ(int8 == int1_2, 0);
 
-  // Test Dialect interface
+  //  Test 2: Test Dialect interfaces
   EXPECT_EQ(ctx == int8.ir_context(), 1);
 
   EXPECT_EQ(int8.dialect().id() == ir::TypeId::get<IntegerDialect>(), 1);

From 04b4bec53954b111627e83ff89e23bbc09973d84 Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Tue, 28 Feb 2023 07:05:43 +0000
Subject: [PATCH 7/8] refine include style

---
 paddle/ir/ir_context.cc      | 3 ++-
 paddle/ir/storage_manager.cc | 4 ++--
 2 files changed, 4 insertions(+), 3 deletions(-)

diff --git a/paddle/ir/ir_context.cc b/paddle/ir/ir_context.cc
index 243d0655af72d..b8725c1ee5069 100644
--- a/paddle/ir/ir_context.cc
+++ b/paddle/ir/ir_context.cc
@@ -12,12 +12,13 @@
 // 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/dialect.h"
-#include "paddle/ir/ir_context.h"
 #include "paddle/ir/spin_lock.h"
 #include "paddle/ir/type_base.h"
 
diff --git a/paddle/ir/storage_manager.cc b/paddle/ir/storage_manager.cc
index 91bc1ba3d360c..4ed630af6c160 100644
--- a/paddle/ir/storage_manager.cc
+++ b/paddle/ir/storage_manager.cc
@@ -12,11 +12,11 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
+#include "paddle/ir/storage_manager.h"
+
 #include <memory>
 #include <unordered_map>
 
-#include "paddle/ir/storage_manager.h"
-
 namespace ir {
 // This is a structure for creating, caching, and looking up Storage of
 // parameteric types.

From ff74640c2938357a2d504b94e873a61ae370915a Mon Sep 17 00:00:00 2001
From: zhangbo9674 <zhangbo54@baidu.com>
Date: Tue, 28 Feb 2023 11:58:52 +0000
Subject: [PATCH 8/8] refine log for debug

---
 paddle/ir/dialect.h          |  2 ++
 paddle/ir/ir_context.cc      | 27 ++++++++++++---------------
 paddle/ir/storage_manager.cc |  8 ++++----
 3 files changed, 18 insertions(+), 19 deletions(-)

diff --git a/paddle/ir/dialect.h b/paddle/ir/dialect.h
index f9051ab89267e..6b5b733f78296 100644
--- a/paddle/ir/dialect.h
+++ b/paddle/ir/dialect.h
@@ -49,11 +49,13 @@ class Dialect {
   ///
   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) << "----------------------------------";
   }
 
   ///
diff --git a/paddle/ir/ir_context.cc b/paddle/ir/ir_context.cc
index b8725c1ee5069..6fca67c16e2fd 100644
--- a/paddle/ir/ir_context.cc
+++ b/paddle/ir/ir_context.cc
@@ -44,7 +44,7 @@ class IrContextImpl {
 
   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);
@@ -54,20 +54,19 @@ class IrContextImpl {
     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 fonund a cached abstract_type of: [TypeId_hash="
+      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)
-        << "IrContext not fonund cached abstract_type of: [TypeId_hash="
-        << std::hash<ir::TypeId>()(type_id) << "].";
+    LOG(WARNING) << "No cache found abstract_type of: [TypeId_hash="
+                 << std::hash<ir::TypeId>()(type_id) << "].";
     return nullptr;
   }
 
   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
+    VLOG(4) << "Register a dialect of: [name=" << name
             << ", dialect_ptr=" << dialect << "].";
     registed_dialect_.emplace(name, dialect);
   }
@@ -76,12 +75,11 @@ class IrContextImpl {
     std::lock_guard<ir::SpinLock> guard(registed_dialect_lock_);
     auto iter = registed_dialect_.find(name);
     if (iter != registed_dialect_.end()) {
-      VLOG(4) << "IrContext fonund a cached dialect of: [name=" << name
+      VLOG(4) << "Fonund a cached dialect of: [name=" << name
               << ", dialect_ptr=" << iter->second << "].";
       return iter->second;
     }
-    LOG(WARNING) << "IrContext not fonund cached dialect of: [name=" << name
-                 << "].";
+    LOG(WARNING) << "No cache fonund dialect of: [name=" << name << "].";
     return nullptr;
   }
 
@@ -111,8 +109,9 @@ IrContext *IrContext::Instance() {
 }
 
 IrContext::IrContext() : impl_(new IrContextImpl()) {
-  VLOG(4) << "BuiltinDialect registered into IrContext.";
+  VLOG(4) << "BuiltinDialect registered into IrContext. ===>";
   GetOrRegisterDialect<BuiltinDialect>();
+  VLOG(4) << "==============================================";
 
   impl_->fp32_type = TypeManager::get<Float32Type>(this);
   impl_->int32_type = TypeManager::get<Int32Type>(this);
@@ -134,11 +133,12 @@ std::unordered_map<TypeId, AbstractType *>
 
 Dialect *IrContext::GetOrRegisterDialect(
     std::string dialect_name, std::function<Dialect *()> constructor) {
-  VLOG(4) << "IrContext get or register a dialect of: [name=" << dialect_name
+  VLOG(4) << "Try to 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.";
+    VLOG(4) << "Create and register a new Dialect of: [name=" << dialect_name
+            << "].";
     dialect = constructor();
     impl().RegisterDialect(dialect_name, dialect);
   }
@@ -164,9 +164,6 @@ Dialect *IrContext::GetRegisteredDialect(const std::string &dialect_name) {
 }
 
 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.GetAbstractType(type_id);
   if (abstract_type) {
diff --git a/paddle/ir/storage_manager.cc b/paddle/ir/storage_manager.cc
index 4ed630af6c160..9cec4a48e1441 100644
--- a/paddle/ir/storage_manager.cc
+++ b/paddle/ir/storage_manager.cc
@@ -72,7 +72,7 @@ StorageManager::StorageBase *StorageManager::GetParametricStorageTypeImpl(
     std::function<bool(const StorageBase *)> equal_func,
     std::function<StorageBase *()> constructor) {
   std::lock_guard<ir::SpinLock> guard(parametric_instance_lock_);
-  VLOG(4) << "StorageManager get parameteretric storage of: [TypeId_hash="
+  VLOG(4) << "Try to get a parameteretric storage of: [TypeId_hash="
           << std::hash<ir::TypeId>()(type_id) << ", param_hash=" << hash_value
           << "].";
   if (parametric_instance_.find(type_id) == parametric_instance_.end())
@@ -84,7 +84,7 @@ StorageManager::StorageBase *StorageManager::GetParametricStorageTypeImpl(
 StorageManager::StorageBase *StorageManager::GetParameterlessStorageTypeImpl(
     TypeId type_id) {
   std::lock_guard<ir::SpinLock> guard(parameterless_instance_lock_);
-  VLOG(4) << "StorageManager get parameterless storage of: [TypeId_hash="
+  VLOG(4) << "Try to get a parameterless storage of: [TypeId_hash="
           << std::hash<ir::TypeId>()(type_id) << "].";
   if (parameterless_instance_.find(type_id) == parameterless_instance_.end())
     throw("TypeId not found in IrContext.");
@@ -94,7 +94,7 @@ StorageManager::StorageBase *StorageManager::GetParameterlessStorageTypeImpl(
 
 void StorageManager::RegisterParametricStorageTypeImpl(TypeId type_id) {
   std::lock_guard<ir::SpinLock> guard(parametric_instance_lock_);
-  VLOG(4) << "StorageManager register parameteric storage of: [TypeId_hash="
+  VLOG(4) << "Register a parameteric storage of: [TypeId_hash="
           << std::hash<ir::TypeId>()(type_id) << "].";
   parametric_instance_.emplace(type_id,
                                std::make_unique<ParametricStorageManager>());
@@ -103,7 +103,7 @@ void StorageManager::RegisterParametricStorageTypeImpl(TypeId type_id) {
 void StorageManager::RegisterParameterlessStorageTypeImpl(
     TypeId type_id, std::function<StorageBase *()> constructor) {
   std::lock_guard<ir::SpinLock> guard(parameterless_instance_lock_);
-  VLOG(4) << "StorageManager register parameterless storage of: [TypeId_hash="
+  VLOG(4) << "Register a parameterless storage of: [TypeId_hash="
           << std::hash<ir::TypeId>()(type_id) << "].";
   if (parameterless_instance_.find(type_id) != parameterless_instance_.end())
     throw("storage class already registered");