diff --git a/c/src/neighbors/all_neighbors.cpp b/c/src/neighbors/all_neighbors.cpp index 22e0929e2b..fa110c4662 100644 --- a/c/src/neighbors/all_neighbors.cpp +++ b/c/src/neighbors/all_neighbors.cpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -16,7 +16,9 @@ #include #include #include -#include +#include +#include +#include #include "../core/exceptions.hpp" #include "../core/interop.hpp" diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index bed7e74084..081179ca46 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -18,7 +18,6 @@ #include #include #include -#include #include "../core/exceptions.hpp" #include "../core/interop.hpp" diff --git a/cpp/include/cuvs/neighbors/all_neighbors.hpp b/cpp/include/cuvs/neighbors/all_neighbors.hpp index 41662d9e31..9f20a28134 100644 --- a/cpp/include/cuvs/neighbors/all_neighbors.hpp +++ b/cpp/include/cuvs/neighbors/all_neighbors.hpp @@ -1,11 +1,14 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ #pragma once -#include +#include +#include +#include + #include namespace cuvs::neighbors::all_neighbors { diff --git a/cpp/include/cuvs/neighbors/brute_force.hpp b/cpp/include/cuvs/neighbors/brute_force.hpp index 724ffcb5a9..b57c2e7a0f 100644 --- a/cpp/include/cuvs/neighbors/brute_force.hpp +++ b/cpp/include/cuvs/neighbors/brute_force.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2024, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -926,3 +926,11 @@ void deserialize(raft::resources const& handle, */ } // namespace cuvs::neighbors::brute_force + +/** Specialized parameters utilizing brute force to build knn graph */ +namespace cuvs::neighbors::graph_build_params { +struct brute_force_params { + cuvs::neighbors::brute_force::index_params build_params; + cuvs::neighbors::brute_force::search_params search_params; +}; +} // namespace cuvs::neighbors::graph_build_params diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 6fd734064c..811861af90 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -8,7 +8,6 @@ #include "common.hpp" #include #include -#include #include #include #include @@ -31,6 +30,73 @@ #include #include +namespace cuvs::neighbors::graph_build_params { +using iterative_search_params = cuvs::neighbors::search_params; + +/** Specialized parameters for ACE (Augmented Core Extraction) graph build */ +struct ace_params { + /** + * Number of partitions for ACE (Augmented Core Extraction) partitioned build. + * + * When set to 0 (default), the number of partitions is automatically derived + * based on available host and GPU memory to maximize partition size while + * ensuring the build fits in memory. + * + * Small values might improve recall but potentially degrade performance and + * increase memory usage. Partitions should not be too small to prevent issues + * in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * + * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the + * partition sizes (up to 3x in our tests). + * + * If the specified number of partitions results in partitions that exceed + * available memory, the value will be automatically increased to fit memory + * constraints and a warning will be issued. + */ + size_t npartitions = 0; + /** + * The index quality for the ACE build. + * + * Bigger values increase the index quality. At some point, increasing this will no longer improve + * the quality. + */ + size_t ef_construction = 120; + /** + * Directory to store ACE build artifacts (e.g., KNN graph, optimized graph). + * + * Used when `use_disk` is true or when the graph does not fit in host and GPU + * memory. This should be the fastest disk in the system and hold enough space + * for twice the dataset, final graph, and label mapping. + */ + std::string build_dir = "/tmp/ace_build"; + /** + * Whether to use disk-based storage for ACE build. + * + * When true, enables disk-based operations for memory-efficient graph construction. + */ + bool use_disk = false; + + /** + * Maximum host memory to use for ACE build in GiB. + * + * When set to 0 (default), uses available host memory. + * When set to a positive value, limits host memory usage to the specified amount. + * Useful for testing or when running alongside other memory-intensive processes. + */ + double max_host_memory_gb = 0; + /** + * Maximum GPU memory to use for ACE build in GiB. + * + * When set to 0 (default), uses available GPU memory. + * When set to a positive value, limits GPU memory usage to the specified amount. + * Useful for testing or when running alongside other memory-intensive processes. + */ + double max_gpu_memory_gb = 0; + + ace_params() = default; +}; + +} // namespace cuvs::neighbors::graph_build_params + namespace cuvs::neighbors::cagra { // For re-exporting into cagra namespace namespace graph_build_params = cuvs::neighbors::graph_build_params; diff --git a/cpp/include/cuvs/neighbors/common.hpp b/cpp/include/cuvs/neighbors/common.hpp index 39967999ed..a118f8180e 100644 --- a/cpp/include/cuvs/neighbors/common.hpp +++ b/cpp/include/cuvs/neighbors/common.hpp @@ -35,6 +35,9 @@ namespace cuvs::neighbors { * @{ */ +/* Graph build algo used in cagra and all_neighbors */ +enum GRAPH_BUILD_ALGO { BRUTE_FORCE = 0, IVF_PQ = 1, NN_DESCENT = 2, ACE = 3 }; + /** Parameters for VPQ compression. */ struct vpq_params { /** diff --git a/cpp/include/cuvs/neighbors/graph_build_types.hpp b/cpp/include/cuvs/neighbors/graph_build_types.hpp deleted file mode 100644 index 5faa08c1d1..0000000000 --- a/cpp/include/cuvs/neighbors/graph_build_types.hpp +++ /dev/null @@ -1,163 +0,0 @@ -/* - * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. - * SPDX-License-Identifier: Apache-2.0 - */ - -#pragma once - -#include -#include -#include - -namespace cuvs::neighbors { - -/** - * @defgroup neighbors_build_algo Graph build algorithm types - * @{ - */ - -enum GRAPH_BUILD_ALGO { BRUTE_FORCE = 0, IVF_PQ = 1, NN_DESCENT = 2, ACE = 3 }; - -namespace graph_build_params { - -/** Specialized parameters utilizing IVF-PQ to build knn graph */ -struct ivf_pq_params { - cuvs::neighbors::ivf_pq::index_params build_params; - cuvs::neighbors::ivf_pq::search_params search_params; - float refinement_rate = 1.0; - - ivf_pq_params() = default; - - /** - * Set default parameters based on shape of the input dataset. - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * raft::resources res; - * // create index_params for a [N. D] dataset - * auto dataset = raft::make_device_matrix(res, N, D); - * auto pq_params = - * graph_build_params::ivf_pq_params(dataset.extents()); - * // modify/update index_params as needed - * pq_params.kmeans_trainset_fraction = 0.1; - * @endcode - */ - ivf_pq_params(raft::matrix_extent dataset_extents, - cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) - { - build_params = ivf_pq::index_params::from_dataset(dataset_extents, metric); - auto n_rows = dataset_extents.extent(0); - auto n_features = dataset_extents.extent(1); - if (n_features <= 32) { - build_params.pq_dim = 16; - build_params.pq_bits = 8; - } else { - build_params.pq_bits = 4; - if (n_features <= 64) { - build_params.pq_dim = 32; - } else if (n_features <= 128) { - build_params.pq_dim = 64; - } else if (n_features <= 192) { - build_params.pq_dim = 96; - } else { - build_params.pq_dim = raft::round_up_safe(n_features / 2, 128); - } - } - - build_params.n_lists = std::max(1, n_rows / 2000); - build_params.kmeans_n_iters = 10; - - const double kMinPointsPerCluster = 32; - const double min_kmeans_trainset_points = kMinPointsPerCluster * build_params.n_lists; - const double max_kmeans_trainset_fraction = 1.0; - const double min_kmeans_trainset_fraction = - std::min(max_kmeans_trainset_fraction, min_kmeans_trainset_points / n_rows); - build_params.kmeans_trainset_fraction = std::clamp( - 1.0 / std::sqrt(n_rows * 1e-5), min_kmeans_trainset_fraction, max_kmeans_trainset_fraction); - build_params.codebook_kind = ivf_pq::codebook_gen::PER_SUBSPACE; - - search_params = cuvs::neighbors::ivf_pq::search_params{}; - search_params.n_probes = std::round(std::sqrt(build_params.n_lists) / 20 + 4); - search_params.lut_dtype = CUDA_R_16F; - search_params.internal_distance_dtype = CUDA_R_16F; - search_params.coarse_search_dtype = CUDA_R_16F; - search_params.max_internal_batch_size = 128 * 1024; - - refinement_rate = 1; - } -}; - -using nn_descent_params = cuvs::neighbors::nn_descent::index_params; - -struct brute_force_params { - cuvs::neighbors::brute_force::index_params build_params; - cuvs::neighbors::brute_force::search_params search_params; -}; - -/** Specialized parameters for ACE (Augmented Core Extraction) graph build */ -struct ace_params { - /** - * Number of partitions for ACE (Augmented Core Extraction) partitioned build. - * - * When set to 0 (default), the number of partitions is automatically derived - * based on available host and GPU memory to maximize partition size while - * ensuring the build fits in memory. - * - * Small values might improve recall but potentially degrade performance and - * increase memory usage. Partitions should not be too small to prevent issues - * in KNN graph construction. The partition size is on average 2 * (n_rows / npartitions) * dim * - * sizeof(T). 2 is because of the core and augmented vectors. Please account for imbalance in the - * partition sizes (up to 3x in our tests). - * - * If the specified number of partitions results in partitions that exceed - * available memory, the value will be automatically increased to fit memory - * constraints and a warning will be issued. - */ - size_t npartitions = 0; - /** - * The index quality for the ACE build. - * - * Bigger values increase the index quality. At some point, increasing this will no longer improve - * the quality. - */ - size_t ef_construction = 120; - /** - * Directory to store ACE build artifacts (e.g., KNN graph, optimized graph). - * - * Used when `use_disk` is true or when the graph does not fit in host and GPU - * memory. This should be the fastest disk in the system and hold enough space - * for twice the dataset, final graph, and label mapping. - */ - std::string build_dir = "/tmp/ace_build"; - /** - * Whether to use disk-based storage for ACE build. - * - * When true, enables disk-based operations for memory-efficient graph construction. - */ - bool use_disk = false; - /** - * Maximum host memory to use for ACE build in GiB. - * - * When set to 0 (default), uses available host memory. - * When set to a positive value, limits host memory usage to the specified amount. - * Useful for testing or when running alongside other memory-intensive processes. - */ - double max_host_memory_gb = 0; - /** - * Maximum GPU memory to use for ACE build in GiB. - * - * When set to 0 (default), uses available GPU memory. - * When set to a positive value, limits GPU memory usage to the specified amount. - * Useful for testing or when running alongside other memory-intensive processes. - */ - double max_gpu_memory_gb = 0; - - ace_params() = default; -}; - -// **** Experimental **** -using iterative_search_params = cuvs::neighbors::search_params; -} // namespace graph_build_params - -/** @} */ // end group neighbors_build_algo -} // namespace cuvs::neighbors diff --git a/cpp/include/cuvs/neighbors/hnsw.hpp b/cpp/include/cuvs/neighbors/hnsw.hpp index a1be1c58cc..7ee91f18ba 100644 --- a/cpp/include/cuvs/neighbors/hnsw.hpp +++ b/cpp/include/cuvs/neighbors/hnsw.hpp @@ -10,7 +10,7 @@ #include "common.hpp" #include -#include +#include #include "cagra.hpp" #include diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp index f6eef54414..710a08cd0c 100644 --- a/cpp/include/cuvs/neighbors/ivf_pq.hpp +++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp @@ -3358,3 +3358,72 @@ void resize_list(raft::resources const& res, } // namespace helpers } // namespace cuvs::neighbors::ivf_pq + +namespace cuvs::neighbors::graph_build_params { +/** Specialized parameters utilizing IVF-PQ to build knn graph */ +struct ivf_pq_params { + cuvs::neighbors::ivf_pq::index_params build_params; + cuvs::neighbors::ivf_pq::search_params search_params; + float refinement_rate = 1.0; + + ivf_pq_params() = default; + + /** + * Set default parameters based on shape of the input dataset. + * Usage example: + * @code{.cpp} + * using namespace cuvs::neighbors; + * raft::resources res; + * // create index_params for a [N. D] dataset + * auto dataset = raft::make_device_matrix(res, N, D); + * auto pq_params = + * graph_build_params::ivf_pq_params(dataset.extents()); + * // modify/update index_params as needed + * pq_params.kmeans_trainset_fraction = 0.1; + * @endcode + */ + ivf_pq_params(raft::matrix_extent dataset_extents, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) + { + build_params = ivf_pq::index_params::from_dataset(dataset_extents, metric); + auto n_rows = dataset_extents.extent(0); + auto n_features = dataset_extents.extent(1); + if (n_features <= 32) { + build_params.pq_dim = 16; + build_params.pq_bits = 8; + } else { + build_params.pq_bits = 4; + if (n_features <= 64) { + build_params.pq_dim = 32; + } else if (n_features <= 128) { + build_params.pq_dim = 64; + } else if (n_features <= 192) { + build_params.pq_dim = 96; + } else { + build_params.pq_dim = raft::round_up_safe(n_features / 2, 128); + } + } + + build_params.n_lists = std::max(1, n_rows / 2000); + build_params.kmeans_n_iters = 10; + + const double kMinPointsPerCluster = 32; + const double min_kmeans_trainset_points = kMinPointsPerCluster * build_params.n_lists; + const double max_kmeans_trainset_fraction = 1.0; + const double min_kmeans_trainset_fraction = + std::min(max_kmeans_trainset_fraction, min_kmeans_trainset_points / n_rows); + build_params.kmeans_trainset_fraction = std::clamp( + 1.0 / std::sqrt(n_rows * 1e-5), min_kmeans_trainset_fraction, max_kmeans_trainset_fraction); + build_params.codebook_kind = ivf_pq::codebook_gen::PER_SUBSPACE; + + search_params = cuvs::neighbors::ivf_pq::search_params{}; + search_params.n_probes = std::round(std::sqrt(build_params.n_lists) / 20 + 4); + search_params.lut_dtype = CUDA_R_16F; + search_params.internal_distance_dtype = CUDA_R_16F; + search_params.coarse_search_dtype = CUDA_R_16F; + search_params.max_internal_batch_size = 128 * 1024; + + refinement_rate = 1; + } +}; +} // namespace cuvs::neighbors::graph_build_params diff --git a/cpp/include/cuvs/neighbors/nn_descent.hpp b/cpp/include/cuvs/neighbors/nn_descent.hpp index c2f6121303..44fbaed592 100644 --- a/cpp/include/cuvs/neighbors/nn_descent.hpp +++ b/cpp/include/cuvs/neighbors/nn_descent.hpp @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -525,3 +525,7 @@ bool has_enough_device_memory(raft::resources const& res, size_t idx_size = 4); } // namespace cuvs::neighbors::nn_descent + +namespace cuvs::neighbors::graph_build_params { +using nn_descent_params = cuvs::neighbors::nn_descent::index_params; +} diff --git a/cpp/src/neighbors/all_neighbors/all_neighbors.cuh b/cpp/src/neighbors/all_neighbors/all_neighbors.cuh index e89440a222..b6831d1f3c 100644 --- a/cpp/src/neighbors/all_neighbors/all_neighbors.cuh +++ b/cpp/src/neighbors/all_neighbors/all_neighbors.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2025, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION. * SPDX-License-Identifier: Apache-2.0 */ @@ -7,7 +7,6 @@ #include "../detail/reachability.cuh" #include "all_neighbors_batched.cuh" #include -#include #include #include