wip

Author: martygrant

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

@@ -62,7 +62,7 @@ ur_exp_command_buffer_handle_t_::ur_exp_command_buffer_handle_t_(
     bool IsInOrder)
     : handle_base(), Context(Context), Device(Device), IsUpdatable(IsUpdatable),
       IsInOrder(IsInOrder), CudaGraph{nullptr}, CudaGraphExec{nullptr},
-      RefCount{1}, NextSyncPoint{0} {
+      NextSyncPoint{0} {
   urContextRetain(Context);
 }
 

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

@@ -186,7 +186,8 @@ struct ur_exp_command_buffer_handle_t_ : ur::cuda::handle_base {
   CUgraph CudaGraph;
   // Cuda Graph Exec handle
   CUgraphExec CudaGraphExec = nullptr;
-
+  // Atomic variable counting the number of reference to this command_buffer
+  // using std::atomic prevents data race when incrementing/decrementing.
   ur::RefCount RefCount;
 
   // Ordered map of sync_points to ur_events, so that we can find the last

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

@@ -85,10 +85,9 @@ UR_APIEXPORT ur_result_t UR_APICALL urContextGetInfo(
 UR_APIEXPORT ur_result_t UR_APICALL
 urContextRelease(ur_context_handle_t hContext) {
   if (hContext->RefCount.release()) {
-    return UR_RESULT_SUCCESS;
+    hContext->invokeExtendedDeleters();
+    delete hContext;
   }
-  hContext->invokeExtendedDeleters();
-  delete hContext;
 
   return UR_RESULT_SUCCESS;
 }

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

@@ -97,7 +97,7 @@ struct ur_context_handle_t_ : ur::cuda::handle_base {
   umf_memory_pool_handle_t MemoryPoolHost = nullptr;
 
   ur_context_handle_t_(const ur_device_handle_t *Devs, uint32_t NumDevices)
-      : handle_base(), Devices{Devs, Devs + NumDevices}, RefCount(0) {
+      : handle_base(), Devices{Devs, Devs + NumDevices} {
     // Create UMF CUDA memory provider for the host memory
     // (UMF_MEMORY_TYPE_HOST) from any device (Devices[0] is used here, because
     // it is guaranteed to exist).

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

@@ -265,7 +265,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventRelease(ur_event_handle_t hEvent) {
   }
 
   // decrement ref count. If it is 0, delete the event.
-  if (hEvent->release()) {
+  if (hEvent->RefCount.release()) {
     std::unique_ptr<ur_event_handle_t_> event_ptr{hEvent};
     ur_result_t Result = UR_RESULT_ERROR_INVALID_EVENT;
     try {

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

@@ -15,6 +15,7 @@
 #include "common.hpp"
 #include "common/ur_ref_count.hpp"
 #include "queue.hpp"
+#include "common/ur_ref_count.hpp"
 
 /// UR Event mapping to CUevent
 struct ur_event_handle_t_ : ur::cuda::handle_base {

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

@@ -315,6 +315,7 @@ struct ur_mem_handle_t_ : ur::cuda::handle_base {
   // Context where the memory object is accessible
   ur_context_handle_t Context;
 
+  /// Reference counting of the handler
   ur::RefCount RefCount;
 
   // Original mem flags passed

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

@@ -107,7 +107,7 @@ urQueueCreate(ur_context_handle_t hContext, ur_device_handle_t hDevice,
 }
 
 UR_APIEXPORT ur_result_t UR_APICALL urQueueRetain(ur_queue_handle_t hQueue) {
-  assert(hQueue->RefCount.getCount() > 0);
+  assert(hQueue->RefCount.getCount());
 
   hQueue->RefCount.retain();
   return UR_RESULT_SUCCESS;

unified-runtime/source/common/cuda-hip/stream_queue.hpp

@@ -16,6 +16,8 @@
 #include <mutex>
 #include <vector>
 
+#include "common/ur_ref_count.hpp"
+
 using ur_stream_guard = std::unique_lock<std::mutex>;
 
 /// Generic implementation of an out-of-order UR queue based on in-order