[DevMSAN] Skip more unsupported memory access pattern (#16427)

UR: oneapi-src/unified-runtime#2484

---------

Co-authored-by: Aaron Greig <aaron.greig@codeplay.com>

Author: zhaomaosu

libdevice/include/sanitizer_defs.hpp

@@ -7,6 +7,7 @@
 //===----------------------------------------------------------------------===//
 #pragma once
 
+#include "spir_global_var.hpp"
 #include <cstdint>
 
 using uptr = uintptr_t;
@@ -19,6 +20,59 @@ using s16 = int16_t;
 using s32 = int32_t;
 using s64 = int64_t;
 
+enum ADDRESS_SPACE : uint32_t {
+  ADDRESS_SPACE_PRIVATE = 0,
+  ADDRESS_SPACE_GLOBAL = 1,
+  ADDRESS_SPACE_CONSTANT = 2,
+  ADDRESS_SPACE_LOCAL = 3,
+  ADDRESS_SPACE_GENERIC = 4,
+};
+
 #define LIKELY(x) __builtin_expect(!!(x), 1)
 #define UNLIKELY(x) __builtin_expect(!!(x), 0)
 #define NORETURN __declspec(noreturn)
+
+#if defined(__SPIR__) || defined(__SPIRV__)
+
+#if defined(__SYCL_DEVICE_ONLY__)
+
+#define __USE_SPIR_BUILTIN__ 1
+
+#ifndef SYCL_EXTERNAL
+#define SYCL_EXTERNAL
+#endif // SYCL_EXTERNAL
+
+#else // __SYCL_DEVICE_ONLY__
+
+#define __USE_SPIR_BUILTIN__
+
+#endif // __SYCL_DEVICE_ONLY__
+
+#if __USE_SPIR_BUILTIN__
+extern SYCL_EXTERNAL int
+__spirv_ocl_printf(const __SYCL_CONSTANT__ char *Format, ...);
+
+extern SYCL_EXTERNAL __SYCL_GLOBAL__ void *
+__spirv_GenericCastToPtrExplicit_ToGlobal(void *, int);
+extern SYCL_EXTERNAL __SYCL_LOCAL__ void *
+__spirv_GenericCastToPtrExplicit_ToLocal(void *, int);
+extern SYCL_EXTERNAL __SYCL_PRIVATE__ void *
+__spirv_GenericCastToPtrExplicit_ToPrivate(void *, int);
+
+extern SYCL_EXTERNAL __attribute__((convergent)) void
+__spirv_ControlBarrier(uint32_t Execution, uint32_t Memory, uint32_t Semantics);
+
+extern "C" SYCL_EXTERNAL void __devicelib_exit();
+#endif // __USE_SPIR_BUILTIN__
+
+__SYCL_GLOBAL__ void *ToGlobal(void *ptr) {
+  return __spirv_GenericCastToPtrExplicit_ToGlobal(ptr, 5);
+}
+__SYCL_LOCAL__ void *ToLocal(void *ptr) {
+  return __spirv_GenericCastToPtrExplicit_ToLocal(ptr, 4);
+}
+__SYCL_PRIVATE__ void *ToPrivate(void *ptr) {
+  return __spirv_GenericCastToPtrExplicit_ToPrivate(ptr, 7);
+}
+
+#endif // __SPIR__ || __SPIRV__

libdevice/sanitizer/asan_rtl.cpp

@@ -17,37 +17,6 @@ __SYCL_GLOBAL__ uptr *__SYCL_LOCAL__ __AsanLaunchInfo;
 
 #if defined(__SPIR__) || defined(__SPIRV__)
 
-#if defined(__SYCL_DEVICE_ONLY__)
-
-#define __USE_SPIR_BUILTIN__ 1
-
-#ifndef SYCL_EXTERNAL
-#define SYCL_EXTERNAL
-#endif // SYCL_EXTERNAL
-
-#else // __SYCL_DEVICE_ONLY__
-
-#define __USE_SPIR_BUILTIN__
-
-#endif // __SYCL_DEVICE_ONLY__
-
-#if __USE_SPIR_BUILTIN__
-extern SYCL_EXTERNAL int
-__spirv_ocl_printf(const __SYCL_CONSTANT__ char *Format, ...);
-
-extern SYCL_EXTERNAL __SYCL_GLOBAL__ void *
-__spirv_GenericCastToPtrExplicit_ToGlobal(void *, int);
-extern SYCL_EXTERNAL __SYCL_LOCAL__ void *
-__spirv_GenericCastToPtrExplicit_ToLocal(void *, int);
-extern SYCL_EXTERNAL __SYCL_PRIVATE__ void *
-__spirv_GenericCastToPtrExplicit_ToPrivate(void *, int);
-
-extern SYCL_EXTERNAL __attribute__((convergent)) void
-__spirv_ControlBarrier(uint32_t Execution, uint32_t Memory, uint32_t Semantics);
-
-extern "C" SYCL_EXTERNAL void __devicelib_exit();
-#endif // __USE_SPIR_BUILTIN__
-
 static const __SYCL_CONSTANT__ char __asan_shadow_value_start[] =
     "[kernel] %p(%d) -> %p:";
 static const __SYCL_CONSTANT__ char __asan_shadow_value[] = " %02X";
@@ -87,14 +56,6 @@ static __SYCL_CONSTANT__ const char __generic_to[] =
     }                                                                          \
   } while (false)
 
