[SYCL] Cache pi_events instead of native handles and fix pool handling

Author: againull

sycl/plugins/level_zero/pi_level_zero.cpp

@@ -478,10 +478,18 @@ _pi_context::getFreeSlotInExistingOrNewPool(ze_event_pool_handle_t &Pool,
   std::list<ze_event_pool_handle_t> *ZePoolCache =
       getZeEventPoolCache(HostVisible, ProfilingEnabled);
 
-  // Remove full pool from the cache.
   if (!ZePoolCache->empty()) {
     if (NumEventsAvailableInEventPool[ZePoolCache->front()] == 0) {
-      ZePoolCache->erase(ZePoolCache->begin());
+      if (DisableEventsCaching) {
+        // Remove full pool from the cache if events caching is disabled.
+        ZePoolCache->erase(ZePoolCache->begin());
+      } else {
+        // If event caching is enabled then we don't destroy events so there is
+        // no need to remove pool from the cache and add it back when it has
+        // available slots. Just keep it in the tail of the cache so that all
+        // pools can be destroyed during context destruction.
+        ZePoolCache->push_front(nullptr);
+      }
     }
   }
   if (ZePoolCache->empty()) {
@@ -878,16 +886,18 @@ pi_result _pi_context::initialize() {
 pi_result _pi_context::finalize() {
   // This function is called when pi_context is deallocated, piContextRelease.
   // There could be some memory that may have not been deallocated.
-  // For example, event pool caches would be still alive.
-  {
-    std::scoped_lock Lock(ZeEventCacheMutex);
-    for (auto &ZeEventCache : ZeEventCaches) {
-      for (auto &ZeEventAndPool : ZeEventCache)
-        ZE_CALL(zeEventDestroy, (ZeEventAndPool.first));
-      ZeEventCache.clear();
+  // For example, event and event pool caches would be still alive.
+
+  if (!DisableEventsCaching) {
+    std::scoped_lock Lock(EventCacheMutex);
+    for (auto &EventCache : EventCaches) {
+      for (auto Event : EventCache) {
+        ZE_CALL(zeEventDestroy, (Event->ZeEvent));
+        delete Event;
+      }
+      EventCache.clear();
     }
   }
-
   {
     std::scoped_lock Lock(ZeEventPoolCacheMutex);
     for (auto &ZePoolCache : ZeEventPoolCache) {
@@ -5449,26 +5459,40 @@ _pi_event::getOrCreateHostVisibleEvent(ze_event_handle_t &ZeHostVisibleEvent) {
   return PI_SUCCESS;
 }
 
-std::pair<ze_event_handle_t, ze_event_pool_handle_t>
-_pi_context::getZeEventFromCache(bool HostVisible, bool WithProfiling) {
-  std::scoped_lock Lock(ZeEventCacheMutex);
-  auto Cache = getZeEventCache(HostVisible, WithProfiling);
+pi_result _pi_event::reset() {
+  Queue = nullptr;
+  CleanedUp = false;
+  Completed = false;
+  CommandData = nullptr;
+  CommandType = PI_COMMAND_TYPE_USER;
+  WaitList = {};
+  RefCount.reset(1);
+
+  if (!isHostVisible())
+    HostVisibleEvent = nullptr;
+
+  ZE_CALL(zeEventHostReset, (ZeEvent));
+  return PI_SUCCESS;
+}
+
+pi_event _pi_context::getEventFromCache(bool HostVisible, bool WithProfiling) {
+  std::scoped_lock Lock(EventCacheMutex);
+  auto Cache = getEventCache(HostVisible, WithProfiling);
   if (Cache->empty())
-    return std::make_pair<ze_event_handle_t, ze_event_pool_handle_t>(nullptr,
-                                                                     nullptr);
+    return nullptr;
 
   auto It = Cache->begin();
-  std::pair<ze_event_handle_t, ze_event_pool_handle_t> ZeEvent = *It;
+  pi_event Event = *It;
   Cache->erase(It);
-  return ZeEvent;
+  return Event;
 }
 
-void _pi_context::addZeEventToCache(ze_event_handle_t ZeEvent,
-                                    ze_event_pool_handle_t ZePool,
-                                    bool HostVisible, bool WithProfiling) {
-  std::scoped_lock Lock(ZeEventCacheMutex);
-  auto Cache = getZeEventCache(HostVisible, WithProfiling);
-  Cache->emplace_back(ZeEvent, ZePool);
+void _pi_context::addEventToCache(pi_event Event) {
+  std::scoped_lock Lock(EventCacheMutex);
+  auto Cache =
+      getEventCache(Event->isHostVisible(), Event->isProfilingEnabled());
+  Event->reset();
+  Cache->emplace_back(Event);
 }
 
 // Helper function for creating a PI event.
@@ -5481,43 +5505,41 @@ static pi_result EventCreate(pi_context Context, pi_queue Queue,
   bool ProfilingEnabled =
       !Queue || (Queue->Properties & PI_QUEUE_PROFILING_ENABLE) != 0;
 
-  auto CachedEvent =
-      Context->getZeEventFromCache(HostVisible, ProfilingEnabled);
+  if (auto CachedEvent =
+          Context->getEventFromCache(HostVisible, ProfilingEnabled)) {
+    *RetEvent = CachedEvent;
+    return PI_SUCCESS;
+  }
+
   ze_event_handle_t ZeEvent;
   ze_event_pool_handle_t ZeEventPool = {};
 
-  // If no any then check cache of event in the context
-  if (CachedEvent.first) {
-    ZeEvent = CachedEvent.first;
-    ZeEventPool = CachedEvent.second;
-  } else {
-    size_t Index = 0;
-
-    if (auto Res = Context->getFreeSlotInExistingOrNewPool(
-            ZeEventPool, Index, HostVisible, ProfilingEnabled))
-      return Res;
+  size_t Index = 0;
 
-    ZeStruct<ze_event_desc_t> ZeEventDesc;
-    ZeEventDesc.index = Index;
-    ZeEventDesc.wait = 0;
+  if (auto Res = Context->getFreeSlotInExistingOrNewPool(
+          ZeEventPool, Index, HostVisible, ProfilingEnabled))
+    return Res;
 
-    if (HostVisible) {
-      ZeEventDesc.signal = ZE_EVENT_SCOPE_FLAG_HOST;
-    } else {
-      //
-      // Set the scope to "device" for every event. This is sufficient for
-      // global device access and peer device access. If needed to be seen on
-      // the host we are doing special handling, see EventsScope options.
-      //
-      // TODO: see if "sub-device" (ZE_EVENT_SCOPE_FLAG_SUBDEVICE) can better be
-      //       used in some circumstances.
-      //
-      ZeEventDesc.signal = 0;
-    }
+  ZeStruct<ze_event_desc_t> ZeEventDesc;
+  ZeEventDesc.index = Index;
+  ZeEventDesc.wait = 0;
 
-    ZE_CALL(zeEventCreate, (ZeEventPool, &ZeEventDesc, &ZeEvent));
+  if (HostVisible) {
+    ZeEventDesc.signal = ZE_EVENT_SCOPE_FLAG_HOST;
+  } else {
+    //
+    // Set the scope to "device" for every event. This is sufficient for
+    // global device access and peer device access. If needed to be seen on
+    // the host we are doing special handling, see EventsScope options.
+    //
+    // TODO: see if "sub-device" (ZE_EVENT_SCOPE_FLAG_SUBDEVICE) can better be
+    //       used in some circumstances.
+    //
+    ZeEventDesc.signal = 0;
   }
 
+  ZE_CALL(zeEventCreate, (ZeEventPool, &ZeEventDesc, &ZeEvent));
+
   try {
     PI_ASSERT(RetEvent, PI_ERROR_INVALID_VALUE);
 
@@ -5869,12 +5891,7 @@ pi_result piEventRelease(pi_event Event) {
     Event->CommandData = nullptr;
   }
   if (Event->OwnZeEvent) {
-    if (!DisableEventsCaching) {
-      ZE_CALL(zeEventHostReset, (Event->ZeEvent));
-      Event->Context->addZeEventToCache(Event->ZeEvent, Event->ZeEventPool,
-                                        Event->isHostVisible(),
-                                        Event->isProfilingEnabled());
-    } else {
+    if (DisableEventsCaching) {
       ZE_CALL(zeEventDestroy, (Event->ZeEvent));
       auto Context = Event->Context;
       if (auto Res = Context->decrementUnreleasedEventsInPool(Event))
@@ -5896,7 +5913,12 @@ pi_result piEventRelease(pi_event Event) {
   if (Event->Queue) {
     PI_CALL(piQueueReleaseInternal(Event->Queue));
   }
-  delete Event;
+
+  if (DisableEventsCaching || !Event->OwnZeEvent) {
+    delete Event;
+  } else {
+    Event->Context->addEventToCache(Event);
+  }
 
   return PI_SUCCESS;
 }

sycl/plugins/level_zero/pi_level_zero.hpp

@@ -276,6 +276,8 @@ template <class T> struct ZeCache : private T {
 struct ReferenceCounter {
   ReferenceCounter(pi_uint32 InitVal) : RefCount{InitVal} {}
 
+  void reset(pi_uint32 InitVal) { RefCount = InitVal; }
+
   // Used when retaining an object.
   void increment() { RefCount++; }
 
@@ -750,14 +752,11 @@ struct _pi_context : _pi_object {
   // when kernel has finished execution.
   std::unordered_map<void *, MemAllocRecord> MemAllocs;
 
-  // Get a native event and its pool from cache.
-  std::pair<ze_event_handle_t, ze_event_pool_handle_t>
-  getZeEventFromCache(bool HostVisible, bool WithProfiling);
+  // Get pi_event from cache.
+  pi_event getEventFromCache(bool HostVisible, bool WithProfiling);
 
-  // Add a native event and its pool to cache.
-  void addZeEventToCache(ze_event_handle_t ZeEvent,
-                         ze_event_pool_handle_t ZePool, bool HostVisible,
-                         bool WithProfiling);
+  // Add pi_event to cache.
+  void addEventToCache(pi_event);
 
 private:
   // If context contains one device then return this device.
@@ -809,18 +808,17 @@ struct _pi_context : _pi_object {
   pi_mutex ZeEventPoolCacheMutex;
 
   // Mutex to control operations on event caches.
-  pi_mutex ZeEventCacheMutex;
+  pi_mutex EventCacheMutex;
 
-  // Caches for native event handles.
-  std::vector<std::list<std::pair<ze_event_handle_t, ze_event_pool_handle_t>>>
-      ZeEventCaches{4};
+  // Caches for events.
+  std::vector<std::list<pi_event>> EventCaches{4};
 
   // Get the cache of events for a provided scope and profiling mode.
-  auto getZeEventCache(bool HostVisible, bool WithProfiling) {
+  auto getEventCache(bool HostVisible, bool WithProfiling) {
     if (HostVisible)
-      return WithProfiling ? &ZeEventCaches[0] : &ZeEventCaches[1];
+      return WithProfiling ? &EventCaches[0] : &EventCaches[1];
     else
-      return WithProfiling ? &ZeEventCaches[2] : &ZeEventCaches[3];
+      return WithProfiling ? &EventCaches[2] : &EventCaches[3];
   }
 };
 
@@ -1374,6 +1372,9 @@ struct _pi_event : _pi_object {
   // L0 event (if any) is not guranteed to have been signalled, or
   // being visible to the host at all.
   bool Completed = {false};
+
+  // Reset _pi_event object.
+  pi_result reset();
 };
 
 struct _pi_program : _pi_object {