[SYCL][UR] Implement sycl_ext_oneapi_device_wait

This commit implements the UR functionality for device-wide
synchronization and the SYCL APIs using it. The latter implements the
sycl_ext_oneapi_device_wait extension.

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

Author: steffenlarsen

llvm/include/llvm/SYCLLowerIR/DeviceConfigFile.td

@@ -98,6 +98,7 @@ def Aspectext_oneapi_clock_sub_group : Aspect<"ext_oneapi_clock_sub_group">;
 def Aspectext_oneapi_clock_work_group : Aspect<"ext_oneapi_clock_work_group">;
 def Aspectext_oneapi_clock_device : Aspect<"ext_oneapi_clock_device">;
 def Aspectext_oneapi_is_integrated_gpu : Aspect<"ext_oneapi_is_integrated_gpu">;
+def Aspectext_oneapi_device_wait : Aspect<"ext_oneapi_device_wait">;
 
 // Deprecated aspects
 def AspectInt64_base_atomics : Aspect<"int64_base_atomics">;
@@ -176,7 +177,8 @@ def : TargetInfo<"__TestAspectList",
     Aspectext_oneapi_clock_sub_group,
     Aspectext_oneapi_clock_work_group,
     Aspectext_oneapi_clock_device,
-    Aspectext_oneapi_is_integrated_gpu],
+    Aspectext_oneapi_is_integrated_gpu,
+    Aspectext_oneapi_device_wait],
     []>;
 // This definition serves the only purpose of testing whether the deprecated aspect list defined in here and in SYCL RT
 // match.

sycl/include/sycl/device.hpp

@@ -365,6 +365,21 @@ class __SYCL_STANDALONE_DEBUG __SYCL_EXPORT device
     return profile.c_str();
   }
 
+  /// Synchronizes with all queues associated with the device.
+  void ext_oneapi_wait();
+
+  /// Dispatches all unconsumed asynchronous exceptions for all queues or
+  /// contexts associated with the queues.
+  void ext_oneapi_throw_asynchronous();
+
+  /// Synchronizes with all queues associated with the device, then dispatches
+  /// all unconsumed asynchronous exceptions for all queues or contexts
+  /// associated with the queues.
+  void ext_oneapi_wait_and_throw() {
+    ext_oneapi_wait();
+    ext_oneapi_throw_asynchronous();
+  }
+
 // TODO: Remove this diagnostics when __SYCL_WARN_IMAGE_ASPECT is removed.
 #if defined(__clang__)
 #pragma clang diagnostic pop

sycl/include/sycl/info/aspects.def

@@ -84,3 +84,4 @@ __SYCL_ASPECT(ext_oneapi_clock_sub_group, 91)
 __SYCL_ASPECT(ext_oneapi_clock_work_group, 92)
 __SYCL_ASPECT(ext_oneapi_clock_device, 93)
 __SYCL_ASPECT(ext_oneapi_is_integrated_gpu, 94)
+__SYCL_ASPECT(ext_oneapi_device_wait, 95)

sycl/source/detail/device_impl.cpp

@@ -503,6 +503,34 @@ device_impl::getImmediateProgressGuarantee(
   return forward_progress_guarantee::weakly_parallel;
 }
 
+void device_impl::wait() const {
+  // Firstly, all associated queues should be cleaned through of all
+  // not-yet-enqueued commands and host_task.
+  for (const std::weak_ptr<queue_impl> &WQueue : MQueues) {
+    std::shared_ptr<queue_impl> Queue = WQueue.lock();
+    assert(Queue && "Queue should never be dangling in the list of queues "
+                    "associated with the device!");
+    Queue->waitForRuntimeLevelCmdsAndClear();
+  }
+
+  // Then we synchronize the entire device.
+  getAdapter().call<detail::UrApiKind::urDeviceWaitExp>(getHandleRef());
+}
+
+void device_impl::throwAsynchronous() {
+  std::lock_guard<std::mutex> Lock(MAsyncExceptionsMutex);
+  for (auto &ExceptionsEntryIt : MAsyncExceptions) {
+    exception_list Exceptions = std::move(ExceptionsEntryIt.second);
+    std::shared_ptr<queue_impl> Queue = ExceptionsEntryIt.first.lock();
+    if (Queue && Queue->getAsynchHandler()) {
+      Queue->getAsynchHandler()(std::move(Exceptions));
+    } else {
+      // If the queue is dead, use the default handler.
+      defaultAsyncHandler(std::move(Exceptions));
+    }
+  }
+}
+
 #ifndef __INTEL_PREVIEW_BREAKING_CHANGES
 #define EXPORT_GET_INFO(PARAM)                                                 \
   template <>                                                                  \

sycl/source/detail/device_impl.hpp

@@ -1597,6 +1597,10 @@ class device_impl : public std::enable_shared_from_this<device_impl> {
              get_info_impl_nocheck<UR_DEVICE_INFO_IS_INTEGRATED_GPU>().value_or(
                  0);
     }
