Merge pull request #14334 from vespa-engine/arnej/get-memory-usage-from-tensors

Arnej/get memory usage from tensors 

17 files changed, 97 insertions, 47 deletions

diff --git a/eval/src/tests/tensor/direct_dense_tensor_builder/direct_dense_tensor_builder_test.cpp b/eval/src/tests/tensor/direct_dense_tensor_builder/direct_dense_tensor_builder_test.cpp
index 08333fa30f3..75e9c7868ce 100644
--- a/eval/src/tests/tensor/direct_dense_tensor_builder/direct_dense_tensor_builder_test.cpp
+++ b/eval/src/tests/tensor/direct_dense_tensor_builder/direct_dense_tensor_builder_test.cpp
@@ -179,9 +179,9 @@ TEST("require that memory used count is reasonable") {
     const DenseTensorView &full_view = dynamic_cast<const DenseTensorView &>(*full);
     DenseTensorView ref_view(full_view.fast_type(), full_view.cellsRef());
 
-    size_t full_sz = full->count_memory_used();
-    size_t view_sz = full_view.count_memory_used();
-    size_t ref_sz = ref_view.count_memory_used();
+    size_t full_sz = full->get_memory_usage().usedBytes();
+    size_t view_sz = full_view.get_memory_usage().usedBytes();
+    size_t ref_sz = ref_view.get_memory_usage().usedBytes();
 
     EXPECT_EQUAL(ref_sz, sizeof(DenseTensorView));
     EXPECT_LESS(ref_sz, full_sz);
diff --git a/eval/src/tests/tensor/direct_sparse_tensor_builder/direct_sparse_tensor_builder_test.cpp b/eval/src/tests/tensor/direct_sparse_tensor_builder/direct_sparse_tensor_builder_test.cpp
index f901b7775fd..651451d81f1 100644
--- a/eval/src/tests/tensor/direct_sparse_tensor_builder/direct_sparse_tensor_builder_test.cpp
+++ b/eval/src/tests/tensor/direct_sparse_tensor_builder/direct_sparse_tensor_builder_test.cpp
@@ -100,9 +100,14 @@ TEST("Test essential object sizes") {
     EXPECT_EQUAL(24u, sizeof(std::pair<SparseTensorAddressRef, double>));
     EXPECT_EQUAL(32u, sizeof(vespalib::hash_node<std::pair<SparseTensorAddressRef, double>>));
     Tensor::UP tensor = buildTensor();
-    size_t used = tensor->count_memory_used();
+    size_t used = tensor->get_memory_usage().usedBytes();
     EXPECT_GREATER(used, sizeof(SparseTensor));
     EXPECT_LESS(used, 10000u);
+    size_t allocated = tensor->get_memory_usage().allocatedBytes();
+    EXPECT_GREATER(allocated, used);
+    EXPECT_LESS(allocated, 50000u);
+    fprintf(stderr, "tensor using %zu bytes of %zu allocated\n",
+            used, allocated);
 }
 
 TEST_MAIN() { TEST_RUN_ALL(); }
diff --git a/eval/src/vespa/eval/eval/simple_tensor.cpp b/eval/src/vespa/eval/eval/simple_tensor.cpp
index 9ed28d87fee..6ab7df34b2f 100644
--- a/eval/src/vespa/eval/eval/simple_tensor.cpp
+++ b/eval/src/vespa/eval/eval/simple_tensor.cpp
@@ -769,13 +769,20 @@ SimpleTensor::decode(nbostream &input)
     return builder.build();
 }
 
