From d1e2dc0b9405104654e97ca80824907a3f3af0e5 Mon Sep 17 00:00:00 2001 From: James Xia Date: Tue, 18 Aug 2026 12:53:24 -0700 Subject: [PATCH 1/2] Replace O(nrow * node_degree^2) dedup with per-thread bit-packed bitset The final duplicate-removal pass in GNND::build() scanned every candidate against everything already placed in the row -- O(node_degree) per candidate. Replace with a per-thread bit-packed "seen" array (1 bit per dataset row), making the check O(1) and the whole pass O(nrow * node_degree). Bits are cleared immediately after each row finishes, so no per-row reset is needed, and the extra memory is a small, nrow-independent fraction of what the step already allocates. Also adds an idx >= nrow bounds check the original loop didn't have, guarding against a pre-existing edge case where an uninitialized sentinel value could reach the final output graph. Benchmarked up to 20M rows / 1024 degree against the original loop, a sort-based variant, a hashmap variant, and a GPU kernel; this wins every configuration tested (1.8x-8.5x over the next best), though the margin narrows with nrow as the per-thread array outgrows cache -- untested beyond 20M rows. --- cpp/src/neighbors/detail/nn_descent.cuh | 76 ++++++++++++++----------- 1 file changed, 44 insertions(+), 32 deletions(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index aff33040cf..253d933743 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1727,47 +1727,59 @@ 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; + // Copy the output graph while removing duplicates. Each thread keeps a bit-packed "seen" + // array, one bit per dataset row, to test and mark ids in O(1) as it scans a row's + // candidates. Only the ids actually placed for a row are ever set, and they're cleared again + // immediately after that row is done, so beyond the one-time zero-initialization when each + // thread starts, no reset across the full array is ever needed. + const size_t num_dedup_words = (static_cast(nrow_) + 63) / 64; +#pragma omp parallel + { + std::vector seen_bits(num_dedup_words, 0); - size_t out_j = 0; + auto test_and_set = [&](size_t idx) -> bool { + uint64_t mask = uint64_t{1} << (idx & 63); + if (seen_bits[idx >> 6] & mask) { return false; } + seen_bits[idx >> 6] |= mask; + return true; + }; - // 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(); +#pragma omp for + for (size_t i = 0; i < (size_t)nrow_; i++) { + auto output_neighbor_list_ptr = graph_shrink_buffer + i * build_config_.node_degree; - 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) { + 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(); + if (idx >= (size_t)nrow_ || idx == i || !test_and_set(idx)) continue; output_neighbor_list_ptr[out_j] = idx; out_j++; } - } - // 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; - } + // 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 = 0; + bool placed = false; + for (size_t attempts = 0; !placed && attempts < build_config_.node_degree; attempts++) { + rnd = cuvs::neighbors::detail::device::xorshift64(rnd); + idx = rnd % nrow_; + placed = (idx != i) && test_and_set(idx); } + output_neighbor_list_ptr[j] = static_cast(idx); + } + + // Unset every bit this row touched so the array is back to all-zero for the next row this + // thread processes. Safe to run over the full row rather than tracking exactly which + // entries were newly set: clearing an already-clear bit is a no-op, and a duplicate value + // that slipped through via the fallback above (only possible if a row exhausts its retry + // budget) shares its bit with wherever it was legitimately set earlier in this same row. + for (size_t k = 0; k < build_config_.node_degree; k++) { + size_t idx = static_cast(output_neighbor_list_ptr[k]); + seen_bits[idx >> 6] &= ~(uint64_t{1} << (idx & 63)); } - output_neighbor_list_ptr[j] = static_cast(idx); } } graph_.h_graph = nullptr; From 22ccd389619aa70f7bb67c846b94052be2948724 Mon Sep 17 00:00:00 2001 From: James Xia Date: Thu, 27 Aug 2026 03:52:14 -0700 Subject: [PATCH 2/2] Simplify seen_bits cleanup in nn_descent dedup to a direct zero seen_bits is thread-local and zeroed on entry to each row, so any word touched by this row's output list contains only bits this row set. Zeroing the whole word is equivalent to masking off one bit but simpler. --- cpp/src/neighbors/detail/nn_descent.cuh | 11 +++++------ 1 file changed, 5 insertions(+), 6 deletions(-) diff --git a/cpp/src/neighbors/detail/nn_descent.cuh b/cpp/src/neighbors/detail/nn_descent.cuh index e7ffeda620..432cde7ffc 100644 --- a/cpp/src/neighbors/detail/nn_descent.cuh +++ b/cpp/src/neighbors/detail/nn_descent.cuh @@ -1772,13 +1772,12 @@ void GNND::build(Data_t* data, } // Unset every bit this row touched so the array is back to all-zero for the next row this - // thread processes. Safe to run over the full row rather than tracking exactly which - // entries were newly set: clearing an already-clear bit is a no-op, and a duplicate value - // that slipped through via the fallback above (only possible if a row exhausts its retry - // budget) shares its bit with wherever it was legitimately set earlier in this same row. + // thread processes. Since seen_bits is all-zero on entry to this row and thread-local, the + // only bits set anywhere are ones this row's own entries set, so the whole word covering + // idx can be zeroed outright rather than masking off just its one bit. for (size_t k = 0; k < build_config_.node_degree; k++) { - size_t idx = static_cast(output_neighbor_list_ptr[k]); - seen_bits[idx >> 6] &= ~(uint64_t{1} << (idx & 63)); + size_t idx = static_cast(output_neighbor_list_ptr[k]); + seen_bits[idx >> 6] = 0; } } }