[mlir] Refactor DialectRegistry delayed interface support into a general DialectExtension mechanism

[rriddle]

Summary

The current dialect registry allows for attaching delayed interfaces, that are added to attrs/dialects/ops/etc.
when the owning dialect gets loaded. This is clunky for quite a few reasons, e.g. each interface type has a
separate tracking structure, and is also quite limiting. This commit refactors this delayed mutation of
dialect constructs into a more general DialectExtension mechanism. This mechanism is essentially a registration
callback that is invoked when a set of dialects have been loaded. This allows for attaching interfaces directly
on the loaded constructs, and also allows for loading new dependent dialects. The latter of which is
extremely useful as it will now enable dependent dialects to only apply in the contexts in which they
are necessary. For example, a dialect dependency can now be conditional on if a user actually needs the
interface that relies on it.

[mehdi_amini]

This is very elegant! The combination of the trick for converting lambda to bare function pointer and deduce the Dialect dependency is quite interesting :)

mlir/include/mlir/IR/DialectRegistry.h

89

It wasn't obvious to me that this parameter wasn't used: can you comment it out maybe? Or name it unusedDialects?

142

Update the comment? Register all dialects and extensions ...

143

This API does not take a context anymore, the comment is stalled on this as well.

189

This is a neat trick!

204

The dtor declaration does not seem necessary here, is it?

212

I don't find any place where addExtension is called with a lambda with more than a single dialect?

mlir/lib/IR/Dialect.cpp

192

Can't you just have a for-range loop here?

202

You're building the requiredDialects for the general case (DialectExtensionBase), but actually the only implementation DialectExtension will ignore this and rebuild it entirely.
Do you anticipate other implementations of DialectExtensionBase that would use it?

[mehdi_amini]

LG with inline comments!

mlir/include/mlir/IR/Dialect.h

 //===- Dialect.h - IR Dialect Description -----------------------*- C++ -*-===//
 //
 // 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 file defines the 'dialect' abstraction.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_DIALECT_H
 #define MLIR_IR_DIALECT_H
 
+#include "mlir/IR/DialectRegistry.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;
 
-using DialectAllocatorFunction = std::function<Dialect *(MLIRContext *)>;
-using DialectAllocatorFunctionRef = function_ref<Dialect *(MLIRContext *)>;
-using DialectInterfaceAllocatorFunction =
-    std::function<std::unique_ptr<DialectInterface>(Dialect *)>;
-using ObjectInterfaceAllocatorFunction = std::function<void(MLIRContext *)>;
+//===----------------------------------------------------------------------===//
+// 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.
 ///
public:
   }
   template <typename InterfaceT>
   typename InterfaceT::Concept *
   getRegisteredInterfaceForOp(OperationName opName) {
     return static_cast<typename InterfaceT::Concept *>(
         getRegisteredInterfaceForOp(InterfaceT::getInterfaceID(), opName));
   }
 
+  /// Register a dialect interface with this dialect instance.
+  void addInterface(std::unique_ptr<DialectInterface> interface);
+
+  /// Register a set of dialect interfaces with this dialect instance.
+  template <typename... Args>
+  void addInterfaces() {
+    (void)std::initializer_list<int>{
+        0, (addInterface(std::make_unique<Args>(this)), 0)...};
+  }
+
 protected:
   /// The constructor takes a unique namespace for this dialect as well as the
   /// context to bind to.
   /// Note: The namespace must not contain '.' characters.
   /// Note: All operations belonging to this dialect must have names starting
   ///       with the namespace followed by '.'.
   /// Example:
   ///       - "tf" for the TensorFlow ops like "tf.add".
protected:
   }
 
   /// Enable support for unregistered operations.
   void allowUnknownOperations(bool allow = true) { unknownOpsAllowed = allow; }
 
   /// Enable support for unregistered types.
   void allowUnknownTypes(bool allow = true) { unknownTypesAllowed = allow; }
 
