[DevMSAN] Use global __msan_track_origins to indicate if track origin is enabled (#19828)

Always propagate origin value will cause bad performance.

Author: zhaomaosu

libdevice/sanitizer/msan_rtl.cpp

@@ -13,6 +13,8 @@ DeviceGlobal<void *> __MsanLaunchInfo;
 #define GetMsanLaunchInfo                                                      \
   ((__SYCL_GLOBAL__ MsanRuntimeData *)__MsanLaunchInfo.get())
 
+extern "C" __attribute__((weak)) const int __msan_track_origins;
+
 namespace {
 
 constexpr int MSAN_REPORT_NONE = 0;
@@ -64,6 +66,8 @@ const __SYCL_CONSTANT__ char __msan_print_unknown[] = "unknown";
 
 namespace {
 
+inline bool IsTrackOriginsEnabled() { return __msan_track_origins; }
+
 inline void ConvertGenericPointer(uptr &addr, uint32_t &as) {
   auto old = addr;
   if ((addr = (uptr)ToPrivate((void *)old))) {
@@ -337,7 +341,8 @@ inline void CopyShadowAndOrigin(uptr dst, uint32_t dst_as, uptr src,
   }
   auto *shadow_src = (__SYCL_GLOBAL__ char *)MemToShadow(src, src_as);
   Memcpy(shadow_dst, shadow_src, size);
-  CopyOrigin(dst, dst_as, src, src_as, size);
+  if (IsTrackOriginsEnabled())
+    CopyOrigin(dst, dst_as, src, src_as, size);
 
   MSAN_DEBUG(__spirv_ocl_printf(__msan_print_copy_shadow, dst, dst_as, src,
                                 src_as, shadow_dst, shadow_src, size));
@@ -365,7 +370,8 @@ inline void MoveShadowAndOrigin(uptr dst, uint32_t dst_as, uptr src,
   auto *shadow_dst = (__SYCL_GLOBAL__ char *)MemToShadow(dst, dst_as);
   auto *shadow_src = (__SYCL_GLOBAL__ char *)MemToShadow(src, src_as);
   // MoveOrigin transfers origins by refering to their shadows
-  MoveOrigin(dst, dst_as, src, src_as, size);
+  if (IsTrackOriginsEnabled())
+    MoveOrigin(dst, dst_as, src, src_as, size);
   Memmove(shadow_dst, shadow_src, size);
 
   MSAN_DEBUG(__spirv_ocl_printf(__msan_print_move_shadow, dst, dst_as, src,

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

@@ -1216,6 +1216,14 @@ void MemorySanitizerOnSpirv::instrumentKernelsMetadata(int TrackOrigins) {
   SmallVector<Constant *, 8> SpirKernelsMetadata;
   SmallVector<uint8_t, 256> KernelNamesBytes;
 
+  // Insert global __msan_track_origins to indicate if origin track is enabled.
+  M.getOrInsertGlobal("__msan_track_origins", Int32Ty, [&] {
+    return new GlobalVariable(
+        M, Int32Ty, true, GlobalValue::WeakODRLinkage,
+        ConstantInt::get(Int32Ty, TrackOrigins), "__msan_track_origins",
+        nullptr, llvm::GlobalValue::NotThreadLocal, kSpirOffloadGlobalAS);
+  });
+
   // SpirKernelsMetadata only saves fixed kernels, and is described by
   // following structure:
   //  uptr unmangled_kernel_name

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

@@ -3,9 +3,9 @@
 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"
 
+; CHECK: __msan_track_origins
 ; CHECK: @__MsanKernelMetadata{{.*}}i64 8, i64 36
 ; CHECK-SAME: [[ATTR0:#[0-9]+]]
-; CHECK-NOT: __msan_track_origins
 ; CHECK-NOT: _tls
 
 define spir_kernel void @MyKernel(ptr addrspace(1) noundef align 4 %_arg_array) sanitize_memory {

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

@@ -1598,7 +1598,8 @@ ur_result_t urEnqueueUSMMemcpy(
                            (void *)SrcShadow, size, 0, nullptr, &Event));
       Events.push_back(Event);
     }
-    {
+    if (getContext()->Options.MsanCheckHostAndSharedUSM ||
+        (IsDeviceUSM(hContext, pSrc) && IsDeviceUSM(hContext, pDst))) {
       const auto SrcOriginBegin = SrcDI->Shadow->MemToOrigin((uptr)pSrc);
       const auto SrcOriginEnd =
           SrcDI->Shadow->MemToOrigin((uptr)pSrc + size - 1) +
@@ -1764,7 +1765,8 @@ ur_result_t urEnqueueUSMMemcpy2D(
       Events.push_back(Event);
     }
 
-    {
+    if (getContext()->Options.MsanCheckHostAndSharedUSM ||
+        (IsDeviceUSM(hContext, pSrc) && IsDeviceUSM(hContext, pDst))) {
       auto pfnUSMMemcpy = getContext()->urDdiTable.Enqueue.pfnUSMMemcpy;
 
       for (size_t HeightIndex = 0; HeightIndex < height; HeightIndex++) {

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

@@ -163,7 +163,7 @@ ur_result_t MsanShadowMemoryCPU::EnqueuePoisonShadowWithOrigin(
                (void *)(uptr)Value);
       memset((void *)ShadowBegin, Value, ShadowEnd - ShadowBegin + 1);
     }
-    {
+    if (Origin) {
       const uptr OriginBegin = MemToOrigin(Ptr);
       const uptr OriginEnd =
           MemToOrigin(Ptr + Size - 1) + MSAN_ORIGIN_GRANULARITY;
@@ -323,7 +323,7 @@ ur_result_t MsanShadowMemoryGPU::EnqueuePoisonShadowWithOrigin(
                           Events.data(), OutEvent));
   }
 
-  {
+  if (Origin) {
     uptr OriginBegin = MemToOrigin(Ptr);
     uptr OriginEnd = MemToOrigin(Ptr + Size - 1) + sizeof(Origin) - 1;
     UR_CALL(EnqueueVirtualMemMap(OriginBegin, OriginEnd, Events, OutEvent));

unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.cpp

@@ -176,6 +176,11 @@ bool IsHostUSM(ur_context_handle_t Context, const void *MemPtr) {
   return USMType == UR_USM_TYPE_HOST;
 }
 
+bool IsDeviceUSM(ur_context_handle_t Context, const void *MemPtr) {
+  ur_usm_type_t USMType = GetUSMType(Context, MemPtr);
+  return USMType == UR_USM_TYPE_DEVICE;
+}
+
 ur_device_handle_t GetUSMAllocDevice(ur_context_handle_t Context,
                                      const void *MemPtr) {
   assert(IsUSM(Context, MemPtr));

unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.hpp

@@ -54,6 +54,7 @@ size_t GetDeviceLocalMemorySize(ur_device_handle_t Device);
 ur_program_handle_t GetProgram(ur_kernel_handle_t Kernel);
 bool IsUSM(ur_context_handle_t Context, const void *MemPtr);
 bool IsHostUSM(ur_context_handle_t Context, const void *MemPtr);
+bool IsDeviceUSM(ur_context_handle_t Context, const void *MemPtr);
 ur_device_handle_t GetUSMAllocDevice(ur_context_handle_t Context,
                                      const void *MemPtr);
 // Get the device of MemPtr. If MemPtr is host USM, then return the device