[mlir] Make some functions public to use custom TypeIDs

[math-fehr]

Summary

Currently, it is only possible to register an operation or a type
when the TypeID is defined at compile time. Same with InterfaceMaps
which can only be defined with compile-time defined interfaces.

With those changes, it is now possible to register types/operations
with custom TypeIDs. This is necessary to define new operations/types
at runtime.

[rriddle]

Thanks!

mlir/include/mlir/IR/Dialect.h

222

nit: Move this next to the templated addTypes.

mlir/include/mlir/IR/TypeSupport.h

65

nit: Can you wrap this in a non-templated get instead of exposing the constructor? Also, can you place it next to the other get?

[math-fehr]

Address comments by moving functions to better place.

Also, use a 'get' function instead of exposing the AbstractType
constructor.

[math-fehr]

Fix phabricator diff.

I think I erased the first patch with the second one.

[math-fehr]

I addressed the comments.

Also, I forgot to precise that I can't land the patch!

mlir/include/mlir/IR/Dialect.h

(void)std::initializer_list<int>{
         0, (AbstractOperation::insert<Args>(*this), 0)...};
   }
 
   /// Register a set of type classes with this dialect.
   template <typename... Args> void addTypes() {
     (void)std::initializer_list<int>{0, (addType<Args>(), 0)...};
   }
 
+  /// Register a type instance with this dialect.
+  /// The use of this method is in general discouraged in favor of
+  /// 'addTypes<CustomType>()'.
+  void addType(TypeID typeID, AbstractType &&typeInfo);
+
   /// Register a set of attribute classes with this dialect.
   template <typename... Args> void addAttributes() {
     (void)std::initializer_list<int>{0, (addAttribute<Args>(), 0)...};
   }
 
   /// Enable support for unregistered operations.
   void allowUnknownOperations(bool allow = true) { unknownOpsAllowed = allow; }
 
   /// Enable support for unregistered types.
   void allowUnknownTypes(bool allow = true) { unknownTypesAllowed = allow; }
 
   /// Register a dialect interface with this dialect instance.
   void addInterface(std::unique_ptr<DialectInterface> interface);
 
   /// Register a set of dialect interfaces with this dialect instance.
   template <typename... Args> void addInterfaces() {
     (void)std::initializer_list<int>{
         0, (addInterface(std::make_unique<Args>(this)), 0)...};
   }
 
 private:
   Dialect(const Dialect &) = delete;
   void operator=(Dialect &) = delete;
 
   /// Register an attribute instance with this dialect.
   template <typename T> void addAttribute() {
     // Add this attribute to the dialect and register it with the uniquer.
     addAttribute(T::getTypeID(), AbstractAttribute::get<T>(*this));
     detail::AttributeUniquer::registerAttribute<T>(context);
   }
   void addAttribute(TypeID typeID, AbstractAttribute &&attrInfo);
 
   /// Register a type instance with this dialect.
   template <typename T> void addType() {
     // Add this type to the dialect and register it with the uniquer.
     addType(T::getTypeID(), AbstractType::get<T>(*this));
     detail::TypeUniquer::registerType<T>(context);
   }
-  void addType(TypeID typeID, AbstractType &&typeInfo);
 
   /// The namespace of this dialect.
   StringRef name;
 
   /// The unique identifier of the derived Op class, this is used in the context
   /// to allow registering multiple times the same dialect.
   TypeID dialectID;
 

mlir/include/mlir/IR/OperationSupport.h

public:
   static void insert(Dialect &dialect) {
     insert(T::getOperationName(), dialect, TypeID::get<T>(),
            T::getParseAssemblyFn(), T::getPrintAssemblyFn(),
            T::getVerifyInvariantsFn(), T::getFoldHookFn(),
            T::getGetCanonicalizationPatternsFn(), T::getInterfaceMap(),
            T::getHasTraitFn());
   }
 
