Skip to content

Commit

Permalink
CAGRA graph optimizer: clamp rev_graph_count (#1987)
Browse files Browse the repository at this point in the history
This PR fixes a potential out of bounds access. 

While building the reverse graph, we count the number of potential edges in `rev_graph_count`

https://github.com/rapidsai/raft/blob/047bfb2a7a24a97a4bece0b553521533868e7889/cpp/include/raft/neighbors/detail/cagra/graph_core.cuh#L204-L205

While we store only max `output_graph_degree` neighbors, the `rev_graph_count` can be larger than `output_graph_degree`.  The loop that uses `rev_graph` is updated to take this limit into account.

Authors:
  - Tamas Bela Feher (https://github.com/tfeher)

Approvers:
  - tsuki (https://github.com/enp1s0)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #1987
  • Loading branch information
tfeher authored Nov 13, 2023
1 parent e5bea87 commit 9110531
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions cpp/include/raft/neighbors/detail/cagra/graph_core.cuh
Original file line number Diff line number Diff line change
Expand Up @@ -528,8 +528,9 @@ void optimize(raft::resources const& res,
constexpr int _omp_chunk = 1024;
#pragma omp parallel for schedule(dynamic, _omp_chunk)
for (uint64_t j = 0; j < graph_size; j++) {
for (uint64_t _k = 0; _k < rev_graph_count.data_handle()[j]; _k++) {
uint64_t k = rev_graph_count.data_handle()[j] - 1 - _k;
uint64_t k = std::min(rev_graph_count.data_handle()[j], output_graph_degree);
while (k) {
k--;
uint64_t i = rev_graph.data_handle()[k + (output_graph_degree * j)];

uint64_t pos =
Expand Down

0 comments on commit 9110531

Please sign in to comment.