diff --git a/mlir/include/mlir/IR/MLIRContext.h b/mlir/include/mlir/IR/MLIRContext.h --- a/mlir/include/mlir/IR/MLIRContext.h +++ b/mlir/include/mlir/IR/MLIRContext.h @@ -142,6 +142,10 @@ // MLIRContextImpl type. MLIRContextImpl &getImpl() { return *impl; } + /// Create a new TypeID that is guaranteed to be unique for the lifetime + /// of the MLIRContext. + TypeID allocateTypeID(); + /// Returns the diagnostic engine for this context. DiagnosticEngine &getDiagEngine(); diff --git a/mlir/include/mlir/Support/TypeID.h b/mlir/include/mlir/Support/TypeID.h --- a/mlir/include/mlir/Support/TypeID.h +++ b/mlir/include/mlir/Support/TypeID.h @@ -89,6 +89,8 @@ // See TypeIDExported below for an explanation of the trampoline behavior. friend struct detail::TypeIDExported; + + friend class TypeIDAllocator; }; /// Enable hashing TypeID. @@ -135,6 +137,24 @@ return detail::TypeIDExported::get(); } +/// This class provides a way to define new TypeIDs at runtime. +/// When the allocator is destructed, all allocated TypeIDs becode invalid and +/// therefore should not be used. +class TypeIDAllocator { +public: + /// Allocate a new TypeID, that is ensured to be unique for the lifetime + /// of the TypeIDAllocator. + TypeID allocateID() { + ids.emplace_back(new TypeID::Storage()); + return TypeID(ids.back().get()); + } + +private: + /// The TypeIDs allocated are the addresses of the different storages. + /// Keeping those in memory ensure uniqueness of the TypeIDs. + std::vector> ids; +}; + } // end namespace mlir namespace llvm { diff --git a/mlir/lib/IR/MLIRContext.cpp b/mlir/lib/IR/MLIRContext.cpp --- a/mlir/lib/IR/MLIRContext.cpp +++ b/mlir/lib/IR/MLIRContext.cpp @@ -25,6 +25,7 @@ #include "mlir/IR/Types.h" #include "mlir/Support/DebugAction.h" #include "mlir/Support/ThreadLocalCache.h" +#include "mlir/Support/TypeID.h" #include "llvm/ADT/DenseMap.h" #include "llvm/ADT/DenseSet.h" #include "llvm/ADT/SetVector.h" @@ -283,6 +284,9 @@ /// An allocator used for AbstractAttribute and AbstractType objects. llvm::BumpPtrAllocator abstractDialectSymbolAllocator; + /// Manages the creation and lifetime of TypeID defined at runtime. + TypeIDAllocator typeIDAllocator; + //===--------------------------------------------------------------------===// // Affine uniquing //===--------------------------------------------------------------------===// @@ -432,6 +436,14 @@ return getImpl().debugActionManager; } +//===----------------------------------------------------------------------===// +// TypeIDs allocation +//===----------------------------------------------------------------------===// + +TypeID MLIRContext::allocateTypeID() { + return getImpl().typeIDAllocator.allocateID(); +} + //===----------------------------------------------------------------------===// // Diagnostic Handlers //===----------------------------------------------------------------------===//