From 62c394ca9156759cdc2c6f6b281ac858dbb07833 Mon Sep 17 00:00:00 2001 From: Megvii Engine Team Date: Tue, 9 Mar 2021 14:56:05 +0800 Subject: [PATCH] feat(cuda/comp_node): enable to query adjacent free blocks size GitOrigin-RevId: 0cd1d202dd05c65c2231a469b511ca13974d43a6 --- src/core/impl/comp_node/cuda/comp_node.cpp | 6 ++++++ src/core/impl/comp_node/mem_alloc/impl.cpp | 21 +++++++++++++++++++++ src/core/impl/comp_node/mem_alloc/impl.h | 4 ++++ src/core/include/megbrain/comp_node.h | 7 +++++++ src/core/include/megbrain/comp_node/alloc.h | 9 +++++++++ 5 files changed, 47 insertions(+) diff --git a/src/core/impl/comp_node/cuda/comp_node.cpp b/src/core/impl/comp_node/cuda/comp_node.cpp index e89d5670..73e8dfb4 100644 --- a/src/core/impl/comp_node/cuda/comp_node.cpp +++ b/src/core/impl/comp_node/cuda/comp_node.cpp @@ -269,6 +269,12 @@ class CudaCompNode::CompNodeImpl final: public CompNode::Impl { return {tot, free}; } +#if !MGB_BUILD_SLIM_SERVING + std::pair get_free_left_and_right(size_t begin_ptr, size_t end_ptr) override { + return m_mem_alloc->get_free_left_and_right(begin_ptr, end_ptr); + } +#endif + Locator locator() override { return m_locator; } diff --git a/src/core/impl/comp_node/mem_alloc/impl.cpp b/src/core/impl/comp_node/mem_alloc/impl.cpp index 27770086..4f6d44a8 100644 --- a/src/core/impl/comp_node/mem_alloc/impl.cpp +++ b/src/core/impl/comp_node/mem_alloc/impl.cpp @@ -21,6 +21,27 @@ using namespace mem_alloc; /* ===================== MemAllocImplHelper ===================== */ +#if !MGB_BUILD_SLIM_SERVING +std::pair MemAllocImplHelper::get_free_left_and_right(size_t begin_ptr, size_t end_ptr) { + MGB_LOCK_GUARD(m_mutex); + auto iter = m_free_blk_addr.lower_bound(begin_ptr); + size_t left_free = 0, right_free = 0; + if (iter != m_free_blk_addr.begin()) { + auto prev = iter; + prev --; + if (prev->first + prev->second.size == begin_ptr) { + left_free = prev->second.size; + } + } + if (iter != m_free_blk_addr.end()) { + if (iter->first == end_ptr) { + right_free = iter->second.size; + } + } + return {left_free, right_free}; +} +#endif + MemAllocImplHelper::MemAddr MemAllocImplHelper::do_alloc( size_t size, bool allow_from_parent, bool log_stat_on_error) { diff --git a/src/core/impl/comp_node/mem_alloc/impl.h b/src/core/impl/comp_node/mem_alloc/impl.h index c3515a44..da226639 100644 --- a/src/core/impl/comp_node/mem_alloc/impl.h +++ b/src/core/impl/comp_node/mem_alloc/impl.h @@ -115,6 +115,10 @@ class MemAllocImplHelper: virtual public MemAllocBase { //! get free mem for this allocator, without locking FreeMemStat get_free_memory_self_unsafe(); +#if !MGB_BUILD_SLIM_SERVING + std::pair get_free_left_and_right(size_t begin_ptr, size_t end_ptr) override; +#endif + public: void print_memory_state() override; diff --git a/src/core/include/megbrain/comp_node.h b/src/core/include/megbrain/comp_node.h index 99112723..6a9f7fb3 100644 --- a/src/core/include/megbrain/comp_node.h +++ b/src/core/include/megbrain/comp_node.h @@ -352,6 +352,10 @@ class CompNode { } #if !MGB_BUILD_SLIM_SERVING + std::pair get_free_left_and_right(size_t begin_ptr, size_t end_ptr) { + return m_impl->get_free_left_and_right(begin_ptr, end_ptr); + } + size_t get_used_memory() const { return m_impl->get_used_memory(); } @@ -535,6 +539,9 @@ class CompNode { virtual std::pair get_mem_status_bytes() = 0; #if !MGB_BUILD_SLIM_SERVING + virtual std::pair get_free_left_and_right(size_t x, size_t y) { + return {x - x, y - y}; + } virtual size_t get_used_memory() { return 0; } diff --git a/src/core/include/megbrain/comp_node/alloc.h b/src/core/include/megbrain/comp_node/alloc.h index 33f9539d..18993a14 100644 --- a/src/core/include/megbrain/comp_node/alloc.h +++ b/src/core/include/megbrain/comp_node/alloc.h @@ -134,6 +134,15 @@ class MemAllocBase { */ virtual FreeMemStat get_free_memory_dev() = 0; +#if !MGB_BUILD_SLIM_SERVING + /*! + * \brief get free memory adjacent to interval [begin_ptr, end_ptr] + */ + virtual std::pair get_free_left_and_right(size_t begin_ptr, size_t end_ptr) { + return {0, 0}; + } +#endif + virtual ~MemAllocBase() = default; }; -- GitLab