compute memory used by any vespalib::tensor::Tensor

7 files changed, 32 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 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 93dd2dbedeb..52ffec14b18 100644
--- a/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
+++ b/eval/src/vespa/eval/tensor/dense/dense_tensor_view.h
@@ -43,6 +43,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..2963b2b4f53 100644
--- a/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/wrapped_simple_tensor.cpp
@@ -54,6 +54,17 @@ WrappedSimpleTensor::accept(TensorVisitor &visitor) const
     }
 }
 
+size_t
+WrappedSimpleTensor::count_memory_used() const
+{
+    size_t result = sizeof(WrappedSimpleTensor);
+    if (_space) {
+        result += sizeof(SimpleTensor);
+        result += _tensor.cells().size() * sizeof(SimpleTensor::Cell);
+    }
+    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;