RAFT contains fundamental widely-used algorithms and primitives for machine learning and information retrieval. The algorithms are CUDA-accelerated and form building blocks for more easily writing high performance applications.
Replace the size_t type in the AnnBase::search for the output neighbor indices with a common AnnBase::index_type.
This PR stops short of changing the behavior of the benchmarks, since it keeps using index_type = size_t.
The introduction of the new type has couple benefits:
Makes the usage of the index_type more clear in the code, distinguishing it from the extents type, which is usually size_t as well.
Makes it possible to quickly change the alias to uint32_t during development and experiments. This is needed to avoid calling extra linalg::map on the produced results when the algorithm output is not compatible with size_t.
As a small extra change, I've factored out common IVF-PQ - CAGRA-Q refinement code into a separate refine_helper function.
achirkin
commented
1 month ago
Replace the
size_ttype in theAnnBase::searchfor the output neighbor indices with a commonAnnBase::index_type. This PR stops short of changing the behavior of the benchmarks, since it keepsusing index_type = size_t.The introduction of the new type has couple benefits:
index_typemore clear in the code, distinguishing it from the extents type, which is usuallysize_tas well.uint32_tduring development and experiments. This is needed to avoid calling extralinalg::mapon the produced results when the algorithm output is not compatible withsize_t.As a small extra change, I've factored out common IVF-PQ - CAGRA-Q refinement code into a separate
refine_helperfunction.