diff --git a/unified-runtime/source/loader/layers/sanitizer/asan/asan_ddi.cpp b/unified-runtime/source/loader/layers/sanitizer/asan/asan_ddi.cpp index 13a59d0757d1b..04f637fabdbca 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_ddi.cpp +++ b/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(*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(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 struct NotSupportedApi; template @@ -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); diff --git a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp b/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp index 03907ed29e3fe..2dc7474f4060d 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp +++ b/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{Addr, + Addr, + Addr + Size, + Size, + AllocType::DEVICE_USM, + false, + Context, + Device, + GetCurrentBacktrace(), + {}}); + + DI->insertAllocInfo(AI); + + // For memory release + { + std::scoped_lock 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 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, } // 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; diff --git a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp b/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp index d2f7ef198cca7..3d0b85fd634b0 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp +++ b/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); diff --git a/unified-runtime/source/loader/layers/sanitizer/asan/asan_libdevice.hpp b/unified-runtime/source/loader/layers/sanitizer/asan/asan_libdevice.hpp index 1eeff52e8ee0d..5d6fc34a789bc 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_libdevice.hpp +++ b/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"; diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp index 18690b94bf375..4b70ad5da28ca 100644 --- a/unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp +++ b/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 DI = getMsanInterceptor()->getDeviceInfo(hDevice); + ManagedQueue Queue(hContext, hDevice); + UR_CALL(DI->Shadow->EnqueuePoisonShadow(Queue, reinterpret_cast(*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); diff --git a/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_ddi.cpp b/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_ddi.cpp index 4439c7b6f97d5..266639e68fba2 100644 --- a/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_ddi.cpp +++ b/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(*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(pMem); + for (const auto &Device : CI->DeviceList) { + auto DI = getTsanInterceptor()->getDeviceInfo(Device); + std::scoped_lock 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);