InnerProduct Distance Metric for CAGRA search #2260

tarang-jain · 2024-04-10T00:19:49Z

InnerProduct Distance Metric for CAGRA search. InnerProduct in graph building is supported using IVF-PQ for building the graph. NNDescent does not currently support any other metric except L2Expanded.

…cagra-dists

tarang-jain · 2024-04-11T22:07:16Z

/rerun tests

enp1s0

Thank you @tarang-jain for adding support for innerproduct similarity! Most of the changes look good to me. There is one performance concern and some minor suggestions, so can you modify them?

cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp

cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh

cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp

cpp/include/raft/neighbors/detail/cagra/compute_distance_vpq.cuh

cpp/include/raft/neighbors/detail/cagra/search_multi_kernel.cuh

cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh

cpp/include/raft/neighbors/detail/cagra/search_multi_kernel.cuh

Co-authored-by: tsuki <[email protected]>

cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh

cpp/include/raft/neighbors/detail/cagra/compute_distance_vpq.cuh

cpp/include/raft/neighbors/detail/cagra/search_single_cta_kernel-inl.cuh

lowener · 2024-04-12T17:30:36Z

Can you add support for InnerProduct on Python side in this PR as well?

…cagra-dists

tarang-jain · 2024-04-19T20:16:44Z

/ok to test

tfeher

Thanks @tarang-jain for the update. LGTM!

enp1s0

Thank you @tarang-jain for the changes. The code looks good to me.

cjnolet

This looks good Tarang. Mostly small items on my side, but want to make sure we're providing a consistent API experience across all of our ANN APIs (this will become particularly important when these APIs are migrated to cuVS).

cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh

cpp/test/neighbors/ann_cagra.cuh

cpp/include/raft/neighbors/cagra.cuh

cjnolet · 2024-04-22T15:50:05Z

cpp/include/raft/neighbors/cagra.cuh

 *
 * Usage example:
 * @code{.cpp}
 *   using namespace raft::neighbors;
 *   // use default index parameters
- *   ivf_pq::index_params build_params;
+ *   ivf_pq::index_params = cagra::get_default_ivf_pq_build_params(dataset);


I'm not sure I like this approach as it doesn't follow the standard API design flow that we've been using in hte other algos. I wonder if instead of adding this additional helper function if we could instead just add an overloaded constructor to the index_params object that accepts an mdsan with the dataset and sets the defaults automatically. This would then follow the standard build flow but add an option for the user to establish the ressonable defaultsby passing in the dataset.

Can't add constructors to the ivf_pq::index_params struct because the aggregate assertion here. A helper is needed. One thing that can be done is that instead of a constructor within the index_params struct, we can have a helper get_default_index_params outside of the struct, but in ivf_pq_types.hpp

Honestly, I'd just remove the static assertion there. It's not necessary. My problem with exposing helper functions like this is that they make the APIs harder to use as they aren't immediatley obvious and don't provide for a consistent experience. Keeping everything together within the index_params object itself allows there to be a single entrypoint (e.g. index_params) for construction. The other option would be to have factory functions which would always be used for construction, however I think that will also make things even more confusing unless we used that pattern across all the ANN APis.

Here we are constructing CAGRA specific ivf_pq::index_params. I do not think the constructor of ivf_pq::index_params shall have the job to define parameters according to CAGRA's requirements.

For consistent experience, the user just need to fill the metric field of cagra::index_params, which is inherited from ann::index_params, and call cagra::build as usual.

Users of build_knn_graph are advanced users who are using build_knn_graph & optimize in a separate step. Other algos don't have these steps, this is custom for CAGRA. Before this PR, the user either passes a manually filled ivf_pq::index_params, or just uses a default. I think a helper function is an improvement.

…cagra-dists

tarang-jain · 2024-04-23T18:47:00Z

/ok to test

tarang-jain · 2024-04-23T23:42:34Z

/ok to test

tarang-jain · 2024-04-24T00:16:17Z

/ok to test

tfeher · 2024-04-24T07:49:28Z

cpp/include/raft/neighbors/ivf_pq_types.hpp

+  /**
+   * Helper that sets values according to the extents of the dataset mdspan.
+   */
+  template <typename DataT, typename Accessor>
+  static index_params from_dataset(
+    mdspan<const DataT, matrix_extent<int64_t>, row_major, Accessor> dataset,
+    raft::distance::DistanceType metric = raft::distance::L2Expanded)
+  {
+    index_params params;
+    params.n_lists =
+      dataset.extent(0) < 4 * 2500 ? 4 : static_cast<uint32_t>(std::sqrt(dataset.extent(0)));
+    params.pq_dim =
+      round_up_safe(static_cast<uint32_t>(dataset.extent(1) / 4), static_cast<uint32_t>(8));
+    params.pq_bits                  = 8;
+    params.kmeans_trainset_fraction = dataset.extent(0) < 10000 ? 1 : 0.1;
+    params.metric                   = metric;
+    return params;
+  }


