register globals

Author: yingcong-wu

source/loader/layers/sanitizer/asan/asan_interceptor.cpp

@@ -431,6 +431,12 @@ ur_result_t AsanInterceptor::unregisterProgram(ur_program_handle_t Program) {
     auto ProgramInfo = getProgramInfo(Program);
     assert(ProgramInfo != nullptr && "unregistered program!");
 
+    std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
+    for (auto AI : ProgramInfo->AllocInfoForGlobals) {
+        m_AllocationMap.erase(AI->AllocBegin);
+    }
+    ProgramInfo->AllocInfoForGlobals.clear();
+
     ProgramInfo->InstrumentedKernels.clear();
 
     return UR_RESULT_SUCCESS;
@@ -549,6 +555,10 @@ AsanInterceptor::registerDeviceGlobals(ur_program_handle_t Program) {
                           {}});
 
             ContextInfo->insertAllocInfo({Device}, AI);
+            ProgramInfo->AllocInfoForGlobals.emplace(AI);
+
+            std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
+            m_AllocationMap.emplace(AI->AllocBegin, std::move(AI));
         }
     }
 

source/loader/layers/sanitizer/asan/asan_interceptor.hpp

@@ -112,6 +112,7 @@ struct ProgramInfo {
     std::atomic<int32_t> RefCount = 1;
 
     // Program is built only once, so we don't need to lock it
+    std::unordered_set<std::shared_ptr<AllocInfo>> AllocInfoForGlobals;
     std::unordered_set<std::string> InstrumentedKernels;
 
     explicit ProgramInfo(ur_program_handle_t Program) : Handle(Program) {