diff --git a/.github/.licenserc.yaml b/.github/.licenserc.yaml index ad6f7f31..1b5f8976 100644 --- a/.github/.licenserc.yaml +++ b/.github/.licenserc.yaml @@ -46,6 +46,7 @@ header: - '.github/renovate.json' - '.github/CODEOWNERS' - 'cmake/mkl_functions' + - 'cmake/mkl_functions_ivf' - 'cmake/patches/tomlplusplus_v330.patch' - 'docker/x86_64/manylinux2014/oneAPI.repo' - 'docs/cpp/index/loader-compatibility.csv' diff --git a/.gitignore b/.gitignore index 6aa7c701..efaf738f 100644 --- a/.gitignore +++ b/.gitignore @@ -4,7 +4,7 @@ usr/ wheelhouse/ # Bundled test data -/data/ +/data/temp # Misc tool related files *.swp diff --git a/CMakeLists.txt b/CMakeLists.txt index e682b4f7..47b39f49 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -77,6 +77,12 @@ include("cmake/fmt.cmake") include("cmake/spdlog.cmake") include("cmake/toml.cmake") +# IVF requires Intel(R) MKL support +if(SVS_EXPERIMENTAL_ENABLE_IVF) + include("cmake/mkl.cmake") + target_compile_options(${SVS_LIB} INTERFACE "-DSVS_HAVE_MKL=1") +endif() + add_library(svs_x86_options_base INTERFACE) add_library(svs::x86_options_base ALIAS svs_x86_options_base) if(CMAKE_SYSTEM_PROCESSOR MATCHES "(x86)|(X86)|(amd64)|(AMD64)") diff --git a/THIRD-PARTY-PROGRAMS b/THIRD-PARTY-PROGRAMS index 9b976e48..7b1f9557 100644 --- a/THIRD-PARTY-PROGRAMS +++ b/THIRD-PARTY-PROGRAMS @@ -159,7 +159,7 @@ Please also refer to the file .github/CONTRIBUTING.md, which clarifies licensing external contributions to this project including patches, pull requests, etc. -------------------------------------------------------------------------------- -7. MKL (cmake/mkl.cmake, https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl.html) +7. Intel(R) MKL (cmake/mkl.cmake, https://www.intel.com/content/www/us/en/developer/tools/oneapi/onemkl.html) Copyright (c) Intel Corporation, All rights reserved. diff --git a/benchmark/CMakeLists.txt b/benchmark/CMakeLists.txt index 0e4bc916..62995be7 100644 --- a/benchmark/CMakeLists.txt +++ b/benchmark/CMakeLists.txt @@ -48,6 +48,17 @@ set(SHARED_LIBRARY_FILES src/inverted/memory/executables/memory_test.cpp ) +# ivf +if (SVS_EXPERIMENTAL_ENABLE_IVF) + list(APPEND SHARED_LIBRARY_FILES + src/ivf/uncompressed.cpp + src/ivf/search.cpp + src/ivf/build.cpp + src/ivf/test.cpp + ) +endif() + + add_library(svs_benchmark_library SHARED ${SHARED_LIBRARY_FILES}) target_include_directories(svs_benchmark_library PUBLIC ${CMAKE_CURRENT_LIST_DIR}/include) diff --git a/benchmark/include/svs-benchmark/ivf/build.h b/benchmark/include/svs-benchmark/ivf/build.h new file mode 100644 index 00000000..46481bbb --- /dev/null +++ b/benchmark/include/svs-benchmark/ivf/build.h @@ -0,0 +1,280 @@ +/* + * Copyright 2025 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 + +// svs-benchmark +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/build.h" +#include "svs-benchmark/datasets.h" +#include "svs-benchmark/index_traits.h" +#include "svs-benchmark/ivf/search.h" +#include "svs-benchmark/search.h" + +// svs +#include "svs/orchestrators/ivf.h" + +// stl +#include +#include +#include +#include +#include +#include + +namespace svsbenchmark::ivf { + +struct StaticBenchmark {}; + +// Forward declarations +struct BuildJob; + +template struct AssociatedJob; + +template <> struct AssociatedJob { + using type = BuildJob; +}; + +template using associated_job_t = typename AssociatedJob::type; + +// Job names +inline constexpr std::string_view benchmark_name(StaticBenchmark) { + return "ivf_static_build"; +} + +// Entry-point for registering the static index building executable. +std::unique_ptr static_workflow(); + +// Shared struct between the static and dynamic paths. +struct BuildJobBase { + public: + // A descriptive name for this workload. + std::string description_; + + // The dataset to load + Dataset dataset_; + + // Paths + std::filesystem::path data_; + std::filesystem::path queries_; + + // The number of queries (taken form queries) to use in the training set. + size_t queries_in_training_set_; + + // Dataset Parameters + svs::DataType data_type_; + svs::DataType query_type_; + svs::DistanceType distance_; + Extent ndims_; + + // Build Parameters + svs::index::ivf::IVFBuildParameters build_parameters_; + size_t num_threads_; + + public: + ///// Contructor + BuildJobBase( + std::string_view description, + svsbenchmark::Dataset dataset, + std::filesystem::path data, + std::filesystem::path queries, + size_t queries_in_training_set, + svs::DataType data_type, + svs::DataType query_type, + svs::DistanceType distance, + size_t ndims, + const svs::index::ivf::IVFBuildParameters& build_parameters, + size_t num_threads + ) + : description_{description} + , dataset_{dataset} + , data_{std::move(data)} + , queries_{std::move(queries)} + , queries_in_training_set_{queries_in_training_set} + , data_type_{data_type} + , query_type_{query_type} + , distance_{distance} + , ndims_{ndims} + , build_parameters_{build_parameters} + , num_threads_{num_threads} {} + + // Compatibility with `ExpectedResults`. + const svs::index::ivf::IVFBuildParameters& get_build_parameters() const { + return build_parameters_; + } + svs::DistanceType get_distance() const { return distance_; } + + // Return an example BuildJob that can be used to generate sample config files. + static BuildJobBase example() { + return BuildJobBase( + "example index build", + Dataset::example(), + "data.fvecs", + "queries.fvecs", + 5000, + svs::DataType::float32, + svs::DataType::float32, + svs::DistanceType::L2, + svs::Dynamic, + svs::index::ivf::IVFBuildParameters(128, 10000, 10, false, 0.1), + 8 + ); + } + + svs::lib::SaveTable + to_toml(std::string_view schema, const svs::lib::Version& version) const { + return svs::lib::SaveTable( + schema, + version, + {SVS_LIST_SAVE_(description), + SVS_LIST_SAVE_(dataset), + SVS_LIST_SAVE_(data), + SVS_LIST_SAVE_(queries), + SVS_LIST_SAVE_(queries_in_training_set), + SVS_LIST_SAVE_(data_type), + SVS_LIST_SAVE_(query_type), + SVS_LIST_SAVE_(distance), + SVS_LIST_SAVE_(ndims), + SVS_LIST_SAVE_(build_parameters), + SVS_LIST_SAVE_(num_threads)} + ); + } + + static BuildJobBase from_toml( + const svs::lib::ContextFreeLoadTable& table, + const std::optional& root + ) { + namespace lib = svs::lib; + return BuildJobBase( + SVS_LOAD_MEMBER_AT_(table, description), + SVS_LOAD_MEMBER_AT_(table, dataset, root), + svsbenchmark::extract_filename(table, "data", root), + svsbenchmark::extract_filename(table, "queries", root), + SVS_LOAD_MEMBER_AT_(table, queries_in_training_set), + SVS_LOAD_MEMBER_AT_(table, data_type), + SVS_LOAD_MEMBER_AT_(table, query_type), + SVS_LOAD_MEMBER_AT_(table, distance), + SVS_LOAD_MEMBER_AT_(table, ndims), + SVS_LOAD_MEMBER_AT_(table, build_parameters), + SVS_LOAD_MEMBER_AT_(table, num_threads) + ); + } +}; + +// Parsed setup for a static index build job. +struct BuildJob : public BuildJobBase { + public: + // Paths + std::filesystem::path groundtruth_; + // Preset search parameters + std::vector preset_parameters_; + // Post-build validation parameters. + svsbenchmark::search::SearchParameters search_parameters_; + // Directory to save the built index. + // An empty optional implies no saving. + std::optional save_directory_; + + public: + template + BuildJob( + std::filesystem::path groundtruth, + std::vector preset_parameters, + svsbenchmark::search::SearchParameters search_parameters, + std::optional save_directory, + Args&&... args + ) + : BuildJobBase(std::forward(args)...) + , groundtruth_{std::move(groundtruth)} + , preset_parameters_{std::move(preset_parameters)} + , search_parameters_{std::move(search_parameters)} + , save_directory_{std::move(save_directory)} {} + + // Return an example BuildJob that can be used to generate sample config files. + static BuildJob example() { + return BuildJob( + "groundtruth.ivecs", // groundtruth + {{10, 1.0}, {10, 4.0}, {50, 1.0}}, // preset_parameters + svsbenchmark::search::SearchParameters::example(), // search_parameters + std::nullopt, // save_directory + BuildJobBase::example() // base-class + ); + } + + // Compatibility with abstract search-space. + std::vector get_search_configs() const { + return preset_parameters_; + } + const svsbenchmark::search::SearchParameters& get_search_parameters() const { + return search_parameters_; + } + + template + auto invoke(F&& f, const Checkpoint& SVS_UNUSED(checkpoint)) const { + return f(dataset_, query_type_, data_type_, distance_, ndims_, *this); + } + + // Save the index if the `save_directory` field is non-empty. + template void maybe_save_index(Index& index) const { + if (!save_directory_) { + return; + } + const auto& root = save_directory_.value(); + svs::lib::save_to_disk(index, root / "clustering"); + } + + static constexpr svs::lib::Version save_version = svs::lib::Version(0, 0, 0); + static constexpr std::string_view serialization_schema = "benchmark_ivf_build_job"; + + // Save the BuildJob to a TOML table. + svs::lib::SaveTable save() const { + // Get a base table. + auto table = BuildJobBase::to_toml(serialization_schema, save_version); + + // Append the extra information needed by the static BuildJob. + SVS_INSERT_SAVE_(table, groundtruth); + SVS_INSERT_SAVE_(table, preset_parameters); + SVS_INSERT_SAVE_(table, search_parameters); + table.insert("save_directory", svs::lib::save(save_directory_.value_or(""))); + return table; + } + + // Load a BuildJob from a TOML table. + static BuildJob load( + const svs::lib::ContextFreeLoadTable& table, + const std::optional& root, + svsbenchmark::SaveDirectoryChecker& checker + ) { + return BuildJob( + svsbenchmark::extract_filename(table, "groundtruth", root), + SVS_LOAD_MEMBER_AT_(table, preset_parameters), + SVS_LOAD_MEMBER_AT_(table, search_parameters), + checker.extract(table.unwrap(), "save_directory"), + BuildJobBase::from_toml(table, root) + ); + } +}; + +// Dispatchers +using StaticBuildDispatcher = svs::lib::Dispatcher< + toml::table, + svsbenchmark::Dataset, + svs::DataType, + svs::DataType, + svs::DistanceType, + Extent, + const BuildJob&>; + +} // namespace svsbenchmark::ivf diff --git a/benchmark/include/svs-benchmark/ivf/common.h b/benchmark/include/svs-benchmark/ivf/common.h new file mode 100644 index 00000000..2ff38b9a --- /dev/null +++ b/benchmark/include/svs-benchmark/ivf/common.h @@ -0,0 +1,39 @@ +/* + * Copyright 2025 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. + */ + +// svs-benchmark +#include "svs-benchmark/benchmark.h" + +// svs +#include "svs/core/distance.h" +#include "svs/index/ivf/common.h" + +// stl +#include + +namespace svsbenchmark::ivf { + +// Test Routines +SVS_BENCHMARK_FOR_TESTS_ONLY inline search::SearchParameters test_search_parameters() { + return search::SearchParameters{10, {0.5, 0.8, 0.9}}; +} + +SVS_BENCHMARK_FOR_TESTS_ONLY inline std::vector +test_search_configs() { + return std::vector({{{10, 1.0}, {50, 1.0}}}); +} + +} // namespace svsbenchmark::ivf diff --git a/benchmark/include/svs-benchmark/ivf/search.h b/benchmark/include/svs-benchmark/ivf/search.h new file mode 100644 index 00000000..7a0c563a --- /dev/null +++ b/benchmark/include/svs-benchmark/ivf/search.h @@ -0,0 +1,215 @@ +/* + * Copyright 2025 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 + +// svs-benchmark +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/datasets.h" +#include "svs-benchmark/index_traits.h" +#include "svs-benchmark/search.h" + +// svs +#include "svs/core/distance.h" +#include "svs/index/ivf/common.h" +#include "svs/lib/dispatcher.h" + +// stl +#include + +namespace svsbenchmark::ivf { + +inline constexpr std::string_view search_benchmark_name() { return "ivf_static_search"; } + +// Entry point for searching. +std::unique_ptr search_static_workflow(); + +// The current state of the index. +struct IVFState { + public: + svs::index::ivf::IVFSearchParameters search_parameters_; + size_t num_threads_; + + public: + IVFState(svs::index::ivf::IVFSearchParameters search_parameters, size_t num_threads) + : search_parameters_{search_parameters} + , num_threads_{num_threads} {} + + template + explicit IVFState(const Index& index) + : IVFState(index.get_search_parameters(), index.get_num_threads()) {} + + // Saving + static constexpr svs::lib::Version save_version{0, 0, 0}; + static constexpr std::string_view serialization_schema = "benchmark_ivf_state"; + svs::lib::SaveTable save() const { + return svs::lib::SaveTable( + serialization_schema, + save_version, + {SVS_LIST_SAVE_(search_parameters), SVS_LIST_SAVE_(num_threads)} + ); + } +}; + +struct SearchJob { + public: + std::string description_; + svsbenchmark::Dataset dataset_; + std::filesystem::path config_; + std::filesystem::path graph_; + std::filesystem::path data_; + std::filesystem::path queries_; + std::filesystem::path groundtruth_; + size_t queries_in_training_set_; + svs::DataType data_type_; + svs::DataType query_type_; + svs::DistanceType distance_; + Extent ndims_; + size_t num_threads_; + svsbenchmark::search::SearchParameters search_parameters_; + std::vector preset_parameters_; + + public: + SearchJob( + std::string description, + svsbenchmark::Dataset dataset, + std::filesystem::path config, + std::filesystem::path graph, + std::filesystem::path data, + std::filesystem::path queries, + std::filesystem::path groundtruth, + size_t queries_in_training_set, + svs::DataType data_type, + svs::DataType query_type, + svs::DistanceType distance, + Extent ndims, + size_t num_threads, + const svsbenchmark::search::SearchParameters& search_parameters, + std::vector preset_parameters + ) + : description_{std::move(description)} + , dataset_{std::move(dataset)} + , config_{std::move(config)} + , graph_{std::move(graph)} + , data_{std::move(data)} + , queries_{std::move(queries)} + , groundtruth_{std::move(groundtruth)} + , queries_in_training_set_{queries_in_training_set} + , data_type_{data_type} + , query_type_{query_type} + , distance_{distance} + , ndims_{ndims} + , num_threads_{num_threads} + , search_parameters_{search_parameters} + , preset_parameters_{std::move(preset_parameters)} {} + + // Return the benchmark search parameters + const svsbenchmark::search::SearchParameters& get_search_parameters() const { + return search_parameters_; + } + + // Compatbility with `ExpectedResults` + static std::nullopt_t get_build_parameters() { return std::nullopt; } + svs::DistanceType get_distance() const { return distance_; } + + // Return the preset search configurations. + const std::vector& get_search_configs() const { + return preset_parameters_; + } + + static SearchJob example() { + return SearchJob{ + "index search", // description + Dataset::example(), // dataset + "path/to/clustering dir", // clustering dir + "path/not/used", // Not applicable + "path/to/data", // data + "path/to/queries", // queries + "path/to/groundtruth", // groundtruth + 5000, // queries_in_training_set + svs::DataType::float32, // data_type + svs::DataType::float32, // query_type + svs::DistanceType::L2, // distance + Extent{svs::Dynamic}, // ndims + 4, // num_threads + search::SearchParameters::example(), // search_parameters + {{10, 1.0}, {10, 4.0}} // preset_parameters + }; + } + + template + auto invoke(F&& f, const Checkpoint& SVS_UNUSED(checkpoiner)) const { + return f(dataset_, query_type_, data_type_, distance_, ndims_, *this); + } + + ///// Save/Load + static constexpr svs::lib::Version save_version{0, 0, 0}; + static constexpr std::string_view serialization_schema = "benchmark_ivf_search_job"; + svs::lib::SaveTable save() const { + return svs::lib::SaveTable( + serialization_schema, + save_version, + {SVS_LIST_SAVE_(description), + SVS_LIST_SAVE_(dataset), + SVS_LIST_SAVE_(config), + SVS_LIST_SAVE_(graph), + SVS_LIST_SAVE_(data), + SVS_LIST_SAVE_(queries), + SVS_LIST_SAVE_(groundtruth), + SVS_LIST_SAVE_(queries_in_training_set), + SVS_LIST_SAVE_(data_type), + SVS_LIST_SAVE_(query_type), + SVS_LIST_SAVE_(distance), + SVS_LIST_SAVE_(ndims), + SVS_LIST_SAVE_(num_threads), + SVS_LIST_SAVE_(search_parameters), + SVS_LIST_SAVE_(preset_parameters)} + ); + } + + static SearchJob load( + const svs::lib::ContextFreeLoadTable& table, + const std::optional& root = {} + ) { + return SearchJob{ + SVS_LOAD_MEMBER_AT_(table, description), + SVS_LOAD_MEMBER_AT_(table, dataset, root), + svsbenchmark::extract_filename(table, "config", root), + svsbenchmark::extract_filename(table, "graph", root), + svsbenchmark::extract_filename(table, "data", root), + svsbenchmark::extract_filename(table, "queries", root), + svsbenchmark::extract_filename(table, "groundtruth", root), + SVS_LOAD_MEMBER_AT_(table, queries_in_training_set), + SVS_LOAD_MEMBER_AT_(table, data_type), + SVS_LOAD_MEMBER_AT_(table, query_type), + SVS_LOAD_MEMBER_AT_(table, distance), + SVS_LOAD_MEMBER_AT_(table, ndims), + SVS_LOAD_MEMBER_AT_(table, num_threads), + SVS_LOAD_MEMBER_AT_(table, search_parameters), + SVS_LOAD_MEMBER_AT_(table, preset_parameters)}; + } +}; + +using StaticSearchDispatcher = svs::lib::Dispatcher< + toml::table, + Dataset, + svs::DataType, + svs::DataType, + svs::DistanceType, + Extent, + const SearchJob&>; + +} // namespace svsbenchmark::ivf diff --git a/benchmark/include/svs-benchmark/ivf/static_traits.h b/benchmark/include/svs-benchmark/ivf/static_traits.h new file mode 100644 index 00000000..baab997e --- /dev/null +++ b/benchmark/include/svs-benchmark/ivf/static_traits.h @@ -0,0 +1,139 @@ +/* + * Copyright 2025 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 + +// svs-benchmark +#include "svs-benchmark/index_traits.h" + +// svs +#include "svs/orchestrators/ivf.h" + +// stl +#include +#include + +namespace svsbenchmark { + +template <> struct IndexTraits { + using index_type = svs::IVF; + using config_type = svs::index::ivf::IVFSearchParameters; + using state_type = svsbenchmark::ivf::IVFState; + static std::string name() { return "static ivf index"; } + + // Configuration Space. + static void apply_config(svs::IVF& index, const config_type& config) { + index.set_search_parameters(config); + } + + template + static auto search( + svs::IVF& index, + const Queries& queries, + size_t num_neighbors, + const config_type& config + ) { + apply_config(index, config); + return index.search(queries, num_neighbors); + } + + static state_type report_state(const svs::IVF& index) { return state_type(index); } + + template + static config_type calibrate( + index_type& index, + const Queries& queries, + const Groundtruth& groundtruth, + size_t num_neighbors, + double target_recall, + svsbenchmark::CalibrateContext SVS_UNUSED(ctx), + svsbenchmark::Placeholder SVS_UNUSED(placeholder) + ) { + config_type baseline = index.get_search_parameters(); + + auto valid_configurations = std::vector(); + + // Search over epsilon search space first. + auto k_reorders = std::vector({1, 4, 10}); + auto n_probes_range = svs::threads::UnitRange(1, 200); + for (auto k_reorder : k_reorders) { + auto copy = baseline; + copy.k_reorder_ = k_reorder; + + copy.n_probes_ = *std::lower_bound( + n_probes_range.begin(), + n_probes_range.end(), + target_recall, + [&](size_t n_probes, double recall) { + copy.n_probes_ = n_probes; + index.set_search_parameters(copy); + + auto result = index.search(queries, num_neighbors); + auto this_recall = svs::k_recall_at_n(groundtruth, result); + return this_recall < recall; + } + ); + + valid_configurations.push_back(copy); + } + + // Loop through each valid configuration - find the fastest. + size_t best_config = std::numeric_limits::max(); + double lowest_latency = std::numeric_limits::max(); + + size_t config_index = 0; + for (auto& config : valid_configurations) { + apply_config(index, config); + auto latencies = std::vector(); + for (size_t i = 0; i < 5; ++i) { + auto tic = svs::lib::now(); + index.search(queries, num_neighbors); + latencies.push_back(svs::lib::time_difference(tic)); + } + + auto min_latency = *std::min_element(latencies.begin(), latencies.end()); + + std::cout << svs::lib::save_to_table(config) << '\n'; + SVS_SHOW(min_latency); + + if (min_latency < lowest_latency) { + best_config = config_index; + lowest_latency = min_latency; + } + ++config_index; + } + + return valid_configurations[best_config]; + } + + template + static config_type calibrate_with_hint( + index_type& index, + const Queries& queries, + const Groundtruth& groundtruth, + size_t num_neighbors, + double target_recall, + svsbenchmark::CalibrateContext ctx, + const config_type& SVS_UNUSED(preset), + svsbenchmark::Placeholder placeholder + ) { + return calibrate( + index, queries, groundtruth, num_neighbors, target_recall, ctx, placeholder + ); + } +}; + +} // namespace svsbenchmark diff --git a/benchmark/include/svs-benchmark/ivf/test.h b/benchmark/include/svs-benchmark/ivf/test.h new file mode 100644 index 00000000..95c9b307 --- /dev/null +++ b/benchmark/include/svs-benchmark/ivf/test.h @@ -0,0 +1,152 @@ +/* + * Copyright 2025 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 + +// svs-benchmark +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/ivf/build.h" +#include "svs-benchmark/ivf/search.h" +#include "svs-benchmark/ivf/static_traits.h" +#include "svs-benchmark/test.h" + +// svs +#include "svs/orchestrators/ivf.h" + +// stl +#include +#include +#include +#include + +namespace svsbenchmark::ivf { + +inline constexpr std::string_view test_benchmark_name() { return "ivf_test_generator"; } + +// A benchmark that generates reference inputs for unit tests. +std::unique_ptr test_generator(); + +///// Test Runner +struct IVFTest { + public: + std::vector groundtruths_; + std::filesystem::path data_f32_; + std::filesystem::path index_config_; + std::filesystem::path graph_; + std::filesystem::path queries_f32_; + size_t queries_in_training_set_; + // Backend-specific members + std::filesystem::path leanvec_data_matrix_; + std::filesystem::path leanvec_query_matrix_; + // Runtime values + size_t num_threads_; + + public: + IVFTest( + std::vector groundtruths, + std::filesystem::path data_f32, + std::filesystem::path index_config, + std::filesystem::path graph, + std::filesystem::path queries_f32, + size_t queries_in_training_set, + // backend-specific members + std::filesystem::path leanvec_data_matrix, + std::filesystem::path leanvec_query_matrix, + // Runtime values + size_t num_threads + ) + : groundtruths_{std::move(groundtruths)} + , data_f32_{std::move(data_f32)} + , index_config_{std::move(index_config)} + , graph_{std::move(graph)} + , queries_f32_{std::move(queries_f32)} + , queries_in_training_set_{queries_in_training_set} + , leanvec_data_matrix_{std::move(leanvec_data_matrix)} + , leanvec_query_matrix_{std::move(leanvec_query_matrix)} + , num_threads_{num_threads} {} + + static IVFTest example() { + return IVFTest{ + {DistanceAndGroundtruth::example()}, // groundtruths + "path/to/data_f32", // data_f32 + "path/to/config", // index_config + "path/to/graph", // graph + "path/to/queries_f32", // queries_f32 + 10000, // queries_in_training_set + "path/to/leanvec_data_matrix", // LeanVec data matrix + "path/to/leanvec_query_matrix", // LeanVec query matrix + 0, // Num Threads (not-saved) + }; + } + + const std::filesystem::path& groundtruth_for(svs::DistanceType distance) const { + for (const auto& pair : groundtruths_) { + if (pair.distance_ == distance) { + return pair.path_; + } + } + throw ANNEXCEPTION("Could not find a groundtruth for {} distance!", distance); + } + + ///// Save/Load + static constexpr svs::lib::Version save_version{0, 0, 0}; + static constexpr std::string_view serialization_schema = "benchmark_ivf_test"; + svs::lib::SaveTable save() const { + return svs::lib::SaveTable( + serialization_schema, + save_version, + {SVS_LIST_SAVE_(groundtruths), + SVS_LIST_SAVE_(data_f32), + SVS_LIST_SAVE_(index_config), + SVS_LIST_SAVE_(graph), + SVS_LIST_SAVE_(queries_f32), + SVS_LIST_SAVE_(queries_in_training_set), + SVS_LIST_SAVE_(leanvec_data_matrix), + SVS_LIST_SAVE_(leanvec_query_matrix)} + ); + } + + static IVFTest load( + const svs::lib::ContextFreeLoadTable& table, + size_t num_threads, + const std::optional& root = {} + ) { + return IVFTest{ + SVS_LOAD_MEMBER_AT_(table, groundtruths, root), + svsbenchmark::extract_filename(table, "data_f32", root), + svsbenchmark::extract_filename(table, "index_config", root), + svsbenchmark::extract_filename(table, "graph", root), + svsbenchmark::extract_filename(table, "queries_f32", root), + SVS_LOAD_MEMBER_AT_(table, queries_in_training_set), + svsbenchmark::extract_filename(table, "leanvec_data_matrix", root), + svsbenchmark::extract_filename(table, "leanvec_query_matrix", root), + num_threads}; + } +}; + +// Specialize ConfigAndResult for `svs::IVF`. +using ConfigAndResult = + svsbenchmark::ConfigAndResultPrototype; + +// Specialize ExpectedResult for `svs::IVF`. +using ExpectedResult = svsbenchmark::ExpectedResultPrototype< + svs::index::ivf::IVFBuildParameters, + svs::index::ivf::IVFSearchParameters>; + +// Test functions take the test input and returns a `TestFunctionReturn` with the results. +using TestFunction = std::function; + +} // namespace svsbenchmark::ivf diff --git a/benchmark/include/svs-benchmark/ivf/uncompressed.h b/benchmark/include/svs-benchmark/ivf/uncompressed.h new file mode 100644 index 00000000..c8af8afd --- /dev/null +++ b/benchmark/include/svs-benchmark/ivf/uncompressed.h @@ -0,0 +1,39 @@ +/* + * Copyright 2025 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 + +// svs-benchmark +#include "svs-benchmark/ivf/build.h" +#include "svs-benchmark/ivf/search.h" +#include "svs-benchmark/ivf/test.h" + +// stl +#include + +namespace svsbenchmark::ivf { + +///// target-registration +// search +void register_uncompressed_static_search(ivf::StaticSearchDispatcher&); + +// build +void register_uncompressed_static_build(ivf::StaticBuildDispatcher&); + +// test +std::vector register_uncompressed_test_routines(); + +} // namespace svsbenchmark::ivf diff --git a/benchmark/include/svs-benchmark/test.h b/benchmark/include/svs-benchmark/test.h index d11a091f..a8190b99 100644 --- a/benchmark/include/svs-benchmark/test.h +++ b/benchmark/include/svs-benchmark/test.h @@ -247,6 +247,9 @@ inline int8_t convert_to(svs::lib::Type, float x) { inline svs::Float16 convert_to(svs::lib::Type, float x) { return svs::Float16{x}; } +inline svs::BFloat16 convert_to(svs::lib::Type, float x) { + return svs::BFloat16{x}; +} inline float convert_to(svs::lib::Type, float x) { return x; } } // namespace detail diff --git a/benchmark/src/ivf/build.cpp b/benchmark/src/ivf/build.cpp new file mode 100644 index 00000000..20d3d679 --- /dev/null +++ b/benchmark/src/ivf/build.cpp @@ -0,0 +1,152 @@ +/* + * Copyright 2025 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. + */ + +// svs-benchmark +#include "svs-benchmark/ivf/build.h" +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/executable.h" +#include "svs-benchmark/ivf/uncompressed.h" + +// svs +#include "svs/lib/saveload.h" +#include "svs/third-party/toml.h" + +// third-party +#include "fmt/core.h" +#include "fmt/ranges.h" + +// stl +#include +#include +#include +#include +#include + +namespace svsbenchmark::ivf { +namespace { + +///// +///// Dispatchers +///// + +template struct BuildDispatcher; + +// Dispatcher for Uncompressed Data. +template <> struct BuildDispatcher { + using type = ivf::StaticBuildDispatcher; + static type dispatcher() { + auto dispatcher = type{}; + ivf::register_uncompressed_static_build(dispatcher); + return dispatcher; + } +}; + +///// +///// Executable +///// + +constexpr std::string_view HELP_TEMPLATE = R"( +Run a {} benchmark for the IVF index. + +Usage: + (1) src-file.toml (output-file.toml/--validate) [basename] + (2) --help + (3) --example + +1. Run all the benchmarks in the global `{}` array in `src-file.toml`. + All elements in the array must be parseable as a `{}`. + + Results will be saved to `output-file.toml`. + + If `--validate` is given as the second argument, then all pre-run checks will be + performed on the input file and arguments but not benchmark will actually be run. + + Optional third argument `basename` will be used as the root for all file paths parsed. + +2. Print this help message. + +3. Display an example input TOML file to `stdout`. + +Backend specializations are dispatched on the following fields of the input TOML file: +* build_type: The dataset type to use. +* query_type: The element type of the query dataset. +* data_type: The input type of the source dataset. +* distance: The distance function to use. +* ndims: The compile-time dimensionality. + +Compiled specializations are listed below: +{{ build_type, query_type, data_type, distance, ndims }} +)"; + +template struct Exe { + public: + // type alises + using job_type = associated_job_t; + using dispatcher_type = typename BuildDispatcher::type; + + static dispatcher_type dispatcher() { + return BuildDispatcher::dispatcher(); + } + + static std::string_view name() { return benchmark_name(BenchmarkType()); } + + static void print_help() { + if constexpr (std::is_same_v) { + fmt::print( + HELP_TEMPLATE, + "static build and search", + name(), + "svsbenchmark::IVF::BuildJob" + ); + } else { + throw ANNEXCEPTION("Unreachable"); + } + auto f = dispatcher(); + for (size_t i = 0; i < f.size(); ++i) { + auto dispatch_strings = std::array{ + f.description(i, 0), + f.description(i, 1), + f.description(i, 2), + f.description(i, 3), + f.description(i, 4), + }; + fmt::print("{{ {} }}\n", fmt::join(dispatch_strings, ", ")); + } + } + + static job_type example() { return job_type::example(); } + + template + static std::optional> + parse_args_and_invoke(F&& f, std::span args) { + auto root = std::optional(); + if (args.size() == 1) { + root = std::filesystem::path(args[0]); + } + if constexpr (std::is_same_v) { + return f(root, svsbenchmark::SaveDirectoryChecker()); + } else { + return f(root); + } + } +}; +} // namespace + +// Return an executor for this benchmark. +std::unique_ptr static_workflow() { + return std::make_unique>>(); +} +} // namespace svsbenchmark::ivf diff --git a/benchmark/src/ivf/search.cpp b/benchmark/src/ivf/search.cpp new file mode 100644 index 00000000..dfd29fc9 --- /dev/null +++ b/benchmark/src/ivf/search.cpp @@ -0,0 +1,120 @@ +/* + * Copyright 2025 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. + */ + +// svs-benchmark +#include "svs-benchmark/ivf/search.h" +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/executable.h" +#include "svs-benchmark/ivf/uncompressed.h" + +// svs +#include "svs/lib/saveload.h" +#include "svs/third-party/toml.h" + +// third-party +#include "fmt/core.h" +#include "fmt/ranges.h" + +// stl +#include +#include +#include +#include +#include + +namespace svsbenchmark::ivf { +namespace { + +const char* HELP = R"( +Run a search-only benchmark for the IVF index. + +Usage: + (1) src-file.toml (output-file.toml/--validate) [basename] + (2) --help + (3) --example + +1. Run all the benchmarks in the global `search_ivf_static` array in `src-file.toml`. + All elements in the array must be parseable as a ``svsbenchmark::ivf::SearchJob``. + + Results will be saved to `output-file.toml`. + + If `--validate` is given as the second argument, then all pre-run checks will be + performed on the input file and arguments but not benchmark will actually be run. + + Optional third argument `basename` will be used as the root for all file paths parsed. + +2. Print this help message. + +3. Display an example input TOML file to `stdout`. + +Backend specializations are dispatched on the following fields of the input TOML file: +* build_type: The dataset type to use. +* query_type: The element type of the query dataset. +* data_type: The input type of the source dataset. +* distance: The distance function to use. +* ndims: The compile-time dimensionality. + +Compiled specializations are listed below: +{build_type, query_type, data_type, distance, ndims} +)"; + +struct Exe { + public: + using job_type = ivf::SearchJob; + using dispatcher_type = ivf::StaticSearchDispatcher; + + static dispatcher_type dispatcher() { + auto dispatcher = dispatcher_type{}; + ivf::register_uncompressed_static_search(dispatcher); + return dispatcher; + } + + static std::string_view name() { return search_benchmark_name(); } + + static void print_help() { + fmt::print("{}", HELP); + auto f = dispatcher(); + for (size_t i = 0; i < f.size(); ++i) { + auto dispatch_strings = std::array{ + f.description(i, 0), + f.description(i, 1), + f.description(i, 2), + f.description(i, 3), + f.description(i, 4), + }; + fmt::print("{{ {} }}\n", fmt::join(dispatch_strings, ", ")); + } + } + + static job_type example() { return job_type::example(); } + + template + static std::optional> + parse_args_and_invoke(F&& f, std::span args) { + auto root = std::optional(); + if (args.size() == 1) { + root = std::filesystem::path(args[0]); + } + return f(root); + } +}; +} // namespace + +// Return an executor for this benchmark. +std::unique_ptr search_static_workflow() { + return std::make_unique>(); +} +} // namespace svsbenchmark::ivf diff --git a/benchmark/src/ivf/test.cpp b/benchmark/src/ivf/test.cpp new file mode 100644 index 00000000..e39268b9 --- /dev/null +++ b/benchmark/src/ivf/test.cpp @@ -0,0 +1,118 @@ +/* + * Copyright 2025 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. + */ + +// svs-benchmark +#include "svs-benchmark/ivf/test.h" +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/executable.h" +#include "svs-benchmark/ivf/uncompressed.h" + +// svs +#include "svs/lib/saveload.h" +#include "svs/third-party/toml.h" + +// third-party +#include "fmt/core.h" +#include "fmt/ranges.h" + +// stl +#include +#include +#include +#include +#include +#include + +namespace svsbenchmark::ivf { +namespace { + +const char* HELP = R"( +Generate reference results for the IVF index. + +Usage: + (1) src-file.toml output-file.toml num_threads [basename] + (2) --help + (3) --example + +1. Run the test generators using `src-file.toml` as the test driver input. (see (3)) + Store the post-processed results into `output-file.toml`. + Third argument `num_threads` sets the number of worker threads to use for each job. + Optional fourth argument `basename` will be used as the root for all file-paths parsed + from `src-file.toml`. + + The output results will be saved to `output-file.toml` as a dictionary with the following + structure: + + "ivf_test_search" : Array of serialized `svsbenchmark::ivf::ExpectedResult` for + each search-only job registered. None of these entries should have the + `build_parameters` field present. + + "ivf_test_build" : Array of serialized `svsbenchmark::ivf::ExpectedResult` for + each build-job registered. All of these entries should have the `build_parameters` + field present. + +2. Print this message. + +3. Display an example input TOML file to `stdout`. +)"; + +struct TestGenerator { + public: + // Type Aliases + using job_type = ivf::IVFTest; + using test_type = std::vector; + + constexpr TestGenerator() = default; + + static std::string_view name() { return ivf::test_benchmark_name(); } + + static test_type tests() { + auto generator = test_type{}; + svsbenchmark::append_to(generator, register_uncompressed_test_routines()); + return generator; + } + + static job_type example() { return job_type::example(); } + static void print_help() { fmt::print("{}", HELP); } + + template + static std::optional + parse_args_and_invoke(F&& f, std::span args) { + // We should have 1 or 2 additional arguments to parse, corresponding to the + // number of threads and an optional data root. + auto nargs = args.size(); + bool nargs_okay = nargs == 1 || nargs == 2; + if (!nargs_okay) { + fmt::print("Received too few arguments for Inverted test generation!"); + print_help(); + return std::nullopt; + } + + auto num_threads = std::stoull(std::string{args[0]}); + auto data_root = std::optional(); + if (nargs == 2) { + data_root = args[1]; + } + return f(num_threads, data_root); + } +}; +} // namespace + +// Return an executor for this benchmark. +std::unique_ptr test_generator() { + return std::make_unique>(); +} +} // namespace svsbenchmark::ivf diff --git a/benchmark/src/ivf/uncompressed.cpp b/benchmark/src/ivf/uncompressed.cpp new file mode 100644 index 00000000..2ddc7cfd --- /dev/null +++ b/benchmark/src/ivf/uncompressed.cpp @@ -0,0 +1,274 @@ +/* + * Copyright 2025 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. + */ + +// svs-benchmark +#include "svs-benchmark/ivf/uncompressed.h" +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/datasets/uncompressed.h" +#include "svs-benchmark/ivf/build.h" +#include "svs-benchmark/ivf/common.h" +#include "svs-benchmark/ivf/static_traits.h" + +// svs +#include "svs/core/distance.h" +#include "svs/lib/dispatcher.h" +#include "svs/third-party/toml.h" + +// stl +#include + +namespace svsbenchmark::ivf { + +namespace { + +// Specializations +#define X(Q, T, D, N) f.template operator()() +template void for_standard_specializations(F&& f) { + if constexpr (!is_minimal) { + X(float, svs::Float16, svs::distance::DistanceL2, 96); // deep + X(float, svs::Float16, svs::distance::DistanceL2, 100); // msturing + X(float, svs::Float16, svs::distance::DistanceIP, 200); // text2image + X(float, svs::Float16, svs::distance::DistanceIP, 512); // open-images, laion + X(float, svs::Float16, svs::distance::DistanceIP, 768); // dpr, rqa + // // Generic fallbacks + // X(float, svs::Float16, svs::distance::DistanceL2, svs::Dynamic); + // X(float, svs::Float16, svs::distance::DistanceIP, svs::Dynamic); + } +} +#undef X + +// Load and Search +template +toml::table run_static_search( + // dispatch arguments. + svsbenchmark::TypedUncompressed SVS_UNUSED(tag), + DispatchType SVS_UNUSED(query_type), + DispatchType SVS_UNUSED(data_type), + D distance, + svs::lib::ExtentTag SVS_UNUSED(extent), + // feed-forward arguments + const SearchJob& job +) { + auto tic = svs::lib::now(); + auto index = svs::IVF::assemble_from_file( + job.config_, + svs::data::SimpleData>::load(job.data_), + distance, + job.num_threads_ + ); + + double load_time = svs::lib::time_difference(tic); + auto queries = svs::data::SimpleData::load(job.queries_); + auto groundtruth = svs::data::SimpleData::load(job.groundtruth_); + + auto results = svsbenchmark::search::run_search( + index, + job, + search::QuerySet(queries, groundtruth, job.queries_in_training_set_), + svsbenchmark::LoadTime{load_time}, + svsbenchmark::Placeholder{} + ); + return svs::lib::save_to_table(results); +} + +// Static build and search +template +toml::table run_static_uncompressed( + // dispatch arguments. + TypedUncompressed SVS_UNUSED(tag), + DispatchType SVS_UNUSED(query_type), + DispatchType SVS_UNUSED(data_type), + D distance, + svs::lib::ExtentTag SVS_UNUSED(extent), + // feed-forward arguments + const BuildJob& job +) { + auto data = svs::data::SimpleData::load(job.data_); + auto tic = svs::lib::now(); + auto clustering = svs::IVF::build_clustering( + job.build_parameters_, data, distance, job.num_threads_ + ); + double build_time = svs::lib::time_difference(tic); + auto index = + svs::IVF::assemble_from_clustering(clustering, data, distance, job.num_threads_); + + // Save the index if requested by the caller. + job.maybe_save_index(clustering); + + // Load and run queries. + auto queries = svs::data::SimpleData::load(job.queries_); + auto groundtruth = svs::data::SimpleData::load(job.groundtruth_); + auto results = svsbenchmark::search::run_search( + index, + job, + search::QuerySet(queries, groundtruth, job.queries_in_training_set_), + svsbenchmark::BuildTime{build_time}, + svsbenchmark::Placeholder{} + ); + return svs::lib::save_to_table(results); +} + +template +svsbenchmark::TestFunctionReturn test_search(const IVFTest& job) { + // Get the groundtruth for the distance. + constexpr svs::DistanceType distance = svs::distance_type_v; + const auto& groundtruth_path = job.groundtruth_for(distance); + auto kind = svsbenchmark::Uncompressed(svs::datatype_v); + + // Construct a `SearchJob` for this operation. + auto search_job = SearchJob{ + "IVF uncompressed reference search", + kind, + job.index_config_, + job.graph_, + job.data_f32_, + job.queries_f32_, + groundtruth_path, + job.queries_in_training_set_, + svs::datatype_v, + svs::DataType::float32, + svs::distance_type_v, + Extent(svs::Dynamic), + job.num_threads_, + test_search_parameters(), + test_search_configs()}; + + // Load the components for the test. + auto tic = svs::lib::now(); + auto data_loader = svs::lib::Lazy([&]() { + return svsbenchmark::convert_data( + svs::lib::Type(), svs::data::SimpleData::load(job.data_f32_) + ); + }); + auto index = svs::IVF::assemble_from_file( + job.index_config_, data_loader, Distance(), job.num_threads_ + ); + double load_time = svs::lib::time_difference(tic); + auto queries = svs::data::SimpleData::load(job.queries_f32_); + auto groundtruth = svs::data::SimpleData::load(groundtruth_path); + + auto results = svsbenchmark::search::run_search( + index, + search_job, + svsbenchmark::search::QuerySet{ + std::move(queries), std::move(groundtruth), job.queries_in_training_set_}, + svsbenchmark::LoadTime{load_time}, + svsbenchmark::Placeholder{} + ); + + return TestFunctionReturn{ + .key_ = "ivf_test_search", + .results_ = svs::lib::save_to_table(ivf::ExpectedResult(std::move(kind), results))}; +} + +template +svsbenchmark::TestFunctionReturn test_build(const IVFTest& job) { + // Get the groundtruth for the distance. + constexpr svs::DistanceType distance = svs::distance_type_v; + const auto& groundtruth_path = job.groundtruth_for(distance); + + auto build_parameters = + svs::index::ivf::IVFBuildParameters{128, 10000, 10, Hierarchical, 0.1}; + + auto kind = svsbenchmark::Uncompressed(svs::datatype_v); + + // Construct a `SearchJob` for this operation. + auto build_job = BuildJob{ + groundtruth_path, + {{10, 1.0}, {50, 1.0}}, + test_search_parameters(), + std::nullopt, + "IVF uncompressed reference build", + kind, + job.data_f32_, + job.queries_f32_, + job.queries_in_training_set_, + svs::datatype_v, + svs::DataType::float32, + svs::distance_type_v, + Extent(svs::Dynamic), + build_parameters, + job.num_threads_}; + + // Load the components for the test. + auto data = svsbenchmark::convert_data( + svs::lib::Type(), svs::data::SimpleData::load(job.data_f32_) + ); + auto tic = svs::lib::now(); + auto clustering = svs::IVF::build_clustering( + build_parameters, data, Distance(), job.num_threads_ + ); + double build_time = svs::lib::time_difference(tic); + auto index = svs::IVF::assemble_from_clustering( + clustering, data, distance, job.num_threads_ + ); + + auto queries = svs::data::SimpleData::load(job.queries_f32_); + auto groundtruth = svs::data::SimpleData::load(groundtruth_path); + + auto results = svsbenchmark::search::run_search( + index, + build_job, + svsbenchmark::search::QuerySet{ + std::move(queries), std::move(groundtruth), job.queries_in_training_set_}, + svsbenchmark::BuildTime{build_time}, + svsbenchmark::Placeholder{} + ); + + return TestFunctionReturn{ + .key_ = "ivf_test_build", + .results_ = svs::lib::save_to_table(ivf::ExpectedResult(std::move(kind), results))}; +} + +} // namespace + +// target-registration. +void register_uncompressed_static_search(ivf::StaticSearchDispatcher& dispatcher) { + for_standard_specializations( + [&dispatcher]() { + auto method = &run_static_search; + dispatcher.register_target(svs::lib::dispatcher_build_docs, method); + } + ); +} + +void register_uncompressed_static_build(ivf::StaticBuildDispatcher& dispatcher) { + for_standard_specializations( + [&dispatcher]() { + auto method = &run_static_uncompressed; + dispatcher.register_target(svs::lib::dispatcher_build_docs, method); + } + ); +} + +std::vector register_uncompressed_test_routines() { + if constexpr (build_test_generators) { + return std::vector({ + // Searching + &test_search, + &test_search, + // Building + &test_build, + &test_build, + &test_build, + &test_build, + }); + } else { + return std::vector(); + } +} + +} // namespace svsbenchmark::ivf diff --git a/benchmark/src/main.cpp b/benchmark/src/main.cpp index 7c06fb28..c8db6003 100644 --- a/benchmark/src/main.cpp +++ b/benchmark/src/main.cpp @@ -25,6 +25,14 @@ // inverted #include "svs-benchmark/inverted/inverted.h" +// ivf +SVS_VALIDATE_BOOL_ENV(SVS_ENABLE_IVF) +#if SVS_ENABLE_IVF +#include "svs-benchmark/ivf/build.h" +#include "svs-benchmark/ivf/search.h" +#include "svs-benchmark/ivf/test.h" +#endif // SVS_ENABLE_IVF + // stl #include #include @@ -43,6 +51,13 @@ svsbenchmark::ExecutableDispatcher build_dispatcher() { dispatcher.register_executable(svsbenchmark::vamana::iterator_benchmark()); // inverted svsbenchmark::inverted::register_executables(dispatcher); + // ivf + SVS_VALIDATE_BOOL_ENV(SVS_ENABLE_IVF) +#if SVS_ENABLE_IVF + dispatcher.register_executable(svsbenchmark::ivf::search_static_workflow()); + dispatcher.register_executable(svsbenchmark::ivf::static_workflow()); + dispatcher.register_executable(svsbenchmark::ivf::test_generator()); +#endif // SVS_ENABLE_IVF // documentation svsbenchmark::register_dataset_documentation(dispatcher); return dispatcher; diff --git a/bindings/python/CMakeLists.txt b/bindings/python/CMakeLists.txt index 279b51df..81ae4f2a 100644 --- a/bindings/python/CMakeLists.txt +++ b/bindings/python/CMakeLists.txt @@ -39,6 +39,14 @@ set(CPP_FILES src/svs_mkl.cpp ) +# ivf +if (SVS_EXPERIMENTAL_ENABLE_IVF) + list(APPEND CPP_FILES + src/ivf.cpp + ) +endif() + + set(LIB_NAME "_svs") pybind11_add_module(${LIB_NAME} MODULE ${CPP_FILES}) target_link_libraries(${LIB_NAME} PRIVATE pybind11::module) diff --git a/bindings/python/include/svs/python/core.h b/bindings/python/include/svs/python/core.h index ed378375..50864281 100644 --- a/bindings/python/include/svs/python/core.h +++ b/bindings/python/include/svs/python/core.h @@ -104,7 +104,7 @@ class UnspecializedGraphLoader { const std::filesystem::path& path() const { return path_; } const Allocator& allocator() const { return allocator_; } - svs::graphs::SimpleGraph load() const { + auto load() const { using other = std::allocator_traits::rebind_alloc; return svs::graphs::SimpleGraph::load(path_, other(allocator_)); } diff --git a/bindings/python/include/svs/python/ivf.h b/bindings/python/include/svs/python/ivf.h new file mode 100644 index 00000000..5f0dc0a3 --- /dev/null +++ b/bindings/python/include/svs/python/ivf.h @@ -0,0 +1,94 @@ +/* + * Copyright 2025 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 + +// svs python bindings +#include "svs/python/common.h" +#include "svs/python/core.h" + +#include "svs/core/distance.h" +#include "svs/lib/bfloat16.h" +#include "svs/lib/datatype.h" +#include "svs/lib/float16.h" +#include "svs/lib/meta.h" +#include "svs/lib/misc.h" + +#include +#include + +namespace svs::python { +namespace ivf_specializations { +/// +/// Flag to selectively enable index building. +/// +enum class EnableBuild { None, FromFile, FromFileAndArray }; + +template +inline constexpr bool enable_build_from_file = + (B == EnableBuild::FromFile || B == EnableBuild::FromFileAndArray); + +template +inline constexpr bool enable_build_from_array = (B == EnableBuild::FromFileAndArray); + +// Define all desired specializations for searching and building. +template void for_standard_specializations(F&& f) { +#define X(Q, T, N, B) f.template operator()(); +#define XN(Q, T, N) X(Q, T, N, EnableBuild::None) + // Pattern: + // QueryType, DataType, Dimensionality, Enable Building + // clang-format off + X(float, svs::BFloat16, 512, EnableBuild::FromFileAndArray); + + XN(float, float, 512); + XN(float, svs::Float16, 512); + + X(float, svs::BFloat16, Dynamic, EnableBuild::FromFileAndArray); + XN(float, float, Dynamic); + XN(float, svs::Float16, Dynamic); + // clang-format on +#undef XN +#undef X +} +} // namespace ivf_specializations + +namespace ivf { +template void add_interface(pybind11::class_& manager) { + manager.def_property_readonly( + "experimental_backend_string", + &Manager::experimental_backend_string, + R"( + Read Only (str): Get a string identifying the full-type of the backend implementation. + + This property is experimental and subject to change without a deprecation warning.)" + ); + + manager.def_property( + "search_parameters", + &Manager::get_search_parameters, + &Manager::set_search_parameters, + R"( + "Read/Write (svs.IVFSearchParameters): Get/set the current search parameters for the + index. These parameters modify both the algorithmic properties of search (affecting recall) + and non-algorthmic properties of search (affecting queries-per-second). + + See also: `svs.IVFSearchParameters`.)" + ); +} + +void wrap(pybind11::module& m); +} // namespace ivf +} // namespace svs::python diff --git a/bindings/python/src/ivf.cpp b/bindings/python/src/ivf.cpp new file mode 100644 index 00000000..443d358c --- /dev/null +++ b/bindings/python/src/ivf.cpp @@ -0,0 +1,643 @@ +/* + * Copyright 2025 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. + */ + +// svs python bindings +#include "svs/python/ivf.h" +#include "svs/python/common.h" +#include "svs/python/core.h" +#include "svs/python/dispatch.h" +#include "svs/python/manager.h" + +// svs +#include "svs/core/data/simple.h" +#include "svs/core/distance.h" +#include "svs/lib/array.h" +#include "svs/lib/datatype.h" +#include "svs/lib/dispatcher.h" +#include "svs/lib/float16.h" +#include "svs/lib/meta.h" +#include "svs/orchestrators/ivf.h" + +// pybind +#include +#include +#include + +// stl +#include +#include +#include +#include + +///// +///// IVF +///// + +namespace py = pybind11; +using namespace svs::python::ivf_specializations; + +namespace svs::python::ivf { +// The build process in IVF uses Kmeans to get centroids and assignments of data. +// This sparse clustering can be saved with centroids stored as float datatype. +// While assembling, the sparse clustering is used to create DenseClusters and +// centroids datatype can be changed as per the search specializations. +// By default, BFloat16 centroids are used to take advantage of AMX +// template +using Clustering = + svs::index::ivf::Clustering, uint32_t>; + +namespace detail { + +///// +///// Assembly +///// + +template +svs::IVF assemble_uncompressed( + Clustering clustering, + svs::VectorDataLoader> data, + svs::DistanceType distance_type, + size_t num_threads, + size_t intra_query_threads = 1 +) { + return svs::IVF::assemble_from_clustering( + std::move(clustering), + std::move(data), + distance_type, + num_threads, + intra_query_threads + ); +} + +template +void register_uncompressed_ivf_assemble(Dispatcher& dispatcher) { + for_standard_specializations( + [&dispatcher]() { + auto method = &assemble_uncompressed; + dispatcher.register_target(svs::lib::dispatcher_build_docs, method); + } + ); +} + +template void register_ivf_assembly(Dispatcher& dispatcher) { + register_uncompressed_ivf_assemble(dispatcher); +} + +///// +///// Assembly from File +///// +// N.B (IB): quite a bit of repetition in Assemble and AssembleFromFile functions. +// Loading the cluster from file and then using the Assemble from clustering +// shows performance loss, mainly due to the threadpool used for loading. +// This needs to be revisited. +template +svs::IVF assemble_from_file_uncompressed( + const std::filesystem::path& cluster_path, + svs::VectorDataLoader> data, + svs::DistanceType distance_type, + size_t num_threads, + size_t intra_query_threads = 1 +) { + return svs::IVF::assemble_from_file( + cluster_path, std::move(data), distance_type, num_threads, intra_query_threads + ); +} + +template +void register_uncompressed_ivf_assemble_from_file(Dispatcher& dispatcher) { + for_standard_specializations( + [&dispatcher]() { + auto method = &assemble_from_file_uncompressed; + dispatcher.register_target(svs::lib::dispatcher_build_docs, method); + } + ); +} + +template +void register_ivf_assembly_from_file(Dispatcher& dispatcher) { + register_uncompressed_ivf_assemble_from_file(dispatcher); +} + +using IVFAssembleTypes = + std::variant; + +///// +///// Build From File +///// + +template +Clustering build_uncompressed( + const svs::index::ivf::IVFBuildParameters& parameters, + svs::VectorDataLoader> data, + svs::DistanceType distance_type, + size_t num_threads +) { + return svs::IVF::build_clustering( + parameters, std::move(data), distance_type, num_threads + ); +} + +template +void register_uncompressed_ivf_build_from_file(Dispatcher& dispatcher) { + for_standard_specializations( + [&dispatcher]() { + if constexpr (enable_build_from_file) { + auto method = &build_uncompressed; + dispatcher.register_target(svs::lib::dispatcher_build_docs, method); + } + } + ); +} + +template void register_ivf_build_from_file(Dispatcher& dispatcher) { + register_uncompressed_ivf_build_from_file(dispatcher); +} + +using IVFBuildTypes = std::variant; + +///// +///// Build from Array +///// + +template +Clustering uncompressed_build_from_array( + const svs::index::ivf::IVFBuildParameters& parameters, + svs::data::ConstSimpleDataView view, + svs::DistanceType distance_type, + size_t num_threads +) { + auto data = + svs::data::SimpleData>(view.size(), view.dimensions()); + svs::data::copy(view, data); + return svs::IVF::build_clustering( + parameters, std::move(data), distance_type, num_threads + ); +} + +template void register_ivf_build_from_array(Dispatcher& dispatcher) { + for_standard_specializations( + [&dispatcher]() { + if constexpr (enable_build_from_array) { + auto method = &uncompressed_build_from_array; + dispatcher.register_target(svs::lib::dispatcher_build_docs, method); + } + } + ); +} + +///// +///// Dispatch Invocation +///// + +using AssemblyDispatcher = svs::lib:: + Dispatcher; + +AssemblyDispatcher assembly_dispatcher() { + auto dispatcher = AssemblyDispatcher{}; + + // Register available backend methods. + register_ivf_assembly(dispatcher); + return dispatcher; +} + +// Assemble +svs::IVF assemble_from_clustering( + Clustering clustering, + IVFAssembleTypes data_kind, + svs::DistanceType distance_type, + svs::DataType SVS_UNUSED(query_type), + bool SVS_UNUSED(enforce_dims), + size_t num_threads, + size_t intra_query_threads = 1 +) { + return assembly_dispatcher().invoke( + std::move(clustering), + std::move(data_kind), + distance_type, + num_threads, + intra_query_threads + ); +} + +using AssemblyFromFileDispatcher = svs::lib::Dispatcher< + svs::IVF, + const std::filesystem::path&, + IVFAssembleTypes, + svs::DistanceType, + size_t, + size_t>; + +AssemblyFromFileDispatcher assembly_from_file_dispatcher() { + auto dispatcher = AssemblyFromFileDispatcher{}; + + // Register available backend methods. + register_ivf_assembly_from_file(dispatcher); + return dispatcher; +} + +// Assemble from file +svs::IVF assemble_from_file( + const std::string& cluster_path, + IVFAssembleTypes data_kind, + svs::DistanceType distance_type, + svs::DataType SVS_UNUSED(query_type), + bool SVS_UNUSED(enforce_dims), + size_t num_threads, + size_t intra_query_threads = 1 +) { + return assembly_from_file_dispatcher().invoke( + cluster_path, std::move(data_kind), distance_type, num_threads, intra_query_threads + ); +} + +constexpr std::string_view ASSEMBLE_DOCSTRING_PROTO = R"( +Assemble a searchable IVF index from provided clustering and data + +Args: + clustering_path/clustering: Path to the directory where the clustering was generated. + OR directly provide the loaded Clustering. + data_loader: The loader for the dataset. See comment below for accepted types. + distance: The distance function to use. + query_type: The data type of the queries. + enforce_dims: Require that the compiled dimensionality of the returned index matches + the dimensionality provided in the ``data_loader`` argument. If a match is not + found, an exception is thrown. + + This is meant to ensure that specialized dimensionality is provided without falling + back to generic implementations. Leaving the ``dims`` out when constructing the + ``data_loader`` will with `enable_dims = True` will always attempt to use a generic + implementation. + num_threads: The number of threads to use for queries (can't be changed after loading). + intra_query_threads: (default: 1) these many threads work on a single query. + Total number of threads required = ``query_batch_size`` * ``intra_query_threads``. + Where ``query_batch_size`` is the number of queries processed in parallel. + Use this parameter only when the ``query_batch_size`` is smaller and ensure your + system has sufficient threads available. Set ``num_threads`` = ``query_batch_size`` + +The top level type is an abstract type backed by various specialized backends that will +be instantiated based on their applicability to the particular problem instance. + +The arguments upon which specialization is conducted are: + +* `data_loader`: Both kind (type of loader) and inner aspects of the loader like data type, + quantization type, and number of dimensions. +* `distance`: The distance measure being used. + +Specializations compiled into the binary are listed below. + +{} +)"; + +void wrap_assemble(py::class_& ivf) { + auto dispatcher = assembly_dispatcher(); + // Procedurally generate the dispatch string. + auto dynamic = std::string{}; + for (size_t i = 0; i < dispatcher.size(); ++i) { + fmt::format_to( + std::back_inserter(dynamic), + R"( +Method {}: + - data_loader: {} + - distance: {} +)", + i, + dispatcher.description(i, 2), + dispatcher.description(i, 3) + ); + } + + ivf.def_static( + "assemble_from_clustering", + &detail::assemble_from_clustering, + py::arg("clustering"), + py::arg("data_loader"), + py::arg("distance") = svs::L2, + py::arg("query_type") = svs::DataType::float32, + py::arg("enforce_dims") = false, + py::arg("num_threads") = 1, + py::arg("intra_query_threads") = 1, + fmt::format(ASSEMBLE_DOCSTRING_PROTO, dynamic).c_str() + ); + ivf.def_static( + "assemble_from_file", + &detail::assemble_from_file, + py::arg("clustering_path"), + py::arg("data_loader"), + py::arg("distance") = svs::L2, + py::arg("query_type") = svs::DataType::float32, + py::arg("enforce_dims") = false, + py::arg("num_threads") = 1, + py::arg("intra_query_threads") = 1, + fmt::format(ASSEMBLE_DOCSTRING_PROTO, dynamic).c_str() + ); +} + +// Build from file +using BuildFromFileDispatcher = svs::lib::Dispatcher< + Clustering, + const svs::index::ivf::IVFBuildParameters&, + IVFBuildTypes, + svs::DistanceType, + size_t>; + +BuildFromFileDispatcher build_from_file_dispatcher() { + auto dispatcher = BuildFromFileDispatcher{}; + register_ivf_build_from_file(dispatcher); + return dispatcher; +} + +Clustering build_from_file( + const svs::index::ivf::IVFBuildParameters& parameters, + IVFBuildTypes data_source, + svs::DistanceType distance_type, + size_t num_threads +) { + return build_from_file_dispatcher().invoke( + parameters, std::move(data_source), distance_type, num_threads + ); +} + +// Build from array. +using BuildFromArrayDispatcher = svs::lib::Dispatcher< + Clustering, + const svs::index::ivf::IVFBuildParameters&, + AnonymousVectorData, + svs::DistanceType, + size_t>; + +BuildFromArrayDispatcher build_from_array_dispatcher() { + auto dispatcher = BuildFromArrayDispatcher{}; + register_ivf_build_from_array(dispatcher); + return dispatcher; +} + +Clustering build_from_array( + const svs::index::ivf::IVFBuildParameters& parameters, + AnonymousVectorData py_data, + svs::DistanceType distance_type, + size_t num_threads +) { + return build_from_array_dispatcher().invoke( + parameters, py_data, distance_type, num_threads + ); +} + +template +void add_build_specialization(py::class_& clustering) { + clustering.def_static( + "build", + [](const svs::index::ivf::IVFBuildParameters& parameters, + py_contiguous_array_t py_data, + svs::DistanceType distance, + size_t num_threads) { + return build_from_array( + parameters, AnonymousVectorData(py_data), distance, num_threads + ); + }, + py::arg("build_parameters"), + py::arg("py_data"), + py::arg("distance"), + py::arg("num_threads") = 1, + R"( + Build IVF clustering over the given data and return a sparse clustering. + Use this clustering to assemble a searcheable IVF index. + + Args: + parameters: Parameters controlling kmeans clustering. + py_data: The dataset to index. + distance: The distance type to use for this dataset. + num_threads: The number of threads to use for index construction. Default: 1. +)" + ); +} + +void wrap_build_from_file(py::class_& clustering) { + constexpr std::string_view docstring_proto = R"( + Build IVF clustering over the given data and return a sparse clustering. + Use the returned clustering to assemble a searcheable IVF index. + + Args: + build_parameters (:py:class:`svs.IVFBuildParameters`): Hyper-parameters + controlling clustering build. + data_loader: The source of the data on-disk. Can either be + :py:class:`svs.DataFile` to represent a standard uncompressed dataset + distance: The similarity-function to use for this index. + num_threads: The number of threads to use for index construction. Default: 1. + + The top level type is an abstract type backed by various specialized backends that will + be instantiated based on their applicability to the particular problem instance. + + The arguments upon which specialization is conducted are: + +* `data_loader`: Only uncompressed data types are supported for IVF cluster building +* `distance`: The distance measure being used. + + Specializations compiled into the binary are listed below. + +{} +)"; + + auto dispatcher = build_from_file_dispatcher(); + // Procedurally generate the dispatch string. + auto dynamic = std::string{}; + for (size_t i = 0; i < dispatcher.size(); ++i) { + fmt::format_to( + std::back_inserter(dynamic), + R"( + Method {}: + - data_loader: {} + - distance: {} +)", + i, + dispatcher.description(i, 1), + dispatcher.description(i, 2) + ); + } + + clustering.def_static( + "build", + &detail::build_from_file, + py::arg("build_parameters"), + py::arg("data_loader"), + py::arg("distance"), + py::arg("num_threads") = 1, + fmt::format(docstring_proto, dynamic).c_str() + ); +} + +// Save the sparse clustering to a directory +void save_clustering(Clustering& clustering, const std::string& clustering_path) { + svs::lib::save_to_disk(clustering, clustering_path); +} + +// Save the sparse clustering to a directory +auto load_clustering(const std::string& clustering_path, size_t num_threads = 1) { + auto threadpool = threads::as_threadpool(num_threads); + return svs::lib::load_from_disk(clustering_path, threadpool); +} + +} // namespace detail + +void wrap(py::module& m) { + // wrap_common(m); + + /// Build Parameters + using IVFBuildParameters = svs::index::ivf::IVFBuildParameters; + py::class_ parameters( + m, "IVFBuildParameters", "Build parameters for kmeans clustering." + ); + + parameters + .def( + py::init([](size_t num_centroids, + size_t minibatch_size, + size_t num_iterations, + bool is_hierarchical, + float training_fraction, + size_t hierarchical_level1_clusters, + size_t seed) { + return svs::index::ivf::IVFBuildParameters{ + num_centroids, + minibatch_size, + num_iterations, + is_hierarchical, + training_fraction, + hierarchical_level1_clusters, + seed}; + }), + py::arg("num_centroids") = 1000, + py::arg("minibatch_size") = 10'000, + py::arg("num_iterations") = 10, + py::arg("is_hierarchical") = false, + py::arg("training_fraction") = 0.1, + py::arg("hierarchical_level1_clusters") = 0, + py::arg("seed") = 0xc0ffee, + R"( + Construct a new instance from keyword arguments. + + Args: + num_centroids: The target number of clusters in the final result. + minibatch_size: The size of each minibatch used to process data at a time. + num_iterations: The number of iterations used in kmeans training. + is_hierarchical: Use hierarchical Kmeans or not. + training_fraction: Fraction of dataset used for training + hierarchical_level1_clusters: Level1 clusters for hierarchical kmeans. + Use heuristic if 0. + seed: The initial seed for the random number generator. + )" + ) + .def_readwrite("num_centroids", &IVFBuildParameters::num_centroids_) + .def_readwrite("minibatch_size", &IVFBuildParameters::minibatch_size_) + .def_readwrite("num_iterations", &IVFBuildParameters::num_iterations_) + .def_readwrite("is_hierarchical", &IVFBuildParameters::is_hierarchical_) + .def_readwrite("training_fraction", &IVFBuildParameters::training_fraction_) + .def_readwrite( + "hierarchical_level1_clusters", + &IVFBuildParameters::hierarchical_level1_clusters_ + ); + + /// Search Parameters + using IVFSearchParameters = svs::index::ivf::IVFSearchParameters; + auto params = py::class_{ + m, + "IVFSearchParameters", + R"( + Parameters controlling recall and performance of the IVF Index. + Args: + n_probes: The number of nearest clusters to be explored + k_reorder: Level of reordering or reranking done when using compressed datasets + )"}; + + params + .def(py::init(), py::arg("n_probes") = 1, py::arg("k_reorder") = 1.0) + .def_readwrite("n_probes", &IVFSearchParameters::n_probes_) + .def_readwrite("k_reorder", &IVFSearchParameters::k_reorder_); + + /// + /// IVF Static Module + /// + std::string name = "IVF"; + py::class_ ivf(m, name.c_str(), "Top level class for the IVF index."); + + detail::wrap_assemble(ivf); + + // Make the IVF type searchable. + add_search_specialization(ivf); + add_search_specialization(ivf); + + // Add threading layer. + add_threading_interface(ivf); + + // Data layer + add_data_interface(ivf); + + // IVF Specific Extensions. + add_interface(ivf); + + // Reconstruction. + // add_reconstruct_interface(ivf); + + name = "Clustering"; + py::class_ clustering( + m, name.c_str(), "Top level class for sparse IVF clustering" + ); + + /// Index building + // Build from Numpy array. + detail::add_build_specialization(clustering); + detail::add_build_specialization(clustering); + detail::add_build_specialization(clustering); + + // Build from datasets on file. + detail::wrap_build_from_file(clustering); + + /// Index Saving and Loading. + clustering.def( + "save_clustering", + &detail::save_clustering, + py::arg("clustering_directory"), + R"( + Save a constructed IVF clustering to disk (useful following build). + + Args: + clustering_directory: Directory where clustering will be saved. + + Note: All directories should be separate to avoid accidental name collision + with any auxiliary files that are needed when saving the various components of + the index. + + If the directory does not exist, it will be created if its parent exists. + + It is the caller's responsibilty to ensure that no existing data will be + overwritten when saving the index to this directory. + )" + ); + clustering.def_static( + "load_clustering", + &detail::load_clustering, + py::arg("clustering_directory"), + py::arg("num_threads") = 1, + R"( + Load IVF clustering from disk (maybe used before assembling). + + Args: + clustering_directory: Directory from where to load the clustering. + num_threads: Number of threads to use when loading (default: 1). + )" + ); +} + +} // namespace svs::python::ivf diff --git a/bindings/python/src/python_bindings.cpp b/bindings/python/src/python_bindings.cpp index 7d3b3988..ae79638a 100644 --- a/bindings/python/src/python_bindings.cpp +++ b/bindings/python/src/python_bindings.cpp @@ -20,6 +20,12 @@ #include "svs/python/core.h" #include "svs/python/dynamic_vamana.h" #include "svs/python/flat.h" + +SVS_VALIDATE_BOOL_ENV(SVS_ENABLE_IVF) +#if SVS_ENABLE_IVF +#include "svs/python/ivf.h" +#endif // SVS_ENABLE_IVF + #include "svs/python/svs_mkl.h" #include "svs/python/vamana.h" @@ -27,6 +33,7 @@ #include "svs/core/distance.h" #include "svs/core/io.h" #include "svs/lib/array.h" +#include "svs/lib/bfloat16.h" #include "svs/lib/datatype.h" #include "svs/lib/float16.h" #include "svs/third-party/toml.h" @@ -59,6 +66,17 @@ void convert_fvecs_to_float16( } } +// Convert fvecs to bfloat16 +void convert_fvecs_to_bfloat16( + const std::string& filename_f32, const std::string& filename_bf16 +) { + auto reader = svs::io::vecs::VecsReader{filename_f32}; + auto writer = svs::io::vecs::VecsWriter{filename_bf16, reader.ndims()}; + for (auto i : reader) { + writer << i; + } +} + // Convert fvecs to svs - typed template void convert_vecs_to_svs_impl(const std::string& vecs_file, const std::string& svs_file) { @@ -91,13 +109,14 @@ void wrap_conversion(py::module& m) { Convert the vecs file (containing the specified element types) to the svs native format. Args: - vecs_file: The source [f/h/i/b]vecs file. + vecs_file: The source [f/h/bfi/b]vecs file. svs_file: The destination native file. dtype: The svs.DataType of the vecs file. Supported types: ({}). File extension type map: * fvecs = svs.DataType.float32 * hvecs = svs.DataType.float16 +* bfvecs = svs.DataType.bfloat16 * ivecs = svs.DataType.uint32 * bvecs = svs.DataType.uint8 )"; @@ -165,6 +184,7 @@ PYBIND11_MODULE(_svs, m) { .value("int32", svs::DataType::int32, "32-bit signed integer.") .value("int64", svs::DataType::int64, "64-bit signed integer.") .value("float16", svs::DataType::float16, "16-bit IEEE floating point.") + .value("bfloat16", svs::DataType::bfloat16, "16-bit brain floating point.") .value("float32", svs::DataType::float32, "32-bit IEEE floating point.") .value("float64", svs::DataType::float64, "64-bit IEEE floating point.") .export_values(); @@ -179,6 +199,21 @@ PYBIND11_MODULE(_svs, m) { Convert the `fvecs` file on disk with 32-bit floating point entries to a `fvecs` file with 16-bit floating point entries. +Args: + source_file: The source file path to convert. + destination_file: The destination file to generate. + )" + ); + + m.def( + "convert_fvecs_to_bfloat16", + &convert_fvecs_to_bfloat16, + py::arg("source_file"), + py::arg("destination_file"), + R"( +Convert the `fvecs` file on disk with 32-bit floating point entries to a `fvecs` file with +16-bit brain floating (bfloat16) point entries. + Args: source_file: The source file path to convert. destination_file: The destination file to generate. @@ -213,4 +248,10 @@ Convert the `fvecs` file on disk with 32-bit floating point entries to a `fvecs` // Vamana svs::python::vamana::wrap(m); svs::python::dynamic_vamana::wrap(m); + + // IVF + SVS_VALIDATE_BOOL_ENV(SVS_ENABLE_IVF) +#if SVS_ENABLE_IVF + svs::python::ivf::wrap(m); +#endif // SVS_ENABLE_IVF } diff --git a/cmake/mkl.cmake b/cmake/mkl.cmake index 13af27fc..088cfbea 100644 --- a/cmake/mkl.cmake +++ b/cmake/mkl.cmake @@ -45,6 +45,11 @@ if (SVS_EXPERIMENTAL_BUILD_CUSTOM_MKL) NO_DEFAULT_PATH REQUIRED ) + if (SVS_EXPERIMENTAL_ENABLE_IVF) + set(mkl_functions_file ${CMAKE_CURRENT_LIST_DIR}/mkl_functions_ivf) + else() + set(mkl_functions_file ${CMAKE_CURRENT_LIST_DIR}/mkl_functions) + endif() # This command creates the custom shared-object that will be linked to. add_custom_command( @@ -56,7 +61,7 @@ if (SVS_EXPERIMENTAL_BUILD_CUSTOM_MKL) "libintel64" "interface=ilp64" "threading=sequential" - "export=${CMAKE_CURRENT_LIST_DIR}/mkl_functions" + "export=${mkl_functions_file}" "MKLROOT=${MKL_ROOT}" "name=${CMAKE_CURRENT_BINARY_DIR}/${SVS_MKL_CUSTOM_LIBRARY_NAME}" ) diff --git a/cmake/mkl_functions_ivf b/cmake/mkl_functions_ivf new file mode 100644 index 00000000..951760fa --- /dev/null +++ b/cmake/mkl_functions_ivf @@ -0,0 +1,7 @@ +cblas_sgemm +mkl_simatcopy +LAPACKE_sgesvd +mkl_get_max_threads +MKL_Free_Buffers +cblas_gemm_bf16bf16f32 +cblas_gemm_f16f16f32 diff --git a/cmake/options.cmake b/cmake/options.cmake index e374a548..2f9dd416 100644 --- a/cmake/options.cmake +++ b/cmake/options.cmake @@ -92,6 +92,11 @@ option(SVS_EXPERIMENTAL_ENABLE_NUMA OFF # disabled by default ) +option(SVS_EXPERIMENTAL_ENABLE_IVF + "Enable IVF implementation. Requires Intel(R) MKL support" + OFF # disabled by default +) + ##### ##### svsbenchmark ##### @@ -140,6 +145,12 @@ else() target_compile_options(${SVS_LIB} INTERFACE -DSVS_INITIALIZE_LOGGER=0) endif() +if (SVS_EXPERIMENTAL_ENABLE_IVF) + target_compile_options(${SVS_LIB} INTERFACE -DSVS_ENABLE_IVF=1) +else() + target_compile_options(${SVS_LIB} INTERFACE -DSVS_ENABLE_IVF=0) +endif() + ##### ##### Helper target to apply relevant compiler optimizations. ##### @@ -161,6 +172,7 @@ target_compile_options( -Wextra -Wpedantic -Wno-gnu-zero-variadic-macro-arguments + -Wno-address-of-packed-member # When calling Intel(R) MKL GEMMs with BFloat16/Float16 -Wno-parentheses # GCC in CI has issues without it ) diff --git a/data/test_dataset/ivf_clustering/clusters_0.bin b/data/test_dataset/ivf_clustering/clusters_0.bin new file mode 100644 index 00000000..9e5f9566 Binary files /dev/null and b/data/test_dataset/ivf_clustering/clusters_0.bin differ diff --git a/data/test_dataset/ivf_clustering/data_1.svs b/data/test_dataset/ivf_clustering/data_1.svs new file mode 100644 index 00000000..5c760e8d Binary files /dev/null and b/data/test_dataset/ivf_clustering/data_1.svs differ diff --git a/data/test_dataset/ivf_clustering/svs_config.toml b/data/test_dataset/ivf_clustering/svs_config.toml new file mode 100644 index 00000000..1cf94f94 --- /dev/null +++ b/data/test_dataset/ivf_clustering/svs_config.toml @@ -0,0 +1,34 @@ +# Copyright 2025 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. + +__version__ = 'v0.0.2' + +[object] +__schema__ = 'IVF clustering' +__version__ = 'v0.0.0' +data_type = 'bfloat16' +filepath = 'clusters_0.bin' +filesize = 41032 +integer_type = 'uint32' +num_clusters = 128 + + [object.centroids] + __schema__ = 'uncompressed_data' + __version__ = 'v0.0.0' + binary_file = 'data_1.svs' + dims = 128 + eltype = 'bfloat16' + name = 'uncompressed' + num_vectors = 128 + uuid = '0cd42b81-8e7a-4fdd-b4d6-81d6d5880fb0' diff --git a/data/test_dataset/reference/ivf_reference.toml b/data/test_dataset/reference/ivf_reference.toml new file mode 100644 index 00000000..4edf53c9 --- /dev/null +++ b/data/test_dataset/reference/ivf_reference.toml @@ -0,0 +1,570 @@ +# Copyright 2025 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. + +start_time = 2024-09-18T19:44:08 +stop_time = 2024-09-18T19:44:52 + +[[ivf_test_build]] +__schema__ = 'benchmark_expected_result' +__version__ = 'v0.0.0' +distance = 'L2' + + [ivf_test_build.build_parameters] + __schema__ = 'ivf_build_parameters' + __version__ = 'v0.0.0' + hierarchical_level1_clusters = 0 + is_hierarchical = true + minibatch_size = 10000 + num_centroids = 128 + num_iterations = 10 + seed = 12648430 + training_fraction = 0.10000000149011612 + + [ivf_test_build.dataset] + __schema__ = 'benchmark_dataset_abstract' + __version__ = 'v0.0.0' + kind = 'uncompressed' + + [ivf_test_build.dataset.dataset] + __schema__ = 'benchmark_dataset_uncompressed' + __version__ = 'v0.0.0' + data_type = 'float32' + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.55700000000000005 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 10 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.98044444444444445 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 50 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.52144444444444449 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 9 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.80588888888888888 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 22 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.90688888888888886 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 32 + +[[ivf_test_build]] +__schema__ = 'benchmark_expected_result' +__version__ = 'v0.0.0' +distance = 'MIP' + + [ivf_test_build.build_parameters] + __schema__ = 'ivf_build_parameters' + __version__ = 'v0.0.0' + hierarchical_level1_clusters = 0 + is_hierarchical = true + minibatch_size = 10000 + num_centroids = 128 + num_iterations = 10 + seed = 12648430 + training_fraction = 0.10000000149011612 + + [ivf_test_build.dataset] + __schema__ = 'benchmark_dataset_abstract' + __version__ = 'v0.0.0' + kind = 'uncompressed' + + [ivf_test_build.dataset.dataset] + __schema__ = 'benchmark_dataset_uncompressed' + __version__ = 'v0.0.0' + data_type = 'float16' + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.24455555555555555 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 10 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.67688888888888887 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 50 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.50755555555555554 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 31 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.80544444444444441 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 70 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.90300000000000002 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 91 + +[[ivf_test_build]] +__schema__ = 'benchmark_expected_result' +__version__ = 'v0.0.0' +distance = 'L2' + + [ivf_test_build.build_parameters] + __schema__ = 'ivf_build_parameters' + __version__ = 'v0.0.0' + hierarchical_level1_clusters = 0 + is_hierarchical = true + minibatch_size = 10000 + num_centroids = 128 + num_iterations = 10 + seed = 12648430 + training_fraction = 0.10000000149011612 + + [ivf_test_build.dataset] + __schema__ = 'benchmark_dataset_abstract' + __version__ = 'v0.0.0' + kind = 'uncompressed' + + [ivf_test_build.dataset.dataset] + __schema__ = 'benchmark_dataset_uncompressed' + __version__ = 'v0.0.0' + data_type = 'bfloat16' + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.54955555555555557 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 10 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.97633333333333339 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 50 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.51855555555555555 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 9 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.81177777777777782 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 22 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.90322222222222226 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 31 + +[[ivf_test_build]] +__schema__ = 'benchmark_expected_result' +__version__ = 'v0.0.0' +distance = 'MIP' + + [ivf_test_build.build_parameters] + __schema__ = 'ivf_build_parameters' + __version__ = 'v0.0.0' + hierarchical_level1_clusters = 0 + is_hierarchical = false + minibatch_size = 10000 + num_centroids = 128 + num_iterations = 10 + seed = 12648430 + training_fraction = 0.10000000149011612 + + [ivf_test_build.dataset] + __schema__ = 'benchmark_dataset_abstract' + __version__ = 'v0.0.0' + kind = 'uncompressed' + + [ivf_test_build.dataset.dataset] + __schema__ = 'benchmark_dataset_uncompressed' + __version__ = 'v0.0.0' + data_type = 'bfloat16' + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.26800000000000002 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 10 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.70766666666666667 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 50 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.504 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 27 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.80444444444444441 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 65 + + [[ivf_test_build.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.90288888888888885 + recall_k = 10 + + [ivf_test_build.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 86 + +[[ivf_test_search]] +__schema__ = 'benchmark_expected_result' +__version__ = 'v0.0.0' +distance = 'L2' + + [ivf_test_search.dataset] + __schema__ = 'benchmark_dataset_abstract' + __version__ = 'v0.0.0' + kind = 'uncompressed' + + [ivf_test_search.dataset.dataset] + __schema__ = 'benchmark_dataset_uncompressed' + __version__ = 'v0.0.0' + data_type = 'float32' + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.48399999999999999 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 10 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.95644444444444443 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 50 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.51244444444444448 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 11 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.80844444444444441 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 27 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.90000000000000002 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 38 + +[[ivf_test_search]] +__schema__ = 'benchmark_expected_result' +__version__ = 'v0.0.0' +distance = 'MIP' + + [ivf_test_search.dataset] + __schema__ = 'benchmark_dataset_abstract' + __version__ = 'v0.0.0' + kind = 'uncompressed' + + [ivf_test_search.dataset.dataset] + __schema__ = 'benchmark_dataset_uncompressed' + __version__ = 'v0.0.0' + data_type = 'float16' + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.066000000000000003 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 10 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.57833333333333337 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 50 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.50155555555555553 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 44 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.80577777777777781 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 72 + + [[ivf_test_search.config_and_recall]] + __schema__ = 'benchmark_config_and_result' + __version__ = 'v0.0.0' + num_neighbors = 10 + num_queries = 900 + recall = 0.90044444444444449 + recall_k = 10 + + [ivf_test_search.config_and_recall.search_parameters] + __schema__ = 'ivf_search_parameters' + __version__ = 'v0.0.0' + k_reorder = 1.0 + n_probes = 85 diff --git a/include/svs/core/data/simple.h b/include/svs/core/data/simple.h index bde2786a..3ceb8126 100644 --- a/include/svs/core/data/simple.h +++ b/include/svs/core/data/simple.h @@ -245,6 +245,9 @@ class SimpleData { /// The type used to return a constant handle to stored vectors. using const_value_type = std::span; + /// Data wrapped in the library allocator. + using lib_alloc_data_type = SimpleData>; + /// Return the underlying allocator. const allocator_type& get_allocator() const { return data_.get_allocator(); } @@ -600,6 +603,8 @@ class SimpleData> { using value_type = std::span; using const_value_type = std::span; + using lib_alloc_data_type = SimpleData>>; + ///// Constructors SimpleData(size_t n_elements, size_t n_dimensions, const Blocked& alloc) : blocksize_{lib::prevpow2( diff --git a/include/svs/core/data/view.h b/include/svs/core/data/view.h index 21e34ce9..e9aa0c9e 100644 --- a/include/svs/core/data/view.h +++ b/include/svs/core/data/view.h @@ -66,6 +66,7 @@ template class DataViewImpl { using raw_reference = Data&; using raw_const_reference = const Data&; + using element_type = typename raw_data_type::element_type; using value_type = typename raw_data_type::value_type; using const_value_type = typename raw_data_type::const_value_type; diff --git a/include/svs/index/ivf/clustering.h b/include/svs/index/ivf/clustering.h new file mode 100644 index 00000000..8f055517 --- /dev/null +++ b/include/svs/index/ivf/clustering.h @@ -0,0 +1,356 @@ +/* + * Copyright 2025 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 + +// svs +#include "svs/concepts/data.h" +#include "svs/core/data.h" +#include "svs/core/logging.h" +#include "svs/index/ivf/extensions.h" +#include "svs/lib/meta.h" +#include "svs/lib/misc.h" +#include "svs/lib/saveload.h" +#include "svs/lib/threads.h" +#include "svs/lib/timing.h" + +namespace svs::index::ivf { + +struct ClusteringStats { + public: + size_t min_size_ = std::numeric_limits::max(); + size_t max_size_ = std::numeric_limits::min(); + size_t empty_clusters_ = 0; + size_t num_clusters_ = 0; + size_t num_leaves_ = 0; + double mean_size_ = 0; + double std_size_ = 0; + + public: + template ClusteringStats(Iter begin, Iter end) { + for (auto it = begin; it != end; ++it) { + const auto& list = *it; + num_clusters_++; + auto these_leaves = list.size(); + num_leaves_ += these_leaves; + min_size_ = std::min(min_size_, these_leaves); + max_size_ = std::max(max_size_, these_leaves); + if (these_leaves == 0) { + empty_clusters_ += 1; + } + } + mean_size_ = static_cast(num_leaves_) / num_clusters_; + + // Compute the standard deviation + double accum = 0; + for (auto it = begin; it != end; ++it) { + const auto& list = *it; + auto x = static_cast(list.size()) - mean_size_; + accum += x * x; + } + std_size_ = std::sqrt(accum / static_cast(num_clusters_)); + } + + std::vector prepare_report() const { + return std::vector({ + SVS_SHOW_STRING_(min_size), + SVS_SHOW_STRING_(max_size), + SVS_SHOW_STRING_(empty_clusters), + SVS_SHOW_STRING_(num_clusters), + SVS_SHOW_STRING_(num_leaves), + SVS_SHOW_STRING_(mean_size), + SVS_SHOW_STRING_(std_size), + }); + } + + [[nodiscard]] std::string report() const { return report(", "); } + [[nodiscard]] std::string report(std::string_view separator) const { + return fmt::format("{}", fmt::join(prepare_report(), separator)); + } +}; + +template class Clustering { + public: + Data centroids_; + std::vector> clusters_; + + public: + using vector_type = std::vector; + using T = typename Data::element_type; + + // Type Aliases + using iterator = typename std::vector::iterator; + using const_iterator = typename std::vector::const_iterator; + + public: + Clustering() = default; + + Clustering(size_t n_clusters, size_t n_dims) + : centroids_{n_clusters, n_dims} + , clusters_(n_clusters) {} + + Clustering(Data centroids, std::vector clusters) + : centroids_{std::move(centroids)} + , clusters_{std::move(clusters)} {} + + size_t size() const { return clusters_.size(); } + + void check_valid(size_t cluster_id) const { + if (cluster_id >= size()) { + throw ANNEXCEPTION( + "Cluster id {} can't be higher than the number of clusters!", + cluster_id, + size() + ); + } + } + + size_t size(size_t id) const { + check_valid(id); + return clusters_[id].size(); + } + + const vector_type& cluster(size_t id) const { + check_valid(id); + return clusters_[id]; + } + + vector_type& cluster(size_t id) { + check_valid(id); + return clusters_[id]; + } + + Data centroids() { return centroids_; } + + // Iterators + iterator begin() { return clusters_.begin(); } + iterator end() { return clusters_.end(); } + + const_iterator cbegin() const { return clusters_.cbegin(); } + const_iterator cend() const { return clusters_.cend(); } + + template void for_each_cluster(F&& f) const { + for (size_t i = 0; i < size(); i++) { + f(cluster(i)); + } + } + + template + void for_each_cluster_parallel(F&& f, Pool& threadpool) const { + threads::parallel_for( + threadpool, + threads::StaticPartition{size()}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + f(cluster(i), i); + } + } + ); + } + + ClusteringStats statistics() const { return ClusteringStats(cbegin(), cend()); } + + // Serializing and Deserializing. + size_t serialize_clusters(std::ostream& stream) const { + size_t bytes = lib::write_binary(stream, size()); + for (size_t i = 0; i < size(); i++) { + bytes += lib::write_binary(stream, size(i)); + bytes += lib::write_binary(stream, clusters_[i]); + } + return bytes; + } + + static std::vector deserialize_clusters(std::istream& stream) { + size_t n_clusters = lib::read_binary(stream); + std::vector clusters(n_clusters); + for (size_t i = 0; i < n_clusters; i++) { + size_t cluster_size = lib::read_binary(stream); + clusters[i].resize(cluster_size); + lib::read_binary(stream, clusters[i]); + } + return clusters; + } + + // Saving and Loading. + static constexpr lib::Version save_version{0, 0, 0}; + static constexpr std::string_view serialization_schema = "IVF clustering"; + lib::SaveTable save(const lib::SaveContext& ctx) const { + // Serialize all clusters into an auxiliary file. + auto fullpath = ctx.generate_name("clusters", "bin"); + size_t filesize = 0; + { + auto io = lib::open_write(fullpath); + filesize += serialize_clusters(io); + } + + return lib::SaveTable( + serialization_schema, + save_version, + {SVS_LIST_SAVE_(centroids, ctx), + {"filepath", lib::save(fullpath.filename())}, + SVS_LIST_SAVE(filesize), + {"data_type", lib::save(datatype_v)}, + {"integer_type", lib::save(datatype_v)}, + {"num_clusters", lib::save(size())}} + ); + } + + template + static Clustering load(const lib::LoadTable& table, Pool& threadpool) { + auto saved_data_type = lib::load_at(table, "data_type"); + + // Ensure we have the correct integer type when decoding. + auto saved_integer_type = lib::load_at(table, "integer_type"); + if (saved_integer_type != datatype_v) { + auto type = datatype_v; + throw ANNEXCEPTION( + "Clustering was saved using {} but we're trying to reload it using {}!", + saved_integer_type, + type + ); + } + + auto expected_filesize = lib::load_at(table, "filesize"); + + auto file = table.resolve_at("filepath"); + size_t actual_filesize = std::filesystem::file_size(file); + if (actual_filesize != expected_filesize) { + throw ANNEXCEPTION( + "Expected cluster file size to be {}. Instead, it is {}!", + actual_filesize, + expected_filesize + ); + } + + auto io = lib::open_read(file); + if (saved_data_type != datatype_v) { + auto centroids_orig = + lib::load_at>(table, "centroids"); + if constexpr (std::is_same_v || std::is_same_v) { + auto centroids = convert_data(centroids_orig, threadpool); + return Clustering{centroids, deserialize_clusters(io)}; + } else { + throw ANNEXCEPTION("Centroids datatype {} not supported!", datatype_v); + } + } + + return Clustering{ + SVS_LOAD_MEMBER_AT_(table, centroids), deserialize_clusters(io)}; + } +}; + +template struct DenseCluster { + public: + DenseCluster(Data data, std::vector ids) + : data_{std::move(data)} + , ids_{std::move(ids)} { + if (data_.size() != ids_.size()) { + throw ANNEXCEPTION("Size mismatch!"); + } + } + + size_t size() const { return data_.size(); } + + template + void on_leaves(Callback&& f, size_t prefetch_offset) const { + size_t p = 0; + size_t clustersize = size(); + auto accessor = extensions::accessor(data_); + + for (size_t pmax = std::min(prefetch_offset, clustersize); p < pmax; ++p) { + accessor.prefetch(data_, p); + } + + for (size_t i = 0; i < clustersize; ++i) { + if (p < clustersize) { + accessor.prefetch(data_, p); + ++p; + } + f(accessor(data_, i), ids_[i], i); + } + } + + auto get_datum(size_t id) const { return data_.get_datum(id); } + auto get_secondary(size_t id) const { return data_.get_secondary(id); } + auto get_global_id(size_t local_id) const { return ids_[local_id]; } + const Data& view_cluster() const { return data_; } + + public: + Data data_; + std::vector ids_; +}; + +template < + data::ImmutableMemoryDataset Centroids, + std::integral I, + data::ImmutableMemoryDataset Data> +class DenseClusteredDataset { + public: + // Type aliases + using index_type = I; + using data_type = Data; + + // Constructor + template + DenseClusteredDataset( + const Clustering& clustering, + const Original& original, + Pool& threadpool, + const Alloc& allocator + ) + : clusters_{} { + clustering.for_each_cluster([&](const auto& cluster) { + size_t cluster_size = cluster.size(); + clusters_.emplace_back( + extensions::create_dense_cluster(original, cluster_size, allocator), + std::vector(cluster_size) + ); + }); + + clustering.for_each_cluster_parallel( + [&](const auto& cluster, size_t cluster_id) { + auto& leaf = clusters_[cluster_id]; + extensions::set_dense_cluster(original, leaf.data_, cluster, leaf.ids_); + }, + threadpool + ); + } + + template void on_leaves(Callback&& f, size_t cluster) const { + clusters_.at(cluster).on_leaves(SVS_FWD(f), prefetch_offset_); + } + + size_t get_prefetch_offset() const { return prefetch_offset_; } + void set_prefetch_offset(size_t offset) { prefetch_offset_ = offset; } + auto get_datum(size_t cluster, size_t id) const { + return clusters_.at(cluster).get_datum(id); + } + auto get_secondary(size_t cluster, size_t id) const { + return clusters_.at(cluster).get_secondary(id); + } + auto get_global_id(size_t cluster, size_t id) const { + return clusters_.at(cluster).get_global_id(id); + } + const Data& view_cluster(size_t cluster) const { + return clusters_.at(cluster).view_cluster(); + } + + private: + std::vector> clusters_; + size_t prefetch_offset_ = 8; +}; + +} // namespace svs::index::ivf diff --git a/include/svs/index/ivf/common.h b/include/svs/index/ivf/common.h new file mode 100644 index 00000000..e6e756c0 --- /dev/null +++ b/include/svs/index/ivf/common.h @@ -0,0 +1,699 @@ +/* + * Copyright 2025 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 + +// svs +#include "svs/concepts/data.h" +#include "svs/concepts/distance.h" +#include "svs/core/data/simple.h" +#include "svs/core/data/view.h" +#include "svs/core/distance.h" +#include "svs/core/logging.h" +#include "svs/index/ivf/sorted_buffer.h" +#include "svs/lib/exception.h" +#include "svs/lib/neighbor.h" +#include "svs/lib/threads/threadpool.h" +#include "svs/lib/timing.h" +#include "svs/lib/type_traits.h" + +// external +#include "tsl/robin_set.h" + +// Intel(R) MKL +#include + +// stl +#include + +// Common definitions. +namespace svs::index::ivf { + +// Small epsilon value used for floating-point comparisons to avoid precision +// issues. The value 1/1024 (approximately 0.0009765625) is chosen as a reasonable +// threshold for numerical stability in algorithms such as k-means clustering, where exact +constexpr double EPSILON = 1.0 / 1024.0; + +/// @brief Parameters controlling the IVF build/k-means algortihm. +struct IVFBuildParameters { + public: + /// The target number of clusters in the final result. + size_t num_centroids_ = 1000; + /// The size of each minibatch. + size_t minibatch_size_ = 10'000; + /// The number of iterations used in kmeans training. + size_t num_iterations_ = 10; + /// Use hierarchical Kmeans + bool is_hierarchical_ = true; + /// Fraction of dataset used for training + float training_fraction_ = 0.1; + // Level1 clusters for hierarchical kmeans (use heuristic when 0) + size_t hierarchical_level1_clusters_ = 0; + /// The initial seed for the random number generator. + size_t seed_ = 0xc0ffee; + + public: + IVFBuildParameters() = default; + IVFBuildParameters( + size_t num_centroids, + size_t minibatch_size = 10'000, + size_t num_iterations = 10, + bool is_hierarchical = true, + float training_fraction = 0.1, + size_t hierarchical_level1_clusters = 0, + size_t seed = 0xc0ffee + ) + : num_centroids_{num_centroids} + , minibatch_size_{minibatch_size} + , num_iterations_{num_iterations} + , is_hierarchical_{is_hierarchical} + , training_fraction_{training_fraction} + , hierarchical_level1_clusters_{hierarchical_level1_clusters} + , seed_{seed} {} + + // Chain setters to help with construction. + SVS_CHAIN_SETTER_(IVFBuildParameters, num_centroids); + SVS_CHAIN_SETTER_(IVFBuildParameters, minibatch_size); + SVS_CHAIN_SETTER_(IVFBuildParameters, num_iterations); + SVS_CHAIN_SETTER_(IVFBuildParameters, is_hierarchical); + SVS_CHAIN_SETTER_(IVFBuildParameters, training_fraction); + SVS_CHAIN_SETTER_(IVFBuildParameters, hierarchical_level1_clusters); + SVS_CHAIN_SETTER_(IVFBuildParameters, seed); + + // Comparison + friend constexpr bool + operator==(const IVFBuildParameters&, const IVFBuildParameters&) = default; + + // Saving and Loading. + static constexpr svs::lib::Version save_version{0, 0, 0}; + static constexpr std::string_view serialization_schema = "ivf_build_parameters"; + lib::SaveTable save() const { + return lib::SaveTable( + serialization_schema, + save_version, + { + SVS_LIST_SAVE_(num_centroids), + SVS_LIST_SAVE_(minibatch_size), + SVS_LIST_SAVE_(num_iterations), + SVS_LIST_SAVE_(is_hierarchical), + SVS_LIST_SAVE_(training_fraction), + SVS_LIST_SAVE_(hierarchical_level1_clusters), + {"seed", lib::save(lib::FullUnsigned(seed_))}, + } + ); + } + + static IVFBuildParameters load(const lib::ContextFreeLoadTable& table) { + return IVFBuildParameters( + SVS_LOAD_MEMBER_AT_(table, num_centroids), + SVS_LOAD_MEMBER_AT_(table, minibatch_size), + SVS_LOAD_MEMBER_AT_(table, num_iterations), + SVS_LOAD_MEMBER_AT_(table, is_hierarchical), + SVS_LOAD_MEMBER_AT_(table, training_fraction), + SVS_LOAD_MEMBER_AT_(table, hierarchical_level1_clusters), + lib::load_at(table, "seed") + ); + } +}; + +/// @brief Parameters controlling the IVF search algorithm. +struct IVFSearchParameters { + public: + /// The number of nearest clusters to be explored + size_t n_probes_ = 1; + /// Level of reordering or reranking done when using compressed datasets + float k_reorder_ = 1.0; + + public: + IVFSearchParameters() = default; + + IVFSearchParameters(size_t n_probes, float k_reorder) + : n_probes_{n_probes} + , k_reorder_{k_reorder} {} + + SVS_CHAIN_SETTER_(IVFSearchParameters, n_probes); + SVS_CHAIN_SETTER_(IVFSearchParameters, k_reorder); + + // Saving and Loading. + static constexpr lib::Version save_version{0, 0, 0}; + static constexpr std::string_view serialization_schema = "ivf_search_parameters"; + lib::SaveTable save() const { + return lib::SaveTable( + serialization_schema, + save_version, + {SVS_LIST_SAVE_(n_probes), SVS_LIST_SAVE_(k_reorder)} + ); + } + + static IVFSearchParameters load(const lib::ContextFreeLoadTable& table) { + return IVFSearchParameters{ + SVS_LOAD_MEMBER_AT_(table, n_probes), SVS_LOAD_MEMBER_AT_(table, k_reorder)}; + } + + constexpr friend bool + operator==(const IVFSearchParameters&, const IVFSearchParameters&) = default; +}; + +// Helper functions to convert data from one type to another using threadpool +template +void convert_data(D1& src, D2& dst, Pool& threadpool) { + // Note: Destination size can be bigger than the source as we preallocate a bigger + // buffer and reuse it to reduce the cost of frequent allocations + if (src.size() > dst.size() || src.dimensions() != dst.dimensions()) { + throw ANNEXCEPTION( + "Unexpected data shapes sizes: {}, {}; dims: {}, {}!", + src.size(), + dst.size(), + src.dimensions(), + dst.dimensions() + ); + } + + threads::parallel_for( + threadpool, + threads::StaticPartition{src.size()}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + dst.set_datum(i, src.get_datum(i)); + } + } + ); +} + +template +auto convert_data(Data& src, Pool& threadpool) { + auto dst = svs::data::SimpleData(src.size(), src.dimensions()); + convert_data(src, dst, threadpool); + return dst; +} + +// Partial specialization to preserve the dimensionality and Allocation type +template +auto convert_data(svs::data::SimpleData& src, Pool& threadpool) { + using allocator_type = svs::lib::rebind_allocator_t; + allocator_type rebound_alloctor = {}; + + auto dst = svs::data::SimpleData( + src.size(), src.dimensions(), rebound_alloctor + ); + convert_data(src, dst, threadpool); + return dst; +} + +template auto convert_data(Data& src) { + auto dst = svs::data::SimpleData(src.size(), src.dimensions()); + auto threadpool = threads::as_threadpool(1); + convert_data(src, dst, threadpool); + return dst; +} + +template +void compute_matmul( + const T* data, const T* centroids, float* results, size_t m, size_t n, size_t k +) { + if constexpr (std::is_same_v) { + cblas_sgemm( + CblasRowMajor, // CBLAS_LAYOUT layout + CblasNoTrans, // CBLAS_TRANSPOSE TransA + CblasTrans, // CBLAS_TRANSPOSE TransB + m, // const int M + n, // const int N + k, // const int K + 1.0, // float alpha + data, // const float* A + k, // const int lda + centroids, // const float* B + k, // const int ldb + 0.0, // const float beta + results, // float* c + n // const int ldc + ); + } else if constexpr (std::is_same_v) { + cblas_gemm_bf16bf16f32( + CblasRowMajor, // CBLAS_LAYOUT layout + CblasNoTrans, // CBLAS_TRANSPOSE TransA + CblasTrans, // CBLAS_TRANSPOSE TransB + m, // const int M + n, // const int N + k, // const int K + 1.0, // float alpha + (const uint16_t*)data, // const *uint16_t A + k, // const int lda + (const uint16_t*)centroids, // const uint16_t* B + k, // const int ldb + 0.0, // const float beta + results, // float* c + n // const int ldc + ); + } else if constexpr (std::is_same_v) { + cblas_gemm_f16f16f32( + CblasRowMajor, // CBLAS_LAYOUT layout + CblasNoTrans, // CBLAS_TRANSPOSE TransA + CblasTrans, // CBLAS_TRANSPOSE TransB + m, // const int M + n, // const int N + k, // const int K + 1.0, // float alpha + (const uint16_t*)data, // const *uint16_t A + k, // const int lda + (const uint16_t*)centroids, // const uint16_t* B + k, // const int ldb + 0.0, // const float beta + results, // float* c + n // const int ldc + ); + } else { + throw ANNEXCEPTION("GEMM type not supported!"); + } +} + +inline static void +generate_unique_ids(std::vector& ids, size_t id_range, std::mt19937& rng) { + size_t n = ids.size(); + tsl::robin_set seen; + seen.reserve(n); + + while (seen.size() < n) { + auto j = rng() % id_range; + seen.insert(j); + } + + size_t i = 0; + for (auto it = seen.begin(); it != seen.end(); ++it, ++i) { + ids[i] = *it; + } +} + +template +void normalize_centroids( + data::SimpleData& centroids, Pool& threadpool, lib::Timer& timer +) { + auto normalize_centroids_t = timer.push_back("normalize centroids"); + threads::parallel_for( + threadpool, + threads::StaticPartition{centroids.size()}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + auto datum = centroids.get_datum(i); + float norm = distance::norm(datum); + if (norm != 0.0) { + float norm_inv = 1.0 / norm; + for (size_t j = 0; j < datum.size(); j++) { + datum[j] = datum[j] * norm_inv; + } + } + } + } + ); + normalize_centroids_t.finish(); +} + +template < + data::ImmutableMemoryDataset Data, + typename T, + typename Distance, + threads::ThreadPool Pool> +void centroid_assignment( + Data& data, + std::vector& data_norm, + threads::UnitRange batch_range, + Distance& SVS_UNUSED(distance), + data::SimpleData& centroids, + std::vector& centroids_norm, + std::vector& assignments, + data::SimpleData& matmul_results, + Pool& threadpool, + lib::Timer& timer +) { + auto generate_assignments = timer.push_back("generate assignments"); + threads::parallel_for( + threadpool, + threads::StaticPartition{batch_range.size()}, + [&](auto indices, auto /*tid*/) { + auto range = threads::UnitRange(indices); + compute_matmul( + data.get_datum(range.start()).data(), + centroids.data(), + matmul_results.get_datum(range.start()).data(), + range.size(), + centroids.size(), + data.dimensions() + ); + if constexpr (std::is_same_v) { + for (auto i : indices) { + auto nearest = + type_traits::sentinel_v, std::greater<>>; + auto dists = matmul_results.get_datum(i); + for (size_t j = 0; j < centroids.size(); j++) { + nearest = std::max(nearest, Neighbor(j, dists[j])); + } + assignments[batch_range.start() + i] = nearest.id(); + } + } else if constexpr (std::is_same_v) { + for (auto i : indices) { + auto nearest = type_traits::sentinel_v, std::less<>>; + auto dists = matmul_results.get_datum(i); + for (size_t j = 0; j < centroids.size(); j++) { + auto dist = data_norm[batch_range.start() + i] + centroids_norm[j] - + 2 * dists[j]; + nearest = std::min(nearest, Neighbor(j, dist)); + } + assignments[batch_range.start() + i] = nearest.id(); + } + } else { + throw ANNEXCEPTION("Only L2 and MIP distances supported in IVF build!"); + } + } + ); + generate_assignments.finish(); +} + +template +void centroid_adjustment( + Data& data, + data::SimpleData& centroids, + std::vector& assignments, + std::vector& counts, + Pool& threadpool, + lib::Timer& timer +) { + auto adjust_centroids = timer.push_back("adjust centroids"); + size_t n_threads = threadpool.size(); + size_t n_centroids = centroids.size(); + + threads::parallel_for( + threadpool, + threads::StaticPartition{n_threads}, + [&](auto /*indices*/, auto tid) { + size_t centroid_start = (n_centroids * tid) / n_threads; + size_t centroid_end = (n_centroids * (tid + 1)) / n_threads; + for (auto i : data.eachindex()) { + auto assignment = assignments[i]; + if (assignment >= centroid_start && assignment < centroid_end) { + counts.at(assignment)++; + auto datum = data.get_datum(i); + auto this_centroid = centroids.get_datum(assignment); + for (size_t p = 0, pmax = this_centroid.size(); p < pmax; ++p) { + this_centroid[p] += lib::narrow_cast(datum[p]); + } + } + } + } + ); + + threads::parallel_for( + threadpool, + threads::StaticPartition{n_centroids}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + if (counts.at(i) != 0) { + auto this_centroid = centroids.get_datum(i); + float norm = 1.0 / (counts.at(i) + 1); + // float norm = 1.0 / counts.at(i); + for (size_t p = 0, pmax = this_centroid.size(); p < pmax; ++p) { + this_centroid[p] *= norm; + } + } + } + } + ); + adjust_centroids.finish(); +} + +template +void centroid_split( + Data& data, + data::SimpleData& centroids, + std::vector& counts, + std::mt19937& rng, + Pool& SVS_UNUSED(threadpool), + lib::Timer& timer +) { + auto split_centroids = timer.push_back("split centroids"); + + auto num_centroids = centroids.size(); + auto dims = centroids.dimensions(); + auto num_data = data.size(); + + auto distribution = std::uniform_real_distribution(0.0, 1.0); + for (size_t i = 0; i < num_centroids; i++) { + if (counts.at(i) == 0) { + size_t j; + for (j = 0; true; j = (j + 1) % num_centroids) { + if (counts.at(j) == 0) { + continue; + } + float p = counts.at(j) / float(num_data); + float r = distribution(rng); + if (r < p) { + break; + } + } + centroids.set_datum(i, centroids.get_datum(j)); + for (size_t k = 0; k < dims; k++) { + if (k % 2 == 0) { + centroids.get_datum(i)[k] *= 1 + EPSILON; + centroids.get_datum(j)[k] *= 1 - EPSILON; + } else { + centroids.get_datum(i)[k] *= 1 - EPSILON; + centroids.get_datum(j)[k] *= 1 + EPSILON; + } + } + counts.at(i) = counts.at(j) / 2; + counts.at(j) -= counts.at(i); + } + } + split_centroids.finish(); +} + +template +void generate_norms(Data& data, std::vector& norms, Pool& threadpool) { + norms.resize(data.size()); + threads::parallel_for( + threadpool, + threads::StaticPartition{data.size()}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + norms[i] = distance::norm_square(data.get_datum(i)); + } + } + ); +} + +template < + data::ImmutableMemoryDataset Data, + typename T, + typename Distance, + threads::ThreadPool Pool> +auto kmeans_training( + const IVFBuildParameters& parameters, + Data& data, + Distance& distance, + data::SimpleData& centroids, + data::SimpleData& matmul_results, + std::mt19937& rng, + Pool& threadpool, + lib::Timer& timer +) { + auto training_timer = timer.push_back("Kmeans training"); + data::SimpleData centroids_fp32 = convert_data(centroids, threadpool); + + if constexpr (std::is_same_v) { + normalize_centroids(centroids_fp32, threadpool, timer); + } + + auto assignments = std::vector(data.size()); + std::vector data_norm; + if constexpr (std::is_same_v) { + generate_norms(data, data_norm, threadpool); + } + std::vector centroids_norm; + + for (size_t iter = 0; iter < parameters.num_iterations_; ++iter) { + auto iter_timer = timer.push_back("iteration"); + auto batchsize = parameters.minibatch_size_; + auto num_batches = lib::div_round_up(data.size(), batchsize); + if constexpr (std::is_same_v) { + generate_norms(centroids_fp32, centroids_norm, threadpool); + } + + // Convert from fp32 to fp16/bf16 + convert_data(centroids_fp32, centroids, threadpool); + + for (size_t batch = 0; batch < num_batches; ++batch) { + auto this_batch = threads::UnitRange{ + batch * batchsize, std::min((batch + 1) * batchsize, data.size())}; + auto data_batch = data::make_view(data, this_batch); + centroid_assignment( + data_batch, + data_norm, + this_batch, + distance, + centroids, + centroids_norm, + assignments, + matmul_results, + threadpool, + timer + ); + } + + // Convert back to fp32 + convert_data(centroids, centroids_fp32, threadpool); + + auto counts = std::vector(centroids.size()); + centroid_adjustment(data, centroids_fp32, assignments, counts, threadpool, timer); + + centroid_split(data, centroids_fp32, counts, rng, threadpool, timer); + + if constexpr (std::is_same_v) { + normalize_centroids(centroids_fp32, threadpool, timer); + } + } + + // Finally call the conversion to get the updated centroids + // after adjustment and split + convert_data(centroids_fp32, centroids, threadpool); + training_timer.finish(); + return centroids_fp32; +} + +template < + data::ImmutableMemoryDataset Queries, + typename Centroids, + typename MatMulResults, + threads::ThreadPool Pool> +void compute_centroid_distances( + const Queries& queries, + const Centroids& centroids, + MatMulResults& matmul_results, + Pool& threadpool +) { + using TQ = typename Queries::element_type; + using TC = typename Centroids::element_type; + size_t num_centroids = centroids.size(); + size_t num_queries = queries.size(); + data::SimpleData queries_conv; + + // Convert if Queries and Centroids datatypes are not same + if constexpr (!std::is_same_v) { + queries_conv = convert_data(queries, threadpool); + } + + threads::parallel_for( + threadpool, + threads::StaticPartition(num_centroids), + [&](auto is, auto tid) { + auto batch = threads::UnitRange{is}; + if constexpr (!std::is_same_v) { + compute_matmul( + queries_conv.data(), + centroids.get_datum(batch.start()).data(), + matmul_results[tid].data(), + num_queries, + batch.size(), + queries.dimensions() + ); + } else { + compute_matmul( + queries.get_datum(0).data(), + centroids.get_datum(batch.start()).data(), + matmul_results[tid].data(), + num_queries, + batch.size(), + queries.dimensions() + ); + } + } + ); +} + +template +void search_centroids( + const Query& query, + Dist& SVS_UNUSED(dist), + const MatMulResults& matmul_results, + Buffer& buffer, + size_t query_id, + const std::vector& centroids_norm, + size_t num_threads +) { + unsigned int count = 0; + buffer.clear(); + if constexpr (std::is_same_v) { + for (size_t j = 0; j < num_threads; j++) { + auto distance = matmul_results[j].get_datum(query_id); + for (size_t k = 0; k < distance.size(); k++) { + buffer.insert({count, distance[k]}); + count++; + } + } + } else if constexpr (std::is_same_v) { + float query_norm = distance::norm_square(query); + for (size_t j = 0; j < num_threads; j++) { + auto distance = matmul_results[j].get_datum(query_id); + for (size_t k = 0; k < distance.size(); k++) { + float dist = query_norm + centroids_norm[count] - 2 * distance[k]; + buffer.insert({count, dist}); + count++; + } + } + } else { + throw ANNEXCEPTION("Only L2 and MIP distances supported in IVF search!"); + } +} + +template < + typename Query, + typename Dist, + typename Cluster, + typename BufferCentroids, + typename BufferLeaves, + threads::ThreadPool Pool> +void search_leaves( + const Query& query, + Dist& dist, + const Cluster& cluster, + const BufferCentroids& buffer_centroids, + BufferLeaves& buffer_leaves, + Pool& threadpool_inner +) { + for (size_t j = 0; j < buffer_leaves.size(); j++) { + buffer_leaves[j].clear(); + } + distance::maybe_fix_argument(dist, query); + threads::parallel_for( + threadpool_inner, + threads::DynamicPartition(buffer_centroids.size(), 1), + [&](auto js, auto tid_inner) { + for (auto j : js) { + auto candidate = buffer_centroids[j]; + auto cluster_id = candidate.id(); + + // Compute the distance between the query and each leaf element. + cluster.on_leaves( + [&](const auto& datum, unsigned int /*gid*/, unsigned int lid) { + auto distance = distance::compute(dist, query, datum); + buffer_leaves[tid_inner].insert({cluster_id, distance, lid}); + }, + cluster_id + ); + } + } + ); +} + +} // namespace svs::index::ivf diff --git a/include/svs/index/ivf/extensions.h b/include/svs/index/ivf/extensions.h new file mode 100644 index 00000000..0ef83587 --- /dev/null +++ b/include/svs/index/ivf/extensions.h @@ -0,0 +1,220 @@ +/* + * Copyright 2025 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 "svs/core/distance.h" +#include "svs/core/query_result.h" +#include "svs/index/ivf/common.h" +#include "svs/lib/invoke.h" + +namespace svs::index::ivf::extensions { + +struct IVFAccessor { + template + svs::svs_invoke_result_t operator()(const Data& data) const { + return svs::svs_invoke(*this, data); + } +}; + +inline constexpr IVFAccessor accessor{}; + +template +data::GetDatumAccessor svs_invoke(svs::tag_t, const Data& SVS_UNUSED(data)) { + return data::GetDatumAccessor{}; +} + +struct IVFPerThreadBatchSearchSetupType { + template + using result_t = svs:: + svs_invoke_result_t; + + template + result_t + operator()(const Data& dataset, const Distance& distance) const { + return svs::svs_invoke(*this, dataset, distance); + } +}; + +inline constexpr IVFPerThreadBatchSearchSetupType per_thread_batch_search_setup{}; + +// Default Implementation. +template +Distance svs_invoke( + svs::tag_t, + const Data& SVS_UNUSED(dataset), + const Distance& distance +) { + return threads::shallow_copy(distance); +} + +/// +/// @brief Customization point for working with a batch of threads. +/// +struct IVFPerThreadBatchSearchType { + template < + data::ImmutableMemoryDataset Data, + typename Cluster, + typename BufferCentroids, + typename BufferLeaves, + typename Scratch, + data::ImmutableMemoryDataset Queries, + std::integral I, + typename SearchCentroids, + typename SearchLeaves> + SVS_FORCE_INLINE void operator()( + const Data& data, + const Cluster& cluster, + BufferCentroids& buffer_centroids, + BufferLeaves& buffer_leaves, + Scratch& scratch, + const Queries& queries, + QueryResultView& result, + threads::UnitRange thread_indices, + size_t tid, + const SearchCentroids& search_centroids, + const SearchLeaves& search_leaves + ) const { + svs::svs_invoke( + *this, + data, + cluster, + buffer_centroids, + buffer_leaves, + scratch, + queries, + result, + thread_indices, + tid, + search_centroids, + search_leaves + ); + } +}; + +/// @brief Customization point object for batch search. +inline constexpr IVFPerThreadBatchSearchType per_thread_batch_search{}; + +// Default Implementation +template < + typename Data, + typename Cluster, + typename BufferCentroids, + typename BufferLeaves, + typename Distance, + typename Queries, + std::integral I, + typename SearchCentroids, + typename SearchLeaves> +void svs_invoke( + svs::tag_t, + const Data& SVS_UNUSED(dataset), + const Cluster& cluster, + BufferCentroids& buffer_centroids, + BufferLeaves& buffer_leaves, + Distance& distance, + const Queries& queries, + QueryResultView& result, + threads::UnitRange thread_indices, + const size_t tid, + const SearchCentroids& search_centroids, + const SearchLeaves& search_leaves +) { + size_t n_inner_threads = buffer_leaves.size(); + size_t buffer_leaves_size = buffer_leaves[0].capacity(); + size_t num_neighbors = result.n_neighbors(); + + // Fallback implementation + for (auto i : thread_indices) { + const auto& query = queries.get_datum(i); + search_centroids(query, buffer_centroids, i); + search_leaves(query, distance, buffer_centroids, buffer_leaves, tid); + + // Accumulate results from intra-query threads + for (size_t j = 1; j < n_inner_threads; ++j) { + for (size_t k = 0; k < buffer_leaves_size; ++k) { + buffer_leaves[0].insert(buffer_leaves[j][k]); + } + } + + for (size_t j = 0; j < buffer_leaves_size; ++j) { + auto& neighbor = buffer_leaves[0][j]; + auto cluster_id = neighbor.id(); + auto local_id = neighbor.get_local_id(); + auto global_id = cluster.get_global_id(cluster_id, local_id); + neighbor.set_id(global_id); + } + + // Store results + for (size_t j = 0; j < num_neighbors; ++j) { + result.set(buffer_leaves[0][j], i, j); + } + } +} + +struct CreateDenseCluster { + using This = CreateDenseCluster; + + template + using return_type = svs::svs_invoke_result_t; + + template + return_type + operator()(const Data& original, size_t new_size, const Alloc& allocator) const { + return svs::svs_invoke(*this, original, new_size, allocator); + } +}; + +inline constexpr CreateDenseCluster create_dense_cluster{}; + +template +svs::data::SimpleData svs_invoke( + svs::tag_t, + const svs::data::SimpleData& original, + size_t new_size, + const NewAlloc& SVS_UNUSED(allocator) +) { + return svs::data::SimpleData(new_size, original.dimensions()); +} + +struct SetDenseCluster { + template + void operator()( + const Src& src, Dst& dst, const std::vector& src_ids, std::vector& dst_ids + ) const { + return svs::svs_invoke(*this, src, dst, src_ids, dst_ids); + } +}; + +inline constexpr SetDenseCluster set_dense_cluster{}; + +template +void svs_invoke( + svs::tag_t, + const Src& src, + Dst& dst, + const std::vector& src_ids, + std::vector& dst_ids +) { + size_t i = 0; + for (auto id : src_ids) { + dst.set_datum(i, src.get_datum(id)); + dst_ids[i] = id; + ++i; + } +} + +} // namespace svs::index::ivf::extensions diff --git a/include/svs/index/ivf/hierarchical_kmeans.h b/include/svs/index/ivf/hierarchical_kmeans.h new file mode 100644 index 00000000..fca610c9 --- /dev/null +++ b/include/svs/index/ivf/hierarchical_kmeans.h @@ -0,0 +1,370 @@ +/* + * Copyright 2025 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 "svs/index/ivf/common.h" + +// stdlib +#include + +namespace svs::index::ivf { + +template +auto calc_level2_clusters( + size_t num_clusters, + size_t num_level1_clusters, + std::vector>& clusters_level1, + size_t num_training_data, + std::mt19937& rng +) { + auto num_level2_clusters = std::vector(num_level1_clusters); + size_t total_centroids_l2 = 0; + for (size_t cluster = 0; cluster < num_level1_clusters; cluster++) { + num_level2_clusters[cluster] = + (1.0 * clusters_level1[cluster].size()) / num_training_data * num_clusters; + total_centroids_l2 += num_level2_clusters[cluster]; + } + while (total_centroids_l2 < num_clusters) { + size_t j = rng() % num_level1_clusters; + if (clusters_level1[j].size() != 0) { + num_level2_clusters[j]++; + total_centroids_l2++; + } + } + return num_level2_clusters; +} + +template < + typename BuildType, + data::ImmutableMemoryDataset Data, + typename Distance, + threads::ThreadPool Pool, + std::integral I = uint32_t> +auto hierarchical_kmeans_clustering_impl( + const IVFBuildParameters& parameters, + Data& data, + Distance& distance, + Pool& threadpool, + lib::Type SVS_UNUSED(integer_type) = {} +) { + auto timer = lib::Timer(); + auto kmeans_timer = timer.push_back("Hierarchical kmeans clustering"); + auto init_timer = timer.push_back("init"); + auto init_trainset_centroid = timer.push_back("create trainset and centroids"); + + constexpr size_t Dims = Data::extent; + using Alloc = svs::HugepageAllocator; + size_t ndims = data.dimensions(); + auto num_clusters = parameters.num_centroids_; + + size_t num_level1_clusters = parameters.hierarchical_level1_clusters_; + if (num_level1_clusters == 0) { + num_level1_clusters = std::sqrt(num_clusters); + } + + fmt::print("Level1 clusters: {}\n", num_level1_clusters); + + size_t num_training_data = + lib::narrow(std::ceil(data.size() * parameters.training_fraction_)); + if (num_training_data < num_clusters || num_training_data > data.size()) { + throw ANNEXCEPTION( + "Invalid number of training data: {}, num_clusters: {}, total data size: " + "{}\n", + num_training_data, + num_clusters, + data.size() + ); + } + + auto centroids_level1 = data::SimpleData{num_level1_clusters, ndims}; + auto data_train = data::SimpleData{num_training_data, ndims}; + auto matmul_results_level1 = + data::SimpleData{parameters.minibatch_size_, num_level1_clusters}; + auto rng = std::mt19937(parameters.seed_); + + std::vector v(num_training_data); + generate_unique_ids(v, data.size(), rng); + threads::parallel_for( + threadpool, + threads::StaticPartition{num_training_data}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + data_train.set_datum(i, data.get_datum(v[i])); + } + } + ); + + v.resize(num_level1_clusters); + generate_unique_ids(v, data_train.size(), rng); + threads::parallel_for( + threadpool, + threads::StaticPartition{num_level1_clusters}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + centroids_level1.set_datum(i, data_train.get_datum(v[i])); + } + } + ); + init_trainset_centroid.finish(); + init_timer.finish(); + + auto level1_training_time = timer.push_back("Level1 training"); + auto centroids_level1_fp32 = kmeans_training( + parameters, + data_train, + distance, + centroids_level1, + matmul_results_level1, + rng, + threadpool, + timer + ); + auto assignments_level1 = std::vector(data_train.size()); + auto batchsize = parameters.minibatch_size_; + auto num_batches = lib::div_round_up(data_train.size(), batchsize); + + std::vector data_norm; + if constexpr (std::is_same_v) { + generate_norms(data_train, data_norm, threadpool); + } + std::vector centroids_level1_norm; + if constexpr (std::is_same_v) { + generate_norms(centroids_level1_fp32, centroids_level1_norm, threadpool); + } + + for (size_t batch = 0; batch < num_batches; ++batch) { + auto this_batch = threads::UnitRange{ + batch * batchsize, std::min((batch + 1) * batchsize, data_train.size())}; + auto data_batch = data::make_view(data_train, this_batch); + centroid_assignment( + data_batch, + data_norm, + this_batch, + distance, + centroids_level1, + centroids_level1_norm, + assignments_level1, + matmul_results_level1, + threadpool, + timer + ); + } + auto clusters_level1 = std::vector>(num_level1_clusters); + for (auto i : data_train.eachindex()) { + clusters_level1[assignments_level1[i]].push_back(i); + } + + auto all_assignments_time = timer.push_back("level1 all assignments"); + auto all_assignments_alloc = timer.push_back("level1 all assignments alloc"); + auto assignments_level1_all = std::vector(data.size()); + all_assignments_alloc.finish(); + + batchsize = parameters.minibatch_size_; + num_batches = lib::div_round_up(data.size(), batchsize); + + if constexpr (std::is_same_v) { + generate_norms(data, data_norm, threadpool); + } + + auto data_batch = data::SimpleData{batchsize, ndims}; + for (size_t batch = 0; batch < num_batches; ++batch) { + auto this_batch = threads::UnitRange{ + batch * batchsize, std::min((batch + 1) * batchsize, data.size())}; + auto data_batch_view = data::make_view(data, this_batch); + auto all_assignments_convert = timer.push_back("level1 all assignments convert"); + convert_data(data_batch_view, data_batch, threadpool); + all_assignments_convert.finish(); + centroid_assignment( + data_batch, + data_norm, + this_batch, + distance, + centroids_level1, + centroids_level1_norm, + assignments_level1_all, + matmul_results_level1, + threadpool, + timer + ); + } + auto all_assignments_cluster = timer.push_back("level1 all assignments clusters"); + auto clusters_level1_all = std::vector>(num_level1_clusters); + for (auto i : data.eachindex()) { + clusters_level1_all[assignments_level1_all[i]].push_back(i); + } + all_assignments_cluster.finish(); + + all_assignments_time.finish(); + level1_training_time.finish(); + + auto level2_training_time = timer.push_back("Level2 training"); + auto num_level2_clusters = calc_level2_clusters( + num_clusters, num_level1_clusters, clusters_level1, num_training_data, rng + ); + + auto centroids_final = data::SimpleData{num_clusters, ndims}; + auto clusters_final = std::vector>(num_clusters); + + size_t max_data_per_cluster = 0; + for (size_t cluster = 0; cluster < num_level1_clusters; cluster++) { + max_data_per_cluster = clusters_level1_all[cluster].size() > max_data_per_cluster + ? clusters_level1_all[cluster].size() + : max_data_per_cluster; + } + auto data_level2 = + data::SimpleData{max_data_per_cluster, ndims}; + auto assignments_level2_all = std::vector(max_data_per_cluster); + + size_t cluster_start = 0; + for (size_t cluster = 0; cluster < num_level1_clusters; cluster++) { + size_t num_clusters_l2 = num_level2_clusters[cluster]; + size_t num_assignments_l2 = clusters_level1[cluster].size(); + size_t num_assignments_l2_all = clusters_level1_all[cluster].size(); + + auto centroids_level2 = data::SimpleData{num_clusters_l2, ndims}; + auto matmul_results_level2 = + data::SimpleData{parameters.minibatch_size_, num_clusters_l2}; + auto data_train_level2 = data::SimpleData{num_assignments_l2, ndims}; + + threads::parallel_for( + threadpool, + threads::StaticPartition{num_assignments_l2}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + data_train_level2.set_datum( + i, data_train.get_datum(clusters_level1[cluster][i]) + ); + } + } + ); + + v.resize(num_clusters_l2); + generate_unique_ids(v, data_train_level2.size(), rng); + threads::parallel_for( + threadpool, + threads::StaticPartition{num_clusters_l2}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + centroids_level2.set_datum(i, data_train_level2.get_datum(v[i])); + } + } + ); + + auto centroids_level2_fp32 = kmeans_training( + parameters, + data_train_level2, + distance, + centroids_level2, + matmul_results_level2, + rng, + threadpool, + timer + ); + + auto all_assignments_level2 = timer.push_back("level2 all assignments"); + threads::parallel_for( + threadpool, + threads::StaticPartition{num_assignments_l2_all}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + data_level2.set_datum( + i, data.get_datum(clusters_level1_all[cluster][i]) + ); + } + } + ); + + batchsize = parameters.minibatch_size_; + num_batches = lib::div_round_up(num_assignments_l2_all, batchsize); + + if constexpr (std::is_same_v) { + generate_norms(data_level2, data_norm, threadpool); + } + std::vector centroids_level2_norm; + if constexpr (std::is_same_v) { + generate_norms(centroids_level2_fp32, centroids_level2_norm, threadpool); + } + for (size_t batch = 0; batch < num_batches; ++batch) { + auto this_batch = threads::UnitRange{ + batch * batchsize, + std::min((batch + 1) * batchsize, num_assignments_l2_all)}; + auto data_batch = data::make_view(data_level2, this_batch); + centroid_assignment( + data_batch, + data_norm, + this_batch, + distance, + centroids_level2, + centroids_level2_norm, + assignments_level2_all, + matmul_results_level2, + threadpool, + timer + ); + } + + for (size_t i = 0; i < num_assignments_l2_all; i++) { + clusters_final[cluster_start + assignments_level2_all[i]].push_back( + clusters_level1_all[cluster][i] + ); + } + + threads::parallel_for( + threadpool, + threads::StaticPartition{num_clusters_l2}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + centroids_final.set_datum( + cluster_start + i, centroids_level2_fp32.get_datum(i) + ); + } + } + ); + + cluster_start += num_clusters_l2; + all_assignments_level2.finish(); + } + level2_training_time.finish(); + + kmeans_timer.finish(); + svs::logging::debug("{}", timer); + fmt::print( + "hierarchical kmeans clustering time: {}\n", lib::as_seconds(timer.elapsed()) + ); + + return std::make_tuple(std::move(centroids_final), std::move(clusters_final)); +} + +template < + typename BuildType, + data::ImmutableMemoryDataset Data, + typename Distance, + threads::ThreadPool Pool, + std::integral I = uint32_t> +auto hierarchical_kmeans_clustering( + const IVFBuildParameters& parameters, + Data& data, + Distance& distance, + Pool& threadpool, + lib::Type integer_type = {} +) { + return hierarchical_kmeans_clustering_impl( + parameters, data, distance, threadpool, integer_type + ); +} + +} // namespace svs::index::ivf diff --git a/include/svs/index/ivf/index.h b/include/svs/index/ivf/index.h new file mode 100644 index 00000000..eaafe10e --- /dev/null +++ b/include/svs/index/ivf/index.h @@ -0,0 +1,348 @@ +/* + * Copyright 2025 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 + +// svs +#include "svs/concepts/data.h" +#include "svs/core/loading.h" +#include "svs/core/query_result.h" +#include "svs/index/ivf/clustering.h" +#include "svs/index/ivf/extensions.h" +#include "svs/index/ivf/hierarchical_kmeans.h" +#include "svs/index/ivf/kmeans.h" +#include "svs/lib/timing.h" + +// format +#include "fmt/core.h" + +// stl +#include +#include + +namespace svs::index::ivf { + +// The maximum batch size for queries is set to 10,000 to balance memory usage and +// performance. This value was chosen based on empirical testing to avoid excessive memory +// allocation while supporting large batch operations typical in high-throughput +// environments. +const size_t MAX_QUERY_BATCH_SIZE = 10000; + +template +class IVFIndex { + public: + using Idx = typename Cluster::index_type; + using Data = typename Cluster::data_type; + using search_parameters_type = IVFSearchParameters; + + IVFIndex( + Centroids centroids, + Cluster cluster, + Dist distance_function, + ThreadPoolProto threadpool_proto, + size_t n_inner_threads = 1 + ) + : centroids_{std::move(centroids)} + , cluster_{std::move(cluster)} + , cluster0_{cluster_.view_cluster(0)} + , distance_{std::move(distance_function)} + , threadpool_{threads::as_threadpool(std::move(threadpool_proto))} + , n_inner_threads_{n_inner_threads} { + // Initialize threadpools for intra-query parallelism + for (size_t i = 0; i < threadpool_.size(); i++) { + threadpool_inner_.push_back(threads::as_threadpool(n_inner_threads_)); + } + + // The first level of search to find the n_probes nearest centroids is done + // using matmul (single thread) and thread batching on the number of centroids. + // Matmul results of each batch with appropriate sizes are initialized here, + // assuming that the static partition will keep the batching same. + auto batches = std::vector>(threadpool_.size()); + threads::parallel_for( + threadpool_, + threads::StaticPartition(centroids_.size()), + [&](auto is, auto tid) { batches[tid] = threads::UnitRange{is}; } + ); + for (size_t i = 0; i < threadpool_.size(); i++) { + matmul_results_.emplace_back(MAX_QUERY_BATCH_SIZE, batches[i].size()); + } + + // Precalculate centroid norms (required in L2 distances) + if constexpr (std::is_same_v) { + for (size_t i = 0; i < centroids_.size(); i++) { + centroids_norm_.push_back(distance::norm_square(centroids_.get_datum(i))); + } + } + } + + ///// Threading Interface + + static constexpr bool can_change_threads() { return false; } + /// + /// @brief Return the current number of threads used for search. + /// + size_t get_num_threads() const { return threadpool_.size(); } + + void set_threadpool(threads::ThreadPoolHandle threadpool) { + if (threadpool.size() != threadpool_.size()) { + throw ANNEXCEPTION("Threadpool change not supported for IVFIndex!"); + } + + threadpool_ = std::move(threadpool); + } + + /// + /// @brief Destroy the original thread pool and set to the provided one. + /// + /// @param threadpool An acceptable thread pool. + /// + /// @copydoc threadpool_requirements + /// + template + void set_threadpool(Pool threadpool) + requires(!std::is_same_v) + { + set_threadpool(threads::ThreadPoolHandle(std::move(threadpool))); + } + + /// + /// @brief Return the current thread pool handle. + /// + threads::ThreadPoolHandle& get_threadpool_handle() { return threadpool_; } + + size_t size() const { return centroids_.size(); } + size_t dimensions() const { return centroids_.dimensions(); } + + ///// Search Parameter Setting + search_parameters_type get_search_parameters() const { return search_parameters_; } + + void set_search_parameters(const search_parameters_type& search_parameters) { + search_parameters_ = search_parameters; + } + + auto search_centroids_closure() { + return [&](const auto& query, auto& buffer, size_t id) { + search_centroids( + query, + distance_, + matmul_results_, + buffer, + id, + centroids_norm_, + get_num_threads() + ); + }; + } + + auto search_leaves_closure() { + return [&](const auto& query, + auto& distance, + const auto& buffer_centroids, + auto& buffer_leaves, + size_t tid) { + search_leaves( + query, + distance, + cluster_, + buffer_centroids, + buffer_leaves, + threadpool_inner_[tid] + ); + }; + } + + // Search + template + void search( + QueryResultView results, + const Queries& queries, + const search_parameters_type& search_parameters, + const lib::DefaultPredicate& SVS_UNUSED(cancel) = lib::Returns(lib::Const()) + ) { + if (queries.size() > MAX_QUERY_BATCH_SIZE) { + throw ANNEXCEPTION( + "Number of queries {} higher than expected {}, increase value of " + "MAX_QUERY_BATCH_SIZE", + queries.size(), + MAX_QUERY_BATCH_SIZE + ); + } + + size_t num_neighbors = results.n_neighbors(); + size_t buffer_leaves_size = search_parameters.k_reorder_ * num_neighbors; + compute_centroid_distances(queries, centroids_, matmul_results_, threadpool_); + + threads::parallel_for( + threadpool_, + threads::StaticPartition(queries.size()), + [&](auto is, auto tid) { + auto buffer_centroids = SortedBuffer>( + search_parameters.n_probes_, distance::comparator(distance_) + ); + std::vector>> buffer_leaves; + for (size_t j = 0; j < n_inner_threads_; j++) { + buffer_leaves.push_back(SortedBuffer>( + buffer_leaves_size, distance::comparator(distance_) + )); + } + + // Pre-allocate scratch space needed by the dataset implementation. + auto scratch = + extensions::per_thread_batch_search_setup(cluster0_, distance_); + + // Perform a search over the batch of queries. + extensions::per_thread_batch_search( + cluster0_, + cluster_, + buffer_centroids, + buffer_leaves, + scratch, + queries, + results, + threads::UnitRange{is}, + tid, + search_centroids_closure(), + search_leaves_closure() + ); + } + ); + } + + std::string name() const { return "IVFIndex"; } + + private: + Centroids centroids_; + Cluster cluster_; + Data cluster0_; + Dist distance_; + threads::ThreadPoolHandle threadpool_; + const size_t n_inner_threads_ = 1; + std::vector threadpool_inner_; + std::vector> matmul_results_; + std::vector centroids_norm_; + + // Tunable Parameters + search_parameters_type search_parameters_{}; +}; + +template < + typename BuildType, + typename DataProto, + typename Distance, + typename ThreadpoolProto> +auto build_clustering( + const IVFBuildParameters& parameters, + const DataProto& data_proto, + Distance distance, + ThreadpoolProto threadpool_proto +) { + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto data = svs::detail::dispatch_load(data_proto, threadpool); + + auto tic = svs::lib::now(); + data::SimpleData centroids; + std::vector> clusters; + if (parameters.is_hierarchical_) { + std::tie(centroids, clusters) = hierarchical_kmeans_clustering( + parameters, data, distance, threadpool + ); + } else { + std::tie(centroids, clusters) = + kmeans_clustering(parameters, data, distance, threadpool); + } + + Clustering clustering(std::move(centroids), std::move(clusters)); + auto build_time = svs::lib::time_difference(tic); + fmt::print("IVF build time: {} seconds\n", build_time); + + auto logger = svs::logging::get(); + svs::logging::debug( + logger, "IVF Clustering Stats: {}", clustering.statistics().report("\n") + ); + + return clustering; +} + +template < + typename Clustering, + typename DataProto, + typename Distance, + typename ThreadpoolProto> +auto assemble_from_clustering( + Clustering clustering, + const DataProto& data_proto, + Distance distance, + ThreadpoolProto threadpool_proto, + const size_t n_inner_threads = 1 +) { + auto timer = lib::Timer(); + auto assemble_timer = timer.push_back("Total Assembling time"); + auto data_load_timer = timer.push_back("Data loading"); + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto data = svs::detail::dispatch_load(data_proto, threadpool); + data_load_timer.finish(); + + auto dense_cluster_timer = timer.push_back("Dense clustering"); + using centroids_type = data::SimpleData; + using data_type = typename decltype(data)::lib_alloc_data_type; + + auto dense_clusters = DenseClusteredDataset( + clustering, data, threadpool, lib::Allocator() + ); + dense_cluster_timer.finish(); + + auto index_build_timer = timer.push_back("IVF index construction"); + auto ivf_index = IVFIndex( + std::move(clustering.centroids()), + std::move(dense_clusters), + std::move(distance), + std::move(threadpool), + n_inner_threads + ); + index_build_timer.finish(); + assemble_timer.finish(); + svs::logging::debug("{}", timer); + return ivf_index; +} + +template < + typename Centroids, + typename DataProto, + typename Distance, + typename ThreadpoolProto> +auto assemble_from_file( + const std::filesystem::path& clustering_path, + const DataProto& data_proto, + Distance distance, + ThreadpoolProto threadpool_proto, + const size_t n_inner_threads = 1 +) { + using centroids_type = data::SimpleData; + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto clustering = svs::lib::load_from_disk>( + clustering_path, threadpool + ); + + return assemble_from_clustering( + std::move(clustering), + data_proto, + std::move(distance), + std::move(threadpool), + n_inner_threads + ); +} + +} // namespace svs::index::ivf diff --git a/include/svs/index/ivf/kmeans.h b/include/svs/index/ivf/kmeans.h new file mode 100644 index 00000000..17e8747b --- /dev/null +++ b/include/svs/index/ivf/kmeans.h @@ -0,0 +1,154 @@ +/* + * Copyright 2025 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 "svs/index/ivf/common.h" + +namespace svs::index::ivf { + +template < + typename BuildType, + data::ImmutableMemoryDataset Data, + typename Distance, + threads::ThreadPool Pool, + std::integral I = uint32_t> +auto kmeans_clustering_impl( + const IVFBuildParameters& parameters, + Data& data, + Distance& distance, + Pool& threadpool, + lib::Type SVS_UNUSED(integer_type) = {} +) { + auto timer = lib::Timer(); + auto kmeans_timer = timer.push_back("Non-hierarchical kmeans clustering"); + auto init_timer = timer.push_back("init"); + + constexpr size_t Dims = Data::extent; + using Alloc = svs::HugepageAllocator; + size_t ndims = data.dimensions(); + auto num_centroids = parameters.num_centroids_; + size_t num_training_data = + lib::narrow(std::ceil(data.size() * parameters.training_fraction_)); + if (num_training_data < num_centroids || num_training_data > data.size()) { + throw ANNEXCEPTION( + "Invalid number of training data: {}, num_centroids: {}, total data size: " + "{}\n", + num_training_data, + num_centroids, + data.size() + ); + } + + // The cluster centroids + auto centroids = data::SimpleData{num_centroids, ndims}; + auto data_train = data::SimpleData{num_training_data, ndims}; + auto matmul_results = + data::SimpleData{parameters.minibatch_size_, num_centroids}; + auto rng = std::mt19937(parameters.seed_); + + std::vector v(num_training_data); + generate_unique_ids(v, data.size(), rng); + threads::parallel_for( + threadpool, + threads::StaticPartition{num_training_data}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + data_train.set_datum(i, data.get_datum(v[i])); + } + } + ); + + v.resize(num_centroids); + generate_unique_ids(v, data_train.size(), rng); + threads::parallel_for( + threadpool, + threads::StaticPartition{num_centroids}, + [&](auto indices, auto /*tid*/) { + for (auto i : indices) { + centroids.set_datum(i, data_train.get_datum(v[i])); + } + } + ); + init_timer.finish(); + + auto centroids_fp32 = kmeans_training( + parameters, data_train, distance, centroids, matmul_results, rng, threadpool, timer + ); + + auto final_assignments_time = timer.push_back("final assignments"); + auto assignments = std::vector(data.size()); + auto batchsize = parameters.minibatch_size_; + auto num_batches = lib::div_round_up(data.size(), batchsize); + + std::vector data_norm; + if constexpr (std::is_same_v) { + generate_norms(data, data_norm, threadpool); + } + std::vector centroids_norm; + if constexpr (std::is_same_v) { + generate_norms(centroids_fp32, centroids_norm, threadpool); + } + + auto data_batch = data::SimpleData{batchsize, ndims}; + for (size_t batch = 0; batch < num_batches; ++batch) { + auto this_batch = threads::UnitRange{ + batch * batchsize, std::min((batch + 1) * batchsize, data.size())}; + auto data_batch_view = data::make_view(data, this_batch); + convert_data(data_batch_view, data_batch, threadpool); + centroid_assignment( + data_batch, + data_norm, + this_batch, + distance, + centroids, + centroids_norm, + assignments, + matmul_results, + threadpool, + timer + ); + } + + auto clusters = std::vector>(num_centroids); + for (auto i : data.eachindex()) { + clusters[assignments[i]].push_back(i); + } + final_assignments_time.finish(); + kmeans_timer.finish(); + svs::logging::debug("{}", timer); + fmt::print("kmeans clustering time: {}\n", lib::as_seconds(timer.elapsed())); + return std::make_tuple(centroids, std::move(clusters)); +} + +template < + typename BuildType, + data::ImmutableMemoryDataset Data, + typename Distance, + threads::ThreadPool Pool, + std::integral I = uint32_t> +auto kmeans_clustering( + const IVFBuildParameters& parameters, + Data& data, + Distance& distance, + Pool& threadpool, + lib::Type integer_type = {} +) { + return kmeans_clustering_impl( + parameters, data, distance, threadpool, integer_type + ); +} +} // namespace svs::index::ivf diff --git a/include/svs/index/ivf/sorted_buffer.h b/include/svs/index/ivf/sorted_buffer.h new file mode 100644 index 00000000..85939ef6 --- /dev/null +++ b/include/svs/index/ivf/sorted_buffer.h @@ -0,0 +1,202 @@ +/* + * Copyright 2025 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 "svs/index/vamana/filter.h" +#include "svs/lib/datatype.h" +#include "svs/lib/neighbor.h" +#include "svs/lib/prefetch.h" +#include "svs/lib/threads/threadlocal.h" + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace svs::index::ivf { +/// +/// @brief Class used to store search results for static greedy search. +/// +/// @tparam Idx Type used to uniquely identify DB vectors +/// @tparam Cmp Type of the comparison function used to sort neighbors by distance. +/// +template > class SortedBuffer { + public: + // External type aliases + using value_type = IVFNeighbor; + using reference = value_type&; + using const_reference = const value_type&; + using compare_type = Cmp; + + using vector_type = std::vector>; + using iterator = typename vector_type::iterator; + using const_iterator = typename vector_type::const_iterator; + + /// A visited filter with 65,535 entries with a memory footpring of 128 kiB. + using set_type = index::vamana::VisitedFilter; + + /// + /// @brief Initialize a buffer with zero capacity. + /// + /// In order to use a buffer that has been default constructed, use the + /// @ref change_maxsize(size_t) method. + /// + SortedBuffer() = default; + + /// + /// @brief Construct a search buffer with the target capacity and comparison function. + /// + /// @param size The number of valid elements to return from a search operation. + /// @param compare The functor used to compare two ``Neighbor``s together. + /// + explicit SortedBuffer(size_t size, Cmp compare = {}) + : compare_{std::move(compare)} + , capacity_{size} + , candidates_{capacity_ + 1} {} + + /// + /// @brief Perform an efficient copy. + /// + /// Copy the portions of the SortedBuffer that matter for the purposes of scratch + /// space. + /// + /// Perserves the sizes of various containers but not necessarily the values. + /// + SortedBuffer shallow_copy() const { return SortedBuffer{capacity_, compare_}; } + + /// + /// @brief Change the target number of elements to return after search. + /// + /// @param new_size The new number of elements to return. + /// + /// Post conditions + /// - The capacity of the search buffer will be set to the new size. + /// - The actual size (number of valid elements) will be the minimum of the current + /// size and the new size. + /// + void change_maxsize(size_t new_size) { + capacity_ = new_size; + candidates_.resize(new_size + 1); + size_ = std::min(size_, new_size); + } + + /// + /// @brief Prepare the buffer for a new search operation. + /// + void clear() { size_ = 0; } + + /// @brief Return the current number of valid elements in the buffer. + size_t size() const { return size_; } + + /// @brief Return the maximum number of neighbors that can be held by the buffer. + size_t capacity() const { return capacity_; } + + /// @brief Return whether or not the buffer is full of valid elements. + bool full() const { return size() == capacity(); } + + /// @brief Access the neighbor at position `i`. + reference operator[](size_t i) { return candidates_[i]; } + + /// @brief Access the neighbor at position `i`. + const_reference operator[](size_t i) const { return candidates_[i]; } + + /// @brief Return the furtherst valid neighbor. + reference back() { return candidates_[size_ - 1]; } + + /// @brief Return the furtherst valid neighbor. + const_reference back() const { return candidates_[size_ - 1]; } + + // Define iterators. + constexpr const_iterator begin() const noexcept { return candidates_.begin(); } + constexpr const_iterator end() const noexcept { return begin() + size(); } + constexpr iterator begin() noexcept { return candidates_.begin(); } + constexpr iterator end() noexcept { return begin() + size(); } + + void unsafe_insert(value_type neighbor, iterator index) { + std::copy_backward(index, end(), end() + 1); + (*index) = neighbor; + } + + /// + /// @brief Return ``true`` if a neighbor with the given distance can be skipped. + /// + bool can_skip(float distance) const { + return compare_(back().distance(), distance) && full(); + } + + /// + /// @brief Insert the neighbor into the buffer. + /// + /// @param neighbor The neighbor to insert. + /// + /// @returns The position where the neighbor was inserted. + /// + size_t insert(value_type neighbor) { + if (can_skip(neighbor.distance())) { + return size(); + } + return insert_inner(neighbor); + } + + size_t insert_inner(value_type neighbor) { + const auto start = begin(); + // Binary search to the first location where `distance` is less than the stored + // neighbor. + auto pos = std::lower_bound( + start, + end(), + neighbor.distance(), + [&](const value_type& other, const float& d) { + return !compare_(d, other.distance()); + } + ); + + size_t i = pos - start; + unsafe_insert(neighbor, pos); + size_ = std::min(size_ + 1, capacity()); + return i; + } + + /// + /// @brief Sort the elements in the buffer according to the internal comparison functor. + /// + void sort() { std::sort(begin(), end(), compare_); } + + /// + /// @brief Return ``true`` if the visited set is enabled. Otherwise, return ``false``. + /// + bool visited_set_enabled() const { return visited_.has_value(); } + + private: + // The comparison functor. + [[no_unique_address]] Cmp compare_ = Cmp{}; + // The current number of valid neighbors. + size_t size_ = 0; + // The maximum capacity of the buffer. + size_t capacity_ = 0; + // Storage for the neighbors. + vector_type candidates_ = {}; + // The visited set. Implemented as a `std::optional` to allow enablind and disabling + // without always requiring allocation of the data structure. + std::optional visited_{std::nullopt}; +}; + +} // namespace svs::index::ivf diff --git a/include/svs/lib/bfloat16.h b/include/svs/lib/bfloat16.h new file mode 100644 index 00000000..a1c257a6 --- /dev/null +++ b/include/svs/lib/bfloat16.h @@ -0,0 +1,143 @@ +/* + * Copyright 2025 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 "svs/lib/narrow.h" +#include "svs/lib/type_traits.h" +#include "svs/third-party/fmt.h" + +#include +#include +#include +#include +#include +#if defined(__AVX512_BF16__) +#include +#endif + +namespace svs { +namespace bfloat16 { +namespace detail { + +// TODO: Update to `bitcast` if its available in the standard library. +inline uint32_t bitcast_float_to_uint32(const float x) { + static_assert(sizeof(float) == sizeof(uint32_t)); + uint32_t u; + memcpy(&u, &x, sizeof(x)); + return u; +} + +inline float bitcast_uint32_to_float(const uint32_t x) { + static_assert(sizeof(float) == sizeof(uint32_t)); + float f; + memcpy(&f, &x, sizeof(x)); + return f; +} + +inline float bfloat16_to_float_untyped(const uint16_t x) { + const uint32_t u = x << 16; + return bitcast_uint32_to_float(u); +} + +inline uint16_t float_to_bfloat16_untyped(const float x) { + const uint32_t u = bitcast_float_to_uint32(x); + return lib::narrow(u >> 16); +} + +} // namespace detail + +// On GCC, we need to add this attribute so that BFloat16 members can appear inside +// packed structs. +class __attribute__((packed)) BFloat16 { + public: + BFloat16() = default; + + // converting constructors + explicit BFloat16(float x) + : value_{detail::float_to_bfloat16_untyped(x)} {} + explicit BFloat16(double x) + : BFloat16(lib::narrow_cast(x)) {} + explicit BFloat16(size_t x) + : BFloat16(lib::narrow(x)) {} + explicit BFloat16(int x) + : BFloat16(lib::narrow(x)) {} + + // conversion functions + operator float() const { return detail::bfloat16_to_float_untyped(value_); } + BFloat16& operator=(float x) { + value_ = detail::float_to_bfloat16_untyped(x); + return *this; + } + + // Allow users to set and expect the contents of the class as a uint16_t using an + // explicit API. + static BFloat16 from_raw(uint16_t value) { return BFloat16{value, FromRawTag{}}; } + uint16_t raw() const { return value_; } + + private: + // Use a tag to construct from a raw value in order to still allow a constructor + // for a lone `uint16_t`. + struct FromRawTag {}; + explicit BFloat16(uint16_t value, FromRawTag /*unused*/) + : value_{value} {} + uint16_t value_; +}; +static_assert(std::is_trivial_v); +static_assert(std::is_standard_layout_v); + +///// +///// Operators +///// + +// For equality, still use `float` rather than the underlying bit pattern to handle cases +// like signed zeros. +inline bool operator==(BFloat16 x, BFloat16 y) { return float{x} == float{y}; } + +///// +///// Pretty Printing +///// + +} // namespace bfloat16 + +using BFloat16 = bfloat16::BFloat16; + +// SVS local arithmetric trait. +template <> inline constexpr bool is_arithmetic_v = true; +template <> inline constexpr bool is_signed_v = true; +template <> inline constexpr bool allow_lossy_conversion = true; + +} // namespace svs + +// Apply hashing to `BFloat16` +namespace std { +template <> struct hash { + inline std::size_t operator()(const svs::BFloat16& x) const noexcept { + return std::hash()(x); + } +}; +} // namespace std + +// Formatting and Printing +template <> struct fmt::formatter : svs::format_empty { + auto format(svs::BFloat16 x, auto& ctx) const { + return fmt::format_to(ctx.out(), "{}f16", float{x}); + } +}; + +inline std::ostream& operator<<(std::ostream& stream, svs::BFloat16 x) { + return stream << fmt::format("{}", x); +} diff --git a/include/svs/lib/datatype.h b/include/svs/lib/datatype.h index d0b94d49..49955707 100644 --- a/include/svs/lib/datatype.h +++ b/include/svs/lib/datatype.h @@ -22,6 +22,7 @@ /// // local deps +#include "svs/lib/bfloat16.h" #include "svs/lib/exception.h" #include "svs/lib/float16.h" #include "svs/third-party/fmt.h" @@ -50,6 +51,7 @@ enum class DataType { int32, int64, float16, + bfloat16, float32, float64, byte, @@ -74,6 +76,9 @@ template <> inline constexpr std::string_view name() { return " template <> inline constexpr std::string_view name() { return "float16"; } +template <> inline constexpr std::string_view name() { + return "bfloat16"; +} template <> inline constexpr std::string_view name() { return "float32"; } @@ -101,6 +106,7 @@ inline constexpr std::string_view name(DataType type) { case DataType::int64: { return name(); } case DataType::float16: { return name(); } + case DataType::bfloat16: { return name(); } case DataType::float32: { return name(); } case DataType::float64: { return name(); } @@ -126,6 +132,7 @@ inline constexpr size_t element_size(DataType type) { case DataType::int64: { return sizeof(int64_t); } case DataType::float16: { return sizeof(svs::Float16); } + case DataType::bfloat16: { return sizeof(svs::BFloat16); } case DataType::float32: { return sizeof(float); } case DataType::float64: { return sizeof(double); } @@ -140,6 +147,9 @@ inline constexpr DataType parse_datatype_floating(std::string_view name) { if (name == "float16") { return DataType::float16; } + if (name == "bfloat16") { + return DataType::bfloat16; + } if (name == "float32") { return DataType::float32; } @@ -191,7 +201,7 @@ inline constexpr DataType parse_datatype(std::string_view name) { } // Floating point. - if (name.starts_with("float")) { + if (name.starts_with("float") || name.starts_with("bfloat")) { return parse_datatype_floating(name); } if (name.starts_with("uint")) { @@ -263,6 +273,7 @@ template <> struct CppType { using type = int32_t; }; template <> struct CppType { using type = int64_t; }; template <> struct CppType { using type = Float16; }; +template <> struct CppType { using type = BFloat16; }; template <> struct CppType { using type = float; }; template <> struct CppType { using type = double; }; @@ -289,6 +300,7 @@ template<> inline constexpr DataType datatype_v = DataType::int32; template<> inline constexpr DataType datatype_v = DataType::int64; template<> inline constexpr DataType datatype_v = DataType::float16; +template<> inline constexpr DataType datatype_v = DataType::bfloat16; template<> inline constexpr DataType datatype_v = DataType::float32; template<> inline constexpr DataType datatype_v = DataType::float64; diff --git a/include/svs/lib/file_iterator.h b/include/svs/lib/file_iterator.h index 0ef8b79e..63f7bcd3 100644 --- a/include/svs/lib/file_iterator.h +++ b/include/svs/lib/file_iterator.h @@ -54,9 +54,10 @@ template io_convert_type_t io_convert(U u) { return narrow>(u); } -// We expect the conversion to `Float16` to be lossy. +// We expect the conversion to `Float16` and `BFloat16` to be lossy. // Hence, don't require precise narrowing. template <> inline Float16 io_convert(float u) { return Float16(u); } +template <> inline BFloat16 io_convert(float u) { return BFloat16(u); } ///// ///// Iterator Helpers diff --git a/include/svs/lib/neighbor.h b/include/svs/lib/neighbor.h index f31d498a..885d918a 100644 --- a/include/svs/lib/neighbor.h +++ b/include/svs/lib/neighbor.h @@ -62,6 +62,7 @@ template struct Neighbor : public Meta { constexpr float distance() const { return distance_; } constexpr void set_distance(float new_distance) { distance_ = new_distance; } constexpr Idx id() const { return id_; } + constexpr void set_id(Idx new_id) { id_ = new_id; } // Members Idx id_; @@ -82,7 +83,9 @@ template struct Neighbor { : Neighbor(lib::narrow(other.id()), other.distance()) {} constexpr float distance() const { return distance_; } + constexpr void set_distance(float new_distance) { distance_ = new_distance; } constexpr Idx id() const { return id_; } + constexpr void set_id(Idx new_id) { id_ = new_id; } // Members Idx id_; @@ -246,4 +249,24 @@ class ValidVisit { /// Type alias for skippable neighbor. template using PredicatedSearchNeighbor = Neighbor; +///// +///// IVF Neighbor +///// + +template struct LocalId { + constexpr LocalId(Idx local_id = 0) + : local_id_{local_id} {} + constexpr Idx get_local_id() const { return local_id_; } + constexpr void set_local_id(Idx local_id) { local_id_ = local_id; } + friend constexpr bool operator==(LocalId, LocalId) = default; + + // members + Idx local_id_; +}; + +/// +/// Type alias for the IVF neighbor +/// +template using IVFNeighbor = Neighbor>; + } // namespace svs diff --git a/include/svs/orchestrators/ivf.h b/include/svs/orchestrators/ivf.h new file mode 100644 index 00000000..7c035f11 --- /dev/null +++ b/include/svs/orchestrators/ivf.h @@ -0,0 +1,180 @@ +/* + * 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 "svs/index/ivf/index.h" +#include "svs/orchestrators/manager.h" + +namespace svs { + +class IVFInterface { + public: + using search_parameters_type = svs::index::ivf::IVFSearchParameters; + + ///// Backend information interface + virtual std::string experimental_backend_string() const = 0; +}; + +template +class IVFImpl : public manager::ManagerImpl { + private: + // Null-terminated array of characters. + static constexpr auto typename_impl = lib::generate_typename(); + + public: + using base_type = manager::ManagerImpl; + using base_type::impl; + using search_parameters_type = typename IFace::search_parameters_type; + + explicit IVFImpl(Impl impl) + : base_type{std::move(impl)} {} + + ///// Parameter Interface + [[nodiscard]] search_parameters_type get_search_parameters() const override { + return impl().get_search_parameters(); + } + + void set_search_parameters(const search_parameters_type& search_parameters) override { + impl().set_search_parameters(search_parameters); + } + + ///// Backend Information Interface + [[nodiscard]] std::string experimental_backend_string() const override { + return std::string{typename_impl.begin(), typename_impl.end() - 1}; + } +}; + +///// +///// IVFManager +///// + +class IVF : public manager::IndexManager { + // Type Alises + public: + using base_type = manager::IndexManager; + using search_parameters_type = typename IVFInterface::search_parameters_type; + + // Constructors + IVF(std::unique_ptr> impl) + : base_type{std::move(impl)} {} + + template + IVF(std::in_place_t, QueryTypes SVS_UNUSED(type), Impl&& impl) + : base_type{std::make_unique>(SVS_FWD(impl))} {} + + ///// Backend String + std::string experimental_backend_string() const { + return impl_->experimental_backend_string(); + } + + ///// Assembling + template < + manager::QueryTypeDefinition QueryTypes, + typename Clustering, + typename DataProto, + typename Distance, + typename ThreadpoolProto> + static IVF assemble_from_clustering( + Clustering clustering, + const DataProto& data_proto, + const Distance& distance, + ThreadpoolProto threadpool_proto, + size_t intra_query_threads = 1 + ) { + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + if constexpr (std::is_same_v, DistanceType>) { + auto dispatcher = DistanceDispatcher(distance); + return dispatcher([&](auto distance_function) { + return IVF( + std::in_place, + manager::as_typelist{}, + index::ivf::assemble_from_clustering( + std::move(clustering), + data_proto, + std::move(distance_function), + std::move(threadpool), + intra_query_threads + ) + ); + }); + } else { + return IVF( + std::in_place, + manager::as_typelist{}, + index::ivf::assemble_from_clustering( + std::move(clustering), + data_proto, + distance, + std::move(threadpool), + intra_query_threads + ) + ); + } + } + + template < + manager::QueryTypeDefinition QueryTypes, + typename Centroids, + typename DataProto, + typename Distance, + typename ThreadpoolProto> + static IVF assemble_from_file( + const std::filesystem::path& clustering_path, + const DataProto& data_proto, + const Distance& distance, + ThreadpoolProto threadpool_proto, + size_t intra_query_threads = 1 + ) { + using centroids_type = data::SimpleData; + auto threadpool = threads::as_threadpool(std::move(threadpool_proto)); + auto clustering = + svs::lib::load_from_disk>( + clustering_path, threadpool + ); + return assemble_from_clustering( + std::move(clustering), + data_proto, + distance, + std::move(threadpool), + intra_query_threads + ); + } + + ///// Building + template + static auto build_clustering( + const index::ivf::IVFBuildParameters& build_parameters, + const DataProto& data_proto, + const Distance& distance, + size_t num_threads + ) { + if constexpr (std::is_same_v, DistanceType>) { + auto dispatcher = DistanceDispatcher(distance); + return dispatcher([&](auto distance_function) { + return index::ivf::build_clustering( + build_parameters, data_proto, std::move(distance_function), num_threads + ); + }); + } else { + return index::ivf::build_clustering( + build_parameters, data_proto, distance, num_threads + ); + } + } +}; + +} // namespace svs diff --git a/tests/CMakeLists.txt b/tests/CMakeLists.txt index 8bda49ac..bab566f2 100644 --- a/tests/CMakeLists.txt +++ b/tests/CMakeLists.txt @@ -65,6 +65,7 @@ set(TEST_SOURCES ${TEST_DIR}/svs/lib/invoke.cpp ${TEST_DIR}/svs/lib/file.cpp ${TEST_DIR}/svs/lib/file_iterator.cpp + ${TEST_DIR}/svs/lib/bfloat16.cpp ${TEST_DIR}/svs/lib/float16.cpp ${TEST_DIR}/svs/lib/memory.cpp ${TEST_DIR}/svs/lib/meta.cpp @@ -169,6 +170,14 @@ SET(INTEGRATION_TESTS # ${TEST_DIR}/integration/numa_search.cpp -- requires SVS_EXPERIMENTAL_ENABLE_NUMA ) +if (SVS_EXPERIMENTAL_ENABLE_IVF) + message("Enabling IVF Integration tests!") + list(APPEND INTEGRATION_TESTS + ${TEST_DIR}/integration/ivf/index_build.cpp + ${TEST_DIR}/integration/ivf/index_search.cpp + ) +endif() + if ((NOT CMAKE_BUILD_TYPE STREQUAL "Debug") OR SVS_FORCE_INTEGRATION_TESTS) message("Enabling Integration Tests!") list(APPEND TEST_SOURCES ${INTEGRATION_TESTS}) @@ -188,6 +197,16 @@ if (SVS_EXPERIMENTAL_ENABLE_NUMA) list(APPEND TEST_SOURCES ${NUMA_TESTS}) endif() +if (SVS_EXPERIMENTAL_ENABLE_IVF) + message("Enabling IVF tests!") + list(APPEND TEST_SOURCES + ${TEST_DIR}/utils/ivf_reference.cpp + ${TEST_DIR}/svs/index/ivf/kmeans.cpp + ${TEST_DIR}/svs/index/ivf/hierarchical_kmeans.cpp + ${TEST_DIR}/svs/index/ivf/common.cpp + ) +endif() + add_executable(tests ${TEST_SOURCES}) # Path to the test dataset. diff --git a/tests/integration/ivf/index_build.cpp b/tests/integration/ivf/index_build.cpp new file mode 100644 index 00000000..3e362203 --- /dev/null +++ b/tests/integration/ivf/index_build.cpp @@ -0,0 +1,114 @@ +/* + * Copyright 2025 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. + */ + +// svs +#include "svs/core/data/simple.h" +#include "svs/core/recall.h" +#include "svs/lib/timing.h" +#include "svs/orchestrators/ivf.h" + +// svsbenchmark +#include "svs-benchmark/benchmark.h" + +// fmt +#include "fmt/core.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" + +// tests +#include "tests/utils/ivf_reference.h" +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" + +// stl +#include +#include +#include +#include + +namespace { + +template +auto build_index( + const svs::index::ivf::IVFBuildParameters parameters, + const std::filesystem::path& data_path, + size_t num_threads, + size_t num_inner_threads, + const Distance& dist_type +) { + auto data = svs::data::SimpleData::load(data_path); + auto clustering = + svs::IVF::build_clustering(parameters, data, dist_type, num_threads); + + return svs::IVF::assemble_from_clustering( + std::move(clustering), std::move(data), dist_type, num_threads, num_inner_threads + ); +} + +template +void test_build(const Distance& distance, size_t num_inner_threads = 1) { + const double epsilon = 0.005; + const auto queries = svs::data::SimpleData::load(test_dataset::query_file()); + CATCH_REQUIRE(svs_test::prepare_temp_directory()); + size_t num_threads = 2; + + auto expected_result = test_dataset::ivf::expected_build_results( + svs::distance_type_v, svsbenchmark::Uncompressed(svs::datatype_v) + ); + auto index = build_index( + expected_result.build_parameters_.value(), + test_dataset::data_svs_file(), + num_threads, + num_inner_threads, + distance + ); + + auto groundtruth = test_dataset::load_groundtruth(svs::distance_type_v); + for (const auto& expected : expected_result.config_and_recall_) { + auto these_queries = test_dataset::get_test_set(queries, expected.num_queries_); + auto these_groundtruth = + test_dataset::get_test_set(groundtruth, expected.num_queries_); + index.set_search_parameters(expected.search_parameters_); + auto results = index.search(these_queries, expected.num_neighbors_); + double recall = svs::k_recall_at_n( + these_groundtruth, results, expected.num_neighbors_, expected.recall_k_ + ); + + fmt::print( + "n_probes: {}, Expected Recall: {}, Actual Recall: {}\n", + index.get_search_parameters().n_probes_, + expected.recall_, + recall + ); + CATCH_REQUIRE(recall > expected.recall_ - epsilon); + CATCH_REQUIRE(recall < expected.recall_ + epsilon); + } +} + +} // namespace + +CATCH_TEST_CASE("IVF Build/Clustering", "[integration][build][ivf]") { + test_build(svs::DistanceL2()); + test_build(svs::DistanceIP()); + + test_build(svs::DistanceL2()); + test_build(svs::DistanceIP()); + + // With 4 inner threads + test_build(svs::DistanceL2(), 4); + test_build(svs::DistanceIP(), 4); +} diff --git a/tests/integration/ivf/index_search.cpp b/tests/integration/ivf/index_search.cpp new file mode 100644 index 00000000..cec5fcda --- /dev/null +++ b/tests/integration/ivf/index_search.cpp @@ -0,0 +1,141 @@ +/* + * Copyright 2025 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. + */ + +// stl +#include +#include +#include +#include +#include +#include + +// svs +#include "svs/core/recall.h" +#include "svs/lib/saveload.h" +#include "svs/orchestrators/ivf.h" + +// svsbenchmark +#include "svs-benchmark/benchmark.h" +#include "svs-benchmark/test.h" + +// catch2 +#include "catch2/catch_test_macros.hpp" +#include "catch2/matchers/catch_matchers_string.hpp" + +// fmt +#include "fmt/core.h" + +// tests +#include "tests/utils/ivf_reference.h" +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" + +namespace { + +void run_search( + svs::IVF& index, + const svs::data::SimpleData& queries_all, + const svs::data::SimpleData& groundtruth_all, + const std::vector& expected_results +) { + // If we make a change that somehow improves accuracy, we'll want to know. + // Use `epsilon` to add to the expected results to set an upper bound on the + // achieved accuracy. + double epsilon = 0.005; + + // Ensure we have at least one entry in the expected results. + CATCH_REQUIRE(!expected_results.empty()); + + const auto queries_in_test_set = expected_results.at(0).num_queries_; + + auto queries = test_dataset::get_test_set(queries_all, queries_in_test_set); + auto groundtruth = test_dataset::get_test_set(groundtruth_all, queries_in_test_set); + + for (const auto& expected : expected_results) { + // Update the query set if needed. + auto num_queries = expected.num_queries_; + if (num_queries != queries.size()) { + queries = test_dataset::get_test_set(queries_all, num_queries); + groundtruth = test_dataset::get_test_set(groundtruth_all, num_queries); + } + + // Configure the index with the current parameters. + // Ensure that the result sticks. + index.set_search_parameters(expected.search_parameters_); + CATCH_REQUIRE(index.get_search_parameters() == expected.search_parameters_); + + // Float32 + auto results = index.search(queries, expected.num_neighbors_); + auto recall = svs::k_recall_at_n( + groundtruth, results, expected.num_neighbors_, expected.recall_k_ + ); + fmt::print( + "n_probes: {}, Expected Recall: {}, Actual Recall: {}\n", + index.get_search_parameters().n_probes_, + expected.recall_, + recall + ); + + CATCH_REQUIRE(recall > expected.recall_ - epsilon); + CATCH_REQUIRE(recall < expected.recall_ + epsilon); + } +} + +template +void test_search( + svs::data::SimpleData data, + const Distance& distance, + const svs::data::SimpleData& queries, + const svs::data::SimpleData& groundtruth, + const size_t num_inner_threads = 1 +) { + size_t num_threads = 2; + + // Find the expected results for this dataset. + auto expected_result = test_dataset::ivf::expected_search_results( + svs::distance_type_v, svsbenchmark::Uncompressed(svs::datatype_v) + ); + + auto index = svs::IVF::assemble_from_file( + test_dataset::clustering_directory(), data, distance, num_threads, num_inner_threads + ); + CATCH_REQUIRE(index.get_num_threads() == num_threads); + + run_search(index, queries, groundtruth, expected_result.config_and_recall_); + CATCH_REQUIRE(index.dimensions() == test_dataset::NUM_DIMENSIONS); +} + +} // namespace + +CATCH_TEST_CASE("IVF Search", "[integration][search][ivf]") { + namespace ivf = svs::index::ivf; + + auto datafile = test_dataset::data_svs_file(); + auto queries = test_dataset::queries(); + auto gt_l2 = test_dataset::groundtruth_euclidean(); + auto gt_ip = test_dataset::groundtruth_mip(); + + auto dist_l2 = svs::distance::DistanceL2(); + auto dist_ip = svs::distance::DistanceIP(); + + auto data = svs::data::SimpleData::load(datafile); + auto data_f16 = svs::index::ivf::convert_data(data); + test_search(data, dist_l2, queries, gt_l2); + test_search(data, dist_l2, queries, gt_l2, 2); + + test_search(data_f16, dist_ip, queries, gt_ip); + test_search(data_f16, dist_ip, queries, gt_ip, 2); +} diff --git a/tests/svs/core/data/data.h b/tests/svs/core/data/data.h index 4ebe2a07..72cd9821 100644 --- a/tests/svs/core/data/data.h +++ b/tests/svs/core/data/data.h @@ -25,6 +25,7 @@ namespace svs_test::data { // A mock dataset with integer entries. class MockDataset { public: + using element_type = int64_t; using value_type = int64_t; using const_value_type = int64_t; diff --git a/tests/svs/index/ivf/common.cpp b/tests/svs/index/ivf/common.cpp new file mode 100644 index 00000000..39df8503 --- /dev/null +++ b/tests/svs/index/ivf/common.cpp @@ -0,0 +1,71 @@ +/* + * Copyright 2025 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. + */ + +// header under test +#include "svs/index/ivf/common.h" + +// tests +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" + +// catch +#include "catch2/catch_test_macros.hpp" + +CATCH_TEST_CASE("Kmeans Clustering", "[ivf][parameters]") { + namespace ivf = svs::index::ivf; + CATCH_SECTION("IVF Build Parameters") { + auto p = ivf::IVFBuildParameters(); + + // Test the setter methods. +#define XX(name, v) \ + CATCH_REQUIRE(p.name##_ != v); \ + CATCH_REQUIRE(p.name(v).name##_ == v); \ + CATCH_REQUIRE(p.name##_ == v); + + XX(num_centroids, 10); + XX(minibatch_size, 100); + XX(num_iterations, 1000); + XX(is_hierarchical, false); + XX(training_fraction, 0.05F); + XX(hierarchical_level1_clusters, 10); + XX(seed, 0x1234); +#undef XX + + // Saving and loading + svs_test::prepare_temp_directory(); + auto dir = svs_test::temp_directory(); + CATCH_REQUIRE(svs::lib::test_self_save_load(p, dir)); + } + + CATCH_SECTION("IVF Search Parameters") { + auto p = ivf::IVFSearchParameters(); + + // Test the setter methods. +#define XX(name, v) \ + CATCH_REQUIRE(p.name##_ != v); \ + CATCH_REQUIRE(p.name(v).name##_ == v); \ + CATCH_REQUIRE(p.name##_ == v); + + XX(n_probes, 10); + XX(k_reorder, 100); +#undef XX + + // Saving and loading + svs_test::prepare_temp_directory(); + auto dir = svs_test::temp_directory(); + CATCH_REQUIRE(svs::lib::test_self_save_load(p, dir)); + } +} diff --git a/tests/svs/index/ivf/hierarchical_kmeans.cpp b/tests/svs/index/ivf/hierarchical_kmeans.cpp new file mode 100644 index 00000000..db15940c --- /dev/null +++ b/tests/svs/index/ivf/hierarchical_kmeans.cpp @@ -0,0 +1,75 @@ +/* + * Copyright 2025 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. + */ + +// header under test +#include "svs/index/ivf/hierarchical_kmeans.h" + +// tests +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" + +// catch +#include "catch2/catch_test_macros.hpp" + +// stl +#include + +namespace { + +template +void test_hierarchical_kmeans_clustering(const Data& data, Distance distance) { + namespace ivf = svs::index::ivf; + + for (size_t n_centroids : {155}) { + for (size_t minibatch : {25}) { + for (size_t iters : {3}) { + for (float training_fraction : {0.55}) { + for (size_t l1_clusters : {0, 9}) { + auto params = ivf::IVFBuildParameters() + .num_centroids(n_centroids) + .minibatch_size(minibatch) + .num_iterations(iters) + .is_hierarchical(true) + .training_fraction(training_fraction) + .hierarchical_level1_clusters(l1_clusters); + + auto threadpool = svs::threads::as_threadpool(10); + auto [centroids, clusters] = + hierarchical_kmeans_clustering( + params, data, distance, threadpool + ); + + CATCH_REQUIRE(centroids.size() == n_centroids); + CATCH_REQUIRE(centroids.dimensions() == data.dimensions()); + CATCH_REQUIRE(clusters.size() == n_centroids); + } + } + } + } + } +} + +} // namespace + +CATCH_TEST_CASE("Hierarchical Kmeans Param Check", "[ivf][hierarchial_parameter_check]") { + CATCH_SECTION("Uncompressed Data") { + auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); + test_hierarchical_kmeans_clustering(data, svs::DistanceIP()); + test_hierarchical_kmeans_clustering(data, svs::DistanceIP()); + test_hierarchical_kmeans_clustering(data, svs::DistanceIP()); + test_hierarchical_kmeans_clustering(data, svs::DistanceL2()); + } +} diff --git a/tests/svs/index/ivf/kmeans.cpp b/tests/svs/index/ivf/kmeans.cpp new file mode 100644 index 00000000..49e20a10 --- /dev/null +++ b/tests/svs/index/ivf/kmeans.cpp @@ -0,0 +1,70 @@ +/* + * Copyright 2025 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. + */ + +// header under test +#include "svs/index/ivf/kmeans.h" + +// tests +#include "tests/utils/test_dataset.h" +#include "tests/utils/utils.h" + +// catch +#include "catch2/catch_test_macros.hpp" + +// stl +#include + +namespace { + +template +void test_kmeans_clustering(const Data& data, Distance distance) { + namespace ivf = svs::index::ivf; + + for (size_t n_centroids : {1, 99}) { + for (size_t minibatch : {25}) { + for (size_t iters : {3}) { + for (float training_fraction : {0.55}) { + auto params = ivf::IVFBuildParameters() + .num_centroids(n_centroids) + .minibatch_size(minibatch) + .num_iterations(iters) + .is_hierarchical(false) + .training_fraction(training_fraction); + auto threadpool = svs::threads::as_threadpool(10); + auto [centroids, clusters] = ivf::kmeans_clustering( + params, data, distance, threadpool + ); + + CATCH_REQUIRE(centroids.size() == n_centroids); + CATCH_REQUIRE(centroids.dimensions() == data.dimensions()); + CATCH_REQUIRE(clusters.size() == n_centroids); + } + } + } + } +} + +} // namespace + +CATCH_TEST_CASE("Build Kmeans Param Check", "[ivf][parameter_check]") { + CATCH_SECTION("Uncompressed Data") { + auto data = svs::data::SimpleData::load(test_dataset::data_svs_file()); + test_kmeans_clustering(data, svs::DistanceIP()); + test_kmeans_clustering(data, svs::DistanceIP()); + test_kmeans_clustering(data, svs::DistanceIP()); + test_kmeans_clustering(data, svs::DistanceL2()); + } +} diff --git a/tests/svs/lib/bfloat16.cpp b/tests/svs/lib/bfloat16.cpp new file mode 100644 index 00000000..ae011b77 --- /dev/null +++ b/tests/svs/lib/bfloat16.cpp @@ -0,0 +1,79 @@ +/* + * Copyright 2025 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 + +#include "svs/lib/bfloat16.h" +#include "svs/lib/narrow.h" + +#include "catch2/catch_test_macros.hpp" + +CATCH_TEST_CASE("Testing BFloat16", "[core][bfloat16]") { + CATCH_SECTION("Implicit Conversion") { + svs::BFloat16 x{1.0f}; + float y = x; + CATCH_REQUIRE(y == 1.0f); + + x = svs::BFloat16{-1}; + CATCH_REQUIRE(float{x} == -1.0f); + + // Construct from `size_t` + x = svs::BFloat16{size_t{100}}; + CATCH_REQUIRE(float{x} == 100.0f); + + // Default Construction. + CATCH_REQUIRE(svs::BFloat16{} == svs::BFloat16(float{0})); + } + + CATCH_SECTION("Arithmetic") { + CATCH_REQUIRE(svs::is_arithmetic_v); + CATCH_REQUIRE(svs::is_signed_v); + + auto x = svs::BFloat16{1}; + auto y = svs::BFloat16{2}; + CATCH_REQUIRE(x + y == 3); + CATCH_REQUIRE(x != y); + CATCH_REQUIRE(x < y); + CATCH_REQUIRE(!(y < x)); + CATCH_REQUIRE(y - x == svs::BFloat16(1)); + } + + CATCH_SECTION("Narrow") { + float x_good{1.0f}; + auto y_good = svs::lib::narrow(x_good); + CATCH_REQUIRE(float{y_good} == x_good); + + // Use low precision outside the bounds of `BFloat16`. + float x_bad{0.000012f}; + CATCH_REQUIRE_THROWS_AS( + svs::lib::narrow(x_bad), svs::lib::narrowing_error + ); + + // Fail when constructing from typemax integers + CATCH_REQUIRE_THROWS_AS( + svs::BFloat16(std::numeric_limits::max() - 1), svs::lib::narrowing_error + ); + + CATCH_REQUIRE_THROWS_AS( + svs::BFloat16(std::numeric_limits::max() - int{1}), + svs::lib::narrowing_error + ); + CATCH_REQUIRE_THROWS_AS( + svs::BFloat16(std::numeric_limits::min() + int{1}), + svs::lib::narrowing_error + ); + } +} diff --git a/tests/svs/lib/datatype.cpp b/tests/svs/lib/datatype.cpp index a0e7ca22..a3573531 100644 --- a/tests/svs/lib/datatype.cpp +++ b/tests/svs/lib/datatype.cpp @@ -65,6 +65,7 @@ CATCH_TEST_CASE("Data Type", "[core][datatype]") { CATCH_REQUIRE(svs::parse_datatype("int128") == DataType::undef); test("float16"); + test("bfloat16"); test("float32"); test("float64"); CATCH_REQUIRE(svs::parse_datatype("float128") == DataType::undef); @@ -86,12 +87,15 @@ CATCH_TEST_CASE("Data Type", "[core][datatype]") { // Use in a hash table. std::unordered_map table{}; table[DataType::float16] = 5; + table[DataType::bfloat16] = 6; table[DataType::float32] = 10; CATCH_REQUIRE(!table.contains(DataType::int8)); CATCH_REQUIRE(table.contains(DataType::float16)); + CATCH_REQUIRE(table.contains(DataType::bfloat16)); CATCH_REQUIRE(table.contains(DataType::float32)); CATCH_REQUIRE(table[DataType::float16] == 5); + CATCH_REQUIRE(table[DataType::bfloat16] == 6); CATCH_REQUIRE(table[DataType::float32] == 10); } diff --git a/tests/utils/ivf_reference.cpp b/tests/utils/ivf_reference.cpp new file mode 100644 index 00000000..5bb932f8 --- /dev/null +++ b/tests/utils/ivf_reference.cpp @@ -0,0 +1,46 @@ +/* + * Copyright 2025 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. + */ + +// svstest +#include "tests/utils/test_dataset.h" + +// svsbenchmark +#include "svs-benchmark/ivf/test.h" + +// svs +#include "svs/core/distance.h" +#include "svs/lib/timing.h" +#include "svs/third-party/toml.h" + +// stl +#include +#include +#include + +namespace test_dataset::ivf { +namespace { +std::filesystem::path reference_path() { + return test_dataset::reference_directory() / "ivf_reference.toml"; +} +} // namespace + +const toml::table& parse_expected() { + // Make the expected results a magic static variable. + // This shaves off a bit of run time as we only need to parse the toml file once. + static toml::table expected = toml::parse_file(reference_path().native()); + return expected; +} +} // namespace test_dataset::ivf diff --git a/tests/utils/ivf_reference.h b/tests/utils/ivf_reference.h new file mode 100644 index 00000000..64674740 --- /dev/null +++ b/tests/utils/ivf_reference.h @@ -0,0 +1,81 @@ +/* + * Copyright 2025 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 + +// svsbenchmark +#include "svs-benchmark/datasets.h" +#include "svs-benchmark/ivf/test.h" + +// svs +#include "svs/core/distance.h" + +// stl +#include +#include +#include + +namespace test_dataset::ivf { + +// Implemented in CPP file. +const toml::table& parse_expected(); + +template +std::vector +expected_results(std::string_view key, svs::DistanceType distance, const T& dataset) { + const auto& table = parse_expected(); + auto v = svs::lib::load>( + svs::lib::node_view_at(table, key), std::nullopt + ); + auto output = std::vector(); + for (const auto& i : v) { + if ((i.distance_ == distance) && i.dataset_.match(dataset)) { + output.push_back(i); + } + } + return output; +} + +/// Return the only reference build for the requested parameters. +/// Throws ANNException if the number of matches is not equal to one. +template +svsbenchmark::ivf::ExpectedResult +expected_build_results(svs::DistanceType distance, const T& dataset) { + auto results = ivf::expected_results("ivf_test_build", distance, dataset); + if (results.size() != 1) { + throw ANNEXCEPTION("Got {} results when only one was expected!", results.size()); + } + // Make sure the only result has build parameters. + auto result = results[0]; + if (!result.build_parameters_.has_value()) { + throw ANNEXCEPTION("Expected build result does not have build parameters!"); + } + return result; +} + +/// Return the only reference search for the requested parameters. +/// Throws ANNException if the number of dataset is not equal to one. +template +svsbenchmark::ivf::ExpectedResult +expected_search_results(svs::DistanceType distance, const T& dataset) { + auto results = ivf::expected_results("ivf_test_search", distance, dataset); + if (results.size() != 1) { + throw ANNEXCEPTION("Got {} results when only one was expected!", results.size()); + } + return results[0]; +} + +} // namespace test_dataset::ivf diff --git a/tests/utils/test_dataset.cpp b/tests/utils/test_dataset.cpp index a89aa042..d1ff3aba 100644 --- a/tests/utils/test_dataset.cpp +++ b/tests/utils/test_dataset.cpp @@ -78,6 +78,10 @@ std::filesystem::path groundtruth_cosine_file() { return dataset_directory() / "groundtruth_cosine.ivecs"; } +std::filesystem::path clustering_directory() { + return dataset_directory() / "ivf_clustering"; +} + svs::data::SimpleData queries() { return svs::load_data(query_file()); } svs::data::SimpleData groundtruth_euclidean() { return svs::load_data(groundtruth_euclidean_file()); diff --git a/tests/utils/test_dataset.h b/tests/utils/test_dataset.h index 51bf913f..aa55f316 100644 --- a/tests/utils/test_dataset.h +++ b/tests/utils/test_dataset.h @@ -58,6 +58,9 @@ std::filesystem::path groundtruth_mip_file(); // Groundtruth of for the queries with respect to the dataset using cosine similarity. std::filesystem::path groundtruth_cosine_file(); +// The directory containing the IVF clustering. +std::filesystem::path clustering_directory(); + ///// Helper Functions svs::data::SimpleData queries(); svs::data::SimpleData groundtruth_euclidean(); diff --git a/tools/benchmark_inputs/ivf/test-generator.toml b/tools/benchmark_inputs/ivf/test-generator.toml new file mode 100644 index 00000000..963ee7b0 --- /dev/null +++ b/tools/benchmark_inputs/ivf/test-generator.toml @@ -0,0 +1,34 @@ +# Copyright 2025 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. + +[ivf_test_generator] +__schema__ = 'benchmark_ivf_test' +__version__ = 'v0.0.0' +data_f32 = 'data_f32.fvecs' +graph = '' +index_config = 'ivf_clustering/' +queries_f32 = 'queries_f32.fvecs' +queries_in_training_set = 100 + + [[ivf_test_generator.groundtruths]] + __schema__ = 'benchmark_distance_and_groundtruth' + __version__ = 'v0.0.0' + distance = 'L2' + path = 'groundtruth_euclidean.ivecs' + + [[ivf_test_generator.groundtruths]] + __schema__ = 'benchmark_distance_and_groundtruth' + __version__ = 'v0.0.0' + distance = 'MIP' + path = 'groundtruth_mip.ivecs' diff --git a/utils/CMakeLists.txt b/utils/CMakeLists.txt index d27bb33c..888b32f4 100644 --- a/utils/CMakeLists.txt +++ b/utils/CMakeLists.txt @@ -46,9 +46,17 @@ create_utility(upgrade_vamana_index_parameters upgrade_vamana_index_parameters.c create_utility(build_index build_index.cpp) create_utility(compute_recall compute_recall.cpp) create_utility(convert_data_to_float16 convert_data_to_float16.cpp) +create_utility(convert_data_to_bfloat16 convert_data_to_bfloat16.cpp) create_utility(search_index search_index.cpp) create_utility(consolidate characterization/consolidate.cpp) create_utility(logging logging.cpp) + +# IVF build and search +if (SVS_EXPERIMENTAL_ENABLE_IVF) + create_utility(build_ivf build_ivf.cpp) + create_utility(search_ivf search_ivf.cpp) +endif() + # create_utility(mutable characterization/mutable.cpp) # # Benchmark diff --git a/utils/build_ivf.cpp b/utils/build_ivf.cpp new file mode 100644 index 00000000..58dbe93f --- /dev/null +++ b/utils/build_ivf.cpp @@ -0,0 +1,130 @@ +/* + * Copyright 2025 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 "svs/orchestrators/ivf.h" +#include "svsmain.h" + +// stl +#include +#include +#include +#include +#include +#include +#include + +using namespace svs::index::ivf; + +template +void build_ivf_clustering( + Params params, + Data data, + Dist dist, + size_t n_threads, + const std::string& clustering_directory +) { + auto clustering = svs::IVF::build_clustering( + std::move(params), std::move(data), std::move(dist), n_threads + ); + svs::lib::save_to_disk(clustering, clustering_directory); +} + +const std::string HELP = + R"( +The required arguments are as follows: + +(1) Data Element Type (string). Options: (int8, uint8, float, float16, bfloat16) +(2) Path to vector dataset (.vecs format) (string). +(3) Number of clusters to be built +(4) Number of threads to use for index construction (integer). +(5) Should use hierarchical Kmeans? (0/1) +(6) Clustering directory for saving. +(7) Distance type (string - distance type) +)"; + +int svs_main(std::vector args) { + if (args.size() != 8) { + std::cout << "Expected 7 arguments. Instead, got " << args.size() - 1 << ". " + << "The required positional arguments are given below." << std::endl + << std::endl + << HELP << std::endl; + return 1; + } + + size_t i = 1; + const auto& data_type(args[i++]); + const auto& vecs_filename(args[i++]); + const size_t n_clusters = std::stoull(args[i++]); + const size_t n_threads = std::stoull(args[i++]); + const size_t is_hierarchical = std::stoull(args[i++]); + const std::string& clustering_directory(args[i++]); + const auto& distance_type = args[i++]; + + const size_t D = svs::Dynamic; + using Alloc = svs::HugepageAllocator; + + auto data = svs::VectorDataLoader(vecs_filename); + + auto dist_disp = [&](dist_type dist) { + auto params = svs::index::ivf::IVFBuildParameters(n_clusters, 10000, 10, false, .1); + if (is_hierarchical) { + params.is_hierarchical_ = true; + } + if (data_type == "float") { + build_ivf_clustering( + std::move(params), + std::move(data), + std::move(dist), + n_threads, + clustering_directory + ); + } else if (data_type == "float16") { + build_ivf_clustering( + std::move(params), + std::move(data), + std::move(dist), + n_threads, + clustering_directory + ); + } else if (data_type == "bfloat16") { + build_ivf_clustering( + std::move(params), + std::move(data), + std::move(dist), + n_threads, + clustering_directory + ); + } else { + throw ANNEXCEPTION("Unsupported data type: ", data_type, '.'); + } + }; + + if (distance_type == std::string("L2")) { + dist_disp(svs::distance::DistanceL2{}); + } else if (distance_type == std::string("MIP")) { + dist_disp(svs::distance::DistanceIP{}); + } else { + throw ANNEXCEPTION( + "Unsupported distance type. Valid values: L2/MIP. Received: ", + distance_type, + '!' + ); + } + + return 0; +} + +SVS_DEFINE_MAIN(); diff --git a/utils/convert_data_to_bfloat16.cpp b/utils/convert_data_to_bfloat16.cpp new file mode 100644 index 00000000..34e1d5ac --- /dev/null +++ b/utils/convert_data_to_bfloat16.cpp @@ -0,0 +1,67 @@ +/* + * Copyright 2025 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 +#include +#include + +#include "svs/core/io.h" +#include "svs/lib/bfloat16.h" +#include "svsmain.h" + +int svs_main(std::vector args) { + if (args.size() != 4) { + std::cout << "Specify the right parameters: input index, output index, " + "vector_type: 0 for SVS data, 1 for fvecs, 2 for fbin" + << std::endl; + return 1; + } + const std::string& filename_f32 = args[1]; + const std::string& filename_bf16 = args[2]; + const size_t file_type = std::stoull(args[3]); + + if (file_type == 0) { + std::cout << "Converting SVS data!" << std::endl; + auto reader = svs::io::v1::NativeFile{filename_f32}.reader(svs::lib::Type()); + auto writer = svs::io::NativeFile{filename_bf16}.writer( + svs::lib::Type(), reader.ndims() + ); + + for (auto i : reader) { + writer << i; + } + } else if (file_type == 1) { + std::cout << "Converting Vecs data!" << std::endl; + auto reader = svs::io::vecs::VecsReader{filename_f32}; + auto writer = + svs::io::vecs::VecsWriter{filename_bf16, reader.ndims()}; + for (auto i : reader) { + writer << i; + } + } else if (file_type == 2) { + std::cout << "Converting Bin data!" << std::endl; + auto reader = svs::io::binary::BinaryReader{filename_f32}; + auto writer = svs::io::binary::BinaryWriter{ + filename_bf16, reader.nvectors(), reader.ndims()}; + for (auto i : reader) { + writer << i; + } + } + + return 0; +} + +SVS_DEFINE_MAIN(); diff --git a/utils/search_ivf.cpp b/utils/search_ivf.cpp new file mode 100644 index 00000000..f1c42d4c --- /dev/null +++ b/utils/search_ivf.cpp @@ -0,0 +1,241 @@ +/* + * Copyright 2025 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 "svs/core/recall.h" +#include "svs/orchestrators/ivf.h" +#include "svsmain.h" + +// stl +#include +#include +#include +#include +#include +#include +#include + +using namespace svs::index::ivf; + +template +using SearchIndexFunction = std::function; + +auto std_dev(std::vector v) { + double sum = std::accumulate(std::begin(v), std::end(v), 0.0); + double m = sum / v.size(); + + double accum = 0.0; + std::for_each(std::begin(v), std::end(v), [&](const double d) { + accum += (d - m) * (d - m); + }); + + double stdev = sqrt(accum / (v.size() - 1)); + return stdev; +} + +template +auto batch_queries( + svs::data::SimpleData query_data, size_t num_batches, size_t batchsize +) { + std::vector> query_batch; + for (size_t batch = 0; batch < num_batches; ++batch) { + auto this_batch = svs::threads::UnitRange{ + batch * batchsize, std::min((batch + 1) * batchsize, query_data.size())}; + query_batch.push_back( + svs::data::SimpleData(this_batch.size(), query_data.dimensions()) + ); + for (size_t i = 0; i < this_batch.size(); i++) { + query_batch[batch].set_datum(i, query_data.get_datum(this_batch[i])); + } + } + return query_batch; +} + +template +void search_index( + const std::string& query_filename, + const std::string& gt_filename, + const size_t n_probes, + const size_t n_neighbors, + const size_t batch_size, + const size_t n_threads, + const size_t n_inner_threads, + const std::filesystem::path& clustering_path, + const std::filesystem::path& data_path, + const Distance dist_type, + const int n_reps +) { + const size_t dim = svs::Dynamic; + + auto data = svs::VectorDataLoader>(data_path); + + auto ivf_index = svs::IVF::assemble_from_file( + clustering_path, std::move(data), dist_type, n_threads, n_inner_threads + ); + + const auto query_data = svs::load_data(query_filename); + const auto groundtruth = svs::load_data(gt_filename); + + ivf_index.set_search_parameters(IVFSearchParameters(n_probes, 1.0)); + + size_t batchsize = query_data.size(); + if (batch_size != 0) { + batchsize = batch_size; + } + + auto query_results = + svs::Matrix{svs::make_dims(query_data.size(), n_neighbors)}; + auto num_batches = svs::lib::div_round_up(query_data.size(), batchsize); + auto query_batch = batch_queries(query_data, num_batches, batchsize); + + auto tic = svs::lib::now(); + for (size_t batch = 0; batch < num_batches; ++batch) { + auto query_result = ivf_index.search(query_batch[batch], n_neighbors); + for (size_t i = 0; i < query_result.n_queries(); i++) { + for (size_t j = 0; j < n_neighbors; j++) { + query_results.at(batch * batchsize + i, j) = query_result.index(i, j); + } + } + } + auto search_time = svs::lib::time_difference(tic); + + std::vector qps; + for (int i = 0; i < n_reps; i++) { + tic = svs::lib::now(); + for (size_t batch = 0; batch < num_batches; ++batch) { + ivf_index.search(query_batch[batch], n_neighbors); + } + search_time = svs::lib::time_difference(tic); + qps.push_back(query_data.size() / search_time); + } + + fmt::print("Raw QPS: {:7.3f} \n", fmt::join(qps, ", ")); + fmt::print( + "Batch Size: {}, Recall: {:.4f}, QPS (Avg: {:7.3f}, Max: {:7.3f}, StdDev: {:7.3f} " + ") " + "\n", + batchsize, + svs::k_recall_at_n(groundtruth, query_results, 10, 10), + std::reduce(qps.begin(), qps.end()) / qps.size(), + *std::max_element(qps.begin(), qps.end()), + std_dev(qps) + ); +} + +constexpr std::string_view HELP = + R"( +The required arguments are as follows: +(1) Query Element Type (string). Options: (int8, uint8, float) +(2) Data Element Type (string). Options: (int8, uint8, float, float16, bfloat16) +(3) Query File Path (string). Supported extentions: (.vecs, .bin) +(4) Groundtruth File Path (string). Supported extentions: (.vecs, .bin) +(5) n_probes (number of clusters to search) (integer) +(6) Number of neighbors to recall (integer) +(7) Batch size (integer) +(8) Number of threads (integer) +(9) Number of intra-query threads (integer) +(10) Clustering directory (string) +(11) Data directory (string) +(12) Number of repetitions to be run for benchmarking purposes (integer) +(13) Distance type (string - distance type) +)"; + +int svs_main(std::vector&& args) { + if (args.size() != 14) { + std::cout << "Expected 13 arguments. Instead, got " << args.size() - 1 << ". " + << "The required positional arguments are given below." << std::endl + << std::endl + << HELP << std::endl; + return 1; + } + + size_t i = 1; + const auto& query_data_type = args[i++]; + const auto& db_data_type = args[i++]; + const auto& query_filename = args[i++]; + const auto& gt_filename = args[i++]; + const size_t n_probes = std::stoull(args[i++]); + const size_t n_neighbors = std::stoull(args[i++]); + const size_t batch_size = std::stoull(args[i++]); + const size_t n_threads = std::stoull(args[i++]); + const size_t n_inner_threads = std::stoull(args[i++]); + const auto& clustering_path = args[i++]; + const auto& data_path = args[i++]; + const size_t nreps = std::stoull(args[i++]); + const auto& distance_type = args[i++]; + + auto dist_disp = [&](dist_type dist) { + using KeyType = std::pair; + using ValueType = SearchIndexFunction; + const auto dispatcher = std::map{ + {{"float", "float16"}, search_index}, + {{"float", "bfloat16"}, search_index}, + {{"float", "float"}, search_index}}; + + auto it = dispatcher.find({query_data_type, db_data_type}); + if (it == dispatcher.end()) { + throw ANNEXCEPTION( + "Unsupported Query and Data type pair: (", + query_data_type, + ", ", + db_data_type, + ")!" + ); + } + + // Unpack and call + const auto& f = it->second; + f(query_filename, + gt_filename, + n_probes, + n_neighbors, + batch_size, + n_threads, + n_inner_threads, + clustering_path, + data_path, + dist, + nreps); + }; + + if (distance_type == std::string("L2")) { + dist_disp(svs::distance::DistanceL2{}); + } else if (distance_type == std::string("MIP")) { + dist_disp(svs::distance::DistanceIP{}); + } else { + throw ANNEXCEPTION( + "Unsupported distance type. Valid values: L2/MIP. Received: ", + distance_type, + '!' + ); + } + + return 0; +} + +// Include the helper main function. +SVS_DEFINE_MAIN();