diff --git a/libdevice/sanitizer/msan_rtl.cpp b/libdevice/sanitizer/msan_rtl.cpp index 9be95fafc27a3..9adc84a64539d 100644 --- a/libdevice/sanitizer/msan_rtl.cpp +++ b/libdevice/sanitizer/msan_rtl.cpp @@ -13,6 +13,8 @@ DeviceGlobal __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, diff --git a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp index f14c538ec8637..6cecdca8b3ad6 100644 --- a/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp +++ b/llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp @@ -1216,6 +1216,14 @@ void MemorySanitizerOnSpirv::instrumentKernelsMetadata(int TrackOrigins) { SmallVector SpirKernelsMetadata; SmallVector 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 diff --git a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/track_origins.ll b/llvm/test/Instrumentation/MemorySanitizer/SPIRV/track_origins.ll index a5b383ba36932..d20b133efff2a 100644 --- a/llvm/test/Instrumentation/MemorySanitizer/SPIRV/track_origins.ll +++ b/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 { diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp index e2bbb166a5feb..810fd76120b2d 100644 --- a/unified-runtime/source/loader/layers/sanitizer/msan/msan_ddi.cpp +++ b/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++) { diff --git a/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp b/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp index 12ee3154cc7fa..826960d6e29ff 100644 --- a/unified-runtime/source/loader/layers/sanitizer/msan/msan_shadow.cpp +++ b/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)); diff --git a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.cpp b/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.cpp index b81cf55b4d584..edfc4f8a56f3d 100644 --- a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.cpp +++ b/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)); diff --git a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.hpp b/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.hpp index 0199cc6f24b9a..1f0f250a9453d 100644 --- a/unified-runtime/source/loader/layers/sanitizer/sanitizer_common/sanitizer_utils.hpp +++ b/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