[L0 v2] introduce raii wrapper for UR handles

Some entities (e.g. devices) do not need to be retained
as they are owned by the platform. For such cases, only
validate RefCount instead of acutally increasing/decreasing it.

Author: igchor

source/adapters/level_zero/v2/common.hpp

@@ -14,6 +14,7 @@
 #include <ze_api.h>
 
 #include "../common.hpp"
+#include "../ur_interface_loader.hpp"
 #include "logger/ur_logger.hpp"
 
 namespace v2 {
@@ -54,8 +55,7 @@ struct ze_handle_wrapper {
     try {
       reset();
     } catch (...) {
-      // TODO: add appropriate logging or pass the error
-      // to the caller (make the dtor noexcept(false) or use tls?)
+      // logging already done in reset
     }
   }
 
@@ -104,5 +104,95 @@ using ze_context_handle_t =
 using ze_command_list_handle_t =
     ze_handle_wrapper<::ze_command_list_handle_t, zeCommandListDestroy>;
 
+template <typename URHandle, ur_result_t (*retain)(URHandle),
+          ur_result_t (*release)(URHandle)>
+struct ref_counted {
+  ref_counted(URHandle handle) : handle(handle) {
+    if (handle) {
+      retain(handle);
+    }
+  }
+
+  ~ref_counted() {
+    if (handle) {
+      release(handle);
+    }
+  }
+
+  operator URHandle() const { return handle; }
+  URHandle operator->() const { return handle; }
+
+  ref_counted(const ref_counted &) = delete;
+  ref_counted &operator=(const ref_counted &) = delete;
+
+  ref_counted(ref_counted &&other) {
+    handle = other.handle;
+    other.handle = nullptr;
+  }
+
+  ref_counted &operator=(ref_counted &&other) {
+    if (this == &other) {
+      return *this;
+    }
+
+    if (handle) {
+      release(handle);
+    }
+
+    handle = other.handle;
+    other.handle = nullptr;
+    return *this;
+  }
+
+  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;                                                \
+    }                                                                          \
+  };
+
+// This version of ref_counted calls retain/release functions.
+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.
+// devices).
+template <typename URHandle>
+using rc_val_only =
+    ref_counted<URHandle, ref_counted_traits<URHandle>::validate,
+                ref_counted_traits<URHandle>::validate>;
+
+DECLARE_REF_COUNTER_TRAITS(::ur_device_handle_t, urDeviceRetain,
+                           urDeviceRelease);
+DECLARE_REF_COUNTER_TRAITS(::ur_context_handle_t, urContextRetain,
+                           urContextRelease);
+DECLARE_REF_COUNTER_TRAITS(::ur_mem_handle_t, urMemRetain, urMemRelease);
+DECLARE_REF_COUNTER_TRAITS(::ur_program_handle_t, urProgramRetain,
+                           urProgramRelease);
+
 } // namespace raii
+
 } // namespace v2

source/adapters/level_zero/v2/event.cpp