+    CASE(ext_oneapi_device_wait) {
+      return get_info_impl_nocheck<UR_DEVICE_INFO_DEVICE_WAIT_SUPPORT_EXP>()
+          .value_or(0);
+    }
     else {
       return false; // This device aspect has not been implemented yet.
     }
@@ -2292,6 +2296,22 @@ class device_impl : public std::enable_shared_from_this<device_impl> {
     return Exceptions;
   }
 
+  /// Synchronizes with all queues on the device.
+  void wait() const;
+
+  // Dispatch all unconsumed asynchronous exception to the appropriate handlers.
+  void throwAsynchronous();
+
+  void registerQueue(const std::weak_ptr<queue_impl> &Q) {
+    std::lock_guard<std::mutex> Lock(MQueuesMutex);
+    MQueues.insert(Q);
+  }
+
+  void unregisterQueue(const std::weak_ptr<queue_impl> &Q) {
+    std::lock_guard<std::mutex> Lock(MQueuesMutex);
+    MQueues.erase(Q);
+  }
+
 private:
   ur_device_handle_t MDevice = 0;
   // This is used for getAdapter so should be above other properties.
@@ -2302,6 +2322,13 @@ class device_impl : public std::enable_shared_from_this<device_impl> {
 
   const ur_device_handle_t MRootDevice;
 
+  // Devices track a list of active queues on it, to allow for synchronization
+  // with host_task and not-yet-enqueued commands.
+  std::mutex MQueuesMutex;
+  std::set<std::weak_ptr<queue_impl>,
+           std::owner_less<std::weak_ptr<queue_impl>>>
+      MQueues;
+
   // Asynchronous exceptions are captured at device-level until flushed, either
   // by queues, events or a synchronization on the device itself.
   std::mutex MAsyncExceptionsMutex;

sycl/source/detail/queue_impl.cpp

@@ -889,32 +889,7 @@ void queue_impl::wait(const detail::code_location &CodeLoc) {
       LastEvent->wait();
     }
   } else if (!isInOrder()) {
-    std::vector<std::weak_ptr<event_impl>> WeakEvents;
-    {
-      std::lock_guard<std::mutex> Lock(MMutex);
-      WeakEvents.swap(MEventsWeak);
-      MMissedCleanupRequests.unset(
-          [&](MissedCleanupRequestsType &MissedCleanupRequests) {
-            for (auto &UpdatedGraph : MissedCleanupRequests)
-              doUnenqueuedCommandCleanup(UpdatedGraph);
-            MissedCleanupRequests.clear();
-          });
-    }
-
-    // Wait for unenqueued or host task events, starting
-    // from the latest submitted task in order to minimize total amount of
-    // calls, then handle the rest with urQueueFinish.
-    for (auto EventImplWeakPtrIt = WeakEvents.rbegin();
-         EventImplWeakPtrIt != WeakEvents.rend(); ++EventImplWeakPtrIt) {
-      if (std::shared_ptr<event_impl> EventImplSharedPtr =
-              EventImplWeakPtrIt->lock()) {
-        // A nullptr UR event indicates that urQueueFinish will not cover it,
-        // either because it's a host task event or an unenqueued one.
-        if (nullptr == EventImplSharedPtr->getHandle()) {
-          EventImplSharedPtr->wait();
-        }
-      }
-    }
+    waitForRuntimeLevelCmdsAndClear();
   }
 
   getAdapter().call<UrApiKind::urQueueFinish>(getHandleRef());
@@ -1127,6 +1102,47 @@ void queue_impl::verifyProps(const property_list &Props) const {
                                                 CheckPropertiesWithData);
 }
 
+void queue_impl::waitForRuntimeLevelCmdsAndClear() {
+  if (isInOrder() && !MNoLastEventMode.load(std::memory_order_relaxed)) {
+    // if MLastEvent is not null and has no associated handle, we need to wait
+    // for it. We do not clear it however.
+    EventImplPtr LastEvent;
+    {
+      std::lock_guard<std::mutex> Lock(MMutex);
+      LastEvent = MDefaultGraphDeps.LastEventPtr;
+    }
+    if (LastEvent && nullptr == LastEvent->getHandle())
+      LastEvent->wait();
+  } else if (!isInOrder()) {
+    std::vector<std::weak_ptr<event_impl>> WeakEvents;
+    {
+      std::lock_guard<std::mutex> Lock(MMutex);
+      WeakEvents.swap(MEventsWeak);
+      MMissedCleanupRequests.unset(
+          [&](MissedCleanupRequestsType &MissedCleanupRequests) {
+            for (auto &UpdatedGraph : MissedCleanupRequests)
+              doUnenqueuedCommandCleanup(UpdatedGraph);
+            MissedCleanupRequests.clear();
+          });
+    }
+
+    // Wait for unenqueued or host task events, starting
+    // from the latest submitted task in order to minimize total amount of
+    // calls, then handle the rest with urQueueFinish.
+    for (auto EventImplWeakPtrIt = WeakEvents.rbegin();
+         EventImplWeakPtrIt != WeakEvents.rend(); ++EventImplWeakPtrIt) {
+      if (std::shared_ptr<event_impl> EventImplSharedPtr =
+              EventImplWeakPtrIt->lock()) {
+        // A nullptr UR event indicates that urQueueFinish will not cover it,
+        // either because it's a host task event or an unenqueued one.
+        if (nullptr == EventImplSharedPtr->getHandle()) {
+          EventImplSharedPtr->wait();
+        }
+      }
+    }
+  }
+}
+
 } // namespace detail
 } // namespace _V1
 } // namespace sycl

