Made the metric be a part of the index struct and set the greater_ = …

…true for similarity metrics
rapidsai · rapids-bot · Aug 24, 2022 · May 13, 2022 · May 16, 2022 · May 16, 2022
commit 38b3cec4986ab5fc7fb04347e4a5997a3f3005e3
diff --git a/cpp/include/raft/spatial/knn/detail/ann_ivf_flat.cuh b/cpp/include/raft/spatial/knn/detail/ann_ivf_flat.cuh
@@ -45,18 +45,23 @@ namespace raft::spatial::knn::detail {
 template <typename T>
 class cuivflHandle {
  public:
-  cuivflHandle(const handle_t& handle,
-               raft::distance::DistanceType metric_type,
-               const ivf_flat_params& params);
+  cuivflHandle(const handle_t& handle, const ivf_flat_params& params)
+    : handle_(handle), stream_(handle_.get_stream()), params_(params), grid_dim_x_(0)
+  {
+  }
 
   /**
    * @brief Build the index from the dataset for efficient search.
    *
    * @param[in] dataset a device pointer to a row-major matrix [n_rows, dim]
    * @param n_rows number of samples
    * @param dim the dimensionality of the data
+   * @param metric distance type
    */
-  void cuivflBuildIndex(const T* dataset, uint32_t n_rows, uint32_t dim);
+  void cuivflBuildIndex(const T* dataset,
+                        uint32_t n_rows,
+                        uint32_t dim,
+                        raft::distance::DistanceType metric);
 
   /**
    * @brief Set the search parameters. Must be called before `cuivflSearch`
@@ -90,7 +95,6 @@ class cuivflHandle {
   const rmm::cuda_stream_view stream_;
   ivf_flat_params params_;
 
-  const raft::distance::DistanceType metric_type_;
   bool greater_;
   uint32_t grid_dim_x_;  // The number of blocks launched across n_probes.
   // The built index
@@ -107,18 +111,6 @@ class cuivflHandle {
   void queryIVFFlatGridSize(const uint32_t n_probes, const uint32_t n_queries, const uint32_t k);
 };
 
-template <typename T>
-cuivflHandle<T>::cuivflHandle(const handle_t& handle,
-                              raft::distance::DistanceType metric_type,
-                              const ivf_flat_params& params)
-  : handle_(handle),
-    stream_(handle_.get_stream()),
-    params_(params),
-    grid_dim_x_(0),
-    metric_type_(metric_type)
-{
-}
-
 /**
  * @brief Record the dataset into the index, one source row at a time.
  *
@@ -187,7 +179,10 @@ __global__ void build_index_kernel(const uint32_t* labels,
 }
 
 template <typename T>
-void cuivflHandle<T>::cuivflBuildIndex(const T* dataset, uint32_t n_rows, uint32_t dim)
+void cuivflHandle<T>::cuivflBuildIndex(const T* dataset,
+                                       uint32_t n_rows,
+                                       uint32_t dim,
+                                       raft::distance::DistanceType metric)
 {
   static_assert(std::is_same_v<T, float> || std::is_same_v<T, uint8_t> || std::is_same_v<T, int8_t>,
                 "unsupported data type");
@@ -218,7 +213,7 @@ void cuivflHandle<T>::cuivflBuildIndex(const T* dataset, uint32_t n_rows, uint32
                                  centers.data(),
                                  n_lists,
                                  params_.kmeans_trainset_fraction,
-                                 metric_type_,
+                                 metric,
                                  stream_);
 
   // Calculate offsets into cluster data using exclusive scan
@@ -263,13 +258,13 @@ void cuivflHandle<T>::cuivflBuildIndex(const T* dataset, uint32_t n_rows, uint32
     RAFT_LOG_TRACE_VEC(r.data(), 20);
     return r;
   };
-  auto&& center_norms = metric_type_ == raft::distance::DistanceType::L2Expanded
+  auto&& center_norms = metric == raft::distance::DistanceType::L2Expanded
                           ? std::optional(compute_norms())
                           : std::nullopt;
 
   // assemble the index
-  index_.emplace(
-    ivf_flat_index<T>{veclen, data, indices, list_sizes, list_offsets, centers, center_norms});
+  index_.emplace(ivf_flat_index<T>{
+    veclen, metric, data, indices, list_sizes, list_offsets, centers, center_norms});
 
   // check index invariants
   index_->check_consistency();
@@ -285,7 +280,7 @@ void cuivflHandle<T>::queryIVFFlatGridSize(const uint32_t n_probes,
                                                                   nullptr,
                                                                   nullptr,
                                                                   n_queries,
-                                                                  metric_type_,
+                                                                  index_->metric,
                                                                   n_probes,
                                                                   k,
                                                                   greater_,
@@ -303,13 +298,17 @@ void cuivflHandle<T>::cuivflSetSearchParameters(const uint32_t n_probes,
   RAFT_EXPECTS(n_probes > 0,
                "n_probes (number of clusters to probe in the search) must be positive.");
   params_.nprobe = n_probes;
-  // Set the greater_
-  if (metric_type_ == raft::distance::DistanceType::L2Expanded ||
-      metric_type_ == raft::distance::DistanceType::L2Unexpanded) {
-    greater_ = false;
-  } else {
-    // Need to set this to true for inner product if need FAISS like behavior for inner product
-    greater_ = false;
+
+  switch (index_->metric) {
+    case raft::distance::DistanceType::InnerProduct:
+    case raft::distance::DistanceType::CosineExpanded:
+    case raft::distance::DistanceType::CorrelationExpanded:
+      // Similarity metrics have the opposite meaning, i.e. nearest neigbours are those with larger
+      // similarity (See the same logic at cpp/include/raft/sparse/selection/detail/knn.cuh:362
+      // {perform_k_selection})
+      greater_ = true;
+      break;
+    default: greater_ = false;
   }
 
   // Set memory buffer to be reused across searches
@@ -376,7 +375,7 @@ void cuivflHandle<T>::cuivflSearchImpl(const T* queries,  // [numQueries, dim]
   float alpha = 1.0f;
   float beta  = 0.0f;
 
-  if (metric_type_ == raft::distance::DistanceType::L2Expanded) {
+  if (index_->metric == raft::distance::DistanceType::L2Expanded) {
     alpha = -2.0f;
     beta  = 1.0f;
     utils::dots_along_rows(
@@ -447,7 +446,7 @@ void cuivflHandle<T>::cuivflSearchImpl(const T* queries,  // [numQueries, dim]
                                                                   queries,
                                                                   coarse_indices_dev.data(),
                                                                   n_queries,
-                                                                  metric_type_,
+                                                                  index_->metric,
                                                                   n_probes,
                                                                   k,
                                                                   greater_,

diff --git a/cpp/include/raft/spatial/knn/detail/ann_ivf_flat_kernel.cuh b/cpp/include/raft/spatial/knn/detail/ann_ivf_flat_kernel.cuh
@@ -27,6 +27,7 @@
 #include <raft/cuda_utils.cuh>
 #include <raft/cudart_utils.h>
 #include <raft/distance/distance.cuh>
+#include <raft/distance/distance_type.hpp>
 #include <raft/pow2_utils.cuh>
 
 #ifdef USE_FAISS
@@ -55,6 +56,8 @@ struct ivf_flat_index {
    * Vectorized load/store size in elements, determines the size of interleaved data chunks.
    */
   uint32_t veclen;
+  /** Distance metric used for clustering. */
+  raft::distance::DistanceType metric;
 
   /**
    * Inverted list data [size, dim].
@@ -82,15 +85,13 @@ struct ivf_flat_index {
   device_mdarray<T, extent_2d, row_major> data;
   /** Inverted list indices: ids of items in the source data [size] */
   device_mdarray<uint32_t, extent_1d, row_major> indices;
-
   /** Sizes of the lists (clusters) [n_lists] */
   device_mdarray<uint32_t, extent_1d, row_major> list_sizes;
   /**
    * Offsets into the lists [n_lists + 1].
    * The last value contains the total length of the index.
    */
   device_mdarray<uint32_t, extent_1d, row_major> list_offsets;
-
   /** k-means cluster centers corresponding to the lists [n_lists, dim] */
   device_mdarray<float, extent_2d, row_major> centers;
   /** (Optional) Precomputed norms of the `centers` w.r.t. the chosen distance metrix [n_lists]  */

diff --git a/cpp/include/raft/spatial/knn/detail/ann_quantized.cuh b/cpp/include/raft/spatial/knn/detail/ann_quantized.cuh
@@ -97,8 +97,8 @@ void approx_knn_cuivfl_ivfflat_build_index(const raft::handle_t& handle,
                                            IntType n,
                                            IntType D)
 {
-  index->ivf_flat<T>() = std::make_unique<detail::cuivflHandle<T>>(handle, metric, *params);
-  index->ivf_flat<T>()->cuivflBuildIndex(dataset, n, D);
+  index->ivf_flat<T>() = std::make_unique<detail::cuivflHandle<T>>(handle, *params);
+  index->ivf_flat<T>()->cuivflBuildIndex(dataset, n, D, metric);
 }
 
 template <typename IntType = int>

diff --git a/cpp/test/spatial/ann_base_kernel.cuh b/cpp/test/spatial/ann_base_kernel.cuh
@@ -93,7 +93,7 @@ void naiveBfKnn(float* dist_topk,
                              input_len,
                              dist_topk + offset * k,
                              indices_topk + offset * k,
-                             true,
+                             type != raft::distance::DistanceType::InnerProduct,
                              static_cast<int>(k),
                              stream,
                              SelectKAlgo::WARP_SORT);