[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.
 
@@ -1394,7 +1402,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) {
@@ -1425,9 +1434,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
@@ -1615,6 +1636,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();
@@ -1693,19 +1725,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;
+}

unified-runtime/source/loader/layers/sanitizer/asan/asan_interceptor.cpp

@@ -267,11 +267,23 @@ ur_result_t AsanInterceptor::preLaunchKernel(ur_kernel_handle_t Kernel,
 
   ur_queue_handle_t InternalQueue = ContextInfo->getInternalQueue(Device);
 
+  // To get right shadow boundary, shadow memory should be updated before
+  // prepareLaunch
+  {
+    // Force to allocate membuffer before prepareLaunch
+    auto &KernelInfo = getOrCreateKernelInfo(Kernel);
+    std::shared_lock<ur_shared_mutex> Guard(KernelInfo.Mutex);
+    for (const auto &[ArgIndex, MemBuffer] : KernelInfo.BufferArgs) {
+      char *ArgPointer = nullptr;
+      UR_CALL(MemBuffer->getHandle(DeviceInfo->Handle, ArgPointer));
+      (void)ArgPointer;
+    }
+  }
+  UR_CALL(updateShadowMemory(ContextInfo, DeviceInfo, InternalQueue));
+
   UR_CALL(prepareLaunch(ContextInfo, DeviceInfo, InternalQueue, Kernel,
                         LaunchInfo));
 
-  UR_CALL(updateShadowMemory(ContextInfo, DeviceInfo, InternalQueue));
-
   UR_CALL(getContext()->urDdiTable.Queue.pfnFinish(InternalQueue));
 
   return UR_RESULT_SUCCESS;
@@ -453,7 +465,7 @@ ur_result_t AsanInterceptor::unregisterProgram(ur_program_handle_t Program) {
   }
   ProgramInfo->AllocInfoForGlobals.clear();
 
-  ProgramInfo->InstrumentedKernels.clear();
+  ProgramInfo->KernelMetadataMap.clear();
 
   return UR_RESULT_SUCCESS;
 }
@@ -508,14 +520,18 @@ ur_result_t AsanInterceptor::registerSpirKernels(ur_program_handle_t Program) {
 
       std::string KernelName =
           std::string(KernelNameV.begin(), KernelNameV.end());
+      bool CheckShadowBounds =
+          SKI.Flags & SanitizedKernelFlags::ASAN_CHECK_SHADOW_BOUNDS;
 
       UR_LOG_L(getContext()->logger, INFO,
-               "SpirKernel(name='{}', isInstrumented={})", KernelName, true);
+               "SpirKernel(name='{}', isInstrumented={}, checkShadowBounds={})",
+               KernelName, true, CheckShadowBounds);
 
-      PI->InstrumentedKernels.insert(std::move(KernelName));
+      PI->KernelMetadataMap[KernelName] =
+          ProgramInfo::KernelMetadata{CheckShadowBounds};
     }
     UR_LOG_L(getContext()->logger, INFO, "Number of sanitized kernel: {}",
-             PI->InstrumentedKernels.size());
+             PI->KernelMetadataMap.size());
   }
 
   return UR_RESULT_SUCCESS;
@@ -666,11 +682,16 @@ KernelInfo &AsanInterceptor::getOrCreateKernelInfo(ur_kernel_handle_t Kernel) {
   auto Program = GetProgram(Kernel);
   auto PI = getProgramInfo(Program);
   assert(PI != nullptr && "unregistered program!");
-  bool IsInstrumented = PI->isKernelInstrumented(Kernel);
+
+  auto KI = std::make_unique<KernelInfo>(Kernel);
+  KI->IsInstrumented = PI->isKernelInstrumented(Kernel);
+  if (KI->IsInstrumented) {
+    auto &KM = PI->getKernelMetadata(Kernel);
+    KI->IsCheckShadowBounds = KM.CheckShadowBounds;
+  }
 
   std::scoped_lock<ur_shared_mutex> Guard(m_KernelMapMutex);
-  m_KernelMap.emplace(Kernel,
-                      std::make_unique<KernelInfo>(Kernel, IsInstrumented));
+  m_KernelMap.emplace(Kernel, std::move(KI));
   return *m_KernelMap[Kernel].get();
 }
 