sycl/source/detail/queue_impl.hpp

@@ -241,8 +241,10 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
   // `std::shared_ptr` allocations.
   template <typename... Ts>
   static std::shared_ptr<queue_impl> create(Ts &&...args) {
-    return std::make_shared<queue_impl>(std::forward<Ts>(args)...,
-                                        private_tag{});
+    auto ImplPtr =
+        std::make_shared<queue_impl>(std::forward<Ts>(args)..., private_tag{});
+    ImplPtr->getDeviceImpl().registerQueue(ImplPtr);
+    return ImplPtr;
   }
 
   ~queue_impl() {
@@ -253,6 +255,7 @@ class queue_impl : public std::enable_shared_from_this<queue_impl> {
       // notification and destroy the trace event for this queue.
       destructorNotification();
 #endif
+      MDevice.unregisterQueue(weak_from_this());
       auto status =
           getAdapter().call_nocheck<UrApiKind::urQueueRelease>(MQueue);
       // If loader is already closed, it'll return a not-initialized status
@@ -704,6 +707,17 @@ 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;
+  }
+
+  /// Waits for all not-yet-enqueued and host_task commands in the queue and
+  /// clears the events associated with the queue (if out-of-order.)
+  /// Note: This should only be called if the queue is guaranteed to be
+  /// synchronized by the caller.
+  void waitForRuntimeLevelCmdsAndClear();
+
 protected:
   template <typename HandlerType = handler>
   EventImplPtr insertHelperBarrier(const HandlerType &Handler) {

sycl/source/detail/ur_device_info_ret_types.inc

@@ -197,4 +197,5 @@ MAP(UR_DEVICE_INFO_CLOCK_SUB_GROUP_SUPPORT_EXP, ur_bool_t)
 MAP(UR_DEVICE_INFO_CLOCK_WORK_GROUP_SUPPORT_EXP, ur_bool_t)
 MAP(UR_DEVICE_INFO_CLOCK_DEVICE_SUPPORT_EXP, ur_bool_t)
 MAP(UR_DEVICE_INFO_IS_INTEGRATED_GPU, ur_bool_t)
+MAP(UR_DEVICE_INFO_DEVICE_WAIT_SUPPORT_EXP, ur_bool_t)
 // clang-format on

sycl/source/device.cpp

@@ -344,5 +344,9 @@ detail::string device::ext_oneapi_cl_profile_impl() const {
   return detail::string{profile};
 }
 
+void device::ext_oneapi_wait() { impl->wait(); }
+
+void device::ext_oneapi_throw_asynchronous() { impl->throwAsynchronous(); }
+
 } // namespace _V1
 } // namespace sycl

sycl/test-e2e/DeviceWait/basic.cpp

@@ -0,0 +1,48 @@
+// RUN: %{build} -o %t.out
+// RUN: %{run} %t.out
+
+#include <sycl/detail/core.hpp>
+#include <sycl/properties/all_properties.hpp>
+
+#include <array>
+#include <vector>
+
+constexpr size_t NContexts = 2;
+constexpr size_t NQueues = 6;
+
+int main() {
+  sycl::device D;
+  std::array<sycl::context, NContexts> Contexts{sycl::context{D},
+                                                sycl::context{D}};
+  std::array<sycl::queue, NQueues> Queues{
+      sycl::queue{Contexts[0], D},
+      sycl::queue{Contexts[0], D, sycl::property::queue::in_order()},
+      sycl::queue{Contexts[0], D},
+      sycl::queue{Contexts[1], D, sycl::property::queue::in_order()},
+      sycl::queue{Contexts[1], D},
+      sycl::queue{Contexts[1], D, sycl::property::queue::in_order()}};
+
+  std::vector<sycl::event> Events;
+  Events.reserve(NQueues);
+  for (sycl::queue &Q : Queues) {
+    sycl::event E = Q.single_task([]() {
+      volatile int value = 1024 * 1024;
+      while (--value)
+        ;
+    });
+    Events.push_back(std::move(E));
+  }
+
+  D.ext_oneapi_wait();
+
+  int Failed = 0;
+  for (size_t I = 0; I < Events.size(); ++I) {
+    sycl::info::event_command_status EventStatus =
+        Events[I].get_info<sycl::info::event::command_execution_status>();
+    if (EventStatus != sycl::info::event_command_status::complete) {
+      std::cout << "Unexpected event status for event at " << I << std::endl;
+      ++Failed;
+    }
+  }
+  return Failed;
+}