feat(cuda/comp_node): enable to query adjacent free blocks size

GitOrigin-RevId: 0cd1d202dd

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<size_t, size_t> 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;
         }

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

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<size_t, size_t> get_free_left_and_right(size_t begin_ptr, size_t end_ptr) override;
+#endif
+
     public:
         void print_memory_state() override;
 

src/core/include/megbrain/comp_node.h

@@ -352,6 +352,10 @@ class CompNode {
         }
 
 #if !MGB_BUILD_SLIM_SERVING
+        std::pair<size_t, size_t> 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<size_t, size_t> get_mem_status_bytes() = 0;
 
 #if !MGB_BUILD_SLIM_SERVING
+                virtual std::pair<size_t, size_t> get_free_left_and_right(size_t x, size_t y) {
+                    return {x - x, y - y};
+                }
                 virtual size_t get_used_memory() {
                     return 0;
                 }

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<size_t, size_t> get_free_left_and_right(size_t begin_ptr, size_t end_ptr) {
+            return {0, 0};
+        }
+#endif
+
         virtual ~MemAllocBase() = default;
 };