Merge pull request #1097 from PietroGhg/pietro/refcounting

[NATIVECPU] Implement correct reference counting for Native CPU adapter
oneapi-src · Nov 24, 2023 · b4e665f · b4e665f
2 parents e197941 + b7731bc
commit b4e665f
Show file tree

Hide file tree

Showing 7 changed files with 21 additions and 12 deletions.
diff --git a/source/adapters/native_cpu/common.hpp b/source/adapters/native_cpu/common.hpp
@@ -61,10 +61,16 @@ struct _ur_object {
   ur_shared_mutex Mutex;
 };
 
+// Todo: replace this with a common helper once it is available
 struct RefCounted {
   std::atomic_uint32_t _refCount;
-  void incrementReferenceCount() { _refCount++; }
-  void decrementReferenceCount() { _refCount--; }
+  uint32_t incrementReferenceCount() { return ++_refCount; }
+  uint32_t decrementReferenceCount() { return --_refCount; }
   RefCounted() : _refCount{1} {}
   uint32_t getReferenceCount() const { return _refCount; }
 };
+
+template <typename T> inline void decrementOrDelete(T *refC) {
+  if (refC->decrementReferenceCount() == 0)
+    delete refC;
+}
diff --git a/source/adapters/native_cpu/context.cpp b/source/adapters/native_cpu/context.cpp
@@ -32,13 +32,13 @@ urContextCreate(uint32_t DeviceCount, const ur_device_handle_t *phDevices,
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urContextRetain(ur_context_handle_t hContext) {
-  std::ignore = hContext;
-  DIE_NO_IMPLEMENTATION
+  hContext->incrementReferenceCount();
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urContextRelease(ur_context_handle_t hContext) {
-  delete hContext;
+  decrementOrDelete(hContext);
   return UR_RESULT_SUCCESS;
 }
 

diff --git a/source/adapters/native_cpu/context.hpp b/source/adapters/native_cpu/context.hpp
@@ -12,9 +12,10 @@
 
 #include <ur_api.h>
 
+#include "common.hpp"
 #include "device.hpp"
 
-struct ur_context_handle_t_ {
+struct ur_context_handle_t_ : RefCounted {
   ur_context_handle_t_(ur_device_handle_t_ *phDevices) : _device{phDevices} {}
 
   ur_device_handle_t _device;

diff --git a/source/adapters/native_cpu/kernel.cpp b/source/adapters/native_cpu/kernel.cpp
@@ -182,7 +182,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urKernelRetain(ur_kernel_handle_t hKernel) {
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urKernelRelease(ur_kernel_handle_t hKernel) {
-  delete hKernel;
+  decrementOrDelete(hKernel);
 
   return UR_RESULT_SUCCESS;
 }

diff --git a/source/adapters/native_cpu/program.cpp b/source/adapters/native_cpu/program.cpp
@@ -115,8 +115,7 @@ urProgramRetain(ur_program_handle_t hProgram) {
 
 UR_APIEXPORT ur_result_t UR_APICALL
 urProgramRelease(ur_program_handle_t hProgram) {
-  delete hProgram;
-
+  decrementOrDelete(hProgram);
   return UR_RESULT_SUCCESS;
 }
 

diff --git a/source/adapters/native_cpu/queue.cpp b/source/adapters/native_cpu/queue.cpp
@@ -43,12 +43,14 @@ UR_APIEXPORT ur_result_t UR_APICALL urQueueCreate(
 
 UR_APIEXPORT ur_result_t UR_APICALL urQueueRetain(ur_queue_handle_t hQueue) {
   std::ignore = hQueue;
+  hQueue->incrementReferenceCount();
 
-  DIE_NO_IMPLEMENTATION;
+  return UR_RESULT_SUCCESS;
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urQueueRelease(ur_queue_handle_t hQueue) {
-  delete hQueue;
+  decrementOrDelete(hQueue);
+
   return UR_RESULT_SUCCESS;
 }
 

diff --git a/source/adapters/native_cpu/queue.hpp b/source/adapters/native_cpu/queue.hpp
@@ -8,5 +8,6 @@
 //
 //===----------------------------------------------------------------------===//
 #pragma once
+#include "common.hpp"
 
-struct ur_queue_handle_t_ {};
+struct ur_queue_handle_t_ : RefCounted {};