[SYCL][DeviceSanitizer] Checking out-of-bounds error on sycl::local_accessor (#13503)

UR: oneapi-src/unified-runtime#1532

To check sycl::local_accessor(aka, dynamic local memory), we need to
extend a new argument in spir kernel, this is because:
- ASan needs to know some size information of local buffer, like its
size and size with redzone, so that it can poison its shadow memory
- By using this new argument, we can also pass some per-launch
information (that is, it is different in each launch of kernel). One
obvious example is SanitizerReport, which saves the error message, so
that we can store and print multiple error reports for one kernel with
different arguments. Another example is the shadow memory of local
memory, this should be different per-launch as well, since one kernel
can be launched multiple times and executed in parallel.

I named this argument as "__asan_launch", which is a pointer pointed to
"LaunchInfo" structure and allocated it in shared USM. To make this
pointer can be used in spir_func w/o extending their argument, I created
a global external local memory (external, so that it can be shared with
other translation units, and its instance is defined in libdevice), and
save the "__asan_launch" into this local memory immediately at the entry
of kernel.

UR can't check the name of kernel arguments, so it can't know if the
kernel has "__asan_launch". So I assume the "__asan_launch" is always
there, and added a check to prevent DAE pass from removing it.

---------

Co-authored-by: Maosu Zhao <[email protected]>
Co-authored-by: Kenneth Benzie (Benie) <[email protected]>

Author: 3 people

libdevice/cmake/modules/SYCLLibdevice.cmake

@@ -159,7 +159,7 @@ 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)
-  set(sanitizer_obj_deps device.h atomic.hpp spirv_vars.h include/sanitizer_device_utils.hpp include/spir_global_var.hpp sycl-compiler)
+  set(sanitizer_obj_deps device.h atomic.hpp spirv_vars.h include/sanitizer_utils.hpp include/spir_global_var.hpp sycl-compiler)
 endif()
 
 add_devicelib(libsycl-itt-stubs SRC itt_stubs.cpp DEP ${itt_obj_deps})

libdevice/include/asan_libdevice.hpp

@@ -59,7 +59,6 @@ struct DeviceSanitizerReport {
 };
 
 struct LocalArgsInfo {
-  uint32_t ArgIndex = 0;
   uint64_t Size = 0;
   uint64_t SizeWithRedZone = 0;
 };

libdevice/include/sanitizer_device_utils.hpp renamed to libdevice/include/sanitizer_utils.hpp

@@ -10,4 +10,10 @@
 #include "spir_global_var.hpp"
 #include <cstdint>
 
+// Treat this header as system one to workaround frontend's restriction
+#pragma clang system_header
+
 enum DeviceType : uint64_t { UNKNOWN, CPU, GPU_PVC, GPU_DG2 };
+
+extern SPIR_GLOBAL_VAR __SYCL_GLOBAL__ uint64_t *__SYCL_LOCAL__
+    __AsanLaunchInfo;

libdevice/include/spir_global_var.hpp

@@ -42,3 +42,11 @@ class
 #define __SYCL_LOCAL__ __attribute__((opencl_local))
 #define __SYCL_PRIVATE__ __attribute__((opencl_private))
 #define __SYCL_CONSTANT__ __attribute__((opencl_constant))
+
+#ifndef SPIR_GLOBAL_VAR
+#ifdef __SYCL_DEVICE_ONLY__
+#define SPIR_GLOBAL_VAR __attribute__((sycl_global_var))
+#else
+#define SPIR_GLOBAL_VAR
+#endif
+#endif

libdevice/sanitizer_utils.cpp

@@ -11,9 +11,7 @@
 #include "spirv_vars.h"
 
 #include "include/asan_libdevice.hpp"
-#include "include/sanitizer_device_utils.hpp"
-#include <cstddef>
-#include <cstdint>
+#include "include/sanitizer_utils.hpp"
 
 using uptr = uintptr_t;
 using s8 = char;
@@ -23,11 +21,10 @@ using u16 = unsigned short;
 
 DeviceGlobal<uptr> __AsanShadowMemoryGlobalStart;
 DeviceGlobal<uptr> __AsanShadowMemoryGlobalEnd;
-DeviceGlobal<uptr> __AsanShadowMemoryLocalStart;
-DeviceGlobal<uptr> __AsanShadowMemoryLocalEnd;
 DeviceGlobal<DeviceType> __DeviceType;
 DeviceGlobal<uint64_t> __AsanDebug;
-DeviceGlobal<DeviceSanitizerReport> __DeviceSanitizerReportMem;
+// Save the pointer to LaunchInfo
+__SYCL_GLOBAL__ uptr *__SYCL_LOCAL__ __AsanLaunchInfo;
 
 #if defined(__SPIR__) || defined(__SPIRV__)
 
