Add ANN refinement method

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