[IR] Type system stage2: add class Type, type uniquer utils, class IRContext

paddle/ir/CMakeLists.txt

@@ -2,4 +2,12 @@ if(NOT WITH_NEWIR)
   return()
 endif()
 
-add_subdirectory(type)
+set(NEWIR_SOURCE_DIR "${PADDLE_SOURCE_DIR}/paddle/ir")
+set(NEWIR_BINARY_DIR "${PADDLE_BINARY_DIR}/paddle/ir")
+
+# ir tests
+add_subdirectory(tests)
+
+file(GLOB IR_SRCS "*.cc")
+
+cc_library(new_ir SRCS ${IR_SRCS})

paddle/ir/builtin_type.h

@@ -0,0 +1,42 @@
+// 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/type.h"
+
+namespace ir {
+///
+/// \brief Definitions of built-in type classes. The built-in type object get
+/// method is as follows: Type fp32 = Float32Type::get(ctx);
+///
+class Float32Type : public ir::Type {
+ public:
+  using Type::Type;
+
+  DECLARE_TYPE_UTILITY_FUNCTOR(Float32Type, ir::TypeStorage);
+
+  static Float32Type get(ir::IrContext *context);
+};
+
+class Int32Type : public ir::Type {
+ public:
+  using Type::Type;
+
+  DECLARE_TYPE_UTILITY_FUNCTOR(Int32Type, ir::TypeStorage);
+
+  static Int32Type get(ir::IrContext *context);
+};
+
+}  // namespace ir

paddle/ir/ir_context.cc

@@ -0,0 +1,118 @@
+// 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 <unordered_map>
+
+#include "paddle/ir/builtin_type.h"
+#include "paddle/ir/ir_context.h"
+#include "paddle/ir/spin_lock.h"
+#include "paddle/ir/type_base.h"
+
+namespace ir {
+// The implementation class of the IrContext class
+class IrContextImpl {
+ public:
+  IrContextImpl() {}
+
+  ~IrContextImpl() {
+    std::lock_guard<ir::SpinLock> guard(registed_abstract_types_lock_);
+    for (auto abstract_type_map : registed_abstract_types_) {
+      delete abstract_type_map.second;
+    }
+    registed_abstract_types_.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="
+            << std::hash<ir::TypeId>()(type_id)
+            << ", AbstractType_ptr=" << abstract_type << "].";
+    registed_abstract_types_.emplace(type_id, abstract_type);
+  }
+
+  AbstractType *lookup(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="
+              << std::hash<ir::TypeId>()(type_id)
+              << ", AbstractType_ptr=" << iter->second << "].";
+      return iter->second;
+    }
+  }
+
+  ir::SpinLock registed_abstract_types_lock_;
+
+  // Cached AbstractType instances.
+  std::unordered_map<TypeId, AbstractType *> registed_abstract_types_;
+
+  // TypeStorage uniquer and cache instances.
+  StorageManager registed_storage_manager_;
+
+  // Some built-in type.
+  Float32Type fp32_type;
+  Int32Type int32_type;
+};
+
+IrContext *IrContext::Instance() {
+  static IrContext context;
+  return &context;
+}
+
+IrContext::IrContext() : impl_(new IrContextImpl()) {
+  VLOG(4) << "IrContext register built-in type...";
+  REGISTER_TYPE_2_IRCONTEXT(Float32Type, this);
+  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,
+                                     AbstractType *abstract_type) {
+  impl().RegisterAbstractType(type_id, abstract_type);
+}
+
+StorageManager &IrContext::storage_manager() {
+  return impl().registed_storage_manager_;
+}
+
+std::unordered_map<TypeId, AbstractType *>
+    &IrContext::registed_abstracted_type() {
+  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
+          << "].";
+  auto &impl = ctx->impl();
+  AbstractType *abstract_type = impl.lookup(type_id);
+  if (abstract_type) {
+    return *abstract_type;
+  } else {
+    throw("Abstract type not found in IrContext.");
+  }
+}
+
+Float32Type Float32Type::get(IrContext *ctx) { return ctx->impl().fp32_type; }
+
+Int32Type Int32Type::get(IrContext *ctx) { return ctx->impl().int32_type; }
+
+}  // namespace ir

paddle/ir/ir_context.h

@@ -0,0 +1,82 @@
+// 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 <glog/logging.h>
+#include <memory>
+#include <unordered_map>
+
+namespace ir {
+class IrContextImpl;
+class StorageManager;
+class AbstractType;
+class TypeId;
+
+///
+/// \brief IrContext is a global parameterless class used to store and manage
+/// Type and its related data structures.
+///
+class IrContext {
+ public:
+  ///
+  /// \brief Initializes a new instance of IrContext.
+  ///
+  static IrContext *Instance();
+
+  ///
+  /// \brief Get an instance of IrContextImpl, a private member of IrContext.
+  /// For the specific definition of IrContextImpl, see ir_context.cc.
+  ///
+  /// \return The instance of IrContextImpl.
+  ///
+  IrContextImpl &impl() { return *impl_; }
+
+  ///
+  /// \brief Register an AbstractType to IrContext
+  ///
+  /// \param type_id The type id of the AbstractType.
+  /// \param abstract_type AbstractType* provided by user.
+  ///
+  void RegisterAbstractType(ir::TypeId type_id, AbstractType *abstract_type);
+
+  ///
+  /// \brief Returns the storage uniquer used for constructing TypeStorage
+  /// instances.
+  ///
+  /// \return The storage uniquer used for constructing TypeStorage
+  /// instances.
+  ///
+  StorageManager &storage_manager();
+
+  ///
+  /// \brief Returns the storage uniquer used for constructing TypeStorage
+  /// instances.
+  ///
+  /// \return The storage uniquer used for constructing TypeStorage
+  /// instances.
+  ///
+  std::unordered_map<TypeId, AbstractType *> &registed_abstracted_type();
+
+  IrContext(const IrContext &) = delete;
+
+  void operator=(const IrContext &) = delete;
+
+ private:
+  IrContext();
+
+  const std::unique_ptr<IrContextImpl> impl_;
+};
+
+}  // namespace ir

paddle/ir/spin_lock.h

@@ -0,0 +1,66 @@
+// 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 <atomic>
+#if defined(_M_X64) || defined(__x86_64__) || defined(_M_IX86) || \
+    defined(__i386__)
+#define __PADDLE_x86__
+#include <immintrin.h>
+#endif
+#include <mutex>
+#include <thread>
+
+namespace ir {
+static inline void CpuRelax() {
+#if defined(__PADDLE_x86__)
+  _mm_pause();
+#endif
+}
+
+class SpinLock {
+ public:
+  SpinLock() : mlock_(false) {}
+
+  void lock() {
+    for (;;) {
+      if (!mlock_.exchange(true, std::memory_order_acquire)) {
+        break;
+      }
+      constexpr int kMaxLoop = 32;
+      for (int loop = 1; mlock_.load(std::memory_order_relaxed);) {
+        if (loop <= kMaxLoop) {
+          for (int i = 1; i <= loop; ++i) {
+            CpuRelax();
+          }
+          loop *= 2;
+        } else {
+          std::this_thread::yield();
+        }
+      }
+    }
+  }
+
+  void unlock() { mlock_.store(false, std::memory_order_release); }
+
+ private:
+  SpinLock(const SpinLock&) = delete;
+  SpinLock(SpinLock&&) = delete;
+  SpinLock& operator=(const SpinLock&) = delete;
+  SpinLock& operator=(SpinLock&&) = delete;
+  std::atomic<bool> mlock_;
+};
+
+}  // namespace ir