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) {