[mlir] Move symbol loading from mlir-cpu-runner to ExecutionEngine.

[ingomueller-net]

Summary

Both the mlir-cpu-runner and the execution engine allow to provide a
list of shared libraries that should be loaded into the process such
that the jitted code can use the symbols from those libraries. The
runner had implemented a protocol that allowed libraries to control
which symbols it wants to provide in that context (with a function
called __mlir_runner_init). In absence of that, the runner would rely on
the loading mechanism of the execution engine, which didn't do anything
particular with the symbols, i.e., only symbols with public visibility
were visible to jitted code.

Libraries used a mix of the two mechanisms: while the runner utils and C
runner utils libs (and potentially others) used public visibility, the
async runtime lib (as the only one in the monorepo) used the loading
protocol. As a consequence, the async runtime library could not be used
through the Python bindings of the execution engine.

This patch moves the loading protocol from the runner to the execution
engine. For the runner, this should not change anything: it lets the
execution engine handle the loading which now implements the same
protocol that the runner had implemented before. However, the Python
binding now get to benefit from the loading protocol as well, so the
async runtime library (and potentially other out-of-tree libraries) can
now be used in that context.

[ingomueller-net]

If accepted, this should make https://reviews.llvm.org/D153009 unnecessary.

[ingomueller-net]

I am really unsure about this patch. I don't understand how the underlying ORC works and have moved the code quite mechanically, so please review carefully.

The test suite seems to work, but I get occasional errors for test/mlir-cpu-runner/async* tests, which is suspicious. Running each test in isolation has never failed, though, and rerunning the whole suite always resulted in different failures, which isn't a pattern I would expect as a consequence to this patch, though.

Also, as potential follow up, I could implement the loading protocol for some of the other libraries (runner utils and C runner utils).

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

320

I am not 100% sure why this is necessary. libPath is a StringRef, so it may not be null terminated. In one test of mine, it apparently wasn't, so I saw extra characters. However, both now and before, that StringRef is/was produced by make_absolute. Maybe it was broken before and nobody noticed...

372

These are actually never called in the patch! I don't know where the right place would be...

403

Is this the right place to do this (and all of the above)?

[mehdi_amini]

This is what I had in mind! Thanks.

If we could document this somewhere as well?

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

372

In the destructor of the ExecutionEngine?

[ingomueller-net]

Implement calling of destroy functions in ~ExecutionEngine and document the loading/unloading protocol.

[ingomueller-net]

In D153029#4424860, @mehdi_amini wrote:

This is what I had in mind! Thanks.

If we could document this somewhere as well?

Thanks for the guidance! I added some documentation in the source code of the corresponding option and implemented calling the destroy functions. Let me know what you think!

[ingomueller-net]

@mehdi_amini: There is an effect/issue that I don't yet fully understand, which I am running into while porting the runner utils library to this mechanism. The issue is that I get doubly defined symbols if there is a forward declaration of a function with the emit_c_interface attribute. For example, if I add the following line to one of the unit tests that use the mlir_async_runtime:

func.func private @mlirAsyncRuntimeAddRef(!llvm.ptr, i64) -> () attributes { llvm.emit_c_interface }

then I get the following output and a stackdump:

Failure value returned from cantFail wrapped call
Duplicate definition of symbol 'printMemrefF32'
UNREACHABLE executed at /usr/local/google/home/ingomueller/git/iree-llvm-sandbox/third_party/llvm-project/llvm/include/llvm/Support/Error.h:755!

If the function does not have that attribute, it works fine. For the libraries that currently do *not* use the init mechanism of this patch yet, there is no problem. Many tests contain something like the following for symbols from the runner utils lib:

func.func private @printMemrefF32(memref<*xf32>)
  attributes { llvm.emit_c_interface }

I dug into this a bit and while I don't understand everything involved yet, I have the feeling that the underlying problem is that the emitted C wrapper uses the symbol from the forward declaration, so ORC doesn't want to override it anymore. More precisely, JITDylib::defineImpl(MaterializationUnit &MU) (from llvm/.../Orc/Core.cpp) seems to find the symbol in not being in SymbolState::NeverSearched and thus considers the symbol as not overridable.

