Merge pull request #2325 from AllanZyne/review/yang/restructure_asan_msan

[DeviceMSAN] Support MemorySanitizer for device offloading

Author: kbenzie

source/loader/CMakeLists.txt

@@ -151,6 +151,21 @@ if(UR_ENABLE_SANITIZER)
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan/asan_statistics.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan/asan_validator.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/asan/asan_validator.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_allocator.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_allocator.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_buffer.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_buffer.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_ddi.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_ddi.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_interceptor.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_interceptor.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_libdevice.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_options.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_options.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_report.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_report.hpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_shadow.cpp
+        ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/msan/msan_shadow.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/sanitizer_common/linux/backtrace.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/sanitizer_common/linux/sanitizer_utils.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/sanitizer_common/sanitizer_allocator.hpp
@@ -160,6 +175,7 @@ if(UR_ENABLE_SANITIZER)
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/sanitizer_common/sanitizer_stacktrace.hpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/sanitizer_common/sanitizer_utils.cpp
         ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/sanitizer_common/sanitizer_utils.hpp
+        ${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
     )

source/loader/layers/sanitizer/asan/asan_ddi.cpp

@@ -2053,26 +2053,10 @@ __urdlllocal ur_result_t UR_APICALL urGetVirtualMemProcAddrTable(
 }
 } // namespace asan
 
