[UR][OpenCL] Refactor UR OpenCL reference counting

unified-runtime/source/adapters/opencl/adapter.cpp

@@ -78,7 +78,7 @@ urAdapterGet(uint32_t NumEntries, ur_adapter_handle_t *phAdapters,
     }
 
     auto &adapter = *phAdapters;
-    adapter->RefCount++;
+    adapter->RefCount.retain();
   }
 
   if (pNumAdapters) {
@@ -90,13 +90,13 @@ urAdapterGet(uint32_t NumEntries, ur_adapter_handle_t *phAdapters,
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urAdapterRetain(ur_adapter_handle_t hAdapter) {
-  ++hAdapter->RefCount;
+  hAdapter->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urAdapterRelease(ur_adapter_handle_t hAdapter) {
-  if (--hAdapter->RefCount == 0) {
+  if (hAdapter->RefCount.release()) {
     delete hAdapter;
   }
   return UR_RESULT_SUCCESS;
@@ -119,7 +119,7 @@ urAdapterGetInfo(ur_adapter_handle_t hAdapter, ur_adapter_info_t propName,
   case UR_ADAPTER_INFO_BACKEND:
     return ReturnValue(UR_BACKEND_OPENCL);
   case UR_ADAPTER_INFO_REFERENCE_COUNT:
-    return ReturnValue(hAdapter->RefCount.load());
+    return ReturnValue(hAdapter->RefCount.getCount());
   case UR_ADAPTER_INFO_VERSION:
     return ReturnValue(uint32_t{1});
   default:

unified-runtime/source/adapters/opencl/adapter.hpp

@@ -15,6 +15,7 @@
 
 #include "CL/cl.h"
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "logger/ur_logger.hpp"
 
 struct ur_adapter_handle_t_ : ur::opencl::handle_base {
@@ -24,13 +25,14 @@ struct ur_adapter_handle_t_ : ur::opencl::handle_base {
   ur_adapter_handle_t_(ur_adapter_handle_t_ &) = delete;
   ur_adapter_handle_t_ &operator=(const ur_adapter_handle_t_ &) = delete;
 
-  std::atomic<uint32_t> RefCount = 0;
   logger::Logger &log = logger::get_logger("opencl");
   cl_ext::ExtFuncPtrCacheT fnCache{};
 
   std::vector<std::unique_ptr<ur_platform_handle_t_>> URPlatforms;
   uint32_t NumPlatforms = 0;
 
+  ur::RefCount RefCount;
+
   // Function pointers to core OpenCL entry points which may not exist in older
   // versions of the OpenCL-ICD-Loader are tracked here and initialized by
   // dynamically loading the symbol by name.

unified-runtime/source/adapters/opencl/command_buffer.cpp

@@ -108,13 +108,13 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferCreateExp(
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urCommandBufferRetainExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
-  hCommandBuffer->incrementReferenceCount();
+  hCommandBuffer->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t hCommandBuffer) {
-  if (hCommandBuffer->decrementReferenceCount() == 0) {
+  if (hCommandBuffer->RefCount.release()) {
     delete hCommandBuffer;
   }
 
@@ -783,7 +783,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urCommandBufferGetInfoExp(
 
   switch (propName) {
   case UR_EXP_COMMAND_BUFFER_INFO_REFERENCE_COUNT:
-    return ReturnValue(hCommandBuffer->getReferenceCount());
+    return ReturnValue(hCommandBuffer->RefCount.getCount());
   case UR_EXP_COMMAND_BUFFER_INFO_DESCRIPTOR: {
     ur_exp_command_buffer_desc_t Descriptor{};
     Descriptor.stype = UR_STRUCTURE_TYPE_EXP_COMMAND_BUFFER_DESC;

unified-runtime/source/adapters/opencl/command_buffer.hpp

@@ -9,6 +9,7 @@
 //===----------------------------------------------------------------------===//
 
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include <ur/ur.hpp>
 
 /// Handle to a kernel command.
@@ -53,11 +54,11 @@ struct ur_exp_command_buffer_handle_t_ : ur::opencl::handle_base {
   /// List of commands in the command-buffer.
   std::vector<std::unique_ptr<ur_exp_command_buffer_command_handle_t_>>
       CommandHandles;
-  /// Object reference count
-  std::atomic_uint32_t RefCount;
   /// Track last submission of the command-buffer
   cl_event LastSubmission;
 
+  ur::RefCount RefCount;
+
   ur_exp_command_buffer_handle_t_(ur_queue_handle_t hQueue,
                                   ur_context_handle_t hContext,
                                   ur_device_handle_t hDevice,
@@ -66,11 +67,7 @@ struct ur_exp_command_buffer_handle_t_ : ur::opencl::handle_base {
       : handle_base(), hInternalQueue(hQueue), hContext(hContext),
         hDevice(hDevice), CLCommandBuffer(CLCommandBuffer),
         IsUpdatable(IsUpdatable), IsInOrder(IsInOrder), IsFinalized(false),
-        RefCount(0), LastSubmission(nullptr) {}
+        LastSubmission(nullptr) {}
 
   ~ur_exp_command_buffer_handle_t_();
-
-  uint32_t incrementReferenceCount() noexcept { return ++RefCount; }
-  uint32_t decrementReferenceCount() noexcept { return --RefCount; }
-  uint32_t getReferenceCount() const noexcept { return RefCount; }
 };

unified-runtime/source/adapters/opencl/context.cpp

@@ -108,7 +108,7 @@ urContextGetInfo(ur_context_handle_t hContext, ur_context_info_t propName,
     return ReturnValue(&hContext->Devices[0], hContext->DeviceCount);
   }
   case UR_CONTEXT_INFO_REFERENCE_COUNT: {
-    return ReturnValue(hContext->getReferenceCount());
+    return ReturnValue(hContext->RefCount.getCount());
   }
   default:
     return UR_RESULT_ERROR_INVALID_ENUMERATION;
@@ -117,7 +117,7 @@ urContextGetInfo(ur_context_handle_t hContext, ur_context_info_t propName,
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urContextRelease(ur_context_handle_t hContext) {
-  if (hContext->decrementReferenceCount() == 0) {
+  if (hContext->RefCount.release()) {
     delete hContext;
   }
 
@@ -126,7 +126,7 @@ urContextRelease(ur_context_handle_t hContext) {
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urContextRetain(ur_context_handle_t hContext) {
-  hContext->incrementReferenceCount();
+  hContext->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 

unified-runtime/source/adapters/opencl/context.hpp

@@ -11,6 +11,7 @@
 
 #include "adapter.hpp"
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "device.hpp"
 
 #include <vector>
@@ -20,8 +21,8 @@ struct ur_context_handle_t_ : ur::opencl::handle_base {
   native_type CLContext;
   std::vector<ur_device_handle_t> Devices;
   uint32_t DeviceCount;
-  std::atomic<uint32_t> RefCount = 0;
   bool IsNativeHandleOwned = true;
+  ur::RefCount RefCount;
 
   ur_context_handle_t_(native_type Ctx, uint32_t DevCount,
                        const ur_device_handle_t *phDevices)
@@ -30,15 +31,8 @@ struct ur_context_handle_t_ : ur::opencl::handle_base {
       Devices.emplace_back(phDevices[i]);
       urDeviceRetain(phDevices[i]);
     }
-    RefCount = 1;
   }
 
-  uint32_t incrementReferenceCount() noexcept { return ++RefCount; }
-
-  uint32_t decrementReferenceCount() noexcept { return --RefCount; }
-
-  uint32_t getReferenceCount() const noexcept { return RefCount; }
-
   static ur_result_t makeWithNative(native_type Ctx, uint32_t DevCount,
                                     const ur_device_handle_t *phDevices,
                                     ur_context_handle_t &Context);

unified-runtime/source/adapters/opencl/device.cpp

@@ -1019,7 +1019,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceGetInfo(ur_device_handle_t hDevice,
     return UR_RESULT_SUCCESS;
   }
   case UR_DEVICE_INFO_REFERENCE_COUNT: {
-    return ReturnValue(hDevice->getReferenceCount());
+    return ReturnValue(hDevice->RefCount.getCount());
   }
   case UR_DEVICE_INFO_PARTITION_MAX_SUB_DEVICES: {
     CL_RETURN_ON_FAILURE(clGetDeviceInfo(hDevice->CLDevice,
@@ -1561,7 +1561,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDevicePartition(
 // Root devices ref count are unchanged through out the program lifetime.
 UR_APIEXPORT ur_result_t UR_APICALL urDeviceRetain(ur_device_handle_t hDevice) {
   if (hDevice->ParentDevice) {
-    hDevice->incrementReferenceCount();
+    hDevice->RefCount.retain();
   }
 
   return UR_RESULT_SUCCESS;
@@ -1571,7 +1571,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urDeviceRetain(ur_device_handle_t hDevice) {
 UR_APIEXPORT ur_result_t UR_APICALL
 urDeviceRelease(ur_device_handle_t hDevice) {
   if (hDevice->ParentDevice) {
-    if (hDevice->decrementReferenceCount() == 0) {
+    if (hDevice->RefCount.release()) {
       delete hDevice;
     }
   }

unified-runtime/source/adapters/opencl/device.hpp

@@ -10,6 +10,7 @@
 #pragma once
 
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "device.hpp"
 #include "platform.hpp"
 
@@ -19,13 +20,12 @@ struct ur_device_handle_t_ : ur::opencl::handle_base {
   ur_platform_handle_t Platform;
   cl_device_type Type = 0;
   ur_device_handle_t ParentDevice = nullptr;
-  std::atomic<uint32_t> RefCount = 0;
   bool IsNativeHandleOwned = true;
+  ur::RefCount RefCount;
 
   ur_device_handle_t_(native_type Dev, ur_platform_handle_t Plat,
                       ur_device_handle_t Parent)
       : handle_base(), CLDevice(Dev), Platform(Plat), ParentDevice(Parent) {
-    RefCount = 1;
     if (Parent) {
       Type = Parent->Type;
       [[maybe_unused]] auto Res = clRetainDevice(CLDevice);
@@ -51,12 +51,6 @@ struct ur_device_handle_t_ : ur::opencl::handle_base {
     }
   }
 
-  uint32_t incrementReferenceCount() noexcept { return ++RefCount; }
-
-  uint32_t decrementReferenceCount() noexcept { return --RefCount; }
-
-  uint32_t getReferenceCount() const noexcept { return RefCount; }
-
   ur_result_t getDeviceVersion(oclv::OpenCLVersion &Version) {
     size_t DevVerSize = 0;
     CL_RETURN_ON_FAILURE(

unified-runtime/source/adapters/opencl/event.cpp

@@ -136,14 +136,14 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventGetNativeHandle(
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urEventRelease(ur_event_handle_t hEvent) {
-  if (hEvent->decrementReferenceCount() == 0) {
+  if (hEvent->RefCount.release()) {
     delete hEvent;
   }
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urEventRetain(ur_event_handle_t hEvent) {
-  hEvent->incrementReferenceCount();
+  hEvent->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
@@ -188,7 +188,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hEvent,
     return ReturnValue(hEvent->Queue);
   }
   case UR_EVENT_INFO_REFERENCE_COUNT: {
-    return ReturnValue(hEvent->getReferenceCount());
+    return ReturnValue(hEvent->RefCount.getCount());
   }
   default: {
     size_t CheckPropSize = 0;

unified-runtime/source/adapters/opencl/event.hpp

@@ -10,6 +10,7 @@
 #pragma once
 
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "queue.hpp"
 
 #include <vector>
@@ -19,13 +20,12 @@ struct ur_event_handle_t_ : ur::opencl::handle_base {
   native_type CLEvent;
   ur_context_handle_t Context;
   ur_queue_handle_t Queue;
-  std::atomic<uint32_t> RefCount = 0;
   bool IsNativeHandleOwned = true;
+  ur::RefCount RefCount;
 
   ur_event_handle_t_(native_type Event, ur_context_handle_t Ctx,
                      ur_queue_handle_t Queue)
       : handle_base(), CLEvent(Event), Context(Ctx), Queue(Queue) {
-    RefCount = 1;
     urContextRetain(Context);
     if (Queue) {
       urQueueRetain(Queue);
@@ -42,12 +42,6 @@ struct ur_event_handle_t_ : ur::opencl::handle_base {
     }
   }
 
-  uint32_t incrementReferenceCount() noexcept { return ++RefCount; }
-
-  uint32_t decrementReferenceCount() noexcept { return --RefCount; }
-
-  uint32_t getReferenceCount() const noexcept { return RefCount; }
-
   ur_result_t ensureQueue() {
     if (!Queue) {
       cl_command_queue native_queue;

unified-runtime/source/adapters/opencl/kernel.cpp

@@ -152,7 +152,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
     return ReturnValue(hKernel->Context);
   }
   case UR_KERNEL_INFO_REFERENCE_COUNT: {
-    return ReturnValue(hKernel->getReferenceCount());
+    return ReturnValue(hKernel->RefCount.getCount());
   }
   default: {
     size_t CheckPropSize = 0;
@@ -343,13 +343,13 @@ urKernelGetSubGroupInfo(ur_kernel_handle_t hKernel, ur_device_handle_t hDevice,
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelRetain(ur_kernel_handle_t hKernel) {
-  hKernel->incrementReferenceCount();
+  hKernel->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelRelease(ur_kernel_handle_t hKernel) {
-  if (hKernel->decrementReferenceCount() == 0) {
+  if (hKernel->RefCount.release()) {
     delete hKernel;
   }
   return UR_RESULT_SUCCESS;

unified-runtime/source/adapters/opencl/kernel.hpp

@@ -11,6 +11,7 @@
 
 #include "adapter.hpp"
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "context.hpp"
 #include "program.hpp"
 
@@ -21,14 +22,13 @@ struct ur_kernel_handle_t_ : ur::opencl::handle_base {
   native_type CLKernel;
   ur_program_handle_t Program;
   ur_context_handle_t Context;
-  std::atomic<uint32_t> RefCount = 0;
   bool IsNativeHandleOwned = true;
   clSetKernelArgMemPointerINTEL_fn clSetKernelArgMemPointerINTEL = nullptr;
+  ur::RefCount RefCount;
 
   ur_kernel_handle_t_(native_type Kernel, ur_program_handle_t Program,
                       ur_context_handle_t Context)
       : handle_base(), CLKernel(Kernel), Program(Program), Context(Context) {
-    RefCount = 1;
     urProgramRetain(Program);
     urContextRetain(Context);
 
@@ -46,12 +46,6 @@ struct ur_kernel_handle_t_ : ur::opencl::handle_base {
     }
   }
 
-  uint32_t incrementReferenceCount() noexcept { return ++RefCount; }
-
-  uint32_t decrementReferenceCount() noexcept { return --RefCount; }
-
-  uint32_t getReferenceCount() const noexcept { return RefCount; }
-
   static ur_result_t makeWithNative(native_type NativeKernel,
                                     ur_program_handle_t Program,
                                     ur_context_handle_t Context,

unified-runtime/source/adapters/opencl/memory.cpp

@@ -521,7 +521,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemGetInfo(ur_mem_handle_t hMemory,
     return ReturnValue(hMemory->Context);
   }
   case UR_MEM_INFO_REFERENCE_COUNT: {
-    return ReturnValue(hMemory->getReferenceCount());
+    return ReturnValue(hMemory->RefCount.getCount());
   }
   default: {
     size_t CheckPropSize = 0;
@@ -569,12 +569,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urMemImageGetInfo(ur_mem_handle_t hMemory,
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urMemRetain(ur_mem_handle_t hMem) {
-  hMem->incrementReferenceCount();
+  hMem->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urMemRelease(ur_mem_handle_t hMem) {
-  if (hMem->decrementReferenceCount() == 0) {
+  if (hMem->RefCount.release()) {
     delete hMem;
   }
   return UR_RESULT_SUCCESS;