ASAN: Remap first page of allocations to host mem

When compiling in ASAN mode, remap the first page of device allocations
to system memory. ASAN's memory allocator uses a small amount of extra
memory to store data for housekeeping purpose. But because this memory
is from the GPU memory pool, it might have uncommon memory type for host
to access. Mapping this section of memory to the host makes this memory
accessible to ASAN.

Change-Id: I36f659d616a4d15558372592439a8723c5c84a69
Signed-off-by: Bing Ma <Bing.Ma@amd.com>

runtime/hsa-runtime/core/inc/memory_region.h

@@ -91,6 +91,7 @@ class MemoryRegion : public Checked<0x9C961F19EE175BB3> {
     AllocateIPC = (1 << 4),         // System memory that can be IPC-shared
     AllocateNonPaged = (1 << 4),    // Non-paged system memory (AllocateIPC alias)
     AllocatePCIeRW = (1 << 5),      // Enforce pseudo fine grain/RW memory
+    AllocateAsan = (1 << 6),        // ASAN - First page of allocation remapped to system memory
   };
 
   typedef uint32_t AllocateFlags;

runtime/hsa-runtime/core/inc/runtime.h

@@ -70,6 +70,12 @@
 #include "core/inc/amd_loader_context.hpp"
 #include "core/inc/amd_hsa_code.hpp"
 
+#if defined(__clang__)
+#if __has_feature(address_sanitizer)
+#define SANITIZER_AMDGPU 1
+#endif
+#endif
+
 //---------------------------------------------------------------------------//
 //    Constants                                                              //
 //---------------------------------------------------------------------------//