@@ -815,6 +836,13 @@ ur_result_t AsanInterceptor::prepareLaunch(
   LaunchInfo.Data.Host.DeviceTy = DeviceInfo->Type;
   LaunchInfo.Data.Host.Debug = getContext()->Options.Debug ? 1 : 0;
 
+  if (KernelInfo.IsCheckShadowBounds) {
+    LaunchInfo.Data.Host.GlobalShadowLowerBound =
+        DeviceInfo->Shadow->ShadowLowerBound;
+    LaunchInfo.Data.Host.GlobalShadowUpperBound =
+        DeviceInfo->Shadow->ShadowUpperBound;
+  }
+
   // Write shadow memory offset for local memory
   if (getContext()->Options.DetectLocals) {
     if (DeviceInfo->Shadow->AllocLocalShadow(
@@ -872,18 +900,20 @@ ur_result_t AsanInterceptor::prepareLaunch(
   // sync asan runtime data to device side
   UR_CALL(LaunchInfo.Data.syncToDevice(Queue));
 
-  UR_LOG_L(getContext()->logger, INFO,
-           "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, PrivateBase={}, "
-           "PrivateShadow={}, LocalArgs={}, NumLocalArgs={}, "
-           "Device={}, Debug={})",
-           (void *)LaunchInfo.Data.getDevicePtr(),
-           (void *)LaunchInfo.Data.Host.GlobalShadowOffset,
-           (void *)LaunchInfo.Data.Host.LocalShadowOffset,
-           (void *)LaunchInfo.Data.Host.PrivateBase,
-           (void *)LaunchInfo.Data.Host.PrivateShadowOffset,
-           (void *)LaunchInfo.Data.Host.LocalArgs,
-           LaunchInfo.Data.Host.NumLocalArgs,
-           ToString(LaunchInfo.Data.Host.DeviceTy), LaunchInfo.Data.Host.Debug);
+  UR_LOG_L(
+      getContext()->logger, INFO,
+      "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, PrivateBase={}, "
+      "PrivateShadow={}, GlobalShadowLowerBound={}, GlobalShadowUpperBound={}, "
+      "LocalArgs={}, NumLocalArgs={}, Device={}, Debug={})",
+      (void *)LaunchInfo.Data.getDevicePtr(),
+      (void *)LaunchInfo.Data.Host.GlobalShadowOffset,
+      (void *)LaunchInfo.Data.Host.LocalShadowOffset,
+      (void *)LaunchInfo.Data.Host.PrivateBase,
+      (void *)LaunchInfo.Data.Host.PrivateShadowOffset,
+      (void *)LaunchInfo.Data.Host.GlobalShadowLowerBound,
+      (void *)LaunchInfo.Data.Host.GlobalShadowUpperBound,
+      (void *)LaunchInfo.Data.Host.LocalArgs, LaunchInfo.Data.Host.NumLocalArgs,
+      ToString(LaunchInfo.Data.Host.DeviceTy), LaunchInfo.Data.Host.Debug);
 
   return UR_RESULT_SUCCESS;
 }
@@ -920,7 +950,14 @@ AsanInterceptor::findAllocInfoByContext(ur_context_handle_t Context) {
 
 bool ProgramInfo::isKernelInstrumented(ur_kernel_handle_t Kernel) const {
   const auto Name = GetKernelName(Kernel);
-  return InstrumentedKernels.find(Name) != InstrumentedKernels.end();
+  return KernelMetadataMap.find(Name) != KernelMetadataMap.end();
+}
+
+const ProgramInfo::KernelMetadata &
+ProgramInfo::getKernelMetadata(ur_kernel_handle_t Kernel) const {
+  const auto Name = GetKernelName(Kernel);
+  assert(KernelMetadataMap.find(Name) != KernelMetadataMap.end());
+  return KernelMetadataMap.at(Name);
 }
 
 ContextInfo::~ContextInfo() {