Add urEnqueueNativeCommand impl for HIP adapter

Add urEnqueueNativeCommand impl for HIP adapter. Also cahnge typo
'quard' to 'guard'.

Author: Hugh Delaney

source/adapters/hip/command_buffer.cpp

@@ -884,7 +884,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferEnqueueExp(
     std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
     ScopedContext Active(hQueue->getDevice());
     uint32_t StreamToken;
-    ur_stream_quard Guard;
+    ur_stream_guard Guard;
     hipStream_t HIPStream = hQueue->getNextComputeStream(
         numEventsInWaitList, phEventWaitList, Guard, &StreamToken);
 

source/adapters/hip/enqueue.cpp

@@ -293,7 +293,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
     ScopedContext Active(Dev);
 
     uint32_t StreamToken;
-    ur_stream_quard Guard;
+    ur_stream_guard Guard;
     hipStream_t HIPStream = hQueue->getNextComputeStream(
         numEventsInWaitList, phEventWaitList, Guard, &StreamToken);
 
@@ -380,7 +380,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueEventsWaitWithBarrier(
   try {
     ScopedContext Active(hQueue->getDevice());
     uint32_t StreamToken;
-    ur_stream_quard Guard;
+    ur_stream_guard Guard;
     hipStream_t HIPStream = hQueue->getNextComputeStream(
         numEventsInWaitList,
         reinterpret_cast<const ur_event_handle_t *>(phEventWaitList), Guard,
@@ -1243,7 +1243,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueUSMFill(
   try {
     ScopedContext Active(hQueue->getDevice());
     uint32_t StreamToken;
-    ur_stream_quard Guard;
+    ur_stream_guard Guard;
     hipStream_t HIPStream = hQueue->getNextComputeStream(
         numEventsInWaitList, phEventWaitList, Guard, &StreamToken);
     Result = enqueueEventsWait(hQueue, HIPStream, numEventsInWaitList,
@@ -1893,7 +1893,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueTimestampRecordingExp(
     ScopedContext Active(hQueue->getDevice());
 
     uint32_t StreamToken;
-    ur_stream_quard Guard;
+    ur_stream_guard Guard;
     hipStream_t HIPStream = hQueue->getNextComputeStream(
         numEventsInWaitList, phEventWaitList, Guard, &StreamToken);
     UR_CHECK_ERROR(enqueueEventsWait(hQueue, HIPStream, numEventsInWaitList,

source/adapters/hip/enqueue_native.cpp

@@ -10,9 +10,44 @@
 
 #include <ur_api.h>
 
-UR_APICALL UR_APIEXPORT ur_result_t urEnqueueNativeCommandExp(
-    ur_queue_handle_t, ur_exp_enqueue_native_command_function_t, void *,
-    const ur_exp_enqueue_native_command_properties_t *, uint32_t,
-    const ur_event_handle_t *, ur_event_handle_t *) {
-  return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
+#include "context.hpp"
+#include "event.hpp"
+#include "queue.hpp"
+
+UR_APIEXPORT ur_result_t UR_APICALL urEnqueueNativeCommandExp(
+    ur_queue_handle_t hQueue,
+    ur_exp_enqueue_native_command_function_t pfnNativeEnqueue, void *data,
+    const ur_exp_enqueue_native_command_properties_t *,
+    uint32_t NumEventsInWaitList, const ur_event_handle_t *phEventWaitList,
+    ur_event_handle_t *phEvent) {
+  // TODO: how should mem migration work across a context here?
+  // Perhaps we will need to add a phMemObjArgs so that we are able to make
+  // sure memory migration happens across devices in the same context
+
+  try {
+    ScopedContext ActiveContext(hQueue->getDevice());
+    ScopedStream ActiveStream(hQueue, NumEventsInWaitList, phEventWaitList);
+    std::unique_ptr<ur_event_handle_t_> RetImplEvent{nullptr};
+
+    if (phEvent) {
+      RetImplEvent =
+          std::unique_ptr<ur_event_handle_t_>(ur_event_handle_t_::makeNative(
+              UR_COMMAND_ENQUEUE_NATIVE_EXP, hQueue, ActiveStream.getStream()));
+      UR_CHECK_ERROR(RetImplEvent->start());
+    }
+
+    pfnNativeEnqueue(hQueue, data); // This is using urQueueGetNativeHandle to
+                                    // get the CUDA stream. It must be the
+                                    // same stream as is used before and after
+    if (phEvent) {
+      UR_CHECK_ERROR(RetImplEvent->record());
+      *phEvent = RetImplEvent.release();
+    }
+
+  } catch (ur_result_t Err) {
+    return Err;
+  } catch (hipError_t hipErr) {
+    return mapErrorUR(hipErr);
+  }
+  return UR_RESULT_SUCCESS;
 }

source/adapters/hip/queue.cpp

@@ -29,6 +29,8 @@ void ur_queue_handle_t_::transferStreamWaitForBarrierIfNeeded(
 }
 
 hipStream_t ur_queue_handle_t_::getNextComputeStream(uint32_t *StreamToken) {
+  if (getThreadLocalStream() != hipStream_t{0})
+    return getThreadLocalStream();
   uint32_t Stream_i;
   uint32_t Token;
   while (true) {
@@ -63,7 +65,9 @@ hipStream_t ur_queue_handle_t_::getNextComputeStream(uint32_t *StreamToken) {
 
 hipStream_t ur_queue_handle_t_::getNextComputeStream(
     uint32_t NumEventsInWaitList, const ur_event_handle_t *EventWaitList,
-    ur_stream_quard &Guard, uint32_t *StreamToken) {
+    ur_stream_guard &Guard, uint32_t *StreamToken) {
+  if (getThreadLocalStream() != hipStream_t{0})
+    return getThreadLocalStream();
   for (uint32_t i = 0; i < NumEventsInWaitList; i++) {
     uint32_t Token = EventWaitList[i]->getComputeStreamToken();
     if (EventWaitList[i]->getQueue() == this && canReuseStream(Token)) {
@@ -76,7 +80,7 @@ hipStream_t ur_queue_handle_t_::getNextComputeStream(
         if (StreamToken) {
           *StreamToken = Token;
         }
-        Guard = ur_stream_quard{std::move(ComputeSyncGuard)};
+        Guard = ur_stream_guard{std::move(ComputeSyncGuard)};
         hipStream_t Res = EventWaitList[i]->getStream();
         computeStreamWaitForBarrierIfNeeded(Res, Stream_i);
         return Res;
@@ -88,6 +92,8 @@ hipStream_t ur_queue_handle_t_::getNextComputeStream(
 }
 
 hipStream_t ur_queue_handle_t_::getNextTransferStream() {
+  if (getThreadLocalStream() != hipStream_t{0})
+    return getThreadLocalStream();
   if (TransferStreams.empty()) { // for example in in-order queue
     return getNextComputeStream();
   }

source/adapters/hip/queue.hpp

@@ -14,7 +14,7 @@
 #include <mutex>
 #include <vector>
 
-using ur_stream_quard = std::unique_lock<std::mutex>;
+using ur_stream_guard = std::unique_lock<std::mutex>;
 
 /// UR queue mapping on to hipStream_t objects.
 ///
@@ -97,7 +97,7 @@ struct ur_queue_handle_t_ {
   // returns a lock that needs to remain locked as long as the stream is in use
   native_type getNextComputeStream(uint32_t NumEventsInWaitList,
                                    const ur_event_handle_t *EventWaitList,
-                                   ur_stream_quard &Guard,
+                                   ur_stream_guard &Guard,
                                    uint32_t *StreamToken = nullptr);
   native_type getNextTransferStream();
   native_type get() { return getNextComputeStream(); };
@@ -247,6 +247,12 @@ struct ur_queue_handle_t_ {
     }
   }
 
+  // Thread local stream will be used if ScopedStream is active
+  static hipStream_t &getThreadLocalStream() {
+    static thread_local hipStream_t stream{0};
+    return stream;
+  }
+
   ur_context_handle_t getContext() const { return Context; };
 
   ur_device_handle_t getDevice() const { return Device; };
@@ -261,3 +267,26 @@ struct ur_queue_handle_t_ {
 
   bool backendHasOwnership() const noexcept { return HasOwnership; }
 };
+
+// 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}; }
+};