As I am writing this, it seems like changing the order in which the symbols are loaded (provided shared libs vs. the IR module that should be run) could avoid this problem but I *think* I have tried this earlier and it didn't work. I'll try this again on Monday and report back, but I am running out of time now.

Should you have a guess on the top of your head (and be at work before I start on Monday), let me know.

[mehdi_amini]

Nothing on top of my head!

[ingomueller-net]

Quick update: I think the problem is that the loading mechanism from this patch loads the symbols into getMainJITDylib() (which is where the jitted code lives) whereas symbols from "normal" shared libraries are loaded into a new createBareJITDylib(). In the former case, they may collide with symbols from the MLIR input; in the latter case, they get a fresh symbol space. I'll investigate more but my feeling is that the libraries loaded by the loading mechanism should also live in their own JITDylib...

[ingomueller-net]

Another update: I think that the problem is actually with FuncToLLVM and how it emits the C interface wrapper. The following function

func.func private @mlirAsyncRuntimeAddRef(!llvm.ptr, i64) -> ()

gets lowered into

llvm.func @mlirAsyncRuntimeAddRef(!llvm.ptr, i64) attributes {sym_visibility = "private"}

which is external, whereas the if the emit_c_interface attribute is added like this:

func.func private @mlirAsyncRuntimeAddRef(!llvm.ptr, i64) -> ()
  attributes { llvm.emit_c_interface }

then it is lowered to the following, which is internal:

llvm.func @mlirAsyncRuntimeAddRef(%arg0: !llvm.ptr, %arg1: i64) attributes {llvm.emit_c_interface, sym_visibility = "private"} {
  llvm.call @_mlir_ciface_mlirAsyncRuntimeAddRef(%arg0, %arg1) : (!llvm.ptr, i64) -> ()
  llvm.return
}
llvm.func @_mlir_ciface_mlirAsyncRuntimeAddRef(!llvm.ptr, i64) attributes {llvm.emit_c_interface, sym_visibility = "private"}

Adding this attribute turns an external function into an internal one. That doesn't sound right, does it?

@ftynse, @Mogball, @Dinistro, @gysit: You also have worked on FuncToLLVM. Any opinions? (I think you can ignore the rest of the patch; just look at this comment.)

[mehdi_amini]

In D153029#4431759, @ingomueller-net wrote:

Another update: I think that the problem is actually with FuncToLLVM and how it emits the C interface wrapper. The following function

func.func private @mlirAsyncRuntimeAddRef(!llvm.ptr, i64) -> ()

gets lowered into

llvm.func @mlirAsyncRuntimeAddRef(!llvm.ptr, i64) attributes {sym_visibility = "private"}

which is external,

It says sym_visibility = "private", what do you mean by "external" here?

[ingomueller-net]

I must be using the wrong terms. In C++ terminology, it turns a declaration (just the signature) into a definition (including a body).

[mehdi_amini]

In D153029#4431809, @ingomueller-net wrote:

I must be using the wrong terms. In C++ terminology, it turns a declaration (just the signature) into a definition (including a body).

Ah got it!

Seems like the issue is that the runtime library should always expose symbols by prefixing them with _mlir_ciface_ and it's not the case here?

[ingomueller-net]

In D153029#4431819, @mehdi_amini wrote:

Seems like the issue is that the runtime library should always expose symbols by prefixing them with _mlir_ciface_ and it's not the case here?

It's not the case but I can't really speak to whether it should be or not.

Independently, I find it problematic that adding this attribute turns a "declaration" into a "definition". Don't you agree? Independently of prefixes, this may cause the problem that I am running into here: IR that wanted to refer to a function provided by some other component suddenly overrides that function...

[mehdi_amini]

Independently, I find it problematic that adding this attribute turns a "declaration" into a "definition". Don't you agree? Independently of prefixes, this may cause the problem that I am running into here: IR that wanted to refer to a function provided by some other component suddenly overrides that function...

I don't think this is necessarily a problem, there is a convention: at a high-level of the IR you're saying "I want to call an external function, for example: call @printMemref1dI32(%x0) : (memref<?xi32>) -> ().

