diff --git a/source/loader/layers/sanitizer/asan_interceptor.cpp b/source/loader/layers/sanitizer/asan_interceptor.cpp index 4a315588fd..129731193f 100644 --- a/source/loader/layers/sanitizer/asan_interceptor.cpp +++ b/source/loader/layers/sanitizer/asan_interceptor.cpp @@ -39,9 +39,11 @@ SanitizerInterceptor::~SanitizerInterceptor() { m_Quarantine = nullptr; m_MemBufferMap.clear(); - m_AllocationMap.clear(); m_KernelMap.clear(); m_ContextMap.clear(); + // AllocationMap need to be cleared after ContextMap because memory leak + // detection depends on it. + m_AllocationMap.clear(); for (auto Adapter : m_Adapters) { getContext()->urDdiTable.Global.pfnAdapterRelease(Adapter); @@ -290,7 +292,7 @@ ur_result_t SanitizerInterceptor::postLaunchKernel(ur_kernel_handle_t Kernel, ReportFatalError(AH); } if (!AH.IsRecover) { - exit(1); + exitWithErrors(); } } } @@ -616,7 +618,7 @@ ur_result_t SanitizerInterceptor::prepareLaunch( ContextInfo->Handle, DeviceInfo->Handle, (uptr)Ptr)) { ReportInvalidKernelArgument(Kernel, ArgIndex, (uptr)Ptr, ValidateResult, PtrPair.second); - exit(1); + exitWithErrors(); } } } @@ -838,12 +840,14 @@ ContextInfo::~ContextInfo() { assert(Result == UR_RESULT_SUCCESS); // check memory leaks - std::vector AllocInfos = - getContext()->interceptor->findAllocInfoByContext(Handle); - for (const auto &It : AllocInfos) { - const auto &[_, AI] = *It; - if (!AI->IsReleased) { - ReportMemoryLeak(AI); + if (getContext()->interceptor->isNormalExit()) { + std::vector AllocInfos = + getContext()->interceptor->findAllocInfoByContext(Handle); + for (const auto &It : AllocInfos) { + const auto &[_, AI] = *It; + if (!AI->IsReleased) { + ReportMemoryLeak(AI); + } } } } diff --git a/source/loader/layers/sanitizer/asan_interceptor.hpp b/source/loader/layers/sanitizer/asan_interceptor.hpp index e5429acd56..c1bf710425 100644 --- a/source/loader/layers/sanitizer/asan_interceptor.hpp +++ b/source/loader/layers/sanitizer/asan_interceptor.hpp @@ -271,6 +271,13 @@ class SanitizerInterceptor { const AsanOptions &getOptions() { return m_Options; } + void exitWithErrors() { + m_NormalExit = false; + exit(1); + } + + bool isNormalExit() { return m_NormalExit; } + private: ur_result_t updateShadowMemory(std::shared_ptr &ContextInfo, std::shared_ptr &DeviceInfo, @@ -320,6 +327,8 @@ class SanitizerInterceptor { std::unordered_set m_Adapters; ur_shared_mutex m_AdaptersMutex; + + bool m_NormalExit = true; }; } // namespace ur_sanitizer_layer