[DeviceMSAN] Fix false negative report due to __spirv_GroupAsyncCopy

libdevice/sanitizer/msan_rtl.cpp

@@ -220,6 +220,14 @@ inline void __msan_exit() {
     __devicelib_exit();
 }
 
+template <typename dataT>
+void GroupAsyncCopy(dataT *Dest, const dataT *Src, size_t NumElements,
+                    size_t Stride) {
+  for (size_t i = 0; i < NumElements; i++) {
+    Dest[i] = Src[i * Stride];
+  }
+}
+
 } // namespace
 
 #define MSAN_MAYBE_WARNING(type, size)                                         \
@@ -576,4 +584,46 @@ DEVICE_EXTERN_C_NOINLINE void __msan_unpoison_stack(__SYCL_PRIVATE__ void *ptr,
       __spirv_ocl_printf(__msan_print_func_end, "__msan_unpoison_stack"));
 }
 
+static __SYCL_CONSTANT__ const char __msan_print_strided_copy_unsupport_type[] =
+    "[kernel] __msan_unpoison_strided_copy: unsupport type(%d)\n";
+
+DEVICE_EXTERN_C_NOINLINE void
+__msan_unpoison_strided_copy(uptr dest, uint32_t dest_as, uptr src,
+                             uint32_t src_as, uint32_t element_size,
+                             uptr counts, uptr stride) {
+  if (!GetMsanLaunchInfo)
+    return;
+
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_func_beg,
+                                "__msan_unpoison_strided_copy"));
+
+  uptr shadow_dest = __msan_get_shadow(dest, dest_as);
+  uptr shadow_src = __msan_get_shadow(src, src_as);
+
+  switch (element_size) {
+  case 1:
+    GroupAsyncCopy<int8_t>((int8_t *)shadow_dest, (int8_t *)shadow_src, counts,
+                           stride);
+    break;
+  case 2:
+    GroupAsyncCopy<int16_t>((int16_t *)shadow_dest, (int16_t *)shadow_src,
+                            counts, stride);
+    break;
+  case 4:
+    GroupAsyncCopy<int32_t>((int32_t *)shadow_dest, (int32_t *)shadow_src,
+                            counts, stride);
+    break;
+  case 8:
+    GroupAsyncCopy<int64_t>((int64_t *)shadow_dest, (int64_t *)shadow_src,
+                            counts, stride);
+    break;
+  default:
+    MSAN_DEBUG(__spirv_ocl_printf(__msan_print_strided_copy_unsupport_type,
+                                  element_size));
+  }
+
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_func_end,
+                                "__msan_unpoison_strided_copy"));
+}
+
 #endif // __SPIR__ || __SPIRV__

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -804,6 +804,8 @@ class MemorySanitizerOnSpirv {
   void initializeKernelCallerMap(Function *F);
 
 private:
+  friend struct MemorySanitizerVisitor;
+
   Module &M;
   LLVMContext &C;
   const DataLayout &DL;
@@ -830,6 +832,7 @@ class MemorySanitizerOnSpirv {
   FunctionCallee MsanUnpoisonShadowDynamicLocalFunc;
   FunctionCallee MsanBarrierFunc;
   FunctionCallee MsanUnpoisonStackFunc;
+  FunctionCallee MsanUnpoisonStridedCopyFunc;
 };
 
 } // end anonymous namespace
@@ -920,6 +923,11 @@ void MemorySanitizerOnSpirv::initializeCallbacks() {
   // )
   MsanUnpoisonStackFunc = M.getOrInsertFunction(
       "__msan_unpoison_stack", IRB.getVoidTy(), PtrTy, IntptrTy);
+
+  MsanUnpoisonStridedCopyFunc = M.getOrInsertFunction(
+      "__msan_unpoison_strided_copy", IRB.getVoidTy(), IntptrTy,
+      IRB.getInt32Ty(), IntptrTy, IRB.getInt32Ty(), IRB.getInt32Ty(),
+      IRB.getInt64Ty(), IRB.getInt64Ty());
 }
 
 // Handle global variables:
@@ -1810,7 +1818,7 @@ static void setNoSanitizedMetadataSPIR(Instruction &I) {
         }
       } else {
         auto FuncName = Func->getName();
-        if (FuncName.contains("__spirv_"))
+        if (FuncName.contains("__spirv_") && !FuncName.contains("__spirv_GroupAsyncCopy"))
           I.setNoSanitizeMetadata();
       }
     }
@@ -1820,6 +1828,50 @@ static void setNoSanitizedMetadataSPIR(Instruction &I) {
     I.setNoSanitizeMetadata();
 }
 
+// This is not a gerneral-purpose function, but a helper for demangling
+// "__spirv_GroupAsyncCopy" function name
+static int getTypeSizeFromManglingName(StringRef Name) {
+  auto GetTypeSize = [](const char C) {
+    switch (C) {
+    case 'c':
+      return 1;
+    case 'd':
+      return 8;
+    case 'f':
+      return 4;
+    case 'i':
+      return 4;
+    case 'l':
+      return 8;
+    case 's':
+      return 2;
+    default:
+      return 0;
+    }
+  };
+
+  // 1. Basic type
+  if (Name[0] != 'D')
+    return GetTypeSize(Name[0]);
+
+  // 2. Vector type
+
+  // Drop "Dv"
+  assert(Name[0] == 'D' && Name[1] == 'v' &&
+         "Invalid mangling name for vector type");
+  Name = Name.drop_front(2);
+
+  // Vector length
+  int Len = std::stoi(Name.str());
+  Name = Name.drop_front(Len >= 10 ? 2 : 1);
+
+  assert(Name[0] == '_' && "Invalid mangling name for vector type");
+  Name = Name.drop_front(1);
+
+  int Size = GetTypeSize(Name[0]);
+  return Len * Size;
+}
+
 namespace {
 
 /// Helper class to attach debug information of the given instruction onto new
@@ -6000,6 +6052,7 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
 
   void visitCallBase(CallBase &CB) {
     assert(!CB.getMetadata(LLVMContext::MD_nosanitize));
+
     if (CB.isInlineAsm()) {
       // For inline asm (either a call to asm function, or callbr instruction),
       // do the usual thing: check argument shadow and mark all outputs as
@@ -6157,6 +6210,34 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
       VAHelper->visitCallBase(CB, IRB);
     }
 
+    if (SpirOrSpirv) {
+      auto *Func = CB.getCalledFunction();
+      if (Func) {
+        auto FuncName = Func->getName();
+        outs() << "!!!! " << FuncName << "\n";
+        if (FuncName.contains("__spirv_GroupAsyncCopy")) {
+          // Like "_Z22__spirv_GroupAsyncCopyiPU3AS3dPU3AS1dllP13__spirv_Event"
+          // "__spirv_GroupAsyncCopy(int, double AS3*, double AS1*, long, long, __spirv_Event*)"
+
+          auto *Dest = CB.getArgOperand(1);
+          auto *Src = CB.getArgOperand(2);
+          auto *NumElements = CB.getArgOperand(3);
+          auto *Stride = CB.getArgOperand(4);
+
+          // Skip "_Z22__spirv_GroupAsyncCopyiPU3AS3", get the size of parameter type directly
+          int ElementSize = getTypeSizeFromManglingName(FuncName.substr(33));
+
+          IRB.CreateCall(
+              MS.Spirv.MsanUnpoisonStridedCopyFunc,
+              {IRB.CreatePointerCast(Dest, MS.Spirv.IntptrTy),
+               IRB.getInt32(Dest->getType()->getPointerAddressSpace()),
+               IRB.CreatePointerCast(Src, MS.Spirv.IntptrTy),
+               IRB.getInt32(Src->getType()->getPointerAddressSpace()),
+               IRB.getInt32(ElementSize), NumElements, Stride});
+        }
+      }
+    }
+
     // Now, get the shadow for the RetVal.
     if (!CB.getType()->isSized())
       return;