Now when you have a declaration func.func private @printMemref1dI32(%ptr : memref<?xi32>) attributes { llvm.emit_c_interface }, what the llvm.emit_c_interface attribute is saying is "this function is implemented in runtime library that exposes a C API calling convention and the method will be prefixed with _mlir_ciface_ ".

That does not say anything about the name to use for the wrapper we insert though, but it is a private function so its name should not matter in any way?

[ingomueller-net]

In D153029#4431845, @mehdi_amini wrote:

Independently, I find it problematic that adding this attribute turns a "declaration" into a "definition". Don't you agree? Independently of prefixes, this may cause the problem that I am running into here: IR that wanted to refer to a function provided by some other component suddenly overrides that function...

...
That does not say anything about the name to use for the wrapper we insert though, but it is a private function so its name should not matter in any way?

Aha, I see. So that "definition" that is inserted is only private, hence it should not clash with any function exported by any other component (but transparently override it, in this case by calling a function with a different name).

I guess then my previous guess went into a better direction: the symbols exported through the loading mechanism should live in a different "symbol space" than the jitted IR, no? In that case, the private functions from the jitted IR module would only be visible there and not collide with those from the libraries. However, currently, the symbols from the libraries are loaded into the same place as those from the jitted module, hence the collisions I currently observe.

If you agree, I'll look into the linking mechanism more and see how I can load the symbols of each library into their own place.

[Dinistro]

Seems like the issue is that the runtime library should always expose symbols by prefixing them with _mlir_ciface_ and it's not the case here?

The definition that calls the _mlir_ciface_* function should only be generated once, ideally in the same module that holds the original definition. So when the definition is not available anywhere, the entity that defines the original symbol should also provided the _mlir_ciface_ symbol, I guess.

Note that this might be a bit messy if one has no control over the external parts, though.

That does not say anything about the name to use for the wrapper we insert though, but it is a private function so its name should not matter in any way?

Unfortunately, that's not correct. Consider the following the function @ingomueller-net showed before:

llvm.func @mlirAsyncRuntimeAddRef(%arg0: !llvm.ptr, %arg1: i64) attributes {llvm.emit_c_interface, sym_visibility = "private"} {
  llvm.call @_mlir_ciface_mlirAsyncRuntimeAddRef(%arg0, %arg1) : (!llvm.ptr, i64) -> ()
  llvm.return
}

While the symbol has private visibility, the linkage of this function is external (default visibility, thus not printed).

[ingomueller-net]

In D153029#4431855, @Dinistro wrote:

Unfortunately, that's not correct. Consider the following the function @ingomueller-net showed before:

llvm.func @mlirAsyncRuntimeAddRef(%arg0: !llvm.ptr, %arg1: i64) attributes {llvm.emit_c_interface, sym_visibility = "private"} {
  llvm.call @_mlir_ciface_mlirAsyncRuntimeAddRef(%arg0, %arg1) : (!llvm.ptr, i64) -> ()
  llvm.return
}

While the symbol has private visibility, the linkage of this function is external (default visibility, thus not printed).

True! If I convert this to LLVM proper, I get:

define void @mlirAsyncRuntimeAddRef(ptr %0, i64 %1) {
  call void @_mlir_ciface_mlirAsyncRuntimeAddRef(ptr %0, i64 %1)
  ret void
}
declare void @_mlir_ciface_mlirAsyncRuntimeAddRef(ptr, i64)

Both symbols have thus external linkage and default visibility, right?

[Dinistro]

In D153029#4431903, @ingomueller-net wrote:

True! If I convert this to LLVM proper, I get:

define void @mlirAsyncRuntimeAddRef(ptr %0, i64 %1) {
  call void @_mlir_ciface_mlirAsyncRuntimeAddRef(ptr %0, i64 %1)
  ret void
}
declare void @_mlir_ciface_mlirAsyncRuntimeAddRef(ptr, i64)

Both symbols have thus external linkage and default visibility, right?

Yeah, I'm just not sure if this is what you really want. If you have two modules that declared @ mlirAsyncRuntimeAddRef before the lowering, you will have the same definition in both modules. This will cause linking issues eventually.

