Change implementation to a new URRefCount class and add it as member to any relevant handles that need ref counting.

Author: martygrant

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

@@ -301,7 +301,7 @@ ur_adapter_handle_t_::ur_adapter_handle_t_()
   ZeInitResult = ZE_RESULT_ERROR_UNINITIALIZED;
   ZesResult = ZE_RESULT_ERROR_UNINITIALIZED;
 
-  resetRefCount(0);
+  RefCount.reset(0);
 
 #ifdef UR_STATIC_LEVEL_ZERO
   // Given static linking of the L0 Loader, we must delay the loader's
@@ -677,7 +677,7 @@ ur_result_t urAdapterGet(
     }
     *Adapters = GlobalAdapter;
 
-    if (GlobalAdapter->incrementRefCount() == 0) {
+    if (GlobalAdapter->getRefCount().increment() == 0) {
       adapterStateInit();
     }
   }
@@ -694,7 +694,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->decrementRefCount() == 0) {
+  if (GlobalAdapter->getRefCount().decrementAndTest()) {
     auto result = adapterStateTeardown();
 #ifdef UR_STATIC_LEVEL_ZERO
     // Given static linking of the L0 Loader, we must delay the loader's
@@ -711,9 +711,9 @@ ur_result_t urAdapterRelease([[maybe_unused]] ur_adapter_handle_t Adapter) {
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t urAdapterRetain(ur_adapter_handle_t) {
+ur_result_t urAdapterRetain([[maybe_unused]] ur_adapter_handle_t Adapter) {
   assert(GlobalAdapter && GlobalAdapter == Adapter);
-  GlobalAdapter->incrementRefCount();
+  GlobalAdapter->getRefCount().increment();
 
   return UR_RESULT_SUCCESS;
 }
@@ -742,12 +742,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->getRefCount());
+    return ReturnValue(GlobalAdapter->getRefCount().getCount());
   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

@@ -9,6 +9,7 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include "common/ur_ref_count.hpp"
 #include "logger/ur_logger.hpp"
 #include "ur_interface_loader.hpp"
 #include <loader/ur_loader.hpp>
@@ -43,6 +44,11 @@ struct ur_adapter_handle_t_ : ur::handle_base<ur::level_zero::ddi_getter> {
   ZeCache<Result<PlatformVec>> PlatformCache;
   logger::Logger &logger;
   HMODULE processHandle = nullptr;
+
+  URRefCount &getRefCount() noexcept { return RefCount; }
+
+private:
+  URRefCount RefCount;
 };
 
 extern ur_adapter_handle_t_ *GlobalAdapter;

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

@@ -247,7 +247,7 @@ ur_result_t urEnqueueUSMFreeExp(
   }
 
   size_t size = umfPoolMallocUsableSize(hPool, Mem);
-  (*Event)->incrementRefCount();
+  (*Event)->getRefCount().increment();
   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->incrementRefCount();
+  CommandBuffer->getRefCount().increment();
   return UR_RESULT_SUCCESS;
 }
 
 ur_result_t
 urCommandBufferReleaseExp(ur_exp_command_buffer_handle_t CommandBuffer) {
-  if (!CommandBuffer->decrementAndTest())
+  if (!CommandBuffer->getRefCount().decrementAndTest())
     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)->incrementRefCount();
+  (*Event)->getRefCount().increment();
   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)->incrementRefCount();
+  (*Event)->getRefCount().increment();
   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->getRefCount()});
+    return ReturnValue(uint32_t{hCommandBuffer->getRefCount().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/level_zero/command_buffer.hpp

@@ -17,6 +17,7 @@
 
 #include "common.hpp"
 
+#include "common/ur_ref_count.hpp"
 #include "context.hpp"
 #include "kernel.hpp"
 #include "queue.hpp"
@@ -149,4 +150,9 @@ struct ur_exp_command_buffer_handle_t_ : public ur_object {
   // Track handle objects to free when command-buffer is destroyed.
   std::vector<std::unique_ptr<ur_exp_command_buffer_command_handle_t_>>
       CommandHandles;
+
+  URRefCount &getRefCount() noexcept { return RefCount; }
+
+private:
+  URRefCount RefCount;
 };

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

@@ -34,6 +34,7 @@
 #include <level_zero/ze_intel_gpu.h>
 #include <umf_pools/disjoint_pool_config_parser.hpp>
 
+#include "common/ur_ref_count.hpp"
 #include "logger/ur_logger.hpp"
 #include "ur_interface_loader.hpp"
 
@@ -257,6 +258,11 @@ struct MemAllocRecord : ur_object {
   // TODO: this should go away when memory isolation issue is fixed in the Level
   // Zero runtime.
   ur_context_handle_t Context;
+
+  URRefCount &getRefCount() noexcept { return RefCount; }
+
+private:
+  URRefCount RefCount;
 };
 
 extern usm::DisjointPoolAllConfigs DisjointPoolConfigInstance;

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->incrementRefCount();
+  Context->getRefCount().increment();
   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->getRefCount()});
+    return ReturnValue(uint32_t{Context->getRefCount().getCount()});
   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->decrementAndTest())
+  if (!Context->getRefCount().decrementAndTest())
     return UR_RESULT_SUCCESS;
 
   if (IndirectAccessTrackingEnabled) {

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

@@ -26,6 +26,7 @@
 #include "queue.hpp"
 #include "usm.hpp"
 
+#include "common/ur_ref_count.hpp"
 #include <umf_helpers.hpp>
 
 struct l0_command_list_cache_info {
@@ -358,6 +359,8 @@ struct ur_context_handle_t_ : ur_object {
   // Get handle to the L0 context
   ze_context_handle_t getZeHandle() const;
 
+  URRefCount &getRefCount() noexcept { return RefCount; }
+
 private:
   enum EventFlags {
     EVENT_FLAG_HOST_VISIBLE = UR_BIT(0),
@@ -404,6 +407,8 @@ struct ur_context_handle_t_ : ur_object {
 
     return &EventCaches[index];
   }
+
+  URRefCount RefCount;
 };
 
 // Helper function to release the context, a caller must lock the platform-level

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->getRefCount()});
+    return ReturnValue(uint32_t{Device->getRefCount().getCount()});
   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->incrementRefCount();
+    Device->getRefCount().increment();
   }
   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->decrementAndTest()) {
+    if (Device->getRefCount().decrementAndTest()) {
       delete Device;
     }
   }

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

@@ -20,6 +20,7 @@
 
 #include "adapters/level_zero/platform.hpp"
 #include "common.hpp"
+#include "common/ur_ref_count.hpp"
 #include <ur/ur.hpp>
 #include <ur_ddi.h>
 #include <ze_api.h>
@@ -242,6 +243,11 @@ struct ur_device_handle_t_ : ur_object {
 
   // unique ephemeral identifer of the device in the adapter
   std::optional<DeviceId> Id;
+
+  URRefCount &getRefCount() noexcept { return RefCount; }
+
+private:
+  URRefCount RefCount;
 };
 
 inline std::vector<ur_device_handle_t>