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) {