Replace O(N*degree^2) dedup with per-thread bit-packed bitset - #2437
Replace O(N*degree^2) dedup with per-thread bit-packed bitset#2437jamxia155 wants to merge 2 commits into
Conversation
|
Auto-sync is disabled for draft pull requests in this repository. Workflows must be run manually. Contributors can view more details about this message here. |
jamxia155
added
improvement
| 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)); | ||
|
|
There was a problem hiding this comment.
This PR brings the entire set of host candidates into device memory. The number of candidates per host can be very large. Earlier we were only capping the degree of on device graph buffers to DEGREE_ON_DEVICE (=32) it seems.
Secondly, I don't see why we were even doing it in this way on host in the first place i.e. comparing every single pair. We already call sort_lists(), so the list of candidates is sorted by distances, right? Since we are sorting pairs, the same candidate should technically appear adjacent. Then we should just do a single scan through each candidate list. Or are you saying that there could be differences (for example the epilog is asymmetric maybe?) which means the same candidate can appear at non-adjacent places in the sorted list?
There was a problem hiding this comment.
Thanks for taking a look, Tarang. I'll look into a streaming approach to keep the memory footprint bounded.
As for the de-dup implementation, the original implementation sorted the candidates by distance so a de-dup based on ID cannot be achieved with a linear scan. However, it is possible to first sort by ID, dedup, then sort again by distance. I will run some tests and see if that closes the gap with the GPU port.
There was a problem hiding this comment.
the original implementation sorted the candidates by distance
It sorts the pairs. So for candidates with identical distance, the same candidate would be placed adjacent, right? I think std::sort does this for tuples by moving on to the next element if the first one is identical.
There was a problem hiding this comment.
But we need to think about cases when the same candidate can have two different distances for example small numerical differences in distances (while adding reverse edges). Is that possible? I would be quite surprised though if there was asymmetry like that.
There was a problem hiding this comment.
The possibility for the same candidate with two different distances arises from a somewhat different scenario (all line numbers refer to cpp/src/neighbors/detail/nn_descent.cuh (permalink):
- 1239-1241:
init_random_graph()seeds each row's initial candidate list segment-by-segment, independently, with no cross-segment uniqueness check. The sameidvalue from line 1240 can happen to end up in multiple segments. - 1246: the distances are initialized as
std::numeric_limits::max(). - 1312:
update_graph()computesseg_idx = new_neighb_id.id() % num_segments; 1316:insert_to_ordered_list()is called using this oneseg_idxonly. If the sameidexists in another segment, its initial value never gets overwritten.
There was a problem hiding this comment.
I benchmarked a few different implementations (original, sort-based, hashmap-based, bitset-based, GPU __ballot_sync-based): the bitset-based approach on CPU won in all configurations tested. So I'm switching this PR to a CPU implementation after all.
Of course, the perf comparison depends on the relative capabilities of the CPU and GPU (node I tested on had 48 logical cores + H100 PCI-E) but the burden of proof to switch to GPU is also higher than the current CPU optimization proposed.
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.
fd92c36 to
8ff98cd
Compare
jamxia155
changed the title
2 participants
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 thenode_degreecandidates was checked against all already-placed entries, givingO(N*degree^2)CPU work that scales poorly with graph degree and dataset size.Replace with a per-thread bit-packed "seen" array (1 bit per dataset row), making the check O(1) and the whole pass
O(N*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.Tested on 48-core AMD + H100, up to 20M rows / 1024 degrees 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 vs. GPU narrows with
nrowas the per-thread array outgrows cache.