Make block_size flexible

bindings/cpp/CMakeLists.txt

@@ -22,6 +22,7 @@ set(SVS_RUNTIME_HEADERS
     include/svs/runtime/training.h
     include/svs/runtime/vamana_index.h
     include/svs/runtime/dynamic_vamana_index.h
+    include/svs/runtime/index_blocksize.h
     include/svs/runtime/flat_index.h
 )
 

bindings/cpp/include/svs/runtime/dynamic_vamana_index.h

@@ -16,6 +16,7 @@
 
 #pragma once
 #include <svs/runtime/api_defs.h>
+#include <svs/runtime/index_blocksize.h>
 #include <svs/runtime/training.h>
 #include <svs/runtime/vamana_index.h>
 
@@ -29,6 +30,9 @@ namespace v0 {
 
 // Abstract interface for Dynamic Vamana-based indexes.
 struct SVS_RUNTIME_API DynamicVamanaIndex : public VamanaIndex {
+    virtual Status
+    add(size_t n, const size_t* labels, const float* x, IndexBlockSize blocksize
+    ) noexcept = 0;
     virtual Status add(size_t n, const size_t* labels, const float* x) noexcept = 0;
     virtual Status
     remove_selected(size_t* num_removed, const IDFilter& selector) noexcept = 0;
@@ -58,6 +62,8 @@ struct SVS_RUNTIME_API DynamicVamanaIndex : public VamanaIndex {
         MetricType metric,
         StorageKind storage_kind
     ) noexcept;
+
+    virtual lib::PowerOfTwo blocksize_bytes() const noexcept = 0;
 };
 
 struct SVS_RUNTIME_API DynamicVamanaIndexLeanVec : public DynamicVamanaIndex {

bindings/cpp/include/svs/runtime/index_blocksize.h

@@ -0,0 +1,46 @@
+/*
+ * Copyright 2023 Intel Corporation
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <cstddef>
+
+#include <svs/lib/exception.h>
+#include <svs/lib/misc.h>
+
+namespace svs::runtime::v0 {
+
+class IndexBlockSize {
+    constexpr static size_t kMaxBlockSizeExp = 30; // 1GB
+    constexpr static size_t kMinBlockSizeExp = 12; // 4KB
+
+    svs::lib::PowerOfTwo blocksize_bytes_;
+
+  public:
+    explicit IndexBlockSize(size_t blocksize_exp) {
+        if (blocksize_exp > kMaxBlockSizeExp) {
+            throw ANNEXCEPTION("Blocksize is too large!");
+        } else if (blocksize_exp < kMinBlockSizeExp) {
+            throw ANNEXCEPTION("Blocksize is too small!");
+        }
+
+        blocksize_bytes_ = svs::lib::PowerOfTwo(blocksize_exp);
+    }
+
+    svs::lib::PowerOfTwo BlockSizeBytes() const { return blocksize_bytes_; }
+};
+
+} // namespace svs::runtime::v0

bindings/cpp/src/dynamic_vamana_index.cpp

@@ -55,6 +55,16 @@ struct DynamicVamanaIndexManagerBase : public DynamicVamanaIndex {
     DynamicVamanaIndexManagerBase& operator=(DynamicVamanaIndexManagerBase&&) = default;
     ~DynamicVamanaIndexManagerBase() override = default;
 
+    Status
+    add(size_t n, const size_t* labels, const float* x, IndexBlockSize blocksize
+    ) noexcept override {
+        return runtime_error_wrapper([&] {
+            svs::data::ConstSimpleDataView<float> data{x, n, impl_->dimensions()};
+            std::span<const size_t> lbls(labels, n);
+            impl_->add(data, lbls, blocksize);
+        });
+    }
+
     Status add(size_t n, const size_t* labels, const float* x) noexcept override {
         return runtime_error_wrapper([&] {
             svs::data::ConstSimpleDataView<float> data{x, n, impl_->dimensions()};
@@ -77,6 +87,8 @@ struct DynamicVamanaIndexManagerBase : public DynamicVamanaIndex {
         });
     }
 
+    lib::PowerOfTwo blocksize_bytes() const noexcept { return impl_->blocksize_bytes(); }
+
     Status search(
         size_t n,
         const float* x,

bindings/cpp/src/dynamic_vamana_index_impl.h

@@ -63,15 +63,30 @@ class DynamicVamanaIndexImpl {
 
     size_t size() const { return impl_ ? impl_->size() : 0; }
 
+    lib::PowerOfTwo blocksize_bytes() const { return impl_->blocksize_bytes(); }
+
     size_t dimensions() const { return dim_; }
 
     MetricType metric_type() const { return metric_type_; }
 
     StorageKind get_storage_kind() const { return storage_kind_; }
 
+    void
+    add(data::ConstSimpleDataView<float> data,
+        std::span<const size_t> labels,
+        IndexBlockSize blocksize) {
+        if (!impl_) {
+            return init_impl(data, labels, blocksize);
+        }
+
+        impl_->add_points(data, labels);
+    }
+
     void add(data::ConstSimpleDataView<float> data, std::span<const size_t> labels) {
         if (!impl_) {
-            return init_impl(data, labels);
+            IndexBlockSize blocksize(data::BlockingParameters::default_blocksize_bytes.raw()
+            );
+            return init_impl(data, labels, blocksize);
         }
 
         impl_->add_points(data, labels);
@@ -389,6 +404,7 @@ class DynamicVamanaIndexImpl {
         const index::vamana::VamanaBuildParameters& parameters,
         const svs::data::ConstSimpleDataView<float>& data,
         std::span<const size_t> labels,
+        IndexBlockSize blocksize,
         StorageArgs&&... storage_args
     ) {
         auto threadpool = default_threadpool();
@@ -397,6 +413,7 @@ class DynamicVamanaIndexImpl {
             std::forward<Tag>(tag),
             data,
             threadpool,
+            blocksize.BlockSizeBytes(),
             std::forward<StorageArgs>(storage_args)...
         );
 
@@ -412,26 +429,32 @@ class DynamicVamanaIndexImpl {
         });
     }
 
-    virtual void
-    init_impl(data::ConstSimpleDataView<float> data, std::span<const size_t> labels) {
+    virtual void init_impl(
+        data::ConstSimpleDataView<float> data,
+        std::span<const size_t> labels,
+        IndexBlockSize blocksize
+    ) {
         impl_.reset(storage::dispatch_storage_kind(
             get_storage_kind(),
             [this](
                 auto&& tag,
                 data::ConstSimpleDataView<float> data,
-                std::span<const size_t> labels
+                std::span<const size_t> labels,
+                IndexBlockSize blocksize
             ) {
                 using Tag = std::decay_t<decltype(tag)>;
                 return build_impl(
                     std::forward<Tag>(tag),
                     this->metric_type_,
                     this->vamana_build_parameters(),
                     data,
-                    labels
+                    labels,
+                    blocksize
                 );
             },
             data,
-            labels
+            labels,
+            blocksize
         ));
     }
 

bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h

@@ -91,15 +91,19 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
         }
     }
 
-    void init_impl(data::ConstSimpleDataView<float> data, std::span<const size_t> labels)
-        override {
+    void init_impl(
+        data::ConstSimpleDataView<float> data,
+        std::span<const size_t> labels,
+        IndexBlockSize blocksize
+    ) override {
         assert(storage::is_leanvec_storage(this->storage_kind_));
         impl_.reset(dispatch_leanvec_storage_kind(
             this->storage_kind_,
             [this](
                 auto&& tag,
                 data::ConstSimpleDataView<float> data,
-                std::span<const size_t> labels
+                std::span<const size_t> labels,
+                IndexBlockSize blocksize
             ) {
                 using Tag = std::decay_t<decltype(tag)>;
                 return DynamicVamanaIndexImpl::build_impl(
@@ -108,12 +112,14 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl {
                     this->vamana_build_parameters(),
                     data,
                     labels,
+                    blocksize,
                     this->leanvec_dims_,
                     this->leanvec_matrices_
                 );
             },
             data,
-            labels
+            labels,
+            blocksize
         ));
     }
 

bindings/cpp/src/flat_index_impl.h

@@ -19,6 +19,7 @@
 #include "svs_runtime_utils.h"
 
 #include <svs/runtime/flat_index.h>
+#include <svs/runtime/index_blocksize.h>
 
 #include <algorithm>
 #include <memory>

bindings/cpp/src/svs_runtime_utils.h

@@ -178,7 +178,8 @@ template <> struct StorageFactory<UnsupportedStorageType> {
     template <svs::threads::ThreadPool Pool>
     static StorageType init(
         const svs::data::ConstSimpleDataView<float>& SVS_UNUSED(data),
-        Pool& SVS_UNUSED(pool)
+        Pool& SVS_UNUSED(pool),
+        svs::lib::PowerOfTwo SVS_UNUSED(blocksize_bytes)
     ) {
         throw StatusException(
             ErrorCode::NOT_IMPLEMENTED, "Requested storage kind is not supported"
@@ -198,8 +199,14 @@ template <typename ElementType> struct StorageFactory<SimpleDatasetType<ElementT
     using StorageType = SimpleDatasetType<ElementType>;
 
     template <svs::threads::ThreadPool Pool>
-    static StorageType init(const svs::data::ConstSimpleDataView<float>& data, Pool& pool) {
-        StorageType result(data.size(), data.dimensions());
+    static StorageType init(
+        const svs::data::ConstSimpleDataView<float>& data,
+        Pool& pool,
+        svs::lib::PowerOfTwo blocksize_bytes
+    ) {
+        auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes};
+        typename StorageType::allocator_type alloc(parameters);
+        StorageType result(data.size(), data.dimensions(), alloc);
         svs::threads::parallel_for(
             pool,
             svs::threads::StaticPartition(result.size()),
@@ -212,6 +219,11 @@ template <typename ElementType> struct StorageFactory<SimpleDatasetType<ElementT
         return result;
     }
 
+    template <svs::threads::ThreadPool Pool>
+    static StorageType init(const svs::data::ConstSimpleDataView<float>& data, Pool& pool) {
+        return init(data, pool, svs::data::BlockingParameters::default_blocksize_bytes);
+    }
+
     template <typename... Args>
     static StorageType load(const std::filesystem::path& path, Args&&... args) {
         return svs::lib::load_from_disk<StorageType>(path, SVS_FWD(args)...);
@@ -224,8 +236,14 @@ struct StorageFactory<SQStorageType> {
     using StorageType = SQStorageType;
 
     template <svs::threads::ThreadPool Pool>
-    static StorageType init(const svs::data::ConstSimpleDataView<float>& data, Pool& pool) {
-        return SQStorageType::compress(data, pool);
+    static StorageType init(
+        const svs::data::ConstSimpleDataView<float>& data,
+        Pool& pool,
+        svs::lib::PowerOfTwo blocksize_bytes
+    ) {
+        auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes};
+        typename StorageType::allocator_type alloc(parameters);
+        return SQStorageType::compress(data, pool, alloc);
     }
 
     template <typename... Args>
@@ -255,8 +273,14 @@ struct StorageFactory<LVQStorageType> {
     using StorageType = LVQStorageType;
 
     template <svs::threads::ThreadPool Pool>
-    static StorageType init(const svs::data::ConstSimpleDataView<float>& data, Pool& pool) {
-        return LVQStorageType::compress(data, pool, 0);
+    static StorageType init(
+        const svs::data::ConstSimpleDataView<float>& data,
+        Pool& pool,
+        svs::lib::PowerOfTwo blocksize_bytes
+    ) {
+        auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes};
+        typename LVQStorageType::allocator_type alloc(parameters);
+        return LVQStorageType::compress(data, pool, 0, alloc);
     }
 
     template <typename... Args>
@@ -289,14 +313,17 @@ struct StorageFactory<LeanVecStorageType> {
     static StorageType init(
         const svs::data::ConstSimpleDataView<float>& data,
         Pool& pool,
+        svs::lib::PowerOfTwo blocksize_bytes,
         size_t leanvec_d = 0,
         std::optional<svs::leanvec::LeanVecMatrices<svs::Dynamic>> matrices = std::nullopt
     ) {
         if (leanvec_d == 0) {
             leanvec_d = (data.dimensions() + 1) / 2;
         }
+        auto parameters = svs::data::BlockingParameters{.blocksize_bytes = blocksize_bytes};
+        typename LeanVecStorageType::allocator_type alloc(parameters);
         return LeanVecStorageType::reduce(
-            data, std::move(matrices), pool, 0, svs::lib::MaybeStatic{leanvec_d}
+            data, std::move(matrices), pool, 0, svs::lib::MaybeStatic{leanvec_d}, alloc
         );
     }
 

bindings/cpp/tests/CMakeLists.txt

@@ -37,6 +37,7 @@ FetchContent_Declare(
 )
 
 FetchContent_MakeAvailable(Catch2)
+
 set(CMAKE_CXX_STANDARD ${PRESET_CMAKE_CXX_STANDARD})
 
 # Add test executable
@@ -50,6 +51,7 @@ add_executable(svs_runtime_test ${TEST_SOURCES})
 target_link_libraries(svs_runtime_test PRIVATE
     svs_runtime
     Catch2::Catch2WithMain
+    fmt::fmt
 )
 
 # Set C++ standard
@@ -64,6 +66,7 @@ set_target_properties(svs_runtime_test PROPERTIES
 target_include_directories(svs_runtime_test PRIVATE
     ${CMAKE_CURRENT_SOURCE_DIR}
     ${CMAKE_CURRENT_SOURCE_DIR}/../include
+    ${CMAKE_CURRENT_SOURCE_DIR}/../../../include
 )
 
 # Enable testing with CTest