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
|
// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#pragma once
#include "empty_subspace.h"
#include "serialized_tensor_ref.h"
#include "subspace_type.h"
#include "vector_bundle.h"
#include <vespa/vespalib/datastore/aligner.h>
#include <vespa/vespalib/util/arrayref.h>
#include <vespa/vespalib/util/string_id.h>
#include <cstddef>
#include <memory>
namespace vespalib { class nbostream; }
namespace vespalib::eval {
struct Value;
class ValueType;
}
namespace search::tensor {
/*
* Class used to store a tensor in a buffer and make tensor views based on
* buffer content.
*
* Layout of buffer is:
*
* num_subspaces_and_flag
* 31 least significant bits is num_subspaces - number of subspaces
* 1 bit to signal that reclaim_labels should be a noop (i.e. buffer has been copied
* as part of compaction or fallback copy (due to datastore buffer fallback resize)).
* labels[num_subspaces * _num_mapped_dimensions] - array of labels for sparse dimensions
* padding - to align start of cells
* cells[num_subspaces * _dense_subspaces_size] - array of tensor cell values
* padding - to align start of next buffer
*
* Alignment is dynamic, based on cell type, memory used by tensor cell values and
* alignment required for reading num_subspaces and labels array.
*/
class TensorBufferOperations
{
SubspaceType _subspace_type;
uint32_t _num_mapped_dimensions;
uint32_t _min_alignment;
std::vector<vespalib::string_id> _addr;
std::vector<vespalib::string_id*> _addr_refs;
EmptySubspace _empty;
using Aligner = vespalib::datastore::Aligner<vespalib::datastore::dynamic_alignment>;
static constexpr size_t CELLS_ALIGNMENT = 16;
static constexpr size_t CELLS_ALIGNMENT_MEM_SIZE_MIN = 32;
static constexpr uint32_t num_subspaces_mask = ((1u << 31) - 1);
static constexpr uint32_t skip_reclaim_labels_mask = (1u << 31);
static constexpr size_t get_num_subspaces_size() noexcept { return sizeof(uint32_t); }
static constexpr size_t get_labels_offset() noexcept { return get_num_subspaces_size(); }
size_t get_cells_mem_size(uint32_t num_subspaces) const noexcept {
return _subspace_type.mem_size() * num_subspaces;
}
auto select_aligner(size_t cells_mem_size) const noexcept {
return Aligner((cells_mem_size < CELLS_ALIGNMENT_MEM_SIZE_MIN) ? _min_alignment : CELLS_ALIGNMENT);
}
size_t get_labels_mem_size(uint32_t num_subspaces) const noexcept {
return sizeof(vespalib::string_id) * _num_mapped_dimensions * num_subspaces;
}
size_t get_cells_offset(uint32_t num_subspaces, auto aligner) const noexcept {
return aligner.align(get_labels_offset() + get_labels_mem_size(num_subspaces));
}
uint32_t get_num_subspaces_and_flag(vespalib::ConstArrayRef<char> buf) const noexcept;
void set_skip_reclaim_labels(vespalib::ArrayRef<char> buf, uint32_t num_subspaces_and_flag) const noexcept;
static uint32_t get_num_subspaces(uint32_t num_subspaces_and_flag) noexcept {
return num_subspaces_and_flag & num_subspaces_mask;
}
static bool get_skip_reclaim_labels(uint32_t num_subspaces_and_flag) noexcept {
return (num_subspaces_and_flag & skip_reclaim_labels_mask) != 0;
}
uint32_t get_num_subspaces(vespalib::ConstArrayRef<char> buf) const noexcept {
return get_num_subspaces(get_num_subspaces_and_flag(buf));
}
public:
// Size (in bytes) used to serialize tensor with num_subspaces.
size_t get_buffer_size(uint32_t num_subspaces) const noexcept {
auto cells_mem_size = get_cells_mem_size(num_subspaces);
auto aligner = select_aligner(cells_mem_size);
return get_cells_offset(num_subspaces, aligner) + aligner.align(cells_mem_size);
}
TensorBufferOperations(const vespalib::eval::ValueType& tensor_type);
~TensorBufferOperations();
TensorBufferOperations(const TensorBufferOperations&) = delete;
TensorBufferOperations(TensorBufferOperations&&) = delete;
TensorBufferOperations& operator=(const TensorBufferOperations&) = delete;
TensorBufferOperations& operator=(TensorBufferOperations&&) = delete;
void store_tensor(vespalib::ArrayRef<char> buf, const vespalib::eval::Value& tensor);
std::unique_ptr<vespalib::eval::Value> make_fast_view(vespalib::ConstArrayRef<char> buf, const vespalib::eval::ValueType& tensor_type) const;
// Mark that reclaim_labels should be skipped for old buffer after copying tensor buffer
void copied_labels(vespalib::ArrayRef<char> buf) const;
// Decrease reference counts for labels and set skip flag unless skip flag is set.
void reclaim_labels(vespalib::ArrayRef<char> buf) const;
// Serialize stored tensor to target (used when saving attribute)
void encode_stored_tensor(vespalib::ConstArrayRef<char> buf, const vespalib::eval::ValueType& type, vespalib::nbostream& target) const;
vespalib::eval::TypedCells get_empty_subspace() const noexcept {
return _empty.cells();
}
VectorBundle get_vectors(vespalib::ConstArrayRef<char> buf) const {
auto num_subspaces = get_num_subspaces(buf);
auto cells_mem_size = get_cells_mem_size(num_subspaces);
auto aligner = select_aligner(cells_mem_size);
return VectorBundle(buf.data() + get_cells_offset(num_subspaces, aligner), num_subspaces, _subspace_type);
}
SerializedTensorRef get_serialized_tensor_ref(vespalib::ConstArrayRef<char> buf) const {
auto num_subspaces = get_num_subspaces(buf);
auto cells_mem_size = get_cells_mem_size(num_subspaces);
auto aligner = select_aligner(cells_mem_size);
vespalib::ConstArrayRef<vespalib::string_id> labels(reinterpret_cast<const vespalib::string_id*>(buf.data() + get_labels_offset()), num_subspaces * _num_mapped_dimensions);
return SerializedTensorRef(VectorBundle(buf.data() + get_cells_offset(num_subspaces, aligner), num_subspaces, _subspace_type), _num_mapped_dimensions, labels);
}
bool is_dense() const noexcept { return _num_mapped_dimensions == 0; }
};
}
|