-  /// Register a dialect interface with this dialect instance.
-  void addInterface(std::unique_ptr<DialectInterface> interface);
-
-  /// Register a set of dialect interfaces with this dialect instance.
-  template <typename... Args> void addInterfaces() {
-    (void)std::initializer_list<int>{
-        0, (addInterface(std::make_unique<Args>(this)), 0)...};
-  }
-
 private:
   Dialect(const Dialect &) = delete;
   void operator=(Dialect &) = delete;
 
   /// Register an attribute instance with this dialect.
   template <typename T> void addAttribute() {
     // Add this attribute to the dialect and register it with the uniquer.
     addAttribute(T::getTypeID(), AbstractAttribute::get<T>(*this));
private:
   /// 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 {
-  /// 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::tuple<TypeID, TypeID, ObjectInterfaceAllocatorFunction>, 2>
-        objectInterfaces;
-  };
-
-  using MapTy =
-      std::map<std::string, std::pair<TypeID, DialectAllocatorFunction>>;
-  using InterfaceMapTy = DenseMap<TypeID, DelayedInterfaces>;
-
-public:
-  explicit DialectRegistry();
-
-  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. 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,
-              const DialectAllocatorFunction &ctor);
-
-  /// Return an allocation function for constructing the dialect identified by
-  /// its namespace, or nullptr if the namespace is not in this registry.
-  DialectAllocatorFunctionRef getDialectAllocator(StringRef name) const;
-
-  // Register all dialects available in the current registry with the registry
-  // in the provided context.
-  void appendTo(DialectRegistry &destination) const {
-    for (const auto &nameAndRegistrationIt : registry)
-      destination.insert(nameAndRegistrationIt.second.first,
-                         nameAndRegistrationIt.first,
-                         nameAndRegistrationIt.second.second);
-    // Merge interfaces.
-    for (auto it : interfaces) {
-      TypeID dialect = it.first;
-      auto destInterfaces = destination.interfaces.find(dialect);
-      if (destInterfaces == destination.interfaces.end()) {
-        destination.interfaces[dialect] = it.second;
-        continue;
-      }
-      // The destination already has delayed interface registrations for this
-      // dialect. Merge registrations into the destination registry.
-      destInterfaces->second.dialectInterfaces.append(
-          it.second.dialectInterfaces.begin(),
-          it.second.dialectInterfaces.end());
-      destInterfaces->second.objectInterfaces.append(
-          it.second.objectInterfaces.begin(), it.second.objectInterfaces.end());
-    }
-  }
-
-  /// Return the names of dialects known to this registry.
-  auto getDialectNames() const {
-    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(TypeID interfaceTypeID,
-                           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, 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.
-  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,
-                           const 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 objectID,
-                          TypeID interfaceTypeID,
-                          const 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, TypeID::get<ObjectTy>(),
-                       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,
                 std::enable_if_t<std::is_base_of<::mlir::Dialect, T>::value>> {
   static inline bool doit(const ::mlir::Dialect &dialect) {

mlir/include/mlir/IR/DialectRegistry.h

+//===- DialectRegistry.h - Dialect Registration and Extension ---*- C++ -*-===//
+//
+// 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 file defines functionality for registring and extending dialects.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef MLIR_IR_DIALECTREGISTRY_H
+#define MLIR_IR_DIALECTREGISTRY_H
+
+#include "mlir/IR/MLIRContext.h"
+#include "llvm/ADT/ArrayRef.h"
+#include "llvm/ADT/SmallVector.h"
+#include "llvm/ADT/StringRef.h"
+
+#include <map>
+#include <tuple>
+
+namespace mlir {
+class Dialect;
+
+using DialectAllocatorFunction = std::function<Dialect *(MLIRContext *)>;
+using DialectAllocatorFunctionRef = function_ref<Dialect *(MLIRContext *)>;
+
+//===----------------------------------------------------------------------===//
+// DialectExtension
+//===----------------------------------------------------------------------===//
+
+/// This class represents an opaque dialect extension. It contains a set of
+/// required dialects and an application function. The required dialects control
+/// when the extension is applied, i.e. the extension is applied when all
+/// required dialects are loaded. The application function can be used to attach
+/// additional functionality to attributes, dialects, operations, types, etc.,
+/// and may also load additional necessary dialects.
+class DialectExtensionBase {
+public:
+  virtual ~DialectExtensionBase();
+
+  /// Return the dialects that our required by this extension to be loaded
+  /// before applying.
+  ArrayRef<StringRef> getRequiredDialects() const { return dialectNames; }
+
+  /// Apply this extension to the given context and the required dialects.
+  virtual void apply(MLIRContext *context,
+                     MutableArrayRef<Dialect *> dialects) const = 0;
+
+  /// Return a copy of this extension.
+  virtual std::unique_ptr<DialectExtensionBase> clone() const = 0;
+
+protected:
+  /// Initialize the extension with a set of required dialects. Note that there
+  /// should always be at least one affected dialect.
+  DialectExtensionBase(ArrayRef<StringRef> dialectNames)
+      : dialectNames(dialectNames.begin(), dialectNames.end()) {
+    assert(!dialectNames.empty() && "expected at least one affected dialect");
+  }
+
+private:
+  /// The names of the dialects affected by this extension.
+  SmallVector<StringRef> dialectNames;
+};
+
+/// This class represents a dialect extension anchored on the given set of
+/// dialects. When all of the specified dialects have been loaded, the
+/// application function of this extension will be executed.
+template <typename DerivedT, typename... DialectsT>
+class DialectExtension : public DialectExtensionBase {
+public:
+  /// Applies this extension to the given context and set of required dialects.
+  virtual void apply(MLIRContext *context, DialectsT *...dialects) const = 0;
+
+  /// Return a copy of this extension.
+  std::unique_ptr<DialectExtensionBase> clone() const final {
+    return std::make_unique<DerivedT>(static_cast<const DerivedT &>(*this));
+  }
+
+protected:
+  DialectExtension()
+      : DialectExtensionBase(
+            ArrayRef<StringRef>({DialectsT::getDialectNamespace()...})) {}
+
+  /// Override the base apply method to allow providing the exact dialect types.
+  void apply(MLIRContext *context,
+             MutableArrayRef<Dialect *> dialects) const final {
+    unsigned dialectIdx = 0;
+    auto derivedDialects = std::tuple<DialectsT *...>{
+        static_cast<DialectsT *>(dialects[dialectIdx++])...};
+    llvm::apply_tuple(
+        [&](DialectsT *...dialect) { apply(context, dialect...); },
+        derivedDialects);
+  }
+};
+
+//===----------------------------------------------------------------------===//
+// DialectRegistry
+//===----------------------------------------------------------------------===//
+
+/// 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>>;
+
+public:
+  explicit DialectRegistry();
+
+  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. 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,
+              const DialectAllocatorFunction &ctor);
+
+  /// Return an allocation function for constructing the dialect identified by
+  /// its namespace, or nullptr if the namespace is not in this registry.
+  DialectAllocatorFunctionRef getDialectAllocator(StringRef name) const;
+
+  // Register all dialects available in the current registry with the registry
+  // in the provided context.
+  void appendTo(DialectRegistry &destination) const {
+    for (const auto &nameAndRegistrationIt : registry)
+      destination.insert(nameAndRegistrationIt.second.first,
+                         nameAndRegistrationIt.first,
+                         nameAndRegistrationIt.second.second);
+    // Merge the extensions.
+    for (const auto &extension : extensions)
+      destination.extensions.push_back(extension->clone());
+  }
+
+  /// Return the names of dialects known to this registry.
+  auto getDialectNames() const {
+    return llvm::map_range(
+        registry,
+        [](const MapTy::value_type &item) -> StringRef { return item.first; });
+  }
+
+  /// Apply any held extensions that require the given dialect. Users are not
+  /// expected to call this directly.
+  void applyExtensions(Dialect *dialect) const;
+
+  /// Apply any applicable extensions to the given context. Users are not
+  /// expected to call this directly.
+  void applyExtensions(MLIRContext *ctx) const;
+
+  /// Add the given extension to the registry.
+  void addExtension(std::unique_ptr<DialectExtensionBase> extension) {
+    extensions.push_back(std::move(extension));
+  }
+
+  /// Add the given extensions to the registry.
+  template <typename... ExtensionsT>
+  void addExtensions() {
+    (void)std::initializer_list<int>{
+        addExtension(std::make_unique<ExtensionsT>())...};
+  }
+
+  /// Add an extension function that requires the given dialects.
+  /// Note: This bare functor overload is provided in addition to the
+  /// std::function variant to enable dialect type deduction, e.g.:
+  ///  registry.addExtension(+[](MLIRContext *ctx, MyDialect *dialect) { ... })
+  ///
+  /// is equivalent to:
+  ///  registry.addExtension<MyDialect>(
+  ///     [](MLIRContext *ctx, MyDialect *dialect){ ... }
+  ///  )
+  template <typename... DialectsT>
+  void addExtension(void (*extensionFn)(MLIRContext *, DialectsT *...)) {
+    addExtension<DialectsT...>(
+        std::function<void(MLIRContext *, DialectsT * ...)>(extensionFn));
+  }
+  template <typename... DialectsT>
+  void
+  addExtension(std::function<void(MLIRContext *, DialectsT *...)> extensionFn) {
+    using ExtensionFnT = std::function<void(MLIRContext *, DialectsT * ...)>;
+
+    struct Extension : public DialectExtension<Extension, DialectsT...> {
+      Extension(const Extension &) = default;
+      Extension(ExtensionFnT extensionFn)
+          : extensionFn(std::move(extensionFn)) {}
+      ~Extension() override = default;
+
+      void apply(MLIRContext *context, DialectsT *...dialects) const final {
+        extensionFn(context, dialects...);
+      }
+      ExtensionFnT extensionFn;
+    };
+    addExtension(std::make_unique<Extension>(std::move(extensionFn)));
+  }
+
+private:
+  MapTy registry;
+  std::vector<std::unique_ptr<DialectExtensionBase>> extensions;
+};
+
+} // namespace mlir
+
+#endif // MLIR_IR_DIALECTREGISTRY_H

mlir/lib/Dialect/Arithmetic/Transforms/BufferizableOpInterfaceImpl.cpp

BufferRelation bufferRelation(Operation *op, OpResult opResult,
     return BufferRelation::None;
   }
 };
 
 } // namespace
 
 void mlir::arith::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<ConstantOp, ConstantOpInterface>();