-private:
+  /// Register a new operation in a Dialect object.
+  /// The use of this method is in general discouraged in favor of
+  /// 'insert<CustomOp>(dialect)'.
   static void insert(StringRef name, Dialect &dialect, TypeID typeID,
                      ParseAssemblyFn parseAssembly,
                      PrintAssemblyFn printAssembly,
                      VerifyInvariantsFn verifyInvariants, FoldHookFn foldHook,
                      GetCanonicalizationPatternsFn getCanonicalizationPatterns,
                      detail::InterfaceMap &&interfaceMap, HasTraitFn hasTrait);
 
+private:
   AbstractOperation(StringRef name, Dialect &dialect, TypeID typeID,
                     ParseAssemblyFn parseAssembly,
                     PrintAssemblyFn printAssembly,
                     VerifyInvariantsFn verifyInvariants, FoldHookFn foldHook,
                     GetCanonicalizationPatternsFn getCanonicalizationPatterns,
                     detail::InterfaceMap &&interfaceMap, HasTraitFn hasTrait);
 
   /// A map of interfaces that were registered to this operation.

mlir/include/mlir/IR/TypeSupport.h

public:
   static const AbstractType &lookup(TypeID typeID, MLIRContext *context);
 
   /// This method is used by Dialect objects when they register the list of
   /// types they contain.
   template <typename T> static AbstractType get(Dialect &dialect) {
     return AbstractType(dialect, T::getInterfaceMap(), T::getTypeID());
   }
 
+  /// This method is used by Dialect objects to register types with
+  /// custom TypeIDs.
+  /// The use of this method is in general discouraged in favor of
+  /// 'get<CustomType>(dialect)';
+  static AbstractType get(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
+                          TypeID typeID) {
+    return AbstractType(dialect, std::move(interfaceMap), typeID);
+  }
+
   /// Return the dialect this type was registered to.
   Dialect &getDialect() const { return const_cast<Dialect &>(dialect); }
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this type, null otherwise. This should not be used
   /// directly.
   template <typename T> typename T::Concept *getInterface() const {
     return interfaceMap.lookup<T>();
   }
 
   /// Return the unique identifier representing the concrete type class.
   TypeID getTypeID() const { return typeID; }
 
 private:
   AbstractType(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
                TypeID typeID)
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
         typeID(typeID) {}
 
   /// This is the dialect that this type was registered to.
   Dialect &dialect;
 
   /// This is a collection of the interfaces registered to this type.

mlir/include/mlir/Support/InterfaceSupport.h

public:
   /// was registered to this map, null otherwise.
   template <typename T> typename T::Concept *lookup() const {
     return reinterpret_cast<typename T::Concept *>(lookup(T::getInterfaceID()));
   }
 
   /// Returns true if the interface map contains an interface for the given id.
   bool contains(TypeID interfaceID) const { return lookup(interfaceID); }
 
+  /// Create an InterfaceMap given with the implementation of the interfaces.
+  /// The use of this constructor is in general discouraged in favor of
+  /// 'InterfaceMap::get<InterfaceA, ...>()'.
+  InterfaceMap(MutableArrayRef<std::pair<TypeID, void *>> elements)
+      : interfaces(elements.begin(), elements.end()) {
+    llvm::sort(interfaces, [](const auto &lhs, const auto &rhs) {
+      return compare(lhs.first, rhs.first);
+    });
+  }
+
 private:
   /// Compare two TypeID instances by comparing the underlying pointer.
   static bool compare(TypeID lhs, TypeID rhs) {
     return lhs.getAsOpaquePointer() < rhs.getAsOpaquePointer();
   }
 
   InterfaceMap() = default;
-  InterfaceMap(MutableArrayRef<std::pair<TypeID, void *>> elements)
-      : interfaces(elements.begin(), elements.end()) {
-    llvm::sort(interfaces, [](const auto &lhs, const auto &rhs) {
-      return compare(lhs.first, rhs.first);
-    });
-  }
 
   template <typename... Ts>
   static InterfaceMap getImpl(std::tuple<Ts...> *) {
     std::pair<TypeID, void *> elements[] = {std::make_pair(
         Ts::getInterfaceID(),
         new (malloc(sizeof(typename Ts::ModelT))) typename Ts::ModelT())...};
     return InterfaceMap(elements);
   }