These are specific default params that are useful for constructing a CAGRA index, but possibly not ideal as default params for an ivf_pq index. The documentation has to make this clear. I think it would be also preferred to keep this helper in the CAGRA namespace.

Actually, I don't know that I agree with this. The sqrt(n) and compression ratio are reasonable defaults for ivfpq in general, irrespective of CAGRA.

Ok, maybe you are right. It is specific for CAGRA only in the sense that it leads to good enough recall for a building a CAGRA graph. But we do not give any recall guarantees (we cannot since it is dataset dependent) so we can treat this general defaults.

tarang-jain · 2024-04-24T17:18:12Z

/ok to test

cpp/include/raft/neighbors/ivf_pq_types.hpp

tfeher · 2024-04-24T19:51:32Z

cpp/include/raft/neighbors/ivf_pq_types.hpp

+  /**
+   * Helper that sets values according to the extents of the dataset mdspan.
+   */
+  template <typename DataT, typename Accessor>
+  static index_params from_dataset(
+    mdspan<const DataT, matrix_extent<int64_t>, row_major, Accessor> dataset,
+    raft::distance::DistanceType metric = raft::distance::L2Expanded)
+  {
+    index_params params;
+    params.n_lists =
+      dataset.extent(0) < 4 * 2500 ? 4 : static_cast<uint32_t>(std::sqrt(dataset.extent(0)));
+    params.pq_dim =
+      round_up_safe(static_cast<uint32_t>(dataset.extent(1) / 4), static_cast<uint32_t>(8));
+    params.pq_bits                  = 8;
+    params.kmeans_trainset_fraction = dataset.extent(0) < 10000 ? 1 : 0.1;
+    params.metric                   = metric;
+    return params;
+  }


Ok, maybe you are right. It is specific for CAGRA only in the sense that it leads to good enough recall for a building a CAGRA graph. But we do not give any recall guarantees (we cannot since it is dataset dependent) so we can treat this general defaults.

…cagra-dists

tarang-jain · 2024-04-24T20:07:24Z

/ok to test

cjnolet · 2024-04-24T22:09:40Z

Ok, maybe you are right. It is specific for CAGRA only in the sense that it leads to good enough recall for a building a CAGRA graph. But we do not give any recall guarantees (we cannot since it is dataset dependent) so we can treat this general defaults.

I went back and forth about this a couple times too, and I ended up at the conclusion that taking the dataset's shape into account should add tiny bit more generalization to the default params, even if they aren't guaranteed to yield the highest recall. IMO, in the worst case, a user doesn't use the factory method and they get the same default params we've always provided and in the best case, they are able to use the factory method and they are happy w/ the resultling recall.

cjnolet · 2024-04-26T00:32:40Z

/merge

tarang-jain · 2024-04-29T23:22:02Z

/ok to test

tarang-jain · 2024-04-30T00:23:04Z

/ok to test

tarang-jain added 7 commits April 5, 2024 11:13

apply updates to 24.06

b6d9980

Merge branch 'branch-24.06' of https://github.com/rapidsai/raft into …

fac65b8

…cagra-dists

remove build errors

ad48b03

search inputs

2b8d898

inner product in compute_distance_vpq.cuh

e44ab17

Merge branch 'branch-24.06' of https://github.com/rapidsai/raft into …

c506216

…cagra-dists

inner product in index build; debug statements

cb7fbba

github-actions bot added the cpp label Apr 10, 2024

tarang-jain mentioned this pull request Apr 10, 2024

Additional Distance Metrics for CAGRA #2187

Closed

tarang-jain added 2 - In Progress Currenty a work in progress non-breaking Non-breaking change feature request New feature or request labels Apr 10, 2024

cjnolet assigned tarang-jain Apr 10, 2024

tarang-jain added 9 commits April 10, 2024 13:26

tests passing

0029bba

Merge branch 'branch-24.06' of https://github.com/rapidsai/raft into …

293bc8f

…cagra-dists

style

c91d895

update testing

7a5f876

rm log statements

c37aa27

pass CagraSort

890372b

tests passing

810ddd1

remove dbg statements

f92e68b

update docs

5bbbc70

tarang-jain marked this pull request as ready for review April 11, 2024 21:08

tarang-jain requested a review from a team as a code owner April 11, 2024 21:08

metric assertions

7febd73

enp1s0 suggested changes Apr 12, 2024

View reviewed changes

add metric as const arg

7e19937

Co-authored-by: tsuki <[email protected]>

lowener reviewed Apr 12, 2024

View reviewed changes

tarang-jain added 2 commits April 19, 2024 10:33

update docs, passing gtests

c5cd0e7

Merge branch 'branch-24.06' of https://github.com/rapidsai/raft into …