-  registry.addOpInterface<IndexCastOp, IndexCastOpInterface>();
-  registry.addOpInterface<SelectOp, SelectOpInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, ArithmeticDialect *dialect) {
+    ConstantOp::attachInterface<ConstantOpInterface>(*ctx);
+    IndexCastOp::attachInterface<IndexCastOpInterface>(*ctx);
+    SelectOp::attachInterface<SelectOpInterface>(*ctx);
+  });
 }

mlir/lib/Dialect/Bufferization/Transforms/BufferDeallocation.cpp

LogicalResult bufferization::deallocateBuffers(Operation *op) {
   if (failed(deallocation.deallocate()))
     return failure();
 
   return success();
 }
 
 void bufferization::registerAllocationOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<memref::AllocOp, DefaultAllocationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, memref::MemRefDialect *dialect) {
+    memref::AllocOp::attachInterface<DefaultAllocationInterface>(*ctx);
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // BufferDeallocationPass construction
 //===----------------------------------------------------------------------===//
 
 std::unique_ptr<Pass> mlir::bufferization::createBufferDeallocationPass() {
   return std::make_unique<BufferDeallocationPass>();
 }

mlir/lib/Dialect/Linalg/ComprehensiveBufferize/ModuleBufferization.cpp

 
 } // namespace std_ext
 } // namespace comprehensive_bufferize
 } // namespace linalg
 } // namespace mlir
 
 void mlir::linalg::comprehensive_bufferize::std_ext::
     registerModuleBufferizationExternalModels(DialectRegistry &registry) {
-  registry.addOpInterface<func::CallOp, std_ext::CallOpInterface>();
-  registry.addOpInterface<func::ReturnOp, std_ext::ReturnOpInterface>();
-  registry.addOpInterface<FuncOp, std_ext::FuncOpInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, func::FuncDialect *dialect) {
+    func::CallOp::attachInterface<std_ext::CallOpInterface>(*ctx);
+    func::ReturnOp::attachInterface<std_ext::ReturnOpInterface>(*ctx);
+    func::FuncOp::attachInterface<std_ext::FuncOpInterface>(*ctx);
+  });
 }
 
 /// Set the attribute that triggers inplace bufferization on a FuncOp argument
 /// `bbArg`.
 static void setInPlaceFuncArgument(BlockArgument bbArg, bool inPlace) {
   auto funcOp = cast<FuncOp>(bbArg.getOwner()->getParentOp());
   funcOp.setArgAttr(bbArg.getArgNumber(),
                     BufferizableOpInterface::kInplaceableAttrName,

mlir/lib/Dialect/Linalg/Transforms/BufferizableOpInterfaceImpl.cpp

if (failed(alloc))
       return failure();
     replaceOpWithBufferizedValues(rewriter, op, *alloc);
     return success();
   }
 };
 
 /// Helper structure that iterates over all LinalgOps in `OpTys` and registers
 /// the `BufferizableOpInterface` with each of them.
