From f1e47113ce41de9ad5ef931492ff25770ab021d5 Mon Sep 17 00:00:00 2001 From: "Zhao, Maosu" Date: Wed, 8 Jan 2025 22:23:06 -0800 Subject: [PATCH] Record usm allocated size (cherry picked from commit 5c5e9b95bc8f259bb4ff81279acaa51fa29276fe) --- .../layers/sanitizer/msan/msan_buffer.cpp | 6 +-- .../loader/layers/sanitizer/msan/msan_ddi.cpp | 41 ++++++++++++++++++- .../sanitizer/msan/msan_interceptor.cpp | 35 ++++++++++++++-- .../sanitizer/msan/msan_interceptor.hpp | 3 +- .../layers/sanitizer/msan/msan_libdevice.hpp | 2 +- 5 files changed, 76 insertions(+), 11 deletions(-) diff --git a/source/loader/layers/sanitizer/msan/msan_buffer.cpp b/source/loader/layers/sanitizer/msan/msan_buffer.cpp index 66ebb10326..b93fd5ed9a 100644 --- a/source/loader/layers/sanitizer/msan/msan_buffer.cpp +++ b/source/loader/layers/sanitizer/msan/msan_buffer.cpp @@ -93,7 +93,7 @@ ur_result_t MemBuffer::getHandle(ur_device_handle_t Device, char *&Handle) { USMDesc.align = getAlignment(); ur_usm_pool_handle_t Pool{}; URes = getMsanInterceptor()->allocateMemory( - Context, Device, &USMDesc, Pool, Size, + Context, Device, &USMDesc, Pool, Size, AllocType::DEVICE_USM, ur_cast(&Allocation)); if (URes != UR_RESULT_SUCCESS) { getContext()->logger.error( @@ -130,8 +130,8 @@ ur_result_t MemBuffer::getHandle(ur_device_handle_t Device, char *&Handle) { ur_usm_desc_t USMDesc{}; USMDesc.align = getAlignment(); ur_usm_pool_handle_t Pool{}; - URes = getMsanInterceptor()->allocateMemory( - Context, nullptr, &USMDesc, Pool, Size, + URes = getContext()->urDdiTable.USM.pfnHostAlloc( + Context, &USMDesc, Pool, Size, ur_cast(&HostAllocation)); if (URes != UR_RESULT_SUCCESS) { getContext()->logger.error("Failed to allocate {} bytes host " diff --git a/source/loader/layers/sanitizer/msan/msan_ddi.cpp b/source/loader/layers/sanitizer/msan/msan_ddi.cpp index 2dfeadc358..3dc9630b38 100644 --- a/source/loader/layers/sanitizer/msan/msan_ddi.cpp +++ b/source/loader/layers/sanitizer/msan/msan_ddi.cpp @@ -99,8 +99,45 @@ ur_result_t urUSMDeviceAlloc( ) { getContext()->logger.debug("==== urUSMDeviceAlloc"); - return getMsanInterceptor()->allocateMemory(hContext, hDevice, pUSMDesc, - pool, size, ppMem); + return getMsanInterceptor()->allocateMemory( + hContext, hDevice, pUSMDesc, pool, size, AllocType::DEVICE_USM, ppMem); +} + +/////////////////////////////////////////////////////////////////////////////// +/// @brief Intercept function for urUSMHostAlloc +__urdlllocal ur_result_t UR_APICALL urUSMHostAlloc( + ur_context_handle_t hContext, ///< [in] handle of the context object + const ur_usm_desc_t + *pUSMDesc, ///< [in][optional] USM memory allocation descriptor + ur_usm_pool_handle_t + pool, ///< [in][optional] Pointer to a pool created using urUSMPoolCreate + size_t + size, ///< [in] size in bytes of the USM memory object to be allocated + void **ppMem ///< [out] pointer to USM host memory object +) { + getContext()->logger.debug("==== urUSMHostAlloc"); + + return getMsanInterceptor()->allocateMemory( + hContext, nullptr, pUSMDesc, pool, size, AllocType::HOST_USM, ppMem); +} + +/////////////////////////////////////////////////////////////////////////////// +/// @brief Intercept function for urUSMSharedAlloc +__urdlllocal ur_result_t UR_APICALL urUSMSharedAlloc( + ur_context_handle_t hContext, ///< [in] handle of the context object + ur_device_handle_t hDevice, ///< [in] handle of the device object + const ur_usm_desc_t * + pUSMDesc, ///< [in][optional] Pointer to USM memory allocation descriptor. + ur_usm_pool_handle_t + pool, ///< [in][optional] Pointer to a pool created using urUSMPoolCreate + size_t + size, ///< [in] size in bytes of the USM memory object to be allocated + void **ppMem ///< [out] pointer to USM shared memory object +) { + getContext()->logger.debug("==== urUSMSharedAlloc"); + + return getMsanInterceptor()->allocateMemory( + hContext, hDevice, pUSMDesc, pool, size, AllocType::SHARED_USM, ppMem); } /////////////////////////////////////////////////////////////////////////////// diff --git a/source/loader/layers/sanitizer/msan/msan_interceptor.cpp b/source/loader/layers/sanitizer/msan/msan_interceptor.cpp index 8f2855343e..8dc474eb65 100644 --- a/source/loader/layers/sanitizer/msan/msan_interceptor.cpp +++ b/source/loader/layers/sanitizer/msan/msan_interceptor.cpp @@ -46,18 +46,36 @@ ur_result_t MsanInterceptor::allocateMemory(ur_context_handle_t Context, ur_device_handle_t Device, const ur_usm_desc_t *Properties, ur_usm_pool_handle_t Pool, - size_t Size, void **ResultPtr) { + size_t Size, AllocType Type, + void **ResultPtr) { auto ContextInfo = getContextInfo(Context); - std::shared_ptr DeviceInfo = getDeviceInfo(Device); + std::shared_ptr DeviceInfo = + Device ? getDeviceInfo(Device) : nullptr; void *Allocated = nullptr; - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, Properties, Pool, Size, &Allocated)); + if (Type == AllocType::DEVICE_USM) { + UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( + Context, Device, Properties, Pool, Size, &Allocated)); + } else if (Type == AllocType::HOST_USM) { + UR_CALL(getContext()->urDdiTable.USM.pfnHostAlloc( + Context, Properties, Pool, Size, &Allocated)); + } else if (Type == AllocType::SHARED_USM) { + UR_CALL(getContext()->urDdiTable.USM.pfnSharedAlloc( + Context, Device, Properties, Pool, Size, &Allocated)); + } *ResultPtr = Allocated; + ContextInfo->MaxAllocatedSize = + std::max(ContextInfo->MaxAllocatedSize, Size); + + // For host/shared usm, we only record the alloc size. + if (Type != AllocType::DEVICE_USM) { + return UR_RESULT_SUCCESS; + } + auto AI = std::make_shared(MsanAllocInfo{(uptr)Allocated, Size, @@ -432,10 +450,14 @@ ur_result_t MsanInterceptor::prepareLaunch( } // Set LaunchInfo + auto ContextInfo = getContextInfo(LaunchInfo.Context); LaunchInfo.Data->GlobalShadowOffset = DeviceInfo->Shadow->ShadowBegin; LaunchInfo.Data->GlobalShadowOffsetEnd = DeviceInfo->Shadow->ShadowEnd; LaunchInfo.Data->DeviceTy = DeviceInfo->Type; LaunchInfo.Data->Debug = getOptions().Debug ? 1 : 0; + UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( + ContextInfo->Handle, DeviceInfo->Handle, nullptr, nullptr, + ContextInfo->MaxAllocatedSize, &LaunchInfo.Data->CleanShadow)); getContext()->logger.info( "launch_info {} (GlobalShadow={}, Device={}, Debug={})", @@ -518,6 +540,11 @@ ur_result_t USMLaunchInfo::initialize() { USMLaunchInfo::~USMLaunchInfo() { [[maybe_unused]] ur_result_t Result; if (Data) { + if (Data->CleanShadow) { + Result = getContext()->urDdiTable.USM.pfnFree(Context, + Data->CleanShadow); + assert(Result == UR_RESULT_SUCCESS); + } Result = getContext()->urDdiTable.USM.pfnFree(Context, (void *)Data); assert(Result == UR_RESULT_SUCCESS); } diff --git a/source/loader/layers/sanitizer/msan/msan_interceptor.hpp b/source/loader/layers/sanitizer/msan/msan_interceptor.hpp index 7252f1e870..fea52741f3 100644 --- a/source/loader/layers/sanitizer/msan/msan_interceptor.hpp +++ b/source/loader/layers/sanitizer/msan/msan_interceptor.hpp @@ -121,6 +121,7 @@ struct ProgramInfo { struct ContextInfo { ur_context_handle_t Handle; + size_t MaxAllocatedSize = 1024; std::atomic RefCount = 1; std::vector DeviceList; @@ -179,7 +180,7 @@ class MsanInterceptor { ur_device_handle_t Device, const ur_usm_desc_t *Properties, ur_usm_pool_handle_t Pool, size_t Size, - void **ResultPtr); + AllocType Type, void **ResultPtr); ur_result_t releaseMemory(ur_context_handle_t Context, void *Ptr); ur_result_t registerProgram(ur_program_handle_t Program); diff --git a/source/loader/layers/sanitizer/msan/msan_libdevice.hpp b/source/loader/layers/sanitizer/msan/msan_libdevice.hpp index 32e8f36552..0888c9dc75 100644 --- a/source/loader/layers/sanitizer/msan/msan_libdevice.hpp +++ b/source/loader/layers/sanitizer/msan/msan_libdevice.hpp @@ -53,7 +53,7 @@ struct MsanLaunchInfo { MsanErrorReport Report; - uint8_t CleanShadow[128] = {}; + void *CleanShadow = nullptr; }; // Based on the observation, only the last 24 bits of the address of the private