eabb031

…cagra-dists

tfeher approved these changes Apr 22, 2024

View reviewed changes

enp1s0 approved these changes Apr 22, 2024

View reviewed changes

cjnolet requested changes Apr 22, 2024

View reviewed changes

tarang-jain added 2 commits April 23, 2024 11:38

add ivf_pq::index_params helper

77ff0c2

Merge branch 'branch-24.06' of https://github.com/rapidsai/raft into …

b61c1c3

…cagra-dists

tarang-jain mentioned this pull request Apr 23, 2024

[FEA] NN Descent + CAGRA should support additional distance metrics #2276

Open

tracking issue; styling

54061d0

make helper static

dda1810

tarang-jain force-pushed the cagra-dists branch from 0e498a8 to dda1810 Compare April 24, 2024 00:15

tfeher requested changes Apr 24, 2024

View reviewed changes

cjnolet reviewed Apr 24, 2024

View reviewed changes

cpp/include/raft/neighbors/ivf_pq_types.hpp Outdated Show resolved Hide resolved

tfeher approved these changes Apr 24, 2024

View reviewed changes

tarang-jain added 2 commits April 24, 2024 13:05

update from_dataset verbiage

743d5bc

Merge branch 'branch-24.06' of https://github.com/rapidsai/raft into …

5791743

…cagra-dists

tarang-jain force-pushed the cagra-dists branch from 893e09d to 5791743 Compare April 24, 2024 20:06

cjnolet approved these changes Apr 26, 2024

View reviewed changes

Merge branch 'branch-24.06' into cagra-dists

f5f30b7

rapids-bot bot merged commit e720de7 into rapidsai:branch-24.06 Apr 30, 2024
70 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

InnerProduct Distance Metric for CAGRA search #2260

InnerProduct Distance Metric for CAGRA search #2260

tarang-jain commented Apr 10, 2024 •

edited

Loading

tarang-jain commented Apr 11, 2024 •

edited

Loading

enp1s0 left a comment

lowener commented Apr 12, 2024

tarang-jain commented Apr 19, 2024

tfeher left a comment

enp1s0 left a comment

cjnolet left a comment

cjnolet Apr 22, 2024

tarang-jain Apr 23, 2024 •

edited

Loading

cjnolet Apr 23, 2024 •

edited

Loading

tfeher Apr 24, 2024 •

edited

Loading

tarang-jain commented Apr 23, 2024

tarang-jain commented Apr 23, 2024

tarang-jain commented Apr 24, 2024

tfeher Apr 24, 2024

cjnolet Apr 24, 2024

tfeher Apr 24, 2024

tarang-jain commented Apr 24, 2024

tfeher Apr 24, 2024

tarang-jain commented Apr 24, 2024

cjnolet commented Apr 24, 2024

cjnolet commented Apr 26, 2024

tarang-jain commented Apr 29, 2024

tarang-jain commented Apr 30, 2024

InnerProduct Distance Metric for CAGRA search #2260

InnerProduct Distance Metric for CAGRA search #2260

Conversation

tarang-jain commented Apr 10, 2024 • edited Loading

tarang-jain commented Apr 11, 2024 • edited Loading

enp1s0 left a comment

Choose a reason for hiding this comment

lowener commented Apr 12, 2024

tarang-jain commented Apr 19, 2024

tfeher left a comment

Choose a reason for hiding this comment

enp1s0 left a comment

Choose a reason for hiding this comment

cjnolet left a comment

Choose a reason for hiding this comment

cjnolet Apr 22, 2024

Choose a reason for hiding this comment

tarang-jain Apr 23, 2024 • edited Loading

Choose a reason for hiding this comment

cjnolet Apr 23, 2024 • edited Loading

Choose a reason for hiding this comment

tfeher Apr 24, 2024 • edited Loading

Choose a reason for hiding this comment

tarang-jain commented Apr 23, 2024

tarang-jain commented Apr 23, 2024

tarang-jain commented Apr 24, 2024

tfeher Apr 24, 2024

Choose a reason for hiding this comment

cjnolet Apr 24, 2024

Choose a reason for hiding this comment

tfeher Apr 24, 2024

Choose a reason for hiding this comment

tarang-jain commented Apr 24, 2024

tfeher Apr 24, 2024

Choose a reason for hiding this comment

tarang-jain commented Apr 24, 2024

cjnolet commented Apr 24, 2024

cjnolet commented Apr 26, 2024

tarang-jain commented Apr 29, 2024

tarang-jain commented Apr 30, 2024

tarang-jain commented Apr 10, 2024 •

edited

Loading

tarang-jain commented Apr 11, 2024 •

edited

Loading

tarang-jain Apr 23, 2024 •

edited

Loading

cjnolet Apr 23, 2024 •

edited

Loading

tfeher Apr 24, 2024 •

edited

Loading