-template <typename... OpTys>
-struct LinalgOpInterfaceHelper;
-
-template <typename First, typename... Others>
-struct LinalgOpInterfaceHelper<First, Others...> {
-  static void registerOpInterface(DialectRegistry &registry) {
-    registry.addOpInterface<First, LinalgOpInterface<First>>();
-    LinalgOpInterfaceHelper<Others...>::registerOpInterface(registry);
+template <typename... Ops>
+struct LinalgOpInterfaceHelper {
+  static void registerOpInterface(MLIRContext *ctx) {
+    (void)std::initializer_list<int>{
+        0, (Ops::template attachInterface<LinalgOpInterface<Ops>>(*ctx), 0)...};
   }
 };
-
-template <>
-struct LinalgOpInterfaceHelper<> {
-  static void registerOpInterface(DialectRegistry &registry) {}
-};
-
 } // namespace
 
 /// Return true if all `neededValues` are in scope at the given
 /// `insertionPoint`.
 static bool
 neededValuesDominateInsertionPoint(const DominanceInfo &domInfo,
                                    Operation *insertionPoint,
                                    const SmallVector<Value> &neededValues) {
return eliminateInitTensors(
             loc, t, insertOp.dest(), mixedOffsets, mixedSizes, mixedStrides);
         return extractOp.result();
       },
       newOps);
 }
 
 void mlir::linalg::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<linalg::InitTensorOp, InitTensorOpInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, linalg::LinalgDialect *dialect) {
+    linalg::InitTensorOp::attachInterface<InitTensorOpInterface>(*ctx);
 
     // Register all Linalg structured ops. `LinalgOp` is an interface and it is
     // not possible to attach an external interface to an existing interface.
     // Therefore, attach the `BufferizableOpInterface` to all ops one-by-one.
     LinalgOpInterfaceHelper<
 #define GET_OP_LIST
 #include "mlir/Dialect/Linalg/IR/LinalgStructuredOps.cpp.inc"
-      >::registerOpInterface(registry);
+        >::registerOpInterface(ctx);
+  });
 }

mlir/lib/Dialect/SCF/Transforms/BufferizableOpInterfaceImpl.cpp

 };
 
 } // namespace
 } // namespace scf
 } // namespace mlir
 
 void mlir::scf::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<ExecuteRegionOp, ExecuteRegionOpInterface>();
-  registry.addOpInterface<ForOp, ForOpInterface>();
-  registry.addOpInterface<IfOp, IfOpInterface>();
-  registry.addOpInterface<YieldOp, YieldOpInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, scf::SCFDialect *dialect) {
+    ExecuteRegionOp::attachInterface<ExecuteRegionOpInterface>(*ctx);
+    ForOp::attachInterface<ForOpInterface>(*ctx);
+    IfOp::attachInterface<IfOpInterface>(*ctx);
+    YieldOp::attachInterface<YieldOpInterface>(*ctx);
+  });
 }

mlir/lib/Dialect/Shape/Transforms/BufferizableOpInterfaceImpl.cpp

 };
 
 } // namespace
 } // namespace shape
 } // namespace mlir
 
 void mlir::shape::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<shape::AssumingOp, AssumingOpInterface>();
-  registry.addOpInterface<shape::AssumingYieldOp, AssumingYieldOpInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, shape::ShapeDialect *dialect) {
+    shape::AssumingOp::attachInterface<AssumingOpInterface>(*ctx);
+    shape::AssumingYieldOp::attachInterface<AssumingYieldOpInterface>(*ctx);
+  });
 }

mlir/lib/Dialect/Tensor/IR/TensorInferTypeOpInterfaceImpl.cpp

reifyResultShapes(Operation *op, OpBuilder &b,
     return success();
   }
 };
 
 } // namespace
 
 void mlir::tensor::registerInferTypeOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry
-      .addOpInterface<tensor::ExpandShapeOp,
-                      ReifyExpandOrCollapseShapeOp<tensor::ExpandShapeOp>>();
-  registry
-      .addOpInterface<tensor::CollapseShapeOp,
-                      ReifyExpandOrCollapseShapeOp<tensor::CollapseShapeOp>>();
-  registry.addOpInterface<tensor::PadOp, ReifyPadOp>();
+  registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
+    ExpandShapeOp::attachInterface<
+        ReifyExpandOrCollapseShapeOp<tensor::ExpandShapeOp>>(*ctx);
+    CollapseShapeOp::attachInterface<
+        ReifyExpandOrCollapseShapeOp<tensor::CollapseShapeOp>>(*ctx);
+    PadOp::attachInterface<ReifyPadOp>(*ctx);
+  });
 }

