usearch icon indicating copy to clipboard operation
usearch copied to clipboard

Feature: More accurate memory statistics

Open mapleFU opened this issue 6 months ago • 3 comments
trafficstars

Describe what you are looking for

In memory_mapping_allocator_gt, memory is allocated in last_arena_, the scaling code is listed below:

https://github.com/unum-cloud/usearch/blob/e414f46b440fc2b989aca648fecfc4db35e22332/include/usearch/index_plugins.hpp#L960

And memory used is not stored in allocator, instead it's calculated:

https://github.com/unum-cloud/usearch/blob/e414f46b440fc2b989aca648fecfc4db35e22332/include/usearch/index_plugins.hpp#L981-L990

Since memory scaling is not simply * 2

Can you contribute to the implementation?

  • [x] I can contribute

Is your feature request specific to a certain interface?

It applies to everything

Contact Details

[email protected]

Is there an existing issue for this?

  • [x] I have searched the existing issues

Code of Conduct

  • [x] I agree to follow this project's Code of Conduct

mapleFU avatar May 16 '25 12:05 mapleFU

Strategy one: add a variable

diff --git a/include/usearch/index_plugins.hpp b/include/usearch/index_plugins.hpp
index e13703d..2318b0b 100644
--- a/include/usearch/index_plugins.hpp
+++ b/include/usearch/index_plugins.hpp
@@ -889,6 +889,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
     std::size_t last_usage_ = head_size();
     std::size_t last_capacity_ = min_capacity();
     std::size_t wasted_space_ = 0;
+    std::size_t total_allocated_ = 0; // Track total allocated memory
 
   public:
     using value_type = byte_t;
@@ -899,13 +900,15 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
     memory_mapping_allocator_gt() = default;
     memory_mapping_allocator_gt(memory_mapping_allocator_gt&& other) noexcept
         : last_arena_(exchange(other.last_arena_, nullptr)), last_usage_(exchange(other.last_usage_, 0)),
-          last_capacity_(exchange(other.last_capacity_, 0)), wasted_space_(exchange(other.wasted_space_, 0)) {}
+          last_capacity_(exchange(other.last_capacity_, 0)), wasted_space_(exchange(other.wasted_space_, 0)),
+          total_allocated_(exchange(other.total_allocated_, 0)) {}
 
     memory_mapping_allocator_gt& operator=(memory_mapping_allocator_gt&& other) noexcept {
         std::swap(last_arena_, other.last_arena_);
         std::swap(last_usage_, other.last_usage_);
         std::swap(last_capacity_, other.last_capacity_);
         std::swap(wasted_space_, other.wasted_space_);
+        std::swap(total_allocated_, other.total_allocated_);
         return *this;
     }
 
@@ -930,6 +933,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
         last_usage_ = head_size();
         last_capacity_ = min_capacity();
         wasted_space_ = 0;
+        total_allocated_ = 0;
     }
 
     /**
@@ -968,6 +972,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
             last_arena_ = new_arena;
             last_capacity_ = new_cap;
             last_usage_ = head_size();
+            total_allocated_ += new_cap;
         }
 
         wasted_space_ += extended_bytes - count_bytes;
@@ -979,15 +984,7 @@ template <std::size_t alignment_ak = 1> class memory_mapping_allocator_gt {
      *  @return The amount of space in bytes.
      */
     std::size_t total_allocated() const noexcept {
-        if (!last_arena_)
-            return 0;
-        std::size_t total_used = 0;
-        std::size_t last_capacity = last_capacity_;
-        do {
-            total_used += last_capacity;
-            last_capacity /= capacity_multiplier();
-        } while (last_capacity >= min_capacity());
-        return total_used;
+        return total_allocated_; // Return tracked total allocated memory
     }

mapleFU avatar May 16 '25 12:05 mapleFU

Another method:

    /**
     *  @brief Returns the amount of memory used by the allocator across all arenas.
     *  @return The amount of space in bytes.
     */
    std::size_t total_allocated() const noexcept {
        if (!last_arena_)
            return 0;
            
        std::size_t total_used = 0;
        byte_t* current_arena = last_arena_;
        
        while (current_arena) {
            std::size_t current_cap = 0;
            std::memcpy(&current_cap, current_arena + sizeof(byte_t*), sizeof(std::size_t));
            total_used += current_cap;
            
            byte_t* previous_arena = nullptr;
            std::memcpy(&previous_arena, current_arena, sizeof(byte_t*));
            current_arena = previous_arena;
        }
        
        return total_used;
    }

mapleFU avatar May 16 '25 12:05 mapleFU

@ashvardanian would you mind take a look?

mapleFU avatar May 16 '25 12:05 mapleFU