-ur_result_t context_t::init(ur_dditable_t *dditable,
-                            const std::set<std::string> &enabledLayerNames,
-                            [[maybe_unused]] codeloc_data codelocData) {
+ur_result_t initAsanDDITable(ur_dditable_t *dditable) {
     ur_result_t result = UR_RESULT_SUCCESS;
 
-    if (enabledLayerNames.count("UR_LAYER_ASAN")) {
-        enabledType = SanitizerType::AddressSanitizer;
-        initAsanInterceptor();
-    } else if (enabledLayerNames.count("UR_LAYER_MSAN")) {
-        enabledType = SanitizerType::MemorySanitizer;
-    } else if (enabledLayerNames.count("UR_LAYER_TSAN")) {
-        enabledType = SanitizerType::ThreadSanitizer;
-    }
-
-    // Only support AddressSanitizer now
-    if (enabledType != SanitizerType::AddressSanitizer) {
-        return result;
-    }
-
-    urDdiTable = *dditable;
+    getContext()->logger.always("==== DeviceSanitizer: ASAN");
 
     if (UR_RESULT_SUCCESS == result) {
         result = ur_sanitizer_layer::asan::urGetGlobalProcAddrTable(
@@ -2134,6 +2118,11 @@ ur_result_t context_t::init(ur_dditable_t *dditable,
             UR_API_VERSION_CURRENT, &dditable->VirtualMem);
     }
 
+    if (result != UR_RESULT_SUCCESS) {
+        getContext()->logger.error("Initialize ASAN DDI table failed: {}",
+                                   result);
+    }
+
     return result;
 }
 

source/loader/layers/sanitizer/asan/asan_ddi.hpp

@@ -17,4 +17,6 @@ namespace ur_sanitizer_layer {
 void initAsanInterceptor();
 void destroyAsanInterceptor();
 
+ur_result_t initAsanDDITable(ur_dditable_t *dditable);
+
 } // namespace ur_sanitizer_layer

source/loader/layers/sanitizer/asan/asan_interceptor.cpp

@@ -676,162 +676,155 @@ ur_result_t AsanInterceptor::prepareLaunch(
     std::shared_ptr<DeviceInfo> &DeviceInfo, ur_queue_handle_t Queue,
     ur_kernel_handle_t Kernel, LaunchInfo &LaunchInfo) {
 
-    do {
-        auto KernelInfo = getKernelInfo(Kernel);
-        assert(KernelInfo && "Kernel should be instrumented");
-
-        // Validate pointer arguments
-        if (getOptions().DetectKernelArguments) {
-            for (const auto &[ArgIndex, PtrPair] : KernelInfo->PointerArgs) {
-                auto Ptr = PtrPair.first;
-                if (Ptr == nullptr) {
-                    continue;
-                }
-                if (auto ValidateResult = ValidateUSMPointer(
-                        ContextInfo->Handle, DeviceInfo->Handle, (uptr)Ptr)) {
-                    ReportInvalidKernelArgument(Kernel, ArgIndex, (uptr)Ptr,
-                                                ValidateResult, PtrPair.second);
-                    exitWithErrors();
-                }
+    auto KernelInfo = getKernelInfo(Kernel);
+    assert(KernelInfo && "Kernel should be instrumented");
+
+    // Validate pointer arguments
+    if (getOptions().DetectKernelArguments) {
+        for (const auto &[ArgIndex, PtrPair] : KernelInfo->PointerArgs) {
+            auto Ptr = PtrPair.first;
+            if (Ptr == nullptr) {
+                continue;
+            }
+            if (auto ValidateResult = ValidateUSMPointer(
+                    ContextInfo->Handle, DeviceInfo->Handle, (uptr)Ptr)) {
+                ReportInvalidKernelArgument(Kernel, ArgIndex, (uptr)Ptr,
+                                            ValidateResult, PtrPair.second);
+                exitWithErrors();
             }
         }
+    }
 
-        // Set membuffer arguments
-        for (const auto &[ArgIndex, MemBuffer] : KernelInfo->BufferArgs) {
-            char *ArgPointer = nullptr;
-            UR_CALL(MemBuffer->getHandle(DeviceInfo->Handle, ArgPointer));
-            ur_result_t URes = getContext()->urDdiTable.Kernel.pfnSetArgPointer(
-                Kernel, ArgIndex, nullptr, ArgPointer);
-            if (URes != UR_RESULT_SUCCESS) {
-                getContext()->logger.error(
-                    "Failed to set buffer {} as the {} arg to kernel {}: {}",
-                    ur_cast<ur_mem_handle_t>(MemBuffer.get()), ArgIndex, Kernel,
-                    URes);
-            }
+    // Set membuffer arguments
+    for (const auto &[ArgIndex, MemBuffer] : KernelInfo->BufferArgs) {
+        char *ArgPointer = nullptr;
+        UR_CALL(MemBuffer->getHandle(DeviceInfo->Handle, ArgPointer));
+        ur_result_t URes = getContext()->urDdiTable.Kernel.pfnSetArgPointer(
+            Kernel, ArgIndex, nullptr, ArgPointer);
+        if (URes != UR_RESULT_SUCCESS) {
+            getContext()->logger.error(
+                "Failed to set buffer {} as the {} arg to kernel {}: {}",
+                ur_cast<ur_mem_handle_t>(MemBuffer.get()), ArgIndex, Kernel,
+                URes);
         }
+    }
 
-        auto ArgNums = GetKernelNumArgs(Kernel);
-        // We must prepare all kernel args before call
-        // urKernelGetSuggestedLocalWorkSize, otherwise the call will fail on
-        // CPU device.
-        if (ArgNums) {
-            ur_result_t URes = getContext()->urDdiTable.Kernel.pfnSetArgPointer(
-                Kernel, ArgNums - 1, nullptr, LaunchInfo.Data.getDevicePtr());
-            if (URes != UR_RESULT_SUCCESS) {
-                getContext()->logger.error("Failed to set launch info: {}",
-                                           URes);
-                return URes;
-            }
+    auto ArgNums = GetKernelNumArgs(Kernel);
+    // We must prepare all kernel args before call
+    // urKernelGetSuggestedLocalWorkSize, otherwise the call will fail on
+    // CPU device.
+    if (ArgNums) {
+        ur_result_t URes = getContext()->urDdiTable.Kernel.pfnSetArgPointer(
+            Kernel, ArgNums - 1, nullptr, LaunchInfo.Data.getDevicePtr());
+        if (URes != UR_RESULT_SUCCESS) {
+            getContext()->logger.error("Failed to set launch info: {}", URes);
+            return URes;
         }
+    }
 
-        if (LaunchInfo.LocalWorkSize.empty()) {
-            LaunchInfo.LocalWorkSize.resize(LaunchInfo.WorkDim);
-            auto URes =
-                getContext()->urDdiTable.Kernel.pfnGetSuggestedLocalWorkSize(
-                    Kernel, Queue, LaunchInfo.WorkDim,
-                    LaunchInfo.GlobalWorkOffset, LaunchInfo.GlobalWorkSize,
-                    LaunchInfo.LocalWorkSize.data());
-            if (URes != UR_RESULT_SUCCESS) {
-                if (URes != UR_RESULT_ERROR_UNSUPPORTED_FEATURE) {
-                    return URes;
-                }
-                // If urKernelGetSuggestedLocalWorkSize is not supported by driver, we fallback
-                // to inefficient implementation
-                for (size_t Dim = 0; Dim < LaunchInfo.WorkDim; ++Dim) {
-                    LaunchInfo.LocalWorkSize[Dim] = 1;
-                }
+    if (LaunchInfo.LocalWorkSize.empty()) {
+        LaunchInfo.LocalWorkSize.resize(LaunchInfo.WorkDim);
+        auto URes =
+            getContext()->urDdiTable.Kernel.pfnGetSuggestedLocalWorkSize(
+                Kernel, Queue, LaunchInfo.WorkDim, LaunchInfo.GlobalWorkOffset,
+                LaunchInfo.GlobalWorkSize, LaunchInfo.LocalWorkSize.data());
+        if (URes != UR_RESULT_SUCCESS) {
+            if (URes != UR_RESULT_ERROR_UNSUPPORTED_FEATURE) {
+                return URes;
+            }
+            // If urKernelGetSuggestedLocalWorkSize is not supported by driver, we fallback
+            // to inefficient implementation
+            for (size_t Dim = 0; Dim < LaunchInfo.WorkDim; ++Dim) {
+                LaunchInfo.LocalWorkSize[Dim] = 1;
             }
         }
+    }
 
-        const size_t *LocalWorkSize = LaunchInfo.LocalWorkSize.data();
-        uint32_t NumWG = 1;
-        for (uint32_t Dim = 0; Dim < LaunchInfo.WorkDim; ++Dim) {
-            NumWG *= (LaunchInfo.GlobalWorkSize[Dim] + LocalWorkSize[Dim] - 1) /
-                     LocalWorkSize[Dim];
-        }
+    const size_t *LocalWorkSize = LaunchInfo.LocalWorkSize.data();
+    uint32_t NumWG = 1;
+    for (uint32_t Dim = 0; Dim < LaunchInfo.WorkDim; ++Dim) {
+        NumWG *= (LaunchInfo.GlobalWorkSize[Dim] + LocalWorkSize[Dim] - 1) /
+                 LocalWorkSize[Dim];
+    }
 
-        // 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.Debug = getOptions().Debug ? 1 : 0;
-
-        auto LocalMemoryUsage =
-            GetKernelLocalMemorySize(Kernel, DeviceInfo->Handle);
-        auto PrivateMemoryUsage =
-            GetKernelPrivateMemorySize(Kernel, DeviceInfo->Handle);
-
-        getContext()->logger.info(
-            "KernelInfo {} (LocalMemory={}, PrivateMemory={})", (void *)Kernel,
-            LocalMemoryUsage, PrivateMemoryUsage);
-
-        // Write shadow memory offset for local memory
-        if (getOptions().DetectLocals) {
-            if (DeviceInfo->Shadow->AllocLocalShadow(
-                    Queue, NumWG, LaunchInfo.Data.Host.LocalShadowOffset,
-                    LaunchInfo.Data.Host.LocalShadowOffsetEnd) !=
-                UR_RESULT_SUCCESS) {
-                getContext()->logger.warning(
-                    "Failed to allocate shadow memory for local "
-                    "memory, maybe the number of workgroup ({}) is too "
-                    "large",
-                    NumWG);
-                getContext()->logger.warning(
-                    "Skip checking local memory of kernel <{}>",
-                    GetKernelName(Kernel));
-            } else {
-                getContext()->logger.info(
-                    "ShadowMemory(Local, WorkGroup{}, {} - {})", NumWG,
-                    (void *)LaunchInfo.Data.Host.LocalShadowOffset,
-                    (void *)LaunchInfo.Data.Host.LocalShadowOffsetEnd);
-            }
+    // 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.Debug = getOptions().Debug ? 1 : 0;
+
+    auto LocalMemoryUsage =
+        GetKernelLocalMemorySize(Kernel, DeviceInfo->Handle);
+    auto PrivateMemoryUsage =
+        GetKernelPrivateMemorySize(Kernel, DeviceInfo->Handle);
+
+    getContext()->logger.info(
+        "KernelInfo {} (LocalMemory={}, PrivateMemory={})", (void *)Kernel,
+        LocalMemoryUsage, PrivateMemoryUsage);
+
+    // Write shadow memory offset for local memory
+    if (getOptions().DetectLocals) {
+        if (DeviceInfo->Shadow->AllocLocalShadow(
+                Queue, NumWG, LaunchInfo.Data.Host.LocalShadowOffset,
+                LaunchInfo.Data.Host.LocalShadowOffsetEnd) !=
+            UR_RESULT_SUCCESS) {
+            getContext()->logger.warning(
+                "Failed to allocate shadow memory for local "
+                "memory, maybe the number of workgroup ({}) is too "
+                "large",
+                NumWG);
+            getContext()->logger.warning(
+                "Skip checking local memory of kernel <{}>",
+                GetKernelName(Kernel));
+        } else {
+            getContext()->logger.info(
+                "ShadowMemory(Local, WorkGroup{}, {} - {})", NumWG,
+                (void *)LaunchInfo.Data.Host.LocalShadowOffset,
+                (void *)LaunchInfo.Data.Host.LocalShadowOffsetEnd);
         }
+    }
 
-        // Write shadow memory offset for private memory
-        if (getOptions().DetectPrivates) {
-            if (DeviceInfo->Shadow->AllocPrivateShadow(
-                    Queue, NumWG, LaunchInfo.Data.Host.PrivateShadowOffset,
-                    LaunchInfo.Data.Host.PrivateShadowOffsetEnd) !=
-                UR_RESULT_SUCCESS) {
-                getContext()->logger.warning(
-                    "Failed to allocate shadow memory for private "
-                    "memory, maybe the number of workgroup ({}) is too "
-                    "large",
-                    NumWG);
-                getContext()->logger.warning(
-                    "Skip checking private memory of kernel <{}>",
-                    GetKernelName(Kernel));
-            } else {
-                getContext()->logger.info(
-                    "ShadowMemory(Private, WorkGroup{}, {} - {})", NumWG,
-                    (void *)LaunchInfo.Data.Host.PrivateShadowOffset,
-                    (void *)LaunchInfo.Data.Host.PrivateShadowOffsetEnd);
-            }
+    // Write shadow memory offset for private memory
+    if (getOptions().DetectPrivates) {
+        if (DeviceInfo->Shadow->AllocPrivateShadow(
+                Queue, NumWG, LaunchInfo.Data.Host.PrivateShadowOffset,
+                LaunchInfo.Data.Host.PrivateShadowOffsetEnd) !=
+            UR_RESULT_SUCCESS) {
+            getContext()->logger.warning(
+                "Failed to allocate shadow memory for private "
+                "memory, maybe the number of workgroup ({}) is too "
+                "large",
+                NumWG);
+            getContext()->logger.warning(
+                "Skip checking private memory of kernel <{}>",
+                GetKernelName(Kernel));
+        } else {
+            getContext()->logger.info(
+                "ShadowMemory(Private, WorkGroup{}, {} - {})", NumWG,
+                (void *)LaunchInfo.Data.Host.PrivateShadowOffset,
+                (void *)LaunchInfo.Data.Host.PrivateShadowOffsetEnd);
         }
+    }
 
-        // Write local arguments info
-        if (!KernelInfo->LocalArgs.empty()) {
-            std::vector<LocalArgsInfo> LocalArgsInfo;
-            for (auto [ArgIndex, ArgInfo] : KernelInfo->LocalArgs) {
-                LocalArgsInfo.push_back(ArgInfo);
-                getContext()->logger.debug(
-                    "local_args (argIndex={}, size={}, sizeWithRZ={})",
-                    ArgIndex, ArgInfo.Size, ArgInfo.SizeWithRedZone);
-            }
-            UR_CALL(LaunchInfo.Data.importLocalArgsInfo(Queue, LocalArgsInfo));
+    // Write local arguments info
+    if (!KernelInfo->LocalArgs.empty()) {
+        std::vector<LocalArgsInfo> LocalArgsInfo;
+        for (auto [ArgIndex, ArgInfo] : KernelInfo->LocalArgs) {
+            LocalArgsInfo.push_back(ArgInfo);
+            getContext()->logger.debug(
+                "local_args (argIndex={}, size={}, sizeWithRZ={})", ArgIndex,
+                ArgInfo.Size, ArgInfo.SizeWithRedZone);
         }
+        UR_CALL(LaunchInfo.Data.importLocalArgsInfo(Queue, LocalArgsInfo));
+    }
 
-        // sync asan runtime data to device side
-        UR_CALL(LaunchInfo.Data.syncToDevice(Queue));
+    // sync asan runtime data to device side
+    UR_CALL(LaunchInfo.Data.syncToDevice(Queue));
 
-        getContext()->logger.debug(
-            "launch_info {} (numLocalArgs={}, localArgs={})",
-            (void *)LaunchInfo.Data.getDevicePtr(),
-            LaunchInfo.Data.Host.NumLocalArgs,
-            (void *)LaunchInfo.Data.Host.LocalArgs);
-    } while (false);
+    getContext()->logger.debug("launch_info {} (numLocalArgs={}, localArgs={})",
+                               (void *)LaunchInfo.Data.getDevicePtr(),
+                               LaunchInfo.Data.Host.NumLocalArgs,
+                               (void *)LaunchInfo.Data.Host.LocalArgs);
 
     return UR_RESULT_SUCCESS;
 }

source/loader/layers/sanitizer/asan/asan_shadow.hpp

@@ -12,7 +12,7 @@
 
 #pragma once
 
-#include "asan/asan_allocator.hpp"
+#include "asan_allocator.hpp"
 #include "sanitizer_common/sanitizer_libdevice.hpp"
 
 #include <unordered_set>
@@ -134,7 +134,7 @@ struct ShadowMemoryPVC final : public ShadowMemoryGPU {
     size_t GetShadowSize() override { return 0x180000000000ULL; }
 };
 
-/// Shadow Memory layout of GPU PVC device
+/// Shadow Memory layout of GPU DG2 device
 ///
 /// USM Allocation Range (48 bits)
 ///   Host/Shared USM : 0x0000_0000_0000_0000 ~ 0x0000_7fff_ffff_ffff