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 f5862efe58f12..8165482e78ca0 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_ddi.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_ddi.cpp @@ -29,34 +29,57 @@ ur_result_t setupContext(ur_context_handle_t Context, uint32_t numDevices, std::shared_ptr CI; UR_CALL(getAsanInterceptor()->insertContext(Context, CI)); - if (numDevices > 0) { - auto DeviceType = GetDeviceType(Context, phDevices[0]); - auto ShadowMemory = - getAsanInterceptor()->getOrCreateShadowMemory(phDevices[0], DeviceType); - - for (uint32_t i = 0; i < numDevices; ++i) { - auto hDevice = phDevices[i]; - std::shared_ptr DI; - UR_CALL(getAsanInterceptor()->insertDevice(hDevice, DI)); - DI->Type = GetDeviceType(Context, hDevice); - if (DI->Type == DeviceType::UNKNOWN) { + if (!numDevices) + return UR_RESULT_SUCCESS; + + for (uint32_t i = 0; i < numDevices; ++i) { + ur_device_handle_t hDevice = phDevices[i]; + DeviceInfo &DI = getAsanInterceptor()->getDeviceInfo(hDevice); + + if (!DI.Handle) { + DI.Handle = hDevice; + + DI.Type = GetDeviceType(Context, hDevice); + if (DI.Type == DeviceType::UNKNOWN) { getContext()->logger.error("Unsupport device"); return UR_RESULT_ERROR_INVALID_DEVICE; } - if (DI->Type != DeviceType) { + + DI.IsSupportSharedSystemUSM = GetDeviceUSMCapability( + hDevice, UR_DEVICE_INFO_USM_SYSTEM_SHARED_SUPPORT); + + // Query alignment + getContext()->urDdiTable.Device.pfnGetInfo( + hDevice, UR_DEVICE_INFO_MEM_BASE_ADDR_ALIGN, sizeof(DI.Alignment), + &DI.Alignment, nullptr); + + DI.Shadow = + getAsanInterceptor()->getOrCreateShadowMemory(hDevice, DI.Type); + + getContext()->logger.info( + "DeviceInfo {} (Type={}, IsSupportSharedSystemUSM={}, Alignment={})", + (void *)DI.Handle, ToString(DI.Type), DI.IsSupportSharedSystemUSM, + DI.Alignment); + } + + getContext()->logger.info("Add device({}) into context({})", + (void *)DI.Handle, (void *)Context); + CI->DeviceList.emplace_back(hDevice); + CI->AllocInfosMap[hDevice]; + } + + // Make sure all device type are same + { + DeviceInfo &DI1 = getAsanInterceptor()->getDeviceInfo(phDevices[0]); + for (uint32_t i = 1; i < numDevices; ++i) { + DeviceInfo &DI2 = getAsanInterceptor()->getDeviceInfo(phDevices[i]); + if (DI1.Type != DI2.Type) { getContext()->logger.error("Different device type in the same context"); return UR_RESULT_ERROR_INVALID_DEVICE; } - getContext()->logger.info( - "DeviceInfo {} (Type={}, IsSupportSharedSystemUSM={})", - (void *)DI->Handle, ToString(DI->Type), DI->IsSupportSharedSystemUSM); - getContext()->logger.info("Add {} into context {}", (void *)DI->Handle, - (void *)Context); - DI->Shadow = ShadowMemory; - CI->DeviceList.emplace_back(hDevice); - CI->AllocInfosMap[hDevice]; } } + return UR_RESULT_SUCCESS; } 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 7447887e95741..6699b9cefe5fb 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp @@ -21,6 +21,7 @@ #include "sanitizer_common/sanitizer_options.hpp" #include "sanitizer_common/sanitizer_stacktrace.hpp" #include "sanitizer_common/sanitizer_utils.hpp" +#include "ur_api.h" namespace ur_sanitizer_layer { namespace asan { @@ -35,8 +36,8 @@ AsanInterceptor::AsanInterceptor() { AsanInterceptor::~AsanInterceptor() { // 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 = nullptr; + for (auto &[_, DeviceInfo] : m_DeviceMap) { + DeviceInfo.Shadow = nullptr; } m_Quarantine = nullptr; @@ -72,8 +73,7 @@ ur_result_t AsanInterceptor::allocateMemory(ur_context_handle_t Context, void **ResultPtr) { auto ContextInfo = getContextInfo(Context); - std::shared_ptr DeviceInfo = - Device ? getDeviceInfo(Device) : nullptr; + DeviceInfo *DeviceInfo = Device ? &getDeviceInfo(Device) : nullptr; /// Modified from llvm/compiler-rt/lib/asan/asan_allocator.cpp uint32_t Alignment = Properties ? Properties->align : 0; @@ -267,7 +267,7 @@ ur_result_t AsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel, auto Context = GetContext(Queue); auto Device = GetDevice(Queue); auto ContextInfo = getContextInfo(Context); - auto DeviceInfo = getDeviceInfo(Device); + auto &DeviceInfo = getDeviceInfo(Device); ManagedQueue InternalQueue(Context, Device); if (!InternalQueue) { @@ -319,23 +319,15 @@ ur_result_t AsanInterceptor::postLaunchKernel(ur_kernel_handle_t Kernel, } std::shared_ptr -AsanInterceptor::getOrCreateShadowMemory(ur_device_handle_t Device, - DeviceType Type) { +AsanInterceptor::getOrCreateShadowMemory(ur_device_handle_t Device) { std::scoped_lock Guard(m_ShadowMapMutex); - if (m_ShadowMap.find(Type) == m_ShadowMap.end()) { - ur_context_handle_t InternalContext; - auto Res = getContext()->urDdiTable.Context.pfnCreate(1, &Device, nullptr, - &InternalContext); - if (Res != UR_RESULT_SUCCESS) { - getContext()->logger.error("Failed to create shadow context"); - return nullptr; - } - std::shared_ptr CI; - insertContext(InternalContext, CI); - m_ShadowMap[Type] = GetShadowMemory(InternalContext, Device, Type); - m_ShadowMap[Type]->Setup(); + auto Shadow = GetShadowMemory(Device); + auto URes = Shadow->Setup(); + if (URes != UR_RESULT_SUCCESS) { + die("Failed to setup shadow memory"); } - return m_ShadowMap[Type]; + m_ShadowMap[Device] = Shadow; + return Shadow; } /// Each 8 bytes of application memory are mapped into one byte of shadow memory @@ -346,10 +338,9 @@ AsanInterceptor::getOrCreateShadowMemory(ur_device_handle_t Device, /// /// ref: /// https://github.com/google/sanitizers/wiki/AddressSanitizerAlgorithm#mapping -ur_result_t -AsanInterceptor::enqueueAllocInfo(std::shared_ptr &DeviceInfo, - ur_queue_handle_t Queue, - std::shared_ptr &AI) { +ur_result_t AsanInterceptor::enqueueAllocInfo(DeviceInfo &DeviceInfo, + ur_queue_handle_t Queue, + std::shared_ptr &AI) { if (AI->IsReleased) { int ShadowByte; switch (AI->Type) { @@ -369,14 +360,14 @@ AsanInterceptor::enqueueAllocInfo(std::shared_ptr &DeviceInfo, ShadowByte = 0xff; assert(false && "Unknow AllocInfo Type"); } - UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(Queue, AI->AllocBegin, - AI->AllocSize, ShadowByte)); + UR_CALL(DeviceInfo.Shadow->EnqueuePoisonShadow(Queue, AI->AllocBegin, + AI->AllocSize, ShadowByte)); return UR_RESULT_SUCCESS; } // Init zero - UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(Queue, AI->AllocBegin, - AI->AllocSize, 0)); + UR_CALL(DeviceInfo.Shadow->EnqueuePoisonShadow(Queue, AI->AllocBegin, + AI->AllocSize, 0)); uptr TailBegin = RoundUpTo(AI->UserEnd, ASAN_SHADOW_GRANULARITY); uptr TailEnd = AI->AllocBegin + AI->AllocSize; @@ -385,8 +376,8 @@ AsanInterceptor::enqueueAllocInfo(std::shared_ptr &DeviceInfo, if (TailBegin != AI->UserEnd) { auto Value = AI->UserEnd - RoundDownTo(AI->UserEnd, ASAN_SHADOW_GRANULARITY); - UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow(Queue, AI->UserEnd, 1, - static_cast(Value))); + UR_CALL(DeviceInfo.Shadow->EnqueuePoisonShadow(Queue, AI->UserEnd, 1, + static_cast(Value))); } int ShadowByte; @@ -412,11 +403,11 @@ AsanInterceptor::enqueueAllocInfo(std::shared_ptr &DeviceInfo, } // Left red zone - UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow( + UR_CALL(DeviceInfo.Shadow->EnqueuePoisonShadow( Queue, AI->AllocBegin, AI->UserBegin - AI->AllocBegin, ShadowByte)); // Right red zone - UR_CALL(DeviceInfo->Shadow->EnqueuePoisonShadow( + UR_CALL(DeviceInfo.Shadow->EnqueuePoisonShadow( Queue, TailBegin, TailEnd - TailBegin, ShadowByte)); return UR_RESULT_SUCCESS; @@ -424,9 +415,9 @@ AsanInterceptor::enqueueAllocInfo(std::shared_ptr &DeviceInfo, ur_result_t AsanInterceptor::updateShadowMemory(std::shared_ptr &ContextInfo, - std::shared_ptr &DeviceInfo, + DeviceInfo &DeviceInfo, ur_queue_handle_t Queue) { - auto &AllocInfos = ContextInfo->AllocInfosMap[DeviceInfo->Handle]; + auto &AllocInfos = ContextInfo->AllocInfosMap[DeviceInfo.Handle]; std::scoped_lock Guard(AllocInfos.Mutex); for (auto &AI : AllocInfos.List) { @@ -615,32 +606,12 @@ ur_result_t AsanInterceptor::eraseContext(ur_context_handle_t Context) { return UR_RESULT_SUCCESS; } -ur_result_t AsanInterceptor::insertDevice(ur_device_handle_t Device, - std::shared_ptr &DI) { +DeviceInfo &AsanInterceptor::getDeviceInfo(ur_device_handle_t Device) { std::scoped_lock Guard(m_DeviceMapMutex); - - if (m_DeviceMap.find(Device) != m_DeviceMap.end()) { - DI = m_DeviceMap.at(Device); - return UR_RESULT_SUCCESS; - } - - DI = std::make_shared(Device); - - DI->IsSupportSharedSystemUSM = - GetDeviceUSMCapability(Device, UR_DEVICE_INFO_USM_SYSTEM_SHARED_SUPPORT); - - // Query alignment - UR_CALL(getContext()->urDdiTable.Device.pfnGetInfo( - Device, UR_DEVICE_INFO_MEM_BASE_ADDR_ALIGN, sizeof(DI->Alignment), - &DI->Alignment, nullptr)); - - // Don't move DI, since it's a return value as well - m_DeviceMap.emplace(Device, DI); - - return UR_RESULT_SUCCESS; + return m_DeviceMap[Device]; } -ur_result_t AsanInterceptor::eraseDevice(ur_device_handle_t Device) { +ur_result_t AsanInterceptor::eraseDeviceInfo(ur_device_handle_t Device) { std::scoped_lock Guard(m_DeviceMapMutex); assert(m_DeviceMap.find(Device) != m_DeviceMap.end()); m_DeviceMap.erase(Device); @@ -716,17 +687,18 @@ AsanInterceptor::getMemBuffer(ur_mem_handle_t MemHandle) { return nullptr; } -ur_result_t AsanInterceptor::prepareLaunch( - std::shared_ptr &ContextInfo, - std::shared_ptr &DeviceInfo, ur_queue_handle_t Queue, - ur_kernel_handle_t Kernel, LaunchInfo &LaunchInfo) { +ur_result_t +AsanInterceptor::prepareLaunch(std::shared_ptr &ContextInfo, + DeviceInfo &DeviceInfo, ur_queue_handle_t Queue, + ur_kernel_handle_t Kernel, + LaunchInfo &LaunchInfo) { auto &KernelInfo = getOrCreateKernelInfo(Kernel); std::shared_lock Guard(KernelInfo.Mutex); auto ArgNums = GetKernelNumArgs(Kernel); - auto LocalMemoryUsage = GetKernelLocalMemorySize(Kernel, DeviceInfo->Handle); + auto LocalMemoryUsage = GetKernelLocalMemorySize(Kernel, DeviceInfo.Handle); auto PrivateMemoryUsage = - GetKernelPrivateMemorySize(Kernel, DeviceInfo->Handle); + GetKernelPrivateMemorySize(Kernel, DeviceInfo.Handle); getContext()->logger.info( "KernelInfo {} (Name={}, ArgNums={}, IsInstrumented={}, " @@ -742,7 +714,7 @@ ur_result_t AsanInterceptor::prepareLaunch( continue; } if (auto ValidateResult = ValidateUSMPointer( - ContextInfo->Handle, DeviceInfo->Handle, (uptr)Ptr)) { + ContextInfo->Handle, DeviceInfo.Handle, (uptr)Ptr)) { ReportInvalidKernelArgument(Kernel, ArgIndex, (uptr)Ptr, ValidateResult, PtrPair.second); if (ValidateResult.Type != ValidateUSMResult::MAYBE_HOST_POINTER) { @@ -755,7 +727,7 @@ ur_result_t AsanInterceptor::prepareLaunch( // Set membuffer arguments for (const auto &[ArgIndex, MemBuffer] : KernelInfo.BufferArgs) { char *ArgPointer = nullptr; - UR_CALL(MemBuffer->getHandle(DeviceInfo->Handle, ArgPointer)); + UR_CALL(MemBuffer->getHandle(DeviceInfo.Handle, ArgPointer)); ur_result_t URes = getContext()->urDdiTable.Kernel.pfnSetArgPointer( Kernel, ArgIndex, nullptr, ArgPointer); if (URes != UR_RESULT_SUCCESS) { @@ -809,14 +781,14 @@ ur_result_t AsanInterceptor::prepareLaunch( } // Prepare asan runtime data - LaunchInfo.Data.Host.GlobalShadowOffset = DeviceInfo->Shadow->ShadowBegin; - LaunchInfo.Data.Host.GlobalShadowOffsetEnd = DeviceInfo->Shadow->ShadowEnd; - LaunchInfo.Data.Host.DeviceTy = DeviceInfo->Type; + LaunchInfo.Data.Host.GlobalShadowOffset = DeviceInfo.Shadow->ShadowBegin; + LaunchInfo.Data.Host.GlobalShadowOffsetEnd = DeviceInfo.Shadow->ShadowEnd; + LaunchInfo.Data.Host.DeviceTy = DeviceInfo.Type; LaunchInfo.Data.Host.Debug = getContext()->Options.Debug ? 1 : 0; // Write shadow memory offset for local memory if (getContext()->Options.DetectLocals) { - if (DeviceInfo->Shadow->AllocLocalShadow( + if (DeviceInfo.Shadow->AllocLocalShadow( Queue, NumWG, LaunchInfo.Data.Host.LocalShadowOffset, LaunchInfo.Data.Host.LocalShadowOffsetEnd) != UR_RESULT_SUCCESS) { getContext()->logger.warning( @@ -836,7 +808,7 @@ ur_result_t AsanInterceptor::prepareLaunch( // Write shadow memory offset for private memory if (getContext()->Options.DetectPrivates) { - if (DeviceInfo->Shadow->AllocPrivateShadow( + if (DeviceInfo.Shadow->AllocPrivateShadow( Queue, NumWG, LaunchInfo.Data.Host.PrivateShadowOffset, LaunchInfo.Data.Host.PrivateShadowOffsetEnd) != UR_RESULT_SUCCESS) { getContext()->logger.warning( 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 53aac416f2567..03acebf9960ad 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp @@ -40,7 +40,7 @@ struct AllocInfoList { }; struct DeviceInfo { - ur_device_handle_t Handle; + ur_device_handle_t Handle = nullptr; DeviceType Type = DeviceType::UNKNOWN; size_t Alignment = 0; @@ -53,10 +53,6 @@ struct DeviceInfo { ur_mutex Mutex; std::queue> Quarantine; size_t QuarantineSize = 0; - - // Device handles are special and alive in the whole process lifetime, - // so we needn't retain&release here. - explicit DeviceInfo(ur_device_handle_t Device) : Handle(Device) {} }; struct QueueInfo { @@ -293,10 +289,6 @@ class AsanInterceptor { std::shared_ptr &CI); ur_result_t eraseContext(ur_context_handle_t Context); - ur_result_t insertDevice(ur_device_handle_t Device, - std::shared_ptr &CI); - ur_result_t eraseDevice(ur_device_handle_t Device); - ur_result_t insertProgram(ur_program_handle_t Program); ur_result_t eraseProgram(ur_program_handle_t Program); @@ -325,11 +317,8 @@ class AsanInterceptor { return m_ContextMap[Context]; } - std::shared_ptr getDeviceInfo(ur_device_handle_t Device) { - std::shared_lock Guard(m_DeviceMapMutex); - assert(m_DeviceMap.find(Device) != m_DeviceMap.end()); - return m_DeviceMap[Device]; - } + DeviceInfo &getDeviceInfo(ur_device_handle_t Device); + ur_result_t eraseDeviceInfo(ur_device_handle_t Device); std::shared_ptr getProgramInfo(ur_program_handle_t Program) { std::shared_lock Guard(m_ProgramMapMutex); @@ -350,24 +339,22 @@ class AsanInterceptor { bool isNormalExit() { return m_NormalExit; } std::shared_ptr - getOrCreateShadowMemory(ur_device_handle_t Device, DeviceType Type); + getOrCreateShadowMemory(ur_device_handle_t Device); ur_shared_mutex KernelLaunchMutex; private: ur_result_t updateShadowMemory(std::shared_ptr &ContextInfo, - std::shared_ptr &DeviceInfo, + DeviceInfo &DeviceInfo, ur_queue_handle_t Queue); - ur_result_t enqueueAllocInfo(std::shared_ptr &DeviceInfo, - ur_queue_handle_t Queue, + ur_result_t enqueueAllocInfo(DeviceInfo &DeviceInfo, ur_queue_handle_t Queue, std::shared_ptr &AI); /// Initialize Global Variables & Kernel Name at first Launch ur_result_t prepareLaunch(std::shared_ptr &ContextInfo, - std::shared_ptr &DeviceInfo, - ur_queue_handle_t Queue, ur_kernel_handle_t Kernel, - LaunchInfo &LaunchInfo); + DeviceInfo &DeviceInfo, ur_queue_handle_t Queue, + ur_kernel_handle_t Kernel, LaunchInfo &LaunchInfo); ur_result_t registerDeviceGlobals(ur_program_handle_t Program); ur_result_t registerSpirKernels(ur_program_handle_t Program); @@ -376,8 +363,7 @@ class AsanInterceptor { std::unordered_map> m_ContextMap; ur_shared_mutex m_ContextMapMutex; - std::unordered_map> - m_DeviceMap; + std::unordered_map m_DeviceMap; ur_shared_mutex m_DeviceMapMutex; std::unordered_map> @@ -403,7 +389,8 @@ class AsanInterceptor { bool m_NormalExit = true; - std::unordered_map> m_ShadowMap; + std::unordered_map> + m_ShadowMap; ur_shared_mutex m_ShadowMapMutex; }; diff --git a/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.cpp b/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.cpp index e3c77989a1899..d31dea0fb0c87 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.cpp @@ -20,9 +20,13 @@ namespace ur_sanitizer_layer { namespace asan { -std::shared_ptr GetShadowMemory(ur_context_handle_t Context, - ur_device_handle_t Device, - DeviceType Type) { +std::shared_ptr GetShadowMemory(ur_device_handle_t Device) { + ur_context_handle_t Context; + [[maybe_unused]] auto URes = + getContext()->urDdiTable.Context.pfnCreate(1, &Device, nullptr, &Context); + assert(URes == UR_RESULT_SUCCESS); + + DeviceType Type = GetDeviceType(Context, Device); if (Type == DeviceType::CPU) { static std::shared_ptr ShadowCPU = std::make_shared(Context, Device); diff --git a/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.hpp b/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.hpp index 575210291684f..81cbadb129825 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.hpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.hpp @@ -24,11 +24,7 @@ namespace asan { struct ShadowMemory { ShadowMemory(ur_context_handle_t Context, ur_device_handle_t Device) - : Context(Context), Device(Device) { - [[maybe_unused]] ur_result_t URes = - getContext()->urDdiTable.Device.pfnRetain(Device); - assert(URes == UR_RESULT_SUCCESS); - } + : Context(Context), Device(Device) {} virtual ~ShadowMemory() { [[maybe_unused]] ur_result_t URes = @@ -163,9 +159,7 @@ struct ShadowMemoryDG2 final : public ShadowMemoryGPU { size_t GetShadowSize() override { return 0x100000000000ULL; } }; -std::shared_ptr GetShadowMemory(ur_context_handle_t Context, - ur_device_handle_t Device, - DeviceType Type); +std::shared_ptr GetShadowMemory(ur_device_handle_t Device); } // namespace asan } // namespace ur_sanitizer_layer diff --git a/unified-runtime/source/loader/layers/sanitizer/asan/asan_validator.cpp b/unified-runtime/source/loader/layers/sanitizer/asan/asan_validator.cpp index 2ff16ae0bfd8e..f6a6dd574c71f 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_validator.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_validator.cpp @@ -42,8 +42,8 @@ ValidateUSMResult ValidateUSMPointer(ur_context_handle_t Context, auto AllocInfoItOp = getAsanInterceptor()->findAllocInfoByAddress(Ptr); if (!AllocInfoItOp) { - auto DI = getAsanInterceptor()->getDeviceInfo(Device); - bool IsSupportSharedSystemUSM = DI->IsSupportSharedSystemUSM; + auto &DI = getAsanInterceptor()->getDeviceInfo(Device); + bool IsSupportSharedSystemUSM = DI.IsSupportSharedSystemUSM; if (IsSupportSharedSystemUSM) { // maybe it's host pointer return ValidateUSMResult::success();