Separate the Registration from Loading dialects in the Context

[mehdi_amini]

Summary

This changes the behavior of constructing MLIRContext to no longer load globally registered dialects on construction. Instead Dialects are only loaded explicitly on demand:

• the Parser is lazily loading Dialects in the context as it encounters them during parsing. This is the only purpose for registering dialects and not load them in the context.
• Passes are expected to declare the dialects they will create entity from (Operations, Attributes, or Types), and the PassManager is loading Dialects into the Context when starting a pipeline.

This changes simplifies the configuration of the registration: a compiler only need to load the dialect for the IR it will emit, and the optimizer is self-contained and load the required Dialects. For example in the Toy tutorial, the compiler only needs to load the Toy dialect in the Context, all the others (linalg, affine, std, LLVM, ...) are automatically loaded depending on the optimization pipeline enabled.

[jpienaar]

Passes are expected to declare the dialects they will create entity from (Operations, Attributes, or Types), and the PassManager is loading Dialects into the Context when starting a pipeline.

Why put this into PassManager? Why not let passes load lazily too only when needed and not do it up front for if it may be needed? What is the overhead of lazily loading and is the fear that it is too high? When is it needed? (E.g., does one need to have dialect loaded to create an op from that dialect? And if the input already has ops from that dialect it isn't really needed?).

[mehdi_amini]

In D85622#2206379, @jpienaar wrote:

Passes are expected to declare the dialects they will create entity from (Operations, Attributes, or Types), and the PassManager is loading Dialects into the Context when starting a pipeline.

Why put this into PassManager? Why not let passes load lazily too only when needed and not do it up front for if it may be needed? What is the overhead of lazily loading and is the fear that it is too high? When is it needed? (E.g., does one need to have dialect loaded to create an op from that dialect? And if the input already has ops from that dialect it isn't really needed?).

Excellent question! The reason is threading: front-loading this at the beginning of the pipeline is required to be lock-free.
We moved away from locking the context during registration, so these methods aren't thread-safe anymore (River had measurable perf improvements).

[rriddle]

Looks good, did a first pass.

mlir/include/mlir/Conversion/Passes.td

69

nit: Wrap this at 80 characters.

83

nit: Move after constructor.

216

Move this before options.

306

Same here.

337

Same here and a few other places.

mlir/include/mlir/ExecutionEngine/JitRunner.h

24 ↗

(On Diff #284254)

Looks unused.

mlir/include/mlir/IR/Dialect.h

233

nit: allows for decoupling

237

The ordering of this container should be stable.

241

nit: add -> insert?

258

Can you change this pass in another dialect registry instead of the context? That would be more composable, and less confusing when looking at this compared to loadAll.

270

nit: What is the benefit of the non-const iterators here?

mlir/include/mlir/IR/MLIRContext.h

26

This is unnecessary.

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

52

This seems like something the dialect registry would provide, why add this to the context?

162

This comment seems old.

mlir/include/mlir/Pass/Pass.h

63

nit: other than dialects that exist in the input. For example,

mlir/include/mlir/Pass/PassManager.h

112

nit: Add a comment here.

mlir/include/mlir/Transforms/Passes.td

105

This looks like a functional change?

mlir/lib/Conversion/PassDetail.h

18

nit: I would add that this is a forward declaration from DialectRegistry.h.

mlir/lib/Dialect/Affine/Transforms/PassDetail.h

16

Same here.

mlir/lib/Dialect/LLVMIR/IR/LLVMDialect.cpp

1690 ↗

(On Diff #284254)

Can we autogen this with the dialect tablegen backend?

mlir/lib/IR/MLIRContext.cpp

447

nit: Use array_pod_sort here.

467

Why the extra ()?

mlir/lib/Parser/AttributeParser.cpp

250

nit: I would add a comment here to specifically call out that this is lazy loading a dialect.

mlir/lib/Parser/Parser.cpp

1473

nit: Check to see if the load was successful first. If the load fails, the second call is guaranteed to fail.

mlir/tools/mlir-opt/mlir-opt.cpp

134

nit: Drop trivial braces.

You could also use llvm::interleave here.

[mehdi_amini]

Address most comments

[mehdi_amini]

(remove spurious new line)

[mehdi_amini]

mlir/include/mlir/IR/MLIRContext.h

26

Leftover from previous iterations...

52

"Someone" needs to own the DialectRegistry: ultimately I'm trying to remove entirely the concept of global registry.

I originally didn't have a registry in the context, but that forces the caller of the parser to thread through separately the registry.
Instead we can register the dialects with the context (without loading them) and load them on-demand when parsing.

mlir/include/mlir/Transforms/Passes.td

105

Maybe a bad rebase? I don't remember changing this

mlir/lib/Conversion/PassDetail.h

18

From Dialect.h you mean?

mlir/lib/IR/MLIRContext.cpp

467

I can't read ternary well without these :)

I can remove it if it hurts your eyes

[mehdi_amini]

(revert unrelated change in mlir/Transforms/Passes.td)

[mehdi_amini]

Restore the change for bufferplacement: it is needed after all

[mehdi_amini]

Thanks for the review River! That was very much a draft...

mlir/include/mlir/Transforms/Passes.td

105

Actually the BufferPlacement pass was just not deriving from the generated base class but directly from FunctionPass, so this change is needed.

[mehdi_amini]

Rebase

[mehdi_amini]

Update ctor initialization

[mehdi_amini]

mlir/include/mlir/IR/MLIRContext.h

43

@rriddle : FYI I added this ctor parameter for the transition. This ease landing this while clients have to update all their uses.

[rriddle]

Please fix the clang-tidy failures.

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

260

Can we update insert to be variadic?

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

25

nit: Drop the extra spaces surrounding "omp..."

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

26

Same here and a few other places.

mlir/include/mlir/IR/Dialect.h

238

You could also have used MapVector, but either is fine.

260

nit: name is in a different style.

267

Same here.

mlir/include/mlir/IR/MLIRContext.h

52

Is the registry on the context intended to be temporary then? i.e. will it be removed when the global registry is gone?

mlir/include/mlir/Pass/Pass.h

62

I think we should remove operations here, a dialect dependency can also arise from depending on attributes or types.

mlir/include/mlir/Pass/PassBase.td

81

Same here.

mlir/include/mlir/Support/MlirOptMain.h

28

nit: an option upfront -> an upfront?

mlir/lib/CAPI/IR/IR.cpp

54

nit: Parameter comment on the false.

mlir/lib/Conversion/PassDetail.h

18

Yeah, I just guessed.

mlir/lib/IR/Dialect.cpp

33

nit: name.str()?

mlir/lib/IR/MLIRContext.cpp

447

Unresolved?

447

nit: Drop trivial braces here.

mlir/lib/Pass/Pass.cpp

747

nit: Fix the naming style here.

mlir/tools/mlir-tblgen/PassGen.cpp

111

nit: Naming style here.

[mehdi_amini]

Address comments

[mehdi_amini]

mlir/include/mlir/IR/Dialect.h

238

I thought about it, but MapVector depends on the order of insertion, and right now this is used with global initializer whose order isn't defined.
It would be deterministic across runs of the same binary (hopefully), but not across platform/environment or across revisions.

mlir/include/mlir/IR/MLIRContext.h

52

Addressed offline: it is only for allowing lazy loading in the parser

mlir/include/mlir/Pass/Pass.h

62

/// Register dependent dialects for the current pass.
/// A pass is expected to register the dialects it will create entities for
/// (Operations, Types, Attributes), other than dialect that exists in the
/// input. For example, a pass that converts from Linalg to Affine would
/// register the Affine dialect but does not need to register Linalg.

[mehdi_amini]

rebase

[tskeith]

This change is causing some flang test to fail. See http://lab.llvm.org:8011/builders/flang-aarch64-ubuntu/builds/2928

[CarolineConcatto]

@mehdi_amini this patch is breaking flang buildbot.
http://lab.llvm.org:8011/builders/flang-aarch64-ubuntu/builds/2928
Do you mind to revert?

[mehdi_amini]

I fixed this earlier, tracking this build right now to see if flang tests
are passing:
http://lab.llvm.org:8011/builders/flang-aarch64-ubuntu/builds/2949

[mehdi_amini]

Build is back green!
http://lab.llvm.org:8011/builders/flang-aarch64-ubuntu/builds/2949

Sorry for the breakage: I tested with flang, but then I had a last minute
change of an API and only tested MLIR...
In general if a patch is breaking flang, feel free to revert directly and
notify in the revision!

mlir/examples/standalone/standalone-opt/standalone-opt.cpp

 #include "llvm/Support/SourceMgr.h"
 #include "llvm/Support/ToolOutputFile.h"
 
 #include "Standalone/StandaloneDialect.h"
 
 int main(int argc, char **argv) {
   mlir::registerAllDialects();
   mlir::registerAllPasses();
-
-  mlir::registerDialect<mlir::standalone::StandaloneDialect>();
   // TODO: Register standalone passes here.
 
-  return failed(mlir::MlirOptMain(argc, argv, "Standalone optimizer driver\n"));
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::standalone::StandaloneDialect>();
+  registry.insert<mlir::StandardOpsDialect>();
+  // Add the following to include *all* MLIR Core dialects, or selectively
+  // include what you need like above. You only need to register dialects that
+  // will be *parsed* by the tool, not the one generated
+  // registerAllDialects(registry);
+
+  return failed(
+      mlir::MlirOptMain(argc, argv, "Standalone optimizer driver\n", registry));
 }

mlir/examples/standalone/test/Standalone/standalone-opt.mlir

 // RUN: standalone-opt --show-dialects | FileCheck %s
-// CHECK: Registered Dialects:
+// CHECK: Available Dialects:
 // CHECK: standalone

mlir/examples/toy/Ch2/toyc.cpp

std::unique_ptr<toy::ModuleAST> parseInputFile(llvm::StringRef filename) {
   }
   auto buffer = fileOrErr.get()->getBuffer();
   LexerBuffer lexer(buffer.begin(), buffer.end(), std::string(filename));
   Parser parser(lexer);
   return parser.parseModule();
 }
 
 int dumpMLIR() {
-  // Register our Dialect with MLIR.
-  mlir::registerDialect<mlir::toy::ToyDialect>();
-
-  mlir::MLIRContext context;
+  mlir::MLIRContext context(/*loadAllDialects=*/false);
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
 
   // Handle '.toy' input to the compiler.
   if (inputType != InputType::MLIR &&
       !llvm::StringRef(inputFilename).endswith(".mlir")) {
     auto moduleAST = parseInputFile(inputFilename);
     if (!moduleAST)
       return 6;
     mlir::OwningModuleRef module = mlirGen(context, *moduleAST);

mlir/examples/toy/Ch3/toyc.cpp

int loadMLIR(llvm::SourceMgr &sourceMgr, mlir::MLIRContext &context,
   if (!module) {
     llvm::errs() << "Error can't load file " << inputFilename << "\n";
     return 3;
   }
   return 0;
 }
 
 int dumpMLIR() {
-  // Register our Dialect with MLIR.
-  mlir::registerDialect<mlir::toy::ToyDialect>();
+  mlir::MLIRContext context(/*loadAllDialects=*/false);
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
 
-  mlir::MLIRContext context;
   mlir::OwningModuleRef module;
   llvm::SourceMgr sourceMgr;
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
   if (int error = loadMLIR(sourceMgr, context, module))
     return error;
 
   if (enableOpt) {
     mlir::PassManager pm(&context);

mlir/examples/toy/Ch4/toyc.cpp

int loadMLIR(llvm::SourceMgr &sourceMgr, mlir::MLIRContext &context,
   if (!module) {
     llvm::errs() << "Error can't load file " << inputFilename << "\n";
     return 3;
   }
   return 0;
 }
 
 int dumpMLIR() {
-  // Register our Dialect with MLIR.
-  mlir::registerDialect<mlir::toy::ToyDialect>();
+  mlir::MLIRContext context(/*loadAllDialects=*/false);
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
 
-  mlir::MLIRContext context;
   mlir::OwningModuleRef module;
   llvm::SourceMgr sourceMgr;
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
   if (int error = loadMLIR(sourceMgr, context, module))
     return error;
 
   if (enableOpt) {
     mlir::PassManager pm(&context);

mlir/examples/toy/Ch5/mlir/LowerToAffineLoops.cpp

 //===----------------------------------------------------------------------===//
 
 /// This is a partial lowering to affine loops of the toy operations that are
 /// computationally intensive (like matmul for example...) while keeping the
 /// rest of the code in the Toy dialect.
 namespace {
 struct ToyToAffineLoweringPass
     : public PassWrapper<ToyToAffineLoweringPass, FunctionPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect, StandardOpsDialect>();
+  }
   void runOnFunction() final;
 };
 } // end anonymous namespace.
 
 void ToyToAffineLoweringPass::runOnFunction() {
   auto function = getFunction();
 
   // We only lower the main function as we expect that all other functions have

mlir/examples/toy/Ch5/toyc.cpp

int loadMLIR(llvm::SourceMgr &sourceMgr, mlir::MLIRContext &context,
   if (!module) {
     llvm::errs() << "Error can't load file " << inputFilename << "\n";
     return 3;
   }
   return 0;
 }
 
 int dumpMLIR() {
-  // Register our Dialect with MLIR.
-  mlir::registerDialect<mlir::toy::ToyDialect>();
+  mlir::MLIRContext context(/*loadAllDialects=*/false);
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
 
-  mlir::MLIRContext context;
   mlir::OwningModuleRef module;
   llvm::SourceMgr sourceMgr;
   mlir::SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
   if (int error = loadMLIR(sourceMgr, context, module))
     return error;
 
   mlir::PassManager pm(&context);
   // Apply any generic pass manager command line options and run the pipeline.

mlir/examples/toy/Ch6/mlir/LowerToAffineLoops.cpp

 //===----------------------------------------------------------------------===//
 
 /// This is a partial lowering to affine loops of the toy operations that are
 /// computationally intensive (like matmul for example...) while keeping the
 /// rest of the code in the Toy dialect.
 namespace {
 struct ToyToAffineLoweringPass
     : public PassWrapper<ToyToAffineLoweringPass, FunctionPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect, StandardOpsDialect>();
+  }
   void runOnFunction() final;
 };
 } // end anonymous namespace.
 
 void ToyToAffineLoweringPass::runOnFunction() {
   auto function = getFunction();
 
   // We only lower the main function as we expect that all other functions have

mlir/examples/toy/Ch6/mlir/LowerToLLVM.cpp

 
 //===----------------------------------------------------------------------===//
 // ToyToLLVMLoweringPass
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct ToyToLLVMLoweringPass
     : public PassWrapper<ToyToLLVMLoweringPass, OperationPass<ModuleOp>> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<LLVM::LLVMDialect, scf::SCFDialect>();
+  }
   void runOnOperation() final;
 };
 } // end anonymous namespace
 
 void ToyToLLVMLoweringPass::runOnOperation() {
   // The first thing to define is the conversion target. This will define the
   // final target for this lowering. For this lowering, we are only targeting
   // the LLVM dialect.

mlir/examples/toy/Ch6/toyc.cpp

int main(int argc, char **argv) {
 
   cl::ParseCommandLineOptions(argc, argv, "toy compiler\n");
 
   if (emitAction == Action::DumpAST)
     return dumpAST();
 
   // If we aren't dumping the AST, then we are compiling with/to MLIR.
 
-  // Register our Dialect with MLIR.
-  mlir::registerDialect<mlir::toy::ToyDialect>();
+  mlir::MLIRContext context(/*loadAllDialects=*/false);
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
 
-  mlir::MLIRContext context;
   mlir::OwningModuleRef module;
   if (int error = loadAndProcessMLIR(context, module))
     return error;
 
   // If we aren't exporting to non-mlir, then we are done.
   bool isOutputingMLIR = emitAction <= Action::DumpMLIRLLVM;
   if (isOutputingMLIR) {
     module->dump();

mlir/examples/toy/Ch7/mlir/LowerToAffineLoops.cpp

 //===----------------------------------------------------------------------===//
 
 /// This is a partial lowering to affine loops of the toy operations that are
 /// computationally intensive (like matmul for example...) while keeping the
 /// rest of the code in the Toy dialect.
 namespace {
 struct ToyToAffineLoweringPass
     : public PassWrapper<ToyToAffineLoweringPass, FunctionPass> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect, StandardOpsDialect>();
+  }
   void runOnFunction() final;
 };
 } // end anonymous namespace.
 
 void ToyToAffineLoweringPass::runOnFunction() {
   auto function = getFunction();
 
   // We only lower the main function as we expect that all other functions have

mlir/examples/toy/Ch7/mlir/LowerToLLVM.cpp

 
 //===----------------------------------------------------------------------===//
 // ToyToLLVMLoweringPass
 //===----------------------------------------------------------------------===//
 
 namespace {
 struct ToyToLLVMLoweringPass
     : public PassWrapper<ToyToLLVMLoweringPass, OperationPass<ModuleOp>> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<LLVM::LLVMDialect, scf::SCFDialect>();
+  }
   void runOnOperation() final;
 };
 } // end anonymous namespace
 
 void ToyToLLVMLoweringPass::runOnOperation() {
   // The first thing to define is the conversion target. This will define the
   // final target for this lowering. For this lowering, we are only targeting
   // the LLVM dialect.

mlir/examples/toy/Ch7/toyc.cpp

int main(int argc, char **argv) {
 
   cl::ParseCommandLineOptions(argc, argv, "toy compiler\n");
 
   if (emitAction == Action::DumpAST)
     return dumpAST();
 
   // If we aren't dumping the AST, then we are compiling with/to MLIR.
 
-  // Register our Dialect with MLIR.
-  mlir::registerDialect<mlir::toy::ToyDialect>();
+  mlir::MLIRContext context(/*loadAllDialects=*/false);
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
 
-  mlir::MLIRContext context;
   mlir::OwningModuleRef module;
   if (int error = loadAndProcessMLIR(context, module))
     return error;
 
   // If we aren't exporting to non-mlir, then we are done.
   bool isOutputingMLIR = emitAction <= Action::DumpMLIRLLVM;
   if (isOutputingMLIR) {
     module->dump();

mlir/include/mlir-c/Registration.h

 /*===-- mlir-c/Registration.h - Registration functions for MLIR ---*- 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                    *|
 |*                                                                            *|
 \*===----------------------------------------------------------------------===*/
 
 #ifndef MLIR_C_REGISTRATION_H
 #define MLIR_C_REGISTRATION_H
 
+#include "mlir-c/IR.h"
+
 #ifdef __cplusplus
 extern "C" {
 #endif
 
-/** Registers all dialects known to core MLIR with the system. This must be
- * called before creating an MlirContext if it needs access to the registered
- * dialects. */
-void mlirRegisterAllDialects();
+/** Registers all dialects known to core MLIR with the provided Context.
+ * This is needed before creating IR for these Dialects.
+ */
+void mlirRegisterAllDialects(MlirContext context);
 
 #ifdef __cplusplus
 }
 #endif
 
 #endif // MLIR_C_REGISTRATION_H

mlir/include/mlir/Conversion/Passes.td

let description = [{
 
     -   Functions without a body are not modified.
     -   The semantics of the other functions is preserved.
     -   Individual operations other than those mentioned above are not modified
         if they do not depend on the loop iterator value or on the result of
         `affine.apply`.
   }];
   let constructor = "mlir::createLowerAffinePass()";
+  let dependentDialects = [
+    "scf::SCFDialect",
+    "StandardOpsDialect",
+    "vector::VectorDialect"
+  ];
 }
 
 //===----------------------------------------------------------------------===//
 // AVX512ToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertAVX512ToLLVM : Pass<"convert-avx512-to-llvm", "ModuleOp"> {
   let summary = "Convert the operations from the avx512 dialect into the LLVM "
                 "dialect";
   let constructor = "mlir::createConvertAVX512ToLLVMPass()";
+  let dependentDialects = ["LLVM::LLVMDialect", "LLVM::LLVMAVX512Dialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUCommon
 //===----------------------------------------------------------------------===//
 
 def GpuToLLVMConversionPass : Pass<"gpu-to-llvm", "ModuleOp"> {
   let summary = "Convert GPU dialect to LLVM dialect with GPU runtime calls";
   let constructor = "mlir::createGpuToLLVMConversionPass()";
   let options = [
     Option<"gpuBinaryAnnotation", "gpu-binary-annotation", "std::string",
            "", "Annotation attribute string for GPU binary">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToNVVM
 //===----------------------------------------------------------------------===//
 
 def ConvertGpuOpsToNVVMOps : Pass<"convert-gpu-to-nvvm", "gpu::GPUModuleOp"> {
   let summary = "Generate NVVM operations for gpu operations";
   let constructor = "mlir::createLowerGpuOpsToNVVMOpsPass()";
+  let dependentDialects = ["NVVM::NVVMDialect"];
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToROCDL
 //===----------------------------------------------------------------------===//
 
 def ConvertGpuOpsToROCDLOps : Pass<"convert-gpu-to-rocdl", "gpu::GPUModuleOp"> {
   let summary = "Generate ROCDL operations for gpu operations";
   let constructor = "mlir::createLowerGpuOpsToROCDLOpsPass()";
+  let dependentDialects = ["ROCDL::ROCDLDialect"];
   let options = [
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToSPIRV
 //===----------------------------------------------------------------------===//
 
 def ConvertGPUToSPIRV : Pass<"convert-gpu-to-spirv", "ModuleOp"> {
   let summary = "Convert GPU dialect to SPIR-V dialect";
   let constructor = "mlir::createConvertGPUToSPIRVPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // GPUToVulkan
 //===----------------------------------------------------------------------===//
 
 def ConvertGpuLaunchFuncToVulkanLaunchFunc
     : Pass<"convert-gpu-launch-to-vulkan-launch", "ModuleOp"> {
   let summary = "Convert gpu.launch_func to vulkanLaunch external call";
   let constructor = "mlir::createConvertGpuLaunchFuncToVulkanLaunchFuncPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 def ConvertVulkanLaunchFuncToVulkanCalls
     : Pass<"launch-func-to-vulkan", "ModuleOp"> {
   let summary = "Convert vulkanLaunch external call to Vulkan runtime external "
                 "calls";
   let constructor = "mlir::createConvertVulkanLaunchFuncToVulkanCallsPass()";
+  let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // LinalgToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertLinalgToLLVM : Pass<"convert-linalg-to-llvm", "ModuleOp"> {
   let summary = "Convert the operations from the linalg dialect into the LLVM "
                 "dialect";
   let constructor = "mlir::createConvertLinalgToLLVMPass()";
+  let dependentDialects = ["scf::SCFDialect", "LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // LinalgToStandard
 //===----------------------------------------------------------------------===//
 
 def ConvertLinalgToStandard : Pass<"convert-linalg-to-std", "ModuleOp"> {
   let summary = "Convert the operations from the linalg dialect into the "
                 "Standard dialect";
   let constructor = "mlir::createConvertLinalgToStandardPass()";
+  let dependentDialects = ["StandardOpsDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // LinalgToSPIRV
 //===----------------------------------------------------------------------===//
 
 def ConvertLinalgToSPIRV : Pass<"convert-linalg-to-spirv", "ModuleOp"> {
   let summary = "Convert Linalg ops to SPIR-V ops";
   let constructor = "mlir::createLinalgToSPIRVPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // SCFToStandard
 //===----------------------------------------------------------------------===//
 
 def SCFToStandard : Pass<"convert-scf-to-std"> {
   let summary = "Convert SCF dialect to Standard dialect, replacing structured"
                 " control flow with a CFG";
   let constructor = "mlir::createLowerToCFGPass()";
+  let dependentDialects = ["StandardOpsDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // SCFToGPU
 //===----------------------------------------------------------------------===//
 
 def ConvertAffineForToGPU : FunctionPass<"convert-affine-for-to-gpu"> {
   let summary = "Convert top-level AffineFor Ops to GPU kernels";
   let constructor = "mlir::createAffineForToGPUPass()";
+  let dependentDialects = ["gpu::GPUDialect"];
   let options = [
     Option<"numBlockDims", "gpu-block-dims", "unsigned", /*default=*/"1u",
            "Number of GPU block dimensions for mapping">,
     Option<"numThreadDims", "gpu-thread-dims", "unsigned", /*default=*/"1u",
            "Number of GPU thread dimensions for mapping">
   ];
 }
 
 def ConvertParallelLoopToGpu : Pass<"convert-parallel-loops-to-gpu"> {
   let summary = "Convert mapped scf.parallel ops to gpu launch operations";
   let constructor = "mlir::createParallelLoopToGpuPass()";
+  let dependentDialects = ["AffineDialect", "gpu::GPUDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // ShapeToStandard
 //===----------------------------------------------------------------------===//
 
 def ConvertShapeToStandard : Pass<"convert-shape-to-std", "ModuleOp"> {
   let summary = "Convert operations from the shape dialect into the standard "
                 "dialect";
   let constructor = "mlir::createConvertShapeToStandardPass()";
+  let dependentDialects = ["StandardOpsDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // ShapeToSCF
 //===----------------------------------------------------------------------===//
 
 def ConvertShapeToSCF : FunctionPass<"convert-shape-to-scf"> {
   let summary = "Convert operations from the shape dialect to the SCF dialect";
   let constructor = "mlir::createConvertShapeToSCFPass()";
+  let dependentDialects = ["scf::SCFDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // SPIRVToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertSPIRVToLLVM : Pass<"convert-spirv-to-llvm", "ModuleOp"> {
   let summary = "Convert SPIR-V dialect to LLVM dialect";
   let constructor = "mlir::createConvertSPIRVToLLVMPass()";
+  let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // StandardToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertStandardToLLVM : Pass<"convert-std-to-llvm", "ModuleOp"> {
   let summary = "Convert scalar and vector operations from the Standard to the "
let description = [{
 
     Functions converted to LLVM IR. Function arguments types are converted
     one-to-one. Function results are converted one-to-one and, in case more than
     1 value is returned, packed into an LLVM IR struct type. Function calls and
     returns are updated accordingly. Block argument types are updated to use
     LLVM IR types.
   }];
   let constructor = "mlir::createLowerToLLVMPass()";
+  let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"useAlignedAlloc", "use-aligned-alloc", "bool", /*default=*/"false",
            "Use aligned_alloc in place of malloc for heap allocations">,
     Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
            /*default=*/"false",
            "Replace FuncOp's MemRef arguments with bare pointers to the MemRef "
            "element types">,
     Option<"emitCWrappers", "emit-c-wrappers", "bool", /*default=*/"false",
            "Emit wrappers for C-compatible pointer-to-struct memref "
            "descriptors">,
     Option<"indexBitwidth", "index-bitwidth", "unsigned",
            /*default=kDeriveIndexBitwidthFromDataLayout*/"0",
            "Bitwidth of the index type, 0 to use size of machine word">,
     Option<"dataLayout", "data-layout", "std::string",
            /*default=*/"\"\"",
            "String description (LLVM format) of the data layout that is "
            "expected on the produced module">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // StandardToSPIRV
 //===----------------------------------------------------------------------===//
 
 def LegalizeStandardForSPIRV : Pass<"legalize-std-for-spirv"> {
   let summary = "Legalize standard ops for SPIR-V lowering";
   let constructor = "mlir::createLegalizeStdOpsForSPIRVLoweringPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 def ConvertStandardToSPIRV : Pass<"convert-std-to-spirv", "ModuleOp"> {
   let summary = "Convert Standard Ops to SPIR-V dialect";
   let constructor = "mlir::createConvertStandardToSPIRVPass()";
+  let dependentDialects = ["spirv::SPIRVDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToSCF
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToSCF : FunctionPass<"convert-vector-to-scf"> {
   let summary = "Lower the operations from the vector dialect into the SCF "
                 "dialect";
   let constructor = "mlir::createConvertVectorToSCFPass()";
+  let dependentDialects = ["AffineDialect", "scf::SCFDialect"];
   let options = [
     Option<"fullUnroll", "full-unroll", "bool", /*default=*/"false",
            "Perform full unrolling when converting vector transfers to SCF">,
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToLLVM
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToLLVM : Pass<"convert-vector-to-llvm", "ModuleOp"> {
   let summary = "Lower the operations from the vector dialect into the LLVM "
                 "dialect";
   let constructor = "mlir::createConvertVectorToLLVMPass()";
+  let dependentDialects = ["LLVM::LLVMDialect"];
   let options = [
     Option<"reassociateFPReductions", "reassociate-fp-reductions",
            "bool", /*default=*/"false",
            "Allows llvm to reassociate floating-point reductions for speed">
   ];
 }
 
 //===----------------------------------------------------------------------===//
 // VectorToROCDL
 //===----------------------------------------------------------------------===//
 
 def ConvertVectorToROCDL : Pass<"convert-vector-to-rocdl", "ModuleOp"> {
   let summary = "Lower the operations from the vector dialect into the ROCDL "
                 "dialect";
   let constructor = "mlir::createConvertVectorToROCDLPass()";
+  let dependentDialects = ["ROCDL::ROCDLDialect"];
 }
 
 #endif // MLIR_CONVERSION_PASSES

mlir/include/mlir/Dialect/Affine/Passes.td

Option<"unrollJamFactor", "unroll-jam-factor", "unsigned",
            /*default=*/"4",
            "Use this unroll jam factor for all loops (default 4)">,
   ];
 }
 
 def AffineVectorize : FunctionPass<"affine-super-vectorize"> {
   let summary = "Vectorize to a target independent n-D vector abstraction";
   let constructor = "mlir::createSuperVectorizePass()";
+  let dependentDialects = ["vector::VectorDialect"];
   let options = [
     ListOption<"vectorSizes", "virtual-vector-size", "int64_t",
                "Specify an n-D virtual vector size for vectorization",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">,
     // Optionally, the fixed mapping from loop to fastest varying MemRef
     // dimension for all the MemRefs within a loop pattern:
     //   the index represents the loop depth, the value represents the k^th
     //   fastest varying memory dimension.

mlir/include/mlir/Dialect/LLVMIR/LLVMDialect.h

 // LLVM type system.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_LLVMDIALECT_H_
 #define MLIR_DIALECT_LLVMIR_LLVMDIALECT_H_
 
 #include "mlir/Dialect/LLVMIR/LLVMTypes.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/IR/OpImplementation.h"
 #include "mlir/IR/TypeSupport.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Interfaces/ControlFlowInterfaces.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"

mlir/include/mlir/Dialect/LLVMIR/LLVMOpBase.td

 #ifndef LLVMIR_OP_BASE
 #define LLVMIR_OP_BASE
 
 include "mlir/IR/OpBase.td"
 
 def LLVM_Dialect : Dialect {
   let name = "llvm";
   let cppNamespace = "LLVM";
+
+  /// FIXME: at the moment this is a dependency of the translation to LLVM IR,
+  /// not really one of this dialect per-se.
+  let dependentDialects = ["omp::OpenMPDialect"];
+
   let hasRegionArgAttrVerify = 1;
   let hasOperationAttrVerify = 1;
   let extraClassDeclaration = [{
     /// Name of the data layout attributes.
     static StringRef getDataLayoutAttrName() { return "llvm.data_layout"; }
 
     /// Verifies if the given string is a well-formed data layout descriptor.
     /// Uses `reportError` to report errors.

mlir/include/mlir/Dialect/LLVMIR/NVVMDialect.h

 //===- NVVMDialect.h - MLIR NVVM IR dialect ---------------------*- 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 NVVM IR dialect in MLIR, containing NVVM operations and
 // NVVM specific extensions to the LLVM type system.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_NVVMDIALECT_H_
 #define MLIR_DIALECT_LLVMIR_NVVMDIALECT_H_
 
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 namespace NVVM {
 
 ///// Ops /////

mlir/include/mlir/Dialect/LLVMIR/NVVMOps.td

 
 //===----------------------------------------------------------------------===//
 // NVVM dialect definitions
 //===----------------------------------------------------------------------===//
 
 def NVVM_Dialect : Dialect {
   let name = "nvvm";
   let cppNamespace = "NVVM";
+  let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // NVVM op definitions
 //===----------------------------------------------------------------------===//
 
 class NVVM_Op<string mnemonic, list<OpTrait> traits = []> :
   LLVM_OpBase<NVVM_Dialect, mnemonic, traits> {

mlir/include/mlir/Dialect/LLVMIR/ROCDLDialect.h

 // https://github.com/RadeonOpenCompute/ROCm-Device-Libs/blob/master/doc/OCKL.md
 // https://llvm.org/docs/AMDGPUUsage.html
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_DIALECT_LLVMIR_ROCDLDIALECT_H_
 #define MLIR_DIALECT_LLVMIR_ROCDLDIALECT_H_
 
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/IR/Dialect.h"
 #include "mlir/IR/OpDefinition.h"
 #include "mlir/Interfaces/SideEffectInterfaces.h"
 
 namespace mlir {
 namespace ROCDL {
 
 ///// Ops /////

mlir/include/mlir/Dialect/LLVMIR/ROCDLOps.td

 
 //===----------------------------------------------------------------------===//
 // ROCDL dialect definitions
 //===----------------------------------------------------------------------===//
 
 def ROCDL_Dialect : Dialect {
   let name = "rocdl";
   let cppNamespace = "ROCDL";
+  let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // ROCDL op definitions
 //===----------------------------------------------------------------------===//
 
 class ROCDL_Op<string mnemonic, list<OpTrait> traits = []> :
   LLVM_OpBase<ROCDL_Dialect, mnemonic, traits> {

mlir/include/mlir/Dialect/Linalg/Passes.td

 def LinalgFusion : FunctionPass<"linalg-fusion"> {
   let summary = "Fuse operations in the linalg dialect";
   let constructor = "mlir::createLinalgFusionPass()";
 }
 
 def LinalgFusionOfTensorOps : Pass<"linalg-fusion-for-tensor-ops"> {
   let summary = "Fuse operations on RankedTensorType in linalg dialect";
   let constructor = "mlir::createLinalgFusionOfTensorOpsPass()";
+  let dependentDialects = ["AffineDialect"];
 }
 
 def LinalgLowerToAffineLoops : FunctionPass<"convert-linalg-to-affine-loops"> {
   let summary = "Lower the operations from the linalg dialect into affine "
                 "loops";
   let constructor = "mlir::createConvertLinalgToAffineLoopsPass()";
+  let dependentDialects = ["AffineDialect"];
 }
 
 def LinalgLowerToLoops : FunctionPass<"convert-linalg-to-loops"> {
   let summary = "Lower the operations from the linalg dialect into loops";
   let constructor = "mlir::createConvertLinalgToLoopsPass()";
+  let dependentDialects = ["scf::SCFDialect", "AffineDialect"];
 }
 
 def LinalgOnTensorsToBuffers : Pass<"convert-linalg-on-tensors-to-buffers", "ModuleOp"> {
   let summary = "Convert the Linalg operations which work on tensor-type "
                 "operands or results to use buffers instead";
   let constructor = "mlir::createConvertLinalgOnTensorsToBuffersPass()";
 }
 
 def LinalgLowerToParallelLoops
     : FunctionPass<"convert-linalg-to-parallel-loops"> {
   let summary = "Lower the operations from the linalg dialect into parallel "
                 "loops";
   let constructor = "mlir::createConvertLinalgToParallelLoopsPass()";
+  let dependentDialects = ["AffineDialect", "scf::SCFDialect"];
 }
 
 def LinalgPromotion : FunctionPass<"linalg-promote-subviews"> {
   let summary = "Promote subview ops to local buffers";
   let constructor = "mlir::createLinalgPromotionPass()";
   let options = [
     Option<"dynamicBuffers", "test-promote-dynamic", "bool",
            /*default=*/"false", "Test generation of dynamic promoted buffers">,
     Option<"useAlloca", "test-use-alloca", "bool",
            /*default=*/"false", "Test generation of alloca'ed buffers.">
   ];
 }
 
 def LinalgTiling : FunctionPass<"linalg-tile"> {
   let summary = "Tile operations in the linalg dialect";
   let constructor = "mlir::createLinalgTilingPass()";
+  let dependentDialects = [
+    "AffineDialect", "scf::SCFDialect"
+  ];
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t",
                "Test generation of dynamic promoted buffers",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
 }
 
 def LinalgTilingToParallelLoops
     : FunctionPass<"linalg-tile-to-parallel-loops"> {
   let summary = "Tile operations in the linalg dialect to parallel loops";
   let constructor = "mlir::createLinalgTilingToParallelLoopsPass()";
   let options = [
     ListOption<"tileSizes", "linalg-tile-sizes", "int64_t",
                "Test generation of dynamic promoted buffers",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
+  let dependentDialects = ["AffineDialect", "scf::SCFDialect"];
 }
 
 #endif // MLIR_DIALECT_LINALG_PASSES

mlir/include/mlir/Dialect/SCF/Passes.td

 def SCFParallelLoopTiling : FunctionPass<"parallel-loop-tiling"> {
   let summary = "Tile parallel loops";
   let constructor = "mlir::createParallelLoopTilingPass()";
   let options = [
     ListOption<"tileSizes", "parallel-loop-tile-sizes", "int64_t",
                "Factors to tile parallel loops by",
                "llvm::cl::ZeroOrMore, llvm::cl::MiscFlags::CommaSeparated">
   ];
+  let dependentDialects = ["AffineDialect"];
 }
 
 #endif // MLIR_DIALECT_SCF_PASSES

mlir/include/mlir/IR/Dialect.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_IR_DIALECT_H
 #define MLIR_IR_DIALECT_H
 
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/TypeID.h"
 
+#include <map>
+
 namespace mlir {
 class DialectAsmParser;
 class DialectAsmPrinter;
 class DialectInterface;
 class OpBuilder;
 class Type;
 
-using DialectAllocatorFunction = std::function<void(MLIRContext *)>;
+using DialectAllocatorFunction = std::function<Dialect *(MLIRContext *)>;
 
 /// Dialects are groups of MLIR operations and behavior associated with the
 /// entire group.  For example, hooks into other systems for constant folding,
 /// default named types for asm printing, etc.
 ///
 /// Instances of the dialect object are global across all MLIRContext's that may
 /// be active in the process.
 ///
private:
   /// Flag that specifies whether this dialect allows unregistered types, i.e.
   /// types prefixed with the dialect namespace but not registered with addType.
   /// These types are represented with OpaqueType.
   bool unknownTypesAllowed = false;
 
   /// A collection of registered dialect interfaces.
   DenseMap<TypeID, std::unique_ptr<DialectInterface>> registeredInterfaces;
 
-  /// Registers a specific dialect creation function with the global registry.
-  /// Used through the registerDialect template.
-  /// Registrations are deduplicated by dialect TypeID and only the first
-  /// registration will be used.
-  static void
-  registerDialectAllocator(TypeID typeID,
-                           const DialectAllocatorFunction &function);
-  template <typename ConcreteDialect>
   friend void registerDialect();
   friend class MLIRContext;
 };
 
-/// Registers all dialects and hooks from the global registries with the
-/// specified 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 in the Context as operations are encountered.
+class DialectRegistry {
+  using MapTy =
+      std::map<std::string, std::pair<TypeID, DialectAllocatorFunction>>;
+
+public:
+  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.
+  void insert(TypeID typeID, StringRef name, DialectAllocatorFunction ctor);
+
+  /// Load a dialect for this namespace in the provided context.
+  Dialect *loadByName(StringRef name, MLIRContext *context);
+
+  // Register all dialects available in the current registry with the registry
+  // in the provided context.
+  void appendTo(DialectRegistry &destination) {
+    for (const auto &nameAndRegistrationIt : registry)
+      destination.insert(nameAndRegistrationIt.second.first,
+                         nameAndRegistrationIt.first,
+                         nameAndRegistrationIt.second.second);
+  }
+  // Load all dialects available in the registry in the provided context.
+  void loadAll(MLIRContext *context) {
+    for (const auto &nameAndRegistrationIt : registry)
+      nameAndRegistrationIt.second.second(context);
+  }
+
+  MapTy::const_iterator begin() const { return registry.begin(); }
+  MapTy::const_iterator end() const { return registry.end(); }
+
+private:
+  MapTy registry;
+};
+
+/// Deprecated: this provides a global registry for convenience, while we're
+/// transitionning the registration mechanism to a stateless approach.
+DialectRegistry &getGlobalDialectRegistry();
+
+/// Registers all dialects from the global registries with the
+/// specified MLIRContext. This won't load the dialects in the context,
+/// but only make them available for lazy loading by name.
 /// Note: This method is not thread-safe.
 void registerAllDialects(MLIRContext *context);
 
+/// Register and return the dialect with the given namespace in the provided
+/// context. Returns nullptr is there is no constructor registered for this
+/// dialect.
+inline Dialect *registerDialect(StringRef name, MLIRContext *context) {
+  return getGlobalDialectRegistry().loadByName(name, context);
+}
+
 /// Utility to register a dialect. Client can register their dialect with the
 /// global registry by calling registerDialect<MyDialect>();
 /// Note: This method is not thread-safe.
 template <typename ConcreteDialect> void registerDialect() {
-  Dialect::registerDialectAllocator(
-      TypeID::get<ConcreteDialect>(),
-      [](MLIRContext *ctx) { ctx->getOrCreateDialect<ConcreteDialect>(); });
+  getGlobalDialectRegistry().insert<ConcreteDialect>();
 }
 
 /// DialectRegistration provides a global initializer that registers a Dialect
 /// allocation routine.
 ///
 /// Usage:
 ///
 ///   // At namespace scope.

mlir/include/mlir/IR/FunctionSupport.h

LogicalResult FunctionLike<ConcreteType>::verifyTrait(Operation *op) {
   for (unsigned i = 0, e = funcOp.getNumArguments(); i != e; ++i) {
     // Verify that all of the argument attributes are dialect attributes, i.e.
     // that they contain a dialect prefix in their name.  Call the dialect, if
     // registered, to verify the attributes themselves.
     for (auto attr : funcOp.getArgAttrs(i)) {
       if (!attr.first.strref().contains('.'))
         return funcOp.emitOpError("arguments may only have dialect attributes");
       auto dialectNamePair = attr.first.strref().split('.');
-      if (auto *dialect = ctx->getRegisteredDialect(dialectNamePair.first)) {
+      if (auto *dialect = ctx->getLoadedDialect(dialectNamePair.first)) {
         if (failed(dialect->verifyRegionArgAttribute(op, /*regionIndex=*/0,
                                                      /*argIndex=*/i, attr)))
           return failure();
       }
     }
   }
 
   for (unsigned i = 0, e = funcOp.getNumResults(); i != e; ++i) {
     // Verify that all of the result attributes are dialect attributes, i.e.
     // that they contain a dialect prefix in their name.  Call the dialect, if
     // registered, to verify the attributes themselves.
     for (auto attr : funcOp.getResultAttrs(i)) {
       if (!attr.first.strref().contains('.'))
         return funcOp.emitOpError("results may only have dialect attributes");
       auto dialectNamePair = attr.first.strref().split('.');
-      if (auto *dialect = ctx->getRegisteredDialect(dialectNamePair.first)) {
+      if (auto *dialect = ctx->getLoadedDialect(dialectNamePair.first)) {
         if (failed(dialect->verifyRegionResultAttribute(op, /*regionIndex=*/0,
                                                         /*resultIndex=*/i,
                                                         attr)))
           return failure();
       }
     }
   }
 

mlir/include/mlir/IR/MLIRContext.h

 #include <functional>
 #include <memory>
 #include <vector>
 
 namespace mlir {
 class AbstractOperation;
 class DiagnosticEngine;
 class Dialect;
+class DialectRegistry;
 class InFlightDiagnostic;
 class Location;
 class MLIRContextImpl;
 class StorageUniquer;
+DialectRegistry &getGlobalDialectRegistry();
 
 /// MLIRContext is the top-level object for a collection of MLIR modules.  It
 /// holds immortal uniqued objects like types, and the tables used to unique
 /// them.
 ///
 /// 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:
-  explicit MLIRContext();
+  /// 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(bool loadAllDialects = true);
   ~MLIRContext();
 
-  /// Return information about all registered IR dialects.
-  std::vector<Dialect *> getRegisteredDialects();
+  /// Return information about all IR dialects loaded in the context.
+  std::vector<Dialect *> getLoadedDialects();
+
+  /// Return the dialect registry associated with this context.
+  DialectRegistry &getDialectRegistry();
+
+  /// 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 *getRegisteredDialect(StringRef name);
+  Dialect *getLoadedDialect(StringRef name);
 
   /// Get a registered IR dialect for the given derived dialect type. The
   /// derived type must provide a static 'getDialectNamespace' method.
-  template <typename T> T *getRegisteredDialect() {
-    return static_cast<T *>(getRegisteredDialect(T::getDialectNamespace()));
+  template <typename T>
+  T *getLoadedDialect() {
+    return static_cast<T *>(getLoadedDialect(T::getDialectNamespace()));
   }
 
   /// Get (or create) a dialect for the given derived dialect type. The derived
   /// type must provide a static 'getDialectNamespace' method.
   template <typename T>
-  T *getOrCreateDialect() {
-    return static_cast<T *>(getOrCreateDialect(
-        T::getDialectNamespace(), TypeID::get<T>(), [this]() {
+  T *getOrLoadDialect() {
+    return static_cast<T *>(
+        getOrLoadDialect(T::getDialectNamespace(), TypeID::get<T>(), [this]() {
           std::unique_ptr<T> dialect(new T(this));
-          dialect->dialectID = TypeID::get<T>();
           return dialect;
         }));
   }
 
+  /// Load a dialect in the context.
+  template <typename Dialect>
+  void loadDialect() {
+    getOrLoadDialect<Dialect>();
+  }
+
+  /// Load a list dialects in the context.
+  template <typename Dialect, typename OtherDialect, typename... MoreDialects>
+  void loadDialect() {
+    getOrLoadDialect<Dialect>();
+    loadDialect<OtherDialect, MoreDialects...>();
+  }
+
+  /// Deprecated: load all globally registered dialects into this context.
+  /// This method will be removed soon, it can be used temporarily as we're
+  /// phasing out the global registry.
+  void loadAllGloballyRegisteredDialects();
+
+  /// 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();
 
   /// Enables creating operations in unregistered dialects.
   void allowUnregisteredDialects(bool allow = true);
 
   /// Return true if multi-threading is enabled by the context.
   bool isMultithreadingEnabled();
public:
 
   /// Returns the storage uniquer used for constructing attribute storage
   /// instances. This should not be used directly.
   StorageUniquer &getAttributeUniquer();
 
 private:
   const std::unique_ptr<MLIRContextImpl> impl;
 
   /// 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 *getOrCreateDialect(StringRef dialectNamespace, TypeID dialectID,
-                              function_ref<std::unique_ptr<Dialect>()> ctor);
+  /// a dialect exist for this namespace with different TypeID. If a dialect has
+  /// not been loaded for this namespace/TypeID yet, use the provided ctor to
+  /// create one on the fly and load it. Returns a pointer to the dialect owned
+  /// by the context.
+  Dialect *getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
 
   MLIRContext(const MLIRContext &) = delete;
   void operator=(const MLIRContext &) = delete;
 };
 
 //===----------------------------------------------------------------------===//
 // MLIRContext CommandLine Options
 //===----------------------------------------------------------------------===//

mlir/include/mlir/IR/OpBase.td

class Dialect {
   string name = ?;
 
   // Short summary of the dialect.
   string summary = ?;
 
   // The description of the dialect.
   string description = ?;
 
+  // A list of dialects this dialect will load on construction as dependencies.
+  // These are dialects that this dialect may involved in canonicalization
+  // pattern or interfaces.
+  list<string> dependentDialects = [];
+
   // The C++ namespace that ops of this dialect should be placed into.
   //
   // By default, uses the name of the dialect as the only namespace. To avoid
   // placing in any namespace, use "". To specify nested namespaces, use "::"
   // as the delimiter, e.g., given "A::B", ops will be placed in
   // `namespace A { namespace B { <ops> } }`.
   //
   // Note that this works in conjunction with dialect C++ code. Depending on how

mlir/include/mlir/InitAllDialects.h

 #include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/Shape/IR/Shape.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/Dialect.h"
 
 namespace mlir {
 
+// Add all the MLIR dialects to the provided registry.
+inline void registerAllDialects(DialectRegistry &registry) {
+  // clang-format off
+  registry.insert<acc::OpenACCDialect,
+                  AffineDialect,
+                  avx512::AVX512Dialect,
+                  gpu::GPUDialect,
+                  LLVM::LLVMAVX512Dialect,
+                  LLVM::LLVMDialect,
+                  linalg::LinalgDialect,
+                  scf::SCFDialect,
+                  omp::OpenMPDialect,
+                  quant::QuantizationDialect,
+                  spirv::SPIRVDialect,
+                  StandardOpsDialect,
+                  vector::VectorDialect,
+                  NVVM::NVVMDialect,
+                  ROCDL::ROCDLDialect,
+                  SDBMDialect,
+                  shape::ShapeDialect>();
+  // clang-format on
+}
+
 // This function should be called before creating any MLIRContext if one expect
 // all the possible dialects to be made available to the context automatically.
 inline void registerAllDialects() {
-  static bool init_once = []() {
-    registerDialect<acc::OpenACCDialect>();
-    registerDialect<AffineDialect>();
-    registerDialect<avx512::AVX512Dialect>();
-    registerDialect<gpu::GPUDialect>();
-    registerDialect<LLVM::LLVMAVX512Dialect>();
-    registerDialect<LLVM::LLVMDialect>();
-    registerDialect<linalg::LinalgDialect>();
-    registerDialect<scf::SCFDialect>();
-    registerDialect<omp::OpenMPDialect>();
-    registerDialect<quant::QuantizationDialect>();
-    registerDialect<spirv::SPIRVDialect>();
-    registerDialect<StandardOpsDialect>();
-    registerDialect<vector::VectorDialect>();
-    registerDialect<NVVM::NVVMDialect>();
-    registerDialect<ROCDL::ROCDLDialect>();
-    registerDialect<SDBMDialect>();
-    registerDialect<shape::ShapeDialect>();
-    return true;
-  }();
-  (void)init_once;
+  static bool initOnce =
+      ([]() { registerAllDialects(getGlobalDialectRegistry()); }(), true);
+  (void)initOnce;
 }
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_H_

mlir/include/mlir/InitAllTranslations.h

 void registerToNVVMIRTranslation();
 void registerToROCDLIRTranslation();
 void registerAVX512ToLLVMIRTranslation();
 
 // This function should be called before creating any MLIRContext if one
 // expects all the possible translations to be made available to the context
 // automatically.
 inline void registerAllTranslations() {
-  static bool init_once = []() {
+  static bool initOnce = []() {
     registerFromLLVMIRTranslation();
     registerFromSPIRVTranslation();
     registerToLLVMIRTranslation();
     registerToSPIRVTranslation();
     registerToNVVMIRTranslation();
     registerToROCDLIRTranslation();
     registerAVX512ToLLVMIRTranslation();
     return true;
   }();
-  (void)init_once;
+  (void)initOnce;
 }
 } // namespace mlir
 
 #endif // MLIR_INITALLTRANSLATIONS_H

mlir/include/mlir/Pass/Pass.h

 //===- Pass.h - Base classes for compiler passes ----------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_PASS_PASS_H
 #define MLIR_PASS_PASS_H
 
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Function.h"
 #include "mlir/Pass/AnalysisManager.h"
 #include "mlir/Pass/PassRegistry.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/PointerIntPair.h"
 #include "llvm/ADT/Statistic.h"
 
 namespace mlir {
public:
   }
 
   /// Returns the pass info for this pass.
   const PassInfo *lookupPassInfo() const { return lookupPassInfo(getTypeID()); }
 
   /// Returns the derived pass name.
   virtual StringRef getName() const = 0;
 
+  /// Register dependent dialects for the current pass.
+  /// A pass is expected to register the dialects it will create entities for
+  /// (Operations, Types, Attributes), other than dialect that exists in the
+  /// input. For example, a pass that converts from Linalg to Affine would
+  /// register the Affine dialect but does not need to register Linalg.
+  virtual void getDependentDialects(DialectRegistry &registry) const {}
+
   /// Returns the command line argument used when registering this pass. Return
   /// an empty string if one does not exist.
   virtual StringRef getArgument() const {
     if (const PassInfo *passInfo = lookupPassInfo())
       return passInfo->getPassArgument();
     return "";
   }
 

mlir/include/mlir/Pass/PassBase.td

class PassBase<string passArg, string base> {
   string summary = "";
 
   // A human readable description of the pass.
   string description = "";
 
   // A C++ constructor call to create an instance of this pass.
   code constructor = [{}];
 
+  // A list of dialects this pass may produce entities in.
+  list<string> dependentDialects = [];
+
   // A set of options provided by this pass.
   list<Option> options = [];
 
   // A set of statistics provided by this pass.
   list<Statistic> statistics = [];
 }
 
 // This class represents an mlir::OperationPass.
 class Pass<string passArg, string operation = "">
   : PassBase<passArg, "::mlir::OperationPass<" # operation # ">">;
 
 // This class represents an mlir::FunctionPass.
 class FunctionPass<string passArg> : PassBase<passArg, "::mlir::FunctionPass">;
 
 #endif // MLIR_PASS_PASSBASE

mlir/include/mlir/Pass/PassManager.h

 //===- PassManager.h - Pass Management Interface ----------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_PASS_PASSMANAGER_H
 #define MLIR_PASS_PASSMANAGER_H
 
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/OperationSupport.h"
 #include "mlir/Support/LogicalResult.h"
 #include "llvm/ADT/Optional.h"
 #include "llvm/ADT/SmallVector.h"
 #include "llvm/ADT/iterator.h"
 #include "llvm/Support/raw_ostream.h"
 
 #include <functional>
public:
 
   /// Iterator over the passes in this pass manager.
   using pass_iterator =
       llvm::pointee_iterator<std::vector<std::unique_ptr<Pass>>::iterator>;
   pass_iterator begin();
   pass_iterator end();
   iterator_range<pass_iterator> getPasses() { return {begin(), end()}; }
 
+  using const_pass_iterator = llvm::pointee_iterator<
+      std::vector<std::unique_ptr<Pass>>::const_iterator>;
+  const_pass_iterator begin() const;
+  const_pass_iterator end() const;
+  iterator_range<const_pass_iterator> getPasses() const {
+    return {begin(), end()};
+  }
+
   /// Run the held passes over the given operation.
   LogicalResult run(Operation *op, AnalysisManager am);
 
   /// Nest a new operation pass manager for the given operation kind under this
   /// pass manager.
   OpPassManager &nest(const OperationName &nestedName);
   OpPassManager &nest(StringRef nestedName);
   template <typename OpT> OpPassManager &nest() {
public:
   /// of pipelines.
   /// Note: The quality of the string representation depends entirely on the
   /// the correctness of per-pass overrides of Pass::printAsTextualPipeline.
   void printAsTextualPipeline(raw_ostream &os);
 
   /// Merge the pass statistics of this class into 'other'.
   void mergeStatisticsInto(OpPassManager &other);
 
+  /// Register dependent dialects for the current pass manager.
+  /// This is forwarding to every pass in this PassManager, see the
+  /// documentation for the same method on the Pass class.
+  void getDependentDialects(DialectRegistry &dialects) const;
+
 private:
   OpPassManager(OperationName name, bool verifyPasses);
 
   /// A pointer to an internal implementation instance.
   std::unique_ptr<detail::OpPassManagerImpl> impl;
 
   /// Allow access to the constructor.
   friend class PassManager;

mlir/include/mlir/Support/MlirOptMain.h

 #include <memory>
 
 namespace llvm {
 class raw_ostream;
 class MemoryBuffer;
 } // end namespace llvm
 
 namespace mlir {
+class DialectRegistry;
 class PassPipelineCLParser;
 
 /// Perform the core processing behind `mlir-opt`:
 /// - outputStream is the stream where the resulting IR is printed.
 /// - buffer is the in-memory file to parser and process.
 /// - passPipeline is the specification of the pipeline that will be applied.
+/// - registry should contain all the dialects that can be parsed in the source.
 /// - splitInputFile will look for a "-----" marker in the input file, and load
 /// each chunk in an individual ModuleOp processed separately.
 /// - verifyDiagnostics enables a verification mode where comments starting with
 /// "expected-(error|note|remark|warning)" are parsed in the input and matched
 /// against emitted diagnostics.
 /// - verifyPasses enables the IR verifier in-between each pass in the pipeline.
 /// - allowUnregisteredDialects allows to parse and create operation without
 /// registering the Dialect in the MLIRContext.
+/// - preloadDialectsInContext will trigger the upfront loading of all
+///   dialects from the global registry in the MLIRContext. This option is
+///   deprecated and will be removed soon.
 LogicalResult MlirOptMain(llvm::raw_ostream &outputStream,
                           std::unique_ptr<llvm::MemoryBuffer> buffer,
                           const PassPipelineCLParser &passPipeline,
-                          bool splitInputFile, bool verifyDiagnostics,
-                          bool verifyPasses, bool allowUnregisteredDialects);
+                          DialectRegistry &registry, bool splitInputFile,
+                          bool verifyDiagnostics, bool verifyPasses,
+                          bool allowUnregisteredDialects,
+                          bool preloadDialectsInContext = true);
 
 /// Implementation for tools like `mlir-opt`.
-LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName);
+/// - toolName is used for the header displayed by `--help`.
+/// - registry should contain all the dialects that can be parsed in the source.
+/// - preloadDialectsInContext will trigger the upfront loading of all
+///   dialects from the global registry in the MLIRContext. This option is
+///   deprecated and will be removed soon.
+LogicalResult MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
+                          DialectRegistry &registry,
+                          bool preloadDialectsInContext = true);
 
 } // end namespace mlir

mlir/include/mlir/TableGen/Dialect.h

 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Dialect wrapper to simplify using TableGen Record defining a MLIR dialect.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_TABLEGEN_DIALECT_H_
 #define MLIR_TABLEGEN_DIALECT_H_
 
 #include "mlir/Support/LLVM.h"
 #include <string>
+#include <vector>
 
 namespace llvm {
 class Record;
 } // end namespace llvm
 
 namespace mlir {
 namespace tblgen {
 // Wrapper class that contains a MLIR dialect's information defined in TableGen
 // and provides helper methods for accessing them.
 class Dialect {
 public:
-  explicit Dialect(const llvm::Record *def) : def(def) {}
+  explicit Dialect(const llvm::Record *def);
 
   // Returns the name of this dialect.
   StringRef getName() const;
 
   // Returns the C++ namespaces that ops of this dialect should be placed into.
   StringRef getCppNamespace() const;
 
   // Returns this dialect's C++ class name.
   std::string getCppClassName() const;
 
   // Returns the summary description of the dialect. Returns empty string if
   // none.
   StringRef getSummary() const;
 
   // Returns the description of the dialect. Returns empty string if none.
   StringRef getDescription() const;
 
+  // Returns the list of dialect (class names) that this dialect depends on.
+  // These are dialects that will be loaded on construction of this dialect.
+  ArrayRef<StringRef> getDependentDialects() const;
+
   // Returns the dialects extra class declaration code.
   llvm::Optional<StringRef> getExtraClassDeclaration() const;
 
   // Returns if this dialect has a constant materializer or not.
   bool hasConstantMaterializer() const;
 
   /// Returns if this dialect has an operation attribute verifier.
   bool hasOperationAttrVerify() const;
public:
   // Compares two dialects by comparing the names of the dialects.
   bool operator<(const Dialect &other) const;
 
   // Returns whether the dialect is defined.
   operator bool() const { return def != nullptr; }
 
 private:
   const llvm::Record *def;
+  std::vector<StringRef> dependentDialects;
 };
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_DIALECT_H_

mlir/include/mlir/TableGen/Pass.h

public:
   StringRef getSummary() const;
 
   /// Return the description of the pass.
   StringRef getDescription() const;
 
   /// Return the C++ constructor call to create an instance of this pass.
   StringRef getConstructor() const;
 
+  /// Return the dialects this pass needs to be registered.
+  ArrayRef<StringRef> getDependentDialects() const;
+
   /// Return the options provided by this pass.
   ArrayRef<PassOption> getOptions() const;
 
   /// Return the statistics provided by this pass.
   ArrayRef<PassStatistic> getStatistics() const;
 
   const llvm::Record *getDef() const { return def; }
 
 private:
   const llvm::Record *def;
+  std::vector<StringRef> dependentDialects;
   std::vector<PassOption> options;
   std::vector<PassStatistic> statistics;
 };
 
 } // end namespace tblgen
 } // end namespace mlir
 
 #endif // MLIR_TABLEGEN_PASS_H_

mlir/include/mlir/Transforms/Passes.td

module {
         return
       }
     }
     ```
   }];
   let constructor = "mlir::createPipelineDataTransferPass()";
 }
 
 def BufferPlacement : FunctionPass<"buffer-placement"> {
   let summary = "Optimizes placement of alloc and dealloc operations";
   let description = [{
     This pass implements an algorithm to optimize the placement of alloc and
     dealloc operations. This pass also inserts missing dealloc operations
     automatically to reclaim memory.
 
 
     Input
module {
         dealloc %0 : memref<2xf32>
         return
       }
     }
     ```
 
   }];
   let constructor = "mlir::createBufferPlacementPass()";
+  // TODO: this pass likely shouldn't depend on Linalg?
+  let dependentDialects = ["linalg::LinalgDialect"];
 }
 
 def Canonicalizer : Pass<"canonicalize"> {
   let summary = "Canonicalize operations";
   let description = [{
     This pass performs various types of canonicalizations over a set of
     operations. See [Operation Canonicalization](Canonicalization.md) for more
     details.

mlir/lib/CAPI/IR/IR.cpp

 //===- IR.cpp - C Interface for Core MLIR APIs ----------------------------===//
 //
 // 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/IR.h"
 
 #include "mlir/CAPI/IR.h"
 #include "mlir/IR/Attributes.h"
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/IR/Types.h"
 #include "mlir/Parser.h"
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 

 };
 } // end namespace
 
 /* ========================================================================== */
 /* Context API.                                                               */
 /* ========================================================================== */
 
 MlirContext mlirContextCreate() {
-  auto *context = new MLIRContext;
+  auto *context = new MLIRContext(/*loadAllDialects=*/false);
   return wrap(context);
 }
 
 void mlirContextDestroy(MlirContext context) { delete unwrap(context); }
 
 /* ========================================================================== */
 /* Location API.                                                              */
 /* ========================================================================== */

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() { mlir::registerAllDialects(); }
+void mlirRegisterAllDialects(MlirContext context) {
+  registerAllDialects(unwrap(context)->getDialectRegistry());
+  // TODO: we may not want to eagerly load here.
+  unwrap(context)->getDialectRegistry().loadAll(unwrap(context));
+}

mlir/lib/Conversion/GPUToVulkan/ConvertGPULaunchFuncToVulkanLaunchFunc.cpp

 // using `spirv::serialize` function, attaches binary data and entry point name
 // as an attributes to vulkan launch call op.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
+#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Serialization.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Function.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/StandardTypes.h"

mlir/lib/Conversion/LinalgToLLVM/LinalgToLLVM.cpp

 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/VectorToLLVM/ConvertVectorToLLVM.h"
 #include "mlir/Conversion/VectorToSCF/VectorToSCF.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgTypes.h"
 #include "mlir/Dialect/Linalg/Passes.h"
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/EDSC/Intrinsics.h"
 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/AffineMap.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/MLIRContext.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Operation.h"

mlir/lib/Conversion/PassDetail.h

 //===- PassDetail.h - Conversion Pass class details -------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef CONVERSION_PASSDETAIL_H_
 #define CONVERSION_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+class AffineDialect;
+class StandardOpsDialect;
+
+// Forward declaration from Dialect.h
+template <typename ConcreteDialect>
+void registerDialect(DialectRegistry &registry);
 
 namespace gpu {
+class GPUDialect;
 class GPUModuleOp;
 } // end namespace gpu
 
+namespace LLVM {
+class LLVMDialect;
+class LLVMAVX512Dialect;
+} // end namespace LLVM
+
+namespace NVVM {
+class NVVMDialect;
+} // end namespace NVVM
+
+namespace ROCDL {
+class ROCDLDialect;
+} // end namespace ROCDL
+
+namespace scf {
+class SCFDialect;
+} // end namespace scf
+
+namespace spirv {
+class SPIRVDialect;
+} // end namespace spirv
+
+namespace vector {
+class VectorDialect;
+} // end namespace vector
+
 #define GEN_PASS_CLASSES
 #include "mlir/Conversion/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // CONVERSION_PASSDETAIL_H_

mlir/lib/Conversion/StandardToLLVM/StandardToLLVM.cpp

 
 /// Create an LLVMTypeConverter using default LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx)
     : LLVMTypeConverter(ctx, LowerToLLVMOptions::getDefaultOptions()) {}
 
 /// Create an LLVMTypeConverter using custom LowerToLLVMOptions.
 LLVMTypeConverter::LLVMTypeConverter(MLIRContext *ctx,
                                      const LowerToLLVMOptions &options)
-    : llvmDialect(ctx->getRegisteredDialect<LLVM::LLVMDialect>()),
+    : llvmDialect(ctx->getOrLoadDialect<LLVM::LLVMDialect>()),
       options(options) {
   assert(llvmDialect && "LLVM IR dialect is not registered");
   if (options.indexBitwidth == kDeriveIndexBitwidthFromDataLayout)
     this->options.indexBitwidth = options.dataLayout.getPointerSizeInBits();
 
   // Register conversions for the standard types.
   addConversion([&](ComplexType type) { return convertComplexType(type); });
   addConversion([&](FloatType type) { return convertFloatType(type); });

mlir/lib/Conversion/StandardToSPIRV/LegalizeStandardForSPIRV.cpp

 //===- LegalizeStandardForSPIRV.cpp - Legalize ops for SPIR-V lowering ----===//
 //
 // 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 transformation pass legalizes operations before the conversion to SPIR-V
 // dialect to handle ops that cannot be lowered directly.
 //
 //===----------------------------------------------------------------------===//
 
 #include "../PassDetail.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRV.h"
 #include "mlir/Conversion/StandardToSPIRV/ConvertStandardToSPIRVPass.h"
+#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/StandardTypes.h"
 
 using namespace mlir;
 
 namespace {

mlir/lib/Dialect/Affine/Transforms/PassDetail.h

 //===- PassDetail.h - Affine Pass class details -----------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_AFFINE_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_AFFINE_TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+// Forward declaration from Dialect.h
+template <typename ConcreteDialect>
+void registerDialect(DialectRegistry &registry);
+
+namespace linalg {
+class LinalgDialect;
+} // end namespace linalg
+namespace vector {
+class VectorDialect;
+} // end namespace vector
 
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/Affine/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // DIALECT_AFFINE_TRANSFORMS_PASSDETAIL_H_

mlir/lib/Dialect/Linalg/IR/LinalgOps.cpp

static void printNamedStructuredOp(OpAsmPrinter &p, NamedStructuredOpType op) {
   if (!outputTensorTypes.empty())
     p << " -> (" << outputTensorTypes << ")";
 }
 
 template <typename NamedStructuredOpType>
 static ParseResult parseNamedStructuredOp(OpAsmParser &parser,
                                           OperationState &result) {
   SmallVector<OpAsmParser::OperandType, 8> operandsInfo;
+  result.getContext()->getOrLoadDialect<StandardOpsDialect>();
 
   // Optional attributes may be added.
   if (parser.parseOperandList(operandsInfo) ||
       parser.parseOptionalAttrDict(result.attributes))
     return failure();
 
   SmallVector<Type, 8> operandTypes;
   if (parser.parseColon() || parser.parseLParen() ||

mlir/lib/Dialect/Linalg/Transforms/PassDetail.h

 //===- PassDetail.h - Linalg Pass class details -----------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_LINALG_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_LINALG_TRANSFORMS_PASSDETAIL_H_
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/IR/Dialect.h"
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+// Forward declaration from Dialect.h
+template <typename ConcreteDialect>
+void registerDialect(DialectRegistry &registry);
+
+namespace scf {
+class SCFDialect;
+} // end namespace scf
 
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/Linalg/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // DIALECT_LINALG_TRANSFORMS_PASSDETAIL_H_

mlir/lib/Dialect/SCF/Transforms/PassDetail.h

 //===- PassDetail.h - Loop Pass class details -------------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef DIALECT_LOOPOPS_TRANSFORMS_PASSDETAIL_H_
 #define DIALECT_LOOPOPS_TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+// Forward declaration from Dialect.h
+template <typename ConcreteDialect>
+void registerDialect(DialectRegistry &registry);
+
+class AffineDialect;
 
 #define GEN_PASS_CLASSES
 #include "mlir/Dialect/SCF/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // DIALECT_LOOPOPS_TRANSFORMS_PASSDETAIL_H_

mlir/lib/Dialect/SDBM/SDBMExpr.cpp

SDBMExpr visitDimExpr(AffineDimExpr expr) {
       return SDBMDimExpr::get(dialect, expr.getPosition());
     }
     SDBMExpr visitSymbolExpr(AffineSymbolExpr expr) {
       return SDBMSymbolExpr::get(dialect, expr.getPosition());
     }
 
     SDBMDialect *dialect;
   } converter;
-  converter.dialect = affine.getContext()->getRegisteredDialect<SDBMDialect>();
+  converter.dialect = affine.getContext()->getOrLoadDialect<SDBMDialect>();
 
   if (auto result = converter.visit(affine))
     return result;
   return None;
 }
 
 //===----------------------------------------------------------------------===//
 // SDBMDiffExpr

mlir/lib/ExecutionEngine/JitRunner.cpp

int mlir::JitRunnerMain(
   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;
+  MLIRContext context(/*loadAllDialects=*/false);
+  registerAllDialects(&context);
+
   auto m = parseMLIRInput(options.inputFilename, &context);
   if (!m) {
     llvm::errs() << "could not parse the input IR\n";
     return 1;
   }
 
   if (mlirTransformer)
     if (failed(mlirTransformer(m.get())))

mlir/lib/IR/Dialect.cpp

 
 DialectAsmParser::~DialectAsmParser() {}
 
 //===----------------------------------------------------------------------===//
 // Dialect Registration
 //===----------------------------------------------------------------------===//
 
 /// Registry for all dialect allocation functions.
-static llvm::ManagedStatic<llvm::MapVector<TypeID, DialectAllocatorFunction>>
-    dialectRegistry;
+static llvm::ManagedStatic<DialectRegistry> dialectRegistry;
+DialectRegistry &mlir::getGlobalDialectRegistry() { return *dialectRegistry; }
 
-void Dialect::registerDialectAllocator(
-    TypeID typeID, const DialectAllocatorFunction &function) {
-  assert(function &&
-         "Attempting to register an empty dialect initialize function");
-  dialectRegistry->insert({typeID, function});
+void mlir::registerAllDialects(MLIRContext *context) {
+  dialectRegistry->appendTo(context->getDialectRegistry());
 }
 
-/// Registers all dialects and hooks from the global registries with the
-/// specified MLIRContext.
-void mlir::registerAllDialects(MLIRContext *context) {
-  for (const auto &it : *dialectRegistry)
-    it.second(context);
+Dialect *DialectRegistry::loadByName(StringRef name, MLIRContext *context) {
+  auto it = registry.find(name.str());
+  if (it == registry.end())
+    return nullptr;
+  return it->second.second(context);
+}
+
+void DialectRegistry::insert(TypeID typeID, StringRef name,
+                             DialectAllocatorFunction ctor) {
+  auto inserted =
+      registry.insert(std::make_pair(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);
+  }
 }
 
 //===----------------------------------------------------------------------===//
 // Dialect
 //===----------------------------------------------------------------------===//
 
 Dialect::Dialect(StringRef name, MLIRContext *context, TypeID id)
     : name(name), dialectID(id), context(context) {

 //===----------------------------------------------------------------------===//
 // Dialect Interface
 //===----------------------------------------------------------------------===//
 
 DialectInterface::~DialectInterface() {}
 
 DialectInterfaceCollectionBase::DialectInterfaceCollectionBase(
     MLIRContext *ctx, TypeID interfaceKind) {
-  for (auto *dialect : ctx->getRegisteredDialects()) {
+  for (auto *dialect : ctx->getLoadedDialects()) {
     if (auto *interface = dialect->getRegisteredInterface(interfaceKind)) {
       interfaces.insert(interface);
       orderedInterfaces.push_back(interface);
     }
   }
 }
 
 DialectInterfaceCollectionBase::~DialectInterfaceCollectionBase() {}
 
 /// 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());
 }

mlir/lib/IR/MLIRContext.cpp

 #include "mlir/IR/Types.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/DenseSet.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/Twine.h"
 #include "llvm/Support/Allocator.h"
 #include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Debug.h"
 #include "llvm/Support/RWMutex.h"
 #include "llvm/Support/raw_ostream.h"
 #include <memory>
 
+#define DEBUG_TYPE "mlircontext"
+
 using namespace mlir;
 using namespace mlir::detail;
 
 using llvm::hash_combine;
 using llvm::hash_combine_range;
 
 //===----------------------------------------------------------------------===//
 // MLIRContext CommandLine Options
public:
   bool printStackTraceOnDiagnostic = false;
 
   //===--------------------------------------------------------------------===//
   // Other
   //===--------------------------------------------------------------------===//
 
   /// This is a list of dialects that are created referring to this context.
   /// The MLIRContext owns the objects.
-  std::vector<std::unique_ptr<Dialect>> dialects;
+  DenseMap<StringRef, std::unique_ptr<Dialect>> loadedDialects;
+  DialectRegistry dialectsRegistry;
 
   /// This is a mapping from operation name to AbstractOperation for registered
   /// operations.
   llvm::StringMap<AbstractOperation> registeredOperations;
 
   /// These are identifiers uniqued into this MLIRContext.
   llvm::StringSet<llvm::BumpPtrAllocator &> identifiers;
 
~MLIRContextImpl() {
     for (auto typeMapping : registeredTypes)
       typeMapping.second->~AbstractType();
     for (auto attrMapping : registeredAttributes)
       attrMapping.second->~AbstractAttribute();
   }
 };
 } // end namespace mlir
 
-MLIRContext::MLIRContext() : impl(new MLIRContextImpl()) {
+MLIRContext::MLIRContext(bool loadAllDialects) : 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);
   }
 
   // Register dialects with this context.
-  getOrCreateDialect<BuiltinDialect>();
-  registerAllDialects(this);
+  getOrLoadDialect<BuiltinDialect>();
+  if (loadAllDialects)
+    loadAllGloballyRegisteredDialects();
 
   // 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);

 
 /// Returns the diagnostic engine for this context.
 DiagnosticEngine &MLIRContext::getDiagEngine() { return getImpl().diagEngine; }
 
 //===----------------------------------------------------------------------===//
 // Dialect and Operation Registration
 //===----------------------------------------------------------------------===//
 
+DialectRegistry &MLIRContext::getDialectRegistry() {
+  return impl->dialectsRegistry;
+}
+
 /// Return information about all registered IR dialects.
-std::vector<Dialect *> MLIRContext::getRegisteredDialects() {
+std::vector<Dialect *> MLIRContext::getLoadedDialects() {
   std::vector<Dialect *> result;
-  result.reserve(impl->dialects.size());
-  for (auto &dialect : impl->dialects)
-    result.push_back(dialect.get());
+  result.reserve(impl->loadedDialects.size());
+  for (auto &dialect : impl->loadedDialects)
+    result.push_back(dialect.second.get());
+  llvm::array_pod_sort(result.begin(), result.end(),
+                       [](Dialect *const *lhs, Dialect *const *rhs) -> int {
+                         return (*lhs)->getNamespace() < (*rhs)->getNamespace();
+                       });
+  return result;
+}
+std::vector<StringRef> MLIRContext::getAvailableDialects() {
+  std::vector<StringRef> result;
+  for (auto &dialect : impl->dialectsRegistry)
+    result.push_back(dialect.first);
   return result;
 }
 
 /// Get a registered IR dialect with the given namespace. If none is found,
 /// then return nullptr.
-Dialect *MLIRContext::getRegisteredDialect(StringRef name) {
+Dialect *MLIRContext::getLoadedDialect(StringRef name) {
   // Dialects are sorted by name, so we can use binary search for lookup.
-  auto it = llvm::lower_bound(
-      impl->dialects, name,
-      [](const auto &lhs, StringRef rhs) { return lhs->getNamespace() < rhs; });
-  return (it != impl->dialects.end() && (*it)->getNamespace() == name)
-             ? (*it).get()
-             : nullptr;
+  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);
 }
 
 /// 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::getOrCreateDialect(StringRef dialectNamespace, TypeID dialectID,
-                                function_ref<std::unique_ptr<Dialect>()> ctor) {
+MLIRContext::getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
   auto &impl = getImpl();
   // Get the correct insertion position sorted by namespace.
-  auto insertPt =
-      llvm::lower_bound(impl.dialects, nullptr,
-                        [&](const std::unique_ptr<Dialect> &lhs,
-                            const std::unique_ptr<Dialect> &rhs) {
-                          if (!lhs)
-                            return dialectNamespace < rhs->getNamespace();
-                          return lhs->getNamespace() < dialectNamespace;
-                        });
+  std::unique_ptr<Dialect> &dialect = impl.loadedDialects[dialectNamespace];
+
+  if (!dialect) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Load new dialect in Context" << dialectNamespace);
+    dialect = ctor();
+    assert(dialect && "dialect ctor failed");
+    return dialect.get();
+  }
 
   // Abort if dialect with namespace has already been registered.
-  if (insertPt != impl.dialects.end() &&
-      (*insertPt)->getNamespace() == dialectNamespace) {
-    if ((*insertPt)->getTypeID() == dialectID)
-      return insertPt->get();
+  if (dialect->getTypeID() != dialectID)
     llvm::report_fatal_error("a dialect with namespace '" + dialectNamespace +
                              "' has already been registered");
-  }
-  auto it = impl.dialects.insert(insertPt, ctor());
-  return &**it;
+
+  return dialect.get();
+
+void MLIRContext::loadAllGloballyRegisteredDialects() {
+  getGlobalDialectRegistry().loadAll(this);
 }
 
 bool MLIRContext::allowsUnregisteredDialects() {
   return impl->allowUnregisteredDialects;
 }
 
 void MLIRContext::allowUnregisteredDialects(bool allowing) {
   impl->allowUnregisteredDialects = allowing;

mlir/lib/IR/Operation.cpp

 /// Return the dialect this operation is associated with, or nullptr if the
 /// associated dialect is not registered.
 Dialect *Operation::getDialect() {
   if (auto *abstractOp = getAbstractOperation())
     return &abstractOp->dialect;
 
   // If this operation hasn't been registered or doesn't have abstract
   // operation, try looking up the dialect name in the context.
-  return getContext()->getRegisteredDialect(getName().getDialect());
+  return getContext()->getLoadedDialect(getName().getDialect());
 }
 
 Region *Operation::getParentRegion() {
   return block ? block->getParent() : nullptr;
 }
 
 Operation *Operation::getParentOp() {
   return block ? block->getParentOp() : nullptr;

mlir/lib/IR/Verifier.cpp

public:
 
   /// Verify the given operation.
   LogicalResult verify(Operation &op);
 
   /// Returns the registered dialect for a dialect-specific attribute.
   Dialect *getDialectForAttribute(const NamedAttribute &attr) {
     assert(attr.first.strref().contains('.') && "expected dialect attribute");
     auto dialectNamePair = attr.first.strref().split('.');
-    return ctx->getRegisteredDialect(dialectNamePair.first);
+    return ctx->getLoadedDialect(dialectNamePair.first);
   }
 
 private:
   /// Verify the given potentially nested region or block.
   LogicalResult verifyRegion(Region &region);
   LogicalResult verifyBlock(Block &block);
   LogicalResult verifyOperation(Operation &op);
 
LogicalResult OperationVerifier::verifyOperation(Operation &op) {
 
   // Otherwise, verify that the parent dialect allows un-registered operations.
   auto dialectPrefix = op.getName().getDialect();
 
   // Check for an existing answer for the operation dialect.
   auto it = dialectAllowsUnknownOps.find(dialectPrefix);
   if (it == dialectAllowsUnknownOps.end()) {
     // If the operation dialect is registered, query it directly.
-    if (auto *dialect = ctx->getRegisteredDialect(dialectPrefix))
+    if (auto *dialect = ctx->getLoadedDialect(dialectPrefix))
       it = dialectAllowsUnknownOps
                .try_emplace(dialectPrefix, dialect->allowsUnknownOperations())
                .first;
     // Otherwise, unregistered dialects (when allowed by the context)
     // conservatively allow unknown operations.
     else {
       if (!op.getContext()->allowsUnregisteredDialects() && !op.getDialect())
         return op.emitOpError()

mlir/lib/Parser/AttributeParser.cpp

 //===- AttributeParser.cpp - MLIR Attribute Parser 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the parser for the MLIR Types.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Parser.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/IntegerSet.h"
 #include "mlir/IR/StandardTypes.h"
 #include "llvm/ADT/StringExtras.h"
 
 using namespace mlir;
 using namespace mlir::detail;
 
 /// Parse an arbitrary attribute.
else if (getToken().isAny(Token::bare_identifier, Token::inttype) ||
              getToken().isKeyword())
       nameId = builder.getIdentifier(getTokenSpelling());
     else
       return emitError("expected attribute name");
     if (!seenKeys.insert(*nameId).second)
       return emitError("duplicate key in dictionary attribute");
     consumeToken();
 
+    // Lazy load a dialect in the context if there is a possible namespace.
+    auto splitName = nameId->strref().split('.');
+    if (!splitName.second.empty())
+      getContext()->getOrLoadDialect(splitName.first);
+
     // Try to parse the '=' for the attribute value.
     if (!consumeIf(Token::equal)) {
       // If there is no '=', we treat this as a unit attribute.
       attributes.push_back({*nameId, builder.getUnitAttr()});
       return success();
     }
 
     auto attr = parseAttribute();
Attribute Parser::parseOpaqueElementsAttr(Type attrType) {
   consumeToken(Token::kw_opaque);
   if (parseToken(Token::less, "expected '<' after 'opaque'"))
     return nullptr;
 
   if (getToken().isNot(Token::string))
     return (emitError("expected dialect namespace"), nullptr);
 
   auto name = getToken().getStringValue();
-  auto *dialect = builder.getContext()->getRegisteredDialect(name);
+  // Lazy load a dialect in the context if there is a possible namespace.
+  Dialect *dialect = builder.getContext()->getOrLoadDialect(name);
+
   // TODO: Allow for having an unknown dialect on an opaque
   // attribute. Otherwise, it can't be roundtripped without having the dialect
   // registered.
   if (!dialect)
     return (emitError("no registered dialect with namespace '" + name + "'"),
             nullptr);
   consumeToken(Token::string);
 

mlir/lib/Parser/DialectSymbolParser.cpp

Attribute attr = parseExtendedSymbol<Attribute>(
       [&](StringRef dialectName, StringRef symbolData,
           llvm::SMLoc loc) -> Attribute {
         // Parse an optional trailing colon type.
         Type attrType = type;
         if (consumeIf(Token::colon) && !(attrType = parseType()))
           return Attribute();
 
         // If we found a registered dialect, then ask it to parse the attribute.
-        if (auto *dialect = state.context->getRegisteredDialect(dialectName)) {
+        if (Dialect *dialect =
+                builder.getContext()->getOrLoadDialect(dialectName)) {
           return parseSymbol<Attribute>(
               symbolData, state.context, state.symbols, [&](Parser &parser) {
                 CustomDialectAsmParser customParser(symbolData, parser);
                 return dialect->parseAttribute(customParser, attrType);
               });
         }
 
         // Otherwise, form a new opaque attribute.

 ///   type-alias    ::= `!` alias-name
 ///
 Type Parser::parseExtendedType() {
   return parseExtendedSymbol<Type>(
       *this, Token::exclamation_identifier, state.symbols.typeAliasDefinitions,
       [&](StringRef dialectName, StringRef symbolData,
           llvm::SMLoc loc) -> Type {
         // If we found a registered dialect, then ask it to parse the type.
-        if (auto *dialect = state.context->getRegisteredDialect(dialectName)) {
+        auto *dialect = state.context->getOrLoadDialect(dialectName);
+
+        if (dialect) {
           return parseSymbol<Type>(
               symbolData, state.context, state.symbols, [&](Parser &parser) {
                 CustomDialectAsmParser customParser(symbolData, parser);
                 return dialect->parseType(customParser);
               });
         }
 
         // Otherwise, form a new opaque type.

mlir/lib/Parser/Parser.cpp

 //===- Parser.cpp - MLIR Parser 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
 //
 //===----------------------------------------------------------------------===//
 //
 // This file implements the parser for the MLIR textual form.
 //
 //===----------------------------------------------------------------------===//
 
 #include "Parser.h"
 #include "mlir/IR/AffineMap.h"
+#include "mlir/IR/Dialect.h"
 #include "mlir/IR/Module.h"
 #include "mlir/IR/Verifier.h"
 #include "mlir/Parser.h"
 #include "llvm/ADT/DenseMap.h"
 #include "llvm/ADT/StringSet.h"
 #include "llvm/ADT/bit.h"
 #include "llvm/Support/PrettyStackTrace.h"
 #include "llvm/Support/SourceMgr.h"
struct CleanupOpStateRegions {
   OperationState &state;
 };
 } // namespace
 
 Operation *OperationParser::parseGenericOperation() {
   // Get location information for the operation.
   auto srcLocation = getEncodedSourceLocation(getToken().getLoc());
 
-  auto name = getToken().getStringValue();
+  std::string name = getToken().getStringValue();
   if (name.empty())
     return (emitError("empty operation name is invalid"), nullptr);
   if (name.find('\0') != StringRef::npos)
     return (emitError("null character not allowed in operation name"), nullptr);
 
   consumeToken(Token::string);
 
   OperationState result(srcLocation, name);
 
+  // Lazy load dialects in the context as needed.
+  if (!result.name.getAbstractOperation()) {
+    StringRef dialectName = StringRef(name).split('.').first;
+    if (!getContext()->getLoadedDialect(dialectName) &&
+        getContext()->getOrLoadDialect(dialectName)) {
+      result.name = OperationName(name, getContext());
+    }
+  }
+
   // Parse the operand list.
   SmallVector<SSAUseInfo, 8> operandInfos;
   if (parseToken(Token::l_paren, "expected '(' to start operand list") ||
       parseOptionalSSAUseList(operandInfos) ||
       parseToken(Token::r_paren, "expected ')' to end operand list")) {
     return nullptr;
   }
 
private:
 
   /// A flag that indicates if any errors were emitted during parsing.
   bool emittedError = false;
 };
 } // end anonymous namespace.
 
 Operation *
 OperationParser::parseCustomOperation(ArrayRef<ResultRecord> resultIDs) {
-  auto opLoc = getToken().getLoc();
-  auto opName = getTokenSpelling();
+  llvm::SMLoc opLoc = getToken().getLoc();
+  StringRef opName = getTokenSpelling();
 
   auto *opDefinition = AbstractOperation::lookup(opName, getContext());
-  if (!opDefinition && !opName.contains('.')) {
+  if (!opDefinition) {
+    if (opName.contains('.')) {
+      // This op has a dialect, we try to check if we can register it in the
+      // context on the fly.
+      StringRef dialectName = opName.split('.').first;
+      if (!getContext()->getLoadedDialect(dialectName) &&
+          getContext()->getOrLoadDialect(dialectName)) {
+        opDefinition = AbstractOperation::lookup(opName, getContext());
+      }
+    } else {
       // If the operation name has no namespace prefix we treat it as a standard
       // operation and prefix it with "std".
       // TODO: Would it be better to just build a mapping of the registered
       // operations in the standard dialect?
-    opDefinition =
-        AbstractOperation::lookup(Twine("std." + opName).str(), getContext());
+      if (getContext()->getOrLoadDialect("std"))
+        opDefinition = AbstractOperation::lookup(Twine("std." + opName).str(),
+                                                 getContext());
+    }
   }
 
   if (!opDefinition) {
     emitError(opLoc) << "custom op '" << opName << "' is unknown";
     return nullptr;
   }
 
   consumeToken();

mlir/lib/Pass/Pass.cpp

 
 OpPassManager::~OpPassManager() {}
 
 OpPassManager::pass_iterator OpPassManager::begin() {
   return impl->passes.begin();
 }
 OpPassManager::pass_iterator OpPassManager::end() { return impl->passes.end(); }
 
+OpPassManager::const_pass_iterator OpPassManager::begin() const {
+  return impl->passes.begin();
+}
+OpPassManager::const_pass_iterator OpPassManager::end() const {
+  return impl->passes.end();
+}
+
 /// Run all of the passes in this manager over the current operation.
 LogicalResult OpPassManager::run(Operation *op, AnalysisManager am) {
   // Run each of the held passes.
   for (auto &pass : impl->passes)
     if (failed(pass->run(op, am)))
       return failure();
   return success();
 }

 }
 
 /// Prints out the passes of the pass manager as the textual representation
 /// of pipelines.
 void OpPassManager::printAsTextualPipeline(raw_ostream &os) {
   ::printAsTextualPipeline(impl->passes, os);
 }
 
+static void registerDialectsForPipeline(const OpPassManager &pm,
+                                        DialectRegistry &dialects) {
+  for (const Pass &pass : pm.getPasses())
+    pass.getDependentDialects(dialects);
+}
+
+void OpPassManager::getDependentDialects(DialectRegistry &dialects) const {
+  registerDialectsForPipeline(*this, dialects);
+}
+
 //===----------------------------------------------------------------------===//
 // OpToOpPassAdaptor
 //===----------------------------------------------------------------------===//
 
 /// Utility to run the given operation and analysis manager on a provided op
 /// pass manager.
 static LogicalResult runPipeline(OpPassManager &pm, Operation *op,
                                  AnalysisManager am) {
auto it = llvm::find_if(
       mgrs, [&](OpPassManager &mgr) { return mgr.getOpName() == name; });
   return it == mgrs.end() ? nullptr : &*it;
 }
 
 OpToOpPassAdaptor::OpToOpPassAdaptor(OpPassManager &&mgr) {
   mgrs.emplace_back(std::move(mgr));
 }
 
+void OpToOpPassAdaptor::getDependentDialects(DialectRegistry &dialects) const {
+  for (auto &pm : mgrs)
+    pm.getDependentDialects(dialects);
+}
+
 /// Merge the current pass adaptor into given 'rhs'.
 void OpToOpPassAdaptor::mergeInto(OpToOpPassAdaptor &rhs) {
   for (auto &pm : mgrs) {
     // If an existing pass manager exists, then merge the given pass manager
     // into it.
     if (auto *existingPM = findPassManagerFor(rhs.mgrs, pm.getOpName())) {
       pm.getImpl().mergeInto(existingPM->getImpl());
     } else {

 PassManager::~PassManager() {}
 
 /// Run the passes within this manager on the provided module.
 LogicalResult PassManager::run(ModuleOp module) {
   // 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(module.getContext());
+
   // Construct an analysis manager for the pipeline.
   ModuleAnalysisManager am(module, instrumentor.get());
 
   // If reproducer generation is enabled, run the pass manager with crash
   // handling enabled.
   LogicalResult result = crashReproducerFileName
                              ? runWithCrashRecovery(module, am)
                              : OpPassManager::run(module, am);

mlir/lib/Pass/PassDetail.h

public:
   void runOnOperation() override;
 
   /// Merge the current pass adaptor into given 'rhs'.
   void mergeInto(OpToOpPassAdaptor &rhs);
 
   /// Returns the pass managers held by this adaptor.
   MutableArrayRef<OpPassManager> getPassManagers() { return mgrs; }
 
+  /// Populate the set of dependent dialects for the passes in the current
+  /// adaptor.
+  void getDependentDialects(DialectRegistry &dialects) const override;
+
   /// Return the async pass managers held by this parallel adaptor.
   MutableArrayRef<SmallVector<OpPassManager, 1>> getParallelPassManagers() {
     return asyncExecutors;
   }
 
   /// Returns the adaptor pass name.
   std::string getAdaptorName();
 

mlir/lib/Support/MlirOptMain.cpp

 }
 
 /// Parses the memory buffer.  If successfully, run a series of passes against
 /// it and print the result.
 static LogicalResult processBuffer(raw_ostream &os,
                                    std::unique_ptr<MemoryBuffer> ownedBuffer,
                                    bool verifyDiagnostics, bool verifyPasses,
                                    bool allowUnregisteredDialects,
-                                   const PassPipelineCLParser &passPipeline) {
+                                   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;
+  MLIRContext context(/*loadAllDialects=*/preloadDialectsInContext);
+  registry.appendTo(context.getDialectRegistry());
+  if (preloadDialectsInContext)
+    registry.loadAll(&context);
   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,
static LogicalResult processBuffer(raw_ostream &os,
   // Verify the diagnostic handler to make sure that each of the diagnostics
   // matched.
   return sourceMgrHandler.verify();
 }
 
 LogicalResult mlir::MlirOptMain(raw_ostream &outputStream,
                                 std::unique_ptr<MemoryBuffer> buffer,
                                 const PassPipelineCLParser &passPipeline,
-                                bool splitInputFile, bool verifyDiagnostics,
-                                bool verifyPasses,
-                                bool allowUnregisteredDialects) {
+                                DialectRegistry &registry, bool splitInputFile,
+                                bool verifyDiagnostics, bool verifyPasses,
+                                bool allowUnregisteredDialects,
+                                bool preloadDialectsInContext) {
   // The split-input-file mode is a very specific mode that slices the file
   // up into small pieces and checks each independently.
   if (splitInputFile)
     return splitAndProcessBuffer(
         std::move(buffer),
         [&](std::unique_ptr<MemoryBuffer> chunkBuffer, raw_ostream &os) {
           return processBuffer(os, std::move(chunkBuffer), verifyDiagnostics,
                                verifyPasses, allowUnregisteredDialects,
-                               passPipeline);
+                               preloadDialectsInContext, passPipeline,
+                               registry);
         },
         outputStream);
 
   return processBuffer(outputStream, std::move(buffer), verifyDiagnostics,
-                       verifyPasses, allowUnregisteredDialects, passPipeline);
+                       verifyPasses, allowUnregisteredDialects,
+                       preloadDialectsInContext, passPipeline, registry);
 }
 
-LogicalResult mlir::MlirOptMain(int argc, char **argv, StringRef toolName) {
+LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
+                                DialectRegistry &registry,
+                                bool preloadDialectsInContext) {
   static cl::opt<std::string> inputFilename(
       cl::Positional, cl::desc("<input file>"), cl::init("-"));
 
   static cl::opt<std::string> outputFilename("o", cl::desc("Output filename"),
                                              cl::value_desc("filename"),
                                              cl::init("-"));
 
   static cl::opt<bool> splitInputFile(
LogicalResult mlir::MlirOptMain(int argc, char **argv, llvm::StringRef toolName,
   registerPassManagerCLOptions();
   PassPipelineCLParser passPipeline("", "Compiler passes to run");
 
   // Build the list of dialects as a header for the --help message.
   std::string helpHeader = (toolName + "\nAvailable Dialects: ").str();
   {
     llvm::raw_string_ostream os(helpHeader);
     MLIRContext context;
-    interleaveComma(context.getRegisteredDialects(), os, [&](Dialect *dialect) {
-      StringRef name = dialect->getNamespace();
-      // filter the builtin dialect.
-      if (name.empty())
-        os << "<builtin>";
-      else
-        os << name;
+    interleaveComma(registry, os, [&](auto &registryEntry) {
+      StringRef name = registryEntry.first;
     });
   }
   // Parse pass names in main to ensure static initialization completed.
   cl::ParseCommandLineOptions(argc, argv, helpHeader);
 
   if (showDialects) {
-    llvm::outs() << "Registered Dialects:\n";
-    MLIRContext context;
+    llvm::outs() << "Available Dialects:\n";
     interleave(
-        context.getRegisteredDialects(), llvm::outs(),
-        [](Dialect *dialect) { llvm::outs() << dialect->getNamespace(); },
-        "\n");
+        registry, llvm::outs(),
+        [](auto &registryEntry) { llvm::outs() << registryEntry.first; }, "\n");
     return success();
   }
 
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
   auto output = openOutputFile(outputFilename, &errorMessage);
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return failure();
   }
 
-  if (failed(MlirOptMain(output->os(), std::move(file), passPipeline,
+  if (failed(MlirOptMain(output->os(), std::move(file), passPipeline, registry,
                          splitInputFile, verifyDiagnostics, verifyPasses,
-                         allowUnregisteredDialects)))
+                         allowUnregisteredDialects, preloadDialectsInContext)))
     return failure();
 
   // Keep the output file if the invocation of MlirOptMain was successful.
   output->keep();
   return success();
 }

mlir/lib/TableGen/Dialect.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/TableGen/Dialect.h"
 #include "llvm/TableGen/Record.h"
 
 using namespace mlir;
 using namespace mlir::tblgen;
+Dialect::Dialect(const llvm::Record *def) : def(def) {
+  for (StringRef dialect : def->getValueAsListOfStrings("dependentDialects"))
+    dependentDialects.push_back(dialect);
+}
 
 StringRef Dialect::getName() const { return def->getValueAsString("name"); }
 
 StringRef Dialect::getCppNamespace() const {
   return def->getValueAsString("cppNamespace");
 }
 
 std::string Dialect::getCppClassName() const {

 StringRef Dialect::getSummary() const {
   return getAsStringOrEmpty(*def, "summary");
 }
 
 StringRef Dialect::getDescription() const {
   return getAsStringOrEmpty(*def, "description");
 }
 
+ArrayRef<StringRef> Dialect::getDependentDialects() const {
+  return dependentDialects;
+}
+
 llvm::Optional<StringRef> Dialect::getExtraClassDeclaration() const {
   auto value = def->getValueAsString("extraClassDeclaration");
   return value.empty() ? llvm::Optional<StringRef>() : value;
 }
 
 bool Dialect::hasConstantMaterializer() const {
   return def->getValueAsBit("hasConstantMaterializer");
 }

mlir/lib/TableGen/Pass.cpp

 // Pass
 //===----------------------------------------------------------------------===//
 
 Pass::Pass(const llvm::Record *def) : def(def) {
   for (auto *init : def->getValueAsListOfDefs("options"))
     options.push_back(PassOption(init));
   for (auto *init : def->getValueAsListOfDefs("statistics"))
     statistics.push_back(PassStatistic(init));
+  for (StringRef dialect : def->getValueAsListOfStrings("dependentDialects"))
+    dependentDialects.push_back(dialect);
 }
 
 StringRef Pass::getArgument() const {
   return def->getValueAsString("argument");
 }
 
 StringRef Pass::getBaseClass() const {
   return def->getValueAsString("baseClass");
 }
 
 StringRef Pass::getSummary() const { return def->getValueAsString("summary"); }
 
 StringRef Pass::getDescription() const {
   return def->getValueAsString("description");
 }
 
 StringRef Pass::getConstructor() const {
   return def->getValueAsString("constructor");
 }
+ArrayRef<StringRef> Pass::getDependentDialects() const {
+  return dependentDialects;
+}
 
 ArrayRef<PassOption> Pass::getOptions() const { return options; }
 
 ArrayRef<PassStatistic> Pass::getStatistics() const { return statistics; }

mlir/lib/Target/LLVMIR/ConvertFromLLVMIR.cpp

if (failed(processInstruction(&inst)))
       return failure();
   }
   return success();
 }
 
 OwningModuleRef
 mlir::translateLLVMIRToModule(std::unique_ptr<llvm::Module> llvmModule,
                               MLIRContext *context) {
+  context->loadDialect<LLVMDialect>();
   OwningModuleRef module(ModuleOp::create(
       FileLineColLoc::get("", /*line=*/0, /*column=*/0, context)));
 
   Importer deserializer(context, module.get());
   for (llvm::GlobalVariable &gv : llvmModule->globals()) {
     if (!deserializer.processGlobal(&gv))
       return {};
   }

mlir/lib/Target/LLVMIR/ModuleTranslation.cpp

static llvm::AtomicOrdering getLLVMAtomicOrdering(AtomicOrdering ordering) {
   llvm_unreachable("incorrect atomic ordering");
 }
 
 ModuleTranslation::ModuleTranslation(Operation *module,
                                      std::unique_ptr<llvm::Module> llvmModule)
     : mlirModule(module), llvmModule(std::move(llvmModule)),
       debugTranslation(
           std::make_unique<DebugTranslation>(module, *this->llvmModule)),
-      ompDialect(
-          module->getContext()->getRegisteredDialect<omp::OpenMPDialect>()),
+      ompDialect(module->getContext()->getOrLoadDialect<omp::OpenMPDialect>()),
       typeTranslator(this->llvmModule->getContext()) {
   assert(satisfiesLLVMModule(mlirModule) &&
          "mlirModule should honor LLVM's module semantics.");
 }
 ModuleTranslation::~ModuleTranslation() {
   if (ompBuilder)
     ompBuilder->finalize();
 }
for (Value v : values) {
     assert(valueMapping.count(v) && "referencing undefined value");
     remapped.push_back(valueMapping.lookup(v));
   }
   return remapped;
 }
 
 std::unique_ptr<llvm::Module> ModuleTranslation::prepareLLVMModule(
     Operation *m, llvm::LLVMContext &llvmContext, StringRef name) {
+  m->getContext()->getOrLoadDialect<LLVM::LLVMDialect>();
   auto llvmModule = std::make_unique<llvm::Module>(name, llvmContext);
-
   if (auto dataLayoutAttr =
           m->getAttr(LLVM::LLVMDialect::getDataLayoutAttrName()))
     llvmModule->setDataLayout(dataLayoutAttr.cast<StringAttr>().getValue());
 
   // Inject declarations for `malloc` and `free` functions that can be used in
   // memref allocation/deallocation coming from standard ops lowering.
   llvm::IRBuilder<> builder(llvmContext);
   llvmModule->getOrInsertFunction("malloc", builder.getInt8PtrTy(),
                                   builder.getInt64Ty());
   llvmModule->getOrInsertFunction("free", builder.getVoidTy(),
                                   builder.getInt8PtrTy());
 
   return llvmModule;
 }

mlir/lib/Transforms/PassDetail.h

 //===- PassDetail.h - Transforms Pass class details -------------*- 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
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef TRANSFORMS_PASSDETAIL_H_
 #define TRANSFORMS_PASSDETAIL_H_
 
 #include "mlir/Pass/Pass.h"
 
 namespace mlir {
+// Forward declaration from Dialect.h
+template <typename ConcreteDialect>
+void registerDialect(DialectRegistry &registry);
+
+namespace linalg {
+class LinalgDialect;
+} // end namespace linalg
 
 #define GEN_PASS_CLASSES
 #include "mlir/Transforms/Passes.h.inc"
 
 } // end namespace mlir
 
 #endif // TRANSFORMS_PASSDETAIL_H_

mlir/test/CAPI/ir.c

if (!mlirTypeIsATuple(tuple) || mlirTupleTypeGetNumTypes(tuple) != 2 ||
     return 20;
   mlirTypeDump(tuple);
   fprintf(stderr, "\n");
 
   return 0;
 }
 
 int main() {
-  mlirRegisterAllDialects();
   MlirContext ctx = mlirContextCreate();
+  mlirRegisterAllDialects(ctx);
   MlirLocation location = mlirLocationUnknownGet(ctx);
 
   MlirModule moduleOp = makeAdd(ctx, location);
   MlirOperation module = mlirModuleGetOperation(moduleOp);
   mlirOperationDump(module);
   // clang-format off
   // CHECK: module {
   // CHECK:   func @add(%[[ARG0:.*]]: memref<?xf32>, %[[ARG1:.*]]: memref<?xf32>) {

mlir/test/EDSC/builder-api-test.cpp

 
 #include "llvm/Support/raw_ostream.h"
 
 using namespace mlir;
 using namespace mlir::edsc;
 using namespace mlir::edsc::intrinsics;
 
 static MLIRContext &globalContext() {
-  static bool init_once = []() {
-    registerDialect<AffineDialect>();
-    registerDialect<linalg::LinalgDialect>();
-    registerDialect<scf::SCFDialect>();
-    registerDialect<StandardOpsDialect>();
-    registerDialect<vector::VectorDialect>();
+  static thread_local MLIRContext context(/*loadAllDialects=*/false);
+  static thread_local bool initOnce = [&]() {
+    // clang-format off
+    context.loadDialect<AffineDialect,
+                        scf::SCFDialect,
+                        linalg::LinalgDialect,
+                        StandardOpsDialect,
+                        vector::VectorDialect>();
+    // clang-format on
     return true;
   }();
-  (void)init_once;
-  static thread_local MLIRContext context;
+  (void)initOnce;
   context.allowUnregisteredDialects();
   return context;
 }
 
 static FuncOp makeFunction(StringRef name, ArrayRef<Type> results = {},
                            ArrayRef<Type> args = {}) {
   auto &ctx = globalContext();
   auto function = FuncOp::create(UnknownLoc::get(&ctx), name,

mlir/test/SDBM/sdbm-api-test.cpp

 #include "mlir/IR/MLIRContext.h"
 
 #include "llvm/Support/raw_ostream.h"
 
 #include "APITest.h"
 
 using namespace mlir;
 
-// Load the SDBM dialect
-static DialectRegistration<SDBMDialect> SDBMRegistration;
 
 static MLIRContext *ctx() {
-  static thread_local MLIRContext context;
+  static thread_local MLIRContext context(/*loadAllDialects=*/false);
+  static thread_local bool once =
+      (context.getOrLoadDialect<SDBMDialect>(), true);
+  (void)once;
   return &context;
 }
 
 static SDBMDialect *dialect() {
   static thread_local SDBMDialect *d = nullptr;
   if (!d) {
-    d = ctx()->getRegisteredDialect<SDBMDialect>();
+    d = ctx()->getOrLoadDialect<SDBMDialect>();
   }
   return d;
 }
 
 static SDBMExpr dim(unsigned pos) { return SDBMDimExpr::get(dialect(), pos); }
 
 static SDBMExpr symb(unsigned pos) {
   return SDBMSymbolExpr::get(dialect(), pos);

mlir/test/lib/Dialect/Affine/TestVectorizationUtils.cpp

 //===- VectorizerTestPass.cpp - VectorizerTestPass Pass Impl --------------===//
 //
 // 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 implements a simple testing pass for vectorization functionality.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Analysis/AffineAnalysis.h"
 #include "mlir/Analysis/NestedMatcher.h"
 #include "mlir/Analysis/SliceAnalysis.h"
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorUtils.h"
 #include "mlir/IR/Builders.h"
 #include "mlir/IR/Diagnostics.h"
 #include "mlir/IR/StandardTypes.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 
 #include "llvm/ADT/STLExtras.h"
llvm::cl::desc(
         "operation."),
     llvm::cl::cat(clOptionsCategory));
 
 namespace {
 struct VectorizerTestPass
     : public PassWrapper<VectorizerTestPass, FunctionPass> {
   static constexpr auto kTestAffineMapOpName = "test_affine_map";
   static constexpr auto kTestAffineMapAttrName = "affine_map";
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<vector::VectorDialect>();
+  }
 
   void runOnFunction() override;
   void testVectorShapeRatio(llvm::raw_ostream &outs);
   void testForwardSlicing(llvm::raw_ostream &outs);
   void testBackwardSlicing(llvm::raw_ostream &outs);
   void testSlicing(llvm::raw_ostream &outs);
   void testComposeMaps(llvm::raw_ostream &outs);
   void testNormalizeMaps();

mlir/test/lib/Dialect/SPIRV/TestAvailability.cpp

struct PrintOpAvailability
   void runOnFunction() override;
 };
 } // end anonymous namespace
 
 void PrintOpAvailability::runOnFunction() {
   auto f = getFunction();
   llvm::outs() << f.getName() << "\n";
 
-  Dialect *spvDialect = getContext().getRegisteredDialect("spv");
+  Dialect *spvDialect = getContext().getLoadedDialect("spv");
 
   f.getOperation()->walk([&](Operation *op) {
     if (op->getDialect() != spvDialect)
       return WalkResult::advance();
 
     auto opName = op->getName();
     auto &os = llvm::outs();
 

mlir/test/lib/Dialect/Test/TestDialect.h

 namespace mlir {
 
 #include "TestOpStructs.h.inc"
 #include "TestOpsDialect.h.inc"
 
 #define GET_OP_CLASSES
 #include "TestOps.h.inc"
 
+void registerTestDialect(DialectRegistry &registry);
+
 } // end namespace mlir
 
 #endif // MLIR_TESTDIALECT_H

mlir/test/lib/Dialect/Test/TestDialect.cpp

 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/Transforms/FoldUtils.h"
 #include "mlir/Transforms/InliningUtils.h"
 #include "llvm/ADT/SetVector.h"
 #include "llvm/ADT/StringSwitch.h"
 
 using namespace mlir;
 
+void mlir::registerTestDialect(DialectRegistry &registry) {
+  registry.insert<TestDialect>();
+}
+
 //===----------------------------------------------------------------------===//
 // TestDialect Interfaces
 //===----------------------------------------------------------------------===//
 
 namespace {
 
 // Test support for interacting with the AsmPrinter.
 struct TestOpAsmInterface : public OpAsmDialectInterface {

mlir/test/lib/Dialect/Test/TestPatterns.cpp

matchAndRewrite(TestTypeProducerOp op, ArrayRef<Value> operands,
 
     rewriter.replaceOpWithNewOp<TestTypeProducerOp>(op, resultType);
     return success();
   }
 };
 
 struct TestTypeConversionDriver
     : public PassWrapper<TestTypeConversionDriver, OperationPass<ModuleOp>> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<TestDialect>();
+  }
+
   void runOnOperation() override {
     // Initialize the type converter.
     TypeConverter converter;
 
     /// Add the legal set of type conversions.
     converter.addConversion([](Type type) -> Type {
       // Treat F64 as legal.
       if (type.isF64())

mlir/test/lib/Transforms/TestAllReduceLowering.cpp

 //===- TestAllReduceLowering.cpp - Test gpu.all_reduce lowering -----------===//
 //
 // 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 contains test passes for lowering the gpu.all_reduce op.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 struct TestAllReduceLoweringPass
     : public PassWrapper<TestAllReduceLoweringPass, OperationPass<ModuleOp>> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<StandardOpsDialect>();
+  }
   void runOnOperation() override {
     OwningRewritePatternList patterns;
     populateGpuRewritePatterns(&getContext(), patterns);
     applyPatternsAndFoldGreedily(getOperation(), patterns);
   }
 };
 } // namespace
 
 namespace mlir {
 void registerTestAllReduceLoweringPass() {
   PassRegistration<TestAllReduceLoweringPass> pass(
       "test-all-reduce-lowering",
       "Lowers gpu.all-reduce ops within the GPU dialect.");
 }
 } // namespace mlir

mlir/test/lib/Transforms/TestBufferPlacement.cpp

void populateTensorLinalgToBufferLinalgConversionPattern(
       MLIRContext *context, BufferAssignmentPlacer *placer,
       TypeConverter *converter, OwningRewritePatternList *patterns) {
     populateWithBufferAssignmentOpConversionPatterns<
         mlir::ReturnOp, mlir::ReturnOp, linalg::CopyOp,
         allowMemrefFunctionResults>(context, placer, converter, patterns);
     patterns->insert<GenericOpConverter>(context, placer, converter);
   }
 
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<linalg::LinalgDialect>();
+  }
+
   void runOnOperation() override {
     MLIRContext &context = this->getContext();
     ConversionTarget target(context);
     BufferAssignmentTypeConverter converter;
 
     // Mark all Standard operations legal.
     target.addLegalDialect<StandardOpsDialect>();
 

mlir/test/lib/Transforms/TestGpuMemoryPromotion.cpp

 //===- TestGPUMemoryPromotionPass.cpp - Test pass for GPU promotion -------===//
 //
 // 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 implements the pass testing the utilities for moving data across
 // different levels of the GPU memory hierarchy.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/MemoryPromotion.h"
+#include "mlir/Dialect/SCF/SCF.h"
+#include "mlir/Dialect/SPIRV/SPIRVDialect.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/IR/Attributes.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 
 namespace {
 /// Simple pass for testing the promotion to workgroup memory in GPU functions.
 /// Promotes all arguments with "gpu.test_promote_workgroup" attribute. This
 /// does not check whether the promotion is legal (e.g., amount of memory used)
 /// or beneficial (e.g., makes previously uncoalesced loads coalesced).
 class TestGpuMemoryPromotionPass
     : public PassWrapper<TestGpuMemoryPromotionPass,
                          OperationPass<gpu::GPUFuncOp>> {
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<StandardOpsDialect, scf::SCFDialect>();
+  }
+
   void runOnOperation() override {
     gpu::GPUFuncOp op = getOperation();
     for (unsigned i = 0, e = op.getNumArguments(); i < e; ++i) {
       if (op.getArgAttrOfType<UnitAttr>(i, "gpu.test_promote_workgroup"))
         promoteToWorkgroupMemory(op, i);
     }
   }
 };

mlir/test/lib/Transforms/TestLinalgHoisting.cpp

 //===- TestLinalgHoisting.cpp - Test Linalg hoisting functions ------------===//
 //
 // 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 implements logic for testing Linalg hoisting functions.
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Hoisting.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 namespace {
 struct TestLinalgHoisting
     : public PassWrapper<TestLinalgHoisting, FunctionPass> {
   TestLinalgHoisting() = default;
   TestLinalgHoisting(const TestLinalgHoisting &pass) {}
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect>();
+  }
 
   void runOnFunction() override;
 
   Option<bool> testHoistViewAllocs{
       *this, "test-hoist-view-allocs",
       llvm::cl::desc("Test hoisting alloc used by view"),
       llvm::cl::init(false)};
   Option<bool> testHoistRedundantTransfers{

mlir/test/lib/Transforms/TestLinalgTransforms.cpp

 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/Dialect/Linalg/Transforms/Transforms.h"
 #include "mlir/Dialect/Linalg/Utils/Utils.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 #include "llvm/ADT/SetVector.h"
 
 using namespace mlir;
 using namespace mlir::linalg;
 
 namespace {
 struct TestLinalgTransforms
     : public PassWrapper<TestLinalgTransforms, FunctionPass> {
   TestLinalgTransforms() = default;
   TestLinalgTransforms(const TestLinalgTransforms &pass) {}
 
+  void getDependentDialects(DialectRegistry &registry) const override {
+    // clang-format off
+    registry.insert<AffineDialect,
+                    scf::SCFDialect,
+                    StandardOpsDialect,
+                    vector::VectorDialect,
+                    gpu::GPUDialect>();
+    // clang-format on
+  }
+
   void runOnFunction() override;
 
   Option<bool> testPatterns{*this, "test-patterns",
                             llvm::cl::desc("Test a mixed set of patterns"),
                             llvm::cl::init(false)};
   Option<bool> testMatmulToVectorPatterns1dTiling{
       *this, "test-matmul-to-vector-patterns-tile-1d",
       llvm::cl::desc(

mlir/test/lib/Transforms/TestVectorTransforms.cpp

 //===- TestVectorToVectorConversion.cpp - Test VectorTransfers lowering ---===//
 //
 // 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 <type_traits>
 
+#include "mlir/Dialect/Affine/IR/AffineOps.h"
+#include "mlir/Dialect/Linalg/IR/LinalgOps.h"
+#include "mlir/Dialect/SCF/SCF.h"
 #include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/Dialect/Vector/VectorOps.h"
 #include "mlir/Dialect/Vector/VectorTransforms.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/Pass/Pass.h"
 
 using namespace mlir;
 using namespace mlir::vector;

 };
 
 struct TestVectorTransferFullPartialSplitPatterns
     : public PassWrapper<TestVectorTransferFullPartialSplitPatterns,
                          FunctionPass> {
   TestVectorTransferFullPartialSplitPatterns() = default;
   TestVectorTransferFullPartialSplitPatterns(
       const TestVectorTransferFullPartialSplitPatterns &pass) {}
+
+  void getDependentDialects(DialectRegistry &registry) const override {
+    registry.insert<AffineDialect, linalg::LinalgDialect, scf::SCFDialect>();
+  }
+
   Option<bool> useLinalgOps{
       *this, "use-linalg-copy",
       llvm::cl::desc("Split using a unmasked vector.transfer + linalg.fill + "
                      "linalg.copy operations."),
       llvm::cl::init(false)};
   void runOnFunction() override {
     MLIRContext *ctx = &getContext();
     OwningRewritePatternList patterns;

mlir/test/mlir-opt/commandline.mlir

 // RUN: mlir-opt --show-dialects | FileCheck %s
-// CHECK: Registered Dialects:
+// CHECK: Available Dialects:
 // CHECK: affine
 // CHECK: gpu
 // CHECK: linalg
 // CHECK: llvm
 // CHECK: nvvm
 // CHECK: omp
 // CHECK: quant
 // CHECK: rocdl
 // CHECK: scf
 // CHECK: sdbm
 // CHECK: spv
 // CHECK: std
 // CHECK: test
 // CHECK: vector

mlir/tools/mlir-linalg-ods-gen/mlir-linalg-ods-gen.cpp

if (!output) {
     exit(1);
   }
 
   // Include the proper Linalg header for end-to-end tblgen testing without
   // resorting to non-portable shgell manipulations.
   if (testEmitIncludeTdHeader)
     output->os() << "include \"mlir/Dialect/Linalg/IR/LinalgStructuredOps.td\"";
 
-  MLIRContext context;
+  MLIRContext context(/*loadAllDialects=*/false);
   llvm::SourceMgr mgr;
   mgr.AddNewSourceBuffer(std::move(file), llvm::SMLoc());
   Parser parser(mgr, &context);
   parseAndEmitAllTensorComprehensions(output->os(), parser);
   output->keep();
 
   return 0;
 }

mlir/tools/mlir-opt/mlir-opt.cpp

 void registerTestAffineLoopUnswitchingPass();
 void registerTestAllReduceLoweringPass();
 void registerTestBufferPlacementPreparationPass();
 void registerTestCallGraphPass();
 void registerTestConstantFold();
 void registerTestConvertGPUKernelToCubinPass();
 void registerTestConvertGPUKernelToHsacoPass();
 void registerTestDominancePass();
+void registerTestDialect(DialectRegistry &);
 void registerTestExpandTanhPass();
 void registerTestFunc();
 void registerTestGpuMemoryPromotionPass();
 void registerTestGpuParallelLoopMappingPass();
 void registerTestInterfaces();
 void registerTestLinalgHoisting();
 void registerTestLinalgTransforms();
 void registerTestLivenessPass();

 #endif
 
 int main(int argc, char **argv) {
   registerAllDialects();
   registerAllPasses();
 #ifdef MLIR_INCLUDE_TESTS
   registerTestPasses();
 #endif
-  return failed(MlirOptMain(argc, argv, "MLIR modular optimizer driver"));
+  DialectRegistry registry;
+  registerAllDialects(registry);
+  registerTestDialect(registry);
+  return failed(MlirOptMain(argc, argv, "MLIR modular optimizer driver\n",
+                            registry,
+                            /*preloadDialectsInContext=*/false));
 }

mlir/tools/mlir-tblgen/DialectGen.cpp

 //===----------------------------------------------------------------------===//
 // GEN: Dialect declarations
 //===----------------------------------------------------------------------===//
 
 /// The code block for the start of a dialect class declaration.
 ///
 /// {0}: The name of the dialect class.
 /// {1}: The dialect namespace.
+/// {2}: initialization code that is emitted in the ctor body before calling
+/// initialize()
 static const char *const dialectDeclBeginStr = R"(
 class {0} : public ::mlir::Dialect {
   explicit {0}(::mlir::MLIRContext *context)
     : ::mlir::Dialect(getDialectNamespace(), context,
       ::mlir::TypeID::get<{0}>()) {{
+    {2}
     initialize();
   }
   void initialize();
   friend class ::mlir::MLIRContext;
 public:
   static ::llvm::StringRef getDialectNamespace() { return "{1}"; }
 )";
 
+/// Registration for a single dependent dialect: to be inserted in the ctor
+/// above for each dependent dialect.
+const char *const dialectRegistrationTemplate = R"(
+    getContext()->getOrLoadDialect<{0}>();
+)";
+
 /// The code block for the attribute parser/printer hooks.
 static const char *const attrParserDecl = R"(
   /// Parse an attribute registered to this dialect.
   ::mlir::Attribute parseAttribute(::mlir::DialectAsmParser &parser,
                                    ::mlir::Type type) const override;
 
   /// Print an attribute registered to this dialect.
   void printAttribute(::mlir::Attribute attr,
::mlir::LogicalResult verifyRegionResultAttribute(
         ::mlir::NamedAttribute attribute) override;
 )";
 
 /// Generate the declaration for the given dialect class.
 static void emitDialectDecl(Dialect &dialect,
                             iterator_range<DialectFilterIterator> dialectAttrs,
                             iterator_range<DialectFilterIterator> dialectTypes,
                             raw_ostream &os) {
+  /// Build the list of dependent dialects
+  std::string dependentDialectRegistrations;
+  {
+    llvm::raw_string_ostream dialectsOs(dependentDialectRegistrations);
+    for (StringRef dependentDialect : dialect.getDependentDialects())
+      dialectsOs << llvm::formatv(dialectRegistrationTemplate,
+                                  dependentDialect);
+  }
   // Emit the start of the decl.
   std::string cppName = dialect.getCppClassName();
-  os << llvm::formatv(dialectDeclBeginStr, cppName, dialect.getName());
+  os << llvm::formatv(dialectDeclBeginStr, cppName, dialect.getName(),
+                      dependentDialectRegistrations);
 
   // Check for any attributes/types registered to this dialect.  If there are,
   // add the hooks for parsing/printing.
   if (!dialectAttrs.empty())
     os << attrParserDecl;
   if (!dialectTypes.empty())
     os << typeParserDecl;
 

mlir/tools/mlir-tblgen/PassGen.cpp

 // GEN: Pass base class generation
 //===----------------------------------------------------------------------===//
 
 /// The code snippet used to generate the start of a pass base class.
 ///
 /// {0}: The def name of the pass record.
 /// {1}: The base class for the pass.
 /// {2): The command line argument for the pass.
+/// {3}: The dependent dialects registration.
 const char *const passDeclBegin = R"(
 //===----------------------------------------------------------------------===//
 // {0}
 //===----------------------------------------------------------------------===//
 
 template <typename DerivedT>
 class {0}Base : public {1} {
 public:
static bool classof(const ::mlir::Pass *pass) {{
     return pass->getTypeID() == ::mlir::TypeID::get<DerivedT>();
   }
 
   /// A clone method to create a copy of this pass.
   std::unique_ptr<::mlir::Pass> clonePass() const override {{
     return std::make_unique<DerivedT>(*static_cast<const DerivedT *>(this));
   }
 
+  /// Return the dialect that must be loaded in the context before this pass.
+  void getDependentDialects(::mlir::DialectRegistry &registry) const override {
+    {3}
+  }
+
 protected:
 )";
 
+/// Registration for a single dependent dialect, to be inserted for each
+/// dependent dialect in the `getDependentDialects` above.
+const char *const dialectRegistrationTemplate = R"(
+  registry.insert<{0}>();
+)";
+
 /// Emit the declarations for each of the pass options.
 static void emitPassOptionDecls(const Pass &pass, raw_ostream &os) {
   for (const PassOption &opt : pass.getOptions()) {
     os.indent(2) << "::mlir::Pass::"
                  << (opt.isListOption() ? "ListOption" : "Option");
 
     os << llvm::formatv("<{0}> {1}{{*this, \"{2}\", ::llvm::cl::desc(\"{3}\")",
                         opt.getType(), opt.getCppVariableName(),
for (const PassStatistic &stat : pass.getStatistics()) {
     os << llvm::formatv(
         "  ::mlir::Pass::Statistic {0}{{this, \"{1}\", \"{2}\"};\n",
         stat.getCppVariableName(), stat.getName(), stat.getDescription());
   }
 }
 
 static void emitPassDecl(const Pass &pass, raw_ostream &os) {
   StringRef defName = pass.getDef()->getName();
+  std::string dependentDialectRegistrations;
+  {
+    llvm::raw_string_ostream dialectsOs(dependentDialectRegistrations);
+    for (StringRef dependentDialect : pass.getDependentDialects())
+      dialectsOs << llvm::formatv(dialectRegistrationTemplate,
+                                  dependentDialect);
+  }
   os << llvm::formatv(passDeclBegin, defName, pass.getBaseClass(),
-                      pass.getArgument());
+                      pass.getArgument(), dependentDialectRegistrations);
   emitPassOptionDecls(pass, os);
   emitPassStatisticDecls(pass, os);
   os << "};\n";
 }
 
 /// Emit the code for registering each of the given passes with the global
 /// PassRegistry.
 static void emitPassDecls(ArrayRef<Pass> passes, raw_ostream &os) {

mlir/tools/mlir-translate/mlir-translate.cpp

int main(int argc, char **argv) {
   if (!output) {
     llvm::errs() << errorMessage << "\n";
     return 1;
   }
 
   // Processes the memory buffer with a new MLIRContext.
   auto processBuffer = [&](std::unique_ptr<llvm::MemoryBuffer> ownedBuffer,
                            raw_ostream &os) {
-    MLIRContext context;
+    MLIRContext context(false);
+    registerAllDialects(&context);
     context.allowUnregisteredDialects();
     context.printOpOnDiagnostic(!verifyDiagnostics);
     llvm::SourceMgr sourceMgr;
     sourceMgr.AddNewSourceBuffer(std::move(ownedBuffer), llvm::SMLoc());
 
     if (!verifyDiagnostics) {
       SourceMgrDiagnosticHandler sourceMgrHandler(sourceMgr, &context);
       return (*translationRequested)(sourceMgr, os, &context);

mlir/unittests/Dialect/Quant/QuantizationUtilsTest.cpp

 #include "mlir/IR/Attributes.h"
 #include "mlir/IR/StandardTypes.h"
 #include "gmock/gmock.h"
 #include "gtest/gtest.h"
 
 using namespace mlir;
 using namespace mlir::quant;
 
-// Load the quant dialect
-static DialectRegistration<QuantizationDialect> QuantOpsRegistration;
-
 namespace {
 
 // Test UniformQuantizedValueConverter converts all APFloat to a magic number 5.
 class TestUniformQuantizedValueConverter
     : public UniformQuantizedValueConverter {
 public:
   TestUniformQuantizedValueConverter(UniformQuantizedType type)
       : UniformQuantizedValueConverter(type), qtype(type) {}

 UniformQuantizedType getTestQuantizedType(Type storageType, MLIRContext *ctx) {
   return UniformQuantizedType::get(/*flags=*/false, storageType,
                                    FloatType::getF32(ctx), /*scale=*/1.0,
                                    /*zeroPoint=*/0, /*storageTypeMin=*/0,
                                    /*storageTypeMax=*/255);
 }
 
 TEST(QuantizationUtilsTest, convertFloatAttrUniform) {
-  MLIRContext ctx;
+  MLIRContext ctx(/*loadAllDialects=*/false);
+  ctx.getOrLoadDialect<QuantizationDialect>();
   IntegerType convertedType = IntegerType::get(8, &ctx);
   auto quantizedType = getTestQuantizedType(convertedType, &ctx);
   TestUniformQuantizedValueConverter converter(quantizedType);
 
   auto realValue = getTestFloatAttr(1.0, &ctx);
   Type typeResult;
   auto valueResult =
       quantizeAttrUniform(realValue, quantizedType, converter, typeResult);
 
   EXPECT_EQ(valueResult.cast<IntegerAttr>().getInt(), 5);
   EXPECT_EQ(
       valueResult.cast<IntegerAttr>().getType().cast<IntegerType>().getWidth(),
       convertedType.getWidth());
 }
 
 TEST(QuantizationUtilsTest, convertRankedDenseAttrUniform) {
-  MLIRContext ctx;
+  MLIRContext ctx(/*loadAllDialects=*/false);
+  ctx.getOrLoadDialect<QuantizationDialect>();
   IntegerType convertedType = IntegerType::get(8, &ctx);
   auto quantizedType = getTestQuantizedType(convertedType, &ctx);
   TestUniformQuantizedValueConverter converter(quantizedType);
   auto realValue = getTestElementsAttr<DenseElementsAttr, ArrayRef<Attribute>>(
       &ctx, {1, 2}, {getTestFloatAttr(1.0, &ctx), getTestFloatAttr(2.0, &ctx)});
 
   Type returnedType;
   auto returnedValue =
       quantizeAttrUniform(realValue, quantizedType, converter, returnedType);
 
   // Check Elements attribute shape and kind are not changed.
   auto tensorType = returnedType.cast<TensorType>();
   auto expectedTensorType = realValue.getType().cast<TensorType>();
   EXPECT_EQ(tensorType.getShape(), expectedTensorType.getShape());
   EXPECT_EQ(tensorType.getElementType(), convertedType);
   EXPECT_TRUE(returnedValue.isa<DenseIntElementsAttr>());
 
   // Check Elements attribute element value is expected.
   auto firstValue = returnedValue.cast<ElementsAttr>().getValue({0, 0});
   EXPECT_EQ(firstValue.cast<IntegerAttr>().getInt(), 5);
 }
 
 TEST(QuantizationUtilsTest, convertRankedSplatAttrUniform) {
-  MLIRContext ctx;
+  MLIRContext ctx(/*loadAllDialects=*/false);
+  ctx.getOrLoadDialect<QuantizationDialect>();
   IntegerType convertedType = IntegerType::get(8, &ctx);
   auto quantizedType = getTestQuantizedType(convertedType, &ctx);
   TestUniformQuantizedValueConverter converter(quantizedType);
   auto realValue = getTestElementsAttr<DenseElementsAttr, Attribute>(
       &ctx, {1, 2}, getTestFloatAttr(1.0, &ctx));
 
   Type returnedType;
   auto returnedValue =
       quantizeAttrUniform(realValue, quantizedType, converter, returnedType);
 
   // Check Elements attribute shape and kind are not changed.
   auto tensorType = returnedType.cast<TensorType>();
   auto expectedTensorType = realValue.getType().cast<TensorType>();
   EXPECT_EQ(tensorType.getShape(), expectedTensorType.getShape());
   EXPECT_EQ(tensorType.getElementType(), convertedType);
   EXPECT_TRUE(returnedValue.isa<SplatElementsAttr>());
 
   // Check Elements attribute element value is expected.
   auto firstValue = returnedValue.cast<ElementsAttr>().getValue({0, 0});
   EXPECT_EQ(firstValue.cast<IntegerAttr>().getInt(), 5);
 }
 
 TEST(QuantizationUtilsTest, convertRankedSparseAttrUniform) {
-  MLIRContext ctx;
+  MLIRContext ctx(/*loadAllDialects=*/false);
+  ctx.getOrLoadDialect<QuantizationDialect>();
   IntegerType convertedType = IntegerType::get(8, &ctx);
   auto quantizedType = getTestQuantizedType(convertedType, &ctx);
   TestUniformQuantizedValueConverter converter(quantizedType);
   auto realValue = getTestSparseElementsAttr(&ctx, {1, 2});
 
   Type returnedType;
   auto returnedValue =
       quantizeAttrUniform(realValue, quantizedType, converter, returnedType);

mlir/unittests/Dialect/SPIRV/DeserializationTest.cpp

 //===----------------------------------------------------------------------===//
 // Test Fixture
 //===----------------------------------------------------------------------===//
 
 /// A deserialization test fixture providing minimal SPIR-V building and
 /// diagnostic checking utilities.
 class DeserializationTest : public ::testing::Test {
 protected:
-  DeserializationTest() {
+  DeserializationTest() : context(/*loadAllDialects=*/false) {
+    context.getOrLoadDialect<mlir::spirv::SPIRVDialect>();
     // Register a diagnostic handler to capture the diagnostic so that we can
     // check it later.
     context.getDiagEngine().registerHandler([&](Diagnostic &diag) {
       diagnostic.reset(new Diagnostic(std::move(diag)));
     });
   }
 
   /// Performs deserialization and returns the constructed spv.module op.

mlir/unittests/Dialect/SPIRV/SerializationTest.cpp

 using namespace mlir;
 
 //===----------------------------------------------------------------------===//
 // Test Fixture
 //===----------------------------------------------------------------------===//
 
 class SerializationTest : public ::testing::Test {
 protected:
-  SerializationTest() { createModuleOp(); }
+  SerializationTest() : context(/*loadAllDialects=*/false) {
+    context.getOrLoadDialect<mlir::spirv::SPIRVDialect>();
+    createModuleOp();
+  }
 
   void createModuleOp() {
     OpBuilder builder(&context);
     OperationState state(UnknownLoc::get(&context),
                          spirv::ModuleOp::getOperationName());
     state.addAttribute("addressing_model",
                        builder.getI32IntegerAttr(static_cast<uint32_t>(
                            spirv::AddressingModel::Logical)));

mlir/unittests/IR/AttributeTest.cpp

static void testSplat(Type eltType, const EltTy &splatElt) {
   EXPECT_EQ(detectedSplat, splat);
 
   for (auto newValue : detectedSplat.template getValues<EltTy>())
     EXPECT_TRUE(newValue == splatElt);
 }
 
 namespace {
 TEST(DenseSplatTest, BoolSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   IntegerType boolTy = IntegerType::get(1, &context);
   RankedTensorType shape = RankedTensorType::get({2, 2}, boolTy);
 
   // Check that splat is automatically detected for boolean values.
   /// True.
   DenseElementsAttr trueSplat = DenseElementsAttr::get(shape, true);
   EXPECT_TRUE(trueSplat.isSplat());
   /// False.
   DenseElementsAttr falseSplat = DenseElementsAttr::get(shape, false);
   EXPECT_TRUE(falseSplat.isSplat());
   EXPECT_NE(falseSplat, trueSplat);
 
   /// Detect and handle splat within 8 elements (bool values are bit-packed).
   /// True.
   auto detectedSplat = DenseElementsAttr::get(shape, {true, true, true, true});
   EXPECT_EQ(detectedSplat, trueSplat);
   /// False.
   detectedSplat = DenseElementsAttr::get(shape, {false, false, false, false});
   EXPECT_EQ(detectedSplat, falseSplat);
 }
 
 TEST(DenseSplatTest, LargeBoolSplat) {
   constexpr int64_t boolCount = 56;
 
-  MLIRContext context;
+  MLIRContext context(false);
   IntegerType boolTy = IntegerType::get(1, &context);
   RankedTensorType shape = RankedTensorType::get({boolCount}, boolTy);
 
   // Check that splat is automatically detected for boolean values.
   /// True.
   DenseElementsAttr trueSplat = DenseElementsAttr::get(shape, true);
   DenseElementsAttr falseSplat = DenseElementsAttr::get(shape, false);
   EXPECT_TRUE(trueSplat.isSplat());
   EXPECT_TRUE(falseSplat.isSplat());
 
   /// Detect that the large boolean arrays are properly splatted.
   /// True.
   SmallVector<bool, 64> trueValues(boolCount, true);
   auto detectedSplat = DenseElementsAttr::get(shape, trueValues);
   EXPECT_EQ(detectedSplat, trueSplat);
   /// False.
   SmallVector<bool, 64> falseValues(boolCount, false);
   detectedSplat = DenseElementsAttr::get(shape, falseValues);
   EXPECT_EQ(detectedSplat, falseSplat);
 }
 
 TEST(DenseSplatTest, BoolNonSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   IntegerType boolTy = IntegerType::get(1, &context);
   RankedTensorType shape = RankedTensorType::get({6}, boolTy);
 
   // Check that we properly handle non-splat values.
   DenseElementsAttr nonSplat =
       DenseElementsAttr::get(shape, {false, false, true, false, false, true});
   EXPECT_FALSE(nonSplat.isSplat());
 }
 
 TEST(DenseSplatTest, OddIntSplat) {
   // Test detecting a splat with an odd(non 8-bit) integer bitwidth.
-  MLIRContext context;
+  MLIRContext context(false);
   constexpr size_t intWidth = 19;
   IntegerType intTy = IntegerType::get(intWidth, &context);
   APInt value(intWidth, 10);
 
   testSplat(intTy, value);
 }
 
 TEST(DenseSplatTest, Int32Splat) {
-  MLIRContext context;
+  MLIRContext context(false);
   IntegerType intTy = IntegerType::get(32, &context);
   int value = 64;
 
   testSplat(intTy, value);
 }
 
 TEST(DenseSplatTest, IntAttrSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   IntegerType intTy = IntegerType::get(85, &context);
   Attribute value = IntegerAttr::get(intTy, 109);
 
   testSplat(intTy, value);
 }
 
 TEST(DenseSplatTest, F32Splat) {
-  MLIRContext context;
+  MLIRContext context(false);
   FloatType floatTy = FloatType::getF32(&context);
   float value = 10.0;
 
   testSplat(floatTy, value);
 }
 
 TEST(DenseSplatTest, F64Splat) {
-  MLIRContext context;
+  MLIRContext context(false);
   FloatType floatTy = FloatType::getF64(&context);
   double value = 10.0;
 
   testSplat(floatTy, APFloat(value));
 }
 
 TEST(DenseSplatTest, FloatAttrSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   FloatType floatTy = FloatType::getF32(&context);
   Attribute value = FloatAttr::get(floatTy, 10.0);
 
   testSplat(floatTy, value);
 }
 
 TEST(DenseSplatTest, BF16Splat) {
-  MLIRContext context;
+  MLIRContext context(false);
   FloatType floatTy = FloatType::getBF16(&context);
   Attribute value = FloatAttr::get(floatTy, 10.0);
 
   testSplat(floatTy, value);
 }
 
 TEST(DenseSplatTest, StringSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   Type stringType =
       OpaqueType::get(Identifier::get("test", &context), "string", &context);
   StringRef value = "test-string";
   testSplat(stringType, value);
 }
 
 TEST(DenseSplatTest, StringAttrSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   Type stringType =
       OpaqueType::get(Identifier::get("test", &context), "string", &context);
   Attribute stringAttr = StringAttr::get("test-string", stringType);
   testSplat(stringType, stringAttr);
 }
 
 TEST(DenseComplexTest, ComplexFloatSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   ComplexType complexType = ComplexType::get(FloatType::getF32(&context));
   std::complex<float> value(10.0, 15.0);
   testSplat(complexType, value);
 }
 
 TEST(DenseComplexTest, ComplexIntSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   ComplexType complexType = ComplexType::get(IntegerType::get(64, &context));
   std::complex<int64_t> value(10, 15);
   testSplat(complexType, value);
 }
 
 TEST(DenseComplexTest, ComplexAPFloatSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   ComplexType complexType = ComplexType::get(FloatType::getF32(&context));
   std::complex<APFloat> value(APFloat(10.0f), APFloat(15.0f));
   testSplat(complexType, value);
 }
 
 TEST(DenseComplexTest, ComplexAPIntSplat) {
-  MLIRContext context;
+  MLIRContext context(false);
   ComplexType complexType = ComplexType::get(IntegerType::get(64, &context));
   std::complex<APInt> value(APInt(64, 10), APInt(64, 15));
   testSplat(complexType, value);
 }
 
 } // end namespace

mlir/unittests/IR/DialectTest.cpp

 struct AnotherTestDialect : public Dialect {
   static StringRef getDialectNamespace() { return "test"; };
   AnotherTestDialect(MLIRContext *context)
       : Dialect(getDialectNamespace(), context,
                 TypeID::get<AnotherTestDialect>()) {}
 };
 
 TEST(DialectDeathTest, MultipleDialectsWithSameNamespace) {
-  MLIRContext context;
+  MLIRContext context(false);
 
   // Registering a dialect with the same namespace twice should result in a
   // failure.
-  context.getOrCreateDialect<TestDialect>();
-  ASSERT_DEATH(context.getOrCreateDialect<AnotherTestDialect>(), "");
+  context.loadDialect<TestDialect>();
+  ASSERT_DEATH(context.loadDialect<AnotherTestDialect>(), "");
 }
 
 } // end namespace

mlir/unittests/IR/OperationSupportTest.cpp

static Operation *createOp(MLIRContext *context,
   context->allowUnregisteredDialects();
   return Operation::create(UnknownLoc::get(context),
                            OperationName("foo.bar", context), resultTypes,
                            operands, llvm::None, llvm::None, 0);
 }
 
 namespace {
 TEST(OperandStorageTest, NonResizable) {
-  MLIRContext context;
+  MLIRContext context(false);
   Builder builder(&context);
 
   Operation *useOp =
       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
   Value operand = useOp->getResult(0);
 
   // Create a non-resizable operation with one operand.
   Operation *user = createOp(&context, operand);
 
   // The same number of operands is okay.
   user->setOperands(operand);
   EXPECT_EQ(user->getNumOperands(), 1u);
 
   // Removing is okay.
   user->setOperands(llvm::None);
   EXPECT_EQ(user->getNumOperands(), 0u);
 
   // Destroy the operations.
   user->destroy();
   useOp->destroy();
 }
 
 TEST(OperandStorageTest, Resizable) {
-  MLIRContext context;
+  MLIRContext context(false);
   Builder builder(&context);
 
   Operation *useOp =
       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
   Value operand = useOp->getResult(0);
 
   // Create a resizable operation with one operand.
   Operation *user = createOp(&context, operand);
TEST(OperandStorageTest, Resizable) {
   EXPECT_EQ(user->getNumOperands(), 3u);
 
   // Destroy the operations.
   user->destroy();
   useOp->destroy();
 }
 
 TEST(OperandStorageTest, RangeReplace) {
-  MLIRContext context;
+  MLIRContext context(false);
   Builder builder(&context);
 
   Operation *useOp =
       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
   Value operand = useOp->getResult(0);
 
   // Create a resizable operation with one operand.
   Operation *user = createOp(&context, operand);
TEST(OperandStorageTest, RangeReplace) {
   EXPECT_EQ(user->getNumOperands(), 0u);
 
   // Destroy the operations.
   user->destroy();
   useOp->destroy();
 }
 
 TEST(OperandStorageTest, MutableRange) {
-  MLIRContext context;
+  MLIRContext context(false);
   Builder builder(&context);
 
   Operation *useOp =
       createOp(&context, /*operands=*/llvm::None, builder.getIntegerType(16));
   Value operand = useOp->getResult(0);
 
   // Create a resizable operation with one operand.
   Operation *user = createOp(&context, operand);

mlir/unittests/Pass/AnalysisManagerTest.cpp

 struct OtherAnalysis {
   OtherAnalysis(Operation *) {}
 };
 struct OpSpecificAnalysis {
   OpSpecificAnalysis(ModuleOp) {}
 };
 
 TEST(AnalysisManagerTest, FineGrainModuleAnalysisPreservation) {
-  MLIRContext context;
+  MLIRContext context(false);
 
   // Test fine grain invalidation of the module analysis manager.
   OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
   ModuleAnalysisManager mam(*module, /*passInstrumentor=*/nullptr);
   AnalysisManager am = mam;
 
   // Query two different analyses, but only preserve one before invalidating.
   am.getAnalysis<MyAnalysis>();
   am.getAnalysis<OtherAnalysis>();
 
   detail::PreservedAnalyses pa;
   pa.preserve<MyAnalysis>();
   am.invalidate(pa);
 
   // Check that only MyAnalysis is preserved.
   EXPECT_TRUE(am.getCachedAnalysis<MyAnalysis>().hasValue());
   EXPECT_FALSE(am.getCachedAnalysis<OtherAnalysis>().hasValue());
 }
 
 TEST(AnalysisManagerTest, FineGrainFunctionAnalysisPreservation) {
-  MLIRContext context;
+  MLIRContext context(false);
   Builder builder(&context);
 
   // Create a function and a module.
   OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
   FuncOp func1 =
       FuncOp::create(builder.getUnknownLoc(), "foo",
                      builder.getFunctionType(llvm::None, llvm::None));
   module->push_back(func1);
TEST(AnalysisManagerTest, FineGrainFunctionAnalysisPreservation) {
   fam.invalidate(pa);
 
   // Check that only MyAnalysis is preserved.
   EXPECT_TRUE(fam.getCachedAnalysis<MyAnalysis>().hasValue());
   EXPECT_FALSE(fam.getCachedAnalysis<OtherAnalysis>().hasValue());
 }
 
 TEST(AnalysisManagerTest, FineGrainChildFunctionAnalysisPreservation) {
-  MLIRContext context;
+  MLIRContext context(false);
   Builder builder(&context);
 
   // Create a function and a module.
   OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
   FuncOp func1 =
       FuncOp::create(builder.getUnknownLoc(), "foo",
                      builder.getFunctionType(llvm::None, llvm::None));
   module->push_back(func1);
struct CustomInvalidatingAnalysis {
   CustomInvalidatingAnalysis(Operation *) {}
 
   bool isInvalidated(const AnalysisManager::PreservedAnalyses &pa) {
     return !pa.isPreserved<TestAnalysisSet>();
   }
 };
 
 TEST(AnalysisManagerTest, CustomInvalidation) {
-  MLIRContext context;
+  MLIRContext context(false);
   Builder builder(&context);
 
   // Create a function and a module.
   OwningModuleRef module(ModuleOp::create(UnknownLoc::get(&context)));
   ModuleAnalysisManager mam(*module, /*passInstrumentor=*/nullptr);
   AnalysisManager am = mam;
 
   detail::PreservedAnalyses pa;

mlir/unittests/SDBM/SDBMTest.cpp

 #include "mlir/IR/AffineExpr.h"
 #include "mlir/IR/MLIRContext.h"
 #include "gtest/gtest.h"
 
 #include "llvm/ADT/DenseSet.h"
 
 using namespace mlir;
 
-/// Load the SDBM dialect.
-static DialectRegistration<SDBMDialect> SDBMRegistration;
 
 static MLIRContext *ctx() {
-  static thread_local MLIRContext context;
+  static thread_local MLIRContext context(false);
+  context.getOrLoadDialect<SDBMDialect>();
   return &context;
 }
 
 static SDBMDialect *dialect() {
   static thread_local SDBMDialect *d = nullptr;
   if (!d) {
-    d = ctx()->getRegisteredDialect<SDBMDialect>();
+    d = ctx()->getOrLoadDialect<SDBMDialect>();
   }
   return d;
 }
 
 static SDBMExpr dim(unsigned pos) { return SDBMDimExpr::get(dialect(), pos); }
 
 static SDBMExpr symb(unsigned pos) {
   return SDBMSymbolExpr::get(dialect(), pos);

mlir/unittests/TableGen/OpBuildGen.cpp

 #include <vector>
 
 namespace mlir {
 
 //===----------------------------------------------------------------------===//
 // Test Fixture
 //===----------------------------------------------------------------------===//
 
+static MLIRContext &getContext() {
+  static MLIRContext ctx(false);
+  ctx.getOrLoadDialect<TestDialect>();
+  return ctx;
+}
 /// Test fixture for providing basic utilities for testing.
 class OpBuildGenTest : public ::testing::Test {
 protected:
   OpBuildGenTest()
-      : ctx{}, builder(&ctx), loc(builder.getUnknownLoc()),
+      : ctx(getContext()), builder(&ctx), loc(builder.getUnknownLoc()),
         i32Ty(builder.getI32Type()), f32Ty(builder.getF32Type()),
         cstI32(builder.create<TableGenConstant>(loc, i32Ty)),
         cstF32(builder.create<TableGenConstant>(loc, f32Ty)),
         noAttrs(), attrStorage{builder.getNamedAttr("attr0",
                                                     builder.getBoolAttr(true)),
                                builder.getNamedAttr(
                                    "attr1", builder.getI32IntegerAttr(33))},
         attrs(attrStorage) {}
void testSingleVariadicInputInferredType() {
     verifyOp(std::move(op), {i32Ty}, {cstI32, cstI32}, noAttrs);
 
     // Test build method with no result types and supplied attributes.
     op = builder.create<OpTy>(loc, ArrayRef<Value>{cstI32, cstI32}, attrs);
     verifyOp(std::move(op), {i32Ty}, {cstI32, cstI32}, attrs);
   }
 
 protected:
-  MLIRContext ctx;
+  MLIRContext &ctx;
   OpBuilder builder;
   Location loc;
   Type i32Ty;
   Type f32Ty;
   TableGenConstant cstI32;
   TableGenConstant cstF32;
 
   ArrayRef<NamedAttribute> noAttrs;

mlir/unittests/TableGen/StructsGenTest.cpp

static test::TestStruct getTestStruct(mlir::MLIRContext *context) {
 
   return test::TestStruct::get(integerAttr, floatAttr, elementsAttr,
                                optionalAttr, context);
 }
 
 /// Validates that test::TestStruct::classof correctly identifies a valid
 /// test::TestStruct.
 TEST(StructsGenTest, ClassofTrue) {
-  mlir::MLIRContext context;
+  mlir::MLIRContext context(false);
   auto structAttr = getTestStruct(&context);
   ASSERT_TRUE(test::TestStruct::classof(structAttr));
 }
 
 /// Validates that test::TestStruct::classof fails when an extra attribute is in
 /// the class.
 TEST(StructsGenTest, ClassofExtraFalse) {
   mlir::MLIRContext context;