[ingomueller-net]

In D153029#4431905, @Dinistro wrote:

Yeah, I'm just not sure if this is what you really want. If you have two modules that declared @ mlirAsyncRuntimeAddRef before the lowering, you will have the same definition in both modules. This will cause linking issues eventually.

In the context of this discussion, one of the two comes from a shared library written in C++ that defines a function with this name. Some unit tests want to call this function, so they add a declaration to an external function. However, they also specify the emit_c_interface attribute, so after lowering they suddenly contain a definition with external visibility of that function, which then clashes with that of the shared library.

[ingomueller-net]

In D153029#4431855, @Dinistro wrote:

llvm.func @mlirAsyncRuntimeAddRef(%arg0: !llvm.ptr, %arg1: i64) attributes {llvm.emit_c_interface, sym_visibility = "private"} {
  llvm.call @_mlir_ciface_mlirAsyncRuntimeAddRef(%arg0, %arg1) : (!llvm.ptr, i64) -> ()
  llvm.return
}

While the symbol has private visibility, the linkage of this function is external (default visibility, thus not printed).

While that seems to be true currently, maybe it shouldn't be the case. The IR module is declaring function it knows exists externally, so the lowering should not emit a wrapper function with the same name, default visibility, and external linkage. That's guaranteed to clash, no? If, instead, the lowering emitted a private function that calls the C interface equivalent, then only this IR module would see it as intended, no?

[mehdi_amini]

While the symbol has private visibility, the linkage of this function is external (default visibility, thus not printed).

Something does not line up here, unless we're using a different terminology for these attributes, I assumed we're mapping these attributes to LLVM: https://llvm.org/docs/LangRef.html#linkage-types
"private" in LLVM is not an external linkage (there is inconsistency with MLIR visibility here when the MLIR symbol is public but the LLVM visibility would be private...).

[mehdi_amini]

In D153029#4431958, @ingomueller-net wrote:

While the symbol has private visibility, the linkage of this function is external (default visibility, thus not printed).

While that seems to be true currently, maybe it shouldn't be the case. The IR module is declaring function it knows exists externally, so the lowering should not emit a wrapper function with the same name, default visibility, and external linkage. That's guaranteed to clash, no?

We actually don't know that it exists externally! (it may or may not... the runtime library should only provide an __mlir_ciface_ symbol: we have not reason to expect another one).

If, instead, the lowering emitted a private function that calls the C interface equivalent, then only this IR module would see it as intended, no?

This is what we should be always doing I believe: if we emit a wrapper, regardless of the name, the only way to avoid accidental collision is to have it private.

[mehdi_amini]

