[mlir] mlir-cpu-runner: Register shutdown callbacks with the runner

[ezhulenev]

Summary

By registering shutdown callback with a mlir-cpu-runner dynamic libraries can prepare for the shutdown and do not depend on the global destructors execution order.

[ezhulenev]

Remove extra space

[ezhulenev]

Fix CMakeLists

[mehdi_amini]

This seems to get in the direction of *more* "implicit global state" sharing instead of less.

What I'd rather see as a direction is to not dlopen libraries in Global mode, and dlsym specific mlir_runner_init/mlir_runner_destroy entry points in the libraries we load. These would be able to populate a list of symbols available to the JIT explicitly instead. This is something we already did for the ModelBuilder wrapper for example, see the extra_symbols here: https://github.com/google/iree/blob/main/experimental/ModelBuilder/ModelRunner.cpp#L57

[ezhulenev]

Yeah, makes sense, also weak linking does not seem to work with MSVC, will try to the dlsym approach.

mlir/lib/ExecutionEngine/AsyncRuntime.cpp

 #include <functional>
 #include <iostream>
 #include <mutex>
 #include <thread>
 #include <vector>
 
 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.
 //===----------------------------------------------------------------------===//
 
 namespace mlir {
 namespace runtime {
 namespace {
 

 private:
   AsyncRuntime *runtime;
   std::atomic<int32_t> refCount;
 };
 
 } // namespace
 
 // Returns the default per-process instance of an async runtime.
+static std::unique_ptr<AsyncRuntime> &getDefaultAsyncRuntime() {
+  static auto runtime = []() {
+    __mlir_cpu_runner_register_shutdown_callback(
+        []() { getDefaultAsyncRuntime().reset(); });
+    return std::make_unique<AsyncRuntime>();
+  }();
+  return runtime;
+}
+
 static AsyncRuntime *getDefaultAsyncRuntimeInstance() {
-  static auto runtime = std::make_unique<AsyncRuntime>();
-  return runtime.get();
+  return getDefaultAsyncRuntime().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

mlir/tools/mlir-cpu-runner/CMakeLists.txt

+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
   nativecodegen
   native
   )
 
 add_llvm_tool(mlir-cpu-runner

mlir/tools/mlir-cpu-runner/mlir-cpu-runner.cpp

 // latter.
 //
 //===----------------------------------------------------------------------===//
 
 #include "mlir/ExecutionEngine/JitRunner.h"
 #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<ShutdownCallback, 4> callbacks;
+};
+
+static llvm::ManagedStatic<ShutdownCallbacks> 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();
   llvm::InitializeNativeTargetAsmPrinter();
   llvm::InitializeNativeTargetAsmParser();
   mlir::initializeLLVMPasses();
 
   return mlir::JitRunnerMain(argc, argv);
 }