[mlir] Enable delayed registration of attribute/operation/type interfaces

This functionality is similar to delayed registration of dialect interfaces. It
allows external interface models to be registered before the dialect containing
the attribute/operation/type interface is loaded, or even before the context is
created.

Reviewed By: rriddle

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

Author: ftynse

mlir/include/mlir/IR/AttributeSupport.h

@@ -50,6 +50,12 @@ class AbstractAttribute {
     return interfaceMap.lookup<T>();
   }
 
+  /// Returns true if the attribute has the interface with the given ID
+  /// registered.
+  bool hasInterface(TypeID interfaceID) const {
+    return interfaceMap.contains(interfaceID);
+  }
+
   /// Return the unique identifier representing the concrete attribute class.
   TypeID getTypeID() const { return typeID; }
 

mlir/include/mlir/IR/Dialect.h

@@ -27,8 +27,9 @@ class Type;
 
 using DialectAllocatorFunction = std::function<Dialect *(MLIRContext *)>;
 using DialectAllocatorFunctionRef = function_ref<Dialect *(MLIRContext *)>;
-using InterfaceAllocatorFunction =
+using DialectInterfaceAllocatorFunction =
     std::function<std::unique_ptr<DialectInterface>(Dialect *)>;
+using ObjectInterfaceAllocatorFunction = std::function<void(MLIRContext *)>;
 
 /// Dialects are groups of MLIR operations, types and attributes, as well as
 /// behavior associated with the entire group.  For example, hooks into other
