Replace O(N*degree^2) CPU dedup with GPU warp-ballot kernel

cpp/src/neighbors/detail/nn_descent.cuh

@@ -1496,6 +1496,94 @@ void GNND<Data_t, Index_t>::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 <typename Index_t>
+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<int>(blockIdx.x) * warps_per_block + warp_id;
+  if (node_id >= nrow) { return; }
+
+  const int* in_row = d_in + static_cast<ptrdiff_t>(node_id) * graph_degree;
+  Index_t* out_row  = d_out + static_cast<ptrdiff_t>(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<uint64_t>(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<int>(rnd % static_cast<uint64_t>(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<Index_t>(warp_buf[j]);
+  }
+}
+
 template <typename Data_t, typename Index_t>
 template <typename DistEpilogue_t>
 void GNND<Data_t, Index_t>::build(Data_t* data,
@@ -1727,48 +1815,45 @@ void GNND<Data_t, Index_t>::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<decltype(idx)>(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<int>(build_config_.node_degree);
+    const int graph_degree = static_cast<int>(graph_.node_degree);
+
+    // Upload h_graph (InternalID_t = int32, same size as Index_t) to device.
+    static_assert(sizeof(InternalID_t<Index_t>) == sizeof(int),
+                  "dedup_graph_kernel assumes InternalID_t is int-sized");
+    const size_t h_graph_elems = static_cast<size_t>(nrow_) * graph_degree;
+    rmm::device_uvector<int> 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<size_t>(nrow_) * node_degree;
+    rmm::device_uvector<Index_t> 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<int>(nrow_) + warps_per_block - 1) / warps_per_block;
+    const size_t smem_bytes   = static_cast<size_t>(warps_per_block) * node_degree * sizeof(int);
+
+    dedup_graph_kernel<<<grid_size, block_size, smem_bytes, stream>>>(
+      d_in.data(), d_out.data(), static_cast<int>(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<uint64_t>(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<decltype(idx)>(output_neighbor_list_ptr[exi_j]) == idx || i == idx) {
-            dup = true;
-            break;
-          }
-        }
-      }
-      output_neighbor_list_ptr[j] = static_cast<int>(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;