CosineExpanded Distance Metric for CAGRA (rapidsai#197)

Currently only IVF-PQ can be used as the graph building algorithm (NN Descent does not support Cosine). As a result, we are limited by IVF-PQ's restriction of data to be of float / half type for the Cosine metric. This PR also fixes an in-place data modification that was being done by IVF-PQ.
Opportunities for optimization:
NN Descent to support Cosine and compute dataset norms only once -- during NN Descent. Re-use those for CAGRA.

[UPDATE 08/21/2025]: NN Descent now support Cosine. This PR allows the initial CAGRA graph to be built by both methods -- IVF_PQ, NN_DESCENT. The IVF_PQ restriction on data types holds, but uint8 and int8 can be supported with NN Descent as the graph building algorithm. ITERATIVE CAGRA SEARCH is currently disabled for Cosine.

[UPDATE 09/23/2025]: This PR also adds Cosine support for IVF_PQ with uint8 / int8 inputs. The above mentioned restriction with IVF_PQ has been removed. So with this PR CAGRA supports Cosine wholly, for float, uint8 and int8 inputs. ITERATIVE_SEARCH however still has some issues as the graph building method with the Cosine metric and has been disabled.

[UPDATE 09/25/2025]:
Binary size comparison for libcuvs.so (CUDA 12.9, x86):
branch-25.10: 1154.42 MB
This PR: 1160.73 MB

Total CAGRA testing time:
branch-25.10:
```
Start 10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST
19/37 Test rapidsai#10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST ...   Passed  825.43 sec
      Start 11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST
20/37 Test rapidsai#11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST ........   Passed    0.58 sec
      Start 12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST
21/37 Test rapidsai#12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST ....   Passed  663.97 sec
      Start 13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST
22/37 Test rapidsai#13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST ....   Passed  397.57 sec
      Start 14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST
23/37 Test rapidsai#14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST ...   Passed  408.16 sec
```
This PR:
```
Start 10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST
19/37 Test rapidsai#10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST ...   Passed  1830.34 sec
      Start 11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST
20/37 Test rapidsai#11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST ........   Passed    0.45 sec
      Start 12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST
21/37 Test rapidsai#12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST ....   Passed  1444.14 sec
      Start 13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST
22/37 Test rapidsai#13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST ....   Passed  973.64 sec
      Start 14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST
23/37 Test rapidsai#14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST ...   Passed  1010.46 sec
```

[UPDATE 09/30/2025]:
Updates to CAGRA C++ tests according to the latest PR reviews.
New total CAGRA testing time:
branch-25.10:
```
      Start  9: NEIGHBORS_ANN_CAGRA_TEST_BUGS
18/37 Test  rapidsai#9: NEIGHBORS_ANN_CAGRA_TEST_BUGS ...........   Passed   16.99 sec
      Start 10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST
19/37 Test rapidsai#10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST ...   Passed  803.64 sec
      Start 11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST
20/37 Test rapidsai#11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST ........   Passed    0.49 sec
      Start 12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST
21/37 Test rapidsai#12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST ....   Passed  667.89 sec
      Start 13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST
22/37 Test rapidsai#13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST ....   Passed  420.49 sec
      Start 14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST
23/37 Test rapidsai#14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST ...   Passed  429.57 sec
```
This PR:
```
      Start  9: NEIGHBORS_ANN_CAGRA_TEST_BUGS
18/37 Test  rapidsai#9: NEIGHBORS_ANN_CAGRA_TEST_BUGS ...........   Passed   26.62 sec
      Start 10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST
19/37 Test rapidsai#10: NEIGHBORS_ANN_CAGRA_FLOAT_UINT32_TEST ...   Passed  973.23 sec
      Start 11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST
20/37 Test rapidsai#11: NEIGHBORS_ANN_CAGRA_HELPERS_TEST ........   Passed    0.43 sec
      Start 12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST
21/37 Test rapidsai#12: NEIGHBORS_ANN_CAGRA_HALF_UINT32_TEST ....   Passed  702.02 sec
      Start 13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST
22/37 Test rapidsai#13: NEIGHBORS_ANN_CAGRA_INT8_UINT32_TEST ....   Passed  491.65 sec
      Start 14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST
23/37 Test rapidsai#14: NEIGHBORS_ANN_CAGRA_UINT8_UINT32_TEST ...   Passed  541.43 sec
```
Fixes rapidsai#1288
Fixes rapidsai#389

Authors:
  - Tarang Jain (https://github.com/tarang-jain)
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Tamas Bela Feher (https://github.com/tfeher)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: rapidsai#197

Author: tarang-jain

cpp/CMakeLists.txt

@@ -231,6 +231,9 @@ if(NOT BUILD_CPU_ONLY)
     src/neighbors/detail/cagra/compute_distance_standard_InnerProduct_float_uint32_dim128_t8.cu
     src/neighbors/detail/cagra/compute_distance_standard_InnerProduct_float_uint32_dim256_t16.cu
     src/neighbors/detail/cagra/compute_distance_standard_InnerProduct_float_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_float_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_float_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_float_uint32_dim512_t32.cu
     src/neighbors/detail/cagra/compute_distance_standard_InnerProduct_half_uint32_dim128_t8.cu
     src/neighbors/detail/cagra/compute_distance_standard_InnerProduct_half_uint32_dim256_t16.cu
     src/neighbors/detail/cagra/compute_distance_standard_InnerProduct_half_uint32_dim512_t32.cu
@@ -246,6 +249,15 @@ if(NOT BUILD_CPU_ONLY)
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_half_uint32_dim128_t8.cu
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_half_uint32_dim256_t16.cu
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_half_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_half_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_half_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_half_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_int8_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_int8_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_int8_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_uint8_uint32_dim128_t8.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_uint8_uint32_dim256_t16.cu
+    src/neighbors/detail/cagra/compute_distance_standard_CosineExpanded_uint8_uint32_dim512_t32.cu
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_int8_uint32_dim128_t8.cu
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_int8_uint32_dim256_t16.cu
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_int8_uint32_dim512_t32.cu

cpp/include/cuvs/neighbors/cagra.hpp

@@ -339,6 +339,14 @@ struct index : cuvs::neighbors::index {
     return graph_view_;
   }
 
+  /** Dataset norms for cosine distance [size] */
+  [[nodiscard]] inline auto dataset_norms() const noexcept
+    -> std::optional<raft::device_vector_view<const float, int64_t>>
+  {
+    if (dataset_norms_.has_value()) { return raft::make_const_mdspan(dataset_norms_->view()); }
+    return std::nullopt;
+  }
+
   // Don't allow copying the index for performance reasons (try avoiding copying data)
   /** \cond */
   index(const index&)                    = delete;
@@ -354,7 +362,8 @@ struct index : cuvs::neighbors::index {
     : cuvs::neighbors::index(),
       metric_(metric),
       graph_(raft::make_device_matrix<IdxT, int64_t>(res, 0, 0)),
-      dataset_(new cuvs::neighbors::empty_dataset<int64_t>(0))
+      dataset_(new cuvs::neighbors::empty_dataset<int64_t>(0)),
+      dataset_norms_(std::nullopt)
   {
   }
 
@@ -420,12 +429,21 @@ struct index : cuvs::neighbors::index {
     : cuvs::neighbors::index(),
       metric_(metric),
       graph_(raft::make_device_matrix<IdxT, int64_t>(res, 0, 0)),
-      dataset_(make_aligned_dataset(res, dataset, 16))
+      dataset_(make_aligned_dataset(res, dataset, 16)),
+      dataset_norms_(std::nullopt)
   {
     RAFT_EXPECTS(dataset.extent(0) == knn_graph.extent(0),
                  "Dataset and knn_graph must have equal number of rows");
     update_graph(res, knn_graph);
 
+    if (metric_ == cuvs::distance::DistanceType::CosineExpanded) {
+      auto p = dynamic_cast<strided_dataset<T, int64_t>*>(dataset_.get());
+      if (p) {
+        auto dataset_view = p->view();
+        if (dataset_view.extent(0) > 0) { compute_dataset_norms_(res); }
+      }
+    }
+
     raft::resource::sync_stream(res);
   }
 
@@ -435,48 +453,81 @@ struct index : cuvs::neighbors::index {
    * If the new dataset rows are aligned on 16 bytes, then only a reference is stored to the
    * dataset. It is the caller's responsibility to ensure that dataset stays alive as long as the
    * index. It is expected that the same set of vectors are used for update_dataset and index build.
+   *
+   * Note: This will clear any precomputed dataset norms.
    */
   void update_dataset(raft::resources const& res,
                       raft::device_matrix_view<const T, int64_t, raft::row_major> dataset)
   {
     dataset_ = make_aligned_dataset(res, dataset, 16);
+    dataset_norms_.reset();
+
+    if (metric() == cuvs::distance::DistanceType::CosineExpanded) {
+      if (dataset.extent(0) > 0) { compute_dataset_norms_(res); }
+    }
   }
 
   /** Set the dataset reference explicitly to a device matrix view with padding. */
   void update_dataset(raft::resources const& res,
                       raft::device_matrix_view<const T, int64_t, raft::layout_stride> dataset)
   {
     dataset_ = make_aligned_dataset(res, dataset, 16);
+    dataset_norms_.reset();
+
+    if (metric() == cuvs::distance::DistanceType::CosineExpanded) {
+      if (dataset.extent(0) > 0) { compute_dataset_norms_(res); }
+    }
   }
 
   /**
    * Replace the dataset with a new dataset.
    *
    * We create a copy of the dataset on the device. The index manages the lifetime of this copy. It
    * is expected that the same set of vectors are used for update_dataset and index build.
+   *
+   * Note: This will clear any precomputed dataset norms.
    */
   void update_dataset(raft::resources const& res,
                       raft::host_matrix_view<const T, int64_t, raft::row_major> dataset)
   {
     dataset_ = make_aligned_dataset(res, dataset, 16);
+    dataset_norms_.reset();
+    if (metric() == cuvs::distance::DistanceType::CosineExpanded) {
+      if (dataset.extent(0) > 0) { compute_dataset_norms_(res); }
+    }
   }
 
   /**
    * Replace the dataset with a new dataset. It is expected that the same set of vectors are used
    * for update_dataset and index build.
+   *
+   * Note: This will clear any precomputed dataset norms.
    */
   template <typename DatasetT>
   auto update_dataset(raft::resources const& res, DatasetT&& dataset)
     -> std::enable_if_t<std::is_base_of_v<cuvs::neighbors::dataset<dataset_index_type>, DatasetT>>
   {
     dataset_ = std::make_unique<DatasetT>(std::move(dataset));
+    dataset_norms_.reset();
+    if (metric() == cuvs::distance::DistanceType::CosineExpanded) {
+      auto p = dynamic_cast<strided_dataset<T, int64_t>*>(dataset_.get());
+      if (p) {
+        auto dataset_view = p->view();
+        if (dataset_view.extent(0) > 0) { compute_dataset_norms_(res); }
+      }
+    }
   }
 
   template <typename DatasetT>
   auto update_dataset(raft::resources const& res, std::unique_ptr<DatasetT>&& dataset)
     -> std::enable_if_t<std::is_base_of_v<neighbors::dataset<dataset_index_type>, DatasetT>>
   {
     dataset_ = std::move(dataset);
+    dataset_norms_.reset();
+    if (metric() == cuvs::distance::DistanceType::CosineExpanded) {
+      auto dataset_view = this->dataset();
+      if (dataset_view.extent(0) > 0) { compute_dataset_norms_(res); }
+    }
   }
 
   /**
@@ -519,6 +570,10 @@ struct index : cuvs::neighbors::index {
   raft::device_matrix<IdxT, int64_t, raft::row_major> graph_;
   raft::device_matrix_view<const IdxT, int64_t, raft::row_major> graph_view_;
   std::unique_ptr<neighbors::dataset<dataset_index_type>> dataset_;
+  // only float distances supported at the moment
+  std::optional<raft::device_vector<float, int64_t>> dataset_norms_;
+
+  void compute_dataset_norms_(raft::resources const& res);
 };
 /**
  * @}
@@ -539,6 +594,7 @@ struct index : cuvs::neighbors::index {
  * The following distance metrics are supported:
  * - L2
  * - InnerProduct (currently only supported with IVF-PQ as the build algorithm)
+ * - CosineExpanded
  *
  * Usage example:
  * @code{.cpp}
@@ -576,6 +632,7 @@ auto build(raft::resources const& res,
  * The following distance metrics are supported:
  * - L2
  * - InnerProduct (currently only supported with IVF-PQ as the build algorithm)
+ * - CosineExpanded
  *
  * Usage example:
  * @code{.cpp}
@@ -613,6 +670,7 @@ auto build(raft::resources const& res,
  * The following distance metrics are supported:
  * - L2
  * - InnerProduct (currently only supported with IVF-PQ as the build algorithm)
+ * - CosineExpanded (dataset norms are computed as float regardless of input data type)
  *
  * Usage example:
  * @code{.cpp}
@@ -649,6 +707,7 @@ auto build(raft::resources const& res,
  *
  * The following distance metrics are supported:
  * - L2
+ * - CosineExpanded (dataset norms are computed as float regardless of input data type)
  *
  * Usage example:
  * @code{.cpp}
@@ -685,6 +744,7 @@ auto build(raft::resources const& res,
  *
  * The following distance metrics are supported:
  * - L2
+ * - CosineExpanded (dataset norms are computed as float regardless of input data type)
  *
  * Usage example:
  * @code{.cpp}
@@ -722,6 +782,7 @@ auto build(raft::resources const& res,
  * The following distance metrics are supported:
  * - L2
  * - InnerProduct (currently only supported with IVF-PQ as the build algorithm)
+ * - CosineExpanded (dataset norms are computed as float regardless of input data type)
  *
  * Usage example:
  * @code{.cpp}
@@ -759,6 +820,7 @@ auto build(raft::resources const& res,
  * The following distance metrics are supported:
  * - L2
  * - InnerProduct (currently only supported with IVF-PQ as the build algorithm)
+ * - CosineExpanded (dataset norms are computed as float regardless of input data type)
  *
  * Usage example:
  * @code{.cpp}
@@ -796,6 +858,7 @@ auto build(raft::resources const& res,
  * The following distance metrics are supported:
  * - L2
  * - InnerProduct (currently only supported with IVF-PQ as the build algorithm)
+ * - CosineExpanded (dataset norms are computed as float regardless of input data type)
  *
  * Usage example:
  * @code{.cpp}

cpp/src/neighbors/cagra.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023-2024, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2025, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -22,10 +22,12 @@
 #include "detail/cagra/cagra_search.cuh"
 #include "detail/cagra/graph_core.cuh"
 
+#include "detail/ann_utils.cuh"
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/host_device_accessor.hpp>
 #include <raft/core/mdspan.hpp>
 #include <raft/core/resources.hpp>
+#include <raft/linalg/norm.cuh>
 
 #include <cuvs/distance/distance.hpp>
 #include <cuvs/neighbors/cagra.hpp>
@@ -36,6 +38,37 @@
 
 namespace cuvs::neighbors::cagra {
 
+// Member function implementations for cagra::index
+template <typename T, typename IdxT>
+void index<T, IdxT>::compute_dataset_norms_(raft::resources const& res)
+{
+  // Get the dataset view
+  auto dataset_view = this->dataset();
+
+  // Allocate norms vector if not already allocated
+  if (!dataset_norms_.has_value() || dataset_norms_->extent(0) != dataset_view.extent(0)) {
+    dataset_norms_.reset();
+    dataset_norms_ = raft::make_device_vector<float, int64_t>(res, dataset_view.extent(0));
+  }
+
+  constexpr float kScale = cuvs::spatial::knn::detail::utils::config<T>::kDivisor /
+                           cuvs::spatial::knn::detail::utils::config<float>::kDivisor;
+
+  // first scale the dataset and then compute norms
+  auto scaled_sq_op = raft::compose_op(
+    raft::sq_op{}, raft::div_const_op<float>{float(kScale)}, raft::cast_op<float>());
+  raft::linalg::reduce<true, true, T, float, int64_t>(dataset_norms_->data_handle(),
+                                                      dataset_view.data_handle(),
+                                                      dataset_view.stride(0),
+                                                      dataset_view.extent(0),
+                                                      (float)0,
+                                                      raft::resource::get_cuda_stream(res),
+                                                      false,
+                                                      scaled_sq_op,
+                                                      raft::add_op(),
+                                                      raft::sqrt_op{});
+}
+
 /**
  * @defgroup cagra CUDA ANN Graph-based nearest neighbor search
  * @{

cpp/src/neighbors/detail/cagra/cagra_build.cuh

@@ -127,8 +127,9 @@ void build_knn_graph(
   cuvs::neighbors::cagra::graph_build_params::ivf_pq_params pq)
 {
   RAFT_EXPECTS(pq.build_params.metric == cuvs::distance::DistanceType::L2Expanded ||
-                 pq.build_params.metric == cuvs::distance::DistanceType::InnerProduct,
-               "Currently only L2Expanded or InnerProduct metric are supported");
+                 pq.build_params.metric == cuvs::distance::DistanceType::InnerProduct ||
+                 pq.build_params.metric == cuvs::distance::DistanceType::CosineExpanded,
+               "Currently only L2Expanded, InnerProduct and CosineExpanded metrics are supported");
 
   uint32_t node_degree = knn_graph.extent(1);
   raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope(
@@ -710,6 +711,11 @@ index<T, IdxT> build(
       std::holds_alternative<cagra::graph_build_params::nn_descent_params>(knn_build_params),
     "IVF_PQ for CAGRA graph build does not support BitwiseHamming as a metric. Please "
     "use nn-descent or the iterative CAGRA search build.");
+  RAFT_EXPECTS(
+    params.metric != cuvs::distance::DistanceType::CosineExpanded ||
+      std::holds_alternative<cagra::graph_build_params::ivf_pq_params>(knn_build_params) ||
+      std::holds_alternative<cagra::graph_build_params::nn_descent_params>(knn_build_params),
+    "CosineExpanded distance is not supported for iterative CAGRA graph build.");
 
   // Validate data type for BitwiseHamming metric
   RAFT_EXPECTS(params.metric != cuvs::distance::DistanceType::BitwiseHamming ||