[DeviceASAN] Enhance out-of-bound by checking shadow bounds on global memory

libdevice/sanitizer/asan_rtl.cpp

@@ -9,6 +9,8 @@
 #include "include/asan_rtl.hpp"
 #include "asan/asan_libdevice.hpp"
 
+extern "C" __attribute__((weak)) const int __asan_check_shadow_bounds;
+
 // Save the pointer to LaunchInfo
 __SYCL_GLOBAL__ uptr *__SYCL_LOCAL__ __AsanLaunchInfo;
 
@@ -40,6 +42,9 @@ static const __SYCL_CONSTANT__ char __asan_print_shadow_value2[] =
 static __SYCL_CONSTANT__ const char __generic_to[] =
     "[kernel] %p(4) - %p(%d)\n";
 
+static __SYCL_CONSTANT__ const char __asan_print_shadow_bound[] =
+    "[kernel] addr: %p, shadow: %p, lower: %p, upper: %p\n";
+
 #define ASAN_REPORT_NONE 0
 #define ASAN_REPORT_START 1
 #define ASAN_REPORT_FINISH 2
@@ -65,8 +70,12 @@ struct DebugInfo {
   uint32_t line;
 };
 
+inline bool IsCheckShadowBounds() { return __asan_check_shadow_bounds; }
+
 void ReportUnknownDevice(const DebugInfo *debug);
 void PrintShadowMemory(uptr addr, uptr shadow_address, uint32_t as);
