diff --git a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h index 57b47d97db..dd8ebc0abb 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h +++ b/cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -423,6 +423,37 @@ void parse_build_param(const nlohmann::json& conf, throw std::runtime_error("invalid value for merge_type"); } } + if (conf.contains("merge_algo")) { + std::string algo = conf.at("merge_algo"); + if (algo == "AUTO") { + param.merge_params.algo = cuvs::neighbors::cagra::merge_algo::AUTO; + } else if (algo == "FASTENER") { + param.merge_params.algo = cuvs::neighbors::cagra::merge_algo::FASTENER; + } else if (algo == "REBUILD") { + param.merge_params.algo = cuvs::neighbors::cagra::merge_algo::REBUILD; + } else { + throw std::runtime_error("invalid value for merge_algo"); + } + } + if (conf.contains("fastener_levels")) { param.merge_params.levels = conf.at("fastener_levels"); } + if (conf.contains("fastener_root_fanout")) { + param.merge_params.root_fanout = conf.at("fastener_root_fanout"); + } + if (conf.contains("fastener_lower_fanout")) { + param.merge_params.lower_fanout = conf.at("fastener_lower_fanout"); + } + if (conf.contains("fastener_leader_fraction")) { + param.merge_params.leader_fraction = conf.at("fastener_leader_fraction"); + } + if (conf.contains("fastener_max_leaders")) { + param.merge_params.max_leaders = conf.at("fastener_max_leaders"); + } + if (conf.contains("fastener_leaf_size")) { + param.merge_params.leaf_size = conf.at("fastener_leaf_size"); + } + if (conf.contains("fastener_leaf_degree")) { + param.merge_params.leaf_degree = conf.at("fastener_leaf_degree"); + } param.cagra_params = [conf](raft::matrix_extent extents, cuvs::distance::DistanceType dist_type) { // Delayed parsing/initialization of cagra_params - it's called once the dataset shape is known diff --git a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h index 87111e4761..7c4f4e6170 100644 --- a/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h +++ b/cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -101,6 +101,7 @@ class cuvs_cagra : public algo, public algo_gpu { dataset_dependent_params cagra_params; size_t num_dataset_splits = 1; CagraMergeType merge_type = CagraMergeType::kPhysical; + cuvs::neighbors::cagra::merge_params merge_params; }; cuvs_cagra(Metric metric, int dim, const build_param& param, int concurrent_searches = 1) @@ -228,21 +229,10 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) auto sub_dev = raft::make_device_matrix_view(sub_ptr, rows, dim_); - auto sub_index = cuvs::neighbors::cagra::index(handle_, params.metric); - if (index_params_.merge_type == CagraMergeType::kPhysical) { - if (dataset_is_on_host) { - sub_index.update_dataset(handle_, sub_host); - } else { - sub_index.update_dataset(handle_, sub_dev); - } - } - if (index_params_.merge_type == CagraMergeType::kLogical) { - if (dataset_is_on_host) { - sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_host); - } else { - sub_index = cuvs::neighbors::cagra::build(handle_, params, sub_dev); - } - } + // Build every partition before the merge so FASTENER and REBUILD consume identical prepared + // inputs. Partition construction remains outside cagra::merge's NVTX range. + auto sub_index = dataset_is_on_host ? cuvs::neighbors::cagra::build(handle_, params, sub_host) + : cuvs::neighbors::cagra::build(handle_, params, sub_dev); auto sub_index_shared = std::make_shared>(std::move(sub_index)); sub_indices_.push_back(std::move(sub_index_shared)); @@ -254,8 +244,8 @@ void cuvs_cagra::build(const T* dataset, size_t nrow) indices.push_back(ptr.get()); } - index_ = std::make_shared>( - std::move(cuvs::neighbors::cagra::merge(handle_, params, indices))); + index_ = std::make_shared>(std::move( + cuvs::neighbors::cagra::merge(handle_, params, indices, index_params_.merge_params))); } } } diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 98f3a96e87..fac6614204 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -697,6 +697,16 @@ struct CUVS_EXPORT index : cuvs::neighbors::index { graph_view_ = knn_graph; } + /** + * Replace the graph by taking ownership of an existing device matrix. + */ + void update_graph(raft::resources const&, + raft::device_matrix&& knn_graph) + { + graph_ = std::move(knn_graph); + graph_view_ = graph_.view(); + } + /** * Replace the graph with a new graph. * @@ -2447,42 +2457,65 @@ void serialize_to_hnswlib( * @{ */ -/** @brief Merge multiple CAGRA indices into a single index. - * - * This function merges multiple CAGRA indices into one, combining both the datasets and graph - * structures. - * - * @note: When device memory is sufficient, the dataset attached to the returned index is allocated - * in device memory by default; otherwise, host memory is used automatically. - * - * @note: This API only supports physical merge (`merge_strategy = MERGE_STRATEGY_PHYSICAL`), and - * attempting a logical merge here will throw an error. - * - * Usage example: - * @code{.cpp} - * using namespace cuvs::neighbors; - * auto dataset0 = raft::make_host_matrix(handle, size0, dim); - * auto dataset1 = raft::make_host_matrix(handle, size1, dim); +enum class merge_algo { + /** Select Fastener when its complete preflight succeeds, otherwise preserve rebuild behavior. */ + AUTO, + /** Require Fastener and reject unsupported configurations. */ + FASTENER, + /** Concatenate datasets and rebuild the graph. */ + REBUILD, +}; + +/** C++ controls for physical CAGRA index merge. */ +struct merge_params { + merge_algo algo = merge_algo::AUTO; + uint32_t levels = 2; + uint32_t root_fanout = 2; + uint32_t lower_fanout = 3; + double leader_fraction = 0.02; + uint32_t max_leaders = 1000; + uint32_t leaf_size = 256; + uint32_t leaf_degree = 4; +}; + +/** @brief Merge multiple physical CAGRA indices into one. * - * auto index0 = cagra::build(res, index_params, dataset0); - * auto index1 = cagra::build(res, index_params, dataset1); + * The overload without `merge_params` uses `merge_algo::AUTO`. AUTO runs Fastener only after a + * non-mutating preflight validates every input and option; otherwise it calls the existing rebuild + * implementation. REBUILD always preserves every input dataset. * - * std::vector*> indices{&index0, &index1}; + * Fastener supports unfiltered, uncompressed `float`, `half`, `int8_t`, and `uint8_t` + * indices using L2Expanded and `uint32_t` graph IDs. Fastener uses `root_fanout` at the first + * split and `lower_fanout` at every later split. Fanouts from 1 through 32, leader fractions in + * (0, 1], leader caps through 8192, leaf sizes from 1 through 256, and leaf degrees from 1 + * through 8 are supported. The configured spill width times `leaf_degree` must not exceed 255. + * `index_params::graph_degree` is the final output degree. * - * auto merged_index = cagra::merge(res, index_params, indices); - * @endcode + * Fastener copies the input datasets into one contiguous device allocation and never mutates the + * input indices. When `attach_dataset_on_build` is true, the result owns a 16-byte-row-aligned + * copy of that consolidated dataset (matching `cagra::build` / `update_dataset`). * - * @param[in] res RAFT resources used for the merge operation. - * @param[in] params Parameters that control the merging process. - * @param[in] indices A vector of pointers to the CAGRA indices to merge. All indices must: - * - Have attached datasets with the same dimension. - * @param[in] row_filter an optional device filter function object that greenlights rows - * to include in the merged index (none_sample_filter for no filtering) - * @return A new CAGRA index containing the merged indices, graph, and dataset. + * @param[in] res RAFT resources used for the merge. + * @param[in] params Parameters for the returned CAGRA index. + * @param[in] indices CAGRA indices to merge. + * @param[in] row_filter Optional row filter. Any filter selects rebuild in AUTO and is rejected by + * explicit FASTENER. + * @return The merged physical CAGRA index. + */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) + -> cuvs::neighbors::cagra::index; + +/** @copydoc merge + * @param[in] merge_params Parameters for the merge. */ auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, std::vector*>& indices, + const cuvs::neighbors::cagra::merge_params& merge_params, const cuvs::neighbors::filtering::base_filter& row_filter = cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; @@ -2495,6 +2528,15 @@ auto merge(raft::resources const& res, cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; +/** @copydoc merge */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::cagra::merge_params& merge_params, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) + -> cuvs::neighbors::cagra::index; + /** @copydoc merge */ auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, @@ -2503,10 +2545,28 @@ auto merge(raft::resources const& res, cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; +/** @copydoc merge */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::cagra::merge_params& merge_params, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) + -> cuvs::neighbors::cagra::index; + +/** @copydoc merge */ +auto merge(raft::resources const& res, + const cuvs::neighbors::cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter = + cuvs::neighbors::filtering::none_sample_filter{}) + -> cuvs::neighbors::cagra::index; + /** @copydoc merge */ auto merge(raft::resources const& res, const cuvs::neighbors::cagra::index_params& params, std::vector*>& indices, + const cuvs::neighbors::cagra::merge_params& merge_params, const cuvs::neighbors::filtering::base_filter& row_filter = cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index ee87c2c0ab..484e2b599d 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -427,6 +427,16 @@ index merge(raft::resources const& handle, return cagra::detail::merge(handle, params, indices, row_filter); } +template +index merge(raft::resources const& handle, + const cagra::index_params& params, + std::vector*>& indices, + const cagra::merge_params& merge_params, + const cuvs::neighbors::filtering::base_filter& row_filter) +{ + return cagra::detail::merge(handle, params, indices, merge_params, row_filter); +} + /** @} */ // end group cagra } // namespace cuvs::neighbors::cagra @@ -439,4 +449,14 @@ index merge(raft::resources const& handle, -> cuvs::neighbors::cagra::index \ { \ return cuvs::neighbors::cagra::merge(handle, params, indices, row_filter); \ + } \ + auto merge(raft::resources const& handle, \ + const cuvs::neighbors::cagra::index_params& params, \ + std::vector*>& indices, \ + const cuvs::neighbors::cagra::merge_params& merge_params, \ + const cuvs::neighbors::filtering::base_filter& row_filter) \ + -> cuvs::neighbors::cagra::index \ + { \ + return cuvs::neighbors::cagra::merge( \ + handle, params, indices, merge_params, row_filter); \ } diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 1dd4cbe075..82a9fe280e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -1,9 +1,13 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once +#include "../../../core/nvtx.hpp" +#include "cagra_merge_scaffold.cuh" +#include "graph_core.cuh" + #include #include @@ -21,19 +25,44 @@ #include #include -#include - -#include -#include +#include +#include +#include +#include #include namespace cuvs::neighbors::cagra::detail { +/** Copy every attached input dataset into its row range in one contiguous device matrix. */ +template +void copy_input_datasets(raft::resources const& handle, + std::vector*> const& indices, + std::vector const& offsets, + int64_t dim, + T* destination) +{ + using cagra_index_t = cuvs::neighbors::cagra::index; + using ds_idx_type = typename cagra_index_t::dataset_index_type; + + for (std::size_t i = 0; i < indices.size(); ++i) { + auto const* source_dataset = + dynamic_cast*>(&indices[i]->data()); + auto source = source_dataset->view(); + raft::copy_matrix(destination + offsets[i] * dim, + static_cast(dim), + source.data_handle(), + static_cast(source_dataset->stride()), + static_cast(dim), + static_cast(source.extent(0)), + raft::resource::get_cuda_stream(handle)); + } +} + template -index merge(raft::resources const& handle, - const cagra::index_params& params, - std::vector*>& indices, - const cuvs::neighbors::filtering::base_filter& row_filter) +index merge_rebuild(raft::resources const& handle, + const cagra::index_params& params, + std::vector*>& indices, + const cuvs::neighbors::filtering::base_filter& row_filter) { using cagra_index_t = cuvs::neighbors::cagra::index; using ds_idx_type = typename cagra_index_t::dataset_index_type; @@ -175,4 +204,236 @@ index merge(raft::resources const& handle, } } +struct fastener_preflight_result { + bool eligible = false; + int64_t rows = 0; + int64_t dim = 0; + std::vector offsets; + std::string reason; +}; + +template +auto preflight_fastener(cagra::index_params const& params, + cagra::merge_params const& merge_params, + std::vector*> const& indices, + cuvs::neighbors::filtering::base_filter const& row_filter) + -> fastener_preflight_result +{ + fastener_preflight_result result; + auto reject = [&](std::string reason) { + result.reason = std::move(reason); + return result; + }; + + if constexpr (!(std::is_same_v || std::is_same_v || + std::is_same_v || std::is_same_v) || + !std::is_same_v) { + return reject("the scalar or graph index type is unsupported"); + } + if (indices.size() < 2) { return reject("at least two input indices are required"); } + if (row_filter.get_filter_type() != cuvs::neighbors::filtering::FilterType::None) { + return reject("row filters are not supported"); + } + if (params.compression.has_value()) { return reject("compressed output is not supported"); } + if (params.metric != cuvs::distance::DistanceType::L2Expanded) { + return reject("only L2Expanded is supported"); + } + if (merge_params.levels == 0) { return reject("levels must be positive"); } + if (merge_params.root_fanout < 1 || merge_params.root_fanout > merge_scaffold::MAX_FANOUT || + merge_params.lower_fanout < 1 || merge_params.lower_fanout > merge_scaffold::MAX_FANOUT) { + return reject("root_fanout and lower_fanout must be between 1 and " + + std::to_string(merge_scaffold::MAX_FANOUT)); + } + if (!(merge_params.leader_fraction > 0.0 && merge_params.leader_fraction <= 1.0)) { + return reject("leader_fraction must be in (0, 1]"); + } + if (merge_params.max_leaders == 0 || merge_params.max_leaders > merge_scaffold::MAX_LEADERS) { + return reject("max_leaders must be between 1 and " + + std::to_string(merge_scaffold::MAX_LEADERS)); + } + if (merge_params.max_leaders < std::max(merge_params.root_fanout, merge_params.lower_fanout)) { + return reject("max_leaders must cover both configured fanouts"); + } + if (merge_params.leaf_size == 0 || merge_params.leaf_size > merge_scaffold::MAX_LEAF_SIZE) { + return reject("leaf_size must be between 1 and " + + std::to_string(merge_scaffold::MAX_LEAF_SIZE)); + } + if (merge_params.leaf_degree == 0 || + merge_params.leaf_degree > static_cast(merge_scaffold::MAX_LEAF_DEGREE)) { + return reject("leaf_degree must be between 1 and " + + std::to_string(merge_scaffold::MAX_LEAF_DEGREE)); + } + + uint64_t const max_spill = std::numeric_limits::max() / merge_params.leaf_degree; + uint64_t spill = merge_params.root_fanout; + auto const candidate_width_limit = + "root_fanout * lower_fanout^(levels - 1) * leaf_degree must not exceed " + + std::to_string(std::numeric_limits::max()); + if (spill > max_spill) { return reject(candidate_width_limit); } + if (merge_params.lower_fanout > 1) { + for (uint32_t level = 1; level < merge_params.levels; ++level) { + if (spill > max_spill / merge_params.lower_fanout) { return reject(candidate_width_limit); } + spill *= merge_params.lower_fanout; + } + } + uint64_t const scaffold_degree = spill * merge_params.leaf_degree; + + using index_type = cuvs::neighbors::cagra::index; + using ds_idx_type = typename index_type::dataset_index_type; + uint64_t rows = 0; + uint64_t max_input_degree = 0; + result.offsets.reserve(indices.size() + 1); + result.offsets.push_back(0); + + for (auto const* index : indices) { + if (index == nullptr) { return reject("all input index pointers must be non-null"); } + auto const* dataset = + dynamic_cast const*>(&index->data()); + if (dataset == nullptr || dataset->n_rows() != static_cast(index->size())) { + return reject("every input must have an attached, uncompressed dataset"); + } + if (index->metric() != params.metric) { + return reject("every input metric must match index_params.metric"); + } + if (result.offsets.size() == 1) { + result.dim = static_cast(index->dim()); + } else if (result.dim != static_cast(index->dim())) { + return reject("all input dimensions must match"); + } + auto graph = index->graph(); + if (graph.extent(0) <= 0 || graph.extent(1) <= 0 || + graph.extent(0) != static_cast(index->size())) { + return reject("every input must have a nonempty device graph"); + } + + auto const input_rows = static_cast(index->size()); + if (rows > std::numeric_limits::max() - input_rows) { + return reject("the combined row count must fit in uint32_t"); + } + rows += input_rows; + max_input_degree = std::max(max_input_degree, static_cast(graph.extent(1))); + result.offsets.push_back(static_cast(rows)); + } + + if (result.dim <= 0 || result.dim > std::numeric_limits::max()) { + return reject("dataset dimension must be positive and fit cuBLAS int dimensions"); + } + if (!merge_scaffold::leaf_gemm_supported(result.dim, merge_params.leaf_size)) { + return reject("dataset dimension exceeds the leaf GEMM workspace limit"); + } + if (rows > std::numeric_limits::max() / spill) { + return reject("combined rows times the configured spill width must fit in uint32_t"); + } + if (params.graph_degree == 0 || static_cast(params.graph_degree) >= rows) { + return reject("graph_degree must be positive and smaller than the combined row count"); + } + if (static_cast(params.graph_degree) > max_input_degree + scaffold_degree) { + return reject("graph_degree exceeds the input graph plus scaffold capacity"); + } + + result.rows = static_cast(rows); + result.eligible = true; + return result; +} + +template +auto merge_fastener(raft::resources const& handle, + cagra::index_params const& params, + cagra::merge_params const& merge_params, + std::vector*>& indices, + fastener_preflight_result const& preflight) -> index +{ + auto dataset = [&] { + raft::common::nvtx::range scope("cagra::merge/consolidate"); + auto combined = raft::make_device_matrix(handle, preflight.rows, preflight.dim); + copy_input_datasets(handle, indices, preflight.offsets, preflight.dim, combined.data_handle()); + return combined; + }(); + + merge_scaffold::build_params scaffold_params; + scaffold_params.levels = merge_params.levels; + scaffold_params.root_fanout = merge_params.root_fanout; + scaffold_params.lower_fanout = merge_params.lower_fanout; + scaffold_params.leader_fraction = merge_params.leader_fraction; + scaffold_params.max_leaders = merge_params.max_leaders; + scaffold_params.leaf_size = merge_params.leaf_size; + scaffold_params.leaf_degree = merge_params.leaf_degree; + + auto merged_graph = [&] { + raft::common::nvtx::range scope("cagra::merge/scaffold"); + auto scaffold = merge_scaffold::build( + handle, raft::make_const_mdspan(dataset.view()), preflight.offsets, scaffold_params); + return merge_scaffold::append_to_input_graphs( + handle, indices, preflight.offsets, raft::make_const_mdspan(scaffold.view())); + }(); + + { + raft::common::nvtx::range scope("cagra::merge/append/sort"); + cagra::detail::graph::sort_knn_graph_device_inplace( + handle, params.metric, raft::make_const_mdspan(dataset.view()), merged_graph.view()); + merged_graph = merge_scaffold::cap_sorted_graph( + handle, raft::make_const_mdspan(merged_graph.view()), params.graph_degree); + } + + auto optimized_graph = + raft::make_device_matrix(handle, preflight.rows, params.graph_degree); + { + raft::common::nvtx::range scope("cagra::merge/optimize"); + cagra::detail::graph::optimize( + handle, merged_graph.view(), optimized_graph.view(), params.guarantee_connectivity); + } + + index merged_index(handle, params.metric); + merged_index.update_graph(handle, std::move(optimized_graph)); + if (params.attach_dataset_on_build) { + // CAGRA search assumes 16-byte row alignment (vectorized loads). Consolidate remains dense for + // the scaffold; attach through make_aligned_dataset so unaligned dims are padded. + merged_index.update_dataset(handle, make_aligned_dataset(handle, std::move(dataset), 16)); + } else { + using ds_idx_type = typename index::dataset_index_type; + merged_index.update_dataset( + handle, std::make_unique>(preflight.dim)); + } + + return merged_index; +} + +template +index merge(raft::resources const& handle, + cagra::index_params const& params, + std::vector*>& indices, + cagra::merge_params const& merge_params, + cuvs::neighbors::filtering::base_filter const& row_filter) +{ + raft::common::nvtx::range merge_scope( + "cagra::merge(algo=%d,parts=%zu)", static_cast(merge_params.algo), indices.size()); + + RAFT_EXPECTS(merge_params.algo == cagra::merge_algo::AUTO || + merge_params.algo == cagra::merge_algo::FASTENER || + merge_params.algo == cagra::merge_algo::REBUILD, + "Unknown cagra::merge algorithm"); + if (merge_params.algo == cagra::merge_algo::REBUILD) { + return merge_rebuild(handle, params, indices, row_filter); + } + + auto preflight = preflight_fastener(params, merge_params, indices, row_filter); + if (!preflight.eligible) { + if (merge_params.algo == cagra::merge_algo::AUTO) { + return merge_rebuild(handle, params, indices, row_filter); + } + RAFT_FAIL("FASTENER cagra::merge is unsupported: %s", preflight.reason.c_str()); + } + + return merge_fastener(handle, params, merge_params, indices, preflight); +} + +template +index merge(raft::resources const& handle, + cagra::index_params const& params, + std::vector*>& indices, + cuvs::neighbors::filtering::base_filter const& row_filter) +{ + return merge(handle, params, indices, cagra::merge_params{}, row_filter); +} + } // namespace cuvs::neighbors::cagra::detail diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh new file mode 100644 index 0000000000..e75db0ca33 --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -0,0 +1,1299 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ +#pragma once + +#include "../neighbors_device_intrinsics.cuh" + +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra::detail::merge_scaffold { + +// ---------------------------------------------------------------------------- +// Constants and size limits +// ---------------------------------------------------------------------------- + +inline constexpr uint32_t MAX_FANOUT = 32; +inline constexpr uint32_t MAX_LEADERS = 8192; +inline constexpr uint32_t MAX_LEAF_SIZE = 256; +inline constexpr int MAX_LEAF_DEGREE = 8; +inline constexpr int ASSIGNMENT_TILE_ROWS = 2048; +inline constexpr uint64_t DETERMINISTIC_SEED = 0x4c616e6472756d; +inline constexpr size_t GEMM_WORKSPACE_BYTES = size_t{2} * 1024 * 1024 * 1024; +inline constexpr int THREADS_PER_BLOCK = 256; +inline constexpr int MAX_STRIDED_GRID_BLOCKS = 1 << 20; +/** Warps per block in the one-row-per-warp kernels; their launch geometry derives from this. */ +inline constexpr int ROW_WARPS_PER_BLOCK = 4; + +/** Grid size for the grid-stride kernels: blocks of THREADS_PER_BLOCK covering `items`, capped so + * oversized workloads loop within resident threads instead. */ +inline auto strided_grid_size(int64_t items) -> int +{ + return static_cast(std::min((items + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK, + MAX_STRIDED_GRID_BLOCKS)); +} + +// ---------------------------------------------------------------------------- +// Build parameters and partition data structures +// ---------------------------------------------------------------------------- + +/** Internal controls for deterministic multi-level ball carving and leaf neighbor construction. */ +struct build_params { + uint32_t levels = 2; + uint32_t root_fanout = 2; + uint32_t lower_fanout = 3; + double leader_fraction = 0.02; + uint32_t max_leaders = 1000; + uint32_t leaf_size = MAX_LEAF_SIZE; + uint32_t leaf_degree = 4; +}; + +/** Controls for one invocation of the reusable many-way split boundary. */ +struct split_params { + uint32_t fanout = 1; + double leader_fraction = 0.02; + uint32_t max_leaders = 1000; + uint32_t leaf_size = MAX_LEAF_SIZE; + uint32_t level = 0; + uint32_t occurrence_stride = 1; +}; + +/** Device state and tuning knobs shared by every split level: the precomputed row norms, the + * tiling and workspace capacities, and the seed feeding the deterministic leader samples. */ +struct split_context { + split_context(rmm::cuda_stream_view stream, int64_t rows) + : norms(static_cast(rows), stream) + { + } + + rmm::device_uvector norms; + int assignment_tile_rows = ASSIGNMENT_TILE_ROWS; + size_t gemm_workspace_bytes = GEMM_WORKSPACE_BYTES; + uint64_t seed = DETERMINISTIC_SEED; +}; + +struct partition_membership { + uint32_t id = 0; + uint16_t occurrence = 0; + uint16_t padding = 0; +}; + +struct partition_range { + uint32_t key = 0; + int64_t start = 0; + int64_t end = 0; +}; + +struct partition_set { + rmm::device_uvector memberships; + std::vector ranges; +}; + +// ---------------------------------------------------------------------------- +// Many-way partitioning +// ---------------------------------------------------------------------------- + +/** Describes a contiguous tile of a parent partition to assign to child leaders. */ +struct assignment_tile { + int64_t input_start = 0; + int64_t group_start = 0; + int64_t group_size = 0; + int64_t output_start = 0; + int32_t rows = 0; + int32_t leader_count = 0; + uint32_t child_key_base = 0; + uint32_t leader_offset = 0; +}; + +/** Describes a completed parent partition copied unchanged into the next split level. */ +struct carry_span { + int64_t input_start = 0; + int64_t output_start = 0; + int32_t rows = 0; + uint32_t child_key = 0; +}; + +/** Leader count selection logic */ +inline int select_leader_count(int64_t rows, split_params const& params) +{ + auto sampled = + static_cast(std::ceil(params.leader_fraction * static_cast(rows))); + sampled = std::max(sampled, params.fanout); + sampled = std::min(sampled, params.max_leaders); + sampled = std::min(sampled, rows); + return static_cast(sampled); +} + +/** Return the bucket index for a leader count: ceil(log2(leaders)), so that + * `1 << leader_bucket_index(leaders)` is the smallest power of two covering it. Requires + * leaders >= 1. */ +inline int leader_bucket_index(int leaders) +{ + return static_cast(std::bit_width(static_cast(leaders - 1))); +} + +/** Host-side decision for one parent partition: where its children and output rows land. */ +struct parent_plan { + partition_range parent; + int64_t output_start = 0; + uint32_t child_key_base = 0; + int32_t leader_count = 0; // 0 => carried unchanged as one child + uint32_t leader_offset = 0; // deterministic leader sample start; meaningful only for splits + + auto carried() const -> bool { return leader_count == 0; } + auto size() const -> int64_t { return parent.end - parent.start; } + auto child_count() const -> uint32_t + { + return carried() ? 1 : static_cast(leader_count); + } + auto output_rows(uint32_t fanout) const -> int64_t + { + return carried() ? size() : size() * static_cast(fanout); + } +}; + +struct split_plan { + std::vector parents; + int64_t output_rows = 0; + uint32_t child_count = 0; +}; + +/** + * Decide on host how every parent partition maps into the next level. + * + * Parents within the leaf size are carried: their memberships pass through unchanged under one + * child key. Larger parents are split: a deterministic leader sample of `leader_fraction` of + * their rows, clamped to [fanout, max_leaders], is taken at evenly strided member positions + * starting from an offset hashed from the seed, level, and parent identity, so reruns select the + * same leaders. Child keys and output rows are dense in parent order. + */ +inline auto plan_split(std::vector const& ranges, + int64_t membership_count, + split_params const& params, + uint64_t seed) -> split_plan +{ + RAFT_EXPECTS(params.fanout >= 1 && params.fanout <= MAX_FANOUT, + "Fastener split fanout must be between 1 and %d", + MAX_FANOUT); + RAFT_EXPECTS(params.leader_fraction > 0.0 && params.leader_fraction <= 1.0, + "Fastener leader fraction must be in (0, 1]"); + RAFT_EXPECTS(params.max_leaders >= params.fanout && params.max_leaders <= MAX_LEADERS, + "Fastener leader cap must cover the fanout and not exceed %d", + MAX_LEADERS); + + split_plan plan; + plan.parents.reserve(ranges.size()); + int64_t covered = 0; + int64_t output_rows = 0; + int64_t child_keys = 0; + + for (size_t parent_index = 0; parent_index < ranges.size(); ++parent_index) { + auto const& parent = ranges[parent_index]; + RAFT_EXPECTS(parent.start == covered && parent.end > parent.start, + "Fastener parent ranges must compactly cover all memberships"); + covered = parent.end; + + parent_plan entry{.parent = parent, + .output_start = output_rows, + .child_key_base = static_cast(child_keys)}; + if (entry.size() > params.leaf_size) { + entry.leader_count = select_leader_count(entry.size(), params); + entry.leader_offset = static_cast( + cuvs::neighbors::detail::device::xorshift64( + seed ^ (static_cast(params.level) << 48) ^ + (static_cast(parent.key) << 1) ^ static_cast(parent_index)) % + static_cast(entry.size())); // mixing in all the relevant state + } + + child_keys += entry.child_count(); + output_rows += entry.output_rows(params.fanout); + plan.parents.push_back(entry); + } + RAFT_EXPECTS(covered == membership_count, + "Fastener parent ranges must compactly cover all memberships"); + RAFT_EXPECTS(output_rows <= static_cast(std::numeric_limits::max()), + "Fastener membership count must fit in uint32_t"); + RAFT_EXPECTS(child_keys <= static_cast(std::numeric_limits::max()), + "Fastener child key count must fit in uint32_t"); + plan.output_rows = output_rows; + plan.child_count = static_cast(child_keys); + return plan; +} + +/** Render the device-facing descriptor for one tile of a split parent. */ +inline auto make_tile(parent_plan const& entry, int64_t start, uint32_t fanout, int tile_rows) + -> assignment_tile +{ + return assignment_tile{ + .input_start = start, + .group_start = entry.parent.start, + .group_size = entry.size(), + .output_start = + entry.output_start + (start - entry.parent.start) * static_cast(fanout), + .rows = static_cast(std::min(tile_rows, entry.parent.end - start)), + .leader_count = entry.leader_count, + .child_key_base = entry.child_key_base, + .leader_offset = entry.leader_offset}; +} + +/** Group the split parents by padded leader count, so each bucket shares one GEMM shape. */ +inline auto bucket_split_parents(split_plan const& plan, split_params const& params) + -> std::vector> +{ + std::vector> buckets( + leader_bucket_index(static_cast(params.max_leaders)) + 1); + for (auto const& entry : plan.parents) { + if (!entry.carried()) { buckets[leader_bucket_index(entry.leader_count)].push_back(&entry); } + } + return buckets; +} + +/** Write the identity membership for each dataset row. */ +static __global__ void initialize_root_memberships_kernel(partition_membership* memberships, + int64_t rows) +{ + int64_t row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row < rows) { memberships[row] = {static_cast(row), uint16_t{0}, uint16_t{0}}; } +} + +/** Represent the full dataset as one identity partition without copying any vectors. */ +inline auto make_root_partition(raft::resources const& res, int64_t rows) -> partition_set +{ + auto stream = raft::resource::get_cuda_stream(res); + partition_set root{rmm::device_uvector(static_cast(rows), stream), + {{uint32_t{0}, int64_t{0}, rows}}}; + int blocks = static_cast((rows + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK); + initialize_root_memberships_kernel<<>>( + root.memberships.data(), rows); + RAFT_CUDA_TRY(cudaGetLastError()); + return root; +} + +/** Copy the memberships of each completed parent to the output and write its one child key. */ +static __global__ void carry_completed_parents_kernel(partition_membership const* input, + carry_span const* spans, + partition_membership* output, + uint32_t* output_keys) +{ + auto span = spans[blockIdx.x]; + for (int row = threadIdx.x; row < span.rows; row += blockDim.x) { + output[span.output_start + row] = input[span.input_start + row]; + output_keys[span.output_start + row] = span.child_key; + } +} + +/** Copy the vectors of each tile row into a dense float buffer. Unused rows become zero. */ +template +__global__ void manyway_gather_tile_points_kernel(T const* dataset, + int64_t dim, + partition_membership const* memberships, + assignment_tile const* tiles, + int batch_size, + int tile_rows, + float* output) +{ + int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + int64_t total = static_cast(batch_size) * tile_rows * dim; + for (; linear < total; linear += stride) { + int64_t d = linear % dim; + int64_t row = (linear / dim) % tile_rows; + int batch = static_cast(linear / (dim * tile_rows)); + auto tile = tiles[batch]; + float value = 0.0f; + if (row < tile.rows) { + uint32_t id = memberships[tile.input_start + row].id; + value = static_cast(dataset[static_cast(id) * dim + d]); + } + output[linear] = value; + } +} + +/** Copy the leader vectors of each tile into a dense float buffer and record the leader IDs. */ +template +__global__ void manyway_gather_tile_leaders_kernel(T const* dataset, + int64_t dim, + partition_membership const* memberships, + assignment_tile const* tiles, + int batch_size, + int padded_leaders, + float* output, + uint32_t* leader_ids) +{ + int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + int64_t total = static_cast(batch_size) * padded_leaders * dim; + for (; linear < total; linear += stride) { + int64_t d = linear % dim; + int leader = static_cast((linear / dim) % padded_leaders); + int batch = static_cast(linear / (dim * padded_leaders)); + auto tile = tiles[batch]; + float value = 0.0f; + if (leader < tile.leader_count) { + int64_t relative = (tile.leader_offset + + (static_cast(leader) * tile.group_size) / tile.leader_count) % + tile.group_size; + uint32_t id = memberships[tile.group_start + relative].id; + value = static_cast(dataset[static_cast(id) * dim + d]); + if (d == 0) { leader_ids[static_cast(batch) * padded_leaders + leader] = id; } + } + output[linear] = value; + } +} + +/** Convert the batched point-leader dot products into squared L2 distances in place. Invalid + * padded rows and leaders become infinity so the generic selection primitive ignores them. */ +static __global__ void manyway_materialize_tile_distances_kernel( + float* dots, + int batch_size, + int tile_rows, + int padded_leaders, + float const* norms, + uint32_t const* leader_ids, + partition_membership const* input_memberships, + assignment_tile const* tiles) +{ + int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + int64_t total = static_cast(batch_size) * tile_rows * padded_leaders; + for (; linear < total; linear += stride) { + int leader = static_cast(linear % padded_leaders); + int row = static_cast((linear / padded_leaders) % tile_rows); + int batch = static_cast(linear / (static_cast(padded_leaders) * tile_rows)); + auto tile = tiles[batch]; + float distance = std::numeric_limits::infinity(); + if (row < tile.rows && leader < tile.leader_count) { + auto membership = input_memberships[tile.input_start + row]; + auto leader_id = + leader_ids[static_cast(batch) * padded_leaders + static_cast(leader)]; + // This could omit the point norm because it is constant across every leader in a row and + // select_k only depends on their relative ordering. Doing so would also require removing the + // nonnegative clamp because the resulting ranking scores may legitimately be negative. + // I'm not doing this for the sake of simplicity and because it would be unlikely to have a + // performance impact; we still need those norms for sorting the edges by distance. + distance = fmaxf(0.0f, norms[membership.id] + norms[leader_id] - 2.0f * dots[linear]); + } + dots[linear] = distance; + } +} + +/** Write selected leader offsets as child keys and copy their point memberships. */ +static __global__ void manyway_emit_tile_assignments_kernel( + int const* selected_leaders, + int batch_size, + int tile_rows, + int fanout, + int occurrence_stride, + partition_membership const* input_memberships, + assignment_tile const* tiles, + uint32_t* output_keys, + partition_membership* output_memberships) +{ + int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + int64_t total = static_cast(batch_size) * tile_rows; + for (; linear < total; linear += stride) { + int row = static_cast(linear % tile_rows); + int batch = static_cast(linear / tile_rows); + auto tile = tiles[batch]; + if (row >= tile.rows) { continue; } + + auto membership = input_memberships[tile.input_start + row]; + int64_t output_base = tile.output_start + static_cast(row) * fanout; + int64_t selected_base = linear * fanout; + for (int rank = 0; rank < fanout; ++rank) { + auto leader = static_cast(selected_leaders[selected_base + rank]); + output_keys[output_base + rank] = tile.child_key_base + leader; + output_memberships[output_base + rank] = { + membership.id, + static_cast(membership.occurrence + rank * occurrence_stride), + uint16_t{0}}; + } + } +} + +/** Read the sorted keys and make one contiguous range for each unique partition key. */ +inline auto collect_partition_ranges(raft::resources const& res, + uint32_t const* keys, + int64_t count, + uint32_t key_count) -> std::vector +{ + auto stream = raft::resource::get_cuda_stream(res); + auto output_capacity = static_cast(std::min(count, key_count)); + rmm::device_uvector device_unique_keys(output_capacity, stream); + rmm::device_uvector device_counts(output_capacity, stream); + auto reduced_end = thrust::reduce_by_key(thrust::cuda::par.on(stream), + thrust::device_pointer_cast(keys), + thrust::device_pointer_cast(keys + count), + thrust::make_constant_iterator(uint32_t{1}), + thrust::device_pointer_cast(device_unique_keys.data()), + thrust::device_pointer_cast(device_counts.data())); + auto group_count = + static_cast(reduced_end.first - thrust::device_pointer_cast(device_unique_keys.data())); + + std::vector unique_keys(group_count); + std::vector counts(group_count); + raft::copy(unique_keys.data(), device_unique_keys.data(), group_count, stream); + raft::copy(counts.data(), device_counts.data(), group_count, stream); + raft::resource::sync_stream(res); + + std::vector groups; + groups.reserve(group_count); + int64_t cursor = 0; + for (size_t index = 0; index < group_count; ++index) { + RAFT_EXPECTS(unique_keys[index] < key_count, "Many-way group key is out of range"); + groups.push_back({unique_keys[index], cursor, cursor + counts[index]}); + cursor += counts[index]; + } + RAFT_EXPECTS(cursor == count, "Many-way partition membership histogram lost entries"); + return groups; +} + +/** Copy every carried parent's memberships unchanged. */ +inline void carry_parents(raft::resources const& res, + partition_set const& parents, + split_plan const& plan, + rmm::device_uvector& keys, + rmm::device_uvector& memberships) +{ + std::vector carries; + for (auto const& entry : plan.parents) { + if (entry.carried()) { + carries.push_back({entry.parent.start, + entry.output_start, + static_cast(entry.size()), + entry.child_key_base}); + } + } + if (carries.empty()) { return; } + + auto stream = raft::resource::get_cuda_stream(res); + rmm::device_uvector device_carries(carries.size(), stream); + raft::copy(device_carries.data(), carries.data(), carries.size(), stream); + carry_completed_parents_kernel<<(carries.size()), + THREADS_PER_BLOCK, + 0, + stream>>>( + parents.memberships.data(), device_carries.data(), memberships.data(), keys.data()); + RAFT_CUDA_TRY(cudaGetLastError()); +} + +/** Compute every pairwise dot product between the rows of A_i and the rows of B_i for a strided + * batch of row-major matrices with rows of length `row_width`: + * out_i[b * a_rows + a] = A_i[a] . B_i[b]. + * + * Pretty much a wrapper around `cublasGemmStridedBatchedEx`. */ +/** Float strided-batched GEMM for row-wise dots. Integer datasets are gathered to float first; + * native INT8 cuBLAS paths are not portable across architectures (e.g. Ada returns + * CUBLAS_STATUS_NOT_SUPPORTED for CUDA_R_8I / CUBLAS_COMPUTE_32I strided-batched GEMM). */ +inline void batched_row_dot_products(raft::resources const& res, + float const* a, + int a_rows, + long long a_stride, + float const* b, + int b_rows, + long long b_stride, + float* out, + long long out_stride, + int row_width, + int batch_count) +{ + float alpha = 1.0f; + float beta = 0.0f; + auto cublas_handle = raft::resource::get_cublas_handle(res); + RAFT_CUBLAS_TRY(cublasSetPointerMode(cublas_handle, CUBLAS_POINTER_MODE_HOST)); + RAFT_CUBLAS_TRY(cublasGemmStridedBatchedEx(cublas_handle, + CUBLAS_OP_T, + CUBLAS_OP_N, + a_rows, + b_rows, + row_width, + &alpha, + a, + CUDA_R_32F, + row_width, + a_stride, + b, + CUDA_R_32F, + row_width, + b_stride, + &beta, + out, + CUDA_R_32F, + a_rows, + out_stride, + batch_count, + CUBLAS_COMPUTE_32F, + CUBLAS_GEMM_DEFAULT)); +} + +/** + * Assign every row of one bucket's split parents to its `fanout` nearest leaders. + * + * Parents are cut into tiles of `assignment_tile_rows` rows; every tile in a bucket shares the + * same padded leader count, so one strided batched GEMM per batch produces all point-leader dot + * products. Tiles are processed in batches sized to the GEMM workspace: each batch gathers its + * tile vectors and its parents' leader vectors into dense float buffers, converts dots to + * distances with the precomputed row norms (|x|^2 + |l|^2 - 2 x.l), uses `select_k` to keep the + * `fanout` nearest leaders per row, and writes the child keys and memberships. + */ +template +void assign_bucket(raft::resources const& res, + raft::device_matrix_view dataset, + partition_set const& parents, + std::vector const& bucket, + int padded_leaders, + split_params const& params, + split_context& context, + rmm::device_uvector& keys, + rmm::device_uvector& memberships) +{ + auto stream = raft::resource::get_cuda_stream(res); + auto dim = dataset.extent(1); + + // Cut every parent in the bucket into fixed-height tiles + int tile_rows = context.assignment_tile_rows; + std::vector tiles; + for (auto const* entry : bucket) { + for (int64_t start = entry->parent.start; start < entry->parent.end; start += tile_rows) { + tiles.push_back(make_tile(*entry, start, params.fanout, tile_rows)); + } + } + + // Size the batch so all per-tile buffers fit in the GEMM workspace budget + size_t point_elements = static_cast(tile_rows) * dim; + size_t leader_elements = static_cast(padded_leaders) * dim; + size_t dot_elements = static_cast(tile_rows) * padded_leaders; + size_t selected_elements = static_cast(tile_rows) * params.fanout; + size_t bytes_per_batch = + (point_elements + leader_elements + dot_elements + selected_elements) * sizeof(float) + + static_cast(padded_leaders) * sizeof(uint32_t) + selected_elements * sizeof(int) + + sizeof(assignment_tile); + RAFT_EXPECTS(bytes_per_batch <= context.gemm_workspace_bytes, + "Fastener assignment workspace is too small for one tile"); + size_t batch_capacity = std::max( + 1, std::min(tiles.size(), context.gemm_workspace_bytes / bytes_per_batch)); + + rmm::device_uvector device_tiles(batch_capacity, stream); + rmm::device_uvector tile_points(batch_capacity * point_elements, stream); + rmm::device_uvector tile_leaders(batch_capacity * leader_elements, stream); + rmm::device_uvector tile_dots(batch_capacity * dot_elements, stream); + rmm::device_uvector tile_leader_ids(batch_capacity * padded_leaders, stream); + rmm::device_uvector selected_distances(batch_capacity * selected_elements, stream); + rmm::device_uvector selected_leaders(batch_capacity * selected_elements, stream); + + auto point_stride = static_cast(point_elements); + auto leader_stride = static_cast(leader_elements); + auto dot_stride = static_cast(dot_elements); + + for (size_t tile_offset = 0; tile_offset < tiles.size(); tile_offset += batch_capacity) { + size_t batch_size = std::min(batch_capacity, tiles.size() - tile_offset); + raft::copy(device_tiles.data(), tiles.data() + tile_offset, batch_size, stream); + + // Gather the batch's tile rows and leader vectors + int point_blocks = strided_grid_size(static_cast(batch_size * point_elements)); + manyway_gather_tile_points_kernel<<>>( + dataset.data_handle(), + dim, + parents.memberships.data(), + device_tiles.data(), + static_cast(batch_size), + tile_rows, + tile_points.data()); + RAFT_CUDA_TRY(cudaGetLastError()); + + int leader_blocks = strided_grid_size(static_cast(batch_size * leader_elements)); + manyway_gather_tile_leaders_kernel<<>>( + dataset.data_handle(), + dim, + parents.memberships.data(), + device_tiles.data(), + static_cast(batch_size), + padded_leaders, + tile_leaders.data(), + tile_leader_ids.data()); + RAFT_CUDA_TRY(cudaGetLastError()); + + // point-leader dot products for the batch + batched_row_dot_products(res, + tile_leaders.data(), + padded_leaders, + leader_stride, + tile_points.data(), + tile_rows, + point_stride, + tile_dots.data(), + dot_stride, + static_cast(dim), + static_cast(batch_size)); + + // Materialize distances, keep each row's nearest leaders, and emit their memberships + int64_t selection_rows = static_cast(batch_size) * tile_rows; + int distance_blocks = strided_grid_size(selection_rows * padded_leaders); + manyway_materialize_tile_distances_kernel<<>>( + tile_dots.data(), + static_cast(batch_size), + tile_rows, + padded_leaders, + context.norms.data(), + tile_leader_ids.data(), + parents.memberships.data(), + device_tiles.data()); + RAFT_CUDA_TRY(cudaGetLastError()); + + cuvs::selection::select_k(res, + raft::make_device_matrix_view( + tile_dots.data(), selection_rows, padded_leaders), + std::nullopt, + raft::make_device_matrix_view( + selected_distances.data(), selection_rows, params.fanout), + raft::make_device_matrix_view( + selected_leaders.data(), selection_rows, params.fanout), + true, + true); + + int assignment_blocks = strided_grid_size(selection_rows); + manyway_emit_tile_assignments_kernel<<>>( + selected_leaders.data(), + static_cast(batch_size), + tile_rows, + static_cast(params.fanout), + static_cast(params.occurrence_stride), + parents.memberships.data(), + device_tiles.data(), + keys.data(), + memberships.data()); + RAFT_CUDA_TRY(cudaGetLastError()); + } +} + +/** + * Split every oversized parent into overlapping nearest-leader children. + * + * All levels, including the root, traverse this boundary. `plan_split` decides on host how every + * parent maps into the next level, `carry_parents` copies completed parents unchanged, and + * `assign_bucket` assigns each bucket of split parents with tiled batched GEMMs. One input row of + * a split parent emits `fanout` memberships, and each membership's occurrence advances by + * `rank * occurrence_stride` so that across levels every copy of a row lands in a distinct + * scaffold slot. + * + * Child keys are dense and sequential across the whole output, and the final stable sort by key + * followed by a reduce-by-key compaction yields one contiguous range per child, which is the next + * level's partition set. + */ +template +auto split_manyway(raft::resources const& res, + raft::device_matrix_view dataset, + partition_set const& parents, + split_params const& params, + split_context& context) -> partition_set +{ + auto stream = raft::resource::get_cuda_stream(res); + auto plan = plan_split( + parents.ranges, static_cast(parents.memberships.size()), params, context.seed); + + rmm::device_uvector keys(static_cast(plan.output_rows), stream); + rmm::device_uvector memberships(static_cast(plan.output_rows), + stream); + + carry_parents(res, parents, plan, keys, memberships); + auto buckets = bucket_split_parents(plan, params); + for (size_t bucket_index = 0; bucket_index < buckets.size(); ++bucket_index) { + if (buckets[bucket_index].empty()) { continue; } + assign_bucket(res, + dataset, + parents, + buckets[bucket_index], + 1 << bucket_index, + params, + context, + keys, + memberships); + } + + thrust::stable_sort_by_key(thrust::cuda::par.on(stream), + thrust::device_pointer_cast(keys.data()), + thrust::device_pointer_cast(keys.data() + keys.size()), + thrust::device_pointer_cast(memberships.data())); + auto ranges = collect_partition_ranges(res, keys.data(), plan.output_rows, plan.child_count); + return {std::move(memberships), std::move(ranges)}; +} + +// ---------------------------------------------------------------------------- +// Leaf processing: bounded leaf slicing and cross-input leaf KNN +// ---------------------------------------------------------------------------- + +/** Return true if one leaf of this dimension and leaf size fits the float GEMM workspace. */ +template +inline auto leaf_gemm_supported(int64_t dimension, uint32_t leaf_size) -> bool +{ + // All scalar types (including integers) gather into float leaf buffers before the GEMM. + (void)sizeof(T); + if (dimension <= 0 || dimension > std::numeric_limits::max()) { return false; } + + size_t vector_elements = static_cast(leaf_size) * static_cast(dimension); + size_t gram_elements = static_cast(leaf_size) * leaf_size; + return vector_elements * sizeof(float) + gram_elements * sizeof(float) <= GEMM_WORKSPACE_BYTES; +} + +struct leaf_set { + partition_set const* partitions = nullptr; + std::vector starts_host; + std::vector ends_host; + rmm::device_uvector starts; + rmm::device_uvector ends; +}; + +/** Divide final grouped partitions into consecutive bounded leaves, without geometric resplitting. + * + * This is really just a fallback for when the tree isn't deep enough, and will produce obviously + * inferior leaves but at no additional cost. + */ +inline auto make_leaves(raft::resources const& res, + partition_set const& partitions, + uint32_t leaf_size) -> leaf_set +{ + auto stream = raft::resource::get_cuda_stream(res); + std::vector starts_host; + std::vector ends_host; + starts_host.reserve((partitions.memberships.size() + leaf_size - 1) / leaf_size); + ends_host.reserve(starts_host.capacity()); + + int64_t covered = 0; + for (auto const& range : partitions.ranges) { + RAFT_EXPECTS(range.start == covered && range.end > range.start && + range.end <= static_cast(partitions.memberships.size()), + "Fastener partition ranges must compactly cover all memberships"); + for (int64_t start = range.start; start < range.end; start += leaf_size) { + starts_host.push_back(static_cast(start)); + ends_host.push_back(static_cast( + std::min(range.end, start + static_cast(leaf_size)))); + } + covered = range.end; + } + RAFT_EXPECTS( + covered == static_cast(partitions.memberships.size()) && !starts_host.empty(), + "Fastener partition ranges did not cover all memberships"); + + rmm::device_uvector starts(starts_host.size(), stream); + rmm::device_uvector ends(ends_host.size(), stream); + raft::copy(starts.data(), starts_host.data(), starts.size(), stream); + raft::copy(ends.data(), ends_host.data(), ends.size(), stream); + return { + &partitions, std::move(starts_host), std::move(ends_host), std::move(starts), std::move(ends)}; +} + +/** Fill every unwritten scaffold slot with its row ID; final prefix deduplication removes it. */ +static __global__ void initialize_self_scaffold_kernel(uint32_t* graph, + int64_t rows, + int graph_degree) +{ + int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t total = rows * graph_degree; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + for (; linear < total; linear += stride) { + graph[linear] = static_cast(linear / graph_degree); + } +} + +/** Copy the vectors of each leaf into a dense buffer of OutT, zero-padding rows past the leaf + * end and dimensions past `input_dim`. Every scalar type is promoted to OutT (float) as-is. */ +template +__global__ void manyway_gather_leaf_vectors_kernel(T const* dataset, + int64_t input_dim, + int64_t output_dim, + int leaf_size, + partition_membership const* memberships, + uint32_t const* leaf_starts, + uint32_t const* leaf_ends, + int64_t leaf_offset, + int64_t leaf_count, + OutT* leaf_vectors) +{ + int64_t linear = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + int64_t stride = static_cast(blockDim.x) * gridDim.x; + int64_t total = leaf_count * leaf_size * output_dim; + for (; linear < total; linear += stride) { + int64_t d = linear % output_dim; + int64_t local_row = (linear / output_dim) % leaf_size; + int64_t local_leaf = linear / (output_dim * leaf_size); + int64_t leaf = leaf_offset + local_leaf; + int64_t leaf_n = static_cast(leaf_ends[leaf] - leaf_starts[leaf]); + OutT value = 0; + if (local_row < leaf_n && d < input_dim) { + uint32_t point = memberships[leaf_starts[leaf] + local_row].id; + auto input = dataset[static_cast(point) * input_dim + d]; + value = static_cast(input); + } + leaf_vectors[linear] = value; + } +} + +/** Select the nearest cross-input neighbors of each point from the leaf Gram matrix. One block + * does one leaf. One thread does one point. */ +template +static __global__ void manyway_leaf_gram_knn_kernel(GramT const* gram, + partition_membership const* memberships, + uint32_t const* origins, + uint32_t const* leaf_starts, + uint32_t const* leaf_ends, + int64_t leaf_offset, + int64_t leaf_count, + int leaf_size, + int leaf_degree, + int union_degree, + uint32_t* graph) +{ + int64_t local_leaf = blockIdx.x; + if (local_leaf >= leaf_count) { return; } + int64_t leaf = leaf_offset + local_leaf; + uint32_t start = leaf_starts[leaf]; + int leaf_n = static_cast(leaf_ends[leaf] - start); + if (leaf_n <= 1 || leaf_n > MAX_LEAF_SIZE) { return; } + + // Stage the leaf's memberships and origin labels once per block. + __shared__ partition_membership records[MAX_LEAF_SIZE]; + __shared__ uint32_t leaf_origins[MAX_LEAF_SIZE]; + for (int i = threadIdx.x; i < leaf_n; i += blockDim.x) { + records[i] = memberships[start + i]; + leaf_origins[i] = origins[records[i].id]; + } + __syncthreads(); + + int u = threadIdx.x; + if (u >= leaf_n) { return; } + double top_d[MAX_LEAF_DEGREE]; + uint16_t top_v[MAX_LEAF_DEGREE]; + for (int t = 0; t < leaf_degree; ++t) { + top_d[t] = std::numeric_limits::max(); + top_v[t] = std::numeric_limits::max(); + } + + // Scan this point's Gram row, keeping the `leaf_degree` nearest cross-origin neighbors. + int64_t gram_base = local_leaf * leaf_size * leaf_size; + double norm_u = static_cast(gram[gram_base + u * leaf_size + u]); + for (int v = 0; v < leaf_n; ++v) { + if (u == v || leaf_origins[u] == leaf_origins[v] || records[u].id == records[v].id) { + continue; + } + double norm_v = static_cast(gram[gram_base + v * leaf_size + v]); + double dot = static_cast(gram[gram_base + v * leaf_size + u]); + double distance = norm_u + norm_v - 2.0 * dot; + if (isfinite(distance)) { distance = fmax(0.0, distance); } + int worst = 0; + for (int t = 1; t < leaf_degree; ++t) { + if (top_d[t] > top_d[worst] || (top_d[t] == top_d[worst] && top_v[t] > top_v[worst])) { + worst = t; + } + } + if (distance < top_d[worst] || + (distance == top_d[worst] && records[v].id < records[top_v[worst]].id)) { + top_d[worst] = distance; + top_v[worst] = static_cast(v); + } + } + + // Pack the valid selections into this occurrence's slots of the scaffold graph. + int selected = 0; + for (int t = 0; t < leaf_degree; ++t) { + bool valid = top_v[t] != std::numeric_limits::max() && isfinite(top_d[t]); + if (valid) { + int64_t output = static_cast(records[u].id) * union_degree + + static_cast(records[u].occurrence) * leaf_degree + selected; + graph[output] = records[top_v[t]].id; + ++selected; + } + } +} + +/** Build directed cross-input nearest neighbors for every leaf occurrence. */ +template +auto build_leaf_neighbors(raft::resources const& res, + raft::device_matrix_view dataset, + leaf_set const& leaves, + uint32_t const* origins, + int union_degree, + build_params const& params) -> raft::device_matrix +{ + auto stream = raft::resource::get_cuda_stream(res); + int64_t rows = dataset.extent(0); + int leaf_size = static_cast(params.leaf_size); + int leaf_degree = static_cast(params.leaf_degree); + auto const& memberships = leaves.partitions->memberships; + + // Prefill every slot with its own row id, leaf KNN overwrites the slots it fills + auto graph = raft::make_device_matrix(res, rows, union_degree); + int scaffold_blocks = strided_grid_size(rows * union_degree); + initialize_self_scaffold_kernel<<>>( + graph.data_handle(), rows, union_degree); + RAFT_CUDA_TRY(cudaGetLastError()); + + // Gather every scalar type into float leaf buffers. Native INT8 cuBLAS is not portable across + // architectures (Ada returns CUBLAS_STATUS_NOT_SUPPORTED for the strided-batched INT8 path). + int64_t input_dimension = dataset.extent(1); + int dimension = static_cast(input_dimension); + size_t vector_elements_per_leaf = static_cast(leaf_size) * dimension; + size_t gram_elements_per_leaf = static_cast(leaf_size) * leaf_size; + size_t bytes_per_leaf = + vector_elements_per_leaf * sizeof(float) + gram_elements_per_leaf * sizeof(float); + size_t batch_capacity = std::max( + 1, std::min(leaves.starts_host.size(), GEMM_WORKSPACE_BYTES / bytes_per_leaf)); + rmm::device_uvector leaf_vectors(batch_capacity * vector_elements_per_leaf, stream); + rmm::device_uvector gram(batch_capacity * gram_elements_per_leaf, stream); + auto vector_stride = static_cast(vector_elements_per_leaf); + auto gram_stride = static_cast(gram_elements_per_leaf); + + for (size_t leaf_offset = 0; leaf_offset < leaves.starts_host.size(); + leaf_offset += batch_capacity) { + size_t batch_size = std::min(batch_capacity, leaves.starts_host.size() - leaf_offset); + int gather_blocks = + strided_grid_size(static_cast(batch_size * vector_elements_per_leaf)); + manyway_gather_leaf_vectors_kernel<<>>( + dataset.data_handle(), + input_dimension, + input_dimension, + leaf_size, + memberships.data(), + leaves.starts.data(), + leaves.ends.data(), + static_cast(leaf_offset), + static_cast(batch_size), + leaf_vectors.data()); + RAFT_CUDA_TRY(cudaGetLastError()); + batched_row_dot_products(res, + leaf_vectors.data(), + leaf_size, + vector_stride, + leaf_vectors.data(), + leaf_size, + vector_stride, + gram.data(), + gram_stride, + dimension, + static_cast(batch_size)); + manyway_leaf_gram_knn_kernel + <<(batch_size), leaf_size, 0, stream>>>(gram.data(), + memberships.data(), + origins, + leaves.starts.data(), + leaves.ends.data(), + static_cast(leaf_offset), + static_cast(batch_size), + leaf_size, + leaf_degree, + union_degree, + graph.data_handle()); + RAFT_CUDA_TRY(cudaGetLastError()); + } + + return graph; +} + +// ---------------------------------------------------------------------------- +// Scaffold build driver +// ---------------------------------------------------------------------------- + +/** Initialize the source-index label for every dataset row. */ +static __global__ void initialize_origins_kernel(uint32_t* origins, + int64_t start, + int64_t rows, + uint32_t origin) +{ + int64_t local_row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (local_row >= rows) { return; } + origins[start + local_row] = origin; +} + +/** Calculate the squared L2 norm of each dataset row. + * + * This could be `raft::linalg::norm`, but half precision vectors might square to inf without a + * wider accumulator. + */ +template +__global__ void manyway_l2_norms_kernel(T const* dataset, int64_t rows, int64_t dim, float* norms) +{ + int lane = threadIdx.x % raft::WarpSize; + int warp = threadIdx.x / raft::WarpSize; + int64_t row = static_cast(blockIdx.x) * ROW_WARPS_PER_BLOCK + warp; + if (row >= rows) { return; } + + // Each lane accumulates a strided slice of the row's squared values. + float sum = 0.0f; + T const* point = dataset + row * dim; + for (int64_t d = lane; d < dim; d += raft::WarpSize) { + float value = static_cast(point[d]); + sum = fmaf(value, value, sum); + } + // Shuffle-reduce the partial sums; lane 0 holds the total. + for (int offset = raft::WarpSize / 2; offset > 0; offset /= 2) { + sum += __shfl_down_sync(0xffffffffu, sum, offset); + } + if (lane == 0) { norms[row] = sum; } +} + +/** Build the many-way scaffold through splitting and leaf construction. */ +template +auto build(raft::resources const& res, + raft::device_matrix_view dataset, + std::vector const& offsets, + build_params const& params = {}) -> raft::device_matrix +{ + auto stream = raft::resource::get_cuda_stream(res); + int64_t rows = dataset.extent(0); + + RAFT_EXPECTS(offsets.size() >= 3, "Fastener requires at least two input datasets"); + RAFT_EXPECTS(rows > 0, "Fastener row count must be positive"); + RAFT_EXPECTS(params.levels > 0, "Fastener levels must be positive"); + RAFT_EXPECTS(params.root_fanout >= 1 && params.root_fanout <= MAX_FANOUT && + params.lower_fanout >= 1 && params.lower_fanout <= MAX_FANOUT, + "Fastener fanouts must be between 1 and %u", + MAX_FANOUT); + RAFT_EXPECTS(params.leader_fraction > 0.0 && params.leader_fraction <= 1.0, + "Fastener leader fraction must be in (0, 1]"); + RAFT_EXPECTS(params.max_leaders >= std::max(params.root_fanout, params.lower_fanout) && + params.max_leaders <= MAX_LEADERS, + "Fastener leader cap must cover both fanouts and not exceed %u", + MAX_LEADERS); + RAFT_EXPECTS(params.leaf_size >= 1 && params.leaf_size <= MAX_LEAF_SIZE, + "Fastener leaf size must be between 1 and %d", + MAX_LEAF_SIZE); + RAFT_EXPECTS(params.leaf_degree >= 1 && params.leaf_degree <= MAX_LEAF_DEGREE, + "Fastener leaf degree must be between 1 and %d", + MAX_LEAF_DEGREE); + RAFT_EXPECTS(leaf_gemm_supported(dataset.extent(1), params.leaf_size), + "Fastener dataset dimension exceeds the leaf GEMM limits"); + + // the number of leaf partitions that a given point will end up in + uint64_t spill = params.root_fanout; + for (uint32_t level = 1; level < params.levels; ++level) { + RAFT_EXPECTS(spill <= std::numeric_limits::max() / params.lower_fanout, + "Fastener spill width overflow"); + spill *= params.lower_fanout; + } + // this is because the degree of the candidate list is stored in uint8_t + RAFT_EXPECTS(spill * params.leaf_degree <= std::numeric_limits::max(), + "Fastener candidate width must not exceed %u", + static_cast(std::numeric_limits::max())); + RAFT_EXPECTS(static_cast(rows) <= std::numeric_limits::max() / spill, + "Fastener total partition memberships (rows=%ld * spill=%lu) must fit in uint32_t", + static_cast(rows), + spill); + int union_degree = static_cast(spill * params.leaf_degree); + + // Per-row index of the input partition this point came from, used to skip same-origin pairs + rmm::device_uvector origins(static_cast(rows), stream); + + for (size_t part = 0; part + 1 < offsets.size(); ++part) { + int64_t part_rows = offsets[part + 1] - offsets[part]; + int blocks = static_cast((part_rows + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK); + initialize_origins_kernel<<>>( + origins.data(), offsets[part], part_rows, static_cast(part)); + RAFT_CUDA_TRY(cudaGetLastError()); + } + + split_context context(stream, rows); + int norm_blocks = static_cast((rows + ROW_WARPS_PER_BLOCK - 1) / ROW_WARPS_PER_BLOCK); + manyway_l2_norms_kernel<<>>( + dataset.data_handle(), rows, dataset.extent(1), context.norms.data()); + RAFT_CUDA_TRY(cudaGetLastError()); + + // the actual splitting + auto partitions = make_root_partition(res, rows); + uint32_t occurrence_stride = 1; + for (uint32_t level = 0; level < params.levels; ++level) { + uint32_t fanout = level == 0 ? params.root_fanout : params.lower_fanout; + + partitions = split_manyway(res, + dataset, + partitions, + split_params{.fanout = fanout, + .leader_fraction = params.leader_fraction, + .max_leaders = params.max_leaders, + .leaf_size = params.leaf_size, + .level = level, + .occurrence_stride = occurrence_stride}, + context); + occurrence_stride *= fanout; + } + + // Leaf construction: only consecutive range slicing occurs after configured geometric depth. + auto leaves = make_leaves(res, partitions, params.leaf_size); + auto scaffold = build_leaf_neighbors(res, dataset, leaves, origins.data(), union_degree, params); + raft::resource::sync_stream(res); + return scaffold; +} + +// ---------------------------------------------------------------------------- +// Candidate graph assembly: combine input graphs with the scaffold +// ---------------------------------------------------------------------------- + +/** + * Offset-copy one input partition graph and append its global scaffold neighbors. + */ +static __global__ void copy_partition_with_scaffold_kernel(uint32_t const* source, + uint32_t const* scaffold, + int64_t source_rows, + int64_t source_degree, + uint32_t* destination, + int64_t destination_degree, + int64_t base_degree, + int64_t scaffold_degree, + uint32_t offset) +{ + int64_t row = static_cast(blockIdx.x) * blockDim.x + threadIdx.x; + if (row >= source_rows) { return; } + int64_t source_base = row * source_degree; + int64_t global_row = row + offset; + int64_t destination_base = global_row * destination_degree; + + // Rows from lower-degree input graphs are cyclically repeated to the common base width, which is + // not ideal (adds work to the sorting by distance) but practical applications for mismatched + // degrees seem hard to imagine. + for (int64_t j = 0; j < base_degree; ++j) { + destination[destination_base + j] = source[source_base + (j % source_degree)] + offset; + } + for (int64_t j = 0; j < scaffold_degree; ++j) { + destination[destination_base + base_degree + j] = scaffold[global_row * scaffold_degree + j]; + } +} + +/** + * @brief Form the pre-optimization candidate graph from input graphs and scaffold edges. + * + * The maximum input graph degree defines the base width, allowing partitions with mixed degrees. + */ +template +auto append_to_input_graphs( + raft::resources const& res, + std::vector*> const& indices, + std::vector const& offsets, + raft::device_matrix_view scaffold) + -> raft::device_matrix +{ + auto stream = raft::resource::get_cuda_stream(res); + int64_t base_degree = 0; + for (auto const* index : indices) { + base_degree = std::max(base_degree, index->graph_degree()); + } + int64_t scaffold_degree = scaffold.extent(1); + int64_t graph_degree = base_degree + scaffold_degree; + auto graph = raft::make_device_matrix(res, scaffold.extent(0), graph_degree); + + for (size_t part = 0; part < indices.size(); ++part) { + auto source = indices[part]->graph(); + RAFT_EXPECTS(source.extent(1) > 0, "Input CAGRA graphs must have nonzero degree"); + int blocks = static_cast((source.extent(0) + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK); + copy_partition_with_scaffold_kernel<<>>( + source.data_handle(), + scaffold.data_handle(), + source.extent(0), + source.extent(1), + graph.data_handle(), + graph_degree, + base_degree, + scaffold_degree, + static_cast(offsets[part])); + RAFT_CUDA_TRY(cudaGetLastError()); + } + raft::resource::sync_stream(res); + return graph; +} + +/** + * Remove invalid, self, and duplicate candidates from each sorted row, retain the nearest prefix, + * and cyclically pad short rows to the requested output width. + * + * Scaffold padding is measured to be negligible compared to the input graph edges. + */ +static __global__ void deduplicate_graph_prefix_kernel(uint32_t const* input, + int64_t rows, + int64_t input_degree, + uint32_t* output, + int64_t output_degree) +{ + constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / raft::WarpSize; + int lane = threadIdx.x % raft::WarpSize; + int warp = threadIdx.x / raft::WarpSize; + int64_t row = static_cast(blockIdx.x) * WARPS_PER_BLOCK + warp; + if (row >= rows) { return; } + + int64_t input_base = row * input_degree; + int64_t output_base = row * output_degree; + int selected = 0; + for (int64_t tile = 0; tile < input_degree && selected < output_degree; tile += raft::WarpSize) { + int64_t column = tile + lane; + bool first = column < input_degree; + uint32_t candidate = first ? input[input_base + column] : uint32_t{0}; + first = first && candidate < rows && candidate != static_cast(row); + for (int64_t prior = 0; prior < column && first; ++prior) { + if (input[input_base + prior] == candidate) { first = false; } + } + + unsigned first_mask = __ballot_sync(0xffffffffu, first); + unsigned lower_mask = lane == 0 ? 0u : (0xffffffffu >> (raft::WarpSize - lane)); + int output_column = selected + __popc(first_mask & lower_mask); + if (first && output_column < output_degree) { output[output_base + output_column] = candidate; } + selected += __popc(first_mask); + } + + if (selected == 0) { + if (lane == 0) { output[output_base] = static_cast((row + 1) % rows); } + selected = 1; + } + if (selected > output_degree) { selected = static_cast(output_degree); } + for (int64_t column = selected + lane; column < output_degree; column += raft::WarpSize) { + output[output_base + column] = output[output_base + (column % selected)]; + } +} + +/** Deduplicate a metric-sorted graph and retain a fixed-width nearest-candidate prefix. */ +inline auto cap_sorted_graph( + raft::resources const& res, + raft::device_matrix_view graph, + int64_t output_degree) -> raft::device_matrix +{ + RAFT_EXPECTS(output_degree > 0 && output_degree <= graph.extent(1), + "Pre-optimize graph degree cap must be within the sorted graph degree"); + auto output = raft::make_device_matrix(res, graph.extent(0), output_degree); + constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / raft::WarpSize; + int blocks = static_cast((graph.extent(0) + WARPS_PER_BLOCK - 1) / WARPS_PER_BLOCK); + deduplicate_graph_prefix_kernel<<>>( + graph.data_handle(), graph.extent(0), graph.extent(1), output.data_handle(), output_degree); + RAFT_CUDA_TRY(cudaGetLastError()); + return output; +} + +} // namespace cuvs::neighbors::cagra::detail::merge_scaffold diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 52b4542798..5f6c1c499a 100644 --- a/cpp/src/neighbors/detail/cagra/graph_core.cuh +++ b/cpp/src/neighbors/detail/cagra/graph_core.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2023-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once @@ -37,6 +37,7 @@ #include #include +#include #include #include @@ -174,6 +175,30 @@ __global__ void kern_sort(const DATA_T* const dataset, // [dataset_chunk_size, } } +// kern_sort capacity is WarpSize * numElementsPerThread; largest specialization uses 32. +constexpr int kMaxSortElementsPerThread = 32; +constexpr uint64_t kMaxSortDegree = + static_cast(raft::WarpSize) * kMaxSortElementsPerThread; + +template +using sort_kernel_type = + void (*)(DataT const*, IdxT, uint32_t, IdxT*, uint32_t, uint32_t, cuvs::distance::DistanceType); + +template +auto select_sort_kernel(uint64_t degree) -> sort_kernel_type +{ + if (degree <= raft::WarpSize * 1) { return kern_sort; } + if (degree <= raft::WarpSize * 2) { return kern_sort; } + if (degree <= raft::WarpSize * 4) { return kern_sort; } + if (degree <= raft::WarpSize * 8) { return kern_sort; } + if (degree <= raft::WarpSize * 16) { return kern_sort; } + if (degree <= kMaxSortDegree) { return kern_sort; } + RAFT_FAIL( + "The degree of input knn graph is too large (%lu). It must be equal to or smaller than %lu.", + degree, + kMaxSortDegree); +} + template __global__ void kern_make_rev_graph_k( OutputMatrixView output_graph, // [graph_size, degree] @@ -971,6 +996,52 @@ void make_reverse_graph_gpu( } } +/** Sort a device-resident graph in place without staging the dataset or graph. */ +template +void sort_knn_graph_device_inplace( + raft::resources const& res, + cuvs::distance::DistanceType metric, + raft::device_matrix_view dataset, + raft::device_matrix_view knn_graph) +{ + RAFT_EXPECTS(dataset.extent(0) == knn_graph.extent(0), + "dataset size is expected to have the same number of graph index size"); + RAFT_EXPECTS( + metric == cuvs::distance::DistanceType::InnerProduct || + metric == cuvs::distance::DistanceType::CosineExpanded || + metric == cuvs::distance::DistanceType::L2Expanded || + metric == cuvs::distance::DistanceType::BitwiseHamming || + metric == cuvs::distance::DistanceType::L1, + "Unsupported metric. Only InnerProduct, CosineExpanded, L2Expanded, BitwiseHamming and L1 are " + "supported"); + + auto const rows = static_cast(dataset.extent(0)); + auto const dim = static_cast(dataset.extent(1)); + auto const degree = static_cast(knn_graph.extent(1)); + RAFT_EXPECTS(rows <= static_cast(std::numeric_limits::max()), + "Dataset size must fit in the graph index type"); + RAFT_EXPECTS(rows <= static_cast(std::numeric_limits::max()) && + dim <= static_cast(std::numeric_limits::max()), + "Dataset extents must fit in uint32_t"); + RAFT_EXPECTS( + degree > 0 && degree <= kMaxSortDegree, "Graph degree must be in [1, %lu]", kMaxSortDegree); + + auto kernel = select_sort_kernel(degree); + + constexpr uint32_t block_size = 256; + auto const warps = block_size / raft::WarpSize; + auto const blocks = static_cast((rows + warps - 1) / warps); + kernel<<>>( + dataset.data_handle(), + static_cast(rows), + static_cast(dim), + knn_graph.data_handle(), + static_cast(rows), + static_cast(degree), + metric); + RAFT_CUDA_TRY(cudaGetLastError()); +} + template , @@ -985,21 +1056,11 @@ void sort_knn_graph( { RAFT_EXPECTS(dataset.extent(0) == knn_graph.extent(0), "dataset size is expected to have the same number of graph index size"); - RAFT_EXPECTS( - metric == cuvs::distance::DistanceType::InnerProduct || - metric == cuvs::distance::DistanceType::CosineExpanded || - metric == cuvs::distance::DistanceType::L2Expanded || - metric == cuvs::distance::DistanceType::BitwiseHamming || - metric == cuvs::distance::DistanceType::L1, - "Unsupported metric. Only InnerProduct, CosineExpanded, L2Expanded, BitwiseHamming and L1 are " - "supported"); const uint64_t dataset_size = dataset.extent(0); const uint64_t dataset_dim = dataset.extent(1); - const DataT* dataset_ptr = dataset.data_handle(); const IdxT graph_size = dataset_size; const uint64_t input_graph_degree = knn_graph.extent(1); - IdxT* const input_graph_ptr = knn_graph.data_handle(); auto large_tmp_mr = raft::resource::get_large_workspace_resource_ref(res); @@ -1018,51 +1079,9 @@ void sort_knn_graph( raft::copy(res, d_input_graph.view(), knn_graph); - void (*kernel_sort)(const DataT* const, - const IdxT, - const uint32_t, - IdxT* const, - const uint32_t, - const uint32_t, - const cuvs::distance::DistanceType); - if (input_graph_degree <= 32) { - constexpr int numElementsPerThread = 1; - kernel_sort = kern_sort; - } else if (input_graph_degree <= 64) { - constexpr int numElementsPerThread = 2; - kernel_sort = kern_sort; - } else if (input_graph_degree <= 128) { - constexpr int numElementsPerThread = 4; - kernel_sort = kern_sort; - } else if (input_graph_degree <= 256) { - constexpr int numElementsPerThread = 8; - kernel_sort = kern_sort; - } else if (input_graph_degree <= 512) { - constexpr int numElementsPerThread = 16; - kernel_sort = kern_sort; - } else if (input_graph_degree <= 1024) { - constexpr int numElementsPerThread = 32; - kernel_sort = kern_sort; - } else { - RAFT_FAIL( - "The degree of input knn graph is too large (%lu). " - "It must be equal to or smaller than %d.", - input_graph_degree, - 1024); - } - const auto block_size = 256; - const auto num_warps_per_block = block_size / raft::WarpSize; - const auto grid_size = (graph_size + num_warps_per_block - 1) / num_warps_per_block; - RAFT_LOG_DEBUG("."); - kernel_sort<<>>( - d_dataset.data_handle(), - dataset_size, - dataset_dim, - d_input_graph.data_handle(), - graph_size, - input_graph_degree, - metric); + sort_knn_graph_device_inplace( + res, metric, raft::make_const_mdspan(d_dataset.view()), d_input_graph.view()); raft::resource::sync_stream(res); RAFT_LOG_DEBUG("."); raft::copy(res, knn_graph, raft::make_const_mdspan(d_input_graph.view())); diff --git a/cpp/tests/CMakeLists.txt b/cpp/tests/CMakeLists.txt index c0decb2243..c13af55e16 100644 --- a/cpp/tests/CMakeLists.txt +++ b/cpp/tests/CMakeLists.txt @@ -212,6 +212,13 @@ ConfigureTest( PERCENT 100 ) +ConfigureTest( + NAME NEIGHBORS_ANN_CAGRA_MERGE_TEST + PATH neighbors/ann_cagra/test_merge_fastener.cu + GPUS 1 + PERCENT 100 +) + ConfigureTest( NAME NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST PATH neighbors/ann_cagra/test_half_uint32_t.cu diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 1e969ec5fe..daeaba8816 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -278,6 +278,8 @@ struct AnnCagraInputs { cuvs::neighbors::MergeStrategy merge_strategy = cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL; cuvs::neighbors::cagra::internal_dtype smem_dtype = cuvs::neighbors::cagra::internal_dtype::F16; + /** When set, physical merge uses this overload instead of the default-params merge. */ + std::optional physical_merge_params = std::nullopt; }; inline ::std::ostream& operator<<(::std::ostream& os, const AnnCagraInputs& p) @@ -1448,7 +1450,10 @@ class AnnCagraIndexMergeTest : public ::testing::TestWithParam { std::vector*> indices_to_merge{&index0, &index1}; if (ps.merge_strategy == cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL) { - auto merged = cagra::merge(handle_, index_params, indices_to_merge); + auto merged = + ps.physical_merge_params.has_value() + ? cagra::merge(handle_, index_params, indices_to_merge, *ps.physical_merge_params) + : cagra::merge(handle_, index_params, indices_to_merge); cagra::search( handle_, search_params, merged, search_queries_view, indices_out_view, dists_out_view); } else { diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu new file mode 100644 index 0000000000..50c44f95d1 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -0,0 +1,954 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +/** + * @file + * @brief Component and end-to-end coverage for CAGRA Fastener merges. + */ + +#include "../../../src/neighbors/detail/cagra/cagra_merge.cuh" +#include "../ann_cagra.cuh" + +#include +#include + +#include +#include +#include +#include +#include +#include + +#include + +#include +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { +namespace { + +template +auto make_dataset(raft::resources const& res, + int64_t rows, + int64_t dim, + uint64_t seed = 1234ULL, + cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded) +{ + auto device = raft::make_device_matrix(res, rows, dim); + raft::random::RngState rng(seed); + InitDataset(res, + device.data_handle(), + static_cast(rows), + static_cast(dim), + metric, + rng); + auto host = raft::make_host_matrix(res, rows, dim); + auto stream = raft::resource::get_cuda_stream(res); + raft::copy(host.data_handle(), device.data_handle(), host.size(), stream); + raft::resource::sync_stream(res); + return host; +} + +template +auto concat_host_datasets(raft::resources const& res, + raft::host_matrix_view first, + raft::host_matrix_view second) +{ + EXPECT_EQ(first.extent(1), second.extent(1)); + auto out = + raft::make_host_matrix(res, first.extent(0) + second.extent(0), first.extent(1)); + std::copy_n(first.data_handle(), first.size(), out.data_handle()); + std::copy_n(second.data_handle(), second.size(), out.data_handle() + first.size()); + return out; +} + +auto make_ring_graph(raft::resources const& res, int64_t rows, int64_t degree) +{ + auto graph = raft::make_host_matrix(res, rows, degree); + for (int64_t i = 0; i < rows; ++i) { + for (int64_t j = 0; j < degree; ++j) { + graph(i, j) = static_cast((i + j + 1) % rows); + } + } + return graph; +} + +template +void expect_valid_graph(index const& merged, int64_t rows, int64_t degree) +{ + ASSERT_EQ(merged.size(), rows); + ASSERT_EQ(merged.graph().extent(0), rows); + ASSERT_EQ(merged.graph().extent(1), degree); + auto host = raft::make_host_matrix(rows, degree); + raft::resources res; + raft::copy(host.data_handle(), + merged.graph().data_handle(), + host.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < degree; ++column) { + EXPECT_LT(host(row, column), static_cast(rows)); + EXPECT_NE(host(row, column), static_cast(row)); + } + } +} + +template +void expect_dataset_order(raft::resources const& res, + index const& merged, + raft::host_matrix_view expected) +{ + auto view = merged.dataset(); + ASSERT_EQ(view.extent(0), expected.extent(0)); + ASSERT_EQ(view.extent(1), expected.extent(1)); + // The merged owning dataset is attached through make_aligned_dataset, which pads each row to a + // 16-byte stride. For int8/uint8 (and any dim whose byte width is not a multiple of 16) the + // stride exceeds the logical dimension, so the rows must be copied honoring that stride rather + // than as one contiguous block. + auto host = raft::make_host_matrix(res, expected.extent(0), expected.extent(1)); + auto stream = raft::resource::get_cuda_stream(res); + int64_t row_stride = view.stride(0); + for (int64_t row = 0; row < view.extent(0); ++row) { + raft::copy(host.data_handle() + row * view.extent(1), + view.data_handle() + row * row_stride, + view.extent(1), + stream); + } + raft::resource::sync_stream(res); + for (int64_t row = 0; row < expected.extent(0); ++row) { + for (int64_t column = 0; column < expected.extent(1); ++column) { + EXPECT_EQ(static_cast(host(row, column)), static_cast(expected(row, column))); + } + } +} + +template +void run_explicit_fastener(cuvs::distance::DistanceType metric) +{ + raft::resources res; + constexpr int64_t rows = 48; + constexpr int64_t dim = 8; + constexpr int64_t degree = 4; + auto dataset0 = make_dataset(res, rows, dim, 1234ULL, metric); + auto dataset1 = make_dataset(res, rows, dim, 5678ULL, metric); + auto expected = concat_host_datasets( + res, raft::make_const_mdspan(dataset0.view()), raft::make_const_mdspan(dataset1.view())); + auto graph0 = make_ring_graph(res, rows, degree); + auto graph1 = make_ring_graph(res, rows, degree); + index index0( + res, metric, raft::make_const_mdspan(dataset0.view()), raft::make_const_mdspan(graph0.view())); + index index1( + res, metric, raft::make_const_mdspan(dataset1.view()), raft::make_const_mdspan(graph1.view())); + + index_params params; + params.metric = metric; + params.graph_degree = degree; + params.intermediate_graph_degree = degree; + params.attach_dataset_on_build = true; + params.guarantee_connectivity = false; + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + fastener.leaf_size = 64; + fastener.leaf_degree = 4; + std::vector*> indices{&index0, &index1}; + + auto merged = merge(res, params, indices, fastener); + expect_valid_graph(merged, rows * 2, degree); + expect_dataset_order(res, merged, raft::make_const_mdspan(expected.view())); + EXPECT_TRUE(merged.data().is_owning()); + EXPECT_EQ(index0.dataset().extent(0), rows); + EXPECT_EQ(index1.dataset().extent(0), rows); + EXPECT_TRUE(index0.data().is_owning()); + EXPECT_TRUE(index1.data().is_owning()); + EXPECT_EQ(index0.size(), rows); + EXPECT_EQ(index1.size(), rows); + EXPECT_EQ(index0.dim(), dim); + EXPECT_EQ(index1.dim(), dim); +} + +/** + * Verifies that split planning is deterministic, lays out child keys and output ranges densely, + * carries completed parents forward, and rejects overlapping parent ranges. + */ +TEST(CagraMergeFastener, PlanSplitIsDeterministicDenseAndCompact) +{ + using namespace detail::merge_scaffold; + std::vector ranges{{0, 0, 40}, {1, 40, 300}, {2, 300, 1000}}; + split_params params{.fanout = 2, + .leader_fraction = 0.05, + .max_leaders = 16, + .leaf_size = 64, + .level = 1, + .occurrence_stride = 2}; + + auto plan = plan_split(ranges, 1000, params, DETERMINISTIC_SEED); + auto replay = plan_split(ranges, 1000, params, DETERMINISTIC_SEED); + + ASSERT_EQ(plan.parents.size(), ranges.size()); + EXPECT_TRUE(plan.parents[0].carried()); + EXPECT_FALSE(plan.parents[1].carried()); + EXPECT_FALSE(plan.parents[2].carried()); + + int64_t output_cursor = 0; + uint32_t key_cursor = 0; + for (size_t i = 0; i < plan.parents.size(); ++i) { + auto const& entry = plan.parents[i]; + EXPECT_EQ(entry.output_start, output_cursor); + EXPECT_EQ(entry.child_key_base, key_cursor); + output_cursor += entry.output_rows(params.fanout); + key_cursor += entry.child_count(); + if (!entry.carried()) { + EXPECT_GE(entry.leader_count, static_cast(params.fanout)); + EXPECT_LE(entry.leader_count, static_cast(params.max_leaders)); + EXPECT_LT(entry.leader_offset, static_cast(entry.size())); + } + EXPECT_EQ(entry.leader_offset, replay.parents[i].leader_offset); + } + EXPECT_EQ(plan.output_rows, output_cursor); + EXPECT_EQ(plan.child_count, key_cursor); + EXPECT_EQ(plan.output_rows, 40 + 260 * 2 + 700 * 2); + + std::vector overlapping{{0, 0, 40}, {1, 30, 100}}; + EXPECT_THROW(plan_split(overlapping, 100, params, DETERMINISTIC_SEED), raft::exception); +} + +/** + * Verifies that repeated levels use the same many-way split operation, emit valid memberships and + * contiguous ranges, and preserve every membership of parents already within the leaf-size limit. + */ +TEST(CagraMergeFastener, SplitManywayCarriesSmallParentsAndSupportsRepeatedLevels) +{ + using namespace detail::merge_scaffold; + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + constexpr int64_t rows = 96; + constexpr int64_t dim = 4; + auto host_dataset = make_dataset(res, rows, dim); + auto dataset = raft::make_device_matrix(res, rows, dim); + raft::copy(dataset.data_handle(), host_dataset.data_handle(), dataset.size(), stream); + + split_context context(stream, rows); + manyway_l2_norms_kernel<<((rows + 3) / 4), 128, 0, stream>>>( + dataset.data_handle(), rows, dim, context.norms.data()); + RAFT_CUDA_TRY(cudaGetLastError()); + + auto root = make_root_partition(res, rows); + auto first = split_manyway(res, + raft::make_const_mdspan(dataset.view()), + root, + split_params{.fanout = 2, + .leader_fraction = 0.01, + .max_leaders = 2, + .leaf_size = 64, + .level = 0, + .occurrence_stride = 1}, + context); + auto second = split_manyway(res, + raft::make_const_mdspan(dataset.view()), + first, + split_params{.fanout = 2, + .leader_fraction = 0.01, + .max_leaders = 2, + .leaf_size = 64, + .level = 1, + .occurrence_stride = 2}, + context); + EXPECT_EQ(second.memberships.size(), static_cast(rows * 4)); + + int64_t cursor = 0; + for (auto const& range : second.ranges) { + EXPECT_EQ(range.start, cursor); + EXPECT_GT(range.end, range.start); + EXPECT_LE(range.end, static_cast(second.memberships.size())); + cursor = range.end; + } + EXPECT_EQ(cursor, static_cast(second.memberships.size())); + + std::vector records(second.memberships.size()); + raft::copy(records.data(), second.memberships.data(), records.size(), stream); + raft::resource::sync_stream(res); + for (auto const& record : records) { + EXPECT_LT(record.id, static_cast(rows)); + EXPECT_LT(record.occurrence, uint16_t{4}); + } + + constexpr int64_t small_rows = 32; + auto small_root = make_root_partition(res, small_rows); + split_context small_context(stream, small_rows); + auto carried = split_manyway(res, + raft::make_const_mdspan(dataset.view()), + small_root, + split_params{.fanout = 3, + .leader_fraction = 0.5, + .max_leaders = 8, + .leaf_size = 64, + .level = 0, + .occurrence_stride = 1}, + small_context); + ASSERT_EQ(carried.ranges.size(), 1); + EXPECT_EQ(carried.memberships.size(), static_cast(small_rows)); + std::vector carried_records(carried.memberships.size()); + raft::copy(carried_records.data(), carried.memberships.data(), carried_records.size(), stream); + raft::resource::sync_stream(res); + for (int64_t row = 0; row < small_rows; ++row) { + EXPECT_EQ(carried_records[row].id, static_cast(row)); + EXPECT_EQ(carried_records[row].occurrence, uint16_t{0}); + } +} + +/** + * Verifies that scaffold slots not populated by a leaf neighbor are initialized with the source + * row, providing a valid sentinel for later graph combination and compaction. + */ +TEST(CagraMergeFastener, InitializesUnwrittenScaffoldSlotsWithSelf) +{ + using namespace detail::merge_scaffold; + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + constexpr int64_t rows = 3; + constexpr int64_t degree = 4; + auto graph = raft::make_device_matrix(res, rows, degree); + + initialize_self_scaffold_kernel<<<1, 256, 0, stream>>>(graph.data_handle(), rows, degree); + RAFT_CUDA_TRY(cudaGetLastError()); + + std::vector initialized(graph.size()); + raft::copy(initialized.data(), graph.data_handle(), initialized.size(), stream); + raft::resource::sync_stream(res); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < degree; ++column) { + EXPECT_EQ(initialized[row * degree + column], static_cast(row)); + } + } +} + +/** + * Builds a many-way scaffold over two input ranges that share one exact duplicate vector and + * checks that partitioning plus intra-leaf KNN place each duplicate in the other's scaffold + * neighbors. Identical points select the same leaders, so they co-occur in the spilled children; + * zero distance then makes them mutual cross-origin nearest neighbors inside those leaves. + */ +TEST(CagraMergeFastener, DuplicateAcrossInputsAppearInEachOthersScaffoldNeighbors) +{ + using namespace detail::merge_scaffold; + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + constexpr int64_t part0_rows = 8; + constexpr int64_t part1_rows = 8; + constexpr int64_t rows = part0_rows + part1_rows; + constexpr int64_t dim = 8; + constexpr uint32_t duplicate0 = 3; + constexpr uint32_t duplicate1 = static_cast(part0_rows + 5); + constexpr uint32_t leaf_degree = 2; + + auto host = raft::make_host_matrix(res, rows, dim); + for (int64_t row = 0; row < rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + host(row, column) = static_cast(row * 17 + column); + } + } + for (int64_t column = 0; column < dim; ++column) { + host(duplicate1, column) = host(duplicate0, column); + } + + auto dataset = raft::make_device_matrix(res, rows, dim); + raft::copy(dataset.data_handle(), host.data_handle(), dataset.size(), stream); + + // Force a real root split (rows > leaf_size). Dense leader sampling keeps each spilled child + // small enough that make_leaves does not separate the co-assigned duplicates. + build_params params; + params.levels = 1; + params.root_fanout = 2; + params.lower_fanout = 2; + params.leader_fraction = 1.0; + params.max_leaders = 16; + params.leaf_size = 8; + params.leaf_degree = leaf_degree; + + std::vector offsets{0, part0_rows, rows}; + auto scaffold = build(res, raft::make_const_mdspan(dataset.view()), offsets, params); + ASSERT_EQ(scaffold.extent(0), rows); + ASSERT_EQ(scaffold.extent(1), params.root_fanout * leaf_degree); + + auto scaffold_host = + raft::make_host_matrix(res, scaffold.extent(0), scaffold.extent(1)); + raft::copy(scaffold_host.data_handle(), scaffold.data_handle(), scaffold.size(), stream); + raft::resource::sync_stream(res); + + auto row_contains = [&](uint32_t row, uint32_t neighbor) { + for (int64_t column = 0; column < scaffold_host.extent(1); ++column) { + if (scaffold_host(row, column) == neighbor) { return true; } + } + return false; + }; + + EXPECT_TRUE(row_contains(duplicate0, duplicate1)); + EXPECT_TRUE(row_contains(duplicate1, duplicate0)); +} + +/** + * Verifies the leaf-GEMM eligibility boundary for every scalar type and confirms that a dimension + * that cannot fit one float leaf in the GEMM workspace is rejected before either input is mutated. + */ +TEST(CagraMergeFastener, LeafGemmLimitsRejectOversizedWorkspaceDimensions) +{ + using namespace detail::merge_scaffold; + EXPECT_TRUE(leaf_gemm_supported(4096, 256)); + EXPECT_TRUE(leaf_gemm_supported(4096, 256)); + EXPECT_TRUE(leaf_gemm_supported(4096, 256)); + EXPECT_TRUE(leaf_gemm_supported(4096, 256)); + + // One float leaf of this width exceeds GEMM_WORKSPACE_BYTES (~2 GiB). + constexpr uint32_t leaf_size = 256; + constexpr int64_t oversized_dim = + static_cast(GEMM_WORKSPACE_BYTES / sizeof(float) / leaf_size) + 1; + EXPECT_FALSE(leaf_gemm_supported(oversized_dim, leaf_size)); + EXPECT_FALSE(leaf_gemm_supported(oversized_dim, leaf_size)); + EXPECT_FALSE(leaf_gemm_supported(oversized_dim, leaf_size)); + + raft::resources res; + constexpr int64_t rows = 2; + constexpr int64_t dim = oversized_dim; + auto dataset0 = make_dataset(res, rows, dim, 1234ULL); + auto dataset1 = make_dataset(res, rows, dim, 5678ULL); + auto graph0 = make_ring_graph(res, rows, 1); + auto graph1 = make_ring_graph(res, rows, 1); + auto metric = cuvs::distance::DistanceType::L2Expanded; + index index0( + res, metric, raft::make_const_mdspan(dataset0.view()), raft::make_const_mdspan(graph0.view())); + index index1( + res, metric, raft::make_const_mdspan(dataset1.view()), raft::make_const_mdspan(graph1.view())); + + index_params params; + params.metric = metric; + params.graph_degree = 1; + params.intermediate_graph_degree = 1; + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + fastener.leaf_size = leaf_size; + std::vector*> indices{&index0, &index1}; + + auto result = detail::preflight_fastener( + params, fastener, indices, cuvs::neighbors::filtering::none_sample_filter{}); + EXPECT_FALSE(result.eligible); + EXPECT_EQ(result.reason, "dataset dimension exceeds the leaf GEMM workspace limit"); + EXPECT_EQ(index0.dataset().extent(0), rows); + EXPECT_EQ(index1.dataset().extent(0), rows); +} + +/** + * Exercises an explicit Fastener merge for float, half, int8, and uint8 data, including graph + * validity, concatenated dataset order, output ownership, and preservation of both input indices. + */ +TEST(CagraMergeFastener, SupportsAllScalarTypes) +{ + auto metric = cuvs::distance::DistanceType::L2Expanded; + run_explicit_fastener(metric); + run_explicit_fastener(metric); + run_explicit_fastener(metric); + run_explicit_fastener(metric); +} + +/** + * Verifies that merging one owning and one non-owning input preserves both inputs and their + * ownership modes while producing an owning output with the expected dataset order. + */ +TEST(CagraMergeFastener, MixedDatasetOwnershipPreservesInputs) +{ + raft::resources res; + constexpr int64_t rows = 48; + constexpr int64_t dim = 8; + constexpr int64_t degree = 4; + auto dataset0 = make_dataset(res, rows, dim, 1234ULL); + auto dataset1 = make_dataset(res, rows, dim, 5678ULL); + auto expected = concat_host_datasets( + res, raft::make_const_mdspan(dataset0.view()), raft::make_const_mdspan(dataset1.view())); + auto device_dataset1 = raft::make_device_matrix(res, rows, dim); + raft::copy(device_dataset1.data_handle(), + dataset1.data_handle(), + dataset1.size(), + raft::resource::get_cuda_stream(res)); + auto graph0 = make_ring_graph(res, rows, degree); + auto graph1 = make_ring_graph(res, rows, degree); + auto metric = cuvs::distance::DistanceType::L2Expanded; + index index0( + res, metric, raft::make_const_mdspan(dataset0.view()), raft::make_const_mdspan(graph0.view())); + index index1(res, + metric, + raft::make_const_mdspan(device_dataset1.view()), + raft::make_const_mdspan(graph1.view())); + ASSERT_TRUE(index0.data().is_owning()); + ASSERT_FALSE(index1.data().is_owning()); + + index_params params; + params.metric = metric; + params.graph_degree = degree; + params.intermediate_graph_degree = degree; + params.attach_dataset_on_build = true; + params.guarantee_connectivity = false; + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + fastener.leaf_size = 64; + fastener.leaf_degree = 4; + std::vector*> indices{&index0, &index1}; + + auto merged = merge(res, params, indices, fastener); + expect_valid_graph(merged, rows * 2, degree); + expect_dataset_order(res, merged, raft::make_const_mdspan(expected.view())); + EXPECT_TRUE(merged.data().is_owning()); + EXPECT_EQ(index0.dataset().extent(0), rows); + EXPECT_EQ(index1.dataset().extent(0), rows); + EXPECT_TRUE(index0.data().is_owning()); + EXPECT_FALSE(index1.data().is_owning()); +} + +/** + * Reuses AnnCagraIndexMergeTest on the primary float L2 shapes from generate_inputs(): build two + * CAGRA graphs over halves of the same InitDataset vectors, merge with default Fastener knobs, and + * require at least 95% recall against brute force. + */ +namespace { +std::vector generate_fastener_merge_recall_inputs() +{ + // Same core shape as the leading AnnCagra float L2 cases: 1000x16 database, 100 queries, k=16. + auto inputs = raft::util::itertools::product( + {100}, + {1000}, + {16}, + {16}, // k + {32}, // graph_degree + {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT}, + {search_algo::AUTO}, + {10}, + {0}, + {256}, + {1}, + {cuvs::distance::DistanceType::L2Expanded}, + {false}, + {true}, + {false}, + {0.95}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {std::optional{std::nullopt}}, + {cuvs::neighbors::MergeStrategy::MERGE_STRATEGY_PHYSICAL}); + for (auto& input : inputs) { + // Default Fastener controls; only the algorithm is forced so AUTO cannot fall back to rebuild. + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + input.physical_merge_params = fastener; + } + return inputs; +} +} // namespace + +typedef AnnCagraIndexMergeTest AnnCagraFastenerMergeRecallTest; +TEST_P(AnnCagraFastenerMergeRecallTest, DefaultFloatMergeSearchRecallAgainstBruteForce) +{ + this->testCagra(); +} +INSTANTIATE_TEST_CASE_P(CagraMergeFastener, + AnnCagraFastenerMergeRecallTest, + ::testing::ValuesIn(generate_fastener_merge_recall_inputs())); + +/** + * Exercises a three-level, non-default uint8 merge with different input graph degrees and verifies + * that every output row has the requested degree and contains only valid neighbor identifiers. + */ +TEST(CagraMergeFastener, MixedDegreesAndThreeLevelUint8OptionsProduceExactDegree) +{ + raft::resources res; + constexpr int64_t rows = 64; + constexpr int64_t dim = 16; + auto dataset0 = make_dataset(res, rows, dim, 1234ULL); + auto dataset1 = make_dataset(res, rows, dim, 5678ULL); + auto graph0 = make_ring_graph(res, rows, 6); + auto graph1 = make_ring_graph(res, rows, 8); + index index0(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset0.view()), + raft::make_const_mdspan(graph0.view())); + index index1(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset1.view()), + raft::make_const_mdspan(graph1.view())); + + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = 8; + params.attach_dataset_on_build = true; + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + fastener.levels = 3; + fastener.root_fanout = 2; + fastener.lower_fanout = 2; + fastener.leader_fraction = 0.25; + fastener.max_leaders = 32; + fastener.leaf_size = 64; + fastener.leaf_degree = 8; + std::vector*> indices{&index0, &index1}; + + auto merged = merge(res, params, indices, fastener); + expect_valid_graph(merged, rows * 2, 8); +} + +/** + * Verifies that input graphs with different degrees are shifted to global identifiers, cyclically + * padded to a common width, and appended to the already-global scaffold columns in row order. + */ +TEST(CagraMergeFastener, AppendCyclicallyPadsMixedInputDegrees) +{ + raft::resources res; + auto dataset0 = make_dataset(res, 2, 1, 1234ULL); + auto dataset1 = make_dataset(res, 2, 1, 5678ULL); + auto graph0 = make_ring_graph(res, 2, 2); + auto graph1 = make_ring_graph(res, 2, 4); + index index0(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset0.view()), + raft::make_const_mdspan(graph0.view())); + index index1(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset1.view()), + raft::make_const_mdspan(graph1.view())); + std::vector*> indices{&index0, &index1}; + std::vector offsets{0, 2, 4}; + + auto scaffold_host = raft::make_host_matrix(4, 1); + uint32_t scaffold_rows[4] = {2, 3, 0, 1}; + for (int64_t row = 0; row < 4; ++row) { + scaffold_host(row, 0) = scaffold_rows[row]; + } + auto scaffold = raft::make_device_matrix(res, 4, 1); + raft::copy(scaffold.data_handle(), + scaffold_host.data_handle(), + scaffold.size(), + raft::resource::get_cuda_stream(res)); + + auto appended = detail::merge_scaffold::append_to_input_graphs( + res, indices, offsets, raft::make_const_mdspan(scaffold.view())); + ASSERT_EQ(appended.extent(0), 4); + ASSERT_EQ(appended.extent(1), 5); + + auto host = raft::make_host_matrix(4, 5); + raft::copy( + host.data_handle(), appended.data_handle(), host.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + // The degree-2 input repeats cyclically to the base width of 4; the degree-4 input copies + // through shifted by its offset; the scaffold column is already global. + uint32_t expected[4][5] = {{1, 0, 1, 0, 2}, {0, 1, 0, 1, 3}, {3, 2, 3, 2, 0}, {2, 3, 2, 3, 1}}; + for (int64_t row = 0; row < 4; ++row) { + for (int64_t column = 0; column < 5; ++column) { + EXPECT_EQ(host(row, column), expected[row][column]) << row << "," << column; + } + } +} + +/** + * Verifies that Fastener can merge higher-degree input graphs into a valid graph whose requested + * output degree is smaller than either input degree. + */ +TEST(CagraMergeFastener, MergeSupportsOutputDegreeBelowInputDegree) +{ + raft::resources res; + constexpr int64_t rows = 48; + constexpr int64_t dim = 8; + constexpr int64_t input_degree = 8; + auto dataset0 = make_dataset(res, rows, dim, 1234ULL); + auto dataset1 = make_dataset(res, rows, dim, 5678ULL); + auto graph0 = make_ring_graph(res, rows, input_degree); + auto graph1 = make_ring_graph(res, rows, input_degree); + index index0(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset0.view()), + raft::make_const_mdspan(graph0.view())); + index index1(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset1.view()), + raft::make_const_mdspan(graph1.view())); + + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = 4; + params.attach_dataset_on_build = true; + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + fastener.leaf_size = 64; + fastener.leaf_degree = 4; + std::vector*> indices{&index0, &index1}; + + auto merged = merge(res, params, indices, fastener); + expect_valid_graph(merged, rows * 2, params.graph_degree); +} + +auto make_float_indices(raft::resources const& res, + cuvs::distance::DistanceType metric, + raft::host_matrix& dataset0, + raft::host_matrix& dataset1, + raft::host_matrix& graph0, + raft::host_matrix& graph1) +{ + std::vector> output; + output.emplace_back( + res, metric, raft::make_const_mdspan(dataset0.view()), raft::make_const_mdspan(graph0.view())); + output.emplace_back( + res, metric, raft::make_const_mdspan(dataset1.view()), raft::make_const_mdspan(graph1.view())); + return output; +} + +/** + * Exercises invalid levels, fanouts, leader settings, leaf settings, and spill width, and verifies + * that preflight rejects each configuration without changing either input dataset. + */ +TEST(CagraMergeFastener, InvalidManywayOptionsFailPreflightWithoutMutation) +{ + raft::resources res; + constexpr int64_t rows = 32; + constexpr int64_t dim = 8; + auto dataset0 = make_dataset(res, rows, dim, 1234ULL); + auto dataset1 = make_dataset(res, rows, dim, 5678ULL); + auto graph0 = make_ring_graph(res, rows, 4); + auto graph1 = make_ring_graph(res, rows, 4); + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = 4; + params.intermediate_graph_degree = 8; + params.attach_dataset_on_build = true; + auto owned = make_float_indices(res, params.metric, dataset0, dataset1, graph0, graph1); + std::vector*> indices{&owned[0], &owned[1]}; + + std::vector invalid; + auto add_invalid = [&](auto update) { + merge_params candidate; + candidate.algo = merge_algo::FASTENER; + update(candidate); + invalid.push_back(candidate); + }; + add_invalid([](auto& value) { value.levels = 0; }); + add_invalid([](auto& value) { value.root_fanout = 0; }); + add_invalid([](auto& value) { value.root_fanout = 33; }); + add_invalid([](auto& value) { value.lower_fanout = 0; }); + add_invalid([](auto& value) { value.lower_fanout = 33; }); + add_invalid([](auto& value) { value.leader_fraction = 0.0; }); + add_invalid([](auto& value) { value.leader_fraction = 1.1; }); + add_invalid( + [](auto& value) { value.leader_fraction = std::numeric_limits::quiet_NaN(); }); + add_invalid([](auto& value) { value.max_leaders = 0; }); + add_invalid([](auto& value) { value.max_leaders = 8193; }); + add_invalid([](auto& value) { value.max_leaders = 2; }); + add_invalid([](auto& value) { value.leaf_size = 0; }); + add_invalid([](auto& value) { value.leaf_size = 512; }); + add_invalid([](auto& value) { value.leaf_degree = 0; }); + add_invalid([](auto& value) { value.leaf_degree = 16; }); + add_invalid([](auto& value) { + value.levels = 3; + value.root_fanout = 32; + value.lower_fanout = 32; + }); + + for (auto const& candidate : invalid) { + auto result = detail::preflight_fastener( + params, candidate, indices, cuvs::neighbors::filtering::none_sample_filter{}); + EXPECT_FALSE(result.eligible) << result.reason; + } + EXPECT_EQ(owned[0].dataset().extent(0), rows); + EXPECT_EQ(owned[1].dataset().extent(0), rows); +} + +/** + * Verifies dispatch semantics: explicit Fastener rejects unsupported settings and metrics, AUTO + * falls back successfully, REBUILD remains available, and every path preserves its input indices. + */ +TEST(CagraMergeFastener, DispatchRejectsOrFallsBackBeforeMutation) +{ + raft::resources res; + constexpr int64_t rows = 32; + constexpr int64_t dim = 8; + auto dataset0 = make_dataset(res, rows, dim, 1234ULL); + auto dataset1 = make_dataset(res, rows, dim, 5678ULL); + auto graph0 = make_ring_graph(res, rows, 4); + auto graph1 = make_ring_graph(res, rows, 4); + index_params params; + params.metric = cuvs::distance::DistanceType::L2Expanded; + params.graph_degree = 4; + params.intermediate_graph_degree = 8; + params.attach_dataset_on_build = true; + + { + auto owned = make_float_indices(res, params.metric, dataset0, dataset1, graph0, graph1); + std::vector*> indices{&owned[0], &owned[1]}; + merge_params unsupported; + unsupported.algo = merge_algo::FASTENER; + unsupported.leaf_size = 512; + EXPECT_ANY_THROW(merge(res, params, indices, unsupported)); + EXPECT_EQ(owned[0].dataset().extent(0), rows); + EXPECT_EQ(owned[1].dataset().extent(0), rows); + } + { + auto inner_product_params = params; + inner_product_params.metric = cuvs::distance::DistanceType::InnerProduct; + auto owned = + make_float_indices(res, inner_product_params.metric, dataset0, dataset1, graph0, graph1); + std::vector*> indices{&owned[0], &owned[1]}; + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + fastener.leaf_size = 64; + EXPECT_ANY_THROW(merge(res, inner_product_params, indices, fastener)); + EXPECT_EQ(owned[0].dataset().extent(0), rows); + EXPECT_EQ(owned[1].dataset().extent(0), rows); + } + { + auto owned = make_float_indices(res, params.metric, dataset0, dataset1, graph0, graph1); + std::vector*> indices{&owned[0], &owned[1]}; + merge_params automatic; + automatic.algo = merge_algo::AUTO; + automatic.leaf_size = 512; + auto merged = merge(res, params, indices, automatic); + EXPECT_EQ(merged.size(), rows * 2); + EXPECT_EQ(owned[0].dataset().extent(0), rows); + EXPECT_EQ(owned[1].dataset().extent(0), rows); + } + { + auto owned = make_float_indices(res, params.metric, dataset0, dataset1, graph0, graph1); + std::vector*> indices{&owned[0], &owned[1]}; + merge_params rebuild{merge_algo::REBUILD}; + auto merged = merge(res, params, indices, rebuild); + EXPECT_EQ(merged.size(), rows * 2); + EXPECT_EQ(owned[0].dataset().extent(0), rows); + EXPECT_EQ(owned[1].dataset().extent(0), rows); + } +} + +/** + * Verifies that the shared graph sorter applies metric-specific ordering by producing different + * first neighbors for L2 distance and inner-product similarity on the same candidate graph. + */ +TEST(CagraMergeFastener, InnerProductOrderingDiffersFromL2) +{ + raft::resources res; + auto dataset = raft::make_host_matrix(3, 2); + dataset(0, 0) = 1.0f; + dataset(0, 1) = 0.0f; + dataset(1, 0) = 2.0f; + dataset(1, 1) = 0.0f; + dataset(2, 0) = 100.0f; + dataset(2, 1) = 100.0f; + auto device_dataset = raft::make_device_matrix(res, 3, 2); + raft::copy(device_dataset.data_handle(), + dataset.data_handle(), + dataset.size(), + raft::resource::get_cuda_stream(res)); + + auto source = raft::make_host_matrix(3, 2); + source(0, 0) = 1; + source(0, 1) = 2; + source(1, 0) = 0; + source(1, 1) = 2; + source(2, 0) = 0; + source(2, 1) = 1; + auto l2 = raft::make_device_matrix(res, 3, 2); + auto ip = raft::make_device_matrix(res, 3, 2); + raft::copy( + l2.data_handle(), source.data_handle(), source.size(), raft::resource::get_cuda_stream(res)); + raft::copy( + ip.data_handle(), source.data_handle(), source.size(), raft::resource::get_cuda_stream(res)); + + detail::graph::sort_knn_graph_device_inplace(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(device_dataset.view()), + l2.view()); + detail::graph::sort_knn_graph_device_inplace(res, + cuvs::distance::DistanceType::InnerProduct, + raft::make_const_mdspan(device_dataset.view()), + ip.view()); + auto l2_host = raft::make_host_matrix(3, 2); + auto ip_host = raft::make_host_matrix(3, 2); + raft::copy( + l2_host.data_handle(), l2.data_handle(), l2.size(), raft::resource::get_cuda_stream(res)); + raft::copy( + ip_host.data_handle(), ip.data_handle(), ip.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + EXPECT_EQ(l2_host(0, 0), 1); + EXPECT_EQ(ip_host(0, 0), 2); +} + +/** + * Verifies the graph-combination pipeline end to end: local identifiers are offset, scaffold edges + * are appended, rows are metric-sorted, duplicates are removed, and short rows are padded. + */ +TEST(CagraMergeFastener, AppendSortAndDedupHandleOffsetsAndPadding) +{ + raft::resources res; + auto dataset0 = raft::make_host_matrix(2, 1); + auto dataset1 = raft::make_host_matrix(2, 1); + dataset0(0, 0) = 0.0f; + dataset0(1, 0) = 2.0f; + dataset1(0, 0) = 10.0f; + dataset1(1, 0) = 12.0f; + auto graph0 = make_ring_graph(res, 2, 1); + auto graph1 = make_ring_graph(res, 2, 1); + index index0(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset0.view()), + raft::make_const_mdspan(graph0.view())); + index index1(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(dataset1.view()), + raft::make_const_mdspan(graph1.view())); + std::vector*> indices{&index0, &index1}; + std::vector offsets{0, 2, 4}; + + auto scaffold_host = raft::make_host_matrix(4, 2); + uint32_t scaffold_rows[4][2] = {{1, 3}, {1, 2}, {2, 0}, {3, 1}}; + for (int64_t row = 0; row < 4; ++row) { + for (int64_t column = 0; column < 2; ++column) { + scaffold_host(row, column) = scaffold_rows[row][column]; + } + } + auto scaffold = raft::make_device_matrix(res, 4, 2); + raft::copy(scaffold.data_handle(), + scaffold_host.data_handle(), + scaffold.size(), + raft::resource::get_cuda_stream(res)); + auto appended = detail::merge_scaffold::append_to_input_graphs( + res, indices, offsets, raft::make_const_mdspan(scaffold.view())); + + auto combined_host = raft::make_host_matrix(4, 1); + combined_host(0, 0) = 0.0f; + combined_host(1, 0) = 2.0f; + combined_host(2, 0) = 10.0f; + combined_host(3, 0) = 12.0f; + auto combined = raft::make_device_matrix(res, 4, 1); + raft::copy(combined.data_handle(), + combined_host.data_handle(), + combined.size(), + raft::resource::get_cuda_stream(res)); + detail::graph::sort_knn_graph_device_inplace(res, + cuvs::distance::DistanceType::L2Expanded, + raft::make_const_mdspan(combined.view()), + appended.view()); + auto output = + detail::merge_scaffold::cap_sorted_graph(res, raft::make_const_mdspan(appended.view()), 3); + auto host = raft::make_host_matrix(4, 3); + raft::copy( + host.data_handle(), output.data_handle(), host.size(), raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + + uint32_t expected[4][3] = {{1, 3, 1}, {0, 2, 0}, {3, 0, 3}, {2, 1, 2}}; + for (int64_t row = 0; row < 4; ++row) { + for (int64_t column = 0; column < 3; ++column) { + EXPECT_EQ(host(row, column), expected[row][column]); + } + } +} + +} // namespace +} // namespace cuvs::neighbors::cagra diff --git a/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py b/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py index f22852f0ae..bf3b1f8aea 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py +++ b/python/cuvs_bench/cuvs_bench/config/algos/constraints/__init__.py @@ -1,5 +1,5 @@ # -# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION. +# SPDX-FileCopyrightText: Copyright (c) 2024-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. # SPDX-License-Identifier: Apache-2.0 @@ -20,9 +20,35 @@ def cuvs_cagra_build(params, dims): + valid = True if "graph_degree" in params and "intermediate_graph_degree" in params: - return params["graph_degree"] <= params["intermediate_graph_degree"] - return True + valid = params["graph_degree"] <= params["intermediate_graph_degree"] + if params.get("merge_algo") == "FASTENER": + levels = params.get("fastener_levels", 2) + root_fanout = params.get("fastener_root_fanout", 2) + lower_fanout = params.get("fastener_lower_fanout", 3) + leader_fraction = params.get("fastener_leader_fraction", 0.02) + max_leaders = params.get("fastener_max_leaders", 1000) + leaf_size = params.get("fastener_leaf_size", 256) + leaf_degree = params.get("fastener_leaf_degree", 4) + valid = valid and levels > 0 + valid = valid and 1 <= root_fanout <= 32 + valid = valid and 1 <= lower_fanout <= 32 + valid = valid and 0.0 < leader_fraction <= 1.0 + valid = valid and max(root_fanout, lower_fanout) <= max_leaders <= 8192 + valid = valid and 1 <= leaf_size <= 256 + valid = valid and 1 <= leaf_degree <= 8 + if valid: + max_spill = 255 // leaf_degree + spill = root_fanout + if lower_fanout > 1: + for _ in range(1, levels): + if spill > max_spill // lower_fanout: + valid = False + break + spill *= lower_fanout + valid = valid and spill <= max_spill + return valid def cuvs_ivf_pq_build(params, dims): diff --git a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml index cf59c3ce22..e36f606ad2 100644 --- a/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml +++ b/python/cuvs_bench/cuvs_bench/config/algos/cuvs_cagra.yaml @@ -31,3 +31,50 @@ groups: itopk: [32, 64, 128, 256, 512] max_iterations: [0, 16] search_width: [1, 2, 4, 8] + merge_fastener_default: + build: + graph_degree: [64] + intermediate_graph_degree: [64] + graph_build_algo: ["NN_DESCENT"] + num_dataset_splits: [2, 8, 32] + merge_type: ["PHYSICAL"] + merge_algo: ["FASTENER"] + fastener_levels: [2] + fastener_root_fanout: [2] + fastener_lower_fanout: [3] + fastener_leader_fraction: [0.02] + fastener_max_leaders: [1000] + fastener_leaf_size: [256] + fastener_leaf_degree: [4] + search: + itopk: [64] + search_width: [1, 4, 16] + merge_fastener_high_quality: + build: + graph_degree: [64] + intermediate_graph_degree: [64] + graph_build_algo: ["NN_DESCENT"] + num_dataset_splits: [2, 8, 32] + merge_type: ["PHYSICAL"] + merge_algo: ["FASTENER"] + fastener_levels: [2] + fastener_root_fanout: [4] + fastener_lower_fanout: [2] + fastener_leader_fraction: [0.01] + fastener_max_leaders: [1000] + fastener_leaf_size: [256] + fastener_leaf_degree: [4] + search: + itopk: [64] + search_width: [1, 4, 16] + merge_rebuild: + build: + graph_degree: [64] + intermediate_graph_degree: [64] + graph_build_algo: ["NN_DESCENT"] + num_dataset_splits: [2, 8, 32] + merge_type: ["PHYSICAL"] + merge_algo: ["REBUILD"] + search: + itopk: [64] + search_width: [1, 4, 16]