mlir/lib/Dialect/Tensor/IR/TensorTilingInterfaceImpl.cpp

auto result = b.create<scf::IfOp>(
         });
     return result;
   }
   return createPadOfExtractSlice();
 }
 
 void mlir::tensor::registerTilingOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<tensor::PadOp, PadOpTiling>();
+  registry.addExtension(+[](MLIRContext *ctx, TensorDialect *dialect) {
+    tensor::PadOp::attachInterface<PadOpTiling>(*ctx);
+  });
 }

mlir/lib/Dialect/Tensor/Transforms/BufferizableOpInterfaceImpl.cpp

 };
 
 } // namespace
 } // namespace tensor
 } // namespace mlir
 
 void mlir::tensor::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<CastOp, CastOpInterface>();
-  registry.addOpInterface<CollapseShapeOp, CollapseShapeOpInterface>();
-  registry.addOpInterface<DimOp, DimOpInterface>();
-  registry.addOpInterface<ExpandShapeOp, ExpandShapeOpInterface>();
-  registry.addOpInterface<ExtractSliceOp, ExtractSliceOpInterface>();
-  registry.addOpInterface<ExtractOp, ExtractOpInterface>();
-  registry.addOpInterface<FromElementsOp, FromElementsOpInterface>();
-  registry.addOpInterface<GenerateOp, GenerateOpInterface>();
-  registry.addOpInterface<InsertOp, InsertOpInterface>();
-  registry.addOpInterface<InsertSliceOp, InsertSliceOpInterface>();
-  registry.addOpInterface<RankOp, RankOpInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, tensor::TensorDialect *dialect) {
+    CastOp::attachInterface<CastOpInterface>(*ctx);
+    CollapseShapeOp::attachInterface<CollapseShapeOpInterface>(*ctx);
+    DimOp::attachInterface<DimOpInterface>(*ctx);
+    ExpandShapeOp::attachInterface<ExpandShapeOpInterface>(*ctx);
+    ExtractSliceOp::attachInterface<ExtractSliceOpInterface>(*ctx);
+    ExtractOp::attachInterface<ExtractOpInterface>(*ctx);
+    FromElementsOp::attachInterface<FromElementsOpInterface>(*ctx);
+    GenerateOp::attachInterface<GenerateOpInterface>(*ctx);
+    InsertOp::attachInterface<InsertOpInterface>(*ctx);
+    InsertSliceOp::attachInterface<InsertSliceOpInterface>(*ctx);
+    RankOp::attachInterface<RankOpInterface>(*ctx);
+  });
 }

mlir/lib/Dialect/Vector/Transforms/BufferizableOpInterfaceImpl.cpp

 };
 
 } // namespace
 } // namespace vector
 } // namespace mlir
 
 void mlir::vector::registerBufferizableOpInterfaceExternalModels(
     DialectRegistry &registry) {
-  registry.addOpInterface<TransferReadOp, TransferReadOpInterface>();
-  registry.addOpInterface<TransferWriteOp, TransferWriteOpInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, vector::VectorDialect *dialect) {
+    TransferReadOp::attachInterface<TransferReadOpInterface>(*ctx);
+    TransferWriteOp::attachInterface<TransferWriteOpInterface>(*ctx);
+  });
 }

mlir/lib/IR/Dialect.cpp

 #include "llvm/Support/Regex.h"
 
 #define DEBUG_TYPE "dialect"
 
 using namespace mlir;
 using namespace detail;
 
 //===----------------------------------------------------------------------===//
-// DialectRegistry
-//===----------------------------------------------------------------------===//
-
-DialectRegistry::DialectRegistry() { insert<BuiltinDialect>(); }
-
-void DialectRegistry::addDialectInterface(
-    StringRef dialectName, TypeID interfaceTypeID,
-    const 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 &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;
-    }
-  }
-
-  ifaces.dialectInterfaces.emplace_back(interfaceTypeID, allocator);
-}
-
-void DialectRegistry::addObjectInterface(
-    StringRef dialectName, TypeID objectID, TypeID interfaceTypeID,
-    const 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 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(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::insert(TypeID typeID, StringRef name,
-                             const 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) 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->getSecond().dialectInterfaces) {
-    if (dialect->getRegisteredInterface(kvp.first))
-      continue;
-    dialect->addInterface(kvp.second(dialect));
-  }
-
-  // Add attribute, operation and type interfaces.
-  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) {
   assert(isValidNamespace(name) && "invalid dialect namespace");
 }
 
