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
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
|
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#include "tensor_buffer_operations.h"
#include "fast_value_view.h"
#include <vespa/eval/eval/value_codec.h>
#include <vespa/eval/eval/value_type.h>
#include <vespa/eval/streamed/streamed_value_view.h>
#include <vespa/vespalib/util/atomic.h>
#include <vespa/vespalib/util/shared_string_repo.h>
#include <algorithm>
using vespalib::ArrayRef;
using vespalib::ConstArrayRef;
using vespalib::MemoryUsage;
using vespalib::SharedStringRepo;
using vespalib::StringIdVector;
using vespalib::eval::FastAddrMap;
using vespalib::eval::FastValueIndex;
using vespalib::eval::StreamedValueView;
using vespalib::eval::TypedCells;
using vespalib::eval::Value;
using vespalib::eval::ValueType;
using vespalib::eval::self_memory_usage;
using vespalib::string_id;
namespace search::tensor {
namespace {
uint32_t
adjust_min_alignment(size_t min_alignment)
{
// Also apply alignment for num_subspaces and labels
return std::max(std::max(sizeof(uint32_t), sizeof(string_id)), min_alignment);
}
}
TensorBufferOperations::TensorBufferOperations(const vespalib::eval::ValueType& tensor_type)
: _subspace_type(tensor_type),
_num_mapped_dimensions(tensor_type.count_mapped_dimensions()),
_min_alignment(adjust_min_alignment(vespalib::eval::CellTypeUtils::alignment(_subspace_type.cell_type()))),
_addr(_num_mapped_dimensions),
_addr_refs(),
_empty(_subspace_type)
{
_addr_refs.reserve(_addr.size());
for (auto& label : _addr) {
_addr_refs.push_back(&label);
}
}
TensorBufferOperations::~TensorBufferOperations() = default;
uint32_t
TensorBufferOperations::get_num_subspaces_and_flag(ConstArrayRef<char> buf) const noexcept
{
assert(buf.size() >= get_num_subspaces_size());
const uint32_t& num_subspaces_and_flag_ref = *reinterpret_cast<const uint32_t*>(buf.data());
return vespalib::atomic::load_ref_relaxed(num_subspaces_and_flag_ref);
}
void
TensorBufferOperations::set_skip_reclaim_labels(ArrayRef<char> buf, uint32_t num_subspaces_and_flag) const noexcept
{
uint32_t& num_subspaces_and_flag_ref = *reinterpret_cast<uint32_t*>(buf.data());
vespalib::atomic::store_ref_relaxed(num_subspaces_and_flag_ref, (num_subspaces_and_flag | skip_reclaim_labels_mask));
}
void
TensorBufferOperations::store_tensor(ArrayRef<char> buf, const vespalib::eval::Value& tensor)
{
uint32_t num_subspaces = tensor.index().size();
assert(num_subspaces <= num_subspaces_mask);
auto labels_end_offset = get_labels_offset() + get_labels_mem_size(num_subspaces);
auto cells_size = num_subspaces * _subspace_type.size();
auto cells_mem_size = num_subspaces * _subspace_type.mem_size(); // Size measured in bytes
auto aligner = select_aligner(cells_mem_size);
auto cells_start_offset = aligner.align(labels_end_offset);
auto cells_end_offset = cells_start_offset + cells_mem_size;
auto store_end = aligner.align(cells_end_offset);
assert(store_end == get_buffer_size(num_subspaces));
assert(buf.size() >= store_end);
*reinterpret_cast<uint32_t*>(buf.data()) = num_subspaces;
auto labels = reinterpret_cast<string_id*>(buf.data() + get_labels_offset());
size_t subspace = 0;
size_t num_subspaces_visited = 0;
auto view = tensor.index().create_view({});
view->lookup({});
while (view->next_result(_addr_refs, subspace)) {
assert(subspace < num_subspaces);
auto subspace_labels = labels + subspace * _num_mapped_dimensions;
for (auto& label : _addr) {
SharedStringRepo::unsafe_copy(label); // tensor has an existing ref
*subspace_labels = label;
++subspace_labels;
}
++num_subspaces_visited;
}
assert(num_subspaces_visited == num_subspaces);
if (labels_end_offset != cells_start_offset) {
memset(buf.data() + labels_end_offset, 0, cells_start_offset - labels_end_offset);
}
auto cells = tensor.cells();
assert(cells_size == cells.size);
if (cells_mem_size > 0) {
memcpy(buf.data() + cells_start_offset, cells.data, cells_mem_size);
}
if (cells_end_offset != buf.size()) {
memset(buf.data() + cells_end_offset, 0, buf.size() - cells_end_offset);
}
}
std::unique_ptr<vespalib::eval::Value>
TensorBufferOperations::make_fast_view(ConstArrayRef<char> buf, const vespalib::eval::ValueType& tensor_type) const
{
auto num_subspaces = get_num_subspaces(buf);
assert(buf.size() >= get_buffer_size(num_subspaces));
ConstArrayRef<string_id> labels(reinterpret_cast<const string_id*>(buf.data() + get_labels_offset()), num_subspaces * _num_mapped_dimensions);
auto cells_size = num_subspaces * _subspace_type.size();
auto cells_mem_size = num_subspaces * _subspace_type.mem_size(); // Size measured in bytes
auto aligner = select_aligner(cells_mem_size);
auto cells_start_offset = get_cells_offset(num_subspaces, aligner);
TypedCells cells(buf.data() + cells_start_offset, _subspace_type.cell_type(), cells_size);
assert(cells_start_offset + cells_mem_size <= buf.size());
return std::make_unique<FastValueView>(tensor_type, labels, cells, _num_mapped_dimensions, num_subspaces);
}
void
TensorBufferOperations::copied_labels(ArrayRef<char> buf) const
{
auto num_subspaces_and_flag = get_num_subspaces_and_flag(buf);
if (!get_skip_reclaim_labels(num_subspaces_and_flag)) {
set_skip_reclaim_labels(buf, num_subspaces_and_flag);
}
}
void
TensorBufferOperations::reclaim_labels(ArrayRef<char> buf) const
{
auto num_subspaces_and_flag = get_num_subspaces_and_flag(buf);
if (get_skip_reclaim_labels(num_subspaces_and_flag)) {
return;
}
set_skip_reclaim_labels(buf, num_subspaces_and_flag);
auto num_subspaces = get_num_subspaces(num_subspaces_and_flag);
ArrayRef<string_id> labels(reinterpret_cast<string_id*>(buf.data() + get_labels_offset()), num_subspaces * _num_mapped_dimensions);
for (auto& label : labels) {
SharedStringRepo::unsafe_reclaim(label);
}
}
void
TensorBufferOperations::encode_stored_tensor(ConstArrayRef<char> buf, const vespalib::eval::ValueType& tensor_type, vespalib::nbostream& target) const
{
auto num_subspaces = get_num_subspaces(buf);
assert(buf.size() >= get_buffer_size(num_subspaces));
ConstArrayRef<string_id> labels(reinterpret_cast<const string_id*>(buf.data() + get_labels_offset()), num_subspaces * _num_mapped_dimensions);
auto cells_size = num_subspaces * _subspace_type.size();
auto cells_mem_size = num_subspaces * _subspace_type.mem_size(); // Size measured in bytes
auto aligner = select_aligner(cells_mem_size);
auto cells_start_offset = get_cells_offset(num_subspaces, aligner);
TypedCells cells(buf.data() + cells_start_offset, _subspace_type.cell_type(), cells_size);
assert(cells_start_offset + cells_mem_size <= buf.size());
StringIdVector labels_copy(labels.begin(), labels.end());
StreamedValueView streamed_value_view(tensor_type, _num_mapped_dimensions, cells, num_subspaces, labels_copy);
vespalib::eval::encode_value(streamed_value_view, target);
}
}
|