llvm.func @mlirAsyncRuntimeAddRef(%arg0: !llvm.ptr, %arg1: i64) attributes {llvm.emit_c_interface, sym_visibility = "private"} {

Seems like "sym_visibility" here is the MLIR annotation for symbols (should be moved to properties...), and nothing specific to LLVM. It is weirdly printed here with the attributes instead of prefixed like func.func

[Dinistro]

In D153029#4431971, @mehdi_amini wrote:

While the symbol has private visibility, the linkage of this function is external (default visibility, thus not printed).

Something does not line up here, unless we're using a different terminology for these attributes, I assumed we're mapping these attributes to LLVM: https://llvm.org/docs/LangRef.html#linkage-types
"private" in LLVM is not an external linkage (there is inconsistency with MLIR visibility here when the MLIR symbol is public but the LLVM visibility would be private...).

That's the source of all confusion. llvm.func has a separate linkage attribute which is independent of the symbol visibility. https://mlir.llvm.org/docs/Dialects/LLVM/#llvmfunc-llvmllvmfuncop
I was not involved in the original design, but I assume the following reasons.
Symbol visibility is an MLIR concept that defines where symbols can be seen, this does not necessarily correspond to how these symbols are linked. For example, a func.func declaration is always private, as it's symbol could otherwise collide with the corresponding definition.
Furthermore, there are only 3 kinds of symbol visibility, while there are many more linkage types.

mlir/include/mlir/ExecutionEngine/ExecutionEngine.h

struct ExecutionEngineOptions {
   /// JIT-compilation and can be used, e.g., for reporting or optimization.
   llvm::function_ref<llvm::Error(llvm::Module *)> transformer = {};
 
   /// `jitCodeGenOptLevel`, when provided, is used as the optimization level for
   /// target code generation.
   std::optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
 
   /// If `sharedLibPaths` are provided, the underlying JIT-compilation will
-  /// open and link the shared libraries for symbol resolution.
+  /// open and link the shared libraries for symbol resolution. Libraries that
+  /// are designed to be used with the `ExecutionEngine` may implement a
+  /// loading and unloading protocol: if they implement the two functions with
+  /// the names defined in `kLibraryInitFnName` and `kLibraryDestroyFnName`,
+  /// these functions will be called upon loading the library and upon
+  /// destruction of the `ExecutionEngine`. In the init function, the library
+  /// may provide a list of symbols that it wants to make available to code
+  /// run by the `ExecutionEngine`. If the two functions are not defined, only
+  /// symbols with public visibility are available to the executed code.
   ArrayRef<StringRef> sharedLibPaths = {};
 
   /// Specifies an existing `sectionMemoryMapper` to be associated with the
   /// compiled code. If none is provided, a default memory mapper that directly
   /// calls into the operating system is used.
   llvm::SectionMemoryManager::MemoryMapper *sectionMemoryMapper = nullptr;
 
   /// If `enableObjectCache` is set, the JIT compiler will create one to store

 ///     void _mlir_funcName(void **)
 ///
 /// where the only argument is interpreted as a list of pointers to the actual
 /// 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:
+  /// Name of init functions of shared libraries. If a library provides a
+  /// function with this name and the one of the destroy function, this function
+  /// is called upon loading the library.
+  static constexpr const char *const kLibraryInitFnName =
+      "__mlir_execution_engine_init";
+
+  /// Name of destroy functions of shared libraries. If a library provides a
+  /// function with this name and the one of the init function, this function is
+  /// called upon destructing the `ExecutionEngine`.
+  static constexpr const char *const kLibraryDestroyFnName =
+      "__mlir_execution_engine_destroy";
+
+  /// Function type for init functions of shared libraries. The library may
+  /// provide a list of symbols that it wants to make available to code run by
+  /// the `ExecutionEngine`. If the two functions are not defined, only symbols
+  /// with public visibility are available to the executed code.
+  using LibraryInitFn = void (*)(llvm::StringMap<void *> &);
+
+  /// Function type for destroy functions of shared libraries.
+  using LibraryDestroyFn = void (*)();
+
   ExecutionEngine(bool enableObjectDump, bool enableGDBNotificationListener,
                   bool enablePerfNotificationListener);
 
+  ~ExecutionEngine();
+
   /// Creates an execution engine for the given MLIR IR. If TargetMachine is
   /// not provided, default TM is created (i.e. ignoring any command line flags
   /// that could affect the set-up).
   static llvm::Expected<std::unique_ptr<ExecutionEngine>>
   create(Operation *op, const ExecutionEngineOptions &options = {},
          std::unique_ptr<llvm::TargetMachine> tm = nullptr);
 
   /// Looks up a packed-argument function wrapping the function with the given
private:
   /// Names of functions that may be looked up.
   std::vector<std::string> functionNames;
 
   /// GDB notification listener.
   llvm::JITEventListener *gdbListener;
 
   /// Perf notification listener.
   llvm::JITEventListener *perfListener;
+
+  /// Destroy functions in the libraries loaded by the ExecutionEngine that are
+  /// called when this ExecutionEngine is destructed.
+  SmallVector<LibraryDestroyFn> destroyFns;
 };
 
 } // namespace mlir
 
 #endif // MLIR_EXECUTIONENGINE_EXECUTIONENGINE_H_

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

 #endif
 
 // Visual Studio had a bug that fails to compile nested generic lambdas
 // inside an `extern "C"` function.
 //   https://developercommunity.visualstudio.com/content/problem/475494/clexe-error-with-lambda-inside-function-templates.html
 // The bug is fixed in VS2019 16.1. Separating the declaration and definition is
 // a work around for older versions of Visual Studio.
 // NOLINTNEXTLINE(*-identifier-naming): externally called.
