Expose search function with pre-filter for ANN #302
Conversation
…eparable-compilation
…e place as compute_distance components
lowener
changed the title
…ble compute_distance_impl and being more explicit about the memory spaces (using lds/ldg)
…sed to the compute_distance function
|
Closes #319 |
cpp/include/cuvs/neighbors/cagra.hpp
Outdated
| */ | ||
|
|
||
| void search(raft::resources const& res, | ||
| cuvs::neighbors::cagra::search_params const& params, | ||
| const cuvs::neighbors::cagra::index<float, uint32_t>& index, | ||
| raft::device_matrix_view<const float, int64_t, raft::row_major> queries, | ||
| raft::device_matrix_view<uint32_t, int64_t, raft::row_major> neighbors, | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances); | ||
| raft::device_matrix_view<float, int64_t, raft::row_major> distances, | ||
| std::optional<cuvs::neighbors::filtering::bitset_filter<uint32_t, int64_t>> |
There was a problem hiding this comment.
We should be able to support both bitset and bitmap using inheritance (and without having to instantiate two explicit templates), right?
There was a problem hiding this comment.
This could work on the host API side if we do dynamic casts on the host side down the line (before passing it to device) and do not add any virtual methods to the filter struct (so that the objects we pass to device don't have virtual tables).
(otherwise it gets very complicated and potentially slow)
There was a problem hiding this comment.
Yes, inheritance doesn't work as well on device side so the filter can't dynamically choose the right virtual function. Doing a dynamic cast on the host side can work so I will try that approach
…n' into 24.10-search-filter
lowener
changed the title
There was a problem hiding this comment.
Overall the changes look good, but I'd like to avoid using pointers on the public APIs where at all possible on the public APIs and either pass by const reference or std::optional.
Also, can we consolidate the search API on brute-force also so that we can remove that search_with_filter function? We can do that in a follow-up if needed.
|
@cjnolet Let's do the brute-force part in a follow-up, as it requires additional changes in the python, C and Rust wrappers code. |
|
/merge |
Related to #136
This PR:
search()function with optional filters in the public API of ANNs (CAGRA, IVF-PQ, IVF-Flat)