[mlir:JitRunner] Use custom shared library init/destroy functions if available

[ezhulenev]

Summary

Use custom mlir runner init/destroy functions to safely init and destroy shared libraries loaded by the JitRunner.

This mechanism is ignored for Windows builds (for now) because init/destroy functions are not exported, and library unloading relies on static destructors.

Re-submit https://reviews.llvm.org/D94270 with a temporary workaround for windows

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

 #include <cassert>
 #include <condition_variable>
 #include <functional>
 #include <iostream>
 #include <mutex>
 #include <thread>
 #include <vector>
 
+#include "llvm/ADT/StringMap.h"
+
 using namespace mlir::runtime;
 
 //===----------------------------------------------------------------------===//
 // Async runtime API.
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace runtime {

 private:
   AsyncRuntime *runtime;
   std::atomic<int32_t> refCount;
 };
 
 } // namespace
 
 // Returns the default per-process instance of an async runtime.
-static AsyncRuntime *getDefaultAsyncRuntimeInstance() {
+static std::unique_ptr<AsyncRuntime> &getDefaultAsyncRuntimeInstance() {
   static auto runtime = std::make_unique<AsyncRuntime>();
-  return runtime.get();
+  return runtime;
+}
+
+static void resetDefaultAsyncRuntime() {
+  return getDefaultAsyncRuntimeInstance().reset();
+}
+
+static AsyncRuntime *getDefaultAsyncRuntime() {
+  return getDefaultAsyncRuntimeInstance().get();
 }
 
 // Async token provides a mechanism to signal asynchronous operation completion.
 struct AsyncToken : public RefCounted {
   // AsyncToken created with a reference count of 2 because it will be returned
   // to the `async.execute` caller and also will be later on emplaced by the
   // asynchronously executed task. If the caller immediately will drop its
   // reference we must ensure that the token will be alive until the

 // Drops references from reference counted runtime object.
 extern "C" void mlirAsyncRuntimeDropRef(RefCountedObjPtr ptr, int32_t count) {
   RefCounted *refCounted = static_cast<RefCounted *>(ptr);
   refCounted->dropRef(count);
 }
 
 // Creates a new `async.token` in not-ready state.
 extern "C" AsyncToken *mlirAsyncRuntimeCreateToken() {
-  AsyncToken *token = new AsyncToken(getDefaultAsyncRuntimeInstance());
+  AsyncToken *token = new AsyncToken(getDefaultAsyncRuntime());
   return token;
 }
 
 // Creates a new `async.value` in not-ready state.
 extern "C" AsyncValue *mlirAsyncRuntimeCreateValue(int32_t size) {
-  AsyncValue *value = new AsyncValue(getDefaultAsyncRuntimeInstance(), size);
+  AsyncValue *value = new AsyncValue(getDefaultAsyncRuntime(), size);
   return value;
 }
 
 // Create a new `async.group` in empty state.
 extern "C" AsyncGroup *mlirAsyncRuntimeCreateGroup() {
-  AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntimeInstance());
+  AsyncGroup *group = new AsyncGroup(getDefaultAsyncRuntime());
   return group;
 }
 
 extern "C" int64_t mlirAsyncRuntimeAddTokenToGroup(AsyncToken *token,
                                                    AsyncGroup *group) {
   std::unique_lock<std::mutex> lockToken(token->mu);
   std::unique_lock<std::mutex> lockGroup(group->mu);
 

 // Small async runtime support library for testing.
 //===----------------------------------------------------------------------===//
 
 extern "C" void mlirAsyncRuntimePrintCurrentThreadId() {
   static thread_local std::thread::id thisId = std::this_thread::get_id();
   std::cout << "Current thread id: " << thisId << std::endl;
 }
 
+//===----------------------------------------------------------------------===//
+// MLIR Runner (JitRunner) dynamic library integration.
+//===----------------------------------------------------------------------===//
+
+// Export symbols for the MLIR runner integration. All other symbols are hidden.
+#ifndef _WIN32
+#define API __attribute__((visibility("default")))
+#else
+#define API
+#endif // _WIN32
+
+extern "C" API void __mlir_runner_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",
+               &mlir::runtime::mlirAsyncRuntimeDropRef);
+  exportSymbol("mlirAsyncRuntimeExecute",
+               &mlir::runtime::mlirAsyncRuntimeExecute);
+  exportSymbol("mlirAsyncRuntimeGetValueStorage",
+               &mlir::runtime::mlirAsyncRuntimeGetValueStorage);
+  exportSymbol("mlirAsyncRuntimeCreateToken",
+               &mlir::runtime::mlirAsyncRuntimeCreateToken);
+  exportSymbol("mlirAsyncRuntimeCreateValue",
+               &mlir::runtime::mlirAsyncRuntimeCreateValue);
+  exportSymbol("mlirAsyncRuntimeEmplaceToken",
+               &mlir::runtime::mlirAsyncRuntimeEmplaceToken);
+  exportSymbol("mlirAsyncRuntimeEmplaceValue",
+               &mlir::runtime::mlirAsyncRuntimeEmplaceValue);
+  exportSymbol("mlirAsyncRuntimeAwaitToken",
+               &mlir::runtime::mlirAsyncRuntimeAwaitToken);
+  exportSymbol("mlirAsyncRuntimeAwaitValue",
+               &mlir::runtime::mlirAsyncRuntimeAwaitValue);
+  exportSymbol("mlirAsyncRuntimeAwaitTokenAndExecute",
+               &mlir::runtime::mlirAsyncRuntimeAwaitTokenAndExecute);
+  exportSymbol("mlirAsyncRuntimeAwaitValueAndExecute",
+               &mlir::runtime::mlirAsyncRuntimeAwaitValueAndExecute);
+  exportSymbol("mlirAsyncRuntimeCreateGroup",
+               &mlir::runtime::mlirAsyncRuntimeCreateGroup);
+  exportSymbol("mlirAsyncRuntimeAddTokenToGroup",
+               &mlir::runtime::mlirAsyncRuntimeAddTokenToGroup);
+  exportSymbol("mlirAsyncRuntimeAwaitAllInGroup",
+               &mlir::runtime::mlirAsyncRuntimeAwaitAllInGroup);
+  exportSymbol("mlirAsyncRuntimeAwaitAllInGroupAndExecute",
+               &mlir::runtime::mlirAsyncRuntimeAwaitAllInGroupAndExecute);
+  exportSymbol("mlirAsyncRuntimePrintCurrentThreadId",
+               &mlir::runtime::mlirAsyncRuntimePrintCurrentThreadId);
+}
+
+extern "C" API void __mlir_runner_destroy() { resetDefaultAsyncRuntime(); }
+
 #endif // MLIR_ASYNCRUNTIME_DEFINE_FUNCTIONS

