[mlir] separable registration of attribute and type interfaces

It may be desirable to provide an interface implementation for an attribute or
a type without modifying the definition of said attribute or type. Notably,
this allows to implement interfaces for attributes and types outside of the
dialect that defines them and, in particular, provide interfaces for built-in
types. Provide the mechanism to do so.

Currently, separable registration requires the attribute or type to have been
registered with the context, i.e. for the dialect containing the attribute or
type to be loaded. This can be relaxed in the future using a mechanism similar
to delayed dialect interface registration.

See https://llvm.discourse.group/t/rfc-separable-attribute-type-interfaces/3637

Depends On D104233

Reviewed By: rriddle

Differential Revision: https://reviews.llvm.org/D104234

Author: ftynse

mlir/docs/Interfaces.md

@@ -207,6 +207,91 @@ if (ExampleOpInterface example = dyn_cast<ExampleOpInterface>(op))
   llvm::errs() << "hook returned = " << example.exampleInterfaceHook() << "\n";
 ```
 
+#### External Models for Attribute/Type Interfaces
+
+It may be desirable to provide an interface implementation for an attribute or a
+type without modifying the definition of said attribute or type. Notably, this
+allows to implement interfaces for attributes and types outside of the dialect
+that defines them and, in particular, provide interfaces for built-in types.
+
+This is achieved by extending the concept-based polymorphism model with two more
+classes derived from `Concept` as follows.
+
+```c++
+struct ExampleTypeInterfaceTraits {
+  struct Concept {
+    virtual unsigned exampleInterfaceHook(Type type) const = 0;
+    virtual unsigned exampleStaticInterfaceHook() const = 0;
+  };
+
+  template <typename ConcreteType>
+  struct Model : public Concept { /*...*/ };
+
+  /// Unlike `Model`, `FallbackModel` passes the type object through to the
+  /// hook, making it accessible in the method body even if the method is not
+  /// defined in the class itself and thus has no `this` access. ODS
+  /// automatically generates this class for all interfaces.
+  template <typename ConcreteType>
+  struct FallbackModel : public Concept {
+    unsigned exampleInterfaceHook(Type type) const override {
+      getImpl()->exampleInterfaceHook(type);
+    }
+    unsigned exampleStaticInterfaceHook() const override {
+      ConcreteType::exampleStaticInterfaceHook();
+    }
+  };
+
+  /// `ExternalModel` provides a place for default implementations of interface
+  /// methods by explicitly separating the model class, which implements the
+  /// interface, from the type class, for which the interface is being
+  /// implemented. Default implementations can be then defined generically
+  /// making use of `cast<ConcreteType>`. If `ConcreteType` does not provide
+  /// the APIs required by the default implementation, custom implementations
+  /// may use `FallbackModel` directly to override the default implementation.
+  /// Being located in a class template, it never gets instantiated and does not
+  /// lead to compilation errors. ODS automatically generates this class and
+  /// places default method implementations in it.
+  template <typename ConcreteModel, typename ConcreteType>
+  struct ExternalModel : public FallbackModel<ConcreteModel> {
+    unsigned exampleInterfaceHook(Type type) const override {
+      // Default implementation can be provided here.
+      return type.cast<ConcreteType>().callSomeTypeSpecificMethod();
+    }
+  };
+};
+```
+
+External models can be provided for attribute and type interfaces by deriving
+either `FallbackModel` or `ExternalModel` and by registering the model class
+with the attribute or type class in a given context. Other contexts will not see
+the interface unless registered.
+
+```c++
+/// External interface implementation for a concrete class. This does not
+/// require modifying the definition of the type class itself.
+struct ExternalModelExample
+    : public ExampleTypeInterface::ExternalModel<ExternalModelExample,
+                                                 IntegerType> {
+  static unsigned exampleStaticInterfaceHook() {
+    // Implementation is provided here.
+    return IntegerType::someStaticMethod();
+  }
+
+  // No need to define `exampleInterfaceHook` that has a default implementation
+  // in `ExternalModel`. But it can be overridden if desired.
+}
+
+int main() {
+  MLIRContext context;
+  /* ... */;
+
+  // Register the interface model with the type in the given context before
+  // using it. The dialect contaiing the type is expected to have been loaded
+  // at this point.
+  IntegerType::registerInterface<ExternalModelExample>(context);
+}
+```
+
 #### Dialect Fallback for OpInterface
 
 Some dialects have an open ecosystem and don't register all of the possible
@@ -215,9 +300,9 @@ implementing an `OpInterface` for these operation. When an operation isn't
 registered or does not provide an implementation for an interface, the query
 will fallback to the dialect itself.
 
-A second model is used for such cases and automatically generated when
-using ODS (see below) with the name `FallbackModel`. This model can be implemented
-for a particular dialect:
+A second model is used for such cases and automatically generated when using ODS
+(see below) with the name `FallbackModel`. This model can be implemented for a
+particular dialect:
 
 ```c++
 // This is the implementation of a dialect fallback for `ExampleOpInterface`.

