[mlir] enable delayed registration of dialect interfaces

[ftynse]

Summary

This introduces a mechanism to register interfaces for a dialect without making
the dialect itself depend on the interface. The registration request happens on
DialectRegistry and, if the dialect has not been loaded yet, the actual
registration is delayed until the dialect is loaded. It requires
DialectRegistry to become aware of the context that contains it and the context
to expose methods for querying if a dialect is loaded.

This mechanism will enable a simple extension mechanism for dialects that can
have interfaces defined outside of the dialect code. It is particularly helpful
for, e.g., translation to LLVM IR where we don't want the dialect itself to
depend on LLVM IR libraries.

[mehdi_amini]

I was discussing a bit with River and we think we can avoid the back pointer to the context by restoring proper encapsulation by removing the getDialectRegistry() from MLIRContext, let me try to prepare a patch that does that.

mlir/include/mlir/IR/Dialect.h

280

I'd expect to see an "append" of the interfaces, it seems like you're replacing the ones in the destination here.

mlir/include/mlir/IR/MLIRContext.h

63 ↗

(On Diff #321745)

I'm not sure these API bear their weight here, can we just rely on getLoadedDialect?

This lead to pattern of double lookup, like you do in DialectRegistry:

if (owningContext->isDialectLoaded(dialectName))
  Dialect *dialect = owningContext->getLoadedDialect(dialectName);

[ftynse]

Address review

[ftynse]

I was discussing a bit with River and we think we can avoid the back pointer to the context by restoring proper encapsulation by removing the getDialectRegistry() from MLIRContext, let me try to prepare a patch that does that.

It looks quite widely used, so may be worth having a separate cleanup patch for that.

[ftynse]

The child commit https://reviews.llvm.org/D96331 makes getDialectRegistry return a const reference instead.

[mehdi_amini]

LG with the child.

mlir/include/mlir/IR/Dialect.h

 namespace mlir {
 class DialectAsmParser;
 class DialectAsmPrinter;
 class DialectInterface;
 class OpBuilder;
 class Type;
 
 using DialectAllocatorFunction = std::function<Dialect *(MLIRContext *)>;
+using InterfaceAllocatorFunction =
+    std::function<std::unique_ptr<DialectInterface>(Dialect *)>;
 
 /// Dialects are groups of MLIR operations, types and attributes, as well as
 /// behavior associated with the entire group.  For example, hooks into other
 /// systems for constant folding, interfaces, default named types for asm
 /// printing, etc.
 ///
 /// Instances of the dialect object are loaded in a specific MLIRContext.
 ///
private:
   /// Flag that specifies whether this dialect allows unregistered types, i.e.
   /// types prefixed with the dialect namespace but not registered with addType.
   /// These types are represented with OpaqueType.
   bool unknownTypesAllowed = false;
 
   /// A collection of registered dialect interfaces.
   DenseMap<TypeID, std::unique_ptr<DialectInterface>> registeredInterfaces;
 
+  friend class DialectRegistry;
   friend void registerDialect();
   friend class MLIRContext;
 };
 
 /// The DialectRegistry maps a dialect namespace to a constructor for the
 /// matching dialect.
 /// This allows for decoupling the list of dialects "available" from the
 /// dialects loaded in the Context. The parser in particular will lazily load
 /// dialects in the Context as operations are encountered.
 class DialectRegistry {
   using MapTy =
       std::map<std::string, std::pair<TypeID, DialectAllocatorFunction>>;
+  using InterfaceMapTy =
+      DenseMap<TypeID, SmallVector<InterfaceAllocatorFunction, 2>>;
 
 public:
+  explicit DialectRegistry(MLIRContext *context = nullptr)
+      : owningContext(context) {}
+
   template <typename ConcreteDialect>
   void insert() {
     insert(TypeID::get<ConcreteDialect>(),
            ConcreteDialect::getDialectNamespace(),
            static_cast<DialectAllocatorFunction>(([](MLIRContext *ctx) {
              // Just allocate the dialect, the context
              // takes ownership of it.
              return ctx->getOrLoadDialect<ConcreteDialect>();
            })));
   }
 
   template <typename ConcreteDialect, typename OtherDialect,
             typename... MoreDialects>
   void insert() {
     insert<ConcreteDialect>();
     insert<OtherDialect, MoreDialects...>();
   }
 
-  /// Add a new dialect constructor to the registry.
+  /// Add a new dialect constructor to the registry. The constructor must be
+  /// calling MLIRContext::getOrLoadDialect in order for the context to take
+  /// ownership of the dialect and for delayed interface registration to happen.
   void insert(TypeID typeID, StringRef name, DialectAllocatorFunction ctor);
 
   /// Load a dialect for this namespace in the provided context.
   Dialect *loadByName(StringRef name, MLIRContext *context);
 
   // Register all dialects available in the current registry with the registry
   // in the provided context.
   void appendTo(DialectRegistry &destination) {
     for (const auto &nameAndRegistrationIt : registry)
       destination.insert(nameAndRegistrationIt.second.first,
                          nameAndRegistrationIt.first,
                          nameAndRegistrationIt.second.second);
+    destination.interfaces.insert(interfaces.begin(), interfaces.end());
   }
   // Load all dialects available in the registry in the provided context.
   void loadAll(MLIRContext *context) {
     for (const auto &nameAndRegistrationIt : registry)
       nameAndRegistrationIt.second.second(context);
   }
 
-  MapTy::const_iterator begin() const { return registry.begin(); }
-  MapTy::const_iterator end() const { return registry.end(); }
+  /// Return the names of dialects known to this registry.
+  auto getDialectNames() {
+    return llvm::map_range(
+        registry, [](const MapTy::value_type &item) { return item.first; });
+  }
+
+  /// Add an interface constructed with the given allocation function to the
+  /// dialect provided as template parameter. The dialect must be present in
+  /// the registry, but may or may not be loaded. If it is not loaded, the
+  /// interface registration is delayed until the loading.
+  template <typename DialectTy>
+  void addDialectInterface(InterfaceAllocatorFunction allocator) {
+    addDialectInterface(DialectTy::getDialectNamespace(), allocator);
+  }
+
+  /// Add an interface to the dialect, both provided as template parameter. The
+  /// dialect must be present in the registry, but may or may not be loaded. If
+  /// it is not loaded, the interface registration is delayed until the loading.
+  template <typename DialectTy, typename InterfaceTy>
+  void addDialectInterface() {
+    addDialectInterface<DialectTy>([](Dialect *dialect) {
+      return std::make_unique<InterfaceTy>(dialect);
+    });
+  }
+
+  /// Register any interfaces required for the given dialect (based on its
+  /// TypeID). Users are not expected to call this directly.
+  void registerDelayedInterfaces(Dialect *dialect);
 
 private:
+  /// Add an interface constructed with the given allocation function to the
+  /// dialect identified by its namespace.
+  void addDialectInterface(StringRef dialectName,
+                           InterfaceAllocatorFunction allocator);
+
   MapTy registry;
+  InterfaceMapTy interfaces;
+
+  /// If this registry belongs to a context, this points back to the context.
+  /// Useful for checking if a dialect is loaded in the context.
+  MLIRContext *owningContext;
 };
 
 } // namespace mlir
 
 namespace llvm {
 /// Provide isa functionality for Dialects.
 template <typename T>
 struct isa_impl<T, ::mlir::Dialect> {
   static inline bool doit(const ::mlir::Dialect &dialect) {
     return mlir::TypeID::get<T>() == dialect.getTypeID();
   }
 };
 } // namespace llvm
 
 #endif

