[mlir] separable registration of attribute and type interfaces

[ftynse]

Summary

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

[ftynse]

Drop some irrelevant changes

[rriddle]

LG thanks! This has been on the TODO for a while, but just never found the time.

mlir/docs/Interfaces.md

230

Extra newline here.

248

typo: implemenations

265

typo: attirbute

mlir/include/mlir/IR/AttributeSupport.h

67–70

Why return null as opposed to aborting in the lookup(like the other method does)? To provide a better error?

mlir/include/mlir/IR/Attributes.h

34

mlir/include/mlir/IR/StorageUniquerSupport.h

90–91

The documentation here should mention that the type must be registered (i.e. its parent dialect must be loaded) in the context (given that otherwise this aborts).

97

The style of this message doesn't match the other fatal errors, e.g. no capital etc. I can't remember which we prefer for report_fatal_error, but we should be consistent.

mlir/include/mlir/IR/TypeSupport.h

68–78

mlir/include/mlir/IR/Types.h

79–84

mlir/include/mlir/Support/InterfaceSupport.h

204–222

An important invariant about this interface map is that the interfaces can't have state, given that the destructors don't get called. Can we add a validation (I guess is_trivially_destructible) for
the provided interfaces?

mlir/test/lib/Dialect/Test/TestInterfaces.td

66

69

78

88

Is the (ins) here necessary? I would have thought it be defaulted. (same above)

mlir/unittests/IR/InterfaceAttachmentTest.cpp

96–100

I can foresee this being a bit tricky if we want some methods defaulted and some not, but I don't really see a good way to avoid that. (Given that the implementer here generally doesn't have control over the interface implementation.

[ftynse]

A bit more cmake.

[ftynse]

Address feedback.

mlir/include/mlir/IR/AttributeSupport.h

67–70

Provisioning for delayed registration of interfaces like we do for dialects. If lookup returns null, we will be able to store an interface construction lambda for later. In the meantime, this also gives a better error message, the original one refers to attribute/type construction.

mlir/include/mlir/IR/StorageUniquerSupport.h

97

We are really inconsistent about this, I see on average the same amount of capitalized and non-capitalized error message, and only some capitalized error messages have trailing periods. Changed here, but I would consider using the same style as for emitError (no capitalization, no trailing period) for simplicity.

mlir/unittests/IR/InterfaceAttachmentTest.cpp

96–100

It should be possible to provide an ODS mechanism that emits classes with default implementations including/excluding the methods from the list, something like `def
ConcreteName : DefineExternalInterfaceModel<InterfaceName> { let excludeDefaults = ["method1", "method2"]; }` if this becomes common.

Otherwise, I need to dig deeper into both the C++ spec and the specific compiler implementations we care about, but I suspect the method of a template may not be instantiated by the compiler if it is never called. That is, if there is a method in the derived class. Not sure about it, but this would let the user redefine only the methods they need directly in the class deriving from ExternalModel.

