Skip to content

Commit b8f20f0

Browse files
yingcong-wuCopilotkbenzie
authored
[DeviceSanitizer] Fix private base leak and reuse private base if possible (#19066)
Bug1: Leak, when private shadow failed to allocate, the already allocated private base would not be freed. Bug2: Leak, the old private base is never freed. Improve: try to reuse private base just like we try to reuse private shadow. --------- Co-authored-by: Copilot <[email protected]> Co-authored-by: Kenneth Benzie (Benie) <[email protected]>
1 parent a85b4ce commit b8f20f0

File tree

4 files changed

+133
-57
lines changed

4 files changed

+133
-57
lines changed

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

Lines changed: 66 additions & 29 deletions
Original file line numberDiff line numberDiff line change
@@ -134,6 +134,12 @@ ur_result_t ShadowMemoryGPU::Destory() {
134134
PrivateShadowOffset = 0;
135135
}
136136

137+
if (PrivateBasePtr != 0) {
138+
UR_CALL(
139+
getContext()->urDdiTable.USM.pfnFree(Context, (void *)PrivateBasePtr));
140+
PrivateBasePtr = 0;
141+
}
142+
137143
if (LocalShadowOffset != 0) {
138144
UR_CALL(getContext()->urDdiTable.USM.pfnFree(Context,
139145
(void *)LocalShadowOffset));
@@ -282,51 +288,82 @@ ur_result_t ShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue,
282288
uint64_t NumWI, uint32_t NumWG,
283289
uptr *&Base, uptr &Begin,
284290
uptr &End) {
285-
{
286-
const size_t Size = NumWI * sizeof(uptr);
287-
ur_usm_desc_t Properties{UR_STRUCTURE_TYPE_USM_DESC, nullptr,
288-
UR_USM_ADVICE_FLAG_DEFAULT, sizeof(uptr)};
289-
UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc(
290-
Context, Device, &Properties, nullptr, Size, (void **)&Base));
291-
}
291+
// Trying to allocate private base array and private shadow, and any one of
292+
// them fail to allocate would be a failure
293+
static size_t LastPrivateBaseAllocedSize = 0;
294+
static size_t LastPrivateShadowAllocedSize = 0;
295+
296+
try {
297+
const size_t NewPrivateBaseSize = NumWI * sizeof(uptr);
298+
if (NewPrivateBaseSize > LastPrivateBaseAllocedSize) {
299+
if (PrivateBasePtr) {
300+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnFree(
301+
Context, (void *)PrivateBasePtr));
302+
PrivateBasePtr = 0;
303+
LastPrivateBaseAllocedSize = 0;
304+
}
292305

293-
{
294-
const size_t RequiredShadowSize =
306+
ur_usm_desc_t PrivateBaseProps{UR_STRUCTURE_TYPE_USM_DESC, nullptr,
307+
UR_USM_ADVICE_FLAG_DEFAULT, sizeof(uptr)};
308+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc(
309+
Context, Device, &PrivateBaseProps, nullptr, NewPrivateBaseSize,
310+
(void **)&PrivateBasePtr));
311+
312+
// No need to clean the shadow base, their should be set by work item on
313+
// launch
314+
315+
// FIXME: we should add private base to statistic
316+
LastPrivateBaseAllocedSize = NewPrivateBaseSize;
317+
}
318+
319+
const size_t NewPrivateShadowSize =
295320
(NumWG * ASAN_PRIVATE_SIZE) >> ASAN_SHADOW_SCALE;
296-
static size_t LastAllocedSize = 0;
297-
if (RequiredShadowSize > LastAllocedSize) {
321+
if (NewPrivateShadowSize > LastPrivateShadowAllocedSize) {
298322
ur_context_handle_t QueueContext = GetContext(Queue);
299323
auto ContextInfo = getAsanInterceptor()->getContextInfo(QueueContext);
324+
300325
if (PrivateShadowOffset) {
301-
UR_CALL(getContext()->urDdiTable.USM.pfnFree(
326+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnFree(
302327
Context, (void *)PrivateShadowOffset));
303-
ContextInfo->Stats.UpdateShadowFreed(LastAllocedSize);
328+
ContextInfo->Stats.UpdateShadowFreed(LastPrivateShadowAllocedSize);
304329
PrivateShadowOffset = 0;
305-
LastAllocedSize = 0;
330+
LastPrivateShadowAllocedSize = 0;
306331
}
307332

308-
UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc(
309-
Context, Device, nullptr, nullptr, RequiredShadowSize,
333+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc(
334+
Context, Device, nullptr, nullptr, NewPrivateShadowSize,
310335
(void **)&PrivateShadowOffset));
336+
LastPrivateShadowAllocedSize = NewPrivateShadowSize;
337+
UR_CALL_THROWS(EnqueueUSMSet(Queue, (void *)PrivateShadowOffset, (char)0,
338+
NewPrivateShadowSize));
339+
ContextInfo->Stats.UpdateShadowMalloced(NewPrivateShadowSize);
340+
}
311341

312-
// Initialize shadow memory
313-
ur_result_t URes = EnqueueUSMSet(Queue, (void *)PrivateShadowOffset,
314-
(char)0, RequiredShadowSize);
315-
if (URes != UR_RESULT_SUCCESS) {
316-
UR_CALL(getContext()->urDdiTable.USM.pfnFree(
317-
Context, (void *)PrivateShadowOffset));
318-
PrivateShadowOffset = 0;
319-
LastAllocedSize = 0;
320-
}
342+
Base = (uptr *)PrivateBasePtr;
343+
Begin = PrivateShadowOffset;
344+
End = PrivateShadowOffset + NewPrivateShadowSize - 1;
321345

322-
ContextInfo->Stats.UpdateShadowMalloced(RequiredShadowSize);
346+
} catch (ur_result_t &UrRes) {
347+
assert(UrRes != UR_RESULT_SUCCESS);
323348

324-
LastAllocedSize = RequiredShadowSize;
349+
if (PrivateBasePtr) {
350+
351+
UR_CALL_NOCHECK(getContext()->urDdiTable.USM.pfnFree(
352+
Context, (void *)PrivateBasePtr));
353+
PrivateBasePtr = 0;
354+
LastPrivateBaseAllocedSize = 0;
325355
}
326356

327-
Begin = PrivateShadowOffset;
328-
End = PrivateShadowOffset + RequiredShadowSize - 1;
357+
if (PrivateShadowOffset) {
358+
UR_CALL_NOCHECK(getContext()->urDdiTable.USM.pfnFree(
359+
Context, (void *)PrivateShadowOffset));
360+
PrivateShadowOffset = 0;
361+
LastPrivateShadowAllocedSize = 0;
362+
}
363+
364+
return UrRes;
329365
}
366+
330367
return UR_RESULT_SUCCESS;
331368
}
332369

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,7 @@ struct ShadowMemoryGPU : public ShadowMemory {
121121
uptr LocalShadowOffset = 0;
122122

123123
uptr PrivateShadowOffset = 0;
124+
uptr PrivateBasePtr = 0;
124125
};
125126

126127
/// Shadow Memory layout of GPU PVC device

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

Lines changed: 65 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -224,11 +224,19 @@ ur_result_t MsanShadowMemoryGPU::Destory() {
224224
static ur_result_t Result = [this]() {
225225
auto Result = getContext()->urDdiTable.VirtualMem.pfnFree(
226226
Context, (const void *)ShadowBegin, GetShadowSize());
227+
227228
if (PrivateShadowOffset != 0) {
228229
UR_CALL(getContext()->urDdiTable.USM.pfnFree(
229230
Context, (void *)PrivateShadowOffset));
230231
PrivateShadowOffset = 0;
231232
}
233+
234+
if (PrivateBasePtr != 0) {
235+
UR_CALL(getContext()->urDdiTable.USM.pfnFree(Context,
236+
(void *)PrivateBasePtr));
237+
PrivateBasePtr = 0;
238+
}
239+
232240
if (LocalShadowOffset != 0) {
233241
UR_CALL(getContext()->urDdiTable.USM.pfnFree(Context,
234242
(void *)LocalShadowOffset));
@@ -409,45 +417,74 @@ ur_result_t MsanShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue,
409417
uint64_t NumWI,
410418
uint32_t NumWG, uptr *&Base,
411419
uptr &Begin, uptr &End) {
412-
{
413-
const size_t Size = NumWI * sizeof(uptr);
414-
ur_usm_desc_t Properties{UR_STRUCTURE_TYPE_USM_DESC, nullptr,
415-
UR_USM_ADVICE_FLAG_DEFAULT, sizeof(uptr)};
416-
UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc(
417-
Context, Device, &Properties, nullptr, Size, (void **)&Base));
418-
}
420+
// Trying to allocate private base array and private shadow, and any one of
421+
// them fail to allocate would be a failure
422+
static size_t LastPrivateBaseAllocedSize = 0;
423+
static size_t LastPrivateShadowAllocedSize = 0;
424+
425+
try {
426+
const size_t NewPrivateBaseSize = NumWI * sizeof(uptr);
427+
if (NewPrivateBaseSize > LastPrivateBaseAllocedSize) {
428+
if (PrivateBasePtr) {
429+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnFree(
430+
Context, (void *)PrivateBasePtr));
431+
PrivateBasePtr = 0;
432+
LastPrivateBaseAllocedSize = 0;
433+
}
434+
435+
ur_usm_desc_t PrivateBaseProps{UR_STRUCTURE_TYPE_USM_DESC, nullptr,
436+
UR_USM_ADVICE_FLAG_DEFAULT, sizeof(uptr)};
437+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc(
438+
Context, Device, &PrivateBaseProps, nullptr, NewPrivateBaseSize,
439+
(void **)&PrivateBasePtr));
440+
441+
// No need to clean the shadow base, their should be set by work item on
442+
// launch
443+
444+
// FIXME: we should add private base to statistic
445+
LastPrivateBaseAllocedSize = NewPrivateBaseSize;
446+
}
447+
448+
const size_t NewPrivateShadowSize = NumWG * MSAN_PRIVATE_SIZE;
449+
if (NewPrivateShadowSize > LastPrivateShadowAllocedSize) {
419450

420-
{
421-
const size_t RequiredShadowSize = NumWG * MSAN_PRIVATE_SIZE;
422-
static size_t LastAllocedSize = 0;
423-
if (RequiredShadowSize > LastAllocedSize) {
424-
auto ContextInfo = getMsanInterceptor()->getContextInfo(Context);
425451
if (PrivateShadowOffset) {
426-
UR_CALL(getContext()->urDdiTable.USM.pfnFree(
452+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnFree(
427453
Context, (void *)PrivateShadowOffset));
428454
PrivateShadowOffset = 0;
429-
LastAllocedSize = 0;
455+
LastPrivateShadowAllocedSize = 0;
430456
}
431457

432-
UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc(
433-
Context, Device, nullptr, nullptr, RequiredShadowSize,
458+
UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc(
459+
Context, Device, nullptr, nullptr, NewPrivateShadowSize,
434460
(void **)&PrivateShadowOffset));
461+
LastPrivateShadowAllocedSize = NewPrivateShadowSize;
462+
UR_CALL_THROWS(EnqueueUSMSet(Queue, (void *)PrivateShadowOffset, (char)0,
463+
NewPrivateShadowSize));
464+
}
435465

436-
// Initialize shadow memory
437-
ur_result_t URes = EnqueueUSMSet(Queue, (void *)PrivateShadowOffset,
438-
(char)0, RequiredShadowSize);
439-
if (URes != UR_RESULT_SUCCESS) {
440-
UR_CALL(getContext()->urDdiTable.USM.pfnFree(
441-
Context, (void *)PrivateShadowOffset));
442-
PrivateShadowOffset = 0;
443-
LastAllocedSize = 0;
444-
}
466+
Base = (uptr *)PrivateBasePtr;
467+
Begin = PrivateShadowOffset;
468+
End = PrivateShadowOffset + NewPrivateShadowSize - 1;
445469

446-
LastAllocedSize = RequiredShadowSize;
470+
} catch (ur_result_t &UrRes) {
471+
assert(UrRes != UR_RESULT_SUCCESS);
472+
473+
if (PrivateBasePtr) {
474+
UR_CALL_NOCHECK(getContext()->urDdiTable.USM.pfnFree(
475+
Context, (void *)PrivateBasePtr));
476+
PrivateBasePtr = 0;
477+
LastPrivateBaseAllocedSize = 0;
447478
}
448479

449-
Begin = PrivateShadowOffset;
450-
End = PrivateShadowOffset + RequiredShadowSize - 1;
480+
if (PrivateShadowOffset) {
481+
UR_CALL_NOCHECK(getContext()->urDdiTable.USM.pfnFree(
482+
Context, (void *)PrivateShadowOffset));
483+
PrivateShadowOffset = 0;
484+
LastPrivateShadowAllocedSize = 0;
485+
}
486+
487+
return UrRes;
451488
}
452489

453490
return UR_RESULT_SUCCESS;

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

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -169,6 +169,7 @@ struct MsanShadowMemoryGPU : public MsanShadowMemory {
169169

170170
uptr LocalShadowOffset = 0;
171171
uptr PrivateShadowOffset = 0;
172+
uptr PrivateBasePtr = 0;
172173
};
173174

174175
// clang-format off

0 commit comments

Comments
 (0)