[mlir] JitRunner: add a config option to register symbols with ExecutionEngine at runtime

[ezhulenev]

Update mlir-spirv-cpu-runner.cpp

mlir/include/mlir/ExecutionEngine/JitRunner.h

 // The translation can be customized by providing an MLIR to MLIR
 // transformation.
 //
 //===----------------------------------------------------------------------===//
 
 #ifndef MLIR_SUPPORT_JITRUNNER_H_
 #define MLIR_SUPPORT_JITRUNNER_H_
 
-#include "mlir/IR/Module.h"
-
 #include "llvm/ADT/STLExtras.h"
-#include "llvm/IR/Module.h"
+#include "llvm/ExecutionEngine/Orc/Core.h"
 
-namespace mlir {
+namespace llvm {
+class Module;
+class LLVMContext;
+
+namespace orc {
+class MangleAndInterner;
+} // namespace orc
+} // namespace llvm
 
-using TranslationCallback = llvm::function_ref<std::unique_ptr<llvm::Module>(
-    ModuleOp, llvm::LLVMContext &)>;
+namespace mlir {
 
 class ModuleOp;
 struct LogicalResult;
 
-// Entry point for all CPU runners. Expects the common argc/argv arguments for
-// standard C++ main functions, `mlirTransformer` and `llvmModuleBuilder`.
-/// `mlirTransformer` is applied after parsing the input into MLIR IR and before
+struct JitRunnerConfig {
+  /// MLIR transformer applied after parsing the input into MLIR IR and before
   /// passing the MLIR module to the ExecutionEngine.
-/// `llvmModuleBuilder` is a custom function that is passed to ExecutionEngine.
-/// It processes MLIR module and creates LLVM IR module.
-int JitRunnerMain(
-    int argc, char **argv,
-    llvm::function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer,
-    TranslationCallback llvmModuleBuilder = nullptr);
+  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 = {});
 
 } // namespace mlir
 
 #endif // MLIR_SUPPORT_JITRUNNER_H_

mlir/lib/ExecutionEngine/JitRunner.cpp

llvm::cl::opt<bool> dumpObjectFile{
       "dump-object-file",
       llvm::cl::desc("Dump JITted-compiled object to file specified with "
                      "-object-filename (<input file>.o by default).")};
 
   llvm::cl::opt<std::string> objectFilename{
       "object-filename",
       llvm::cl::desc("Dump JITted-compiled object to file <input file>.o")};
 };
+
+struct CompileAndExecuteConfig {
+  /// LLVM module transformer that is passed to ExecutionEngine.
+  llvm::function_ref<llvm::Error(llvm::Module *)> transformer;
+
+  /// 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;
+
+  /// A custom function that is passed to ExecutinEngine to register symbols at
+  /// runtime.
+  llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
+      runtimeSymbolMap;
+};
+
 } // end anonymous namespace
 
 static OwningModuleRef parseMLIRInput(StringRef inputFilename,
                                       MLIRContext *context) {
   // Set up the input file.
   std::string errorMessage;
   auto file = openInputFile(inputFilename, &errorMessage);
   if (!file) {
if (flag) {
       optLevel = j;
       break;
     }
   }
   return optLevel;
 }
 
 // JIT-compile the given module and run "entryPoint" with "args" as arguments.
-static Error
-compileAndExecute(Options &options, ModuleOp module,
-                  TranslationCallback llvmModuleBuilder, StringRef entryPoint,
-                  std::function<llvm::Error(llvm::Module *)> transformer,
-                  void **args) {
+static Error compileAndExecute(Options &options, ModuleOp module,
+                               StringRef entryPoint,
+                               CompileAndExecuteConfig config, void **args) {
   Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
   if (auto clOptLevel = getCommandLineOptLevel(options))
     jitCodeGenOptLevel =
         static_cast<llvm::CodeGenOpt::Level>(clOptLevel.getValue());
   SmallVector<StringRef, 4> libs(options.clSharedLibs.begin(),
                                  options.clSharedLibs.end());
   auto expectedEngine = mlir::ExecutionEngine::create(
-      module, llvmModuleBuilder, transformer, jitCodeGenOptLevel, libs);
+      module, config.llvmModuleBuilder, config.transformer, jitCodeGenOptLevel,
+      libs);
   if (!expectedEngine)
     return expectedEngine.takeError();
 
   auto engine = std::move(*expectedEngine);
+  if (config.runtimeSymbolMap)
+    engine->registerSymbols(config.runtimeSymbolMap);
+
   auto expectedFPtr = engine->lookup(entryPoint);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
 
   if (options.dumpObjectFile)
     engine->dumpToObjectFile(options.objectFilename.empty()
                                  ? options.inputFilename + ".o"
                                  : options.objectFilename);
 
   void (*fptr)(void **) = *expectedFPtr;
   (*fptr)(args);
 
   return Error::success();
 }
 