mlir/include/mlir/IR/AttributeSupport.h

@@ -59,14 +59,24 @@ class AbstractAttribute {
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
         typeID(typeID) {}
 
+  /// Give StorageUserBase access to the mutable lookup.
+  template <typename ConcreteT, typename BaseT, typename StorageT,
+            typename UniquerT, template <typename T> class... Traits>
+  friend class detail::StorageUserBase;
+
+  /// Look up the specified abstract attribute in the MLIRContext and return a
+  /// (mutable) pointer to it. Return a null pointer if the attribute could not
+  /// be found in the context.
+  static AbstractAttribute *lookupMutable(TypeID typeID, MLIRContext *context);
+
   /// This is the dialect that this attribute was registered to.
-  Dialect &dialect;
+  const Dialect &dialect;
 
   /// This is a collection of the interfaces registered to this attribute.
   detail::InterfaceMap interfaceMap;
 
   /// The unique identifier of the derived Attribute class.
-  TypeID typeID;
+  const TypeID typeID;
 };
 
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/Attributes.h

@@ -31,6 +31,7 @@ class Attribute {
 
   using ImplType = AttributeStorage;
   using ValueType = void;
+  using AbstractType = AbstractAttribute;
 
   constexpr Attribute() : impl(nullptr) {}
   /* implicit */ Attribute(const ImplType *impl)

mlir/include/mlir/IR/StorageUniquerSupport.h

@@ -87,6 +87,21 @@ class StorageUserBase : public BaseT, public Traits<ConcreteT>... {
     return detail::InterfaceMap::template get<Traits<ConcreteT>...>();
   }
 
+  /// Attach the given models as implementations of the corresponding interfaces
+  /// for the concrete storage user class. The type must be registered with the
+  /// context, i.e. the dialect to which the type belongs must be loaded. The
+  /// call will abort otherwise.
+  template <typename... IfaceModels>
+  static void attachInterface(MLIRContext &context) {
+    typename ConcreteT::AbstractType *abstract =
+        ConcreteT::AbstractType::lookupMutable(TypeID::get<ConcreteT>(),
+                                               &context);
+    if (!abstract)
+      llvm::report_fatal_error("Registering an interface for an attribute/type "
+                               "that is not itself registered.");
+    abstract->interfaceMap.template insert<IfaceModels...>();
+  }
+
   /// Get or create a new ConcreteT instance within the ctx. This
   /// function is guaranteed to return a non null object and will assert if
   /// the arguments provided are invalid.

mlir/include/mlir/IR/TypeSupport.h

@@ -67,14 +67,24 @@ class AbstractType {
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
         typeID(typeID) {}
 
+  /// Give StorageUserBase access to the mutable lookup.
+  template <typename ConcreteT, typename BaseT, typename StorageT,
+            typename UniquerT, template <typename T> class... Traits>
+  friend class detail::StorageUserBase;
+
+  /// Look up the specified abstract type in the MLIRContext and return a
+  /// (mutable) pointer to it. Return a null pointer if the type could not
+  /// be found in the context.
+  static AbstractType *lookupMutable(TypeID typeID, MLIRContext *context);
+
   /// This is the dialect that this type was registered to.
-  Dialect &dialect;
+  const Dialect &dialect;
 
   /// This is a collection of the interfaces registered to this type.
   detail::InterfaceMap interfaceMap;
 
   /// The unique identifier of the derived Type class.
-  TypeID typeID;
+  const TypeID typeID;
 };
 
 //===----------------------------------------------------------------------===//
@@ -105,11 +115,11 @@ class TypeStorage : public StorageUniquer::BaseStorage {
   /// Set the abstract type for this storage instance. This is used by the
   /// TypeUniquer when initializing a newly constructed type storage object.
   void initialize(const AbstractType &abstractTy) {
-    abstractType = &abstractTy;
+    abstractType = const_cast<AbstractType *>(&abstractTy);
   }
 
   /// The abstract description for this type.
-  const AbstractType *abstractType;
+  AbstractType *abstractType;
 };
 
 /// Default storage type for types that require no additional initialization or

mlir/include/mlir/IR/Types.h

@@ -79,6 +79,8 @@ class Type {
 
   using ImplType = TypeStorage;
 
+  using AbstractType = AbstractType;
+
   constexpr Type() : impl(nullptr) {}
   /* implicit */ Type(const ImplType *impl)
       : impl(const_cast<ImplType *>(impl)) {}

mlir/include/mlir/Support/InterfaceSupport.h

@@ -75,6 +75,8 @@ class Interface : public BaseType {
   using FallbackModel = typename Traits::template FallbackModel<T>;
   using InterfaceBase =
       Interface<ConcreteType, ValueT, Traits, BaseType, BaseTrait>;
+  template <typename T, typename U>
+  using ExternalModel = typename Traits::template ExternalModel<T, U>;
 
   /// This is a special trait that registers a given interface with an object.
   template <typename ConcreteT>
@@ -139,6 +141,25 @@ struct FilterTypes {
           typename FilterTypeT<Pred<Es>::value>::template type<Es>>()...));
 };
 
+namespace {
+/// Type trait indicating whether all template arguments are
+/// trivially-destructible.
+template <typename... Args>
+struct all_trivially_destructible;
+
+template <typename Arg, typename... Args>
+struct all_trivially_destructible<Arg, Args...> {
+  static constexpr const bool value =
+      std::is_trivially_destructible<Arg>::value &&
+      all_trivially_destructible<Args...>::value;
+};
+
+template <>
+struct all_trivially_destructible<> {
+  static constexpr const bool value = true;
+};
+} // namespace
+
 /// This class provides an efficient mapping between a given `Interface` type,
 /// and a particular implementation of its concept.
 class InterfaceMap {
@@ -198,6 +219,28 @@ class InterfaceMap {
     });
   }
 
