[mlir] Make JitRunnerMain main take a DialectRegistry

[ftynse]

Summary

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.

mlir/include/mlir/ExecutionEngine/JitRunner.h

 
 namespace orc {
 class MangleAndInterner;
 } // namespace orc
 } // namespace llvm
 
 namespace mlir {
 
+class DialectRegistry;
 class ModuleOp;
 struct LogicalResult;
 
 struct JitRunnerConfig {
   /// MLIR transformer applied after parsing the input into MLIR IR and before
   /// passing the MLIR module to the ExecutionEngine.
   llvm::function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer = nullptr;
 
   /// A custom function that is passed to ExecutionEngine. It processes MLIR
   /// module and creates LLVM IR module.
   llvm::function_ref<std::unique_ptr<llvm::Module>(ModuleOp,
                                                    llvm::LLVMContext &)>
       llvmModuleBuilder = nullptr;
 
   /// A callback to register symbols with ExecutionEngine at runtime.
   llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
       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
 
 #endif // MLIR_SUPPORT_JITRUNNER_H_

mlir/lib/ExecutionEngine/JitRunner.cpp

 
 #include "mlir/ExecutionEngine/JitRunner.h"
 
 #include "mlir/Dialect/LLVMIR/LLVMDialect.h"
 #include "mlir/ExecutionEngine/ExecutionEngine.h"
 #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"
 
 #include "llvm/ADT/STLExtras.h"
 #include "llvm/ExecutionEngine/Orc/JITTargetMachineBuilder.h"
 #include "llvm/IR/IRBuilder.h"
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/LegacyPassNameParser.h"
Error compileAndExecuteSingleReturnFunction(Options &options, ModuleOp module,
   // Intentional printing of the output so we can test.
   llvm::outs() << res << '\n';
 
   return Error::success();
 }
 
 /// 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;
   llvm::cl::ParseCommandLineOptions(argc, argv, "MLIR CPU execution driver\n");
 
   Optional<unsigned> optLevel = getCommandLineOptLevel(options);
   SmallVector<std::reference_wrapper<llvm::cl::opt<bool>>, 4> optFlags{
       options.optO0, options.optO1, options.optO2, options.optO3};
int mlir::JitRunnerMain(int argc, char **argv, const DialectRegistry &registry,
   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);
+  MLIRContext context(registry);
 
   auto m = parseMLIRInput(options.inputFilename, &context);
   if (!m) {
     llvm::errs() << "could not parse the input IR\n";
     return 1;
   }
 
   if (config.mlirTransformer)

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

 set(LLVM_LINK_COMPONENTS
   Core
   Support
   nativecodegen
   native
   )
 
 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

 //===- mlir-cpu-runner.cpp - MLIR CPU Execution Driver---------------------===//
 //
 // 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
 //
 //===----------------------------------------------------------------------===//
 //
 // Main entry point to a command line utility that executes an MLIR file on the
 // CPU by  translating MLIR to LLVM IR before JIT-compiling and executing the
 // latter.
 //
 //===----------------------------------------------------------------------===//
 
+#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"
 
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   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

target_include_directories(cuda-runtime-wrappers
     PRIVATE
     ${CMAKE_CUDA_TOOLKIT_INCLUDE_DIRECTORIES}
   )
   target_link_libraries(cuda-runtime-wrappers
     PRIVATE
     ${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
     MLIRTransforms
     MLIRTranslation
     ${CUDA_RUNTIME_LIBRARY}
   )
 

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

 
 #include "llvm/ADT/STLExtras.h"
 
 #include "mlir/Conversion/AsyncToLLVM/AsyncToLLVM.h"
 #include "mlir/Conversion/GPUCommon/GPUCommonPass.h"
 #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"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
int main(int argc, char **argv) {
   LLVMInitializeNVPTXTargetMC();
   LLVMInitializeNVPTXAsmPrinter();
 
   mlir::initializeLLVMPasses();
 
   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

target_include_directories(rocm-runtime-wrappers
     "${HIP_PATH}/../include"
     "${HIP_PATH}/include"
   )
   target_link_libraries(rocm-runtime-wrappers
     PRIVATE
     ${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
     lldELF
     MLIRJitRunner
     MLIRAnalysis
     MLIREDSC
     MLIRExecutionEngine
+    MLIRGPU
     MLIRIR
+    MLIRLLVMIR
     MLIRParser
     MLIRROCDLIR
+    MLIRStandard
     MLIRSupport
     MLIRTargetLLVMIR
     MLIRTargetROCDLIR
     MLIRTransforms
     MLIRTranslation
     ${ROCM_RUNTIME_LIBRARY}
   )
 

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

 #include "mlir/Conversion/GPUToROCDL/GPUToROCDLPass.h"
 #include "mlir/Conversion/SCFToStandard/SCFToStandard.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVM.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/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"
 #include "mlir/Target/ROCDLIR.h"
 #include "mlir/Transforms/DialectConversion.h"
 #include "mlir/Transforms/Passes.h"
 #include "llvm/Support/ErrorOr.h"
 #include "llvm/Support/FileUtilities.h"
int main(int argc, char **argv) {
   LLVMInitializeAMDGPUTargetMC();
   LLVMInitializeAMDGPUAsmPrinter();
 
   mlir::initializeLLVMPasses();
 
   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

 set(LLVM_LINK_COMPONENTS
   Linker
 )
 
 if (MLIR_SPIRV_CPU_RUNNER_ENABLED)
   message(STATUS "Building SPIR-V CPU runner")
 
   add_llvm_tool(mlir-spirv-cpu-runner
     mlir-spirv-cpu-runner.cpp
   )
 
   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
     MLIRSupport
   )
 endif()

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

 // CPU by translating MLIR GPU module and host part to LLVM IR before
 // JIT-compiling and executing.
 //
 //===----------------------------------------------------------------------===//
 
 #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"
 
 #include "llvm/IR/LLVMContext.h"
 #include "llvm/IR/Module.h"
 #include "llvm/Linker/Linker.h"
 #include "llvm/Support/InitLLVM.h"
int main(int argc, char **argv) {
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   mlir::initializeLLVMPasses();
 
   mlir::JitRunnerConfig jitRunnerConfig;
   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);
 }

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

if (MLIR_VULKAN_RUNNER_ENABLED)
   # it may cause LLVM version conflict when used together with other shared
   # libraries depending on LLVM. Notably, Mesa, who implements Vulkan
   # drivers on Linux, depends on the system libLLVM.so.
   target_link_libraries(vulkan-runtime-wrappers
     PUBLIC
     ${Vulkan_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}
     MLIRAnalysis
     MLIREDSC
     MLIRExecutionEngine
+    MLIRGPU
     MLIRIR
     MLIRJitRunner
     MLIRLLVMIR
     MLIRParser
+    MLIRSPIRV
     MLIRSPIRVTransforms
+    MLIRStandard
     MLIRSupport
     MLIRTargetLLVMIR
     MLIRTransforms
     MLIRTranslation
     ${Vulkan_LIBRARY}
   )
 
   # Manually expand the target library, since our MLIR libraries

mlir/tools/mlir-vulkan-runner/mlir-vulkan-runner.cpp

 // JIT-compiling and executing the latter.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/Conversion/GPUToSPIRV/GPUToSPIRVPass.h"
 #include "mlir/Conversion/GPUToVulkan/ConvertGPUToVulkanPass.h"
 #include "mlir/Conversion/StandardToLLVM/ConvertStandardToLLVMPass.h"
 #include "mlir/Conversion/StandardToSPIRV/StandardToSPIRVPass.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 "llvm/Support/InitLLVM.h"
 #include "llvm/Support/TargetSelect.h"
 
 using namespace mlir;
 
 static LogicalResult runMLIRPasses(ModuleOp module) {
int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   mlir::initializeLLVMPasses();
 
   mlir::JitRunnerConfig jitRunnerConfig;
   jitRunnerConfig.mlirTransformer = runMLIRPasses;
 
-  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);
 }