From 0b783c515baeeabd0bcd90af935dac3ecf03bb0e Mon Sep 17 00:00:00 2001 From: "Wu, Yingcong" Date: Wed, 8 Jan 2025 09:01:06 +0100 Subject: [PATCH] save --- source/loader/CMakeLists.txt | 2 + .../loader/layers/sanitizer/asan/asan_ddi.cpp | 10 +- .../sanitizer/asan/asan_interceptor.cpp | 6 +- .../sanitizer/asan/asan_interceptor.hpp | 7 +- .../layers/sanitizer/asan/asan_shadow.cpp | 13 +- .../layers/sanitizer/ur_obj_handler.cpp | 0 .../layers/sanitizer/ur_obj_handler.hpp | 116 ++++++++++++++++++ .../layers/sanitizer/ur_sanitizer_layer.hpp | 2 + 8 files changed, 145 insertions(+), 11 deletions(-) create mode 100644 source/loader/layers/sanitizer/ur_obj_handler.cpp create mode 100644 source/loader/layers/sanitizer/ur_obj_handler.hpp diff --git a/source/loader/CMakeLists.txt b/source/loader/CMakeLists.txt index a10e99f422..2f0efcd6f2 100644 --- a/source/loader/CMakeLists.txt +++ b/source/loader/CMakeLists.txt @@ -179,6 +179,8 @@ if(UR_ENABLE_SANITIZER) ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/ur_sanddi.cpp ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/ur_sanitizer_layer.cpp ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/ur_sanitizer_layer.hpp + ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/ur_obj_handler.cpp + ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/ur_obj_handler.hpp ) if(UR_ENABLE_SYMBOLIZER) diff --git a/source/loader/layers/sanitizer/asan/asan_ddi.cpp b/source/loader/layers/sanitizer/asan/asan_ddi.cpp index bf4dff157a..2c6a9667ae 100644 --- a/source/loader/layers/sanitizer/asan/asan_ddi.cpp +++ b/source/loader/layers/sanitizer/asan/asan_ddi.cpp @@ -509,6 +509,7 @@ __urdlllocal ur_result_t UR_APICALL urContextCreate( pfnCreate(numDevices, phDevices, pProperties, phContext); if (result == UR_RESULT_SUCCESS) { + getContext()->objectHandler.add(*phContext); UR_CALL(setupContext(*phContext, numDevices, phDevices)); } @@ -543,6 +544,7 @@ __urdlllocal ur_result_t UR_APICALL urContextCreateWithNativeHandle( phDevices, pProperties, phContext); if (result == UR_RESULT_SUCCESS) { + getContext()->objectHandler.add(*phContext); UR_CALL(setupContext(*phContext, numDevices, phDevices)); } @@ -563,7 +565,8 @@ __urdlllocal ur_result_t UR_APICALL urContextRetain( getContext()->logger.debug("==== urContextRetain"); - UR_CALL(pfnRetain(hContext)); + // UR_CALL(pfnRetain(hContext)); + UR_CALL(getContext()->objectHandler.retain(hContext)); auto ContextInfo = getAsanInterceptor()->getContextInfo(hContext); UR_ASSERT(ContextInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE); @@ -585,7 +588,8 @@ __urdlllocal ur_result_t UR_APICALL urContextRelease( getContext()->logger.debug("==== urContextRelease"); - UR_CALL(pfnRelease(hContext)); + // UR_CALL(pfnRelease(hContext)); + UR_CALL(getContext()->objectHandler.release(hContext)); auto ContextInfo = getAsanInterceptor()->getContextInfo(hContext); UR_ASSERT(ContextInfo != nullptr, UR_RESULT_ERROR_INVALID_VALUE); @@ -2037,6 +2041,8 @@ ur_result_t initAsanDDITable(ur_dditable_t *dditable) { getContext()->logger.always("==== DeviceSanitizer: ASAN"); + getContext()->objectHandler.installDdiTable(dditable); + if (UR_RESULT_SUCCESS == result) { result = ur_sanitizer_layer::asan::urGetGlobalProcAddrTable( UR_API_VERSION_CURRENT, &dditable->Global); diff --git a/source/loader/layers/sanitizer/asan/asan_interceptor.cpp b/source/loader/layers/sanitizer/asan/asan_interceptor.cpp index 31c35201de..596c3154c5 100644 --- a/source/loader/layers/sanitizer/asan/asan_interceptor.cpp +++ b/source/loader/layers/sanitizer/asan/asan_interceptor.cpp @@ -896,7 +896,8 @@ ContextInfo::~ContextInfo() { assert(URes == UR_RESULT_SUCCESS); } - URes = getContext()->urDdiTable.Context.pfnRelease(Handle); + // URes = getContext()->urDdiTable.Context.pfnRelease(Handle); + URes = getContext()->objectHandler.release(Handle); assert(URes == UR_RESULT_SUCCESS); // check memory leaks @@ -944,7 +945,8 @@ AsanRuntimeDataWrapper::~AsanRuntimeDataWrapper() { LaunchInfo::~LaunchInfo() { [[maybe_unused]] ur_result_t Result; - Result = getContext()->urDdiTable.Context.pfnRelease(Context); + // Result = getContext()->urDdiTable.Context.pfnRelease(Context); + Result = getContext()->objectHandler.release(Context); assert(Result == UR_RESULT_SUCCESS); Result = getContext()->urDdiTable.Device.pfnRelease(Device); assert(Result == UR_RESULT_SUCCESS); diff --git a/source/loader/layers/sanitizer/asan/asan_interceptor.hpp b/source/loader/layers/sanitizer/asan/asan_interceptor.hpp index 27d5e37532..6772d9b6e5 100644 --- a/source/loader/layers/sanitizer/asan/asan_interceptor.hpp +++ b/source/loader/layers/sanitizer/asan/asan_interceptor.hpp @@ -148,8 +148,7 @@ struct ContextInfo { AsanStatsWrapper Stats; explicit ContextInfo(ur_context_handle_t Context) : Handle(Context) { - [[maybe_unused]] auto Result = - getContext()->urDdiTable.Context.pfnRetain(Context); + [[maybe_unused]] auto Result = getContext()->objectHandler.retain(Handle); assert(Result == UR_RESULT_SUCCESS); } @@ -252,9 +251,9 @@ struct LaunchInfo { this->LocalWorkSize = std::vector(LocalWorkSize, LocalWorkSize + WorkDim); } - [[maybe_unused]] auto Result = - getContext()->urDdiTable.Context.pfnRetain(Context); + [[maybe_unused]] auto Result = getContext()->objectHandler.retain(Context); assert(Result == UR_RESULT_SUCCESS); + Result = getContext()->urDdiTable.Device.pfnRetain(Device); assert(Result == UR_RESULT_SUCCESS); } diff --git a/source/loader/layers/sanitizer/asan/asan_shadow.cpp b/source/loader/layers/sanitizer/asan/asan_shadow.cpp index 145fd232c1..cf42322d04 100644 --- a/source/loader/layers/sanitizer/asan/asan_shadow.cpp +++ b/source/loader/layers/sanitizer/asan/asan_shadow.cpp @@ -22,6 +22,7 @@ namespace asan { std::shared_ptr GetShadowMemory(ur_context_handle_t Context, ur_device_handle_t Device, DeviceType Type) { + getContext()->objectHandler.use(Context); if (Type == DeviceType::CPU) { static std::shared_ptr ShadowCPU = std::make_shared(Context, Device); @@ -109,6 +110,8 @@ ur_result_t ShadowMemoryGPU::Setup() { // the SVM range, so that GFX driver will automatically switch to reservation on the GPU // heap. const void *StartAddress = (void *)(0x100'0000'0000'0000ULL); + + getContext()->objectHandler.use(Context); // TODO: Protect Bad Zone auto Result = getContext()->urDdiTable.VirtualMem.pfnReserve( Context, StartAddress, ShadowSize, (void **)&ShadowBegin); @@ -120,7 +123,8 @@ ur_result_t ShadowMemoryGPU::Setup() { } ShadowEnd = ShadowBegin + ShadowSize; // Retain the context which reserves shadow memory - getContext()->urDdiTable.Context.pfnRetain(Context); + // getContext()->urDdiTable.Context.pfnRetain(Context); + getContext()->objectHandler.retain(Context); // Set shadow memory for null pointer // For GPU, wu use up to 1 page of shadow memory @@ -147,6 +151,7 @@ ur_result_t ShadowMemoryGPU::Destory() { } static ur_result_t Result = [this]() { + getContext()->objectHandler.use(Context); const size_t PageSize = GetVirtualMemGranularity(Context, Device); for (auto [MappedPtr, PhysicalMem] : VirtualMemMaps) { UR_CALL(getContext()->urDdiTable.VirtualMem.pfnUnmap( @@ -156,7 +161,8 @@ ur_result_t ShadowMemoryGPU::Destory() { } UR_CALL(getContext()->urDdiTable.VirtualMem.pfnFree( Context, (const void *)ShadowBegin, GetShadowSize())); - UR_CALL(getContext()->urDdiTable.Context.pfnRelease(Context)); + // UR_CALL(getContext()->urDdiTable.Context.pfnRelease(Context)); + UR_CALL(getContext()->objectHandler.release(Context)); return UR_RESULT_SUCCESS; }(); if (!Result) { @@ -171,7 +177,8 @@ ur_result_t ShadowMemoryGPU::Destory() { if (ShadowBegin != 0) { UR_CALL(getContext()->urDdiTable.VirtualMem.pfnFree( Context, (const void *)ShadowBegin, GetShadowSize())); - UR_CALL(getContext()->urDdiTable.Context.pfnRelease(Context)); + // UR_CALL(getContext()->urDdiTable.Context.pfnRelease(Context)); + UR_CALL(getContext()->objectHandler.release(Context)); ShadowBegin = ShadowEnd = 0; } return UR_RESULT_SUCCESS; diff --git a/source/loader/layers/sanitizer/ur_obj_handler.cpp b/source/loader/layers/sanitizer/ur_obj_handler.cpp new file mode 100644 index 0000000000..e69de29bb2 diff --git a/source/loader/layers/sanitizer/ur_obj_handler.hpp b/source/loader/layers/sanitizer/ur_obj_handler.hpp new file mode 100644 index 0000000000..15d948aa5a --- /dev/null +++ b/source/loader/layers/sanitizer/ur_obj_handler.hpp @@ -0,0 +1,116 @@ +/* +The UrObjectHandler is intend to provide a global maps for all UrObjects and corresponding XXXInfo objects that used with sanitizer layers. +Also, it provides a checker that checks for UrObjects' status to avoid any use-after-released cases. +*/ + +/** + * 20250107: first we impl this as a checker to check for use-after-released cases. + */ + +#include "ur/ur.hpp" +#include "ur_api.h" + +#include +#include +#include +#include + +#pragma once + +namespace ur_sanitizer_layer { + +typedef std::variant UrObjectT; + +class UrObjectHandler { + public: + void add(UrObjectT urObject) { + std::scoped_lock Guard(urObjectStatusMapMutex); + // if (urObjectStatusMap.find(urObject) != urObjectStatusMap.end()) { + // if (urObjectStatusMap[urObject].refCount > 0) { + // assert(false && "Add of a exist object"); + // } else { + // // remove an old object + // ; // Nothing to do for now as we only do ref-counting + // } + // } + assert(urObjectStatusMap.find(urObject) == urObjectStatusMap.end() && + "Add of a exist object"); + std::ignore = urObjectStatusMap[urObject]; + } + + ur_result_t retain(UrObjectT urObject) { + assert(ddiTableInstalled && "DdiTable is not installed"); + assert(urObjectStatusMap.find(urObject) != urObjectStatusMap.end() && + "Retain of a nonexistent object"); + urObjectStatusMap[urObject].retain(); + + if (std::holds_alternative(urObject)) { + return urDdiTable.Context.pfnRetain( + std::get(urObject)); + } else if (std::holds_alternative(urObject)) { + return urDdiTable.Device.pfnRetain( + std::get(urObject)); + } + assert(false && "Abonomal object type"); + return UR_RESULT_SUCCESS; + } + + ur_result_t release(UrObjectT urObject) { + assert(ddiTableInstalled && "DdiTable is not installed"); + assert(urObjectStatusMap.find(urObject) != urObjectStatusMap.end() && + "Release of a nonexistent object"); + urObjectStatusMap[urObject].release(); + if (urObjectStatusMap[urObject].refCount == 0) { + std::scoped_lock Guard(urObjectStatusMapMutex); + urObjectStatusMap.erase(urObject); + } + + if (std::holds_alternative(urObject)) { + return urDdiTable.Context.pfnRelease( + std::get(urObject)); + } else if (std::holds_alternative(urObject)) { + return urDdiTable.Device.pfnRelease( + std::get(urObject)); + } + assert(false && "Abonomal object type"); + return UR_RESULT_SUCCESS; + } + + bool use(UrObjectT urObject) { + assert(urObjectStatusMap.find(urObject) != urObjectStatusMap.end() && + "Use of a nonexistent object"); + return urObjectStatusMap[urObject].use(); + } + + void installDdiTable(ur_dditable_t *dditable) { + urDdiTable = *dditable; + ddiTableInstalled = true; + } + + ~UrObjectHandler() { + // Check for not released objects + } + + private: + struct UrObjectInfo { + UrObjectInfo() : refCount(1) {} + std::atomic refCount; + + void retain() { refCount += 1; } + void release() { + assert(refCount > 0 && "Release of a invalid object"); + refCount -= 1; + } + bool use() { + assert(refCount > 0 && "Use of a invalid object"); + return refCount > 0; + } + }; + + ur_dditable_t urDdiTable; + bool ddiTableInstalled = false; + ur_shared_mutex urObjectStatusMapMutex; + std::unordered_map urObjectStatusMap; +}; + +} // namespace ur_sanitizer_layer diff --git a/source/loader/layers/sanitizer/ur_sanitizer_layer.hpp b/source/loader/layers/sanitizer/ur_sanitizer_layer.hpp index 55291db82a..5c12bfccf0 100644 --- a/source/loader/layers/sanitizer/ur_sanitizer_layer.hpp +++ b/source/loader/layers/sanitizer/ur_sanitizer_layer.hpp @@ -14,6 +14,7 @@ #include "logger/ur_logger.hpp" #include "ur_proxy_layer.hpp" +#include "ur_obj_handler.hpp" #define SANITIZER_COMP_NAME "sanitizer layer" @@ -33,6 +34,7 @@ class __urdlllocal context_t : public proxy_layer_context_t, ur_dditable_t urDdiTable = {}; logger::Logger logger; SanitizerType enabledType = SanitizerType::None; + UrObjectHandler objectHandler; context_t(); ~context_t();