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(); }
 };
 
 }