Replace O(N*degree^2) CPU dedup with GPU warp-ballot kernel - #2437
Replace O(N*degree^2) CPU dedup with GPU warp-ballot kernel#2437jamxia155 wants to merge 1 commit into
Conversation
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.
|
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.
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 GPU
dedup_graph_kernelthat runs one warp per node. The warp scans neighbors in original order, using__ballot_syncforO(warp-width)duplicate detection, and fills any remaining slots withxorshift64random nodes. The H2D/D2H transfers areO(N*degree)-- the same order as a single pass over the graph -- while the replaced CPU work isO(N*degree^2), so the transfers are dominated by the savings at any practical degree.Tested on 32-core AMD + H100 with
n_rows=2.5Mandgraph_degree=56, this change reduces runtime by about 6 seconds.