From 916b1e71668a143d925780a5a7fb5ec35dd6cc65 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 16:46:26 +0900 Subject: [PATCH 1/8] Add brute-force kNN graph construction for CAGRA --- c/src/neighbors/cagra.cpp | 6 +- cpp/include/cuvs/neighbors/cagra.hpp | 55 ++++++++++++++- cpp/src/neighbors/cagra.cuh | 10 +++ cpp/src/neighbors/cagra_build_inst.cu.in | 16 +++++ .../neighbors/detail/cagra/cagra_build.cuh | 65 +++++++++++++++++- cpp/tests/neighbors/ann_cagra.cuh | 68 ++++++++++++++++++- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 3 +- 7 files changed, 213 insertions(+), 10 deletions(-) diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 97cc9d8bac..3ff1e8d623 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -762,11 +762,7 @@ static void update_dataset(raft::resources* res_ptr, } static void _set_graph_build_params( - std::variant& out_params, + cuvs::neighbors::cagra::graph_build_params_t& out_params, cuvsCagraIndexParams& params, cuvsCagraGraphBuildAlgo algo, int64_t n_rows, diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 418db9d88f..6fd2c0cb0a 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -8,6 +8,7 @@ #include #include #include +#include #include #include #include @@ -254,7 +255,8 @@ using graph_build_params_t = std::variant; + graph_build_params::iterative_search_params, + graph_build_params::brute_force_params>; /** * @brief A strategy for selecting the graph build parameters based on similar HNSW index @@ -317,6 +319,9 @@ struct index_params : cuvs::neighbors::index_params { * // 4. Choose iterative graph building using CAGRA's search() and optimize() [Experimental] * params.graph_build_params = * cagra::graph_build_params::iterative_search_params(); + * + * // 5. Choose exact brute-force kNN construction (float and half datasets) + * params.graph_build_params = cagra::graph_build_params::brute_force_params{}; * @endcode */ graph_build_params_t graph_build_params; @@ -4374,6 +4379,54 @@ void distribute(const raft::resources& clique, const std::string& filename, cuvs::neighbors::mg_index, T, IdxT>* index); +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major host dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::host_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major device dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::device_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major host dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::host_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + +/** + * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. + * @param res RAFT resources. + * @param dataset Row-major device dataset of float or half values. + * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. + * @param build_params Brute-force index and search parameters. + */ +void build_knn_graph(raft::resources const& res, + raft::device_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params build_params); + /** * @brief Build a kNN graph using IVF-PQ. * diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 0e1454d987..b82b259fd4 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -191,6 +191,16 @@ void build_knn_graph( detail::build_knn_graph(res, dataset, knn_graph, build_params); } +template +void build_knn_graph( + raft::resources const& res, + raft::mdspan, raft::row_major, Accessor> dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + detail::build_knn_graph(res, dataset, knn_graph, params); +} + /** * @brief Sort a KNN graph index. * Preprocessing step for `cagra::optimize`: If a KNN graph is not built using diff --git a/cpp/src/neighbors/cagra_build_inst.cu.in b/cpp/src/neighbors/cagra_build_inst.cu.in index 2a3e2b43c1..80e0e33fa3 100644 --- a/cpp/src/neighbors/cagra_build_inst.cu.in +++ b/cpp/src/neighbors/cagra_build_inst.cu.in @@ -43,6 +43,22 @@ void build_knn_graph(raft::resources const& handle, cuvs::neighbors::cagra::build_knn_graph(handle, dataset, knn_graph, params); } +void build_knn_graph(raft::resources const& res, + raft::host_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + detail::build_knn_graph(res, dataset, knn_graph, params); +} + +void build_knn_graph(raft::resources const& res, + raft::device_matrix_view dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + detail::build_knn_graph(res, dataset, knn_graph, params); +} + #define CUVS_DEFINE_CAGRA_BUILD_OVERLOAD(DatasetViewT, ...) \ auto build(raft::resources const& res, \ const cuvs::neighbors::cagra::index_params& params, \ diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 9ccd177b53..4eb400543d 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -12,6 +12,8 @@ #include "graph_core.cuh" #include +#include + #include #include #include @@ -82,6 +84,59 @@ void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_ } } +// Exact kNN construction. Request one extra +// neighbor and remove self by ID, including when duplicate vectors tie with self. +template +void build_knn_graph(raft::resources const& res, + raft::mdspan, Layout, Accessor> dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + auto const partition_size = static_cast(dataset.extent(0)); + auto const intermediate_graph_degree = static_cast(knn_graph.extent(1)); + RAFT_EXPECTS(knn_graph.extent(0) == dataset.extent(0), + "CAGRA: kNN graph and dataset must have the same number of rows"); + RAFT_EXPECTS(intermediate_graph_degree > 0 && partition_size > intermediate_graph_degree, + "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); + if constexpr (std::is_same_v || std::is_same_v) { + auto device_queries = + raft::make_device_matrix(res, partition_size, dataset.extent(1)); + raft::copy(res, device_queries.view(), dataset); + auto brute_force_index = cuvs::neighbors::brute_force::build( + res, params.build_params, raft::make_const_mdspan(device_queries.view())); + auto device_neighbors = raft::make_device_matrix( + res, partition_size, intermediate_graph_degree + 1); + auto device_distances = + raft::make_device_matrix(res, partition_size, intermediate_graph_degree + 1); + cuvs::neighbors::brute_force::search(res, + params.search_params, + brute_force_index, + raft::make_const_mdspan(device_queries.view()), + device_neighbors.view(), + device_distances.view()); + + auto host_neighbors = + raft::make_host_matrix(partition_size, intermediate_graph_degree + 1); + raft::copy(res, host_neighbors.view(), device_neighbors.view()); + raft::resource::sync_stream(res); + for (size_t row = 0; row < partition_size; ++row) { + size_t output_rank = 0; + for (size_t rank = 0; rank < intermediate_graph_degree + 1; ++rank) { + auto const neighbor = host_neighbors(row, rank); + if (neighbor == static_cast(row)) { continue; } + RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, + "CAGRA: brute-force kNN returned an invalid neighbor"); + knn_graph(row, output_rank++) = static_cast(neighbor); + if (output_rank == intermediate_graph_degree) { break; } + } + RAFT_EXPECTS(output_rank == intermediate_graph_degree, + "CAGRA: brute-force kNN did not return enough non-self neighbors"); + } + } else { + RAFT_FAIL("CAGRA: brute-force kNN build supports only float and half datasets"); + } +} + template void write_to_graph(raft::host_matrix_view knn_graph, raft::host_matrix_view neighbors_host_view, @@ -972,11 +1027,12 @@ inline void validate_cagra_knn_graph_build_constraints(index_params const& param std::holds_alternative( knn_build_params) || std::holds_alternative(knn_build_params), - "IVF_PQ for CAGRA graph build does not support BitwiseHamming as a metric. Please " + "The selected CAGRA graph builder does not support BitwiseHamming as a metric. Please " "use nn-descent or the iterative CAGRA search build."); RAFT_EXPECTS( params.metric != cuvs::distance::DistanceType::CosineExpanded || std::holds_alternative(knn_build_params) || + std::holds_alternative(knn_build_params) || std::holds_alternative(knn_build_params), "CosineExpanded distance is not supported for iterative CAGRA graph build."); @@ -987,7 +1043,8 @@ inline void validate_cagra_knn_graph_build_constraints(index_params const& param } /** - * Iterative / IVF-PQ / NN-descent KNN graph construction and `optimize` → final host CAGRA graph. + * Iterative / IVF-PQ / NN-descent / brute-force KNN graph construction and `optimize` → final host + * CAGRA graph. * * @param knn_graph_dataset mdspan passed to IVF-PQ / NN-descent `build_knn_graph` (any stride). */ @@ -1016,6 +1073,10 @@ auto build_cagra_host_graph_from_knn_params(raft::resources const& res, ivf_pq_params.build_params.metric = params.metric; } build_knn_graph(res, knn_graph_dataset, knn_graph->view(), ivf_pq_params); + } else if (std::holds_alternative(knn_build_params)) { + auto brute_force_params = std::get(knn_build_params); + brute_force_params.build_params.metric = params.metric; + build_knn_graph(res, knn_graph_dataset, knn_graph->view(), brute_force_params); } else { auto nn_descent_params = std::get(knn_build_params); diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 716dfd3cbe..70ccfa0009 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -283,6 +283,8 @@ enum class graph_build_algo { NN_DESCENT, /* Experimental, iteratively execute CAGRA's search() and optimize() */ ITERATIVE_CAGRA_SEARCH, + /* Use exact brute-force kNN graph construction */ + BRUTE_FORCE, /* Choose default automatically */ AUTO }; @@ -336,7 +338,8 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p) {search_algo::MULTI_KERNEL, "multi_kernel"}, // {search_algo::AUTO, "auto"} // }; - std::vector build_algo = {"IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "AUTO"}; + std::vector build_algo = { + "IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "AUTO", "BRUTE_FORCE"}; std::vector merge_strategy = {"PHYSICAL", "LOGICAL"}; os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim << ", degree=" << p.graph_degree << ", k=" << p.k << ", " << algo_name[p.algo] @@ -366,6 +369,11 @@ class AnnCagraTest : public ::testing::TestWithParam { template void testCagra() { + if constexpr (!std::is_same_v && !std::is_same_v) { + if (ps.build_algo == graph_build_algo::BRUTE_FORCE) { + GTEST_SKIP() << "Brute-force kNN graph construction supports only float and half"; + } + } // IVF_PQ graph build does not support BitwiseHamming if (ps.metric == cuvs::distance::DistanceType::BitwiseHamming && ((!std::is_same_v) || (ps.build_algo == graph_build_algo::IVF_PQ))) @@ -443,6 +451,9 @@ class AnnCagraTest : public ::testing::TestWithParam { index_params.graph_build_params = graph_build_params::iterative_search_params(); break; } + case graph_build_algo::BRUTE_FORCE: + index_params.graph_build_params = graph_build_params::brute_force_params{}; + break; case graph_build_algo::AUTO: // do nothing break; @@ -588,6 +599,11 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam { protected: void testCagra() { + if constexpr (!std::is_same_v && !std::is_same_v) { + if (ps.build_algo == graph_build_algo::BRUTE_FORCE) { + GTEST_SKIP() << "Brute-force kNN graph construction supports only float and half"; + } + } if (ps.metric == cuvs::distance::DistanceType::L1 && ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH) GTEST_SKIP(); @@ -666,6 +682,9 @@ class AnnCagraAddNodesTest : public ::testing::TestWithParam { index_params.graph_build_params = graph_build_params::iterative_search_params(); break; } + case graph_build_algo::BRUTE_FORCE: + index_params.graph_build_params = graph_build_params::brute_force_params{}; + break; case graph_build_algo::AUTO: // do nothing break; @@ -792,6 +811,11 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { protected: void testCagra() { + if constexpr (!std::is_same_v && !std::is_same_v) { + if (ps.build_algo == graph_build_algo::BRUTE_FORCE) { + GTEST_SKIP() << "Brute-force kNN graph construction supports only float and half"; + } + } if (ps.metric == cuvs::distance::DistanceType::L1 && ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH) GTEST_SKIP(); @@ -876,6 +900,9 @@ class AnnCagraFilterTest : public ::testing::TestWithParam { index_params.graph_build_params = graph_build_params::iterative_search_params(); break; } + case graph_build_algo::BRUTE_FORCE: + index_params.graph_build_params = graph_build_params::brute_force_params{}; + break; case graph_build_algo::AUTO: // do nothing break; @@ -1174,6 +1201,11 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam void testCagra() { + if constexpr (!std::is_same_v && !std::is_same_v) { + if (ps.build_algo == graph_build_algo::BRUTE_FORCE) { + GTEST_SKIP() << "Brute-force kNN graph construction supports only float and half"; + } + } if (ps.metric == cuvs::distance::DistanceType::L1 && ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH) GTEST_SKIP(); @@ -1283,6 +1315,9 @@ class AnnCagraIndexFilteredMergeTest : public ::testing::TestWithParam { template void testCagra() { + if constexpr (!std::is_same_v && !std::is_same_v) { + if (ps.build_algo == graph_build_algo::BRUTE_FORCE) { + GTEST_SKIP() << "Brute-force kNN graph construction supports only float and half"; + } + } if (ps.metric == cuvs::distance::DistanceType::L1 && ps.build_algo != graph_build_algo::ITERATIVE_CAGRA_SEARCH) GTEST_SKIP(); @@ -1516,6 +1556,9 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { index_params.graph_build_params = graph_build_params::iterative_search_params(); break; } + case graph_build_algo::BRUTE_FORCE: + index_params.graph_build_params = graph_build_params::brute_force_params{}; + break; case graph_build_algo::AUTO: // do nothing break; @@ -1955,6 +1998,26 @@ inline std::vector generate_inputs() {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); + // Add two small exact-builder cases instead of expanding the parameter sweep. + inputs2 = raft::util::itertools::product( + {16}, // queries + {256}, // rows + {16}, // dimensions + {8}, // k + {32}, // graph degree + {graph_build_algo::BRUTE_FORCE}, + {search_algo::SINGLE_CTA}, + {0}, // max queries + {0}, // team size + {64}, // itopk size + {1}, // search width + {cuvs::distance::DistanceType::L2Expanded, cuvs::distance::DistanceType::InnerProduct}, + {false}, // device dataset + {true}, // include serialized dataset + {false}, // source indices + {0.99}); + inputs.insert(inputs.end(), inputs2.begin(), inputs2.end()); + return inputs; } @@ -2182,6 +2245,9 @@ class AnnCagraMultiPartitionTest : public ::testing::TestWithParam; +graph_build_params::iterative_search_params, +graph_build_params::brute_force_params>; ``` From 956a51c21b1da6aea59be50b0ed872af29c891a5 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 18:21:58 +0900 Subject: [PATCH 2/8] Fix CAGRA test output bug --- cpp/tests/neighbors/ann_cagra.cuh | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 70ccfa0009..473bd956a7 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -339,7 +339,7 @@ inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p) {search_algo::AUTO, "auto"} // }; std::vector build_algo = { - "IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "AUTO", "BRUTE_FORCE"}; + "IVF_PQ", "NN_DESCENT", "ITERATIVE_CAGRA_SEARCH", "BRUTE_FORCE", "AUTO"}; std::vector merge_strategy = {"PHYSICAL", "LOGICAL"}; os << "{n_queries=" << p.n_queries << ", dataset shape=" << p.n_rows << "x" << p.dim << ", degree=" << p.graph_degree << ", k=" << p.k << ", " << algo_name[p.algo] From 45d79477f5c98a0681a77b7e27d0630692422cca Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 18:36:34 +0900 Subject: [PATCH 3/8] Dataset split for large dataset --- .../neighbors/detail/cagra/cagra_build.cuh | 56 +--------- .../detail/cagra/cagra_build_brute_force.cuh | 103 ++++++++++++++++++ 2 files changed, 104 insertions(+), 55 deletions(-) create mode 100644 cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh index 4eb400543d..95218db467 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -8,12 +8,11 @@ #include "../../../util/kvikio_io.hpp" #include "../../ivf_pq/ivf_pq_fp16_overflow.cuh" #include "cagra_ace.cuh" +#include "cagra_build_brute_force.cuh" #include "cagra_search.cuh" #include "graph_core.cuh" #include -#include - #include #include #include @@ -84,59 +83,6 @@ void check_graph_degree(size_t& intermediate_degree, size_t& graph_degree, size_ } } -// Exact kNN construction. Request one extra -// neighbor and remove self by ID, including when duplicate vectors tie with self. -template -void build_knn_graph(raft::resources const& res, - raft::mdspan, Layout, Accessor> dataset, - raft::host_matrix_view knn_graph, - graph_build_params::brute_force_params params) -{ - auto const partition_size = static_cast(dataset.extent(0)); - auto const intermediate_graph_degree = static_cast(knn_graph.extent(1)); - RAFT_EXPECTS(knn_graph.extent(0) == dataset.extent(0), - "CAGRA: kNN graph and dataset must have the same number of rows"); - RAFT_EXPECTS(intermediate_graph_degree > 0 && partition_size > intermediate_graph_degree, - "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); - if constexpr (std::is_same_v || std::is_same_v) { - auto device_queries = - raft::make_device_matrix(res, partition_size, dataset.extent(1)); - raft::copy(res, device_queries.view(), dataset); - auto brute_force_index = cuvs::neighbors::brute_force::build( - res, params.build_params, raft::make_const_mdspan(device_queries.view())); - auto device_neighbors = raft::make_device_matrix( - res, partition_size, intermediate_graph_degree + 1); - auto device_distances = - raft::make_device_matrix(res, partition_size, intermediate_graph_degree + 1); - cuvs::neighbors::brute_force::search(res, - params.search_params, - brute_force_index, - raft::make_const_mdspan(device_queries.view()), - device_neighbors.view(), - device_distances.view()); - - auto host_neighbors = - raft::make_host_matrix(partition_size, intermediate_graph_degree + 1); - raft::copy(res, host_neighbors.view(), device_neighbors.view()); - raft::resource::sync_stream(res); - for (size_t row = 0; row < partition_size; ++row) { - size_t output_rank = 0; - for (size_t rank = 0; rank < intermediate_graph_degree + 1; ++rank) { - auto const neighbor = host_neighbors(row, rank); - if (neighbor == static_cast(row)) { continue; } - RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, - "CAGRA: brute-force kNN returned an invalid neighbor"); - knn_graph(row, output_rank++) = static_cast(neighbor); - if (output_rank == intermediate_graph_degree) { break; } - } - RAFT_EXPECTS(output_rank == intermediate_graph_degree, - "CAGRA: brute-force kNN did not return enough non-self neighbors"); - } - } else { - RAFT_FAIL("CAGRA: brute-force kNN build supports only float and half datasets"); - } -} - template void write_to_graph(raft::host_matrix_view knn_graph, raft::host_matrix_view neighbors_host_view, diff --git a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh new file mode 100644 index 0000000000..a525dcd1b7 --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh @@ -0,0 +1,103 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include +#include + +#include +#include +#include +#include +#include + +#include +#include +#include + +namespace cuvs::neighbors::cagra::detail { + +// Exact kNN construction. Request one extra +// neighbor and remove self by ID, including when duplicate vectors tie with self. +template +void build_knn_graph(raft::resources const& res, + raft::mdspan, Layout, Accessor> dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + auto const partition_size = static_cast(dataset.extent(0)); + auto const intermediate_graph_degree = static_cast(knn_graph.extent(1)); + RAFT_EXPECTS(knn_graph.extent(0) == dataset.extent(0), + "CAGRA: kNN graph and dataset must have the same number of rows"); + RAFT_EXPECTS(intermediate_graph_degree > 0 && partition_size > intermediate_graph_degree, + "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); + if constexpr (std::is_same_v || std::is_same_v) { + // Borrow contiguous device data. Host and strided inputs need a contiguous device copy + // for the brute-force index, which must stay alive across all query batches. + std::optional> device_dataset; + auto dataset_view = [&]() { + if constexpr (raft::is_device_mdspan_v) { + if (dataset.stride(1) == 1 && dataset.stride(0) == dataset.extent(1)) { + return raft::make_device_matrix_view( + dataset.data_handle(), dataset.extent(0), dataset.extent(1)); + } + } + device_dataset.emplace( + raft::make_device_matrix(res, partition_size, dataset.extent(1))); + raft::copy(res, device_dataset->view(), dataset); + return raft::make_const_mdspan(device_dataset->view()); + }(); + auto brute_force_index = + cuvs::neighbors::brute_force::build(res, params.build_params, dataset_view); + + // Size result buffers from the full remaining workspace after building the index. + // Allocate them from the large workspace resource, as in the other graph builders. + auto const free_bytes = raft::resource::get_workspace_free_bytes(res); + auto const top_k = intermediate_graph_degree + 1; + auto const bytes_per_row = top_k * (sizeof(int64_t) + sizeof(float)); + auto const max_batch_size = std::min(partition_size, free_bytes / bytes_per_row); + RAFT_EXPECTS(max_batch_size > 0, + "CAGRA: insufficient workspace for a brute-force kNN query batch"); + auto workspace_mr = raft::resource::get_large_workspace_resource_ref(res); + auto device_neighbors = raft::make_device_mdarray( + res, workspace_mr, raft::make_extents(max_batch_size, top_k)); + auto device_distances = raft::make_device_mdarray( + res, workspace_mr, raft::make_extents(max_batch_size, top_k)); + auto host_neighbors = raft::make_host_matrix(max_batch_size, top_k); + + for (size_t offset = 0; offset < partition_size; offset += max_batch_size) { + auto const batch_size = std::min(max_batch_size, partition_size - offset); + auto queries = raft::make_device_matrix_view( + dataset_view.data_handle() + offset * dataset.extent(1), batch_size, dataset.extent(1)); + auto neighbors = raft::make_device_matrix_view( + device_neighbors.data_handle(), batch_size, top_k); + auto distances = raft::make_device_matrix_view( + device_distances.data_handle(), batch_size, top_k); + auto neighbors_host = raft::make_host_matrix_view( + host_neighbors.data_handle(), batch_size, top_k); + cuvs::neighbors::brute_force::search( + res, params.search_params, brute_force_index, queries, neighbors, distances); + raft::copy(res, neighbors_host, neighbors); + raft::resource::sync_stream(res); + for (size_t row = 0; row < batch_size; ++row) { + size_t output_rank = 0; + for (size_t rank = 0; rank < top_k; ++rank) { + auto const neighbor = neighbors_host(row, rank); + if (neighbor == static_cast(offset + row)) { continue; } + RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, + "CAGRA: brute-force kNN returned an invalid neighbor"); + knn_graph(offset + row, output_rank++) = static_cast(neighbor); + if (output_rank == intermediate_graph_degree) { break; } + } + RAFT_EXPECTS(output_rank == intermediate_graph_degree, + "CAGRA: brute-force kNN did not return enough non-self neighbors"); + } + } + } else { + RAFT_FAIL("CAGRA: brute-force kNN build supports only float and half datasets"); + } +} + +} // namespace cuvs::neighbors::cagra::detail From 96105dcfdc4d9a7aac404a366104fd3be96f9dda Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 19:29:37 +0900 Subject: [PATCH 4/8] Update exact kNN build --- .../detail/cagra/cagra_build_brute_force.cuh | 341 +++++++++++++++--- 1 file changed, 286 insertions(+), 55 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh index a525dcd1b7..6ce1d7b0a2 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh @@ -6,19 +6,246 @@ #include #include +#include +#include #include #include #include +#include #include #include +#include #include #include #include +#include namespace cuvs::neighbors::cagra::detail { +// Choose equal query/reference chunks from the full free workspace. In addition to the +// two data chunks and results, budget norms and a conservative upper bound for the +// pairwise distances and intermediate top-k buffers used by brute-force search. +inline size_t brute_force_host_chunk_size( + size_t free_bytes, size_t rows, size_t dim, size_t element_size, size_t top_k) +{ + // Two translation offsets for the direct partition merge. + free_bytes -= std::min(free_bytes, 2 * sizeof(int64_t)); + size_t low = 0, high = rows; + while (low < high) { + auto const mid = low + (high - low + 1) / 2; + // Local search results, two 2*k merge buffers, and room for + // selection/sorting scratch. All explicit buffers are allocated once. + auto const result_bytes = + (std::min(mid, top_k) + 6 * top_k) * (sizeof(int64_t) + sizeof(float)); + // The tiled search uses at most 1024 query rows at a time, and at most `mid` + // reference columns for both pairwise distances and intermediate top-k results. + auto const scratch_bytes = std::min(mid, size_t{1024}) * (sizeof(int64_t) + 2 * sizeof(float)); + auto const bytes_per_row = + 2 * dim * element_size + result_bytes + 2 * sizeof(float) + scratch_bytes; + if (bytes_per_row <= free_bytes / mid) { + low = mid; + } else { + high = mid - 1; + } + } + return low; +} + +template +void build_brute_force_host_knn_graph( + raft::resources const& res, + raft::mdspan, Layout, Accessor> dataset, + raft::host_matrix_view knn_graph, + graph_build_params::brute_force_params params) +{ + auto const rows = static_cast(dataset.extent(0)); + auto const dim = dataset.extent(1); + auto const degree = static_cast(knn_graph.extent(1)); + auto const top_k = degree + 1; + auto const chunk_size = brute_force_host_chunk_size( + raft::resource::get_workspace_free_bytes(res), rows, dim, sizeof(T), top_k); + RAFT_EXPECTS(chunk_size >= top_k, + "CAGRA: brute-force host chunk size must be at least graph degree + 1; " + "increase the available workspace"); + auto mr = raft::resource::get_large_workspace_resource_ref(res); + auto queries = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, dim)); + auto references = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, dim)); + auto neighbors = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, top_k)); + auto distances = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, top_k)); + auto candidate_neighbors = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, 2 * top_k)); + auto candidate_distances = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, 2 * top_k)); + auto merged_neighbors = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, 2 * top_k)); + auto merged_distances = + raft::make_device_mdarray(res, mr, raft::make_extents(chunk_size, 2 * top_k)); + auto translations = raft::make_device_mdarray(res, mr, raft::make_extents(2)); + auto host_neighbors = raft::make_host_matrix(chunk_size, top_k); + bool const select_min = cuvs::distance::is_min_close(params.build_params.metric); + + bool const needs_packing = dataset.stride(1) != 1 || dataset.stride(0) != dim; + std::optional> packed_host_chunk; + if (needs_packing) { + packed_host_chunk.emplace(raft::make_host_matrix(chunk_size, dim)); + } + auto copy_chunk = [&](size_t offset, size_t count, auto destination) { + if (!needs_packing) { + auto source = raft::make_host_matrix_view( + dataset.data_handle() + offset * dim, count, dim); + raft::copy(res, destination, source); + } else { + // Pack on the host so arbitrary strides never require device staging storage. + auto source = + raft::make_host_matrix_view(packed_host_chunk->data_handle(), count, dim); + for (size_t row = 0; row < count; ++row) { + for (int64_t col = 0; col < dim; ++col) { + source(row, col) = dataset(offset + row, col); + } + } + raft::copy(res, destination, source); + // The same packing buffer is reused by the next chunk transfer. + raft::resource::sync_stream(res); + } + }; + + for (size_t query_offset = 0; query_offset < rows; query_offset += chunk_size) { + auto const query_rows = std::min(chunk_size, rows - query_offset); + auto query_view = + raft::make_device_matrix_view(queries.data_handle(), query_rows, dim); + copy_chunk(query_offset, query_rows, query_view); + size_t merged_k = 0; + + for (size_t reference_offset = 0; reference_offset < rows; reference_offset += chunk_size) { + auto const reference_rows = std::min(chunk_size, rows - reference_offset); + auto const local_k = std::min(top_k, reference_rows); + auto reference_view = + raft::make_device_matrix_view(references.data_handle(), reference_rows, dim); + copy_chunk(reference_offset, reference_rows, reference_view); + auto index = cuvs::neighbors::brute_force::build( + res, params.build_params, raft::make_const_mdspan(reference_view)); + // knn_merge_parts supports ascending distances and k <= 1024. Write the + // local search directly after the accumulated list for its partitioned layout. + bool const direct_merge = + select_min && top_k <= 1024 && merged_k == top_k && local_k == top_k; + auto* local_neighbors = direct_merge ? merged_neighbors.data_handle() + query_rows * top_k + : neighbors.data_handle(); + auto* local_distances = direct_merge ? merged_distances.data_handle() + query_rows * top_k + : distances.data_handle(); + auto neighbor_view = + raft::make_device_matrix_view(local_neighbors, query_rows, local_k); + auto distance_view = + raft::make_device_matrix_view(local_distances, query_rows, local_k); + cuvs::neighbors::brute_force::search(res, + params.search_params, + index, + raft::make_const_mdspan(query_view), + neighbor_view, + distance_view); + if (direct_merge) { + auto const offset = static_cast(reference_offset); + raft::linalg::map_offset(res, translations.view(), [offset] __device__(int64_t part) { + return part == 0 ? int64_t{0} : offset; + }); + cuvs::neighbors::knn_merge_parts(res, + raft::make_device_matrix_view( + merged_distances.data_handle(), 2 * query_rows, top_k), + raft::make_device_matrix_view( + merged_neighbors.data_handle(), 2 * query_rows, top_k), + raft::make_device_matrix_view( + candidate_distances.data_handle(), query_rows, top_k), + raft::make_device_matrix_view( + candidate_neighbors.data_handle(), query_rows, top_k), + translations.view()); + // Both buffers have capacity for two full lists; swapping never allocates. + std::swap(merged_distances, candidate_distances); + std::swap(merged_neighbors, candidate_neighbors); + continue; + } + if (merged_k == 0) { + // The first chunk already has global IDs (offset zero), so select directly. + cuvs::selection::select_k(res, + raft::make_const_mdspan(distance_view), + std::optional{raft::make_const_mdspan(neighbor_view)}, + raft::make_device_matrix_view( + merged_distances.data_handle(), query_rows, local_k), + raft::make_device_matrix_view( + merged_neighbors.data_handle(), query_rows, local_k), + select_min, + true); + merged_k = local_k; + continue; + } + auto const candidate_k = merged_k + local_k; + auto const next_k = std::min(top_k, candidate_k); + // Translate the local IDs and concatenate using existing RAFT operations. + raft::linalg::map(res, + neighbor_view, + raft::add_const_op{static_cast(reference_offset)}, + raft::make_const_mdspan(neighbor_view)); + auto append_candidates = [&](auto source, auto* destination, size_t column_offset) { + using value_type = std::remove_pointer_t; + auto target = raft::make_device_strided_matrix_view( + destination + column_offset, query_rows, source.extent(1), candidate_k); + raft::copy(res, target, source); + }; + append_candidates(raft::make_device_matrix_view( + merged_distances.data_handle(), query_rows, merged_k), + candidate_distances.data_handle(), + 0); + append_candidates(raft::make_device_matrix_view( + merged_neighbors.data_handle(), query_rows, merged_k), + candidate_neighbors.data_handle(), + 0); + append_candidates( + raft::make_const_mdspan(distance_view), candidate_distances.data_handle(), merged_k); + append_candidates( + raft::make_const_mdspan(neighbor_view), candidate_neighbors.data_handle(), merged_k); + auto candidate_distance_view = raft::make_device_matrix_view( + candidate_distances.data_handle(), query_rows, candidate_k); + auto candidate_neighbor_view = raft::make_device_matrix_view( + candidate_neighbors.data_handle(), query_rows, candidate_k); + auto merged_distance_view = raft::make_device_matrix_view( + merged_distances.data_handle(), query_rows, next_k); + auto merged_neighbor_view = raft::make_device_matrix_view( + merged_neighbors.data_handle(), query_rows, next_k); + cuvs::selection::select_k(res, + candidate_distance_view, + std::optional{candidate_neighbor_view}, + merged_distance_view, + merged_neighbor_view, + select_min, + true); + merged_k = next_k; + } + RAFT_EXPECTS(merged_k == top_k, "CAGRA: brute-force kNN did not return enough candidates"); + auto host_neighbor_view = raft::make_host_matrix_view( + host_neighbors.data_handle(), query_rows, top_k); + auto merged_neighbor_view = raft::make_device_matrix_view( + merged_neighbors.data_handle(), query_rows, top_k); + raft::copy(res, host_neighbor_view, merged_neighbor_view); + raft::resource::sync_stream(res); + for (size_t row = 0; row < query_rows; ++row) { + size_t output_rank = 0; + for (size_t rank = 0; rank < top_k && output_rank < degree; ++rank) { + auto const neighbor = host_neighbor_view(row, rank); + RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < rows, + "CAGRA: brute-force kNN returned an invalid neighbor"); + if (neighbor == static_cast(query_offset + row)) { continue; } + knn_graph(query_offset + row, output_rank++) = static_cast(neighbor); + } + RAFT_EXPECTS(output_rank == degree, + "CAGRA: brute-force kNN did not return enough non-self neighbors"); + } + } +} + // Exact kNN construction. Request one extra // neighbor and remove self by ID, including when duplicate vectors tie with self. template @@ -34,65 +261,69 @@ void build_knn_graph(raft::resources const& res, RAFT_EXPECTS(intermediate_graph_degree > 0 && partition_size > intermediate_graph_degree, "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); if constexpr (std::is_same_v || std::is_same_v) { - // Borrow contiguous device data. Host and strided inputs need a contiguous device copy - // for the brute-force index, which must stay alive across all query batches. - std::optional> device_dataset; - auto dataset_view = [&]() { - if constexpr (raft::is_device_mdspan_v) { - if (dataset.stride(1) == 1 && dataset.stride(0) == dataset.extent(1)) { - return raft::make_device_matrix_view( - dataset.data_handle(), dataset.extent(0), dataset.extent(1)); + if constexpr (raft::is_host_mdspan_v) { + build_brute_force_host_knn_graph(res, dataset, knn_graph, params); + } else { + // Borrow contiguous device data. Strided inputs need a contiguous device copy + // for the brute-force index, which must stay alive across all query batches. + std::optional> device_dataset; + auto dataset_view = [&]() { + if constexpr (raft::is_device_mdspan_v) { + if (dataset.stride(1) == 1 && dataset.stride(0) == dataset.extent(1)) { + return raft::make_device_matrix_view( + dataset.data_handle(), dataset.extent(0), dataset.extent(1)); + } } - } - device_dataset.emplace( - raft::make_device_matrix(res, partition_size, dataset.extent(1))); - raft::copy(res, device_dataset->view(), dataset); - return raft::make_const_mdspan(device_dataset->view()); - }(); - auto brute_force_index = - cuvs::neighbors::brute_force::build(res, params.build_params, dataset_view); + device_dataset.emplace( + raft::make_device_matrix(res, partition_size, dataset.extent(1))); + raft::copy(res, device_dataset->view(), dataset); + return raft::make_const_mdspan(device_dataset->view()); + }(); + auto brute_force_index = + cuvs::neighbors::brute_force::build(res, params.build_params, dataset_view); - // Size result buffers from the full remaining workspace after building the index. - // Allocate them from the large workspace resource, as in the other graph builders. - auto const free_bytes = raft::resource::get_workspace_free_bytes(res); - auto const top_k = intermediate_graph_degree + 1; - auto const bytes_per_row = top_k * (sizeof(int64_t) + sizeof(float)); - auto const max_batch_size = std::min(partition_size, free_bytes / bytes_per_row); - RAFT_EXPECTS(max_batch_size > 0, - "CAGRA: insufficient workspace for a brute-force kNN query batch"); - auto workspace_mr = raft::resource::get_large_workspace_resource_ref(res); - auto device_neighbors = raft::make_device_mdarray( - res, workspace_mr, raft::make_extents(max_batch_size, top_k)); - auto device_distances = raft::make_device_mdarray( - res, workspace_mr, raft::make_extents(max_batch_size, top_k)); - auto host_neighbors = raft::make_host_matrix(max_batch_size, top_k); + // Size result buffers from the full remaining workspace after building the index. + // Allocate them from the large workspace resource, as in the other graph builders. + auto const free_bytes = raft::resource::get_workspace_free_bytes(res); + auto const top_k = intermediate_graph_degree + 1; + auto const bytes_per_row = top_k * (sizeof(int64_t) + sizeof(float)); + auto const max_batch_size = std::min(partition_size, free_bytes / bytes_per_row); + RAFT_EXPECTS(max_batch_size > 0, + "CAGRA: insufficient workspace for a brute-force kNN query batch"); + auto workspace_mr = raft::resource::get_large_workspace_resource_ref(res); + auto device_neighbors = raft::make_device_mdarray( + res, workspace_mr, raft::make_extents(max_batch_size, top_k)); + auto device_distances = raft::make_device_mdarray( + res, workspace_mr, raft::make_extents(max_batch_size, top_k)); + auto host_neighbors = raft::make_host_matrix(max_batch_size, top_k); - for (size_t offset = 0; offset < partition_size; offset += max_batch_size) { - auto const batch_size = std::min(max_batch_size, partition_size - offset); - auto queries = raft::make_device_matrix_view( - dataset_view.data_handle() + offset * dataset.extent(1), batch_size, dataset.extent(1)); - auto neighbors = raft::make_device_matrix_view( - device_neighbors.data_handle(), batch_size, top_k); - auto distances = raft::make_device_matrix_view( - device_distances.data_handle(), batch_size, top_k); - auto neighbors_host = raft::make_host_matrix_view( - host_neighbors.data_handle(), batch_size, top_k); - cuvs::neighbors::brute_force::search( - res, params.search_params, brute_force_index, queries, neighbors, distances); - raft::copy(res, neighbors_host, neighbors); - raft::resource::sync_stream(res); - for (size_t row = 0; row < batch_size; ++row) { - size_t output_rank = 0; - for (size_t rank = 0; rank < top_k; ++rank) { - auto const neighbor = neighbors_host(row, rank); - if (neighbor == static_cast(offset + row)) { continue; } - RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, - "CAGRA: brute-force kNN returned an invalid neighbor"); - knn_graph(offset + row, output_rank++) = static_cast(neighbor); - if (output_rank == intermediate_graph_degree) { break; } + for (size_t offset = 0; offset < partition_size; offset += max_batch_size) { + auto const batch_size = std::min(max_batch_size, partition_size - offset); + auto queries = raft::make_device_matrix_view( + dataset_view.data_handle() + offset * dataset.extent(1), batch_size, dataset.extent(1)); + auto neighbors = raft::make_device_matrix_view( + device_neighbors.data_handle(), batch_size, top_k); + auto distances = raft::make_device_matrix_view( + device_distances.data_handle(), batch_size, top_k); + auto neighbors_host = raft::make_host_matrix_view( + host_neighbors.data_handle(), batch_size, top_k); + cuvs::neighbors::brute_force::search( + res, params.search_params, brute_force_index, queries, neighbors, distances); + raft::copy(res, neighbors_host, neighbors); + raft::resource::sync_stream(res); + for (size_t row = 0; row < batch_size; ++row) { + size_t output_rank = 0; + for (size_t rank = 0; rank < top_k; ++rank) { + auto const neighbor = neighbors_host(row, rank); + if (neighbor == static_cast(offset + row)) { continue; } + RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, + "CAGRA: brute-force kNN returned an invalid neighbor"); + knn_graph(offset + row, output_rank++) = static_cast(neighbor); + if (output_rank == intermediate_graph_degree) { break; } + } + RAFT_EXPECTS(output_rank == intermediate_graph_degree, + "CAGRA: brute-force kNN did not return enough non-self neighbors"); } - RAFT_EXPECTS(output_rank == intermediate_graph_degree, - "CAGRA: brute-force kNN did not return enough non-self neighbors"); } } } else { From 43b93b6b61b6c2dfbe7739770a8dd3140860a667 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 19:37:18 +0900 Subject: [PATCH 5/8] Update chunk size calc --- .../detail/cagra/cagra_build_brute_force.cuh | 31 +++++++------------ 1 file changed, 12 insertions(+), 19 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh index 6ce1d7b0a2..f786d37b28 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh @@ -32,25 +32,18 @@ inline size_t brute_force_host_chunk_size( { // Two translation offsets for the direct partition merge. free_bytes -= std::min(free_bytes, 2 * sizeof(int64_t)); - size_t low = 0, high = rows; - while (low < high) { - auto const mid = low + (high - low + 1) / 2; - // Local search results, two 2*k merge buffers, and room for - // selection/sorting scratch. All explicit buffers are allocated once. - auto const result_bytes = - (std::min(mid, top_k) + 6 * top_k) * (sizeof(int64_t) + sizeof(float)); - // The tiled search uses at most 1024 query rows at a time, and at most `mid` - // reference columns for both pairwise distances and intermediate top-k results. - auto const scratch_bytes = std::min(mid, size_t{1024}) * (sizeof(int64_t) + 2 * sizeof(float)); - auto const bytes_per_row = - 2 * dim * element_size + result_bytes + 2 * sizeof(float) + scratch_bytes; - if (bytes_per_row <= free_bytes / mid) { - low = mid; - } else { - high = mid - 1; - } - } - return low; + auto const data_bytes_per_row = 2 * dim * element_size; + // Per query, 7 = 1 top_k list for local results + 2 lists in each of the two + // merge buffers + 2 lists reserved for selection/sorting scratch. The first + // 5 lists are explicit allocations; the final 2 are a conservative allowance, + // not an exact scratch requirement of the selection library. + auto const result_bytes_per_row = 7 * top_k * (sizeof(int64_t) + sizeof(float)); + // Bound the query tile by the whole dataset (up to the search limit of 1024), + // rather than the unknown chunk size. This may overestimate scratch, never underestimate it. + auto const scratch_bytes_per_row = + std::min(rows, size_t{1024}) * (sizeof(int64_t) + 2 * sizeof(float)) + 2 * sizeof(float); + auto const bytes_per_row = data_bytes_per_row + result_bytes_per_row + scratch_bytes_per_row; + return std::min(rows, free_bytes / bytes_per_row); } template From c0bac99b65b1222d7b414a84b1a6a0bbbbcd634d Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Tue, 20 Oct 2026 23:52:23 +0900 Subject: [PATCH 6/8] Update API docs --- cpp/include/cuvs/neighbors/cagra.hpp | 72 ++++++++---- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 105 ++++++++++++++++++ 2 files changed, 157 insertions(+), 20 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 6fd2c0cb0a..94af0420fc 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -4380,11 +4380,19 @@ void distribute(const raft::resources& clique, cuvs::neighbors::mg_index, T, IdxT>* index); /** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major host dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. + * @brief Build an exact kNN graph using brute-force search. + * @ingroup cagra_cpp_index_build + * + * Each row of the output contains the nearest neighbor indices for the corresponding dataset + * point, ordered by distance and excluding the point itself. The graph is stored in host memory. + * The dataset is processed in chunks, allowing graphs larger than available device memory to be + * built. + * + * @param[in] res RAFT resources. + * @param[in] dataset Row-major host dataset of float values, with shape [n_rows, dim]. + * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - + * 1]. + * @param[in] build_params Brute-force index and search parameters. */ void build_knn_graph(raft::resources const& res, raft::host_matrix_view dataset, @@ -4392,11 +4400,19 @@ void build_knn_graph(raft::resources const& res, graph_build_params::brute_force_params build_params); /** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major device dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. + * @brief Build an exact kNN graph using brute-force search. + * @ingroup cagra_cpp_index_build + * + * Each row of the output contains the nearest neighbor indices for the corresponding dataset + * point, ordered by distance and excluding the point itself. The graph is stored in host memory. + * The dataset is processed in chunks, allowing graphs larger than available device memory to be + * built. + * + * @param[in] res RAFT resources. + * @param[in] dataset Row-major device dataset of float values, with shape [n_rows, dim]. + * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - + * 1]. + * @param[in] build_params Brute-force index and search parameters. */ void build_knn_graph(raft::resources const& res, raft::device_matrix_view dataset, @@ -4404,11 +4420,19 @@ void build_knn_graph(raft::resources const& res, graph_build_params::brute_force_params build_params); /** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major host dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. + * @brief Build an exact kNN graph using brute-force search. + * @ingroup cagra_cpp_index_build + * + * Each row of the output contains the nearest neighbor indices for the corresponding dataset + * point, ordered by distance and excluding the point itself. The graph is stored in host memory. + * The dataset is processed in chunks, allowing graphs larger than available device memory to be + * built. + * + * @param[in] res RAFT resources. + * @param[in] dataset Row-major host dataset of half values, with shape [n_rows, dim]. + * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - + * 1]. + * @param[in] build_params Brute-force index and search parameters. */ void build_knn_graph(raft::resources const& res, raft::host_matrix_view dataset, @@ -4416,11 +4440,19 @@ void build_knn_graph(raft::resources const& res, graph_build_params::brute_force_params build_params); /** - * @brief Build an exact, distance-ordered kNN graph, excluding self neighbors. - * @param res RAFT resources. - * @param dataset Row-major device dataset of float or half values. - * @param knn_graph Host output with matching rows and degree in [1, n_rows - 1]. - * @param build_params Brute-force index and search parameters. + * @brief Build an exact kNN graph using brute-force search. + * @ingroup cagra_cpp_index_build + * + * Each row of the output contains the nearest neighbor indices for the corresponding dataset + * point, ordered by distance and excluding the point itself. The graph is stored in host memory. + * The dataset is processed in chunks, allowing graphs larger than available device memory to be + * built. + * + * @param[in] res RAFT resources. + * @param[in] dataset Row-major device dataset of half values, with shape [n_rows, dim]. + * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - + * 1]. + * @param[in] build_params Brute-force index and search parameters. */ void build_knn_graph(raft::resources const& res, raft::device_matrix_view dataset, diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index aca906cc12..7faf260902 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -1207,6 +1207,111 @@ cuvs::neighbors::device_bbq_dataset_view const& dataset) `cuvs::neighbors::cagra::device_bbq_index` + +### neighbors::cagra::build_knn_graph + +Build an exact kNN graph using brute-force search. + +```cpp +void build_knn_graph(raft::resources const& res, +raft::host_matrix_view dataset, +raft::host_matrix_view knn_graph, +graph_build_params::brute_force_params build_params); +``` + +Each row of the output contains the nearest neighbor indices for the corresponding dataset point, ordered by distance and excluding the point itself. The graph is stored in host memory. The dataset is processed in chunks, allowing graphs larger than available device memory to be built. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | `raft::resources const&` | RAFT resources. | +| `dataset` | in | `raft::host_matrix_view` | Row-major host dataset of float values, with shape [n_rows, dim]. | +| `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | + +**Returns** + +`void` + +**Additional overload:** `neighbors::cagra::build_knn_graph` + +Build an exact kNN graph using brute-force search. + +```cpp +void build_knn_graph(raft::resources const& res, +raft::device_matrix_view dataset, +raft::host_matrix_view knn_graph, +graph_build_params::brute_force_params build_params); +``` + +Each row of the output contains the nearest neighbor indices for the corresponding dataset point, ordered by distance and excluding the point itself. The graph is stored in host memory. The dataset is processed in chunks, allowing graphs larger than available device memory to be built. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | `raft::resources const&` | RAFT resources. | +| `dataset` | in | `raft::device_matrix_view` | Row-major device dataset of float values, with shape [n_rows, dim]. | +| `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | + +**Returns** + +`void` + +**Additional overload:** `neighbors::cagra::build_knn_graph` + +Build an exact kNN graph using brute-force search. + +```cpp +void build_knn_graph(raft::resources const& res, +raft::host_matrix_view dataset, +raft::host_matrix_view knn_graph, +graph_build_params::brute_force_params build_params); +``` + +Each row of the output contains the nearest neighbor indices for the corresponding dataset point, ordered by distance and excluding the point itself. The graph is stored in host memory. The dataset is processed in chunks, allowing graphs larger than available device memory to be built. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | `raft::resources const&` | RAFT resources. | +| `dataset` | in | `raft::host_matrix_view` | Row-major host dataset of half values, with shape [n_rows, dim]. | +| `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | + +**Returns** + +`void` + +**Additional overload:** `neighbors::cagra::build_knn_graph` + +Build an exact kNN graph using brute-force search. + +```cpp +void build_knn_graph(raft::resources const& res, +raft::device_matrix_view dataset, +raft::host_matrix_view knn_graph, +graph_build_params::brute_force_params build_params); +``` + +Each row of the output contains the nearest neighbor indices for the corresponding dataset point, ordered by distance and excluding the point itself. The graph is stored in host memory. The dataset is processed in chunks, allowing graphs larger than available device memory to be built. + +**Parameters** + +| Name | Direction | Type | Description | +| --- | --- | --- | --- | +| `res` | in | `raft::resources const&` | RAFT resources. | +| `dataset` | in | `raft::device_matrix_view` | Row-major device dataset of half values, with shape [n_rows, dim]. | +| `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | + +**Returns** + +`void` + ## CAGRA extend functions From 52cf51e7df39db348103c9ec4e5f20b900784b43 Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 00:42:05 +0900 Subject: [PATCH 7/8] Fix the device working memory size calculation --- .../detail/cagra/cagra_build_brute_force.cuh | 31 ++++++++++++++----- 1 file changed, 23 insertions(+), 8 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh index f786d37b28..6613b89c74 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh @@ -247,11 +247,11 @@ void build_knn_graph(raft::resources const& res, raft::host_matrix_view knn_graph, graph_build_params::brute_force_params params) { - auto const partition_size = static_cast(dataset.extent(0)); - auto const intermediate_graph_degree = static_cast(knn_graph.extent(1)); + auto const partition_size = static_cast(dataset.extent(0)); + auto const graph_degree = static_cast(knn_graph.extent(1)); RAFT_EXPECTS(knn_graph.extent(0) == dataset.extent(0), "CAGRA: kNN graph and dataset must have the same number of rows"); - RAFT_EXPECTS(intermediate_graph_degree > 0 && partition_size > intermediate_graph_degree, + RAFT_EXPECTS(graph_degree > 0 && partition_size > graph_degree, "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); if constexpr (std::is_same_v || std::is_same_v) { if constexpr (raft::is_host_mdspan_v) { @@ -277,9 +277,24 @@ void build_knn_graph(raft::resources const& res, // Size result buffers from the full remaining workspace after building the index. // Allocate them from the large workspace resource, as in the other graph builders. - auto const free_bytes = raft::resource::get_workspace_free_bytes(res); - auto const top_k = intermediate_graph_degree + 1; - auto const bytes_per_row = top_k * (sizeof(int64_t) + sizeof(float)); + auto free_bytes = raft::resource::get_workspace_free_bytes(res); + auto const top_k = graph_degree + 1; + // Fused L2 search needs norm workspace for both the index and query rows. + free_bytes -= std::min(free_bytes, partition_size * sizeof(float)); + auto bytes_per_row = top_k * (sizeof(int64_t) + sizeof(float)) + sizeof(float); + auto const metric = params.build_params.metric; + bool const fused_l2 = + top_k <= 64 && (metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::L2Unexpanded || + metric == cuvs::distance::DistanceType::L2SqrtUnexpanded); + if (!fused_l2) { + // Tiled search (including inner product for any k) also allocates pairwise + // distances, intermediate top-k lists, and selection/sorting scratch. + // Bound each tile/list by a full index row to avoid depending on tile sizing. + bytes_per_row += partition_size * (2 * sizeof(int64_t) + 3 * sizeof(float)) + + 2 * top_k * (sizeof(int64_t) + sizeof(float)); + } auto const max_batch_size = std::min(partition_size, free_bytes / bytes_per_row); RAFT_EXPECTS(max_batch_size > 0, "CAGRA: insufficient workspace for a brute-force kNN query batch"); @@ -312,9 +327,9 @@ void build_knn_graph(raft::resources const& res, RAFT_EXPECTS(neighbor >= 0 && static_cast(neighbor) < partition_size, "CAGRA: brute-force kNN returned an invalid neighbor"); knn_graph(offset + row, output_rank++) = static_cast(neighbor); - if (output_rank == intermediate_graph_degree) { break; } + if (output_rank == graph_degree) { break; } } - RAFT_EXPECTS(output_rank == intermediate_graph_degree, + RAFT_EXPECTS(output_rank == graph_degree, "CAGRA: brute-force kNN did not return enough non-self neighbors"); } } From 1f3494109c02538a83f078eb3c278154500dd20e Mon Sep 17 00:00:00 2001 From: enp1s0 Date: Wed, 21 Oct 2026 01:06:31 +0900 Subject: [PATCH 8/8] Fix dorking memory size calc --- cpp/include/cuvs/neighbors/cagra.hpp | 12 ++++++--- .../detail/cagra/cagra_build_brute_force.cuh | 26 +++++++++++-------- fern/pages/cpp_api/cpp-api-neighbors-cagra.md | 8 +++--- 3 files changed, 27 insertions(+), 19 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 94af0420fc..d951926553 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -4392,7 +4392,8 @@ void distribute(const raft::resources& clique, * @param[in] dataset Row-major host dataset of float values, with shape [n_rows, dim]. * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - * 1]. - * @param[in] build_params Brute-force index and search parameters. + * @param[in] build_params Brute-force index and search parameters. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. */ void build_knn_graph(raft::resources const& res, raft::host_matrix_view dataset, @@ -4412,7 +4413,8 @@ void build_knn_graph(raft::resources const& res, * @param[in] dataset Row-major device dataset of float values, with shape [n_rows, dim]. * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - * 1]. - * @param[in] build_params Brute-force index and search parameters. + * @param[in] build_params Brute-force index and search parameters. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. */ void build_knn_graph(raft::resources const& res, raft::device_matrix_view dataset, @@ -4432,7 +4434,8 @@ void build_knn_graph(raft::resources const& res, * @param[in] dataset Row-major host dataset of half values, with shape [n_rows, dim]. * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - * 1]. - * @param[in] build_params Brute-force index and search parameters. + * @param[in] build_params Brute-force index and search parameters. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. */ void build_knn_graph(raft::resources const& res, raft::host_matrix_view dataset, @@ -4452,7 +4455,8 @@ void build_knn_graph(raft::resources const& res, * @param[in] dataset Row-major device dataset of half values, with shape [n_rows, dim]. * @param[out] knn_graph Host output with shape [n_rows, degree], where degree is in [1, n_rows - * 1]. - * @param[in] build_params Brute-force index and search parameters. + * @param[in] build_params Brute-force index and search parameters. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. */ void build_knn_graph(raft::resources const& res, raft::device_matrix_view dataset, diff --git a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh index 6613b89c74..063f20d58a 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh @@ -253,6 +253,13 @@ void build_knn_graph(raft::resources const& res, "CAGRA: kNN graph and dataset must have the same number of rows"); RAFT_EXPECTS(graph_degree > 0 && partition_size > graph_degree, "CAGRA: kNN graph degree must be positive and smaller than the dataset size"); + auto const metric = params.build_params.metric; + RAFT_EXPECTS(metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded || + metric == cuvs::distance::DistanceType::InnerProduct || + metric == cuvs::distance::DistanceType::CosineExpanded, + "CAGRA: brute-force kNN graph build supports only L2Expanded, L2SqrtExpanded, " + "InnerProduct and CosineExpanded metrics"); if constexpr (std::is_same_v || std::is_same_v) { if constexpr (raft::is_host_mdspan_v) { build_brute_force_host_knn_graph(res, dataset, knn_graph, params); @@ -277,17 +284,14 @@ void build_knn_graph(raft::resources const& res, // Size result buffers from the full remaining workspace after building the index. // Allocate them from the large workspace resource, as in the other graph builders. - auto free_bytes = raft::resource::get_workspace_free_bytes(res); - auto const top_k = graph_degree + 1; - // Fused L2 search needs norm workspace for both the index and query rows. - free_bytes -= std::min(free_bytes, partition_size * sizeof(float)); - auto bytes_per_row = top_k * (sizeof(int64_t) + sizeof(float)) + sizeof(float); - auto const metric = params.build_params.metric; - bool const fused_l2 = - top_k <= 64 && (metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::L2SqrtExpanded || - metric == cuvs::distance::DistanceType::L2Unexpanded || - metric == cuvs::distance::DistanceType::L2SqrtUnexpanded); + auto free_bytes = raft::resource::get_workspace_free_bytes(res); + auto const top_k = graph_degree + 1; + auto bytes_per_row = top_k * (sizeof(int64_t) + sizeof(float)); + bool const fused_l2 = top_k <= 64 && (metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::L2SqrtExpanded); + // Fused L2 allocates norm workspace for both inputs; tiled search reuses index norms. + if (fused_l2) { free_bytes -= std::min(free_bytes, partition_size * sizeof(float)); } + if (metric != cuvs::distance::DistanceType::InnerProduct) { bytes_per_row += sizeof(float); } if (!fused_l2) { // Tiled search (including inner product for any k) also allocates pairwise // distances, intermediate top-k lists, and selection/sorting scratch. diff --git a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md index 7faf260902..2d482c7737 100644 --- a/fern/pages/cpp_api/cpp-api-neighbors-cagra.md +++ b/fern/pages/cpp_api/cpp-api-neighbors-cagra.md @@ -1228,7 +1228,7 @@ Each row of the output contains the nearest neighbor indices for the correspondi | `res` | in | `raft::resources const&` | RAFT resources. | | `dataset` | in | `raft::host_matrix_view` | Row-major host dataset of float values, with shape [n_rows, dim]. | | `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | -| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | **Returns** @@ -1254,7 +1254,7 @@ Each row of the output contains the nearest neighbor indices for the correspondi | `res` | in | `raft::resources const&` | RAFT resources. | | `dataset` | in | `raft::device_matrix_view` | Row-major device dataset of float values, with shape [n_rows, dim]. | | `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | -| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | **Returns** @@ -1280,7 +1280,7 @@ Each row of the output contains the nearest neighbor indices for the correspondi | `res` | in | `raft::resources const&` | RAFT resources. | | `dataset` | in | `raft::host_matrix_view` | Row-major host dataset of half values, with shape [n_rows, dim]. | | `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | -| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | **Returns** @@ -1306,7 +1306,7 @@ Each row of the output contains the nearest neighbor indices for the correspondi | `res` | in | `raft::resources const&` | RAFT resources. | | `dataset` | in | `raft::device_matrix_view` | Row-major device dataset of half values, with shape [n_rows, dim]. | | `knn_graph` | out | `raft::host_matrix_view` | Host output with shape [n_rows, degree], where degree is in [1, n_rows - 1]. | -| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. | +| `build_params` | in | `graph_build_params::brute_force_params` | Brute-force index and search parameters. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | **Returns**