@@ -154,6 +154,9 @@ ur_result_t ur_event_handle_t_::release() {
   if (isTimestamped() && !getEventEndTimestamp()) {
     // L0 will write end timestamp to this event some time in the future,
     // so we can't release it yet.
+
+    // If this code is being executed, queue has to be valid (queue cannot
+    // be released before all operations complete).
     assert(hQueue);
     hQueue->deferEventFree(this);
     return UR_RESULT_SUCCESS;

source/adapters/level_zero/v2/event.hpp

@@ -98,13 +98,13 @@ struct ur_event_handle_t_ : _ur_object {
   uint64_t getEventEndTimestamp();
 
 protected:
-  ur_context_handle_t hContext;
+  v2::raii::weak<ur_context_handle_t> hContext;
 
   // non-owning handle to the L0 event
   const ze_event_handle_t hZeEvent;
 
   // queue and commandType that this event is associated with, set by enqueue
-  // commands
+  // commands. DO NOT ref count hQueue here to avoid circular references.
   ur_queue_handle_t hQueue = nullptr;
   ur_command_t commandType = UR_COMMAND_FORCE_UINT32;
 

source/adapters/level_zero/v2/event_pool.hpp

@@ -50,7 +50,7 @@ class event_pool {
   event_flags_t getFlags() const;
 
 private:
-  ur_context_handle_t hContext;
+  raii::weak<ur_context_handle_t> hContext;
   std::unique_ptr<event_provider> provider;
 
   std::deque<ur_pooled_event_t> events;

source/adapters/level_zero/v2/event_pool_cache.hpp

@@ -41,7 +41,7 @@ class event_pool_cache {
   raii::cache_borrowed_event_pool borrow(DeviceId, event_flags_t flags);
 
 private:
-  ur_context_handle_t hContext;
+  raii::weak<ur_context_handle_t> hContext;
   ur_mutex mutex;
   ProviderCreateFunc providerCreate;
 

source/adapters/level_zero/v2/event_provider_normal.hpp

@@ -56,7 +56,7 @@ class provider_normal : public event_provider {
 
 private:
   queue_type queueType;
-  ur_context_handle_t urContext;
+  raii::weak<ur_context_handle_t> urContext;
   event_flags_t flags;
 
   std::unique_ptr<provider_pool> createProviderPool();

source/adapters/level_zero/v2/kernel.cpp

@@ -40,8 +40,6 @@ ur_kernel_handle_t_::ur_kernel_handle_t_(ur_program_handle_t hProgram,
                                          const char *kernelName)
     : hProgram(hProgram),
       deviceKernels(hProgram->Context->getPlatform()->getNumDevices()) {
-  ur::level_zero::urProgramRetain(hProgram);
-
   for (auto &Dev : hProgram->AssociatedDevices) {
     auto zeDevice = Dev->ZeDevice;
     // Program may be associated with all devices from the context but built
@@ -75,8 +73,6 @@ ur_kernel_handle_t_::ur_kernel_handle_t_(
     const ur_kernel_native_properties_t *pProperties)
     : hProgram(hProgram),
       deviceKernels(context ? context->getPlatform()->getNumDevices() : 0) {
-  ur::level_zero::urProgramRetain(hProgram);
-
   auto ownZeHandle = pProperties ? pProperties->isNativeHandleOwned : false;
 
   ze_kernel_handle_t zeKernel = ur_cast<ze_kernel_handle_t>(hNativeKernel);
@@ -94,19 +90,6 @@ ur_kernel_handle_t_::ur_kernel_handle_t_(
   completeInitialization();
 }
 
-ur_result_t ur_kernel_handle_t_::release() {
-  // manually release kernels to allow errors to be propagated
-  for (auto &singleDeviceKernelOpt : deviceKernels) {
-    if (singleDeviceKernelOpt.has_value()) {
-      singleDeviceKernelOpt.value().hKernel.reset();
-    }
-  }
-
-  UR_CALL_THROWS(ur::level_zero::urProgramRelease(hProgram));
-
-  return UR_RESULT_SUCCESS;
-}
-
 void ur_kernel_handle_t_::completeInitialization() {
   // Cache kernel name. Should be the same for all devices
   assert(deviceKernels.size() > 0);
@@ -365,7 +348,6 @@ ur_result_t urKernelRelease(
   if (!hKernel->RefCount.decrementAndTest())
     return UR_RESULT_SUCCESS;
 
-  hKernel->release();
   delete hKernel;
 
   return UR_RESULT_SUCCESS;

source/adapters/level_zero/v2/kernel.hpp

@@ -20,7 +20,7 @@ struct ur_single_device_kernel_t {
                             ze_kernel_handle_t hKernel, bool ownZeHandle);
   ur_result_t release();
 
-  ur_device_handle_t hDevice;
+  v2::raii::rc_val_only<ur_device_handle_t> hDevice;
   v2::raii::ze_kernel_handle_t hKernel;
   mutable ZeCache<ZeStruct<ze_kernel_properties_t>> zeKernelProperties;
 };
@@ -74,9 +74,6 @@ struct ur_kernel_handle_t_ : _ur_object {
 
   std::vector<char> getSourceAttributes() const;
 
-  // Perform cleanup.
-  ur_result_t release();
-
   // Add a pending memory allocation for which device is not yet known.
   ur_result_t
   addPendingMemoryAllocation(pending_memory_allocation_t allocation);
@@ -92,7 +89,7 @@ struct ur_kernel_handle_t_ : _ur_object {
 
 private:
   // Keep the program of the kernel.
-  const ur_program_handle_t hProgram;
+  const v2::raii::rc<ur_program_handle_t> hProgram;
 
   // Vector of ur_single_device_kernel_t indexed by deviceIndex().
   std::vector<std::optional<ur_single_device_kernel_t>> deviceKernels;

source/adapters/level_zero/v2/memory.cpp

@@ -178,7 +178,7 @@ void *ur_discrete_mem_handle_t::allocateOnDevice(ur_device_handle_t hDevice,
       hContext, hDevice, nullptr, UR_USM_TYPE_DEVICE, size, &ptr));
 
   deviceAllocations[id] =
-      usm_unique_ptr_t(ptr, [hContext = this->hContext](void *ptr) {
+      usm_unique_ptr_t(ptr, [hContext = this->hContext.get()](void *ptr) {
         auto ret = hContext->getDefaultUSMPool()->free(ptr);
         if (ret != UR_RESULT_SUCCESS) {
           logger::error("Failed to free device memory: {}", ret);
@@ -230,7 +230,7 @@ ur_discrete_mem_handle_t::ur_discrete_mem_handle_t(
     devicePtr = allocateOnDevice(hDevice, size);
   } else {
     deviceAllocations[hDevice->Id.value()] = usm_unique_ptr_t(
-        devicePtr, [hContext = this->hContext, ownZePtr](void *ptr) {
+        devicePtr, [hContext = this->hContext.get(), ownZePtr](void *ptr) {
           if (!ownZePtr) {
             return;
           }
@@ -361,22 +361,11 @@ static bool useHostBuffer(ur_context_handle_t hContext) {
              ZE_DEVICE_PROPERTY_FLAG_INTEGRATED;
 }
 
-namespace ur::level_zero {
-ur_result_t urMemRetain(ur_mem_handle_t hMem);
-ur_result_t urMemRelease(ur_mem_handle_t hMem);
-} // namespace ur::level_zero
-
 ur_mem_sub_buffer_t::ur_mem_sub_buffer_t(ur_mem_handle_t hParent, size_t offset,
                                          size_t size,
                                          device_access_mode_t accessMode)
     : ur_mem_handle_t_(hParent->getContext(), size, accessMode),
-      hParent(hParent), offset(offset), size(size) {
-  ur::level_zero::urMemRetain(hParent);
-}
-
-ur_mem_sub_buffer_t::~ur_mem_sub_buffer_t() {
-  ur::level_zero::urMemRelease(hParent);
-}
+      hParent(hParent), offset(offset), size(size) {}
 
 void *ur_mem_sub_buffer_t::getDevicePtr(
     ur_device_handle_t hDevice, device_access_mode_t access, size_t offset,

source/adapters/level_zero/v2/memory.hpp

@@ -48,7 +48,7 @@ struct ur_mem_handle_t_ : private _ur_object {
 
 protected:
   const device_access_mode_t accessMode;
-  const ur_context_handle_t hContext;
+  const v2::raii::rc<ur_context_handle_t> hContext;
   const size_t size;
 };
 
@@ -141,7 +141,7 @@ struct ur_discrete_mem_handle_t : public ur_mem_handle_t_ {
 
   // Specifies device on which the latest allocation resides.
   // If null, there is no allocation.
-  ur_device_handle_t activeAllocationDevice = nullptr;
+  v2::raii::rc_val_only<ur_device_handle_t> activeAllocationDevice = nullptr;
 
   // If not null, copy the buffer content back to this memory on release.
   void *writeBackPtr = nullptr;
@@ -157,7 +157,6 @@ struct ur_discrete_mem_handle_t : public ur_mem_handle_t_ {
 struct ur_mem_sub_buffer_t : public ur_mem_handle_t_ {
   ur_mem_sub_buffer_t(ur_mem_handle_t hParent, size_t offset, size_t size,
                       device_access_mode_t accesMode);
-  ~ur_mem_sub_buffer_t();
 
   void *
   getDevicePtr(ur_device_handle_t, device_access_mode_t, size_t offset,
@@ -172,7 +171,7 @@ struct ur_mem_sub_buffer_t : public ur_mem_handle_t_ {
   ur_shared_mutex &getMutex() override;
 
 private:
-  ur_mem_handle_t hParent;
+  v2::raii::rc<ur_mem_handle_t> hParent;
   size_t offset;
   size_t size;
 };