Inner Product for CAGRA-Q

cpp/CMakeLists.txt

@@ -227,6 +227,30 @@ if(BUILD_SHARED_LIBS)
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_uint8_uint32_dim128_t8.cu
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_uint8_uint32_dim256_t16.cu
     src/neighbors/detail/cagra/compute_distance_standard_L2Expanded_uint8_uint32_dim512_t32.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_float_uint32_dim128_t8_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_float_uint32_dim128_t8_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_float_uint32_dim256_t16_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_float_uint32_dim256_t16_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_float_uint32_dim512_t32_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_float_uint32_dim512_t32_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_half_uint32_dim128_t8_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_half_uint32_dim128_t8_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_half_uint32_dim256_t16_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_half_uint32_dim256_t16_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_half_uint32_dim512_t32_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_half_uint32_dim512_t32_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_int8_uint32_dim128_t8_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_int8_uint32_dim128_t8_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_int8_uint32_dim256_t16_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_int8_uint32_dim256_t16_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_int8_uint32_dim512_t32_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_int8_uint32_dim512_t32_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_uint8_uint32_dim128_t8_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_uint8_uint32_dim128_t8_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_uint8_uint32_dim256_t16_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_uint8_uint32_dim256_t16_8pq_4subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_uint8_uint32_dim512_t32_8pq_2subd_half.cu
+    src/neighbors/detail/cagra/compute_distance_vpq_InnerProduct_uint8_uint32_dim512_t32_8pq_4subd_half.cu
     src/neighbors/detail/cagra/compute_distance_vpq_L2Expanded_float_uint32_dim128_t8_8pq_2subd_half.cu
     src/neighbors/detail/cagra/compute_distance_vpq_L2Expanded_float_uint32_dim128_t8_8pq_4subd_half.cu
     src/neighbors/detail/cagra/compute_distance_vpq_L2Expanded_float_uint32_dim256_t16_8pq_2subd_half.cu

cpp/include/cuvs/neighbors/common.hpp

@@ -40,8 +40,24 @@
 
 namespace cuvs::neighbors {
 
+/**
+ * @defgroup neighbors_index Approximate Nearest Neighbors Types
+ * @{
+ */
+
+/** The base for approximate KNN index structures. */
+struct index {};
+
+/** The base for KNN index parameters. */
+struct index_params {
+  /** Distance type. */
+  cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded;
+  /** The argument used by some distance metrics. */
+  float metric_arg = 2.0f;
+};
+
 /** Parameters for VPQ compression. */
-struct vpq_params {
+struct vpq_params : index_params {
   /**
    * The bit length of the vector element after compression by PQ.
    *
@@ -77,22 +93,6 @@ struct vpq_params {
   double pq_kmeans_trainset_fraction = 0;
 };
 
-/**
- * @defgroup neighbors_index Approximate Nearest Neighbors Types
- * @{
- */
-
-/** The base for approximate KNN index structures. */
-struct index {};
-
-/** The base for KNN index parameters. */
-struct index_params {
-  /** Distance type. */
-  cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded;
-  /** The argument used by some distance metrics. */
-  float metric_arg = 2.0f;
-};
-
 struct search_params {};
 
 /** @} */  // end group neighbors_index

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

@@ -486,15 +486,19 @@ index<T, IdxT> build(
 
   // Construct an index from dataset and optimized knn graph.
   if (params.compression.has_value()) {
-    RAFT_EXPECTS(params.metric == cuvs::distance::DistanceType::L2Expanded,
-                 "VPQ compression is only supported with L2Expanded distance mertric");
+    RAFT_EXPECTS(
+      params.metric == cuvs::distance::DistanceType::L2Expanded ||
+        params.metric == cuvs::distance::DistanceType::InnerProduct,
+      "VPQ compression is only supported with L2Expanded and InnerProduct distance mertric");
     index<T, IdxT> idx(res, params.metric);
     idx.update_graph(res, raft::make_const_mdspan(cagra_graph.view()));
+    auto compression_params   = *params.compression;
+    compression_params.metric = params.metric;
     idx.update_dataset(
       res,
       // TODO: hardcoding codebook math to `half`, we can do runtime dispatching later
       cuvs::neighbors::vpq_build<decltype(dataset), half, int64_t>(
-        res, *params.compression, dataset));
+        res, compression_params, dataset));
 
     return idx;
   }