-extern "C" API void __mlir_runner_init(llvm::StringMap<void *> &exportSymbols);
+extern "C" API void
+__mlir_execution_engine_init(llvm::StringMap<void *> &exportSymbols);
 
 // NOLINTNEXTLINE(*-identifier-naming): externally called.
-void __mlir_runner_init(llvm::StringMap<void *> &exportSymbols) {
+void __mlir_execution_engine_init(llvm::StringMap<void *> &exportSymbols) {
   auto exportSymbol = [&](llvm::StringRef name, auto ptr) {
     assert(exportSymbols.count(name) == 0 && "symbol already exists");
     exportSymbols[name] = reinterpret_cast<void *>(ptr);
   };
 
   exportSymbol("mlirAsyncRuntimeAddRef",
                &mlir::runtime::mlirAsyncRuntimeAddRef);
   exportSymbol("mlirAsyncRuntimeDropRef",
exportSymbol("mlirAsyncRuntimeAwaitAllInGroupAndExecute",
                &mlir::runtime::mlirAsyncRuntimeAwaitAllInGroupAndExecute);
   exportSymbol("mlirAsyncRuntimGetNumWorkerThreads",
                &mlir::runtime::mlirAsyncRuntimGetNumWorkerThreads);
   exportSymbol("mlirAsyncRuntimePrintCurrentThreadId",
                &mlir::runtime::mlirAsyncRuntimePrintCurrentThreadId);
 }
 
 // NOLINTNEXTLINE(*-identifier-naming): externally called.
-extern "C" API void __mlir_runner_destroy() { resetDefaultAsyncRuntime(); }
+extern "C" API void __mlir_execution_engine_destroy() {
+  resetDefaultAsyncRuntime();
+}
 
 } // namespace runtime
 } // namespace mlir
 
 #endif // MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS

mlir/lib/ExecutionEngine/ExecutionEngine.cpp

if (enablePerfNotificationListener) {
     if (auto *listener = llvm::JITEventListener::createPerfJITEventListener())
       perfListener = listener;
     else if (auto *listener =
                  llvm::JITEventListener::createIntelJITEventListener())
       perfListener = listener;
   }
 }
 
