better accounting in SimpleTensor

3 files changed, 19 insertions, 10 deletions

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/wrapped_simple_tensor.cpp b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
index fa01869c754..08fbb945830 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
@@ -59,9 +59,11 @@ WrappedSimpleTensor::get_memory_usage() const
 {
     size_t used = sizeof(WrappedSimpleTensor);
     if (_space) {
-        used += _space->count_memory_used();
+        auto plus = _space->get_memory_usage();
+        plus.incUsedBytes(used);
+        plus.incAllocatedBytes(used);
+        return plus;
     }
-    // XXX - does not account for extra allocated memory
     return MemoryUsage(used, used, 0, 0);
 }