[DeviceSanitizer] Add IPC memory support for ASAN, MSAN, and TSAN

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

@@ -1670,6 +1670,49 @@ __urdlllocal ur_result_t UR_APICALL urMemoryExportExportMemoryHandleExp(
   return UR_RESULT_SUCCESS;
 }
 
+/// @brief Intercept function for urIPCOpenMemHandleExp
+__urdlllocal ur_result_t UR_APICALL urIPCOpenMemHandleExp(
+    /// [in] handle of the context object
+    ur_context_handle_t hContext,
+    /// [in] handle of the device object the corresponding USM device memory
+    /// was allocated on
+    ur_device_handle_t hDevice,
+    /// [in] the IPC memory handle data
+    void *pIPCMemHandleData,
+    /// [in] size of the IPC memory handle data
+    size_t ipcMemHandleDataSize,
+    /// [out] pointer to a pointer to device USM memory
+    void **ppMem) {
+  UR_LOG_L(getContext()->logger, DEBUG, "==== urIPCOpenMemHandleExp");
+
+  UR_CALL(getContext()->urDdiTable.IPCExp.pfnOpenMemHandleExp(
+      hContext, hDevice, pIPCMemHandleData, ipcMemHandleDataSize, ppMem));
+
+  size_t MemSize;
+  UR_CALL(getContext()->urDdiTable.USM.pfnGetMemAllocInfo(
+      hContext, *ppMem, UR_USM_ALLOC_INFO_SIZE, sizeof(MemSize), &MemSize,
+      nullptr));
+  UR_CALL(getAsanInterceptor()->registerIPCMemory(
+      hContext, hDevice, reinterpret_cast<uptr>(*ppMem), MemSize));
+
+  return UR_RESULT_SUCCESS;
+}
+
+/// @brief Intercept function for urIPCCloseMemHandleExp
+__urdlllocal ur_result_t UR_APICALL urIPCCloseMemHandleExp(
+    /// [in] handle of the context object
+    ur_context_handle_t hContext,
+    /// [in] pointer to device USM memory opened through urIPCOpenMemHandleExp
+    void *pMem) {
+  UR_LOG_L(getContext()->logger, DEBUG, "==== urIPCCloseMemHandleExp");
+
+  UR_CALL(getContext()->urDdiTable.IPCExp.pfnCloseMemHandleExp(hContext, pMem));
+  UR_CALL(
+      getAsanInterceptor()->unregisterIPCMemory(reinterpret_cast<uptr>(pMem)));
+
+  return UR_RESULT_SUCCESS;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for filling application's Adapter table
 ///        with current process' addresses
@@ -1903,6 +1946,25 @@ ur_result_t urGetEnqueueExpProcAddrTable(
   return result;
 }
 
+/// @brief Exported function for filling application's IPCExp table
+///        with current process' addresses
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+ur_result_t urGetIPCExpProcAddrTable(
+    /// [in,out] pointer to table of DDI function pointers
+    ur_ipc_exp_dditable_t *pDdiTable) {
+  ur_result_t result = UR_RESULT_SUCCESS;
+
+  pDdiTable->pfnOpenMemHandleExp =
+      ur_sanitizer_layer::asan::urIPCOpenMemHandleExp;
+  pDdiTable->pfnCloseMemHandleExp =
+      ur_sanitizer_layer::asan::urIPCCloseMemHandleExp;
+
+  return result;
+}
+
 template <class A, class B> struct NotSupportedApi;
 
 template <class MsgType, class R, class... A>
@@ -2120,6 +2182,11 @@ ur_result_t initAsanDDITable(ur_dditable_t *dditable) {
         &dditable->MemoryExportExp);
   }
 
