[PATCH] D79812: [MLIR] Add symbol map to mlir ExecutionEngine

[Eugene Zhulenev via Phabricator via llvm-commits]

ezhulenev created this revision.
ezhulenev added a reviewer: rriddle.
Herald added subscribers: llvm-commits, Kayjukh, frgossen, grosul1, Joonsoo, stephenneuendorffer, liufengdb, lucyrfox, mgester, arpith-jacob, nicolasvasilache, antiagainst, shauheen, jpienaar, mehdi_amini.
Herald added a project: LLVM.
ezhulenev added a reviewer: ftynse.
ezhulenev updated this revision to Diff 263534.
ezhulenev added a comment.

Fix typo.


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


Repository:
  rG LLVM Github Monorepo

https://reviews.llvm.org/D79812

Files:
  mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
  mlir/lib/ExecutionEngine/ExecutionEngine.cpp


Index: mlir/lib/ExecutionEngine/ExecutionEngine.cpp
===================================================================
--- mlir/lib/ExecutionEngine/ExecutionEngine.cpp
+++ mlir/lib/ExecutionEngine/ExecutionEngine.cpp
@@ -196,8 +196,9 @@
 Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
     ModuleOp m, std::function<Error(llvm::Module *)> transformer,
     Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel,
-    ArrayRef<StringRef> sharedLibPaths, bool enableObjectCache,
-    bool enableGDBNotificationListener, bool enablePerfNotificationListener) {
+    ArrayRef<StringRef> sharedLibPaths, SymbolMap symbolMap,
+    bool enableObjectCache, bool enableGDBNotificationListener,
+    bool enablePerfNotificationListener) {
   auto engine = std::make_unique<ExecutionEngine>(
       enableObjectCache, enableGDBNotificationListener,
       enablePerfNotificationListener);
@@ -288,6 +289,14 @@
       cantFail(DynamicLibrarySearchGenerator::GetForCurrentProcess(
           dataLayout.getGlobalPrefix())));
 
+  // Add additional symbols provided via the symbol map.
+  llvm::orc::MangleAndInterner Mangle(mainJD.getExecutionSession(), dataLayout);
+  llvm::orc::SymbolMap mangledSymbolMap;
+  for (auto &kv : symbolMap) {
+    mangledSymbolMap[Mangle(kv.first)] = std::move(kv.second);
+  }
+  cantFail(mainJD.define(absoluteSymbols(mangledSymbolMap)));
+
   return std::move(engine);
 }
 
Index: mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
===================================================================
--- mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
+++ mlir/include/mlir/ExecutionEngine/ExecutionEngine.h
@@ -63,23 +63,38 @@
   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
+  /// Unmangled symbol names to JIT evaluated symbol mapping. Symbols will be
+  /// mangled before adding to the jit engine.
+  using SymbolMap = llvm::DenseMap<StringRef, llvm::JITEvaluatedSymbol>;
+
+  /// 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 `symbolMap` is not empty, the underlying JIT-compilation will use it
+  /// 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 = {},
          Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel = llvm::None,
-         ArrayRef<StringRef> sharedLibPaths = {}, bool enableObjectCache = true,
+         ArrayRef<StringRef> sharedLibPaths = {},
+         SymbolMap symbolMap = SymbolMap(), bool enableObjectCache = true,
          bool enableGDBNotificationListener = true,
          bool enablePerfNotificationListener = true);
 


-------------- next part --------------
A non-text attachment was scrubbed...
Name: D79812.263534.patch
Type: text/x-patch
Size: 4406 bytes
Desc: not available
URL: <http://lists.llvm.org/pipermail/llvm-commits/attachments/20200512/69af6cb9/attachment.bin>