bool Dialect::isValidNamespace(StringRef str) {
   return dialectNameRegex.match(str);
 }
 
 /// Register a set of dialect interfaces with this dialect instance.
 void Dialect::addInterface(std::unique_ptr<DialectInterface> interface) {
   auto it = registeredInterfaces.try_emplace(interface->getID(),
                                              std::move(interface));
   (void)it;
-  assert(it.second && "interface kind has already been registered");
+  LLVM_DEBUG({
+    if (!it.second) {
+      llvm::dbgs() << "[" DEBUG_TYPE
+                      "] repeated interface registration for dialect "
+                   << getNamespace();
+    }
+  });
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect Interface
 //===----------------------------------------------------------------------===//
 
 DialectInterface::~DialectInterface() = default;
 

 DialectInterfaceCollectionBase::~DialectInterfaceCollectionBase() = default;
 
 /// Get the interface for the dialect of given operation, or null if one
 /// is not registered.
 const DialectInterface *
 DialectInterfaceCollectionBase::getInterfaceFor(Operation *op) const {
   return getInterfaceFor(op->getDialect());
 }
+
+//===----------------------------------------------------------------------===//
+// DialectExtension
+//===----------------------------------------------------------------------===//
+
+DialectExtensionBase::~DialectExtensionBase() = default;
+
+//===----------------------------------------------------------------------===//
+// DialectRegistry
+//===----------------------------------------------------------------------===//
+
+DialectRegistry::DialectRegistry() { insert<BuiltinDialect>(); }
+
+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,
+                             const 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::applyExtensions(Dialect *dialect) const {
+  MLIRContext *ctx = dialect->getContext();
+  StringRef dialectName = dialect->getNamespace();
+
+  // Functor used to try to apply the given extension.
+  auto applyExtension = [&](const DialectExtensionBase &extension) {
+    ArrayRef<StringRef> dialectNames = extension.getRequiredDialects();
+
+    // Handle the simple case of a single dialect name. In this case, the
+    // required dialect should be the current dialect.
+    if (dialectNames.size() == 1) {
+      if (dialectNames.front() == dialectName)
+        extension.apply(ctx, dialect);
+      return;
+    }
+
+    // Otherwise, check to see if this extension requires this dialect.
+    const StringRef *nameIt = llvm::find(dialectNames, dialectName);
+    if (nameIt == dialectNames.end())
+      return;
+
+    // If it does, ensure that all of the other required dialects have been
+    // loaded.
+    SmallVector<Dialect *> requiredDialects;
+    requiredDialects.reserve(dialectNames.size());
+    for (auto it = dialectNames.begin(), e = dialectNames.end(); it != e;
+         ++it) {
+      // The current dialect is known to be loaded.
+      if (it == nameIt) {
+        requiredDialects.push_back(dialect);
+        continue;
+      }
+      // Otherwise, check if it is loaded.
+      Dialect *loadedDialect = ctx->getLoadedDialect(*it);
+      if (!loadedDialect)
+        return;
+      requiredDialects.push_back(loadedDialect);
+    }
+    extension.apply(ctx, requiredDialects);
+  };
+
+  for (const auto &extension : extensions)
+    applyExtension(*extension);
+}
+
+void DialectRegistry::applyExtensions(MLIRContext *ctx) const {
+  // Functor used to try to apply the given extension.
+  auto applyExtension = [&](const DialectExtensionBase &extension) {
+    ArrayRef<StringRef> dialectNames = extension.getRequiredDialects();
+
+    // Check to see if all of the dialects for this extension are loaded.
+    SmallVector<Dialect *> requiredDialects;
+    requiredDialects.reserve(dialectNames.size());
+    for (StringRef dialectName : dialectNames) {
+      Dialect *loadedDialect = ctx->getLoadedDialect(dialectName);
+      if (!loadedDialect)
+        return;
+      requiredDialects.push_back(loadedDialect);
+    }
+    extension.apply(ctx, requiredDialects);
+  };
+
+  for (const auto &extension : extensions)
+    applyExtension(*extension);
+}

mlir/lib/IR/MLIRContext.cpp

 
 //===----------------------------------------------------------------------===//
 // Dialect and Operation Registration
 //===----------------------------------------------------------------------===//
 
 void MLIRContext::appendDialectRegistry(const DialectRegistry &registry) {
   registry.appendTo(impl->dialectsRegistry);
 
-  // For the already loaded dialects, register the interfaces immediately.
-  for (const auto &kvp : impl->loadedDialects)
-    registry.registerDelayedInterfaces(kvp.second.get());
+  // For the already loaded dialects, apply any possible extensions immediately.
+  registry.applyExtensions(this);
 }
 
 const DialectRegistry &MLIRContext::getDialectRegistry() {
   return impl->dialectsRegistry;
 }
 
 /// Return information about all registered IR dialects.
 std::vector<Dialect *> MLIRContext::getLoadedDialects() {
#endif
     // were already created.
     auto stringAttrsIt = impl.dialectReferencingStrAttrs.find(dialectNamespace);
     if (stringAttrsIt != impl.dialectReferencingStrAttrs.end()) {
       for (StringAttrStorage *storage : stringAttrsIt->second)
         storage->referencedDialect = dialect.get();
       impl.dialectReferencingStrAttrs.erase(stringAttrsIt);
     }
 
-    // Actually register the interfaces with delayed registration.
-    impl.dialectsRegistry.registerDelayedInterfaces(dialect.get());
+    // Apply any extensions to this newly loaded dialect.
+    impl.dialectsRegistry.applyExtensions(dialect.get());
     return dialect.get();
   }
 
   // Abort if dialect with namespace has already been registered.
   std::unique_ptr<Dialect> &dialect = dialectIt->second;
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");

mlir/lib/Target/LLVMIR/Dialect/AMX/AMXToLLVMIRTranslation.cpp

#include "mlir/Dialect/AMX/AMXConversions.inc"
 
     return failure();
   }
 };
 } // namespace
 
 void mlir::registerAMXDialectTranslation(DialectRegistry &registry) {
   registry.insert<amx::AMXDialect>();
-  registry.addDialectInterface<amx::AMXDialect,
-                               AMXDialectLLVMIRTranslationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, amx::AMXDialect *dialect) {
+    dialect->addInterfaces<AMXDialectLLVMIRTranslationInterface>();
+  });
 }
 
 void mlir::registerAMXDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerAMXDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/ArmNeon/ArmNeonToLLVMIRTranslation.cpp

