blob: c8d96adc220375cff330c95d3ee710040da40509 (
plain) (
blame)
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
|
// Copyright Vespa.ai. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root.
#pragma once
#include "tensor_store.h"
#include "tensor_buffer_operations.h"
#include "tensor_buffer_type_mapper.h"
#include "large_subspaces_buffer_type.h"
#include "small_subspaces_buffer_type.h"
#include <vespa/eval/eval/value_type.h>
#include <vespa/vespalib/datastore/array_store.h>
namespace search::tensor {
/**
* Class for storing tensor buffers in memory and making tensor views
* based on stored tensor buffer.
*/
class TensorBufferStore : public TensorStore
{
using RefType = vespalib::datastore::EntryRefT<19>;
using ArrayStoreType = vespalib::datastore::ArrayStore<char, RefType, TensorBufferTypeMapper>;
vespalib::eval::ValueType _tensor_type;
TensorBufferOperations _ops;
ArrayStoreType _array_store;
public:
static constexpr double array_store_grow_factor = 1.03;
static constexpr uint32_t array_store_max_type_id = 300;
TensorBufferStore(const vespalib::eval::ValueType& tensor_type, std::shared_ptr<vespalib::alloc::MemoryAllocator> allocator, uint32_t max_small_subspaces_type_id);
~TensorBufferStore();
void holdTensor(EntryRef ref) override;
EntryRef move_on_compact(EntryRef ref) override;
vespalib::MemoryUsage update_stat(const vespalib::datastore::CompactionStrategy& compaction_strategy) override;
std::unique_ptr<vespalib::datastore::ICompactionContext> start_compact(const vespalib::datastore::CompactionStrategy& compaction_strategy) override;
EntryRef store_tensor(const vespalib::eval::Value& tensor) override;
EntryRef store_encoded_tensor(vespalib::nbostream& encoded) override;
std::unique_ptr<vespalib::eval::Value> get_tensor(EntryRef ref) const override;
bool encode_stored_tensor(EntryRef ref, vespalib::nbostream& target) const override;
vespalib::eval::TypedCells get_empty_subspace() const noexcept {
return _ops.get_empty_subspace();
}
VectorBundle get_vectors(EntryRef ref) const {
if (!ref.valid()) {
return VectorBundle();
}
auto buf = _array_store.get(ref);
return _ops.get_vectors(buf);
}
SerializedTensorRef get_serialized_tensor_ref(EntryRef ref) const {
if (!ref.valid()) {
return SerializedTensorRef();
}
auto buf = _array_store.get(ref);
return _ops.get_serialized_tensor_ref(buf);
}
// Used by unit test
static constexpr uint32_t get_offset_bits() noexcept { return RefType::offset_bits; }
};
}
|
