diff --git a/source/loader/layers/sanitizer/asan_shadow.cpp b/source/loader/layers/sanitizer/asan_shadow.cpp
index 629ce3a491..7ffe917b49 100644
--- a/source/loader/layers/sanitizer/asan_shadow.cpp
+++ b/source/loader/layers/sanitizer/asan_shadow.cpp
@@ -256,21 +256,14 @@ ur_result_t ShadowMemoryGPU::ReleaseShadow(std::shared_ptr<AllocInfo> AI) {
 }
 
 uptr ShadowMemoryPVC::MemToShadow(uptr Ptr) {
-    if (Ptr & 0xFF00000000000000ULL) { // Device USM
-        return ShadowBegin + 0x80000000000ULL +
-               ((Ptr & 0xFFFFFFFFFFFFULL) >> ASAN_SHADOW_SCALE);
-    } else { // Only consider 47bit VA
-        return ShadowBegin + ((Ptr & 0x7FFFFFFFFFFFULL) >> ASAN_SHADOW_SCALE);
-    }
+    return ShadowBegin + (((((Ptr & 0x8000'0000'0000'0000) >> 16) + Ptr) &
+                           0xf'ffff'ffff'ffff) >>
+                          ASAN_SHADOW_SCALE);
 }
 
 uptr ShadowMemoryDG2::MemToShadow(uptr Ptr) {
-    if (Ptr & 0xFFFF000000000000ULL) { // Device USM
-        return ShadowBegin + 0x80000000000ULL +
-               ((Ptr & 0x7FFFFFFFFFFFULL) >> ASAN_SHADOW_SCALE);
-    } else { // Host/Shared USM
-        return ShadowBegin + (Ptr >> ASAN_SHADOW_SCALE);
-    }
+    return ShadowBegin +
+           ((Ptr & 0x0000'FFFF'FFFF'FFFFULL) >> ASAN_SHADOW_SCALE);
 }
 
 } // namespace ur_sanitizer_layer
diff --git a/source/loader/layers/sanitizer/asan_shadow.hpp b/source/loader/layers/sanitizer/asan_shadow.hpp
index 7ae095062a..35e16d6e7e 100644
--- a/source/loader/layers/sanitizer/asan_shadow.hpp
+++ b/source/loader/layers/sanitizer/asan_shadow.hpp
@@ -109,7 +109,7 @@ struct ShadowMemoryPVC final : public ShadowMemoryGPU {
     size_t GetShadowSize() override { return 0x180000000000ULL; }
 };
 
-/// Shadow Memory layout of GPU PVC device
+/// Shadow Memory layout of GPU DG2 device
 ///
 /// USM Allocation Range (48 bits)
 ///   Host/Shared USM : 0x0000_0000_0000_0000 ~ 0x0000_7fff_ffff_ffff
diff --git a/source/loader/layers/sanitizer/ur_sanddi.cpp b/source/loader/layers/sanitizer/ur_sanddi.cpp
index 95b1649691..f544dfaaea 100644
--- a/source/loader/layers/sanitizer/ur_sanddi.cpp
+++ b/source/loader/layers/sanitizer/ur_sanddi.cpp
@@ -73,6 +73,12 @@ __urdlllocal ur_result_t UR_APICALL urAdapterGet(
         return UR_RESULT_ERROR_UNSUPPORTED_FEATURE;
     }
 
+    // Set env AllocateHostAllocationsInHeapExtendedHost to 0, then host and
+    // shared usm will have same allocation range on GPU device. This will
+    // simplify shadow memory mapping logic on GPU device.
+    setenv("NEOReadDebugKeys", "1", 1);
+    setenv("AllocateHostAllocationsInHeapExtendedHost", "0", 1);
+
     ur_result_t result = pfnAdapterGet(NumEntries, phAdapters, pNumAdapters);
     if (result == UR_RESULT_SUCCESS && phAdapters) {
         const uint32_t NumAdapters = pNumAdapters ? *pNumAdapters : NumEntries;