[DeviceAsan][NFC] Code Restructure (#15843)

UR: oneapi-src/unified-runtime#2232 Code refactoring for implementing MemorySanitizer --------- Co-authored-by: Piotr Balcer <[email protected]>
intel · Nov 22, 2024 · 1e1ffc0 · 1e1ffc0
1 parent db59b4f
commit 1e1ffc0
Show file tree

Hide file tree

Showing 7 changed files with 119 additions and 93 deletions.
diff --git a/.github/CODEOWNERS b/.github/CODEOWNERS
@@ -169,12 +169,12 @@ sycl/test-e2e/LLVMIntrinsicLowering/ @intel/dpcpp-spirv-reviewers
 
 # Sanitizer
 clang/lib/Driver/SanitizerArgs.cpp @intel/dpcpp-sanitizers-review
-libdevice/sanitizer_utils.cpp @intel/dpcpp-sanitizers-review
-libdevice/include/asan_libdevice.hpp @intel/dpcpp-sanitizers-review
-libdevice/include/sanitizer_utils.hpp @intel/dpcpp-sanitizers-review
+libdevice/include/asan_rtl.hpp @intel/dpcpp-sanitizers-review
+libdevice/include/sanitizer_defs.hpp @intel/dpcpp-sanitizers-review
+libdevice/sanitizer/ @intel/dpcpp-sanitizers-review
+llvm/include/llvm/Transforms/Instrumentation/AddressSanitizer.h @intel/dpcpp-sanitizers-review
+llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h @intel/dpcpp-sanitizers-review
+llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerOptions.h @intel/dpcpp-sanitizers-review
 llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp @intel/dpcpp-sanitizers-review
-sycl/test-e2e/AddressSanitizer/ @intel/dpcpp-sanitizers-review
 llvm/test/Instrumentation/AddressSanitizer/ @intel/dpcpp-sanitizers-review
-llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerOptions.h @intel/dpcpp-sanitizers-review
-llvm/include/llvm/Transforms/Instrumentation/AddressSanitizerCommon.h @intel/dpcpp-sanitizers-review
-llvm/include/llvm/Transforms/Instrumentation/AddressSanitizer.h @intel/dpcpp-sanitizers-review
+sycl/test-e2e/AddressSanitizer/ @intel/dpcpp-sanitizers-review
diff --git a/clang/lib/Driver/ToolChains/SYCL.cpp b/clang/lib/Driver/ToolChains/SYCL.cpp
@@ -658,43 +658,44 @@ SYCL::getDeviceLibraries(const Compilation &C, const llvm::Triple &TargetTriple,
     addLibraries(SYCLDeviceAnnotationLibs);
 
 #if !defined(_WIN32)
+  std::string SanitizeVal;
   size_t sanitizer_lib_idx = getSingleBuildTarget();
   if (Arg *A = Args.getLastArg(options::OPT_fsanitize_EQ,
                                options::OPT_fno_sanitize_EQ)) {
     if (A->getOption().matches(options::OPT_fsanitize_EQ) &&
-        A->getValues().size() == 1) {
-      std::string SanitizeVal = A->getValue();
-      if (SanitizeVal == "address")
-        addSingleLibrary(SYCLDeviceAsanLibs[sanitizer_lib_idx]);
-    }
+        A->getValues().size() == 1)
+      SanitizeVal = A->getValue();
   } else {
     // User can pass -fsanitize=address to device compiler via
     // -Xsycl-target-frontend, sanitize device library must be
     // linked with user's device image if so.
-    bool IsDeviceAsanEnabled = false;
-    auto SyclFEArg = Args.getAllArgValues(options::OPT_Xsycl_frontend);
-    IsDeviceAsanEnabled = (std::count(SyclFEArg.begin(), SyclFEArg.end(),
-                                      "-fsanitize=address") > 0);
-    if (!IsDeviceAsanEnabled) {
-      auto SyclFEArgEq = Args.getAllArgValues(options::OPT_Xsycl_frontend_EQ);
-      IsDeviceAsanEnabled = (std::count(SyclFEArgEq.begin(), SyclFEArgEq.end(),
-                                        "-fsanitize=address") > 0);
-    }
-
-    // User can also enable asan for SYCL device via -Xarch_device option.
-    if (!IsDeviceAsanEnabled) {
-      auto DeviceArchVals = Args.getAllArgValues(options::OPT_Xarch_device);
-      for (auto DArchVal : DeviceArchVals) {
-        if (DArchVal.find("-fsanitize=address") != std::string::npos) {
-          IsDeviceAsanEnabled = true;
-          break;
-        }
+    std::vector<std::string> EnabledDeviceSanitizers;
+
+    // NOTE: "-fsanitize=" applies to all device targets
+    auto SyclFEArgVals = Args.getAllArgValues(options::OPT_Xsycl_frontend);
+    auto SyclFEEQArgVals = Args.getAllArgValues(options::OPT_Xsycl_frontend_EQ);
+    auto ArchDeviceVals = Args.getAllArgValues(options::OPT_Xarch_device);
+
+    std::vector<std::string> ArgVals(
+        SyclFEArgVals.size() + SyclFEEQArgVals.size() + ArchDeviceVals.size());
+    ArgVals.insert(ArgVals.end(), SyclFEArgVals.begin(), SyclFEArgVals.end());
+    ArgVals.insert(ArgVals.end(), SyclFEEQArgVals.begin(),
+                   SyclFEEQArgVals.end());
+    ArgVals.insert(ArgVals.end(), ArchDeviceVals.begin(), ArchDeviceVals.end());
+
+    // Driver will report error if address sanitizer and memory sanitizer are
+    // both enabled, so we only need to check first one here.
+    for (const std::string &Arg : ArgVals) {
+      if (Arg.find("-fsanitize=address") != std::string::npos) {
+        SanitizeVal = "address";
+        break;
       }
     }
-
-    if (IsDeviceAsanEnabled)
-      addSingleLibrary(SYCLDeviceAsanLibs[sanitizer_lib_idx]);
   }
+
+  if (SanitizeVal == "address")
+    addSingleLibrary(SYCLDeviceAsanLibs[sanitizer_lib_idx]);
+
 #endif
 
   if (isNativeCPU)

diff --git a/libdevice/cmake/modules/SYCLLibdevice.cmake b/libdevice/cmake/modules/SYCLLibdevice.cmake
@@ -197,7 +197,7 @@ function(add_devicelibs filename)
   cmake_parse_arguments(ARG
     ""
     ""
-    "SRC;EXTRA_OPTS;DEPENDENCIES"
+    "SRC;EXTRA_OPTS;DEPENDENCIES;SKIP_ARCHS"
     ${ARGN})
 
   foreach(filetype IN LISTS filetypes)
@@ -209,6 +209,9 @@ function(add_devicelibs filename)
   endforeach()
 
   foreach(arch IN LISTS devicelib_arch)
+    if(arch IN_LIST ARG_SKIP_ARCHS)
+      continue()
+    endif()
     compile_lib(${filename}-${arch}
       FILETYPE bc
       SRC ${ARG_SRC}
@@ -229,16 +232,17 @@ set(imf_obj_deps device_imf.hpp imf_half.hpp imf_bf16.hpp imf_rounding_op.hpp im
 set(itt_obj_deps device_itt.h spirv_vars.h device.h sycl-compiler)
 set(bfloat16_obj_deps sycl-headers sycl-compiler)
 if (NOT MSVC AND UR_SANITIZER_INCLUDE_DIR)
-  set(sanitizer_obj_deps
+  set(asan_obj_deps
     device.h atomic.hpp spirv_vars.h
-    ${UR_SANITIZER_INCLUDE_DIR}/asan_libdevice.hpp
-    include/sanitizer_utils.hpp
+    ${UR_SANITIZER_INCLUDE_DIR}/asan/asan_libdevice.hpp
+    include/asan_rtl.hpp
     include/spir_global_var.hpp
     sycl-compiler)
 
   set(sanitizer_generic_compile_opts ${compile_opts}
                             -fno-sycl-instrument-device-code
-                            -I${UR_SANITIZER_INCLUDE_DIR})
+                            -I${UR_SANITIZER_INCLUDE_DIR}
+                            -I${CMAKE_CURRENT_SOURCE_DIR})
 
   set(asan_pvc_compile_opts_obj -fsycl -c
                                 ${sanitizer_generic_compile_opts}
@@ -346,19 +350,27 @@ if(MSVC)
     DEPENDENCIES ${cmath_obj_deps})
 else()
   if(UR_SANITIZER_INCLUDE_DIR)
+    # asan jit
     add_devicelibs(libsycl-asan
-      SRC sanitizer_utils.cpp
-      DEPENDENCIES ${sanitizer_obj_deps}
-      EXTRA_OPTS -fno-sycl-instrument-device-code -I${UR_SANITIZER_INCLUDE_DIR})
+      SRC sanitizer/asan_rtl.cpp
+      DEPENDENCIES ${asan_obj_deps}
+      SKIP_ARCHS nvptx64-nvidia-cuda
+                 amdgcn-amd-amdhsa
+      EXTRA_OPTS -fno-sycl-instrument-device-code
+                 -I${UR_SANITIZER_INCLUDE_DIR}
+                 -I${CMAKE_CURRENT_SOURCE_DIR})
+
+    # asan aot
     set(asan_filetypes obj obj-new-offload bc)
     set(asan_devicetypes pvc cpu dg2)
+
     foreach(asan_ft IN LISTS asan_filetypes)
       foreach(asan_device IN LISTS asan_devicetypes)
         compile_lib_ext(libsycl-asan-${asan_device}
-        SRC sanitizer_utils.cpp
-        FILETYPE ${asan_ft}
-        DEPENDENCIES ${sanitizer_obj_deps}
-        OPTS ${asan_${asan_device}_compile_opts_${asan_ft}})
+                        SRC sanitizer/asan_rtl.cpp
+                        FILETYPE ${asan_ft}
+                        DEPENDENCIES ${asan_obj_deps}
+                        OPTS ${asan_${asan_device}_compile_opts_${asan_ft}})
       endforeach()
     endforeach()
   endif()

diff --git a/libdevice/include/sanitizer_utils.hpp → libdevice/include/asan_rtl.hpp b/libdevice/include/sanitizer_utils.hpp → libdevice/include/asan_rtl.hpp
@@ -1,4 +1,4 @@
-//==-- sanitizer_device_utils.hpp - Declaration for sanitizer global var ---==//
+//==-- asan_rtl.hpp - Declaration for sanitizer global var ---==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include "sanitizer_defs.hpp"
 #include "spir_global_var.hpp"
 #include <cstdint>
 

diff --git a/libdevice/include/sanitizer_defs.hpp b/libdevice/include/sanitizer_defs.hpp
@@ -0,0 +1,24 @@
+//==-- sanitizer_defs.hpp - common macros shared by sanitizers ---==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+#pragma once
+
+#include <cstdint>
+
+using uptr = uintptr_t;
+using u8 = uint8_t;
+using u16 = uint16_t;
+using u32 = uint32_t;
+using u64 = uint64_t;
+using s8 = int8_t;
+using s16 = int16_t;
+using s32 = int32_t;
+using s64 = int64_t;
+
+#define LIKELY(x) __builtin_expect(!!(x), 1)
+#define UNLIKELY(x) __builtin_expect(!!(x), 0)
+#define NORETURN __declspec(noreturn)
diff --git a/libdevice/sanitizer_utils.cpp → libdevice/sanitizer/asan_rtl.cpp b/libdevice/sanitizer_utils.cpp → libdevice/sanitizer/asan_rtl.cpp
@@ -1,24 +1,17 @@
-//==--- sanitizer_utils.cpp - device sanitizer util inserted by compiler ---==//
+//==--- asan_rtl.cpp - device address sanitizer runtime library ------------==//
 //
 // Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
 // See https://llvm.org/LICENSE.txt for license information.
 // SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
 //
 //===----------------------------------------------------------------------===//
 
-#include "asan_libdevice.hpp"
+#include "include/asan_rtl.hpp"
+#include "asan/asan_libdevice.hpp"
 #include "atomic.hpp"
 #include "device.h"
 #include "spirv_vars.h"
 
-#include "include/sanitizer_utils.hpp"
-
-using uptr = uintptr_t;
-using s8 = char;
-using u8 = unsigned char;
-using s16 = short;
-using u16 = unsigned short;
-
 // Save the pointer to LaunchInfo
 __SYCL_GLOBAL__ uptr *__SYCL_LOCAL__ __AsanLaunchInfo;
 
@@ -375,7 +368,7 @@ bool MemIsZero(__SYCL_GLOBAL__ const char *beg, uptr size) {
 static __SYCL_CONSTANT__ const char __mem_sanitizer_report[] =
     "[kernel] SanitizerReport (ErrorType=%d, IsRecover=%d)\n";
 
-void __asan_internal_report_save(DeviceSanitizerErrorType error_type) {
+void __asan_internal_report_save(ErrorType error_type) {
   const int Expected = ASAN_REPORT_NONE;
   int Desired = ASAN_REPORT_START;
 
@@ -387,21 +380,21 @@ void __asan_internal_report_save(DeviceSanitizerErrorType error_type) {
       __spirv_BuiltInWorkgroupId.z;
 
   auto &SanitizerReport = ((__SYCL_GLOBAL__ LaunchInfo *)__AsanLaunchInfo)
-                              ->SanitizerReport[WG_LID % ASAN_MAX_NUM_REPORTS];
+                              ->Report[WG_LID % ASAN_MAX_NUM_REPORTS];
 
   if (atomicCompareAndSet(
           &(((__SYCL_GLOBAL__ LaunchInfo *)__AsanLaunchInfo)->ReportFlag), 1,
           0) == 0 &&
       atomicCompareAndSet(&SanitizerReport.Flag, Desired, Expected) ==
           Expected) {
-    SanitizerReport.ErrorType = error_type;
+    SanitizerReport.ErrorTy = error_type;
     SanitizerReport.IsRecover = false;
 
     // Show we've done copying
     atomicStore(&SanitizerReport.Flag, ASAN_REPORT_FINISH);
 
     ASAN_DEBUG(__spirv_ocl_printf(__mem_sanitizer_report,
-                                  SanitizerReport.ErrorType,
+                                  SanitizerReport.ErrorTy,
                                   SanitizerReport.IsRecover));
   }
   __devicelib_exit();
@@ -410,8 +403,7 @@ void __asan_internal_report_save(DeviceSanitizerErrorType error_type) {
 void __asan_internal_report_save(
     uptr ptr, uint32_t as, const char __SYCL_CONSTANT__ *file, uint32_t line,
     const char __SYCL_CONSTANT__ *func, bool is_write, uint32_t access_size,
-    DeviceSanitizerMemoryType memory_type, DeviceSanitizerErrorType error_type,
-    bool is_recover = false) {
+    MemoryType memory_type, ErrorType error_type, bool is_recover = false) {
 
   const int Expected = ASAN_REPORT_NONE;
   int Desired = ASAN_REPORT_START;
@@ -424,7 +416,7 @@ void __asan_internal_report_save(
       __spirv_BuiltInWorkgroupId.z;
 
   auto &SanitizerReport = ((__SYCL_GLOBAL__ LaunchInfo *)__AsanLaunchInfo)
-                              ->SanitizerReport[WG_LID % ASAN_MAX_NUM_REPORTS];
+                              ->Report[WG_LID % ASAN_MAX_NUM_REPORTS];
 
   if ((is_recover ||
        atomicCompareAndSet(
@@ -470,15 +462,15 @@ void __asan_internal_report_save(
     SanitizerReport.Address = ptr;
     SanitizerReport.IsWrite = is_write;
     SanitizerReport.AccessSize = access_size;
-    SanitizerReport.ErrorType = error_type;
-    SanitizerReport.MemoryType = memory_type;
+    SanitizerReport.ErrorTy = error_type;
+    SanitizerReport.MemoryTy = memory_type;
     SanitizerReport.IsRecover = is_recover;
 
     // Show we've done copying
     atomicStore(&SanitizerReport.Flag, ASAN_REPORT_FINISH);
 
     ASAN_DEBUG(__spirv_ocl_printf(__mem_sanitizer_report,
-                                  SanitizerReport.ErrorType,
+                                  SanitizerReport.ErrorTy,
                                   SanitizerReport.IsRecover));
   }
   __devicelib_exit();
@@ -488,29 +480,29 @@ void __asan_internal_report_save(
 /// ASAN Error Reporters
 ///
 
-DeviceSanitizerMemoryType GetMemoryTypeByShadowValue(int shadow_value) {
+MemoryType GetMemoryTypeByShadowValue(int shadow_value) {
   switch (shadow_value) {
   case kUsmDeviceRedzoneMagic:
   case kUsmDeviceDeallocatedMagic:
-    return DeviceSanitizerMemoryType::USM_DEVICE;
+    return MemoryType::USM_DEVICE;
   case kUsmHostRedzoneMagic:
   case kUsmHostDeallocatedMagic:
-    return DeviceSanitizerMemoryType::USM_HOST;
+    return MemoryType::USM_HOST;
   case kUsmSharedRedzoneMagic:
   case kUsmSharedDeallocatedMagic:
-    return DeviceSanitizerMemoryType::USM_SHARED;
+    return MemoryType::USM_SHARED;
   case kPrivateLeftRedzoneMagic:
   case kPrivateMidRedzoneMagic:
   case kPrivateRightRedzoneMagic:
-    return DeviceSanitizerMemoryType::PRIVATE;
+    return MemoryType::PRIVATE;
   case kMemBufferRedzoneMagic:
-    return DeviceSanitizerMemoryType::MEM_BUFFER;
+    return MemoryType::MEM_BUFFER;
   case kSharedLocalRedzoneMagic:
-    return DeviceSanitizerMemoryType::LOCAL;
+    return MemoryType::LOCAL;
   case kDeviceGlobalRedzoneMagic:
-    return DeviceSanitizerMemoryType::DEVICE_GLOBAL;
+    return MemoryType::DEVICE_GLOBAL;
   default:
-    return DeviceSanitizerMemoryType::UNKNOWN;
+    return MemoryType::UNKNOWN;
   }
 }
 
@@ -528,9 +520,8 @@ void __asan_report_access_error(uptr addr, uint32_t as, size_t size,
   }
   // FIXME: check if shadow_address out-of-bound
 
-  DeviceSanitizerMemoryType memory_type =
-      GetMemoryTypeByShadowValue(shadow_value);
-  DeviceSanitizerErrorType error_type;
+  MemoryType memory_type = GetMemoryTypeByShadowValue(shadow_value);
+  ErrorType error_type;
 
   switch (shadow_value) {
   case kUsmDeviceRedzoneMagic:
@@ -542,18 +533,18 @@ void __asan_report_access_error(uptr addr, uint32_t as, size_t size,
   case kMemBufferRedzoneMagic:
   case kSharedLocalRedzoneMagic:
   case kDeviceGlobalRedzoneMagic:
-    error_type = DeviceSanitizerErrorType::OUT_OF_BOUNDS;
+    error_type = ErrorType::OUT_OF_BOUNDS;
     break;
   case kUsmDeviceDeallocatedMagic:
   case kUsmHostDeallocatedMagic:
   case kUsmSharedDeallocatedMagic:
-    error_type = DeviceSanitizerErrorType::USE_AFTER_FREE;
+    error_type = ErrorType::USE_AFTER_FREE;
     break;
   case kNullPointerRedzoneMagic:
-    error_type = DeviceSanitizerErrorType::NULL_POINTER;
+    error_type = ErrorType::NULL_POINTER;
     break;
   default:
-    error_type = DeviceSanitizerErrorType::UNKNOWN;
+    error_type = ErrorType::UNKNOWN;
   }
 
   __asan_internal_report_save(addr, as, file, line, func, is_write, size,
@@ -573,16 +564,15 @@ void __asan_report_misalign_error(uptr addr, uint32_t as, size_t size,
   }
   int shadow_value = *shadow;
 
-  DeviceSanitizerErrorType error_type = DeviceSanitizerErrorType::MISALIGNED;
-  DeviceSanitizerMemoryType memory_type =
-      GetMemoryTypeByShadowValue(shadow_value);
+  ErrorType error_type = ErrorType::MISALIGNED;
+  MemoryType memory_type = GetMemoryTypeByShadowValue(shadow_value);
 
   __asan_internal_report_save(addr, as, file, line, func, is_write, size,
                               memory_type, error_type, is_recover);
 }
 
 void __asan_report_unknown_device() {
-  __asan_internal_report_save(DeviceSanitizerErrorType::UNKNOWN_DEVICE);
+  __asan_internal_report_save(ErrorType::UNKNOWN_DEVICE);
 }
 
 ///