@@ -271,11 +272,19 @@ class Dialect {
 /// dialects loaded in the Context. The parser in particular will lazily load
 /// dialects in the Context as operations are encountered.
 class DialectRegistry {
+  /// Lists of interfaces that need to be registered when the dialect is loaded.
+  struct DelayedInterfaces {
+    /// Dialect interfaces.
+    SmallVector<std::pair<TypeID, DialectInterfaceAllocatorFunction>, 2>
+        dialectInterfaces;
+    /// Attribute/Operation/Type interfaces.
+    SmallVector<std::pair<TypeID, ObjectInterfaceAllocatorFunction>, 2>
+        objectInterfaces;
+  };
+
   using MapTy =
       std::map<std::string, std::pair<TypeID, DialectAllocatorFunction>>;
-  using InterfaceMapTy =
-      DenseMap<TypeID,
-               SmallVector<std::pair<TypeID, InterfaceAllocatorFunction>, 2>>;
+  using InterfaceMapTy = DenseMap<TypeID, DelayedInterfaces>;
 
 public:
   explicit DialectRegistry() {}
@@ -329,7 +338,7 @@ class DialectRegistry {
   /// the registry.
   template <typename DialectTy>
   void addDialectInterface(TypeID interfaceTypeID,
-                           InterfaceAllocatorFunction allocator) {
+                           DialectInterfaceAllocatorFunction allocator) {
     addDialectInterface(DialectTy::getDialectNamespace(), interfaceTypeID,
                         allocator);
   }
@@ -344,6 +353,36 @@ class DialectRegistry {
         });
   }
 
+  /// Add an external op interface model for an op that belongs to a dialect,
+  /// both provided as template parameters. The dialect must be present in the
+  /// registry.
+  template <typename OpTy, typename ModelTy>
+  void addOpInterface() {
+    StringRef opName = OpTy::getOperationName();
+    StringRef dialectName = opName.split('.').first;
+    addObjectInterface(dialectName == opName ? "" : dialectName,
+                       ModelTy::Interface::getInterfaceID(),
+                       [](MLIRContext *context) {
+                         OpTy::template attachInterface<ModelTy>(*context);
+                       });
+  }
+
+  /// Add an external attribute interface model for an attribute type `AttrTy`
+  /// that is going to belong to `DialectTy`. The dialect must be present in the
+  /// registry.
+  template <typename DialectTy, typename AttrTy, typename ModelTy>
+  void addAttrInterface() {
+    addStorageUserInterface<AttrTy, ModelTy>(DialectTy::getDialectNamespace());
+  }
+
+  /// Add an external type interface model for an type class `TypeTy` that is
+  /// going to belong to `DialectTy`. The dialect must be present in the
+  /// registry.
+  template <typename DialectTy, typename TypeTy, typename ModelTy>
+  void addTypeInterface() {
+    addStorageUserInterface<TypeTy, ModelTy>(DialectTy::getDialectNamespace());
+  }
+
   /// Register any interfaces required for the given dialect (based on its
   /// TypeID). Users are not expected to call this directly.
   void registerDelayedInterfaces(Dialect *dialect) const;
@@ -352,7 +391,22 @@ class DialectRegistry {
   /// Add an interface constructed with the given allocation function to the
   /// dialect identified by its namespace.
   void addDialectInterface(StringRef dialectName, TypeID interfaceTypeID,
-                           InterfaceAllocatorFunction allocator);
+                           DialectInterfaceAllocatorFunction allocator);
+
+  /// Add an attribute/operation/type interface constructible with the given
+  /// allocation function to the dialect identified by its namespace.
+  void addObjectInterface(StringRef dialectName, TypeID interfaceTypeID,
+                          ObjectInterfaceAllocatorFunction allocator);
+
+  /// Add an external model for an attribute/type interface to the dialect
+  /// identified by its namespace.
+  template <typename ObjectTy, typename ModelTy>
+  void addStorageUserInterface(StringRef dialectName) {
+    addObjectInterface(dialectName, ModelTy::Interface::getInterfaceID(),
+                       [](MLIRContext *context) {
+                         ObjectTy::template attachInterface<ModelTy>(*context);
+                       });
+  }
 
   MapTy registry;
   InterfaceMapTy interfaces;

mlir/include/mlir/IR/TypeSupport.h

@@ -58,6 +58,11 @@ class AbstractType {
     return interfaceMap.lookup<T>();
   }
 
+  /// Returns true if the type has the interface with the given ID.
+  bool hasInterface(TypeID interfaceID) const {
+    return interfaceMap.contains(interfaceID);
+  }
+
   /// Return the unique identifier representing the concrete type class.
   TypeID getTypeID() const { return typeID; }
 

mlir/include/mlir/Support/InterfaceSupport.h

@@ -15,6 +15,7 @@
 
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/TypeName.h"
 
 namespace mlir {
@@ -235,8 +236,10 @@ class InterfaceMap {
           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");
+      if (it != interfaces.end() && it->first == id) {
+        LLVM_DEBUG(llvm::dbgs() << "Ignoring repeated interface registration");
+        continue;
+      }
       interfaces.insert(it, element);
     }
   }

mlir/lib/IR/Dialect.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/BuiltinDialect.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/DialectInterface.h"
@@ -31,7 +32,7 @@ DialectAsmParser::~DialectAsmParser() {}
 
 void DialectRegistry::addDialectInterface(
     StringRef dialectName, TypeID interfaceTypeID,
-    InterfaceAllocatorFunction allocator) {
+    DialectInterfaceAllocatorFunction allocator) {
   assert(allocator && "unexpected null interface allocation function");
   auto it = registry.find(dialectName.str());
   assert(it != registry.end() &&
@@ -40,8 +41,8 @@ void DialectRegistry::addDialectInterface(
   // Bail out if the interface with the given ID is already in the registry for
   // the given dialect. We expect a small number (dozens) of interfaces so a
   // linear search is fine here.
-  auto &dialectInterfaces = interfaces[it->second.first];
-  for (const auto &kvp : dialectInterfaces) {
+  auto &ifaces = interfaces[it->second.first];
+  for (const auto &kvp : ifaces.dialectInterfaces) {
     if (kvp.first == interfaceTypeID) {
       LLVM_DEBUG(llvm::dbgs()
                  << "[" DEBUG_TYPE
@@ -51,7 +52,36 @@ void DialectRegistry::addDialectInterface(
     }
   }
 
-  dialectInterfaces.emplace_back(interfaceTypeID, allocator);
+  ifaces.dialectInterfaces.emplace_back(interfaceTypeID, allocator);
+}
+
+void DialectRegistry::addObjectInterface(
+    StringRef dialectName, TypeID interfaceTypeID,
+    ObjectInterfaceAllocatorFunction allocator) {
+  assert(allocator && "unexpected null interface allocation function");
+
+  // Builtin dialect has an empty prefix and is always registered.
+  TypeID dialectTypeID;
+  if (!dialectName.empty()) {
+    auto it = registry.find(dialectName.str());
+    assert(it != registry.end() &&
+           "adding an interface for an op from an unregistered dialect");
+    dialectTypeID = it->second.first;
+  } else {
+    dialectTypeID = TypeID::get<BuiltinDialect>();
+  }
+
+  auto &ifaces = interfaces[dialectTypeID];
+  for (const auto &kvp : ifaces.objectInterfaces) {
+    if (kvp.first == interfaceTypeID) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "[" DEBUG_TYPE
+                    "] repeated interface object interface registration");
+      return;
+    }
+  }
+
+  ifaces.objectInterfaces.emplace_back(interfaceTypeID, allocator);
 }
 
 DialectAllocatorFunctionRef
@@ -79,11 +109,15 @@ void DialectRegistry::registerDelayedInterfaces(Dialect *dialect) const {
     return;
 
   // Add an interface if it is not already present.
-  for (const auto &kvp : it->second) {
+  for (const auto &kvp : it->getSecond().dialectInterfaces) {
     if (dialect->getRegisteredInterface(kvp.first))
       continue;
     dialect->addInterface(kvp.second(dialect));
   }
+
+  // Add attribute, operation and type interfaces.
+  for (const auto &kvp : it->getSecond().objectInterfaces)
+    kvp.second(dialect->getContext());
 }
 
 //===----------------------------------------------------------------------===//

mlir/lib/IR/MLIRContext.cpp

@@ -355,12 +355,12 @@ MLIRContext::MLIRContext(const DialectRegistry &registry)
     printStackTraceOnDiagnostic(clOptions->printStackTraceOnDiagnostic);
   }
 
-  // Ensure the builtin dialect is always pre-loaded.
-  getOrLoadDialect<BuiltinDialect>();
-
   // Pre-populate the registry.
   registry.appendTo(impl->dialectsRegistry);
 
+  // Ensure the builtin dialect is always pre-loaded.
+  getOrLoadDialect<BuiltinDialect>();
+
   // Initialize several common attributes and types to avoid the need to lock
   // the context when accessing them.