-size_t
-SimpleTensor::count_memory_used() const {
-    size_t result = sizeof(SimpleTensor);
-    size_t addr_size = sizeof(Label) * _type.dimensions().size();
-    size_t cell_size = sizeof(Cell) + addr_size;
-    result += _cells.size() * cell_size;
-    return result;
+vespalib::MemoryUsage
+SimpleTensor::get_memory_usage() const {
+    size_t addr_use = sizeof(Label) * _type.dimensions().size();
+    size_t cell_use = sizeof(Cell) + addr_use;
+    size_t cells_use = _cells.size() * cell_use;
+
+    size_t addr_alloc = sizeof(Label) * _type.dimensions().capacity();
+    size_t cell_alloc = sizeof(Cell) + addr_alloc;
+    size_t cells_alloc = _cells.capacity() * cell_alloc;
+
+    size_t mine_sz = sizeof(SimpleTensor);
+    size_t used = mine_sz + cells_use;
+    size_t allocated = mine_sz + cells_alloc;
+    return MemoryUsage(allocated, used, 0, 0);
 }
 
 } // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/eval/simple_tensor.h b/eval/src/vespa/eval/eval/simple_tensor.h
index 052d7cb70bd..d717838176b 100644
--- a/eval/src/vespa/eval/eval/simple_tensor.h
+++ b/eval/src/vespa/eval/eval/simple_tensor.h
@@ -93,7 +93,7 @@ public:
     static std::unique_ptr<SimpleTensor> concat(const SimpleTensor &a, const SimpleTensor &b, const vespalib::string &dimension);
     static void encode(const SimpleTensor &tensor, nbostream &output);
     static std::unique_ptr<SimpleTensor> decode(nbostream &input);
-    size_t count_memory_used() const;
+    MemoryUsage get_memory_usage() const;
 };
 
 } // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor.h b/eval/src/vespa/eval/tensor/dense/dense_tensor.h
index 4114661a074..f4b2b0a584f 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor.h
@@ -22,9 +22,12 @@ public:
     template <typename RCT>
     bool operator==(const DenseTensor<RCT> &rhs) const;
 
-    size_t count_memory_used() const override {
-        return sizeof(DenseTensor) + (sizeof(CT) * _cells.size());
-    }
+    MemoryUsage get_memory_usage() const override {
+        size_t alloc = sizeof(DenseTensor) + (sizeof(CT) * _cells.capacity());
+        size_t used = sizeof(DenseTensor) + (sizeof(CT) * _cells.size());
+        // missing: extra memory used by _type
+        return MemoryUsage(alloc, used, 0, 0);
+    } 
     
 private:
     eval::ValueType _type;
diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h b/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
index a07a3eede77..cf3e2864a30 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
@@ -44,8 +44,9 @@ public:
     Tensor::UP clone() const override;
     eval::TensorSpec toSpec() const override;
     void accept(TensorVisitor &visitor) const override;
-    size_t count_memory_used() const override {
-        return sizeof(DenseTensorView);
+    MemoryUsage get_memory_usage() const override {
+        size_t sz = sizeof(DenseTensorView);
+        return MemoryUsage(sz, sz, 0, 0);
     }
 
     template <typename T> static ConstArrayRef<T> typify_cells(const eval::Value &self) {
diff --git a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
index b2bd9330160..87ab80c2a8e 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
@@ -248,13 +248,13 @@ SparseTensor::remove(const CellValues &cellAddresses) const
     return remover.build();
 }
 
-size_t
-SparseTensor::count_memory_used() const
+MemoryUsage
+SparseTensor::get_memory_usage() const
 {
-    size_t result = sizeof(SparseTensor) + _cells.getMemoryConsumption();
-    for (const auto &cell : _cells) {
-        result += cell.first.size();
-    }
+    MemoryUsage result = _stash.get_memory_usage();
+    size_t plus = sizeof(SparseTensor) + _cells.getMemoryConsumption();
+    result.incUsedBytes(plus);
+    result.incAllocatedBytes(plus); // should probably be even more
     return result;
 }
 
diff --git a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
index 6bd181e1895..34a6622bdd6 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
@@ -53,7 +53,7 @@ public:
     Tensor::UP clone() const override;
     eval::TensorSpec toSpec() const override;
     void accept(TensorVisitor &visitor) const override;
-    size_t count_memory_used() const override;
+    MemoryUsage get_memory_usage() const override;
 };
 
 }
