[DevMSAN] Unpoison sret argument for builtin function to get spec constant (#19800)

* For builtin func like "__sycl_getComposite2020SpecConstantValue", if
structs which are larger than 64b will be returned via sret arguments
and will be initialized inside the function. So we need to unpoison the
sret arguments.
* Always set ZE_AFFINITY_MAST to 0 when running device sanitizer tests
since device sanitizer only support one GPU card now.

Author: zhaomaosu

libdevice/sanitizer/msan_rtl.cpp

@@ -671,7 +671,7 @@ __msan_unpoison_shadow_dynamic_local(uptr ptr, uint32_t num_args) {
                                 "__msan_unpoison_shadow_dynamic_local"));
 }
 
-static __SYCL_CONSTANT__ const char __msan_print_set_shadow_private[] =
+static __SYCL_CONSTANT__ const char __msan_print_set_shadow[] =
     "[kernel] __msan_set_value(beg=%p, end=%p, val=%02X)\n";
 
 // We outline the function of setting shadow memory of private memory, because
@@ -684,8 +684,7 @@ DEVICE_EXTERN_C_NOINLINE void __msan_poison_stack(__SYCL_PRIVATE__ void *ptr,
   MSAN_DEBUG(__spirv_ocl_printf(__msan_print_func_beg, "__msan_poison_stack"));
 
   auto shadow_address = MemToShadow((uptr)ptr, ADDRESS_SPACE_PRIVATE);
-  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow_private,
-                                (void *)shadow_address,
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow, (void *)shadow_address,
                                 (void *)(shadow_address + size), 0xff));
 
   if (shadow_address != GetMsanLaunchInfo->CleanShadow) {
@@ -704,8 +703,7 @@ DEVICE_EXTERN_C_NOINLINE void __msan_unpoison_stack(__SYCL_PRIVATE__ void *ptr,
       __spirv_ocl_printf(__msan_print_func_beg, "__msan_unpoison_stack"));
 
   auto shadow_address = MemToShadow((uptr)ptr, ADDRESS_SPACE_PRIVATE);
-  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow_private,
-                                (void *)shadow_address,
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow, (void *)shadow_address,
                                 (void *)(shadow_address + size), 0x0));
 
   if (shadow_address != GetMsanLaunchInfo->CleanShadow) {
@@ -716,6 +714,26 @@ DEVICE_EXTERN_C_NOINLINE void __msan_unpoison_stack(__SYCL_PRIVATE__ void *ptr,
       __spirv_ocl_printf(__msan_print_func_end, "__msan_unpoison_stack"));
 }
 
