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.

79

nit: Move after constructor.

213

Move this before options.

299

Same here.

330

Same here and a few other places.

mlir/include/mlir/ExecutionEngine/JitRunner.h

24

Looks unused.

mlir/include/mlir/IR/Dialect.h

223

nit: allows for decoupling

227

The ordering of this container should be stable.

231

nit: add -> insert?

248

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.

260

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

mlir/include/mlir/IR/MLIRContext.h

25

This is unnecessary.

[rriddle]

mlir/include/mlir/IR/MLIRContext.h

47

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

139

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

Can we autogen this with the dialect tablegen backend?

mlir/lib/IR/MLIRContext.cpp

451

nit: Use array_pod_sort here.

469

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

178

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

25

Leftover from previous iterations...

47

"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

469

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

42

@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 ↗

(On Diff #285583)

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

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

26 ↗

(On Diff #285583)

Same here and a few other places.

mlir/include/mlir/IR/Dialect.h

228

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

250

nit: name is in a different style.

257

Same here.

mlir/include/mlir/IR/MLIRContext.h

47

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

26 ↗

(On Diff #285583)

nit: an option upfront -> an upfront?

mlir/lib/CAPI/IR/IR.cpp

66

nit: Parameter comment on the false.

mlir/lib/Conversion/PassDetail.h

18

Yeah, I just guessed.

mlir/lib/IR/Dialect.cpp

34

nit: name.str()?

mlir/lib/IR/MLIRContext.cpp

451

Unresolved?

451

nit: Drop trivial braces here.

mlir/lib/Pass/Pass.cpp

747

nit: Fix the naming style here.

mlir/tools/mlir-tblgen/PassGen.cpp

108

nit: Naming style here.

[mehdi_amini]

Address comments

[mehdi_amini]

mlir/include/mlir/IR/Dialect.h

228

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

47

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/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;
+  // 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;
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
+
   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;
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
+
   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.add<AffineDialect>();
+    registry.add<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;
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
+
   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.add<AffineDialect>();
+    registry.add<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.add<LLVM::LLVMDialect>();
+    registry.add<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;
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
+
   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.add<AffineDialect>();
+    registry.add<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.add<LLVM::LLVMDialect>();
+    registry.add<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;
+  // Load our Dialect in this MLIR Context.
+  context.getOrLoadDialect<mlir::toy::ToyDialect>();
+
   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/IR.h

 /*============================================================================*/
 
 /** Creates an MLIR context and transfers its ownership to the caller. */
 MlirContext mlirContextCreate();
 
 /** Takes an MLIR context owned by the caller and destroys it. */
 void mlirContextDestroy(MlirContext context);
 
+/** Load all the globally registered dialects in the provided context.
+ *  TODO: remove the concept of globally registered dialect by exposing the
+ *  DialectRegistry.
+ */
+void mlirContextLoadAllDialects(MlirContext context);
+
 /*============================================================================*/
 /* Location API.                                                              */
 /*============================================================================*/
 
 /** Creates an File/Line/Column location owned by the given context. */
 MlirLocation mlirLocationFileLineColGet(MlirContext context,
                                         const char *filename, unsigned line,
                                         unsigned col);

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 dependentDialects = ["LLVM::LLVMDialect", "LLVM::LLVMAVX512Dialect"];
   let constructor = "mlir::createConvertAVX512ToLLVMPass()";
 }
 
 //===----------------------------------------------------------------------===//
 // GPUCommon
 //===----------------------------------------------------------------------===//
 
 def ConvertGpuLaunchFuncToGpuRuntimeCalls : Pass<"launch-func-to-gpu-runtime",

 
 //===----------------------------------------------------------------------===//
 // 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 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">
   ];
+  let dependentDialects = ["gpu::GPUDialect"];
 }
 
 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 "
Option<"useBarePtrCallConv", "use-bare-ptr-memref-call-conv", "bool",
            "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">,
   ];
+  let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // 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 options = [
     Option<"fullUnroll", "full-unroll", "bool", /*default=*/"false",
            "Perform full unrolling when converting vector transfers to SCF">,
   ];
+  let dependentDialects = ["AffineDialect", "scf::SCFDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // 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 options = [
     Option<"reassociateFPReductions", "reassociate-fp-reductions",
            "bool", /*default=*/"false",
            "Allows llvm to reassociate floating-point reductions for speed">
   ];