@@ -407,9 +413,19 @@ class Runtime {
   static void AsyncEventsLoop(void*);
 
   struct AllocationRegion {
-    AllocationRegion() : region(NULL), size(0), size_requested(0), user_ptr(nullptr) {}
-    AllocationRegion(const MemoryRegion* region_arg, size_t size_arg, size_t size_requested)
-        : region(region_arg), size(size_arg), size_requested(size_requested), user_ptr(nullptr) {}
+    AllocationRegion()
+        : region(NULL),
+          size(0),
+          size_requested(0),
+          alloc_flags(core::MemoryRegion::AllocateNoFlags),
+          user_ptr(nullptr) {}
+    AllocationRegion(const MemoryRegion* region_arg, size_t size_arg, size_t size_requested,
+                     MemoryRegion::AllocateFlags alloc_flags)
+        : region(region_arg),
+          size(size_arg),
+          size_requested(size_requested),
+          alloc_flags(alloc_flags),
+          user_ptr(nullptr) {}
 
     struct notifier_t {
       void* ptr;
@@ -420,6 +436,7 @@ class Runtime {
     const MemoryRegion* region;
     size_t size;           /* actual size = align_up(size_requested, granularity) */
     size_t size_requested; /* size requested by user */
+    MemoryRegion::AllocateFlags alloc_flags;
     void* user_ptr;
     std::unique_ptr<std::vector<notifier_t>> notifiers;
   };

runtime/hsa-runtime/core/runtime/amd_memory_region.cpp

@@ -199,6 +199,13 @@ hsa_status_t MemoryRegion::AllocateImpl(size_t& size, AllocateFlags alloc_flags,
     useSubAlloc &= ((alloc_flags & (~AllocateRestrict)) == 0);
     if (useSubAlloc) {
       *address = fragment_allocator_.alloc(size);
+
+      if ((alloc_flags & AllocateAsan) &&
+          hsaKmtReplaceAsanHeaderPage(*address) != HSAKMT_STATUS_SUCCESS) {
+        fragment_allocator_.free(*address);
+        *address = NULL;
+        return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+      }
       return HSA_STATUS_SUCCESS;
     }
   }
@@ -253,6 +260,12 @@ hsa_status_t MemoryRegion::AllocateImpl(size_t& size, AllocateFlags alloc_flags,
       return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
     }
 
+    if ((alloc_flags & AllocateAsan) &&
+        hsaKmtReplaceAsanHeaderPage(*address) != HSAKMT_STATUS_SUCCESS) {
+      FreeKfdMemory(*address, size);
+      *address = NULL;
+      return HSA_STATUS_ERROR_OUT_OF_RESOURCES;
+    }
     return HSA_STATUS_SUCCESS;
   }
 

runtime/hsa-runtime/core/runtime/hsa.cpp

@@ -1084,6 +1084,8 @@ hsa_status_t
   TRY;
   IS_OPEN();
 
+  core::MemoryRegion::AllocateFlags alloc_flag = core::MemoryRegion::AllocateNoFlags;
+
   if (size == 0 || ptr == NULL) {
     return HSA_STATUS_ERROR_INVALID_ARGUMENT;
   }
@@ -1091,8 +1093,7 @@ hsa_status_t
   const core::MemoryRegion* mem_region = core::MemoryRegion::Convert(region);
   IS_VALID(mem_region);
 
-  return core::Runtime::runtime_singleton_->AllocateMemory(
-      mem_region, size, core::MemoryRegion::AllocateNoFlags, ptr);
+  return core::Runtime::runtime_singleton_->AllocateMemory(mem_region, size, alloc_flag, ptr);
   CATCH;
 }
 

runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp

@@ -761,6 +761,10 @@ hsa_status_t hsa_amd_memory_pool_allocate(hsa_amd_memory_pool_t memory_pool, siz
 
   if (flags == HSA_AMD_MEMORY_POOL_PCIE_FLAG) alloc_flag |= core::MemoryRegion::AllocatePCIeRW;
 
+#ifdef SANITIZER_AMDGPU
+  alloc_flag |= core::MemoryRegion::AllocateAsan;
+#endif
+
   return core::Runtime::runtime_singleton_->AllocateMemory(mem_region, size, alloc_flag, ptr);
   CATCH;
 }

runtime/hsa-runtime/core/runtime/runtime.cpp

@@ -291,7 +291,7 @@ hsa_status_t Runtime::AllocateMemory(const MemoryRegion* region, size_t size,
   // Track the allocation result so that it could be freed properly.
   if (status == HSA_STATUS_SUCCESS) {
     ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
-    allocation_map_[*address] = AllocationRegion(region, size, size_requested);
+    allocation_map_[*address] = AllocationRegion(region, size, size_requested, alloc_flags);
   }
 
   return status;
@@ -305,6 +305,7 @@ hsa_status_t Runtime::FreeMemory(void* ptr) {
   const MemoryRegion* region = nullptr;
   size_t size = 0;
   std::unique_ptr<std::vector<AllocationRegion::notifier_t>> notifiers;
+  MemoryRegion::AllocateFlags alloc_flags = core::MemoryRegion::AllocateNoFlags;
 
   {
     ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
@@ -317,6 +318,7 @@ hsa_status_t Runtime::FreeMemory(void* ptr) {
     }
     region = it->second.region;
     size = it->second.size;
+    alloc_flags = it->second.alloc_flags;
 
     // Imported fragments can't be released with FreeMemory.
     if (region == nullptr) {
@@ -338,6 +340,9 @@ hsa_status_t Runtime::FreeMemory(void* ptr) {
     }
   }
 
+  if (alloc_flags & core::MemoryRegion::AllocateAsan)
+    assert(hsaKmtReturnAsanHeaderPage(ptr) == HSAKMT_STATUS_SUCCESS);
+
   return region->Free(ptr, size);
 }
 
@@ -1025,7 +1030,8 @@ hsa_status_t Runtime::IPCAttach(const hsa_amd_ipc_memory_t* handle, size_t len,
       len = Min(len, importSize - fragOffset);
     }
     ScopedAcquire<KernelSharedMutex> lock(&memory_lock_);
-    allocation_map_[importAddress] = AllocationRegion(nullptr, len, len);
+    allocation_map_[importAddress] =
+        AllocationRegion(nullptr, len, len, core::MemoryRegion::AllocateNoFlags);
   };
 
   if ((importHandle.handle[6] & 0x80000000) != 0) {