+  if (UR_RESULT_SUCCESS == result) {
+    result =
+        ur_sanitizer_layer::asan::urGetIPCExpProcAddrTable(&dditable->IPCExp);
+  }
+
   if (result != UR_RESULT_SUCCESS) {
     UR_LOG_L(getContext()->logger, ERR, "Initialize ASAN DDI table failed: {}",
              result);

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

@@ -260,6 +260,49 @@ ur_result_t AsanInterceptor::releaseMemory(ur_context_handle_t Context,
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t AsanInterceptor::registerIPCMemory(ur_context_handle_t Context,
+                                               ur_device_handle_t Device,
+                                               uptr Addr, size_t Size) {
+  auto CI = getContextInfo(Context);
+  auto DI = getDeviceInfo(Device);
+
+  auto AI = std::make_shared<AllocInfo>(AllocInfo{Addr,
+                                                  Addr,
+                                                  Addr + Size,
+                                                  Size,
+                                                  AllocType::DEVICE_USM,
+                                                  false,
+                                                  Context,
+                                                  Device,
+                                                  GetCurrentBacktrace(),
+                                                  {}});
+
+  DI->insertAllocInfo(AI);
+
+  // For memory release
+  {
+    std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
+    m_AllocationMap.emplace(AI->AllocBegin, std::move(AI));
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
+ur_result_t AsanInterceptor::unregisterIPCMemory(uptr Addr) {
+  auto AllocInfoItOp = findAllocInfoByAddress(Addr);
+  if (AllocInfoItOp.has_value()) {
+    auto AllocInfo = AllocInfoItOp.value()->second;
+    AllocInfo->IsReleased = true;
+    AllocInfo->ReleaseStack = GetCurrentBacktrace();
+    getDeviceInfo(AllocInfo->Device)->insertAllocInfo(AllocInfo);
+
+    std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
+    m_AllocationMap.erase(*AllocInfoItOp);
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t AsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
                                              ur_queue_handle_t Queue,
                                              LaunchInfo &LaunchInfo) {
@@ -369,9 +412,12 @@ AsanInterceptor::enqueueAllocInfo(std::shared_ptr<DeviceInfo> &DeviceInfo,
   }
 
   // Init zero
-  static const int8_t Zero = 0;
   UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(Queue, AI->AllocBegin,
-                                                  AI->AllocSize, &Zero));
+                                                  AI->AllocSize, &kZeroMagic));
+
+  // If no redzones to poison, so we're done.
+  if (AI->UserEnd == (AI->AllocBegin + AI->AllocSize))
+    return UR_RESULT_SUCCESS;
 
   uptr TailBegin = RoundUpTo(AI->UserEnd, ASAN_SHADOW_GRANULARITY);
   uptr TailEnd = AI->AllocBegin + AI->AllocSize;

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

@@ -305,6 +305,12 @@ class AsanInterceptor {
 
   ur_result_t releaseMemory(ur_context_handle_t Context, void *Ptr);
 
+  ur_result_t registerIPCMemory(ur_context_handle_t Context,
+                                ur_device_handle_t Device, uptr Addr,
+                                size_t Size);
+
+  ur_result_t unregisterIPCMemory(uptr Addr);
+
   ur_result_t registerProgram(ur_program_handle_t Program);
 
   ur_result_t unregisterProgram(ur_program_handle_t Program);

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

@@ -111,6 +111,8 @@ const int8_t kPrivateRightRedzoneMagic = (int8_t)0xf3;
 // Unknown shadow value
 constexpr int8_t kUnknownMagic = (int8_t)0xff;
 
+constexpr int8_t kZeroMagic = 0;
+
 constexpr auto kSPIR_AsanDeviceGlobalMetadata = "__AsanDeviceGlobalMetadata";
 constexpr auto kSPIR_AsanSpirKernelMetadata = "__AsanKernelMetadata";
 

unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp

@@ -1829,6 +1829,37 @@ __urdlllocal ur_result_t UR_APICALL urMemoryExportAllocExportableMemoryExp(
   return UR_RESULT_SUCCESS;
 }
 
+/// @brief Intercept function for urIPCOpenMemHandleExp
+__urdlllocal ur_result_t UR_APICALL urIPCOpenMemHandleExp(
+    /// [in] handle of the context object
+    ur_context_handle_t hContext,
+    /// [in] handle of the device object the corresponding USM device memory
+    /// was allocated on
+    ur_device_handle_t hDevice,
+    /// [in] the IPC memory handle data
+    void *pIPCMemHandleData,
+    /// [in] size of the IPC memory handle data
+    size_t ipcMemHandleDataSize,
+    /// [out] pointer to a pointer to device USM memory
+    void **ppMem) {
+  UR_LOG_L(getContext()->logger, DEBUG, "==== urIPCOpenMemHandleExp");
+
+  UR_CALL(getContext()->urDdiTable.IPCExp.pfnOpenMemHandleExp(
+      hContext, hDevice, pIPCMemHandleData, ipcMemHandleDataSize, ppMem));
+
+  // Mark IPC memory as initialized in shadow memory
+  size_t MemSize;
+  UR_CALL(getContext()->urDdiTable.USM.pfnGetMemAllocInfo(
+      hContext, *ppMem, UR_USM_ALLOC_INFO_SIZE, sizeof(MemSize), &MemSize,
+      nullptr));
+  std::shared_ptr<DeviceInfo> DI = getMsanInterceptor()->getDeviceInfo(hDevice);
+  ManagedQueue Queue(hContext, hDevice);
+  UR_CALL(DI->Shadow->EnqueuePoisonShadow(Queue, reinterpret_cast<uptr>(*ppMem),
+                                          MemSize, &kMemInitializedMagic));
+
+  return UR_RESULT_SUCCESS;
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Exported function for filling application's Adapter table
 ///        with current process' addresses
@@ -2042,6 +2073,23 @@ urGetMemoryExportExpProcAddrTable(ur_memory_export_exp_dditable_t *pDdiTable) {
   return UR_RESULT_SUCCESS;
 }
 
+/// @brief Exported function for filling application's IPCExp table
+///        with current process' addresses
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+ur_result_t urGetIPCExpProcAddrTable(
+    /// [in,out] pointer to table of DDI function pointers
+    ur_ipc_exp_dditable_t *pDdiTable) {
+  ur_result_t result = UR_RESULT_SUCCESS;
+
+  pDdiTable->pfnOpenMemHandleExp =
+      ur_sanitizer_layer::msan::urIPCOpenMemHandleExp;
+
+  return result;
+}
+
 ur_result_t urCheckVersion(ur_api_version_t version) {
   if (UR_MAJOR_VERSION(ur_sanitizer_layer::getContext()->version) !=
           UR_MAJOR_VERSION(version) ||
@@ -2116,6 +2164,11 @@ ur_result_t initMsanDDITable(ur_dditable_t *dditable) {
         &dditable->MemoryExportExp);
   }
 
+  if (UR_RESULT_SUCCESS == result) {
+    result =
+        ur_sanitizer_layer::msan::urGetIPCExpProcAddrTable(&dditable->IPCExp);
+  }
+
   if (result != UR_RESULT_SUCCESS) {
     UR_LOG_L(getContext()->logger, ERR, "Initialize MSAN DDI table failed: {}",
              result);

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

@@ -1381,6 +1381,58 @@ __urdlllocal ur_result_t UR_APICALL urMemoryExportFreeExportableMemoryExp(
       hContext, hDevice, pMem);
 }
 
+/// @brief Intercept function for urIPCOpenMemHandleExp
+__urdlllocal ur_result_t UR_APICALL urIPCOpenMemHandleExp(
+    /// [in] handle of the context object
+    ur_context_handle_t hContext,
+    /// [in] handle of the device object the corresponding USM device memory
+    /// was allocated on
+    ur_device_handle_t hDevice,
+    /// [in] the IPC memory handle data
+    void *pIPCMemHandleData,
+    /// [in] size of the IPC memory handle data
+    size_t ipcMemHandleDataSize,
+    /// [out] pointer to a pointer to device USM memory
+    void **ppMem) {
+  UR_LOG_L(getContext()->logger, DEBUG, "==== urIPCOpenMemHandleExp");
+
+  UR_CALL(getContext()->urDdiTable.IPCExp.pfnOpenMemHandleExp(
+      hContext, hDevice, pIPCMemHandleData, ipcMemHandleDataSize, ppMem));
+
+  size_t MemSize;
+  UR_CALL(getContext()->urDdiTable.USM.pfnGetMemAllocInfo(
+      hContext, *ppMem, UR_USM_ALLOC_INFO_SIZE, sizeof(MemSize), &MemSize,
+      nullptr));
+  auto DI = getTsanInterceptor()->getDeviceInfo(hDevice);
+  DI->insertAllocInfo(TsanAllocInfo{reinterpret_cast<uptr>(*ppMem), MemSize});
+
+  return UR_RESULT_SUCCESS;
+}
+
+/// @brief Intercept function for urIPCCloseMemHandleExp
+__urdlllocal ur_result_t UR_APICALL urIPCCloseMemHandleExp(
+    /// [in] handle of the context object
+    ur_context_handle_t hContext,
+    /// [in] pointer to device USM memory opened through urIPCOpenMemHandleExp
+    void *pMem) {
+  UR_LOG_L(getContext()->logger, DEBUG, "==== urIPCCloseMemHandleExp");
+
+  UR_CALL(getContext()->urDdiTable.IPCExp.pfnCloseMemHandleExp(hContext, pMem));
+
+  auto CI = getTsanInterceptor()->getContextInfo(hContext);
+  auto Addr = reinterpret_cast<uptr>(pMem);
+  for (const auto &Device : CI->DeviceList) {
+    auto DI = getTsanInterceptor()->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; });
+    if (It != DI->AllocInfos.end())
+      DI->AllocInfos.erase(It);
+  }
+
+  return UR_RESULT_SUCCESS;
+}
+
 ur_result_t urCheckVersion(ur_api_version_t version) {
   if (UR_MAJOR_VERSION(ur_sanitizer_layer::getContext()->version) !=
           UR_MAJOR_VERSION(version) ||
@@ -1627,6 +1679,26 @@ urGetMemoryExportExpProcAddrTable(ur_memory_export_exp_dditable_t *pDdiTable) {
 
   return UR_RESULT_SUCCESS;
 }
+
+/// @brief Exported function for filling application's IPCExp table
+///        with current process' addresses
+///
+/// @returns
+///     - ::UR_RESULT_SUCCESS
+///     - ::UR_RESULT_ERROR_INVALID_NULL_POINTER
+ur_result_t urGetIPCExpProcAddrTable(
+    /// [in,out] pointer to table of DDI function pointers
+    ur_ipc_exp_dditable_t *pDdiTable) {
+  ur_result_t result = UR_RESULT_SUCCESS;
+
+  pDdiTable->pfnOpenMemHandleExp =
+      ur_sanitizer_layer::tsan::urIPCOpenMemHandleExp;
+  pDdiTable->pfnCloseMemHandleExp =
+      ur_sanitizer_layer::tsan::urIPCCloseMemHandleExp;
+
+  return result;
+}
+
 } // namespace tsan
 
 ur_result_t initTsanDDITable(ur_dditable_t *dditable) {
@@ -1686,6 +1758,11 @@ ur_result_t initTsanDDITable(ur_dditable_t *dditable) {
         &dditable->MemoryExportExp);
   }
 
+  if (UR_RESULT_SUCCESS == result) {
+    result =
+        ur_sanitizer_layer::tsan::urGetIPCExpProcAddrTable(&dditable->IPCExp);
+  }
+
   if (result != UR_RESULT_SUCCESS) {
     UR_LOG_L(getContext()->logger, ERR, "Initialize TSAN DDI table failed: {}",
              result);