+ExecutionEngine::~ExecutionEngine() {
+  // Run all dynamic library destroy callbacks to prepare for the shutdown.
+  for (LibraryDestroyFn destroy : destroyFns)
+    destroy();
+}
+
 Expected<std::unique_ptr<ExecutionEngine>>
 ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
                         std::unique_ptr<llvm::TargetMachine> tm) {
   auto engine = std::make_unique<ExecutionEngine>(
       options.enableObjectDump, options.enableGDBNotificationListener,
       options.enablePerfNotificationListener);
 
   // Remember all entry-points if object dumping is enabled.
ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   // either based on the host defaults or command line arguments when specified
   // (set-up by callers of this method). It could also be passed to the
   // translation or dialect conversion instead of this.
   setupTargetTripleAndDataLayout(llvmModule.get(), tm.get());
   packFunctionArguments(llvmModule.get());
 
   auto dataLayout = llvmModule->getDataLayout();
 
+  // Use absolute library path so that gdb can find the symbol table.
+  SmallVector<SmallString<256>, 4> sharedLibPaths;
+  transform(
+      options.sharedLibPaths, std::back_inserter(sharedLibPaths),
+      [](StringRef libPath) {
+        SmallString<256> absPath(libPath.begin(), libPath.end());
+        cantFail(llvm::errorCodeToError(llvm::sys::fs::make_absolute(absPath)));
+        return absPath;
+      });
+
   // Callback to create the object layer with symbol resolution to current
   // process and dynamically linked libraries.
   auto objectLinkingLayerCreator = [&](ExecutionSession &session,
                                        const Triple &tt) {
     auto objectLayer = std::make_unique<RTDyldObjectLinkingLayer>(
         session, [sectionMemoryMapper = options.sectionMemoryMapper]() {
           return std::make_unique<SectionMemoryManager>(sectionMemoryMapper);
         });
auto objectLinkingLayerCreator = [&](ExecutionSession &session,
     // cf llvm/lib/ExecutionEngine/Orc/LLJIT.cpp LLJIT::createObjectLinkingLayer
     llvm::Triple targetTriple(llvm::Twine(llvmModule->getTargetTriple()));
     if (targetTriple.isOSBinFormatCOFF()) {
       objectLayer->setOverrideObjectFlagsWithResponsibilityFlags(true);
       objectLayer->setAutoClaimResponsibilityForObjectSymbols(true);
     }
 
     // Resolve symbols from shared libraries.
-    for (auto libPath : options.sharedLibPaths) {
+    for (auto &libPath : sharedLibPaths) {
       auto mb = llvm::MemoryBuffer::getFile(libPath);
       if (!mb) {
         errs() << "Failed to create MemoryBuffer for: " << libPath
                << "\nError: " << mb.getError().message() << "\n";
         continue;
       }
       auto &jd = session.createBareJITDylib(std::string(libPath));
       auto loaded = DynamicLibrarySearchGenerator::Load(
-          libPath.data(), dataLayout.getGlobalPrefix());
+          libPath.str().str().c_str(), dataLayout.getGlobalPrefix());
       if (!loaded) {
         errs() << "Could not load " << libPath << ":\n  " << loaded.takeError()
                << "\n";
         continue;
       }
       jd.addGenerator(std::move(*loaded));
       cantFail(objectLayer->add(jd, std::move(mb.get())));
     }
ExecutionEngine::create(Operation *m, const ExecutionEngineOptions &options,
   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())));
 
+  // If shared library implements custom execution layer library init and
+  // destroy functions, we'll use them to register the library.
+
+  llvm::StringMap<void *> exportSymbols;
+  SmallVector<LibraryDestroyFn> destroyFns;
+
+  for (auto &libPath : sharedLibPaths) {
+    auto lib = llvm::sys::DynamicLibrary::getPermanentLibrary(
+        libPath.str().str().c_str());
+    void *initSym = lib.getAddressOfSymbol(kLibraryInitFnName);
+    void *destroySim = lib.getAddressOfSymbol(kLibraryDestroyFnName);
+
+    // Library does not provide call backs, rely on symbol visiblity.
+    if (!initSym || !destroySim) {
+      continue;
+    }
+
+    auto initFn = reinterpret_cast<LibraryInitFn>(initSym);
+    initFn(exportSymbols);
+
+    auto destroyFn = reinterpret_cast<LibraryDestroyFn>(destroySim);
+    destroyFns.push_back(destroyFn);
+  }
+  engine->destroyFns = std::move(destroyFns);
+
+  // Build a runtime symbol map from the exported symbols and register them.
+  auto runtimeSymbolMap = [&](llvm::orc::MangleAndInterner interner) {
+    auto symbolMap = llvm::orc::SymbolMap();
+    for (auto &exportSymbol : exportSymbols)
+      symbolMap[interner(exportSymbol.getKey())] = {
+          llvm::orc::ExecutorAddr::fromPtr(exportSymbol.getValue()),
+          llvm::JITSymbolFlags::Exported};
+    return symbolMap;
+  };
+  engine->registerSymbols(runtimeSymbolMap);
+
   return std::move(engine);
 }
 
 Expected<void (*)(void **)>
 ExecutionEngine::lookupPacked(StringRef name) const {
   auto result = lookup(makePackedFunctionName(name));
   if (!result)
     return result.takeError();
   return reinterpret_cast<void (*)(void **)>(result.get());
 }
 

