add get_memory_usage for tensor::Tensor

7 files changed, 36 insertions, 0 deletions

diff --git a/eval/src/vespa/eval/tensor/dense/dense_tensor.h b/eval/src/vespa/eval/tensor/dense/dense_tensor.h
index 4114661a074..54025e47d2e 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor.h
@@ -25,6 +25,12 @@ public:
     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..4c855cff25b 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
@@ -47,6 +47,10 @@ public:
     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) {
         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 b2bd9330160..62780a9cfcc 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.cpp
@@ -258,6 +258,16 @@ SparseTensor::count_memory_used() const
     return result;
 }
 
+MemoryUsage
+SparseTensor::get_memory_usage() const
+{
+    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;
+}
+
 }
 
 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 6bd181e1895..e14ce0daea0 100644
--- a/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
+++ b/eval/src/vespa/eval/tensor/sparse/sparse_tensor.h
@@ -54,6 +54,7 @@ public:
     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..b75e8b14197 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>
@@ -60,6 +61,7 @@ public:
     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..914a3422161 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
@@ -64,6 +64,18 @@ WrappedSimpleTensor::count_memory_used() const
     return result;
 }
 
+
+MemoryUsage
+WrappedSimpleTensor::get_memory_usage() const
+{
+    size_t used = sizeof(WrappedSimpleTensor);
+    if (_space) {
+        used += _space->count_memory_used();
+    }
+    // XXX - does not account for extra allocated memory
+    return MemoryUsage(used, used, 0, 0);
+}
+
 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 6b549718a29..1544b1615db 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.h
@@ -34,6 +34,7 @@ public:
     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;