use MemoryUsage from Tensor

2 files changed, 13 insertions, 12 deletions

diff --git a/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp b/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
index a77cf8070e5..d64ce935fb1 100644
--- a/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
+++ b/searchlib/src/tests/tensor/direct_tensor_store/direct_tensor_store_test.cpp
@@ -61,11 +61,12 @@ TEST_F(DirectTensorStoreTest, heap_allocated_memory_is_tracked)
     store.store_tensor(make_tensor(5));
     auto mem_1 = store.getMemoryUsage();
     auto ref = store.store_tensor(make_tensor(10));
-    size_t tensor_memory_used = store.get_tensor(ref)->count_memory_used();
+    auto tensor_memory_used = store.get_tensor(ref)->get_memory_usage();
     auto mem_2 = store.getMemoryUsage();
-    EXPECT_GT(tensor_memory_used, 100);
-    EXPECT_GE(mem_2.allocatedBytes(), mem_1.allocatedBytes() + tensor_memory_used);
-    EXPECT_GT(mem_2.usedBytes(), mem_1.usedBytes() + tensor_memory_used);
+    EXPECT_GT(tensor_memory_used.usedBytes(), 100);
+    EXPECT_GT(tensor_memory_used.allocatedBytes(), 500);
+    EXPECT_GE(mem_2.allocatedBytes(), mem_1.allocatedBytes() + tensor_memory_used.allocatedBytes());
+    EXPECT_GT(mem_2.usedBytes(), mem_1.usedBytes() + tensor_memory_used.allocatedBytes());
 }
 
 TEST_F(DirectTensorStoreTest, invalid_ref_returns_nullptr)
@@ -77,18 +78,18 @@ TEST_F(DirectTensorStoreTest, invalid_ref_returns_nullptr)
 TEST_F(DirectTensorStoreTest, hold_adds_entry_to_hold_list)
 {
     auto ref = store.store_tensor(make_tensor(5));
-    size_t tensor_memory_used = store.get_tensor(ref)->count_memory_used();
+    auto tensor_memory_used = store.get_tensor(ref)->get_memory_usage();
     auto mem_1 = store.getMemoryUsage();
     store.holdTensor(ref);
     auto mem_2 = store.getMemoryUsage();
-    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_memory_used);
+    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_memory_used.allocatedBytes());
 }
 
 TEST_F(DirectTensorStoreTest, move_allocates_new_entry_and_puts_old_entry_on_hold)
 {
     auto t = make_tensor(5);
     auto* exp = t.get();
-    size_t tensor_memory_used = exp->count_memory_used();
+    auto tensor_memory_used = exp->get_memory_usage();
     auto ref_1 = store.store_tensor(std::move(t));
     auto mem_1 = store.getMemoryUsage();
 
@@ -97,7 +98,7 @@ TEST_F(DirectTensorStoreTest, move_allocates_new_entry_and_puts_old_entry_on_hol
     EXPECT_NE(ref_1, ref_2);
     expect_tensor(exp, ref_1);
     expect_tensor(exp, ref_2);
-    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_memory_used);
+    EXPECT_GT(mem_2.allocatedBytesOnHold(), mem_1.allocatedBytesOnHold() + tensor_memory_used.allocatedBytes());
 }
 
 GTEST_MAIN_RUN_ALL_TESTS()
diff --git a/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp b/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp
index 9b8659bf576..90be394227c 100644
--- a/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp
+++ b/searchlib/src/vespa/searchlib/tensor/direct_tensor_store.cpp
@@ -20,7 +20,7 @@ DirectTensorStore::TensorBufferType::cleanHold(void* buffer, size_t offset, size
 {
     TensorSP* elem = static_cast<TensorSP*>(buffer) + offset;
     for (size_t i = 0; i < num_elems; ++i) {
-        clean_ctx.extraBytesCleaned((*elem)->count_memory_used());
+        clean_ctx.extraBytesCleaned((*elem)->get_memory_usage().allocatedBytes());
         *elem = _emptyEntry;
         ++elem;
     }
@@ -31,7 +31,7 @@ DirectTensorStore::add_entry(TensorSP tensor)
 {
     auto ref = _tensor_store.addEntry(tensor);
     auto& state = _tensor_store.getBufferState(RefType(ref).bufferId());
-    state.incExtraUsedBytes(tensor->count_memory_used());
+    state.incExtraUsedBytes(tensor->get_memory_usage().allocatedBytes());
     return ref;
 }
 
@@ -70,7 +70,7 @@ DirectTensorStore::holdTensor(EntryRef ref)
     }
     const auto& tensor = _tensor_store.getEntry(ref);
     assert(tensor);
-    _tensor_store.holdElem(ref, 1, tensor->count_memory_used());
+    _tensor_store.holdElem(ref, 1, tensor->get_memory_usage().allocatedBytes());
 }
 
 EntryRef
@@ -82,7 +82,7 @@ DirectTensorStore::move(EntryRef ref)
     const auto& old_tensor = _tensor_store.getEntry(ref);
     assert(old_tensor);
     auto new_ref = add_entry(old_tensor);
-    _tensor_store.holdElem(ref, 1, old_tensor->count_memory_used());
+    _tensor_store.holdElem(ref, 1, old_tensor->get_memory_usage().allocatedBytes());
     return new_ref;
 }