diff options
author | Tor Egge <Tor.Egge@online.no> | 2023-06-26 14:01:55 +0200 |
---|---|---|
committer | Tor Egge <Tor.Egge@online.no> | 2023-06-26 14:01:55 +0200 |
commit | b0eda0e4a90d6abb9a8a7ce7831dee4cf19a2f3e (patch) | |
tree | 7b1dc6109bfda259c5f523514ebc66174be93e22 | |
parent | 0c341f8ed39b3edcd1938d964cbdf9ce7c179411 (diff) |
Cap number of entries in a dense tensor store buffer to avoid very
large buffers.
3 files changed, 35 insertions, 3 deletions
diff --git a/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp b/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp index 26ef57aab65..87420e8939f 100644 --- a/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp +++ b/searchlib/src/tests/tensor/dense_tensor_store/dense_tensor_store_test.cpp @@ -9,6 +9,7 @@ LOG_SETUP("dense_tensor_store_test"); #include <vespa/eval/eval/value.h> #include <vespa/eval/eval/value_type.h> #include <vespa/eval/eval/test/value_compare.h> +#include <vespa/vespalib/util/size_literals.h> using search::tensor::DenseTensorStore; using vespalib::eval::SimpleValue; @@ -90,5 +91,26 @@ TEST("require that array size is calculated correctly") TEST_DO(assertArraySize("tensor<int8>(x[65])", 96)); } +void +assert_max_buffer_entries(const vespalib::string& tensor_type, uint32_t exp_entries) +{ + Fixture f(tensor_type); + EXPECT_EQUAL(exp_entries, f.store.get_max_buffer_entries()); +} + +TEST("require that max entries is calculated correctly") +{ + TEST_DO(assert_max_buffer_entries("tensor(x[1])", 1_Mi)); + TEST_DO(assert_max_buffer_entries("tensor(x[32])", 1_Mi)); + TEST_DO(assert_max_buffer_entries("tensor(x[64])", 512_Ki)); + TEST_DO(assert_max_buffer_entries("tensor(x[1024])", 32_Ki)); + TEST_DO(assert_max_buffer_entries("tensor(x[1024])", 32_Ki)); + TEST_DO(assert_max_buffer_entries("tensor(x[16777216])", 2)); + TEST_DO(assert_max_buffer_entries("tensor(x[33554428])", 2)); + TEST_DO(assert_max_buffer_entries("tensor(x[33554429])", 1)); + TEST_DO(assert_max_buffer_entries("tensor(x[33554432])", 1)); + TEST_DO(assert_max_buffer_entries("tensor(x[303554432])", 1)); +} + TEST_MAIN() { TEST_RUN_ALL(); } diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp index c51d0ec7fd3..638c254aac1 100644 --- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp +++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.cpp @@ -34,6 +34,13 @@ size_t my_align(size_t size, size_t alignment) { return (size - (size % alignment)); } +size_t +cap_max_entries(size_t max_entries, size_t max_buffer_size, size_t entry_size) +{ + size_t dynamic_max_entries = (max_buffer_size + (entry_size - 1)) / entry_size; + return std::min(max_entries, dynamic_max_entries); +} + } DenseTensorStore::TensorSizeCalc::TensorSizeCalc(const ValueType &type) @@ -55,7 +62,7 @@ DenseTensorStore::TensorSizeCalc::TensorSizeCalc(const ValueType &type) } DenseTensorStore::BufferType::BufferType(const TensorSizeCalc &tensorSizeCalc, std::shared_ptr<vespalib::alloc::MemoryAllocator> allocator) - : vespalib::datastore::BufferType<char>(tensorSizeCalc.alignedSize(), MIN_BUFFER_ARRAYS, RefType::offsetSize()), + : vespalib::datastore::BufferType<char>(tensorSizeCalc.alignedSize(), MIN_BUFFER_ARRAYS, cap_max_entries(RefType::offsetSize(), max_dense_tensor_buffer_size, tensorSizeCalc.alignedSize())), _allocator(std::move(allocator)) {} diff --git a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h index 0dd483e7f08..b7375c4d2c3 100644 --- a/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h +++ b/searchlib/src/vespa/searchlib/tensor/dense_tensor_store.h @@ -20,9 +20,11 @@ namespace search::tensor { class DenseTensorStore : public TensorStore { public: - using RefType = vespalib::datastore::EntryRefT<22>; + // 4 Ki buffers of 256 MiB each is 1 TiB. + using RefType = vespalib::datastore::EntryRefT<20>; using DataStoreType = vespalib::datastore::DataStoreT<RefType>; using ValueType = vespalib::eval::ValueType; + static constexpr size_t max_dense_tensor_buffer_size = 256_Mi; struct TensorSizeCalc { @@ -90,8 +92,9 @@ public: return VectorBundle(getRawBuffer(ref), 1, _subspace_type); } const SubspaceType& get_subspace_type() const noexcept { return _subspace_type; } - // The following method is meant to be used only for unit tests. + // The following methods are meant to be used only for unit tests. uint32_t getArraySize() const { return _bufferType.getArraySize(); } + uint32_t get_max_buffer_entries() const noexcept { return _bufferType.get_max_entries(); } }; } |