[SYCL] Keep kernel compiler resources alive during static destruction

sycl-jit/jit-compiler/lib/translation/JITContext.cpp

@@ -10,6 +10,20 @@
 
 using namespace jit_compiler;
 
+#ifndef _WIN32
+JITContext *JITContext::Instance = nullptr;
+
+void JITContext::destroyInstance() {
+  delete Instance;
+  Instance = nullptr;
+}
+
+// Match the runtime's late unload priority, after user static destructors.
+__attribute__((destructor(110))) static void jitUnload() {
+  JITContext::destroyInstance();
+}
+#endif
+
 JITBinary::JITBinary(std::string &&Binary, BinaryFormat Fmt)
     : Blob{std::move(Binary)}, Format{Fmt} {}
 

sycl-jit/jit-compiler/lib/translation/JITContext.h

@@ -51,10 +51,21 @@ class JITContext {
 
 public:
   static JITContext &getInstance() {
+#ifdef _WIN32
     static JITContext Instance{};
     return Instance;
+#else
+    // Runtime-compiled bundles can outlive ordinary static destructors. Keep
+    // their binary storage alive until the JIT library's late unload handler.
+    static JITContext &Context = *(Instance = new JITContext{});
+    return Context;
+#endif
   }
 
+#ifndef _WIN32
+  static void destroyInstance();
+#endif
+
   template <typename... Ts> JITBinary &emplaceBinary(Ts &&...Args) {
     WriteLockT WriteLock{BinariesMutex};
     auto JBUPtr = std::make_unique<JITBinary>(std::forward<Ts>(Args)...);
@@ -69,6 +80,9 @@ class JITContext {
   }
 
 private:
+#ifndef _WIN32
+  static JITContext *Instance;
+#endif
   JITContext() = default;
   ~JITContext() = default;
   JITContext(const JITContext &) = delete;

sycl/source/detail/global_handler.cpp

@@ -15,6 +15,7 @@
 #include <detail/config.hpp>
 #include <detail/device_kernel_info.hpp>
 #include <detail/global_handler.hpp>
+#include <detail/jit_compiler.hpp>
 #include <detail/platform_impl.hpp>
 #include <detail/program_manager/program_manager.hpp>
 #include <detail/scheduler/scheduler.hpp>
@@ -206,6 +207,15 @@ ProgramManager &GlobalHandler::getProgramManager() {
   return PM;
 }
 
+#if SYCL_EXT_JIT_ENABLE
+jit_compiler &GlobalHandler::getJITCompiler() {
+  const LockGuard Lock{MJITCompiler.Lock};
+  if (!MJITCompiler.Inst)
+    MJITCompiler.Inst.reset(new jit_compiler());
+  return *MJITCompiler.Inst;
+}
+#endif
+
 std::unordered_map<platform_impl *, std::shared_ptr<context_impl>> &
 GlobalHandler::getPlatformToDefaultContextCache() {
   // The optimization with static reference is not done because
@@ -447,6 +457,9 @@ void shutdown_late() {
   GlobalHandler::RTGlobalObjHandler->MScheduler.Inst.reset(nullptr);
   GlobalHandler::RTGlobalObjHandler->MProgramManager.Inst.reset(nullptr);
 
+  // The JIT compiler owns image tables referenced by the program manager.
+  GlobalHandler::RTGlobalObjHandler->MJITCompiler.Inst.reset(nullptr);
+
   // Clear the adapters and reset the instance if it was there.
   GlobalHandler::RTGlobalObjHandler->unloadAdapters();
   if (GlobalHandler::RTGlobalObjHandler->MAdapters.Inst)

sycl/source/detail/global_handler.hpp

@@ -23,6 +23,7 @@ class platform_impl;
 class context_impl;
 class Scheduler;
 class ProgramManager;
+class jit_compiler;
 class Sync;
 class adapter_impl;
 class ods_target_list;
@@ -77,6 +78,7 @@ class GlobalHandler {
   /// returned access object keeps the Scheduler alive until it is destroyed.
   static SchedulerAccess getSchedulerAccess() noexcept;
   ProgramManager &getProgramManager();
+  jit_compiler &getJITCompiler();
   Sync &getSync();
   std::vector<std::shared_ptr<platform_impl>> &getPlatformCache();
 
@@ -168,6 +170,7 @@ class GlobalHandler {
   // count accesses that must finish before the Scheduler can be destroyed.
   static std::atomic<uint64_t> MSchedulerAccessState;
   InstWithLock<ProgramManager> MProgramManager;
+  InstWithLock<jit_compiler> MJITCompiler;
   InstWithLock<Sync> MSync;
   InstWithLock<std::vector<std::shared_ptr<platform_impl>>> MPlatformCache;
   InstWithLock<

sycl/source/detail/jit_compiler.cpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 #include <sycl/feature_test.hpp>
 #if SYCL_EXT_JIT_ENABLE
+#include <detail/global_handler.hpp>
 #include <detail/jit_compiler.hpp>
 #include <detail/persistent_device_code_cache.hpp>
 #include <sycl/detail/os_util.hpp>
@@ -16,6 +17,10 @@ namespace sycl {
 inline namespace _V1 {
 namespace detail {
 
+jit_compiler &jit_compiler::get_instance() {
+  return GlobalHandler::instance().getJITCompiler();
+}
+
 std::function<void(void *)> jit_compiler::CustomDeleterForLibHandle =
     [](void *StoredPtr) {
       if (!StoredPtr)

sycl/source/detail/jit_compiler.hpp

@@ -54,10 +54,7 @@ class jit_compiler {
 
   bool isAvailable() { return Available; }
 
-  static jit_compiler &get_instance() {
-    static jit_compiler instance{};
-    return instance;
-  }
+  static jit_compiler &get_instance();
 
 private:
   jit_compiler();
@@ -67,6 +64,9 @@ class jit_compiler {
   jit_compiler &operator=(const jit_compiler &) = delete;
   jit_compiler &operator=(const jit_compiler &&) = delete;
 
+  friend class GlobalHandler;
+  friend struct std::default_delete<jit_compiler>;
+
   sycl_device_binaries
   createDeviceBinaries(const ::jit_compiler::RTCBundleInfo &BundleInfo,
                        const std::string &Prefix);

sycl/test-e2e/KernelCompiler/sycl_static_lifetime.cpp

@@ -0,0 +1,65 @@
+// REQUIRES: sycl-jit, aspect-usm_shared_allocations
+// UNSUPPORTED: target-native_cpu
+// UNSUPPORTED-TRACKER: https://github.com/intel/llvm/issues/20142
+
+// RUN: %{build} -o %t.out
+// RUN: %if hip %{ env SYCL_JIT_AMDGCN_PTX_TARGET_CPU=%{amd_arch} %} %{run} %t.out
+// RUN: %if hip %{ env SYCL_JIT_AMDGCN_PTX_TARGET_CPU=%{amd_arch} %} %{run} %t.out keep
+
+// Cached source-built bundles must support execution and cleanup during static
+// destruction, after ordinary runtime static destructors would have run.
+#include <sycl/detail/core.hpp>
+#include <sycl/kernel_bundle.hpp>
+#include <sycl/usm.hpp>
+
+#include <cassert>
+#include <optional>
+#include <vector>
+
+namespace syclex = sycl::ext::oneapi::experimental;
+using ExecutableBundle = sycl::kernel_bundle<sycl::bundle_state::executable>;
+
+constexpr const char *Source = R"(
+#include <sycl/detail/core.hpp>
+#include <sycl/ext/oneapi/free_function_kernel_properties.hpp>
+extern "C" SYCL_EXTERNAL
+SYCL_EXT_ONEAPI_FUNCTION_PROPERTY((sycl::ext::oneapi::experimental::single_task_kernel))
+void increment(int *Value) { ++*Value; }
+)";
+
+struct KernelCache {
+  std::vector<ExecutableBundle> Bundles;
+  std::optional<sycl::queue> Queue;
+
+  ~KernelCache() {
+    for (auto &Bundle : Bundles) {
+      auto Kernel = Bundle.ext_oneapi_get_kernel("increment");
+      int *Value = sycl::malloc_shared<int>(1, *Queue);
+      *Value = 0;
+      Queue->submit([&](sycl::handler &CGH) {
+        CGH.set_arg(0, Value);
+        CGH.single_task(Kernel);
+      });
+      Queue->wait_and_throw();
+      assert(*Value == 1);
+      sycl::free(Value, *Queue);
+    }
+  }
+};
+
+int main(int Argc, char **) {
+  sycl::queue Q;
+  static KernelCache Cache;
+  assert(Q.get_device().ext_oneapi_can_build(syclex::source_language::sycl));
+  auto SourceBundle = syclex::create_kernel_bundle_from_source(
+      Q.get_context(), syclex::source_language::sycl, Source);
+  for (int I = 0; I < 2; ++I) {
+    auto Bundle = syclex::build(SourceBundle);
+    assert(Bundle.ext_oneapi_has_kernel("increment"));
+    auto Kernel = Bundle.ext_oneapi_get_kernel("increment");
+    if (Argc > 1) {
+      Cache.Queue = Q;
+      Cache.Bundles.push_back(Bundle);
+    }
+  }
+}