14 files changed, 111 insertions, 3 deletions

diff --git a/eval/src/vespa/eval/eval/fast_sparse_map.h b/eval/src/vespa/eval/eval/fast_sparse_map.h
index 666a9db4140..2f3409a45f6 100644
--- a/eval/src/vespa/eval/eval/fast_sparse_map.h
+++ b/eval/src/vespa/eval/eval/fast_sparse_map.h
@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include "memory_usage_stuff.h"
 #include <vespa/vespalib/util/arrayref.h>
 #include <vespa/vespalib/stllike/string.h>
 #include <vespa/vespalib/stllike/hash_map.h>
@@ -93,6 +94,20 @@ public:
         _labels.reserve(_num_dims * expected_subspaces);
     }
     ~FastSparseMap();
+
+    MemoryUsage estimate_extra_memory_usage() const {
+        MemoryUsage extra_usage;
+        size_t map_self_size = sizeof(_map);
+        size_t map_used = _map.getMemoryUsed();
+        size_t map_allocated = _map.getMemoryConsumption();
+        // avoid double-counting the map itself
+        map_used = std::min(map_used, map_used - map_self_size);
+        map_allocated = std::min(map_allocated, map_allocated - map_self_size);
+        extra_usage.merge(MemoryUsage(map_used, map_allocated, 0, 0));
+        extra_usage.merge(vector_extra_memory_usage(_labels));
+        return extra_usage;
+    }
+
     size_t size() const { return _map.size(); }
     size_t num_dims() const { return _num_dims; }
     static constexpr size_t npos() { return -1; }
diff --git a/eval/src/vespa/eval/eval/fast_value.hpp b/eval/src/vespa/eval/eval/fast_value.hpp
index e0b401b800c..8914d63167c 100644
--- a/eval/src/vespa/eval/eval/fast_value.hpp
+++ b/eval/src/vespa/eval/eval/fast_value.hpp
@@ -191,6 +191,12 @@ struct FastValue final : Value, ValueBuilder<T> {
         self.release();
         return std::unique_ptr<Value>(this);
     }
+    MemoryUsage get_memory_usage() const override {
+        MemoryUsage usage = self_memory_usage<FastValue<T>>();
+        usage.merge(vector_extra_memory_usage(my_cells));
+        usage.merge(my_index.map.estimate_extra_memory_usage());
+        return usage;
+    }
 };
 template <typename T> FastValue<T>::~FastValue() = default;
 
diff --git a/eval/src/vespa/eval/eval/memory_usage_stuff.h b/eval/src/vespa/eval/eval/memory_usage_stuff.h
new file mode 100644
index 00000000000..cfc680b342f
--- /dev/null
+++ b/eval/src/vespa/eval/eval/memory_usage_stuff.h
@@ -0,0 +1,18 @@
+// Copyright 2017 Yahoo Holdings. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
+
+#pragma once
+
+#include <vespa/vespalib/util/memoryusage.h>
+#include <vector>
+
+namespace vespalib::eval {
+
+template <typename T>
+MemoryUsage self_memory_usage() { return MemoryUsage(sizeof(T), sizeof(T), 0, 0); }
+
+template <typename T>
+MemoryUsage vector_extra_memory_usage(const std::vector<T> &vec) {
+    return MemoryUsage(sizeof(T) * vec.size(), sizeof(T) * vec.capacity(), 0, 0);
+}
+
+}
diff --git a/eval/src/vespa/eval/eval/simple_tensor.h b/eval/src/vespa/eval/eval/simple_tensor.h
index 50229070cad..7ab59199af2 100644
--- a/eval/src/vespa/eval/eval/simple_tensor.h
+++ b/eval/src/vespa/eval/eval/simple_tensor.h
@@ -96,7 +96,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);
-    MemoryUsage get_memory_usage() const;
+    MemoryUsage get_memory_usage() const override;
 };
 
 } // namespace vespalib::eval
diff --git a/eval/src/vespa/eval/eval/simple_value.cpp b/eval/src/vespa/eval/eval/simple_value.cpp
index ec843ce320e..734a92be788 100644
--- a/eval/src/vespa/eval/eval/simple_value.cpp
+++ b/eval/src/vespa/eval/eval/simple_value.cpp
@@ -182,6 +182,17 @@ SimpleValue::add_mapping(ConstArrayRef<vespalib::stringref> addr)
     assert(was_inserted);
 }
 
