diff --git a/src/core/impl/comp_node/cuda/comp_node.cpp b/src/core/impl/comp_node/cuda/comp_node.cpp index e89d56708b2b082e7b8270d7fbd08b6d59052004..73e8dfb49633f592781ad77799a4ac22090e2cea 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 277700860f6c602682d2ccf8ef292a74a7340e10..4f6d44a87bb84e57f7ac0e225ac22fb6c22c57f5 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 c3515a44ab2b5434f13bc1e235ea4fb268adc350..da226639abe788282871e7afcd1de5ec2dd99d70 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 99112723918782625edd7494ae4407db1c1a6693..6a9f7fb3e52c92d387c5df536529b7be1076c570 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 33f9539dc00f026764e2714a15853c5b1f77afdd..18993a149c79401e1442133b4ac3c9d080d7e572 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; };