-static Error compileAndExecuteVoidFunction(
-    Options &options, ModuleOp module, TranslationCallback llvmModuleBuilder,
+static Error compileAndExecuteVoidFunction(Options &options, ModuleOp module,
                                            StringRef entryPoint,
-    std::function<llvm::Error(llvm::Module *)> transformer) {
+                                           CompileAndExecuteConfig config) {
   auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
   if (!mainFunction || mainFunction.empty())
     return make_string_error("entry point not found");
   void *empty = nullptr;
-  return compileAndExecute(options, module, llvmModuleBuilder, entryPoint,
-                           transformer, &empty);
+  return compileAndExecute(options, module, entryPoint, config, &empty);
 }
 
 template <typename Type>
 Error checkCompatibleReturnType(LLVM::LLVMFuncOp mainFunction);
 template <>
 Error checkCompatibleReturnType<int32_t>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getType().getFunctionResultType().isIntegerTy(32))
     return make_string_error("only single llvm.i32 function result supported");
   return Error::success();
 }
 template <>
 Error checkCompatibleReturnType<int64_t>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getType().getFunctionResultType().isIntegerTy(64))
     return make_string_error("only single llvm.i64 function result supported");
   return Error::success();
 }
 template <>
 Error checkCompatibleReturnType<float>(LLVM::LLVMFuncOp mainFunction) {
   if (!mainFunction.getType().getFunctionResultType().isFloatTy())
     return make_string_error("only single llvm.f32 function result supported");
   return Error::success();
 }
 template <typename Type>
-Error compileAndExecuteSingleReturnFunction(
-    Options &options, ModuleOp module, TranslationCallback llvmModuleBuilder,
+Error compileAndExecuteSingleReturnFunction(Options &options, ModuleOp module,
                                             StringRef entryPoint,
-    std::function<llvm::Error(llvm::Module *)> transformer) {
+                                            CompileAndExecuteConfig config) {
   auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
   if (!mainFunction || mainFunction.isExternal())
     return make_string_error("entry point not found");
 
   if (mainFunction.getType().getFunctionNumParams() != 0)
     return make_string_error("function inputs not supported");
 
   if (Error error = checkCompatibleReturnType<Type>(mainFunction))
     return error;
 
   Type res;
   struct {
     void *data;
   } data;
   data.data = &res;
-  if (auto error = compileAndExecute(options, module, llvmModuleBuilder,
-                                     entryPoint, transformer, (void **)&data))
+  if (auto error = compileAndExecute(options, module, entryPoint, config,
+                                     (void **)&data))
     return error;
 
   // 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, `mlirTransformer` and `llvmModuleBuilder`.
-/// `mlirTransformer` is applied after parsing the input into MLIR IR and before
-/// passing the MLIR module to the ExecutionEngine.
-/// `llvmModuleBuilder` is a custom function that is passed to ExecutionEngine.
-/// It processes MLIR module and creates LLVM IR module.
-int mlir::JitRunnerMain(
-    int argc, char **argv,
-    function_ref<LogicalResult(mlir::ModuleOp)> mlirTransformer,
-    TranslationCallback llvmModuleBuilder) {
+/// standard C++ main functions.
+int mlir::JitRunnerMain(int argc, char **argv, 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, JitRunnerConfig config) {
   registerAllDialects(context.getDialectRegistry());
 
   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())))
+  if (config.mlirTransformer)
+    if (failed(config.mlirTransformer(m.get())))
       return EXIT_FAILURE;
 
   auto tmBuilderOrError = llvm::orc::JITTargetMachineBuilder::detectHost();
   if (!tmBuilderOrError) {
     llvm::errs() << "Failed to create a JITTargetMachineBuilder for the host\n";
     return EXIT_FAILURE;
   }
   auto tmOrError = tmBuilderOrError->createTargetMachine();
   if (!tmOrError) {
     llvm::errs() << "Failed to create a TargetMachine for the host\n";
     return EXIT_FAILURE;
   }
 
   auto transformer = mlir::makeLLVMPassesTransformer(
       passes, optLevel, /*targetMachine=*/tmOrError->get(), optPosition);
 
+  CompileAndExecuteConfig compileAndExecuteConfig;
+  compileAndExecuteConfig.transformer = transformer;
+  compileAndExecuteConfig.llvmModuleBuilder = config.llvmModuleBuilder;
+  compileAndExecuteConfig.runtimeSymbolMap = config.runtimesymbolMap;
+
   // Get the function used to compile and execute the module.
   using CompileAndExecuteFnT =
-      Error (*)(Options &, ModuleOp, TranslationCallback, StringRef,
-                std::function<llvm::Error(llvm::Module *)>);
+      Error (*)(Options &, ModuleOp, StringRef, CompileAndExecuteConfig);
   auto compileAndExecuteFn =
       StringSwitch<CompileAndExecuteFnT>(options.mainFuncType.getValue())
           .Case("i32", compileAndExecuteSingleReturnFunction<int32_t>)
           .Case("i64", compileAndExecuteSingleReturnFunction<int64_t>)
           .Case("f32", compileAndExecuteSingleReturnFunction<float>)
           .Case("void", compileAndExecuteVoidFunction)
           .Default(nullptr);
 
-  Error error =
-      compileAndExecuteFn
-          ? compileAndExecuteFn(options, m.get(), llvmModuleBuilder,
-                                options.mainFuncName.getValue(), transformer)
+  Error error = compileAndExecuteFn
+                    ? compileAndExecuteFn(options, m.get(),
+                                          options.mainFuncName.getValue(),
+                                          compileAndExecuteConfig)
                     : make_string_error("unsupported function type");
 
   int exitCode = EXIT_SUCCESS;
   llvm::handleAllErrors(std::move(error),
                         [&exitCode](const llvm::ErrorInfoBase &info) {
                           llvm::errs() << "Error: ";
                           info.log(llvm::errs());
                           llvm::errs() << '\n';
                           exitCode = EXIT_FAILURE;
                         });
 
   return exitCode;
 }

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

 #include "llvm/Support/TargetSelect.h"
 
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   mlir::initializeLLVMPasses();
 
