From 2d99519df574984ab0c2a62583b006e3e9f651bb Mon Sep 17 00:00:00 2001
From: Yang Zhao <yang2.zhao@intel.com>
Date: Sat, 28 Dec 2024 05:51:11 +0800
Subject: [PATCH] [DeviceMSAN] Fix MemToShadow algorithm and VA reservation
 (#16475)

UR: https://github.com/oneapi-src/unified-runtime/pull/2507
---
 clang/lib/Driver/ToolChains/SYCL.cpp        |  4 ----
 libdevice/cmake/modules/SYCLLibdevice.cmake |  2 ++
 libdevice/sanitizer/msan_rtl.cpp            | 13 +++++++++----
 sycl/cmake/modules/UnifiedRuntimeTag.cmake  | 12 ++++++------
 4 files changed, 17 insertions(+), 14 deletions(-)

diff --git a/clang/lib/Driver/ToolChains/SYCL.cpp b/clang/lib/Driver/ToolChains/SYCL.cpp
index ed254e01bd41e..adaf227b1abcb 100644
--- a/clang/lib/Driver/ToolChains/SYCL.cpp
+++ b/clang/lib/Driver/ToolChains/SYCL.cpp
@@ -671,14 +671,12 @@ SYCL::getDeviceLibraries(const Compilation &C, const llvm::Triple &TargetTriple,
   };
 
   std::string SanitizeVal;
-  std::string SanitizeArg;
   size_t sanitizer_lib_idx = getSingleBuildTarget();
   if (Arg *A = Args.getLastArg(options::OPT_fsanitize_EQ,
                                options::OPT_fno_sanitize_EQ)) {
     if (A->getOption().matches(options::OPT_fsanitize_EQ) &&
         A->getValues().size() == 1) {
       SanitizeVal = A->getValue();
-      SanitizeArg = A->getAsString(Args);
     }
   } else {
     // User can pass -fsanitize=address to device compiler via
@@ -703,12 +701,10 @@ SYCL::getDeviceLibraries(const Compilation &C, const llvm::Triple &TargetTriple,
     for (const std::string &Arg : ArgVals) {
       if (Arg.find("-fsanitize=address") != std::string::npos) {
         SanitizeVal = "address";
-        SanitizeArg = Arg;
         break;
       }
       if (Arg.find("-fsanitize=memory") != std::string::npos) {
         SanitizeVal = "memory";
-        SanitizeArg = Arg;
         break;
       }
     }
diff --git a/libdevice/cmake/modules/SYCLLibdevice.cmake b/libdevice/cmake/modules/SYCLLibdevice.cmake
index 63a079fdc0d40..64979f2b65e4a 100644
--- a/libdevice/cmake/modules/SYCLLibdevice.cmake
+++ b/libdevice/cmake/modules/SYCLLibdevice.cmake
@@ -239,6 +239,7 @@ if (NOT MSVC AND UR_SANITIZER_INCLUDE_DIR)
     device.h atomic.hpp spirv_vars.h
     ${UR_SANITIZER_INCLUDE_DIR}/asan/asan_libdevice.hpp
     include/asan_rtl.hpp
+    include/sanitizer_defs.hpp
     include/spir_global_var.hpp
     ${sycl-compiler_deps})
 
@@ -296,6 +297,7 @@ if (NOT MSVC AND UR_SANITIZER_INCLUDE_DIR)
     device.h atomic.hpp spirv_vars.h
     ${UR_SANITIZER_INCLUDE_DIR}/msan/msan_libdevice.hpp
     include/msan_rtl.hpp
+    include/sanitizer_defs.hpp
     include/spir_global_var.hpp
     sycl-compiler)
 endif()
diff --git a/libdevice/sanitizer/msan_rtl.cpp b/libdevice/sanitizer/msan_rtl.cpp
index ab02f4d0662e5..218aef9aafe3c 100644
--- a/libdevice/sanitizer/msan_rtl.cpp
+++ b/libdevice/sanitizer/msan_rtl.cpp
@@ -132,10 +132,15 @@ inline uptr __msan_get_shadow_cpu(uptr addr) {
 
 inline uptr __msan_get_shadow_pvc(uptr addr, uint32_t as) {
   // Device USM only
-  uptr shadow_ptr = ((__SYCL_GLOBAL__ MsanLaunchInfo *)__MsanLaunchInfo.get())
-                        ->GlobalShadowOffset +
-                    (addr & 0x3FFF'FFFF'FFFFULL);
-  return shadow_ptr;
+  auto shadow_begin = ((__SYCL_GLOBAL__ MsanLaunchInfo *)__MsanLaunchInfo.get())
+                          ->GlobalShadowOffset;
+  auto shadow_end = ((__SYCL_GLOBAL__ MsanLaunchInfo *)__MsanLaunchInfo.get())
+                        ->GlobalShadowOffsetEnd;
+  if (addr < shadow_begin) {
+    return addr + (shadow_begin - 0xff00'0000'0000'0000ULL);
+  } else {
+    return addr - (0xff00'ffff'ffff'ffffULL - shadow_end);
+  }
 }
 
 } // namespace
diff --git a/sycl/cmake/modules/UnifiedRuntimeTag.cmake b/sycl/cmake/modules/UnifiedRuntimeTag.cmake
index ed520578b22f0..cc8d3050b790e 100644
--- a/sycl/cmake/modules/UnifiedRuntimeTag.cmake
+++ b/sycl/cmake/modules/UnifiedRuntimeTag.cmake
@@ -1,7 +1,7 @@
-# commit 93fc1331edfd5519ea1fab42a8b96abe9e53672f
-# Merge: 259b67db d2c04518
+# commit fe6c83a3a40e69eaa0c4829c8ba99b03e05bb4e2
+# Merge: 93fc1331 cf0c9482
 # Author: aarongreig <aaron.greig@codeplay.com>
-# Date:   Thu Dec 26 16:52:33 2024 +0000
-#     Merge pull request #2402 from yingcong-wu/yc/1202-mmap-failure
-#     [DevASAN][CPU] bugfix for mmap return value check.
-set(UNIFIED_RUNTIME_TAG 93fc1331edfd5519ea1fab42a8b96abe9e53672f)
+# Date:   Fri Dec 27 17:39:52 2024 +0000
+#     Merge pull request #2507 from AllanZyne/review/yang/fix_msan_shadow
+#     [DeviceMSAN] Fix MemToShadow algorithm and VA reservation
+set(UNIFIED_RUNTIME_TAG fe6c83a3a40e69eaa0c4829c8ba99b03e05bb4e2)