[DevMSAN] Fix missed symbols __msan_memset & __msan_warning (#16477)

zhaomaosu · web-flow · commit 600069d5dd4b · 2024-12-30T08:48:25.000-08:00
diff --git a/libdevice/sanitizer/msan_rtl.cpp b/libdevice/sanitizer/msan_rtl.cpp
@@ -159,6 +159,19 @@ MSAN_MAYBE_WARNING(u16, 2)
 MSAN_MAYBE_WARNING(u32, 4)
 MSAN_MAYBE_WARNING(u64, 8)
 
+DEVICE_EXTERN_C_NOINLINE void
+__msan_warning(const char __SYCL_CONSTANT__ *file, uint32_t line,
+               const char __SYCL_CONSTANT__ *func) {
+  __msan_report_error(1, file, line, func);
+}
+
+DEVICE_EXTERN_C_NOINLINE void
+__msan_warning_noreturn(const char __SYCL_CONSTANT__ *file, uint32_t line,
+                        const char __SYCL_CONSTANT__ *func) {
+  __msan_internal_report_save(1, file, line, func);
+  __devicelib_exit();
+}
+
 DEVICE_EXTERN_C_NOINLINE uptr __msan_get_shadow(uptr addr, uint32_t as) {
   // Return clean shadow (0s) by default
   uptr shadow_ptr = (uptr)CleanShadow;
@@ -187,4 +200,19 @@ DEVICE_EXTERN_C_NOINLINE uptr __msan_get_shadow(uptr addr, uint32_t as) {
   return shadow_ptr;
 }
 
+#define MSAN_MEMSET(as)                                                        \
+  DEVICE_EXTERN_C_NOINLINE void __msan_memset_p##as(                           \
+      __attribute__((address_space(as))) char *dest, int val, size_t size) {   \
+    uptr shadow = __msan_get_shadow((uptr)dest, as);                           \
+    for (size_t i = 0; i < size; i++) {                                        \
+      dest[i] = val;                                                           \
+      ((__SYCL_GLOBAL__ char *)shadow)[i] = 0;                                 \
+    }                                                                          \
+  }
+
+MSAN_MEMSET(0)
+MSAN_MEMSET(1)
+MSAN_MEMSET(3)
+MSAN_MEMSET(4)
+
 #endif // __SPIR__ || __SPIRV__
diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
@@ -233,6 +233,8 @@ static const unsigned kRetvalTLSSize = 800;
 // Accesses sizes are powers of two: 1, 2, 4, 8.
 static const size_t kNumberOfAccessSizes = 4;
 
+static constexpr unsigned kNumberOfAddressSpace = 5;
+
 /// Track origins of uninitialized values.
 ///
 /// Adds a section to MemorySanitizer report that points to the allocation
@@ -678,6 +680,9 @@ class MemorySanitizer {
   /// MSan runtime replacements for memmove, memcpy and memset.
   FunctionCallee MemmoveFn, MemcpyFn, MemsetFn;
 
+  /// MSan runtime replacements for memset with address space.
+  FunctionCallee MemsetOffloadFn[kNumberOfAddressSpace];
+
   /// KMSAN callback for task-local function argument shadow.
   StructType *MsanContextStateTy;
   FunctionCallee MsanGetContextStateFn;
@@ -964,7 +969,19 @@ void MemorySanitizer::createUserspaceApi(Module &M,
   } else {
     StringRef WarningFnName =
         Recover ? "__msan_warning" : "__msan_warning_noreturn";
-    WarningFn = M.getOrInsertFunction(WarningFnName, IRB.getVoidTy());
+    if (!TargetTriple.isSPIROrSPIRV()) {
+      WarningFn = M.getOrInsertFunction(WarningFnName, IRB.getVoidTy());
+    } else {
+      // __msan_warning[_noreturn](
+      //   char* file,
+      //   unsigned int line,
+      //   char* func
+      // )
+      WarningFn = M.getOrInsertFunction(
+          WarningFnName, IRB.getVoidTy(),
+          IRB.getInt8PtrTy(kSpirOffloadConstantAS), IRB.getInt32Ty(),
+          IRB.getInt8PtrTy(kSpirOffloadConstantAS));
+    }
   }
 
   // Create the global TLS variables.
@@ -1050,13 +1067,24 @@ void MemorySanitizer::initializeCallbacks(Module &M,
   MsanSetOriginFn = M.getOrInsertFunction(
       "__msan_set_origin", TLI.getAttrList(C, {2}, /*Signed=*/false),
       IRB.getVoidTy(), PtrTy, IntptrTy, IRB.getInt32Ty());
-  MemmoveFn =
-      M.getOrInsertFunction("__msan_memmove", PtrTy, PtrTy, PtrTy, IntptrTy);
-  MemcpyFn =
-      M.getOrInsertFunction("__msan_memcpy", PtrTy, PtrTy, PtrTy, IntptrTy);
-  MemsetFn = M.getOrInsertFunction("__msan_memset",
-                                   TLI.getAttrList(C, {1}, /*Signed=*/true),
-                                   PtrTy, PtrTy, IRB.getInt32Ty(), IntptrTy);
+  if (!TargetTriple.isSPIROrSPIRV()) {
+    MemmoveFn =
+        M.getOrInsertFunction("__msan_memmove", PtrTy, PtrTy, PtrTy, IntptrTy);
+    MemcpyFn =
+        M.getOrInsertFunction("__msan_memcpy", PtrTy, PtrTy, PtrTy, IntptrTy);
+    MemsetFn = M.getOrInsertFunction("__msan_memset",
+                                     TLI.getAttrList(C, {1}, /*Signed=*/true),
+                                     PtrTy, PtrTy, IRB.getInt32Ty(), IntptrTy);
+  } else {
+    for (unsigned FirstArgAS = 0; FirstArgAS < kNumberOfAddressSpace;
+         FirstArgAS++) {
+      const std::string Suffix = "_p" + itostr(FirstArgAS);
+      PointerType *FirstArgPtrTy = IRB.getPtrTy(FirstArgAS);
+      MemsetOffloadFn[FirstArgAS] = M.getOrInsertFunction(
+          "__msan_memset" + Suffix, TLI.getAttrList(C, {1}, /*Signed=*/true),
+          FirstArgPtrTy, FirstArgPtrTy, IRB.getInt32Ty(), IntptrTy);
+    }
+  }
 
   MsanInstrumentAsmStoreFn = M.getOrInsertFunction(
       "__msan_instrument_asm_store", IRB.getVoidTy(), PtrTy, IntptrTy);
@@ -1560,6 +1588,35 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
     return LazyWarningDebugLocationCount[DebugLoc] >= ClDisambiguateWarning;
   }
 
+  void appendDebugInfoToArgs(IRBuilder<> &IRB, SmallVectorImpl<Value *> &Args) {
+    auto *M = F.getParent();
+    auto &C = IRB.getContext();
+    auto DebugLoc = IRB.getCurrentDebugLocation();
+
+    // SPIR constant address space
+    auto *ConstASPtrTy =
+        PointerType::get(Type::getInt8Ty(C), kSpirOffloadConstantAS);
+
+    // file name and line number
+    if (DebugLoc) {
+      llvm::SmallString<128> Source = DebugLoc->getDirectory();
+      sys::path::append(Source, DebugLoc->getFilename());
+      auto *FileNameGV = getOrCreateGlobalString(*M, "__msan_file", Source,
+                                                 kSpirOffloadConstantAS);
+      Args.push_back(ConstantExpr::getPointerCast(FileNameGV, ConstASPtrTy));
+      Args.push_back(ConstantInt::get(Type::getInt32Ty(C), DebugLoc.getLine()));
+    } else {
+      Args.push_back(ConstantPointerNull::get(ConstASPtrTy));
+      Args.push_back(ConstantInt::get(Type::getInt32Ty(C), 0));
+    }
+
+    // function name
+    auto FuncName = F.getName();
+    auto *FuncNameGV = getOrCreateGlobalString(
+        *M, "__msan_func", demangle(FuncName), kSpirOffloadConstantAS);
+    Args.push_back(ConstantExpr::getPointerCast(FuncNameGV, ConstASPtrTy));
+  }
+
   /// Helper function to insert a warning at IRB's current insert point.
   void insertWarningFn(IRBuilder<> &IRB, Value *Origin) {
     if (!Origin)
@@ -1584,10 +1641,16 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       }
     }
 
-    if (MS.CompileKernel || MS.TrackOrigins)
-      IRB.CreateCall(MS.WarningFn, Origin)->setCannotMerge();
-    else
-      IRB.CreateCall(MS.WarningFn)->setCannotMerge();
+    if (!SpirOrSpirv) {
+      if (MS.CompileKernel || MS.TrackOrigins)
+        IRB.CreateCall(MS.WarningFn, Origin)->setCannotMerge();
+      else
+        IRB.CreateCall(MS.WarningFn)->setCannotMerge();
+    } else { // SPIR or SPIR-V
+      SmallVector<Value *, 3> Args;
+      appendDebugInfoToArgs(IRB, Args);
+      IRB.CreateCall(MS.WarningFn, Args)->setCannotMerge();
+    }
     // FIXME: Insert UnreachableInst if !MS.Recover?
     // This may invalidate some of the following checks and needs to be done
     // at the very end.
@@ -1617,43 +1680,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
             ConvertedShadow2,
             MS.TrackOrigins && Origin ? Origin : (Value *)IRB.getInt32(0)};
 