+DEVICE_EXTERN_C_NOINLINE void __msan_unpoison_shadow(uptr ptr, uint32_t as,
+                                                     uptr size) {
+  if (!GetMsanLaunchInfo)
+    return;
+
+  MSAN_DEBUG(
+      __spirv_ocl_printf(__msan_print_func_beg, "__msan_unpoison_shadow"));
+
+  auto shadow_address = MemToShadow(ptr, as);
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow, (void *)shadow_address,
+                                (void *)(shadow_address + size), 0x0));
+
+  if (shadow_address != GetMsanLaunchInfo->CleanShadow) {
+    Memset((__SYCL_GLOBAL__ char *)shadow_address, 0, size);
+  }
+
+  MSAN_DEBUG(
+      __spirv_ocl_printf(__msan_print_func_end, "__msan_unpoison_shadow"));
+}
+
 static __SYCL_CONSTANT__ const char __msan_print_private_base[] =
     "[kernel] __msan_set_private_base(sid=%llu): %p\n";
 

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -813,6 +813,8 @@ class MemorySanitizerOnSpirv {
   Constant *getOrCreateGlobalString(StringRef Name, StringRef Value,
                                     unsigned AddressSpace);
 
+  static bool isSupportedBuiltIn(StringRef Name);
+
   operator bool() const { return IsSPIRV; }
 
 private:
@@ -823,7 +825,6 @@ class MemorySanitizerOnSpirv {
   void instrumentKernelsMetadata(int TrackOrigins);
   void instrumentPrivateArguments(Function &F, Instruction *FnPrologueEnd);
   void instrumentPrivateBase(Function &F);
-
   void initializeRetVecMap(Function *F);
   void initializeKernelCallerMap(Function *F);
 
@@ -856,6 +857,7 @@ class MemorySanitizerOnSpirv {
   FunctionCallee MsanUnpoisonShadowDynamicLocalFunc;
   FunctionCallee MsanBarrierFunc;
   FunctionCallee MsanUnpoisonStackFunc;
+  FunctionCallee MsanUnpoisonShadowFunc;
   FunctionCallee MsanSetPrivateBaseFunc;
   FunctionCallee MsanUnpoisonStridedCopyFunc;
 };
@@ -949,6 +951,14 @@ void MemorySanitizerOnSpirv::initializeCallbacks() {
   MsanUnpoisonStackFunc = M.getOrInsertFunction(
       "__msan_unpoison_stack", IRB.getVoidTy(), PtrTy, IntptrTy);
 
+  // __msan_unpoison_(
+  //   uptr ptr,
+  //   uint32_t as,
+  //   size_t size
+  // )
+  MsanUnpoisonShadowFunc = M.getOrInsertFunction(
+      "__msan_unpoison_shadow", IRB.getVoidTy(), IntptrTy, Int32Ty, IntptrTy);
+
   // __msan_set_private_base(
   //   as(0) void *  ptr
   // )
@@ -987,9 +997,16 @@ void MemorySanitizerOnSpirv::instrumentGlobalVariables() {
       G.setName("nameless_global");
 
     if (isUnsupportedDeviceGlobal(&G)) {
-      for (auto *User : G.users())
-        if (auto *Inst = dyn_cast<Instruction>(User))
-          Inst->setNoSanitizeMetadata();
+      for (auto *User : G.users()) {
+        if (!isa<Instruction>(User))
+          continue;
+        if (auto *CI = dyn_cast<CallInst>(User)) {
+          Function *Callee = CI->getCalledFunction();
+          if (Callee && isSupportedBuiltIn(Callee->getName()))
+            continue;
+        }
+        cast<Instruction>(User)->setNoSanitizeMetadata();
+      }
       continue;
     }
 
@@ -1150,6 +1167,10 @@ void MemorySanitizerOnSpirv::instrumentPrivateBase(Function &F) {
   IRB.CreateCall(MsanSetPrivateBaseFunc, {PrivateBase});
 }
 
+bool MemorySanitizerOnSpirv::isSupportedBuiltIn(StringRef Name) {
+  return Name.contains("__sycl_getComposite2020SpecConstantValue");
+}
+
 void MemorySanitizerOnSpirv::instrumentPrivateArguments(
     Function &F, Instruction *FnPrologueEnd) {
   if (!ClSpirOffloadPrivates)
@@ -6994,6 +7015,25 @@ struct MemorySanitizerVisitor : public InstVisitor<MemorySanitizerVisitor> {
                IRB.CreatePointerCast(Src, MS.Spirv.IntptrTy),
                IRB.getInt32(Src->getType()->getPointerAddressSpace()),
                IRB.getInt32(ElementSize), NumElements, Stride});
+        } else if (FuncName.contains(
+                       "__sycl_getComposite2020SpecConstantValue")) {
+          // clang-format off
+          // Handle builtin functions like "_Z40__sycl_getComposite2020SpecConstantValue"
+          // Structs which are larger than 64b will be returned via sret arguments
+          // and will be initialized inside the function. So we need to unpoison
+          // the sret arguments.
+          // clang-format on
+          if (Func->hasStructRetAttr()) {
+            Type *SCTy = Func->getParamStructRetType(0);
+            unsigned Size = Func->getDataLayout().getTypeStoreSize(SCTy);
+            auto *Addr = CB.getArgOperand(0);
+            IRB.CreateCall(
+                MS.Spirv.MsanUnpoisonShadowFunc,
+                {IRB.CreatePointerCast(Addr, MS.Spirv.IntptrTy),
+                 ConstantInt::get(MS.Spirv.Int32Ty,
+                                  Addr->getType()->getPointerAddressSpace()),
+                 ConstantInt::get(MS.Spirv.IntptrTy, Size)});
+          }
         }
       }
     }

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

@@ -0,0 +1,22 @@
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-poison-stack-with-call=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"
+
+%"class.sycl::_V1::specialization_id" = type { %"struct.user_def_types::no_cnstr" }
+%"struct.user_def_types::no_cnstr" = type { float, i32, i8 }
+
+@__usid_str = external addrspace(4) constant [57 x i8]
+@_Z19spec_const_externalIN14user_def_types8no_cnstrELi1EE = external addrspace(1) constant %"class.sycl::_V1::specialization_id"
+
+define spir_func i1 @_Z50check_kernel_handler_by_reference_external_handlerRN4sycl3_V114kernel_handlerEN14user_def_types8no_cnstrE() {
+entry:
+  %ref.tmp.i = alloca %"struct.user_def_types::no_cnstr", align 4
+  %ref.tmp.ascast.i = addrspacecast ptr %ref.tmp.i to ptr addrspace(4)
+; CHECK: [[REG1:%[0-9]+]] = ptrtoint ptr addrspace(4) %ref.tmp.ascast.i to i64
+; CHECK: call void @__msan_unpoison_shadow(i64 [[REG1]], i32 4, i64 12)
+  call spir_func void @_Z40__sycl_getComposite2020SpecConstantValueIN14user_def_types8no_cnstrEET_PKcPKvS6_(ptr addrspace(4) dead_on_unwind writable sret(%"struct.user_def_types::no_cnstr") align 4 %ref.tmp.ascast.i, ptr addrspace(4) noundef @__usid_str, ptr addrspace(4) noundef addrspacecast (ptr addrspace(1) @_Z19spec_const_externalIN14user_def_types8no_cnstrELi1EE to ptr addrspace(4)), ptr addrspace(4) noundef null)
+  ret i1 false
+}
+
+declare spir_func void @_Z40__sycl_getComposite2020SpecConstantValueIN14user_def_types8no_cnstrEET_PKcPKvS6_(ptr addrspace(4) sret(%"struct.user_def_types::no_cnstr"), ptr addrspace(4), ptr addrspace(4), ptr addrspace(4))

sycl/test-e2e/AddressSanitizer/lit.local.cfg

@@ -28,3 +28,5 @@ unsupported_san_flags = [
 ]
 if any(flag in config.cxx_flags for flag in unsupported_san_flags):
     config.unsupported=True
+
+config.environment["ZE_AFFINITY_MASK"] = "0"

sycl/test-e2e/MemorySanitizer/lit.local.cfg

@@ -35,3 +35,5 @@ unsupported_san_flags = [
 ]
 if any(flag in config.cxx_flags for flag in unsupported_san_flags):
     config.unsupported=True
+
+config.environment["ZE_AFFINITY_MASK"] = "0"

sycl/test-e2e/ThreadSanitizer/lit.local.cfg

@@ -33,3 +33,5 @@ unsupported_san_flags = [
 ]
 if any(flag in config.cxx_flags for flag in unsupported_san_flags):
     config.unsupported=True
+
+config.environment["ZE_AFFINITY_MASK"] = "0"