[mlir] avoid exposing mutable DialectRegistry from MLIRContext

[ftynse]

Summary

MLIRContext allows its users to access directly to the DialectRegistry it
contains. While sometimes useful for registering additional dialects on an
already existing context, this breaks the encapsulation by essentially giving
raw accesses to a part of the context's internal state. Remove this mutable
access and instead provide a method to append a given DialectRegistry to the
one already contained in the context. Also provide a shortcut mechanism to
construct a context from an already existing registry, which seems to be a
common use case in the wild. Keep read-only access to the registry contained in
the context in case it needs to be copied or used for constructing another
context.

With this change, DialectRegistry is no longer concerned with loading the
dialects and deciding whether to invoke delayed interface registration. Loading
is concentrated in the MLIRContext, and the functionality of the registry
better reflects its name.

Depends On D96137

[mehdi_amini]

mlir/include/mlir/IR/MLIRContext.h

49

Append for consistency with everywhere else here (this isn't google3)

92

mlir/lib/CAPI/Registration/Registration.cpp

17

Could we provide a registerAllDialects(MLIRContext *) helper that replace these three lines? I suspect this'll be a common pattern

mlir/lib/IR/MLIRContext.cpp

330

Trying to follow all the diff is really made harder by the order of the patches, you're undoing the change made in the parent in too many places.

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

140

Can you add a TODO, we shouldn't need this and only load the dialect we need.

[ftynse]

Address review

[ftynse]

mlir/lib/IR/MLIRContext.cpp

330

Sorry about that.

[rriddle]

mlir/include/mlir/IR/Dialect.h

271

Do we need to return this by-value?

mlir/include/mlir/IR/MLIRContext.h

47

What is this function used for now? Can we remove it?

mlir/lib/Pass/Pass.cpp

870

nit: Spell out auto here. It isn't clear if this is copying something expensive or not.

[ftynse]

Address review

[ftynse]

mlir/include/mlir/IR/Dialect.h

271

Good point, changed to function_ref.

mlir/include/mlir/IR/MLIRContext.h

47

SPIR-V translation constructs a new context with the same registry as the old context, so it needs readonly access to the registry. An alternative could be to provide a function to append the registry of one context into another context without exposing it.

mlir/include/mlir/CAPI/Registration.h

 //===----------------------------------------------------------------------===//
 // Corrolary to MLIR_DECLARE_CAPI_DIALECT_REGISTRATION that defines an impl.
 // Takes the same name passed to the above and the fully qualified class name
 // of the dialect class.
 //===----------------------------------------------------------------------===//
 
 #define MLIR_DEFINE_CAPI_DIALECT_REGISTRATION(Name, Namespace, ClassName)      \
   void mlirContextRegister##Name##Dialect(MlirContext context) {               \
-    unwrap(context)->getDialectRegistry().insert<ClassName>();                 \
+    mlir::DialectRegistry registry;                                            \
+    registry.insert<ClassName>();                                              \
+    unwrap(context)->appendDialectRegistry(registry);                          \
   }                                                                            \
   MlirDialect mlirContextLoad##Name##Dialect(MlirContext context) {            \
     return wrap(unwrap(context)->getOrLoadDialect<ClassName>());               \
   }                                                                            \
   MlirStringRef mlir##Name##DialectGetNamespace() {                            \
     return wrap(ClassName::getDialectNamespace());                             \
   }                                                                            \
   const MlirDialectRegistrationHooks *mlirGetDialectHooks__##Namespace##__() { \
     static MlirDialectRegistrationHooks hooks = {                              \
         mlirContextRegister##Name##Dialect, mlirContextLoad##Name##Dialect,    \
         mlir##Name##DialectGetNamespace};                                      \
     return &hooks;                                                             \
   }
 
 #endif // MLIR_CAPI_REGISTRATION_H

mlir/include/mlir/IR/Dialect.h

 /// dialects in the Context as operations are encountered.
 class DialectRegistry {
   using MapTy =
       std::map<std::string, std::pair<TypeID, DialectAllocatorFunction>>;
   using InterfaceMapTy =
       DenseMap<TypeID, SmallVector<InterfaceAllocatorFunction, 2>>;
 
 public:
-  explicit DialectRegistry(MLIRContext *context = nullptr)
-      : owningContext(context) {}
+  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, DialectAllocatorFunction ctor);
 
