diff --git a/bindings/c/CMakeLists.txt b/bindings/c/CMakeLists.txt index be466bd7..60c0ec27 100644 --- a/bindings/c/CMakeLists.txt +++ b/bindings/c/CMakeLists.txt @@ -35,6 +35,7 @@ set(SVS_C_API_SOURCES src/svs_c.cpp src/dispatcher_vamana.cpp src/dispatcher_dynamic_vamana.cpp + src/data_builder.cpp ) add_library(${TARGET_NAME} SHARED diff --git a/bindings/c/include/svs/c_api/svs_c.h b/bindings/c/include/svs/c_api/svs_c.h index dccccb97..2577b994 100644 --- a/bindings/c/include/svs/c_api/svs_c.h +++ b/bindings/c/include/svs/c_api/svs_c.h @@ -355,6 +355,37 @@ SVS_API bool svs_index_builder_set_threadpool_custom( svs_index_builder_h builder, svs_threadpool_i pool, svs_error_h out_err /*=NULL*/ ); +/// @brief Estimate the memory usage of an index based on the builder configuration and +/// number of vectors +/// @param builder The index builder handle +/// @param num_vectors The number of vectors to be indexed +/// @param out_breakdown Pointer to a structure to hold the memory breakdown +/// @param out_err An optional error handle to capture errors +/// @return true on success, false on failure +SVS_API bool svs_index_builder_estimate_memory( + svs_index_builder_h builder, + size_t num_vectors, + svs_memory_breakdown_t* out_breakdown, + svs_error_h out_err /*=NULL*/ +); + +/// @brief Estimate the memory usage of a dynamic index based on the builder configuration, +/// number of vectors, and block size +/// @param builder The index builder handle +/// @param num_vectors The number of vectors to be indexed +/// @param blocksize_bytes The block size in bytes for dynamic index building (0 for +/// default) +/// @param out_breakdown Pointer to a structure to hold the memory breakdown +/// @param out_err An optional error handle to capture errors +/// @return true on success, false on failure +SVS_API bool svs_index_builder_estimate_memory_dynamic( + svs_index_builder_h builder, + size_t num_vectors, + size_t blocksize_bytes, + svs_memory_breakdown_t* out_breakdown, + svs_error_h out_err /*=NULL*/ +); + /// @brief Build an index from the provided data /// @param builder The index builder handle /// @param data Pointer to the vector data (float array) diff --git a/bindings/c/src/data_builder.cpp b/bindings/c/src/data_builder.cpp new file mode 100644 index 00000000..176fb99e --- /dev/null +++ b/bindings/c/src/data_builder.cpp @@ -0,0 +1,133 @@ +/* + * Copyright 2026 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. + */ + +#include "data_builder.hpp" + +#include "storage.hpp" + +#include +#include +#include + +#include + +namespace svs::c_runtime { +// namespace { +template +size_t +estimate_size(DataBuilder builder, size_t num_vectors, size_t dimension, svs::lib::Empty) { + using allocator_type = typename DataBuilder::allocator_type; + static_assert( + !svs::data::is_blocked_v, + "estimate_size requires a non-blocked allocator type." + ); + return builder.estimate_size(num_vectors, dimension, allocator_type{}); +} + +template +size_t estimate_blocked_size( + DataBuilder builder, size_t num_vectors, size_t dimension, size_t blocksize_bytes +) { + using allocator_type = typename DataBuilder::allocator_type; + static_assert( + svs::data::is_blocked_v, + "estimate_blocked_size requires a blocked allocator type." + ); + svs::data::BlockingParameters block_params; + if (blocksize_bytes != 0) { + block_params.blocksize_bytes = svs::lib::prevpow2(blocksize_bytes); + } + auto allocator = allocator_type{block_params}; + return builder.estimate_size(num_vectors, dimension, allocator); +} + +template +void register_data_size_specializations(Dispatcher& dispatcher) { + auto size_closure = [&dispatcher]() { + dispatcher.register_target(&estimate_size); + }; + + for_simple_specializations(size_closure); + for_leanvec_specializations(size_closure); + for_lvq_specializations(size_closure); + for_sq_specializations(size_closure); + + auto blocked_size_closure = [&dispatcher]() { + dispatcher.register_target(&estimate_blocked_size); + }; + + for_simple_specializations(blocked_size_closure); + for_leanvec_specializations(blocked_size_closure); + for_lvq_specializations(blocked_size_closure); + for_sq_specializations(blocked_size_closure); +} + +using BlocksizeArg = std::variant; + +using EstimateSizeDispatcher = + svs::lib::Dispatcher; + +const EstimateSizeDispatcher& build_data_size_dispatcher() { + static EstimateSizeDispatcher dispatcher = [] { + EstimateSizeDispatcher d{}; + register_data_size_specializations(d); + return d; + }(); + return dispatcher; +} + +size_t dispatch_data_size_estimation( + const Storage* storage, + size_t num_vectors, + size_t dimension, + BlocksizeArg blocksize_bytes +) { + return build_data_size_dispatcher().invoke( + storage, num_vectors, dimension, blocksize_bytes + ); +} +//} // namespace + +size_t estimate_data_size(const Storage* storage, size_t num_vectors, size_t dimension) { + if (storage == nullptr) { + throw std::invalid_argument("Storage pointer cannot be null."); + } + if (num_vectors == 0) { + throw std::invalid_argument("Number of vectors must be greater than zero."); + } + if (dimension == 0) { + throw std::invalid_argument("Dimension must be greater than zero."); + } + return dispatch_data_size_estimation( + storage, num_vectors, dimension, svs::lib::Empty{} + ); +} + +size_t estimate_data_size_blocked( + const Storage* storage, size_t num_vectors, size_t dimension, size_t blocksize_bytes +) { + if (storage == nullptr) { + throw std::invalid_argument("Storage pointer cannot be null."); + } + if (num_vectors == 0) { + throw std::invalid_argument("Number of vectors must be greater than zero."); + } + if (dimension == 0) { + throw std::invalid_argument("Dimension must be greater than zero."); + } + return dispatch_data_size_estimation(storage, num_vectors, dimension, blocksize_bytes); +} +} // namespace svs::c_runtime diff --git a/bindings/c/src/data_builder.hpp b/bindings/c/src/data_builder.hpp index f4dfafd9..c92abbd4 100644 --- a/bindings/c/src/data_builder.hpp +++ b/bindings/c/src/data_builder.hpp @@ -19,3 +19,11 @@ #include "data_builder/lvq.hpp" #include "data_builder/simple.hpp" #include "data_builder/sq.hpp" +#include "storage.hpp" + +namespace svs::c_runtime { +size_t estimate_data_size(const Storage* storage, size_t num_vectors, size_t dimension); +size_t estimate_data_size_blocked( + const Storage* storage, size_t num_vectors, size_t dimension, size_t blocksize_bytes +); +} // namespace svs::c_runtime diff --git a/bindings/c/src/data_builder/leanvec.hpp b/bindings/c/src/data_builder/leanvec.hpp index 13ba58a5..3c51f4fc 100644 --- a/bindings/c/src/data_builder/leanvec.hpp +++ b/bindings/c/src/data_builder/leanvec.hpp @@ -19,6 +19,7 @@ #include "svs/c_api/svs_c.h" +#include "data_builder/lvq.hpp" #include "storage.hpp" #include "types_support.hpp" @@ -75,6 +76,42 @@ class LeanVecDataBuilder { load(const std::filesystem::path& path, const allocator_type& allocator = {}) { return svs::lib::load_from_disk(path, allocator); } + + size_t estimate_size( + size_t num_vectors, size_t dimension, const allocator_type& allocator = {} + ) const { + // Current version of LeanVecDataBuilder supports LVQ-only datasets, so we can + // directly reuse LVQDataBuilder::estimate_size() + // + // LeanDataset uses primary-only LVQ (ResidualBits == 0), so we can use + // LVQDataBuilder and LVQDataBuilder to estimate sizes for primary and + // secondary datasets. + + // Estimate primary size + using primary_data_builder = LVQDataBuilder; + const auto primary_size = + primary_data_builder{}.estimate_size(num_vectors, leanvec_dims_, allocator); + + // Estimate secondary size + using secondary_data_builder = LVQDataBuilder; + const auto secondary_size = + secondary_data_builder{}.estimate_size(num_vectors, dimension, allocator); + + // TODO: Fix the actual memory breakdown reported by index by implementing + // dataset_allocated_bytes() specialization for LeanDataset. LeanVec matrices are 2 + // SimpleData matrices of float, each of size (dimension x leanvec_dims) + const size_t matrices_size = 0; // 2 * dimension * leanvec_dims_ * sizeof(float); + + // LeanVec means is the vector of double of size (dimension) + const size_t means_size = 0; // dimension * sizeof(double); + + // is_pca_ flag is a boolean, so it takes 1 byte + const size_t is_pca_size = 0; // sizeof(bool); + + const auto total_size = + primary_size + secondary_size + matrices_size + means_size + is_pca_size; + return total_size; + } }; template diff --git a/bindings/c/src/data_builder/lvq.hpp b/bindings/c/src/data_builder/lvq.hpp index d7d5912a..7235b14d 100644 --- a/bindings/c/src/data_builder/lvq.hpp +++ b/bindings/c/src/data_builder/lvq.hpp @@ -52,11 +52,18 @@ class LVQDataBuilder { public: LVQDataBuilder() {} + // Follow the logic of svs::leanvec::detail::PickContainer which looks like: + // "Use Turbo-encoding for 4-bit LVQ." + using Sequential = svs::quantization::lvq::Sequential; + using Turbo16x8 = svs::quantization::lvq::Turbo<16, 8>; + template + using AutoStrategy = std::conditional_t<(Primary == 4), Turbo16x8, Sequential>; + using data_type = svs::quantization::lvq::LVQDataset< PrimaryBits, ResidualBits, svs::Dynamic, - svs::quantization::lvq::Sequential, + AutoStrategy, Allocator>; using allocator_type = Allocator; @@ -73,6 +80,52 @@ class LVQDataBuilder { load(const std::filesystem::path& path, const allocator_type& allocator = {}) { return svs::lib::load_from_disk(path, allocator); } + + static constexpr size_t primary_element_size(size_t dimension, size_t alignment = 0) { + using primary_type = typename data_type::primary_type; + using layout_type = typename primary_type::helper_type; + using layout_dims_type = svs::lib::MaybeStatic; + const auto layout_dims = layout_dims_type{dimension}; + return primary_type::compute_data_dimensions(layout_type{layout_dims}, alignment); + } + + static constexpr size_t residual_element_size(size_t dims) { + if constexpr (ResidualBits == 0) { + return 0; + } else { + using residual_type = typename data_type::residual_type; + using dims_type = svs::lib::MaybeStatic; + auto residual_dims = dims_type{dims}; + return residual_type::total_bytes(residual_dims); + } + } + + size_t estimate_size( + size_t num_vectors, size_t dimension, const allocator_type& allocator = {} + ) const { + const size_t alignment = 0; // Assuming no specific alignment for estimation + + const auto primary_element_sz = primary_element_size(dimension, alignment); + const auto primary_size = + svs::c_runtime::adjust_blocked_size(num_vectors, primary_element_sz, allocator); + + const auto residual_element_sz = residual_element_size(dimension); + const auto residual_size = svs::c_runtime::adjust_blocked_size( + num_vectors, residual_element_sz, allocator + ); + + // Assuming a single centroid for estimation purposes + const size_t num_centroids = 1; // Assuming 1 centroid for estimation + // TODO: Fix the actual memory breakdown reported by index by implementing + // dataset_allocated_bytes() specialization for LVQDataset. + const size_t centroid_size = 0; // Skipping centroids for estimation + // const auto centroid_size = + // sizeof(typename data_type::centroid_type::element_type) * dimension; + + const auto total_size = + primary_size + residual_size + num_centroids * centroid_size; + return total_size; + } }; template diff --git a/bindings/c/src/data_builder/simple.hpp b/bindings/c/src/data_builder/simple.hpp index 1f016241..48f1125a 100644 --- a/bindings/c/src/data_builder/simple.hpp +++ b/bindings/c/src/data_builder/simple.hpp @@ -59,6 +59,15 @@ class SimpleDataBuilder { load(const std::filesystem::path& path, const allocator_type& allocator = {}) { return svs::lib::load_from_disk(path, allocator); } + + size_t estimate_size( + size_t num_vectors, size_t dimension, const allocator_type& allocator = {} + ) const { + const auto element_size = sizeof(typename data_type::element_type) * dimension; + const auto total_size = + svs::c_runtime::adjust_blocked_size(num_vectors, element_size, allocator); + return total_size; + } }; template diff --git a/bindings/c/src/data_builder/sq.hpp b/bindings/c/src/data_builder/sq.hpp index f2fd9559..0b38d3c3 100644 --- a/bindings/c/src/data_builder/sq.hpp +++ b/bindings/c/src/data_builder/sq.hpp @@ -57,6 +57,18 @@ template > class SQDat load(const std::filesystem::path& path, const allocator_type& allocator = {}) { return svs::lib::load_from_disk(path, allocator); } + + size_t estimate_size( + size_t num_vectors, size_t dimension, const allocator_type& allocator = {} + ) const { + const auto element_size = sizeof(typename data_type::element_type) * dimension; + const auto data_size = + svs::c_runtime::adjust_blocked_size(num_vectors, element_size, allocator); + // TODO: Fix the actual memory breakdown reported by index by implementing + // dataset_allocated_bytes() specialization for SQDataset. + const size_t scale_bias_size = 0; // sizeof(float) * 2; + return data_size + scale_bias_size; + } }; template diff --git a/bindings/c/src/dispatcher_dynamic_vamana.cpp b/bindings/c/src/dispatcher_dynamic_vamana.cpp index 3d5669fb..283b1f75 100644 --- a/bindings/c/src/dispatcher_dynamic_vamana.cpp +++ b/bindings/c/src/dispatcher_dynamic_vamana.cpp @@ -165,4 +165,50 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_load( blocksize_bytes ); } + +svs::index::vamana::MemoryBreakdown dispatch_dynamic_vamana_memory_estimate( + const svs::index::vamana::VamanaBuildParameters& build_params, + size_t num_vectors, + size_t dimension, + const Storage* storage, + svs::DistanceType SVS_UNUSED(distance_type), + size_t blocksize_bytes +) { + svs::index::vamana::MemoryBreakdown breakdown{}; + // Graph: SimpleBlockedData with num_vectors rows and (max_degree + 1) + // cols; the +1 slot stores the per-node neighbor count. + using index_type = uint32_t; + const size_t max_degree = build_params.graph_max_degree; + + // TODO Fix/refactor DynamicVamana index builder to use proper allocator type and + // blocking parameters for graph, so that the memory estimate can be accurate for + // blocked data. For now, we use the default blocking parameters. + // There is MutableVamanaIndex deduction guides for index building defined in + // dynamic_index.h which set SimpleBlockedGraph as default graph type. + using graph_type = graphs::SimpleBlockedGraph; + using graph_data_type = typename graph_type::data_type; + using allocator_type = graph_data_type::allocator_type; + using graph_builder_type = svs::SimpleDataBuilder; + + breakdown.graph_bytes = + graph_builder_type{}.estimate_size(num_vectors, (max_degree + 1)); + + breakdown.data_bytes = + estimate_data_size_blocked(storage, num_vectors, dimension, blocksize_bytes); + + // Metadata: single entry point held as Idx, plus the SlotMetadata vector, plus the + // IDTranslator maps. + size_t metadata_bytes = + sizeof(index_type) + sizeof(svs::index::vamana::SlotMetadata) * num_vectors; + // The IDTranslator holds two tsl::robin_map instances (external->internal and + // internal->external), neither of which exposes its allocated byte count. We + // approximate the storage as the id pair held in each of the two directions. This + // ignores the maps' load-factor slack and control bytes, so it is an estimate of + // the hash-map overhead that is accurate to within a few percent. + metadata_bytes += + 2 * num_vectors * + (sizeof(IDTranslator::external_id_type) + sizeof(IDTranslator::internal_id_type)); + breakdown.metadata_bytes = metadata_bytes; + return breakdown; +} } // namespace svs::c_runtime diff --git a/bindings/c/src/dispatcher_dynamic_vamana.hpp b/bindings/c/src/dispatcher_dynamic_vamana.hpp index 41ac71da..8994eca4 100644 --- a/bindings/c/src/dispatcher_dynamic_vamana.hpp +++ b/bindings/c/src/dispatcher_dynamic_vamana.hpp @@ -49,4 +49,13 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_load( size_t blocksize_bytes ); +svs::index::vamana::MemoryBreakdown dispatch_dynamic_vamana_memory_estimate( + const svs::index::vamana::VamanaBuildParameters& build_params, + size_t num_vectors, + size_t dimension, + const Storage* storage, + svs::DistanceType distance_type, + size_t blocksize_bytes +); + } // namespace svs::c_runtime diff --git a/bindings/c/src/dispatcher_vamana.cpp b/bindings/c/src/dispatcher_vamana.cpp index 1c9b873b..47ef0340 100644 --- a/bindings/c/src/dispatcher_vamana.cpp +++ b/bindings/c/src/dispatcher_vamana.cpp @@ -129,4 +129,27 @@ svs::Vamana dispatch_vamana_index_load( build_params, VamanaSource{directory}, storage, distance_type, std::move(pool) ); } + +svs::index::vamana::MemoryBreakdown dispatch_vamana_memory_estimate( + const svs::index::vamana::VamanaBuildParameters& build_params, + size_t num_vectors, + size_t dimension, + const Storage* storage, + svs::DistanceType SVS_UNUSED(distance_type) +) { + svs::index::vamana::MemoryBreakdown breakdown{}; + + // Graph: SimpleData with num_vectors rows and (max_degree + 1) cols; + // the +1 slot stores the per-node neighbor count. + using index_type = uint32_t; + const size_t max_degree = build_params.graph_max_degree; + auto graph_data_builder = SimpleDataBuilder{}; + breakdown.graph_bytes = graph_data_builder.estimate_size(num_vectors, (max_degree + 1)); + + // Data: SimpleData with num_vectors rows and `dimension` cols. + breakdown.data_bytes = estimate_data_size(storage, num_vectors, dimension); + // Metadata: single entry point held as Idx. + breakdown.metadata_bytes = sizeof(index_type); + return breakdown; +} } // namespace svs::c_runtime diff --git a/bindings/c/src/dispatcher_vamana.hpp b/bindings/c/src/dispatcher_vamana.hpp index 90174dfe..457c77d7 100644 --- a/bindings/c/src/dispatcher_vamana.hpp +++ b/bindings/c/src/dispatcher_vamana.hpp @@ -44,4 +44,12 @@ svs::Vamana dispatch_vamana_index_load( svs::threads::ThreadPoolHandle pool ); +svs::index::vamana::MemoryBreakdown dispatch_vamana_memory_estimate( + const svs::index::vamana::VamanaBuildParameters& build_params, + size_t num_vectors, + size_t dimension, + const Storage* storage, + svs::DistanceType distance_type +); + } // namespace svs::c_runtime diff --git a/bindings/c/src/index_builder.hpp b/bindings/c/src/index_builder.hpp index 4b221397..322a9806 100644 --- a/bindings/c/src/index_builder.hpp +++ b/bindings/c/src/index_builder.hpp @@ -18,6 +18,7 @@ #include "svs/c_api/svs_c.h" #include "algorithm.hpp" +#include "data_builder.hpp" #include "dispatcher_dynamic_vamana.hpp" #include "dispatcher_vamana.hpp" #include "index.hpp" @@ -29,6 +30,8 @@ #include #include #include +#include +#include #include #include @@ -152,5 +155,44 @@ struct IndexBuilder { } return nullptr; } + + // Estimate the memory a built static Vamana index would consume + // for `num_vectors` vectors. Mirrors the accounting done by + // svs::index::vamana::VamanaIndex::get_memory_breakdown(). + svs::index::vamana::MemoryBreakdown estimate_memory_breakdown(size_t num_vectors + ) const { + NOT_IMPLEMENTED_IF( + algorithm->type != SVS_ALGORITHM_TYPE_VAMANA, + "Memory estimation is currently supported only for Vamana algorithm" + ); + auto vamana_algorithm = std::static_pointer_cast(algorithm); + return dispatch_vamana_memory_estimate( + vamana_algorithm->build_parameters(), + num_vectors, + dimension, + storage.get(), + to_distance_type(distance_metric) + ); + } + + // Estimate the memory a built dynamic Vamana index would consume + // for `num_vectors` vectors. Mirrors the accounting done by + // svs::index::vamana::MutableVamanaIndex::get_memory_breakdown(). + svs::index::vamana::MemoryBreakdown + estimate_memory_breakdown_dynamic(size_t num_vectors, size_t blocksize_bytes) const { + NOT_IMPLEMENTED_IF( + algorithm->type != SVS_ALGORITHM_TYPE_VAMANA, + "Memory estimation is currently supported only for Vamana algorithm" + ); + auto vamana_algorithm = std::static_pointer_cast(algorithm); + return dispatch_dynamic_vamana_memory_estimate( + vamana_algorithm->build_parameters(), + num_vectors, + dimension, + storage.get(), + to_distance_type(distance_metric), + blocksize_bytes + ); + } }; } // namespace svs::c_runtime diff --git a/bindings/c/src/svs_c.cpp b/bindings/c/src/svs_c.cpp index 40baf635..ef664e56 100644 --- a/bindings/c/src/svs_c.cpp +++ b/bindings/c/src/svs_c.cpp @@ -423,6 +423,56 @@ extern "C" bool svs_index_builder_set_threadpool_custom( ); } +extern "C" bool svs_index_builder_estimate_memory( + svs_index_builder_h builder, + size_t num_vectors, + svs_memory_breakdown_t* out_breakdown, + svs_error_h out_err +) { + using namespace svs::c_runtime; + return wrap_exceptions( + [&]() { + EXPECT_ARG_NOT_NULL(builder); + EXPECT_ARG_NOT_NULL(out_breakdown); + EXPECT_ARG_GT_THAN(num_vectors, 0); + auto breakdown = builder->impl->estimate_memory_breakdown(num_vectors); + out_breakdown->graph_bytes = breakdown.graph_bytes; + out_breakdown->data_bytes = breakdown.data_bytes; + out_breakdown->metadata_bytes = breakdown.metadata_bytes; + return true; + }, + out_err, + false + ); +} + +extern "C" bool svs_index_builder_estimate_memory_dynamic( + svs_index_builder_h builder, + size_t num_vectors, + size_t blocksize_bytes, + svs_memory_breakdown_t* out_breakdown, + svs_error_h out_err +) { + using namespace svs::c_runtime; + return wrap_exceptions( + [&]() { + EXPECT_ARG_NOT_NULL(builder); + EXPECT_ARG_NOT_NULL(out_breakdown); + EXPECT_ARG_GT_THAN(num_vectors, 0); + EXPECT_ARG_GT_THAN(blocksize_bytes, 0); + auto breakdown = builder->impl->estimate_memory_breakdown_dynamic( + num_vectors, blocksize_bytes + ); + out_breakdown->graph_bytes = breakdown.graph_bytes; + out_breakdown->data_bytes = breakdown.data_bytes; + out_breakdown->metadata_bytes = breakdown.metadata_bytes; + return true; + }, + out_err, + false + ); +} + extern "C" svs_index_h svs_index_build( svs_index_builder_h builder, const float* data, size_t num_vectors, svs_error_h out_err ) { diff --git a/bindings/c/src/types_support.hpp b/bindings/c/src/types_support.hpp index 5b261324..b7bc8daa 100644 --- a/bindings/c/src/types_support.hpp +++ b/bindings/c/src/types_support.hpp @@ -97,5 +97,22 @@ struct IDFilterAdapter : public IDFilterInterface { } }; +template +size_t +adjust_blocked_size(size_t num_vectors, size_t element_size, const Alloc& allocator) { + if constexpr (svs::data::is_blocked_v) { + // If using blocked allocator, account for block size overhead + // following the same logic as in SimpleData .ctor for Blocked allocators + assert(element_size > 0); + const auto blocksize = + lib::prevpow2(allocator.parameters().blocksize_bytes.value() / element_size); + size_t elements_per_block = blocksize.value(); + size_t num_blocks = lib::div_round_up(num_vectors, elements_per_block); + return num_blocks * blocksize.value() * element_size; + } else { + return num_vectors * element_size; + } +} + } // namespace c_runtime } // namespace svs diff --git a/bindings/c/tests/c_api_dynamic_index.cpp b/bindings/c/tests/c_api_dynamic_index.cpp index bc1a9565..5059a49a 100644 --- a/bindings/c/tests/c_api_dynamic_index.cpp +++ b/bindings/c/tests/c_api_dynamic_index.cpp @@ -354,6 +354,37 @@ CATCH_TEST_CASE("C API Dynamic Index", "[c_api][index][dynamic]") { svs_index_free(loaded_index); svs_index_free(index); } +} + +CATCH_TEST_CASE("C API Dynamic Index Memory", "[c_api][index][memory][dynamic]") { + // TODO: fix the blocked memory breakdown reported by index for LVQ, LeanVec storages. + // For now, we will: + // * test only the default simple and SQ storages. + // * Align graph and data sizes to BLOCK_SIZE to avoid test failures. + const size_t BLOCK_SIZE = 8 * 1024; // 8 KB block size for testing + const size_t DIMENSION = 32; + const size_t GRAPH_DEGREE = 16; + const size_t NUM_VECTORS = BLOCK_SIZE / DIMENSION; // full blocks of data + const size_t K = 5; + + std::vector data; + std::vector ids(NUM_VECTORS); + generate_test_data(data, NUM_VECTORS, DIMENSION); + + // Generate sequential IDs + for (size_t i = 0; i < NUM_VECTORS; ++i) { + ids[i] = i; + } + + svs_error_h error = svs_error_create(); + + svs_algorithm_h algorithm = svs_algorithm_create_vamana(GRAPH_DEGREE, 100, 100, error); + CATCH_REQUIRE(algorithm != nullptr); + + svs_index_builder_h builder = svs_index_builder_create( + SVS_DISTANCE_METRIC_EUCLIDEAN, DIMENSION, algorithm, error + ); + CATCH_REQUIRE(builder != nullptr); CATCH_SECTION("Memory Accounting Functions") { // Build dynamic index @@ -365,7 +396,7 @@ CATCH_TEST_CASE("C API Dynamic Index", "[c_api][index][dynamic]") { // Test get_memory_usage size_t memory_usage = 0; - success = svs_index_get_memory_usage(index, &memory_usage, error); + bool success = svs_index_get_memory_usage(index, &memory_usage, error); CATCH_REQUIRE(success); CATCH_REQUIRE(svs_error_ok(error)); CATCH_REQUIRE(memory_usage > 0); @@ -387,6 +418,101 @@ CATCH_TEST_CASE("C API Dynamic Index", "[c_api][index][dynamic]") { svs_index_free(index); } + CATCH_SECTION("Estimate Memory vs Actual Breakdown") { + // Build a dynamic index and compare its actual memory breakdown against + // the pre-build estimate produced by + // svs_index_builder_estimate_memory_dynamic(). `storage` may be nullptr + // to exercise the default (simple float32) storage. + auto estimate_and_verify = [&](svs_storage_h storage) { + svs_algorithm_h local_algorithm = + svs_algorithm_create_vamana(16, 32, 50, error); + CATCH_REQUIRE(local_algorithm != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + svs_index_builder_h local_builder = svs_index_builder_create( + SVS_DISTANCE_METRIC_EUCLIDEAN, DIMENSION, local_algorithm, error + ); + CATCH_REQUIRE(local_builder != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + bool ok = svs_index_builder_set_threadpool( + local_builder, SVS_THREADPOOL_KIND_NATIVE, 4, error + ); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(svs_error_ok(error)); + + if (storage != nullptr) { + ok = svs_index_builder_set_storage(local_builder, storage, error); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(svs_error_ok(error)); + } + + // Estimate before build. + svs_memory_breakdown_t estimated{}; + ok = svs_index_builder_estimate_memory_dynamic( + local_builder, NUM_VECTORS, BLOCK_SIZE, &estimated, error + ); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(svs_error_ok(error)); + CATCH_REQUIRE(estimated.graph_bytes > 0); + CATCH_REQUIRE(estimated.data_bytes > 0); + CATCH_REQUIRE(estimated.metadata_bytes > 0); + + // Build the dynamic index and query the actual breakdown. + svs_index_h index = svs_index_build_dynamic( + local_builder, data.data(), ids.data(), NUM_VECTORS, BLOCK_SIZE, error + ); + CATCH_REQUIRE(index != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + svs_memory_breakdown_t actual{}; + ok = svs_index_get_memory_breakdown(index, &actual, error); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(svs_error_ok(error)); + + // Allow up to 1% deviation between the pre-build estimate and the + // actual allocation (compressed storages may add small per-dataset + // overhead not accounted for by the estimator, and vice versa). + auto within_1pct = [](size_t estimate, size_t actual_val) { + if (estimate == actual_val) { + return true; + } + const auto [smaller, larger] = std::minmax(estimate, actual_val); + return (larger - smaller) * 100 <= larger; + }; + CATCH_REQUIRE(within_1pct(estimated.graph_bytes, actual.graph_bytes)); + CATCH_REQUIRE(within_1pct(estimated.data_bytes, actual.data_bytes)); + CATCH_REQUIRE(within_1pct(estimated.metadata_bytes, actual.metadata_bytes)); + + svs_index_free(index); + svs_index_builder_free(local_builder); + svs_algorithm_free(local_algorithm); + }; + + // Default storage (simple float32). + estimate_and_verify(nullptr); + + // Simple float16 storage. + { + svs_storage_h storage = svs_storage_create_simple(SVS_DATA_TYPE_FLOAT16, error); + CATCH_REQUIRE(check_storage_support(storage, error) == true); + if (storage != nullptr) { + estimate_and_verify(storage); + svs_storage_free(storage); + } + } + + // Scalar quantization storage. + { + svs_storage_h storage = svs_storage_create_sq(SVS_DATA_TYPE_INT8, error); + CATCH_REQUIRE(check_storage_support(storage, error) == true); + if (storage != nullptr) { + estimate_and_verify(storage); + svs_storage_free(storage); + } + } + } + svs_index_builder_free(builder); svs_algorithm_free(algorithm); svs_error_free(error); diff --git a/bindings/c/tests/c_api_index.cpp b/bindings/c/tests/c_api_index.cpp index d9ab2446..a4be8ed7 100644 --- a/bindings/c/tests/c_api_index.cpp +++ b/bindings/c/tests/c_api_index.cpp @@ -765,6 +765,14 @@ CATCH_TEST_CASE("C API Threadpool Management", "[c_api][index][threadpool]") { svs_algorithm_free(algorithm); svs_error_free(error); } +} + +CATCH_TEST_CASE("C API Index Memory", "[c_api][index][memory]") { + const size_t NUM_VECTORS = 1000; + const size_t DIMENSION = 32; + + std::vector data; + generate_test_data(data, NUM_VECTORS, DIMENSION); CATCH_SECTION("Memory Accounting Functions") { svs_error_h error = svs_error_create(); @@ -839,6 +847,136 @@ CATCH_TEST_CASE("C API Threadpool Management", "[c_api][index][threadpool]") { svs_algorithm_free(algorithm); svs_error_free(error); } + + CATCH_SECTION("Estimate Memory vs Actual Breakdown") { + svs_error_h error = svs_error_create(); + + // Build an index and compare its actual memory breakdown against the + // pre-build estimate produced by svs_index_builder_estimate_memory(). + // `storage` may be nullptr to exercise the default (simple float32) storage. + auto estimate_and_verify = [&](svs_storage_h storage) { + svs_algorithm_h algorithm = svs_algorithm_create_vamana(16, 32, 50, error); + CATCH_REQUIRE(algorithm != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + svs_index_builder_h builder = svs_index_builder_create( + SVS_DISTANCE_METRIC_EUCLIDEAN, DIMENSION, algorithm, error + ); + CATCH_REQUIRE(builder != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + bool success = svs_index_builder_set_threadpool( + builder, SVS_THREADPOOL_KIND_NATIVE, 4, error + ); + CATCH_REQUIRE(success); + CATCH_REQUIRE(svs_error_ok(error)); + + if (storage != nullptr) { + success = svs_index_builder_set_storage(builder, storage, error); + CATCH_REQUIRE(success); + CATCH_REQUIRE(svs_error_ok(error)); + } + + // Estimate before build. + svs_memory_breakdown_t estimated{}; + success = + svs_index_builder_estimate_memory(builder, NUM_VECTORS, &estimated, error); + CATCH_REQUIRE(success); + CATCH_REQUIRE(svs_error_ok(error)); + CATCH_REQUIRE(estimated.graph_bytes > 0); + CATCH_REQUIRE(estimated.data_bytes > 0); + CATCH_REQUIRE(estimated.metadata_bytes > 0); + + // Build the index and query the actual breakdown. + svs_index_h index = svs_index_build(builder, data.data(), NUM_VECTORS, error); + CATCH_REQUIRE(index != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + svs_memory_breakdown_t actual{}; + success = svs_index_get_memory_breakdown(index, &actual, error); + CATCH_REQUIRE(success); + CATCH_REQUIRE(svs_error_ok(error)); + + // Allow up to 1% deviation between the pre-build estimate and the + // actual allocation (compressed storages may add small per-dataset + // overhead not accounted for by the estimator, and vice versa). + auto within_1pct = [](size_t estimate, size_t actual_val) { + if (estimate == actual_val) { + return true; + } + const auto [smaller, larger] = std::minmax(estimate, actual_val); + return (larger - smaller) * 100 <= larger; + }; + CATCH_REQUIRE(within_1pct(estimated.graph_bytes, actual.graph_bytes)); + CATCH_REQUIRE(estimated.data_bytes == actual.data_bytes); + CATCH_REQUIRE(within_1pct(estimated.data_bytes, actual.data_bytes)); + CATCH_REQUIRE(within_1pct(estimated.metadata_bytes, actual.metadata_bytes)); + + svs_index_free(index); + svs_index_builder_free(builder); + svs_algorithm_free(algorithm); + }; + + // Default storage (simple float32). + estimate_and_verify(nullptr); + + // Simple float16 storage. + { + svs_storage_h storage = svs_storage_create_simple(SVS_DATA_TYPE_FLOAT16, error); + CATCH_REQUIRE(check_storage_support(storage, error) == true); + if (storage != nullptr) { + estimate_and_verify(storage); + svs_storage_free(storage); + } + } + + // Scalar quantization storage + { + svs_storage_h storage = svs_storage_create_sq(SVS_DATA_TYPE_INT8, error); + CATCH_REQUIRE(check_storage_support(storage, error) == true); + if (storage != nullptr) { + estimate_and_verify(storage); + svs_storage_free(storage); + } + } + + // LVQ: primary = int4, residual = int8. + { + svs_storage_h storage = + svs_storage_create_lvq(SVS_DATA_TYPE_INT4, SVS_DATA_TYPE_INT8, error); + CATCH_REQUIRE(check_storage_support(storage, error) == true); + if (storage != nullptr) { + estimate_and_verify(storage); + svs_storage_free(storage); + } + } + + // LeanVec: leanvec_dims = DIMENSION / 2, primary = int4, secondary = int8. + { + svs_storage_h storage = svs_storage_create_leanvec( + DIMENSION / 2, SVS_DATA_TYPE_INT4, SVS_DATA_TYPE_INT8, error + ); + CATCH_REQUIRE(check_storage_support(storage, error) == true); + if (storage != nullptr) { + estimate_and_verify(storage); + svs_storage_free(storage); + } + } + + // LeanVec: leanvec_dims = DIMENSION / 2, primary = int4, secondary = int4. + { + svs_storage_h storage = svs_storage_create_leanvec( + DIMENSION / 2, SVS_DATA_TYPE_INT4, SVS_DATA_TYPE_INT4, error + ); + CATCH_REQUIRE(check_storage_support(storage, error) == true); + if (storage != nullptr) { + estimate_and_verify(storage); + svs_storage_free(storage); + } + } + + svs_error_free(error); + } } namespace {