[UR][Offload] Improvements to offload event handling

unified-runtime/source/adapters/offload/enqueue.cpp

@@ -74,7 +74,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueKernelLaunch(
                      hKernel->Args.getStorageSize(), &LaunchArgs, &EventOut));
 
   if (phEvent) {
-    auto *Event = new ur_event_handle_t_();
+    auto *Event = new ur_event_handle_t_(UR_COMMAND_KERNEL_LAUNCH, hQueue);
     Event->OffloadEvent = EventOut;
     *phEvent = Event;
   }
@@ -117,7 +117,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferRead(
   }
 
   if (phEvent) {
-    auto *Event = new ur_event_handle_t_();
+    auto *Event = new ur_event_handle_t_(UR_COMMAND_MEM_BUFFER_READ, hQueue);
     Event->OffloadEvent = EventOut;
     *phEvent = Event;
   }
@@ -149,21 +149,22 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferWrite(
   }
 
   if (phEvent) {
-    auto *Event = new ur_event_handle_t_();
+    auto *Event = new ur_event_handle_t_(UR_COMMAND_MEM_BUFFER_WRITE, hQueue);
     Event->OffloadEvent = EventOut;
     *phEvent = Event;
   }
 
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t enqueueNoOp(ur_queue_handle_t hQueue, ur_event_handle_t *phEvent) {
+ur_result_t enqueueNoOp(ur_command_t Type, ur_queue_handle_t hQueue,
+                        ur_event_handle_t *phEvent) {
   // This path is a no-op, but we can't output a real event because
   // Offload doesn't currently support creating arbitrary events, and we
   // don't know the last real event in the queue. Instead we just have to
   // wait on the whole queue and then return an empty (implicitly
   // finished) event.
-  *phEvent = ur_event_handle_t_::createEmptyEvent();
+  *phEvent = ur_event_handle_t_::createEmptyEvent(Type, hQueue);
   return urQueueFinish(hQueue);
 }
 
@@ -197,7 +198,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemBufferMap(
     }
 
     if (phEvent) {
-      enqueueNoOp(hQueue, phEvent);
+      enqueueNoOp(UR_COMMAND_MEM_BUFFER_MAP, hQueue, phEvent);
     }
   }
   *ppRetMap = MapPtr;
@@ -231,7 +232,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEnqueueMemUnmap(
     }
 
     if (phEvent) {
-      enqueueNoOp(hQueue, phEvent);
+      enqueueNoOp(UR_COMMAND_MEM_UNMAP, hQueue, phEvent);
     }
   }
   BufferImpl.unmap(pMappedPtr);

unified-runtime/source/adapters/offload/event.cpp

@@ -12,6 +12,7 @@
 #include <ur_api.h>
 
 #include "event.hpp"
+#include "queue.hpp"
 #include "ur2offload.hpp"
 
 UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hEvent,
@@ -22,6 +23,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hEvent,
   UrReturnHelper ReturnValue(propSize, pPropValue, pPropSizeRet);
 
   switch (propName) {
+  case UR_EVENT_INFO_CONTEXT:
+    return ReturnValue(hEvent->UrQueue->UrContext);
+  case UR_EVENT_INFO_COMMAND_QUEUE:
+    return ReturnValue(hEvent->UrQueue);
+  case UR_EVENT_INFO_COMMAND_TYPE:
+    return ReturnValue(hEvent->Type);
   case UR_EVENT_INFO_REFERENCE_COUNT:
     return ReturnValue(hEvent->RefCount.load());
   default:
@@ -63,9 +70,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventRelease(ur_event_handle_t hEvent) {
     // if (Res) {
     //   return offloadResultToUR(Res);
     // }
+    delete hEvent;
   }
 
-  delete hEvent;
   return UR_RESULT_SUCCESS;
 }
 

unified-runtime/source/adapters/offload/event.hpp

@@ -18,9 +18,14 @@
 struct ur_event_handle_t_ : RefCounted {
   ol_event_handle_t OffloadEvent;
   ur_command_t Type;
+  ur_queue_handle_t UrQueue;
 
-  static ur_event_handle_t createEmptyEvent() {
-    auto *Event = new ur_event_handle_t_();
+  ur_event_handle_t_(ur_command_t Type, ur_queue_handle_t Queue)
+      : Type(Type), UrQueue(Queue) {}
+
+  static ur_event_handle_t createEmptyEvent(ur_command_t Type,
+                                            ur_queue_handle_t Queue) {
+    auto *Event = new ur_event_handle_t_(Type, Queue);
     // Null event represents an empty event. Waiting on it is a no-op.
     Event->OffloadEvent = nullptr;
 

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

@@ -17,9 +17,10 @@
 #include "queue.hpp"
 #include "ur2offload.hpp"
 
-UR_APIEXPORT ur_result_t UR_APICALL urQueueCreate(
-    [[maybe_unused]] ur_context_handle_t hContext, ur_device_handle_t hDevice,
-    const ur_queue_properties_t *, ur_queue_handle_t *phQueue) {
+UR_APIEXPORT ur_result_t UR_APICALL urQueueCreate(ur_context_handle_t hContext,
+                                                  ur_device_handle_t hDevice,
+                                                  const ur_queue_properties_t *,
+                                                  ur_queue_handle_t *phQueue) {
 
   assert(hContext->Device == hDevice);
 
@@ -31,6 +32,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueCreate(
   }
 
   Queue->OffloadDevice = hDevice->OffloadDevice;
+  Queue->UrContext = hContext;
 
   *phQueue = Queue;
 

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

@@ -18,4 +18,5 @@
 struct ur_queue_handle_t_ : RefCounted {
   ol_queue_handle_t OffloadQueue;
   ol_device_handle_t OffloadDevice;
+  ur_context_handle_t UrContext;
 };

unified-runtime/test/conformance/event/urEventGetProfilingInfo.cpp

@@ -5,6 +5,7 @@
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 
 #include "fixtures.h"
+#include "uur/checks.h"
 #include "uur/known_failure.h"
 
 using urEventGetProfilingInfoTest = uur::event::urEventTest;
@@ -179,8 +180,9 @@ TEST_P(urEventGetProfilingInfoTest, InvalidValue) {
 
   const ur_profiling_info_t property_name = UR_PROFILING_INFO_COMMAND_QUEUED;
   size_t property_size = 0;
-  ASSERT_SUCCESS(urEventGetProfilingInfo(event, property_name, 0, nullptr,
-                                         &property_size));
+  ASSERT_SUCCESS_OR_OPTIONAL_QUERY(
+      urEventGetProfilingInfo(event, property_name, 0, nullptr, &property_size),
+      property_name);
   ASSERT_NE(property_size, 0);
 
   uint64_t property_value = 0;
@@ -221,8 +223,10 @@ UUR_INSTANTIATE_DEVICE_TEST_SUITE(urEventGetProfilingInfoForWaitWithBarrier);
 
 TEST_P(urEventGetProfilingInfoForWaitWithBarrier, Success) {
   uint64_t submit_value = 0;
-  ASSERT_SUCCESS(urEventGetProfilingInfo(event, UR_PROFILING_INFO_COMMAND_START,
-                                         size, &submit_value, nullptr));
+  ASSERT_SUCCESS_OR_OPTIONAL_QUERY(
+      urEventGetProfilingInfo(event, UR_PROFILING_INFO_COMMAND_START, size,
+                              &submit_value, nullptr),
+      UR_PROFILING_INFO_COMMAND_START);
   ASSERT_NE(submit_value, 0);
 
   uint64_t complete_value = 0;