Support device thread sanitizer for device globals (#17548)

1.Add a global '__TsanDeviceGlobalMetadata' to record device global's
information
2.Read global meta data when build/link program done, and then poison
related
shadow memory

Author: zhaomaosu

source/loader/layers/sanitizer/tsan/tsan_ddi.cpp

@@ -130,6 +130,95 @@ ur_result_t urContextRelease(
   return UR_RESULT_SUCCESS;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urProgramBuild
+ur_result_t urProgramBuild(
+    /// [in] handle of the context object
+    ur_context_handle_t hContext,
+    /// [in] handle of the program object
+    ur_program_handle_t hProgram,
+    /// [in] string of build options
+    const char *pOptions) {
+  getContext()->logger.debug("==== urProgramBuild");
+
+  UR_CALL(
+      getContext()->urDdiTable.Program.pfnBuild(hContext, hProgram, pOptions));
+
+  UR_CALL(getTsanInterceptor()->registerProgram(hProgram));
+
+  return UR_RESULT_SUCCESS;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urProgramBuildExp
+ur_result_t urProgramBuildExp(
+    /// [in] Handle of the program to build.
+    ur_program_handle_t hProgram,
+    /// [in] number of devices
+    uint32_t numDevices,
+    /// [in][range(0, numDevices)] pointer to array of device handles
+    ur_device_handle_t *phDevices,
+    /// [in][optional] pointer to build options null-terminated string.
+    const char *pOptions) {
+  getContext()->logger.debug("==== urProgramBuildExp");
+
+  UR_CALL(getContext()->urDdiTable.ProgramExp.pfnBuildExp(hProgram, numDevices,
+                                                          phDevices, pOptions));
+  UR_CALL(getTsanInterceptor()->registerProgram(hProgram));
+
+  return UR_RESULT_SUCCESS;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urProgramLink
+ur_result_t urProgramLink(
+    /// [in] handle of the context instance.
+    ur_context_handle_t hContext,
+    /// [in] number of program handles in `phPrograms`.
+    uint32_t count,
+    /// [in][range(0, count)] pointer to array of program handles.
+    const ur_program_handle_t *phPrograms,
+    /// [in][optional] pointer to linker options null-terminated string.
+    const char *pOptions,
+    /// [out] pointer to handle of program object created.
+    ur_program_handle_t *phProgram) {
+  getContext()->logger.debug("==== urProgramLink");
+
+  UR_CALL(getContext()->urDdiTable.Program.pfnLink(hContext, count, phPrograms,
+                                                   pOptions, phProgram));
+
+  UR_CALL(getTsanInterceptor()->registerProgram(*phProgram));
+
+  return UR_RESULT_SUCCESS;
+}
+
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urProgramLinkExp
+ur_result_t urProgramLinkExp(
+    /// [in] handle of the context instance.
+    ur_context_handle_t hContext,
+    /// [in] number of devices
+    uint32_t numDevices,
+    /// [in][range(0, numDevices)] pointer to array of device handles
+    ur_device_handle_t *phDevices,
+    /// [in] number of program handles in `phPrograms`.
+    uint32_t count,
+    /// [in][range(0, count)] pointer to array of program handles.
+    const ur_program_handle_t *phPrograms,
+    /// [in][optional] pointer to linker options null-terminated string.
+    const char *pOptions,
+    /// [out] pointer to handle of program object created.
+    ur_program_handle_t *phProgram) {
+  getContext()->logger.debug("==== urProgramLinkExp");
+
+  UR_CALL(getContext()->urDdiTable.ProgramExp.pfnLinkExp(
+      hContext, numDevices, phDevices, count, phPrograms, pOptions, phProgram));
+
+  UR_CALL(getTsanInterceptor()->registerProgram(*phProgram));
+
+  return UR_RESULT_SUCCESS;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urUSMDeviceAlloc
 __urdlllocal ur_result_t UR_APICALL urUSMDeviceAlloc(
@@ -283,6 +372,39 @@ __urdlllocal ur_result_t UR_APICALL urGetContextProcAddrTable(
   return result;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Exported function for filling application's Program table
+///        with current process' addresses
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+ur_result_t urGetProgramProcAddrTable(
+    /// [in,out] pointer to table of DDI function pointers
+    ur_program_dditable_t *pDdiTable) {
+  pDdiTable->pfnBuild = ur_sanitizer_layer::tsan::urProgramBuild;
+  pDdiTable->pfnLink = ur_sanitizer_layer::tsan::urProgramLink;
+
+  return UR_RESULT_SUCCESS;
+}
+
+/// @brief Exported function for filling application's ProgramExp table
+///        with current process' addresses
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+ur_result_t urGetProgramExpProcAddrTable(
+    /// [in,out] pointer to table of DDI function pointers
+    ur_program_exp_dditable_t *pDdiTable) {
+  ur_result_t result = UR_RESULT_SUCCESS;
+
+  pDdiTable->pfnBuildExp = ur_sanitizer_layer::tsan::urProgramBuildExp;
+  pDdiTable->pfnLinkExp = ur_sanitizer_layer::tsan::urProgramLinkExp;
+
+  return result;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for filling application's USM table
 ///        with current process' addresses
@@ -363,6 +485,16 @@ ur_result_t initTsanDDITable(ur_dditable_t *dditable) {
         UR_API_VERSION_CURRENT, &dditable->Context);
   }
 
+  if (UR_RESULT_SUCCESS == result) {
+    result =
+        ur_sanitizer_layer::tsan::urGetProgramProcAddrTable(&dditable->Program);
+  }
+
+  if (UR_RESULT_SUCCESS == result) {
+    result = ur_sanitizer_layer::tsan::urGetProgramExpProcAddrTable(
+        &dditable->ProgramExp);
+  }
+
   if (UR_RESULT_SUCCESS == result) {
     result = ur_sanitizer_layer::tsan::urGetUSMProcAddrTable(
         UR_API_VERSION_CURRENT, &dditable->USM);

source/loader/layers/sanitizer/tsan/tsan_interceptor.cpp

@@ -69,15 +69,14 @@ ur_result_t DeviceInfo::allocShadowMemory() {
   return UR_RESULT_SUCCESS;
 }
 
-void ContextInfo::insertAllocInfo(ur_device_handle_t Device,
-                                  std::shared_ptr<TsanAllocInfo> &AI) {
+void ContextInfo::insertAllocInfo(ur_device_handle_t Device, TsanAllocInfo AI) {
   if (Device) {
     std::scoped_lock<ur_shared_mutex> Guard(AllocInfosMapMutex);
-    AllocInfosMap[Device].emplace_back(AI);
+    AllocInfosMap[Device].emplace_back(std::move(AI));
   } else {
     for (auto Device : DeviceList) {
       std::scoped_lock<ur_shared_mutex> Guard(AllocInfosMapMutex);
-      AllocInfosMap[Device].emplace_back(AI);
+      AllocInfosMap[Device].emplace_back(std::move(AI));
     }
   }
 }
@@ -103,16 +102,62 @@ ur_result_t TsanInterceptor::allocateMemory(ur_context_handle_t Context,
         Context, Device, Properties, Pool, Size, &Allocated));
   }
 
-  auto AI = std::make_shared<TsanAllocInfo>(
-      TsanAllocInfo{reinterpret_cast<uptr>(Allocated), Size});
-
+  auto AI = TsanAllocInfo{reinterpret_cast<uptr>(Allocated), Size};
   // For updating shadow memory
-  CI->insertAllocInfo(Device, AI);
+  CI->insertAllocInfo(Device, std::move(AI));
 
   *ResultPtr = Allocated;
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t TsanInterceptor::registerProgram(ur_program_handle_t Program) {
+  getContext()->logger.info("registerDeviceGlobals");
+  UR_CALL(registerDeviceGlobals(Program));
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t TsanInterceptor::registerDeviceGlobals(ur_program_handle_t Program) {
+  std::vector<ur_device_handle_t> Devices = GetDevices(Program);
+  assert(Devices.size() != 0 && "No devices in registerDeviceGlobals");
+  auto Context = GetContext(Program);
+  auto ContextInfo = getContextInfo(Context);
+
+  for (auto Device : Devices) {
+    ManagedQueue Queue(Context, Device);
+
+    size_t MetadataSize;
+    void *MetadataPtr;
+    auto Result = getContext()->urDdiTable.Program.pfnGetGlobalVariablePointer(
+        Device, Program, kSPIR_TsanDeviceGlobalMetadata, &MetadataSize,
+        &MetadataPtr);
+    if (Result != UR_RESULT_SUCCESS) {
+      getContext()->logger.info("No device globals");
+      continue;
+    }
+
+    const uint64_t NumOfDeviceGlobal = MetadataSize / sizeof(DeviceGlobalInfo);
+    assert((MetadataSize % sizeof(DeviceGlobalInfo) == 0) &&
+           "DeviceGlobal metadata size is not correct");
+    std::vector<DeviceGlobalInfo> GVInfos(NumOfDeviceGlobal);
+    Result = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy(
+        Queue, true, &GVInfos[0], MetadataPtr,
+        sizeof(DeviceGlobalInfo) * NumOfDeviceGlobal, 0, nullptr, nullptr);
+    if (Result != UR_RESULT_SUCCESS) {
+      getContext()->logger.error("Device Global[{}] Read Failed: {}",
+                                 kSPIR_TsanDeviceGlobalMetadata, Result);
+      return Result;
+    }
+
+    for (size_t i = 0; i < NumOfDeviceGlobal; i++) {
+      const auto &GVInfo = GVInfos[i];
+      auto AI = TsanAllocInfo{GVInfo.Addr, GVInfo.Size};
+      ContextInfo->insertAllocInfo(Device, std::move(AI));
+    }
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t TsanInterceptor::insertContext(ur_context_handle_t Context,
                                            std::shared_ptr<ContextInfo> &CI) {
   std::scoped_lock<ur_shared_mutex> Guard(m_ContextMapMutex);
@@ -225,9 +270,10 @@ TsanInterceptor::updateShadowMemory(std::shared_ptr<ContextInfo> &CI,
                                     ur_queue_handle_t Queue) {
   std::scoped_lock<ur_shared_mutex> Guard(CI->AllocInfosMapMutex);
   for (auto &AllocInfo : CI->AllocInfosMap[DI->Handle]) {
-    UR_CALL(DI->Shadow->CleanShadow(Queue, AllocInfo->AllocBegin,
-                                    AllocInfo->AllocSize));
+    UR_CALL(DI->Shadow->CleanShadow(Queue, AllocInfo.AllocBegin,
+                                    AllocInfo.AllocSize));
   }
+  CI->AllocInfosMap[DI->Handle].clear();
   return UR_RESULT_SUCCESS;
 }
 

source/loader/layers/sanitizer/tsan/tsan_interceptor.hpp

@@ -48,8 +48,7 @@ struct ContextInfo {
   std::vector<ur_device_handle_t> DeviceList;
 
   ur_shared_mutex AllocInfosMapMutex;
-  std::unordered_map<ur_device_handle_t,
-                     std::vector<std::shared_ptr<TsanAllocInfo>>>
+  std::unordered_map<ur_device_handle_t, std::vector<TsanAllocInfo>>
       AllocInfosMap;
 
   explicit ContextInfo(ur_context_handle_t Context) : Handle(Context) {
@@ -68,8 +67,12 @@ struct ContextInfo {
 
   ContextInfo &operator=(const ContextInfo &) = delete;
 
-  void insertAllocInfo(ur_device_handle_t Device,
-                       std::shared_ptr<TsanAllocInfo> &AI);
+  void insertAllocInfo(ur_device_handle_t Device, TsanAllocInfo AI);
+};
+
+struct DeviceGlobalInfo {
+  uptr Size;
+  uptr Addr;
 };
 
 struct TsanRuntimeDataWrapper {
@@ -132,6 +135,8 @@ class TsanInterceptor {
                              ur_usm_pool_handle_t Pool, size_t Size,
                              AllocType Type, void **ResultPtr);
 
+  ur_result_t registerProgram(ur_program_handle_t Program);
+
   ur_result_t insertContext(ur_context_handle_t Context,
                             std::shared_ptr<ContextInfo> &CI);
 
@@ -168,6 +173,8 @@ class TsanInterceptor {
                             ur_queue_handle_t Queue, ur_kernel_handle_t Kernel,
                             LaunchInfo &LaunchInfo);
 
+  ur_result_t registerDeviceGlobals(ur_program_handle_t Program);
+
 private:
   std::unordered_map<ur_context_handle_t, std::shared_ptr<ContextInfo>>
       m_ContextMap;

source/loader/layers/sanitizer/tsan/tsan_libdevice.hpp

@@ -94,6 +94,10 @@ struct TsanRuntimeData {
   TsanErrorReport Report[TSAN_MAX_NUM_REPORTS];
 };
 
+constexpr auto kSPIR_TsanDeviceGlobalMetadata = "__TsanDeviceGlobalMetadata";
+
+constexpr auto kSPIR_TsanSpirKernelMetadata = "__TsanKernelMetadata";
+
 #if !defined(__SPIR__) && !defined(__SPIRV__)
 } // namespace ur_sanitizer_layer
 #endif // !__SPIR__ && !__SPIRV__