mlir/docs/Interfaces.md

 
 /// Later, we can query if a specific operation(like 'MyOp') overrides the given
 /// interface.
 Operation *op = ...;
 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 implemenations 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 attirbute 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
 operations. In such cases it is still possible to provide support for
 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`.
 struct FallbackExampleOpInterface
     : public ExampleOpInterface::FallbackModel<
           FallbackExampleOpInterface> {
   static bool classof(Operation *op) { return true; }
 

mlir/include/mlir/IR/AttributeSupport.h

public:
   TypeID getTypeID() const { return typeID; }
 
 private:
   AbstractAttribute(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
                     TypeID typeID)
       : 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;
 };
 
 //===----------------------------------------------------------------------===//
 // AttributeStorage
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 class AttributeUniquer;

mlir/include/mlir/IR/Attributes.h

public:
   /// Utility class for implementing attributes.
   template <typename ConcreteType, typename BaseType, typename StorageType,
             template <typename T> class... Traits>
   using AttrBase = detail::StorageUserBase<ConcreteType, BaseType, StorageType,
                                            detail::AttributeUniquer, Traits...>;
 
   using ImplType = AttributeStorage;
   using ValueType = void;
+  using Abstract = AbstractAttribute;
 
   constexpr Attribute() : impl(nullptr) {}
   /* implicit */ Attribute(const ImplType *impl)
       : impl(const_cast<ImplType *>(impl)) {}
 
   Attribute(const Attribute &other) = default;
   Attribute &operator=(const Attribute &other) = default;
 

mlir/include/mlir/IR/StorageUniquerSupport.h

public:
   }
 
   /// Returns an interface map for the interfaces registered to this storage
   /// user. This should not be used directly.
   static detail::InterfaceMap getInterfaceMap() {
     return detail::InterfaceMap::template get<Traits<ConcreteT>...>();
   }
 
+  /// Attach the given models as implementations of the corresponding interfaces
+  /// for the concrete storage user class.
+  template <typename... IfaceModels>
+  static void attachInterface(MLIRContext &context) {
+    typename ConcreteT::Abstract *abstract =
+        ConcreteT::Abstract::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.
   template <typename... Args>
   static ConcreteT get(MLIRContext *ctx, Args... args) {
     // Ensure that the invariants are correct for construction.
     assert(
         succeeded(ConcreteT::verify(getDefaultDiagnosticEmitFn(ctx), args...)));

mlir/include/mlir/IR/TypeSupport.h

public:
 
   /// 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) {}
+  /// 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;
 };
 
 //===----------------------------------------------------------------------===//
 // TypeStorage
 //===----------------------------------------------------------------------===//
 
 namespace detail {
 struct TypeUniquer;

 protected:
   /// This constructor is used by derived classes as part of the TypeUniquer.
   TypeStorage() : abstractType(nullptr) {}
 
 private:
   /// 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
 /// storage.
 using DefaultTypeStorage = TypeStorage;
 
 //===----------------------------------------------------------------------===//
 // TypeStorageAllocator

mlir/include/mlir/IR/Types.h

 //       the key.
 class Type {
 public:
   /// Utility class for implementing types.
   template <typename ConcreteType, typename BaseType, typename StorageType,
             template <typename T> class... Traits>
   using TypeBase = detail::StorageUserBase<ConcreteType, BaseType, StorageType,
                                            detail::TypeUniquer, Traits...>;
 
   using ImplType = TypeStorage;
 
+  using Abstract = AbstractType;
+
   constexpr Type() : impl(nullptr) {}
   /* implicit */ Type(const ImplType *impl)
       : impl(const_cast<ImplType *>(impl)) {}
 
   Type(const Type &other) = default;
   Type &operator=(const Type &other) = default;
 
   bool operator==(Type other) const { return impl == other.impl; }
   bool operator!=(Type other) const { return !(*this == other); }

mlir/include/mlir/Support/InterfaceSupport.h

 class Interface : public BaseType {
 public:
   using Concept = typename Traits::Concept;
   template <typename T> using Model = typename Traits::template Model<T>;
   template <typename T>
   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>
   struct Trait : public BaseTrait<ConcreteT, Trait> {
     using ModelT = Model<ConcreteT>;
 
     /// Define an accessor for the ID of this interface.
     static TypeID getInterfaceID() { return TypeID::get<ConcreteType>(); }
public:
   /// '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);
     });
   }
 
+  /// Insert the given models as implementations of the corresponding interfaces
+  /// for the concrete attribute class.
+  template <typename... IfaceModels>
+  void insert() {
+    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) {
     return lhs.getAsOpaquePointer() < rhs.getAsOpaquePointer();
   }
 
   InterfaceMap() = default;
 

mlir/lib/IR/MLIRContext.cpp

#endif
 
   // Affine expression uniquing.
   StorageUniquer affineUniquer;
 
   //===--------------------------------------------------------------------===//
   // Type uniquing
   //===--------------------------------------------------------------------===//
 
-  DenseMap<TypeID, const AbstractType *> registeredTypes;
+  DenseMap<TypeID, AbstractType *> registeredTypes;
   StorageUniquer typeUniquer;
 
   /// Cached Type Instances.
   BFloat16Type bf16Ty;
   Float16Type f16Ty;
   Float32Type f32Ty;
   Float64Type f64Ty;
   Float80Type f80Ty;
   Float128Type f128Ty;
   IndexType indexTy;
   IntegerType int1Ty, int8Ty, int16Ty, int32Ty, int64Ty, int128Ty;
   NoneType noneType;
 
   //===--------------------------------------------------------------------===//
   // Attribute uniquing
   //===--------------------------------------------------------------------===//
 
-  DenseMap<TypeID, const AbstractAttribute *> registeredAttributes;
+  DenseMap<TypeID, AbstractAttribute *> registeredAttributes;
   StorageUniquer attributeUniquer;
 
   /// Cached Attribute Instances.
   BoolAttr falseAttr, trueAttr;
   UnitAttr unitAttr;
   UnknownLoc unknownLocAttr;
   DictionaryAttr emptyDictionaryAttr;
   StringAttr emptyStringAttr;

 
 //===----------------------------------------------------------------------===//
 // AbstractAttribute
 //===----------------------------------------------------------------------===//
 
 /// 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;
 }
 
 //===----------------------------------------------------------------------===//
 // AbstractOperation
 //===----------------------------------------------------------------------===//
 
 ParseResult AbstractOperation::parseAssembly(OpAsmParser &parser,
                                              OperationState &result) const {
: name(Identifier::get(name, dialect.getContext())), dialect(dialect),
       printAssemblyFn(std::move(printAssembly)),
       verifyInvariantsFn(std::move(verifyInvariants)) {}
 
 //===----------------------------------------------------------------------===//
 // AbstractType
 //===----------------------------------------------------------------------===//
 
 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;
 }
 
 //===----------------------------------------------------------------------===//
 // Identifier uniquing
 //===----------------------------------------------------------------------===//
 
 /// Return an identifier for the specified string.
 Identifier Identifier::get(const Twine &string, MLIRContext *context) {

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

 set(LLVM_OPTIONAL_SOURCES
   TestDialect.cpp
   TestPatterns.cpp
   TestTraits.cpp
 )
 
 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_tablegen(TestOpInterfaces.cpp.inc -gen-op-interface-defs)
 add_public_tablegen_target(MLIRTestInterfaceIncGen)
 
 set(LLVM_TARGET_DEFINITIONS TestAttrDefs.td)
 mlir_tablegen(TestAttrDefs.h.inc -gen-attrdef-decls)

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

 
 #include <tuple>
 
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/DialectImplementation.h"
 
+#include "TestAttrInterfaces.h.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.h.inc"
 
 #endif // MLIR_TESTATTRIBUTES_H

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

void CompoundAAttr::print(DialectAsmPrinter &printer) const {
   llvm::interleaveComma(getArrayOfInts(), printer);
   printer << "]>";
 }
 
 //===----------------------------------------------------------------------===//
 // Tablegen Generated Definitions
 //===----------------------------------------------------------------------===//
 
+#include "TestAttrInterfaces.cpp.inc"
+
 #define GET_ATTRDEF_CLASSES
 #include "TestAttrDefs.cpp.inc"
 
 //===----------------------------------------------------------------------===//
 // TestDialect
 //===----------------------------------------------------------------------===//
 
 void TestDialect::registerAttributes() {

mlir/test/lib/Dialect/Test/TestInterfaces.td

def TestTypeInterface : TypeInterface<"TestTypeInterface"> {
   let extraTraitClassDeclaration = [{
     /// Prints the type name.
     void printTypeE(Location loc) const {
       emitRemark(loc) << $_type << " - TestE";
     }
   }];
 }
 
+def TestExternalTypeInterface : TypeInterface<"TestExternalTypeInterface"> {
+  let cppNamespace = "::mlir";
+  let methods = [
+    InterfaceMethod<"Returns the bitwidth of the type plus 'arg'.",
+      "unsigned", "getBitwidthPlusArg", (ins "unsigned":$arg)>,
+    StaticInterfaceMethod<"Returns some value plus 'arg'.",
+      "unsigned", "staticGetSomeValuePlusArg", (ins "unsigned":$arg)>,
+    InterfaceMethod<"Returns the argument doubled.",
+      "unsigned", "getBitwidthPlusDoubleArgument", (ins "unsigned":$arg), "",
+      [{return $_type.getIntOrFloatBitWidth() + 2 * arg;}]>,
+    StaticInterfaceMethod<"Returns the argument.",
+      "unsigned", "staticGetArgument", (ins "unsigned":$arg), "",
+      [{return arg;}]>,
+  ];
+}
+
+def TestExternalFallbackTypeInterface
+    : TypeInterface<"TestExternalFallbackTypeInterface"> {
+  let cppNamespace = "::mlir";
+  let methods = [
+    InterfaceMethod<"Returns the bitwidth of the given integer type.",
+      "unsigned", "getBitwidth", (ins), "", [{return $_type.getWidth();}]>,
+  ];
+}
+
+def TestExternalAttrInterface : AttrInterface<"TestExternalAttrInterface"> {
+  let cppNamespace = "::mlir";
+  let methods = [
+    InterfaceMethod<"Gets the dialect pointer.", "const ::mlir::Dialect *",
+      "getDialectPtr", (ins)>,
+    StaticInterfaceMethod<"Returns some number.", "int", "getSomeNumber",
+      (ins)>,
+  ];
+}
+
 def TestEffectOpInterface
     : EffectOpInterfaceBase<"TestEffectOpInterface",
                             "::mlir::TestEffects::Effect"> {
   let cppNamespace = "::mlir";
 }
 
 class TestEffect<string effectName>
     : SideEffect<TestEffectOpInterface, effectName, DefaultResource>;
 
 class TestEffects<list<TestEffect> effects = []>
    : SideEffectsTraitBase<TestEffectOpInterface, effects>;
 
 def TestConcreteEffect : TestEffect<"TestEffects::Concrete">;
 
 #endif // MLIR_TEST_DIALECT_TEST_INTERFACES

mlir/tools/mlir-tblgen/OpInterfacesGen.cpp

llvm::interleaveComma(
         method.getArguments(), os,
         [&](const InterfaceMethod::Argument &arg) { os << arg.type; });
     os << ");\n";
   }
   os << "  };\n";
 }
 
 void InterfaceGenerator::emitModelDecl(Interface &interface) {
+  // Emit the basic model and the fallback model.
   for (const char *modelClass : {"Model", "FallbackModel"}) {
     os << "  template<typename " << valueTemplate << ">\n";
     os << "  class " << modelClass << " : public Concept {\n  public:\n";
+    os << "    using Interface = " << interface.getCppNamespace()
+       << (interface.getCppNamespace().empty() ? "" : "::")
+       << interface.getName() << ";\n";
     os << "    " << modelClass << "() : Concept{";
     llvm::interleaveComma(
         interface.getMethods(), os,
         [&](const InterfaceMethod &method) { os << method.getName(); });
     os << "} {}\n\n";
 
     // Insert each of the virtual method overrides.
     for (auto &method : interface.getMethods()) {
       emitCPPType(method.getReturnType(), os << "    static inline ");
       emitMethodNameAndArgs(method, os, valueType,
                             /*addThisArg=*/!method.isStatic(),
                             /*addConst=*/false);
       os << ";\n";
     }
     os << "  };\n";
   }
+
+  // Emit the template for the external model.
+  os << "  template<typename ConcreteModel, typename " << valueTemplate
+     << ">\n";
+  os << "  class ExternalModel : public FallbackModel<ConcreteModel> {\n";
+  os << "  public:\n";
+
+  // Emit declarations for methods that have default implementations. Other
+  // methods are expected to be implemented by the concrete derived model.
+  for (auto &method : interface.getMethods()) {
+    if (!method.getDefaultImplementation())
+      continue;
+    os << "    ";
+    if (method.isStatic())
+      os << "static ";
+    emitCPPType(method.getReturnType(), os);
+    os << method.getName() << "(";
+    if (!method.isStatic()) {
+      emitCPPType(valueType, os);
+      os << "tablegen_opaque_val";
+      if (!method.arg_empty())
+        os << ", ";
+    }
+    llvm::interleaveComma(method.getArguments(), os,
+                          [&](const InterfaceMethod::Argument &arg) {
+                            emitCPPType(arg.type, os);
+                            os << arg.name;
+                          });
+    os << ")";
+    if (!method.isStatic())
+      os << " const";
+    os << ";\n";
+  }
+  os << "  };\n";
 }
 
 void InterfaceGenerator::emitModelMethodsDef(Interface &interface) {
   for (auto &method : interface.getMethods()) {
     os << "template<typename " << valueTemplate << ">\n";
     emitCPPType(method.getReturnType(), os);
     os << "detail::" << interface.getName() << "InterfaceTraits::Model<"
        << valueTemplate << ">::";
for (auto &method : interface.getMethods()) {
     os << method.getName() << '(';
     if (!method.isStatic())
       os << "tablegen_opaque_val" << (method.arg_empty() ? "" : ", ");
     llvm::interleaveComma(
         method.getArguments(), os,
         [&](const InterfaceMethod::Argument &arg) { os << arg.name; });
     os << ");\n}\n";
   }
+
+  // Emit default implementations for the external model.
+  for (auto &method : interface.getMethods()) {
+    if (!method.getDefaultImplementation())
+      continue;
+    os << "template<typename ConcreteModel, typename " << valueTemplate
+       << ">\n";
+    emitCPPType(method.getReturnType(), os);
+    os << "detail::" << interface.getName()
+       << "InterfaceTraits::ExternalModel<ConcreteModel, " << valueTemplate
+       << ">::";
+
+    os << method.getName() << "(";
+    if (!method.isStatic()) {
+      emitCPPType(valueType, os);
+      os << "tablegen_opaque_val";
+      if (!method.arg_empty())
+        os << ", ";
+    }
+    llvm::interleaveComma(method.getArguments(), os,
+                          [&](const InterfaceMethod::Argument &arg) {
+                            emitCPPType(arg.type, os);
+                            os << arg.name;
+                          });
+    os << ")";
+    if (!method.isStatic())
+      os << " const";
+
+    os << " {\n";
+
+    // Use the empty context for static methods.
+    tblgen::FmtContext ctx;
+    os << tblgen::tgfmt(method.getDefaultImplementation()->trim(),
+                        method.isStatic() ? &ctx : &nonStaticMethodFmt);
+    os << "\n}\n";
+  }
 }
 
 void InterfaceGenerator::emitTraitDecl(Interface &interface,
                                        StringRef interfaceName,
                                        StringRef interfaceTraitsName) {
   os << llvm::formatv("  template <typename {3}>\n"
                       "  struct {0}Trait : public ::mlir::{2}<{0},"
                       " detail::{1}>::Trait<{3}> {{\n",

mlir/tools/mlir-tblgen/SPIRVUtilsGen.cpp

os << "  class " << modelClass << " : public Concept {\n"
        << "(const Concept *impl, Operation *tblgen_opaque_op) const final {\n"
        << "      auto op = llvm::cast<ConcreteOp>(tblgen_opaque_op);\n"
        << "      (void)op;\n"
        // Forward to the method on the concrete operation type.
        << "      return op." << availability.getQueryFnName() << "();\n"
        << "    }\n"
        << "  };\n";
   }
+  os << "  template<typename ConcreteModel, typename ConcreteOp>\n";
+  os << "  class ExternalModel : public FallbackModel<ConcreteOp> {};\n";
 }
 
 static void emitInterfaceDecl(const Availability &availability,
                               raw_ostream &os) {
   StringRef interfaceName = availability.getInterfaceClassName();
   std::string interfaceTraitsName =
       std::string(formatv("{0}Traits", interfaceName));
 

mlir/unittests/IR/CMakeLists.txt

 add_mlir_unittest(MLIRIRTests
   AttributeTest.cpp
   DialectTest.cpp
+  InterfaceAttachmentTest.cpp
   MemRefTypeTest.cpp
   OperationSupportTest.cpp
   ShapedTypeTest.cpp
   SubElementInterfaceTest.cpp
 )
 target_link_libraries(MLIRIRTests
   PRIVATE
   MLIRIR)

mlir/unittests/IR/InterfaceAttachmentTest.cpp

+//===- InterfaceAttachmentTest.cpp - Test attaching interfaces ------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This implements the tests for attaching interfaces to attributes and types
+// without having to specify them on the attribute or type class directly.
+//
+//===----------------------------------------------------------------------===//
+
+#include "mlir/IR/BuiltinAttributes.h"
+#include "mlir/IR/BuiltinTypes.h"
+#include "gtest/gtest.h"
+
+#include "../../test/lib/Dialect/Test/TestAttributes.h"
+#include "../../test/lib/Dialect/Test/TestTypes.h"
+
+using namespace mlir;
+using namespace mlir::test;
+
+namespace {
+
+/// External interface model for the integer type. Only provides non-default
+/// methods.
+struct Model
+    : public TestExternalTypeInterface::ExternalModel<Model, IntegerType> {
+  unsigned getBitwidthPlusArg(Type type, unsigned arg) const {
+    return type.getIntOrFloatBitWidth() + arg;
+  }
+
+  static unsigned staticGetSomeValuePlusArg(unsigned arg) { return 42 + arg; }
+};
+
+/// External interface model for the float type. Provides non-deafult and
+/// overrides default methods.
+struct OverridingModel
+    : public TestExternalTypeInterface::ExternalModel<OverridingModel,
+                                                      FloatType> {
+  unsigned getBitwidthPlusArg(Type type, unsigned arg) const {
+    return type.getIntOrFloatBitWidth() + arg;
+  }
+
+  static unsigned staticGetSomeValuePlusArg(unsigned arg) { return 42 + arg; }
+
+  unsigned getBitwidthPlusDoubleArgument(Type type, unsigned arg) const {
+    return 128;
+  }
+
+  static unsigned staticGetArgument(unsigned arg) { return 420; }
+};
+
+TEST(InterfaceAttachment, Type) {
+  MLIRContext context;
+
+  // Check that the type has no interface.
+  IntegerType i8 = IntegerType::get(&context, 8);
+  ASSERT_FALSE(i8.isa<TestExternalTypeInterface>());
+
+  // Attach an interface and check that the type now has the interface.
+  IntegerType::attachInterface<Model>(context);
+  TestExternalTypeInterface iface = i8.dyn_cast<TestExternalTypeInterface>();
+  ASSERT_TRUE(iface != nullptr);
+  EXPECT_EQ(iface.getBitwidthPlusArg(10), 18u);
+  EXPECT_EQ(iface.staticGetSomeValuePlusArg(0), 42u);
+  EXPECT_EQ(iface.getBitwidthPlusDoubleArgument(2), 12u);
+  EXPECT_EQ(iface.staticGetArgument(17), 17u);
+
+  // Same, but with the default implementation overridden.
+  FloatType flt = Float32Type::get(&context);
+  ASSERT_FALSE(flt.isa<TestExternalTypeInterface>());
+  Float32Type::attachInterface<OverridingModel>(context);
+  iface = flt.dyn_cast<TestExternalTypeInterface>();
+  ASSERT_TRUE(iface != nullptr);
+  EXPECT_EQ(iface.getBitwidthPlusArg(10), 42u);
+  EXPECT_EQ(iface.staticGetSomeValuePlusArg(10), 52u);
+  EXPECT_EQ(iface.getBitwidthPlusDoubleArgument(3), 128u);
+  EXPECT_EQ(iface.staticGetArgument(17), 420u);
+
+  // Other contexts shouldn't have the attribute attached.
+  MLIRContext other;
+  IntegerType i8other = IntegerType::get(&other, 8);
+  EXPECT_FALSE(i8other.isa<TestExternalTypeInterface>());
+}
+
+/// The interface provides a default implementation that expects
+/// ConcreteType::getWidth to exist, which is the case for IntegerType. So this
+/// just derives from the ExternalModel.
+struct TestExternalFallbackTypeIntegerModel
+    : public TestExternalFallbackTypeInterface::ExternalModel<
+          TestExternalFallbackTypeIntegerModel, IntegerType> {};
+
+/// The interface provides a default implementation that expects
+/// ConcreteType::getWidth to exist, which is *not* the case for VectorType. Use
+/// FallbackModel instead to override this and make sure the code still compiles
+/// because we never instantiate the ExternalModel class template with a
+/// template argument that would have led to compilation failures.
+struct TestExternalFallbackTypeVectorModel
+    : public TestExternalFallbackTypeInterface::FallbackModel<
+          TestExternalFallbackTypeVectorModel> {
+  unsigned getBitwidth(Type type) const {
+    IntegerType elementType = type.cast<VectorType>()
+                                  .getElementType()
+                                  .dyn_cast_or_null<IntegerType>();
+    return elementType ? elementType.getWidth() : 0;
+  }
+};
+
+TEST(InterfaceAttachment, Fallback) {
+  MLIRContext context;
+
+  // Just check that we can attach the interface.
+  IntegerType i8 = IntegerType::get(&context, 8);
+  ASSERT_FALSE(i8.isa<TestExternalFallbackTypeInterface>());
+  IntegerType::attachInterface<TestExternalFallbackTypeIntegerModel>(context);
+  ASSERT_TRUE(i8.isa<TestExternalFallbackTypeInterface>());
+
+  // Call the method so it is guaranteed not to be instantiated.
+  VectorType vec = VectorType::get({42}, i8);
+  ASSERT_FALSE(vec.isa<TestExternalFallbackTypeInterface>());
+  VectorType::attachInterface<TestExternalFallbackTypeVectorModel>(context);
+  ASSERT_TRUE(vec.isa<TestExternalFallbackTypeInterface>());
+  EXPECT_EQ(vec.cast<TestExternalFallbackTypeInterface>().getBitwidth(), 8u);
+}
+
+/// External model for attribute interfaces.
+struct TextExternalIntegerAttrModel
+    : public TestExternalAttrInterface::ExternalModel<
+          TextExternalIntegerAttrModel, IntegerAttr> {
+  const Dialect *getDialectPtr(Attribute attr) const {
+    return &attr.cast<IntegerAttr>().getDialect();
+  }
+
+  static int getSomeNumber() { return 42; }
+};
+
+TEST(InterfaceAttachment, Attribute) {
+  MLIRContext context;
+
+  // Attribute interfaces use the exact same mechanism as types, so just check
+  // that the basics work for attributes.
+  IntegerAttr attr = IntegerAttr::get(IntegerType::get(&context, 32), 42);
+  ASSERT_FALSE(attr.isa<TestExternalAttrInterface>());
+  IntegerAttr::attachInterface<TextExternalIntegerAttrModel>(context);
+  auto iface = attr.dyn_cast<TestExternalAttrInterface>();
+  ASSERT_TRUE(iface != nullptr);
+  EXPECT_EQ(iface.getDialectPtr(), &attr.getDialect());
+  EXPECT_EQ(iface.getSomeNumber(), 42);
+}
+
+} // end namespace