diff --git a/c/src/neighbors/cagra.cpp b/c/src/neighbors/cagra.cpp index 30d356101f..436b1370e0 100644 --- a/c/src/neighbors/cagra.cpp +++ b/c/src/neighbors/cagra.cpp @@ -757,11 +757,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 98df8cdadf..9ed658a3ea 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; @@ -4415,6 +4420,90 @@ void distribute(const raft::resources& clique, const std::string& filename, cuvs::neighbors::mg_index, T, IdxT>* index); +/** + * @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. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. + */ +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 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. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. + */ +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 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. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. + */ +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 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. Supported metrics are + * L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. + */ +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 550b317552..4db3361280 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 d3c5b0fbe0..23c3d740b6 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_INST_CAGRA_BUILD(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 8c8766cbd5..bec6a67105 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh @@ -8,6 +8,7 @@ #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 @@ -982,11 +983,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."); @@ -997,7 +999,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). */ @@ -1026,6 +1029,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/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..063f20d58a --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_build_brute_force.cuh @@ -0,0 +1,346 @@ +/* + * 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 +#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)); + 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 +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 +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 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(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); + } 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); + + // 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; + 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. + // 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"); + 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 == graph_degree) { break; } + } + RAFT_EXPECTS(output_rank == 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 diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 3917e16b7e..e3c5936921 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -282,6 +282,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 }; @@ -335,7 +337,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", "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] @@ -365,6 +368,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))) @@ -442,6 +450,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; @@ -587,6 +598,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(); @@ -665,6 +681,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; @@ -795,6 +814,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(); @@ -879,6 +903,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; @@ -1177,6 +1204,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(); @@ -1286,6 +1318,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(); @@ -1518,6 +1558,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; @@ -1956,6 +1999,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; } @@ -2183,6 +2246,9 @@ class AnnCagraMultiPartitionTest : public ::testing::TestWithParam; +graph_build_params::iterative_search_params, +graph_build_params::brute_force_params>; ``` @@ -1206,6 +1207,111 @@ cuvs::neighbors::device_bbq_dataset_view 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. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | + +**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. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | + +**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. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | + +**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. Supported metrics are L2Expanded, L2SqrtExpanded, InnerProduct and CosineExpanded. | + +**Returns** + +`void` + ## CAGRA extend functions