fix event caching

source/adapters/level_zero/context.cpp

@@ -565,18 +565,26 @@ ur_event_handle_t ur_context_handle_t_::getEventFromContextCache(
     bool HostVisible, bool WithProfiling, ur_device_handle_t Device,
     bool CounterBasedEventEnabled) {
   std::scoped_lock<ur_mutex> Lock(EventCacheMutex);
-  auto Cache = getEventCache(HostVisible, WithProfiling, Device);
-  if (Cache->empty())
+  auto Cache = getEventCache(HostVisible, WithProfiling, Device,
+                             CounterBasedEventEnabled);
+  if (Cache->empty()) {
+    logger::info("Cache empty (Host Visible: {}, Profiling: {}, Counter: {}, "
+                 "Device: {})",
+                 HostVisible, WithProfiling, CounterBasedEventEnabled, Device);
     return nullptr;
+  }
 
   auto It = Cache->begin();
   ur_event_handle_t Event = *It;
-  if (Event->CounterBasedEventsEnabled != CounterBasedEventEnabled) {
-    return nullptr;
-  }
   Cache->erase(It);
   // We have to reset event before using it.
   Event->reset();
+
+  logger::info("Using {} event (Host Visible: {}, Profiling: {}, Counter: {}, "
+               "Device: {}) from cache {}",
+               Event, Event->HostVisibleEvent, Event->isProfilingEnabled(),
+               Event->CounterBasedEventsEnabled, Device, Cache);
+
   return Event;
 }
 
@@ -588,8 +596,13 @@ void ur_context_handle_t_::addEventToContextCache(ur_event_handle_t Event) {
     Device = Event->UrQueue->Device;
   }
 
-  auto Cache = getEventCache(Event->isHostVisible(),
-                             Event->isProfilingEnabled(), Device);
+  auto Cache =
+      getEventCache(Event->isHostVisible(), Event->isProfilingEnabled(), Device,
+                    Event->CounterBasedEventsEnabled);
+  logger::info("Inserting {} event (Host Visible: {}, Profiling: {}, Counter: "
+               "{}, Device: {}) into cache {}",
+               Event, Event->HostVisibleEvent, Event->isProfilingEnabled(),
+               Event->CounterBasedEventsEnabled, Device, Cache);
   Cache->emplace_back(Event);
 }
 

source/adapters/level_zero/context.hpp

@@ -169,15 +169,6 @@ struct ur_context_handle_t_ : _ur_object {
   // holding the current pool usage counts.
   ur_mutex ZeEventPoolCacheMutex;
 
-  // Mutex to control operations on event caches.
-  ur_mutex EventCacheMutex;
-
-  // Caches for events.
-  using EventCache = std::vector<std::list<ur_event_handle_t>>;
-  EventCache EventCaches{4};
-  std::vector<std::unordered_map<ur_device_handle_t, size_t>>
-      EventCachesDeviceMap{4};
-
   // Initialize the PI context.
   ur_result_t initialize();
 
@@ -313,36 +304,45 @@ struct ur_context_handle_t_ : _ur_object {
   ze_context_handle_t getZeHandle() const;
 
 private:
+  enum EventFlags {
+    EVENT_FLAG_HOST_VISIBLE = UR_BIT(0),
+    EVENT_FLAG_WITH_PROFILING = UR_BIT(1),
+    EVENT_FLAG_COUNTER = UR_BIT(2),
+    EVENT_FLAG_DEVICE = UR_BIT(3), // if set, subsequent bits are device id
+    MAX_EVENT_FLAG_BITS =
+        4, // this is used as an offset for embedding device id
+  };
+
+  // Mutex to control operations on event caches.
+  ur_mutex EventCacheMutex;
+
+  // Caches for events.
+  using EventCache = std::list<ur_event_handle_t>;
+  std::vector<EventCache> EventCaches;
+
   // Get the cache of events for a provided scope and profiling mode.
-  auto getEventCache(bool HostVisible, bool WithProfiling,
-                     ur_device_handle_t Device) {
+  EventCache *getEventCache(bool HostVisible, bool WithProfiling,
+                            ur_device_handle_t Device, bool Counter) {
+
+    size_t index = 0;
     if (HostVisible) {
-      if (Device) {
-        auto EventCachesMap =
-            WithProfiling ? &EventCachesDeviceMap[0] : &EventCachesDeviceMap[1];
-        if (EventCachesMap->find(Device) == EventCachesMap->end()) {
-          EventCaches.emplace_back();
-          EventCachesMap->insert(
-              std::make_pair(Device, EventCaches.size() - 1));
-        }
-        return &EventCaches[(*EventCachesMap)[Device]];
-      } else {
-        return WithProfiling ? &EventCaches[0] : &EventCaches[1];
-      }
-    } else {
-      if (Device) {
-        auto EventCachesMap =
-            WithProfiling ? &EventCachesDeviceMap[2] : &EventCachesDeviceMap[3];
-        if (EventCachesMap->find(Device) == EventCachesMap->end()) {
-          EventCaches.emplace_back();
-          EventCachesMap->insert(
-              std::make_pair(Device, EventCaches.size() - 1));
-        }
-        return &EventCaches[(*EventCachesMap)[Device]];
-      } else {
-        return WithProfiling ? &EventCaches[2] : &EventCaches[3];
-      }
+      index |= EVENT_FLAG_HOST_VISIBLE;
+    }
+    if (WithProfiling) {
+      index |= EVENT_FLAG_WITH_PROFILING;
     }
+    if (Counter) {
+      index |= EVENT_FLAG_COUNTER;
+    }
+    if (Device) {
+      index |= EVENT_FLAG_DEVICE | (*Device->Id << MAX_EVENT_FLAG_BITS);
+    }
+
+    if (index >= EventCaches.size()) {
+      EventCaches.resize(index + 1);
+    }
+
+    return &EventCaches[index];
   }
 };