+  let dependentDialects = ["LLVM::LLVMDialect"];
 }
 
 //===----------------------------------------------------------------------===//
 // 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/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/ExecutionEngine/JitRunner.h

 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SUPPORT_JITRUNNER_H_
 #define MLIR_SUPPORT_JITRUNNER_H_
 
 #include "llvm/ADT/STLExtras.h"
 
 namespace mlir {
-
+class MLIRContext;
 class ModuleOp;
 struct LogicalResult;
 
 // Entry point for all CPU runners. Expects the common argc/argv arguments for
 // standard C++ main functions and an mlirTransformer.
 // The latter is applied after parsing the input into MLIR IR and before passing
 // the MLIR module to the ExecutionEngine.
 int JitRunnerMain(
     int argc, char **argv,
     llvm::function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer);
 
 } // namespace mlir
 
 #endif // MLIR_SUPPORT_JITRUNNER_H_

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 "llvm/ADT/StringMap.h"
+
 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 to decouple 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 = llvm::StringMap<std::pair<TypeID, DialectAllocatorFunction>>;
+
+public:
+  template <typename ConcreteDialect>
+  void add() {
+    add(TypeID::get<ConcreteDialect>(), ConcreteDialect::getDialectNamespace(),
+        static_cast<DialectAllocatorFunction>(([](MLIRContext *ctx) {
+          // Just allocate the dialect, the context
+          // takes ownership of it.
+          return ctx->getOrLoadDialect<ConcreteDialect>();
+        })));
+  }
+
+  /// Add a new dialect constructor to the registry.
+  void add(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 addAll(MLIRContext *context) {
+    for (const auto &name_and_registration_it : registry)
+      context->getDialectRegistry().add(name_and_registration_it.second.first,
+                                        name_and_registration_it.first(),
+                                        name_and_registration_it.second.second);
+  }
+  // Load all dialects available in the registry in the provided context.
+  void loadAll(MLIRContext *context) {
+    for (const auto &name_and_registration_it : registry)
+      name_and_registration_it.second.second(context);
+  }
+
+  MapTy::iterator begin() { return registry.begin(); }
+  MapTy::iterator end() { return registry.end(); }
+  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.
+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);
+inline void registerAllDialects(MLIRContext *context) {
+  getGlobalDialectRegistry().addAll(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>(); });
+template <typename ConcreteDialect>
+void registerDialect() {
+  getGlobalDialectRegistry().add<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

 //===- MLIRContext.h - MLIR Global Context Class ----------------*- 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_IR_MLIRCONTEXT_H
 #define MLIR_IR_MLIRCONTEXT_H
 
 #include "mlir/Support/LLVM.h"
 #include "mlir/Support/TypeID.h"
 #include <functional>
 #include <memory>
 #include <vector>
-
 namespace mlir {
 class AbstractOperation;
 class DiagnosticEngine;
 class Dialect;
+class DialectRegistry;
 class InFlightDiagnostic;
 class Location;
 class MLIRContextImpl;
+class MLIRContext;
 class StorageUniquer;
 
 /// 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();
   ~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;
         }));
   }
 
+  /// 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);
+  Dialect *getOrLoadDialect(StringRef dialectNamespace, TypeID dialectID,
 
   MLIRContext(const MLIRContext &) = delete;
   void operator=(const MLIRContext &) = delete;
 };
 
 //===----------------------------------------------------------------------===//
 // MLIRContext CommandLine Options
 //===----------------------------------------------------------------------===//

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) {
+  registry.add<AffineDialect>();
+  registry.add<avx512::AVX512Dialect>();
+  registry.add<gpu::GPUDialect>();
+  registry.add<LLVM::LLVMAVX512Dialect>();
+  registry.add<LLVM::LLVMDialect>();
+  registry.add<linalg::LinalgDialect>();
+  registry.add<scf::SCFDialect>();
+  registry.add<omp::OpenMPDialect>();
+  registry.add<quant::QuantizationDialect>();
+  registry.add<spirv::SPIRVDialect>();
+  registry.add<StandardOpsDialect>();
+  registry.add<vector::VectorDialect>();
+  registry.add<NVVM::NVVMDialect>();
+  registry.add<ROCDL::ROCDLDialect>();
+  registry.add<SDBMDialect>();
+  registry.add<shape::ShapeDialect>();
+}
+
 // 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<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;
-  }();
+  static bool init_once =
+      ([]() { registerAllDialects(getGlobalDialectRegistry()); }(), true);
   (void)init_once;
 }
 } // namespace mlir
 
 #endif // MLIR_INITALLDIALECTS_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 operations for,
