Add more explicit error message when creating a type or attribute for an unregistered dialect (NFC)

[rriddle]

mlir/lib/Support/StorageUniquer.cpp

237

A singleton is only one type of storage object, you also need to check for parametric storage objects.

[mehdi_amini]

Add checks for parametric storage as well

[rriddle]

mlir/include/mlir/IR/TypeSupport.h

19

Why is this include necessary here, but not in AttributeSupport?

[mehdi_amini]

Remove spurious include

[GMNGeoffrey]

Thanks for improving error messages Mehdi!

mlir/include/mlir/IR/AttributeSupport.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_ATTRIBUTESUPPORT_H
 #define MLIR_IR_ATTRIBUTESUPPORT_H
 
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/StorageUniquerSupport.h"
 #include "llvm/ADT/PointerIntPair.h"
+#include "llvm/ADT/Twine.h"
 
 namespace mlir {
 class MLIRContext;
 class Type;
 
 //===----------------------------------------------------------------------===//
 // AbstractAttribute
 //===----------------------------------------------------------------------===//

 // MLIRContext. This class manages all creation and uniquing of attributes.
 class AttributeUniquer {
 public:
   /// Get an uniqued instance of a parametric attribute T.
   template <typename T, typename... Args>
   static typename std::enable_if_t<
       !std::is_same<typename T::ImplType, AttributeStorage>::value, T>
   get(MLIRContext *ctx, Args &&...args) {
+#ifndef NDEBUG
+    if (!ctx->getAttributeUniquer().isParametricStorageInitialized(
+            T::getTypeID()))
+      llvm::report_fatal_error(llvm::Twine("can't create Attribute '") +
+                               llvm::getTypeName<T>() +
+                               "' because storage uniquer isn't initialized: "
+                               "the dialect was likely not loaded.");
+#endif
     return ctx->getAttributeUniquer().get<typename T::ImplType>(
         [ctx](AttributeStorage *storage) {
           initializeAttributeStorage(storage, ctx, T::getTypeID());
         },
         T::getTypeID(), std::forward<Args>(args)...);
   }
   /// Get an uniqued instance of a singleton attribute T.
   template <typename T>
   static typename std::enable_if_t<
       std::is_same<typename T::ImplType, AttributeStorage>::value, T>
   get(MLIRContext *ctx) {
+#ifndef NDEBUG
+    if (!ctx->getAttributeUniquer().isSingletonStorageInitialized(
+            T::getTypeID()))
+      llvm::report_fatal_error(llvm::Twine("can't create Attribute '") +
+                               llvm::getTypeName<T>() +
+                               "' because storage uniquer isn't initialized: "
+                               "the dialect was likely not loaded.");
+#endif
     return ctx->getAttributeUniquer().get<typename T::ImplType>(T::getTypeID());
   }
 
   template <typename T, typename... Args>
   static LogicalResult mutate(MLIRContext *ctx, typename T::ImplType *impl,
                               Args &&...args) {
     assert(impl && "cannot mutate null attribute");
     return ctx->getAttributeUniquer().mutate(T::getTypeID(), impl,

mlir/include/mlir/IR/TypeSupport.h

 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_TYPE_SUPPORT_H
 #define MLIR_IR_TYPE_SUPPORT_H
 
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/StorageUniquerSupport.h"
+#include "llvm/ADT/Twine.h"
 
 namespace mlir {
 class Dialect;
 class MLIRContext;
 
 //===----------------------------------------------------------------------===//
 // AbstractType
 //===----------------------------------------------------------------------===//
 

 /// A utility class to get, or create, unique instances of types within an
 /// MLIRContext. This class manages all creation and uniquing of types.
 struct TypeUniquer {
   /// Get an uniqued instance of a parametric type T.
   template <typename T, typename... Args>
   static typename std::enable_if_t<
       !std::is_same<typename T::ImplType, TypeStorage>::value, T>
   get(MLIRContext *ctx, Args &&...args) {
+#ifndef NDEBUG
+    if (!ctx->getTypeUniquer().isParametricStorageInitialized(T::getTypeID()))
+      llvm::report_fatal_error(llvm::Twine("can't create type '") +
+                               llvm::getTypeName<T>() +
+                               "' because storage uniquer isn't initialized: "
+                               "the dialect was likely not loaded.");
+#endif
     return ctx->getTypeUniquer().get<typename T::ImplType>(
         [&](TypeStorage *storage) {
           storage->initialize(AbstractType::lookup(T::getTypeID(), ctx));
         },
         T::getTypeID(), std::forward<Args>(args)...);
   }
   /// Get an uniqued instance of a singleton type T.
   template <typename T>
   static typename std::enable_if_t<
       std::is_same<typename T::ImplType, TypeStorage>::value, T>
   get(MLIRContext *ctx) {
+#ifndef NDEBUG
+    if (!ctx->getTypeUniquer().isSingletonStorageInitialized(T::getTypeID()))
+      llvm::report_fatal_error(llvm::Twine("can't create type '") +
+                               llvm::getTypeName<T>() +
+                               "' because storage uniquer isn't initialized: "
+                               "the dialect was likely not loaded.");
+#endif
     return ctx->getTypeUniquer().get<typename T::ImplType>(T::getTypeID());
   }
 
   /// Change the mutable component of the given type instance in the provided
   /// context.
   template <typename T, typename... Args>
   static LogicalResult mutate(MLIRContext *ctx, typename T::ImplType *impl,
                               Args &&...args) {

mlir/include/mlir/Support/StorageUniquer.h

public:
   template <typename Storage> Storage *get(TypeID id) {
     return static_cast<Storage *>(getSingletonImpl(id));
   }
   /// Utility override when the storage type represents the type id.
   template <typename Storage> Storage *get() {
     return get<Storage>(TypeID::get<Storage>());
   }
 
+  /// Test if there is a singleton storage uniquer initialized for the provided
+  /// TypeID. This is only useful for debugging/diagnostic purpose: the uniquer
+  /// is initialized when a dialect is loaded.
+  bool isSingletonStorageInitialized(TypeID id);
+
+  /// Test if there is a parametric storage uniquer initialized for the provided
+  /// TypeID. This is only useful for debugging/diagnostic purpose: the uniquer
+  /// is initialized when a dialect is loaded.
+  bool isParametricStorageInitialized(TypeID id);
+
   /// Changes the mutable component of 'storage' by forwarding the trailing
   /// arguments to the 'mutate' function of the derived class.
   template <typename Storage, typename... Args>
   LogicalResult mutate(TypeID id, Storage *storage, Args &&...args) {
     auto mutationFn = [&](StorageAllocator &allocator) -> LogicalResult {
       return static_cast<Storage &>(*storage).mutate(
           allocator, std::forward<Args>(args)...);
     };

mlir/lib/Support/StorageUniquer.cpp

 struct StorageUniquerImpl {
   using BaseStorage = StorageUniquer::BaseStorage;
   using StorageAllocator = StorageUniquer::StorageAllocator;
 
   //===--------------------------------------------------------------------===//
   // Parametric Storage
   //===--------------------------------------------------------------------===//
 
+  /// Check if an instance of a parametric storage class exists.
+  bool hasParametricStorage(TypeID id) { return parametricUniquers.count(id); }
+
   /// Get or create an instance of a parametric type.
   BaseStorage *
   getOrCreate(TypeID id, unsigned hashValue,
               function_ref<bool(const BaseStorage *)> isEqual,
               function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
     assert(parametricUniquers.count(id) &&
            "creating unregistered storage instance");
     ParametricStorageUniquer::LookupKey lookupKey{hashValue, isEqual};
struct StorageUniquerImpl {
 
   /// Get or create an instance of a singleton storage class.
   BaseStorage *getSingleton(TypeID id) {
     BaseStorage *singletonInstance = singletonInstances[id];
     assert(singletonInstance && "expected singleton instance to exist");
     return singletonInstance;
   }
 
+  /// Check if an instance of a singleton storage class exists.
+  bool hasSingleton(TypeID id) { return singletonInstances.count(id); }
+
   //===--------------------------------------------------------------------===//
   // Instance Storage
   //===--------------------------------------------------------------------===//
 
   /// Map of type ids to the storage uniquer to use for registered objects.
   DenseMap<TypeID, std::unique_ptr<ParametricStorageUniquer>>
       parametricUniquers;
 

 }
 
 /// Implementation for getting an instance of a derived type with default
 /// storage.
 auto StorageUniquer::getSingletonImpl(TypeID id) -> BaseStorage * {
   return impl->getSingleton(id);
 }
 
+/// Test is the storage singleton is initialized.
+bool StorageUniquer::isSingletonStorageInitialized(TypeID id) {
+  return impl->hasSingleton(id);
+}
+
+/// Test is the parametric storage is initialized.
+bool StorageUniquer::isParametricStorageInitialized(TypeID id) {
+  return impl->hasParametricStorage(id);
+}
+
 /// Implementation for registering an instance of a derived type with default
 /// storage.
 void StorageUniquer::registerSingletonImpl(
     TypeID id, function_ref<BaseStorage *(StorageAllocator &)> ctorFn) {
   assert(!impl->singletonInstances.count(id) &&
          "storage class already registered");
   impl->singletonInstances.try_emplace(id, ctorFn(impl->singletonAllocator));
 }