[DeviceSanitizer] Memory overhead statistics

source/loader/CMakeLists.txt

@@ -137,6 +137,8 @@ if(UR_ENABLE_SANITIZER)
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_report.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_shadow_setup.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_shadow_setup.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_statistics.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan_statistics.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/common.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/stacktrace.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/stacktrace.hpp

source/loader/layers/sanitizer/asan_interceptor.cpp

@@ -16,7 +16,9 @@
 #include "asan_quarantine.hpp"
 #include "asan_report.hpp"
 #include "asan_shadow_setup.hpp"
+#include "asan_statistics.hpp"
 #include "stacktrace.hpp"
+#include "ur_sanitizer_layer.hpp"
 #include "ur_sanitizer_utils.hpp"
 
 namespace ur_sanitizer_layer {
@@ -77,6 +79,8 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
         uptr ShadowEnd =
             MemToShadow_PVC(DeviceInfo->ShadowOffset, Ptr + Size - 1);
         assert(ShadowBegin <= ShadowEnd);
+
+        auto &Stats = getContext()->interceptor->getStats();
         {
             static const size_t PageSize =
                 GetVirtualMemGranularity(Context, DeviceInfo->Handle);
@@ -110,11 +114,13 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
                 if (URes != UR_RESULT_SUCCESS) {
                     getContext()->logger.debug("urVirtualMemMap(): {}", URes);
                 }
+                ++Stats.shadow_mmaps;
 
                 // Initialize to zero
                 if (URes == UR_RESULT_SUCCESS) {
                     // Reset PhysicalMem to null since it's been mapped
                     PhysicalMem = nullptr;
+                    Stats.shadow_mmaped += PageSize;
 
                     auto URes =
                         urEnqueueUSMSet(Queue, (void *)MappedPtr, 0, PageSize);
@@ -146,18 +152,26 @@ ur_result_t enqueueMemSetShadow(ur_context_handle_t Context,
 
 } // namespace
 
+ContextInfo::~ContextInfo() {
+    [[maybe_unused]] auto Result =
+        getContext()->urDdiTable.Context.pfnRelease(Handle);
+    assert(Result == UR_RESULT_SUCCESS);
+}
+
 SanitizerInterceptor::SanitizerInterceptor(logger::Logger &logger)
-    : logger(logger) {
-    if (Options(logger).MaxQuarantineSizeMB) {
+    : m_Options(logger) {
+    if (getOptions().MaxQuarantineSizeMB) {
         m_Quarantine = std::make_unique<Quarantine>(
-            static_cast<uint64_t>(Options(logger).MaxQuarantineSizeMB) * 1024 *
+            static_cast<uint64_t>(getOptions().MaxQuarantineSizeMB) * 1024 *
             1024);
     }
 }
 
 SanitizerInterceptor::~SanitizerInterceptor() {
     DestroyShadowMemoryOnCPU();
     DestroyShadowMemoryOnPVC();
+
+    m_Stats.Print();
 }
 
 /// The memory chunk allocated from the underlying allocator looks like this:
@@ -173,6 +187,7 @@ ur_result_t SanitizerInterceptor::allocateMemory(
     AllocType Type, void **ResultPtr) {
 
     auto ContextInfo = getContextInfo(Context);
+
     std::shared_ptr<DeviceInfo> DeviceInfo =
         Device ? getDeviceInfo(Device) : nullptr;
 
@@ -191,8 +206,8 @@ ur_result_t SanitizerInterceptor::allocateMemory(
         Alignment = MinAlignment;
     }
 
-    uptr RZLog = ComputeRZLog(Size, Options(logger).MinRZSize,
-                              Options(logger).MaxRZSize);
+    uptr RZLog =
+        ComputeRZLog(Size, getOptions().MinRZSize, getOptions().MaxRZSize);
     uptr RZSize = RZLog2Size(RZLog);
     uptr RoundedSize = RoundUpTo(Size, Alignment);
     uptr NeededSize = RoundedSize + RZSize * 2;
@@ -219,6 +234,11 @@ ur_result_t SanitizerInterceptor::allocateMemory(
         return UR_RESULT_ERROR_INVALID_ARGUMENT;
     }
 
+    // Udpate statistics
+    ++m_Stats.usm_mallocs;
+    m_Stats.usm_malloced += NeededSize;
+    m_Stats.usm_malloced_redzones += NeededSize - Size;
+
     uptr AllocBegin = reinterpret_cast<uptr>(Allocated);
     [[maybe_unused]] uptr AllocEnd = AllocBegin + NeededSize;
     uptr UserBegin = AllocBegin + RZSize;
@@ -304,11 +324,18 @@ ur_result_t SanitizerInterceptor::releaseMemory(ur_context_handle_t Context,
         ContextInfo->insertAllocInfo({AllocInfo->Device}, AllocInfo);
     }
 
+    ++m_Stats.usm_frees;
+    m_Stats.usm_freed += AllocInfo->AllocSize;
+
     // If quarantine is disabled, USM is freed immediately
     if (!m_Quarantine) {
         getContext()->logger.debug("Free: {}", (void *)AllocInfo->AllocBegin);
         std::scoped_lock<ur_shared_mutex> Guard(m_AllocationMapMutex);
         m_AllocationMap.erase(AllocInfoIt);
+
+        ++m_Stats.usm_real_frees;
+        m_Stats.usm_really_freed += AllocInfo->AllocSize;
+
         return getContext()->urDdiTable.USM.pfnFree(
             Context, (void *)(AllocInfo->AllocBegin));
     }
@@ -320,6 +347,10 @@ ur_result_t SanitizerInterceptor::releaseMemory(ur_context_handle_t Context,
             getContext()->logger.info("Quarantine Free: {}",
                                       (void *)It->second->AllocBegin);
             m_AllocationMap.erase(It);
+
+            ++m_Stats.usm_real_frees;
+            m_Stats.usm_really_freed += It->second->AllocSize;
+
             UR_CALL(getContext()->urDdiTable.USM.pfnFree(
                 Context, (void *)(It->second->AllocBegin)));
         }
@@ -345,8 +376,8 @@ ur_result_t SanitizerInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
         return UR_RESULT_ERROR_INVALID_QUEUE;
     }
 
-    UR_CALL(
-        prepareLaunch(Context, DeviceInfo, InternalQueue, Kernel, LaunchInfo));
+    UR_CALL(prepareLaunch(ContextInfo, DeviceInfo, InternalQueue, Kernel,
+                          LaunchInfo));
 
     UR_CALL(updateShadowMemory(ContextInfo, DeviceInfo, InternalQueue));
 
@@ -654,9 +685,9 @@ SanitizerInterceptor::getMemBuffer(ur_mem_handle_t MemHandle) {
 }
 
 ur_result_t SanitizerInterceptor::prepareLaunch(
-    ur_context_handle_t Context, std::shared_ptr<DeviceInfo> &DeviceInfo,
-    ur_queue_handle_t Queue, ur_kernel_handle_t Kernel,
-    USMLaunchInfo &LaunchInfo) {
+    std::shared_ptr<ContextInfo> &ContextInfo,
+    std::shared_ptr<DeviceInfo> &DeviceInfo, ur_queue_handle_t Queue,
+    ur_kernel_handle_t Kernel, USMLaunchInfo &LaunchInfo) {
     auto Program = GetProgram(Kernel);
 
     do {
@@ -709,7 +740,7 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
 
         // Write debug
         // We use "uint64_t" here because EnqueueWriteGlobal will fail when it's "uint32_t"
-        uint64_t Debug = Options(logger).Debug ? 1 : 0;
+        uint64_t Debug = getOptions().Debug ? 1 : 0;
         EnqueueWriteGlobal(kSPIR_AsanDebug, &Debug, sizeof(Debug));
 
         // Write shadow memory offset for global memory
@@ -750,12 +781,12 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
                      LocalWorkSize[Dim];
         }
 
-        auto EnqueueAllocateShadowMemory = [Context, &DeviceInfo,
+        auto EnqueueAllocateShadowMemory = [Context = ContextInfo->Handle,
+                                            Device = DeviceInfo->Handle,
                                             Queue](size_t Size, uptr &Ptr) {
             void *Allocated = nullptr;
             auto URes = getContext()->urDdiTable.USM.pfnDeviceAlloc(
-                Context, DeviceInfo->Handle, nullptr, nullptr, Size,
-                &Allocated);
+                Context, Device, nullptr, nullptr, Size, &Allocated);
             if (URes != UR_RESULT_SUCCESS) {
                 return URes;
             }
@@ -782,7 +813,7 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
             LocalMemoryUsage, PrivateMemoryUsage);
 
         // Write shadow memory offset for local memory
-        if (Options(logger).DetectLocals) {
+        if (getOptions().DetectLocals) {
             // CPU needn't this
             if (DeviceInfo->Type == DeviceType::GPU_PVC) {
                 const size_t LocalMemorySize =
@@ -812,6 +843,8 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
                         LaunchInfo.Data->LocalShadowOffset +
                         LocalShadowMemorySize - 1;
 
+                    m_Stats.shadow_malloced += LocalShadowMemorySize;
+
                     getContext()->logger.info(
                         "ShadowMemory(Local, {} - {})",
                         (void *)LaunchInfo.Data->LocalShadowOffset,
@@ -821,7 +854,7 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
         }
 
         // Write shadow memory offset for private memory
-        if (Options(logger).DetectPrivates) {
+        if (getOptions().DetectPrivates) {
             if (DeviceInfo->Type == DeviceType::CPU) {
                 LaunchInfo.Data->PrivateShadowOffset = DeviceInfo->ShadowOffset;
             } else if (DeviceInfo->Type == DeviceType::GPU_PVC) {
@@ -848,6 +881,9 @@ ur_result_t SanitizerInterceptor::prepareLaunch(
                     LaunchInfo.Data->PrivateShadowOffsetEnd =
                         LaunchInfo.Data->PrivateShadowOffset +
                         PrivateShadowMemorySize - 1;
+
+                    m_Stats.shadow_malloced += PrivateShadowMemorySize;
+
                     getContext()->logger.info(
                         "ShadowMemory(Private, {} - {})",
                         (void *)LaunchInfo.Data->PrivateShadowOffset,
@@ -905,14 +941,21 @@ ur_result_t USMLaunchInfo::updateKernelInfo(const KernelInfo &KI) {
 USMLaunchInfo::~USMLaunchInfo() {
     [[maybe_unused]] ur_result_t Result;
     if (Data) {
+        auto ContextInfo = getContext()->interceptor->getContextInfo(Context);
+        auto &Stats = getContext()->interceptor->getStats();
+
         auto Type = GetDeviceType(Device);
         if (Type == DeviceType::GPU_PVC) {
             if (Data->PrivateShadowOffset) {
+                Stats.shadow_freed += Data->PrivateShadowOffsetEnd -
+                                      Data->PrivateShadowOffset + 1;
                 Result = getContext()->urDdiTable.USM.pfnFree(
                     Context, (void *)Data->PrivateShadowOffset);
                 assert(Result == UR_RESULT_SUCCESS);
             }
             if (Data->LocalShadowOffset) {
+                Stats.shadow_freed +=
+                    Data->LocalShadowOffsetEnd - Data->LocalShadowOffset + 1;
                 Result = getContext()->urDdiTable.USM.pfnFree(
                     Context, (void *)Data->LocalShadowOffset);
                 assert(Result == UR_RESULT_SUCCESS);

source/loader/layers/sanitizer/asan_interceptor.hpp

@@ -15,6 +15,8 @@
 #include "asan_allocator.hpp"
 #include "asan_buffer.hpp"
 #include "asan_libdevice.hpp"
+#include "asan_options.hpp"
+#include "asan_statistics.hpp"
 #include "common.hpp"
 #include "ur_sanitizer_layer.hpp"
 
@@ -114,11 +116,7 @@ struct ContextInfo {
         assert(Result == UR_RESULT_SUCCESS);
     }
 
-    ~ContextInfo() {
-        [[maybe_unused]] auto Result =
-            getContext()->urDdiTable.Context.pfnRelease(Handle);
-        assert(Result == UR_RESULT_SUCCESS);
-    }
+    ~ContextInfo();
 
     void insertAllocInfo(const std::vector<ur_device_handle_t> &Devices,
                          std::shared_ptr<AllocInfo> &AI) {
@@ -221,6 +219,9 @@ class SanitizerInterceptor {
         return m_KernelMap[Kernel];
     }
 
+    const AsanOptions &getOptions() { return m_Options; }
+    AsanStats& getStats() { return m_Stats; }
+
   private:
     ur_result_t updateShadowMemory(std::shared_ptr<ContextInfo> &ContextInfo,
                                    std::shared_ptr<DeviceInfo> &DeviceInfo,
@@ -231,7 +232,7 @@ class SanitizerInterceptor {
                                  std::shared_ptr<AllocInfo> &AI);
 
     /// Initialize Global Variables & Kernel Name at first Launch
-    ur_result_t prepareLaunch(ur_context_handle_t Context,
+    ur_result_t prepareLaunch(std::shared_ptr<ContextInfo> &ContextInfo,
                               std::shared_ptr<DeviceInfo> &DeviceInfo,
                               ur_queue_handle_t Queue,
                               ur_kernel_handle_t Kernel,
@@ -261,7 +262,9 @@ class SanitizerInterceptor {
     ur_shared_mutex m_AllocationMapMutex;
 
     std::unique_ptr<Quarantine> m_Quarantine;
-    logger::Logger &logger;
+
+    AsanOptions m_Options;
+    AsanStats m_Stats;
 };
 
 } // namespace ur_sanitizer_layer

source/loader/layers/sanitizer/asan_options.hpp

@@ -27,20 +27,20 @@ struct AsanOptions {
     AsanOptions(AsanOptions &other) = delete;
     void operator=(const AsanOptions &) = delete;
 
-    static AsanOptions &getInstance(logger::Logger &logger) {
-        static AsanOptions instance(logger);
-        return instance;
-    }
+    // static AsanOptions &getInstance(logger::Logger &logger) {
+    //     static AsanOptions instance(logger);
+    //     return instance;
+    // }
 
     bool Debug = false;
     uint64_t MinRZSize = 16;
     uint64_t MaxRZSize = 2048;
     uint32_t MaxQuarantineSizeMB = 0;
     bool DetectLocals = true;
     bool DetectPrivates = true;
+    bool PrintStats = false;
 
-  private:
-    AsanOptions(logger::Logger &logger) {
+    explicit AsanOptions(logger::Logger &logger) {
         auto OptionsEnvMap = getenv_to_map("UR_LAYER_ASAN_OPTIONS");
         if (!OptionsEnvMap.has_value()) {
             return;
@@ -93,6 +93,7 @@ struct AsanOptions {
         SetBoolOption("debug", Debug);
         SetBoolOption("detect_locals", DetectLocals);
         SetBoolOption("detect_privates", DetectPrivates);
+        SetBoolOption("print_stats", PrintStats);
 
         auto KV = OptionsEnvMap->find("quarantine_size_mb");
         if (KV != OptionsEnvMap->end()) {
@@ -142,8 +143,8 @@ struct AsanOptions {
     }
 };
 
-inline const AsanOptions &Options(logger::Logger &logger) {
-    return AsanOptions::getInstance(logger);
-}
+// inline const AsanOptions &Options(logger::Logger &logger) {
+//     return AsanOptions::getInstance(logger);
+// }
 
 } // namespace ur_sanitizer_layer

source/loader/layers/sanitizer/asan_report.cpp

@@ -124,7 +124,7 @@ void ReportUseAfterFree(const DeviceSanitizerReport &Report,
     getContext()->logger.always("  #0 {} {}:{}", Func, File, Report.Line);
     getContext()->logger.always("");
 
-    if (Options(getContext()->logger).MaxQuarantineSizeMB > 0) {
+    if (getContext()->interceptor->getOptions().MaxQuarantineSizeMB > 0) {
         auto AllocInfoItOp =
             getContext()->interceptor->findAllocInfoByAddress(Report.Address);