+  /// Insert the given models as implementations of the corresponding interfaces
+  /// for the concrete attribute class.
+  template <typename... IfaceModels>
+  void insert() {
+    static_assert(all_trivially_destructible<IfaceModels...>::value,
+                  "interface models must be trivially destructible");
+    std::pair<TypeID, void *> elements[] = {
+        std::make_pair(IfaceModels::Interface::getInterfaceID(),
+                       new (malloc(sizeof(IfaceModels))) IfaceModels())...};
+    // Insert directly into the right position to keep the interfaces sorted.
+    for (auto &element : elements) {
+      TypeID id = element.first;
+      auto it =
+          llvm::lower_bound(interfaces, id, [](const auto &it, TypeID id) {
+            return compare(it.first, id);
+          });
+      if (it != interfaces.end() && it->first == id)
+        llvm::report_fatal_error("Interface already registered");
+      interfaces.insert(it, element);
+    }
+  }
+
 private:
   /// Compare two TypeID instances by comparing the underlying pointer.
   static bool compare(TypeID lhs, TypeID rhs) {

mlir/lib/IR/MLIRContext.cpp

@@ -306,7 +306,7 @@ class MLIRContextImpl {
   // Type uniquing
   //===--------------------------------------------------------------------===//
 
-  DenseMap<TypeID, const AbstractType *> registeredTypes;
+  DenseMap<TypeID, AbstractType *> registeredTypes;
   StorageUniquer typeUniquer;
 
   /// Cached Type Instances.
@@ -324,7 +324,7 @@ class MLIRContextImpl {
   // Attribute uniquing
   //===--------------------------------------------------------------------===//
 
-  DenseMap<TypeID, const AbstractAttribute *> registeredAttributes;
+  DenseMap<TypeID, AbstractAttribute *> registeredAttributes;
   StorageUniquer attributeUniquer;
 
   /// Cached Attribute Instances.
@@ -666,12 +666,20 @@ void Dialect::addAttribute(TypeID typeID, AbstractAttribute &&attrInfo) {
 /// Get the dialect that registered the attribute with the provided typeid.
 const AbstractAttribute &AbstractAttribute::lookup(TypeID typeID,
                                                    MLIRContext *context) {
+  const AbstractAttribute *abstract = lookupMutable(typeID, context);
+  if (!abstract)
+    llvm::report_fatal_error("Trying to create an Attribute that was not "
+                             "registered in this MLIRContext.");
+  return *abstract;
+}
+
+AbstractAttribute *AbstractAttribute::lookupMutable(TypeID typeID,
+                                                    MLIRContext *context) {
   auto &impl = context->getImpl();
   auto it = impl.registeredAttributes.find(typeID);
   if (it == impl.registeredAttributes.end())
-    llvm::report_fatal_error("Trying to create an Attribute that was not "
-                             "registered in this MLIRContext.");
-  return *it->second;
+    return nullptr;
+  return it->second;
 }
 
 //===----------------------------------------------------------------------===//
@@ -733,12 +741,19 @@ AbstractOperation::AbstractOperation(
 //===----------------------------------------------------------------------===//
 
 const AbstractType &AbstractType::lookup(TypeID typeID, MLIRContext *context) {
+  const AbstractType *type = lookupMutable(typeID, context);
+  if (!type)
+    llvm::report_fatal_error(
+        "Trying to create a Type that was not registered in this MLIRContext.");
+  return *type;
+}
+
+AbstractType *AbstractType::lookupMutable(TypeID typeID, MLIRContext *context) {
   auto &impl = context->getImpl();
   auto it = impl.registeredTypes.find(typeID);
   if (it == impl.registeredTypes.end())
-    llvm::report_fatal_error(
-        "Trying to create a Type that was not registered in this MLIRContext.");
-  return *it->second;
+    return nullptr;
+  return it->second;
 }
 
 //===----------------------------------------------------------------------===//

mlir/test/lib/Dialect/Test/CMakeLists.txt

@@ -5,6 +5,8 @@ set(LLVM_OPTIONAL_SOURCES
 )
 
 set(LLVM_TARGET_DEFINITIONS TestInterfaces.td)
+mlir_tablegen(TestAttrInterfaces.h.inc -gen-attr-interface-decls)
+mlir_tablegen(TestAttrInterfaces.cpp.inc -gen-attr-interface-defs)
 mlir_tablegen(TestTypeInterfaces.h.inc -gen-type-interface-decls)
 mlir_tablegen(TestTypeInterfaces.cpp.inc -gen-type-interface-defs)
 mlir_tablegen(TestOpInterfaces.h.inc -gen-op-interface-decls)

mlir/test/lib/Dialect/Test/TestAttributes.cpp

@@ -93,6 +93,8 @@ void CompoundAAttr::print(DialectAsmPrinter &printer) const {
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//
 
+#include "TestAttrInterfaces.cpp.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.cpp.inc"