From 88cae728c6a063429c2ed1eebf6c60931b224090 Mon Sep 17 00:00:00 2001 From: "Zhao, Maosu" Date: Mon, 1 Dec 2025 04:25:17 +0100 Subject: [PATCH 1/2] [DevSAN] Replace memory allocation with 'SafeAllocate' Report warning to user if device memory is out of resource. --- .../layers/sanitizer/asan/asan_interceptor.cpp | 18 ++---------------- .../layers/sanitizer/asan/asan_interceptor.hpp | 6 +++--- .../layers/sanitizer/asan/asan_shadow.cpp | 17 ++++++++--------- .../layers/sanitizer/msan/msan_interceptor.cpp | 8 ++++---- .../layers/sanitizer/msan/msan_interceptor.hpp | 6 +++--- .../layers/sanitizer/msan/msan_shadow.cpp | 17 ++++++++--------- .../sanitizer_common/sanitizer_allocator.cpp | 15 ++++++++++++--- .../layers/sanitizer/tsan/tsan_interceptor.cpp | 18 +++++------------- .../layers/sanitizer/tsan/tsan_shadow.cpp | 5 ++--- 9 files changed, 47 insertions(+), 63 deletions(-) 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 75de16d79a8a3..cfc4d31c40617 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp @@ -98,22 +98,8 @@ ur_result_t AsanInterceptor::allocateMemory(ur_context_handle_t Context, Pool = ContextInfo->getUSMPool(); } - if (Type == AllocType::DEVICE_USM) { - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, Properties, Pool, NeededSize, &Allocated)); - } else if (Type == AllocType::HOST_USM) { - UR_CALL(getContext()->urDdiTable.USM.pfnHostAlloc(Context, Properties, Pool, - NeededSize, &Allocated)); - } else if (Type == AllocType::SHARED_USM) { - UR_CALL(getContext()->urDdiTable.USM.pfnSharedAlloc( - Context, Device, Properties, Pool, NeededSize, &Allocated)); - } else if (Type == AllocType::MEM_BUFFER) { - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, Properties, Pool, NeededSize, &Allocated)); - } else { - UR_LOG_L(getContext()->logger, ERR, "Unsupport memory type"); - return UR_RESULT_ERROR_INVALID_ARGUMENT; - } + UR_CALL(SafeAllocate(Context, Device, NeededSize, Properties, Pool, Type, + &Allocated)); // Udpate statistics ContextInfo->Stats.UpdateUSMMalloced(NeededSize, NeededSize - Size); 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 bee76e839e1cc..a29786956e895 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.hpp @@ -232,9 +232,9 @@ struct AsanRuntimeDataWrapper { Host.NumLocalArgs = LocalArgs.size(); const size_t LocalArgsInfoSize = sizeof(LocalArgsInfo) * Host.NumLocalArgs; - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, LocalArgsInfoSize, - ur_cast(&Host.LocalArgs))); + UR_CALL(SafeAllocate(Context, Device, LocalArgsInfoSize, nullptr, nullptr, + AllocType::DEVICE_USM, + ur_cast(&Host.LocalArgs))); UR_CALL(getContext()->urDdiTable.Enqueue.pfnUSMMemcpy( Queue, true, Host.LocalArgs, &LocalArgs[0], LocalArgsInfoSize, 0, 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 b784eb228a975..a6018f37d9487 100644 --- a/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/asan/asan_shadow.cpp @@ -266,9 +266,8 @@ ur_result_t ShadowMemoryGPU::AllocLocalShadow(ur_queue_handle_t Queue, LastAllocedSize = 0; } - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, RequiredShadowSize, - (void **)&LocalShadowOffset)); + UR_CALL(SafeAllocate(Context, Device, RequiredShadowSize, nullptr, nullptr, + AllocType::DEVICE_USM, (void **)&LocalShadowOffset)); // Initialize shadow memory ur_result_t URes = EnqueueUSMSet(Queue, (void *)LocalShadowOffset, (char)0, @@ -312,9 +311,9 @@ ur_result_t ShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue, ur_usm_desc_t PrivateBaseProps{UR_STRUCTURE_TYPE_USM_DESC, nullptr, UR_USM_ADVICE_FLAG_DEFAULT, sizeof(uptr)}; - UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, &PrivateBaseProps, nullptr, NewPrivateBaseSize, - (void **)&PrivateBasePtr)); + UR_CALL_THROWS(SafeAllocate( + Context, Device, NewPrivateBaseSize, &PrivateBaseProps, nullptr, + AllocType::DEVICE_USM, (void **)&PrivateBasePtr)); // No need to clean the shadow base, their should be set by work item on // launch @@ -337,9 +336,9 @@ ur_result_t ShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue, LastPrivateShadowAllocedSize = 0; } - UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, NewPrivateShadowSize, - (void **)&PrivateShadowOffset)); + UR_CALL_THROWS(SafeAllocate(Context, Device, NewPrivateShadowSize, + nullptr, nullptr, AllocType::DEVICE_USM, + (void **)&PrivateShadowOffset)); LastPrivateShadowAllocedSize = NewPrivateShadowSize; UR_CALL_THROWS(EnqueueUSMSet(Queue, (void *)PrivateShadowOffset, (char)0, NewPrivateShadowSize)); diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp index 806fd9f9638b5..cf9ae6c106acf 100644 --- a/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp @@ -498,10 +498,10 @@ ur_result_t MsanInterceptor::prepareLaunch( // Clean shadow // Its content is always zero, and is used for unsupport memory types - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - ContextInfo->Handle, DeviceInfo->Handle, nullptr, nullptr, - ContextInfo->CleanShadowSize, - (void **)&LaunchInfo.Data.Host.CleanShadow)); + UR_CALL(SafeAllocate(ContextInfo->Handle, DeviceInfo->Handle, + ContextInfo->CleanShadowSize, nullptr, nullptr, + AllocType::DEVICE_USM, + (void **)&LaunchInfo.Data.Host.CleanShadow)); UR_CALL(EnqueueUSMSet(Queue, (void *)LaunchInfo.Data.Host.CleanShadow, (char)0, ContextInfo->CleanShadowSize, 0, nullptr, nullptr)); diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.hpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.hpp index 37708e5ca80b7..66fd010df5e7f 100644 --- a/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.hpp +++ b/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.hpp @@ -207,9 +207,9 @@ struct MsanRuntimeDataWrapper { Host.NumLocalArgs = LocalArgs.size(); const size_t LocalArgsInfoSize = sizeof(MsanLocalArgsInfo) * Host.NumLocalArgs; - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, LocalArgsInfoSize, - ur_cast(&Host.LocalArgs))); + UR_CALL(SafeAllocate(Context, Device, LocalArgsInfoSize, nullptr, nullptr, + AllocType::DEVICE_USM, + ur_cast(&Host.LocalArgs))); UR_CALL(getContext()->urDdiTable.Enqueue.pfnUSMMemcpy( Queue, true, Host.LocalArgs, &LocalArgs[0], LocalArgsInfoSize, 0, diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp index c61d2c66fc7d2..1d7bf3e257059 100644 --- a/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp @@ -387,9 +387,8 @@ ur_result_t MsanShadowMemoryGPU::AllocLocalShadow(ur_queue_handle_t Queue, LastAllocedSize = 0; } - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, RequiredShadowSize, - (void **)&LocalShadowOffset)); + UR_CALL(SafeAllocate(Context, Device, RequiredShadowSize, nullptr, nullptr, + AllocType::DEVICE_USM, (void **)&LocalShadowOffset)); // Initialize shadow memory ur_result_t URes = EnqueueUSMSet(Queue, (void *)LocalShadowOffset, (char)0, @@ -433,9 +432,9 @@ ur_result_t MsanShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue, ur_usm_desc_t PrivateBaseProps{UR_STRUCTURE_TYPE_USM_DESC, nullptr, UR_USM_ADVICE_FLAG_DEFAULT, sizeof(uptr)}; - UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, &PrivateBaseProps, nullptr, NewPrivateBaseSize, - (void **)&PrivateBasePtr)); + UR_CALL_THROWS(SafeAllocate( + Context, Device, NewPrivateBaseSize, &PrivateBaseProps, nullptr, + AllocType::DEVICE_USM, (void **)&PrivateBasePtr)); // No need to clean the shadow base, their should be set by work item on // launch @@ -454,9 +453,9 @@ ur_result_t MsanShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue, LastPrivateShadowAllocedSize = 0; } - UR_CALL_THROWS(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, NewPrivateShadowSize, - (void **)&PrivateShadowOffset)); + UR_CALL_THROWS(SafeAllocate(Context, Device, NewPrivateShadowSize, + nullptr, nullptr, AllocType::DEVICE_USM, + (void **)&PrivateShadowOffset)); LastPrivateShadowAllocedSize = NewPrivateShadowSize; UR_CALL_THROWS(EnqueueUSMSet(Queue, (void *)PrivateShadowOffset, (char)0, NewPrivateShadowSize)); diff --git a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp b/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp index 2e6de1eac7f28..584c11770e343 100644 --- a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp @@ -50,11 +50,20 @@ ur_result_t SafeAllocate(ur_context_handle_t Context, ur_device_handle_t Device, Device ? GetDeviceType(Context, Device) : DeviceType::UNKNOWN; switch (Type) { case AllocType::DEVICE_USM: - case AllocType::MEM_BUFFER: - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, Properties, Pool, Size, Allocated)); + case AllocType::MEM_BUFFER: { + auto UrRes = getContext()->urDdiTable.USM.pfnDeviceAlloc( + Context, Device, Properties, Pool, Size, Allocated); + if (UrRes != UR_RESULT_SUCCESS) { + if (UrRes == UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY) + UR_LOG_L(getContext()->logger, WARN, + "Device memory out of resource. Please consider reducing the " + "application or running the application with MPI on multi " + "devices."); + return UrRes; + } validateDeviceUSM((uptr)*Allocated, DevieType); break; + } case AllocType::HOST_USM: UR_CALL(getContext()->urDdiTable.USM.pfnHostAlloc(Context, Properties, Pool, Size, Allocated)); diff --git a/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.cpp b/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.cpp index d872025b7d9a9..90904e404850a 100644 --- a/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_interceptor.cpp @@ -78,9 +78,9 @@ ur_result_t TsanRuntimeDataWrapper::importLocalArgsInfo( Host.NumLocalArgs = LocalArgs.size(); const size_t LocalArgsInfoSize = sizeof(TsanLocalArgsInfo) * Host.NumLocalArgs; - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, LocalArgsInfoSize, - ur_cast(&Host.LocalArgs))); + UR_CALL(SafeAllocate(Context, Device, LocalArgsInfoSize, nullptr, nullptr, + AllocType::DEVICE_USM, + ur_cast(&Host.LocalArgs))); UR_CALL(getContext()->urDdiTable.Enqueue.pfnUSMMemcpy( Queue, true, Host.LocalArgs, &LocalArgs[0], LocalArgsInfoSize, 0, nullptr, @@ -140,16 +140,8 @@ ur_result_t TsanInterceptor::allocateMemory(ur_context_handle_t Context, void *Allocated = nullptr; - 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)); - } + UR_CALL( + SafeAllocate(Context, Device, Size, Properties, Pool, Type, &Allocated)); auto AI = TsanAllocInfo{reinterpret_cast(Allocated), Size}; // For updating shadow memory diff --git a/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_shadow.cpp b/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_shadow.cpp index 236f2765ef223..269c57a86696a 100644 --- a/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_shadow.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/tsan/tsan_shadow.cpp @@ -203,9 +203,8 @@ ur_result_t ShadowMemoryGPU::AllocLocalShadow(ur_queue_handle_t Queue, LastAllocatedSize = 0; } - UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc( - Context, Device, nullptr, nullptr, RequiredShadowSize, - (void **)&LocalShadowOffset)); + UR_CALL(SafeAllocate(Context, Device, RequiredShadowSize, nullptr, nullptr, + AllocType::DEVICE_USM, (void **)&LocalShadowOffset)); // Initialize shadow memory ur_result_t URes = From 2c84d8eb2d33848e0bab2e53c2e9c41d24d53dc6 Mon Sep 17 00:00:00 2001 From: "Zhao, Maosu" Date: Wed, 10 Dec 2025 08:21:55 +0100 Subject: [PATCH 2/2] update the messages --- .../layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp b/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp index 584c11770e343..b0f315627f61b 100644 --- a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp +++ b/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_allocator.cpp @@ -56,7 +56,7 @@ ur_result_t SafeAllocate(ur_context_handle_t Context, ur_device_handle_t Device, if (UrRes != UR_RESULT_SUCCESS) { if (UrRes == UR_RESULT_ERROR_OUT_OF_DEVICE_MEMORY) UR_LOG_L(getContext()->logger, WARN, - "Device memory out of resource. Please consider reducing the " + "Out of device memory. Please consider reducing the " "application or running the application with MPI on multi " "devices."); return UrRes;