-  /// Load a dialect for this namespace in the provided context.
-  Dialect *loadByName(StringRef name, MLIRContext *context);
+  /// Return an allocation function for constructing the dialect identified by
+  /// its namespace, or nullptr if the namespace is not in this registry.
+  DialectAllocatorFunction getDialectAllocator(StringRef name) const;
 
   // Register all dialects available in the current registry with the registry
   // in the provided context.
-  void appendTo(DialectRegistry &destination) {
+  void appendTo(DialectRegistry &destination) const {
     for (const auto &nameAndRegistrationIt : registry)
       destination.insert(nameAndRegistrationIt.second.first,
                          nameAndRegistrationIt.first,
                          nameAndRegistrationIt.second.second);
     destination.interfaces.insert(interfaces.begin(), interfaces.end());
   }
-  // Load all dialects available in the registry in the provided context.
-  void loadAll(MLIRContext *context) {
-    for (const auto &nameAndRegistrationIt : registry)
-      nameAndRegistrationIt.second.second(context);
-  }
 
   /// Return the names of dialects known to this registry.
-  auto getDialectNames() {
+  auto getDialectNames() const {
     return llvm::map_range(
-        registry, [](const MapTy::value_type &item) { return item.first; });
+        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, but may or may not be loaded. If it is not loaded, the
-  /// interface registration is delayed until the loading.
+  /// the registry.
   template <typename DialectTy>
   void addDialectInterface(InterfaceAllocatorFunction allocator) {
     addDialectInterface(DialectTy::getDialectNamespace(), allocator);
   }
 
   /// Add an interface to the dialect, both provided as template parameter. The
-  /// dialect must be present in the registry, but may or may not be loaded. If
-  /// it is not loaded, the interface registration is delayed until the loading.
+  /// dialect must be present in the registry.
   template <typename DialectTy, typename InterfaceTy>
   void addDialectInterface() {
     addDialectInterface<DialectTy>([](Dialect *dialect) {
       return std::make_unique<InterfaceTy>(dialect);
     });
   }
 
   /// Register any interfaces required for the given dialect (based on its
   /// TypeID). Users are not expected to call this directly.
-  void registerDelayedInterfaces(Dialect *dialect);
+  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,
                            InterfaceAllocatorFunction allocator);
 
   MapTy registry;
   InterfaceMapTy interfaces;
-
-  /// If this registry belongs to a context, this points back to the context.
-  /// Useful for checking if a dialect is loaded in the context.
-  MLIRContext *owningContext;
 };
 
 } // namespace mlir
 
 namespace llvm {
 /// Provide isa functionality for Dialects.
 template <typename T>
 struct isa_impl<T, ::mlir::Dialect> {
   static inline bool doit(const ::mlir::Dialect &dialect) {
     return mlir::TypeID::get<T>() == dialect.getTypeID();
   }
 };
 } // namespace llvm
 
 #endif

