[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,83 @@ 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>
+ur_result_t validateRetain([[maybe_unused]] URHandle handle) {
+  assert(reinterpret_cast<_ur_object *>(handle)->RefCount.load() != 0);
+  return UR_RESULT_SUCCESS;
+}
+
+template <typename URHandle>
+ur_result_t validateRelease([[maybe_unused]] URHandle handle) {
+  assert(reinterpret_cast<_ur_object *>(handle)->RefCount.load() != 0);
+  return UR_RESULT_SUCCESS;
+}
+
+// Devices are owned by the platform, so we don't need to retain/release them
+// as long as the platform is alive.
+using ur_device_handle_t =
+    ref_counted<::ur_device_handle_t, validateRetain<::ur_device_handle_t>,
+                validateRelease<::ur_device_handle_t>>;
+
+// Spec requires that the context is not destroyed until all entities
+// using the context are destroyed.
+using ur_context_handle_t =
+    ref_counted<::ur_context_handle_t, validateRetain<::ur_context_handle_t>,
+                validateRelease<::ur_context_handle_t>>;
+
+using ur_mem_handle_t =
+    ref_counted<::ur_mem_handle_t, ur::level_zero::urMemRetain, urMemRelease>;
+
+using ur_program_handle_t =
+    ref_counted<::ur_program_handle_t, ur::level_zero::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,14 @@ struct ur_event_handle_t_ : _ur_object {
   uint64_t getEventEndTimestamp();
 
 protected:
+  // Do not use ref couting on context to avoid circular dependency.
   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,6 +50,7 @@ class event_pool {
   event_flags_t getFlags() const;
 
 private:
+  // Do not use ref couting on context to avoid circular dependency.
   ur_context_handle_t hContext;
   std::unique_ptr<event_provider> provider;
 

source/adapters/level_zero/v2/event_pool_cache.hpp

@@ -41,6 +41,7 @@ class event_pool_cache {
   raii::cache_borrowed_event_pool borrow(DeviceId, event_flags_t flags);
 
 private:
+  // Do not use ref couting on context to avoid circular dependency.
   ur_context_handle_t hContext;
   ur_mutex mutex;
   ProviderCreateFunc providerCreate;

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::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::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::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::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::ur_mem_handle_t hParent;
   size_t offset;
   size_t size;
 };

source/adapters/level_zero/v2/queue_immediate_in_order.cpp

@@ -121,9 +121,9 @@ ur_queue_immediate_in_order_t::queueGetInfo(ur_queue_info_t propName,
   // TODO: consider support for queue properties and size
   switch ((uint32_t)propName) { // cast to avoid warnings on EXT enum values
   case UR_QUEUE_INFO_CONTEXT:
-    return ReturnValue(hContext);
+    return ReturnValue(hContext.get());
   case UR_QUEUE_INFO_DEVICE:
-    return ReturnValue(hDevice);
+    return ReturnValue(hDevice.get());
   case UR_QUEUE_INFO_REFERENCE_COUNT:
     return ReturnValue(uint32_t{RefCount.load()});
   case UR_QUEUE_INFO_FLAGS:

source/adapters/level_zero/v2/queue_immediate_in_order.hpp

@@ -36,8 +36,8 @@ struct ur_command_list_handler_t {
 
 struct ur_queue_immediate_in_order_t : _ur_object, public ur_queue_handle_t_ {
 private:
-  ur_context_handle_t hContext;
-  ur_device_handle_t hDevice;
+  v2::raii::ur_context_handle_t hContext;
+  v2::raii::ur_device_handle_t hDevice;
   ur_queue_flags_t flags;
 
   raii::cache_borrowed_event_pool eventPool;