From fd92c3670bc2dd72281ac30cae486c844b8ba2ba Mon Sep 17 00:00:00 2001 From: James Xia Date: Wed, 12 Aug 2026 07:44:18 -0700 Subject: [PATCH] Replace O(N*degree^2) CPU dedup with GPU warp-ballot kernel The graph shrink step in GNND::build() copied the NN-descent output while removing duplicate and self-referencing neighbor IDs using a nested scan: for each of the N nodes, each of the `node_degree` candidates was checked against all already-placed entries, giving O(N*degree^2) CPU work that scales poorly with graph degree and dataset size. Replace with a GPU kernel (dedup_graph_kernel) that runs one warp per node. The warp scans InternalID_t neighbors in original order, using __ballot_sync for O(warp-width) duplicate detection, and fills any remaining slots with xorshift64 random nodes. Original neighbor order is preserved, avoiding the recall regression seen with sort-based approaches. The H2D/D2H transfers are O(N*degree) -- the same order as a single pass over the graph -- while the replaced CPU work is O(N*degree^2), so the transfers are dominated by the savings at any practical degree. --- cpp/src/neighbors/detail/nn_descent.cuh | 167 ++++++++++++++++++------ 1 file changed, 126 insertions(+), 41 deletions(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index aff33040cf..752ae7fe11 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1496,6 +1496,94 @@ void GNND::local_join(cudaStream_t stream, DistEpilogue_t dist_ } } +// GPU kernel for deduplicating and shrinking the NN-descent output graph. +// +// One warp (32 threads) per node. The warp scans the node's `graph_degree` InternalID_t +// neighbors in the input (stored as raw int32) in order, emitting up to `node_degree` valid, +// non-self, non-duplicate IDs to the output while preserving original order. Duplicate +// detection uses warp-ballot: for each candidate, every earlier warp lane votes on whether it +// already holds that ID; any match means duplicate. Remaining output slots are filled with +// random non-duplicate non-self nodes via xorshift64, again using warp-ballot for O(warp) +// duplicate checks. +// +// grid: ceil(nrow / warps_per_block) blocks +// block: warps_per_block * 32 threads +template +RAFT_KERNEL dedup_graph_kernel(const int* __restrict__ d_in, + Index_t* __restrict__ d_out, + int nrow, + int graph_degree, + int node_degree) +{ + const int warp_id = threadIdx.x / 32; + const int lane = threadIdx.x % 32; + const int warps_per_block = blockDim.x / 32; + const int node_id = static_cast(blockIdx.x) * warps_per_block + warp_id; + if (node_id >= nrow) { return; } + + const int* in_row = d_in + static_cast(node_id) * graph_degree; + Index_t* out_row = d_out + static_cast(node_id) * node_degree; + + // Each thread in the warp holds one output slot (valid when slot < node_degree and filled). + // We process candidates one at a time; the whole warp votes on duplicates. + // For node_degree > 32 we keep a running register array isn't viable, so we use shared + // memory: one int32 array of node_degree per warp slot in the block. + extern __shared__ int smem[]; + int* warp_buf = smem + warp_id * node_degree; // node_degree ints per warp + + int out_count = 0; // number of valid entries placed so far + + // --- Pass 1: copy valid neighbors in original order, skipping duplicates and self --- + for (int j = 0; j < graph_degree && out_count < node_degree; j++) { + // Decode InternalID_t: raw value >= 0 means new (id = raw), < 0 means old (id = -raw-1). + int raw = in_row[j]; + int id = (raw >= 0) ? raw : (-raw - 1); + if (id >= nrow || id == node_id) { continue; } + + // Check for duplicate against already-placed entries in warp_buf. + bool dup = false; + for (int k = 0; k < out_count && !dup; k += 32) { + int check = (k + lane < out_count) ? warp_buf[k + lane] : -1; + unsigned mask = __ballot_sync(0xffffffff, check == id); + if (mask) { dup = true; } + } + if (!dup) { + warp_buf[out_count] = id; + out_count++; + } + } + + // --- Pass 2: fill remaining slots with random non-duplicate non-self nodes --- + // Cap attempts per slot to node_degree to avoid an infinite loop when nrow is small. + uint64_t rnd = static_cast(node_id) * node_degree + out_count + 1; + while (out_count < node_degree) { + bool filled = false; + for (int attempt = 0; attempt < node_degree && !filled; attempt++) { + rnd = cuvs::neighbors::detail::device::xorshift64(rnd); + int idx = static_cast(rnd % static_cast(nrow)); + if (idx == node_id) { continue; } + + bool dup = false; + for (int k = 0; k < out_count && !dup; k += 32) { + int check = (k + lane < out_count) ? warp_buf[k + lane] : -1; + unsigned mask = __ballot_sync(0xffffffff, check == idx); + if (mask) { dup = true; } + } + if (!dup) { + warp_buf[out_count] = idx; + out_count++; + filled = true; + } + } + if (!filled) { break; } + } + + // Write from shared memory to global output (coalesced within the warp). + for (int j = lane; j < node_degree; j += 32) { + out_row[j] = static_cast(warp_buf[j]); + } +} + template template void GNND::build(Data_t* data, @@ -1727,48 +1815,45 @@ void GNND::build(Data_t* data, Index_t* graph_shrink_buffer = (Index_t*)graph_.h_dists.data_handle(); - // Copy the output graph while removing duplicates. -#pragma omp parallel for - for (size_t i = 0; i < (size_t)nrow_; i++) { - auto output_neighbor_list_ptr = graph_shrink_buffer + i * build_config_.node_degree; - - size_t out_j = 0; - - // Copy neighbor list while removing duplicates. - for (size_t in_j = 0; in_j < build_config_.node_degree; in_j++) { - size_t idx = graph_.h_graph[i * graph_.node_degree + in_j].id(); - - bool dup = false; - for (size_t exi_j = 0; exi_j < out_j; exi_j++) { - if (static_cast(output_neighbor_list_ptr[exi_j]) == idx || i == idx) { - dup = true; - break; - } - } - if (!dup) { - output_neighbor_list_ptr[out_j] = idx; - out_j++; - } - } + // Deduplicate and shrink the graph on GPU. Each warp handles one node: it scans the + // InternalID_t neighbor list (pinned host -> device), removes duplicates and self-references + // while preserving original order using warp-ballot, then fills any remaining slots with + // random non-duplicate nodes via xorshift64. The two H2D/D2H transfers (~N*degree*4 bytes + // each) are far cheaper than the O(N * degree^2) CPU nested scan they replace. + { + cudaStream_t stream = raft::resource::get_cuda_stream(res); + const int node_degree = static_cast(build_config_.node_degree); + const int graph_degree = static_cast(graph_.node_degree); + + // Upload h_graph (InternalID_t = int32, same size as Index_t) to device. + static_assert(sizeof(InternalID_t) == sizeof(int), + "dedup_graph_kernel assumes InternalID_t is int-sized"); + const size_t h_graph_elems = static_cast(nrow_) * graph_degree; + rmm::device_uvector d_in(h_graph_elems, stream); + RAFT_CUDA_TRY(cudaMemcpyAsync( + d_in.data(), graph_.h_graph, h_graph_elems * sizeof(int), cudaMemcpyHostToDevice, stream)); + + const size_t out_elems = static_cast(nrow_) * node_degree; + rmm::device_uvector d_out(out_elems, stream); + + // One warp (32 threads) per node. Each thread covers one or more neighbor slots. + // shared memory per block: node_degree ints for the node's working neighbor list. + const int warps_per_block = 4; + const int block_size = warps_per_block * 32; + const int grid_size = (static_cast(nrow_) + warps_per_block - 1) / warps_per_block; + const size_t smem_bytes = static_cast(warps_per_block) * node_degree * sizeof(int); + + dedup_graph_kernel<<>>( + d_in.data(), d_out.data(), static_cast(nrow_), graph_degree, node_degree); + RAFT_CUDA_TRY(cudaPeekAtLastError()); - // Fill with random nodes if the length of the filled neighbor list is less than the degree. - for (size_t j = out_j; j < build_config_.node_degree; j++) { - uint64_t rnd = static_cast(i * build_config_.node_degree + j + 1); - uint64_t idx; - bool dup = true; - for (size_t attempts = 0; dup && attempts < build_config_.node_degree; attempts++) { - rnd = cuvs::neighbors::detail::device::xorshift64(rnd); - idx = rnd % nrow_; - dup = false; - for (size_t exi_j = 0; exi_j < j; exi_j++) { - if (static_cast(output_neighbor_list_ptr[exi_j]) == idx || i == idx) { - dup = true; - break; - } - } - } - output_neighbor_list_ptr[j] = static_cast(idx); - } + // Download the deduplicated graph to graph_shrink_buffer (pinned host memory). + RAFT_CUDA_TRY(cudaMemcpyAsync(graph_shrink_buffer, + d_out.data(), + out_elems * sizeof(Index_t), + cudaMemcpyDeviceToHost, + stream)); + raft::resource::sync_stream(res); } graph_.h_graph = nullptr;