-enum ADDRESS_SPACE : uint32_t {
-  ADDRESS_SPACE_PRIVATE = 0,
-  ADDRESS_SPACE_GLOBAL = 1,
-  ADDRESS_SPACE_CONSTANT = 2,
-  ADDRESS_SPACE_LOCAL = 3,
-  ADDRESS_SPACE_GENERIC = 4,
-};
-
 namespace {
 
 void __asan_report_unknown_device();

libdevice/sanitizer/msan_rtl.cpp

@@ -18,8 +18,6 @@ constexpr int MSAN_REPORT_NONE = 0;
 constexpr int MSAN_REPORT_START = 1;
 constexpr int MSAN_REPORT_FINISH = 2;
 
-static const uint64_t CleanShadow[16] = {};
-
 static const __SYCL_CONSTANT__ char __msan_print_warning_return[] =
     "[kernel] !!! msan warning return\n";
 
@@ -38,17 +36,10 @@ static const __SYCL_CONSTANT__ char __msan_print_report[] =
 static const __SYCL_CONSTANT__ char __msan_print_unsupport_device_type[] =
     "[kernel] Unsupport device type: %d\n";
 
-#if defined(__SPIR__) || defined(__SPIRV__)
-
-#if defined(__SYCL_DEVICE_ONLY__)
-#define __USE_SPIR_BUILTIN__ 1
-#endif
+static const __SYCL_CONSTANT__ char __msan_print_generic_to[] =
+    "[kernel] %p(4) - %p(%d)\n";
 
-#if __USE_SPIR_BUILTIN__
-extern SYCL_EXTERNAL int
-__spirv_ocl_printf(const __SYCL_CONSTANT__ char *Format, ...);
-extern "C" SYCL_EXTERNAL void __devicelib_exit();
-#endif
+#if defined(__SPIR__) || defined(__SPIRV__)
 
 #define MSAN_DEBUG(X)                                                          \
   do {                                                                         \
@@ -61,6 +52,21 @@ extern "C" SYCL_EXTERNAL void __devicelib_exit();
 
 namespace {
 
+inline void ConvertGenericPointer(uptr &addr, uint32_t &as) {
+  auto old = addr;
+  if ((addr = (uptr)ToPrivate((void *)old))) {
+    as = ADDRESS_SPACE_PRIVATE;
+  } else if ((addr = (uptr)ToLocal((void *)old))) {
+    as = ADDRESS_SPACE_LOCAL;
+  } else {
+    // FIXME: I'm not sure if we need to check ADDRESS_SPACE_CONSTANT,
+    // but this can really simplify the generic pointer conversion logic
+    as = ADDRESS_SPACE_GLOBAL;
+    addr = old;
+  }
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_generic_to, old, addr, as));
+}
+
 void __msan_internal_report_save(const uint32_t size,
                                  const char __SYCL_CONSTANT__ *file,
                                  const uint32_t line,
@@ -131,6 +137,13 @@ inline uptr __msan_get_shadow_cpu(uptr addr) {
 }
 
 inline uptr __msan_get_shadow_pvc(uptr addr, uint32_t as) {
+  if (as == ADDRESS_SPACE_GENERIC) {
+    ConvertGenericPointer(addr, as);
+    if (as != ADDRESS_SPACE_GLOBAL)
+      return (uptr)((__SYCL_GLOBAL__ MsanLaunchInfo *)__MsanLaunchInfo.get())
+          ->CleanShadow;
+  }
+
   // Device USM only
   auto shadow_begin = ((__SYCL_GLOBAL__ MsanLaunchInfo *)__MsanLaunchInfo.get())
                           ->GlobalShadowOffset;
@@ -174,7 +187,9 @@ __msan_warning_noreturn(const char __SYCL_CONSTANT__ *file, uint32_t line,
 
 DEVICE_EXTERN_C_NOINLINE uptr __msan_get_shadow(uptr addr, uint32_t as) {
   // Return clean shadow (0s) by default
-  uptr shadow_ptr = (uptr)CleanShadow;
+  uptr shadow_ptr =
+      (uptr)((__SYCL_GLOBAL__ MsanLaunchInfo *)__MsanLaunchInfo.get())
+          ->CleanShadow;
 
   if (UNLIKELY(!__MsanLaunchInfo)) {
     __spirv_ocl_printf(__msan_print_warning_nolaunchinfo);

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -945,11 +945,21 @@ void MemorySanitizer::createKernelApi(Module &M, const TargetLibraryInfo &TLI) {
 }
 
 static Constant *getOrInsertGlobal(Module &M, StringRef Name, Type *Ty) {
-  return M.getOrInsertGlobal(Name, Ty, [&] {
-    return new GlobalVariable(M, Ty, false, GlobalVariable::ExternalLinkage,
-                              nullptr, Name, nullptr,
-                              GlobalVariable::InitialExecTLSModel);
-  });
+  // FIXME: spirv target doesn't support TLS, need to handle it later.
+  if (Triple(M.getTargetTriple()).isSPIROrSPIRV()) {
+    return M.getOrInsertGlobal(Name, Ty, [&] {
+      return new GlobalVariable(M, Ty, false, GlobalVariable::ExternalLinkage,
+                                Constant::getNullValue(Ty), Name, nullptr,
+                                GlobalVariable::NotThreadLocal,
+                                kSpirOffloadGlobalAS);
+    });
+  } else {
+    return M.getOrInsertGlobal(Name, Ty, [&] {
+      return new GlobalVariable(M, Ty, false, GlobalVariable::ExternalLinkage,
+                                nullptr, Name, nullptr,
+                                GlobalVariable::InitialExecTLSModel);
+    });
+  }
 }
 
 /// Insert declarations for userspace-specific functions and globals.
@@ -1275,6 +1285,11 @@ static unsigned TypeSizeToSizeIndex(TypeSize TS) {
 }
 
 static bool isUnsupportedSPIRAccess(const Value *Addr, Instruction *I) {
+  if (isa<Instruction>(Addr) &&
+      cast<Instruction>(Addr)->getMetadata(LLVMContext::MD_nosanitize)) {
+    return true;
+  }
+
   // Skip SPIR-V built-in varibles
   auto *OrigValue = Addr->stripInBoundsOffsets();
   assert(OrigValue != nullptr);
@@ -1285,8 +1300,9 @@ static bool isUnsupportedSPIRAccess(const Value *Addr, Instruction *I) {
   switch (PtrTy->getPointerAddressSpace()) {
   case kSpirOffloadPrivateAS:
   case kSpirOffloadLocalAS:
-  case kSpirOffloadGenericAS:
     return true;
+  case kSpirOffloadGenericAS:
+    return false;
   }
 
   return false;
@@ -1302,8 +1318,12 @@ static void setNoSanitizedMetadataSPIR(Instruction &I) {
     Addr = RMW->getPointerOperand();
   else if (const auto *XCHG = dyn_cast<AtomicCmpXchgInst>(&I))
     Addr = XCHG->getPointerOperand();
+  else if (const auto *ASC = dyn_cast<AddrSpaceCastInst>(&I))
+    Addr = ASC->getPointerOperand();
   else if (isa<MemCpyInst>(&I))
     I.setNoSanitizeMetadata();
+  else if (isa<AllocaInst>(&I))
+    I.setNoSanitizeMetadata();
   else if (const auto *CI = dyn_cast<CallInst>(&I)) {
     auto *Func = CI->getCalledFunction();
     if (Func) {
@@ -1327,6 +1347,11 @@ static void setNoSanitizedMetadataSPIR(Instruction &I) {
           Addr = CI->getOperand(OpOffset);
           break;
         }
+        case Intrinsic::lifetime_start:
+        case Intrinsic::lifetime_end: {
+          I.setNoSanitizeMetadata();
+          break;
+        }
         }
       } else {
         auto FuncName = Func->getName();
@@ -2263,7 +2288,10 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       unsigned ArgOffset = 0;
       const DataLayout &DL = F->getDataLayout();
       for (auto &FArg : F->args()) {
-        if (!FArg.getType()->isSized() || FArg.getType()->isScalableTy()) {
+        // FIXME: Need to find a reasonable way to handle byval arguments for
+        // spirv target.
+        if (!FArg.getType()->isSized() || FArg.getType()->isScalableTy() ||
+            (SpirOrSpirv && FArg.hasByValAttr())) {
           LLVM_DEBUG(dbgs() << (FArg.getType()->isScalableTy()
                                     ? "vscale not fully supported\n"
                                     : "Arg is not sized\n"));
@@ -3185,6 +3213,11 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
   // Same as memcpy.
   void visitMemSetInst(MemSetInst &I) {
+    if (SpirOrSpirv && isa<Instruction>(I.getArgOperand(0)) &&
+        cast<Instruction>(I.getArgOperand(0))
+            ->getMetadata(LLVMContext::MD_nosanitize))
+      return;
+
     IRBuilder<> IRB(&I);
     IRB.CreateCall(
         SpirOrSpirv ? MS.MemsetOffloadFn[cast<PointerType>(
@@ -4874,7 +4907,8 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       bool NoUndef = CB.paramHasAttr(i, Attribute::NoUndef);
       bool EagerCheck = MayCheckCall && !ByVal && NoUndef;
 
-      if (EagerCheck) {
+      // Always do eager check for spirv target
+      if (EagerCheck || SpirOrSpirv) {
         insertShadowCheck(A, &CB);
         Size = DL.getTypeAllocSize(A->getType());
       } else {

llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_unsupported_access.ll

@@ -0,0 +1,21 @@
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -S | FileCheck %s
+target datalayout = "e-i64:64-v16:16-v24:32-v32:32-v48:64-v96:128-v192:256-v256:256-v512:512-v1024:1024-n8:16:32:64-G1"
+target triple = "spir64-unknown-unknown"
+
+define spir_kernel void @MyKernel(ptr %__SYCLKernel) {
+; CHECK-LABEL: @MyKernel
+entry:
+  %_arg_array1.addr = alloca ptr addrspace(1), i32 0, align 8
+; CHECK: %_arg_array1.addr = alloca ptr addrspace(1){{.*!nosanitize}}
+  %_arg_array1.addr.ascast = addrspacecast ptr %_arg_array1.addr to ptr addrspace(4)
+; CHECK: %_arg_array1.addr.ascast = addrspacecast ptr %_arg_array1.addr to ptr addrspace(4){{.*!nosanitize}}
+  call void @llvm.lifetime.start.p0(i64 64, ptr %__SYCLKernel)
+; CHECK: call void @llvm.lifetime.start.p0(i64 64, ptr %__SYCLKernel){{.*!nosanitize}}
+  ret void
+}
+
+; Function Attrs: nocallback nofree nosync nounwind willreturn memory(argmem: readwrite)
+declare void @llvm.lifetime.start.p0(i64 immarg, ptr nocapture) #0
+
+attributes #0 = { nocallback nofree nosync nounwind willreturn memory(argmem: readwrite) }
+

sycl/cmake/modules/UnifiedRuntimeTag.cmake

@@ -1,7 +1,7 @@
-# commit fe6c83a3a40e69eaa0c4829c8ba99b03e05bb4e2
-# Merge: 93fc1331 cf0c9482
+# commit 13f5e52ac047bb81f478f0dbe9416a8bc179130a
+# Merge: fe6c83a3 fa1e678a
 # Author: aarongreig <aaron.greig@codeplay.com>
-# Date:   Fri Dec 27 17:39:52 2024 +0000
-#     Merge pull request #2507 from AllanZyne/review/yang/fix_msan_shadow
-#     [DeviceMSAN] Fix MemToShadow algorithm and VA reservation
-set(UNIFIED_RUNTIME_TAG fe6c83a3a40e69eaa0c4829c8ba99b03e05bb4e2)
+# Date:   Mon Dec 30 15:14:19 2024 +0000
+#     Merge pull request #2484 from zhaomaosu/move-clean-shadow-to-launchinfo
+#     [DevMSAN] Move clean shadow into launch info
+set(UNIFIED_RUNTIME_TAG 13f5e52ac047bb81f478f0dbe9416a8bc179130a)

sycl/test-e2e/MemorySanitizer/check_buffer.cpp

@@ -1,4 +1,6 @@
 // REQUIRES: linux, cpu || (gpu && level_zero)
+// RUN: %{build} %device_msan_flags -O0 -g -o %t2.out
+// RUN: %{run} not %t2.out 2>&1 | FileCheck %s
 // RUN: %{build} %device_msan_flags -O1 -g -o %t2.out
 // RUN: %{run} not %t2.out 2>&1 | FileCheck %s
 // RUN: %{build} %device_msan_flags -O2 -g -o %t3.out

sycl/test-e2e/MemorySanitizer/check_call.cpp

@@ -1,4 +1,6 @@
 // REQUIRES: linux, cpu || (gpu && level_zero)
+// RUN: %{build} %device_msan_flags -O0 -g -o %t2.out
+// RUN: %{run} not %t2.out 2>&1 | FileCheck %s
 // RUN: %{build} %device_msan_flags -O1 -g -o %t2.out
 // RUN: %{run} not %t2.out 2>&1 | FileCheck %s
 // RUN: %{build} %device_msan_flags -O2 -g -o %t3.out

sycl/test-e2e/MemorySanitizer/check_divide.cpp

@@ -1,4 +1,6 @@
 // REQUIRES: linux, cpu || (gpu && level_zero)
+// RUN: %{build} %device_msan_flags -O0 -g -o %t2.out
+// RUN: %{run} not %t2.out 2>&1 | FileCheck %s
 // RUN: %{build} %device_msan_flags -O1 -g -o %t2.out
 // RUN: %{run} not %t2.out 2>&1 | FileCheck %s
 // RUN: %{build} %device_msan_flags -O2 -g -o %t3.out