Add reference counting to handle_base struct to allow all UR handles to inherit common reference counting functionality. Update L0 and L0V2 to use this.

Author: martygrant

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

@@ -675,7 +675,7 @@ ur_result_t urAdapterGet(
     }
     *Adapters = GlobalAdapter;
 
-    if (GlobalAdapter->RefCount++ == 0) {
+    if (GlobalAdapter->incrementRefCount() == 0) {
       adapterStateInit();
     }
   }
@@ -692,7 +692,7 @@ ur_result_t urAdapterRelease([[maybe_unused]] ur_adapter_handle_t Adapter) {
 
   // NOTE: This does not require guarding with a mutex; the instant the ref
   // count hits zero, both Get and Retain are UB.
-  if (--GlobalAdapter->RefCount == 0) {
+  if (GlobalAdapter->decrementRefCount() == 0) {
     auto result = adapterStateTeardown();
 #ifdef UR_STATIC_LEVEL_ZERO
     // Given static linking of the L0 Loader, we must delay the loader's
@@ -709,9 +709,9 @@ ur_result_t urAdapterRelease([[maybe_unused]] ur_adapter_handle_t Adapter) {
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t urAdapterRetain([[maybe_unused]] ur_adapter_handle_t Adapter) {
+ur_result_t urAdapterRetain(ur_adapter_handle_t) {
   assert(GlobalAdapter && GlobalAdapter == Adapter);
-  GlobalAdapter->RefCount++;
+  GlobalAdapter->incrementRefCount();
 
   return UR_RESULT_SUCCESS;
 }
@@ -740,12 +740,12 @@ ur_result_t urAdapterGetInfo(ur_adapter_handle_t, ur_adapter_info_t PropName,
   case UR_ADAPTER_INFO_BACKEND:
     return ReturnValue(UR_BACKEND_LEVEL_ZERO);
   case UR_ADAPTER_INFO_REFERENCE_COUNT:
-    return ReturnValue(GlobalAdapter->RefCount.load());
+    return ReturnValue(GlobalAdapter->getRefCount());
   case UR_ADAPTER_INFO_VERSION: {
 #ifdef UR_ADAPTER_LEVEL_ZERO_V2
     uint32_t adapterVersion = 2;
 #else
-    uint32_t adapterVersion = 1;
+      uint32_t adapterVersion = 1;
 #endif
     return ReturnValue(adapterVersion);
   }

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

@@ -11,7 +11,6 @@
 
 #include "logger/ur_logger.hpp"
 #include "ur_interface_loader.hpp"
-#include <atomic>
 #include <loader/ur_loader.hpp>
 #include <loader/ze_loader.h>
 #include <optional>
@@ -26,7 +25,6 @@ class ur_legacy_sink;
 
 struct ur_adapter_handle_t_ : ur::handle_base<ur::level_zero::ddi_getter> {
   ur_adapter_handle_t_();
-  std::atomic<uint32_t> RefCount = 0;
 
   zes_pfnDriverGetDeviceByUuidExp_t getDeviceByUUIdFunctionPtr = nullptr;
   zes_pfnDriverGet_t getSysManDriversFunctionPtr = nullptr;

unified-runtime/source/adapters/level_zero/async_alloc.cpp

@@ -247,7 +247,7 @@ ur_result_t urEnqueueUSMFreeExp(
   }
 
   size_t size = umfPoolMallocUsableSize(hPool, Mem);
-  (*Event)->RefCount.increment();
+  (*Event)->incrementRefCount();
   usmPool->AsyncPool.insert(Mem, size, *Event, Queue);
 
   // Signal that USM free event was finished

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

@@ -840,13 +840,13 @@ urCommandBufferCreateExp(ur_context_handle_t Context, ur_device_handle_t Device,
 
 ur_result_t
 urCommandBufferRetainExp(ur_exp_command_buffer_handle_t CommandBuffer) {
-  CommandBuffer->RefCount.increment();
+  CommandBuffer->incrementRefCount();
   return UR_RESULT_SUCCESS;
 }
 
 ur_result_t
 urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t CommandBuffer) {
-  if (!CommandBuffer->RefCount.decrementAndTest())
+  if (!CommandBuffer->decrementRefCount() == 0)
     return UR_RESULT_SUCCESS;
 
   UR_CALL(waitForOngoingExecution(CommandBuffer));
@@ -1641,7 +1641,7 @@ ur_result_t enqueueImmediateAppendPath(
   if (CommandBuffer->CurrentSubmissionEvent) {
     UR_CALL(urEventReleaseInternal(CommandBuffer->CurrentSubmissionEvent));
   }
-  (*Event)->RefCount.increment();
+  (*Event)->incrementRefCount();
   CommandBuffer->CurrentSubmissionEvent = *Event;
 
   UR_CALL(Queue->executeCommandList(CommandListHelper, false, false));
@@ -1724,7 +1724,7 @@ ur_result_t enqueueWaitEventPath(ur_exp_command_buffer_handle_t CommandBuffer,
   if (CommandBuffer->CurrentSubmissionEvent) {
     UR_CALL(urEventReleaseInternal(CommandBuffer->CurrentSubmissionEvent));
   }
-  (*Event)->RefCount.increment();
+  (*Event)->incrementRefCount();
   CommandBuffer->CurrentSubmissionEvent = *Event;
 
   UR_CALL(Queue->executeCommandList(SignalCommandList, false /*IsBlocking*/,
@@ -1848,7 +1848,7 @@ urCommandBufferGetInfoExp(ur_exp_command_buffer_handle_t hCommandBuffer,
 
   switch (propName) {
   case UR_EXP_COMMAND_BUFFER_INFO_REFERENCE_COUNT:
-    return ReturnValue(uint32_t{hCommandBuffer->RefCount.load()});
+    return ReturnValue(uint32_t{hCommandBuffer->getRefCount()});
   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/level_zero/common.hpp

@@ -220,55 +220,9 @@ void zeParseError(ze_result_t ZeError, const char *&ErrorString);
 #define ZE_CALL_NOCHECK_NAME(ZeName, ZeArgs, callName)                         \
   ZeCall().doCall(ZeName ZeArgs, callName, #ZeArgs, false)
 
-// This wrapper around std::atomic is created to limit operations with reference
-// counter and to make allowed operations more transparent in terms of
-// thread-safety in the plugin. increment() and load() operations do not need a
-// mutex guard around them since the underlying data is already atomic.
-// decrementAndTest() method is used to guard a code which needs to be
-// executed when object's ref count becomes zero after release. This method also
-// doesn't need a mutex guard because decrement operation is atomic and only one
-// thread can reach ref count equal to zero, i.e. only a single thread can pass
-// through this check.
-struct ReferenceCounter {
-  ReferenceCounter() : RefCount{1} {}
-
-  // Reset the counter to the initial value.
-  void reset() { RefCount = 1; }
-
-  // Used when retaining an object.
-  void increment() { RefCount++; }
-
-  // Supposed to be used in ur*GetInfo* methods where ref count value is
-  // requested.
-  uint32_t load() { return RefCount.load(); }
-
-  // This method allows to guard a code which needs to be executed when object's
-  // ref count becomes zero after release. It is important to notice that only a
-  // single thread can pass through this check. This is true because of several
-  // reasons:
-  //   1. Decrement operation is executed atomically.
-  //   2. It is not allowed to retain an object after its refcount reaches zero.
-  //   3. It is not allowed to release an object more times than the value of
-  //   the ref count.
-  // 2. and 3. basically means that we can't use an object at all as soon as its
-  // refcount reaches zero. Using this check guarantees that code for deleting
-  // an object and releasing its resources is executed once by a single thread
-  // and we don't need to use any mutexes to guard access to this object in the
-  // scope after this check. Of course if we access another objects in this code
-  // (not the one which is being deleted) then access to these objects must be
-  // guarded, for example with a mutex.
-  bool decrementAndTest() { return --RefCount == 0; }
-
-private:
-  std::atomic<uint32_t> RefCount;
-};
-
 // Base class to store common data
 struct ur_object : ur::handle_base<ur::level_zero::ddi_getter> {
-  ur_object() : handle_base(), RefCount{} {}
-
-  // Must be atomic to prevent data race when incrementing/decrementing.
-  ReferenceCounter RefCount;
+  ur_object() : handle_base() {}
 
   // This mutex protects accesses to all the non-const member variables.
   // Exclusive access is required to modify any of these members.

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

@@ -61,7 +61,7 @@ ur_result_t urContextRetain(
 
     /// [in] handle of the context to get a reference of.
     ur_context_handle_t Context) {
-  Context->RefCount.increment();
+  Context->incrementRefCount();
   return UR_RESULT_SUCCESS;
 }
 
@@ -113,7 +113,7 @@ ur_result_t urContextGetInfo(
   case UR_CONTEXT_INFO_NUM_DEVICES:
     return ReturnValue(uint32_t(Context->Devices.size()));
   case UR_CONTEXT_INFO_REFERENCE_COUNT:
-    return ReturnValue(uint32_t{Context->RefCount.load()});
+    return ReturnValue(uint32_t{Context->getRefCount()});
   case UR_CONTEXT_INFO_USM_MEMCPY2D_SUPPORT:
     // 2D USM memcpy is supported.
     return ReturnValue(uint8_t{UseMemcpy2DOperations});
@@ -251,7 +251,7 @@ ur_device_handle_t ur_context_handle_t_::getRootDevice() const {
 // from the list of tracked contexts.
 ur_result_t ContextReleaseHelper(ur_context_handle_t Context) {
 
-  if (!Context->RefCount.decrementAndTest())
+  if (!Context->decrementRefCount() == 0)
     return UR_RESULT_SUCCESS;
 
   if (IndirectAccessTrackingEnabled) {

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

@@ -470,7 +470,7 @@ ur_result_t urDeviceGetInfo(
     return ReturnValue((uint32_t)Device->SubDevices.size());
   }
   case UR_DEVICE_INFO_REFERENCE_COUNT:
-    return ReturnValue(uint32_t{Device->RefCount.load()});
+    return ReturnValue(uint32_t{Device->getRefCount()});
   case UR_DEVICE_INFO_SUPPORTED_PARTITIONS: {
     // SYCL spec says: if this SYCL device cannot be partitioned into at least
     // two sub devices then the returned vector must be empty.
@@ -1666,15 +1666,15 @@ ur_result_t urDeviceGetGlobalTimestamps(
 ur_result_t urDeviceRetain(ur_device_handle_t Device) {
   // The root-device ref-count remains unchanged (always 1).
   if (Device->isSubDevice()) {
-    Device->RefCount.increment();
+    Device->incrementRefCount();
   }
   return UR_RESULT_SUCCESS;
 }
 
 ur_result_t urDeviceRelease(ur_device_handle_t Device) {
   // Root devices are destroyed during the piTearDown process.
   if (Device->isSubDevice()) {
-    if (Device->RefCount.decrementAndTest()) {
+    if (Device->decrementRefCount() == 0) {
       delete Device;
     }
   }

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

@@ -505,7 +505,7 @@ ur_result_t urEventGetInfo(
     return ReturnValue(Result);
   }
   case UR_EVENT_INFO_REFERENCE_COUNT: {
-    return ReturnValue(Event->RefCount.load());
+    return ReturnValue(Event->getRefCount());
   }
   default:
     UR_LOG(ERR, "Unsupported ParamName in urEventGetInfo: ParamName={}(0x{})",
@@ -874,7 +874,7 @@ ur_result_t
 /// [in] handle of the event object
 urEventRetain(/** [in] handle of the event object */ ur_event_handle_t Event) {
   Event->RefCountExternal++;
-  Event->RefCount.increment();
+  Event->incrementRefCount();
 
   return UR_RESULT_SUCCESS;
 }
@@ -1088,7 +1088,7 @@ ur_event_handle_t_::~ur_event_handle_t_() {
 
 ur_result_t urEventReleaseInternal(ur_event_handle_t Event,
                                    bool *isEventDeleted) {
-  if (!Event->RefCount.decrementAndTest())
+  if (!Event->decrementRefCount() == 0)
     return UR_RESULT_SUCCESS;
 
   if (Event->OriginAllocEvent) {
@@ -1429,7 +1429,7 @@ ur_result_t ur_event_handle_t_::reset() {
   CommandType = UR_EXT_COMMAND_TYPE_USER;
   WaitList = {};
   RefCountExternal = 0;
-  RefCount.reset();
+  resetRefCount();
   CommandList = std::nullopt;
   completionBatch = std::nullopt;
   OriginAllocEvent = nullptr;
@@ -1524,7 +1524,7 @@ ur_result_t ur_ze_event_list_t::createAndRetainUrZeEventList(
       std::shared_lock<ur_shared_mutex> Lock(CurQueue->LastCommandEvent->Mutex);
       this->ZeEventList[0] = CurQueue->LastCommandEvent->ZeEvent;
       this->UrEventList[0] = CurQueue->LastCommandEvent;
-      this->UrEventList[0]->RefCount.increment();
+      this->UrEventList[0]->incrementRefCount();
       TmpListLength = 1;
     } else if (EventListLength > 0) {
       this->ZeEventList = new ze_event_handle_t[EventListLength];
@@ -1660,7 +1660,7 @@ ur_result_t ur_ze_event_list_t::createAndRetainUrZeEventList(
               IsInternal, IsMultiDevice));
           MultiDeviceZeEvent = MultiDeviceEvent->ZeEvent;
           const auto &ZeCommandList = CommandList->first;
-          EventList[I]->RefCount.increment();
+          EventList[I]->incrementRefCount();
 
           // Append a Barrier to wait on the original event while signalling the
           // new multi device event.
@@ -1676,11 +1676,11 @@ ur_result_t ur_ze_event_list_t::createAndRetainUrZeEventList(
 
           this->ZeEventList[TmpListLength] = MultiDeviceZeEvent;
           this->UrEventList[TmpListLength] = MultiDeviceEvent;
-          this->UrEventList[TmpListLength]->RefCount.increment();
+          this->UrEventList[TmpListLength]->incrementRefCount();
         } else {
           this->ZeEventList[TmpListLength] = EventList[I]->ZeEvent;
           this->UrEventList[TmpListLength] = EventList[I];
-          this->UrEventList[TmpListLength]->RefCount.increment();
+          this->UrEventList[TmpListLength]->incrementRefCount();
         }
 
         if (QueueLock.has_value()) {

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

@@ -787,7 +787,7 @@ ur_result_t urKernelGetInfo(
     }
   }
   case UR_KERNEL_INFO_REFERENCE_COUNT:
-    return ReturnValue(uint32_t{Kernel->RefCount.load()});
+    return ReturnValue(uint32_t{Kernel->getRefCount()});
   case UR_KERNEL_INFO_ATTRIBUTES:
     try {
       uint32_t Size;
@@ -938,15 +938,15 @@ ur_result_t urKernelGetSubGroupInfo(
 ur_result_t urKernelRetain(
     /// [in] handle for the Kernel to retain
     ur_kernel_handle_t Kernel) {
-  Kernel->RefCount.increment();
+  Kernel->incrementRefCount();
 
   return UR_RESULT_SUCCESS;
 }
 
 ur_result_t urKernelRelease(
     /// [in] handle for the Kernel to release
     ur_kernel_handle_t Kernel) {
-  if (!Kernel->RefCount.decrementAndTest())
+  if (!Kernel->decrementRefCount() == 0)
     return UR_RESULT_SUCCESS;
 
   auto KernelProgram = Kernel->Program;

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

@@ -1052,7 +1052,7 @@ ur_result_t urEnqueueMemBufferMap(
 
     // Add the event to the command list.
     CommandList->second.append(reinterpret_cast<ur_event_handle_t>(*Event));
-    (*Event)->RefCount.increment();
+    (*Event)->incrementRefCount();
 
     const auto &ZeCommandList = CommandList->first;
     const auto &WaitList = (*Event)->WaitList;
@@ -1183,7 +1183,7 @@ ur_result_t urEnqueueMemUnmap(
       nullptr /*ForcedCmdQueue*/));
 
   CommandList->second.append(reinterpret_cast<ur_event_handle_t>(*Event));
-  (*Event)->RefCount.increment();
+  (*Event)->incrementRefCount();
 
   const auto &ZeCommandList = CommandList->first;
 
@@ -1635,14 +1635,14 @@ ur_result_t urMemBufferCreate(
 ur_result_t urMemRetain(
     /// [in] handle of the memory object to get access
     ur_mem_handle_t Mem) {
-  Mem->RefCount.increment();
+  Mem->incrementRefCount();
   return UR_RESULT_SUCCESS;
 }
 
 ur_result_t urMemRelease(
     /// [in] handle of the memory object to release
     ur_mem_handle_t Mem) {
-  if (!Mem->RefCount.decrementAndTest())
+  if (!Mem->decrementRefCount() == 0)
     return UR_RESULT_SUCCESS;
 
   if (Mem->isImage()) {
@@ -1848,7 +1848,7 @@ ur_result_t urMemGetInfo(
     return ReturnValue(size_t{Buffer->Size});
   }
   case UR_MEM_INFO_REFERENCE_COUNT: {
-    return ReturnValue(Buffer->RefCount.load());
+    return ReturnValue(Buffer->getRefCount());
   }
   default: {
     return UR_RESULT_ERROR_INVALID_ENUMERATION;