Merge pull request #14180 from vespa-engine/arnej/count-used-memory

Arnej/count used memory

11 files changed, 61 insertions, 0 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 3f4641ed2ee..08333fa30f3 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
@@ -174,4 +174,20 @@ TEST("require that dense tensor cells iterator works for 2d tensor") {
     EXPECT_FALSE(itr.valid());
 }
 
+TEST("require that memory used count is reasonable") {
+    Tensor::UP full = build2DTensor();
+    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();
+
+    EXPECT_EQUAL(ref_sz, sizeof(DenseTensorView));
+    EXPECT_LESS(ref_sz, full_sz);
+    EXPECT_EQUAL(full_sz, view_sz);
+    EXPECT_LESS(full_sz, 10000u);
+    EXPECT_GREATER(full_sz, sizeof(DenseTensor<double>));
+}
+
 TEST_MAIN() { TEST_RUN_ALL(); }
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 86b6abedd39..f901b7775fd 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
@@ -99,6 +99,10 @@ TEST("Test essential object sizes") {
     EXPECT_EQUAL(16u, sizeof(SparseTensorAddressRef));
     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();
+    EXPECT_GREATER(used, sizeof(SparseTensor));
+    EXPECT_LESS(used, 10000u);
 }
 
 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 a72a24be211..9ed28d87fee 100644
--- a/eval/src/vespa/eval/eval/simple_tensor.cpp
+++ b/eval/src/vespa/eval/eval/simple_tensor.cpp
@@ -769,5 +769,14 @@ 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;
+}
+
 } // namespace vespalib::eval
 } // namespace vespalib
diff --git a/eval/src/vespa/eval/eval/simple_tensor.h b/eval/src/vespa/eval/eval/simple_tensor.h
index cbf1ac99e05..052d7cb70bd 100644
--- a/eval/src/vespa/eval/eval/simple_tensor.h
+++ b/eval/src/vespa/eval/eval/simple_tensor.h
@@ -93,6 +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;
 };
 
 } // 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 d0246fef635..4114661a074 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor.h
@@ -21,6 +21,11 @@ public:
     // for unit tests
     template <typename RCT>
     bool operator==(const DenseTensor<RCT> &rhs) const;
+
+    size_t count_memory_used() const override {
+        return sizeof(DenseTensor) + (sizeof(CT) * _cells.size());
+    }
+    
 private:
     eval::ValueType _type;
     std::vector<CT> _cells;
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 bf2464791f6..a07a3eede77 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
@@ -44,6 +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);
+    }
 
     template <typename T> static ConstArrayRef<T> typify_cells(const eval::Value &self) {
         return static_cast<const DenseTensorView &>(self).cellsRef().typify<T>();
diff --git a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
index d183c33f5cd..db35de6786d 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
@@ -243,6 +243,16 @@ SparseTensor::remove(const CellValues &cellAddresses) const
     return remover.build();
 }
 
+size_t
+SparseTensor::count_memory_used() const
+{
+    size_t result = sizeof(SparseTensor) + _cells.getMemoryConsumption();
+    for (const auto &cell : _cells) {
+        result += cell.first.size();
+    }
+    return result;
+}
+
 }
 
 VESPALIB_HASH_MAP_INSTANTIATE(vespalib::tensor::SparseTensorAddressRef, double);
diff --git a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
index e5ea639b460..6bd181e1895 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
@@ -53,6 +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;
 };
 
 }
diff --git a/eval/src/vespa/eval/tensor/tensor.h b/eval/src/vespa/eval/tensor/tensor.h
index d822c99a6d8..bef7309c609 100644
--- a/eval/src/vespa/eval/tensor/tensor.h
+++ b/eval/src/vespa/eval/tensor/tensor.h
@@ -59,6 +59,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;
 
     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 7c09bc4e4ab..fe73bf92063 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
@@ -54,6 +54,16 @@ WrappedSimpleTensor::accept(TensorVisitor &visitor) const
     }
 }
 
+size_t
+WrappedSimpleTensor::count_memory_used() const
+{
+    size_t result = sizeof(WrappedSimpleTensor);
+    if (_space) {
+        result += _space->count_memory_used();
+    }
+    return result;
+}
+
 Tensor::UP
 WrappedSimpleTensor::clone() const
 {
diff --git a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
index 12ee1237d67..6b549718a29 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
@@ -33,6 +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;
     Tensor::UP clone() const override;
     // functions below should not be used for this implementation
     Tensor::UP apply(const CellFunction &) const override;