#include "mlir/Dialect/ArmNeon/ArmNeonConversions.inc"
 
     return failure();
   }
 };
 } // namespace
 
 void mlir::registerArmNeonDialectTranslation(DialectRegistry &registry) {
   registry.insert<arm_neon::ArmNeonDialect>();
-  registry.addDialectInterface<arm_neon::ArmNeonDialect,
-                               ArmNeonDialectLLVMIRTranslationInterface>();
+  registry.addExtension(
+      +[](MLIRContext *ctx, arm_neon::ArmNeonDialect *dialect) {
+        dialect->addInterfaces<ArmNeonDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerArmNeonDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerArmNeonDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/ArmSVE/ArmSVEToLLVMIRTranslation.cpp

#include "mlir/Dialect/ArmSVE/ArmSVEConversions.inc"
 
     return failure();
   }
 };
 } // namespace
 
 void mlir::registerArmSVEDialectTranslation(DialectRegistry &registry) {
   registry.insert<arm_sve::ArmSVEDialect>();
-  registry.addDialectInterface<arm_sve::ArmSVEDialect,
-                               ArmSVEDialectLLVMIRTranslationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, arm_sve::ArmSVEDialect *dialect) {
+    dialect->addInterfaces<ArmSVEDialectLLVMIRTranslationInterface>();
+  });
 }
 
 void mlir::registerArmSVEDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerArmSVEDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/LLVMIR/LLVMToLLVMIRTranslation.cpp

convertOperation(Operation *op, llvm::IRBuilderBase &builder,
                    LLVM::ModuleTranslation &moduleTranslation) const final {
     return convertOperationImpl(*op, builder, moduleTranslation);
   }
 };
 } // namespace
 
 void mlir::registerLLVMDialectTranslation(DialectRegistry &registry) {
   registry.insert<LLVM::LLVMDialect>();
-  registry.addDialectInterface<LLVM::LLVMDialect,
-                               LLVMDialectLLVMIRTranslationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, LLVM::LLVMDialect *dialect) {
+    dialect->addInterfaces<LLVMDialectLLVMIRTranslationInterface>();
+  });
 }
 
 void mlir::registerLLVMDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerLLVMDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/NVVM/NVVMToLLVMIRTranslation.cpp

amendOperation(Operation *op, NamedAttribute attribute,
     }
     return success();
   }
 };
 } // namespace
 
 void mlir::registerNVVMDialectTranslation(DialectRegistry &registry) {
   registry.insert<NVVM::NVVMDialect>();
-  registry.addDialectInterface<NVVM::NVVMDialect,
-                               NVVMDialectLLVMIRTranslationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, NVVM::NVVMDialect *dialect) {
+    dialect->addInterfaces<NVVMDialectLLVMIRTranslationInterface>();
+  });
 }
 
 void mlir::registerNVVMDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerNVVMDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/OpenACC/OpenACCToLLVMIRTranslation.cpp

return llvm::TypeSwitch<Operation *, LogicalResult>(op)
       .Default([&](Operation *op) {
         return op->emitError("unsupported OpenACC operation: ")
                << op->getName();
       });
 }
 
 void mlir::registerOpenACCDialectTranslation(DialectRegistry &registry) {
   registry.insert<acc::OpenACCDialect>();
-  registry.addDialectInterface<acc::OpenACCDialect,
-                               OpenACCDialectLLVMIRTranslationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, acc::OpenACCDialect *dialect) {
+    dialect->addInterfaces<OpenACCDialectLLVMIRTranslationInterface>();
+  });
 }
 
 void mlir::registerOpenACCDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerOpenACCDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/OpenMP/OpenMPToLLVMIRTranslation.cpp

return llvm::TypeSwitch<Operation *, LogicalResult>(op)
       .Default([&](Operation *inst) {
         return inst->emitError("unsupported OpenMP operation: ")
                << inst->getName();
       });
 }
 
 void mlir::registerOpenMPDialectTranslation(DialectRegistry &registry) {
   registry.insert<omp::OpenMPDialect>();
-  registry.addDialectInterface<omp::OpenMPDialect,
-                               OpenMPDialectLLVMIRTranslationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, omp::OpenMPDialect *dialect) {
+    dialect->addInterfaces<OpenMPDialectLLVMIRTranslationInterface>();
+  });
 }
 
 void mlir::registerOpenMPDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerOpenMPDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/ROCDL/ROCDLToLLVMIRTranslation.cpp

amendOperation(Operation *op, NamedAttribute attribute,
     }
     return success();
   }
 };
 } // namespace
 
 void mlir::registerROCDLDialectTranslation(DialectRegistry &registry) {
   registry.insert<ROCDL::ROCDLDialect>();
-  registry.addDialectInterface<ROCDL::ROCDLDialect,
-                               ROCDLDialectLLVMIRTranslationInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, ROCDL::ROCDLDialect *dialect) {
+    dialect->addInterfaces<ROCDLDialectLLVMIRTranslationInterface>();
+  });
 }
 
 void mlir::registerROCDLDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerROCDLDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/lib/Target/LLVMIR/Dialect/X86Vector/X86VectorToLLVMIRTranslation.cpp

#include "mlir/Dialect/X86Vector/X86VectorConversions.inc"
 
     return failure();
   }
 };
 } // namespace
 
 void mlir::registerX86VectorDialectTranslation(DialectRegistry &registry) {
   registry.insert<x86vector::X86VectorDialect>();
-  registry.addDialectInterface<x86vector::X86VectorDialect,
-                               X86VectorDialectLLVMIRTranslationInterface>();
+  registry.addExtension(
+      +[](MLIRContext *ctx, x86vector::X86VectorDialect *dialect) {
+        dialect->addInterfaces<X86VectorDialectLLVMIRTranslationInterface>();
+      });
 }
 
 void mlir::registerX86VectorDialectTranslation(MLIRContext &context) {
   DialectRegistry registry;
   registerX86VectorDialectTranslation(registry);
   context.appendDialectRegistry(registry);
 }

