Weak ptr only from rc

igchor · igchor · commit 0c845632d3dd · 2025-01-14T21:00:50.000Z
diff --git a/source/adapters/level_zero/v2/common.hpp b/source/adapters/level_zero/v2/common.hpp
@@ -150,13 +150,36 @@ struct ref_counted {
   URHandle handle;
 };
 
+// This version of ref_counted does not call retain/release functions.
+// It is used to avoid circular references, most notably to ur_context_handle_t.
+// This is equivalent to just using URHandle but makes it clear that no ref
+// counting is expected.
+template <typename URHandle> struct weak {
+  template <ur_result_t (*retain)(URHandle), ur_result_t (*release)(URHandle)>
+  weak(const ref_counted<URHandle, retain, release> &handle)
+      : handle(handle.get()) {}
+
+  operator URHandle() const { return handle; }
+  URHandle operator->() const { return handle; }
+
+  weak(const weak &) = default;
+  weak &operator=(const weak &) = default;
+
+  weak(weak &&other) = default;
+  weak &operator=(weak &&other) = default;
+
+  URHandle get() const { return handle; }
+
+private:
+  URHandle handle;
+};
+
 template <typename URHandle> struct ref_counted_traits;
 
 #define DECLARE_REF_COUNTER_TRAITS(URHandle, retainFn, releaseFn)              \
   template <> struct ref_counted_traits<URHandle> {                            \
     static ur_result_t retain(URHandle handle) { return retainFn(handle); }    \
     static ur_result_t release(URHandle handle) { return releaseFn(handle); }  \
-    static ur_result_t nop(URHandle) { return UR_RESULT_SUCCESS; }             \
     static ur_result_t validate([[maybe_unused]] URHandle handle) {            \
       assert(reinterpret_cast<_ur_object *>(handle)->RefCount.load() != 0);    \
       return UR_RESULT_SUCCESS;                                                \
@@ -168,14 +191,6 @@ template <typename URHandle>
 using rc = ref_counted<URHandle, ref_counted_traits<URHandle>::retain,
                        ref_counted_traits<URHandle>::release>;
 
-// This version of ref_counted does not call retain/release functions.
-// It is used to avoid circular references, most notably to ur_context_handle_t.
-// This is equivalent to just using URHandle but makes it clear that no ref
-// counting is expected.
-template <typename URHandle>
-using weak = ref_counted<URHandle, ref_counted_traits<URHandle>::nop,
-                         ref_counted_traits<URHandle>::nop>;
-
 // This version of ref_counted validates that the ref count is not zero on every
 // release and retain in debug mode, and does nothing in the release mode.
 // Used for types that should always be alibe during the adapter lifetime (e.g.
diff --git a/source/adapters/level_zero/v2/context.cpp b/source/adapters/level_zero/v2/context.cpp
@@ -49,7 +49,8 @@ ur_context_handle_t_::ur_context_handle_t_(ze_context_handle_t hContext,
                                            const ur_device_handle_t *phDevices,
                                            bool ownZeContext)
     : commandListCache(hContext),
-      eventPoolCache(this, phDevices[0]->Platform->getNumDevices(),
+      eventPoolCache(v2::raii::rc_val_only<ur_context_handle_t>(this),
+                     phDevices[0]->Platform->getNumDevices(),
                      [context = this, platform = phDevices[0]->Platform](
                          DeviceId deviceId, v2::event_flags_t flags)
                          -> std::unique_ptr<v2::event_provider> {
@@ -60,11 +61,14 @@ ur_context_handle_t_::ur_context_handle_t_(ze_context_handle_t hContext,
 
                        // TODO: just use per-context id?
                        return std::make_unique<v2::provider_normal>(
-                           context, v2::QUEUE_IMMEDIATE, flags);
+                           v2::raii::rc_val_only<ur_context_handle_t>(context),
+                           v2::QUEUE_IMMEDIATE, flags);
                      }),
-      nativeEventsPool(this, std::make_unique<v2::provider_normal>(
-                                 this, v2::QUEUE_IMMEDIATE,
-                                 v2::EVENT_FLAGS_PROFILING_ENABLED)),
+      nativeEventsPool(v2::raii::rc_val_only<ur_context_handle_t>(this),
+                       std::make_unique<v2::provider_normal>(
+                           v2::raii::rc_val_only<ur_context_handle_t>(this),
+                           v2::QUEUE_IMMEDIATE,
+                           v2::EVENT_FLAGS_PROFILING_ENABLED)),
       hContext(hContext, ownZeContext),
       hDevices(phDevices, phDevices + numDevices),
       p2pAccessDevices(populateP2PDevices(
diff --git a/source/adapters/level_zero/v2/event.cpp b/source/adapters/level_zero/v2/event.cpp
@@ -87,9 +87,9 @@ uint64_t *event_profiling_data_t::eventEndTimestampAddr() {
   return &recordEventEndTimestamp;
 }
 
-ur_event_handle_t_::ur_event_handle_t_(ur_context_handle_t hContext,
-                                       ze_event_handle_t hZeEvent,
-                                       v2::event_flags_t flags)
+ur_event_handle_t_::ur_event_handle_t_(
+    v2::raii::weak<ur_context_handle_t> hContext, ze_event_handle_t hZeEvent,
+    v2::event_flags_t flags)
     : hContext(hContext), hZeEvent(hZeEvent), flags(flags),
       profilingData(hZeEvent) {}
 
@@ -189,7 +189,7 @@ ur_context_handle_t ur_event_handle_t_::getContext() const { return hContext; }
 ur_command_t ur_event_handle_t_::getCommandType() const { return commandType; }
 
 ur_pooled_event_t::ur_pooled_event_t(
-    ur_context_handle_t hContext,
+    v2::raii::weak<ur_context_handle_t> hContext,
     v2::raii::cache_borrowed_event eventAllocation, v2::event_pool *pool)
     : ur_event_handle_t_(hContext, eventAllocation.get(), pool->getFlags()),
       zeEvent(std::move(eventAllocation)), pool(pool) {}
@@ -200,7 +200,8 @@ ur_result_t ur_pooled_event_t::forceRelease() {
 }
 
 ur_native_event_t::ur_native_event_t(
-    ur_native_handle_t hNativeEvent, ur_context_handle_t hContext,
+    ur_native_handle_t hNativeEvent,
+    v2::raii::weak<ur_context_handle_t> hContext,
     const ur_event_native_properties_t *pProperties)
     : ur_event_handle_t_(
           hContext,
@@ -390,7 +391,9 @@ urEventCreateWithNativeHandle(ur_native_handle_t hNativeEvent,
     *phEvent = hContext->nativeEventsPool.allocate();
     ZE2UR_CALL(zeEventHostSignal, ((*phEvent)->getZeEvent()));
   } else {
-    *phEvent = new ur_native_event_t(hNativeEvent, hContext, pProperties);
+    *phEvent = new ur_native_event_t(
+        hNativeEvent, v2::raii::rc_val_only<ur_context_handle_t>(hContext),
+        pProperties);
   }
   return UR_RESULT_SUCCESS;
 } catch (...) {
diff --git a/source/adapters/level_zero/v2/event.hpp b/source/adapters/level_zero/v2/event.hpp
@@ -47,8 +47,8 @@ struct event_profiling_data_t {
 
 struct ur_event_handle_t_ : _ur_object {
 public:
-  ur_event_handle_t_(ur_context_handle_t hContext, ze_event_handle_t hZeEvent,
-                     v2::event_flags_t flags);
+  ur_event_handle_t_(v2::raii::weak<ur_context_handle_t> hContext,
+                     ze_event_handle_t hZeEvent, v2::event_flags_t flags);
 
   // Set the queue and command that this event is associated with
   void resetQueueAndCommand(ur_queue_handle_t hQueue, ur_command_t commandType);
@@ -113,7 +113,7 @@ struct ur_event_handle_t_ : _ur_object {
 };
 
 struct ur_pooled_event_t : ur_event_handle_t_ {
-  ur_pooled_event_t(ur_context_handle_t hContext,
+  ur_pooled_event_t(v2::raii::weak<ur_context_handle_t> hContext,
                     v2::raii::cache_borrowed_event eventAllocation,
                     v2::event_pool *pool);
 
@@ -126,7 +126,7 @@ struct ur_pooled_event_t : ur_event_handle_t_ {
 
 struct ur_native_event_t : ur_event_handle_t_ {
   ur_native_event_t(ur_native_handle_t hNativeEvent,
-                    ur_context_handle_t hContext,
+                    v2::raii::weak<ur_context_handle_t> hContext,
                     const ur_event_native_properties_t *pProperties);
 
   ur_result_t forceRelease() override;
diff --git a/source/adapters/level_zero/v2/event_pool.hpp b/source/adapters/level_zero/v2/event_pool.hpp
@@ -28,10 +28,9 @@ namespace v2 {
 
 class event_pool {
 public:
-  // store weak reference to the queue as event_pool is part of the queue
-  event_pool(ur_context_handle_t hContext,
+  event_pool(raii::weak<ur_context_handle_t> hContext,
              std::unique_ptr<event_provider> Provider)
-      : hContext(hContext), provider(std::move(Provider)),
+      : hContext(std::move(hContext)), provider(std::move(Provider)),
         mutex(std::make_unique<std::mutex>()){};
 
   event_pool(event_pool &&other) = default;
diff --git a/source/adapters/level_zero/v2/event_pool_cache.cpp b/source/adapters/level_zero/v2/event_pool_cache.cpp
@@ -13,10 +13,10 @@
 
 namespace v2 {
 
-event_pool_cache::event_pool_cache(ur_context_handle_t hContext,
+event_pool_cache::event_pool_cache(raii::weak<ur_context_handle_t> hContext,
                                    size_t max_devices,
                                    ProviderCreateFunc ProviderCreate)
-    : hContext(hContext), providerCreate(ProviderCreate) {
+    : hContext(std::move(hContext)), providerCreate(ProviderCreate) {
   pools.resize(max_devices * (1ULL << EVENT_FLAGS_USED_BITS));
 }
 
diff --git a/source/adapters/level_zero/v2/event_pool_cache.hpp b/source/adapters/level_zero/v2/event_pool_cache.hpp
@@ -35,7 +35,7 @@ class event_pool_cache {
   using ProviderCreateFunc = std::function<std::unique_ptr<event_provider>(
       DeviceId, event_flags_t flags)>;
 
-  event_pool_cache(ur_context_handle_t hContext, size_t max_devices,
+  event_pool_cache(raii::weak<ur_context_handle_t> hContext, size_t max_devices,
                    ProviderCreateFunc);
 
   raii::cache_borrowed_event_pool borrow(DeviceId, event_flags_t flags);
diff --git a/source/adapters/level_zero/v2/event_provider_counter.cpp b/source/adapters/level_zero/v2/event_provider_counter.cpp
@@ -21,8 +21,9 @@
 namespace v2 {
 
 provider_counter::provider_counter(ur_platform_handle_t platform,
-                                   ur_context_handle_t context,
-                                   ur_device_handle_t device) {
+                                   raii::weak<ur_context_handle_t> context,
+                                   raii::rc_val_only<ur_device_handle_t> device)
+    : device(std::move(device)) {
   ZE2UR_CALL_THROWS(zeDriverGetExtensionFunctionAddress,
                     (platform->ZeDriver, "zexCounterBasedEventCreate",
                      (void **)&this->eventCreateFunc));
diff --git a/source/adapters/level_zero/v2/event_provider_counter.hpp b/source/adapters/level_zero/v2/event_provider_counter.hpp
@@ -34,13 +34,15 @@ typedef ze_result_t (*zexCounterBasedEventCreate)(
 
 class provider_counter : public event_provider {
 public:
-  provider_counter(ur_platform_handle_t platform, ur_context_handle_t,
-                   ur_device_handle_t);
+  provider_counter(ur_platform_handle_t platform,
+                   raii::weak<ur_context_handle_t> context,
+                   raii::rc_val_only<ur_device_handle_t> device);
 
   raii::cache_borrowed_event allocate() override;
   event_flags_t eventFlags() const override;
 
 private:
+  raii::rc_val_only<ur_device_handle_t> device;
   ze_context_handle_t translatedContext;
   ze_device_handle_t translatedDevice;
 
diff --git a/source/adapters/level_zero/v2/event_provider_normal.hpp b/source/adapters/level_zero/v2/event_provider_normal.hpp
@@ -47,9 +47,9 @@ class provider_pool {
 // supplies multi-device events for a given context
 class provider_normal : public event_provider {
 public:
-  provider_normal(ur_context_handle_t context, queue_type qtype,
+  provider_normal(raii::weak<ur_context_handle_t> context, queue_type qtype,
                   event_flags_t flags)
-      : queueType(qtype), urContext(context), flags(flags) {}
+      : queueType(qtype), urContext(std::move(context)), flags(flags) {}
 
   raii::cache_borrowed_event allocate() override;
   event_flags_t eventFlags() const override;
diff --git a/test/adapters/level_zero/v2/event_pool_test.cpp b/test/adapters/level_zero/v2/event_pool_test.cpp
@@ -110,17 +110,20 @@ struct EventPoolTest : public uur::urQueueTestWithParam<ProviderParams> {
         mockVec.push_back(device);
 
         cache = std::unique_ptr<event_pool_cache>(new event_pool_cache(
-            nullptr, MAX_DEVICES,
+            v2::raii::rc_val_only<ur_context_handle_t>(context), MAX_DEVICES,
             [this, params](DeviceId, event_flags_t flags)
                 -> std::unique_ptr<event_provider> {
                 // normally id would be used to find the appropriate device to create the provider
                 switch (params.provider) {
                 case TEST_PROVIDER_COUNTER:
-                    return std::make_unique<provider_counter>(platform, context,
-                                                              device);
+                    return std::make_unique<provider_counter>(
+                        platform,
+                        v2::raii::rc_val_only<ur_context_handle_t>(context),
+                        device);
                 case TEST_PROVIDER_NORMAL:
                     return std::make_unique<provider_normal>(
-                        context, params.queue, flags);
+                        v2::raii::rc_val_only<ur_context_handle_t>(context),
+                        params.queue, flags);
                 }
                 return nullptr;
             }));
