[mlir] Support repeated delayed registration of dialect interfaces

[ftynse]

Summary

Dialects themselves do not support repeated addition of interfaces with the
same TypeID. However, in case of delayed registration, the registry may contain
such an interface, or have the same interface registered several times due to,
e.g., dependencies. Make sure we delayed registration does not attempt to add
an interface with the same TypeID more than once.

[mehdi_amini]

mlir/unittests/IR/DialectTest.cpp

120

I don't see a "death test" on double addition of different implementation of the same interface, can you try to add one?

[ftynse]

Add death test

mlir/include/mlir/IR/Dialect.h

 /// 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>>;
+      DenseMap<TypeID,
+               SmallVector<std::pair<TypeID, InterfaceAllocatorFunction>, 2>>;
 
 public:
   explicit DialectRegistry() {}
 
   template <typename ConcreteDialect>
   void insert() {
     insert(TypeID::get<ConcreteDialect>(),
            ConcreteDialect::getDialectNamespace(),
return llvm::map_range(
         registry,
         [](const MapTy::value_type &item) -> StringRef { 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.
   template <typename DialectTy>
-  void addDialectInterface(InterfaceAllocatorFunction allocator) {
-    addDialectInterface(DialectTy::getDialectNamespace(), allocator);
+  void addDialectInterface(TypeID interfaceTypeID,
+                           InterfaceAllocatorFunction allocator) {
+    addDialectInterface(DialectTy::getDialectNamespace(), interfaceTypeID,
+                        allocator);
   }
 
   /// Add an interface to the dialect, both provided as template parameter. The
   /// dialect must be present in the registry.
   template <typename DialectTy, typename InterfaceTy>
   void addDialectInterface() {
-    addDialectInterface<DialectTy>([](Dialect *dialect) {
+    addDialectInterface<DialectTy>(
+        InterfaceTy::getInterfaceID(), [](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) const;
 
 private:
   /// Add an interface constructed with the given allocation function to the
   /// dialect identified by its namespace.
-  void addDialectInterface(StringRef dialectName,
+  void addDialectInterface(StringRef dialectName, TypeID interfaceTypeID,
                            InterfaceAllocatorFunction allocator);
 
   MapTy registry;
   InterfaceMapTy interfaces;
 };
 
 } // namespace mlir
 

mlir/lib/IR/Dialect.cpp

 //===- Dialect.cpp - Dialect implementation -------------------------------===//
 //
 // 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/Diagnostics.h"
 #include "mlir/IR/DialectImplementation.h"
 #include "mlir/IR/DialectInterface.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Operation.h"
 #include "llvm/ADT/MapVector.h"
 #include "llvm/ADT/Twine.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/ManagedStatic.h"
 #include "llvm/Support/Regex.h"
 
+#define DEBUG_TYPE "dialect"
+
 using namespace mlir;
 using namespace detail;
 
 DialectAsmParser::~DialectAsmParser() {}
 
 //===----------------------------------------------------------------------===//
 // DialectRegistry
 //===----------------------------------------------------------------------===//
 
 void DialectRegistry::addDialectInterface(
-    StringRef dialectName, InterfaceAllocatorFunction allocator) {
+    StringRef dialectName, TypeID interfaceTypeID,
+    InterfaceAllocatorFunction allocator) {
   assert(allocator && "unexpected null interface allocation function");
   auto it = registry.find(dialectName.str());
   assert(it != registry.end() &&
          "adding an interface for an unregistered dialect");
-  interfaces[it->second.first].push_back(allocator);
+
+  // 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) {
+    if (kvp.first == interfaceTypeID) {
+      LLVM_DEBUG(llvm::dbgs()
+                 << "[" DEBUG_TYPE
+                    "] repeated interface registration for dialect "
+                 << dialectName);
+      return;
+    }
+  }
+
+  dialectInterfaces.emplace_back(interfaceTypeID, allocator);
 }
 
 DialectAllocatorFunctionRef
 DialectRegistry::getDialectAllocator(StringRef name) const {
   auto it = registry.find(name.str());
   if (it == registry.end())
     return nullptr;
   return it->second.second;
void DialectRegistry::insert(TypeID typeID, StringRef name,
   }
 }
 
 void DialectRegistry::registerDelayedInterfaces(Dialect *dialect) const {
   auto it = interfaces.find(dialect->getTypeID());
   if (it == interfaces.end())
     return;
 
-  for (const InterfaceAllocatorFunction &createInterface : it->second)
-    dialect->addInterface(createInterface(dialect));
+  // Add an interface if it is not already present.
+  for (const auto &kvp : it->second) {
+    if (dialect->getRegisteredInterface(kvp.first))
+      continue;
+    dialect->addInterface(kvp.second(dialect));
+  }
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)
     : name(name), dialectID(id), context(context) {

mlir/unittests/IR/DialectTest.cpp

TEST(Dialect, DelayedInterfaceRegistration) {
   secondRegistry
       .addDialectInterface<SecondTestDialect, SecondTestDialectInterface>();
   context.appendDialectRegistry(secondRegistry);
   secondTestDialectInterface =
       secondTestDialect->getRegisteredInterface<SecondTestDialectInterface>();
   EXPECT_TRUE(secondTestDialectInterface != nullptr);
 }
 
+TEST(Dialect, RepeatedDelayedRegistration) {
+  // Set up the delayed registration.
+  DialectRegistry registry;
+  registry.insert<TestDialect>();
+  registry.addDialectInterface<TestDialect, TestDialectInterface>();
+  MLIRContext context(registry);
+
+  // 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);
+
+  // Try adding the same dialect interface again and check that we don't crash
+  // on repeated interface registration.
+  DialectRegistry secondRegistry;
+  secondRegistry.insert<TestDialect>();
+  secondRegistry.addDialectInterface<TestDialect, TestDialectInterface>();
+  context.appendDialectRegistry(secondRegistry);
+  testDialectInterface =
+      testDialect->getRegisteredInterface<TestDialectInterfaceBase>();
+  EXPECT_TRUE(testDialectInterface != nullptr);
+}
+
+// A dialect that registers two interfaces with the same InterfaceID, triggering
+// an assertion failure.
+struct RepeatedRegistrationDialect : public Dialect {
+  static StringRef getDialectNamespace() { return "repeatedreg"; }
+  RepeatedRegistrationDialect(MLIRContext *context)
+      : Dialect(getDialectNamespace(), context,
+                TypeID::get<RepeatedRegistrationDialect>()) {
+    addInterfaces<TestDialectInterface>();
+    addInterfaces<SecondTestDialectInterface>();
+  }
+};
+
+TEST(Dialect, RepeatedInterfaceRegistrationDeath) {
+  MLIRContext context;
+  (void)context;
+
+  // This triggers an assertion in debug mode.
+#ifndef NDEBUG
+  ASSERT_DEATH(context.loadDialect<RepeatedRegistrationDialect>(),
+               "interface kind has already been registered");
+#endif
+}
+
 } // end namespace