mlir/lib/ExecutionEngine/JitRunner.cpp

 static Error
 compileAndExecute(Options &options, Operation *module, StringRef entryPoint,
                   CompileAndExecuteConfig config, void **args,
                   std::unique_ptr<llvm::TargetMachine> tm = nullptr) {
   std::optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
   if (auto clOptLevel = getCommandLineOptLevel(options))
     jitCodeGenOptLevel = static_cast<llvm::CodeGenOpt::Level>(*clOptLevel);
 
-  // If shared library implements custom mlir-runner library init and destroy
-  // functions, we'll use them to register the library with the execution
-  // engine. Otherwise we'll pass library directly to the execution engine.
-  SmallVector<SmallString<256>, 4> libPaths;
-
-  // Use absolute library path so that gdb can find the symbol table.
-  transform(
-      options.clSharedLibs, std::back_inserter(libPaths),
-      [](std::string libPath) {
-        SmallString<256> absPath(libPath.begin(), libPath.end());
-        cantFail(llvm::errorCodeToError(llvm::sys::fs::make_absolute(absPath)));
-        return absPath;
-      });
-
-  // Libraries that we'll pass to the ExecutionEngine for loading.
-  SmallVector<StringRef, 4> executionEngineLibs;
-
-  using MlirRunnerInitFn = void (*)(llvm::StringMap<void *> &);
-  using MlirRunnerDestroyFn = void (*)();
-
-  llvm::StringMap<void *> exportSymbols;
-  SmallVector<MlirRunnerDestroyFn> destroyFns;
-
-  // Handle libraries that do support mlir-runner init/destroy callbacks.
-  for (auto &libPath : libPaths) {
-    auto lib = llvm::sys::DynamicLibrary::getPermanentLibrary(libPath.c_str());
-    void *initSym = lib.getAddressOfSymbol("__mlir_runner_init");
-    void *destroySim = lib.getAddressOfSymbol("__mlir_runner_destroy");
-
-    // Library does not support mlir runner, load it with ExecutionEngine.
-    if (!initSym || !destroySim) {
-      executionEngineLibs.push_back(libPath);
-      continue;
-    }
-
-    auto initFn = reinterpret_cast<MlirRunnerInitFn>(initSym);
-    initFn(exportSymbols);
-
-    auto destroyFn = reinterpret_cast<MlirRunnerDestroyFn>(destroySim);
-    destroyFns.push_back(destroyFn);
-  }
-
-  // Build a runtime symbol map from the config and exported symbols.
-  auto runtimeSymbolMap = [&](llvm::orc::MangleAndInterner interner) {
-    auto symbolMap = config.runtimeSymbolMap ? config.runtimeSymbolMap(interner)
-                                             : llvm::orc::SymbolMap();
-    for (auto &exportSymbol : exportSymbols)
-      symbolMap[interner(exportSymbol.getKey())] =
-          { llvm::orc::ExecutorAddr::fromPtr(exportSymbol.getValue()),
-            llvm::JITSymbolFlags::Exported };
-    return symbolMap;
-  };
+  SmallVector<StringRef, 4> sharedLibs(options.clSharedLibs.begin(),
+                                       options.clSharedLibs.end());
 
   mlir::ExecutionEngineOptions engineOptions;
   engineOptions.llvmModuleBuilder = config.llvmModuleBuilder;
   if (config.transformer)
     engineOptions.transformer = config.transformer;
   engineOptions.jitCodeGenOptLevel = jitCodeGenOptLevel;
-  engineOptions.sharedLibPaths = executionEngineLibs;
+  engineOptions.sharedLibPaths = sharedLibs;
   engineOptions.enableObjectDump = true;
   auto expectedEngine =
       mlir::ExecutionEngine::create(module, engineOptions, std::move(tm));
   if (!expectedEngine)
     return expectedEngine.takeError();
 
   auto engine = std::move(*expectedEngine);
-  engine->registerSymbols(runtimeSymbolMap);
 
   auto expectedFPtr = engine->lookupPacked(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);
 
-  // Run all dynamic library destroy callbacks to prepare for the shutdown.
-  for (MlirRunnerDestroyFn destroy : destroyFns)
-    destroy();
-
   return Error::success();
 }
 
 static Error compileAndExecuteVoidFunction(
     Options &options, Operation *module, StringRef entryPoint,
     CompileAndExecuteConfig config, std::unique_ptr<llvm::TargetMachine> tm) {
   auto mainFunction = dyn_cast_or_null<LLVM::LLVMFuncOp>(
       SymbolTable::lookupSymbolIn(module, entryPoint));