allow uninserted program

oneapi-src · Nov 27, 2024 · 6572a04 · 6572a04
1 parent 093a95d
commit 6572a04
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git 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
@@ -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
@@ -266,8 +266,10 @@ class AsanInterceptor {
 
     std::shared_ptr<ProgramInfo> getProgramInfo(ur_program_handle_t Program) {
         std::shared_lock<ur_shared_mutex> 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<KernelInfo> getKernelInfo(ur_kernel_handle_t Kernel) {