@@ -134,6 +131,16 @@ inline uptr MemToShadow_DG2(uptr addr, uint32_t as) {
   return shadow_ptr;
 }
 
+static __SYCL_CONSTANT__ const char __mem_launch_info[] =
+    "[kernel] launch_info: %p (local_shadow=%p~%p, numLocalArgs=%d, "
+    "localArgs=%p)\n";
+
+static __SYCL_CONSTANT__ const char __generic_to[] =
+    "[kernel] %p(4) - %p(%d)\n";
+
+static __SYCL_CONSTANT__ const char __generic_to_fail[] =
+    "[kernel] %p(4) - unknown address space\n";
+
 inline uptr MemToShadow_PVC(uptr addr, uint32_t as) {
 
   if (as == ADDRESS_SPACE_GENERIC) {
@@ -172,9 +179,6 @@ inline uptr MemToShadow_PVC(uptr addr, uint32_t as) {
     }
     return shadow_ptr;
   } else if (as == ADDRESS_SPACE_LOCAL) { // local
-    if (__AsanShadowMemoryLocalStart == 0) {
-      return 0;
-    }
     // The size of SLM is 128KB on PVC
     constexpr unsigned SLM_SIZE = 128 * 1024;
     // work-group linear id
@@ -184,14 +188,28 @@ inline uptr MemToShadow_PVC(uptr addr, uint32_t as) {
         __spirv_BuiltInWorkgroupId.y * __spirv_BuiltInNumWorkgroups.z +
         __spirv_BuiltInWorkgroupId.z;
 
-    uptr shadow_ptr = __AsanShadowMemoryLocalStart +
+    auto launch_info = (__SYCL_GLOBAL__ const LaunchInfo *)__AsanLaunchInfo;
+    const auto shadow_offset = launch_info->LocalShadowOffset;
+    const auto shadow_offset_end = launch_info->LocalShadowOffsetEnd;
+
+    if (shadow_offset == 0) {
+      return 0;
+    }
+
+    if (__AsanDebug)
+      __spirv_ocl_printf(__mem_launch_info, launch_info,
+                         launch_info->LocalShadowOffset,
+                         launch_info->LocalShadowOffsetEnd,
+                         launch_info->NumLocalArgs, launch_info->LocalArgs);
+
+    uptr shadow_ptr = shadow_offset +
                       ((wg_lid * SLM_SIZE) >> ASAN_SHADOW_SCALE) +
                       ((addr & (SLM_SIZE - 1)) >> 3);
 
-    if (shadow_ptr > __AsanShadowMemoryLocalEnd) {
+    if (shadow_ptr > shadow_offset_end) {
       if (__asan_report_out_of_shadow_bounds() && __AsanDebug) {
         __spirv_ocl_printf(__local_shadow_out_of_bound, addr, shadow_ptr,
-                           wg_lid, (uptr)__AsanShadowMemoryLocalStart);
+                           wg_lid, (uptr)shadow_offset);
       }
       return 0;
     }
@@ -268,22 +286,18 @@ bool MemIsZero(__SYCL_GLOBAL__ const char *beg, uptr size) {
 bool __asan_internal_report_save(DeviceSanitizerErrorType error_type) {
   const int Expected = ASAN_REPORT_NONE;
   int Desired = ASAN_REPORT_START;
-  if (atomicCompareAndSet(&__DeviceSanitizerReportMem.get().Flag, Desired,
-                          Expected) == Expected) {
-    __DeviceSanitizerReportMem.get().ErrorType = error_type;
+  auto &SanitizerReport =
+      ((__SYCL_GLOBAL__ LaunchInfo *)__AsanLaunchInfo)->SanitizerReport;
+  if (atomicCompareAndSet(&SanitizerReport.Flag, Desired, Expected) ==
+      Expected) {
+    SanitizerReport.ErrorType = error_type;
     // Show we've done copying
-    atomicStore(&__DeviceSanitizerReportMem.get().Flag, ASAN_REPORT_FINISH);
+    atomicStore(&SanitizerReport.Flag, ASAN_REPORT_FINISH);
     return true;
   }
   return false;
 }
 
-#ifdef __SYCL_DEVICE_ONLY__
-#define __DEVICE_SANITIZER_REPORT_ACCESSOR __DeviceSanitizerReportMem.get()
-#else // __SYCL_DEVICE_ONLY__
-#define __DEVICE_SANITIZER_REPORT_ACCESSOR
-#endif // __SYCL_DEVICE_ONLY__
-
 bool __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,
@@ -292,8 +306,20 @@ bool __asan_internal_report_save(
 
   const int Expected = ASAN_REPORT_NONE;
   int Desired = ASAN_REPORT_START;
-  if (atomicCompareAndSet(&__DEVICE_SANITIZER_REPORT_ACCESSOR.Flag, Desired,
-                          Expected) == Expected) {
+
+  if (__AsanDebug) {
+    auto *launch_info = (__SYCL_GLOBAL__ LaunchInfo *)__AsanLaunchInfo;
+    __spirv_ocl_printf(__mem_launch_info, launch_info,
+                       launch_info->LocalShadowOffset,
+                       launch_info->LocalShadowOffsetEnd,
+                       launch_info->NumLocalArgs, launch_info->LocalArgs);
+  }
+
+  auto &SanitizerReport =
+      ((__SYCL_GLOBAL__ LaunchInfo *)__AsanLaunchInfo)->SanitizerReport;
+
+  if (atomicCompareAndSet(&SanitizerReport.Flag, Desired, Expected) ==
+      Expected) {
 
     int FileLength = 0;
     int FuncLength = 0;
@@ -305,39 +331,40 @@ bool __asan_internal_report_save(
       for (auto *C = func; *C != '\0'; ++C, ++FuncLength)
         ;
 
-    int MaxFileIdx = sizeof(__DEVICE_SANITIZER_REPORT_ACCESSOR.File) - 1;
-    int MaxFuncIdx = sizeof(__DEVICE_SANITIZER_REPORT_ACCESSOR.Func) - 1;
+    int MaxFileIdx = sizeof(SanitizerReport.File) - 1;
+    int MaxFuncIdx = sizeof(SanitizerReport.Func) - 1;
 
     if (FileLength < MaxFileIdx)
       MaxFileIdx = FileLength;
     if (FuncLength < MaxFuncIdx)
       MaxFuncIdx = FuncLength;
 
     for (int Idx = 0; Idx < MaxFileIdx; ++Idx)
-      __DEVICE_SANITIZER_REPORT_ACCESSOR.File[Idx] = file[Idx];
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.File[MaxFileIdx] = '\0';
+      SanitizerReport.File[Idx] = file[Idx];
+    SanitizerReport.File[MaxFileIdx] = '\0';
 
     for (int Idx = 0; Idx < MaxFuncIdx; ++Idx)
-      __DEVICE_SANITIZER_REPORT_ACCESSOR.Func[Idx] = func[Idx];
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.Func[MaxFuncIdx] = '\0';
-
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.Line = line;
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.GID0 = __spirv_GlobalInvocationId_x();
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.GID1 = __spirv_GlobalInvocationId_y();
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.GID2 = __spirv_GlobalInvocationId_z();
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.LID0 = __spirv_LocalInvocationId_x();
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.LID1 = __spirv_LocalInvocationId_y();
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.LID2 = __spirv_LocalInvocationId_z();
-
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.Address = ptr;
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.IsWrite = is_write;
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.AccessSize = access_size;
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.ErrorType = error_type;
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.MemoryType = memory_type;
-    __DEVICE_SANITIZER_REPORT_ACCESSOR.IsRecover = is_recover;
+      SanitizerReport.Func[Idx] = func[Idx];
+    SanitizerReport.Func[MaxFuncIdx] = '\0';
+
+    SanitizerReport.Line = line;
+    SanitizerReport.GID0 = __spirv_GlobalInvocationId_x();
+    SanitizerReport.GID1 = __spirv_GlobalInvocationId_y();
+    SanitizerReport.GID2 = __spirv_GlobalInvocationId_z();
+    SanitizerReport.LID0 = __spirv_LocalInvocationId_x();
+    SanitizerReport.LID1 = __spirv_LocalInvocationId_y();
+    SanitizerReport.LID2 = __spirv_LocalInvocationId_z();
+
+    SanitizerReport.Address = ptr;
+    SanitizerReport.IsWrite = is_write;
+    SanitizerReport.AccessSize = access_size;
+    SanitizerReport.ErrorType = error_type;
+    SanitizerReport.MemoryType = memory_type;
+    SanitizerReport.IsRecover = is_recover;
 
     // Show we've done copying
-    atomicStore(&__DEVICE_SANITIZER_REPORT_ACCESSOR.Flag, ASAN_REPORT_FINISH);
+    atomicStore(&SanitizerReport.Flag, ASAN_REPORT_FINISH);
+    return true;
   }
   return false;
 }
@@ -545,7 +572,7 @@ ASAN_REPORT_ERROR(store, true, 4)
   DEVICE_EXTERN_C_NOINLINE void __asan_##type##size(                           \
       uptr addr, uint32_t as, const char __SYCL_CONSTANT__ *file,              \
       uint32_t line, const char __SYCL_CONSTANT__ *func) {                     \
-    u##size *shadow_address = (u##size *)MemToShadow(addr, as);                \
+    auto *shadow_address = (__SYCL_GLOBAL__ u##size *)MemToShadow(addr, as);   \
     if (shadow_address && *shadow_address) {                                   \
       __asan_report_access_error(addr, as, size, is_write, addr, file, line,   \
                                  func);                                        \
@@ -554,7 +581,7 @@ ASAN_REPORT_ERROR(store, true, 4)
   DEVICE_EXTERN_C_NOINLINE void __asan_##type##size##_noabort(                 \
       uptr addr, uint32_t as, const char __SYCL_CONSTANT__ *file,              \
       uint32_t line, const char __SYCL_CONSTANT__ *func) {                     \
-    u##size *shadow_address = (u##size *)MemToShadow(addr, as);                \
+    auto *shadow_address = (__SYCL_GLOBAL__ u##size *)MemToShadow(addr, as);   \
     if (shadow_address && *shadow_address) {                                   \
       __asan_report_access_error(addr, as, size, is_write, addr, file, line,   \
                                  func, true);                                  \
@@ -595,7 +622,7 @@ static __SYCL_CONSTANT__ const char __mem_set_shadow_local[] =
     "[kernel] set_shadow_local(beg=%p, end=%p, val:%02X)\n";
 
 DEVICE_EXTERN_C_NOINLINE void
-__asan_set_shadow_local_memory(uptr ptr, size_t size,
+__asan_set_shadow_static_local(uptr ptr, size_t size,
                                size_t size_with_redzone) {
   // Since ptr is aligned to ASAN_SHADOW_GRANULARITY,
   // if size != aligned_size, then the buffer tail of ptr is not aligned
@@ -638,4 +665,48 @@ __asan_set_shadow_local_memory(uptr ptr, size_t size,
   }
 }
 
+static __SYCL_CONSTANT__ const char __mem_local_arg[] =
+    "[kernel] local_arg(index=%d, size=%d, size_rz=%d)\n";
+
+static __SYCL_CONSTANT__ const char __mem_set_shadow_dynamic_local_begin[] =
+    "[kernel] BEGIN __asan_set_shadow_dynamic_local\n";
+static __SYCL_CONSTANT__ const char __mem_set_shadow_dynamic_local_end[] =
+    "[kernel] END   __asan_set_shadow_dynamic_local\n";
+static __SYCL_CONSTANT__ const char __mem_report_arg_count_incorrect[] =
+    "[kernel] ERROR: The number of local args is incorrect, expect %d, actual "
+    "%d\n";
+
+DEVICE_EXTERN_C_NOINLINE void
+__asan_set_shadow_dynamic_local(uptr ptr, uint32_t num_args) {
+  if (__AsanDebug)
+    __spirv_ocl_printf(__mem_set_shadow_dynamic_local_begin);
+
+  auto *launch_info = (__SYCL_GLOBAL__ const LaunchInfo *)__AsanLaunchInfo;
+  if (num_args != launch_info->NumLocalArgs) {
+    __spirv_ocl_printf(__mem_report_arg_count_incorrect, num_args,
+                       launch_info->NumLocalArgs);
+    return;
+  }
+
+  uptr *args = (uptr *)ptr;
+  if (__AsanDebug)
+    __spirv_ocl_printf(__mem_launch_info, launch_info,
+                       launch_info->LocalShadowOffset,
+                       launch_info->LocalShadowOffsetEnd,
+                       launch_info->NumLocalArgs, launch_info->LocalArgs);
+
+  for (uint32_t i = 0; i < num_args; ++i) {
+    auto *local_arg = &launch_info->LocalArgs[i];
+    if (__AsanDebug)
+      __spirv_ocl_printf(__mem_local_arg, i, local_arg->Size,
+                         local_arg->SizeWithRedZone);
+
+    __asan_set_shadow_static_local(args[i], local_arg->Size,
+                                   local_arg->SizeWithRedZone);
+  }
+
+  if (__AsanDebug)
+    __spirv_ocl_printf(__mem_set_shadow_dynamic_local_end);
+}
+
 #endif // __SPIR__ || __SPIRV__

llvm/lib/Transforms/IPO/DeadArgumentElimination.cpp

@@ -540,6 +540,14 @@ void DeadArgumentEliminationPass::surveyFunction(const Function &F) {
     return;
   }
 
+  // Don't touch sanitized functions. The "__asan_launch" argument needs to be
+  // present at all times, even if it's not used.
+  if (F.getCallingConv() == CallingConv::SPIR_KERNEL &&
+      F.hasFnAttribute(Attribute::SanitizeAddress)) {
+    markLive(F);
+    return;
+  }
+
   unsigned RetCount = numRetVals(&F);
 
   // Assume all return values are dead