From b4df74f1989a9fe062d57c8f21d302d243420293 Mon Sep 17 00:00:00 2001 From: tiancyin Date: Tue, 22 Oct 2024 15:02:36 +0800 Subject: [PATCH] wsl/hsakmt: add sub-allocator Reviewed-by: Shane Xiao Reviewed-by: lyndonli Signed-off-by: tiancyin --- libhsakmt.h | 24 +++++++ memory.cpp | 151 ++++++++++++++++++++++++++++++++++++++++----- openclose.cpp | 2 + topology.cpp | 1 + util/simple_heap.h | 30 ++++++++- 5 files changed, 190 insertions(+), 18 deletions(-) diff --git a/libhsakmt.h b/libhsakmt.h index 0a2627ced3..aebba33d7e 100644 --- a/libhsakmt.h +++ b/libhsakmt.h @@ -177,6 +177,30 @@ bool is_forked_child(void); void clear_allocation_map(void); +class BlockAllocator { +private: + static const size_t block_size_ = 128 * 1024 * 1024; // 128MB blocks. + +public: + void* alloc(size_t request_size, size_t& allocated_size) const; + void free(void* ptr, size_t length) const; + size_t block_size() const { return block_size_; } +}; + +void reset_suballocator(void); +void trim_suballocator(void); + +HSAKMT_STATUS hsaKmtAllocMemoryAlignInternal(HSAuint32 PreferredNode, + HSAuint64 SizeInBytes, + HSAuint64 Alignment, + HsaMemFlags MemFlags, + void **MemoryAddress, + bool SkipSubAlloc = false); + +HSAKMT_STATUS hsaKmtFreeMemoryInternal(void *MemoryAddress, + HSAuint64 SizeInBytes, + bool SkipSubAlloc = false); + bool queue_acquire_buffer(void *MemoryAddress); bool queue_release_buffer(void *MemoryAddress); /* Calculate VGPR and SGPR register file size per CU */ diff --git a/memory.cpp b/memory.cpp index 04f73ef0fd..13d10f9849 100644 --- a/memory.cpp +++ b/memory.cpp @@ -33,6 +33,7 @@ #include #include #include "inc/wddm/gpu_memory.h" +#include "util/simple_heap.h" struct Allocation { Allocation() @@ -112,11 +113,45 @@ bool isSystemMemoryAvailable(HSAuint64 SizeInBytes) { return SizeInBytes <= info.freeram; } -HSAKMT_STATUS HSAKMTAPI hsaKmtAllocMemoryAlign(HSAuint32 PreferredNode, - HSAuint64 SizeInBytes, - HSAuint64 Alignment, - HsaMemFlags MemFlags, - void **MemoryAddress) { +void* BlockAllocator::alloc(size_t request_size, size_t& allocated_size) const { + void *address; + HsaMemFlags MemFlags; + + MemFlags.Value = 0; + MemFlags.ui32.CoarseGrain = 1; + MemFlags.ui32.NoSubstitute = 1; + allocated_size = wsl::AlignUp(request_size, block_size()); + if (HSAKMT_STATUS_SUCCESS == hsaKmtAllocMemoryAlignInternal(1, allocated_size, 0, MemFlags, &address, true)) + return address; + + return nullptr; +} + +void BlockAllocator::free(void* ptr, size_t length) const { + if (HSAKMT_STATUS_SUCCESS != hsaKmtFreeMemoryInternal(ptr, length, true)) + pr_err("wsl-thunk: BlockAllocator::free() err, address %p, length:%zu\n", ptr, length); +} + +static wsl::SimpleHeap fragment_allocator_; +static std::unique_ptr fragment_allocator_lock_ = std::make_unique(); + +void reset_suballocator(void) { + fragment_allocator_lock_ = std::make_unique(); + std::lock_guard lock(*fragment_allocator_lock_); + fragment_allocator_.reset(); +} + +void trim_suballocator(void) { + std::lock_guard lock(*fragment_allocator_lock_); + fragment_allocator_.trim(); +} + +HSAKMT_STATUS hsaKmtAllocMemoryAlignInternal(HSAuint32 PreferredNode, + HSAuint64 SizeInBytes, + HSAuint64 Alignment, + HsaMemFlags MemFlags, + void **MemoryAddress, + bool SkipSubAlloc) { CHECK_DXG_OPEN(); if (!MemoryAddress) @@ -174,6 +209,33 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtAllocMemoryAlign(HSAuint32 PreferredNode, if (create_info.flags.virtual_alloc || create_info.flags.physical_only) create_info.domain = thunk_proxy::AllocDomain::kLocal; + /* Only allow using the suballocator for ordinary VRAM.*/ + bool trim_safe = false; + if (!SkipSubAlloc && + create_info.domain == thunk_proxy::AllocDomain::kLocal && + !(create_info.flags.virtual_alloc || create_info.flags.physical_only)) { + std::lock_guard gard(*fragment_allocator_lock_); + + /* just quickly skip SA if size is bigger than SA block size.*/ + gpusize real_size; + if (create_info.size > GPU_HUGE_PAGE_SIZE) + real_size = wsl::AlignUp(create_info.size, GPU_HUGE_PAGE_SIZE); + else + real_size = wsl::AlignUp(create_info.size, getpagesize()); + + if (real_size < fragment_allocator_.default_block_size()) { + *MemoryAddress = fragment_allocator_.alloc(real_size); + if (*MemoryAddress) + return HSAKMT_STATUS_SUCCESS; + } + + /* SA might keep a lot of free blocks as *cache*. + * We can trim them if direct allocation fails at first time. + */ + trim_safe = true; + } + +after_trim: auto code = dev->CreateGpuMemory(create_info, &gpu_mem); if (code == ErrorCode::Success) { std::lock_guard gard(*allocation_map_lock_); @@ -189,18 +251,41 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtAllocMemoryAlign(HSAuint32 PreferredNode, create_info.size, false, nullptr, SizeInBytes, MemFlags.ui32.GTTAccess ? 0 : PreferredNode, MemFlags.Value); return HSAKMT_STATUS_SUCCESS; + } else if (trim_safe) { + /* attempt to release memory from the block allocator and retry */ + std::lock_guard gard(*fragment_allocator_lock_); + fragment_allocator_.trim(); + trim_safe = false; + goto after_trim; } return HSAKMT_STATUS_ERROR; } -HSAKMT_STATUS HSAKMTAPI hsaKmtFreeMemory(void *MemoryAddress, - HSAuint64 SizeInBytes) { +HSAKMT_STATUS HSAKMTAPI hsaKmtAllocMemoryAlign(HSAuint32 PreferredNode, + HSAuint64 SizeInBytes, + HSAuint64 Alignment, + HsaMemFlags MemFlags, + void **MemoryAddress) { + return hsaKmtAllocMemoryAlignInternal(PreferredNode, SizeInBytes, + Alignment, MemFlags, + MemoryAddress); +} + +HSAKMT_STATUS hsaKmtFreeMemoryInternal(void *MemoryAddress, + HSAuint64 SizeInBytes, + bool SkipSubAlloc) { CHECK_DXG_OPEN(); if (!MemoryAddress) return HSAKMT_STATUS_INVALID_PARAMETER; + if (!SkipSubAlloc) { + std::lock_guard gard(*fragment_allocator_lock_); + if (fragment_allocator_.free(MemoryAddress)) + return HSAKMT_STATUS_SUCCESS; + } + wsl::thunk::GpuMemory *gpu_mem = nullptr; { std::lock_guard gard(*allocation_map_lock_); @@ -220,23 +305,28 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtFreeMemory(void *MemoryAddress, return HSAKMT_STATUS_SUCCESS; } +HSAKMT_STATUS HSAKMTAPI hsaKmtFreeMemory(void *MemoryAddress, + HSAuint64 SizeInBytes) { + return hsaKmtFreeMemoryInternal(MemoryAddress, SizeInBytes); +} + bool queue_acquire_buffer(void *MemoryAddress) { if (!MemoryAddress) - return false; + return false; wsl::thunk::GpuMemory *gpu_mem = nullptr; { - std::lock_guard gard(*allocation_map_lock_); - auto it = allocation_map_.find(MemoryAddress); - if (it == allocation_map_.end()) { - return HSAKMT_STATUS_ERROR; - } + std::lock_guard gard(*allocation_map_lock_); + auto it = allocation_map_.find(MemoryAddress); + if (it == allocation_map_.end()) { + return HSAKMT_STATUS_ERROR; + } - gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle); - gpu_mem->GetQueueReference(); + gpu_mem = wsl::thunk::GpuMemory::Convert(it->second.handle); + gpu_mem->GetQueueReference(); } if (gpu_mem == nullptr) - return false; + return false; return true; } @@ -478,6 +568,12 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtMapMemoryToGPU(void *MemoryAddress, void *aligned_ptr = (void *)start; size_t aligned_size = end - start; + { + std::lock_guard gard(*fragment_allocator_lock_); + if (nullptr != fragment_allocator_.block_base(aligned_ptr)) + return HSAKMT_STATUS_SUCCESS; + } + { std::lock_guard gard(*allocation_map_lock_); // GTT mem @@ -557,6 +653,12 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtUnmapMemoryToGPU(void *MemoryAddress) { pr_debug("address %p\n", MemoryAddress); + { + std::lock_guard gard(*fragment_allocator_lock_); + if (nullptr != fragment_allocator_.block_base(MemoryAddress)) + return HSAKMT_STATUS_SUCCESS; + } + wsl::thunk::GpuMemoryHandle handle = nullptr; { std::lock_guard gard(*allocation_map_lock_); @@ -623,13 +725,21 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtQueryPointerInfo(const void *Pointer, return HSAKMT_STATUS_INVALID_PARAMETER; pr_debug("pointer %p\n", Pointer); + void *pointer = const_cast(Pointer); memset(PointerInfo, 0, sizeof(HsaPointerInfo)); + void* block_base = nullptr; + { + std::lock_guard gard(*fragment_allocator_lock_); + block_base = fragment_allocator_.block_base(pointer); + if (block_base != nullptr) + pointer = block_base; + } Allocation allocation_info; { std::lock_guard gard(*allocation_map_lock_); - auto it = allocation_map_.find(Pointer); + auto it = allocation_map_.find(pointer); if (it == allocation_map_.end()) { PointerInfo->Type = HSA_POINTER_UNKNOWN; return HSAKMT_STATUS_ERROR; @@ -649,6 +759,13 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtQueryPointerInfo(const void *Pointer, PointerInfo->MemFlags.Value = allocation_info.mem_flags_value; PointerInfo->CPUAddress = allocation_info.cpu_addr; PointerInfo->GPUAddress = allocation_info.gpu_addr; + if (block_base != nullptr) { + uint64_t offset = reinterpret_cast(Pointer) - + reinterpret_cast(block_base); + PointerInfo->CPUAddress = reinterpret_cast( + reinterpret_cast(PointerInfo->CPUAddress) + offset); + PointerInfo->GPUAddress += offset; + } PointerInfo->UserData = allocation_info.user_data; return HSAKMT_STATUS_SUCCESS; diff --git a/openclose.cpp b/openclose.cpp index 8ff4abec65..55f48c1e7e 100644 --- a/openclose.cpp +++ b/openclose.cpp @@ -80,6 +80,7 @@ static void child_fork_handler(void) { * in the child process so it is not cleared */ static void clear_after_fork(void) { + reset_suballocator(); clear_allocation_map(); if (dxg_fd) { close(dxg_fd); @@ -194,6 +195,7 @@ HSAKMT_STATUS HSAKMTAPI hsaKmtOpenKFD(void) { result = HSAKMT_STATUS_KERNEL_ALREADY_OPENED; } + reset_suballocator(); pthread_mutex_unlock(&hsakmt_mutex); return result; topology_sysfs_failed: diff --git a/topology.cpp b/topology.cpp index 906ae6f9cb..35c45d41f4 100644 --- a/topology.cpp +++ b/topology.cpp @@ -1377,6 +1377,7 @@ void topology_drop_snapshot(void) { free(g_system); g_system = NULL; + trim_suballocator(); for (auto device : wdevices_) delete device; wdevices_.clear(); diff --git a/util/simple_heap.h b/util/simple_heap.h index eca19055d8..1181682fa0 100644 --- a/util/simple_heap.h +++ b/util/simple_heap.h @@ -51,7 +51,6 @@ #include #include -#include "core/util/utils.h" namespace wsl { @@ -233,6 +232,35 @@ template class SimpleHeap { return reinterpret_cast(base); } + /* Return block-base the ptr belongs to if the ptr is a valid ptr which is allocated + * from this simpleheap and the block-base is allocated from block_allocator_*/ + void* block_base(void* ptr) { + if (ptr == nullptr) + return nullptr; + + uintptr_t base = reinterpret_cast(ptr); + + // Find fragment and validate. + auto frag_map_it = block_list_.upper_bound(base); + if (frag_map_it == block_list_.begin()) + return nullptr; + frag_map_it--; + auto& frag_map = frag_map_it->second; + auto fragment = frag_map.find(base); + if (fragment == frag_map.end() || isFree(fragment->second)) + return nullptr; + + return reinterpret_cast(frag_map_it->first); + } + + void reset() { + free_list_.clear(); + block_list_.clear(); + block_cache_.clear(); + in_use_size_ = 0; + cache_size_ = 0; + } + bool free(void* ptr) { if (ptr == nullptr) return true;