[mlir] Make JitRunnerMain main take a DialectRegistry

Historically, JitRunner has been registering all available dialects with the
context and depending on them without the real need. Make it take a registry
that contains only the dialects that are expected in the input and stop linking
in all dialects.

Reviewed By: mehdi_amini

Differential Revision: https://reviews.llvm.org/D96436

Author: ftynse

mlir/include/mlir/ExecutionEngine/JitRunner.h

@@ -32,6 +32,7 @@ class MangleAndInterner;
 
 namespace mlir {
 
+class DialectRegistry;
 class ModuleOp;
 struct LogicalResult;
 
@@ -51,9 +52,12 @@ struct JitRunnerConfig {
       runtimesymbolMap = nullptr;
 };
 
-// Entry point for all CPU runners. Expects the common argc/argv arguments for
-// standard C++ main functions.
-int JitRunnerMain(int argc, char **argv, JitRunnerConfig config = {});
+/// Entry point for all CPU runners. Expects the common argc/argv arguments for
+/// standard C++ main functions. The supplied dialect registry is expected to
+/// contain any registers that appear in the input IR, they will be loaded
+/// on-demand by the parser.
+int JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
+                  JitRunnerConfig config = {});
 
 } // namespace mlir
 

mlir/lib/ExecutionEngine/JitRunner.cpp

@@ -21,7 +21,6 @@
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/MLIRContext.h"
-#include "mlir/InitAllDialects.h"
 #include "mlir/Parser.h"
 #include "mlir/Support/FileUtilities.h"
 
