1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
|
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#include "direct_tensor_store.h"
#include "tensor_deserialize.h"
#include <vespa/eval/eval/fast_value.h>
#include <vespa/eval/eval/value.h>
#include <vespa/eval/eval/value_codec.h>
#include <vespa/vespalib/datastore/compacting_buffers.h>
#include <vespa/vespalib/datastore/compaction_context.h>
#include <vespa/vespalib/datastore/compaction_strategy.h>
#include <vespa/vespalib/datastore/datastore.hpp>
#include <vespa/vespalib/datastore/buffer_type.hpp>
#include <vespa/vespalib/util/size_literals.h>
using vespalib::datastore::CompactionContext;
using vespalib::datastore::CompactionSpec;
using vespalib::datastore::CompactionStrategy;
using vespalib::datastore::EntryRef;
using vespalib::datastore::ICompactionContext;
using vespalib::eval::FastValueBuilderFactory;
using vespalib::eval::Value;
namespace search::tensor {
constexpr size_t MIN_BUFFER_ARRAYS = 8_Ki;
DirectTensorStore::TensorBufferType::TensorBufferType()
: ParentType(1, MIN_BUFFER_ARRAYS, TensorStoreType::RefType::offsetSize())
{
}
void
DirectTensorStore::TensorBufferType::cleanHold(void* buffer, size_t offset, ElemCount num_elems, CleanContext clean_ctx)
{
TensorSP* elem = static_cast<TensorSP*>(buffer) + offset;
const auto& empty = empty_entry();
for (size_t i = 0; i < num_elems; ++i) {
clean_ctx.extraBytesCleaned((*elem)->get_memory_usage().allocatedBytes());
*elem = empty;
++elem;
}
}
EntryRef
DirectTensorStore::add_entry(TensorSP tensor)
{
auto ref = _tensor_store.addEntry(tensor);
auto& state = _tensor_store.getBufferState(RefType(ref).bufferId());
state.stats().inc_extra_used_bytes(tensor->get_memory_usage().allocatedBytes());
return ref;
}
DirectTensorStore::DirectTensorStore()
: TensorStore(_tensor_store),
_tensor_store(std::make_unique<TensorBufferType>())
{
_tensor_store.enableFreeLists();
}
DirectTensorStore::~DirectTensorStore() = default;
void
DirectTensorStore::holdTensor(EntryRef ref)
{
if (!ref.valid()) {
return;
}
const auto& tensor = _tensor_store.getEntry(ref);
assert(tensor);
_tensor_store.holdElem(ref, 1, tensor->get_memory_usage().allocatedBytes());
}
EntryRef
DirectTensorStore::move(EntryRef ref)
{
if (!ref.valid()) {
return EntryRef();
}
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->get_memory_usage().allocatedBytes());
return new_ref;
}
vespalib::MemoryUsage
DirectTensorStore::update_stat(const CompactionStrategy& compaction_strategy)
{
auto memory_usage = _store.getMemoryUsage();
_compaction_spec = CompactionSpec(compaction_strategy.should_compact_memory(memory_usage), false);
return memory_usage;
}
std::unique_ptr<ICompactionContext>
DirectTensorStore::start_compact(const CompactionStrategy& compaction_strategy)
{
auto compacting_buffers = _store.start_compact_worst_buffers(_compaction_spec, compaction_strategy);
return std::make_unique<CompactionContext>(*this, std::move(compacting_buffers));
}
EntryRef
DirectTensorStore::store_tensor(std::unique_ptr<Value> tensor)
{
assert(tensor);
return add_entry(std::move(tensor));
}
EntryRef
DirectTensorStore::store_tensor(const Value& tensor)
{
return add_entry(FastValueBuilderFactory::get().copy(tensor));
}
EntryRef
DirectTensorStore::store_encoded_tensor(vespalib::nbostream& encoded)
{
return add_entry(deserialize_tensor(encoded));
}
std::unique_ptr<Value>
DirectTensorStore::get_tensor(EntryRef ref) const
{
if (!ref.valid()) {
return {};
}
return FastValueBuilderFactory::get().copy(*_tensor_store.getEntry(ref));
}
bool
DirectTensorStore::encode_stored_tensor(EntryRef ref, vespalib::nbostream& target) const
{
if (!ref.valid()) {
return false;
}
vespalib::eval::encode_value(*_tensor_store.getEntry(ref), target);
return true;
}
}
|
