[MLIR] Add support for resolving symbols to in-process functions for ExecutionEngine.

[timshen]

Summary

This comes in handy for calling into C/C++ implementations defined by the caller of the jit.

[mehdi_amini]

How does it compare to ExecutionEngine::registerSymbols ?

[timshen]

You are right, registerSymbols works.

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

create(ModuleOp m,
          ArrayRef<StringRef> sharedLibPaths = {}, 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;
 
+  /// Maps a symbol to an in-process address.
+  llvm::Error addInProcessSymbol(StringRef name, void *address);
+
   /// Invokes the function with the given name passing it the list of arguments
   /// as a list of opaque pointers.
   llvm::Error invoke(StringRef name, MutableArrayRef<void *> args = llvm::None);
 
   /// Set the target triple on the module. This is implicitly done when creating
   /// the engine.
   static bool setupTargetTriple(llvm::Module *llvmModule);
 

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

Expected<void (*)(void **)> ExecutionEngine::lookup(StringRef name) const {
 
   auto rawFPtr = expectedSymbol->getAddress();
   auto fptr = reinterpret_cast<void (*)(void **)>(rawFPtr);
   if (!fptr)
     return make_string_error("looked up function is null");
   return fptr;
 }
 
+llvm::Error ExecutionEngine::addInProcessSymbol(StringRef name, void *address) {
+  return jit->getMainJITDylib().define(llvm::orc::absoluteSymbols(
+      {{jit->getExecutionSession().intern(name),
+        llvm::JITEvaluatedSymbol::fromPointer(address)}}));
+}
+
 Error ExecutionEngine::invoke(StringRef name, MutableArrayRef<void *> args) {
   auto expectedFPtr = lookup(name);
   if (!expectedFPtr)
     return expectedFPtr.takeError();
   auto fptr = *expectedFPtr;
 
   (*fptr)(args.data());
 
   return Error::success();
 }