diff --git a/eval/src/vespa/eval/tensor/tensor.h b/eval/src/vespa/eval/tensor/tensor.h
index bef7309c609..0deb36b24a0 100644
--- a/eval/src/vespa/eval/tensor/tensor.h
+++ b/eval/src/vespa/eval/tensor/tensor.h
@@ -5,6 +5,7 @@
 #include "cell_function.h"
 #include "tensor_address.h"
 #include <vespa/vespalib/stllike/string.h>
+#include <vespa/vespalib/util/memoryusage.h>
 #include <vespa/eval/eval/tensor.h>
 #include <vespa/eval/eval/tensor_spec.h>
 #include <vespa/eval/eval/value_type.h>
@@ -59,7 +60,7 @@ public:
     virtual Tensor::UP clone() const = 0; // want to remove, but needed by document
     virtual eval::TensorSpec toSpec() const = 0;
     virtual void accept(TensorVisitor &visitor) const = 0;
-    virtual size_t count_memory_used() const = 0;
+    virtual MemoryUsage get_memory_usage() const = 0;
 
     using TypeList = std::initializer_list<std::reference_wrapper<const eval::ValueType>>;
     static bool supported(TypeList types);
diff --git a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
index fe73bf92063..08fbb945830 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
@@ -54,14 +54,17 @@ WrappedSimpleTensor::accept(TensorVisitor &visitor) const
     }
 }
 
-size_t
-WrappedSimpleTensor::count_memory_used() const
+MemoryUsage
+WrappedSimpleTensor::get_memory_usage() const
 {
-    size_t result = sizeof(WrappedSimpleTensor);
+    size_t used = sizeof(WrappedSimpleTensor);
     if (_space) {
-        result += _space->count_memory_used();
+        auto plus = _space->get_memory_usage();
+        plus.incUsedBytes(used);
+        plus.incAllocatedBytes(used);
+        return plus;
     }
-    return result;
+    return MemoryUsage(used, used, 0, 0);
 }
 
 Tensor::UP
diff --git a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
index 6b549718a29..cf9103a4895 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
@@ -33,7 +33,7 @@ public:
     eval::TensorSpec toSpec() const override;
     double as_double() const override;
     void accept(TensorVisitor &visitor) const override;
-    size_t count_memory_used() const override;
+    MemoryUsage get_memory_usage() const override;
     Tensor::UP clone() const override;
     // functions below should not be used for this implementation
     Tensor::UP apply(const CellFunction &) const override;
diff --git a/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp b/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
index a77cf8070e5..5a815a96dfb 100644
--- a/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
+++ b/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
@@ -61,11 +61,12 @@ TEST_F(DirectTensorStoreTest, heap_allocated_memory_is_tracked)
     store.store_tensor(make_tensor(5));
     auto mem_1 = store.getMemoryUsage();
     auto ref = store.store_tensor(make_tensor(10));
-    size_t tensor_memory_used = store.get_tensor(ref)->count_memory_used();
+    auto tensor_mem_usage = store.get_tensor(ref)->get_memory_usage();
     auto mem_2 = store.getMemoryUsage();
-    EXPECT_GT(tensor_memory_used, 100);
-    EXPECT_GE(mem_2.allocatedBytes(), mem_1.allocatedBytes() + tensor_memory_used);
-    EXPECT_GT(mem_2.usedBytes(), mem_1.usedBytes() + tensor_memory_used);
+    EXPECT_GT(tensor_mem_usage.usedBytes(), 100);
+    EXPECT_GT(tensor_mem_usage.allocatedBytes(), 500);
+    EXPECT_GE(mem_2.allocatedBytes(), mem_1.allocatedBytes() + tensor_mem_usage.allocatedBytes());
+    EXPECT_GT(mem_2.usedBytes(), mem_1.usedBytes() + tensor_mem_usage.allocatedBytes());
 }
 
 TEST_F(DirectTensorStoreTest, invalid_ref_returns_nullptr)
