[Sanitizer] Support only device compilation on CPU device (#10)

* [Sanitizer] Support only device compilation on CPU device
* Refine some logics

source/loader/CMakeLists.txt

@@ -117,12 +117,12 @@ if(UR_ENABLE_SANITIZER)
     if(WIN32)
         target_sources(ur_loader
             PRIVATE
-            ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/windows/check_san.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/windows/san_utils.cpp
         )
     else()
         target_sources(ur_loader
             PRIVATE
-            ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/linux/check_san.cpp
+            ${CMAKE_CURRENT_SOURCE_DIR}/layers/sanitizer/linux/san_utils.cpp
         )
     endif()
 

source/loader/layers/sanitizer/asan_interceptor.cpp

@@ -101,7 +101,14 @@ std::string getKernelName(ur_kernel_handle_t Kernel) {
 } // namespace
 
 SanitizerInterceptor::SanitizerInterceptor()
-    : m_IsInASanContext(IsInASanContext()) {}
+    : m_IsInASanContext(IsInASanContext()),
+      m_ShadowMemInited(m_IsInASanContext) {}
+
+SanitizerInterceptor::~SanitizerInterceptor() {
+    if (!m_IsInASanContext && m_ShadowMemInited && !DestroyShadowMem()) {
+        context.logger.error("Failed to destroy shadow memory");
+    }
+}
 
 /// The memory chunk allocated from the underlying allocator looks like this:
 /// L L L L L L U U U U U U R R
@@ -289,20 +296,20 @@ ur_result_t SanitizerInterceptor::allocShadowMemory(
     ur_context_handle_t Context, std::shared_ptr<DeviceInfo> &DeviceInfo) {
     if (DeviceInfo->Type == DeviceType::CPU) {
         if (!m_IsInASanContext) {
-            context.logger.error("Host AddressSanitizer needs to be enabled");
-            return UR_RESULT_ERROR_INVALID_CONTEXT;
+            static std::once_flag OnceFlag;
+            bool Result = true;
+            std::call_once(OnceFlag, [&]() {
+                Result = m_ShadowMemInited = SetupShadowMem();
+            });
+
+            if (!Result) {
+                context.logger.error("Failed to allocate shadow memory");
+                return UR_RESULT_ERROR_OUT_OF_RESOURCES;
+            }
         }
 
-        // Based on "compiler-rt/lib/asan/asan_mapping.h"
-        // Typical shadow mapping on Linux/x86_64 with SHADOW_OFFSET == 0x00007fff8000:
-        DeviceInfo->ShadowOffset = 0x00007fff8000ULL;
-        DeviceInfo->ShadowOffsetEnd = 0x10007fff7fffULL;
-        // // Default Linux/i386 mapping on x86_64 machine:
-        // DeviceInfo->ShadowOffset = 0x20000000ULL;
-        // DeviceInfo->ShadowOffsetEnd = 0x3fffffffULL;
-        // // Default Linux/i386 mapping on i386 machine
-        // DeviceInfo->ShadowOffset = 0x20000000ULL;
-        // DeviceInfo->ShadowOffsetEnd = 0x37ffffffULL;
+        DeviceInfo->ShadowOffset = LOW_SHADOW_BEGIN;
+        DeviceInfo->ShadowOffsetEnd = HIGH_SHADOW_END;
     } else if (DeviceInfo->Type == DeviceType::GPU_PVC) {
         /// SHADOW MEMORY MAPPING (PVC, with CPU 47bit)
         ///   Host/Shared USM : 0x0              ~ 0x0fff_ffff_ffff

source/loader/layers/sanitizer/asan_interceptor.hpp

@@ -99,6 +99,8 @@ class SanitizerInterceptor {
   public:
     SanitizerInterceptor();
 
+    ~SanitizerInterceptor();
+
     ur_result_t allocateMemory(ur_context_handle_t Context,
                                ur_device_handle_t Device,
                                const ur_usm_desc_t *Properties,
@@ -158,6 +160,7 @@ class SanitizerInterceptor {
     ur_shared_mutex m_ContextMapMutex;
 
     bool m_IsInASanContext;
+    bool m_ShadowMemInited;
 };
 
 } // namespace ur_sanitizer_layer

source/loader/layers/sanitizer/common.hpp

@@ -30,6 +30,18 @@ using u32 = unsigned int;
 constexpr unsigned ASAN_SHADOW_SCALE = 3;
 constexpr unsigned ASAN_SHADOW_GRANULARITY = 1ULL << ASAN_SHADOW_SCALE;
 
+// Based on "compiler-rt/lib/asan/asan_mapping.h"
+// Typical shadow mapping on Linux/x86_64 with SHADOW_OFFSET == 0x00007fff8000:
+constexpr uptr LOW_SHADOW_BEGIN = 0x00007fff8000ULL;
+constexpr uptr LOW_SHADOW_END = 0x00008fff6fffULL;
+constexpr uptr SHADOW_GAP_BEGIN = 0x00008fff7000ULL;
+constexpr uptr SHADOW_GAP_END = 0x02008fff6fffULL;
+constexpr uptr HIGH_SHADOW_BEGIN = 0x02008fff7000ULL;
+constexpr uptr HIGH_SHADOW_END = 0x10007fff7fffULL;
+constexpr uptr LOW_SHADOW_SIZE = LOW_SHADOW_END - LOW_SHADOW_BEGIN;
+constexpr uptr SHADOW_GAP_SIZE = SHADOW_GAP_END - SHADOW_GAP_BEGIN;
+constexpr uptr HIGH_SHADOW_SIZE = HIGH_SHADOW_END - HIGH_SHADOW_BEGIN;
+
 inline constexpr bool IsPowerOfTwo(uptr x) {
     return (x & (x - 1)) == 0 && x != 0;
 }
@@ -89,4 +101,8 @@ inline constexpr uptr ComputeRZLog(uptr user_requested_size) {
 
 bool IsInASanContext();
 
+bool SetupShadowMem();
+
+bool DestroyShadowMem();
+
 } // namespace ur_sanitizer_layer

source/loader/layers/sanitizer/linux/san_utils.cpp

@@ -0,0 +1,57 @@
+#include "common.hpp"
+#include <sys/mman.h>
+#include <asm/param.h>
+
+extern "C" __attribute__((weak)) void __asan_init(void);
+
+namespace ur_sanitizer_layer {
+
+bool IsInASanContext() { return __asan_init != nullptr; }
+
+static bool ReserveShadowMem(uptr Addr, uptr Size) {
+    Size = RoundUpTo(Size, EXEC_PAGESIZE);
+    Addr = RoundDownTo(Addr, EXEC_PAGESIZE);
+    void *P =
+        mmap((void *)Addr, Size, PROT_READ | PROT_WRITE,
+             MAP_PRIVATE | MAP_FIXED | MAP_NORESERVE | MAP_ANONYMOUS, -1, 0);
+    return Addr == (uptr)P;
+}
+
+static bool ProtectShadowGap(uptr Addr, uptr Size) {
+    void *P =
+        mmap((void *)Addr, Size, PROT_NONE,
+             MAP_PRIVATE | MAP_FIXED | MAP_NORESERVE | MAP_ANONYMOUS, -1, 0);
+    return Addr == (uptr)P;
+}
+
+bool SetupShadowMem() {
+    if (!ReserveShadowMem(LOW_SHADOW_BEGIN, LOW_SHADOW_SIZE)) {
+        return false;
+    }
+
+    if (!ReserveShadowMem(HIGH_SHADOW_BEGIN, HIGH_SHADOW_SIZE)) {
+        return false;
+    }
+
+    if (!ProtectShadowGap(SHADOW_GAP_BEGIN, SHADOW_GAP_SIZE)) {
+        return false;
+    }
+    return true;
+}
+
+bool DestroyShadowMem() {
+    if (munmap((void *)LOW_SHADOW_BEGIN, LOW_SHADOW_SIZE) == -1) {
+        return false;
+    }
+
+    if (munmap((void *)HIGH_SHADOW_BEGIN, HIGH_SHADOW_SIZE) == -1) {
+        return false;
+    }
+
+    if (munmap((void *)SHADOW_GAP_BEGIN, SHADOW_GAP_SIZE) == -1) {
+        return false;
+    }
+    return true;
+}
+
+} // namespace ur_sanitizer_layer