mlir/lib/IR/Dialect.cpp

 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Regex.h"
 
 using namespace mlir;
 using namespace detail;
 
 DialectAsmParser::~DialectAsmParser() {}
 
+//===----------------------------------------------------------------------===//
+// DialectRegistry
+//===----------------------------------------------------------------------===//
+
+void DialectRegistry::addDialectInterface(
+    StringRef dialectName, InterfaceAllocatorFunction allocator) {
+  assert(allocator && "unexpected null interface allocation function");
+
+  // If the dialect is already loaded, directly add the interface.
+  if (Dialect *dialect = owningContext
+                             ? owningContext->getLoadedDialect(dialectName)
+                             : nullptr) {
+    dialect->addInterface(allocator(dialect));
+    return;
+  }
+
+  // Otherwise, store it in the interface map for delayed registration.
+  auto it = registry.find(dialectName.str());
+  assert(it != registry.end() &&
+         "adding an interface for an unregistered dialect");
+  interfaces[it->second.first].push_back(allocator);
+}
+
 Dialect *DialectRegistry::loadByName(StringRef name, MLIRContext *context) {
   auto it = registry.find(name.str());
   if (it == registry.end())
     return nullptr;
   return it->second.second(context);
 }
 
 void DialectRegistry::insert(TypeID typeID, StringRef name,
                              DialectAllocatorFunction ctor) {
   auto inserted = registry.insert(
       std::make_pair(std::string(name), std::make_pair(typeID, ctor)));
   if (!inserted.second && inserted.first->second.first != typeID) {
     llvm::report_fatal_error(
         "Trying to register different dialects for the same namespace: " +
         name);
   }
 }
 
