[SYCL] Release deferred resources after event and queue waits

sycl/source/detail/event_impl.cpp

@@ -9,6 +9,7 @@
 #include <detail/adapter_impl.hpp>
 #include <detail/event_impl.hpp>
 #include <detail/event_info.hpp>
+#include <detail/global_handler.hpp>
 #include <detail/queue_impl.hpp>
 #include <detail/scheduler/scheduler.hpp>
 #include <sycl/context.hpp>
@@ -405,12 +406,23 @@ void event_impl::wait(bool *Success) {
 #endif
 
   auto EventHandle = getHandle();
+  bool ScanAuxiliaryResources = false;
   if (EventHandle)
     // presence of the native handle means the command has been enqueued, so no
     // need to go via the slow path event waiting in the scheduler
     waitInternal(Success);
   else if (MCommand)
-    detail::Scheduler::getInstance().waitForEvent(*this, Success);
+    ScanAuxiliaryResources =
+        detail::Scheduler::getInstance().waitForEvent(*this, Success);
+
+  // Deferred resources are otherwise only released by later scheduler
+  // activity, i.e. possibly not before runtime shutdown. The check is
+  // lock-free so that wait() takes no lock when nothing is deferred.
+  auto SchedAccess = detail::GlobalHandler::getSchedulerAccess();
+  if (detail::Scheduler *Sched = SchedAccess.get()) {
+    if (Sched->hasDeferredResources())
+      Sched->releaseResourcesAfterWait(*this, ScanAuxiliaryResources);
+  }
 
 #ifdef XPTI_ENABLE_INSTRUMENTATION
   instrumentationEpilog(TelemetryEvent, Name, StreamID, IId);

sycl/source/detail/global_handler.cpp

@@ -24,6 +24,7 @@
 #include <sycl/detail/device_filter.hpp>
 #include <sycl/detail/spinlock.hpp>
 
+#include <thread>
 #ifdef _WIN32
 #include <windows.h>
 #endif
@@ -36,6 +37,11 @@ namespace detail {
 
 using LockGuard = std::lock_guard<SpinLock>;
 SpinLock GlobalHandler::MSyclGlobalHandlerProtector{};
+std::atomic<Scheduler *> GlobalHandler::MSchedulerPtr{nullptr};
+std::atomic<uint64_t> GlobalHandler::MSchedulerAccessState{0};
+#ifdef _WIN32
+std::atomic<bool> ProcessTerminating{false};
+#endif
 
 // forward decl
 void shutdown_early(bool);
@@ -122,7 +128,11 @@ void GlobalHandler::attachScheduler(Scheduler *Scheduler) {
   // releaseResources will cause dead lock due to host queue release
   if (MScheduler.Inst)
     prepareSchedulerToRelease(true);
+  stopSchedulerAccess();
+  MSchedulerPtr.store(nullptr, std::memory_order_release);
   MScheduler.Inst.reset(Scheduler);
+  MSchedulerPtr.store(Scheduler, std::memory_order_release);
+  MSchedulerAccessState.store(0, std::memory_order_release);
 }
 
 static void enableOnCrashStackPrinting() {
@@ -135,7 +145,9 @@ static void enableOnCrashStackPrinting() {
 }
 
 Scheduler &GlobalHandler::getScheduler() {
-  getOrCreate(MScheduler);
+  Scheduler &Sched = getOrCreate(MScheduler);
+  if (!MSchedulerPtr.load(std::memory_order_relaxed))
+    MSchedulerPtr.store(&Sched, std::memory_order_release);
   registerSchedulerUsage();
   // On Windows the registration of the signal handler before main function
   // (e.g. from DLLMain or from constructors of program scope objects) doesn't
@@ -144,10 +156,46 @@ Scheduler &GlobalHandler::getScheduler() {
   // 2) first call to getScheduler is likely to be done after main starts.
   // The same is done in getAdapters.
   enableOnCrashStackPrinting();
-  return *MScheduler.Inst;
+  return Sched;
+}
+
+bool GlobalHandler::isSchedulerAlive() const {
+  return MSchedulerPtr.load(std::memory_order_acquire) != nullptr;
+}
+
+GlobalHandler::SchedulerAccess GlobalHandler::getSchedulerAccess() noexcept {
+  constexpr uint64_t Closed = uint64_t{1} << 63;
+  uint64_t State = MSchedulerAccessState.load(std::memory_order_relaxed);
+  while (!(State & Closed)) {
+    if (MSchedulerAccessState.compare_exchange_weak(
+            State, State + 1, std::memory_order_acquire,
+            std::memory_order_relaxed)) {
+      if (Scheduler *Sched = MSchedulerPtr.load(std::memory_order_acquire))
+        return SchedulerAccess(Sched);
+      releaseSchedulerAccess();
+      return SchedulerAccess();
+    }
+  }
+  return SchedulerAccess();
 }
 
-bool GlobalHandler::isSchedulerAlive() const { return MScheduler.Inst.get(); }
+GlobalHandler::SchedulerAccess::~SchedulerAccess() {
+  if (MScheduler)
+    GlobalHandler::releaseSchedulerAccess();
+}
+
+void GlobalHandler::releaseSchedulerAccess() noexcept {
+  MSchedulerAccessState.fetch_sub(1, std::memory_order_release);
+}
+
+void GlobalHandler::stopSchedulerAccess(bool WaitForAccess) noexcept {
+  constexpr uint64_t Closed = uint64_t{1} << 63;
+  MSchedulerAccessState.fetch_or(Closed, std::memory_order_acq_rel);
+  if (!WaitForAccess)
+    return;
+  while ((MSchedulerAccessState.load(std::memory_order_acquire) & ~Closed) != 0)
+    std::this_thread::yield();
+}
 
 void GlobalHandler::registerSchedulerUsage(bool ModifyCounter) {
   thread_local ObjectUsageCounter SchedulerCounter(ModifyCounter);
@@ -318,6 +366,13 @@ void shutdown_early(bool CanJoinThreads = true) {
   if (!GlobalHandler::RTGlobalObjHandler)
     return;
 
+#ifdef _WIN32
+  // Threads are terminated without unwinding during process shutdown, so
+  // scheduler-access guards held by those threads cannot be released.
+  if (CanJoinThreads)
+    ProcessTerminating.store(true, std::memory_order_release);
+#endif
+
 #if defined(XPTI_ENABLE_INSTRUMENTATION) && defined(_WIN32)
   if (xptiTraceEnabled())
     return; // When doing xpti tracing, we can't safely shutdown on Win.
@@ -380,7 +435,15 @@ void shutdown_late() {
 #endif
 
   // First, release resources, that may access adapters.
+  bool WaitForSchedulerAccess = true;
+#ifdef _WIN32
+  WaitForSchedulerAccess = !ProcessTerminating.load(std::memory_order_acquire);
+#endif
+  GlobalHandler::RTGlobalObjHandler->stopSchedulerAccess(
+      WaitForSchedulerAccess);
   GlobalHandler::RTGlobalObjHandler->MPlatformCache.Inst.reset(nullptr);
+  GlobalHandler::RTGlobalObjHandler->MSchedulerPtr.store(
+      nullptr, std::memory_order_release);
   GlobalHandler::RTGlobalObjHandler->MScheduler.Inst.reset(nullptr);
   GlobalHandler::RTGlobalObjHandler->MProgramManager.Inst.reset(nullptr);
 

sycl/source/detail/global_handler.hpp

@@ -11,6 +11,8 @@
 #include <sycl/detail/spinlock.hpp>
 #include <sycl/detail/util.hpp>
 
+#include <atomic>
+#include <cstdint>
 #include <memory>
 #include <unordered_map>
 
@@ -41,6 +43,22 @@ class ThreadPool;
 /// construction or destruction is generated anyway.
 class GlobalHandler {
 public:
+  class SchedulerAccess {
+  public:
+    SchedulerAccess() = default;
+    SchedulerAccess(const SchedulerAccess &) = delete;
+    SchedulerAccess &operator=(const SchedulerAccess &) = delete;
+    ~SchedulerAccess();
+
+    Scheduler *get() const noexcept { return MScheduler; }
+
+  private:
+    friend class GlobalHandler;
+    explicit SchedulerAccess(Scheduler *Sched) : MScheduler(Sched) {}
+
+    Scheduler *MScheduler = nullptr;
+  };
+
   static bool isInstanceAlive() { return RTGlobalObjHandler != nullptr; }
   /// \return a reference to a GlobalHandler singleton instance. The reference
   /// is valid as long as runtime library is loaded (i.e. until `DllMain` or
@@ -54,6 +72,10 @@ class GlobalHandler {
   void registerSchedulerUsage(bool ModifyCounter = true);
   Scheduler &getScheduler();
   bool isSchedulerAlive() const;
+  /// \return the Scheduler if it has been created, nullptr otherwise. Unlike
+  /// getScheduler(), takes no lock and never creates the Scheduler. The
+  /// returned access object keeps the Scheduler alive until it is destroyed.
+  static SchedulerAccess getSchedulerAccess() noexcept;
   ProgramManager &getProgramManager();
   Sync &getSync();
   std::vector<std::shared_ptr<platform_impl>> &getPlatformCache();
@@ -89,18 +111,28 @@ class GlobalHandler {
 
   // Used in SYCL unit tests to reset the GlobalHandler instance.
   static void resetGlobalHandler() {
+    stopSchedulerAccess();
+    MSchedulerPtr.store(nullptr, std::memory_order_release);
+    MSchedulerAccessState.store(0, std::memory_order_release);
     RTGlobalObjHandler = new GlobalHandler();
   };
 
   // Used in SYCL unit tests to simulate runtime teardown; pair with
   // restoreGlobalHandler().
   static GlobalHandler *detachGlobalHandler() {
+    stopSchedulerAccess();
+    MSchedulerPtr.store(nullptr, std::memory_order_release);
     GlobalHandler *Old = RTGlobalObjHandler;
     RTGlobalObjHandler = nullptr;
     return Old;
   }
   static void restoreGlobalHandler(GlobalHandler *Handler) {
     RTGlobalObjHandler = Handler;
+    MSchedulerPtr.store(Handler && Handler->MScheduler.Inst
+                            ? Handler->MScheduler.Inst.get()
+                            : nullptr,
+                        std::memory_order_release);
+    MSchedulerAccessState.store(0, std::memory_order_release);
   }
 
 private:
@@ -124,7 +156,17 @@ class GlobalHandler {
   template <typename T, typename... Types>
   T &getOrCreate(InstWithLock<T> &IWL, Types &&...Args);
 
+  static void releaseSchedulerAccess() noexcept;
+  static void stopSchedulerAccess(bool WaitForAccess = true) noexcept;
+
   InstWithLock<Scheduler> MScheduler;
+  // Mirror of MScheduler.Inst for lock-free readers. Published with release
+  // semantics once the Scheduler is fully constructed and cleared before it is
+  // destroyed.
+  static std::atomic<Scheduler *> MSchedulerPtr;
+  // The high bit prevents new accesses during shutdown; the remaining bits
+  // count accesses that must finish before the Scheduler can be destroyed.
+  static std::atomic<uint64_t> MSchedulerAccessState;
   InstWithLock<ProgramManager> MProgramManager;
   InstWithLock<Sync> MSync;
   InstWithLock<std::vector<std::shared_ptr<platform_impl>>> MPlatformCache;

sycl/source/detail/queue_impl.cpp

@@ -8,6 +8,7 @@
 
 #include <detail/event_deps.hpp>
 #include <detail/event_impl.hpp>
+#include <detail/global_handler.hpp>
 #include <detail/memory_manager.hpp>
 #include <detail/queue_impl.hpp>
 #include <detail/scheduler/commands.hpp>
@@ -1456,6 +1457,14 @@ void queue_impl::wait(const detail::code_location &CodeLoc) {
       Event->wait();
   }
 
+  {
+    auto SchedAccess = GlobalHandler::getSchedulerAccess();
+    if (Scheduler *Sched = SchedAccess.get()) {
+      if (Sched->hasDeferredResources())
+        Sched->releaseResources(BlockingT::NON_BLOCKING);
+    }
+  }
+
 #ifdef XPTI_ENABLE_INSTRUMENTATION
   // There is an early return in instrumentationEpilog() if no subscribers are
   // subscribing to queue.wait().