[DeviceMSAN] Check use-of-uninitialized value on dynamic local memory

libdevice/sanitizer/msan_rtl.cpp

@@ -160,7 +160,7 @@ inline uptr __msan_get_shadow_dg2(uptr addr, uint32_t as) {
   if (addr < shadow_begin) {
     return addr + (shadow_begin - DG2_DEVICE_USM_BEGIN);
   } else {
-    return addr - (DG2_DEVICE_USM_END - shadow_end);
+    return addr - (DG2_DEVICE_USM_END - shadow_end + 1);
   }
 }
 
@@ -176,7 +176,7 @@ inline uptr __msan_get_shadow_pvc(uptr addr, uint32_t as) {
     if (addr < shadow_begin) {
       return addr + (shadow_begin - PVC_DEVICE_USM_BEGIN);
     } else {
-      return addr - (PVC_DEVICE_USM_END - shadow_end);
+      return addr - (PVC_DEVICE_USM_END - shadow_end + 1);
     }
   } else if (as == ADDRESS_SPACE_LOCAL) {
     // The size of SLM is 128KB on PVC
@@ -417,4 +417,73 @@ DEVICE_EXTERN_C_INLINE void __msan_barrier() {
                              __spv::MemorySemanticsMask::WorkgroupMemory);
 }
 
+static __SYCL_CONSTANT__ const char __msan_print_local_arg[] =
+    "[kernel] local_arg(index=%d, size=%d)\n";
+
+static __SYCL_CONSTANT__ const char
+    __msan_print_set_shadow_dynamic_local_begin[] =
+        "[kernel] BEGIN __msan_poison_shadow_dynamic_local\n";
+static __SYCL_CONSTANT__ const char
+    __msan_print_set_shadow_dynamic_local_end[] =
+        "[kernel] END   __msan_poison_shadow_dynamic_local\n";
+static __SYCL_CONSTANT__ const char __msan_print_report_arg_count_incorrect[] =
+    "[kernel] ERROR: The number of local args is incorrect, expect %d, actual "
+    "%d\n";
+
+DEVICE_EXTERN_C_NOINLINE void
+__msan_poison_shadow_dynamic_local(uptr ptr, uint32_t num_args) {
+  if (!GetMsanLaunchInfo)
+    return;
+
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow_dynamic_local_begin));
+
+  if (num_args != GetMsanLaunchInfo->NumLocalArgs) {
+    __spirv_ocl_printf(__msan_print_report_arg_count_incorrect, num_args,
+                       GetMsanLaunchInfo->NumLocalArgs);
+    return;
+  }
+
+  uptr *args = (uptr *)ptr;
+
+  for (uint32_t i = 0; i < num_args; ++i) {
+    auto *local_arg = &GetMsanLaunchInfo->LocalArgs[i];
+    MSAN_DEBUG(__spirv_ocl_printf(__msan_print_local_arg, i, local_arg->Size));
+
+    __msan_poison_shadow_static_local(args[i], local_arg->Size);
+  }
+
+  MSAN_DEBUG(__spirv_ocl_printf(__msan_print_set_shadow_dynamic_local_end));
+}
+
+static __SYCL_CONSTANT__ const char
+    __mem_unpoison_shadow_dynamic_local_begin[] =
+        "[kernel] BEGIN __msan_unpoison_shadow_dynamic_local\n";
+static __SYCL_CONSTANT__ const char __mem_unpoison_shadow_dynamic_local_end[] =
+    "[kernel] END   __msan_unpoison_shadow_dynamic_local\n";
+
+DEVICE_EXTERN_C_NOINLINE void
+__msan_unpoison_shadow_dynamic_local(uptr ptr, uint32_t num_args) {
+  if (!GetMsanLaunchInfo)
+    return;
+
+  MSAN_DEBUG(__spirv_ocl_printf(__mem_unpoison_shadow_dynamic_local_begin));
+
+  if (num_args != GetMsanLaunchInfo->NumLocalArgs) {
+    __spirv_ocl_printf(__msan_print_report_arg_count_incorrect, num_args,
+                       GetMsanLaunchInfo->NumLocalArgs);
+    return;
+  }
+
+  uptr *args = (uptr *)ptr;
+
+  for (uint32_t i = 0; i < num_args; ++i) {
+    auto *local_arg = &GetMsanLaunchInfo->LocalArgs[i];
+    MSAN_DEBUG(__spirv_ocl_printf(__msan_print_local_arg, i, local_arg->Size));
+
+    __msan_unpoison_shadow_static_local(args[i], local_arg->Size);
+  }
+
+  MSAN_DEBUG(__spirv_ocl_printf(__mem_unpoison_shadow_dynamic_local_end));
+}
+
 #endif // __SPIR__ || __SPIRV__

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -777,9 +777,11 @@ class MemorySanitizerOnSpirv {
     IsSPIRV = TargetTriple.isSPIROrSPIRV();
 
     IntptrTy = DL.getIntPtrType(C);
+    Int32Ty = Type::getInt32Ty(C);
   }
 
   bool instrumentModule();
