From a55b433e04f933841c05c3b4223b8eecb7ac5ac7 Mon Sep 17 00:00:00 2001 From: Tor Egge Date: Fri, 30 Sep 2022 23:58:26 +0200 Subject: Add TensorBufferStore. --- searchlib/CMakeLists.txt | 1 + .../tensor/tensor_buffer_store/CMakeLists.txt | 9 ++ .../tensor_buffer_store_test.cpp | 164 +++++++++++++++++++++ .../src/vespa/searchlib/tensor/CMakeLists.txt | 4 + .../tensor/large_subspaces_buffer_type.cpp | 86 +++++++++++ .../searchlib/tensor/large_subspaces_buffer_type.h | 40 +++++ .../tensor/small_subspaces_buffer_type.cpp | 67 +++++++++ .../searchlib/tensor/small_subspaces_buffer_type.h | 40 +++++ .../vespa/searchlib/tensor/tensor_buffer_store.cpp | 97 ++++++++++++ .../vespa/searchlib/tensor/tensor_buffer_store.h | 38 +++++ .../searchlib/tensor/tensor_buffer_type_mapper.cpp | 47 ++++++ .../searchlib/tensor/tensor_buffer_type_mapper.h | 35 +++++ .../src/vespa/vespalib/datastore/array_store.h | 3 +- 13 files changed, 630 insertions(+), 1 deletion(-) create mode 100644 searchlib/src/tests/tensor/tensor_buffer_store/CMakeLists.txt create mode 100644 searchlib/src/tests/tensor/tensor_buffer_store/tensor_buffer_store_test.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.h create mode 100644 searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.h create mode 100644 searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h create mode 100644 searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.cpp create mode 100644 searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.h diff --git a/searchlib/CMakeLists.txt b/searchlib/CMakeLists.txt index 76dadc5605e..a7d831aa623 100644 --- a/searchlib/CMakeLists.txt +++ b/searchlib/CMakeLists.txt @@ -224,6 +224,7 @@ vespa_define_module( src/tests/tensor/hnsw_index src/tests/tensor/hnsw_saver src/tests/tensor/tensor_buffer_operations + src/tests/tensor/tensor_buffer_store src/tests/transactionlog src/tests/transactionlogstress src/tests/true diff --git a/searchlib/src/tests/tensor/tensor_buffer_store/CMakeLists.txt b/searchlib/src/tests/tensor/tensor_buffer_store/CMakeLists.txt new file mode 100644 index 00000000000..749d38a1383 --- /dev/null +++ b/searchlib/src/tests/tensor/tensor_buffer_store/CMakeLists.txt @@ -0,0 +1,9 @@ +# Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. +vespa_add_executable(searchlib_tensor_buffer_store_test_app TEST + SOURCES + tensor_buffer_store_test.cpp + DEPENDS + searchlib + GTest::GTest +) +vespa_add_test(NAME searchlib_tensor_buffer_store_test_app COMMAND searchlib_tensor_buffer_store_test_app) diff --git a/searchlib/src/tests/tensor/tensor_buffer_store/tensor_buffer_store_test.cpp b/searchlib/src/tests/tensor/tensor_buffer_store/tensor_buffer_store_test.cpp new file mode 100644 index 00000000000..101b84e01aa --- /dev/null +++ b/searchlib/src/tests/tensor/tensor_buffer_store/tensor_buffer_store_test.cpp @@ -0,0 +1,164 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include +#include +#include +#include +#include + +using search::tensor::TensorBufferStore; +using vespalib::datastore::EntryRef; +using vespalib::eval::SimpleValue; +using vespalib::eval::TensorSpec; +using vespalib::eval::Value; +using vespalib::eval::ValueType; + +const vespalib::string tensor_type_spec("tensor(x{})"); + +class TensorBufferStoreTest : public testing::Test +{ +protected: + ValueType _tensor_type; + TensorBufferStore _store; + TensorBufferStoreTest(); + ~TensorBufferStoreTest() override; + EntryRef store_tensor(const Value& tensor); + EntryRef store_tensor(const TensorSpec& spec); + std::unique_ptr load_tensor(EntryRef ref); + TensorSpec load_tensor_spec(EntryRef ref); + vespalib::nbostream encode_stored_tensor(EntryRef ref); + void assert_store_load(const TensorSpec& tensor_spec); + void assert_store_load_many(const TensorSpec& tensor_spec); + void assert_store_move_load(const TensorSpec& tensor_spec); + void assert_store_encode_store_encoded_load(const TensorSpec& tensor_spec); +}; + +TensorBufferStoreTest::TensorBufferStoreTest() + : testing::Test(), + _tensor_type(ValueType::from_spec(tensor_type_spec)), + _store(_tensor_type, {}, 4) +{ +} + +TensorBufferStoreTest::~TensorBufferStoreTest() = default; + +EntryRef +TensorBufferStoreTest::store_tensor(const Value& tensor) +{ + EXPECT_EQ(_tensor_type, tensor.type()); + return _store.store_tensor(tensor); +} + +EntryRef +TensorBufferStoreTest::store_tensor(const TensorSpec& spec) +{ + auto tensor = SimpleValue::from_spec(spec); + return store_tensor(*tensor); +} + +std::unique_ptr +TensorBufferStoreTest::load_tensor(EntryRef ref) +{ + return _store.get_tensor(ref); +} + +vespalib::nbostream +TensorBufferStoreTest::encode_stored_tensor(EntryRef ref) +{ + vespalib::nbostream out; + _store.encode_stored_tensor(ref, out); + return out; +} + +TensorSpec +TensorBufferStoreTest::load_tensor_spec(EntryRef ref) +{ + auto loaded = load_tensor(ref); + return TensorSpec::from_value(*loaded); +} + +void +TensorBufferStoreTest::assert_store_load(const TensorSpec& tensor_spec) +{ + auto ref = store_tensor(tensor_spec); + auto loaded_spec = load_tensor_spec(ref); + _store.holdTensor(ref); + EXPECT_EQ(tensor_spec, loaded_spec); +} + +void +TensorBufferStoreTest::assert_store_load_many(const TensorSpec& tensor_spec) +{ + constexpr uint32_t cnt = 2000; + std::vector refs; + for (uint32_t i = 0; i < cnt; ++i) { + refs.emplace_back(store_tensor(tensor_spec)); + } + for (auto ref : refs) { + auto loaded_spec = load_tensor_spec(ref); + _store.holdTensor(ref); + EXPECT_EQ(tensor_spec, loaded_spec); + } +} + +void +TensorBufferStoreTest::assert_store_move_load(const TensorSpec& tensor_spec) +{ + auto ref = store_tensor(tensor_spec); + auto ref2 = _store.move(ref); + EXPECT_NE(ref, ref2); + auto loaded_spec = load_tensor_spec(ref2); + _store.holdTensor(ref2); + EXPECT_EQ(tensor_spec, loaded_spec); +} + +void +TensorBufferStoreTest::assert_store_encode_store_encoded_load(const TensorSpec& tensor_spec) +{ + auto ref = store_tensor(tensor_spec); + auto encoded = encode_stored_tensor(ref); + _store.holdTensor(ref); + auto ref2 = _store.store_encoded_tensor(encoded); + EXPECT_NE(ref, ref2); + auto loaded_spec = load_tensor_spec(ref2); + _store.holdTensor(ref2); + EXPECT_EQ(tensor_spec, loaded_spec); +} + +std::vector tensor_specs = { + TensorSpec(tensor_type_spec), + TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5), + TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5).add({{"x", "b"}}, 5.5), + TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5).add({{"x", "b"}}, 5.5).add({{"x", "c"}}, 6.5), + TensorSpec(tensor_type_spec).add({{"x", "a"}}, 4.5).add({{"x", "b"}}, 5.5).add({{"x", "c"}}, 6.5).add({{"x", "d"}}, 7.5) +}; + +TEST_F(TensorBufferStoreTest, tensor_can_be_stored_and_loaded) +{ + for (auto& tensor_spec : tensor_specs) { + assert_store_load(tensor_spec); + } +} + +TEST_F(TensorBufferStoreTest, tensor_can_be_stored_and_loaded_many_times) +{ + for (auto& tensor_spec : tensor_specs) { + assert_store_load_many(tensor_spec); + } +} + +TEST_F(TensorBufferStoreTest, stored_tensor_can_be_copied) +{ + for (auto& tensor_spec : tensor_specs) { + assert_store_move_load(tensor_spec); + } +} + +TEST_F(TensorBufferStoreTest, stored_tensor_can_be_encoded_and_stored_as_encoded_and_loaded) +{ + for (auto& tensor_spec : tensor_specs) { + assert_store_encode_store_encoded_load(tensor_spec); + } +} + +GTEST_MAIN_RUN_ALL_TESTS() diff --git a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt index 7815ef7e770..46bfc0909aa 100644 --- a/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt +++ b/searchlib/src/vespa/searchlib/tensor/CMakeLists.txt @@ -24,13 +24,17 @@ vespa_add_library(searchlib_tensor OBJECT imported_tensor_attribute_vector_read_guard.cpp inner_product_distance.cpp inv_log_level_generator.cpp + large_subspaces_buffer_type.cpp nearest_neighbor_index.cpp nearest_neighbor_index_saver.cpp serialized_fast_value_attribute.cpp + small_subspaces_buffer_type.cpp streamed_value_saver.cpp streamed_value_store.cpp tensor_attribute.cpp tensor_buffer_operations.cpp + tensor_buffer_store.cpp + tensor_buffer_type_mapper.cpp tensor_deserialize.cpp tensor_store.cpp reusable_set_visited_tracker.cpp diff --git a/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.cpp b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.cpp new file mode 100644 index 00000000000..cdd4d35c1df --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.cpp @@ -0,0 +1,86 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "large_subspaces_buffer_type.h" +#include "tensor_buffer_operations.h" +#include "tensor_buffer_type_mapper.h" +#include +#include +#include + +using vespalib::alloc::MemoryAllocator; + +namespace search::tensor { + +LargeSubspacesBufferType::LargeSubspacesBufferType(const AllocSpec& spec, std::shared_ptr memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept + : ParentType(1u, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor), + _memory_allocator(std::move(memory_allocator)), + _ops(type_mapper.get_tensor_buffer_operations()) +{ +} + +LargeSubspacesBufferType::~LargeSubspacesBufferType() = default; + +void +LargeSubspacesBufferType::cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) +{ + auto elem = static_cast(buffer) + offset; + for (size_t i = 0; i < numElems; ++i) { + if (!elem->empty()) { + cleanCtx.extraBytesCleaned(elem->size()); + _ops.reclaim_labels({elem->data(), elem->size()}); + ArrayType().swap(*elem); + } + ++elem; + } +} + +void +LargeSubspacesBufferType::destroyElements(void *buffer, ElemCount numElems) +{ + auto elem = static_cast(buffer); + for (size_t i = 0; i < numElems; ++i) { + if (!elem->empty()) { + _ops.reclaim_labels({elem->data(), elem->size()}); + ArrayType().swap(*elem); + } + ++elem; + } +} + +void +LargeSubspacesBufferType::fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems) +{ + auto old_elems = static_cast(oldBuffer); + auto new_elems = static_cast(newBuffer); + for (size_t i = 0; i < numElems; ++i) { + auto& old_elem = old_elems[i]; + new (new_elems + i) ArrayType(old_elem); + if (!old_elem.empty()) { + _ops.copied_labels({old_elem.data(), old_elem.size()}); + } + } +} + +void +LargeSubspacesBufferType::initializeReservedElements(void *buffer, ElemCount reservedElements) +{ + auto new_elems = static_cast(buffer); + const auto& empty = empty_entry(); + for (size_t i = 0; i < reservedElements; ++i) { + new (new_elems + i) ArrayType(empty); + } +} + +const vespalib::alloc::MemoryAllocator* +LargeSubspacesBufferType::get_memory_allocator() const +{ + return _memory_allocator.get(); +} + +} + +namespace vespalib::datastore { + +template class BufferType>; + +} diff --git a/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.h b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.h new file mode 100644 index 00000000000..cfab8ef20af --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/large_subspaces_buffer_type.h @@ -0,0 +1,40 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +#include +#include +#include +#include + +namespace vespalib::alloc { class MemoryAllocator; } + +namespace search::tensor { + +class TensorBufferOperations; +class TensorBufferTypeMapper; + +/* + * Class representing buffer type for tensors with a large number of + * subspaces in array store. Tensor buffers are externally allocated + * (cf. vespalib::Array). + */ +class LargeSubspacesBufferType : public vespalib::datastore::BufferType> +{ + using AllocSpec = vespalib::datastore::ArrayStoreConfig::AllocSpec; + using ArrayType = vespalib::Array; + using ParentType = vespalib::datastore::BufferType; + using CleanContext = typename ParentType::CleanContext; + std::shared_ptr _memory_allocator; + TensorBufferOperations& _ops; +public: + LargeSubspacesBufferType(const AllocSpec& spec, std::shared_ptr memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept; + ~LargeSubspacesBufferType() override; + void cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override; + void destroyElements(void *buffer, ElemCount numElems) override; + void fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems) override; + void initializeReservedElements(void *buffer, ElemCount reservedElements) override; + const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override; +}; + +} diff --git a/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.cpp b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.cpp new file mode 100644 index 00000000000..adbd3dee2b7 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.cpp @@ -0,0 +1,67 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "small_subspaces_buffer_type.h" +#include "tensor_buffer_operations.h" +#include "tensor_buffer_type_mapper.h" +#include + +using vespalib::alloc::MemoryAllocator; + +namespace search::tensor { + +SmallSubspacesBufferType::SmallSubspacesBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept + : ParentType(array_size, spec.minArraysInBuffer, spec.maxArraysInBuffer, spec.numArraysForNewBuffer, spec.allocGrowFactor), + _memory_allocator(std::move(memory_allocator)), + _ops(type_mapper.get_tensor_buffer_operations()) +{ +} + +SmallSubspacesBufferType::~SmallSubspacesBufferType() = default; + +void +SmallSubspacesBufferType::cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext) +{ + char* elem = static_cast(buffer) + offset; + while (numElems >= getArraySize()) { + _ops.reclaim_labels(vespalib::ArrayRef(elem, getArraySize())); + elem += getArraySize(); + numElems -= getArraySize(); + } +} + +void +SmallSubspacesBufferType::destroyElements(void *buffer, ElemCount numElems) +{ + char* elem = static_cast(buffer); + while (numElems >= getArraySize()) { + _ops.reclaim_labels(vespalib::ArrayRef(elem, getArraySize())); + elem += getArraySize(); + numElems -= getArraySize(); + } +} + +void +SmallSubspacesBufferType::fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems) +{ + memcpy(newBuffer, oldBuffer, numElems); + const char *elem = static_cast(oldBuffer); + while (numElems >= getArraySize()) { + _ops.copied_labels(vespalib::ConstArrayRef(elem, getArraySize())); + elem += getArraySize(); + numElems -= getArraySize(); + } +} + +void +SmallSubspacesBufferType::initializeReservedElements(void *buffer, ElemCount reservedElements) +{ + memset(buffer, 0, reservedElements); +} + +const vespalib::alloc::MemoryAllocator* +SmallSubspacesBufferType::get_memory_allocator() const +{ + return _memory_allocator.get(); +} + +} diff --git a/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.h b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.h new file mode 100644 index 00000000000..a778183c5a2 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/small_subspaces_buffer_type.h @@ -0,0 +1,40 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +#include +#include +#include + +namespace vespalib::alloc { class MemoryAllocator; } + +namespace search::tensor { + +class TensorBufferOperations; +class TensorBufferTypeMapper; + +/* + * Class representing buffer type for tensors with a small number of + * subspaces in array store. Tensor buffers are internal in data store buffer. + */ +class SmallSubspacesBufferType : public vespalib::datastore::BufferType +{ + using AllocSpec = vespalib::datastore::ArrayStoreConfig::AllocSpec; + using ParentType = vespalib::datastore::BufferType; + std::shared_ptr _memory_allocator; + TensorBufferOperations& _ops; +public: + SmallSubspacesBufferType(const SmallSubspacesBufferType&) = delete; + SmallSubspacesBufferType& operator=(const SmallSubspacesBufferType&) = delete; + SmallSubspacesBufferType(SmallSubspacesBufferType&&) noexcept = default; + SmallSubspacesBufferType& operator=(SmallSubspacesBufferType&&) noexcept = default; + SmallSubspacesBufferType(uint32_t array_size, const AllocSpec& spec, std::shared_ptr memory_allocator, TensorBufferTypeMapper& type_mapper) noexcept; + ~SmallSubspacesBufferType() override; + void cleanHold(void* buffer, size_t offset, ElemCount numElems, CleanContext cleanCtx) override; + void destroyElements(void *buffer, ElemCount numElems) override; + void fallbackCopy(void *newBuffer, const void *oldBuffer, ElemCount numElems) override; + void initializeReservedElements(void *buffer, ElemCount reservedElements) override; + const vespalib::alloc::MemoryAllocator* get_memory_allocator() const override; +}; + +} diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp new file mode 100644 index 00000000000..6c1b3bbd1ee --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.cpp @@ -0,0 +1,97 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "tensor_buffer_store.h" +#include +#include +#include +#include +#include +#include + +using vespalib::alloc::MemoryAllocator; +using vespalib::datastore::EntryRef; +using vespalib::eval::StreamedValueBuilderFactory; +using vespalib::eval::Value; +using vespalib::eval::ValueType; + +namespace search::tensor { + +namespace { + +constexpr float ALLOC_GROW_FACTOR = 0.2; + +} + +TensorBufferStore::TensorBufferStore(const ValueType& tensor_type, std::shared_ptr allocator, uint32_t max_small_subspaces_type_id) + : TensorStore(ArrayStoreType::get_data_store_base(_array_store)), + _tensor_type(tensor_type), + _ops(_tensor_type), + _array_store(ArrayStoreType::optimizedConfigForHugePage(max_small_subspaces_type_id, + TensorBufferTypeMapper(max_small_subspaces_type_id, &_ops), + MemoryAllocator::HUGEPAGE_SIZE, 4_Ki, 8_Ki, ALLOC_GROW_FACTOR), + std::move(allocator), TensorBufferTypeMapper(max_small_subspaces_type_id, &_ops)), + _add_buffer() +{ +} + +TensorBufferStore::~TensorBufferStore() = default; + +void +TensorBufferStore::holdTensor(EntryRef ref) +{ + _array_store.remove(ref); +} + +EntryRef +TensorBufferStore::move(EntryRef ref) +{ + if (!ref.valid()) { + return EntryRef(); + } + auto buf = _array_store.get(ref); + auto new_ref = _array_store.add(buf); + _ops.copied_labels(buf); + _array_store.remove(ref); + return new_ref; +} + +EntryRef +TensorBufferStore::store_tensor(const Value &tensor) +{ + uint32_t num_subspaces = tensor.index().size(); + auto array_size = _ops.get_array_size(num_subspaces); + _add_buffer.resize(array_size); + _ops.store_tensor(_add_buffer, tensor); + return _array_store.add(_add_buffer); +} + +EntryRef +TensorBufferStore::store_encoded_tensor(vespalib::nbostream &encoded) +{ + const auto &factory = StreamedValueBuilderFactory::get(); + auto val = vespalib::eval::decode_value(encoded, factory); + return store_tensor(*val); +} + +std::unique_ptr +TensorBufferStore::get_tensor(EntryRef ref) const +{ + if (!ref.valid()) { + return {}; + } + auto buf = _array_store.get(ref); + return _ops.make_fast_view(buf, _tensor_type); +} + +bool +TensorBufferStore::encode_stored_tensor(EntryRef ref, vespalib::nbostream &target) const +{ + if (!ref.valid()) { + return false; + } + auto buf = _array_store.get(ref); + _ops.encode_stored_tensor(buf, _tensor_type, target); + return true; +} + +} diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h new file mode 100644 index 00000000000..14572eb07dc --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_store.h @@ -0,0 +1,38 @@ +// Copyright Yahoo. 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 +#include + +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; + vespalib::eval::ValueType _tensor_type; + TensorBufferOperations _ops; + ArrayStoreType _array_store; + std::vector _add_buffer; +public: + TensorBufferStore(const vespalib::eval::ValueType& tensor_type, std::shared_ptr allocator, uint32_t max_small_subspaces_type_id); + ~TensorBufferStore(); + void holdTensor(EntryRef ref) override; + EntryRef move(EntryRef ref) override; + EntryRef store_tensor(const vespalib::eval::Value &tensor); + EntryRef store_encoded_tensor(vespalib::nbostream &encoded); + std::unique_ptr get_tensor(EntryRef ref) const; + bool encode_stored_tensor(EntryRef ref, vespalib::nbostream &target) const; +}; + +} diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.cpp b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.cpp new file mode 100644 index 00000000000..b4b0b9bbc79 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.cpp @@ -0,0 +1,47 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#include "tensor_buffer_type_mapper.h" +#include "tensor_buffer_operations.h" +#include + +namespace search::tensor { + +TensorBufferTypeMapper::TensorBufferTypeMapper() + : _array_sizes(), + _ops(nullptr) +{ +} + +TensorBufferTypeMapper::TensorBufferTypeMapper(uint32_t max_small_subspaces_type_id, TensorBufferOperations* ops) + : _array_sizes(), + _ops(ops) +{ + _array_sizes.reserve(max_small_subspaces_type_id + 1); + _array_sizes.emplace_back(0); // type id 0 uses LargeSubspacesBufferType + for (uint32_t type_id = 1; type_id <= max_small_subspaces_type_id; ++type_id) { + auto num_subspaces = type_id - 1; + _array_sizes.emplace_back(_ops->get_array_size(num_subspaces)); + } +} + +TensorBufferTypeMapper::~TensorBufferTypeMapper() = default; + +uint32_t +TensorBufferTypeMapper::get_type_id(size_t array_size) const +{ + assert(!_array_sizes.empty()); + auto result = std::lower_bound(_array_sizes.begin() + 1, _array_sizes.end(), array_size); + if (result == _array_sizes.end()) { + return 0; // type id 0 uses LargeSubspacesBufferType + } + return result - _array_sizes.begin(); +} + +size_t +TensorBufferTypeMapper::get_array_size(uint32_t type_id) const +{ + assert(type_id > 0 && type_id < _array_sizes.size()); + return _array_sizes[type_id]; +} + +} diff --git a/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.h b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.h new file mode 100644 index 00000000000..1e02c1cb608 --- /dev/null +++ b/searchlib/src/vespa/searchlib/tensor/tensor_buffer_type_mapper.h @@ -0,0 +1,35 @@ +// Copyright Yahoo. Licensed under the terms of the Apache 2.0 license. See LICENSE in the project root. + +#pragma once + +#include +#include + +namespace search::tensor { + +class LargeSubspacesBufferType; +class SmallSubspacesBufferType; +class TensorBufferOperations; + +/* + * This class provides mapping between type ids and array sizes needed for + * storing a tensor. + */ +class TensorBufferTypeMapper +{ + std::vector _array_sizes; + TensorBufferOperations* _ops; +public: + using SmallBufferType = SmallSubspacesBufferType; + using LargeBufferType = LargeSubspacesBufferType; + + TensorBufferTypeMapper(); + TensorBufferTypeMapper(uint32_t max_small_subspaces_type_id, TensorBufferOperations* ops); + ~TensorBufferTypeMapper(); + + uint32_t get_type_id(size_t array_size) const; + size_t get_array_size(uint32_t type_id) const; + TensorBufferOperations& get_tensor_buffer_operations() const noexcept { return *_ops; } +}; + +} diff --git a/vespalib/src/vespa/vespalib/datastore/array_store.h b/vespalib/src/vespa/vespalib/datastore/array_store.h index fd0271915dd..2c6aaea6b61 100644 --- a/vespalib/src/vespa/vespalib/datastore/array_store.h +++ b/vespalib/src/vespa/vespalib/datastore/array_store.h @@ -110,7 +110,8 @@ public: static vespalib::GenerationHolder &getGenerationHolderLocation(ArrayStore &self) { return DataStoreBase::getGenerationHolderLocation(self._store); } - + // need object location before construction + static DataStoreBase& get_data_store_base(ArrayStore &self) { return self._store; } // Should only be used for unit testing const BufferState &bufferState(EntryRef ref) const; -- cgit v1.2.3