+MemoryUsage
+SimpleValue::estimate_extra_memory_usage() const
+{
+    using Node = std::map<Labels,size_t>::value_type;
+    size_t key_extra_size = sizeof(vespalib::string) * _num_mapped_dims;
+    size_t node_extra_size = 2 * sizeof(Node *); // left/right child ptr
+    size_t entry_size = sizeof(Node) + key_extra_size + node_extra_size;
+    size_t size = entry_size * _index.size();
+    return MemoryUsage(size, size, 0, 0);
+}
+
 std::unique_ptr<Value::Index::View>
 SimpleValue::create_view(const std::vector<size_t> &dims) const
 {
diff --git a/eval/src/vespa/eval/eval/simple_value.h b/eval/src/vespa/eval/eval/simple_value.h
index c838161bbde..a8e84a95419 100644
--- a/eval/src/vespa/eval/eval/simple_value.h
+++ b/eval/src/vespa/eval/eval/simple_value.h
@@ -31,6 +31,7 @@ protected:
     size_t num_mapped_dims() const { return _num_mapped_dims; }
     size_t subspace_size() const { return _subspace_size; }
     void add_mapping(ConstArrayRef<vespalib::stringref> addr);
+    MemoryUsage estimate_extra_memory_usage() const;
 public:
     SimpleValue(const ValueType &type, size_t num_mapped_dims_in, size_t subspace_size_in);
     ~SimpleValue() override;
@@ -63,6 +64,12 @@ public:
         self.release();
         return std::unique_ptr<Value>(this);
     }
+    MemoryUsage get_memory_usage() const override {
+        MemoryUsage usage = self_memory_usage<SimpleValueT<T>>();
+        usage.merge(vector_extra_memory_usage(_cells));
+        usage.merge(estimate_extra_memory_usage());
+        return usage;
+    }
 };
 
 /**
diff --git a/eval/src/vespa/eval/eval/value.h b/eval/src/vespa/eval/eval/value.h
index 76a9002210c..1e656f58e52 100644
--- a/eval/src/vespa/eval/eval/value.h
+++ b/eval/src/vespa/eval/eval/value.h
@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include "memory_usage_stuff.h"
 #include "value_type.h"
 #include "typed_cells.h"
 #include <vespa/vespalib/stllike/string.h>
@@ -61,6 +62,8 @@ struct Value {
     virtual const Index &index() const = 0;
 // --- end of new interface
 
+    virtual MemoryUsage get_memory_usage() const = 0;
+
 // --- old interface that may be (partially) removed in the future
     virtual bool is_double() const { return type().is_double(); }
     virtual bool is_tensor() const { return type().is_tensor(); }
@@ -92,6 +95,7 @@ public:
     DoubleValue(double value) : _value(value) {}
     TypedCells cells() const override { return TypedCells(ConstArrayRef<double>(&_value, 1)); }
     const Index &index() const override { return TrivialIndex::get(); }
+    MemoryUsage get_memory_usage() const override { return self_memory_usage<DoubleValue>(); }
     bool is_double() const override { return true; }
     double as_double() const override { return _value; }
     const ValueType &type() const override { return _type; }
@@ -113,6 +117,7 @@ public:
     const ValueType &type() const final override { return _type; }
     TypedCells cells() const final override { return _cells; }
     const Index &index() const final override { return TrivialIndex::get(); }
+    MemoryUsage get_memory_usage() const override { return self_memory_usage<DenseValueView>(); }
 };
 
 /**
@@ -130,6 +135,7 @@ public:
     const ValueType &type() const final override { return _type; }
     TypedCells cells() const final override { return _cells; }
     const Index &index() const final override { return _index; }
+    MemoryUsage get_memory_usage() const override { return self_memory_usage<ValueView>(); }
 };
 
 /**
diff --git a/eval/src/vespa/eval/tensor/mixed/packed_labels.cpp b/eval/src/vespa/eval/tensor/mixed/packed_labels.cpp
index a05497f6b68..a048986c4ac 100644
--- a/eval/src/vespa/eval/tensor/mixed/packed_labels.cpp
+++ b/eval/src/vespa/eval/tensor/mixed/packed_labels.cpp
@@ -37,6 +37,17 @@ PackedLabels::get_label(uint32_t index) const
     return vespalib::stringref(p, sz);
 }
 
+MemoryUsage
+PackedLabels::estimate_extra_memory_usage() const
+{
+    MemoryUsage extra_usage;
+    size_t offsets_size = _offsets.size() * sizeof(uint32_t);
+    size_t labels_size = _label_store.size() * sizeof(char);
+    extra_usage.merge(MemoryUsage(offsets_size, offsets_size, 0, 0));
+    extra_usage.merge(MemoryUsage(labels_size, labels_size, 0, 0));
+    return extra_usage;
+}
+
 void
 PackedLabels::validate_labels(uint32_t num_labels)
 {
diff --git a/eval/src/vespa/eval/tensor/mixed/packed_labels.h b/eval/src/vespa/eval/tensor/mixed/packed_labels.h
index dac338448fb..ebe8c7bea7b 100644
--- a/eval/src/vespa/eval/tensor/mixed/packed_labels.h
+++ b/eval/src/vespa/eval/tensor/mixed/packed_labels.h
@@ -2,6 +2,7 @@
 
 #pragma once
 
+#include <vespa/eval/eval/memory_usage_stuff.h>
 #include <vespa/vespalib/stllike/string.h>
 #include <vespa/vespalib/util/arrayref.h>
 
@@ -32,6 +33,7 @@ public:
 
     vespalib::stringref get_label(uint32_t index) const;
 
+    MemoryUsage estimate_extra_memory_usage() const;
 private:
     const ConstArrayRef<uint32_t> _offsets;
     const ConstArrayRef<char> _label_store;
diff --git a/eval/src/vespa/eval/tensor/mixed/packed_mappings.cpp b/eval/src/vespa/eval/tensor/mixed/packed_mappings.cpp
index 226ac01ef71..c42a77cf384 100644
--- a/eval/src/vespa/eval/tensor/mixed/packed_mappings.cpp
+++ b/eval/src/vespa/eval/tensor/mixed/packed_mappings.cpp
@@ -99,6 +99,16 @@ PackedMappings::fill_address_by_sortid(uint32_t internal_index, Address &address
     return _int_store[offset];
 }
 
+MemoryUsage
+PackedMappings::estimate_extra_memory_usage() const
+{
+    MemoryUsage extra_usage;
+    size_t store_size = _int_store.size() * sizeof(uint32_t);
+    extra_usage.merge(MemoryUsage(store_size, store_size, 0, 0));
+    extra_usage.merge(_label_store.estimate_extra_memory_usage());
+    return extra_usage;
+}
+
 void
 PackedMappings::validate() const
 {
diff --git a/eval/src/vespa/eval/tensor/mixed/packed_mappings.h b/eval/src/vespa/eval/tensor/mixed/packed_mappings.h
index ce9ff066271..c36bf69bb05 100644
--- a/eval/src/vespa/eval/tensor/mixed/packed_mappings.h
+++ b/eval/src/vespa/eval/tensor/mixed/packed_mappings.h
@@ -47,6 +47,8 @@ public:
 
     // mapping from label enum to stringref (and vice versa)
     const PackedLabels & label_store() const { return _label_store; }
+
+    MemoryUsage estimate_extra_memory_usage() const;
 private:
     PackedMappings(uint32_t num_dims, uint32_t num_mappings,
                    ConstArrayRef<uint32_t> int_store,
diff --git a/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.cpp b/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.cpp
index 660cad82e4f..f1e3b30a026 100644
--- a/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.cpp
+++ b/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.cpp
@@ -1,9 +1,20 @@
 // Copyright Verizon Media. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
 
 #include "packed_mixed_tensor.h"
+#include <vespa/vespalib/util/typify.h>
 
 namespace vespalib::eval::packed_mixed_tensor {
 
+namespace {
+
+struct MySize {
+    template <typename CT> static size_t invoke(TypedCells cells) {
+        return (sizeof(CT) * cells.size);
+    }
+};
+
+} // namespace <unnamed>
+
 /*********************************************************************************/
 
 class PackedMixedTensorIndexView : public Value::Index::View
