Skip to content

Commit b3fbf72

Browse files
martygrantmartygrant
committed
Use new ref count class in sanitizer layers. Fix some compile errors.
1 parent 9ce413c commit b3fbf72

File tree

17 files changed

+101
-56
lines changed

17 files changed

+101
-56
lines changed

‎unified-runtime/source/adapters/hip/kernel.hpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -237,7 +237,7 @@ struct ur_kernel_handle_t_ : ur::hip::handle_base {
237237
ur_context_handle_t Ctxt)
238238
: handle_base(), Function{Func},
239239
FunctionWithOffsetParam{FuncWithOffsetParam}, Name{Name}, Context{Ctxt},
240-
Program{Program}, RefCount{1} {
240+
Program{Program} {
241241
urProgramRetain(Program);
242242
urContextRetain(Context);
243243

‎unified-runtime/source/adapters/hip/physical_mem.hpp

Lines changed: 1 addition & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,9 +20,7 @@
2020
/// TODO: Implement.
2121
///
2222
struct ur_physical_mem_handle_t_ : ur::hip::handle_base {
23-
std::atomic_uint32_t RefCount;
24-
25-
ur_physical_mem_handle_t_() : handle_base(), RefCount(1) {}
23+
ur_physical_mem_handle_t_() : handle_base() {}
2624

2725
UR_ReferenceCounter &getRefCounter() noexcept { return RefCounter; }
2826

‎unified-runtime/source/adapters/level_zero/event.hpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -257,11 +257,11 @@ struct ur_event_handle_t_ : ur_object {
257257
private:
258258
UR_ReferenceCounter RefCounter;
259259

260-
// Besides each PI object keeping a total reference count in
261-
// ur_object::RefCount we keep special track of the event *external*
260+
// Besides each UR object keeping a total reference count in
261+
// RefCounter we keep special track of the event *external*
262262
// references. This way we are able to tell when the event is not referenced
263263
// externally anymore, i.e. it can't be passed as a dependency event to
264-
// piEnqueue* functions and explicitly waited meaning that we can do some
264+
// urEnqueue* functions and explicitly waited meaning that we can do some
265265
// optimizations:
266266
// 1. For in-order queues we can reset and reuse event even if it was not yet
267267
// completed by submitting a reset command to the queue (since there are no

‎unified-runtime/source/adapters/level_zero/queue.hpp

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -691,8 +691,8 @@ struct ur_queue_handle_t_ : ur_object {
691691
private:
692692
UR_ReferenceCounter RefCounter;
693693

694-
// Besides each PI object keeping a total reference count in
695-
// ur_object::RefCount we keep special track of the queue *external*
694+
// Besides each UR object keeping a total reference count in
695+
// RefCounter we keep special track of the queue *external*
696696
// references. This way we are able to tell when the queue is being finished
697697
// externally, and can wait for internal references to complete, and do proper
698698
// cleanup of the queue.

‎unified-runtime/source/adapters/offload/common.hpp

Lines changed: 0 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -19,7 +19,3 @@ struct ddi_getter {
1919
};
2020
using handle_base = ur::handle_base<ur::offload::ddi_getter>;
2121
} // namespace ur::offload
22-
23-
struct RefCounted : ur::offload::handle_base {
24-
std::atomic_uint32_t RefCount = 1;
25-
};

‎unified-runtime/source/adapters/offload/event.cpp

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ UR_APIEXPORT ur_result_t UR_APICALL urEventGetInfo(ur_event_handle_t hKernel,
2323

2424
switch (propName) {
2525
case UR_EVENT_INFO_REFERENCE_COUNT:
26-
return ReturnValue(hKernel->RefCount.load());
26+
return ReturnValue(hKernel->getRefCounter().getCount());
2727
default:
2828
return UR_RESULT_ERROR_UNSUPPORTED_ENUMERATION;
2929
}
@@ -51,13 +51,13 @@ urEventWait(uint32_t numEvents, const ur_event_handle_t *phEventWaitList) {
5151
}
5252

5353
UR_APIEXPORT ur_result_t UR_APICALL urEventRetain(ur_event_handle_t hEvent) {
54-
hEvent->RefCount++;
54+
hEvent->getRefCounter().increment();
5555

5656
return UR_RESULT_SUCCESS;
5757
}
5858

5959
UR_APIEXPORT ur_result_t UR_APICALL urEventRelease(ur_event_handle_t hEvent) {
60-
if (--hEvent->RefCount == 0) {
60+
if (hEvent->getRefCounter().decrement() == 0) {
6161
// There's a small bug in olDestroyEvent that will crash. Leak the event
6262
// in the meantime.
6363
// auto Res = olDestroyEvent(hEvent->OffloadEvent);

‎unified-runtime/source/adapters/offload/event.hpp

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,13 @@
1414
#include <ur_api.h>
1515

1616
#include "common.hpp"
17+
#include "common/ur_ref_counter.hpp"
1718

18-
struct ur_event_handle_t_ : RefCounted {
19+
struct ur_event_handle_t_ {
1920
ol_event_handle_t OffloadEvent;
21+
22+
UR_ReferenceCounter &getRefCounter() noexcept { return RefCounter; }
23+
24+
private:
25+
UR_ReferenceCounter RefCounter;
2026
};

‎unified-runtime/source/adapters/opencl/device.hpp

Lines changed: 0 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -20,13 +20,11 @@ struct ur_device_handle_t_ : ur::opencl::handle_base {
2020
ur_platform_handle_t Platform;
2121
cl_device_type Type = 0;
2222
ur_device_handle_t ParentDevice = nullptr;
23-
std::atomic<uint32_t> RefCount = 0;
2423
bool IsNativeHandleOwned = true;
2524

2625
ur_device_handle_t_(native_type Dev, ur_platform_handle_t Plat,
2726
ur_device_handle_t Parent)
2827
: handle_base(), CLDevice(Dev), Platform(Plat), ParentDevice(Parent) {
29-
RefCount = 1;
3028
if (Parent) {
3129
Type = Parent->Type;
3230
[[maybe_unused]] auto Res = clRetainDevice(CLDevice);

‎unified-runtime/source/loader/layers/sanitizer/asan/asan_buffer.hpp

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include <memory>
1818
#include <optional>
1919

20+
#include "common/ur_ref_counter.hpp"
2021
#include "ur/ur.hpp"
2122

2223
namespace ur_sanitizer_layer {
@@ -68,9 +69,12 @@ struct MemBuffer {
6869

6970
std::optional<SubBuffer_t> SubBuffer;
7071

71-
std::atomic<int32_t> RefCount = 1;
72+
UR_ReferenceCounter &getRefCounter() noexcept { return RefCounter; }
7273

7374
ur_shared_mutex Mutex;
75+
76+
private:
77+
UR_ReferenceCounter RefCounter;
7478
};
7579

7680
ur_result_t EnqueueMemCopyRectHelper(

‎unified-runtime/source/loader/layers/sanitizer/asan/asan_ddi.cpp

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -301,7 +301,7 @@ __urdlllocal ur_result_t UR_APICALL urProgramRetain(
301301

302302
auto ProgramInfo = getAsanInterceptor()->getProgramInfo(hProgram);
303303
if (ProgramInfo != nullptr) {
304-
ProgramInfo->RefCount++;
304+
ProgramInfo->getRefCounter().increment();
305305
}
306306

307307
return UR_RESULT_SUCCESS;
@@ -454,7 +454,7 @@ ur_result_t UR_APICALL urProgramRelease(
454454
UR_CALL(pfnProgramRelease(hProgram));
455455

456456
auto ProgramInfo = getAsanInterceptor()->getProgramInfo(hProgram);
457-
if (ProgramInfo != nullptr && --ProgramInfo->RefCount == 0) {
457+
if (ProgramInfo != nullptr && ProgramInfo->getRefCounter().decrement() == 0) {
458458
UR_CALL(getAsanInterceptor()->unregisterProgram(hProgram));
459459
UR_CALL(getAsanInterceptor()->eraseProgram(hProgram));
460460
}
@@ -608,7 +608,7 @@ __urdlllocal ur_result_t UR_APICALL urContextRetain(
608608

609609
auto ContextInfo = getAsanInterceptor()->getContextInfo(hContext);
610610
UR_ASSERT(ContextInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
611-
ContextInfo->RefCount++;
611+
ContextInfo->getRefCounter().increment();
612612

613613
return UR_RESULT_SUCCESS;
614614
}
@@ -630,7 +630,7 @@ __urdlllocal ur_result_t UR_APICALL urContextRelease(
630630

631631
auto ContextInfo = getAsanInterceptor()->getContextInfo(hContext);
632632
UR_ASSERT(ContextInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
633-
if (--ContextInfo->RefCount == 0) {
633+
if (ContextInfo->getRefCounter().decrement() == 0) {
634634
UR_CALL(getAsanInterceptor()->eraseContext(hContext));
635635
}
636636

@@ -750,7 +750,7 @@ __urdlllocal ur_result_t UR_APICALL urMemRetain(
750750
UR_LOG_L(getContext()->logger, DEBUG, "==== urMemRetain");
751751

752752
if (auto MemBuffer = getAsanInterceptor()->getMemBuffer(hMem)) {
753-
MemBuffer->RefCount++;
753+
MemBuffer->getRefCounter().increment();
754754
} else {
755755
UR_CALL(pfnRetain(hMem));
756756
}
@@ -772,7 +772,7 @@ __urdlllocal ur_result_t UR_APICALL urMemRelease(
772772
UR_LOG_L(getContext()->logger, DEBUG, "==== urMemRelease");
773773

774774
if (auto MemBuffer = getAsanInterceptor()->getMemBuffer(hMem)) {
775-
if (--MemBuffer->RefCount != 0) {
775+
if (MemBuffer->getRefCounter().decrement() != 0) {
776776
return UR_RESULT_SUCCESS;
777777
}
778778
UR_CALL(MemBuffer->free());
@@ -1425,7 +1425,7 @@ __urdlllocal ur_result_t UR_APICALL urKernelRetain(
14251425
UR_CALL(pfnRetain(hKernel));
14261426

14271427
auto &KernelInfo = getAsanInterceptor()->getOrCreateKernelInfo(hKernel);
1428-
KernelInfo.RefCount++;
1428+
KernelInfo.getRefCounter().increment();
14291429

14301430
return UR_RESULT_SUCCESS;
14311431
}
@@ -1444,7 +1444,7 @@ __urdlllocal ur_result_t urKernelRelease(
14441444
UR_LOG_L(getContext()->logger, DEBUG, "==== urKernelRelease");
14451445

14461446
auto &KernelInfo = getAsanInterceptor()->getOrCreateKernelInfo(hKernel);
1447-
if (--KernelInfo.RefCount == 0) {
1447+
if (KernelInfo.getRefCounter().decrement() == 0) {
14481448
UR_CALL(getAsanInterceptor()->eraseKernelInfo(hKernel));
14491449
}
14501450
UR_CALL(pfnRelease(hKernel));

‎unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp

Lines changed: 16 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,7 @@
2121
#include "sanitizer_common/sanitizer_common.hpp"
2222
#include "sanitizer_common/sanitizer_options.hpp"
2323
#include "ur_sanitizer_layer.hpp"
24+
#include "common/ur_ref_counter.hpp"
2425

2526
#include <memory>
2627
#include <optional>
@@ -82,7 +83,6 @@ struct QueueInfo {
8283

8384
struct KernelInfo {
8485
ur_kernel_handle_t Handle;
85-
std::atomic<int32_t> RefCount = 1;
8686

8787
// sanitized kernel
8888
bool IsInstrumented = false;
@@ -107,11 +107,15 @@ struct KernelInfo {
107107
getContext()->urDdiTable.Kernel.pfnRelease(Handle);
108108
assert(Result == UR_RESULT_SUCCESS);
109109
}
110+
111+
UR_ReferenceCounter &getRefCounter() noexcept { return RefCounter; }
112+
113+
private:
114+
UR_ReferenceCounter RefCounter;
110115
};
111116

112117
struct ProgramInfo {
113118
ur_program_handle_t Handle;
114-
std::atomic<int32_t> RefCount = 1;
115119

116120
// Program is built only once, so we don't need to lock it
117121
std::unordered_set<std::shared_ptr<AllocInfo>> AllocInfoForGlobals;
@@ -130,6 +134,11 @@ struct ProgramInfo {
130134
}
131135

132136
bool isKernelInstrumented(ur_kernel_handle_t Kernel) const;
137+
138+
UR_ReferenceCounter &getRefCounter() noexcept { return RefCounter; }
139+
140+
private:
141+
UR_ReferenceCounter RefCounter;
133142
};
134143

135144
struct ContextInfo {
@@ -138,8 +147,6 @@ struct ContextInfo {
138147
ur_usm_pool_handle_t USMPool{};
139148
std::once_flag PoolInit;
140149

141-
std::atomic<int32_t> RefCount = 1;
142-
143150
std::vector<ur_device_handle_t> DeviceList;
144151
std::unordered_map<ur_device_handle_t, AllocInfoList> AllocInfosMap;
145152

@@ -163,6 +170,11 @@ struct ContextInfo {
163170
}
164171

165172
ur_usm_pool_handle_t getUSMPool();
173+
174+
UR_ReferenceCounter &getRefCounter() noexcept { return RefCounter; }
175+
176+
private:
177+
UR_ReferenceCounter RefCounter;
166178
};
167179

168180
struct AsanRuntimeDataWrapper {

‎unified-runtime/source/loader/layers/sanitizer/msan/msan_buffer.hpp

Lines changed: 6 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@
1717
#include <memory>
1818
#include <optional>
1919

20+
#include "common/ur_ref_counter.hpp"
2021
#include "ur/ur.hpp"
2122

2223
namespace ur_sanitizer_layer {
@@ -68,9 +69,12 @@ struct MemBuffer {
6869

6970
std::optional<SubBuffer_t> SubBuffer;
7071

71-
std::atomic<int32_t> RefCount = 1;
72-
7372
ur_shared_mutex Mutex;
73+
74+
UR_ReferenceCounter &getRefCounter() noexcept { return RefCounter; }
75+
76+
private:
77+
UR_ReferenceCounter RefCounter;
7478
};
7579

7680
ur_result_t EnqueueMemCopyRectHelper(

‎unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp

Lines changed: 9 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@
1212
*/
1313

1414
#include "msan_ddi.hpp"
15+
#include "common/ur_ref_counter.hpp"
1516
#include "msan_interceptor.hpp"
1617
#include "sanitizer_common/sanitizer_utils.hpp"
1718
#include "ur_sanitizer_layer.hpp"
@@ -248,7 +249,7 @@ urProgramRetain(ur_program_handle_t
248249

249250
auto ProgramInfo = getMsanInterceptor()->getProgramInfo(hProgram);
250251
UR_ASSERT(ProgramInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
251-
ProgramInfo->RefCount++;
252+
ProgramInfo->getRefCounter().increment();
252253

253254
return UR_RESULT_SUCCESS;
254255
}
@@ -381,7 +382,7 @@ ur_result_t urProgramRelease(
381382

382383
auto ProgramInfo = getMsanInterceptor()->getProgramInfo(hProgram);
383384
UR_ASSERT(ProgramInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
384-
if (--ProgramInfo->RefCount == 0) {
385+
if (ProgramInfo->getRefCounter().decrement() == 0) {
385386
UR_CALL(getMsanInterceptor()->unregisterProgram(hProgram));
386387
UR_CALL(getMsanInterceptor()->eraseProgram(hProgram));
387388
}
@@ -512,7 +513,7 @@ ur_result_t urContextRetain(
512513

513514
auto ContextInfo = getMsanInterceptor()->getContextInfo(hContext);
514515
UR_ASSERT(ContextInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
515-
ContextInfo->RefCount++;
516+
ContextInfo->getRefCounter().increment();
516517

517518
return UR_RESULT_SUCCESS;
518519
}
@@ -530,7 +531,7 @@ ur_result_t urContextRelease(
530531

531532
auto ContextInfo = getMsanInterceptor()->getContextInfo(hContext);
532533
UR_ASSERT(ContextInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE);
533-
if (--ContextInfo->RefCount == 0) {
534+
if (ContextInfo->getRefCounter().decrement() == 0) {
534535
UR_CALL(getMsanInterceptor()->eraseContext(hContext));
535536
}
536537

@@ -642,7 +643,7 @@ ur_result_t urMemRetain(
642643
UR_LOG_L(getContext()->logger, DEBUG, "==== urMemRetain");
643644

644645
if (auto MemBuffer = getMsanInterceptor()->getMemBuffer(hMem)) {
645-
MemBuffer->RefCount++;
646+
MemBuffer->getRefCounter().increment();
646647
} else {
647648
UR_CALL(pfnRetain(hMem));
648649
}
@@ -660,7 +661,7 @@ ur_result_t urMemRelease(
660661
UR_LOG_L(getContext()->logger, DEBUG, "==== urMemRelease");
661662

662663
if (auto MemBuffer = getMsanInterceptor()->getMemBuffer(hMem)) {
663-
if (--MemBuffer->RefCount != 0) {
664+
if (MemBuffer->getRefCounter().decrement() != 0) {
664665
return UR_RESULT_SUCCESS;
665666
}
666667
UR_CALL(MemBuffer->free());
@@ -1331,7 +1332,7 @@ ur_result_t urKernelRetain(
13311332
UR_CALL(pfnRetain(hKernel));
13321333

13331334
auto &KernelInfo = getMsanInterceptor()->getOrCreateKernelInfo(hKernel);
1334-
KernelInfo.RefCount++;
1335+
KernelInfo.getRefCounter().increment();
13351336

13361337
return UR_RESULT_SUCCESS;
13371338
}
@@ -1346,7 +1347,7 @@ ur_result_t urKernelRelease(
13461347
UR_LOG_L(getContext()->logger, DEBUG, "==== urKernelRelease");
13471348

13481349
auto &KernelInfo = getMsanInterceptor()->getOrCreateKernelInfo(hKernel);
1349-
if (--KernelInfo.RefCount == 0) {
1350+
if (KernelInfo.getRefCounter().decrement() == 0) {
13501351
UR_CALL(getMsanInterceptor()->eraseKernelInfo(hKernel));
13511352
}
13521353
UR_CALL(pfnRelease(hKernel));

0 commit comments

Comments
 (0)