[DevTSAN] Move AllocInfo into DeviceInfo to support indirect access (#19634)

If we maintain AllocInfo in ContextInfo, this will cause the pointer
which is allocated by another context can't be poisoned.

Author: zhaomaosu

unified-runtime/source/loader/layers/sanitizer/tsan/tsan_buffer.cpp

@@ -186,8 +186,10 @@ ur_result_t MemBuffer::getHandle(ur_device_handle_t Device, char *&Handle) {
 }
 
 ur_result_t MemBuffer::free() {
-  for (const auto &[_, Ptr] : Allocations) {
-    ur_result_t URes = getTsanInterceptor()->releaseMemory(Context, Ptr);
+  for (const auto &[Device, Ptr] : Allocations) {
+    ur_result_t URes = Device
+                           ? getTsanInterceptor()->releaseMemory(Context, Ptr)
+                           : getContext()->urDdiTable.USM.pfnFree(Context, Ptr);
     if (URes != UR_RESULT_SUCCESS) {
       UR_LOG_L(getContext()->logger, ERR, "Failed to free buffer handle {}",
                (void *)Ptr);

unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.cpp

@@ -101,7 +101,7 @@ ur_result_t DeviceInfo::allocShadowMemory() {
   return UR_RESULT_SUCCESS;
 }
 
-void ContextInfo::insertAllocInfo(TsanAllocInfo AI) {
+void DeviceInfo::insertAllocInfo(TsanAllocInfo AI) {
   std::scoped_lock<ur_shared_mutex> Guard(AllocInfosMutex);
   AllocInfos.insert(std::move(AI));
 }
@@ -153,7 +153,15 @@ ur_result_t TsanInterceptor::allocateMemory(ur_context_handle_t Context,
 
   auto AI = TsanAllocInfo{reinterpret_cast<uptr>(Allocated), Size};
   // For updating shadow memory
-  CI->insertAllocInfo(std::move(AI));
+  if (Device) {
+    auto DI = getDeviceInfo(Device);
+    DI->insertAllocInfo(std::move(AI));
+  } else {
+    for (const auto &Device : CI->DeviceList) {
+      auto DI = getDeviceInfo(Device);
+      DI->insertAllocInfo(AI);
+    }
+  }
 
   *ResultPtr = Allocated;
   return UR_RESULT_SUCCESS;
@@ -163,11 +171,14 @@ ur_result_t TsanInterceptor::releaseMemory(ur_context_handle_t Context,
                                            void *Ptr) {
   auto CI = getContextInfo(Context);
   auto Addr = reinterpret_cast<uptr>(Ptr);
-  {
-    std::scoped_lock<ur_shared_mutex> Guard(CI->AllocInfosMutex);
-    auto It = std::find_if(CI->AllocInfos.begin(), CI->AllocInfos.end(),
+
+  for (const auto &Device : CI->DeviceList) {
+    auto DI = getDeviceInfo(Device);
+    std::scoped_lock<ur_shared_mutex> Guard(DI->AllocInfosMutex);
+    auto It = std::find_if(DI->AllocInfos.begin(), DI->AllocInfos.end(),
                            [&](auto &P) { return P.AllocBegin == Addr; });
-    CI->AllocInfos.erase(It);
+    if (It != DI->AllocInfos.end())
+      DI->AllocInfos.erase(It);
   }
 
   UR_CALL(getContext()->urDdiTable.USM.pfnFree(Context, Ptr));
@@ -343,7 +354,7 @@ ur_result_t TsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
 
   UR_CALL(prepareLaunch(CI, DI, InternalQueue, Kernel, LaunchInfo));
 
-  UR_CALL(updateShadowMemory(CI, DI, Kernel, InternalQueue));
+  UR_CALL(updateShadowMemory(DI, Kernel, InternalQueue));
 
   UR_CALL(getContext()->urDdiTable.Queue.pfnFinish(InternalQueue));
 
@@ -470,12 +481,12 @@ ur_result_t TsanInterceptor::prepareLaunch(std::shared_ptr<ContextInfo> &,
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t TsanInterceptor::updateShadowMemory(
-    std::shared_ptr<ContextInfo> &CI, std::shared_ptr<DeviceInfo> &DI,
-    ur_kernel_handle_t Kernel, ur_queue_handle_t Queue) {
+ur_result_t TsanInterceptor::updateShadowMemory(std::shared_ptr<DeviceInfo> &DI,
+                                                ur_kernel_handle_t Kernel,
+                                                ur_queue_handle_t Queue) {
   auto &PI = getProgramInfo(GetProgram(Kernel));
-  std::scoped_lock<ur_shared_mutex> Guard(CI->AllocInfosMutex);
-  for (auto &AllocInfo : CI->AllocInfos) {
+  std::scoped_lock<ur_shared_mutex> Guard(DI->AllocInfosMutex);
+  for (auto &AllocInfo : DI->AllocInfos) {
     UR_CALL(DI->Shadow->CleanShadow(Queue, AllocInfo.AllocBegin,
                                     AllocInfo.AllocSize));
   }

unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.hpp

@@ -44,9 +44,14 @@ struct DeviceInfo {
 
   std::shared_ptr<ShadowMemory> Shadow;
 
+  ur_shared_mutex AllocInfosMutex;
+  std::set<TsanAllocInfo> AllocInfos;
+
   explicit DeviceInfo(ur_device_handle_t Device) : Handle(Device) {}
 
   ur_result_t allocShadowMemory();
+
+  void insertAllocInfo(TsanAllocInfo AI);
 };
 
 struct ContextInfo {
@@ -56,9 +61,6 @@ struct ContextInfo {
 
   std::vector<ur_device_handle_t> DeviceList;
 
-  ur_shared_mutex AllocInfosMutex;
-  std::set<TsanAllocInfo> AllocInfos;
-
   ur_shared_mutex InternalQueueMapMutex;
   std::unordered_map<ur_device_handle_t, std::optional<ManagedQueue>>
       InternalQueueMap;
@@ -80,8 +82,6 @@ struct ContextInfo {
 
   ContextInfo &operator=(const ContextInfo &) = delete;
 
-  void insertAllocInfo(TsanAllocInfo AI);
-
   ur_queue_handle_t getInternalQueue(ur_device_handle_t);
 };
 
@@ -297,8 +297,7 @@ class TsanInterceptor {
   ur_shared_mutex KernelLaunchMutex;
 
 private:
-  ur_result_t updateShadowMemory(std::shared_ptr<ContextInfo> &CI,
-                                 std::shared_ptr<DeviceInfo> &DI,
+  ur_result_t updateShadowMemory(std::shared_ptr<DeviceInfo> &DI,
                                  ur_kernel_handle_t Kernel,
                                  ur_queue_handle_t Queue);