mlir/include/mlir/IR/MLIRContext.h

 ///
 /// MLIRContext gets a redundant "MLIR" prefix because otherwise it ends up with
 /// a very generic name ("Context") and because it is uncommon for clients to
 /// interact with it.
 ///
 class MLIRContext {
 public:
   /// Create a new Context.
-  /// The loadAllDialects parameters allows to load all dialects from the global
-  /// registry on Context construction. It is deprecated and will be removed
-  /// soon.
   explicit MLIRContext();
+  explicit MLIRContext(const DialectRegistry &registry);
   ~MLIRContext();
 
   /// Return information about all IR dialects loaded in the context.
   std::vector<Dialect *> getLoadedDialects();
 
   /// Return the dialect registry associated with this context.
-  DialectRegistry &getDialectRegistry();
+  const DialectRegistry &getDialectRegistry();
+
+  /// Append the contents of the given dialect registry to the registry
+  /// associated with this context.
+  void appendDialectRegistry(const DialectRegistry &registry);
 
   /// Return information about all available dialects in the registry in this
   /// context.
   std::vector<StringRef> getAvailableDialects();
 
   /// Get a registered IR dialect with the given namespace. If an exact match is
   /// not found, then return nullptr.
   Dialect *getLoadedDialect(StringRef name);
public:
 
   /// Load a list dialects in the context.
   template <typename Dialect, typename OtherDialect, typename... MoreDialects>
   void loadDialect() {
     getOrLoadDialect<Dialect>();
     loadDialect<OtherDialect, MoreDialects...>();
   }
 
+  /// Load all dialects available in the registry in this context.
+  void loadAllAvailableDialects();
+
   /// Get (or create) a dialect for the given derived dialect name.
   /// The dialect will be loaded from the registry if no dialect is found.
   /// If no dialect is loaded for this name and none is available in the
   /// registry, returns nullptr.
   Dialect *getOrLoadDialect(StringRef name);
 
   /// Return true if we allow to create operation for unregistered dialects.
   bool allowsUnregisteredDialects();

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Tensor/IR/Tensor.h"
 #include "mlir/Dialect/Tosa/IR/TosaOps.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Dialect.h"
 
 namespace mlir {
 
-// Add all the MLIR dialects to the provided registry.
+/// Add all the MLIR dialects to the provided registry.
 inline void registerAllDialects(DialectRegistry &registry) {
   // clang-format off
   registry.insert<acc::OpenACCDialect,
                   AffineDialect,
                   arm_neon::ArmNeonDialect,
                   async::AsyncDialect,
                   avx512::AVX512Dialect,
                   complex::ComplexDialect,
registry.insert<acc::OpenACCDialect,
                   ROCDL::ROCDLDialect,
                   SDBMDialect,
                   shape::ShapeDialect,
                   tensor::TensorDialect,
                   tosa::TosaDialect>();
   // clang-format on
 }
 
+/// Append all the MLIR dialects to the registry contained in the given context.
+inline void registerAllDialects(MLIRContext &context) {
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  context.appendDialectRegistry(registry);
+}
+
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

mlir/lib/CAPI/Registration/Registration.cpp

 //===- Registration.cpp - C Interface for MLIR Registration ---------------===//
 //
 // 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-c/Registration.h"
 
 #include "mlir/CAPI/IR.h"
 #include "mlir/InitAllDialects.h"
 
 void mlirRegisterAllDialects(MlirContext context) {
-  registerAllDialects(unwrap(context)->getDialectRegistry());
+  mlir::registerAllDialects(*unwrap(context));
   // TODO: we may not want to eagerly load here.
-  unwrap(context)->getDialectRegistry().loadAll(unwrap(context));
+  unwrap(context)->loadAllAvailableDialects();
 }

mlir/lib/ExecutionEngine/JitRunner.cpp

for (unsigned i = 0, e = options.llvmPasses.size(); i < e; ++i) {
     passes.push_back(options.llvmPasses[i]);
     if (optCLIPosition < options.llvmPasses.getPosition(i)) {
       optPosition = i;
       optCLIPosition = UINT_MAX; // To ensure we never insert again
     }
   }
 
   MLIRContext context;
-  registerAllDialects(context.getDialectRegistry());
+  registerAllDialects(context);
 
   auto m = parseMLIRInput(options.inputFilename, &context);
   if (!m) {
     llvm::errs() << "could not parse the input IR\n";
     return 1;
   }
 
   if (config.mlirTransformer)

mlir/lib/IR/Dialect.cpp

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

mlir/lib/IR/MLIRContext.cpp

#endif
 
   /// Cached Attribute Instances.
   BoolAttr falseAttr, trueAttr;
   UnitAttr unitAttr;
   UnknownLoc unknownLocAttr;
   DictionaryAttr emptyDictionaryAttr;
 
 public:
-  MLIRContextImpl(MLIRContext *ctx)
-      : dialectsRegistry(ctx), identifiers(identifierAllocator) {}
+  MLIRContextImpl() : identifiers(identifierAllocator) {}
   ~MLIRContextImpl() {
     for (auto typeMapping : registeredTypes)
       typeMapping.second->~AbstractType();
     for (auto attrMapping : registeredAttributes)
       attrMapping.second->~AbstractAttribute();
   }
 };
 } // end namespace mlir
 
