[SYCL] Fix asynchronous exception behavior (#20274)

This commit makes the following changes to the behavior of asynchronous
exception handling:

1. The death of a queue should not consume asynchronous exceptions.
2. Calling wait_and_throw on an event after the associated queue has
died should still consume exceptions that were originally associated
with the queue. This should respect the async_handler priority to the
best of its ability.
3. Calling wait_and_throw or throw_asynchronous on a queue without an
async_handler should fall back to using the async_handler of the
associated context, then the default async_handler if none were attached
to the context.
4. Throwing asynchronous exceptions from anywhere will now consume all
unconsumed asynchronous exceptions previously reported, no matter the
event/queue/context/device.

Additionally, this lays the ground work for
#20266 by moving the tracking of
unconsumed asynchronous exception to the devices.

---------

Signed-off-by: Larsen, Steffen <[email protected]>

Author: steffenlarsen

sycl/include/sycl/exception_list.hpp

@@ -23,7 +23,7 @@ inline namespace _V1 {
 
 // Forward declaration
 namespace detail {
-class queue_impl;
+class Scheduler;
 }
 
 /// A list of asynchronous exceptions.
@@ -45,7 +45,7 @@ class __SYCL_EXPORT exception_list {
   iterator end() const;
 
 private:
-  friend class detail::queue_impl;
+  friend class detail::Scheduler;
   void PushBack(const_reference Value);
   void PushBack(value_type &&Value);
   void Clear() noexcept;

sycl/source/detail/event_impl.cpp

@@ -211,8 +211,8 @@ void event_impl::initHostProfilingInfo() {
   MHostProfilingInfo->setDevice(&Device);
 }
 
-void event_impl::setSubmittedQueue(std::weak_ptr<queue_impl> SubmittedQueue) {
-  MSubmittedQueue = std::move(SubmittedQueue);
+void event_impl::setSubmittedQueue(queue_impl *SubmittedQueue) {
+  MSubmittedQueue = SubmittedQueue->weak_from_this();
 }
 
 #ifdef XPTI_ENABLE_INSTRUMENTATION
@@ -307,9 +307,7 @@ void event_impl::wait(bool *Success) {
 
 void event_impl::wait_and_throw() {
   wait();
-
-  if (std::shared_ptr<queue_impl> SubmittedQueue = MSubmittedQueue.lock())
-    SubmittedQueue->throw_asynchronous();
+  Scheduler::getInstance().flushAsyncExceptions();
 }
 
 void event_impl::checkProfilingPreconditions() const {

sycl/source/detail/event_impl.hpp

@@ -267,7 +267,7 @@ class event_impl {
   /// Sets original queue used for submission.
   ///
   /// @return
-  void setSubmittedQueue(std::weak_ptr<queue_impl> SubmittedQueue);
+  void setSubmittedQueue(queue_impl *SubmittedQueue);
 
   /// Indicates if this event is not associated with any command and doesn't
   /// have native handle.

sycl/source/detail/queue_impl.cpp

@@ -514,7 +514,7 @@ EventImplPtr queue_impl::submit_command_to_graph(
     std::unique_ptr<detail::CG> CommandGroup, sycl::detail::CGType CGType,
     sycl::ext::oneapi::experimental::node_type UserFacingNodeType) {
   auto EventImpl = detail::event_impl::create_completed_host_event();
-  EventImpl->setSubmittedQueue(weak_from_this());
+  EventImpl->setSubmittedQueue(this);
   ext::oneapi::experimental::detail::node_impl *NodeImpl = nullptr;
 
   // GraphImpl is read and written in this scope so we lock this graph

sycl/source/detail/queue_impl.hpp

@@ -253,7 +253,6 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
       // notification and destroy the trace event for this queue.
       destructorNotification();
 #endif
-      throw_asynchronous();
       auto status =
           getAdapter().call_nocheck<UrApiKind::urQueueRelease>(MQueue);
       // If loader is already closed, it'll return a not-initialized status
@@ -296,6 +295,8 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
 
   context_impl &getContextImpl() const { return *MContext; }
 
+  std::weak_ptr<context_impl> getContextImplWeakPtr() const { return MContext; }
+
   device_impl &getDeviceImpl() const { return MDevice; }
 
   /// \return an associated SYCL device.
@@ -411,37 +412,18 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
   /// @param Loc is the code location of the submit call (default argument)
   void wait(const detail::code_location &Loc = {});
 
-  /// \return list of asynchronous exceptions occurred during execution.
-  exception_list getExceptionList() const { return MExceptions; }
-
   /// @param Loc is the code location of the submit call (default argument)
   void wait_and_throw(const detail::code_location &Loc = {}) {
     wait(Loc);
     throw_asynchronous();
   }
 
-  /// Performs a blocking wait for the completion of all enqueued tasks in the
-  /// queue.
-  ///
   /// Synchronous errors will be reported through SYCL exceptions.
   /// Asynchronous errors will be passed to the async_handler passed to the
   /// queue on construction. If no async_handler was provided then
-  /// asynchronous exceptions will be lost.
-  void throw_asynchronous() {
-    if (!MAsyncHandler)
-      return;
-
-    exception_list Exceptions;
-    {
-      std::lock_guard<std::mutex> Lock(MMutex);
-      std::swap(Exceptions, MExceptions);
-    }
-    // Unlock the mutex before calling user-provided handler to avoid
-    // potential deadlock if the same queue is somehow referenced in the
-    // handler.
-    if (Exceptions.size())
-      MAsyncHandler(std::move(Exceptions));
-  }
+  /// asynchronous exceptions will be passed to the async_handler associated
+  /// with the context if present, or the default async_handler otherwise.
+  void throw_asynchronous() { Scheduler::getInstance().flushAsyncExceptions(); }
 
   /// Creates UR properties array.
   ///
@@ -588,14 +570,6 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
   event mem_advise(const void *Ptr, size_t Length, ur_usm_advice_flags_t Advice,
                    const std::vector<event> &DepEvents, bool CallerNeedsEvent);
 
-  /// Puts exception to the list of asynchronous ecxeptions.
-  ///
-  /// \param ExceptionPtr is a pointer to exception to be put.
-  void reportAsyncException(const std::exception_ptr &ExceptionPtr) {
-    std::lock_guard<std::mutex> Lock(MMutex);
-    MExceptions.PushBack(ExceptionPtr);
-  }
-
   static ThreadPool &getThreadPool() {
     return GlobalHandler::instance().getHostTaskThreadPool();
   }
@@ -716,6 +690,11 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
   }
 #endif
 
+  /// Returns the async_handler associated with the queue.
+  const async_handler &getAsynchHandler() const noexcept {
+    return MAsyncHandler;
+  }
+
 protected:
   template <typename HandlerType = handler>
   EventImplPtr insertHelperBarrier(const HandlerType &Handler) {
@@ -1003,10 +982,6 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
   /// These events are tracked, but not owned, by the queue.
   std::vector<std::weak_ptr<event_impl>> MEventsWeak;
 
-  /// Events without data dependencies (such as USM) need an owner,
-  /// additionally, USM operations are not added to the scheduler command graph,
-  /// queue is the only owner on the runtime side.
-  exception_list MExceptions;
   const async_handler MAsyncHandler;
   const property_list MPropList;
 

sycl/source/detail/scheduler/commands.cpp

@@ -359,12 +359,14 @@ class DispatchHostTask {
         AdapterWithEvents.first->call<UrApiKind::urEventWait>(RawEvents.size(),
                                                               RawEvents.data());
       } catch (const sycl::exception &) {
-        MThisCmd->MEvent->getSubmittedQueue()->reportAsyncException(
-            std::current_exception());
+        auto QueuePtr = MThisCmd->MEvent->getSubmittedQueue();
+        Scheduler::getInstance().reportAsyncException(QueuePtr,
+                                                      std::current_exception());
         return false;
       } catch (...) {
-        MThisCmd->MEvent->getSubmittedQueue()->reportAsyncException(
-            std::current_exception());
+        auto QueuePtr = MThisCmd->MEvent->getSubmittedQueue();
+        Scheduler::getInstance().reportAsyncException(QueuePtr,
+                                                      std::current_exception());
         return false;
       }
     }
@@ -407,10 +409,12 @@ class DispatchHostTask {
           make_error_code(errc::runtime),
           std::string("Couldn't wait for host-task's dependencies")));
 
-      MThisCmd->MEvent->getSubmittedQueue()->reportAsyncException(EPtr);
+      auto QueuePtr = MThisCmd->MEvent->getSubmittedQueue();
+      auto &SchedulerInst = Scheduler::getInstance();
+      SchedulerInst.reportAsyncException(QueuePtr, EPtr);
       // reset host-task's lambda and quit
       HostTask.MHostTask.reset();
-      Scheduler::getInstance().NotifyHostTaskCompletion(MThisCmd);
+      SchedulerInst.NotifyHostTaskCompletion(MThisCmd);
       return;
     }
 
@@ -469,8 +473,8 @@ class DispatchHostTask {
         }
       }
 #endif
-      MThisCmd->MEvent->getSubmittedQueue()->reportAsyncException(
-          CurrentException);
+      auto QueuePtr = MThisCmd->MEvent->getSubmittedQueue();
+      Scheduler::getInstance().reportAsyncException(QueuePtr, CurrentException);
     }
 
     HostTask.MHostTask.reset();
@@ -487,8 +491,8 @@ class DispatchHostTask {
       Scheduler::getInstance().NotifyHostTaskCompletion(MThisCmd);
     } catch (...) {
       auto CurrentException = std::current_exception();
-      MThisCmd->MEvent->getSubmittedQueue()->reportAsyncException(
-          CurrentException);
+      auto QueuePtr = MThisCmd->MEvent->getSubmittedQueue();
+      Scheduler::getInstance().reportAsyncException(QueuePtr, CurrentException);
     }
   }
 };
@@ -563,7 +567,8 @@ Command::Command(
       MCommandBuffer(CommandBuffer), MSyncPointDeps(SyncPoints) {
   MWorkerQueue = MQueue;
   MEvent->setWorkerQueue(MWorkerQueue);
-  MEvent->setSubmittedQueue(MWorkerQueue);
+  if (Queue)
+    MEvent->setSubmittedQueue(Queue);
   MEvent->setCommand(this);
   if (MQueue)
     MEvent->setContextImpl(MQueue->getContextImpl());
@@ -1958,7 +1963,7 @@ ExecCGCommand::ExecCGCommand(
     assert(SubmitQueue &&
            "Host task command group must have a valid submit queue");
 
-    MEvent->setSubmittedQueue(SubmitQueue->weak_from_this());
+    MEvent->setSubmittedQueue(SubmitQueue);
     // Initialize host profiling info if the queue has profiling enabled.
     if (SubmitQueue->MIsProfilingEnabled)
       MEvent->initHostProfilingInfo();

sycl/source/detail/scheduler/scheduler.cpp

@@ -260,7 +260,7 @@ EventImplPtr Scheduler::addCopyBack(Requirement *Req) {
       auto WorkerQueue = NewCmd->getEvent()->getWorkerQueue();
       assert(WorkerQueue &&
              "WorkerQueue for CopyBack command must be not null");
-      WorkerQueue->reportAsyncException(std::current_exception());
+      reportAsyncException(WorkerQueue, std::current_exception());
     }
   }
   EventImplPtr NewEvent = NewCmd->getEvent();
@@ -706,6 +706,45 @@ bool Scheduler::areEventsSafeForSchedulerBypass(events_range DepEvents,
     return Event.getHandle() != nullptr;
   });
 }
+
+void Scheduler::reportAsyncException(
+    const std::shared_ptr<queue_impl> &QueuePtr,
+    const std::exception_ptr &ExceptionPtr) {
+  std::lock_guard<std::mutex> Lock(MAsyncExceptionsMutex);
+  MAsyncExceptions[AsyncExceptionKey{QueuePtr,
+                                     QueuePtr->getContextImplWeakPtr()}]
+      .PushBack(ExceptionPtr);
+}
+
+void Scheduler::flushAsyncExceptions() {
+  decltype(MAsyncExceptions) AsyncExceptions;
+  {
+    std::lock_guard<std::mutex> Lock(MAsyncExceptionsMutex);
+    std::swap(AsyncExceptions, MAsyncExceptions);
+  }
+  for (auto &ExceptionsEntryIt : AsyncExceptions) {
+    exception_list Exceptions = std::move(ExceptionsEntryIt.second);
+
+    if (Exceptions.size() == 0)
+      continue;
+
+    std::shared_ptr<queue_impl> Queue = ExceptionsEntryIt.first.first.lock();
+    if (Queue && Queue->getAsynchHandler()) {
+      Queue->getAsynchHandler()(std::move(Exceptions));
+      continue;
+    }
+
+    std::shared_ptr<context_impl> Context =
+        ExceptionsEntryIt.first.second.lock();
+    if (Context && Context->get_async_handler()) {
+      Context->get_async_handler()(std::move(Exceptions));
+    } else {
+      // If the queue is dead, use the default handler.
+      defaultAsyncHandler(std::move(Exceptions));
+    }
+  }
+}
+
 } // namespace detail
 } // namespace _V1
 } // namespace sycl