@@ -77,18 +78,18 @@ TEST_F(DirectTensorStoreTest, invalid_ref_returns_nullptr)
 TEST_F(DirectTensorStoreTest, hold_adds_entry_to_hold_list)
 {
     auto ref = store.store_tensor(make_tensor(5));
-    size_t tensor_memory_used = store.get_tensor(ref)->count_memory_used();
+    auto tensor_mem_usage = store.get_tensor(ref)->get_memory_usage();
     auto mem_1 = store.getMemoryUsage();
     store.holdTensor(ref);
     auto mem_2 = store.getMemoryUsage();
-    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_memory_used);
+    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_mem_usage.allocatedBytes());
 }
 
 TEST_F(DirectTensorStoreTest, move_allocates_new_entry_and_puts_old_entry_on_hold)
 {
     auto t = make_tensor(5);
     auto* exp = t.get();
-    size_t tensor_memory_used = exp->count_memory_used();
+    auto tensor_mem_usage = exp->get_memory_usage();
     auto ref_1 = store.store_tensor(std::move(t));
     auto mem_1 = store.getMemoryUsage();
 
@@ -97,7 +98,7 @@ TEST_F(DirectTensorStoreTest, move_allocates_new_entry_and_puts_old_entry_on_hol
     EXPECT_NE(ref_1, ref_2);
     expect_tensor(exp, ref_1);
     expect_tensor(exp, ref_2);
-    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_memory_used);
+    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_mem_usage.allocatedBytes());
 }
 
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp b/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp
index 9b8659bf576..90be394227c 100644
--- a/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp
@@ -20,7 +20,7 @@ DirectTensorStore::TensorBufferType::cleanHold(void* buffer, size_t offset, size
 {
     TensorSP* elem = static_cast<TensorSP*>(buffer) + offset;
     for (size_t i = 0; i < num_elems; ++i) {
-        clean_ctx.extraBytesCleaned((*elem)->count_memory_used());
+        clean_ctx.extraBytesCleaned((*elem)->get_memory_usage().allocatedBytes());
         *elem = _emptyEntry;
         ++elem;
     }
@@ -31,7 +31,7 @@ DirectTensorStore::add_entry(TensorSP tensor)
 {
     auto ref = _tensor_store.addEntry(tensor);
     auto& state = _tensor_store.getBufferState(RefType(ref).bufferId());
-    state.incExtraUsedBytes(tensor->count_memory_used());
+    state.incExtraUsedBytes(tensor->get_memory_usage().allocatedBytes());
     return ref;
 }
 
@@ -70,7 +70,7 @@ DirectTensorStore::holdTensor(EntryRef ref)
     }
     const auto& tensor = _tensor_store.getEntry(ref);
     assert(tensor);
-    _tensor_store.holdElem(ref, 1, tensor->count_memory_used());
+    _tensor_store.holdElem(ref, 1, tensor->get_memory_usage().allocatedBytes());
 }
 
 EntryRef
@@ -82,7 +82,7 @@ DirectTensorStore::move(EntryRef ref)
     const auto& old_tensor = _tensor_store.getEntry(ref);
     assert(old_tensor);
     auto new_ref = add_entry(old_tensor);
-    _tensor_store.holdElem(ref, 1, old_tensor->count_memory_used());
+    _tensor_store.holdElem(ref, 1, old_tensor->get_memory_usage().allocatedBytes());
     return new_ref;
 }
 
diff --git a/vespalib/src/tests/stash/stash.cpp b/vespalib/src/tests/stash/stash.cpp
index 077e721abb7..ebf38a1343a 100644
--- a/vespalib/src/tests/stash/stash.cpp
+++ b/vespalib/src/tests/stash/stash.cpp
@@ -99,7 +99,7 @@ TEST("require that base types have expected size") {
     EXPECT_EQUAL(8u, char_ptr_size());
     EXPECT_EQUAL(16u, chunk_header_size());
     EXPECT_EQUAL(16u, dtor_hook_size());
-    EXPECT_EQUAL(16u, free_hook_size());
+    EXPECT_EQUAL(24u, free_hook_size());
     EXPECT_EQUAL(24u, array_dtor_hook_size());
 }
 
