diff --git a/bindings/c/include/svs/c/svs_c.h b/bindings/c/include/svs/c/svs_c.h index 9a7ae49bc..5692b3b69 100644 --- a/bindings/c/include/svs/c/svs_c.h +++ b/bindings/c/include/svs/c/svs_c.h @@ -189,6 +189,65 @@ struct svs_threadpool_interface { void* self; }; +/// @brief Allocators used for memory management in the SVS C API. +enum svs_allocator_kind { + SVS_ALLOCATOR_KIND_DEFAULT = 0, + SVS_ALLOCATOR_KIND_HUGE_PAGE = 1, + SVS_ALLOCATOR_KIND_CUSTOM = 2 +}; + +/// @brief Operations table for a custom allocator interface +/// @remarks The user must ensure that the allocator implementation is thread-safe and +/// that the provided function pointers remain valid for the lifetime of the allocator +/// interface. +/// @var svs_allocator_interface_ops::version +/// Version of the allocator interface. +/// @var svs_allocator_interface_ops::struct_size +/// Size of the structure, used for versioning and compatibility checks. +/// @var svs_allocator_interface_ops::allocate +/// Function pointer to allocate memory with the specified size and alignment. +/// @param self Pointer to the allocator instance. +/// @param size Size of the memory to allocate in bytes. +/// @param alignment Alignment requirement for the allocated memory in bytes. +/// @param out_err Handle to capture any error that occurs during allocation. User code +/// may call svs_error_set() to set the error code and message if an error occurs. +/// @return Pointer to the allocated memory, or NULL if allocation fails. +/// @var svs_allocator_interface_ops::deallocate +/// Function pointer to deallocate memory previously allocated by the allocator. +/// @param self Pointer to the allocator instance. +/// @param ptr Pointer to the memory to deallocate. +/// @param size Size of the memory to deallocate in bytes. +/// @param alignment Alignment requirement for the allocated memory in bytes. +struct svs_allocator_interface_ops { + uint32_t version; + size_t struct_size; + void* (*allocate)(void* self, size_t size, size_t alignment, svs_error_h out_err); + void (*deallocate)(void* self, void* ptr, size_t size, size_t alignment); +}; + +/// @brief Macro to create a user-defined allocator interface operations structure +/// @param allocate_func Function pointer to allocate memory with the specified size and +/// alignment +/// @param deallocate_func Function pointer to deallocate memory previously allocated by the +/// allocator +#define SVS_INIT_ALLOCATOR_OPS(allocate_func, deallocate_func) \ + { \ + .version = SVS_C_API_VERSION, \ + .struct_size = sizeof(struct svs_allocator_interface_ops), \ + .allocate = &allocate_func, .deallocate = &deallocate_func \ + } + +/// @brief Structure representing a custom allocator interface +/// @var svs_allocator_interface::ops +/// Function pointers for the allocator operations. +/// @var svs_allocator_interface::self +/// Pointer to the user-defined allocator instance. This pointer is passed to the +/// function pointers in @p ops when they are called. +struct svs_allocator_interface { + struct svs_allocator_interface_ops* ops; + void* self; +}; + /// @brief Operations table for a custom ID filter interface /// @remarks The user must ensure that the ID filter implementation is thread-safe and /// that the provided function pointers remain valid for the lifetime of the ID filter @@ -424,10 +483,14 @@ typedef enum svs_algorithm_type svs_algorithm_type_t; typedef enum svs_data_type svs_data_type_t; typedef enum svs_storage_kind svs_storage_kind_t; typedef enum svs_threadpool_kind svs_threadpool_kind_t; +typedef enum svs_allocator_kind svs_allocator_kind_t; typedef struct svs_threadpool_interface_ops svs_threadpool_ops_t; typedef struct svs_threadpool_interface svs_threadpool_t; typedef struct svs_threadpool_interface* svs_threadpool_i; +typedef struct svs_allocator_interface_ops svs_allocator_ops_t; +typedef struct svs_allocator_interface svs_allocator_t; +typedef struct svs_allocator_interface* svs_allocator_i; typedef struct svs_id_filter_interface_ops svs_id_filter_ops_t; typedef struct svs_id_filter_interface svs_id_filter_t; @@ -696,6 +759,28 @@ 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 Set the allocator configuration for the index builder +/// @param builder The index builder handle +/// @param kind The kind of allocator to use +/// @param out_err An optional error handle to capture errors +/// @return true on success, false on failure +SVS_API bool svs_index_builder_set_allocator( + svs_index_builder_h builder, svs_allocator_kind_t kind, svs_error_h out_err /*=NULL*/ +); + +/// @brief Set the custom allocator for the index builder +/// @param builder The index builder handle +/// @param allocator The custom allocator interface +/// @param out_err An optional error handle to capture errors +/// @return true on success, false on failure +/// @remarks The builder copies @p allocator and its ops table by value, so those two +/// objects may be freed or modified after this call returns. The object referenced by +/// @p allocator->self is not copied and must outlive the builder and every index built or +/// loaded with it. +SVS_API bool svs_index_builder_set_allocator_custom( + svs_index_builder_h builder, svs_allocator_i allocator, 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 diff --git a/bindings/c/src/allocator.hpp b/bindings/c/src/allocator.hpp index b6859945d..b8a1889d8 100644 --- a/bindings/c/src/allocator.hpp +++ b/bindings/c/src/allocator.hpp @@ -15,8 +15,15 @@ */ #pragma once +#include "svs/c/svs_c.h" + +#include "error.hpp" + +#include #include +#include #include +#include namespace svs { namespace c_runtime { @@ -25,5 +32,79 @@ template , Allocator>; +template class CustomAllocator : public svs::AllocatorInterface { + public: + using value_type = T; + + static void validate(const svs_allocator_i allocator) { + if (allocator == nullptr) { + throw std::invalid_argument("Custom allocator pointer cannot be null."); + } + if (allocator->ops == nullptr) { + throw std::invalid_argument("Custom allocator interface is not initialized."); + } + if (allocator->ops->version > svs_get_version()) { + throw std::invalid_argument( + "Custom allocator interface version is not supported." + ); + } + if (allocator->ops->struct_size < sizeof(svs_allocator_ops_t)) { + throw std::invalid_argument( + "Incompatible custom allocator interface struct size." + ); + } + if (allocator->ops->allocate == nullptr || allocator->ops->deallocate == nullptr) { + throw std::invalid_argument( + "Custom allocator interface has null function pointers." + ); + } + } + + CustomAllocator(const svs_allocator_ops_t& ops, void* self) + : ops_(ops) + , self_(self) {} + + void* allocate(size_t n) override { + svs_error_desc err{SVS_ERROR_UNKNOWN, "Unknown error in custom allocator allocate"}; + + auto result = ops_.allocate(self_, n * sizeof(T), alignof(T), &err); + if (result == nullptr) { + throw svs::c_runtime::out_of_memory( + "Custom allocator failed to allocate memory: (" + std::to_string(err.code) + + ") " + err.message + ); + } + return result; + } + + void deallocate(void* p, size_t n) override { + ops_.deallocate(self_, p, n * sizeof(T), alignof(T)); + } + + AllocatorInterface* clone() const override { + return new CustomAllocator(ops_, self_); + } + + AllocatorInterface* rebind_to(DataType type) const override { + return svs::lib::match( + AllocatorInterface::rebind_types{}, + type, + [this](svs::lib::Type) -> AllocatorInterface* { + return new CustomAllocator(ops_, self_); + } + ); + } + + private: + svs_allocator_ops_t ops_; + void* self_; +}; + +template +AllocatorHandle make_custom_allocator_handle(const svs_allocator_i allocator) { + CustomAllocator::validate(allocator); + return AllocatorHandle{ + std::make_unique>(*allocator->ops, allocator->self)}; +} } // namespace c_runtime } // namespace svs diff --git a/bindings/c/src/data_builder/leanvec.hpp b/bindings/c/src/data_builder/leanvec.hpp index 0cd5c1a1f..01263e1f8 100644 --- a/bindings/c/src/data_builder/leanvec.hpp +++ b/bindings/c/src/data_builder/leanvec.hpp @@ -140,7 +140,8 @@ struct lib:: }; template void for_leanvec_specializations(F&& f) { - using byte_alloc = svs::c_runtime::MaybeBlockedAlloc; + using byte_alloc = svs::c_runtime:: + MaybeBlockedAlloc>; #define X(P, S, D) f.template operator(), D>(); #define XX(P, S) X(P, S, DistanceL2) X(P, S, DistanceIP) X(P, S, DistanceCosineSimilarity) diff --git a/bindings/c/src/data_builder/lvq.hpp b/bindings/c/src/data_builder/lvq.hpp index 92ea32f9f..c477849c7 100644 --- a/bindings/c/src/data_builder/lvq.hpp +++ b/bindings/c/src/data_builder/lvq.hpp @@ -153,7 +153,8 @@ struct lib::DispatchConverter< }; template void for_lvq_specializations(F&& f) { - using byte_alloc = svs::c_runtime::MaybeBlockedAlloc; + using byte_alloc = svs::c_runtime:: + MaybeBlockedAlloc>; #define X(P, S, D) f.template operator(), D>(); #define XX(P, S) X(P, S, DistanceL2) X(P, S, DistanceIP) X(P, S, DistanceCosineSimilarity) // Pattern: diff --git a/bindings/c/src/data_builder/simple.hpp b/bindings/c/src/data_builder/simple.hpp index b5c6c38ac..e36038031 100644 --- a/bindings/c/src/data_builder/simple.hpp +++ b/bindings/c/src/data_builder/simple.hpp @@ -91,8 +91,10 @@ struct lib::DispatchConverter void for_simple_specializations(F&& f) { - using float_alloc = svs::c_runtime::MaybeBlockedAlloc; - using float16_alloc = svs::c_runtime::MaybeBlockedAlloc; + using float_alloc = + svs::c_runtime::MaybeBlockedAlloc>; + using float16_alloc = svs::c_runtime:: + MaybeBlockedAlloc>; #define X(T, A, D) f.template operator(), D>(); #define XX(T, A) X(T, A, DistanceL2) X(T, A, DistanceIP) X(T, A, DistanceCosineSimilarity) XX(float, float_alloc) diff --git a/bindings/c/src/data_builder/sq.hpp b/bindings/c/src/data_builder/sq.hpp index a5c9bd3c0..4abd2ee4f 100644 --- a/bindings/c/src/data_builder/sq.hpp +++ b/bindings/c/src/data_builder/sq.hpp @@ -93,8 +93,10 @@ struct lib::DispatchConverter }; template void for_sq_specializations(F&& f) { - using int8_alloc = svs::c_runtime::MaybeBlockedAlloc; - using uint8_alloc = svs::c_runtime::MaybeBlockedAlloc; + using int8_alloc = + svs::c_runtime::MaybeBlockedAlloc>; + using uint8_alloc = + svs::c_runtime::MaybeBlockedAlloc>; #define X(T, A, D) f.template operator(), D>(); #define XX(T, A) X(T, A, DistanceL2) X(T, A, DistanceIP) X(T, A, DistanceCosineSimilarity) XX(uint8_t, uint8_alloc) diff --git a/bindings/c/src/dispatcher_dynamic_vamana.cpp b/bindings/c/src/dispatcher_dynamic_vamana.cpp index 283b1f756..1001d5ab7 100644 --- a/bindings/c/src/dispatcher_dynamic_vamana.cpp +++ b/bindings/c/src/dispatcher_dynamic_vamana.cpp @@ -45,6 +45,7 @@ svs::DynamicVamana build_dynamic_vamana_index( DataBuilder builder, Distance D, svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle, size_t blocksize_bytes ) { svs::data::BlockingParameters block_params; @@ -52,14 +53,21 @@ svs::DynamicVamana build_dynamic_vamana_index( block_params.blocksize_bytes = svs::lib::prevpow2(blocksize_bytes); } using allocator_type = typename DataBuilder::allocator_type; - auto allocator = allocator_type{block_params}; - auto data = builder.build(std::move(src_data.first), pool, allocator); + using value_type = typename allocator_type::value_type; + + auto data_allocator_handle = svs::lib::rebind_allocator(allocator_handle); + auto data_allocator = allocator_type{block_params, data_allocator_handle}; + auto data = builder.build(std::move(src_data.first), pool, data_allocator); + + auto graph_allocator_handle = svs::lib::rebind_allocator(allocator_handle); + auto graph_allocator = svs::data::Blocked{block_params, graph_allocator_handle}; return svs::DynamicVamana::build( build_params, std::move(data), std::move(src_data.second), std::move(D), - std::move(pool) + std::move(pool), + graph_allocator ); } @@ -70,6 +78,7 @@ svs::DynamicVamana load_dynamic_vamana_index( DataLoader loader, Distance D, svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle, size_t blocksize_bytes ) { svs::data::BlockingParameters block_params; @@ -77,11 +86,17 @@ svs::DynamicVamana load_dynamic_vamana_index( block_params.blocksize_bytes = svs::lib::prevpow2(blocksize_bytes); } using allocator_type = typename DataLoader::allocator_type; - auto allocator = allocator_type{block_params}; + using value_type = typename allocator_type::value_type; + auto data_allocator_handle = svs::lib::rebind_allocator(allocator_handle); + auto allocator = allocator_type{block_params, data_allocator_handle}; auto data = loader.load(directory / "data", allocator); + + auto graph_allocator_handle = svs::lib::rebind_allocator(allocator_handle); + auto graph_allocator = svs::data::Blocked{block_params, graph_allocator_handle}; + return svs::DynamicVamana::assemble( directory / "config", - svs::GraphLoader{directory / "graph"}, + svs::GraphLoader{directory / "graph", graph_allocator}, std::move(data), std::move(D), std::move(pool) @@ -118,6 +133,7 @@ using BuildDynamicIndexDispatcher = svs::lib::Dispatcher< const Storage*, svs::DistanceType, svs::threads::ThreadPoolHandle, + const AllocatorHandle&, size_t>; const BuildDynamicIndexDispatcher& build_dynamic_vamana_index_dispatcher() { @@ -136,6 +152,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_build( const Storage* storage, svs::DistanceType distance_type, svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle, size_t blocksize_bytes ) { return build_dynamic_vamana_index_dispatcher().invoke( @@ -144,6 +161,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_build( storage, distance_type, std::move(pool), + allocator_handle, blocksize_bytes ); } @@ -154,6 +172,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_load( const Storage* storage, svs::DistanceType distance_type, svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle, size_t blocksize_bytes ) { return build_dynamic_vamana_index_dispatcher().invoke( @@ -162,6 +181,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_load( storage, distance_type, std::move(pool), + allocator_handle, blocksize_bytes ); } @@ -175,24 +195,26 @@ svs::index::vamana::MemoryBreakdown dispatch_dynamic_vamana_memory_estimate( size_t blocksize_bytes ) { svs::index::vamana::MemoryBreakdown breakdown{}; + + // Graph size // 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; + using graph_allocator_type = svs::data::Blocked>; + 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; + svs::data::BlockingParameters block_params; + if (blocksize_bytes != 0) { + block_params.blocksize_bytes = svs::lib::prevpow2(blocksize_bytes); + } + auto graph_allocator = graph_allocator_type{block_params}; + auto graph_data_builder = svs::SimpleDataBuilder{}; breakdown.graph_bytes = - graph_builder_type{}.estimate_size(num_vectors, (max_degree + 1)); + graph_data_builder.estimate_size(num_vectors, (max_degree + 1), graph_allocator); + // Data size breakdown.data_bytes = estimate_data_size_blocked(storage, num_vectors, dimension, blocksize_bytes); diff --git a/bindings/c/src/dispatcher_dynamic_vamana.hpp b/bindings/c/src/dispatcher_dynamic_vamana.hpp index 8994eca4c..148ab2ac5 100644 --- a/bindings/c/src/dispatcher_dynamic_vamana.hpp +++ b/bindings/c/src/dispatcher_dynamic_vamana.hpp @@ -37,6 +37,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_build( const Storage* storage, svs::DistanceType distance_type, svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle, size_t blocksize_bytes ); @@ -46,6 +47,7 @@ svs::DynamicVamana dispatch_dynamic_vamana_index_load( const Storage* storage, svs::DistanceType distance_type, svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle, size_t blocksize_bytes ); diff --git a/bindings/c/src/dispatcher_vamana.cpp b/bindings/c/src/dispatcher_vamana.cpp index 47ef03402..4725e38f9 100644 --- a/bindings/c/src/dispatcher_vamana.cpp +++ b/bindings/c/src/dispatcher_vamana.cpp @@ -41,11 +41,19 @@ svs::Vamana build_vamana_index( svs::data::ConstSimpleDataView src_data, DataBuilder builder, Distance distance, - svs::threads::ThreadPoolHandle pool + svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle ) { - auto data = builder.build(std::move(src_data), pool); + using value_type = typename DataBuilder::allocator_type::value_type; + auto data = builder.build( + std::move(src_data), pool, svs::lib::rebind_allocator(allocator_handle) + ); return svs::Vamana::build( - build_params, std::move(data), distance, std::move(pool) + build_params, + std::move(data), + distance, + std::move(pool), + svs::lib::rebind_allocator(allocator_handle) ); } @@ -55,12 +63,17 @@ svs::Vamana load_vamana_index( const std::filesystem::path& directory, DataLoader loader, Distance distance, - svs::threads::ThreadPoolHandle pool + svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle ) { - auto data = loader.load(directory / "data"); + using value_type = typename DataLoader::allocator_type::value_type; + auto data = loader.load( + directory / "data", svs::lib::rebind_allocator(allocator_handle) + ); return svs::Vamana::assemble( directory / "config", - svs::GraphLoader{directory / "graph"}, + svs::GraphLoader>{ + directory / "graph", svs::lib::rebind_allocator(allocator_handle)}, std::move(data), distance, std::move(pool) @@ -95,7 +108,8 @@ using BuildIndexDispatcher = svs::lib::Dispatcher< VamanaSource, const Storage*, svs::DistanceType, - svs::threads::ThreadPoolHandle>; + svs::threads::ThreadPoolHandle, + const AllocatorHandle&>; const BuildIndexDispatcher& build_vamana_index_dispatcher() { static BuildIndexDispatcher dispatcher = [] { @@ -111,10 +125,16 @@ svs::Vamana dispatch_vamana_index_build( svs::data::ConstSimpleDataView data, const Storage* storage, svs::DistanceType distance_type, - svs::threads::ThreadPoolHandle pool + svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle ) { return build_vamana_index_dispatcher().invoke( - build_params, VamanaSource{std::move(data)}, storage, distance_type, std::move(pool) + build_params, + VamanaSource{std::move(data)}, + storage, + distance_type, + std::move(pool), + allocator_handle ); } @@ -123,10 +143,16 @@ svs::Vamana dispatch_vamana_index_load( const std::filesystem::path& directory, const Storage* storage, svs::DistanceType distance_type, - svs::threads::ThreadPoolHandle pool + svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle ) { return build_vamana_index_dispatcher().invoke( - build_params, VamanaSource{directory}, storage, distance_type, std::move(pool) + build_params, + VamanaSource{directory}, + storage, + distance_type, + std::move(pool), + allocator_handle ); } diff --git a/bindings/c/src/dispatcher_vamana.hpp b/bindings/c/src/dispatcher_vamana.hpp index 457c77d7a..89adc4da6 100644 --- a/bindings/c/src/dispatcher_vamana.hpp +++ b/bindings/c/src/dispatcher_vamana.hpp @@ -33,7 +33,8 @@ svs::Vamana dispatch_vamana_index_build( svs::data::ConstSimpleDataView data, const Storage* storage, svs::DistanceType distance_type, - svs::threads::ThreadPoolHandle pool + svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle ); svs::Vamana dispatch_vamana_index_load( @@ -41,7 +42,8 @@ svs::Vamana dispatch_vamana_index_load( const std::filesystem::path& directory, const Storage* storage, svs::DistanceType distance_type, - svs::threads::ThreadPoolHandle pool + svs::threads::ThreadPoolHandle pool, + const AllocatorHandle& allocator_handle ); svs::index::vamana::MemoryBreakdown dispatch_vamana_memory_estimate( diff --git a/bindings/c/src/error.hpp b/bindings/c/src/error.hpp index a26ce0efc..94eae582e 100644 --- a/bindings/c/src/error.hpp +++ b/bindings/c/src/error.hpp @@ -97,6 +97,11 @@ class unsupported_hw : public std::runtime_error { using std::runtime_error::runtime_error; }; +class out_of_memory : public std::runtime_error { + public: + using std::runtime_error::runtime_error; +}; + // A helper to wrap C++ exceptions and convert them to C error codes/messages. template > Result wrap_exceptions(Callable&& func, svs_error_h err, Result err_res = {}) noexcept { @@ -115,6 +120,9 @@ Result wrap_exceptions(Callable&& func, svs_error_h err, Result err_res = {}) no } catch (const svs::c_runtime::unsupported_hw& ex) { SET_ERROR(err, SVS_ERROR_UNSUPPORTED_HW, ex.what()); return err_res; + } catch (const svs::c_runtime::out_of_memory& ex) { + SET_ERROR(err, SVS_ERROR_OUT_OF_MEMORY, ex.what()); + return err_res; } catch (const svs::lib::ANNException& ex) { SET_ERROR(err, SVS_ERROR_GENERIC, ex.what()); return err_res; diff --git a/bindings/c/src/index_builder.hpp b/bindings/c/src/index_builder.hpp index d9500a182..780034b1e 100644 --- a/bindings/c/src/index_builder.hpp +++ b/bindings/c/src/index_builder.hpp @@ -47,6 +47,7 @@ struct IndexBuilder { std::shared_ptr algorithm; std::shared_ptr storage; ThreadPoolBuilder pool_builder; + AllocatorHandle allocator_handle; IndexBuilder( svs_distance_metric_t distance_metric, @@ -57,7 +58,8 @@ struct IndexBuilder { , dimension(dimension) , algorithm(std::move(algorithm)) , storage(std::make_shared(SVS_DATA_TYPE_FLOAT32)) - , pool_builder{} {} + , pool_builder{} + , allocator_handle{make_allocator_handle(svs::lib::Allocator{})} {} ~IndexBuilder() {} @@ -69,6 +71,10 @@ struct IndexBuilder { std::swap(this->pool_builder, threadpool_builder); } + void set_allocator_handle(AllocatorHandle allocator_handle) { + this->allocator_handle = std::move(allocator_handle); + } + std::shared_ptr build(const svs::data::ConstSimpleDataView& data) { if (algorithm->type == SVS_ALGORITHM_TYPE_VAMANA) { auto vamana_algorithm = std::static_pointer_cast(algorithm); @@ -79,7 +85,8 @@ struct IndexBuilder { data, storage.get(), to_distance_type(distance_metric), - pool_builder.build() + pool_builder.build(), + allocator_handle ), pool_builder ); @@ -99,7 +106,8 @@ struct IndexBuilder { directory, storage.get(), to_distance_type(distance_metric), - pool_builder.build() + pool_builder.build(), + allocator_handle ), pool_builder ); @@ -125,6 +133,7 @@ struct IndexBuilder { storage.get(), to_distance_type(distance_metric), pool_builder.build(), + allocator_handle, blocksize_bytes ), pool_builder @@ -147,6 +156,7 @@ struct IndexBuilder { storage.get(), to_distance_type(distance_metric), pool_builder.build(), + allocator_handle, blocksize_bytes ), pool_builder diff --git a/bindings/c/src/svs_c.cpp b/bindings/c/src/svs_c.cpp index 1c705d014..8dd47efd7 100644 --- a/bindings/c/src/svs_c.cpp +++ b/bindings/c/src/svs_c.cpp @@ -17,6 +17,7 @@ #include "svs/c/svs_c.h" #include "algorithm.hpp" +#include "allocator.hpp" #include "error.hpp" #include "index.hpp" #include "index_builder.hpp" @@ -31,8 +32,10 @@ #include #include +#include #include #include +#include #include // C API implementation @@ -458,6 +461,54 @@ extern "C" bool svs_index_builder_set_threadpool_custom( ); } +SVS_API bool svs_index_builder_set_allocator( + svs_index_builder_h builder, svs_allocator_kind_t kind, svs_error_h out_err +) { + using namespace svs::c_runtime; + return wrap_exceptions( + [&]() { + EXPECT_ARG_NOT_NULL(builder); + switch (kind) { + case SVS_ALLOCATOR_KIND_DEFAULT: { + builder->impl->set_allocator_handle( + svs::make_allocator_handle(svs::lib::Allocator{}) + ); + break; + } + case SVS_ALLOCATOR_KIND_HUGE_PAGE: { + builder->impl->set_allocator_handle( + svs::make_allocator_handle(svs::HugepageAllocator{}) + ); + break; + } + default: + throw std::invalid_argument("Invalid allocator kind"); + } + return true; + }, + out_err, + false + ); +} + +SVS_API bool svs_index_builder_set_allocator_custom( + svs_index_builder_h builder, svs_allocator_i allocator, svs_error_h out_err +) { + using namespace svs::c_runtime; + return wrap_exceptions( + [&]() { + EXPECT_ARG_NOT_NULL(builder); + EXPECT_ARG_NOT_NULL(allocator); + builder->impl->set_allocator_handle( + make_custom_allocator_handle(allocator) + ); + return true; + }, + out_err, + false + ); +} + extern "C" bool svs_index_builder_estimate_memory( svs_index_builder_h builder, size_t num_vectors, diff --git a/bindings/c/tests/c_api_dynamic_index.cpp b/bindings/c/tests/c_api_dynamic_index.cpp index da4f14814..2adc74bd2 100644 --- a/bindings/c/tests/c_api_dynamic_index.cpp +++ b/bindings/c/tests/c_api_dynamic_index.cpp @@ -376,6 +376,7 @@ CATCH_TEST_CASE("C API Dynamic Index Memory", "[c_api][index][memory][dynamic]") const size_t GRAPH_DEGREE = 16; const size_t NUM_VECTORS = BLOCK_SIZE / DIMENSION; // full blocks of data const size_t K = 5; + const size_t NUM_THREADS = 4; std::vector data; std::vector ids(NUM_VECTORS); @@ -388,7 +389,7 @@ CATCH_TEST_CASE("C API Dynamic Index Memory", "[c_api][index][memory][dynamic]") svs_error_h error = svs_error_create(); - svs_algorithm_h algorithm = svs_algorithm_create_vamana(GRAPH_DEGREE, 100, 100, error); + svs_algorithm_h algorithm = svs_algorithm_create_vamana(GRAPH_DEGREE, 32, 50, error); CATCH_REQUIRE(algorithm != nullptr); svs_index_builder_h builder = svs_index_builder_create( @@ -396,6 +397,11 @@ CATCH_TEST_CASE("C API Dynamic Index Memory", "[c_api][index][memory][dynamic]") ); CATCH_REQUIRE(builder != nullptr); + bool success = svs_index_builder_set_threadpool( + builder, SVS_THREADPOOL_KIND_NATIVE, NUM_THREADS, error + ); + CATCH_REQUIRE(success); + CATCH_SECTION("Dynamic Index Memory Accounting") { // Build dynamic index svs_index_h index = svs_index_build_dynamic( @@ -458,7 +464,7 @@ CATCH_TEST_CASE("C API Dynamic Index Memory", "[c_api][index][memory][dynamic]") } // Estimate before build. - svs_memory_breakdown_t estimated{}; + svs_memory_breakdown_t estimated = SVS_INIT_MEMORY_BREAKDOWN(); ok = svs_index_builder_estimate_memory_dynamic( local_builder, NUM_VECTORS, BLOCK_SIZE, &estimated, error ); @@ -475,7 +481,7 @@ CATCH_TEST_CASE("C API Dynamic Index Memory", "[c_api][index][memory][dynamic]") CATCH_REQUIRE(index != nullptr); CATCH_REQUIRE(svs_error_ok(error)); - svs_memory_breakdown_t actual{}; + svs_memory_breakdown_t actual = SVS_INIT_MEMORY_BREAKDOWN(); ok = svs_index_get_memory_breakdown(index, &actual, error); CATCH_REQUIRE(ok); CATCH_REQUIRE(svs_error_ok(error)); @@ -644,6 +650,122 @@ CATCH_TEST_CASE("C API Dynamic Index Memory", "[c_api][index][memory][dynamic]") svs_search_params_free(search_params); } + CATCH_SECTION("Allocator Configuration") { + // Each built-in allocator kind must yield a usable dynamic index. + auto build_with_allocator = [&](svs_allocator_kind_t kind) { + svs_index_builder_h local_builder = svs_index_builder_create( + SVS_DISTANCE_METRIC_EUCLIDEAN, DIMENSION, algorithm, error + ); + CATCH_REQUIRE(local_builder != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + bool success = svs_index_builder_set_threadpool( + local_builder, SVS_THREADPOOL_KIND_NATIVE, NUM_THREADS, error + ); + CATCH_REQUIRE(success); + + bool ok = svs_index_builder_set_allocator(local_builder, kind, error); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(svs_error_ok(error)); + + // Use big blocksize to avoid many HugePage allocations. + svs_index_h index = svs_index_build_dynamic( + local_builder, data.data(), ids.data(), NUM_VECTORS, 1 << 30, error + ); + CATCH_REQUIRE(index != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + size_t memory_usage = 0; + ok = svs_index_get_memory_usage(index, &memory_usage, error); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(memory_usage > 0); + + svs_index_free(index); + svs_index_builder_free(local_builder); + }; + + build_with_allocator(SVS_ALLOCATOR_KIND_DEFAULT); + build_with_allocator(SVS_ALLOCATOR_KIND_HUGE_PAGE); + + // Null builder is rejected. + CATCH_REQUIRE( + svs_index_builder_set_allocator(nullptr, SVS_ALLOCATOR_KIND_DEFAULT, error) == + false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + + // Unknown allocator kind is rejected. + CATCH_REQUIRE( + svs_index_builder_set_allocator( + builder, static_cast(999), error + ) == false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + } + + CATCH_SECTION("Custom Allocator") { + // Wire up a memory-tracking custom allocator. + TrackingAllocator tracker; + struct svs_allocator_interface_ops alloc_ops = SVS_INIT_ALLOCATOR_OPS( + tracking_allocator_allocate, tracking_allocator_deallocate + ); + struct svs_allocator_interface allocator = SVS_MAKE_INTERFACE(&tracker, alloc_ops); + + // Null builder and null allocator are both rejected. + CATCH_REQUIRE( + svs_index_builder_set_allocator_custom(nullptr, &allocator, error) == false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + CATCH_REQUIRE( + svs_index_builder_set_allocator_custom(builder, nullptr, error) == false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + + bool ok = svs_index_builder_set_allocator_custom(builder, &allocator, error); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(svs_error_ok(error)); + + // Build the dynamic index; the vector data is allocated through `tracker`. + svs_index_h index = svs_index_build_dynamic( + builder, data.data(), ids.data(), NUM_VECTORS, BLOCK_SIZE, error + ); + CATCH_REQUIRE(index != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + // The custom allocator must actually have been used. + CATCH_REQUIRE(tracker.alloc_count > 0); + CATCH_REQUIRE(tracker.live_bytes > 0); + CATCH_REQUIRE(tracker.total_bytes >= tracker.live_bytes); + + // Actual memory usage reported by the built index. + svs_memory_breakdown_t breakdown = SVS_INIT_MEMORY_BREAKDOWN(); + ok = svs_index_get_memory_breakdown(index, &breakdown, error); + CATCH_REQUIRE(ok); + CATCH_REQUIRE(svs_error_ok(error)); + + auto within_1pct = [](size_t a, size_t b) { + if (a == b) { + return true; + } + const auto [smaller, larger] = std::minmax(a, b); + return (larger - smaller) * 100 <= larger; + }; + // For the dynamic Vamana index, both vector data and the blocked graph are routed + // through the builder's allocator. ID-translation metadata uses its own internal + // allocator, so tracked live bytes must match the data and graph portions of the + // breakdown. + // TODO: Modify MutableVamanaIndex metadata allocations to use the builder's + // allocator so that this check can be extended to include metadata_bytes. + CATCH_REQUIRE( + within_1pct(tracker.live_bytes, breakdown.data_bytes + breakdown.graph_bytes) + ); + + // Freeing the index returns every tracked byte to the allocator. + svs_index_free(index); + CATCH_REQUIRE(tracker.live_bytes == 0); + CATCH_REQUIRE(tracker.dealloc_count == tracker.alloc_count); + } + 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 2d78dbb53..45b487562 100644 --- a/bindings/c/tests/c_api_index.cpp +++ b/bindings/c/tests/c_api_index.cpp @@ -883,7 +883,7 @@ CATCH_TEST_CASE("C API Index Memory Management", "[c_api][index][memory]") { } // Estimate before build. - svs_memory_breakdown_t estimated{}; + svs_memory_breakdown_t estimated = SVS_INIT_MEMORY_BREAKDOWN(); success = svs_index_builder_estimate_memory(builder, NUM_VECTORS, &estimated, error); CATCH_REQUIRE(success); @@ -897,7 +897,7 @@ CATCH_TEST_CASE("C API Index Memory Management", "[c_api][index][memory]") { CATCH_REQUIRE(index != nullptr); CATCH_REQUIRE(svs_error_ok(error)); - svs_memory_breakdown_t actual{}; + svs_memory_breakdown_t actual = SVS_INIT_MEMORY_BREAKDOWN(); success = svs_index_get_memory_breakdown(index, &actual, error); CATCH_REQUIRE(success); CATCH_REQUIRE(svs_error_ok(error)); @@ -1085,6 +1085,154 @@ CATCH_TEST_CASE("C API Index Memory Management", "[c_api][index][memory]") { svs_algorithm_free(algorithm); svs_error_free(error); } + + CATCH_SECTION("Allocator Configuration") { + svs_error_h error = svs_error_create(); + + // Each built-in allocator kind must yield a usable index. + auto build_with_allocator = [&](svs_allocator_kind_t kind) { + 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)); + + success = svs_index_builder_set_allocator(builder, kind, error); + CATCH_REQUIRE(success); + CATCH_REQUIRE(svs_error_ok(error)); + + svs_index_h index = svs_index_build(builder, data.data(), NUM_VECTORS, error); + CATCH_REQUIRE(index != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + size_t memory_usage = 0; + success = svs_index_get_memory_usage(index, &memory_usage, error); + CATCH_REQUIRE(success); + CATCH_REQUIRE(memory_usage > 0); + + svs_index_free(index); + svs_index_builder_free(builder); + svs_algorithm_free(algorithm); + }; + + build_with_allocator(SVS_ALLOCATOR_KIND_DEFAULT); + build_with_allocator(SVS_ALLOCATOR_KIND_HUGE_PAGE); + + // Null builder is rejected. + CATCH_REQUIRE( + svs_index_builder_set_allocator(nullptr, SVS_ALLOCATOR_KIND_DEFAULT, error) == + false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + + // Unknown allocator kind is rejected. + { + svs_algorithm_h algorithm = svs_algorithm_create_vamana(16, 32, 50, error); + svs_index_builder_h builder = svs_index_builder_create( + SVS_DISTANCE_METRIC_EUCLIDEAN, DIMENSION, algorithm, error + ); + CATCH_REQUIRE(builder != nullptr); + CATCH_REQUIRE( + svs_index_builder_set_allocator( + builder, static_cast(999), error + ) == false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + svs_index_builder_free(builder); + svs_algorithm_free(algorithm); + } + + svs_error_free(error); + } + + CATCH_SECTION("Custom Allocator") { + svs_error_h error = svs_error_create(); + + 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)); + + // Wire up a memory-tracking custom allocator. + TrackingAllocator tracker; + struct svs_allocator_interface_ops alloc_ops = SVS_INIT_ALLOCATOR_OPS( + tracking_allocator_allocate, tracking_allocator_deallocate + ); + struct svs_allocator_interface allocator = SVS_MAKE_INTERFACE(&tracker, alloc_ops); + + // Null builder and null allocator are both rejected. + CATCH_REQUIRE( + svs_index_builder_set_allocator_custom(nullptr, &allocator, error) == false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + CATCH_REQUIRE( + svs_index_builder_set_allocator_custom(builder, nullptr, error) == false + ); + CATCH_REQUIRE(svs_error_get_code(error) == SVS_ERROR_INVALID_ARGUMENT); + + success = svs_index_builder_set_allocator_custom(builder, &allocator, error); + CATCH_REQUIRE(success); + CATCH_REQUIRE(svs_error_ok(error)); + + // Build the index; both data and graph allocations flow through `tracker`. + svs_index_h index = svs_index_build(builder, data.data(), NUM_VECTORS, error); + CATCH_REQUIRE(index != nullptr); + CATCH_REQUIRE(svs_error_ok(error)); + + // The custom allocator must actually have been used. + CATCH_REQUIRE(tracker.alloc_count > 0); + CATCH_REQUIRE(tracker.live_bytes > 0); + CATCH_REQUIRE(tracker.total_bytes >= tracker.live_bytes); + + // Actual memory usage reported by the built index. + size_t memory_usage = 0; + success = svs_index_get_memory_usage(index, &memory_usage, error); + CATCH_REQUIRE(success); + CATCH_REQUIRE(svs_error_ok(error)); + + // For the static Vamana index both the data and the graph are allocated through + // the builder's allocator, so the bytes still held by the custom allocator must + // match the graph + data portion of the breakdown (the tiny metadata entry point + // is not routed through the allocator). The same must hold against the pre-build + // estimate. + auto within_1pct = [](size_t a, size_t b) { + if (a == b) { + return true; + } + const auto [smaller, larger] = std::minmax(a, b); + return (larger - smaller) * 100 <= larger; + }; + CATCH_REQUIRE(within_1pct(tracker.live_bytes, memory_usage)); + + // Freeing the index returns every tracked byte to the allocator. + svs_index_free(index); + CATCH_REQUIRE(tracker.live_bytes == 0); + CATCH_REQUIRE(tracker.dealloc_count == tracker.alloc_count); + + svs_index_builder_free(builder); + svs_algorithm_free(algorithm); + svs_error_free(error); + } } namespace { diff --git a/bindings/c/tests/c_api_test_utils.h b/bindings/c/tests/c_api_test_utils.h index 715de9c59..ad0e826c8 100644 --- a/bindings/c/tests/c_api_test_utils.h +++ b/bindings/c/tests/c_api_test_utils.h @@ -18,6 +18,7 @@ #include "svs/c/svs_c.h" // Standard library +#include #include #include #include @@ -164,3 +165,54 @@ inline bool check_storage_support(svs_storage_h storage, svs_error_h error) { /// True when compressed storage is expected to be usable on this host, so callers /// can skip the build/search portion of a test that cannot run. inline bool storage_usable(svs_storage_h storage) { return storage != nullptr; } + +/// Memory-accounting state for a custom allocator used in tests. A single instance is +/// passed as the `self` pointer of a svs_allocator_interface; because every rebound or +/// cloned copy of the underlying allocator shares that pointer, one TrackingAllocator +/// accounts for every allocation funneled through the interface. The counters are atomic +/// because index builds allocate concurrently from multiple worker threads. +struct TrackingAllocator { + std::atomic live_bytes{0}; ///< Currently held bytes (allocated - freed). + std::atomic total_bytes{0}; ///< Cumulative bytes ever requested. + std::atomic peak_bytes{0}; ///< Maximum simultaneously live bytes. + std::atomic alloc_count{0}; ///< Number of allocate() calls. + std::atomic dealloc_count{0}; ///< Number of deallocate() calls. +}; + +/// svs_allocator_interface_ops::allocate implementation backed by a TrackingAllocator. +inline void* tracking_allocator_allocate( + void* self, size_t size, size_t alignment, svs_error_h out_err +) { + auto* tracker = static_cast(self); + // std::aligned_alloc requires a power-of-two alignment that is at least + // alignof(std::max_align_t) and a size that is a multiple of that alignment. + size_t align = + alignment < alignof(std::max_align_t) ? alignof(std::max_align_t) : alignment; + size_t rounded = ((size + align - 1) / align) * align; + void* ptr = std::aligned_alloc(align, rounded); + if (ptr == nullptr) { + svs_error_set( + out_err, SVS_ERROR_OUT_OF_MEMORY, "TrackingAllocator: allocation failed" + ); + return nullptr; + } + tracker->total_bytes.fetch_add(size, std::memory_order_relaxed); + size_t live = tracker->live_bytes.fetch_add(size, std::memory_order_relaxed) + size; + // Best-effort peak update (racy under contention but only used for diagnostics). + size_t peak = tracker->peak_bytes.load(std::memory_order_relaxed); + while (live > peak && + !tracker->peak_bytes.compare_exchange_weak(peak, live, std::memory_order_relaxed) + ) {} + tracker->alloc_count.fetch_add(1, std::memory_order_relaxed); + return ptr; +} + +/// svs_allocator_interface_ops::deallocate implementation backed by a TrackingAllocator. +inline void tracking_allocator_deallocate( + void* self, void* ptr, size_t size, size_t /*alignment*/ +) { + auto* tracker = static_cast(self); + tracker->live_bytes.fetch_sub(size, std::memory_order_relaxed); + tracker->dealloc_count.fetch_add(1, std::memory_order_relaxed); + std::free(ptr); +} diff --git a/include/svs/core/allocator.h b/include/svs/core/allocator.h index 1e449a7ee..8924aae0e 100644 --- a/include/svs/core/allocator.h +++ b/include/svs/core/allocator.h @@ -32,6 +32,7 @@ /// @defgroup core_allocators_public Public API for Allocators /// +#include "svs/lib/datatype.h" #include "svs/lib/memory.h" #include "svs/lib/misc.h" #include "svs/lib/narrow.h" @@ -556,21 +557,34 @@ concept Allocator = HasValueType && std::is_copy_constructible_v class AllocatorInterface { public: + // The types that allocator handle can rebind to. + // Order here is for performance - most common types are listed first. + using rebind_types = lib::Types< + std::byte, + uint32_t, + float, + svs::Float16, + int8_t, + uint8_t, + int32_t, + double, + svs::BFloat16, + uint64_t, + int64_t, + uint16_t, + int16_t>; + virtual ~AllocatorInterface() = default; virtual void* allocate(size_t n) = 0; virtual void deallocate(void* ptr, size_t n) = 0; // covariant return type virtual AllocatorInterface* clone() const = 0; - virtual AllocatorInterface* rebind_float() const = 0; - virtual AllocatorInterface* rebind_float16() const = 0; + virtual AllocatorInterface* rebind_to(DataType type) const = 0; }; template class AllocatorImpl : public AllocatorInterface { public: - using rebind_allocator_float = lib::rebind_allocator_t; - using rebind_allocator_float16 = lib::rebind_allocator_t; - // pass by value due to clone() explicit AllocatorImpl(Impl impl) : AllocatorInterface{} @@ -584,12 +598,15 @@ template class AllocatorImpl : public AllocatorInterfac AllocatorImpl* clone() const override { return new AllocatorImpl(impl_); } - AllocatorImpl* rebind_float() const override { - return new AllocatorImpl(rebind_allocator_float{impl_}); - } - - AllocatorImpl* rebind_float16() const override { - return new AllocatorImpl(rebind_allocator_float16{impl_}); + AllocatorInterface* rebind_to(DataType type) const override { + return svs::lib::match( + AllocatorInterface::rebind_types{}, + type, + [this](lib::Type) -> AllocatorInterface* { + using rebind_allocator = lib::rebind_allocator_t; + return new AllocatorImpl(rebind_allocator{impl_}); + } + ); } private: @@ -600,6 +617,9 @@ template class AllocatorHandle { public: using value_type = T; + explicit AllocatorHandle(std::unique_ptr impl) + : impl_{std::move(impl)} {} + template explicit AllocatorHandle(Impl&& impl) requires(!std::is_same_v) && @@ -621,27 +641,16 @@ template class AllocatorHandle { // Enable rebinding of allocators. template friend class AllocatorHandle; - template + template + requires(!std::is_same_v) && (lib::in(AllocatorInterface::rebind_types{})) AllocatorHandle(const AllocatorHandle& other) - requires std::is_same_v && (!std::is_same_v) - : impl_{other.impl_->rebind_float()} {} - template - AllocatorHandle(const AllocatorHandle& other) - requires std::is_same_v && (!std::is_same_v) - : impl_{other.impl_->rebind_float16()} {} + : impl_{other.impl_->rebind_to(datatype_v)} {} - template - AllocatorHandle& operator=(const AllocatorHandle& other) - requires std::is_same_v && (!std::is_same_v) - { - impl_.reset(other.impl_->rebind_float()); - return *this; - } - template + template AllocatorHandle& operator=(const AllocatorHandle& other) - requires std::is_same_v && (!std::is_same_v) + requires(!std::is_same_v) && (lib::in(AllocatorInterface::rebind_types{})) { - impl_.reset(other.impl_->rebind_float16()); + impl_.reset(other.impl_->rebind_to(datatype_v)); return *this; } diff --git a/include/svs/core/graph.h b/include/svs/core/graph.h index 779407256..01b286733 100644 --- a/include/svs/core/graph.h +++ b/include/svs/core/graph.h @@ -31,9 +31,10 @@ namespace svs { /// /// @tparam Idx The type used to encode nodes in the graph. /// -template struct GraphLoader { +template > +struct GraphLoader { // Type aliases - using return_type = graphs::SimpleGraph>; + using return_type = graphs::SimpleGraph; /// @brief Construct a new GraphLoader /// @@ -42,14 +43,16 @@ template struct GraphLoader { /// The saved graph directory will generally be created when saving a graph based /// index. The ``path`` argument should be this directory. /// - GraphLoader(const std::filesystem::path& path) - : path_{path} {} + GraphLoader(const std::filesystem::path& path, const Allocator& allocator = {}) + : path_{path} + , allocator_{allocator} {} /// @brief Load the graph into memory. - return_type load() const { return return_type::load(path_); } + return_type load() const { return return_type::load(path_, allocator_); } ///// Members std::filesystem::path path_{}; + Allocator allocator_{}; }; /// diff --git a/include/svs/index/vamana/dynamic_index.h b/include/svs/index/vamana/dynamic_index.h index de4341b8f..4419ead92 100644 --- a/include/svs/index/vamana/dynamic_index.h +++ b/include/svs/index/vamana/dynamic_index.h @@ -201,6 +201,56 @@ class MutableVamanaIndex { translator_.insert(external_ids, threads::UnitRange(0, external_ids.size())); } + template + MutableVamanaIndex( + const VamanaBuildParameters& parameters, + Graph graph, + Data data, + Idx entry_point, + Dist distance_function, + const ExternalIds& external_ids, + ThreadPoolProto threadpool_proto, + // Optional logger parameter + svs::logging::logger_ptr logger = svs::logging::get() + ) + : MutableVamanaIndex{ + std::move(graph), + std::move(data), + entry_point, + std::move(distance_function), + external_ids, + std::move(threadpool_proto), + std::move(logger)} { + if (graph_.n_nodes() != data_.size()) { + throw ANNEXCEPTION("Wrong sizes!"); + } + build_parameters_ = parameters; + // Verify and set defaults before using the parameters to set other member + // variables. + verify_and_set_default_index_parameters(build_parameters_, distance_function); + + assert(build_parameters_.graph_max_degree == graph_.max_degree()); + alpha_ = build_parameters_.alpha; + construction_window_size_ = build_parameters_.window_size; + max_candidates_ = build_parameters_.max_candidate_pool_size; + prune_to_ = build_parameters_.prune_to; + use_full_search_history_ = build_parameters_.use_full_search_history; + // Perform graph construction. + auto builder = VamanaBuilder( + graph_, + data_, + distance_, + build_parameters_, + threadpool_, + extensions::estimate_prefetch_parameters(data_), + logger_ + ); + builder.construct(1.0f, entry_point_[0], logging::Level::Trace, logger_); + builder.construct( + build_parameters_.alpha, entry_point_[0], logging::Level::Trace, logger_ + ); + } + /// /// Build a graph from scratch. /// @@ -1395,6 +1445,45 @@ struct VamanaStateLoader { } // namespace detail +// Build +template < + typename DataProto, + typename Distance, + typename ExternalIdsProto, + typename ThreadPoolProto, + typename GraphAllocator = data::Blocked>> +auto auto_dynamic_build( + const VamanaBuildParameters& parameters, + DataProto&& data_proto, + ExternalIdsProto&& external_ids_proto, + Distance distance, + ThreadPoolProto threadpool_proto, + const GraphAllocator& graph_allocator = {}, + svs::logging::logger_ptr logger = svs::logging::get() +) { + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto data = svs::detail::dispatch_load(SVS_FWD(data_proto), threadpool); + auto external_ids = svs::detail::dispatch_load(SVS_FWD(external_ids_proto), threadpool); + auto entry_point = extensions::compute_entry_point(data, threadpool); + + // Perform graph construction. + auto verified_parameters = parameters; + verify_and_set_default_index_parameters(verified_parameters, distance); + + auto graph = + default_graph(data.size(), verified_parameters.graph_max_degree, graph_allocator); + using Idx = typename decltype(graph)::index_type; + return MutableVamanaIndex{ + verified_parameters, + std::move(graph), + std::move(data), + lib::narrow(entry_point), + std::move(distance), + std::move(external_ids), + std::move(threadpool), + std::move(logger)}; +} + // Assembly template < typename GraphLoader, diff --git a/include/svs/orchestrators/dynamic_vamana.h b/include/svs/orchestrators/dynamic_vamana.h index 6d20476b3..0d9d5dfe6 100644 --- a/include/svs/orchestrators/dynamic_vamana.h +++ b/include/svs/orchestrators/dynamic_vamana.h @@ -119,6 +119,9 @@ class DynamicVamana : public manager::IndexManager { using base_type = manager::IndexManager; using VamanaSearchParameters = index::vamana::VamanaSearchParameters; + /// @private + struct BuildTag {}; + /// @private struct AssembleTag {}; /// @@ -288,13 +291,15 @@ class DynamicVamana : public manager::IndexManager { manager::QueryTypeDefinition QueryTypes, typename DataLoader, typename Distance, - typename ThreadPoolProto> + typename ThreadPoolProto, + typename GraphAllocator = data::Blocked>> static DynamicVamana build( const index::vamana::VamanaBuildParameters& parameters, DataLoader&& data_loader, std::span ids, Distance distance, - ThreadPoolProto threadpool_proto + ThreadPoolProto threadpool_proto, + const GraphAllocator& graph_allocator = {} ) { auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); auto data = @@ -304,16 +309,24 @@ class DynamicVamana : public manager::IndexManager { auto dispatcher = DistanceDispatcher(distance); return dispatcher([&](auto distance_function) { return make_dynamic_vamana>( + BuildTag{}, parameters, std::move(data), ids, std::move(distance_function), - std::move(threadpool) + std::move(threadpool), + graph_allocator ); }); } else { return make_dynamic_vamana>( - parameters, std::move(data), ids, std::move(distance), std::move(threadpool) + BuildTag{}, + parameters, + std::move(data), + ids, + std::move(distance), + std::move(threadpool), + graph_allocator ); } } @@ -508,4 +521,11 @@ DynamicVamana make_dynamic_vamana(Args&&... args) { std::make_unique>(std::forward(args)...)}; } +template +DynamicVamana make_dynamic_vamana(DynamicVamana::BuildTag SVS_UNUSED(tag), Args&&... args) { + return make_dynamic_vamana( + index::vamana::auto_dynamic_build(std::forward(args)...) + ); +} + } // namespace svs