All-neighbors graph on CPU#2313
Conversation
jinsolp
added
improvement
dantegd
left a comment
dantegd
left a comment
There was a problem hiding this comment.
PR looks great! Just had 3 comments and a broader concern: batched inner-product results can still be incorrect when similarities are negative if I'm not mistaken. reset_global_matrices uses std::numeric_limits::min(), which is the smallest positive value, so the sentinel can rank ahead of valid negative inner products during the merge. I think switching this to lowest() and adding a centered-data test might be a good idea.
| auto distances_view = | ||
| direct.second.has_value() | ||
| ? direct.second.value() | ||
| : raft::make_device_matrix<T, IdxT>(this->res, dataset.extent(0), this->k).view(); |
There was a problem hiding this comment.
What do you think about keeping an owning device_matrix alive here when distances weren’t requested? This view points into a temporary destroyed at the end of this statement, so the brute force search below receives a dangling pointer for indices-only builds. Could we retain the matrix in a local owner and add indices-only coverage?
| indices.extent(1) == distances.value().extent(1), | ||
| "indices matrix and distances matrix has to be the same shape."); | ||
| } | ||
| if (core_distances.has_value()) { |
There was a problem hiding this comment.
Should we also require core_distances.size() == dataset.extent(0) here? The host path uses this size as num_rows and indexes the distances matrix with it, while the reachability epilogue is indexed by dataset row. A mismatched C++ view can therefore cause out-of-bounds access. What do you think about validating the extent here?
| * - A host-resident dataset accepts either host- or device-resident outputs (indices, distances, | ||
| * core_distances); a device-resident dataset requires device-resident outputs. All provided | ||
| * outputs must share the same memory space. | ||
| * - Host-resident outputs never materialize the full [num_rows x k] graph on the GPU. |
There was a problem hiding this comment.
full_build_host() still allocates full device indices and optional distances when n_clusters == 1, no? So the current wording overstates the memory reduction a little bit.
What do you think about narrowing this guarantee to batched builds (n_clusters > 1)?
2 participants
Closes #1903
This PR adds support for CPU kNN graph for batched all-neighbors. Passing host matrices for the indices and distances avoids having to fit the entire graph on GPU. Small perf gap and expected peak GPU memory reductions below: