diff --git a/bindings/cpp/CMakeLists.txt b/bindings/cpp/CMakeLists.txt index 80cfe36b..2e633207 100644 --- a/bindings/cpp/CMakeLists.txt +++ b/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 ) diff --git a/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h b/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h index 2e53b805..ad34b27c 100644 --- a/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h +++ b/bindings/cpp/include/svs/runtime/dynamic_vamana_index.h @@ -16,6 +16,7 @@ #pragma once #include +#include #include #include @@ -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 { diff --git a/bindings/cpp/include/svs/runtime/index_blocksize.h b/bindings/cpp/include/svs/runtime/index_blocksize.h new file mode 100644 index 00000000..70273f53 --- /dev/null +++ b/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 + +#include +#include + +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 diff --git a/bindings/cpp/src/dynamic_vamana_index.cpp b/bindings/cpp/src/dynamic_vamana_index.cpp index 10e677da..4ec53f1a 100644 --- a/bindings/cpp/src/dynamic_vamana_index.cpp +++ b/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 data{x, n, impl_->dimensions()}; + std::span 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 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, diff --git a/bindings/cpp/src/dynamic_vamana_index_impl.h b/bindings/cpp/src/dynamic_vamana_index_impl.h index fd91e5bc..8dc1e5a3 100644 --- a/bindings/cpp/src/dynamic_vamana_index_impl.h +++ b/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 data, + std::span labels, + IndexBlockSize blocksize) { + if (!impl_) { + return init_impl(data, labels, blocksize); + } + + impl_->add_points(data, labels); + } + void add(data::ConstSimpleDataView data, std::span 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& data, std::span labels, + IndexBlockSize blocksize, StorageArgs&&... storage_args ) { auto threadpool = default_threadpool(); @@ -397,6 +413,7 @@ class DynamicVamanaIndexImpl { std::forward(tag), data, threadpool, + blocksize.BlockSizeBytes(), std::forward(storage_args)... ); @@ -412,14 +429,18 @@ class DynamicVamanaIndexImpl { }); } - virtual void - init_impl(data::ConstSimpleDataView data, std::span labels) { + virtual void init_impl( + data::ConstSimpleDataView data, + std::span labels, + IndexBlockSize blocksize + ) { impl_.reset(storage::dispatch_storage_kind( get_storage_kind(), [this]( auto&& tag, data::ConstSimpleDataView data, - std::span labels + std::span labels, + IndexBlockSize blocksize ) { using Tag = std::decay_t; return build_impl( @@ -427,11 +448,13 @@ class DynamicVamanaIndexImpl { this->metric_type_, this->vamana_build_parameters(), data, - labels + labels, + blocksize ); }, data, - labels + labels, + blocksize )); } diff --git a/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h b/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h index de998ae7..a97187f9 100644 --- a/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h +++ b/bindings/cpp/src/dynamic_vamana_index_leanvec_impl.h @@ -91,15 +91,19 @@ struct DynamicVamanaIndexLeanVecImpl : public DynamicVamanaIndexImpl { } } - void init_impl(data::ConstSimpleDataView data, std::span labels) - override { + void init_impl( + data::ConstSimpleDataView data, + std::span 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 data, - std::span labels + std::span labels, + IndexBlockSize blocksize ) { using Tag = std::decay_t; 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 )); } diff --git a/bindings/cpp/src/flat_index_impl.h b/bindings/cpp/src/flat_index_impl.h index 1e12dfa4..2f669597 100644 --- a/bindings/cpp/src/flat_index_impl.h +++ b/bindings/cpp/src/flat_index_impl.h @@ -19,6 +19,7 @@ #include "svs_runtime_utils.h" #include +#include #include #include diff --git a/bindings/cpp/src/svs_runtime_utils.h b/bindings/cpp/src/svs_runtime_utils.h index 9e00dcd4..83a03291 100644 --- a/bindings/cpp/src/svs_runtime_utils.h +++ b/bindings/cpp/src/svs_runtime_utils.h @@ -178,7 +178,8 @@ template <> struct StorageFactory { template static StorageType init( const svs::data::ConstSimpleDataView& 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 struct StorageFactory; template - static StorageType init(const svs::data::ConstSimpleDataView& data, Pool& pool) { - StorageType result(data.size(), data.dimensions()); + static StorageType init( + const svs::data::ConstSimpleDataView& 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 struct StorageFactory + static StorageType init(const svs::data::ConstSimpleDataView& data, Pool& pool) { + return init(data, pool, svs::data::BlockingParameters::default_blocksize_bytes); + } + template static StorageType load(const std::filesystem::path& path, Args&&... args) { return svs::lib::load_from_disk(path, SVS_FWD(args)...); @@ -224,8 +236,14 @@ struct StorageFactory { using StorageType = SQStorageType; template - static StorageType init(const svs::data::ConstSimpleDataView& data, Pool& pool) { - return SQStorageType::compress(data, pool); + static StorageType init( + const svs::data::ConstSimpleDataView& 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 @@ -255,8 +273,14 @@ struct StorageFactory { using StorageType = LVQStorageType; template - static StorageType init(const svs::data::ConstSimpleDataView& data, Pool& pool) { - return LVQStorageType::compress(data, pool, 0); + static StorageType init( + const svs::data::ConstSimpleDataView& 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 @@ -289,14 +313,17 @@ struct StorageFactory { static StorageType init( const svs::data::ConstSimpleDataView& data, Pool& pool, + svs::lib::PowerOfTwo blocksize_bytes, size_t leanvec_d = 0, std::optional> 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 ); } diff --git a/bindings/cpp/tests/CMakeLists.txt b/bindings/cpp/tests/CMakeLists.txt index 1fb81ec5..6ca4f580 100644 --- a/bindings/cpp/tests/CMakeLists.txt +++ b/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 diff --git a/bindings/cpp/tests/runtime_test.cpp b/bindings/cpp/tests/runtime_test.cpp index 11b0efd1..e94138fa 100644 --- a/bindings/cpp/tests/runtime_test.cpp +++ b/bindings/cpp/tests/runtime_test.cpp @@ -83,6 +83,7 @@ void write_and_read_index( size_t n, size_t d, svs::runtime::v0::StorageKind storage_kind, + svs::runtime::v0::IndexBlockSize blocksize, svs::runtime::v0::MetricType metric = svs::runtime::v0::MetricType::L2 ) { // Build index @@ -99,7 +100,7 @@ void write_and_read_index( std::vector labels(n); std::iota(labels.begin(), labels.end(), 0); - status = index->add(n, labels.data(), xb.data()); + status = index->add(n, labels.data(), xb.data(), blocksize); CATCH_REQUIRE(status.ok()); svs_test::prepare_temp_directory(); @@ -134,6 +135,7 @@ void write_and_read_index( status = loaded->search(nq, xq, k, distances.data(), result_labels.data()); CATCH_REQUIRE(status.ok()); + CATCH_REQUIRE(index->blocksize_bytes() == blocksize.BlockSizeBytes()); // Clean up svs::runtime::v0::DynamicVamanaIndex::destroy(index); svs::runtime::v0::DynamicVamanaIndex::destroy(loaded); @@ -141,7 +143,9 @@ void write_and_read_index( // Helper that writes and reads and index of requested size // Reports memory usage -UsageInfo run_save_and_load_test(const size_t target_mibytes) { +UsageInfo run_save_and_load_test( + const size_t target_mibytes, svs::runtime::v0::IndexBlockSize blocksize +) { // Generate requested MiB of test data constexpr size_t mem_test_d = 128; const size_t target_bytes = target_mibytes * 1024 * 1024; @@ -171,7 +175,7 @@ UsageInfo run_save_and_load_test(const size_t target_mibytes) { ); CATCH_REQUIRE(status.ok()); CATCH_REQUIRE(index != nullptr); - status = index->add(mem_test_n, labels.data(), large_test_data.data()); + status = index->add(mem_test_n, labels.data(), large_test_data.data(), blocksize); CATCH_REQUIRE(status.ok()); std::ofstream out(filename, std::ios::binary); @@ -224,7 +228,12 @@ CATCH_TEST_CASE("WriteAndReadIndexSVS", "[runtime]") { ); }; write_and_read_index( - build_func, test_data, test_n, test_d, svs::runtime::v0::StorageKind::FP32 + build_func, + test_data, + test_n, + test_d, + svs::runtime::v0::StorageKind::FP32, + svs::runtime::v0::IndexBlockSize(15) ); } @@ -241,7 +250,12 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSFP16", "[runtime]") { ); }; write_and_read_index( - build_func, test_data, test_n, test_d, svs::runtime::v0::StorageKind::FP16 + build_func, + test_data, + test_n, + test_d, + svs::runtime::v0::StorageKind::FP16, + svs::runtime::v0::IndexBlockSize(16) ); } @@ -258,7 +272,12 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSSQI8", "[runtime]") { ); }; write_and_read_index( - build_func, test_data, test_n, test_d, svs::runtime::v0::StorageKind::SQI8 + build_func, + test_data, + test_n, + test_d, + svs::runtime::v0::StorageKind::SQI8, + svs::runtime::v0::IndexBlockSize(17) ); } @@ -275,7 +294,12 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSLVQ4x4", "[runtime]") { ); }; write_and_read_index( - build_func, test_data, test_n, test_d, svs::runtime::v0::StorageKind::LVQ4x4 + build_func, + test_data, + test_n, + test_d, + svs::runtime::v0::StorageKind::LVQ4x4, + svs::runtime::v0::IndexBlockSize(18) ); } @@ -293,7 +317,12 @@ CATCH_TEST_CASE("WriteAndReadIndexSVSVamanaLeanVec4x4", "[runtime]") { ); }; write_and_read_index( - build_func, test_data, test_n, test_d, svs::runtime::v0::StorageKind::LeanVec4x4 + build_func, + test_data, + test_n, + test_d, + svs::runtime::v0::StorageKind::LeanVec4x4, + svs::runtime::v0::IndexBlockSize(19) ); } @@ -330,6 +359,87 @@ CATCH_TEST_CASE("LeanVecWithTrainingData", "[runtime]") { svs::runtime::v0::DynamicVamanaIndex::destroy(index); } +CATCH_TEST_CASE("LeanVecWithTrainingDataCustomBlockSize", "[runtime]") { + const auto& test_data = get_test_data(); + // Build LeanVec index with explicit training + svs::runtime::v0::DynamicVamanaIndex* index = nullptr; + svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::Status status = svs::runtime::v0::DynamicVamanaIndexLeanVec::build( + &index, + test_d, + svs::runtime::v0::MetricType::L2, + svs::runtime::v0::StorageKind::LeanVec4x4, + 32, + build_params + ); + if (!svs::runtime::v0::DynamicVamanaIndex::check_storage_kind( + svs::runtime::v0::StorageKind::LeanVec4x4 + ) + .ok()) { + CATCH_REQUIRE(!status.ok()); + CATCH_SKIP("Storage kind is not supported, skipping test."); + } + CATCH_REQUIRE(status.ok()); + CATCH_REQUIRE(index != nullptr); + + // Add data - should work with provided leanvec dims + std::vector labels(test_n); + std::iota(labels.begin(), labels.end(), 0); + + int block_size_exp = 17; // block_size_bytes = 2^block_size_exp + status = index->add( + test_n, + labels.data(), + test_data.data(), + svs::runtime::v0::IndexBlockSize(block_size_exp) + ); + CATCH_REQUIRE(status.ok()); + + CATCH_REQUIRE(index->blocksize_bytes().raw() == block_size_exp); + + svs::runtime::v0::DynamicVamanaIndex::destroy(index); +} + +CATCH_TEST_CASE("TrainingDataCustomBlockSize", "[runtime]") { + const auto& test_data = get_test_data(); + // Build LeanVec index with explicit training + svs::runtime::v0::DynamicVamanaIndex* index = nullptr; + svs::runtime::v0::VamanaIndex::BuildParams build_params{64}; + svs::runtime::v0::Status status = svs::runtime::v0::DynamicVamanaIndex::build( + &index, + test_d, + svs::runtime::v0::MetricType::L2, + svs::runtime::v0::StorageKind::FP32, + build_params + ); + if (!svs::runtime::v0::DynamicVamanaIndex::check_storage_kind( + svs::runtime::v0::StorageKind::FP32 + ) + .ok()) { + CATCH_REQUIRE(!status.ok()); + CATCH_SKIP("Storage kind is not supported, skipping test."); + } + CATCH_REQUIRE(status.ok()); + CATCH_REQUIRE(index != nullptr); + + // Add data - should work with provided leanvec dims + std::vector labels(test_n); + std::iota(labels.begin(), labels.end(), 0); + + int block_size_exp = 17; // block_size_bytes = 2^block_size_exp + status = index->add( + test_n, + labels.data(), + test_data.data(), + svs::runtime::v0::IndexBlockSize(block_size_exp) + ); + CATCH_REQUIRE(status.ok()); + + CATCH_REQUIRE(index->blocksize_bytes().raw() == block_size_exp); + + svs::runtime::v0::DynamicVamanaIndex::destroy(index); +} + CATCH_TEST_CASE("FlatIndexWriteAndRead", "[runtime]") { const auto& test_data = get_test_data(); svs::runtime::v0::FlatIndex* index = nullptr; @@ -399,7 +509,9 @@ CATCH_TEST_CASE("SearchWithIDFilter", "[runtime]") { // Add data std::vector labels(test_n); std::iota(labels.begin(), labels.end(), 0); - status = index->add(test_n, labels.data(), test_data.data()); + status = index->add( + test_n, labels.data(), test_data.data(), svs::runtime::v0::IndexBlockSize(30) + ); CATCH_REQUIRE(status.ok()); const int nq = 8; @@ -445,7 +557,9 @@ CATCH_TEST_CASE("RangeSearchFunctional", "[runtime]") { // Add data std::vector labels(test_n); std::iota(labels.begin(), labels.end(), 0); - status = index->add(test_n, labels.data(), test_data.data()); + status = index->add( + test_n, labels.data(), test_data.data(), svs::runtime::v0::IndexBlockSize(30) + ); CATCH_REQUIRE(status.ok()); const int nq = 5; @@ -472,19 +586,19 @@ CATCH_TEST_CASE("MemoryUsageOnLoad", "[runtime][memory]") { }; CATCH_SECTION("SmallIndex") { - auto stats = run_save_and_load_test(10); + auto stats = run_save_and_load_test(10, svs::runtime::v0::IndexBlockSize(30)); CATCH_REQUIRE(stats.file_size < 20 * 1024 * 1024); CATCH_REQUIRE(stats.rss_increase < threshold(stats.file_size)); } CATCH_SECTION("MediumIndex") { - auto stats = run_save_and_load_test(50); + auto stats = run_save_and_load_test(50, svs::runtime::v0::IndexBlockSize(30)); CATCH_REQUIRE(stats.file_size < 100 * 1024 * 1024); CATCH_REQUIRE(stats.rss_increase < threshold(stats.file_size)); } CATCH_SECTION("LargeIndex") { - auto stats = run_save_and_load_test(200); + auto stats = run_save_and_load_test(200, svs::runtime::v0::IndexBlockSize(30)); CATCH_REQUIRE(stats.file_size < 400 * 1024 * 1024); CATCH_REQUIRE(stats.rss_increase < threshold(stats.file_size)); } diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index 169be199..da9a4ae5 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -260,6 +260,12 @@ class MutableVamanaIndex { ); } + template