From 69966c77a55011961e248ffc4ffa08556083b79d Mon Sep 17 00:00:00 2001 From: "Wu, Yingcong" Date: Wed, 6 Nov 2024 21:35:06 -0800 Subject: [PATCH 1/3] minor fixes --- .../layers/sanitizer/asan_interceptor.cpp | 21 +++++++++++-------- .../loader/layers/sanitizer/asan_shadow.cpp | 1 + source/loader/layers/sanitizer/ur_sanddi.cpp | 2 ++ 3 files changed, 15 insertions(+), 9 deletions(-) diff --git a/source/loader/layers/sanitizer/asan_interceptor.cpp b/source/loader/layers/sanitizer/asan_interceptor.cpp index 4a315588fd..5523a6d2a8 100644 --- a/source/loader/layers/sanitizer/asan_interceptor.cpp +++ b/source/loader/layers/sanitizer/asan_interceptor.cpp @@ -215,25 +215,28 @@ ur_result_t SanitizerInterceptor::releaseMemory(ur_context_handle_t Context, if (ReleaseList.size()) { std::scoped_lock Guard(m_AllocationMapMutex); for (auto &It : ReleaseList) { + auto ToFreeAllocInfo = It->second; getContext()->logger.info("Quarantine Free: {}", - (void *)It->second->AllocBegin); + (void *)ToFreeAllocInfo->AllocBegin); - ContextInfo->Stats.UpdateUSMRealFreed(AllocInfo->AllocSize, - AllocInfo->getRedzoneSize()); + ContextInfo->Stats.UpdateUSMRealFreed( + ToFreeAllocInfo->AllocSize, ToFreeAllocInfo->getRedzoneSize()); - m_AllocationMap.erase(It); - if (AllocInfo->Type == AllocType::HOST_USM) { + if (ToFreeAllocInfo->Type == AllocType::HOST_USM) { for (auto &Device : ContextInfo->DeviceList) { UR_CALL(getDeviceInfo(Device)->Shadow->ReleaseShadow( - AllocInfo)); + ToFreeAllocInfo)); } } else { - UR_CALL(getDeviceInfo(AllocInfo->Device) - ->Shadow->ReleaseShadow(AllocInfo)); + UR_CALL(getDeviceInfo(ToFreeAllocInfo->Device) + ->Shadow->ReleaseShadow(ToFreeAllocInfo)); } UR_CALL(getContext()->urDdiTable.USM.pfnFree( - Context, (void *)(It->second->AllocBegin))); + Context, (void *)(ToFreeAllocInfo->AllocBegin))); + + // Erase it at last to avoid use-after-free. + m_AllocationMap.erase(It); } } ContextInfo->Stats.UpdateUSMFreed(AllocInfo->AllocSize); diff --git a/source/loader/layers/sanitizer/asan_shadow.cpp b/source/loader/layers/sanitizer/asan_shadow.cpp index 629ce3a491..e0be521ffd 100644 --- a/source/loader/layers/sanitizer/asan_shadow.cpp +++ b/source/loader/layers/sanitizer/asan_shadow.cpp @@ -249,6 +249,7 @@ ur_result_t ShadowMemoryGPU::ReleaseShadow(std::shared_ptr AI) { getContext()->logger.debug("urVirtualMemUnmap: {} ~ {}", (void *)MappedPtr, (void *)(MappedPtr + PageSize - 1)); + VirtualMemMaps.erase(MappedPtr); } } diff --git a/source/loader/layers/sanitizer/ur_sanddi.cpp b/source/loader/layers/sanitizer/ur_sanddi.cpp index 95b1649691..f065a116e4 100644 --- a/source/loader/layers/sanitizer/ur_sanddi.cpp +++ b/source/loader/layers/sanitizer/ur_sanddi.cpp @@ -357,6 +357,7 @@ __urdlllocal ur_result_t UR_APICALL urProgramLink( UR_CALL(pfnProgramLink(hContext, count, phPrograms, pOptions, phProgram)); + UR_CALL(getContext()->interceptor->insertProgram(*phProgram)); UR_CALL(getContext()->interceptor->registerProgram(hContext, *phProgram)); return UR_RESULT_SUCCESS; @@ -388,6 +389,7 @@ ur_result_t UR_APICALL urProgramLinkExp( UR_CALL(pfnProgramLinkExp(hContext, numDevices, phDevices, count, phPrograms, pOptions, phProgram)); + UR_CALL(getContext()->interceptor->insertProgram(*phProgram)); UR_CALL(getContext()->interceptor->registerProgram(hContext, *phProgram)); return UR_RESULT_SUCCESS; From 6572a044d41286949eb4be35bacfbbfdda825daa Mon Sep 17 00:00:00 2001 From: "Wu, Yingcong" Date: Wed, 27 Nov 2024 03:39:29 +0100 Subject: [PATCH 2/3] allow uninserted program --- source/loader/layers/sanitizer/asan/asan_ddi.cpp | 11 +++++++---- .../loader/layers/sanitizer/asan/asan_interceptor.cpp | 3 +++ .../loader/layers/sanitizer/asan/asan_interceptor.hpp | 6 ++++-- 3 files changed, 14 insertions(+), 6 deletions(-) diff --git a/source/loader/layers/sanitizer/asan/asan_ddi.cpp b/source/loader/layers/sanitizer/asan/asan_ddi.cpp index 489c4cc4e4..9695a8461e 100644 --- a/source/loader/layers/sanitizer/asan/asan_ddi.cpp +++ b/source/loader/layers/sanitizer/asan/asan_ddi.cpp @@ -55,6 +55,9 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices, bool isInstrumentedKernel(ur_kernel_handle_t hKernel) { auto hProgram = GetProgram(hKernel); auto PI = getAsanInterceptor()->getProgramInfo(hProgram); + if (PI == nullptr) { + return false; + } return PI->isKernelInstrumented(hKernel); } @@ -290,8 +293,9 @@ __urdlllocal ur_result_t UR_APICALL urProgramRetain( UR_CALL(pfnRetain(hProgram)); auto ProgramInfo = getAsanInterceptor()->getProgramInfo(hProgram); - UR_ASSERT(ProgramInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE); - ProgramInfo->RefCount++; + if (ProgramInfo != nullptr) { + ProgramInfo->RefCount++; + } return UR_RESULT_SUCCESS; } @@ -419,8 +423,7 @@ ur_result_t UR_APICALL urProgramRelease( UR_CALL(pfnProgramRelease(hProgram)); auto ProgramInfo = getAsanInterceptor()->getProgramInfo(hProgram); - UR_ASSERT(ProgramInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE); - if (--ProgramInfo->RefCount == 0) { + if (ProgramInfo != nullptr && --ProgramInfo->RefCount == 0) { UR_CALL(getAsanInterceptor()->unregisterProgram(hProgram)); UR_CALL(getAsanInterceptor()->eraseProgram(hProgram)); } diff --git a/source/loader/layers/sanitizer/asan/asan_interceptor.cpp b/source/loader/layers/sanitizer/asan/asan_interceptor.cpp index af5bd59944..95bde5a5b5 100644 --- a/source/loader/layers/sanitizer/asan/asan_interceptor.cpp +++ b/source/loader/layers/sanitizer/asan/asan_interceptor.cpp @@ -436,6 +436,7 @@ ur_result_t AsanInterceptor::registerProgram(ur_program_handle_t Program) { ur_result_t AsanInterceptor::unregisterProgram(ur_program_handle_t Program) { auto ProgramInfo = getProgramInfo(Program); + assert(ProgramInfo != nullptr && "unregistered program!"); for (auto AI : ProgramInfo->AllocInfoForGlobals) { UR_CALL(getDeviceInfo(AI->Device)->Shadow->ReleaseShadow(AI)); @@ -483,6 +484,7 @@ ur_result_t AsanInterceptor::registerSpirKernels(ur_program_handle_t Program) { } auto PI = getProgramInfo(Program); + assert(PI != nullptr && "unregistered program!"); for (const auto &SKI : SKInfo) { if (SKI.Size == 0) { continue; @@ -519,6 +521,7 @@ AsanInterceptor::registerDeviceGlobals(ur_program_handle_t Program) { auto Context = GetContext(Program); auto ContextInfo = getContextInfo(Context); auto ProgramInfo = getProgramInfo(Program); + assert(ProgramInfo != nullptr && "unregistered program!"); for (auto Device : Devices) { ManagedQueue Queue(Context, Device); diff --git a/source/loader/layers/sanitizer/asan/asan_interceptor.hpp b/source/loader/layers/sanitizer/asan/asan_interceptor.hpp index d8dd11101c..4254fcff15 100644 --- a/source/loader/layers/sanitizer/asan/asan_interceptor.hpp +++ b/source/loader/layers/sanitizer/asan/asan_interceptor.hpp @@ -266,8 +266,10 @@ class AsanInterceptor { std::shared_ptr getProgramInfo(ur_program_handle_t Program) { std::shared_lock Guard(m_ProgramMapMutex); - assert(m_ProgramMap.find(Program) != m_ProgramMap.end()); - return m_ProgramMap[Program]; + if (m_ProgramMap.find(Program) != m_ProgramMap.end()) { + return m_ProgramMap[Program]; + } + return nullptr; } std::shared_ptr getKernelInfo(ur_kernel_handle_t Kernel) { From 8411e2bcc841ee82822362af8b81db0c3d99210a Mon Sep 17 00:00:00 2001 From: "Wu, Yingcong" Date: Wed, 27 Nov 2024 03:50:55 +0100 Subject: [PATCH 3/3] update api usage --- source/loader/layers/sanitizer/asan/asan_ddi.cpp | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/source/loader/layers/sanitizer/asan/asan_ddi.cpp b/source/loader/layers/sanitizer/asan/asan_ddi.cpp index 9695a8461e..741b4d421c 100644 --- a/source/loader/layers/sanitizer/asan/asan_ddi.cpp +++ b/source/loader/layers/sanitizer/asan/asan_ddi.cpp @@ -369,7 +369,7 @@ __urdlllocal ur_result_t UR_APICALL urProgramLink( UR_CALL(pfnProgramLink(hContext, count, phPrograms, pOptions, phProgram)); UR_CALL(getAsanInterceptor()->insertProgram(*phProgram)); - UR_CALL(getAsanInterceptor()->registerProgram(hContext, *phProgram)); + UR_CALL(getAsanInterceptor()->registerProgram(*phProgram)); return UR_RESULT_SUCCESS; } @@ -401,7 +401,7 @@ ur_result_t UR_APICALL urProgramLinkExp( phPrograms, pOptions, phProgram)); UR_CALL(getAsanInterceptor()->insertProgram(*phProgram)); - UR_CALL(getAsanInterceptor()->registerProgram(hContext, *phProgram)); + UR_CALL(getAsanInterceptor()->registerProgram(*phProgram)); return UR_RESULT_SUCCESS; }