-        {
-          auto *M = F.getParent();
-          auto *ConstASPtrTy = IRB.getInt8PtrTy(kSpirOffloadConstantAS);
-
-          // file name and line number
-          {
-            bool HasDebugLoc = false;
-            auto *ConvertedShadowInst = dyn_cast<Instruction>(ConvertedShadow);
-
-            if (ConvertedShadowInst) {
-              if (auto &Loc = ConvertedShadowInst->getDebugLoc()) {
-                llvm::SmallString<128> Source = Loc->getDirectory();
-                sys::path::append(Source, Loc->getFilename());
-                auto *FileNameGV = getOrCreateGlobalString(
-                    *M, "__asan_file", Source, kSpirOffloadConstantAS);
-                Args.push_back(
-                    ConstantExpr::getPointerCast(FileNameGV, ConstASPtrTy));
-                Args.push_back(
-                    ConstantInt::get(IRB.getInt32Ty(), Loc.getLine()));
-
-                HasDebugLoc = true;
-              }
-            }
-
-            if (!HasDebugLoc) {
-              Args.push_back(ConstantPointerNull::get(ConstASPtrTy));
-              Args.push_back(ConstantInt::get(IRB.getInt32Ty(), 0));
-            }
-          }
-
-          // function name
-          auto FuncName = F.getName();
-          auto *FuncNameGV = getOrCreateGlobalString(
-              *M, "__asan_func", demangle(FuncName), kSpirOffloadConstantAS);
-          Args.push_back(
-              ConstantExpr::getPointerCast(FuncNameGV, ConstASPtrTy));
-        }
+        appendDebugInfoToArgs(IRB, Args);
 
         CallBase *CB = IRB.CreateCall(Fn, Args);
         CB->addParamAttr(0, Attribute::ZExt);
