From 50e754d08b6a9021c95019b3eec5d5bc60cfdbe7 Mon Sep 17 00:00:00 2001 From: David Yat Sin Date: Thu, 11 May 2023 21:27:39 +0000 Subject: [PATCH] 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 --- runtime/hsa-runtime/core/inc/memory_region.h | 1 + runtime/hsa-runtime/core/inc/runtime.h | 23 ++++++++++++++++--- .../core/runtime/amd_memory_region.cpp | 13 +++++++++++ runtime/hsa-runtime/core/runtime/hsa.cpp | 5 ++-- .../hsa-runtime/core/runtime/hsa_ext_amd.cpp | 4 ++++ runtime/hsa-runtime/core/runtime/runtime.cpp | 10 ++++++-- 6 files changed, 49 insertions(+), 7 deletions(-) diff --git a/runtime/hsa-runtime/core/inc/memory_region.h b/runtime/hsa-runtime/core/inc/memory_region.h index 48ecb3dcdc..feba80695b 100644 --- a/runtime/hsa-runtime/core/inc/memory_region.h +++ b/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; diff --git a/runtime/hsa-runtime/core/inc/runtime.h b/runtime/hsa-runtime/core/inc/runtime.h index 11e8ce82b1..b5307532c8 100644 --- a/runtime/hsa-runtime/core/inc/runtime.h +++ b/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> notifiers; }; diff --git a/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp b/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp index b6f91efa7d..a8e4c75c4e 100644 --- a/runtime/hsa-runtime/core/runtime/amd_memory_region.cpp +++ b/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; } diff --git a/runtime/hsa-runtime/core/runtime/hsa.cpp b/runtime/hsa-runtime/core/runtime/hsa.cpp index bcc1414573..c509fda5ec 100644 --- a/runtime/hsa-runtime/core/runtime/hsa.cpp +++ b/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; } diff --git a/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp b/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp index a4feac9894..09c5976e42 100644 --- a/runtime/hsa-runtime/core/runtime/hsa_ext_amd.cpp +++ b/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; } diff --git a/runtime/hsa-runtime/core/runtime/runtime.cpp b/runtime/hsa-runtime/core/runtime/runtime.cpp index ee095f0119..dbc2e2cb8f 100644 --- a/runtime/hsa-runtime/core/runtime/runtime.cpp +++ b/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 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> notifiers; + MemoryRegion::AllocateFlags alloc_flags = core::MemoryRegion::AllocateNoFlags; { ScopedAcquire 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 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) {