-MLIRContext::MLIRContext() : impl(new MLIRContextImpl(this)) {
+MLIRContext::MLIRContext() : MLIRContext(DialectRegistry()) {}
+
+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);
+
   // 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);

 
 /// Returns the diagnostic engine for this context.
 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
 
 //===----------------------------------------------------------------------===//
 // Dialect and Operation Registration
 //===----------------------------------------------------------------------===//
 
-DialectRegistry &MLIRContext::getDialectRegistry() {
+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());
+}
+
+const DialectRegistry &MLIRContext::getDialectRegistry() {
   return impl->dialectsRegistry;
 }
 
 /// Return information about all registered IR dialects.
 std::vector<Dialect *> MLIRContext::getLoadedDialects() {
   std::vector<Dialect *> result;
   result.reserve(impl->loadedDialects.size());
   for (auto &dialect : impl->loadedDialects)
Dialect *MLIRContext::getLoadedDialect(StringRef name) {
   auto it = impl->loadedDialects.find(name);
   return (it != impl->loadedDialects.end()) ? it->second.get() : nullptr;
 }
 
 Dialect *MLIRContext::getOrLoadDialect(StringRef name) {
   Dialect *dialect = getLoadedDialect(name);
   if (dialect)
     return dialect;
-  return impl->dialectsRegistry.loadByName(name, this);
+  DialectAllocatorFunction allocator =
+      impl->dialectsRegistry.getDialectAllocator(name);
+  return allocator ? allocator(this) : nullptr;
 }
 
 /// Get a dialect for the provided namespace and TypeID: abort the program if a
 /// dialect exist for this namespace with different TypeID. Returns a pointer to
 /// the dialect owned by the context.
 Dialect *
 MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
                               function_ref<std::unique_ptr<Dialect>()> ctor) {
#endif
   // Abort if dialect with namespace has already been registered.
   if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
 
   return dialect.get();
 }
 