diff --git a/vespalib/src/vespa/vespalib/stllike/string.h b/vespalib/src/vespa/vespalib/stllike/string.h
index aa1e3677532..121a243d403 100644
--- a/vespalib/src/vespa/vespalib/stllike/string.h
+++ b/vespalib/src/vespa/vespalib/stllike/string.h
@@ -528,6 +528,13 @@ public:
     void reserve(size_type newCapacity) {
         reserveBytes(newCapacity + 1);
     }
+
+    size_t count_allocated_memory() const {
+        return sizeof(small_string) + isAllocated() ? _bufferSize : 0;
+    }
+    size_t count_used_memory() const {
+        return sizeof(small_string) - StackSize + size();
+    }
 private:
     void assign_slower(const void * s, size_type sz) __attribute((noinline));
     void init_slower(const void *s) noexcept __attribute((noinline));
diff --git a/vespalib/src/vespa/vespalib/util/stash.cpp b/vespalib/src/vespa/vespalib/util/stash.cpp
index 836654c2fb2..31580e871db 100644
--- a/vespalib/src/vespa/vespalib/util/stash.cpp
+++ b/vespalib/src/vespa/vespalib/util/stash.cpp
@@ -53,8 +53,9 @@ Stash::do_alloc(size_t size)
         _chunks = new (chunk_mem) stash::Chunk(_chunks);
         return _chunks->alloc(size, _chunk_size);
     } else {
-        char *mem = static_cast<char*>(malloc(sizeof(stash::DeleteMemory) + size));
-        _cleanup = new (mem) stash::DeleteMemory(_cleanup);
+        size_t allocate = sizeof(stash::DeleteMemory) + size;
+        char *mem = static_cast<char*>(malloc(allocate));
+        _cleanup = new (mem) stash::DeleteMemory(allocate, _cleanup);
         return (mem + sizeof(stash::DeleteMemory));
     }
 }
@@ -121,4 +122,22 @@ Stash::count_used() const
     return used;
 }
 
+MemoryUsage
+Stash::get_memory_usage() const
+{
+    size_t allocated = 0;
+    size_t used = 0;
+    for (stash::Chunk *chunk = _chunks; chunk != nullptr; chunk = chunk->next) {
+        allocated += _chunk_size;
+        used += chunk->used;
+    }
+    for (auto cleanup = _cleanup; cleanup; cleanup = cleanup->next) {
+        if (auto memory = dynamic_cast<stash::DeleteMemory *>(cleanup)) {
+            allocated += memory->allocated;
+            used += memory->allocated;
+        }
+    }
+    return MemoryUsage(allocated, used, 0, 0);
+};
+
 } // namespace vespalib
diff --git a/vespalib/src/vespa/vespalib/util/stash.h b/vespalib/src/vespa/vespalib/util/stash.h
index c5e8631ca9e..48ac16080c0 100644
--- a/vespalib/src/vespa/vespalib/util/stash.h
+++ b/vespalib/src/vespa/vespalib/util/stash.h
@@ -4,6 +4,7 @@
 
 #include "traits.h"
 #include "arrayref.h"
+#include "memoryusage.h"
 #include <cstdlib>
 
 namespace vespalib {
@@ -19,8 +20,9 @@ protected:
 
 // used as header for memory allocated outside the stash
 struct DeleteMemory : public Cleanup {
-    explicit DeleteMemory(Cleanup *next_in) noexcept : Cleanup(next_in) {}
+    explicit DeleteMemory(size_t sz, Cleanup *next_in) noexcept : Cleanup(next_in), allocated(sz) {}
     void cleanup() override { free((void*)this); }
+    size_t allocated;
 };
 
 // used as prefix for objects to be destructed
@@ -140,6 +142,7 @@ public:
 
     size_t count_used() const;
     size_t get_chunk_size() const { return _chunk_size; }
+    MemoryUsage get_memory_usage() const;
 
     char *alloc(size_t size) {
         char *ret = __builtin_expect(is_small(size) && _chunks != nullptr, true)