From 3fe23a36dbc4b5cea4ad1378d13222f7adebdc3d Mon Sep 17 00:00:00 2001
From: Patryk Kaminski <patryk.kaminski@intel.com>
Date: Wed, 24 Jul 2024 11:13:13 +0200
Subject: [PATCH] Disable memory poisoning in DisjointPool

ASan throws an error whenever the memory poison is called for the memory allocated on GPU:
"AddressSanitizer: CHECK failed: asan_mapping.h:359 "((AddrIsInMem(p))) != (0)" (0x0, 0x0)".

This commit disables poisoning.
---
 src/pool/pool_disjoint.cpp | 35 +++++++++++++++++++++++++++--------
 1 file changed, 27 insertions(+), 8 deletions(-)

diff --git a/src/pool/pool_disjoint.cpp b/src/pool/pool_disjoint.cpp
index d48d430ce..a1375a3ee 100644
--- a/src/pool/pool_disjoint.cpp
+++ b/src/pool/pool_disjoint.cpp
@@ -31,6 +31,27 @@
 #include "utils_math.h"
 #include "utils_sanitizers.h"
 
+// Temporary solution for disabling memory poisoning. This is needed because
+// AddressSanitizer does not support memory poisoning for GPU allocations.
+// More info: https://github.com/oneapi-src/unified-memory-framework/issues/634
+#ifndef POISON_MEMORY
+#define POISON_MEMORY 0
+#endif
+
+static inline void annotate_memory_inaccessible([[maybe_unused]] void *ptr,
+                                                [[maybe_unused]] size_t size) {
+#ifdef POISON_MEMORY
+    utils_annotate_memory_inaccessible(ptr, size);
+#endif
+}
+
+static inline void annotate_memory_undefined([[maybe_unused]] void *ptr,
+                                             [[maybe_unused]] size_t size) {
+#ifdef POISON_MEMORY
+    utils_annotate_memory_undefined(ptr, size);
+#endif
+}
+
 typedef struct umf_disjoint_pool_shared_limits_t {
     size_t MaxSize;
     std::atomic<size_t> TotalSize;
@@ -400,7 +421,7 @@ static void *memoryProviderAlloc(umf_memory_provider_handle_t hProvider,
     if (ret != UMF_RESULT_SUCCESS) {
         throw MemoryProviderError{ret};
     }
-    utils_annotate_memory_inaccessible(ptr, size);
+    annotate_memory_inaccessible(ptr, size);
     return ptr;
 }
 
@@ -822,7 +843,7 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, bool &FromPool) try {
     FromPool = false;
     if (Size > getParams().MaxPoolableSize) {
         Ptr = memoryProviderAlloc(getMemHandle(), Size);
-        utils_annotate_memory_undefined(Ptr, Size);
+        annotate_memory_undefined(Ptr, Size);
         return Ptr;
     }
 
@@ -839,7 +860,7 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, bool &FromPool) try {
     }
 
     VALGRIND_DO_MEMPOOL_ALLOC(this, Ptr, Size);
-    utils_annotate_memory_undefined(Ptr, Bucket.getSize());
+    annotate_memory_undefined(Ptr, Bucket.getSize());
 
     return Ptr;
 } catch (MemoryProviderError &e) {
@@ -877,7 +898,7 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, size_t Alignment,
     FromPool = false;
     if (AlignedSize > getParams().MaxPoolableSize) {
         Ptr = memoryProviderAlloc(getMemHandle(), Size, Alignment);
-        utils_annotate_memory_undefined(Ptr, Size);
+        annotate_memory_undefined(Ptr, Size);
         return Ptr;
     }
 
@@ -894,8 +915,7 @@ void *DisjointPool::AllocImpl::allocate(size_t Size, size_t Alignment,
     }
 
     VALGRIND_DO_MEMPOOL_ALLOC(this, AlignPtrUp(Ptr, Alignment), Size);
-    utils_annotate_memory_undefined(AlignPtrUp(Ptr, Alignment), Size);
-
+    annotate_memory_undefined(AlignPtrUp(Ptr, Alignment), Size);
     return AlignPtrUp(Ptr, Alignment);
 } catch (MemoryProviderError &e) {
     umf::getPoolLastStatusRef<DisjointPool>() = e.code;
@@ -962,8 +982,7 @@ void DisjointPool::AllocImpl::deallocate(void *Ptr, bool &ToPool) {
             }
 
             VALGRIND_DO_MEMPOOL_FREE(this, Ptr);
-            utils_annotate_memory_inaccessible(Ptr, Bucket.getSize());
-
+            annotate_memory_inaccessible(Ptr, Bucket.getSize());
             if (Bucket.getSize() <= Bucket.ChunkCutOff()) {
                 Bucket.freeChunk(Ptr, Slab, ToPool);
             } else {