-  return mlir::JitRunnerMain(argc, argv, nullptr);
+  return mlir::JitRunnerMain(argc, argv);
 }

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

int main(int argc, char **argv) {
 
   // Initialize LLVM NVPTX backend.
   LLVMInitializeNVPTXTarget();
   LLVMInitializeNVPTXTargetInfo();
   LLVMInitializeNVPTXTargetMC();
   LLVMInitializeNVPTXAsmPrinter();
 
   mlir::initializeLLVMPasses();
-  return mlir::JitRunnerMain(argc, argv, &runMLIRPasses);
+
+  mlir::JitRunnerConfig jitRunnerConfig;
+  jitRunnerConfig.mlirTransformer = &runMLIRPasses;
+
+  return mlir::JitRunnerMain(argc, argv, jitRunnerConfig);
 }

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

 
 int main(int argc, char **argv) {
   llvm::InitLLVM y(argc, argv);
 
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   mlir::initializeLLVMPasses();
 
-  return mlir::JitRunnerMain(argc, argv, &runMLIRPasses, &convertMLIRModule);
+  mlir::JitRunnerConfig jitRunnerConfig;
+  jitRunnerConfig.mlirTransformer = &runMLIRPasses;
+  jitRunnerConfig.llvmModuleBuilder = &convertMLIRModule;
+
+  return mlir::JitRunnerMain(argc, argv, jitRunnerConfig);
 }

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

int main(int argc, char **argv) {
   llvm::llvm_shutdown_obj x;
   registerPassManagerCLOptions();
 
   llvm::InitLLVM y(argc, argv);
   llvm::InitializeNativeTarget();
   llvm::InitializeNativeTargetAsmPrinter();
   mlir::initializeLLVMPasses();
 
-  return mlir::JitRunnerMain(argc, argv, &runMLIRPasses);
+  mlir::JitRunnerConfig jitRunnerConfig;
+  jitRunnerConfig.mlirTransformer = &runMLIRPasses;
+
+  return mlir::JitRunnerMain(argc, argv, jitRunnerConfig);
 }