[DeviceMSAN] Check use-of-uninitialized value on private memory (#17309)

AllanZyne · web-flow · commit 9ad2e8777423 · 2025-03-19T10:15:13.000Z
Support check use-of-uninitialized value on private memory
diff --git a/clang/lib/Driver/SanitizerArgs.cpp b/clang/lib/Driver/SanitizerArgs.cpp
@@ -1278,6 +1278,9 @@ void SanitizerArgs::addArgs(const ToolChain &TC, const llvm::opt::ArgList &Args,
 
       CmdArgs.push_back("-mllvm");
       CmdArgs.push_back("-msan-eager-checks=1");
+
+      CmdArgs.push_back("-mllvm");
+      CmdArgs.push_back("-msan-poison-stack-with-call=1");
     } else if (Sanitizers.has(SanitizerKind::Thread)) {
       CmdArgs.push_back("-fsanitize=thread");
       // The tsan function entry/exit builtins are used to record stack
diff --git a/libdevice/sanitizer/asan_rtl.cpp b/libdevice/sanitizer/asan_rtl.cpp
@@ -881,6 +881,8 @@ static __SYCL_CONSTANT__ const char __mem_set_shadow_private_end[] =
 static __SYCL_CONSTANT__ const char __mem_set_shadow_private[] =
     "[kernel] set_shadow_private(beg=%p, end=%p, val:%02X)\n";
 
+// We outline the function of setting shadow memory of private memory, because
+// it may allocate failed on UR
 DEVICE_EXTERN_C_NOINLINE void __asan_set_shadow_private(uptr begin, uptr size,
                                                         char val) {
   if (!__AsanLaunchInfo)
diff --git a/libdevice/sanitizer/msan_rtl.cpp b/libdevice/sanitizer/msan_rtl.cpp
diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp
diff --git a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_large_access_size.ll b/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_large_access_size.ll
diff --git a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_unsupported_access.ll b/llvm/test/Instrumentation/MemorySanitizer/SPIRV/check_unsupported_access.ll
@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -S | FileCheck %s
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-privates=0 -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"
 
diff --git a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/instrument_private_mem.ll b/llvm/test/Instrumentation/MemorySanitizer/SPIRV/instrument_private_mem.ll
@@ -0,0 +1,35 @@
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-privates=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"
+
+define spir_kernel void @MyKernel() sanitize_memory {
+; CHECK-LABEL: @MyKernel
+entry:
+  %array = alloca [4 x i32], align 4
+  ; CHECK: call void @__msan_poison_stack(ptr %array, i64 16)
+  ret void
+}
+
+%"class.sycl::_V1::range" = type { %"class.sycl::_V1::detail::array" }
+%"class.sycl::_V1::detail::array" = type { [1 x i64] }
+
+define spir_func void @ByValFunc(ptr noundef byval(%"class.sycl::_V1::range") align 8 %_arg_array12) sanitize_memory {
+; CHECK-LABEL: @ByValFunc
+entry:
+  ; CHECK: %0 = ptrtoint ptr %_arg_array12 to i64
+  ; CHECK: %1 = call i64 @__msan_get_shadow(i64 %0, i32 0, ptr addrspace(2) null)
+  ; CHECK: %2 = inttoptr i64 %1 to ptr addrspace(1)
+  ; CHECK: call void @llvm.memset.p1.i64(ptr addrspace(1) align 8 %2, i8 0, i64 8, i1 false)
+  %_arg_array12.ascast = addrspacecast ptr %_arg_array12 to ptr addrspace(4)
+  ret void
+}
+
+define spir_kernel void @ByValKernel(ptr noundef byval(%"class.sycl::_V1::range") align 8 %_arg_array12) sanitize_memory {
+; CHECK-LABEL: @ByValKernel
+entry:
+  ; CHECK: %_arg_array12.byval = alloca %"class.sycl::_V1::range", align 8
+  ; CHECK: call void @__msan_unpoison_stack(ptr %_arg_array12.byval, i64 8), !nosanitize
+  ; CHECK: call void @llvm.memcpy.p0.p0.i64(ptr align 8 %_arg_array12.byval, ptr align 8 %_arg_array12, i64 8, i1 false), !nosanitize
+  call void @ByValFunc(ptr %_arg_array12)
+  ret void
+}
diff --git a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/instrument_static_local_mem.ll b/llvm/test/Instrumentation/MemorySanitizer/SPIRV/instrument_static_local_mem.ll
@@ -1,4 +1,4 @@
-; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-locals=1 -S | FileCheck %s
+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-locals=1 -msan-spir-privates=0 -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"
 
diff --git a/sycl/test-e2e/MemorySanitizer/check_buffer.cpp b/sycl/test-e2e/MemorySanitizer/check_buffer.cpp
@@ -8,20 +8,17 @@
 
 #include <sycl/detail/core.hpp>
 
-__attribute__((noinline)) long long foo(int data1, long long data2) {
+__attribute__((noinline)) int foo(int data1, int data2) {
   return data1 + data2;
 }
 
 int main() {
   sycl::queue q;
 
   sycl::buffer<int, 1> buf1(sycl::range<1>(1));
-  sycl::buffer<long long, 1> buf2(sycl::range<1>(1));
   q.submit([&](sycl::handler &h) {
      auto array1 = buf1.get_access<sycl::access::mode::read_write>(h);
-     auto array2 = buf2.get_access<sycl::access::mode::read_write>(h);
-     h.single_task<class MyKernel>(
-         [=]() { array1[0] = foo(array1[0], array2[0]); });
+     h.single_task<class MyKernel>([=]() { foo(array1[0], array1[0]); });
    }).wait();
   // CHECK: use-of-uninitialized-value
   // CHECK: kernel <{{.*MyKernel}}>
diff --git a/sycl/test-e2e/MemorySanitizer/private/single_private.cpp b/sycl/test-e2e/MemorySanitizer/private/single_private.cpp
@@ -0,0 +1,32 @@
+// REQUIRES: linux, cpu || (gpu && level_zero)
+// RUN: %{build} %device_msan_flags -O0 -g -o %t1.out
+// RUN: %{run} not %t1.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
+// RUN: %{run} not %t3.out 2>&1 | FileCheck %s
+
+#include <sycl/detail/core.hpp>
+#include <sycl/usm.hpp>
+
+__attribute__((noinline)) int check(int p) { return p; }
+__attribute__((noinline)) int foo(int *p) { return check(*p); }
+// CHECK-NOT: [kernel]
+// CHECK: DeviceSanitizer: use-of-uninitialized-value
+// CHECK: #0 {{foo.*}} {{.*single_private.cpp}}:[[@LINE-3]]
+
+int main() {
+  sycl::queue Q;
+  auto *array = sycl::malloc_device<int>(1, Q);
+
+  Q.submit([&](sycl::handler &h) {
+    h.single_task<class MyKernel>([=]() {
+      int p[4];
+      *array += foo(p);
+    });
+  });
+  Q.wait();
+
+  sycl::free(array, Q);
+  return 0;
+}
diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_interceptor.cpp
@@ -517,10 +517,30 @@ ur_result_t MsanInterceptor::prepareLaunch(
       getContext()->logger.warning("Skip checking local memory of kernel <{}> ",
                                    GetKernelName(Kernel));
     } else {
-      getContext()->logger.info("ShadowMemory(Local, WorkGroup={}, {} - {})",
-                                NumWG,
-                                (void *)LaunchInfo.Data->LocalShadowOffset,
-                                (void *)LaunchInfo.Data->LocalShadowOffsetEnd);
+      getContext()->logger.debug("ShadowMemory(Local, WorkGroup={}, {} - {})",
+                                 NumWG,
+                                 (void *)LaunchInfo.Data->LocalShadowOffset,
+                                 (void *)LaunchInfo.Data->LocalShadowOffsetEnd);
+    }
+  }
+
+  // Write shadow memory offset for private memory
+  if (KernelInfo.IsCheckPrivates) {
+    if (DeviceInfo->Shadow->AllocPrivateShadow(
+            Queue, NumWG, LaunchInfo.Data->PrivateShadowOffset,
+            LaunchInfo.Data->PrivateShadowOffsetEnd) != UR_RESULT_SUCCESS) {
+      getContext()->logger.warning(
+          "Failed to allocate shadow memory for private "
+          "memory, maybe the number of workgroup ({}) is too "
+          "large",
+          NumWG);
+      getContext()->logger.warning(
+          "Skip checking private memory of kernel <{}>", GetKernelName(Kernel));
+    } else {
+      getContext()->logger.debug(
+          "ShadowMemory(Private, WorkGroup={}, {} - {})", NumWG,
+          (void *)LaunchInfo.Data->PrivateShadowOffset,
+          (void *)LaunchInfo.Data->PrivateShadowOffsetEnd);
     }
     // Write local arguments info
     if (!KernelInfo.LocalArgs.empty()) {
@@ -535,11 +555,11 @@ ur_result_t MsanInterceptor::prepareLaunch(
   }
 
   getContext()->logger.info(
-      "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, CleanShadow={}, "
-      "LocalArgs={}, NumLocalArgs={}, "
-      "Device={}, Debug={})",
+      "LaunchInfo {} (GlobalShadow={}, LocalShadow={}, PrivateShadow={}, "
+      "CleanShadow={}, LocalArgs={}, NumLocalArgs={}, Device={}, Debug={})",
       (void *)LaunchInfo.Data, (void *)LaunchInfo.Data->GlobalShadowOffset,
       (void *)LaunchInfo.Data->LocalShadowOffset,
+      (void *)LaunchInfo.Data->PrivateShadowOffset,
       (void *)LaunchInfo.Data->CleanShadow, (void *)LaunchInfo.Data->LocalArgs,
       LaunchInfo.Data->NumLocalArgs, ToString(LaunchInfo.Data->DeviceTy),
       LaunchInfo.Data->Debug);
diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_libdevice.hpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_libdevice.hpp
@@ -37,6 +37,7 @@ struct MsanErrorReport {
 
   uint32_t AccessSize = 0;
   ErrorType ErrorTy = ErrorType::UNKNOWN;
+  uintptr_t Origin;
 };
 
 struct MsanLocalArgsInfo {
@@ -50,6 +51,9 @@ struct MsanLaunchInfo {
   uintptr_t LocalShadowOffset = 0;
   uintptr_t LocalShadowOffsetEnd = 0;
 
+  uintptr_t PrivateShadowOffset = 0;
+  uintptr_t PrivateShadowOffsetEnd = 0;
+
   uintptr_t CleanShadow = 0;
 
   DeviceType DeviceTy = DeviceType::UNKNOWN;
diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_report.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_report.cpp
@@ -30,8 +30,15 @@ void ReportUsesUninitializedValue(const MsanErrorReport &Report,
   // Try to demangle the kernel name
   KernelName = DemangleName(KernelName);
 
-  getContext()->logger.always(
-      "====WARNING: DeviceSanitizer: use-of-uninitialized-value");
+  if (Report.Origin) {
+    getContext()->logger.always(
+        "====WARNING: DeviceSanitizer: use-of-uninitialized-value (shadow: {})",
+        (void *)Report.Origin);
+  } else {
+    getContext()->logger.always(
+        "====WARNING: DeviceSanitizer: use-of-uninitialized-value)");
+  }
+
   getContext()->logger.always(
       "use of size {} at kernel <{}> LID({}, {}, {}) GID({}, "
       "{}, {})",
diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp
@@ -333,6 +333,42 @@ ur_result_t MsanShadowMemoryGPU::AllocLocalShadow(ur_queue_handle_t Queue,
   return UR_RESULT_SUCCESS;
 }
 
+ur_result_t MsanShadowMemoryGPU::AllocPrivateShadow(ur_queue_handle_t Queue,
+                                                    uint32_t NumWG, uptr &Begin,
+                                                    uptr &End) {
+  const size_t RequiredShadowSize = NumWG * MSAN_PRIVATE_SIZE;
+  static size_t LastAllocedSize = 0;
+  if (RequiredShadowSize > LastAllocedSize) {
+    auto ContextInfo = getMsanInterceptor()->getContextInfo(Context);
+    if (PrivateShadowOffset) {
+      UR_CALL(getContext()->urDdiTable.USM.pfnFree(
+          Context, (void *)PrivateShadowOffset));
+      PrivateShadowOffset = 0;
+      LastAllocedSize = 0;
+    }
+
+    UR_CALL(getContext()->urDdiTable.USM.pfnDeviceAlloc(
+        Context, Device, nullptr, nullptr, RequiredShadowSize,
+        (void **)&PrivateShadowOffset));
+
+    // Initialize shadow memory
+    ur_result_t URes = EnqueueUSMBlockingSet(Queue, (void *)PrivateShadowOffset,
+                                             0, RequiredShadowSize);
+    if (URes != UR_RESULT_SUCCESS) {
+      UR_CALL(getContext()->urDdiTable.USM.pfnFree(
+          Context, (void *)PrivateShadowOffset));
+      PrivateShadowOffset = 0;
+      LastAllocedSize = 0;
+    }
+
+    LastAllocedSize = RequiredShadowSize;
+  }
+
+  Begin = PrivateShadowOffset;
+  End = PrivateShadowOffset + RequiredShadowSize - 1;
+  return UR_RESULT_SUCCESS;
+}
+
 uptr MsanShadowMemoryPVC::MemToShadow(uptr Ptr) {
   assert(MsanShadowMemoryPVC::IsDeviceUSM(Ptr) && "Ptr must be device USM");
   if (Ptr < ShadowBegin) {
diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.hpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.hpp
@@ -46,6 +46,10 @@ struct MsanShadowMemory {
   virtual ur_result_t AllocLocalShadow(ur_queue_handle_t Queue, uint32_t NumWG,
                                        uptr &Begin, uptr &End) = 0;
 
+  virtual ur_result_t AllocPrivateShadow(ur_queue_handle_t Queue,
+                                         uint32_t NumWG, uptr &Begin,
+                                         uptr &End) = 0;
+
   ur_context_handle_t Context{};
 
   ur_device_handle_t Device{};
@@ -95,6 +99,13 @@ struct MsanShadowMemoryCPU final : public MsanShadowMemory {
     End = ShadowEnd;
     return UR_RESULT_SUCCESS;
   }
+
+  ur_result_t AllocPrivateShadow(ur_queue_handle_t, uint32_t, uptr &Begin,
+                                 uptr &End) override {
+    Begin = ShadowBegin;
+    End = ShadowEnd;
+    return UR_RESULT_SUCCESS;
+  }
 };
 
 struct MsanShadowMemoryGPU : public MsanShadowMemory {
@@ -116,6 +127,9 @@ struct MsanShadowMemoryGPU : public MsanShadowMemory {
   ur_result_t AllocLocalShadow(ur_queue_handle_t Queue, uint32_t NumWG,
                                uptr &Begin, uptr &End) override final;
 
+  ur_result_t AllocPrivateShadow(ur_queue_handle_t Queue, uint32_t NumWG,
+                                 uptr &Begin, uptr &End) override final;
+
   virtual size_t GetShadowSize() = 0;
 
   virtual uptr GetStartAddress() { return 0; }
@@ -132,6 +146,8 @@ struct MsanShadowMemoryGPU : public MsanShadowMemory {
   ur_mutex VirtualMemMapsMutex;
 
   uptr LocalShadowOffset = 0;
+
+  uptr PrivateShadowOffset = 0;
 };
 
 // clang-format off

Original file line number	Diff line number	Diff line change
`@@ -1,4 +1,4 @@`
`1`		`-; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -S \| FileCheck %s`
	`1`	`+; RUN: opt < %s -passes=msan -msan-instrumentation-with-call-threshold=0 -msan-eager-checks=1 -msan-spir-privates=0 -S \| FileCheck %s`
`2`	`2`	`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"`
`3`	`3`	`target triple = "spir64-unknown-unknown"`
`4`	`4`