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

[ftynse]

Summary

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.

[rriddle]

mlir/include/mlir/IR/Dialect.h

373–378

Could we just have attachInterface ignore duplicates?

415–419

Same here.

mlir/unittests/IR/InterfaceAttachmentTest.cpp

342

typo: expceted

[ftynse]

Address review

[awarzynski]

Hi @ftynse ,

It looks like this change is causing all our Flang BuildBot workers to fail :( You can take a look here:

• https://lab.llvm.org/buildbot/#/builders/79/builds/11159
• https://lab.llvm.org/buildbot/#/builders/21/builds/17456

Would you be able to either fix or revert? Thank you!

[awarzynski]

mlir/include/mlir/Support/InterfaceSupport.h

241

I think that you are missing DEBUG_TYPE in this file: https://github.com/llvm/llvm-project/blob/main/llvm/include/llvm/Support/Debug.h#L15

LLVM_DEBUG() requires the DEBUG_TYPE macro to be defined

[ftynse]

In D104397#2824440, @awarzynski wrote:

Hi @ftynse ,

It looks like this change is causing all our Flang BuildBot workers to fail :( You can take a look here:

• https://lab.llvm.org/buildbot/#/builders/79/builds/11159
• https://lab.llvm.org/buildbot/#/builders/21/builds/17456

Would you be able to either fix or revert? Thank you!

This was fixed forward.

mlir/include/mlir/IR/AttributeSupport.h

public:
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this attribute, null otherwise. This should not be used
   /// directly.
   template <typename T> typename T::Concept *getInterface() const {
     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; }
 
 private:
   AbstractAttribute(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
                     TypeID typeID)
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
         typeID(typeID) {}

mlir/include/mlir/IR/Dialect.h

 class DialectAsmParser;
 class DialectAsmPrinter;
 class DialectInterface;
 class OpBuilder;
 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
 /// systems for constant folding, interfaces, default named types for asm
 /// printing, etc.
 ///
 /// Instances of the dialect object are loaded in a specific 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 {
+  /// 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() {}
 
   template <typename ConcreteDialect>
   void insert() {
     insert(TypeID::get<ConcreteDialect>(),
            ConcreteDialect::getDialectNamespace(),
return llvm::map_range(
         [](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(TypeID interfaceTypeID,
-                           InterfaceAllocatorFunction allocator) {
+                           DialectInterfaceAllocatorFunction 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>(
         InterfaceTy::getInterfaceID(), [](Dialect *dialect) {
           return std::make_unique<InterfaceTy>(dialect);
         });
   }
 
+  /// 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;
 
 private:
   /// 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;
 };
 
 } // 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/include/mlir/IR/TypeSupport.h

public:
 
   /// Returns an instance of the concept object for the given interface if it
   /// was registered to this type, null otherwise. This should not be used
   /// directly.
   template <typename T> typename T::Concept *getInterface() const {
     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; }
 
 private:
   AbstractType(Dialect &dialect, detail::InterfaceMap &&interfaceMap,
                TypeID typeID)
       : dialect(dialect), interfaceMap(std::move(interfaceMap)),
         typeID(typeID) {}

mlir/include/mlir/Support/InterfaceSupport.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SUPPORT_INTERFACESUPPORT_H
 #define MLIR_SUPPORT_INTERFACESUPPORT_H
 
 #include "mlir/Support/TypeID.h"
 #include "llvm/ADT/ArrayRef.h"
 #include "llvm/ADT/DenseMap.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/TypeName.h"
 
 namespace mlir {
 namespace detail {
 //===----------------------------------------------------------------------===//
 // Interface
 //===----------------------------------------------------------------------===//
 
std::pair<TypeID, void *> elements[] = {
                        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");
+      if (it != interfaces.end() && it->first == id) {
+        LLVM_DEBUG(llvm::dbgs() << "Ignoring repeated interface registration");
+        continue;
+      }
       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();

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/BuiltinDialect.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, TypeID interfaceTypeID,
-    InterfaceAllocatorFunction allocator) {
+    DialectInterfaceAllocatorFunction 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");
 
   // 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
                     "] repeated interface registration for dialect "
                  << dialectName);
       return;
     }
   }
 
-  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
 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 {
   auto it = interfaces.find(dialect->getTypeID());
   if (it == interfaces.end())
     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());
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)
     : name(name), dialectID(id), context(context) {

mlir/lib/IR/MLIRContext.cpp

MLIRContext::MLIRContext(const DialectRegistry &registry)
     : impl(new MLIRContextImpl) {
   // 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.
-  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.
 
   //// Types.
   /// Floating-point Types.
   impl->bf16Ty = TypeUniquer::get<BFloat16Type>(this);
   impl->f16Ty = TypeUniquer::get<Float16Type>(this);
   impl->f32Ty = TypeUniquer::get<Float32Type>(this);

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/BuiltinDialect.h"
 #include "mlir/IR/BuiltinOps.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "gtest/gtest.h"
 
 #include "../../test/lib/Dialect/Test/TestAttributes.h"
 #include "../../test/lib/Dialect/Test/TestDialect.h"
 #include "../../test/lib/Dialect/Test/TestTypes.h"
 
TEST(InterfaceAttachment, Type) {
   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>());
 }
 
+/// External interface model for the test type from the test dialect.
+struct TestTypeModel
+    : public TestExternalTypeInterface::ExternalModel<TestTypeModel,
+                                                      test::TestType> {
+  unsigned getBitwidthPlusArg(Type type, unsigned arg) const { return arg; }
+
+  static unsigned staticGetSomeValuePlusArg(unsigned arg) { return 10 + arg; }
+};
+
+TEST(InterfaceAttachment, TypeDelayedContextConstruct) {
+  // Put the interface in the registry.
+  DialectRegistry registry;
+  registry.insert<test::TestDialect>();
+  registry.addTypeInterface<test::TestDialect, test::TestType, TestTypeModel>();
+
+  // Check that when a context is constructed with the given registry, the type
+  // interface gets registered.
+  MLIRContext context(registry);
+  context.loadDialect<test::TestDialect>();
+  test::TestType testType = test::TestType::get(&context);
+  auto iface = testType.dyn_cast<TestExternalTypeInterface>();
+  ASSERT_TRUE(iface != nullptr);
+  EXPECT_EQ(iface.getBitwidthPlusArg(42), 42u);
+  EXPECT_EQ(iface.staticGetSomeValuePlusArg(10), 20u);
+}
+
+TEST(InterfaceAttachment, TypeDelayedContextAppend) {
+  // Put the interface in the registry.
+  DialectRegistry registry;
+  registry.insert<test::TestDialect>();
+  registry.addTypeInterface<test::TestDialect, test::TestType, TestTypeModel>();
+
+  // Check that when the registry gets appended to the context, the interface
+  // becomes available for objects in loaded dialects.
+  MLIRContext context;
+  context.loadDialect<test::TestDialect>();
+  test::TestType testType = test::TestType::get(&context);
+  EXPECT_FALSE(testType.isa<TestExternalTypeInterface>());
+  context.appendDialectRegistry(registry);
+  EXPECT_TRUE(testType.isa<TestExternalTypeInterface>());
+}
+
+TEST(InterfaceAttachment, RepeatedRegistration) {
+  DialectRegistry registry;
+  registry.addTypeInterface<BuiltinDialect, IntegerType, Model>();
+  MLIRContext context(registry);
+
+  // Should't fail on repeated registration through the dialect registry.
+  context.appendDialectRegistry(registry);
+}
+
+TEST(InterfaceAttachment, TypeBuiltinDelayed) {
+  // Builtin dialect needs to registration or loading, but delayed interface
+  // registration must still work.
+  DialectRegistry registry;
+  registry.addTypeInterface<BuiltinDialect, IntegerType, Model>();
+
+  MLIRContext context(registry);
+  IntegerType i16 = IntegerType::get(&context, 16);
+  EXPECT_TRUE(i16.isa<TestExternalTypeInterface>());
+
+  MLIRContext initiallyEmpty;
+  IntegerType i32 = IntegerType::get(&initiallyEmpty, 32);
+  EXPECT_FALSE(i32.isa<TestExternalTypeInterface>());
+  initiallyEmpty.appendDialectRegistry(registry);
+  EXPECT_TRUE(i32.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
TEST(InterfaceAttachment, Fallback) {
   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
+struct TestExternalIntegerAttrModel
     : public TestExternalAttrInterface::ExternalModel<
-          TextExternalIntegerAttrModel, IntegerAttr> {
+          TestExternalIntegerAttrModel, 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);
+  IntegerAttr::attachInterface<TestExternalIntegerAttrModel>(context);
   auto iface = attr.dyn_cast<TestExternalAttrInterface>();
   ASSERT_TRUE(iface != nullptr);
   EXPECT_EQ(iface.getDialectPtr(), &attr.getDialect());
   EXPECT_EQ(iface.getSomeNumber(), 42);
 }
 
+/// External model for an interface attachable to a non-builtin attribute.
+struct TestExternalSimpleAAttrModel
+    : public TestExternalAttrInterface::ExternalModel<
+          TestExternalSimpleAAttrModel, test::SimpleAAttr> {
+  const Dialect *getDialectPtr(Attribute attr) const {
+    return &attr.getDialect();
+  }
+
+  static int getSomeNumber() { return 21; }
+};
+
+TEST(InterfaceAttachmentTest, AttributeDelayed) {
+  // Attribute interfaces use the exact same mechanism as types, so just check
+  // that the delayed registration work for attributes.
+  DialectRegistry registry;
+  registry.insert<test::TestDialect>();
+  registry.addAttrInterface<test::TestDialect, test::SimpleAAttr,
+                            TestExternalSimpleAAttrModel>();
+
+  MLIRContext context(registry);
+  context.loadDialect<test::TestDialect>();
+  auto attr = test::SimpleAAttr::get(&context);
+  EXPECT_TRUE(attr.isa<TestExternalAttrInterface>());
+
+  MLIRContext initiallyEmpty;
+  initiallyEmpty.loadDialect<test::TestDialect>();
+  attr = test::SimpleAAttr::get(&initiallyEmpty);
+  EXPECT_FALSE(attr.isa<TestExternalAttrInterface>());
+  initiallyEmpty.appendDialectRegistry(registry);
+  EXPECT_TRUE(attr.isa<TestExternalAttrInterface>());
+}
+
 /// External interface model for the module operation. Only provides non-default
 /// methods.
 struct TestExternalOpModel
     : public TestExternalOpInterface::ExternalModel<TestExternalOpModel,
                                                     ModuleOp> {
   unsigned getNameLengthPlusArg(Operation *op, unsigned arg) const {
     return op->getName().getStringRef().size() + arg;
   }
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()));
 }
 
+struct TestExternalTestOpModel
+    : public TestExternalOpInterface::ExternalModel<TestExternalTestOpModel,
+                                                    test::OpJ> {
+  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;
+  }
+};
+
+TEST(InterfaceAttachment, OperationDelayedContextConstruct) {
+  DialectRegistry registry;
+  registry.insert<test::TestDialect>();
+  registry.addOpInterface<ModuleOp, TestExternalOpModel>();
+  registry.addOpInterface<test::OpJ, TestExternalTestOpModel>();
+
+  // 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 =
+      builder.create<test::OpJ>(builder.getUnknownLoc(), builder.getI32Type());
+  EXPECT_TRUE(isa<TestExternalOpInterface>(module.getOperation()));
+  EXPECT_TRUE(isa<TestExternalOpInterface>(op.getOperation()));
+}
+
+TEST(InterfaceAttachment, OperationDelayedContextAppend) {
+  DialectRegistry registry;
+  registry.insert<test::TestDialect>();
+  registry.addOpInterface<ModuleOp, TestExternalOpModel>();
+  registry.addOpInterface<test::OpJ, TestExternalTestOpModel>();
+
+  // 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 =
+      builder.create<test::OpJ>(builder.getUnknownLoc(), builder.getI32Type());
+  EXPECT_FALSE(isa<TestExternalOpInterface>(module.getOperation()));
+  EXPECT_FALSE(isa<TestExternalOpInterface>(op.getOperation()));
+  context.appendDialectRegistry(registry);
+  EXPECT_TRUE(isa<TestExternalOpInterface>(module.getOperation()));
+  EXPECT_TRUE(isa<TestExternalOpInterface>(op.getOperation()));
+}
+
 } // end namespace