[UR][CUDA][HIP] Cleanup and share queue interop code (#17927)

This patch should essentially be NFC. The way it's currently setup is
still very hacky and could use further refactoring, but this patch just
isolates the hack a little bit and makes it more obvious what's going
on, as well as removing some code duplication.

Author: npmiller

unified-runtime/source/adapters/cuda/enqueue_native.cpp

@@ -25,7 +25,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueNativeCommandExp(
 
   try {
     ScopedContext ActiveContext(hQueue->getDevice());
-    ScopedStream ActiveStream(hQueue, NumEventsInWaitList, phEventWaitList);
+    InteropGuard ActiveStream(hQueue, NumEventsInWaitList, phEventWaitList);
     std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
 
     if (hQueue->getContext()->getDevices().size() > 1) {

unified-runtime/source/adapters/cuda/queue.cpp

@@ -180,7 +180,7 @@ urQueueGetNativeHandle(ur_queue_handle_t hQueue, ur_queue_native_desc_t *pDesc,
 
   ScopedContext Active(hQueue->getDevice());
   *phNativeQueue =
-      reinterpret_cast<ur_native_handle_t>(hQueue->getNextComputeStream());
+      reinterpret_cast<ur_native_handle_t>(hQueue->getInteropStream());
   return UR_RESULT_SUCCESS;
 }
 

unified-runtime/source/adapters/cuda/queue.hpp

@@ -22,6 +22,8 @@
 using cuda_stream_queue = stream_queue_t<CUstream, 128, 64, CUevent>;
 struct ur_queue_handle_t_ : public cuda_stream_queue {};
 
+using InteropGuard = cuda_stream_queue::interop_guard;
+
 // Function which creates the profiling stream. Called only from makeNative
 // event when profiling is required.
 template <> inline void cuda_stream_queue::createHostSubmitTimeStream() {
@@ -38,24 +40,3 @@ inline void cuda_stream_queue::createStreamWithPriority(CUstream *Stream,
                                                         int Priority) {
   UR_CHECK_ERROR(cuStreamCreateWithPriority(Stream, Flags, Priority));
 }
-
-// RAII object to make hQueue stream getter methods all return the same stream
-// within the lifetime of this object.
-//
-// This is useful for urEnqueueNativeCommandExp where we want guarantees that
-// the user submitted native calls will be dispatched to a known stream, which
-// must be "got" within the user submitted fuction.
-class ScopedStream {
-  ur_queue_handle_t hQueue;
-
-public:
-  ScopedStream(ur_queue_handle_t hQueue, uint32_t NumEventsInWaitList,
-               const ur_event_handle_t *EventWaitList)
-      : hQueue{hQueue} {
-    ur_stream_guard Guard;
-    hQueue->getThreadLocalStream() =
-        hQueue->getNextComputeStream(NumEventsInWaitList, EventWaitList, Guard);
-  }
-  CUstream getStream() { return hQueue->getThreadLocalStream(); }
-  ~ScopedStream() { hQueue->getThreadLocalStream() = CUstream{0}; }
-};

unified-runtime/source/adapters/hip/enqueue_native.cpp

@@ -28,7 +28,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueNativeCommandExp(
 
   try {
     ScopedDevice ActiveDevice(hQueue->getDevice());
-    ScopedStream ActiveStream(hQueue, NumEventsInWaitList, phEventWaitList);
+    InteropGuard ActiveStream(hQueue, NumEventsInWaitList, phEventWaitList);
     std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
 
     if (hQueue->getContext()->getDevices().size() > 1) {

unified-runtime/source/adapters/hip/queue.cpp

@@ -208,7 +208,7 @@ urQueueGetNativeHandle(ur_queue_handle_t hQueue, ur_queue_native_desc_t *,
                        ur_native_handle_t *phNativeQueue) {
   ScopedDevice Active(hQueue->getDevice());
   *phNativeQueue =
-      reinterpret_cast<ur_native_handle_t>(hQueue->getNextComputeStream());
+      reinterpret_cast<ur_native_handle_t>(hQueue->getInteropStream());
   return UR_RESULT_SUCCESS;
 }
 

unified-runtime/source/adapters/hip/queue.hpp

@@ -19,6 +19,8 @@
 using hip_stream_queue = stream_queue_t<hipStream_t, 64, 16, hipEvent_t>;
 struct ur_queue_handle_t_ : public hip_stream_queue {};
 
+using InteropGuard = hip_stream_queue::interop_guard;
+
 template <>
 inline void hip_stream_queue::createStreamWithPriority(hipStream_t *Stream,
                                                        unsigned int Flags,
@@ -35,26 +37,3 @@ template <> inline void hip_stream_queue::createHostSubmitTimeStream() {
         hipStreamCreateWithFlags(&HostSubmitTimeStream, hipStreamNonBlocking));
   });
 }
-
-// RAII object to make hQueue stream getter methods all return the same stream
-// within the lifetime of this object.
-//
-// This is useful for urEnqueueNativeCommandExp where we want guarantees that
-// the user submitted native calls will be dispatched to a known stream, which
-// must be "got" within the user submitted function.
-//
-// TODO: Add a test that this scoping works
-class ScopedStream {
-  ur_queue_handle_t hQueue;
-
-public:
-  ScopedStream(ur_queue_handle_t hQueue, uint32_t NumEventsInWaitList,
-               const ur_event_handle_t *EventWaitList)
-      : hQueue{hQueue} {
-    ur_stream_guard Guard;
-    hQueue->getThreadLocalStream() =
-        hQueue->getNextComputeStream(NumEventsInWaitList, EventWaitList, Guard);
-  }
-  hipStream_t getStream() { return hQueue->getThreadLocalStream(); }
-  ~ScopedStream() { hQueue->getThreadLocalStream() = hipStream_t{0}; }
-};

unified-runtime/source/common/cuda-hip/stream_queue.hpp

@@ -107,8 +107,6 @@ struct stream_queue_t {
   // get_next_compute/transfer_stream() functions return streams from
   // appropriate pools in round-robin fashion
   native_type getNextComputeStream(uint32_t *StreamToken = nullptr) {
-    if (getThreadLocalStream() != native_type{0})
-      return getThreadLocalStream();
     uint32_t StreamI;
     uint32_t Token;
     while (true) {
@@ -150,8 +148,6 @@ struct stream_queue_t {
                                    const ur_event_handle_t *EventWaitList,
                                    ur_stream_guard &Guard,
                                    uint32_t *StreamToken = nullptr) {
-    if (getThreadLocalStream() != native_type{0})
-      return getThreadLocalStream();
     for (uint32_t i = 0; i < NumEventsInWaitList; i++) {
       uint32_t Token = getEventComputeStreamToken(EventWaitList[i]);
       if (getEventQueue(EventWaitList[i]) == this && canReuseStream(Token)) {
@@ -175,15 +171,7 @@ struct stream_queue_t {
     return getNextComputeStream(StreamToken);
   }
 
-  // Thread local stream will be used if ScopedStream is active
-  static native_type &getThreadLocalStream() {
-    static thread_local native_type stream{0};
-    return stream;
-  }
-
   native_type getNextTransferStream() {
-    if (getThreadLocalStream() != native_type{0})
-      return getThreadLocalStream();
     if (TransferStreams.empty()) { // for example in in-order queue
       return getNextComputeStream();
     }
@@ -354,4 +342,34 @@ struct stream_queue_t {
   uint32_t getNextEventId() noexcept { return ++EventCount; }
 
   bool backendHasOwnership() const noexcept { return HasOwnership; }
+
+  // Interop handling, for regular interop we return the next compute stream,
+  // for native commands we use the interop_guard and return a thread local
+  // stream. Native commands require to only have one in-order stream to work.
+  native_type getInteropStream() {
+    if (getThreadLocalStream() != native_type{0})
+      return getThreadLocalStream();
+
+    return getNextComputeStream();
+  }
+
+  static native_type &getThreadLocalStream() {
+    static thread_local native_type stream{0};
+    return stream;
+  }
+
+  class interop_guard {
+    stream_queue_t *q;
+
+  public:
+    interop_guard(stream_queue_t *q, uint32_t NumEventsInWaitList,
+                  const ur_event_handle_t *EventWaitList)
+        : q{q} {
+      ur_stream_guard Guard;
+      q->getThreadLocalStream() =
+          q->getNextComputeStream(NumEventsInWaitList, EventWaitList, Guard);
+    }
+    native_type getStream() { return q->getThreadLocalStream(); }
+    ~interop_guard() { q->getThreadLocalStream() = native_type{0}; }
+  };
 };

unified-runtime/test/adapters/cuda/urQueueGetNativeHandle.cpp

@@ -30,7 +30,7 @@ TEST_P(urCudaQueueGetNativeHandleTest, OutOfOrder) {
   ASSERT_SUCCESS_CUDA(cuStreamSynchronize(Stream));
 }
 
-TEST_P(urCudaQueueGetNativeHandleTest, ScopedStream) {
+TEST_P(urCudaQueueGetNativeHandleTest, InteropGuard) {
   CUstream Stream1, Stream2;
   ur_queue_properties_t props = {
       /*.stype =*/UR_STRUCTURE_TYPE_QUEUE_PROPERTIES,
@@ -50,7 +50,7 @@ TEST_P(urCudaQueueGetNativeHandleTest, ScopedStream) {
   ASSERT_NE(Stream1, Stream2);
 
   {
-    ScopedStream ActiveStream(OutOfOrderQueue, 0, nullptr);
+    InteropGuard ActiveStream(OutOfOrderQueue, 0, nullptr);
 
     ASSERT_SUCCESS(urQueueGetNativeHandle(OutOfOrderQueue, nullptr,
                                           (ur_native_handle_t *)&Stream1));