+void MLIRContext::loadAllAvailableDialects() {
+  for (StringRef name : getAvailableDialects())
+    getOrLoadDialect(name);
+}
+
 llvm::hash_code MLIRContext::getRegistryHash() {
   llvm::hash_code hash(0);
   // Factor in number of loaded dialects, attributes, operations, types.
   hash = llvm::hash_combine(hash, impl->loadedDialects.size());
   hash = llvm::hash_combine(hash, impl->registeredAttributes.size());
   hash = llvm::hash_combine(hash, impl->registeredOperations.size());
   hash = llvm::hash_combine(hash, impl->registeredTypes.size());
   return hash;

mlir/lib/Pass/Pass.cpp

LogicalResult PassManager::run(Operation *op) {
 
   // Before running, make sure to coalesce any adjacent pass adaptors in the
   // pipeline.
   getImpl().coalesceAdjacentAdaptorPasses();
 
   // Register all dialects for the current pipeline.
   DialectRegistry dependentDialects;
   getDependentDialects(dependentDialects);
-  dependentDialects.loadAll(context);
+  context->appendDialectRegistry(dependentDialects);
+  for (auto name : dependentDialects.getDialectNames())
+    context->getOrLoadDialect(name);
 
   // Initialize all of the passes within the pass manager with a new generation.
   llvm::hash_code newInitKey = context->getRegistryHash();
   if (newInitKey != initializationKey) {
     initialize(context, impl->initializationGeneration + 1);
     initializationKey = newInitKey;
   }
 

mlir/lib/Support/MlirOptMain.cpp

static LogicalResult processBuffer(raw_ostream &os,
                                    bool preloadDialectsInContext,
                                    const PassPipelineCLParser &passPipeline,
                                    DialectRegistry &registry) {
   // Tell sourceMgr about this buffer, which is what the parser will pick up.
   SourceMgr sourceMgr;
   sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), SMLoc());
 
   // Parse the input file.
-  MLIRContext context;
-  registry.appendTo(context.getDialectRegistry());
+  MLIRContext context(registry);
   if (preloadDialectsInContext)
-    registry.loadAll(&context);
+    context.loadAllAvailableDialects();
   context.allowUnregisteredDialects(allowUnregisteredDialects);
   context.printOpOnDiagnostic(!verifyDiagnostics);
 
   // If we are in verify diagnostics mode then we have a lot of work to do,
   // otherwise just perform the actions without worrying about it.
   if (!verifyDiagnostics) {
     SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
     return performActions(os, verifyDiagnostics, verifyPasses, sourceMgr,

mlir/lib/Target/SPIRV/TranslateRegistration.cpp

if (spirvModules.begin() == spirvModules.end())
     return srcModule.emitError("found no 'spv.module' op");
 
   if (std::next(spirvModules.begin()) != spirvModules.end())
     return srcModule.emitError("found more than one 'spv.module' op");
 
   if (failed(spirv::serialize(*spirvModules.begin(), binary, emitDebugInfo)))
     return failure();
 
-  MLIRContext deserializationContext;
-  context->getDialectRegistry().loadAll(&deserializationContext);
+  MLIRContext deserializationContext(context->getDialectRegistry());
+  // TODO: we should only load the required dialects instead of all dialects.
+  deserializationContext.loadAllAvailableDialects();
   // Then deserialize to get back a SPIR-V module.
   spirv::OwningSPIRVModuleRef spirvModule =
       spirv::deserialize(binary, &deserializationContext);
   if (!spirvModule)
     return failure();
 
   // Wrap around in a new MLIR module.
   OwningModuleRef dstModule(ModuleOp::create(FileLineColLoc::get(

mlir/lib/Translation/Translation.cpp

 //===- Translation.cpp - Translation registry -----------------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Definitions of the translation registry.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Translation.h"
 #include "mlir/IR/AsmState.h"
 #include "mlir/IR/BuiltinOps.h"
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
 #include "mlir/Support/ToolUtilities.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 

 //===----------------------------------------------------------------------===//
 
 TranslateFromMLIRRegistration::TranslateFromMLIRRegistration(
     StringRef name, const TranslateFromMLIRFunction &function,
     std::function<void(DialectRegistry &)> dialectRegistration) {
   registerTranslation(name, [function, dialectRegistration](
                                 llvm::SourceMgr &sourceMgr, raw_ostream &output,
                                 MLIRContext *context) {
-    dialectRegistration(context->getDialectRegistry());
+    DialectRegistry registry;
+    dialectRegistration(registry);
+    context->appendDialectRegistry(registry);
     auto module = OwningModuleRef(parseSourceFile(sourceMgr, context));
     if (!module)
       return failure();
     return function(module.get(), output);
   });
 }
 
 //===----------------------------------------------------------------------===//

mlir/unittests/IR/DialectTest.cpp

 
 TEST(Dialect, DelayedInterfaceRegistration) {
   DialectRegistry registry;
   registry.insert<TestDialect, SecondTestDialect>();
 
   // Delayed registration of an interface for TestDialect.
   registry.addDialectInterface<TestDialect, TestDialectInterface>();
 
-  MLIRContext context;
-  registry.appendTo(context.getDialectRegistry());
+  MLIRContext context(registry);
 
   // Load the TestDialect and check that the interface got registered for it.
   auto *testDialect = context.getOrLoadDialect<TestDialect>();
   ASSERT_TRUE(testDialect != nullptr);
   auto *testDialectInterface =
       testDialect->getRegisteredInterface<TestDialectInterfaceBase>();
   EXPECT_TRUE(testDialectInterface != nullptr);
 
   // Load the SecondTestDialect and check that the interface is not registered
   // for it.
   auto *secondTestDialect = context.getOrLoadDialect<SecondTestDialect>();
   ASSERT_TRUE(secondTestDialect != nullptr);
   auto *secondTestDialectInterface =
       secondTestDialect->getRegisteredInterface<SecondTestDialectInterface>();
   EXPECT_TRUE(secondTestDialectInterface == nullptr);
 
   // Use the same mechanism as for delayed registration but for an already
   // loaded dialect and check that the interface is now registered.
-  context.getDialectRegistry()
+  DialectRegistry secondRegistry;
+  secondRegistry.insert<SecondTestDialect>();
+  secondRegistry
       .addDialectInterface<SecondTestDialect, SecondTestDialectInterface>();
+  context.appendDialectRegistry(secondRegistry);
   secondTestDialectInterface =
       secondTestDialect->getRegisteredInterface<SecondTestDialectInterface>();
   EXPECT_TRUE(secondTestDialectInterface != nullptr);
 }
 
 } // end namespace