[DevASAN] Move quarantine into context info for better management (#19602)

Before this change, quarantine is maintained in AsanInterceptor, this
will cause sometimes we free a USM pointer with wrong context.

Author: zhaomaosu

unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp

@@ -25,12 +25,7 @@
 namespace ur_sanitizer_layer {
 namespace asan {
 
-AsanInterceptor::AsanInterceptor() {
-  if (getContext()->Options.MaxQuarantineSizeMB) {
-    m_Quarantine = std::make_unique<Quarantine>(
-        getContext()->Options.MaxQuarantineSizeMB * 1024 * 1024);
-  }
-}
+AsanInterceptor::AsanInterceptor() {}
 
 AsanInterceptor::~AsanInterceptor() {
   // We must release these objects before releasing adapters, since
@@ -39,7 +34,6 @@ AsanInterceptor::~AsanInterceptor() {
     DeviceInfo->Shadow = nullptr;
   }
 
-  m_Quarantine = nullptr;
   m_MemBufferMap.clear();
   m_KernelMap.clear();
   m_ContextMap.clear();
@@ -224,7 +218,7 @@ ur_result_t AsanInterceptor::releaseMemory(ur_context_handle_t Context,
   }
 
   // If quarantine is disabled, USM is freed immediately
-  if (!m_Quarantine) {
+  if (!ContextInfo->m_Quarantine) {
     UR_LOG_L(getContext()->logger, DEBUG, "Free: {}",
              (void *)AllocInfo->AllocBegin);
 
@@ -239,7 +233,8 @@ ur_result_t AsanInterceptor::releaseMemory(ur_context_handle_t Context,
   }
 
   // If quarantine is enabled, cache it
-  auto ReleaseList = m_Quarantine->put(AllocInfo->Device, AllocInfoIt);
+  auto ReleaseList =
+      ContextInfo->m_Quarantine->put(AllocInfo->Device, AllocInfoIt);
   if (ReleaseList.size()) {
     std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
     for (auto &It : ReleaseList) {

unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp

@@ -16,6 +16,7 @@
 #include "asan_allocator.hpp"
 #include "asan_buffer.hpp"
 #include "asan_libdevice.hpp"
+#include "asan_quarantine.hpp"
 #include "asan_shadow.hpp"
 #include "asan_statistics.hpp"
 #include "sanitizer_common/sanitizer_common.hpp"
@@ -33,8 +34,6 @@
 namespace ur_sanitizer_layer {
 namespace asan {
 
-class Quarantine;
-
 struct AllocInfoList {
   std::vector<std::shared_ptr<AllocInfo>> List;
   ur_shared_mutex Mutex;
@@ -148,12 +147,18 @@ struct ContextInfo {
   std::unordered_map<ur_device_handle_t, std::optional<ManagedQueue>>
       InternalQueueMap;
 
+  std::optional<Quarantine> m_Quarantine;
+
   AsanStatsWrapper Stats;
 
   explicit ContextInfo(ur_context_handle_t Context) : Handle(Context) {
     [[maybe_unused]] auto Result =
         getContext()->urDdiTable.Context.pfnRetain(Context);
     assert(Result == UR_RESULT_SUCCESS);
+    if (getContext()->Options.MaxQuarantineSizeMB) {
+      m_Quarantine.emplace(getContext()->Options.MaxQuarantineSizeMB * 1024 *
+                           1024);
+    }
   }
 
   ~ContextInfo();
@@ -404,8 +409,6 @@ class AsanInterceptor {
   AllocationMap m_AllocationMap;
   ur_shared_mutex m_AllocationMapMutex;
 
-  std::unique_ptr<Quarantine> m_Quarantine;
-
   std::unordered_set<ur_adapter_handle_t> m_Adapters;
   ur_shared_mutex m_AdaptersMutex;