[MLIR] Add symbol map to mlir ExecutionEngine

Add additional symbol mapping to be able to provide custom symbols to jitted code at runtime.

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

Author: ezhulenev

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

@@ -63,21 +63,28 @@ class ExecutionEngine {
   ExecutionEngine(bool enableObjectCache, bool enableGDBNotificationListener,
                   bool enablePerfNotificationListener);
 
-  /// Creates an execution engine for the given module.  If `transformer` is
-  /// provided, it will be called on the LLVM module during JIT-compilation and
-  /// can be used, e.g., for reporting or optimization. `jitCodeGenOptLevel`,
-  /// when provided, is used as the optimization level for target code
-  /// generation. If `sharedLibPaths` are provided, the underlying
-  /// JIT-compilation will open and link the shared libraries for symbol
-  /// resolution. If `enableObjectCache` is set, the JIT compiler will create
-  /// one to store the object generated for the given module. If enable
-  /// `enableGDBNotificationListener` is set, the JIT compiler will notify
-  /// the llvm's global GDB notification listener. If
-  /// `enablePerfNotificationListener` is set, the JIT compiler will notify
+  /// Creates an execution engine for the given module.
+  ///
+  /// If `transformer` is provided, it will be called on the LLVM module during
+  /// JIT-compilation and can be used, e.g., for reporting or optimization.
+  ///
+  /// `jitCodeGenOptLevel`, when provided, is used as the optimization level for
+  /// target code generation.
+  ///
+  /// If `sharedLibPaths` are provided, the underlying JIT-compilation will
+  /// open and link the shared libraries for symbol resolution.
+  ///
+  /// If `enableObjectCache` is set, the JIT compiler will create one to store
+  /// the object generated for the given module.
+  ///
+  /// If enable `enableGDBNotificationListener` is set, the JIT compiler will
+  /// notify the llvm's global GDB notification listener.
+  ///
+  /// If `enablePerfNotificationListener` is set, the JIT compiler will notify
   /// the llvm's global Perf notification listener.
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
   create(ModuleOp m,
-         std::function<llvm::Error(llvm::Module *)> transformer = {},
+         llvm::function_ref<llvm::Error(llvm::Module *)> transformer = {},
          Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
          ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = true,
          bool enableGDBNotificationListener = true,
@@ -105,6 +112,11 @@ class ExecutionEngine {
   /// Dump object code to output file `filename`.
   void dumpToObjectFile(StringRef filename);
 
+  /// Register symbols with this ExecutionEngine.
+  void registerSymbols(
+      llvm::function_ref<llvm::orc::SymbolMap(llvm::orc::MangleAndInterner)>
+          symbolMap);
+
 private:
   /// Ordering of llvmContext and jit is important for destruction purposes: the
   /// jit must be destroyed before the context.

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

@@ -51,7 +51,9 @@ using llvm::orc::DynamicLibrarySearchGenerator;
 using llvm::orc::ExecutionSession;
 using llvm::orc::IRCompileLayer;
 using llvm::orc::JITTargetMachineBuilder;
+using llvm::orc::MangleAndInterner;
 using llvm::orc::RTDyldObjectLinkingLayer;
+using llvm::orc::SymbolMap;
 using llvm::orc::ThreadSafeModule;
 using llvm::orc::TMOwningSimpleCompiler;
 
@@ -99,6 +101,14 @@ void ExecutionEngine::dumpToObjectFile(StringRef filename) {
   cache->dumpToObjectFile(filename);
 }
 
+void ExecutionEngine::registerSymbols(
+    llvm::function_ref<SymbolMap(MangleAndInterner)> symbolMap) {
+  auto &mainJitDylib = jit->getMainJITDylib();
+  cantFail(mainJitDylib.define(
+      absoluteSymbols(symbolMap(llvm::orc::MangleAndInterner(
+          mainJitDylib.getExecutionSession(), jit->getDataLayout())))));
+}
+
 // Setup LLVM target triple from the current machine.
 bool ExecutionEngine::setupTargetTriple(Module *llvmModule) {
   // Setup the machine properties from the current architecture.
@@ -194,7 +204,7 @@ ExecutionEngine::ExecutionEngine(bool enableObjectCache,
                        : nullptr) {}
 
 Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
-    ModuleOp m, std::function<Error(llvm::Module *)> transformer,
+    ModuleOp m, llvm::function_ref<Error(llvm::Module *)> transformer,
     Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel,
     ArrayRef<StringRef> sharedLibPaths, bool enableObjectCache,
     bool enableGDBNotificationListener, bool enablePerfNotificationListener) {