From 637e0fc7eb930bf4d95dfd2a4eab64411097387f Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Thu, 16 Jul 2026 18:15:42 +0000 Subject: [PATCH 01/13] factoring out in-place edge sort --- cpp/src/neighbors/detail/cagra/graph_core.cuh | 122 ++++++++++-------- 1 file changed, 67 insertions(+), 55 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/graph_core.cuh b/cpp/src/neighbors/detail/cagra/graph_core.cuh index 52b4542798..00fb0de36b 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,24 @@ __global__ void kern_sort(const DATA_T* const dataset, // [dataset_chunk_size, } } +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 <= 32) { return kern_sort; } + if (degree <= 64) { return kern_sort; } + if (degree <= 128) { return kern_sort; } + if (degree <= 256) { return kern_sort; } + if (degree <= 512) { return kern_sort; } + if (degree <= 1024) { return kern_sort; } + RAFT_FAIL( + "The degree of input knn graph is too large (%lu). It must be equal to or smaller than 1024.", + degree); +} + template __global__ void kern_make_rev_graph_k( OutputMatrixView output_graph, // [graph_size, degree] @@ -971,6 +990,51 @@ 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 <= 1024, "Graph degree must be in [1, 1024]"); + + 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 +1049,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 +1072,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())); From 7972d3a6ada42bdbb0827f8ec65f677ac6e130e3 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 00:03:08 +0000 Subject: [PATCH 02/13] core merge logic --- .../detail/cagra/cagra_merge_scaffold.cuh | 1396 +++++++++++++++++ 1 file changed, 1396 insertions(+) create mode 100644 cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh 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..3342d06530 --- /dev/null +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -0,0 +1,1396 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * 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 + +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 WARP_SIZE = 32; +/** Warps per block in the one-row-per-warp kernels; their launch geometry derives from this. */ +inline constexpr int ROW_WARPS_PER_BLOCK = 4; +// Centered INT8 dot products accumulate into INT32. Overflow requires more than 131,071 +// dimensions even at the worst-case magnitude of 128, which is not realistic. +inline constexpr int64_t MAX_INTEGER_LEAF_DIMENSION = + std::numeric_limits::max() / (128 * 128); + +/** Grid size for the grid-stride kernels: blocks of 256 threads 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 + 255) / 256, 1048576)); +} + +// ---------------------------------------------------------------------------- +// 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 = 256; + 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 = 256; + 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; +} + +/** One candidate assignment of a row to a leader. Defaults to the invalid sentinel, which sorts + * after every real candidate. */ +struct leader_candidate { + float distance = std::numeric_limits::max(); + uint32_t leader = std::numeric_limits::max(); + + __device__ auto valid() const -> bool { return leader != std::numeric_limits::max(); } +}; + +/** Nearer candidates order first. Equal distances order by the lower leader index, keeping the + * selection deterministic. */ +__device__ inline auto operator<(leader_candidate a, leader_candidate b) -> bool +{ + return a.distance < b.distance || (a.distance == b.distance && a.leader < b.leader); +} + +/** Replace the worst entry in the best list if the new candidate is better. */ +__device__ inline void manyway_insert_local(leader_candidate candidate, + int fanout, + leader_candidate* best) +{ + int worst = 0; + for (int j = 1; j < fanout; ++j) { + if (best[worst] < best[j]) { worst = j; } + } + if (candidate < best[worst]) { best[worst] = candidate; } +} + +/** Sort the selected candidates in place from the nearest to the farthest. + * + * This is hand-rolled because thrust::sort has an allocation and I can't get a cuda::std::sort to + * compile. */ +__device__ inline void manyway_sort_selected(int fanout, leader_candidate* selected) +{ + for (int i = 1; i < fanout; ++i) { + leader_candidate candidate = selected[i]; + int j = i; + while (j > 0 && candidate < selected[j - 1]) { + selected[j] = selected[j - 1]; + --j; + } + selected[j] = candidate; + } +} + +/** 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 + 255) / 256); + 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; + } +} + +/** Select the nearest leaders for each tile row from the dot products. Write the child keys and + * memberships. One warp does one row. */ +static __global__ void manyway_select_tiles_kernel(float const* dots, + int batch_size, + int tile_rows, + int padded_leaders, + int fanout, + int occurrence_stride, + float const* norms, + uint32_t const* leader_ids, + partition_membership const* input_memberships, + assignment_tile const* tiles, + uint32_t* output_keys, + partition_membership* output_memberships) +{ + int lane = threadIdx.x % WARP_SIZE; + int warp = threadIdx.x / WARP_SIZE; + int64_t linear_row = static_cast(blockIdx.x) * ROW_WARPS_PER_BLOCK + warp; + int64_t total_rows = static_cast(batch_size) * tile_rows; + if (linear_row >= total_rows) { return; } + int batch = static_cast(linear_row / tile_rows); + int row = static_cast(linear_row % tile_rows); + auto tile = tiles[batch]; + if (row >= tile.rows) { return; } + + int64_t input_index = tile.input_start + row; + auto membership = input_memberships[input_index]; + leader_candidate local_best[MAX_FANOUT]; + + for (int leader = lane; leader < tile.leader_count; leader += WARP_SIZE) { + uint32_t leader_id = leader_ids[static_cast(batch) * padded_leaders + leader]; + float dot = dots[(static_cast(batch) * tile_rows + row) * padded_leaders + leader]; + float distance = fmaxf(0.0f, norms[membership.id] + norms[leader_id] - 2.0f * dot); + manyway_insert_local( + {.distance = distance, .leader = static_cast(leader)}, fanout, local_best); + } + + extern __shared__ unsigned char shared_bytes[]; + auto* shared_candidates = reinterpret_cast(shared_bytes); + int shared_base = (warp * WARP_SIZE + lane) * fanout; + for (int j = 0; j < fanout; ++j) { + shared_candidates[shared_base + j] = local_best[j]; + } + __syncwarp(); + + if (lane == 0) { + leader_candidate selected[MAX_FANOUT]; + int warp_base = warp * WARP_SIZE * fanout; + for (int candidate = 0; candidate < WARP_SIZE * fanout; ++candidate) { + if (shared_candidates[warp_base + candidate].valid()) { + manyway_insert_local(shared_candidates[warp_base + candidate], fanout, selected); + } + } + manyway_sort_selected(fanout, selected); + int64_t output_base = + tile.output_start + (input_index - tile.input_start) * static_cast(fanout); + for (int rank = 0; rank < fanout; ++rank) { + output_keys[output_base + rank] = tile.child_key_base + selected[rank].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()), 256, 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`. */ +template +void batched_row_dot_products(raft::resources const& res, + DataT const* a, + int a_rows, + long long a_stride, + DataT const* b, + int b_rows, + long long b_stride, + AccT* out, + long long out_stride, + int row_width, + int batch_count) +{ + static_assert((std::is_same_v && std::is_same_v) || + (std::is_same_v && std::is_same_v), + "Unsupported batched dot product type combination"); + constexpr auto data_type = std::is_same_v ? CUDA_R_32F : CUDA_R_8I; + constexpr auto out_type = std::is_same_v ? CUDA_R_32F : CUDA_R_32I; + constexpr auto compute_type = + std::is_same_v ? CUBLAS_COMPUTE_32F : CUBLAS_COMPUTE_32I; + AccT alpha = 1; + AccT beta = 0; + 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, + data_type, + row_width, + a_stride, + b, + data_type, + row_width, + b_stride, + &beta, + out, + out_type, + a_rows, + out_stride, + batch_count, + compute_type, + 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, and the selection kernel + * converts dots to distances with the precomputed row norms (|x|^2 + |l|^2 - 2 x.l), keeps the + * `fanout` nearest leaders per row (ties broken by lower leader index), 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 bytes_per_batch = (point_elements + leader_elements + dot_elements) * sizeof(float) + + static_cast(padded_leaders) * sizeof(uint32_t) + + 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); + + auto point_stride = static_cast(point_elements); + auto leader_stride = static_cast(leader_elements); + auto dot_stride = static_cast(dot_elements); + size_t selection_shared = + ROW_WARPS_PER_BLOCK * WARP_SIZE * params.fanout * sizeof(leader_candidate); + + 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)); + + // Keep each row's `fanout` nearest leaders and write child keys and memberships + int selection_blocks = + static_cast((static_cast(batch_size) * tile_rows + ROW_WARPS_PER_BLOCK - 1) / + ROW_WARPS_PER_BLOCK); + manyway_select_tiles_kernel<<>>(tile_dots.data(), + static_cast(batch_size), + tile_rows, + padded_leaders, + static_cast(params.fanout), + static_cast(params.occurrence_stride), + context.norms.data(), + tile_leader_ids.data(), + 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 goes into the GEMM workspace. */ +template +inline auto leaf_gemm_supported(int64_t dimension, uint32_t leaf_size) -> bool +{ + if (dimension <= 0 || dimension > std::numeric_limits::max()) { return false; } + + int64_t storage_dimension = dimension; + size_t vector_element_size = sizeof(float); + size_t gram_element_size = sizeof(float); + if constexpr (std::is_integral_v) { + // The below check is a little paranoid + if (dimension > MAX_INTEGER_LEAF_DIMENSION) { return false; } + storage_dimension = (dimension + 3) & ~int64_t{3}; + vector_element_size = sizeof(int8_t); + gram_element_size = sizeof(int32_t); + } + + size_t vector_elements = static_cast(leaf_size) * storage_dimension; + size_t gram_elements = static_cast(leaf_size) * leaf_size; + return vector_elements * vector_element_size + gram_elements * gram_element_size <= + 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`. Unsigned input values move to a center of zero. */ +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]; + if constexpr (std::is_same_v) { + value = static_cast(static_cast(input) - 128); + } else { + 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()); + + if constexpr (std::is_same_v || std::is_same_v) { + // Size the leaf batch so the vector and Gram buffers fit the GEMM workspace budget + 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)); + // Gather this batch's leaves into the dense vector buffer. + 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()); + // the Gram matrix of every leaf in the batch + 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)); + // select each point's nearest cross-input neighbors from the corresponding Gram matrix + 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()); + } + } else if constexpr (std::is_same_v || std::is_same_v) { + // Same pipeline as the float path, with rows padded to a multiple of 4 and an INT32 Gram. + int64_t input_dimension = dataset.extent(1); + int64_t padded_dimension = (input_dimension + 3) & ~int64_t{3}; + size_t vector_elements_per_leaf = static_cast(leaf_size) * padded_dimension; + size_t gram_elements_per_leaf = static_cast(leaf_size) * leaf_size; + size_t bytes_per_leaf = + vector_elements_per_leaf * sizeof(int8_t) + gram_elements_per_leaf * sizeof(int32_t); + 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); + int dimension = static_cast(padded_dimension); + 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, + padded_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 % WARP_SIZE; + int warp = threadIdx.x / WARP_SIZE; + 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 += WARP_SIZE) { + float value = static_cast(point[d]); + sum = fmaf(value, value, sum); + } + // Shuffle-reduce the partial sums; lane 0 holds the total. + for (int offset = WARP_SIZE / 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 32"); + 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 8192"); + 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 255"); + 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 + 255) / 256); + 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) + 255) / 256); + 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 = 256 / WARP_SIZE; + int lane = threadIdx.x % WARP_SIZE; + int warp = threadIdx.x / WARP_SIZE; + 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 += WARP_SIZE) { + 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 >> (WARP_SIZE - 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 += WARP_SIZE) { + 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 THREADS_PER_BLOCK = 256; + constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE; + 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 From bc6ec7bbdc58cdc72f7b586581dd887c861b233a Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 00:50:10 +0000 Subject: [PATCH 03/13] remove VMM dataset consolidation from fastener merge --- .../src/cuvs/cuvs_ann_bench_param_parser.h | 31 + cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h | 26 +- cpp/include/cuvs/neighbors/cagra.hpp | 113 ++- cpp/src/neighbors/cagra.cuh | 22 +- .../neighbors/detail/cagra/cagra_merge.cuh | 286 +++++- cpp/tests/CMakeLists.txt | 7 + .../ann_cagra/test_merge_fastener.cu | 900 ++++++++++++++++++ .../config/algos/constraints/__init__.py | 32 +- .../cuvs_bench/config/algos/cuvs_cagra.yaml | 47 + 9 files changed, 1405 insertions(+), 59 deletions(-) create mode 100644 cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu 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..3fee950087 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 @@ -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 d1937cba27..fe46ca3132 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,38 +2457,49 @@ 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 that consolidated copy. * - * @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, @@ -2487,6 +2508,15 @@ auto merge(raft::resources const& res, cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; +/** @brief Merge with explicit C++ algorithm and Fastener controls. */ +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, @@ -2495,10 +2525,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; @@ -2510,6 +2558,15 @@ auto merge(raft::resources const& res, 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..490c225367 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, merge_params, indices, 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..147a59c522 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,243 @@ 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 = [&](char const* reason) { + result.reason = 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 > 32 || + merge_params.lower_fanout < 1 || merge_params.lower_fanout > 32) { + return reject("root_fanout and lower_fanout must be between 1 and 32"); + } + 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 > 8192) { + return reject("max_leaders must be between 1 and 8192"); + } + 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 256"); + } + 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 8"); + } + + uint64_t const max_spill = std::numeric_limits::max() / merge_params.leaf_degree; + uint64_t spill = merge_params.root_fanout; + if (spill > max_spill) { + return reject("root_fanout * lower_fanout^(levels - 1) * leaf_degree must not exceed 255"); + } + 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("root_fanout * lower_fanout^(levels - 1) * leaf_degree must not exceed 255"); + } + 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)) { + if constexpr (std::is_integral_v) { + return reject("integer dataset dimension exceeds the INT32 leaf GEMM limit"); + } else { + 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; + + // Graph combination: append shifted input graphs to the cross-input scaffold. + 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); + } + + // Optimization: reduce the sorted candidate prefix to the requested final CAGRA graph. + 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) { + using matrix_t = decltype(dataset); + using layout_t = typename matrix_t::layout_type; + using container_policy_t = typename matrix_t::container_policy_type; + using owning_t = owning_dataset; + auto out_layout = raft::make_strided_layout(dataset.view().extents(), + cuda::std::array{preflight.dim, 1}); + merged_index.update_dataset(handle, owning_t{std::move(dataset), out_layout}); + } 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, + cagra::merge_params const& merge_params, + std::vector*>& indices, + 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, cagra::merge_params{}, indices, row_filter); +} + } // namespace cuvs::neighbors::cagra::detail 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/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu new file mode 100644 index 0000000000..22a33e8238 --- /dev/null +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -0,0 +1,900 @@ +/* + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-License-Identifier: Apache-2.0 + */ + +#include "../../../src/neighbors/detail/cagra/cagra_merge.cuh" +#include "../naive_knn.cuh" + +#include +#include + +#include +#include +#include +#include +#include + +#include +#include + +#include +#include +#include +#include + +namespace cuvs::neighbors::cagra { +namespace { + +template +auto make_value(int64_t row, int64_t column) -> T +{ + int value = static_cast((row * 7 + column * 3) % 23) - 11; + if constexpr (std::is_same_v) { + return __float2half(static_cast(value) / 8.0f); + } else if constexpr (std::is_same_v) { + return static_cast(value) / 8.0f; + } else if constexpr (std::is_same_v) { + return static_cast(value + 32); + } else { + return static_cast(value); + } +} + +template +auto make_dataset(raft::resources const& res, int64_t rows, int64_t dim, int64_t offset) +{ + auto dataset = raft::make_host_matrix(res, rows, dim); + for (int64_t i = 0; i < rows; ++i) { + for (int64_t j = 0; j < dim; ++j) { + dataset(i, j) = make_value(i + offset, j); + } + } + return dataset; +} + +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, + int64_t rows, + int64_t dim) +{ + ASSERT_EQ(merged.dataset().extent(0), rows); + ASSERT_EQ(merged.dataset().extent(1), dim); + auto host = raft::make_host_matrix(res, rows, dim); + raft::copy(host.data_handle(), + merged.dataset().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 < dim; ++column) { + EXPECT_EQ(static_cast(host(row, column)), + static_cast(make_value(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, 0); + auto dataset1 = make_dataset(res, rows, dim, rows); + 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, rows * 2, dim); + 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); +} + +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); +} + +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, 0); + 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}); + } +} + +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)); + } + } +} + +TEST(CagraMergeFastener, LeafGemmLimitsRejectOnlyUnrealisticIntegerDimensions) +{ + using namespace detail::merge_scaffold; + EXPECT_TRUE(leaf_gemm_supported(4096, 256)); + EXPECT_TRUE(leaf_gemm_supported(4096, 256)); + EXPECT_TRUE(leaf_gemm_supported(MAX_INTEGER_LEAF_DIMENSION, 256)); + EXPECT_TRUE(leaf_gemm_supported(MAX_INTEGER_LEAF_DIMENSION, 256)); + EXPECT_FALSE(leaf_gemm_supported(MAX_INTEGER_LEAF_DIMENSION + 1, 256)); + EXPECT_FALSE(leaf_gemm_supported(MAX_INTEGER_LEAF_DIMENSION + 1, 256)); + + raft::resources res; + constexpr int64_t rows = 2; + constexpr int64_t dim = MAX_INTEGER_LEAF_DIMENSION + 1; + auto dataset0 = make_dataset(res, rows, dim, 0); + auto dataset1 = make_dataset(res, rows, dim, rows); + 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 = 64; + 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, "integer dataset dimension exceeds the INT32 leaf GEMM limit"); + EXPECT_EQ(index0.dataset().extent(0), rows); + EXPECT_EQ(index1.dataset().extent(0), rows); +} + +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); +} + +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, 0); + auto dataset1 = make_dataset(res, rows, dim, rows); + 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, rows * 2, dim); + 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()); +} + +TEST(CagraMergeFastener, DefaultFloatMergeSearchRecallAgainstBruteForce) +{ + raft::resources res; + auto stream = raft::resource::get_cuda_stream(res); + constexpr int64_t part_rows = 256; + constexpr int64_t rows = 2 * part_rows; + constexpr int64_t dim = 16; + constexpr int64_t queries = 16; + constexpr int64_t k = 12; + + auto value = [](int64_t row, int64_t column) { + uint64_t bits = (static_cast(row) + 1) * 0x9e3779b97f4a7c15ull; + bits ^= (static_cast(column) + 1) * 0xbf58476d1ce4e5b9ull; + bits ^= bits >> 30; + bits *= 0x94d049bb133111ebull; + bits ^= bits >> 31; + return static_cast(bits & 0xffffu) / 65535.0f; + }; + + auto dataset0 = raft::make_host_matrix(part_rows, dim); + auto dataset1 = raft::make_host_matrix(part_rows, dim); + for (int64_t row = 0; row < part_rows; ++row) { + for (int64_t column = 0; column < dim; ++column) { + dataset0(row, column) = value(row, column); + dataset1(row, column) = value(row + part_rows, column); + } + } + + index_params partition_params; + partition_params.metric = cuvs::distance::DistanceType::L2Expanded; + partition_params.graph_degree = 16; + partition_params.intermediate_graph_degree = 32; + partition_params.attach_dataset_on_build = true; + partition_params.guarantee_connectivity = false; + auto index0 = build(res, partition_params, raft::make_const_mdspan(dataset0.view())); + auto index1 = build(res, partition_params, raft::make_const_mdspan(dataset1.view())); + std::vector*> indices{&index0, &index1}; + + merge_params fastener; + fastener.algo = merge_algo::FASTENER; + auto merged = merge(res, partition_params, indices, fastener); + expect_valid_graph(merged, rows, partition_params.graph_degree); + + auto query_host = raft::make_host_matrix(queries, dim); + for (int64_t query = 0; query < queries; ++query) { + int64_t source = query * 31; + for (int64_t column = 0; column < dim; ++column) { + query_host(query, column) = value(source, column); + } + } + auto query_device = raft::make_device_matrix(res, queries, dim); + raft::copy(query_device.data_handle(), query_host.data_handle(), query_host.size(), stream); + auto actual_indices = raft::make_device_matrix(res, queries, k); + auto actual_distances = raft::make_device_matrix(res, queries, k); + search_params search_params; + search_params.itopk_size = 32; + search_params.search_width = 4; + search_params.max_iterations = 0; + search(res, + search_params, + merged, + raft::make_const_mdspan(query_device.view()), + actual_indices.view(), + actual_distances.view()); + + auto truth_indices = raft::make_device_matrix(res, queries, k); + auto truth_distances = raft::make_device_matrix(res, queries, k); + cuvs::neighbors::naive_knn(res, + truth_distances.data_handle(), + truth_indices.data_handle(), + query_device.data_handle(), + merged.dataset().data_handle(), + static_cast(queries), + static_cast(rows), + static_cast(dim), + static_cast(k), + cuvs::distance::DistanceType::L2Expanded); + + std::vector actual(static_cast(queries * k)); + std::vector truth(static_cast(queries * k)); + raft::copy(actual.data(), actual_indices.data_handle(), actual.size(), stream); + raft::copy(truth.data(), truth_indices.data_handle(), truth.size(), stream); + raft::resource::sync_stream(res); + + size_t matches = 0; + for (int64_t query = 0; query < queries; ++query) { + for (int64_t candidate = 0; candidate < k; ++candidate) { + auto id = actual[query * k + candidate]; + for (int64_t expected = 0; expected < k; ++expected) { + if (id == truth[query * k + expected]) { + ++matches; + break; + } + } + } + } + EXPECT_GE(static_cast(matches) / static_cast(queries * k), 0.60); +} + +TEST(CagraMergeFastener, MixedDegreesAndThreeLevelUint8OptionsProduceExactDegree) +{ + raft::resources res; + constexpr int64_t rows = 64; + constexpr int64_t dim = 16; + auto dataset0 = make_dataset(res, rows, dim, 0); + auto dataset1 = make_dataset(res, rows, dim, rows); + 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); +} + +TEST(CagraMergeFastener, AppendCyclicallyPadsMixedInputDegrees) +{ + raft::resources res; + auto dataset0 = make_dataset(res, 2, 1, 0); + auto dataset1 = make_dataset(res, 2, 1, 2); + 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; + } + } +} + +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, 0); + auto dataset1 = make_dataset(res, rows, dim, rows); + 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; +} + +TEST(CagraMergeFastener, InvalidManywayOptionsFailPreflightWithoutMutation) +{ + raft::resources res; + constexpr int64_t rows = 32; + constexpr int64_t dim = 8; + auto dataset0 = make_dataset(res, rows, dim, 0); + auto dataset1 = make_dataset(res, rows, dim, rows); + 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); +} + +TEST(CagraMergeFastener, DispatchRejectsOrFallsBackBeforeMutation) +{ + raft::resources res; + constexpr int64_t rows = 32; + constexpr int64_t dim = 8; + auto dataset0 = make_dataset(res, rows, dim, 0); + auto dataset1 = make_dataset(res, rows, dim, rows); + 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); + } +} + +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); +} + +TEST(CagraMergeFastener, Uint8InnerProductSortUsesExactWideAccumulation) +{ + raft::resources res; + constexpr int64_t dim = 1024; + auto dataset = raft::make_host_matrix(3, dim); + for (int64_t d = 0; d < dim; ++d) { + dataset(0, d) = 255; + dataset(1, d) = 255; + dataset(2, d) = 255; + } + dataset(0, 0) = 1; + dataset(1, 0) = 1; + dataset(2, 0) = 0; + auto device_dataset = raft::make_device_matrix(res, 3, dim); + raft::copy(device_dataset.data_handle(), + dataset.data_handle(), + dataset.size(), + raft::resource::get_cuda_stream(res)); + + auto graph_host = raft::make_host_matrix(3, 2); + graph_host(0, 0) = 2; + graph_host(0, 1) = 1; + graph_host(1, 0) = 0; + graph_host(1, 1) = 2; + graph_host(2, 0) = 0; + graph_host(2, 1) = 1; + auto graph = raft::make_device_matrix(res, 3, 2); + raft::copy(graph.data_handle(), + graph_host.data_handle(), + graph.size(), + raft::resource::get_cuda_stream(res)); + detail::graph::sort_knn_graph_device_inplace(res, + cuvs::distance::DistanceType::InnerProduct, + raft::make_const_mdspan(device_dataset.view()), + graph.view()); + raft::copy(graph_host.data_handle(), + graph.data_handle(), + graph.size(), + raft::resource::get_cuda_stream(res)); + raft::resource::sync_stream(res); + EXPECT_EQ(graph_host(0, 0), 1); +} + +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] From e2ce96f2104e9f35b68d75db543b36e2f0761d47 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 18:43:27 +0000 Subject: [PATCH 04/13] pedantic docstring change --- cpp/include/cuvs/neighbors/cagra.hpp | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index fe46ca3132..962041571b 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -2508,7 +2508,9 @@ auto merge(raft::resources const& res, cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; -/** @brief Merge with explicit C++ algorithm and Fastener controls. */ +/** @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, From a9149dbce8630b95d4d840f620eecb2f6a85e185 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 19:15:05 +0000 Subject: [PATCH 05/13] assert/validation changes --- .../neighbors/detail/cagra/cagra_merge.cuh | 24 ++++++++++--------- .../detail/cagra/cagra_merge_scaffold.cuh | 6 +++-- 2 files changed, 17 insertions(+), 13 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 147a59c522..cbc50c02fe 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -220,8 +220,8 @@ auto preflight_fastener(cagra::index_params const& params, -> fastener_preflight_result { fastener_preflight_result result; - auto reject = [&](char const* reason) { - result.reason = reason; + auto reject = [&](std::string reason) { + result.reason = std::move(reason); return result; }; @@ -239,25 +239,29 @@ auto preflight_fastener(cagra::index_params const& params, 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 > 32 || - merge_params.lower_fanout < 1 || merge_params.lower_fanout > 32) { - return reject("root_fanout and lower_fanout must be between 1 and 32"); + 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 > 8192) { - return reject("max_leaders must be between 1 and 8192"); + 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 256"); + 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 8"); + 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; @@ -360,7 +364,6 @@ auto merge_fastener(raft::resources const& handle, scaffold_params.leaf_size = merge_params.leaf_size; scaffold_params.leaf_degree = merge_params.leaf_degree; - // Graph combination: append shifted input graphs to the cross-input scaffold. auto merged_graph = [&] { raft::common::nvtx::range scope("cagra::merge/scaffold"); auto scaffold = merge_scaffold::build( @@ -377,7 +380,6 @@ auto merge_fastener(raft::resources const& handle, handle, raft::make_const_mdspan(merged_graph.view()), params.graph_degree); } - // Optimization: reduce the sorted candidate prefix to the requested final CAGRA graph. auto optimized_graph = raft::make_device_matrix(handle, preflight.rows, params.graph_degree); { diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh index 3342d06530..cd228608a1 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -1175,12 +1175,14 @@ auto build(raft::resources const& res, 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 32"); + "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 8192"); + "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); From cf78e323154a87973c3f5174898119ea03e35806 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 21:08:37 +0000 Subject: [PATCH 06/13] test improvements --- cpp/tests/neighbors/ann_cagra.cuh | 7 +- .../ann_cagra/test_merge_fastener.cu | 413 ++++++++++-------- 2 files changed, 232 insertions(+), 188 deletions(-) diff --git a/cpp/tests/neighbors/ann_cagra.cuh b/cpp/tests/neighbors/ann_cagra.cuh index 79bf339827..e5c0c66c46 100644 --- a/cpp/tests/neighbors/ann_cagra.cuh +++ b/cpp/tests/neighbors/ann_cagra.cuh @@ -277,6 +277,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) @@ -1403,7 +1405,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 index 22a33e8238..e80becf20c 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -3,8 +3,13 @@ * 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 "../naive_knn.cuh" +#include "../ann_cagra.cuh" #include #include @@ -14,10 +19,11 @@ #include #include #include +#include -#include #include +#include #include #include #include @@ -27,30 +33,38 @@ namespace cuvs::neighbors::cagra { namespace { template -auto make_value(int64_t row, int64_t column) -> T +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) { - int value = static_cast((row * 7 + column * 3) % 23) - 11; - if constexpr (std::is_same_v) { - return __float2half(static_cast(value) / 8.0f); - } else if constexpr (std::is_same_v) { - return static_cast(value) / 8.0f; - } else if constexpr (std::is_same_v) { - return static_cast(value + 32); - } else { - return static_cast(value); - } + 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 make_dataset(raft::resources const& res, int64_t rows, int64_t dim, int64_t offset) +auto concat_host_datasets(raft::resources const& res, + raft::host_matrix_view first, + raft::host_matrix_view second) { - auto dataset = raft::make_host_matrix(res, rows, dim); - for (int64_t i = 0; i < rows; ++i) { - for (int64_t j = 0; j < dim; ++j) { - dataset(i, j) = make_value(i + offset, j); - } - } - return dataset; + 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) @@ -88,21 +102,20 @@ void expect_valid_graph(index const& merged, int64_t rows, int64_t template void expect_dataset_order(raft::resources const& res, index const& merged, - int64_t rows, - int64_t dim) + raft::host_matrix_view expected) { - ASSERT_EQ(merged.dataset().extent(0), rows); - ASSERT_EQ(merged.dataset().extent(1), dim); - auto host = raft::make_host_matrix(res, rows, dim); + ASSERT_EQ(merged.dataset().extent(0), expected.extent(0)); + ASSERT_EQ(merged.dataset().extent(1), expected.extent(1)); + auto host = raft::make_host_matrix(res, expected.extent(0), expected.extent(1)); raft::copy(host.data_handle(), merged.dataset().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 < dim; ++column) { + 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(make_value(row, column))); + static_cast(expected(row, column))); } } } @@ -114,10 +127,12 @@ void run_explicit_fastener(cuvs::distance::DistanceType metric) constexpr int64_t rows = 48; constexpr int64_t dim = 8; constexpr int64_t degree = 4; - auto dataset0 = make_dataset(res, rows, dim, 0); - auto dataset1 = make_dataset(res, rows, dim, rows); - auto graph0 = make_ring_graph(res, rows, degree); - auto graph1 = make_ring_graph(res, rows, degree); + 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( @@ -137,7 +152,7 @@ void run_explicit_fastener(cuvs::distance::DistanceType metric) auto merged = merge(res, params, indices, fastener); expect_valid_graph(merged, rows * 2, degree); - expect_dataset_order(res, merged, rows * 2, dim); + 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); @@ -149,6 +164,10 @@ void run_explicit_fastener(cuvs::distance::DistanceType metric) 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; @@ -191,6 +210,10 @@ TEST(CagraMergeFastener, PlanSplitIsDeterministicDenseAndCompact) 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; @@ -198,7 +221,7 @@ TEST(CagraMergeFastener, SplitManywayCarriesSmallParentsAndSupportsRepeatedLevel 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, 0); + 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); @@ -271,6 +294,10 @@ TEST(CagraMergeFastener, SplitManywayCarriesSmallParentsAndSupportsRepeatedLevel } } +/** + * 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; @@ -293,6 +320,74 @@ TEST(CagraMergeFastener, InitializesUnwrittenScaffoldSlotsWithSelf) } } +/** + * 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 an integer + * dimension beyond the safe accumulation limit is rejected before either input index is mutated. + */ TEST(CagraMergeFastener, LeafGemmLimitsRejectOnlyUnrealisticIntegerDimensions) { using namespace detail::merge_scaffold; @@ -306,8 +401,8 @@ TEST(CagraMergeFastener, LeafGemmLimitsRejectOnlyUnrealisticIntegerDimensions) raft::resources res; constexpr int64_t rows = 2; constexpr int64_t dim = MAX_INTEGER_LEAF_DIMENSION + 1; - auto dataset0 = make_dataset(res, rows, dim, 0); - auto dataset1 = make_dataset(res, rows, dim, rows); + 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; @@ -333,6 +428,10 @@ TEST(CagraMergeFastener, LeafGemmLimitsRejectOnlyUnrealisticIntegerDimensions) 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; @@ -342,14 +441,20 @@ TEST(CagraMergeFastener, SupportsAllScalarTypes) 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, 0); - auto dataset1 = make_dataset(res, rows, dim, rows); + 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(), @@ -381,7 +486,7 @@ TEST(CagraMergeFastener, MixedDatasetOwnershipPreservesInputs) auto merged = merge(res, params, indices, fastener); expect_valid_graph(merged, rows * 2, degree); - expect_dataset_order(res, merged, rows * 2, dim); + 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); @@ -389,112 +494,65 @@ TEST(CagraMergeFastener, MixedDatasetOwnershipPreservesInputs) EXPECT_FALSE(index1.data().is_owning()); } -TEST(CagraMergeFastener, DefaultFloatMergeSearchRecallAgainstBruteForce) +/** + * 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() { - raft::resources res; - auto stream = raft::resource::get_cuda_stream(res); - constexpr int64_t part_rows = 256; - constexpr int64_t rows = 2 * part_rows; - constexpr int64_t dim = 16; - constexpr int64_t queries = 16; - constexpr int64_t k = 12; - - auto value = [](int64_t row, int64_t column) { - uint64_t bits = (static_cast(row) + 1) * 0x9e3779b97f4a7c15ull; - bits ^= (static_cast(column) + 1) * 0xbf58476d1ce4e5b9ull; - bits ^= bits >> 30; - bits *= 0x94d049bb133111ebull; - bits ^= bits >> 31; - return static_cast(bits & 0xffffu) / 65535.0f; - }; - - auto dataset0 = raft::make_host_matrix(part_rows, dim); - auto dataset1 = raft::make_host_matrix(part_rows, dim); - for (int64_t row = 0; row < part_rows; ++row) { - for (int64_t column = 0; column < dim; ++column) { - dataset0(row, column) = value(row, column); - dataset1(row, column) = value(row + part_rows, column); - } - } - - index_params partition_params; - partition_params.metric = cuvs::distance::DistanceType::L2Expanded; - partition_params.graph_degree = 16; - partition_params.intermediate_graph_degree = 32; - partition_params.attach_dataset_on_build = true; - partition_params.guarantee_connectivity = false; - auto index0 = build(res, partition_params, raft::make_const_mdspan(dataset0.view())); - auto index1 = build(res, partition_params, raft::make_const_mdspan(dataset1.view())); - std::vector*> indices{&index0, &index1}; - - merge_params fastener; - fastener.algo = merge_algo::FASTENER; - auto merged = merge(res, partition_params, indices, fastener); - expect_valid_graph(merged, rows, partition_params.graph_degree); - - auto query_host = raft::make_host_matrix(queries, dim); - for (int64_t query = 0; query < queries; ++query) { - int64_t source = query * 31; - for (int64_t column = 0; column < dim; ++column) { - query_host(query, column) = value(source, column); - } + // 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}, + {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; } - auto query_device = raft::make_device_matrix(res, queries, dim); - raft::copy(query_device.data_handle(), query_host.data_handle(), query_host.size(), stream); - auto actual_indices = raft::make_device_matrix(res, queries, k); - auto actual_distances = raft::make_device_matrix(res, queries, k); - search_params search_params; - search_params.itopk_size = 32; - search_params.search_width = 4; - search_params.max_iterations = 0; - search(res, - search_params, - merged, - raft::make_const_mdspan(query_device.view()), - actual_indices.view(), - actual_distances.view()); - - auto truth_indices = raft::make_device_matrix(res, queries, k); - auto truth_distances = raft::make_device_matrix(res, queries, k); - cuvs::neighbors::naive_knn(res, - truth_distances.data_handle(), - truth_indices.data_handle(), - query_device.data_handle(), - merged.dataset().data_handle(), - static_cast(queries), - static_cast(rows), - static_cast(dim), - static_cast(k), - cuvs::distance::DistanceType::L2Expanded); - - std::vector actual(static_cast(queries * k)); - std::vector truth(static_cast(queries * k)); - raft::copy(actual.data(), actual_indices.data_handle(), actual.size(), stream); - raft::copy(truth.data(), truth_indices.data_handle(), truth.size(), stream); - raft::resource::sync_stream(res); + return inputs; +} +} // namespace - size_t matches = 0; - for (int64_t query = 0; query < queries; ++query) { - for (int64_t candidate = 0; candidate < k; ++candidate) { - auto id = actual[query * k + candidate]; - for (int64_t expected = 0; expected < k; ++expected) { - if (id == truth[query * k + expected]) { - ++matches; - break; - } - } - } - } - EXPECT_GE(static_cast(matches) / static_cast(queries * k), 0.60); +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, 0); - auto dataset1 = make_dataset(res, rows, dim, rows); + 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, @@ -525,11 +583,15 @@ TEST(CagraMergeFastener, MixedDegreesAndThreeLevelUint8OptionsProduceExactDegree 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, 0); - auto dataset1 = make_dataset(res, 2, 1, 2); + 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, @@ -574,14 +636,18 @@ TEST(CagraMergeFastener, AppendCyclicallyPadsMixedInputDegrees) } } +/** + * 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, 0); - auto dataset1 = make_dataset(res, rows, dim, rows); + 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, @@ -622,13 +688,17 @@ auto make_float_indices(raft::resources const& res, 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, 0); - auto dataset1 = make_dataset(res, rows, dim, rows); + 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; @@ -677,13 +747,17 @@ TEST(CagraMergeFastener, InvalidManywayOptionsFailPreflightWithoutMutation) 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, 0); - auto dataset1 = make_dataset(res, rows, dim, rows); + 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; @@ -737,6 +811,10 @@ TEST(CagraMergeFastener, DispatchRejectsOrFallsBackBeforeMutation) } } +/** + * 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; @@ -787,49 +865,10 @@ TEST(CagraMergeFastener, InnerProductOrderingDiffersFromL2) EXPECT_EQ(ip_host(0, 0), 2); } -TEST(CagraMergeFastener, Uint8InnerProductSortUsesExactWideAccumulation) -{ - raft::resources res; - constexpr int64_t dim = 1024; - auto dataset = raft::make_host_matrix(3, dim); - for (int64_t d = 0; d < dim; ++d) { - dataset(0, d) = 255; - dataset(1, d) = 255; - dataset(2, d) = 255; - } - dataset(0, 0) = 1; - dataset(1, 0) = 1; - dataset(2, 0) = 0; - auto device_dataset = raft::make_device_matrix(res, 3, dim); - raft::copy(device_dataset.data_handle(), - dataset.data_handle(), - dataset.size(), - raft::resource::get_cuda_stream(res)); - - auto graph_host = raft::make_host_matrix(3, 2); - graph_host(0, 0) = 2; - graph_host(0, 1) = 1; - graph_host(1, 0) = 0; - graph_host(1, 1) = 2; - graph_host(2, 0) = 0; - graph_host(2, 1) = 1; - auto graph = raft::make_device_matrix(res, 3, 2); - raft::copy(graph.data_handle(), - graph_host.data_handle(), - graph.size(), - raft::resource::get_cuda_stream(res)); - detail::graph::sort_knn_graph_device_inplace(res, - cuvs::distance::DistanceType::InnerProduct, - raft::make_const_mdspan(device_dataset.view()), - graph.view()); - raft::copy(graph_host.data_handle(), - graph.data_handle(), - graph.size(), - raft::resource::get_cuda_stream(res)); - raft::resource::sync_stream(res); - EXPECT_EQ(graph_host(0, 0), 1); -} - +/** + * 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; From 209485c2947d7cff2484ac23954136babd7ac401 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 21:09:19 +0000 Subject: [PATCH 07/13] assert/message changes & aligned signatures --- cpp/src/neighbors/cagra.cuh | 2 +- .../neighbors/detail/cagra/cagra_merge.cuh | 20 +++++++++---------- 2 files changed, 11 insertions(+), 11 deletions(-) diff --git a/cpp/src/neighbors/cagra.cuh b/cpp/src/neighbors/cagra.cuh index 490c225367..484e2b599d 100644 --- a/cpp/src/neighbors/cagra.cuh +++ b/cpp/src/neighbors/cagra.cuh @@ -434,7 +434,7 @@ index merge(raft::resources const& handle, const cagra::merge_params& merge_params, const cuvs::neighbors::filtering::base_filter& row_filter) { - return cagra::detail::merge(handle, params, merge_params, indices, row_filter); + return cagra::detail::merge(handle, params, indices, merge_params, row_filter); } /** @} */ // end group cagra diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index cbc50c02fe..e9c967640f 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -241,27 +241,27 @@ auto preflight_fastener(cagra::index_params const& params, 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)]); + 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)]); + 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)]); + 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)]); + 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; @@ -410,8 +410,8 @@ auto merge_fastener(raft::resources const& handle, template index merge(raft::resources const& handle, cagra::index_params const& params, - cagra::merge_params const& merge_params, std::vector*>& indices, + cagra::merge_params const& merge_params, cuvs::neighbors::filtering::base_filter const& row_filter) { raft::common::nvtx::range merge_scope( @@ -442,7 +442,7 @@ index merge(raft::resources const& handle, std::vector*>& indices, cuvs::neighbors::filtering::base_filter const& row_filter) { - return merge(handle, params, cagra::merge_params{}, indices, row_filter); + return merge(handle, params, indices, cagra::merge_params{}, row_filter); } } // namespace cuvs::neighbors::cagra::detail From da96f2620bb381bd018b12f9fa109211a018fbec Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 21:26:02 +0000 Subject: [PATCH 08/13] pre-commit hooks --- .../ann/src/cuvs/cuvs_ann_bench_param_parser.h | 2 +- cpp/include/cuvs/neighbors/cagra.hpp | 6 +++--- .../detail/cagra/cagra_merge_scaffold.cuh | 2 +- .../neighbors/ann_cagra/test_merge_fastener.cu | 15 +++++++-------- 4 files changed, 12 insertions(+), 13 deletions(-) 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 3fee950087..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 */ diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index 962041571b..e1bc1ac851 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -2508,9 +2508,9 @@ auto merge(raft::resources const& res, cuvs::neighbors::filtering::none_sample_filter{}) -> cuvs::neighbors::cagra::index; -/** @copydoc merge -* @param[in] merge_params Parameters for the merge. -*/ +/** @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, diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh index cd228608a1..2e0d2769c6 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ #pragma once diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu index e80becf20c..64dcfc90de 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -1,5 +1,5 @@ /* - * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. + * SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved. * SPDX-License-Identifier: Apache-2.0 */ @@ -60,8 +60,8 @@ auto concat_host_datasets(raft::resources const& res, 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)); + 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; @@ -114,8 +114,7 @@ void expect_dataset_order(raft::resources const& res, 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))); + EXPECT_EQ(static_cast(host(row, column)), static_cast(expected(row, column))); } } } @@ -129,7 +128,7 @@ void run_explicit_fastener(cuvs::distance::DistanceType metric) 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( + 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); @@ -453,9 +452,9 @@ TEST(CagraMergeFastener, MixedDatasetOwnershipPreservesInputs) 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( + 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); + auto device_dataset1 = raft::make_device_matrix(res, rows, dim); raft::copy(device_dataset1.data_handle(), dataset1.data_handle(), dataset1.size(), From 9a3cabf52d46543fd1c2970df9edc9a90ab72e99 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 22:16:39 +0000 Subject: [PATCH 09/13] clearer check rationale --- .../neighbors/detail/cagra/cagra_merge.cuh | 11 ++- .../detail/cagra/cagra_merge_scaffold.cuh | 85 ++++++++++--------- cpp/src/neighbors/detail/cagra/graph_core.cuh | 25 ++++-- 3 files changed, 67 insertions(+), 54 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index e9c967640f..d6e40adb79 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -266,14 +266,13 @@ auto preflight_fastener(cagra::index_params const& params, uint64_t const max_spill = std::numeric_limits::max() / merge_params.leaf_degree; uint64_t spill = merge_params.root_fanout; - if (spill > max_spill) { - return reject("root_fanout * lower_fanout^(levels - 1) * leaf_degree must not exceed 255"); - } + 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("root_fanout * lower_fanout^(levels - 1) * leaf_degree must not exceed 255"); - } + if (spill > max_spill / merge_params.lower_fanout) { return reject(candidate_width_limit); } spill *= merge_params.lower_fanout; } } diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh index 2e0d2769c6..245b2620b6 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -15,6 +15,7 @@ #include #include #include +#include #include #include @@ -48,7 +49,8 @@ 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 WARP_SIZE = 32; +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; // Centered INT8 dot products accumulate into INT32. Overflow requires more than 131,071 @@ -56,11 +58,12 @@ inline constexpr int ROW_WARPS_PER_BLOCK = 4; inline constexpr int64_t MAX_INTEGER_LEAF_DIMENSION = std::numeric_limits::max() / (128 * 128); -/** Grid size for the grid-stride kernels: blocks of 256 threads covering `items`, capped so +/** 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 + 255) / 256, 1048576)); + return static_cast(std::min((items + THREADS_PER_BLOCK - 1) / THREADS_PER_BLOCK, + MAX_STRIDED_GRID_BLOCKS)); } // ---------------------------------------------------------------------------- @@ -74,7 +77,7 @@ struct build_params { uint32_t lower_fanout = 3; double leader_fraction = 0.02; uint32_t max_leaders = 1000; - uint32_t leaf_size = 256; + uint32_t leaf_size = MAX_LEAF_SIZE; uint32_t leaf_degree = 4; }; @@ -83,7 +86,7 @@ struct split_params { uint32_t fanout = 1; double leader_fraction = 0.02; uint32_t max_leaders = 1000; - uint32_t leaf_size = 256; + uint32_t leaf_size = MAX_LEAF_SIZE; uint32_t level = 0; uint32_t occurrence_stride = 1; }; @@ -337,8 +340,9 @@ inline auto make_root_partition(raft::resources const& res, int64_t rows) -> par 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 + 255) / 256); - initialize_root_memberships_kernel<<>>(root.memberships.data(), 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; } @@ -430,8 +434,8 @@ static __global__ void manyway_select_tiles_kernel(float const* dots, uint32_t* output_keys, partition_membership* output_memberships) { - int lane = threadIdx.x % WARP_SIZE; - int warp = threadIdx.x / WARP_SIZE; + int lane = threadIdx.x % raft::WarpSize; + int warp = threadIdx.x / raft::WarpSize; int64_t linear_row = static_cast(blockIdx.x) * ROW_WARPS_PER_BLOCK + warp; int64_t total_rows = static_cast(batch_size) * tile_rows; if (linear_row >= total_rows) { return; } @@ -444,7 +448,7 @@ static __global__ void manyway_select_tiles_kernel(float const* dots, auto membership = input_memberships[input_index]; leader_candidate local_best[MAX_FANOUT]; - for (int leader = lane; leader < tile.leader_count; leader += WARP_SIZE) { + for (int leader = lane; leader < tile.leader_count; leader += raft::WarpSize) { uint32_t leader_id = leader_ids[static_cast(batch) * padded_leaders + leader]; float dot = dots[(static_cast(batch) * tile_rows + row) * padded_leaders + leader]; float distance = fmaxf(0.0f, norms[membership.id] + norms[leader_id] - 2.0f * dot); @@ -454,7 +458,7 @@ static __global__ void manyway_select_tiles_kernel(float const* dots, extern __shared__ unsigned char shared_bytes[]; auto* shared_candidates = reinterpret_cast(shared_bytes); - int shared_base = (warp * WARP_SIZE + lane) * fanout; + int shared_base = (warp * raft::WarpSize + lane) * fanout; for (int j = 0; j < fanout; ++j) { shared_candidates[shared_base + j] = local_best[j]; } @@ -462,8 +466,8 @@ static __global__ void manyway_select_tiles_kernel(float const* dots, if (lane == 0) { leader_candidate selected[MAX_FANOUT]; - int warp_base = warp * WARP_SIZE * fanout; - for (int candidate = 0; candidate < WARP_SIZE * fanout; ++candidate) { + int warp_base = warp * raft::WarpSize * fanout; + for (int candidate = 0; candidate < raft::WarpSize * fanout; ++candidate) { if (shared_candidates[warp_base + candidate].valid()) { manyway_insert_local(shared_candidates[warp_base + candidate], fanout, selected); } @@ -539,7 +543,10 @@ inline void carry_parents(raft::resources const& res, 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()), 256, 0, 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()); } @@ -654,7 +661,7 @@ void assign_bucket(raft::resources const& res, auto leader_stride = static_cast(leader_elements); auto dot_stride = static_cast(dot_elements); size_t selection_shared = - ROW_WARPS_PER_BLOCK * WARP_SIZE * params.fanout * sizeof(leader_candidate); + ROW_WARPS_PER_BLOCK * raft::WarpSize * params.fanout * sizeof(leader_candidate); 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); @@ -662,7 +669,7 @@ void assign_bucket(raft::resources const& res, // 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<<>>( + manyway_gather_tile_points_kernel<<>>( dataset.data_handle(), dim, parents.memberships.data(), @@ -673,7 +680,7 @@ void assign_bucket(raft::resources const& res, RAFT_CUDA_TRY(cudaGetLastError()); int leader_blocks = strided_grid_size(static_cast(batch_size * leader_elements)); - manyway_gather_tile_leaders_kernel<<>>( + manyway_gather_tile_leaders_kernel<<>>( dataset.data_handle(), dim, parents.memberships.data(), @@ -702,7 +709,7 @@ void assign_bucket(raft::resources const& res, static_cast((static_cast(batch_size) * tile_rows + ROW_WARPS_PER_BLOCK - 1) / ROW_WARPS_PER_BLOCK); manyway_select_tiles_kernel<<>>(tile_dots.data(), static_cast(batch_size), @@ -991,7 +998,7 @@ auto build_leaf_neighbors(raft::resources const& res, // 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<<>>( + initialize_self_scaffold_kernel<<>>( graph.data_handle(), rows, union_degree); RAFT_CUDA_TRY(cudaGetLastError()); @@ -1016,7 +1023,7 @@ auto build_leaf_neighbors(raft::resources const& res, int gather_blocks = strided_grid_size(static_cast(batch_size * vector_elements_per_leaf)); // Gather this batch's leaves into the dense vector buffer. - manyway_gather_leaf_vectors_kernel<<>>( + manyway_gather_leaf_vectors_kernel<<>>( dataset.data_handle(), input_dimension, input_dimension, @@ -1076,7 +1083,7 @@ auto build_leaf_neighbors(raft::resources const& res, 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<<>>( + manyway_gather_leaf_vectors_kernel<<>>( dataset.data_handle(), input_dimension, padded_dimension, @@ -1141,20 +1148,20 @@ static __global__ void initialize_origins_kernel(uint32_t* origins, template __global__ void manyway_l2_norms_kernel(T const* dataset, int64_t rows, int64_t dim, float* norms) { - int lane = threadIdx.x % WARP_SIZE; - int warp = threadIdx.x / WARP_SIZE; + 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 += WARP_SIZE) { + 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 = WARP_SIZE / 2; offset > 0; offset /= 2) { + for (int offset = raft::WarpSize / 2; offset > 0; offset /= 2) { sum += __shfl_down_sync(0xffffffffu, sum, offset); } if (lane == 0) { norms[row] = sum; } @@ -1201,7 +1208,8 @@ auto build(raft::resources const& res, } // 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 255"); + "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), @@ -1213,15 +1221,15 @@ auto build(raft::resources const& res, 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 + 255) / 256); - initialize_origins_kernel<<>>( + 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<<>>( + manyway_l2_norms_kernel<<>>( dataset.data_handle(), rows, dataset.extent(1), context.norms.data()); RAFT_CUDA_TRY(cudaGetLastError()); @@ -1310,8 +1318,8 @@ auto append_to_input_graphs( 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) + 255) / 256); - copy_partition_with_scaffold_kernel<<>>( + 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), @@ -1339,16 +1347,16 @@ static __global__ void deduplicate_graph_prefix_kernel(uint32_t const* input, uint32_t* output, int64_t output_degree) { - constexpr int WARPS_PER_BLOCK = 256 / WARP_SIZE; - int lane = threadIdx.x % WARP_SIZE; - int warp = threadIdx.x / WARP_SIZE; + 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 += WARP_SIZE) { + 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}; @@ -1358,7 +1366,7 @@ static __global__ void deduplicate_graph_prefix_kernel(uint32_t const* input, } unsigned first_mask = __ballot_sync(0xffffffffu, first); - unsigned lower_mask = lane == 0 ? 0u : (0xffffffffu >> (WARP_SIZE - lane)); + 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); @@ -1369,7 +1377,7 @@ static __global__ void deduplicate_graph_prefix_kernel(uint32_t const* input, selected = 1; } if (selected > output_degree) { selected = static_cast(output_degree); } - for (int64_t column = selected + lane; column < output_degree; column += WARP_SIZE) { + for (int64_t column = selected + lane; column < output_degree; column += raft::WarpSize) { output[output_base + column] = output[output_base + (column % selected)]; } } @@ -1383,8 +1391,7 @@ inline auto cap_sorted_graph( 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 THREADS_PER_BLOCK = 256; - constexpr int WARPS_PER_BLOCK = THREADS_PER_BLOCK / WARP_SIZE; + 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<<(raft::WarpSize) * kMaxSortElementsPerThread; + template using sort_kernel_type = void (*)(DataT const*, IdxT, uint32_t, IdxT*, uint32_t, uint32_t, cuvs::distance::DistanceType); @@ -182,15 +187,16 @@ using sort_kernel_type = template auto select_sort_kernel(uint64_t degree) -> sort_kernel_type { - if (degree <= 32) { return kern_sort; } - if (degree <= 64) { return kern_sort; } - if (degree <= 128) { return kern_sort; } - if (degree <= 256) { return kern_sort; } - if (degree <= 512) { return kern_sort; } - if (degree <= 1024) { return kern_sort; } + 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 1024.", - degree); + "The degree of input knn graph is too large (%lu). It must be equal to or smaller than %lu.", + degree, + kMaxSortDegree); } template @@ -1017,7 +1023,8 @@ void sort_knn_graph_device_inplace( 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 <= 1024, "Graph degree must be in [1, 1024]"); + RAFT_EXPECTS( + degree > 0 && degree <= kMaxSortDegree, "Graph degree must be in [1, %lu]", kMaxSortDegree); auto kernel = select_sort_kernel(degree); From af25e86d5da863623eb0984413484f4bb3b6f463 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Wed, 22 Jul 2026 22:18:54 +0000 Subject: [PATCH 10/13] test arg mismatch --- cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu index 64dcfc90de..9bcdd46779 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -506,7 +506,8 @@ std::vector generate_fastener_merge_recall_inputs() {100}, {1000}, {16}, - {16}, + {16}, // k + {32}, // graph_degree {graph_build_algo::IVF_PQ, graph_build_algo::NN_DESCENT}, {search_algo::AUTO}, {10}, From 2eb64fafef22c6274559db43b06d676bcc65bbd4 Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Thu, 23 Jul 2026 21:05:43 +0000 Subject: [PATCH 11/13] refactor to use existing select_k kernel --- .../detail/cagra/cagra_merge_scaffold.cuh | 239 ++++++++---------- 1 file changed, 109 insertions(+), 130 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh index 245b2620b6..18dd2415c6 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -7,6 +7,7 @@ #include "../neighbors_device_intrinsics.cuh" #include +#include #include #include @@ -281,51 +282,6 @@ inline auto bucket_split_parents(split_plan const& plan, split_params const& par return buckets; } -/** One candidate assignment of a row to a leader. Defaults to the invalid sentinel, which sorts - * after every real candidate. */ -struct leader_candidate { - float distance = std::numeric_limits::max(); - uint32_t leader = std::numeric_limits::max(); - - __device__ auto valid() const -> bool { return leader != std::numeric_limits::max(); } -}; - -/** Nearer candidates order first. Equal distances order by the lower leader index, keeping the - * selection deterministic. */ -__device__ inline auto operator<(leader_candidate a, leader_candidate b) -> bool -{ - return a.distance < b.distance || (a.distance == b.distance && a.leader < b.leader); -} - -/** Replace the worst entry in the best list if the new candidate is better. */ -__device__ inline void manyway_insert_local(leader_candidate candidate, - int fanout, - leader_candidate* best) -{ - int worst = 0; - for (int j = 1; j < fanout; ++j) { - if (best[worst] < best[j]) { worst = j; } - } - if (candidate < best[worst]) { best[worst] = candidate; } -} - -/** Sort the selected candidates in place from the nearest to the farthest. - * - * This is hand-rolled because thrust::sort has an allocation and I can't get a cuda::std::sort to - * compile. */ -__device__ inline void manyway_sort_selected(int fanout, leader_candidate* selected) -{ - for (int i = 1; i < fanout; ++i) { - leader_candidate candidate = selected[i]; - int j = i; - while (j > 0 && candidate < selected[j - 1]) { - selected[j] = selected[j - 1]; - --j; - } - selected[j] = candidate; - } -} - /** Write the identity membership for each dataset row. */ static __global__ void initialize_root_memberships_kernel(partition_membership* memberships, int64_t rows) @@ -419,64 +375,69 @@ __global__ void manyway_gather_tile_leaders_kernel(T const* dataset, } } -/** Select the nearest leaders for each tile row from the dot products. Write the child keys and - * memberships. One warp does one row. */ -static __global__ void manyway_select_tiles_kernel(float const* dots, - int batch_size, - int tile_rows, - int padded_leaders, - int fanout, - int occurrence_stride, - float const* norms, - uint32_t const* leader_ids, - partition_membership const* input_memberships, - assignment_tile const* tiles, - uint32_t* output_keys, - partition_membership* output_memberships) +/** 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) { - int lane = threadIdx.x % raft::WarpSize; - int warp = threadIdx.x / raft::WarpSize; - int64_t linear_row = static_cast(blockIdx.x) * ROW_WARPS_PER_BLOCK + warp; - int64_t total_rows = static_cast(batch_size) * tile_rows; - if (linear_row >= total_rows) { return; } - int batch = static_cast(linear_row / tile_rows); - int row = static_cast(linear_row % tile_rows); - auto tile = tiles[batch]; - if (row >= tile.rows) { return; } - - int64_t input_index = tile.input_start + row; - auto membership = input_memberships[input_index]; - leader_candidate local_best[MAX_FANOUT]; - - for (int leader = lane; leader < tile.leader_count; leader += raft::WarpSize) { - uint32_t leader_id = leader_ids[static_cast(batch) * padded_leaders + leader]; - float dot = dots[(static_cast(batch) * tile_rows + row) * padded_leaders + leader]; - float distance = fmaxf(0.0f, norms[membership.id] + norms[leader_id] - 2.0f * dot); - manyway_insert_local( - {.distance = distance, .leader = static_cast(leader)}, fanout, local_best); + 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; } +} - extern __shared__ unsigned char shared_bytes[]; - auto* shared_candidates = reinterpret_cast(shared_bytes); - int shared_base = (warp * raft::WarpSize + lane) * fanout; - for (int j = 0; j < fanout; ++j) { - shared_candidates[shared_base + j] = local_best[j]; - } - __syncwarp(); - - if (lane == 0) { - leader_candidate selected[MAX_FANOUT]; - int warp_base = warp * raft::WarpSize * fanout; - for (int candidate = 0; candidate < raft::WarpSize * fanout; ++candidate) { - if (shared_candidates[warp_base + candidate].valid()) { - manyway_insert_local(shared_candidates[warp_base + candidate], fanout, selected); - } - } - manyway_sort_selected(fanout, selected); - int64_t output_base = - tile.output_start + (input_index - tile.input_start) * static_cast(fanout); +/** 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) { - output_keys[output_base + rank] = tile.child_key_base + selected[rank].leader; + 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), @@ -611,10 +572,9 @@ void batched_row_dot_products(raft::resources const& res, * 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, and the selection kernel - * converts dots to distances with the precomputed row norms (|x|^2 + |l|^2 - 2 x.l), keeps the - * `fanout` nearest leaders per row (ties broken by lower leader index), and writes the child keys - * and memberships. + * 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, @@ -640,12 +600,14 @@ void assign_bucket(raft::resources const& res, } // 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 bytes_per_batch = (point_elements + leader_elements + dot_elements) * sizeof(float) + - static_cast(padded_leaders) * sizeof(uint32_t) + - sizeof(assignment_tile); + 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( @@ -656,12 +618,12 @@ void assign_bucket(raft::resources const& res, 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); - size_t selection_shared = - ROW_WARPS_PER_BLOCK * raft::WarpSize * params.fanout * sizeof(leader_candidate); 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); @@ -704,25 +666,42 @@ void assign_bucket(raft::resources const& res, static_cast(dim), static_cast(batch_size)); - // Keep each row's `fanout` nearest leaders and write child keys and memberships - int selection_blocks = - static_cast((static_cast(batch_size) * tile_rows + ROW_WARPS_PER_BLOCK - 1) / - ROW_WARPS_PER_BLOCK); - manyway_select_tiles_kernel<<>>(tile_dots.data(), - static_cast(batch_size), - tile_rows, - padded_leaders, - static_cast(params.fanout), - static_cast(params.occurrence_stride), - context.norms.data(), - tile_leader_ids.data(), - parents.memberships.data(), - device_tiles.data(), - keys.data(), - memberships.data()); + // 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()); } } From 6f65f878edad0b086417a1a61e135cd534d19dcd Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Fri, 24 Jul 2026 00:05:26 +0000 Subject: [PATCH 12/13] fixed attached dataset alignment issue --- cpp/include/cuvs/neighbors/cagra.hpp | 3 ++- cpp/src/neighbors/detail/cagra/cagra_merge.cuh | 10 +++------- 2 files changed, 5 insertions(+), 8 deletions(-) diff --git a/cpp/include/cuvs/neighbors/cagra.hpp b/cpp/include/cuvs/neighbors/cagra.hpp index e1bc1ac851..0d1867b00f 100644 --- a/cpp/include/cuvs/neighbors/cagra.hpp +++ b/cpp/include/cuvs/neighbors/cagra.hpp @@ -2492,7 +2492,8 @@ struct merge_params { * `index_params::graph_degree` is the final output degree. * * 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 that consolidated copy. + * 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. * @param[in] params Parameters for the returned CAGRA index. diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index d6e40adb79..88db6218bb 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -390,13 +390,9 @@ auto merge_fastener(raft::resources const& handle, index merged_index(handle, params.metric); merged_index.update_graph(handle, std::move(optimized_graph)); if (params.attach_dataset_on_build) { - using matrix_t = decltype(dataset); - using layout_t = typename matrix_t::layout_type; - using container_policy_t = typename matrix_t::container_policy_type; - using owning_t = owning_dataset; - auto out_layout = raft::make_strided_layout(dataset.view().extents(), - cuda::std::array{preflight.dim, 1}); - merged_index.update_dataset(handle, owning_t{std::move(dataset), out_layout}); + // 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( From c8fd5407e86fcb5ec44cd0b1cdc2a56795558dbb Mon Sep 17 00:00:00 2001 From: Ben Landrum Date: Fri, 24 Jul 2026 19:08:35 +0000 Subject: [PATCH 13/13] CI fixes --- .../neighbors/detail/cagra/cagra_merge.cuh | 6 +- .../detail/cagra/cagra_merge_scaffold.cuh | 251 ++++++------------ .../ann_cagra/test_merge_fastener.cu | 49 ++-- 3 files changed, 116 insertions(+), 190 deletions(-) diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh index 88db6218bb..82a9fe280e 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge.cuh @@ -319,11 +319,7 @@ auto preflight_fastener(cagra::index_params const& params, return reject("dataset dimension must be positive and fit cuBLAS int dimensions"); } if (!merge_scaffold::leaf_gemm_supported(result.dim, merge_params.leaf_size)) { - if constexpr (std::is_integral_v) { - return reject("integer dataset dimension exceeds the INT32 leaf GEMM limit"); - } else { - return reject("dataset dimension exceeds the leaf GEMM workspace limit"); - } + 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"); diff --git a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh index 18dd2415c6..e75db0ca33 100644 --- a/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh +++ b/cpp/src/neighbors/detail/cagra/cagra_merge_scaffold.cuh @@ -54,10 +54,6 @@ 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; -// Centered INT8 dot products accumulate into INT32. Overflow requires more than 131,071 -// dimensions even at the worst-case magnitude of 128, which is not realistic. -inline constexpr int64_t MAX_INTEGER_LEAF_DIMENSION = - std::numeric_limits::max() / (128 * 128); /** Grid size for the grid-stride kernels: blocks of THREADS_PER_BLOCK covering `items`, capped so * oversized workloads loop within resident threads instead. */ @@ -517,28 +513,23 @@ inline void carry_parents(raft::resources const& res, * out_i[b * a_rows + a] = A_i[a] . B_i[b]. * * Pretty much a wrapper around `cublasGemmStridedBatchedEx`. */ -template -void batched_row_dot_products(raft::resources const& res, - DataT const* a, - int a_rows, - long long a_stride, - DataT const* b, - int b_rows, - long long b_stride, - AccT* out, - long long out_stride, - int row_width, - int batch_count) +/** 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) { - static_assert((std::is_same_v && std::is_same_v) || - (std::is_same_v && std::is_same_v), - "Unsupported batched dot product type combination"); - constexpr auto data_type = std::is_same_v ? CUDA_R_32F : CUDA_R_8I; - constexpr auto out_type = std::is_same_v ? CUDA_R_32F : CUDA_R_32I; - constexpr auto compute_type = - std::is_same_v ? CUBLAS_COMPUTE_32F : CUBLAS_COMPUTE_32I; - AccT alpha = 1; - AccT beta = 0; + 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, @@ -549,20 +540,20 @@ void batched_row_dot_products(raft::resources const& res, row_width, &alpha, a, - data_type, + CUDA_R_32F, row_width, a_stride, b, - data_type, + CUDA_R_32F, row_width, b_stride, &beta, out, - out_type, + CUDA_R_32F, a_rows, out_stride, batch_count, - compute_type, + CUBLAS_COMPUTE_32F, CUBLAS_GEMM_DEFAULT)); } @@ -762,27 +753,17 @@ auto split_manyway(raft::resources const& res, // Leaf processing: bounded leaf slicing and cross-input leaf KNN // ---------------------------------------------------------------------------- -/** Return true if one leaf of this dimension and leaf size goes into the GEMM workspace. */ +/** 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; } - int64_t storage_dimension = dimension; - size_t vector_element_size = sizeof(float); - size_t gram_element_size = sizeof(float); - if constexpr (std::is_integral_v) { - // The below check is a little paranoid - if (dimension > MAX_INTEGER_LEAF_DIMENSION) { return false; } - storage_dimension = (dimension + 3) & ~int64_t{3}; - vector_element_size = sizeof(int8_t); - gram_element_size = sizeof(int32_t); - } - - size_t vector_elements = static_cast(leaf_size) * storage_dimension; + size_t vector_elements = static_cast(leaf_size) * static_cast(dimension); size_t gram_elements = static_cast(leaf_size) * leaf_size; - return vector_elements * vector_element_size + gram_elements * gram_element_size <= - GEMM_WORKSPACE_BYTES; + return vector_elements * sizeof(float) + gram_elements * sizeof(float) <= GEMM_WORKSPACE_BYTES; } struct leaf_set { @@ -846,7 +827,7 @@ static __global__ void initialize_self_scaffold_kernel(uint32_t* graph, } /** Copy the vectors of each leaf into a dense buffer of OutT, zero-padding rows past the leaf - * end and dimensions past `input_dim`. Unsigned input values move to a center of zero. */ + * 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, @@ -872,11 +853,7 @@ __global__ void manyway_gather_leaf_vectors_kernel(T const* dataset, 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]; - if constexpr (std::is_same_v) { - value = static_cast(static_cast(input) - 128); - } else { - value = static_cast(input); - } + value = static_cast(input); } leaf_vectors[linear] = value; } @@ -981,124 +958,62 @@ auto build_leaf_neighbors(raft::resources const& res, graph.data_handle(), rows, union_degree); RAFT_CUDA_TRY(cudaGetLastError()); - if constexpr (std::is_same_v || std::is_same_v) { - // Size the leaf batch so the vector and Gram buffers fit the GEMM workspace budget - 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)); - // Gather this batch's leaves into the dense vector buffer. - 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()); - // the Gram matrix of every leaf in the batch - 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)); - // select each point's nearest cross-input neighbors from the corresponding Gram matrix - 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()); - } - } else if constexpr (std::is_same_v || std::is_same_v) { - // Same pipeline as the float path, with rows padded to a multiple of 4 and an INT32 Gram. - int64_t input_dimension = dataset.extent(1); - int64_t padded_dimension = (input_dimension + 3) & ~int64_t{3}; - size_t vector_elements_per_leaf = static_cast(leaf_size) * padded_dimension; - size_t gram_elements_per_leaf = static_cast(leaf_size) * leaf_size; - size_t bytes_per_leaf = - vector_elements_per_leaf * sizeof(int8_t) + gram_elements_per_leaf * sizeof(int32_t); - 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); - int dimension = static_cast(padded_dimension); - 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, - padded_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()); - } + // 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; diff --git a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu index 9bcdd46779..50c44f95d1 100644 --- a/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu +++ b/cpp/tests/neighbors/ann_cagra/test_merge_fastener.cu @@ -104,13 +104,22 @@ void expect_dataset_order(raft::resources const& res, index const& merged, raft::host_matrix_view expected) { - ASSERT_EQ(merged.dataset().extent(0), expected.extent(0)); - ASSERT_EQ(merged.dataset().extent(1), expected.extent(1)); - auto host = raft::make_host_matrix(res, expected.extent(0), expected.extent(1)); - raft::copy(host.data_handle(), - merged.dataset().data_handle(), - host.size(), - raft::resource::get_cuda_stream(res)); + 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) { @@ -384,22 +393,28 @@ TEST(CagraMergeFastener, DuplicateAcrossInputsAppearInEachOthersScaffoldNeighbor } /** - * Verifies the leaf-GEMM eligibility boundary for every scalar type and confirms that an integer - * dimension beyond the safe accumulation limit is rejected before either input index is mutated. + * 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, LeafGemmLimitsRejectOnlyUnrealisticIntegerDimensions) +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(MAX_INTEGER_LEAF_DIMENSION, 256)); - EXPECT_TRUE(leaf_gemm_supported(MAX_INTEGER_LEAF_DIMENSION, 256)); - EXPECT_FALSE(leaf_gemm_supported(MAX_INTEGER_LEAF_DIMENSION + 1, 256)); - EXPECT_FALSE(leaf_gemm_supported(MAX_INTEGER_LEAF_DIMENSION + 1, 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 = MAX_INTEGER_LEAF_DIMENSION + 1; + 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); @@ -416,13 +431,13 @@ TEST(CagraMergeFastener, LeafGemmLimitsRejectOnlyUnrealisticIntegerDimensions) params.intermediate_graph_degree = 1; merge_params fastener; fastener.algo = merge_algo::FASTENER; - fastener.leaf_size = 64; + 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, "integer dataset dimension exceeds the INT32 leaf GEMM limit"); + 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); }