[mlir] Fix delayed object interfaces registration

[vinograd47]

Summary

Store both interfaceID and objectID as key for interface registration callback.
Otherwise the implementation allows to register only one external model per one object in the single dialect.

[ftynse]

I wouldn't call this a fix. Having the same class implement the interface for different objects actually sounds like a new functionality.

[vinograd47]

In D107274#2921570, @ftynse wrote:

I wouldn't call this a fix. Having the same class implement the interface for different objects actually sounds like a new functionality.

It's not the same class. In the test I've used template class TestExternalTestOpModel<ConcreteOp>, so for different objects we have different types. It might be different classes (for example, TestExternalTestOpModelForJ and TestExternalTestOpModelForH), it doesn't matter. But DialectRegistration uses base interface type ID as key, so it doesn't distinguish different model classes for the same interface.

[vinograd47]

Or in other words, current implementation doesn't allow to attach the same interface implementation to several objects from one dialect. And it doesn't matter if it is a single interface model for all objects or it is different model classes.

mlir/include/mlir/IR/Dialect.h

 
 #ifndef MLIR_IR_DIALECT_H
 #define MLIR_IR_DIALECT_H
 
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/TypeID.h"
 
 #include <map>
+#include <tuple>
 
 namespace mlir {
 class DialectAsmParser;
 class DialectAsmPrinter;
 class DialectInterface;
 class OpBuilder;
 class Type;
 

 /// 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>
+    SmallVector<std::tuple<TypeID, TypeID, ObjectInterfaceAllocatorFunction>, 2>
         objectInterfaces;
   };
 
   using MapTy =
       std::map<std::string, std::pair<TypeID, DialectAllocatorFunction>>;
   using InterfaceMapTy = DenseMap<TypeID, DelayedInterfaces>;
 
 public:
public:
 
   /// 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, ModelTy::Interface::getInterfaceID(),