+  void instrumentFunction(Function &F);
 
   Constant *getOrCreateGlobalString(StringRef Name, StringRef Value,
                                     unsigned AddressSpace);
@@ -788,6 +790,7 @@ class MemorySanitizerOnSpirv {
   void initializeCallbacks();
   void instrumentGlobalVariables();
   void instrumentStaticLocalMemory();
+  void instrumentDynamicLocalMemory(Function &F);
   void instrumentKernelsMetadata();
 
   void initializeRetVecMap(Function *F);
@@ -799,15 +802,25 @@ class MemorySanitizerOnSpirv {
   const DataLayout &DL;
   bool IsSPIRV;
   Type *IntptrTy;
+  Type *Int32Ty;
 
   StringMap<GlobalVariable *> GlobalStringMap;
 
   DenseMap<Function *, SmallVector<Instruction *, 8>> KernelToRetVecMap;
   DenseMap<Function *, SmallVector<Constant *, 8>> KernelToLocalMemMap;
   DenseMap<Function *, DenseSet<Function *>> FuncToKernelCallerMap;
 
+  // Make sure that we insert barriers only once per function, and the barrier
+  // needs to be inserted after all "MsanPoisonShadowStaticLocalFunc" and
+  // "MsanPoisonShadowDynamicLocalFunc", and before
+  // "MsanUnpoisonShadowStaticLocalFunc" and
+  // "MsanUnpoisonShadowDynamicLocalFunc".
+  DenseMap<Function *, bool> InsertBarrier;
+
   FunctionCallee MsanPoisonShadowStaticLocalFunc;
   FunctionCallee MsanUnpoisonShadowStaticLocalFunc;
+  FunctionCallee MsanPoisonShadowDynamicLocalFunc;
+  FunctionCallee MsanUnpoisonShadowDynamicLocalFunc;
   FunctionCallee MsanBarrierFunc;
 };
 
@@ -874,6 +887,21 @@ void MemorySanitizerOnSpirv::initializeCallbacks() {
       M.getOrInsertFunction("__msan_unpoison_shadow_static_local",
                             IRB.getVoidTy(), IntptrTy, IntptrTy);
 
+  // __asan_poison_shadow_dynamic_local(
+  //   uptr ptr,
+  //   uint32_t num_args
+  // )
+  MsanPoisonShadowDynamicLocalFunc = M.getOrInsertFunction(
+      "__msan_poison_shadow_dynamic_local", IRB.getVoidTy(), IntptrTy, Int32Ty);
+
+  // __asan_unpoison_shadow_dynamic_local(
+  //   uptr ptr,
+  //   uint32_t num_args
+  // )
+  MsanUnpoisonShadowDynamicLocalFunc =
+      M.getOrInsertFunction("__msan_unpoison_shadow_dynamic_local",
+                            IRB.getVoidTy(), IntptrTy, Int32Ty);
+
   // __msan_barrier()
   MsanBarrierFunc = M.getOrInsertFunction("__msan_barrier", IRB.getVoidTy());
 }
@@ -951,16 +979,15 @@ void MemorySanitizerOnSpirv::instrumentStaticLocalMemory() {
   if (!ClSpirOffloadLocals)
     return;
 
-  DenseMap<Function *, bool> InsertBarrier;
-
-  auto Instrument = [this, &InsertBarrier](GlobalVariable *G, Function *F) {
+  auto Instrument = [this](GlobalVariable *G, Function *F) {
     const uint64_t SizeInBytes = DL.getTypeAllocSize(G->getValueType());
 
-    // Poison shadow of static local memory
     if (!InsertBarrier[F]) {
       IRBuilder<> Builder(&F->getEntryBlock().front());
       Builder.CreateCall(MsanBarrierFunc);
     }
+
+    // Poison shadow of static local memory
     IRBuilder<> Builder(&F->getEntryBlock().front());
     Builder.CreateCall(MsanPoisonShadowStaticLocalFunc,
                        {Builder.CreatePointerCast(G, IntptrTy),
@@ -1001,6 +1028,54 @@ void MemorySanitizerOnSpirv::instrumentStaticLocalMemory() {
   }
 }
 
+void MemorySanitizerOnSpirv::instrumentDynamicLocalMemory(Function &F) {
+  if (!ClSpirOffloadLocals)
+    return;
+
+  // Poison shadow of local memory in kernel argument, required by CPU device
+  SmallVector<Argument *> LocalArgs;
+  for (auto &Arg : F.args()) {
+    Type *PtrTy = dyn_cast<PointerType>(Arg.getType()->getScalarType());
+    if (PtrTy && PtrTy->getPointerAddressSpace() == kSpirOffloadLocalAS)
+      LocalArgs.push_back(&Arg);
+  }
+
+  if (LocalArgs.empty())
+    return;
+
+  if (!InsertBarrier[&F]) {
+    IRBuilder<> Builder(&F.getEntryBlock().front());
+    Builder.CreateCall(MsanBarrierFunc);
+  }
+
+  IRBuilder<> IRB(&F.getEntryBlock().front());
+
+  AllocaInst *ArgsArray = IRB.CreateAlloca(
+      IntptrTy, ConstantInt::get(Int32Ty, LocalArgs.size()), "local_args");
+  for (size_t i = 0; i < LocalArgs.size(); i++) {
+    auto *StoreDest =
+        IRB.CreateGEP(IntptrTy, ArgsArray, ConstantInt::get(Int32Ty, i));
+    IRB.CreateStore(IRB.CreatePointerCast(LocalArgs[i], IntptrTy), StoreDest);
+  }
+
+  auto *ArgsArrayAddr = IRB.CreatePointerCast(ArgsArray, IntptrTy);
+  IRB.CreateCall(MsanPoisonShadowDynamicLocalFunc,
+                 {ArgsArrayAddr, ConstantInt::get(Int32Ty, LocalArgs.size())});
+
+  // Unpoison shadow of dynamic local memory, required by CPU device
+  initializeRetVecMap(&F);
+  for (Instruction *Ret : KernelToRetVecMap[&F]) {
+    IRBuilder<> IRBRet(Ret);
+    if (!InsertBarrier[&F])
+      IRBRet.CreateCall(MsanBarrierFunc);
+    IRBRet.CreateCall(
+        MsanUnpoisonShadowDynamicLocalFunc,
+        {ArgsArrayAddr, ConstantInt::get(Int32Ty, LocalArgs.size())});
+  }
+
+  InsertBarrier[&F] = true;
+}
+
 // Instrument __MsanKernelMetadata, which records information of sanitized
 // kernel
 void MemorySanitizerOnSpirv::instrumentKernelsMetadata() {
@@ -1087,6 +1162,14 @@ bool MemorySanitizerOnSpirv::instrumentModule() {
   return true;
 }
 
+void MemorySanitizerOnSpirv::instrumentFunction(Function &F) {
+  if (!IsSPIRV)
+    return;
+
+  if (F.getCallingConv() == CallingConv::SPIR_KERNEL)
+    instrumentDynamicLocalMemory(F);
+}
+
 PreservedAnalyses MemorySanitizerPass::run(Module &M,
                                            ModuleAnalysisManager &AM) {
   // Return early if nosanitize_memory module flag is present for the module.
@@ -1110,6 +1193,7 @@ PreservedAnalyses MemorySanitizerPass::run(Module &M,
     MemorySanitizer Msan(*F.getParent(), MsanSpirv, Options);
     Modified |=
         Msan.sanitizeFunction(F, FAM.getResult<TargetLibraryAnalysis>(F));
+    MsanSpirv.instrumentFunction(F);
   }
 
   if (!Modified)

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

@@ -0,0 +1,25 @@
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-locals=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"
+
+@WGLocal = internal addrspace(3) global i64 zeroinitializer, align 8
+
+define spir_kernel void @MyKernel(ptr addrspace(3) noundef align 4 %_arg_acc) sanitize_memory {
+; CHECK-LABEL: @MyKernel
+entry:
+  ; CHECK: %local_args = alloca i64, align 8
+  ; CHECK-NEXT: %0 = getelementptr i64, ptr %local_args, i32 0
+  ; CHECK-NEXT: %1 = ptrtoint ptr addrspace(3) %_arg_acc to i64
+  ; CHECK-NEXT: store i64 %1, ptr %0, align 8
+  ; CHECK-NEXT: %2 = ptrtoint ptr %local_args to i64
+  ; CHECK-NEXT: call void @__msan_poison_shadow_dynamic_local(i64 %2, i32 1)
+
+  ; CHECK: @__msan_poison_shadow_static_local{{.*}}@WGLocal
+  ; CHECK: @__msan_barrier
+  store i32 0, ptr addrspace(3) @WGLocal, align 8
+  ; CHECK: @__msan_barrier
+  ; CHECK: @__msan_unpoison_shadow_static_local{{.*}}@WGLocal
+
+  ; CHECK: call void @__msan_unpoison_shadow_dynamic_local(i64 %2, i32 1)
+  ret void
+}

sycl/test-e2e/MemorySanitizer/local/local_accessor.cpp

@@ -0,0 +1,40 @@
+// REQUIRES: linux, cpu || (gpu && level_zero)
+// RUN: %{build} %device_msan_flags -g -O0 -o %t1.out
+// RUN: %{run} not %t1.out 2>&1 | FileCheck %s
+// RUN: %{build} %device_msan_flags -g -O1 -o %t2.out
+// RUN: %{run} not %t2.out 2>&1 | FileCheck %s
+// RUN: %{build} %device_msan_flags -g -O2 -o %t3.out
+// RUN: %{run} not %t3.out 2>&1 | FileCheck %s
+
+// XFAIL: spirv-backend && gpu && run-mode
+// XFAIL-TRACKER: https://github.com/llvm/llvm-project/issues/122075
+
+#include <sycl/ext/oneapi/group_local_memory.hpp>
+#include <sycl/usm.hpp>
+
+constexpr std::size_t global_size = 4;
+constexpr std::size_t local_size = 1;
+
+__attribute__((noinline)) int check(int data, sycl::nd_item<1> &item) {
+  auto ptr =
+      sycl::ext::oneapi::group_local_memory<int[global_size]>(item.get_group());
+  auto &ref = *ptr;
+  return data + ref[0];
+}
+
+int main() {
+  sycl::queue Q;
+
+  Q.submit([&](sycl::handler &cgh) {
+    auto acc = sycl::local_accessor<int>(local_size, cgh);
+    cgh.parallel_for<class MyKernel>(
+        sycl::nd_range<1>(global_size, local_size),
+        [=](sycl::nd_item<1> item) { check(acc[item.get_local_id()], item); });
+  });
+  Q.wait();
+  // CHECK-NOT: [kernel]
+  // CHECK: DeviceSanitizer: use-of-uninitialized-value
+  // CHECK: #0 {{.*}} {{.*local_accessor.cpp}}:[[@LINE-5]]
+
+  return 0;
+}

unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp

@@ -1389,6 +1389,31 @@ ur_result_t urKernelSetArgMemObj(
   return UR_RESULT_SUCCESS;
 }
 
+///////////////////////////////////////////////////////////////////////////////
+/// @brief Intercept function for urKernelSetArgLocal
+__urdlllocal ur_result_t UR_APICALL urKernelSetArgLocal(
+    /// [in] handle of the kernel object
+    ur_kernel_handle_t hKernel,
+    /// [in] argument index in range [0, num args - 1]
+    uint32_t argIndex,
+    /// [in] size of the local buffer to be allocated by the runtime
+    size_t argSize,
+    /// [in][optional] pointer to local buffer properties.
+    const ur_kernel_arg_local_properties_t *pProperties) {
+  auto pfnSetArgLocal = getContext()->urDdiTable.Kernel.pfnSetArgLocal;
+
+  getContext()->logger.debug(
+      "==== urKernelSetArgLocal (argIndex={}, argSize={})", argIndex, argSize);
+
+  {
+    auto &KI = getMsanInterceptor()->getOrCreateKernelInfo(hKernel);
+    std::scoped_lock<ur_shared_mutex> Guard(KI.Mutex);
+    KI.LocalArgs[argIndex] = MsanLocalArgsInfo{argSize};
+  }
+
+  return pfnSetArgLocal(hKernel, argIndex, argSize, pProperties);
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 /// @brief Intercept function for urEnqueueUSMFill
 ur_result_t UR_APICALL urEnqueueUSMFill(
@@ -1738,6 +1763,7 @@ ur_result_t urGetKernelProcAddrTable(
   pDdiTable->pfnRelease = ur_sanitizer_layer::msan::urKernelRelease;
   pDdiTable->pfnSetArgValue = ur_sanitizer_layer::msan::urKernelSetArgValue;
   pDdiTable->pfnSetArgMemObj = ur_sanitizer_layer::msan::urKernelSetArgMemObj;
+  pDdiTable->pfnSetArgLocal = ur_sanitizer_layer::msan::urKernelSetArgLocal;
 
   return result;
 }