somewhat better accounting of SimpleTensor memory

3 files changed, 11 insertions, 2 deletions

diff --git a/eval/src/vespa/eval/eval/simple_tensor.cpp b/eval/src/vespa/eval/eval/simple_tensor.cpp
index a72a24be211..a60852ea1e6 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(Address) * _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/wrapped_simple_tensor.cpp b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
index 2963b2b4f53..fe73bf92063 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
@@ -59,8 +59,7 @@ WrappedSimpleTensor::count_memory_used() const
 {
     size_t result = sizeof(WrappedSimpleTensor);
     if (_space) {
-        result += sizeof(SimpleTensor);
-        result += _tensor.cells().size() * sizeof(SimpleTensor::Cell);
+        result += _space->count_memory_used();
     }
     return result;
 }