Skip to content

Commit 22ed65f

Browse files
AllanZynekbenzie
AllanZyne
authored and
kbenzie
committed
Fix multiple contexts in memory overhead statistics (#17897)
- Create the internal context in "ShadowMemory" constructor - Remove "static" in "ShadowMemory::Setup" and "ShadowMemory::Destroy", because `AsanInterceptor::getOrCreateShadowMemory` is single thread
1 parent 4daa3cb commit 22ed65f

File tree

3 files changed

+87
-110
lines changed

3 files changed

+87
-110
lines changed

‎source/loader/layers/sanitizer/asan/asan_interceptor.cpp‎

Lines changed: 1 addition & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -49,7 +49,6 @@ AsanInterceptor::~AsanInterceptor() {
4949

5050
for (auto &[_, ShadowMemory] : m_ShadowMap) {
5151
ShadowMemory->Destory();
52-
getContext()->urDdiTable.Context.pfnRelease(ShadowMemory->Context);
5352
}
5453

5554
for (auto Adapter : m_Adapters) {
@@ -323,16 +322,7 @@ AsanInterceptor::getOrCreateShadowMemory(ur_device_handle_t Device,
323322
DeviceType Type) {
324323
std::scoped_lock<ur_shared_mutex> Guard(m_ShadowMapMutex);
325324
if (m_ShadowMap.find(Type) == m_ShadowMap.end()) {
326-
ur_context_handle_t InternalContext;
327-
auto Res = getContext()->urDdiTable.Context.pfnCreate(1, &Device, nullptr,
328-
&InternalContext);
329-
if (Res != UR_RESULT_SUCCESS) {
330-
getContext()->logger.error("Failed to create shadow context");
331-
return nullptr;
332-
}
333-
std::shared_ptr<ContextInfo> CI;
334-
insertContext(InternalContext, CI);
335-
m_ShadowMap[Type] = CreateShadowMemory(InternalContext, Device, Type);
325+
m_ShadowMap[Type] = CreateShadowMemory(Device, Type);
336326
m_ShadowMap[Type]->Setup();
337327
}
338328
return m_ShadowMap[Type];

‎source/loader/layers/sanitizer/asan/asan_shadow.cpp‎

Lines changed: 70 additions & 83 deletions
Original file line numberDiff line numberDiff line change
@@ -20,44 +20,39 @@
2020
namespace ur_sanitizer_layer {
2121
namespace asan {
2222

23-
std::shared_ptr<ShadowMemory> CreateShadowMemory(ur_context_handle_t Context,
24-
ur_device_handle_t Device,
23+
std::shared_ptr<ShadowMemory> CreateShadowMemory(ur_device_handle_t Device,
2524
DeviceType Type) {
26-
if (Type == DeviceType::CPU) {
27-
return std::make_shared<ShadowMemoryCPU>(Context, Device);
28-
} else if (Type == DeviceType::GPU_PVC) {
29-
return std::make_shared<ShadowMemoryPVC>(Context, Device);
30-
} else if (Type == DeviceType::GPU_DG2) {
31-
return std::make_shared<ShadowMemoryDG2>(Context, Device);
32-
} else {
33-
getContext()->logger.error("Unsupport device type");
34-
return nullptr;
25+
switch (Type) {
26+
case DeviceType::CPU:
27+
return std::make_shared<ShadowMemoryCPU>(Device);
28+
case DeviceType::GPU_PVC:
29+
return std::make_shared<ShadowMemoryPVC>(Device);
30+
case DeviceType::GPU_DG2:
31+
return std::make_shared<ShadowMemoryDG2>(Device);
32+
default:
33+
die("CreateShadowMemory: Unsupport device type");
3534
}
3635
}
3736

3837
ur_result_t ShadowMemoryCPU::Setup() {
39-
static ur_result_t Result = [this]() {
40-
size_t ShadowSize = GetShadowSize();
41-
ShadowBegin = MmapNoReserve(0, ShadowSize);
42-
if (ShadowBegin == 0) {
43-
return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
44-
}
45-
DontCoredumpRange(ShadowBegin, ShadowSize);
46-
ShadowEnd = ShadowBegin + ShadowSize;
47-
48-
// Set shadow memory for null pointer
49-
// For CPU, we use a typical page size of 4K bytes.
50-
constexpr size_t NullptrRedzoneSize = 4096;
51-
auto URes = EnqueuePoisonShadow({}, 0, NullptrRedzoneSize,
52-
kNullPointerRedzoneMagic);
53-
if (URes != UR_RESULT_SUCCESS) {
54-
getContext()->logger.error("EnqueuePoisonShadow(NullPointerRZ): {}",
55-
URes);
56-
return URes;
57-
}
38+
size_t ShadowSize = GetShadowSize();
39+
ShadowBegin = MmapNoReserve(0, ShadowSize);
40+
if (ShadowBegin == 0) {
41+
return UR_RESULT_ERROR_OUT_OF_HOST_MEMORY;
42+
}
43+
DontCoredumpRange(ShadowBegin, ShadowSize);
44+
ShadowEnd = ShadowBegin + ShadowSize;
45+
46+
// Set shadow memory for null pointer
47+
// For CPU, we use a typical page size of 4K bytes.
48+
constexpr size_t NullptrRedzoneSize = 4096;
49+
auto URes =
50+
EnqueuePoisonShadow({}, 0, NullptrRedzoneSize, kNullPointerRedzoneMagic);
51+
if (URes != UR_RESULT_SUCCESS) {
52+
getContext()->logger.error("EnqueuePoisonShadow(NullPointerRZ): {}", URes);
5853
return URes;
59-
}();
60-
return Result;
54+
}
55+
return URes;
6156
}
6257

6358
ur_result_t ShadowMemoryCPU::Destory() {
@@ -99,39 +94,33 @@ ur_result_t ShadowMemoryGPU::Setup() {
9994
// we reserve shadow memory for each contexts, this will cause out-of-resource
10095
// error when user uses multiple contexts. Therefore, we just create one
10196
// shadow memory here.
102-
static ur_result_t Result = [this]() {
103-
const size_t ShadowSize = GetShadowSize();
104-
// To reserve very large amount of GPU virtual memroy, the pStart param
105-
// should be beyond the SVM range, so that GFX driver will automatically
106-
// switch to reservation on the GPU heap.
107-
const void *StartAddress = (void *)(0x100'0000'0000'0000ULL);
108-
// TODO: Protect Bad Zone
109-
auto Result = getContext()->urDdiTable.VirtualMem.pfnReserve(
110-
Context, StartAddress, ShadowSize, (void **)&ShadowBegin);
111-
if (Result != UR_RESULT_SUCCESS) {
112-
getContext()->logger.error(
113-
"Shadow memory reserved failed with size {}: {}", (void *)ShadowSize,
114-
Result);
115-
return Result;
116-
}
117-
ShadowEnd = ShadowBegin + ShadowSize;
118-
// Retain the context which reserves shadow memory
119-
getContext()->urDdiTable.Context.pfnRetain(Context);
120-
121-
// Set shadow memory for null pointer
122-
// For GPU, wu use up to 1 page of shadow memory
123-
const size_t NullptrRedzoneSize = GetVirtualMemGranularity(Context, Device)
124-
<< ASAN_SHADOW_SCALE;
125-
ManagedQueue Queue(Context, Device);
126-
Result = EnqueuePoisonShadow(Queue, 0, NullptrRedzoneSize,
127-
kNullPointerRedzoneMagic);
128-
if (Result != UR_RESULT_SUCCESS) {
129-
getContext()->logger.error("EnqueuePoisonShadow(NullPointerRZ): {}",
130-
Result);
131-
return Result;
132-
}
97+
const size_t ShadowSize = GetShadowSize();
98+
// To reserve very large amount of GPU virtual memroy, the pStart param
99+
// should be beyond the SVM range, so that GFX driver will automatically
100+
// switch to reservation on the GPU heap.
101+
const void *StartAddress = (void *)(0x100'0000'0000'0000ULL);
102+
// TODO: Protect Bad Zone
103+
auto Result = getContext()->urDdiTable.VirtualMem.pfnReserve(
104+
Context, StartAddress, ShadowSize, (void **)&ShadowBegin);
105+
if (Result != UR_RESULT_SUCCESS) {
106+
getContext()->logger.error("Shadow memory reserved failed with size {}: {}",
107+
(void *)ShadowSize, Result);
133108
return Result;
134-
}();
109+
}
110+
ShadowEnd = ShadowBegin + ShadowSize;
111+
112+
// Set shadow memory for null pointer
113+
// For GPU, wu use up to 1 page of shadow memory
114+
const size_t NullptrRedzoneSize = GetVirtualMemGranularity(Context, Device)
115+
<< ASAN_SHADOW_SCALE;
116+
ManagedQueue Queue(Context, Device);
117+
Result = EnqueuePoisonShadow(Queue, 0, NullptrRedzoneSize,
118+
kNullPointerRedzoneMagic);
119+
if (Result != UR_RESULT_SUCCESS) {
120+
getContext()->logger.error("EnqueuePoisonShadow(NullPointerRZ): {}",
121+
Result);
122+
return Result;
123+
}
135124
return Result;
136125
}
137126

@@ -142,7 +131,13 @@ ur_result_t ShadowMemoryGPU::Destory() {
142131
PrivateShadowOffset = 0;
143132
}
144133

145-
static ur_result_t Result = [this]() {
134+
if (LocalShadowOffset != 0) {
135+
UR_CALL(getContext()->urDdiTable.USM.pfnFree(Context,
136+
(void *)LocalShadowOffset));
137+
LocalShadowOffset = 0;
138+
}
139+
140+
{
146141
const size_t PageSize = GetVirtualMemGranularity(Context, Device);
147142
for (auto [MappedPtr, PhysicalMem] : VirtualMemMaps) {
148143
UR_CALL(getContext()->urDdiTable.VirtualMem.pfnUnmap(
@@ -151,24 +146,14 @@ ur_result_t ShadowMemoryGPU::Destory() {
151146
}
152147
UR_CALL(getContext()->urDdiTable.VirtualMem.pfnFree(
153148
Context, (const void *)ShadowBegin, GetShadowSize()));
154-
UR_CALL(getContext()->urDdiTable.Context.pfnRelease(Context));
155-
return UR_RESULT_SUCCESS;
156-
}();
157-
if (!Result) {
158-
return Result;
159-
}
160149

161-
if (LocalShadowOffset != 0) {
162-
UR_CALL(getContext()->urDdiTable.USM.pfnFree(Context,
163-
(void *)LocalShadowOffset));
164-
LocalShadowOffset = 0;
165-
}
166-
if (ShadowBegin != 0) {
167-
UR_CALL(getContext()->urDdiTable.VirtualMem.pfnFree(
168-
Context, (const void *)ShadowBegin, GetShadowSize()));
169-
UR_CALL(getContext()->urDdiTable.Context.pfnRelease(Context));
170-
ShadowBegin = ShadowEnd = 0;
150+
if (ShadowBegin != 0) {
151+
UR_CALL(getContext()->urDdiTable.VirtualMem.pfnFree(
152+
Context, (const void *)ShadowBegin, GetShadowSize()));
153+
ShadowBegin = ShadowEnd = 0;
154+
}
171155
}
156+
172157
return UR_RESULT_SUCCESS;
173158
}
174159

@@ -248,7 +233,8 @@ ur_result_t ShadowMemoryGPU::AllocLocalShadow(ur_queue_handle_t Queue,
248233
(NumWG * LocalMemorySize) >> ASAN_SHADOW_SCALE;
249234
static size_t LastAllocedSize = 0;
250235
if (RequiredShadowSize > LastAllocedSize) {
251-
auto ContextInfo = getAsanInterceptor()->getContextInfo(Context);
236+
ur_context_handle_t QueueContext = GetContext(Queue);
237+
auto ContextInfo = getAsanInterceptor()->getContextInfo(QueueContext);
252238
if (LocalShadowOffset) {
253239
UR_CALL(getContext()->urDdiTable.USM.pfnFree(Context,
254240
(void *)LocalShadowOffset));
@@ -288,7 +274,8 @@ ur_result_t ShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue,
288274
(NumWG * ASAN_PRIVATE_SIZE) >> ASAN_SHADOW_SCALE;
289275
static size_t LastAllocedSize = 0;
290276
if (RequiredShadowSize > LastAllocedSize) {
291-
auto ContextInfo = getAsanInterceptor()->getContextInfo(Context);
277+
ur_context_handle_t QueueContext = GetContext(Queue);
278+
auto ContextInfo = getAsanInterceptor()->getContextInfo(QueueContext);
292279
if (PrivateShadowOffset) {
293280
UR_CALL(getContext()->urDdiTable.USM.pfnFree(
294281
Context, (void *)PrivateShadowOffset));

‎source/loader/layers/sanitizer/asan/asan_shadow.hpp‎

Lines changed: 16 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -13,26 +13,31 @@
1313

1414
#pragma once
1515

16-
#include "asan_allocator.hpp"
16+
#include "sanitizer_common/sanitizer_common.hpp"
1717
#include "sanitizer_common/sanitizer_libdevice.hpp"
1818
#include "ur_sanitizer_layer.hpp"
1919

20-
#include <unordered_set>
21-
2220
namespace ur_sanitizer_layer {
2321
namespace asan {
2422

2523
struct ShadowMemory {
26-
ShadowMemory(ur_context_handle_t Context, ur_device_handle_t Device)
27-
: Context(Context), Device(Device) {
24+
ShadowMemory(ur_device_handle_t Device) : Device(Device) {
2825
[[maybe_unused]] ur_result_t URes =
2926
getContext()->urDdiTable.Device.pfnRetain(Device);
3027
assert(URes == UR_RESULT_SUCCESS);
28+
29+
// Create the internal context used for managing the shadow memory
30+
URes = getContext()->urDdiTable.Context.pfnCreate(1, &Device, nullptr,
31+
&Context);
32+
assert(URes == UR_RESULT_SUCCESS);
3133
}
3234

3335
virtual ~ShadowMemory() {
3436
[[maybe_unused]] ur_result_t URes =
35-
getContext()->urDdiTable.Device.pfnRelease(Device);
37+
getContext()->urDdiTable.Context.pfnRelease(Context);
38+
assert(URes == UR_RESULT_SUCCESS);
39+
40+
URes = getContext()->urDdiTable.Device.pfnRelease(Device);
3641
assert(URes == UR_RESULT_SUCCESS);
3742
}
3843

@@ -64,8 +69,7 @@ struct ShadowMemory {
6469
};
6570

6671
struct ShadowMemoryCPU final : public ShadowMemory {
67-
ShadowMemoryCPU(ur_context_handle_t Context, ur_device_handle_t Device)
68-
: ShadowMemory(Context, Device) {}
72+
ShadowMemoryCPU(ur_device_handle_t Device) : ShadowMemory(Device) {}
6973

7074
ur_result_t Setup() override;
7175

@@ -94,8 +98,7 @@ struct ShadowMemoryCPU final : public ShadowMemory {
9498
};
9599

96100
struct ShadowMemoryGPU : public ShadowMemory {
97-
ShadowMemoryGPU(ur_context_handle_t Context, ur_device_handle_t Device)
98-
: ShadowMemory(Context, Device) {}
101+
ShadowMemoryGPU(ur_device_handle_t Device) : ShadowMemory(Device) {}
99102

100103
ur_result_t Setup() override;
101104

@@ -136,8 +139,7 @@ struct ShadowMemoryGPU : public ShadowMemory {
136139
/// Device USM : 0x0800_0000_0000 ~ 0x17ff_ffff_ffff
137140
///
138141
struct ShadowMemoryPVC final : public ShadowMemoryGPU {
139-
ShadowMemoryPVC(ur_context_handle_t Context, ur_device_handle_t Device)
140-
: ShadowMemoryGPU(Context, Device) {}
142+
ShadowMemoryPVC(ur_device_handle_t Device) : ShadowMemoryGPU(Device) {}
141143

142144
uptr MemToShadow(uptr Ptr) override;
143145

@@ -155,16 +157,14 @@ struct ShadowMemoryPVC final : public ShadowMemoryGPU {
155157
/// Device USM : 0x0800_0000_0000 ~ 0x0fff_ffff_ffff
156158
///
157159
struct ShadowMemoryDG2 final : public ShadowMemoryGPU {
158-
ShadowMemoryDG2(ur_context_handle_t Context, ur_device_handle_t Device)
159-
: ShadowMemoryGPU(Context, Device) {}
160+
ShadowMemoryDG2(ur_device_handle_t Device) : ShadowMemoryGPU(Device) {}
160161

161162
uptr MemToShadow(uptr Ptr) override;
162163

163164
size_t GetShadowSize() override { return 0x100000000000ULL; }
164165
};
165166

166-
std::shared_ptr<ShadowMemory> CreateShadowMemory(ur_context_handle_t Context,
167-
ur_device_handle_t Device,
167+
std::shared_ptr<ShadowMemory> CreateShadowMemory(ur_device_handle_t Device,
168168
DeviceType Type);
169169

170170
} // namespace asan

0 commit comments

Comments
 (0)