sycl/source/detail/scheduler/scheduler.hpp

@@ -485,6 +485,18 @@ class Scheduler {
   static bool areEventsSafeForSchedulerBypass(events_range DepEvents,
                                               context_impl &Context);
 
+  /// Puts exception to the list of asynchronous ecxeptions.
+  ///
+  /// \param QueuePtr is a pointer referring to the queue to report the
+  /// asynchronous exceptions for.
+  /// \param ExceptionPtr is a pointer to exception to be put.
+  void reportAsyncException(const std::shared_ptr<queue_impl> &QueuePtr,
+                            const std::exception_ptr &ExceptionPtr);
+
+  /// Reports all unconsumed asynchronous exceptions to either the queue's
+  /// async_handler, the context's async_handler or the default async_handler.
+  void flushAsyncExceptions();
+
 protected:
   using RWLockT = std::shared_timed_mutex;
   using ReadLockT = std::shared_lock<RWLockT>;
@@ -872,6 +884,23 @@ class Scheduler {
       MAuxiliaryResources;
   std::mutex MAuxiliaryResourcesMutex;
 
+  // Asynchronous exceptions are captured at device-level until flushed, either
+  // by queues, events or a synchronization on the device itself.
+  std::mutex MAsyncExceptionsMutex;
+  using AsyncExceptionKey =
+      std::pair<std::weak_ptr<queue_impl>, std::weak_ptr<context_impl>>;
+  struct AsyncExceptionKeyOwnerLess {
+    bool operator()(const AsyncExceptionKey &LHS,
+                    const AsyncExceptionKey &RHS) const noexcept {
+      return std::owner_less<std::weak_ptr<queue_impl>>{}(LHS.first,
+                                                          RHS.first) ||
+             std::owner_less<std::weak_ptr<context_impl>>{}(LHS.second,
+                                                            RHS.second);
+    }
+  };
+  std::map<AsyncExceptionKey, exception_list, AsyncExceptionKeyOwnerLess>
+      MAsyncExceptions;
+
   friend class Command;
   friend class DispatchHostTask;
   friend class queue_impl;