Add ANN refinement method (#1038)

This PR implements refinement for approximate nearest neighbor search.

Refinement is a post processing step for ANN search, it follows an ANN search that returned `k0` neighbor candidates,
and select `k` out of these candidates. The selection by calculating exact distances from the original dataset.

Refinement can increase accuracy. It is useful for ANN methods that quantize the dataset and therefore loose accuracy during distance calculation (e.g. IVF-PQ).

Authors:
  - Tamas Bela Feher (https://github.com/tfeher)

Approvers:
  - Robert Maynard (https://github.com/robertmaynard)
  - Artem M. Chirkin (https://github.com/achirkin)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #1038

cpp/bench/CMakeLists.txt

@@ -124,6 +124,7 @@ if(BUILD_BENCH)
     bench/neighbors/knn/ivf_pq_float_uint32_t.cu
     bench/neighbors/knn/ivf_pq_int8_t_int64_t.cu
     bench/neighbors/knn/ivf_pq_uint8_t_uint32_t.cu
+    bench/neighbors/refine.cu
     bench/neighbors/selection.cu
     bench/main.cpp
     OPTIONAL

cpp/bench/neighbors/refine.cu

@@ -0,0 +1,122 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include <common/benchmark.hpp>
+
+#include <raft/random/rng.cuh>
+
+#include <raft/core/device_mdspan.hpp>
+#include <raft/core/handle.hpp>
+#include <raft/distance/distance_types.hpp>
+#include <raft/neighbors/detail/refine.cuh>
+#include <raft/neighbors/refine.cuh>
+
+#if defined RAFT_DISTANCE_COMPILED
+#include <raft/distance/specializations.cuh>
+#endif
+
+#if defined RAFT_NN_COMPILED
+#include <raft/spatial/knn/specializations.cuh>
+#endif
+
+#include <rmm/cuda_stream_view.hpp>
+#include <rmm/mr/device/per_device_resource.hpp>
+#include <rmm/mr/device/pool_memory_resource.hpp>
+
+#include "../../test/neighbors/refine_helper.cuh"
+
+#include <iostream>
+#include <sstream>
+
+using namespace raft::neighbors::detail;
+
+namespace raft::bench::neighbors {
+
+template <typename IdxT>
+inline auto operator<<(std::ostream& os, const RefineInputs<IdxT>& p) -> std::ostream&
+{
+  os << p.n_rows << "#" << p.dim << "#" << p.n_queries << "#" << p.k0 << "#" << p.k << "#"
+     << (p.host_data ? "host" : "device");
+  return os;
+}
+
+RefineInputs<int64_t> p;
+
+template <typename DataT, typename DistanceT, typename IdxT>
+class RefineAnn : public fixture {
+ public:
+  RefineAnn(RefineInputs<IdxT> p) : data(handle_, p) {}
+
+  void run_benchmark(::benchmark::State& state) override
+  {
+    std::ostringstream label_stream;
+    label_stream << data.p;
+    state.SetLabel(label_stream.str());
+
+    auto old_mr = rmm::mr::get_current_device_resource();
+    rmm::mr::pool_memory_resource<rmm::mr::device_memory_resource> pool_mr(old_mr);
+    rmm::mr::set_current_device_resource(&pool_mr);
+
+    if (data.p.host_data) {
+      loop_on_state(state, [this]() {
+        raft::neighbors::refine<IdxT, DataT, DistanceT, IdxT>(handle_,
+                                                              data.dataset_host.view(),
+                                                              data.queries_host.view(),
+                                                              data.candidates_host.view(),
+                                                              data.refined_indices_host.view(),
+                                                              data.refined_distances_host.view(),
+                                                              data.p.metric);
+      });
+    } else {
+      loop_on_state(state, [&]() {
+        raft::neighbors::refine<IdxT, DataT, DistanceT, IdxT>(handle_,
+                                                              data.dataset.view(),
+                                                              data.queries.view(),
+                                                              data.candidates.view(),
+                                                              data.refined_indices.view(),
+                                                              data.refined_distances.view(),
+                                                              data.p.metric);
+      });
+    }
+    rmm::mr::set_current_device_resource(old_mr);
+  }
+
+ private:
+  raft::handle_t handle_;
+  RefineHelper<DataT, DistanceT, IdxT> data;
+};
+
+std::vector<RefineInputs<int64_t>> getInputs()
+{
+  std::vector<RefineInputs<int64_t>> out;
+  raft::distance::DistanceType metric = raft::distance::DistanceType::L2Expanded;
+  for (bool host_data : {true, false}) {
+    for (int64_t n_queries : {1000, 10000}) {
+      for (int64_t dim : {128, 512}) {
+        out.push_back(RefineInputs<int64_t>{n_queries, 2000000, dim, 32, 128, metric, host_data});
+        out.push_back(RefineInputs<int64_t>{n_queries, 2000000, dim, 10, 40, metric, host_data});
+      }
+    }
+  }
+  return out;
+}
+
+using refine_float_int64 = RefineAnn<float, float, int64_t>;
+RAFT_BENCH_REGISTER(refine_float_int64, "", getInputs());
+
+using refine_uint8_int64 = RefineAnn<uint8_t, float, int64_t>;
+RAFT_BENCH_REGISTER(refine_uint8_int64, "", getInputs());
+}  // namespace raft::bench::neighbors

cpp/include/raft/neighbors/detail/refine.cuh

@@ -0,0 +1,232 @@
+/*
+ * Copyright (c) 2022, NVIDIA CORPORATION.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#pragma once
+
+#include <raft/core/device_mdarray.hpp>
+#include <raft/core/handle.hpp>
+#include <raft/core/host_mdspan.hpp>
+#include <raft/core/nvtx.hpp>
+#include <raft/spatial/knn/detail/ann_utils.cuh>
+#include <raft/spatial/knn/detail/ivf_flat_build.cuh>
+#include <raft/spatial/knn/detail/ivf_flat_search.cuh>
+
+#include <cstdlib>
+#include <omp.h>
+
+#include <thrust/sequence.h>
+
+namespace raft::neighbors::detail {
+
+/** Checks whether the input data extents are compatible. */
+template <typename extents_t>
+void check_input(extents_t dataset,
+                 extents_t queries,
+                 extents_t candidates,
+                 extents_t indices,
+                 extents_t distances,
+                 distance::DistanceType metric)
+{
+  auto n_queries = queries.extent(0);
+  auto k         = distances.extent(1);
+
+  RAFT_EXPECTS(k <= raft::spatial::knn::detail::topk::kMaxCapacity,
+               "k must be lest than topk::kMaxCapacity (%d).",
+               raft::spatial::knn::detail::topk::kMaxCapacity);
+
+  RAFT_EXPECTS(indices.extent(0) == n_queries && distances.extent(0) == n_queries &&
+                 candidates.extent(0) == n_queries,
+               "Number of rows in output indices and distances matrices must equal number of rows "
+               "in search matrix.");
+
+  RAFT_EXPECTS(indices.extent(1) == k,
+               "Number of columns in output indices and distances matrices must be equal to k");
+
+  RAFT_EXPECTS(queries.extent(1) == dataset.extent(1),
+               "Number of columns must be equal for dataset and queries");
+
+  RAFT_EXPECTS(candidates.extent(1) >= k,
+               "Number of neighbor candidates must not be smaller than k (%d vs %d)",
+               static_cast<int>(candidates.extent(1)),
+               static_cast<int>(k));
+}
+
+/**
+ * See raft::neighbors::refine for docs.
+ */
+template <typename idx_t, typename data_t, typename distance_t, typename matrix_idx>
+void refine_device(raft::handle_t const& handle,
+                   raft::device_matrix_view<const data_t, matrix_idx, row_major> dataset,
+                   raft::device_matrix_view<const data_t, matrix_idx, row_major> queries,
+                   raft::device_matrix_view<const idx_t, matrix_idx, row_major> neighbor_candidates,
+                   raft::device_matrix_view<idx_t, matrix_idx, row_major> indices,
+                   raft::device_matrix_view<distance_t, matrix_idx, row_major> distances,
+                   distance::DistanceType metric = distance::DistanceType::L2Unexpanded)
+{
+  matrix_idx n_candidates = neighbor_candidates.extent(1);
+  matrix_idx n_queries    = queries.extent(0);
+  matrix_idx dim          = dataset.extent(1);
+  uint32_t k              = static_cast<uint32_t>(indices.extent(1));
+
+  common::nvtx::range<common::nvtx::domain::raft> fun_scope(
+    "neighbors::refine(%zu, %u)", size_t(n_queries), uint32_t(n_candidates));
+
+  check_input(dataset.extents(),
+              queries.extents(),
+              neighbor_candidates.extents(),
+              indices.extents(),
+              distances.extents(),
+              metric);
+
+  // The refinement search can be mapped to an IVF flat search:
+  // - We consider that the candidate vectors form a cluster, separately for each query.
+  // - In other words, the n_queries * n_candidates vectors form n_queries clusters, each with
+  //   n_candidates elements.
+  // - We consider that the coarse level search is already performed and assigned a single cluster
+  //   to search for each query (the cluster formed from the corresponding candidates).
+  // - We run IVF flat search with n_probes=1 to select the best k elements of the candidates.
+  rmm::device_uvector<uint32_t> fake_coarse_idx(n_queries, handle.get_stream());
+
+  thrust::sequence(
+    handle.get_thrust_policy(), fake_coarse_idx.data(), fake_coarse_idx.data() + n_queries);
+
+  raft::neighbors::ivf_flat::index<data_t, idx_t> refinement_index(
+    handle, metric, n_queries, false, dim);
+
+  raft::spatial::knn::ivf_flat::detail::fill_refinement_index(handle,
+                                                              &refinement_index,
+                                                              dataset.data_handle(),
+                                                              neighbor_candidates.data_handle(),
+                                                              n_queries,
+                                                              n_candidates);
+
+  uint32_t grid_dim_x = 1;
+  raft::spatial::knn::ivf_flat::detail::ivfflat_interleaved_scan<
+    data_t,
+    typename raft::spatial::knn::detail::utils::config<data_t>::value_t,
+    idx_t>(refinement_index,
+           queries.data_handle(),
+           fake_coarse_idx.data(),
+           static_cast<uint32_t>(n_queries),
+           refinement_index.metric(),
+           1,
+           k,
+           raft::spatial::knn::ivf_flat::detail::is_min_close(metric),
+           indices.data_handle(),
+           distances.data_handle(),
+           grid_dim_x,
+           handle.get_stream());
+}
+
+/** Helper structure for naive CPU implementation of refine. */
+typedef struct {
+  uint64_t id;
+  float distance;
+} struct_for_refinement;
+
+int _postprocessing_qsort_compare(const void* v1, const void* v2)
+{
+  // sort in ascending order
+  if (((struct_for_refinement*)v1)->distance > ((struct_for_refinement*)v2)->distance) {
+    return 1;
+  } else if (((struct_for_refinement*)v1)->distance < ((struct_for_refinement*)v2)->distance) {
+    return -1;
+  } else {
+    return 0;
+  }
+}
+
+/**
+ * Naive CPU implementation of refine operation
+ *
+ * All pointers are expected to be accessible on the host.
+ */
+template <typename idx_t, typename data_t, typename distance_t, typename matrix_idx>
+void refine_host(raft::host_matrix_view<const data_t, matrix_idx, row_major> dataset,
+                 raft::host_matrix_view<const data_t, matrix_idx, row_major> queries,
+                 raft::host_matrix_view<const idx_t, matrix_idx, row_major> neighbor_candidates,
+                 raft::host_matrix_view<idx_t, matrix_idx, row_major> indices,
+                 raft::host_matrix_view<distance_t, matrix_idx, row_major> distances,
+                 distance::DistanceType metric = distance::DistanceType::L2Unexpanded)
+{
+  check_input(dataset.extents(),
+              queries.extents(),
+              neighbor_candidates.extents(),
+              indices.extents(),
+              distances.extents(),
+              metric);
+
+  switch (metric) {
+    case raft::distance::DistanceType::L2Expanded: break;
+    case raft::distance::DistanceType::InnerProduct: break;
+    default: throw raft::logic_error("Unsopported metric");
+  }
+
+  size_t numDataset            = dataset.extent(0);
+  size_t numQueries            = queries.extent(0);
+  size_t dimDataset            = dataset.extent(1);
+  const data_t* dataset_ptr    = dataset.data_handle();
+  const data_t* queries_ptr    = queries.data_handle();
+  const idx_t* neighbors       = neighbor_candidates.data_handle();
+  idx_t topK                   = neighbor_candidates.extent(1);
+  idx_t refinedTopK            = indices.extent(1);
+  idx_t* refinedNeighbors      = indices.data_handle();
+  distance_t* refinedDistances = distances.data_handle();
+
+  common::nvtx::range<common::nvtx::domain::raft> fun_scope(
+    "neighbors::refine_host(%zu, %u)", size_t(numQueries), uint32_t(topK));
+
+#pragma omp parallel
+  {
+    struct_for_refinement* sfr =
+      (struct_for_refinement*)malloc(sizeof(struct_for_refinement) * topK);
+    for (size_t i = omp_get_thread_num(); i < numQueries; i += omp_get_num_threads()) {
+      // compute distance with original dataset vectors
+      const data_t* cur_query = queries_ptr + ((uint64_t)dimDataset * i);
+      for (size_t j = 0; j < (size_t)topK; j++) {
+        idx_t id                  = neighbors[j + (topK * i)];
+        const data_t* cur_dataset = dataset_ptr + ((uint64_t)dimDataset * id);
+        float distance            = 0.0;
+        for (size_t k = 0; k < (size_t)dimDataset; k++) {
+          float val_q = (float)(cur_query[k]);
+          float val_d = (float)(cur_dataset[k]);
+          if (metric == raft::distance::DistanceType::InnerProduct) {
+            distance += -val_q * val_d;  // Negate because we sort in scending order.
+          } else {
+            distance += (val_q - val_d) * (val_q - val_d);
+          }
+        }
+        sfr[j].id       = id;
+        sfr[j].distance = distance;
+      }
+
+      qsort(sfr, topK, sizeof(struct_for_refinement), _postprocessing_qsort_compare);
+
+      for (size_t j = 0; j < (size_t)refinedTopK; j++) {
+        refinedNeighbors[j + (refinedTopK * i)] = sfr[j].id;
+        if (refinedDistances == NULL) continue;
+        if (metric == raft::distance::DistanceType::InnerProduct) {
+          refinedDistances[j + (refinedTopK * i)] = -sfr[j].distance;
+        } else {
+          refinedDistances[j + (refinedTopK * i)] = -sfr[j].distance;
+        }
+      }
+    }
+    free(sfr);
+  }
+}
+
+}  // namespace raft::neighbors::detail