mlir/lib/ExecutionEngine/CMakeLists.txt

add_mlir_library(mlir_async_runtime
   AsyncRuntime.cpp
 
   EXCLUDE_FROM_LIBMLIR
 
   LINK_LIBS PUBLIC
   mlir_c_runner_utils_static
   ${LLVM_PTHREAD_LIB}
 )
+set_property(TARGET mlir_async_runtime PROPERTY CXX_VISIBILITY_PRESET hidden)
 target_compile_definitions(mlir_async_runtime PRIVATE mlir_async_runtime_EXPORTS)

mlir/lib/ExecutionEngine/JitRunner.cpp

 // JIT-compile the given module and run "entryPoint" with "args" as arguments.
 static Error compileAndExecute(Options &options, ModuleOp module,
                                StringRef entryPoint,
                                CompileAndExecuteConfig config, void **args) {
   Optional<llvm::CodeGenOpt::Level> jitCodeGenOptLevel;
   if (auto clOptLevel = getCommandLineOptLevel(options))
     jitCodeGenOptLevel =
         static_cast<llvm::CodeGenOpt::Level>(clOptLevel.getValue());
+
+  // 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<StringRef, 4> libs(options.clSharedLibs.begin(),
                                  options.clSharedLibs.end());
+
+  // 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 : libs) {
+    auto lib = llvm::sys::DynamicLibrary::getPermanentLibrary(libPath.data());
+    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::JITEvaluatedSymbol::fromPointer(exportSymbol.getValue());
+    return symbolMap;
+  };
+
   auto expectedEngine = mlir::ExecutionEngine::create(
       module, config.llvmModuleBuilder, config.transformer, jitCodeGenOptLevel,
-      libs);
+      executionEngineLibs);
   if (!expectedEngine)
     return expectedEngine.takeError();
 
   auto engine = std::move(*expectedEngine);
-  if (config.runtimeSymbolMap)
-    engine->registerSymbols(config.runtimeSymbolMap);
+  engine->registerSymbols(runtimeSymbolMap);
 
   auto expectedFPtr = engine->lookup(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.
+  llvm::for_each(destroyFns, [](MlirRunnerDestroyFn destroy) { destroy(); });
+
   return Error::success();
 }
 
 static Error compileAndExecuteVoidFunction(Options &options, ModuleOp module,
                                            StringRef entryPoint,
                                            CompileAndExecuteConfig config) {
   auto mainFunction = module.lookupSymbol<LLVM::LLVMFuncOp>(entryPoint);
   if (!mainFunction || mainFunction.empty())