[UR][NativeCPU] Refactor UR Native CPU reference counting

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

@@ -10,17 +10,18 @@
 
 #include "adapter.hpp"
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "ur_api.h"
 
 struct ur_adapter_handle_t_ : ur::native_cpu::handle_base {
-  std::atomic<uint32_t> RefCount = 0;
   logger::Logger &logger = logger::get_logger("native_cpu");
+  ur::RefCount RefCount;
 } Adapter;
 
 UR_APIEXPORT ur_result_t UR_APICALL urAdapterGet(
     uint32_t, ur_adapter_handle_t *phAdapters, uint32_t *pNumAdapters) {
   if (phAdapters) {
-    Adapter.RefCount++;
+    Adapter.RefCount.retain();
     *phAdapters = &Adapter;
   }
   if (pNumAdapters) {
@@ -30,12 +31,12 @@ UR_APIEXPORT ur_result_t UR_APICALL urAdapterGet(
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urAdapterRelease(ur_adapter_handle_t) {
-  Adapter.RefCount--;
+  Adapter.RefCount.release();
   return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urAdapterRetain(ur_adapter_handle_t) {
-  Adapter.RefCount++;
+  Adapter.RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
@@ -57,7 +58,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urAdapterGetInfo(ur_adapter_handle_t,
   case UR_ADAPTER_INFO_BACKEND:
     return ReturnValue(UR_BACKEND_NATIVE_CPU);
   case UR_ADAPTER_INFO_REFERENCE_COUNT:
-    return ReturnValue(Adapter.RefCount.load());
+    return ReturnValue(Adapter.RefCount.getCount());
   case UR_ADAPTER_INFO_VERSION:
     return ReturnValue(uint32_t{1});
   default:

unified-runtime/source/adapters/native_cpu/common.hpp

@@ -44,22 +44,13 @@ struct ddi_getter {
 using handle_base = ur::handle_base<ddi_getter>;
 } // namespace ur::native_cpu
 
-// Todo: replace this with a common helper once it is available
-struct RefCounted : ur::native_cpu::handle_base {
-  std::atomic_uint32_t _refCount;
-  uint32_t incrementReferenceCount() { return ++_refCount; }
-  uint32_t decrementReferenceCount() { return --_refCount; }
-  RefCounted() : handle_base(), _refCount{1} {}
-  uint32_t getReferenceCount() const { return _refCount; }
-};
-
 // Base class to store common data
-struct ur_object : RefCounted {
+struct ur_object {
   ur_shared_mutex Mutex;
 };
 
 template <typename T> inline void decrementOrDelete(T *refC) {
-  if (refC->decrementReferenceCount() == 0)
+  if (refC->RefCount.release() == 0)
     delete refC;
 }
 

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

@@ -30,7 +30,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urContextCreate(
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urContextRetain(ur_context_handle_t hContext) {
-  hContext->incrementReferenceCount();
+  hContext->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
@@ -51,7 +51,7 @@ urContextGetInfo(ur_context_handle_t hContext, ur_context_info_t propName,
   case UR_CONTEXT_INFO_DEVICES:
     return returnValue(hContext->_device);
   case UR_CONTEXT_INFO_REFERENCE_COUNT:
-    return returnValue(uint32_t{hContext->getReferenceCount()});
+    return returnValue(uint32_t{hContext->RefCount.getCount()});
   case UR_CONTEXT_INFO_USM_MEMCPY2D_SUPPORT:
     return returnValue(true);
   case UR_CONTEXT_INFO_USM_FILL2D_SUPPORT:

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

@@ -15,6 +15,7 @@
 #include <ur_api.h>
 
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "device.hpp"
 #include "ur/ur.hpp"
 
@@ -83,7 +84,7 @@ static usm_alloc_info get_alloc_info(void *ptr) {
 
 } // namespace native_cpu
 
-struct ur_context_handle_t_ : RefCounted {
+struct ur_context_handle_t_ {
   ur_context_handle_t_(ur_device_handle_t_ *phDevices) : _device{phDevices} {}
 
   ur_device_handle_t _device;
@@ -135,6 +136,8 @@ struct ur_context_handle_t_ : RefCounted {
     return ptr;
   }
 
+  ur::RefCount RefCount;
+
 private:
   std::mutex alloc_mutex;
   std::set<const void *> allocations;

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

@@ -28,7 +28,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hEvent,
   case UR_EVENT_INFO_COMMAND_TYPE:
     return ReturnValue(hEvent->getCommandType());
   case UR_EVENT_INFO_REFERENCE_COUNT:
-    return ReturnValue(hEvent->getReferenceCount());
+    return ReturnValue(hEvent->RefCount.getCount());
   case UR_EVENT_INFO_COMMAND_EXECUTION_STATUS:
     return ReturnValue(hEvent->getExecutionStatus());
   case UR_EVENT_INFO_CONTEXT:
@@ -69,7 +69,7 @@ urEventWait(uint32_t numEvents, const ur_event_handle_t *phEventWaitList) {
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urEventRetain(ur_event_handle_t hEvent) {
-  hEvent->incrementReferenceCount();
+  hEvent->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 

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

@@ -8,14 +8,17 @@
 //
 //===----------------------------------------------------------------------===//
 #pragma once
-#include "common.hpp"
-#include "ur_api.h"
+
 #include <cstdint>
 #include <future>
 #include <mutex>
 #include <vector>
 
-struct ur_event_handle_t_ : RefCounted {
+#include "common.hpp"
+#include "common/ur_ref_count.hpp"
+#include "ur_api.h"
+
+struct ur_event_handle_t_ {
 
   ur_event_handle_t_(ur_queue_handle_t queue, ur_command_t command_type);
 
@@ -55,6 +58,8 @@ struct ur_event_handle_t_ : RefCounted {
 
   uint64_t get_end_timestamp() const { return timestamp_end; }
 
+  ur::RefCount RefCount;
+
 private:
   ur_queue_handle_t queue;
   ur_context_handle_t context;

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

@@ -95,7 +95,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetInfo(ur_kernel_handle_t hKernel,
   case UR_KERNEL_INFO_FUNCTION_NAME:
     return ReturnValue(hKernel->_name.c_str());
   case UR_KERNEL_INFO_REFERENCE_COUNT:
-    return ReturnValue(uint32_t{hKernel->getReferenceCount()});
+    return ReturnValue(uint32_t{hKernel->RefCount.getCount()});
   case UR_KERNEL_INFO_ATTRIBUTES:
     return ReturnValue("");
   case UR_KERNEL_INFO_SPILL_MEM_SIZE:
@@ -194,7 +194,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelGetSubGroupInfo(
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urKernelRetain(ur_kernel_handle_t hKernel) {
-  hKernel->incrementReferenceCount();
+  hKernel->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 

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

@@ -8,13 +8,15 @@
 
 #pragma once
 
+#include <cstring>
+#include <utility>
+
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "memory.hpp"
 #include "nativecpu_state.hpp"
 #include "program.hpp"
-#include <cstring>
 #include <ur_api.h>
-#include <utility>
 
 using nativecpu_kernel_t = void(void *const *, native_cpu::state *);
 using nativecpu_ptr_t = nativecpu_kernel_t *;
@@ -27,7 +29,7 @@ struct local_arg_info_t {
       : argIndex(argIndex), argSize(argSize) {}
 };
 
-struct ur_kernel_handle_t_ : RefCounted {
+struct ur_kernel_handle_t_ {
 
   ur_kernel_handle_t_(ur_program_handle_t hProgram, const char *name,
                       nativecpu_task_t subhandler)
@@ -188,10 +190,12 @@ struct ur_kernel_handle_t_ : RefCounted {
   void addPtrArg(void *Ptr, size_t Index) { Args.addPtrArg(Index, Ptr); }
 
   void addArgReference(ur_mem_handle_t Arg) {
-    Arg->incrementReferenceCount();
+    Arg->RefCount.getCount();
     ReferencedArgs.push_back(Arg);
   }
 
+  ur::RefCount RefCount;
+
 private:
   void removeArgReferences() {
     for (auto arg : ReferencedArgs)

unified-runtime/source/adapters/native_cpu/memory.hpp

@@ -15,6 +15,7 @@
 #include <cstring>
 
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "context.hpp"
 
 struct ur_mem_handle_t_ : ur_object {
@@ -43,6 +44,8 @@ struct ur_mem_handle_t_ : ur_object {
   char *_mem;
   bool _ownsMem;
 
+  ur::RefCount RefCount;
+
 private:
   const bool IsImage;
 };

unified-runtime/source/adapters/native_cpu/program.cpp

@@ -171,7 +171,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urProgramLinkExp(
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urProgramRetain(ur_program_handle_t hProgram) {
-  hProgram->incrementReferenceCount();
+  hProgram->RefCount.retain();
   return UR_RESULT_SUCCESS;
 }
 
@@ -205,7 +205,7 @@ urProgramGetInfo(ur_program_handle_t hProgram, ur_program_info_t propName,
 
   switch (propName) {
   case UR_PROGRAM_INFO_REFERENCE_COUNT:
-    return returnValue(hProgram->getReferenceCount());
+    return returnValue(hProgram->RefCount.getCount());
   case UR_PROGRAM_INFO_CONTEXT:
     return returnValue(nullptr);
   case UR_PROGRAM_INFO_NUM_DEVICES:

unified-runtime/source/adapters/native_cpu/program.hpp

@@ -10,23 +10,22 @@
 
 #pragma once
 
+#include <array>
+#include <map>
+
 #include <ur_api.h>
 
+#include "common/ur_ref_count.hpp"
 #include "context.hpp"
 
-#include <array>
-#include <map>
-
 namespace native_cpu {
 using WGSize_t = std::array<uint32_t, 3>;
 }
 
-struct ur_program_handle_t_ : RefCounted {
+struct ur_program_handle_t_ {
   ur_program_handle_t_(ur_context_handle_t ctx, const unsigned char *pBinary)
       : _ctx{ctx}, _ptr{pBinary} {}
 
-  uint32_t getReferenceCount() const noexcept { return _refCount; }
-
   ur_context_handle_t _ctx;
   const unsigned char *_ptr;
   struct _compare {
@@ -41,6 +40,8 @@ struct ur_program_handle_t_ : RefCounted {
   std::unordered_map<std::string, native_cpu::WGSize_t>
       KernelMaxWorkGroupSizeMD;
   std::unordered_map<std::string, uint64_t> KernelMaxLinearWorkGroupSizeMD;
+
+  ur::RefCount RefCount;
 };
 
 // The nativecpu_entry struct is also defined as LLVM-IR in the

unified-runtime/source/adapters/native_cpu/queue.cpp

@@ -28,7 +28,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueGetInfo(ur_queue_handle_t hQueue,
   case UR_QUEUE_INFO_DEVICE:
     return ReturnValue(hQueue->getDevice());
   case UR_QUEUE_INFO_REFERENCE_COUNT:
-    return ReturnValue(hQueue->getReferenceCount());
+    return ReturnValue(hQueue->RefCount.getCount());
   case UR_QUEUE_INFO_EMPTY:
     return ReturnValue(hQueue->isEmpty());
   default:
@@ -48,7 +48,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueCreate(
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urQueueRetain(ur_queue_handle_t hQueue) {
-  hQueue->incrementReferenceCount();
+  hQueue->RefCount.retain();
 
   return UR_RESULT_SUCCESS;
 }

unified-runtime/source/adapters/native_cpu/queue.hpp

@@ -8,12 +8,15 @@
 //
 //===----------------------------------------------------------------------===//
 #pragma once
+
+#include <set>
+
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include "event.hpp"
 #include "ur_api.h"
-#include <set>
 
-struct ur_queue_handle_t_ : RefCounted {
+struct ur_queue_handle_t_ {
   ur_queue_handle_t_(ur_device_handle_t device, ur_context_handle_t context,
                      const ur_queue_properties_t *pProps)
       : device(device), context(context),
@@ -43,7 +46,7 @@ struct ur_queue_handle_t_ : RefCounted {
       auto ev = *events.begin();
       // ur_event_handle_t_::wait removes itself from the events set in the
       // queue.
-      ev->incrementReferenceCount();
+      ev->RefCount.retain();
       // Unlocking mutex for removeEvent and for event callbacks that may need
       // to acquire it.
       lock.unlock();
@@ -64,6 +67,8 @@ struct ur_queue_handle_t_ : RefCounted {
     return events.size() == 0;
   }
 
+  ur::RefCount RefCount;
+
 private:
   ur_device_handle_t device;
   ur_context_handle_t context;