[mlir] Add classes to define new TypeIDs at runtime

[math-fehr]

Summary

TypeIDAllocator enables the allocation of new TypeIDs at runtime,
that are unique during the lifetime of the allocator.

NonMovableTypeIDOwner is a class used to define a new TypeID for each instance
of a class, using the instance address. This class cannot be copied or moved.

[bondhugula]

mlir/include/mlir/Support/TypeID.h

148

Nit: Since this is a TypeID allocator, you can name this allocate().

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

147

Why does this need to be in the context? Couldn't the TypeID allocator be owned by whatever is using it? (e.g. the extensible dialect)

mlir/include/mlir/Support/TypeID.h

142

156

Why not use a BumpPtrAllocator here instead? It doesn't look like you need the individual instances.

[math-fehr]

Rename function, and use bump pointer allocator

[math-fehr]

Fix spelling error

[math-fehr]

This comment was removed by math-fehr.

mlir/include/mlir/IR/MLIRContext.h

147

We could put a TypeID allocator in the extensible dialects for the operations and types TypeID.
However, we would still have a problem with the Dialects and Traits defined at runtime, since they need to be stored in the MLIRContext (I do not have a better solution for that currently).

There is however another solution we could consider, which is to create a class (let's call it DynamicObject) which hold a TypeID, as well as the storage if necessary. If we would make the DynamicOp/DynamicDialects/... all child classes of this DynamicObject, we wouldn't need any TypeIDAllocator, beside for the operations (since they are stored entirely in the AbstractOperation class). Also, alternatively, we could use the address of the DynamicTypeDefinition for the TypeID (which would give us for free a lookup for the type given the TypeID).

[math-fehr]

Sorry, I think I added myself as a reviewer by mistake.

[math-fehr]

@rriddle, a gentle ping on this.
What do you think would be the best course of action here?
As I said in a previous comment, I could either keep the TypeIDAllocator in the MLIRContext, or move it to the ExtensibleDialect and having the "dynamic" objects owning their TypeIDs.

[math-fehr]

Hey @rriddle, a gentle ping on this ^

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

147

I really like the idea of a DynamicObject of some kind, which could internally use its own address for the TypeID? I think something like this could be placed directly within TypeID.h. The name would need to be different of course.

[math-fehr]

Add NonMovableTypeIDOwner, and remove TypeIDAllocator from MLIRContext

[math-fehr]

Was this class what you add in mind? I can also define another class that owns its TypeID by having a pointer to some place in memory if this can be useful.

[rriddle]

LGTM, nice.

mlir/include/mlir/Support/TypeID.h

158

I'm not sure we need the NonMovable in the name of the type. Something like SelfOwningTypeID or something that better convey the usage of this class seems better.

160

?

166

If you end up removing Owner from the name, this could also just be an implicit conversion to TypeID.

[trilorez]

Hi, I was wondering about the status of this patch. I'd be interested in using it in my project.

[math-fehr]

Address comments, and rename class to SelfOwningTypeID

[math-fehr]

Sorry, I just applied the last changes, and this should be ready to be committed I believe!
@rriddle, could you commit it for me if this is good for you? I do not have commit rights.

[rriddle]

In D104534#3267813, @math-fehr wrote:

Sorry, I just applied the last changes, and this should be ready to be committed I believe!
@rriddle, could you commit it for me if this is good for you? I do not have commit rights.

Yeah, I'll land tomorrow.

[rriddle]

In D104534#3289397, @rriddle wrote:

In D104534#3267813, @math-fehr wrote:

Sorry, I just applied the last changes, and this should be ready to be committed I believe!
@rriddle, could you commit it for me if this is good for you? I do not have commit rights.

Yeah, I'll land tomorrow.

Woah, missed this. Can you reupload with arc? Alternatively, I need an email to land with.

[math-fehr]

No problem!

Could you land it with [e-mail] ? My arc is currently broken :/

mlir/include/mlir/IR/MLIRContext.h

public:
 
   /// Return true if this operation name is registered in this context.
   bool isOperationRegistered(StringRef name);
 
   // This is effectively private given that only MLIRContext.cpp can see the
   // 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();
 
   /// Returns the storage uniquer used for creating affine constructs.
   StorageUniquer &getAffineUniquer();
 
   /// Returns the storage uniquer used for constructing type storage instances.
   /// This should not be used directly.

mlir/include/mlir/Support/TypeID.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SUPPORT_TYPEID_H
 #define MLIR_SUPPORT_TYPEID_H
 
 #include "mlir/Support/LLVM.h"
 #include "llvm/ADT/DenseMapInfo.h"
+#include "llvm/Support/Allocator.h"
 #include "llvm/Support/PointerLikeTypeTraits.h"
 
 namespace mlir {
 
 namespace detail {
 struct TypeIDExported;
 } // namespace detail
 

 private:
   TypeID(const Storage *storage) : storage(storage) {}
 
   /// The storage of this type info object.
   const Storage *storage;
 
   // See TypeIDExported below for an explanation of the trampoline behavior.
   friend struct detail::TypeIDExported;
+
+  friend class TypeIDAllocator;
 };
 
 /// Enable hashing TypeID.
 inline ::llvm::hash_code hash_value(TypeID id) {
   return llvm::hash_value(id.storage);
 }
 
 namespace detail {

 TypeID TypeID::get() {
   return detail::TypeIDExported::get<T>();
 }
 template <template <typename> class Trait>
 TypeID TypeID::get() {
   return detail::TypeIDExported::get<Trait>();
 }
 
+/// 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<TypeID::Storage> ids;
+};
+
 } // end namespace mlir
 
 namespace llvm {
 template <> struct DenseMapInfo<mlir::TypeID> {
   static mlir::TypeID getEmptyKey() {
     void *pointer = llvm::DenseMapInfo<void *>::getEmptyKey();
     return mlir::TypeID::getFromOpaquePointer(pointer);
   }
   static mlir::TypeID getTombstoneKey() {
     void *pointer = llvm::DenseMapInfo<void *>::getTombstoneKey();
     return mlir::TypeID::getFromOpaquePointer(pointer);
   }
   static unsigned getHashValue(mlir::TypeID val) {
     return mlir::hash_value(val);
   }
   static bool isEqual(mlir::TypeID lhs, mlir::TypeID rhs) { return lhs == rhs; }
 };
 
 /// We align TypeID::Storage by 8, so allow LLVM to steal the low bits.

mlir/lib/IR/MLIRContext.cpp

 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Identifier.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/Location.h"
 #include "mlir/IR/OpImplementation.h"
 #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"
 #include "llvm/ADT/SmallString.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
#endif
   /// A thread local cache of identifiers to reduce lock contention.
   ThreadLocalCache<llvm::StringMap<
       llvm::StringMapEntry<PointerUnion<Dialect *, MLIRContext *>> *>>
       localIdentifierCache;
 
   /// 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
   //===--------------------------------------------------------------------===//
 
   // Affine allocator and mutex for thread safety.
   llvm::BumpPtrAllocator affineAllocator;
   llvm::sys::SmartRWMutex<true> affineMutex;
 

 // Debugging
 //===----------------------------------------------------------------------===//
 
 DebugActionManager &MLIRContext::getDebugActionManager() {
   return getImpl().debugActionManager;
 }
 
 //===----------------------------------------------------------------------===//
+// TypeIDs allocation
+//===----------------------------------------------------------------------===//
+
+TypeID MLIRContext::allocateTypeID() {
+  return getImpl().typeIDAllocator.allocate();
+}
+
+//===----------------------------------------------------------------------===//
 // Diagnostic Handlers
 //===----------------------------------------------------------------------===//
 
 /// Returns the diagnostic engine for this context.
 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
 
 //===----------------------------------------------------------------------===//
 // Dialect and Operation Registration