+void DialectRegistry::registerDelayedInterfaces(Dialect *dialect) {
+  auto it = interfaces.find(dialect->getTypeID());
+  if (it == interfaces.end())
+    return;
+
+  for (const InterfaceAllocatorFunction &createInterface : it->second)
+    dialect->addInterface(createInterface(dialect));
+}
+
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)
     : name(name), dialectID(id), context(context) {
   assert(isValidNamespace(name) && "invalid dialect namespace");
 }

mlir/lib/IR/MLIRContext.cpp

#endif
 
   /// Cached Attribute Instances.
   BoolAttr falseAttr, trueAttr;
   UnitAttr unitAttr;
   UnknownLoc unknownLocAttr;
   DictionaryAttr emptyDictionaryAttr;
 
 public:
-  MLIRContextImpl() : identifiers(identifierAllocator) {}
+  MLIRContextImpl(MLIRContext *ctx)
+      : dialectsRegistry(ctx), identifiers(identifierAllocator) {}
   ~MLIRContextImpl() {
     for (auto typeMapping : registeredTypes)
       typeMapping.second->~AbstractType();
     for (auto attrMapping : registeredAttributes)
       attrMapping.second->~AbstractAttribute();
   }
 };
 } // end namespace mlir
 
-MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
+MLIRContext::MLIRContext() : impl(new MLIRContextImpl(this)) {
   // Initialize values based on the command line flags if they were provided.
   if (clOptions.isConstructed()) {
     disableMultithreading(clOptions->disableThreading);
     printOpOnDiagnostic(clOptions->printOpOnDiagnostic);
     printStackTraceOnDiagnostic(clOptions->printStackTraceOnDiagnostic);
   }
 
   // Ensure the builtin dialect is always pre-loaded.
std::vector<Dialect *> MLIRContext::getLoadedDialects() {
   llvm::array_pod_sort(result.begin(), result.end(),
                        [](Dialect *const *lhs, Dialect *const *rhs) -> int {
                          return (*lhs)->getNamespace() < (*rhs)->getNamespace();
                        });
   return result;
 }
 std::vector<StringRef> MLIRContext::getAvailableDialects() {
   std::vector<StringRef> result;
-  for (auto &dialect : impl->dialectsRegistry)
-    result.push_back(dialect.first);
+  for (auto dialect : impl->dialectsRegistry.getDialectNames())
+    result.push_back(dialect);
   return result;
 }
 
 /// Get a registered IR dialect with the given namespace. If none is found,
 /// then return nullptr.
 Dialect *MLIRContext::getLoadedDialect(StringRef name) {
   // Dialects are sorted by name, so we can use binary search for lookup.
   auto it = impl->loadedDialects.find(name);
#endif
     // Refresh all the identifiers dialect field, this catches cases where a
     // dialect may be loaded after identifier prefixed with this dialect name
     // were already created.
     for (auto &identifierEntry : impl.identifiers)
       if (!identifierEntry.second &&
           identifierEntry.first().startswith(dialectNamespace))
         identifierEntry.second = dialect.get();
 
+    // Actually register the interfaces with delayed registration.
+    impl.dialectsRegistry.registerDelayedInterfaces(dialect.get());
     return dialect.get();
   }
 
   // Abort if dialect with namespace has already been registered.
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
 

mlir/lib/Support/MlirOptMain.cpp

LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   registerPassManagerCLOptions();
   PassPipelineCLParser passPipeline("", "Compiler passes to run");
 
   // Build the list of dialects as a header for the --help message.
   std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();
   {
     llvm::raw_string_ostream os(helpHeader);
     MLIRContext context;
-    interleaveComma(registry, os, [&](auto &registryEntry) {
-      StringRef name = registryEntry.first;
-      os << name;
-    });
+    interleaveComma(registry.getDialectNames(), os,
+                    [&](auto name) { os << name; });
   }
   // Parse pass names in main to ensure static initialization completed.
   cl::ParseCommandLineOptions(argc, argv, helpHeader);
 
   if (showDialects) {
     llvm::outs() << "Available Dialects:\n";
     interleave(
-        registry, llvm::outs(),
-        [](auto &registryEntry) { llvm::outs() << registryEntry.first; }, "\n");
+        registry.getDialectNames(), llvm::outs(),
+        [](auto name) { llvm::outs() << name; }, "\n");
     return success();
   }
 
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
     llvm::errs() << errorMessage << "\n";

