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;