[DevTSAN] Fix UR objects leak problem (#18880)

We forgot to destroy tsan interceptor, this will cause objects leak
issue.

Author: zhaomaosu

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

@@ -43,6 +43,32 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices,
 
 } // namespace
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urAdapterGet
+ur_result_t urAdapterGet(
+    /// [in] the number of adapters to be added to phAdapters. If phAdapters
+    /// is not NULL, then NumEntries should be greater than zero, otherwise
+    /// ::UR_RESULT_ERROR_INVALID_SIZE, will be returned.
+    uint32_t NumEntries,
+    /// [out][optional][range(0, NumEntries)] array of handle of adapters. If
+    /// NumEntries is less than the number of adapters available, then
+    /// ::urAdapterGet shall only retrieve that number of platforms.
+    ur_adapter_handle_t *phAdapters,
+    /// [out][optional] returns the total number of adapters available.
+    uint32_t *pNumAdapters) {
+  auto pfnAdapterGet = getContext()->urDdiTable.Adapter.pfnGet;
+
+  ur_result_t result = pfnAdapterGet(NumEntries, phAdapters, pNumAdapters);
+  if (result == UR_RESULT_SUCCESS && phAdapters) {
+    const uint32_t NumAdapters = pNumAdapters ? *pNumAdapters : NumEntries;
+    for (uint32_t i = 0; i < NumAdapters; ++i) {
+      UR_CALL(getTsanInterceptor()->holdAdapter(phAdapters[i]));
+    }
+  }
+
+  return result;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urContextCreate
 __urdlllocal ur_result_t UR_APICALL urContextCreate(
@@ -1201,6 +1227,27 @@ ur_result_t urCheckVersion(ur_api_version_t version) {
   return UR_RESULT_SUCCESS;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Exported function for filling application's Adapter table
+///        with current process' addresses
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_UNINITIALIZED
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+///     - ::UR_RESULT_ERROR_UNSUPPORTED_VERSION
+ur_result_t urGetAdapterProcAddrTable(
+    /// [in,out] pointer to table of DDI function pointers
+    ur_adapter_dditable_t *pDdiTable) {
+  if (nullptr == pDdiTable) {
+    return UR_RESULT_ERROR_INVALID_NULL_POINTER;
+  }
+
+  pDdiTable->pfnGet = ur_sanitizer_layer::tsan::urAdapterGet;
+
+  return UR_RESULT_SUCCESS;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for filling application's Context table
 ///        with current process' addresses
@@ -1381,6 +1428,11 @@ ur_result_t initTsanDDITable(ur_dditable_t *dditable) {
     result = ur_sanitizer_layer::tsan::urCheckVersion(UR_API_VERSION_CURRENT);
   }
 
+  if (UR_RESULT_SUCCESS == result) {
+    result =
+        ur_sanitizer_layer::tsan::urGetAdapterProcAddrTable(&dditable->Adapter);
+  }
+
   if (UR_RESULT_SUCCESS == result) {
     result =
         ur_sanitizer_layer::tsan::urGetContextProcAddrTable(&dditable->Context);

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

@@ -103,6 +103,23 @@ void ContextInfo::insertAllocInfo(ur_device_handle_t Device, TsanAllocInfo AI) {
   }
 }
 
+TsanInterceptor::~TsanInterceptor() {
+  // We must release these objects before releasing adapters, since
+  // they may use the adapter in their destructor
+  for (const auto &[_, DeviceInfo] : m_DeviceMap) {
+    DeviceInfo->Shadow->Destroy();
+    DeviceInfo->Shadow = nullptr;
+  }
+
+  m_MemBufferMap.clear();
+  m_KernelMap.clear();
+  m_ContextMap.clear();
+
+  for (auto Adapter : m_Adapters) {
+    getContext()->urDdiTable.Adapter.pfnRelease(Adapter);
+  }
+}
+
 ur_result_t TsanInterceptor::allocateMemory(ur_context_handle_t Context,
                                             ur_device_handle_t Device,
                                             const ur_usm_desc_t *Properties,

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

@@ -20,6 +20,9 @@
 #include "tsan_shadow.hpp"
 #include "ur_sanitizer_layer.hpp"
 
+#include <unordered_map>
+#include <unordered_set>
+
 namespace ur_sanitizer_layer {
 namespace tsan {
 
@@ -176,6 +179,8 @@ struct LaunchInfo {
 
 class TsanInterceptor {
 public:
+  ~TsanInterceptor();
+
   ur_result_t allocateMemory(ur_context_handle_t Context,
                              ur_device_handle_t Device,
                              const ur_usm_desc_t *Properties,
@@ -202,6 +207,16 @@ class TsanInterceptor {
 
   std::shared_ptr<MemBuffer> getMemBuffer(ur_mem_handle_t MemHandle);
 
+  ur_result_t holdAdapter(ur_adapter_handle_t Adapter) {
+    std::scoped_lock<ur_shared_mutex> Guard(m_AdaptersMutex);
+    if (m_Adapters.find(Adapter) != m_Adapters.end()) {
+      return UR_RESULT_SUCCESS;
+    }
+    UR_CALL(getContext()->urDdiTable.Adapter.pfnRetain(Adapter));
+    m_Adapters.insert(Adapter);
+    return UR_RESULT_SUCCESS;
+  }
+
   ur_result_t preLaunchKernel(ur_kernel_handle_t Kernel,
                               ur_queue_handle_t Queue, LaunchInfo &LaunchInfo);
 
@@ -255,6 +270,9 @@ class TsanInterceptor {
   std::unordered_map<ur_mem_handle_t, std::shared_ptr<MemBuffer>>
       m_MemBufferMap;
   ur_shared_mutex m_MemBufferMapMutex;
+
+  std::unordered_set<ur_adapter_handle_t> m_Adapters;
+  ur_shared_mutex m_AdaptersMutex;
 };
 
 } // namespace tsan

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

@@ -47,7 +47,7 @@ ur_result_t ShadowMemoryCPU::Setup() {
   return URes;
 }
 
-ur_result_t ShadowMemoryCPU::Destory() {
+ur_result_t ShadowMemoryCPU::Destroy() {
   if (ShadowBegin == 0 && ShadowEnd == 0)
     return UR_RESULT_SUCCESS;
   static ur_result_t URes = [this]() {
@@ -98,12 +98,10 @@ ur_result_t ShadowMemoryGPU::Setup() {
     return Result;
   }
   ShadowEnd = ShadowBegin + ShadowSize;
-  // Retain the context which reserves shadow memory
-  getContext()->urDdiTable.Context.pfnRetain(Context);
   return UR_RESULT_SUCCESS;
 }
 
-ur_result_t ShadowMemoryGPU::Destory() {
+ur_result_t ShadowMemoryGPU::Destroy() {
   if (ShadowBegin == 0) {
     return UR_RESULT_SUCCESS;
   }

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

@@ -28,7 +28,7 @@ struct ShadowMemory {
 
   virtual ur_result_t Setup() = 0;
 
-  virtual ur_result_t Destory() = 0;
+  virtual ur_result_t Destroy() = 0;
 
   virtual RawShadow *MemToShadow(uptr Ptr) = 0;
 
@@ -69,7 +69,7 @@ struct ShadowMemoryCPU final : public ShadowMemory {
 
   ur_result_t Setup() override;
 
-  ur_result_t Destory() override;
+  ur_result_t Destroy() override;
 
   RawShadow *MemToShadow(uptr Ptr) override;
 
@@ -92,7 +92,7 @@ struct ShadowMemoryGPU : public ShadowMemory {
 
   ur_result_t Setup() override;
 
-  ur_result_t Destory() override;
+  ur_result_t Destroy() override;
 
   ur_result_t CleanShadow(ur_queue_handle_t Queue, uptr Ptr,
                           uptr Size) override;

unified-runtime/source/loader/layers/sanitizer/ur_sanitizer_layer.cpp

@@ -14,6 +14,7 @@
 #include "ur_sanitizer_layer.hpp"
 #include "asan/asan_ddi.hpp"
 #include "msan/msan_ddi.hpp"
+#include "tsan/tsan_ddi.hpp"
 
 namespace ur_sanitizer_layer {
 context_t *getContext() {
@@ -39,6 +40,9 @@ ur_result_t context_t::tearDown() {
   case SanitizerType::MemorySanitizer:
     destroyMsanInterceptor();
     break;
+  case SanitizerType::ThreadSanitizer:
+    destroyTsanInterceptor();
+    break;
   default:
     break;
   }