@@ -299,7 +298,8 @@ Error compileAndExecuteSingleReturnFunction(Options &options, ModuleOp module,
 
 /// Entry point for all CPU runners. Expects the common argc/argv arguments for
 /// standard C++ main functions.
-int mlir::JitRunnerMain(int argc, char **argv, JitRunnerConfig config) {
+int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
+                        JitRunnerConfig config) {
   // Create the options struct containing the command line options for the
   // runner. This must come before the command line options are parsed.
   Options options;
@@ -330,8 +330,7 @@ int mlir::JitRunnerMain(int argc, char **argv, JitRunnerConfig config) {
     }
   }
 
-  MLIRContext context;
-  registerAllDialects(context);
+  MLIRContext context(registry);
 
   auto m = parseMLIRInput(options.inputFilename, &context);
   if (!m) {

mlir/tools/mlir-cpu-runner/CMakeLists.txt

@@ -9,15 +9,14 @@ add_llvm_tool(mlir-cpu-runner
   mlir-cpu-runner.cpp
   )
 llvm_update_compile_flags(mlir-cpu-runner)
-get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 target_link_libraries(mlir-cpu-runner PRIVATE
-  ${dialect_libs}
   MLIRAnalysis
   MLIREDSC
   MLIRExecutionEngine
   MLIRIR
   MLIRJitRunner
   MLIRLLVMIR
+  MLIROpenMP
   MLIRParser
   MLIRTargetLLVMIR
   MLIRSupport

mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp

@@ -12,9 +12,11 @@
 //
 //===----------------------------------------------------------------------===//
 
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/OpenMP/OpenMPDialect.h"
 #include "mlir/ExecutionEngine/JitRunner.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
-#include "mlir/InitAllDialects.h"
+#include "mlir/IR/Dialect.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
@@ -25,5 +27,7 @@ int main(int argc, char **argv) {
   llvm::InitializeNativeTargetAsmParser();
   mlir::initializeLLVMPasses();
 
-  return mlir::JitRunnerMain(argc, argv);
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::LLVM::LLVMDialect, mlir::omp::OpenMPDialect>();
+  return mlir::JitRunnerMain(argc, argv, registry);
 }

mlir/tools/mlir-cuda-runner/CMakeLists.txt

@@ -42,17 +42,20 @@ if(MLIR_CUDA_RUNNER_ENABLED)
     ${CUDA_RUNTIME_LIBRARY}
   )
 
-  get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
   get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
   set(LIBS
-    ${dialect_libs}
     ${conversion_libs}
     MLIRJitRunner
     MLIRAnalysis
+    MLIRAsync
     MLIREDSC
     MLIRExecutionEngine
+    MLIRGPU
     MLIRIR
+    MLIRLLVMIR
+    MLIRNVVMIR
     MLIRParser
+    MLIRStandard
     MLIRSupport
     MLIRTargetLLVMIR
     MLIRTargetNVVMIR

mlir/tools/mlir-cuda-runner/mlir-cuda-runner.cpp

@@ -19,15 +19,16 @@
 #include "mlir/Conversion/GPUToNVVM/GPUToNVVMPass.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/Async/IR/Async.h"
 #include "mlir/Dialect/Async/Passes.h"
 #include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/NVVMDialect.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/ExecutionEngine/JitRunner.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/InitAllDialects.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Target/NVVMIR.h"
@@ -149,5 +150,10 @@ int main(int argc, char **argv) {
   mlir::JitRunnerConfig jitRunnerConfig;
   jitRunnerConfig.mlirTransformer = runMLIRPasses;
 
-  return mlir::JitRunnerMain(argc, argv, jitRunnerConfig);
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::LLVM::LLVMDialect, mlir::NVVM::NVVMDialect,
+                  mlir::async::AsyncDialect, mlir::gpu::GPUDialect,
+                  mlir::StandardOpsDialect>();
+
+  return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
 }

mlir/tools/mlir-rocm-runner/CMakeLists.txt

@@ -65,10 +65,8 @@ if(MLIR_ROCM_RUNNER_ENABLED)
     ${ROCM_RUNTIME_LIBRARY}
   )
 
-  get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
   get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
   set(LIBS
-    ${dialect_libs}
     ${conversion_libs}
     lldCommon
     lldDriver
@@ -77,9 +75,12 @@ if(MLIR_ROCM_RUNNER_ENABLED)
     MLIRAnalysis
     MLIREDSC
     MLIRExecutionEngine
+    MLIRGPU
     MLIRIR
+    MLIRLLVMIR
     MLIRParser
     MLIRROCDLIR
+    MLIRStandard
     MLIRSupport
     MLIRTargetLLVMIR
     MLIRTargetROCDLIR

mlir/tools/mlir-rocm-runner/mlir-rocm-runner.cpp

@@ -23,10 +23,10 @@
 #include "mlir/Dialect/GPU/Passes.h"
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/Dialect/LLVMIR/ROCDLDialect.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/ExecutionEngine/JitRunner.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
 #include "mlir/IR/BuiltinOps.h"
-#include "mlir/InitAllDialects.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Support/FileUtilities.h"
@@ -337,5 +337,9 @@ int main(int argc, char **argv) {
   mlir::JitRunnerConfig jitRunnerConfig;
   jitRunnerConfig.mlirTransformer = runMLIRPasses;
 
-  return mlir::JitRunnerMain(argc, argv, jitRunnerConfig);
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::LLVM::LLVMDialect, mlir::gpu::GPUDialect,
+                  mlir::ROCDL::ROCDLDialect, mlir::StandardOpsDialect>();
+
+  return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
 }

mlir/tools/mlir-spirv-cpu-runner/CMakeLists.txt

@@ -12,18 +12,19 @@ if (MLIR_SPIRV_CPU_RUNNER_ENABLED)
   llvm_update_compile_flags(mlir-spirv-cpu-runner)
 
   get_property(conversion_libs GLOBAL PROPERTY MLIR_CONVERSION_LIBS)
-  get_property(dialect_libs GLOBAL PROPERTY MLIR_DIALECT_LIBS)
 
   target_link_libraries(mlir-spirv-cpu-runner PRIVATE
     ${conversion_libs}
-    ${dialect_libs}
     MLIRAnalysis
     MLIREDSC
     MLIRExecutionEngine
+    MLIRGPU
     MLIRIR
     MLIRJitRunner
     MLIRLLVMIR
     MLIRParser
+    MLIRSPIRV
+    MLIRStandard
     MLIRTargetLLVMIR
     MLIRTransforms
     MLIRTranslation

mlir/tools/mlir-spirv-cpu-runner/mlir-spirv-cpu-runner.cpp

@@ -15,12 +15,15 @@
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/SPIRVToLLVM/SPIRVToLLVMPass.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
+#include "mlir/Dialect/GPU/GPUDialect.h"
 #include "mlir/Dialect/GPU/Passes.h"
+#include "mlir/Dialect/LLVMIR/LLVMDialect.h"
+#include "mlir/Dialect/SPIRV/IR/SPIRVDialect.h"
 #include "mlir/Dialect/SPIRV/IR/SPIRVOps.h"
 #include "mlir/Dialect/SPIRV/Transforms/Passes.h"
+#include "mlir/Dialect/StandardOps/IR/Ops.h"
 #include "mlir/ExecutionEngine/JitRunner.h"
 #include "mlir/ExecutionEngine/OptUtils.h"
-#include "mlir/InitAllDialects.h"
 #include "mlir/Pass/Pass.h"
 #include "mlir/Pass/PassManager.h"
 #include "mlir/Target/LLVMIR.h"
@@ -90,5 +93,9 @@ int main(int argc, char **argv) {
   jitRunnerConfig.mlirTransformer = runMLIRPasses;
   jitRunnerConfig.llvmModuleBuilder = convertMLIRModule;
 
-  return mlir::JitRunnerMain(argc, argv, jitRunnerConfig);
+  mlir::DialectRegistry registry;
+  registry.insert<mlir::LLVM::LLVMDialect, mlir::gpu::GPUDialect,
+                  mlir::spirv::SPIRVDialect, mlir::StandardOpsDialect>();
+
+  return mlir::JitRunnerMain(argc, argv, registry, jitRunnerConfig);
 }