+  /// other than dialect that exists in its 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 operations 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);
 
+  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/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 : Pass<"buffer-placement"> {
+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/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/InitAllDialects.h"
 #include "mlir/Parser.h"
 
 using namespace mlir;
 
 /* ========================================================================== */
 /* Definitions of methods for non-owning structures used in C API.            */
 /* ========================================================================== */
 
static ArrayRef<CppTy> unwrapList(unsigned size, CTy *first,
   return storage;
 }
 
 /* ========================================================================== */
 /* Context API.                                                               */
 /* ========================================================================== */
 
 MlirContext mlirContextCreate() {
   auto *context = new MLIRContext;
   return wrap(context);
 }
 
 void mlirContextDestroy(MlirContext context) { delete unwrap(context); }
 
+void mlirContextLoadAllDialects(MlirContext context) {
+  registerAllDialects(unwrap(context));
+  getGlobalDialectRegistry().loadAll(unwrap(context));
+}
+
 /* ========================================================================== */
 /* Location API.                                                              */
 /* ========================================================================== */
 
 MlirLocation mlirLocationFileLineColGet(MlirContext context,
                                         const char *filename, unsigned line,
                                         unsigned col) {
   return wrap(FileLineColLoc::get(filename, line, col, 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;
+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 =
         llvmDialect->getDataLayout().getPointerSizeInBits();
 
   // Register conversions for the standard types.
   addConversion([&](ComplexType type) { return convertComplexType(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 {
+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/LLVMIR/IR/LLVMDialect.cpp

struct LLVMDialectImpl {
   // eventually replaced with a proper MLIR data layout.
   llvm::DataLayout layout;
 };
 } // end namespace detail
 } // end namespace LLVM
 } // end namespace mlir
 
 void LLVMDialect::initialize() {
+
+  /// FIXME: at the moment this is a dependency of the translation to LLVM IR,
+  /// not really one of this dialect per-se.
+  getContext()->getOrLoadDialect("omp");
+
   impl = new detail::LLVMDialectImpl();
   // clang-format off
   addTypes<LLVMVoidType,
            LLVMHalfType,
            LLVMBFloatType,
            LLVMFloatType,
            LLVMDoubleType,
            LLVMFP128Type,

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

 
 // TODO: This should be the llvm.nvvm dialect once this is supported.
 void NVVMDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
       >();
 
+  // Register dependent dialect
+  getContext()->getOrLoadDialect<LLVM::LLVMDialect>();
+
   // Support unknown operations because not all NVVM operations are registered.
   allowUnknownOperations();
 }
 
 namespace mlir {
 namespace NVVM {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/NVVMOps.cpp.inc"
 } // namespace NVVM
 } // namespace mlir
 

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

 
 // TODO: This should be the llvm.rocdl dialect once this is supported.
 void ROCDLDialect::initialize() {
   addOperations<
 #define GET_OP_LIST
 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
       >();
 
+  // Register dependent dialect
+  getContext()->getOrLoadDialect<LLVM::LLVMDialect>();
+
   // Support unknown operations because not all ROCDL operations are registered.
   allowUnknownOperations();
 }
 
 namespace mlir {
 namespace ROCDL {
 #define GET_OP_CLASSES
 #include "mlir/Dialect/LLVMIR/ROCDLOps.cpp.inc"
 } // namespace ROCDL
 } // namespace mlir
 

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 {
+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 {
+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

for (unsigned i = 0, e = options.llvmPasses.size(); i < e; ++i) {
     passes.push_back(options.llvmPasses[i]);
     if (optCLIPosition < options.llvmPasses.getPosition(i)) {
       optPosition = i;
       optCLIPosition = UINT_MAX; // To ensure we never insert again
     }
   }
 
   MLIRContext context;
+  registerAllDialects(&context);
+
   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});
+Dialect *DialectRegistry::loadByName(StringRef name, MLIRContext *context) {
+  auto it = registry.find(name);
+  if (it == registry.end())
+    return nullptr;
+  return it->second.second(context);
+}
+
+void DialectRegistry::add(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);
   }
-
-/// 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
 //===----------------------------------------------------------------------===//
 
 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>> dialects;
+  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;
 
MLIRContext::MLIRContext() : 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>();
+  getOrLoadDialect<BuiltinDialect>();
   registerAllDialects(this);
 
   // 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<FloatType>(this, StandardTypes::BF16);

 
 /// 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.push_back(dialect.second.get());
+  llvm::sort(result, [](Dialect *lhs, Dialect *rhs) {
+    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());
+  llvm::sort(result);
   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->dialects.find(name);
+  return (it != impl->dialects.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.dialects[dialectNamespace];
+
+  if (!dialect) {
+    LLVM_DEBUG(llvm::dbgs()
+               << "Load new dialect in Context" << dialectNamespace);
+    dialect = ctor();
+    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();
 }
 
 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();
 
+    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);
+  auto *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)) {
+        auto *dialect = builder.getContext()->getOrLoadDialect(dialectName);
+        if (dialect) {
           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?
+      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 dependent_dialects;
+  getDependentDialects(dependent_dialects);
+  dependent_dialects.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

static LogicalResult processBuffer(raw_ostream &os,
                                    bool allowUnregisteredDialects,
                                    const PassPipelineCLParser &passPipeline) {
   // 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;
+  registerAllDialects(&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,

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->getOrLoadDialect<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) {
-  auto *dialect = m->getContext()->getRegisteredDialect<LLVM::LLVMDialect>();
+  auto *dialect = m->getContext()->getOrLoadDialect<LLVM::LLVMDialect>();
   assert(dialect && "LLVM dialect must be registered");
 
   auto llvmModule = std::make_unique<llvm::Module>(name, llvmContext);
   llvmModule->setDataLayout(dialect->getDataLayout());
 
   // 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/BufferPlacement.cpp

 // be invalid with respect to program semantics. The only thing that is
 // currently missing is a high-level loop analysis that allows us to move allocs
 // and deallocs outside of the loop blocks. Furthermore, it doesn't also accept
 // functions which return buffers already.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Transforms/BufferPlacement.h"
+#include "PassDetail.h"
 #include "mlir/Dialect/Linalg/IR/LinalgOps.h"
 #include "mlir/IR/Operation.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/ADT/SetOperations.h"
 
 using namespace mlir;
 

 
 //===----------------------------------------------------------------------===//
 // BufferPlacementPass
 //===----------------------------------------------------------------------===//
 
 /// The actual buffer placement pass that moves alloc and dealloc nodes into
 /// the right positions. It uses the algorithm described at the top of the
 /// file.
-struct BufferPlacementPass
-    : mlir::PassWrapper<BufferPlacementPass, FunctionPass> {
+struct BufferPlacementPass : BufferPlacementBase<BufferPlacementPass> {
 
   void runOnFunction() override {
     // Place all required alloc, copy and dealloc nodes.
     BufferPlacement placement(getFunction());
     placement.place();
   }
 };
 

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 {
+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

void collectStats(MlirOperation operation) {
   printf("Number of blocks: %u\n", stats.numBlocks);
   printf("Number of regions: %u\n", stats.numRegions);
   printf("Number of values: %u\n", stats.numValues);
 }
 
 int main() {
   mlirRegisterAllDialects();
   MlirContext ctx = mlirContextCreate();
+  mlirContextLoadAllDialects(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;
+  static bool init_once = [&]() {
+    context.getOrLoadDialect<AffineDialect>();
+    context.getOrLoadDialect<scf::SCFDialect>();
+    context.getOrLoadDialect<linalg::LinalgDialect>();
+    context.getOrLoadDialect<StandardOpsDialect>();
+    context.getOrLoadDialect<vector::VectorDialect>();
     return true;
   }();
   (void)init_once;
-  static thread_local MLIRContext context;
   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 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.add<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/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.add<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.add<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.add<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.add<StandardOpsDialect>();
+    registry.add<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.add<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

 //===- TestLinalgTransforms.cpp - Test Linalg transformation patterns -----===//
 //
 // 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 transformations.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Dialect/Affine/IR/AffineOps.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 {
+    registry.add<AffineDialect>();
+    registry.add<scf::SCFDialect>();
+    registry.add<StandardOpsDialect>();
+    registry.add<vector::VectorDialect>();
+  }
+
   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.add<AffineDialect>();
+    registry.add<linalg::LinalgDialect>();
+    registry.add<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-cuda-runner/all-reduce-or.mlir

-// RUN: mlir-cuda-runner %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
+// RUN: mlir-cuda-runner --print-ir-after-all %s --shared-libs=%cuda_wrapper_library_dir/libcuda-runtime-wrappers%shlibext,%linalg_test_lib_dir/libmlir_runner_utils%shlibext --entry-point-result=void | FileCheck %s
 
 func @main() {
   %data = alloc() : memref<2x6xi32>
   %sum = alloc() : memref<2xi32>
   %cst0 = constant 0 : i32
   %cst1 = constant 1 : i32
   %cst2 = constant 2 : i32
   %cst4 = constant 4 : i32

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

#endif
   registerAsmPrinterCLOptions();
   registerMLIRContextCLOptions();
   registerPassManagerCLOptions();
   PassPipelineCLParser passPipeline("", "Compiler passes to run");
 
   // Parse pass names in main to ensure static initialization completed.
   cl::ParseCommandLineOptions(argc, argv, "MLIR modular optimizer driver\n");
 
-  if(showDialects) {
-    llvm::outs() << "Registered Dialects:\n";
+  if (showDialects) {
     MLIRContext context;
-    for(Dialect *dialect : context.getRegisteredDialects()) {
-      llvm::outs() << dialect->getNamespace() << "\n";
+    llvm::outs() << "Available Dialects:\n";
+    for (StringRef name : context.getAvailableDialects()) {
+      llvm::outs() << name << "\n";
     }
     return 0;
   }
 
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
#endif
 
   if (failed(MlirOptMain(output->os(), std::move(file), passPipeline,
                          splitInputFile, verifyDiagnostics, verifyPasses,
                          allowUnregisteredDialects))) {
     return 1;
   }
   // Keep the output file if the invocation of MlirOptMain was successful.
   output->keep();
+
   return 0;
 }

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));
   }
 
+  void getDependentDialects(DialectRegistry &registry) const override {
+    {3}
+  }
+
 protected:
 )";
 
+const char *const dialectRegistrationTemplate = R"(
+  registry.add<{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 dialects_os(dependentDialectRegistrations);
+    for (StringRef dependentDialect : pass.getDependentDialects())
+      dialects_os << 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-vulkan-runner/mlir-vulkan-runner.cpp

 #include "llvm/Support/TargetSelect.h"
 
 using namespace mlir;
 
 static LogicalResult runMLIRPasses(ModuleOp module) {
   PassManager passManager(module.getContext());
   applyPassManagerCLOptions(passManager);
 
+  module.getContext()->getOrLoadDialect<spirv::SPIRVDialect>();
+
   passManager.addPass(createGpuKernelOutliningPass());
   passManager.addPass(createLegalizeStdOpsForSPIRVLoweringPass());
   passManager.addPass(createConvertGPUToSPIRVPass());
   OpPassManager &modulePM = passManager.nest<spirv::ModuleOp>();
   modulePM.addPass(spirv::createLowerABIAttributesPass());
   modulePM.addPass(spirv::createUpdateVersionCapabilityExtensionPass());
   passManager.addPass(createConvertGpuLaunchFuncToVulkanLaunchFuncPass());
   LowerToLLVMOptions llvmOptions = {

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;
+  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;
+  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;
+  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;
+  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() {
+    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.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/DialectTest.cpp

AnotherTestDialect(MLIRContext *context)
                 TypeID::get<AnotherTestDialect>()) {}
 };
 
 TEST(DialectDeathTest, MultipleDialectsWithSameNamespace) {
   MLIRContext context;
 
   // Registering a dialect with the same namespace twice should result in a
   // failure.
-  context.getOrCreateDialect<TestDialect>();
-  ASSERT_DEATH(context.getOrCreateDialect<AnotherTestDialect>(), "");
+  context.getOrLoadDialect<TestDialect>();
+  ASSERT_DEATH(context.getOrLoadDialect<AnotherTestDialect>(), "");
 }
 
 } // end namespace

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;
+  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;
+  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 verifyOp(OpTy &&concreteOp, std::vector<Type> resultTypes,
     EXPECT_EQ(op->getAttrs().size(), attrs.size());
     for (unsigned idx : llvm::seq<unsigned>(0U, attrs.size()))
       EXPECT_EQ(op->getAttr(attrs[idx].first.strref()), attrs[idx].second);
 
     concreteOp.erase();
   }
 
 protected:
-  MLIRContext ctx;
+  MLIRContext &ctx;
   OpBuilder builder;
   Location loc;
   Type i32Ty;
   Type f32Ty;
   TableGenConstant cstI32;
   TableGenConstant cstF32;
 
   ArrayRef<NamedAttribute> noAttrs;