mlir/unittests/IR/DialectTest.cpp

struct SecondTestDialectInterface : public 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>();
+  registry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+    dialect->addInterfaces<TestDialectInterface>();
+  });
 
   MLIRContext context(registry);
 
   // Load the TestDialect and check that the interface got registered for it.
   Dialect *testDialect = context.getOrLoadDialect<TestDialect>();
   ASSERT_TRUE(testDialect != nullptr);
   auto *testDialectInterface = dyn_cast<TestDialectInterfaceBase>(testDialect);
   EXPECT_TRUE(testDialectInterface != nullptr);
 
   // Load the SecondTestDialect and check that the interface is not registered
   // for it.
   Dialect *secondTestDialect = context.getOrLoadDialect<SecondTestDialect>();
   ASSERT_TRUE(secondTestDialect != nullptr);
   auto *secondTestDialectInterface =
       dyn_cast<SecondTestDialectInterface>(secondTestDialect);
   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.
   DialectRegistry secondRegistry;
   secondRegistry.insert<SecondTestDialect>();
-  secondRegistry
-      .addDialectInterface<SecondTestDialect, SecondTestDialectInterface>();
+  secondRegistry.addExtension(
+      +[](MLIRContext *ctx, SecondTestDialect *dialect) {
+        dialect->addInterfaces<SecondTestDialectInterface>();
+      });
   context.appendDialectRegistry(secondRegistry);
   secondTestDialectInterface =
       dyn_cast<SecondTestDialectInterface>(secondTestDialect);
   EXPECT_TRUE(secondTestDialectInterface != nullptr);
 }
 
 TEST(Dialect, RepeatedDelayedRegistration) {
   // Set up the delayed registration.
   DialectRegistry registry;
   registry.insert<TestDialect>();
-  registry.addDialectInterface<TestDialect, TestDialectInterface>();
+  registry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+    dialect->addInterfaces<TestDialectInterface>();
+  });
   MLIRContext context(registry);
 
   // Load the TestDialect and check that the interface got registered for it.
   Dialect *testDialect = context.getOrLoadDialect<TestDialect>();
   ASSERT_TRUE(testDialect != nullptr);
   auto *testDialectInterface = dyn_cast<TestDialectInterfaceBase>(testDialect);
   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>();
+  secondRegistry.addExtension(+[](MLIRContext *ctx, TestDialect *dialect) {
+    dialect->addInterfaces<TestDialectInterface>();
+  });
   context.appendDialectRegistry(secondRegistry);
   testDialectInterface = dyn_cast<TestDialectInterfaceBase>(testDialect);
   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
-}
-
 } // namespace

mlir/unittests/IR/InterfaceAttachmentTest.cpp

struct TestTypeModel
 
   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>();
+  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
+    test::TestType::attachInterface<TestTypeModel>(*ctx);
+  });
 
   // 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>();
+  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
+    test::TestType::attachInterface<TestTypeModel>(*ctx);
+  });
 
   // 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>();
+  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
+    IntegerType::attachInterface<Model>(*ctx);
+  });
   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>();
+  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
+    IntegerType::attachInterface<Model>(*ctx);
+  });
 
   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>());
struct TestExternalSimpleAAttrModel
   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>();
+  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
+    test::SimpleAAttr::attachInterface<TestExternalSimpleAAttrModel>(*ctx);
+  });
 
   MLIRContext context(registry);
   context.loadDialect<test::TestDialect>();
   auto attr = test::SimpleAAttr::get(&context);
   EXPECT_TRUE(attr.isa<TestExternalAttrInterface>());
 
   MLIRContext initiallyEmpty;
   initiallyEmpty.loadDialect<test::TestDialect>();
struct TestExternalTestOpModel
   static unsigned getNameLengthPlusArgTwice(unsigned 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<test::OpJ>>();
-  registry.addOpInterface<test::OpH, TestExternalTestOpModel<test::OpH>>();
+  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
+    ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
+  });
+  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
+    test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
+    test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
+  });
 
-  // Construct the context directly from a registry. The interfaces are expected
-  // to be readily available on operations.
+  // Construct the context directly from a registry. The interfaces are
+  // expected to be readily available on operations.
   MLIRContext context(registry);
   context.loadDialect<test::TestDialect>();
 
   OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
   OpBuilder builder(module->getBody(), module->getBody()->begin());
   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>(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<test::OpJ>>();
-  registry.addOpInterface<test::OpH, TestExternalTestOpModel<test::OpH>>();
+  registry.addExtension(+[](MLIRContext *ctx, BuiltinDialect *dialect) {
+    ModuleOp::attachInterface<TestExternalOpModel>(*ctx);
+  });
+  registry.addExtension(+[](MLIRContext *ctx, test::TestDialect *dialect) {
+    test::OpJ::attachInterface<TestExternalTestOpModel<test::OpJ>>(*ctx);
+    test::OpH::attachInterface<TestExternalTestOpModel<test::OpH>>(*ctx);
+  });
 
   // 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>();
 
   OwningOpRef<ModuleOp> module = ModuleOp::create(UnknownLoc::get(&context));
   OpBuilder builder(module->getBody(), module->getBody()->begin());