mlir/unittests/IR/DialectTest.cpp

 //===- DialectTest.cpp - Dialect unit tests -------------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/IR/Dialect.h"
+#include "mlir/IR/DialectInterface.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 namespace {
 struct TestDialect : public Dialect {
   static StringRef getDialectNamespace() { return "test"; };
TEST(DialectDeathTest, MultipleDialectsWithSameNamespace) {
   MLIRContext context;
 
   // Registering a dialect with the same namespace twice should result in a
   // failure.
   context.loadDialect<TestDialect>();
   ASSERT_DEATH(context.loadDialect<AnotherTestDialect>(), "");
 }
 
+struct SecondTestDialect : public Dialect {
+  static StringRef getDialectNamespace() { return "test2"; }
+  SecondTestDialect(MLIRContext *context)
+      : Dialect(getDialectNamespace(), context,
+                TypeID::get<SecondTestDialect>()) {}
+};
+
+struct TestDialectInterfaceBase
+    : public DialectInterface::Base<TestDialectInterfaceBase> {
+  TestDialectInterfaceBase(Dialect *dialect) : Base(dialect) {}
+  virtual int function() const { return 42; }
+};
+
+struct TestDialectInterface : public TestDialectInterfaceBase {
+  using TestDialectInterfaceBase::TestDialectInterfaceBase;
+  int function() const final { return 56; }
+};
+
+struct SecondTestDialectInterface : public TestDialectInterfaceBase {
+  using TestDialectInterfaceBase::TestDialectInterfaceBase;
+  int function() const final { return 78; }
+};
+
+TEST(Dialect, DelayedInterfaceRegistration) {
+  DialectRegistry registry;
+  registry.insert<TestDialect, SecondTestDialect>();
+
+  // Delayed registration of an interface for TestDialect.
+  registry.addDialectInterface<TestDialect, TestDialectInterface>();
+
+  MLIRContext context;
+  registry.appendTo(context.getDialectRegistry());
+
+  // Load the TestDialect and check that the interface got registered for it.
+  auto *testDialect = context.getOrLoadDialect<TestDialect>();
+  ASSERT_TRUE(testDialect != nullptr);
+  auto *testDialectInterface =
+      testDialect->getRegisteredInterface<TestDialectInterfaceBase>();
+  EXPECT_TRUE(testDialectInterface != nullptr);
+
+  // Load the SecondTestDialect and check that the interface is not registered
+  // for it.
+  auto *secondTestDialect = context.getOrLoadDialect<SecondTestDialect>();
+  ASSERT_TRUE(secondTestDialect != nullptr);
+  auto *secondTestDialectInterface =
+      secondTestDialect->getRegisteredInterface<SecondTestDialectInterface>();
+  EXPECT_TRUE(secondTestDialectInterface == nullptr);
+
+  // Use the same mechanism as for delayed registration but for an already
+  // loaded dialect and check that the interface is now registered.
+  context.getDialectRegistry()
+      .addDialectInterface<SecondTestDialect, SecondTestDialectInterface>();
+  secondTestDialectInterface =
+      secondTestDialect->getRegisteredInterface<SecondTestDialectInterface>();
+  EXPECT_TRUE(secondTestDialectInterface != nullptr);
+}
+
 } // end namespace