@@ -194,6 +205,15 @@ PackedMixedTensorAllMappings::next_result(ConstArrayRef<vespalib::stringref*> ad
 
 PackedMixedTensor::~PackedMixedTensor() = default;
 
+MemoryUsage
+PackedMixedTensor::get_memory_usage() const {
+    MemoryUsage usage = self_memory_usage<PackedMixedTensor>();
+    size_t cells_size = typify_invoke<1,TypifyCellType,MySize>(type().cell_type(), cells());
+    usage.merge(MemoryUsage(cells_size, cells_size, 0, 0));
+    usage.merge(_mappings.estimate_extra_memory_usage());
+    return usage;
+}
+
 std::unique_ptr<Value::Index::View>
 PackedMixedTensor::create_view(const std::vector<size_t> &dims) const
 {
diff --git a/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.h b/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.h
index 604b1c94aeb..215af27d79f 100644
--- a/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.h
+++ b/eval/src/vespa/eval/tensor/mixed/packed_mixed_tensor.h
@@ -43,6 +43,8 @@ public:
     const Value::Index &index() const override { return *this; }
     TypedCells cells() const override { return _cells; }
 
+    MemoryUsage get_memory_usage() const override;
+
     // Value::Index API:
     size_t size() const override { return _mappings.size(); }
     std::unique_ptr<View> create_view(const std::vector<size_t> &dims) const override;
diff --git a/eval/src/vespa/eval/tensor/tensor.h b/eval/src/vespa/eval/tensor/tensor.h
index e56c6299d94..6a8fd9adc45 100644
--- a/eval/src/vespa/eval/tensor/tensor.h
+++ b/eval/src/vespa/eval/tensor/tensor.h
@@ -5,7 +5,6 @@
 #include "cell_function.h"
 #include "tensor_address.h"
 #include <vespa/vespalib/stllike/string.h>
-#include <vespa/vespalib/util/memoryusage.h>
 #include <vespa/eval/eval/operation.h>
 #include <vespa/eval/eval/tensor.h>
 #include <vespa/eval/eval/tensor_spec.h>
@@ -61,7 +60,6 @@ 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 MemoryUsage get_memory_usage() const = 0;
 
     using TypeList = std::initializer_list<std::reference_wrapper<const eval::ValueType>>;
     static bool supported(TypeList types);