@@ -3160,7 +3187,10 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
   void visitMemSetInst(MemSetInst &I) {
     IRBuilder<> IRB(&I);
     IRB.CreateCall(
-        MS.MemsetFn,
+        SpirOrSpirv ? MS.MemsetOffloadFn[cast<PointerType>(
+                                             I.getArgOperand(0)->getType())
+                                             ->getAddressSpace()]
+                    : MS.MemsetFn,
         {I.getArgOperand(0),
          IRB.CreateIntCast(I.getArgOperand(1), IRB.getInt32Ty(), false),
          IRB.CreateIntCast(I.getArgOperand(2), MS.IntptrTy, false)});
diff --git a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_large_access_size.ll b/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_large_access_size.ll
@@ -0,0 +1,17 @@
+; 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"
+
+; Function Attrs: sanitize_memory
+define spir_kernel void @MyKernel(<3 x i32> %extractVec.i8.i.i.i) #0 {
+; CHECK-LABEL: @MyKernel
+entry:
+  br label %for.body.i
+
+; CHECK: call void @__msan_warning_noreturn(ptr addrspace(2) null, i32 0, ptr addrspace(2) @__msan_func_MyKernel)
+for.body.i:                                       ; preds = %for.body.i, %entry
+  %div.i.i.i.i.i.i = sdiv <3 x i32> zeroinitializer, %extractVec.i8.i.i.i
+  br label %for.body.i
+}
+
+attributes #0 = { sanitize_memory }
diff --git a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_memset.ll b/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_memset.ll
@@ -0,0 +1,16 @@
+; 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 %offset.i) {
+; CHECK-LABEL: @MyKernel
+entry:
+  call void @llvm.memset.p0.i64(ptr %offset.i, i8 0, i64 0, i1 false)
+; CHECK: call ptr @__msan_memset_p0
+  ret void
+}
+
+; Function Attrs: nocallback nofree nounwind willreturn memory(argmem: write)
+declare void @llvm.memset.p0.i64(ptr nocapture writeonly, i8, i64, i1 immarg) #0
+
+attributes #0 = { nocallback nofree nounwind willreturn memory(argmem: write) }
diff --git a/sycl/test-e2e/MemorySanitizer/check_large_access.cpp b/sycl/test-e2e/MemorySanitizer/check_large_access.cpp
@@ -0,0 +1,26 @@
+// REQUIRES: linux, cpu || (gpu && level_zero)
+// 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
+// RUN: %{run} not %t3.out 2>&1 | FileCheck %s
+
+#include "sycl/detail/core.hpp"
+#include <sycl/vector.hpp>
+
+int main() {
+  sycl::buffer<sycl::int3, 1> b(sycl::range<1>(2));
+  sycl::queue myQueue;
+  myQueue
+      .submit([&](sycl::handler &cgh) {
+        auto B = b.get_access<sycl::access::mode::read_write>(cgh);
+        cgh.parallel_for<class MyKernel>(
+            sycl::range<1>{2}, [=](sycl::id<1> ID) {
+              B[ID] = sycl::int3{(sycl::int3)ID[0]} / B[ID];
+            });
+      })
+      .wait();
+  // CHECK: use-of-uninitialized-value
+  // CHECK: kernel <{{.*MyKernel}}>
+
+  return 0;
+}
