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 @@ -16,6 +16,7 @@ #include "mlir/Support/LLVM.h" #include "llvm/ADT/DenseMapInfo.h" +#include "llvm/Support/Allocator.h" #include "llvm/Support/PointerLikeTypeTraits.h" namespace mlir { @@ -89,6 +90,8 @@ // See TypeIDExported below for an explanation of the trampoline behavior. friend struct detail::TypeIDExported; + + friend class TypeIDAllocator; }; /// Enable hashing TypeID. @@ -135,6 +138,21 @@ return detail::TypeIDExported::get(); } +/// This class provides a way to define new TypeIDs at runtime. +/// When the allocator is destructed, all allocated TypeIDs become 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 allocate() { return TypeID(ids.Allocate()); } + +private: + /// The TypeIDs allocated are the addresses of the different storages. + /// Keeping those in memory ensure uniqueness of the TypeIDs. + llvm::SpecificBumpPtrAllocator 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.allocate(); +} + //===----------------------------------------------------------------------===// // Diagnostic Handlers //===----------------------------------------------------------------------===//