[L0 v2] Implement deferred event deallocation

and fix timestamp checking in events. Until now,
recordEventEndTimestamp was being checked directly to
see if the WriteGlobalTimestamp completed. This is not
safe, as we didn't use any atomics. To avoid any races,
only read recordEventEndTimestamp when the associated
zeEvent completed.

Author: igchor

scripts/templates/queue_api.hpp.mako

@@ -25,6 +25,9 @@ from templates import helper as th
 
 struct ur_queue_handle_t_ {
     virtual ~ur_queue_handle_t_();
+
+    virtual void deferEventFree(ur_event_handle_t hEvent) = 0;
+
     %for obj in th.get_queue_related_functions(specs, n, tags):
     virtual ${x}_result_t ${th.transform_queue_related_function_name(n, tags, obj, format=["type"])} = 0;
     %endfor

source/adapters/level_zero/v2/event.cpp

@@ -13,6 +13,7 @@
 #include "event.hpp"
 #include "event_pool.hpp"
 #include "event_provider.hpp"
+#include "queue_api.hpp"
 
 #include "../ur_interface_loader.hpp"
 
@@ -24,6 +25,10 @@ ur_event_handle_t_::ur_event_handle_t_(
       zeTimerResolution(getDevice()->ZeDeviceProperties->timerResolution),
       timestampMaxValue(getDevice()->getTimestampMask()) {}
 
+void ur_event_handle_t_::resetQueue(ur_queue_handle_t hQueue) {
+  this->hQueue = hQueue;
+}
+
 void ur_event_handle_t_::reset() {
   // consider make an abstraction for regular/counter based
   // events if there's more of this type of conditions
@@ -33,6 +38,7 @@ void ur_event_handle_t_::reset() {
 }
 
 ze_event_handle_t ur_event_handle_t_::getZeEvent() const {
+  assert(hQueue);
   return zeEvent.get();
 }
 
@@ -41,14 +47,27 @@ ur_result_t ur_event_handle_t_::retain() {
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t ur_event_handle_t_::releaseDeferred() {
+  assert(zeEventQueryStatus(zeEvent.get()) == ZE_RESULT_SUCCESS);
+  assert(RefCount.load() == 0);
+
+  pool->free(this);
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t ur_event_handle_t_::release() {
   if (!RefCount.decrementAndTest())
     return UR_RESULT_SUCCESS;
 
+  // Need to take a lock before checking if the event is timestamped.
+  std::unique_lock<ur_shared_mutex> lock(Mutex);
+
   if (isTimestamped() && adjustedEventEndTimestamp == 0) {
     // L0 will write end timestamp to this event some time in the future,
     // so we can't release it yet.
-    // TODO: delay releasing until the end timestamp is written.
+
+    assert(hQueue);
+    hQueue->deferEventFree(this);
     return UR_RESULT_SUCCESS;
   }
 
@@ -99,17 +118,16 @@ uint64_t ur_event_handle_t_::getEventEndTimestamp() {
   if (adjustedEventEndTimestamp)
     return adjustedEventEndTimestamp;
 
-  // If the result is 0, we have not yet gotten results back and so we just
-  // return it.
-  if (recordEventEndTimestamp == 0)
-    return recordEventEndTimestamp;
+  auto status = zeEventQueryStatus(zeEvent.get());
+  if (status != ZE_RESULT_SUCCESS) {
+    // profiling info not ready
+    return 0;
+  }
 
-  // Now that we have the result, there is no need to keep it in the queue
-  // anymore, so we cache it on the event and evict the record from the
-  // queue.
   adjustedEventEndTimestamp =
       adjustEndEventTimestamp(getEventStartTimestmap(), recordEventEndTimestamp,
                               timestampMaxValue, zeTimerResolution);
+
   return adjustedEventEndTimestamp;
 }
 
@@ -118,11 +136,13 @@ void ur_event_handle_t_::recordStartTimestamp() {
   UR_CALL_THROWS(ur::level_zero::urDeviceGetGlobalTimestamps(
       getDevice(), &deviceStartTimestamp, nullptr));
 
+  assert(adjustedEventStartTimestamp == 0);
   adjustedEventStartTimestamp = deviceStartTimestamp;
 }
 
-uint64_t *ur_event_handle_t_::getEventEndTimestampPtr() {
-  return &recordEventEndTimestamp;
+std::pair<uint64_t *, ze_event_handle_t>
+ur_event_handle_t_::getEventEndTimestampAndHandle() {
+  return {&recordEventEndTimestamp, zeEvent.get()};
 }
 
 namespace ur::level_zero {

source/adapters/level_zero/v2/event.hpp

@@ -27,12 +27,19 @@ struct ur_event_handle_t_ : _ur_object {
   ur_event_handle_t_(v2::raii::cache_borrowed_event eventAllocation,
                      v2::event_pool *pool);
 
+  // Set the queue that this event is associated with
+  void resetQueue(ur_queue_handle_t hQueue);
+
   void reset();
   ze_event_handle_t getZeEvent() const;
 
   ur_result_t retain();
   ur_result_t release();
 
+  // releases a signaled and no longer in-use event, that's on the
+  // deffered events list in the queue
+  ur_result_t releaseDeferred();
+
   // Tells if this event was created as a timestamp event, allowing profiling
   // info even if profiling is not enabled.
   bool isTimestamped() const;
@@ -44,12 +51,16 @@ struct ur_event_handle_t_ : _ur_object {
   ur_device_handle_t getDevice() const;
 
   void recordStartTimestamp();
-  uint64_t *getEventEndTimestampPtr();
+
+  // Get pointer to the end timestamp, and ze event handle.
+  // Caller is responsible for signaling the event once the timestamp is ready.
+  std::pair<uint64_t *, ze_event_handle_t> getEventEndTimestampAndHandle();
 
   uint64_t getEventStartTimestmap() const;
   uint64_t getEventEndTimestamp();
 
 private:
+  ur_queue_handle_t hQueue = nullptr;
   v2::raii::cache_borrowed_event zeEvent;
   v2::event_pool *pool;
 

source/adapters/level_zero/v2/event_pool.cpp

@@ -15,7 +15,7 @@ namespace v2 {
 
 static constexpr size_t EVENTS_BURST = 64;
 
-ur_event_handle_t_ *event_pool::allocate() {
+ur_event_handle_t_ *event_pool::allocate(ur_queue_handle_t hQueue) {
   TRACK_SCOPE_LATENCY("event_pool::allocate");
 
   std::unique_lock<std::mutex> lock(*mutex);
@@ -32,6 +32,8 @@ ur_event_handle_t_ *event_pool::allocate() {
   auto event = freelist.back();
   freelist.pop_back();
 
+  event->resetQueue(hQueue);
+
   return event;
 }
 

source/adapters/level_zero/v2/event_pool.hpp

@@ -41,7 +41,7 @@ class event_pool {
   DeviceId Id() { return provider->device()->Id.value(); };
 
   // Allocate an event from the pool. Thread safe.
-  ur_event_handle_t_ *allocate();
+  ur_event_handle_t_ *allocate(ur_queue_handle_t hQueue);
 
   // Free an event back to the pool. Thread safe.
   void free(ur_event_handle_t_ *event);

source/adapters/level_zero/v2/queue_api.hpp

@@ -16,6 +16,9 @@
 
 struct ur_queue_handle_t_ {
   virtual ~ur_queue_handle_t_();
+
+  virtual void deferEventFree(ur_event_handle_t hEvent) = 0;
+
   virtual ur_result_t queueGetInfo(ur_queue_info_t, size_t, void *,
                                    size_t *) = 0;
   virtual ur_result_t queueRetain() = 0;

source/adapters/level_zero/v2/queue_immediate_in_order.cpp

@@ -104,7 +104,7 @@ ur_queue_immediate_in_order_t::ur_queue_immediate_in_order_t(
 ur_event_handle_t
 ur_queue_immediate_in_order_t::getSignalEvent(ur_event_handle_t *hUserEvent) {
   if (hUserEvent) {
-    *hUserEvent = eventPool->allocate();
+    *hUserEvent = eventPool->allocate(this);
     return *hUserEvent;
   } else {
     return nullptr;
@@ -156,6 +156,11 @@ ur_result_t ur_queue_immediate_in_order_t::queueRelease() {
   return UR_RESULT_SUCCESS;
 }
 
+void ur_queue_immediate_in_order_t::deferEventFree(ur_event_handle_t hEvent) {
+  std::unique_lock<ur_shared_mutex> lock(this->Mutex);
+  deferredEvents.push_back(hEvent);
+}
+
 ur_result_t ur_queue_immediate_in_order_t::queueGetNativeHandle(
     ur_queue_native_desc_t *pDesc, ur_native_handle_t *phNativeQueue) {
   std::ignore = pDesc;
@@ -175,6 +180,12 @@ ur_result_t ur_queue_immediate_in_order_t::queueFinish() {
   ZE2UR_CALL(zeCommandListHostSynchronize,
              (handler.commandList.get(), UINT64_MAX));
 
+  // Free deferred events
+  for (auto &hEvent : deferredEvents) {
+    hEvent->releaseDeferred();
+  }
+  deferredEvents.clear();
+
   return UR_RESULT_SUCCESS;
 }
 
@@ -983,9 +994,12 @@ ur_result_t ur_queue_immediate_in_order_t::enqueueTimestampRecordingExp(
 
   signalEvent->recordStartTimestamp();
 
+  auto [timestampPtr, zeSignalEvent] =
+      signalEvent->getEventEndTimestampAndHandle();
+
   ZE2UR_CALL(zeCommandListAppendWriteGlobalTimestamp,
-             (handler.commandList.get(), signalEvent->getEventEndTimestampPtr(),
-              signalEvent->getZeEvent(), numWaitEvents, pWaitEvents));
+             (handler.commandList.get(), timestampPtr, zeSignalEvent,
+              numWaitEvents, pWaitEvents));
 
   if (blocking) {
     ZE2UR_CALL(zeCommandListHostSynchronize,

source/adapters/level_zero/v2/queue_immediate_in_order.hpp

@@ -46,12 +46,16 @@ struct ur_queue_immediate_in_order_t : _ur_object, public ur_queue_handle_t_ {
 
   std::vector<ze_event_handle_t> waitList;
 
+  std::vector<ur_event_handle_t> deferredEvents;
+
   std::pair<ze_event_handle_t *, uint32_t>
   getWaitListView(const ur_event_handle_t *phWaitEvents,
                   uint32_t numWaitEvents);
 
   ur_event_handle_t getSignalEvent(ur_event_handle_t *hUserEvent);
 
+  void deferEventFree(ur_event_handle_t hEvent) override;
+
   ur_result_t enqueueRegionCopyUnlocked(
       ur_mem_handle_t src, ur_mem_handle_t dst, bool blocking,
       ur_rect_offset_t srcOrigin, ur_rect_offset_t dstOrigin,

test/adapters/level_zero/v2/event_pool_test.cpp

@@ -150,7 +150,7 @@ TEST_P(EventPoolTest, Basic) {
         {
             auto pool = cache->borrow(device->Id.value(), getParam().flags);
 
-            first = pool->allocate();
+            first = pool->allocate(reinterpret_cast<ur_queue_handle_t>(0x1));
             zeFirst = first->getZeEvent();
 
             urEventRelease(first);
@@ -160,7 +160,7 @@ TEST_P(EventPoolTest, Basic) {
         {
             auto pool = cache->borrow(device->Id.value(), getParam().flags);
 
-            second = pool->allocate();
+            second = pool->allocate(reinterpret_cast<ur_queue_handle_t>(0x1));
             zeSecond = second->getZeEvent();
 
             urEventRelease(second);
@@ -179,7 +179,8 @@ TEST_P(EventPoolTest, Threaded) {
                 auto pool = cache->borrow(device->Id.value(), getParam().flags);
                 std::vector<ur_event_handle_t> events;
                 for (int i = 0; i < 100; ++i) {
-                    events.push_back(pool->allocate());
+                    events.push_back(pool->allocate(
+                        reinterpret_cast<ur_queue_handle_t>(0x1)));
                 }
                 for (int i = 0; i < 100; ++i) {
                     urEventRelease(events[i]);
@@ -197,15 +198,16 @@ TEST_P(EventPoolTest, ProviderNormalUseMostFreePool) {
     auto pool = cache->borrow(device->Id.value(), getParam().flags);
     std::list<ur_event_handle_t> events;
     for (int i = 0; i < 128; ++i) {
-        events.push_back(pool->allocate());
+        events.push_back(
+            pool->allocate(reinterpret_cast<ur_queue_handle_t>(0x1)));
     }
     auto frontZeHandle = events.front()->getZeEvent();
     for (int i = 0; i < 8; ++i) {
         urEventRelease(events.front());
         events.pop_front();
     }
     for (int i = 0; i < 8; ++i) {
-        auto e = pool->allocate();
+        auto e = pool->allocate(reinterpret_cast<ur_queue_handle_t>(0x1));
         events.push_back(e);
     }
 

test/conformance/event/event_adapter_level_zero_v2.match

@@ -5,8 +5,8 @@ urEventGetInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.
 urEventGetInfoNegativeTest.InvalidNullHandle/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urEventGetInfoNegativeTest.InvalidSizePropSize/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urEventGetInfoNegativeTest.InvalidSizePropSizeSmall/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
-{{OPT}}urEventGetProfilingInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_PROFILING_INFO_COMMAND_QUEUED
-{{OPT}}urEventGetProfilingInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_PROFILING_INFO_COMMAND_SUBMIT
+urEventGetProfilingInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_PROFILING_INFO_COMMAND_QUEUED
+urEventGetProfilingInfoTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}___UR_PROFILING_INFO_COMMAND_SUBMIT
 urEventGetProfilingInfoWithTimingComparisonTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urEventSetCallbackTest.Success/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_
 urEventSetCallbackTest.ValidateParameters/Intel_R__oneAPI_Unified_Runtime_over_Level_Zero___{{.*}}_