+    addObjectInterface(dialectName, TypeID::get<OpTy>(),
+                       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.

 private:
   /// Add an interface constructed with the given allocation function to the
   /// dialect identified by its namespace.
   void addDialectInterface(StringRef dialectName, TypeID interfaceTypeID,
                            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,
+  void addObjectInterface(StringRef dialectName, TypeID objectID,
+                          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(),
+    addObjectInterface(dialectName, TypeID::get<ObjectTy>(),
+                       ModelTy::Interface::getInterfaceID(),
                        [](MLIRContext *context) {
                          ObjectTy::template attachInterface<ModelTy>(*context);
                        });
   }
 
   MapTy registry;
   InterfaceMapTy interfaces;
 };

mlir/lib/IR/Dialect.cpp

if (kvp.first == interfaceTypeID) {
       return;
     }
   }
 
   ifaces.dialectInterfaces.emplace_back(interfaceTypeID, allocator);
 }
 
 void DialectRegistry::addObjectInterface(
-    StringRef dialectName, TypeID interfaceTypeID,
+    StringRef dialectName, TypeID objectID, TypeID interfaceTypeID,
     ObjectInterfaceAllocatorFunction allocator) {
   assert(allocator && "unexpected null interface allocation function");
+
   auto it = registry.find(dialectName.str());
   assert(it != registry.end() &&
          "adding an interface for an op from an unregistered dialect");
 
-  auto &ifaces = interfaces[it->second.first];
-  for (const auto &kvp : ifaces.objectInterfaces) {
-    if (kvp.first == interfaceTypeID) {
+  auto dialectID = it->second.first;
+  auto &ifaces = interfaces[dialectID];
+
+  for (const auto &info : ifaces.objectInterfaces) {
+    if (std::get<0>(info) == objectID && std::get<1>(info) == interfaceTypeID) {
       LLVM_DEBUG(llvm::dbgs()
                  << "[" DEBUG_TYPE
                     "] repeated interface object interface registration");
       return;
     }
   }
 
-  ifaces.objectInterfaces.emplace_back(interfaceTypeID, allocator);
+  ifaces.objectInterfaces.emplace_back(objectID, 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::registerDelayedInterfaces(Dialect *dialect) const {
   // Add an interface if it is not already present.
   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());
+  for (const auto &info : it->getSecond().objectInterfaces)
+    std::get<2>(info)(dialect->getContext());
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)
     : name(name), dialectID(id), context(context) {

mlir/unittests/IR/InterfaceAttachmentTest.cpp

TEST(InterfaceAttachment, Operation) {
   EXPECT_EQ(iface.getNameLengthMinusArg(1000), 21u);
 
   // Another context doesn't have the interfaces registered.
   MLIRContext other;
   auto otherModuleOp = ModuleOp::create(UnknownLoc::get(&other));
   ASSERT_FALSE(isa<TestExternalOpInterface>(otherModuleOp.getOperation()));
 }
 
+template <class ConcreteOp>
 struct TestExternalTestOpModel
-    : public TestExternalOpInterface::ExternalModel<TestExternalTestOpModel,
-                                                    test::OpJ> {
+    : public TestExternalOpInterface::ExternalModel<
+          TestExternalTestOpModel<ConcreteOp>, ConcreteOp> {
   unsigned getNameLengthPlusArg(Operation *op, unsigned arg) const {
     return op->getName().getStringRef().size() + arg;
   }
 
   static unsigned getNameLengthPlusArgTwice(unsigned arg) {
-    return test::OpJ::getOperationName().size() + 2 * arg;
+    return ConcreteOp::getOperationName().size() + 2 * arg;
   }
 };
 
 TEST(InterfaceAttachment, OperationDelayedContextConstruct) {
   DialectRegistry registry;
   registry.insert<test::TestDialect>();
   registry.addOpInterface<ModuleOp, TestExternalOpModel>();
-  registry.addOpInterface<test::OpJ, TestExternalTestOpModel>();
+  registry.addOpInterface<test::OpJ, TestExternalTestOpModel<test::OpJ>>();
+  registry.addOpInterface<test::OpH, TestExternalTestOpModel<test::OpH>>();
 
   // Construct the context directly from a registry. The interfaces are expected
   // to be readily available on operations.
   MLIRContext context(registry);
   context.loadDialect<test::TestDialect>();
+
   ModuleOp module = ModuleOp::create(UnknownLoc::get(&context));
   OpBuilder builder(module);
-  auto op =
+  auto opJ =
       builder.create<test::OpJ>(builder.getUnknownLoc(), builder.getI32Type());
+  auto opH =
+      builder.create<test::OpH>(builder.getUnknownLoc(), opJ.getResult());
+  auto opI =
+      builder.create<test::OpI>(builder.getUnknownLoc(), opJ.getResult());
+
   EXPECT_TRUE(isa<TestExternalOpInterface>(module.getOperation()));
-  EXPECT_TRUE(isa<TestExternalOpInterface>(op.getOperation()));
+  EXPECT_TRUE(isa<TestExternalOpInterface>(opJ.getOperation()));
+  EXPECT_TRUE(isa<TestExternalOpInterface>(opH.getOperation()));
+  EXPECT_FALSE(isa<TestExternalOpInterface>(opI.getOperation()));
 }
 
 TEST(InterfaceAttachment, OperationDelayedContextAppend) {
   DialectRegistry registry;
   registry.insert<test::TestDialect>();
   registry.addOpInterface<ModuleOp, TestExternalOpModel>();
-  registry.addOpInterface<test::OpJ, TestExternalTestOpModel>();
+  registry.addOpInterface<test::OpJ, TestExternalTestOpModel<test::OpJ>>();
+  registry.addOpInterface<test::OpH, TestExternalTestOpModel<test::OpH>>();
 
   // Construct the context, create ops, and only then append the registry. The
   // interfaces are expected to be available after appending the registry.
   MLIRContext context;
   context.loadDialect<test::TestDialect>();
+
   ModuleOp module = ModuleOp::create(UnknownLoc::get(&context));
   OpBuilder builder(module);
-  auto op =
+  auto opJ =
       builder.create<test::OpJ>(builder.getUnknownLoc(), builder.getI32Type());
+  auto opH =
+      builder.create<test::OpH>(builder.getUnknownLoc(), opJ.getResult());
+  auto opI =
+      builder.create<test::OpI>(builder.getUnknownLoc(), opJ.getResult());
+
   EXPECT_FALSE(isa<TestExternalOpInterface>(module.getOperation()));
-  EXPECT_FALSE(isa<TestExternalOpInterface>(op.getOperation()));
+  EXPECT_FALSE(isa<TestExternalOpInterface>(opJ.getOperation()));
+  EXPECT_FALSE(isa<TestExternalOpInterface>(opH.getOperation()));
+  EXPECT_FALSE(isa<TestExternalOpInterface>(opI.getOperation()));
+
   context.appendDialectRegistry(registry);
+
   EXPECT_TRUE(isa<TestExternalOpInterface>(module.getOperation()));
-  EXPECT_TRUE(isa<TestExternalOpInterface>(op.getOperation()));
+  EXPECT_TRUE(isa<TestExternalOpInterface>(opJ.getOperation()));
+  EXPECT_TRUE(isa<TestExternalOpInterface>(opH.getOperation()));
+  EXPECT_FALSE(isa<TestExternalOpInterface>(opI.getOperation()));
 }
 
 } // end namespace