+void SaveReport(ErrorType error_type, MemoryType memory_type, bool is_recover,
+                const DebugInfo *debug);
 
 __SYCL_GLOBAL__ void *ToGlobal(void *ptr) {
   return __spirv_GenericCastToPtrExplicit_ToGlobal(ptr, 5);
@@ -115,6 +124,16 @@ inline uptr MemToShadow_DG2(uptr addr, uint32_t as,
           launch_info->GlobalShadowOffset + (addr >> ASAN_SHADOW_SCALE);
     }
 
+    if (IsCheckShadowBounds() &&
+        (shadow_ptr < launch_info->GlobalShadowLowerBound ||
+         shadow_ptr > launch_info->GlobalShadowUpperBound)) {
+      ASAN_DEBUG(__spirv_ocl_printf(__asan_print_shadow_bound, addr, shadow_ptr,
+                                    launch_info->GlobalShadowLowerBound,
+                                    launch_info->GlobalShadowUpperBound));
+      SaveReport(ErrorType::OUT_OF_BOUNDS, MemoryType::GLOBAL, false, debug);
+      return 0;
+    }
+
     ASAN_DEBUG(
         const auto shadow_offset_end = launch_info->GlobalShadowOffsetEnd;
         if (shadow_ptr > shadow_offset_end) {
@@ -168,7 +187,7 @@ inline uptr MemToShadow_DG2(uptr addr, uint32_t as,
       __spirv_ocl_printf(__private_shadow_out_of_bound, addr, shadow_ptr, sid,
                          private_base);
       return 0;
-    };
+    }
 
     return shadow_ptr;
   }
@@ -193,6 +212,16 @@ inline uptr MemToShadow_PVC(uptr addr, uint32_t as,
                    ((addr & 0x7FFFFFFFFFFF) >> ASAN_SHADOW_SCALE);
     }
 
+    if (IsCheckShadowBounds() &&
+        (shadow_ptr < launch_info->GlobalShadowLowerBound ||
+         shadow_ptr > launch_info->GlobalShadowUpperBound)) {
+      ASAN_DEBUG(__spirv_ocl_printf(__asan_print_shadow_bound, addr, shadow_ptr,
+                                    launch_info->GlobalShadowLowerBound,
+                                    launch_info->GlobalShadowUpperBound));
+      SaveReport(ErrorType::OUT_OF_BOUNDS, MemoryType::GLOBAL, false, debug);
+      return 0;
+    }
+
     ASAN_DEBUG(
         const auto shadow_offset_end = launch_info->GlobalShadowOffsetEnd;
         if (shadow_ptr > shadow_offset_end) {

llvm/include/llvm/Transforms/Instrumentation/SPIRVSanitizerCommonUtils.h

@@ -50,6 +50,7 @@ enum SanitizedKernelFlags : uint32_t {
   CHECK_PRIVATES = 1U << 3,
   CHECK_GENERICS = 1U << 4,
   MSAN_TRACK_ORIGINS = 1U << 5,
+  ASAN_CHECK_SHADOW_BOUNDS = 1U << 6,
 };
 
 } // namespace llvm

llvm/lib/Transforms/Instrumentation/AddressSanitizer.cpp

@@ -441,27 +441,35 @@ static cl::opt<AsanDtorKind> ClOverrideDestructorKind(
     cl::init(AsanDtorKind::Invalid), cl::Hidden);
 
 // SYCL flags
+static cl::opt<bool> ClSpirOffloadPrivates(
+    "asan-spir-privates",
+    cl::desc("Instrument private pointer on SPIR-V target"), cl::Hidden,
+    cl::init(true));
+
 static cl::opt<bool>
-    ClSpirOffloadPrivates("asan-spir-privates",
-                          cl::desc("instrument private pointer"), cl::Hidden,
-                          cl::init(true));
+    ClSpirOffloadGlobals("asan-spir-globals",
+                         cl::desc("Instrument global pointer on SPIR-V target"),
+                         cl::Hidden, cl::init(true));
 
-static cl::opt<bool> ClSpirOffloadGlobals("asan-spir-globals",
-                                          cl::desc("instrument global pointer"),
-                                          cl::Hidden, cl::init(true));
+static cl::opt<bool>
+    ClSpirOffloadLocals("asan-spir-locals",
+                        cl::desc("Instrument local pointer on SPIR-V target"),
+                        cl::Hidden, cl::init(true));
 
-static cl::opt<bool> ClSpirOffloadLocals("asan-spir-locals",
-                                         cl::desc("instrument local pointer"),
-                                         cl::Hidden, cl::init(true));
+static cl::opt<bool> ClSpirOffloadGenerics(
+    "asan-spir-generics",
+    cl::desc("Instrument generic pointer on SPIR-V target"), cl::Hidden,
+    cl::init(true));
 
 static cl::opt<bool>
-    ClSpirOffloadGenerics("asan-spir-generics",
-                          cl::desc("instrument generic pointer"), cl::Hidden,
-                          cl::init(true));
+    ClDeviceGlobals("asan-device-globals",
+                    cl::desc("Instrument device globals on SPIR-V target"),
+                    cl::Hidden, cl::init(true));
 
-static cl::opt<bool> ClDeviceGlobals("asan-device-globals",
-                                     cl::desc("instrument device globals"),
-                                     cl::Hidden, cl::init(true));
+static cl::opt<bool> ClSpirCheckShadowBounds(
+    "asan-spir-shadow-bounds",
+    cl::desc("Enable checking shadow bounds on SPIR-V target"), cl::Hidden,
+    cl::init(false));
 
 // Debug flags.
 
@@ -1411,7 +1419,8 @@ static void ExtendSpirKernelArgs(Module &M, FunctionAnalysisManager &FAM,
   // following structure:
   //  uptr unmangled_kernel_name
   //  uptr unmangled_kernel_name_size
-  StructType *StructTy = StructType::get(IntptrTy, IntptrTy);
+  //  uptr sanitized_flags
+  StructType *StructTy = StructType::get(IntptrTy, IntptrTy, IntptrTy);
 
   if (!HasESIMD)
     for (Function &F : M) {
@@ -1442,9 +1451,21 @@ static void ExtendSpirKernelArgs(Module &M, FunctionAnalysisManager &FAM,
       KernelNamesBytes.append(KernelName.begin(), KernelName.end());
       auto *KernelNameGV = GetOrCreateGlobalString(
           M, "__asan_kernel", KernelName, kSpirOffloadConstantAS);
+
+      uintptr_t SanitizerFlags = 0;
+      SanitizerFlags |= ClSpirOffloadLocals ? SanitizedKernelFlags::CHECK_LOCALS
+                                            : SanitizedKernelFlags::NO_CHECK;
+      SanitizerFlags |= ClSpirOffloadPrivates
+                            ? SanitizedKernelFlags::CHECK_PRIVATES
+                            : SanitizedKernelFlags::NO_CHECK;
+      SanitizerFlags |= ClSpirCheckShadowBounds != 0
+                            ? SanitizedKernelFlags::ASAN_CHECK_SHADOW_BOUNDS
+                            : SanitizedKernelFlags::NO_CHECK;
+
       SpirKernelsMetadata.emplace_back(ConstantStruct::get(
           StructTy, ConstantExpr::getPointerCast(KernelNameGV, IntptrTy),
-          ConstantInt::get(IntptrTy, KernelName.size())));
+          ConstantInt::get(IntptrTy, KernelName.size()),
+          ConstantInt::get(IntptrTy, SanitizerFlags)));
     }
 
   // Create global variable to record spirv kernels' information
@@ -1632,6 +1653,17 @@ PreservedAnalyses AddressSanitizerPass::run(Module &M,
     ExtendSpirKernelArgs(M, FAM, HasESIMD);
     Modified = true;
 
+    {
+      IRBuilder<> IRB(M.getContext());
+      M.getOrInsertGlobal("__asan_check_shadow_bounds", IRB.getInt32Ty(), [&] {
+        return new GlobalVariable(
+            M, IRB.getInt32Ty(), true, GlobalValue::WeakODRLinkage,
+            ConstantInt::get(IRB.getInt32Ty(), ClSpirCheckShadowBounds),
+            "__asan_check_shadow_bounds", nullptr,
+            llvm::GlobalValue::NotThreadLocal, kSpirOffloadGlobalAS);
+      });
+    }
+
     if (HasESIMD) {
       GlobalStringMap.clear();
       return PreservedAnalyses::none();
@@ -1710,19 +1742,21 @@ static bool isUnsupportedAMDGPUAddrspace(Value *Addr) {
 }
 
 static bool isUnsupportedDeviceGlobal(GlobalVariable *G) {
-  // Non image scope device globals are implemented by device USM, and the
-  // out-of-bounds check for them will be done by sanitizer USM part. So we
-  // exclude them here.
-  if (!G->hasAttribute("sycl-device-image-scope"))
-    return true;
-
   // Skip instrumenting on "__AsanKernelMetadata" etc.
-  if (G->getName().starts_with("__Asan"))
+  if (G->getName().starts_with("__Asan") || G->getName().starts_with("__asan"))
     return true;
 
   if (G->getAddressSpace() == kSpirOffloadLocalAS)
     return !ClSpirOffloadLocals;
 
+  // When shadow bounds check is enabled, we need to instrument all global
+  // variables that user code can access
+  if (ClSpirCheckShadowBounds)
+    return false;
+
+  // Non image scope device globals are implemented by device USM, and the
+  // out-of-bounds check for them will be done by sanitizer USM part. So we
+  // exclude them here.
   Attribute Attr = G->getAttribute("sycl-device-image-scope");
   return (!Attr.isStringAttribute() || Attr.getValueAsString() == "false");
 }

llvm/test/Instrumentation/AddressSanitizer/SPIRV/extend_launch_info_arg.ll

@@ -26,4 +26,4 @@ entry:
 attributes #0 = { sanitize_address }
 ;; sycl-device-global-size = 16 * 2
 ;; sycl-host-access = 0 read-only
-; CHECK: attributes [[ATTR0]] = { "sycl-device-global-size"="32" "sycl-device-image-scope" "sycl-host-access"="0" "sycl-unique-id"="__AsanKernelMetadata833c47834a0b74946e370c23c39607cc" }
+; CHECK: attributes [[ATTR0]] = { "sycl-device-global-size"="48" "sycl-device-image-scope" "sycl-host-access"="0" "sycl-unique-id"="__AsanKernelMetadata833c47834a0b74946e370c23c39607cc" }

sycl/test-e2e/AddressSanitizer/out-of-bounds/USM/arbitrary_access.cpp

@@ -0,0 +1,25 @@
+// REQUIRES: linux, gpu && level_zero
+// RUN: %{build} %device_asan_flags -Xarch_device -mllvm=-asan-spir-shadow-bounds=1 -O0 -g -o %t1.out
+// RUN: %{run} not %t1.out 2>&1 | FileCheck %s
+// RUN: %{build} %device_asan_flags -Xarch_device -mllvm=-asan-spir-shadow-bounds=1 -O2 -g -o %t3.out
+// RUN: %{run} not %t3.out 2>&1 | FileCheck %s
+
+#include <sycl/detail/core.hpp>
+#include <sycl/usm.hpp>
+
+void out_of_bounds_function() { *(int *)0xdeadbeef = 42; }
+// CHECK: out-of-bounds-access
+// CHECK-SAME: 0xdeadbeef
+// CHECK: WRITE of size 4 at kernel {{<.*MyKernel>}}
+// CHECK: {{.*arbitrary_access.cpp}}:[[@LINE-4]]
+
+int main() {
+  sycl::queue Q;
+
+  Q.submit([&](sycl::handler &h) {
+    h.single_task<class MyKernel>([=]() { out_of_bounds_function(); });
+  });
+  Q.wait();
+
+  return 0;
+}