diff --git a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp --- a/mlir/lib/ExecutionEngine/AsyncRuntime.cpp +++ b/mlir/lib/ExecutionEngine/AsyncRuntime.cpp @@ -26,6 +26,12 @@ using namespace mlir::runtime; +// Shutdown hook registration function will be resolved to the symbol defined by +// the mlir-cpu-runner binary when it will dlopen the async runtime library. +using ShutdownCallback = void (*)(); +extern "C" __attribute__((weak)) void + __mlir_cpu_runner_register_shutdown_callback(ShutdownCallback); + //===----------------------------------------------------------------------===// // Async runtime API. //===----------------------------------------------------------------------===// @@ -109,9 +115,17 @@ } // namespace // Returns the default per-process instance of an async runtime. +static std::unique_ptr &getDefaultAsyncRuntime() { + static auto runtime = []() { + __mlir_cpu_runner_register_shutdown_callback( + []() { getDefaultAsyncRuntime().reset(); }); + return std::make_unique(); + }(); + return runtime; +} + static AsyncRuntime *getDefaultAsyncRuntimeInstance() { - static auto runtime = std::make_unique(); - return runtime.get(); + return getDefaultAsyncRuntime().get(); } // Async token provides a mechanism to signal asynchronous operation completion. diff --git a/mlir/tools/mlir-cpu-runner/CMakeLists.txt b/mlir/tools/mlir-cpu-runner/CMakeLists.txt --- a/mlir/tools/mlir-cpu-runner/CMakeLists.txt +++ b/mlir/tools/mlir-cpu-runner/CMakeLists.txt @@ -1,3 +1,7 @@ +set(CMAKE_EXE_LINKER_FLAGS + "${CMAKE_EXE_LINKER_FLAGS} + -Wl,--export-dynamic-symbol=__mlir_cpu_runner_register_shutdown_callback") + set(LLVM_LINK_COMPONENTS Core Support diff --git a/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp b/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp --- a/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp +++ b/mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp @@ -16,8 +16,36 @@ #include "mlir/ExecutionEngine/OptUtils.h" #include "mlir/InitAllDialects.h" #include "llvm/Support/InitLLVM.h" +#include "llvm/Support/ManagedStatic.h" #include "llvm/Support/TargetSelect.h" +namespace mlir { +namespace runtime { + +using ShutdownCallback = void (*)(); + +// Shutdown callbacks that every dynamic library can use to properly destory all +// the static objects when the mlir-cpu-runner calls llvm_shutdown(). +struct ShutdownCallbacks { + ~ShutdownCallbacks() { + llvm::for_each(callbacks, [](ShutdownCallback callback) { callback(); }); + } + llvm::SmallVector callbacks; +}; + +static llvm::ManagedStatic shutdownCallbacks; + +} // namespace runtime +} // namespace mlir + +using namespace mlir::runtime; + +// Register a shutdown callback function with the mlir-cpu-runner. +extern "C" void +__mlir_cpu_runner_register_shutdown_callback(ShutdownCallback callback) { + shutdownCallbacks->callbacks.push_back(callback); +} + int main(int argc, char **argv) { llvm::InitLLVM y(argc, argv); llvm::InitializeNativeTarget();