[MLIR] Add symbol map to mlir ExecutionEngine

[ezhulenev]

Summary

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

[ezhulenev]

Fix typo.

[mehdi_amini]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

Copying an entire DenseMap just to read it isn't ideal. Can we pass is as const ref?
Or better allow the client of the API to not be forced into this narrow contract by using a callback to populate the internal llvm::orc::SymbolMap or something like this?

297

Nit: no trivial brace please.

[mehdi_amini]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

Also: isn't something can can be done after creation? With a specific API to add symbols?

[ezhulenev]

Pass symbolMap as a std::function callback.s

[ezhulenev]

Fix typo.

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

Done. I changed DenseMap to std::function callback. In theory it can be done later, after the creation of ExecutionEngine, though I think it makes sense to pass these details at the time of "create".

[mehdi_amini]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

Nit: please use function_ref for callbacks when you don't keep a ref to it.

I think it makes sense to pass these details at the time of "create".

Why? Seems like the kind of thing that makes more sense to me to add "as needed" instead of forcing it on the creation if there is no reason for it.
It also makes the creation API simpler.

[ezhulenev]

Use llvm::function_ref instead of std::function.

[ezhulenev]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

Changed function parameters to llvm::function_ref.

It has kind of the same purpose as sharedLibPaths (register symbols), and that's why it feels right to keep them next to each other, but I agree that create signature grows too big, and something has to be done. It doesn't make sense to me to make this change just for the symbol map, it has to be a new creation API with options struct, or builder, or something else.

[mehdi_amini]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

Right but if we agree that a refactoring of this API would be better, I rather see this *before* adding more to it (incidentally it'll reduce the churn on the users of the API)

[ftynse]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

Alternatively, we can have an instance method ExecutionEngine::registerSymbols that can be called when necessary without touching the ExecutionEngine::create interface (it should be updated regardless, but this patch will be unblocked).

[ezhulenev]

Add registerSymbols() function to mlir ExecutionEngine.

[ezhulenev]

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

199

I've added registerSymbols, not to do a full builder-style rewrite in this patch, because it's the one thing that does not have to be a part of create.

[mehdi_amini]

LG, thanks!

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

123

Nit comment should read "with this ExecutionEngine", "an" is a bit undefined / seems to refer to another object.

[ezhulenev]

Update comment.

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

 /// arguments of the function, followed by a pointer to the result.  This allows
 /// the engine to provide the caller with a generic function pointer that can
 /// be used to invoke the JIT-compiled function.
 class ExecutionEngine {
 public:
   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);
 
   /// Looks up a packed-argument function with the given name and returns a
   /// pointer to it.  Propagates errors in case of failure.
   llvm::Expected<void (*)(void **)> lookup(StringRef name) const;
 
   /// Invokes the function with the given name passing it the list of arguments.
public:
 
   /// Set the target triple on the module. This is implicitly done when creating
   /// the engine.
   static bool setupTargetTriple(llvm::Module *llvmModule);
 
   /// Dump object code to output file `filename`.
   void dumpToObjectFile(StringRef filename);
 
 private:
   /// Ordering of llvmContext and jit is important for destruction purposes: the
   /// jit must be destroyed before the context.
   llvm::LLVMContext llvmContext;
 
   /// Underlying LLJIT.
   std::unique_ptr<llvm::orc::LLJIT> jit;
 
   /// Underlying cache.

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

: cache(enableObjectCache ? new SimpleObjectCache() : nullptr),
                       : nullptr),
       perfListener(enablePerfNotificationListener
                        ? llvm::JITEventListener::createPerfJITEventListener()
                        : nullptr) {}
 
 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);
 
   std::unique_ptr<llvm::LLVMContext> ctx(new llvm::LLVMContext);
   auto llvmModule = translateModuleToLLVMIR(m);
   if (!llvmModule)
     return make_string_error("could not convert to LLVM IR");
Expected<std::unique_ptr<ExecutionEngine>> ExecutionEngine::create(
   engine->jit = std::move(jit);
 
   // Resolve symbols that are statically linked in the current process.
   llvm::orc::JITDylib &mainJD = engine->jit->getMainJITDylib();
   mainJD.addGenerator(
       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);
 }
 
 Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
   auto expectedSymbol = jit->lookup(makePackedFunctionName(name));
 
   // JIT lookup may return an Error referring to strings stored internally by
   // the JIT. If the Error outlives the ExecutionEngine, it would want have a