From 1be5e3acfc7849568314c81298dcf774dbee6d27 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@users.noreply.github.com>
Date: Wed, 16 Jun 2021 10:48:00 -0400
Subject: [PATCH] Moving FAISS ANN wrapper to raft (#3963)

This accompanies the RAFT PR: https://github.com/rapidsai/raft/pull/265

Closes #3954

Authors:
  - Corey J. Nolet (https://github.com/cjnolet)

Approvers:
  - Victor Lafargue (https://github.com/viclafargue)
  - Dante Gama Dessavre (https://github.com/dantegd)

URL: https://github.com/rapidsai/cuml/pull/3963
---
 cpp/cmake/thirdparty/get_raft.cmake         |   1 -
 cpp/include/cuml/neighbors/knn.hpp          |  56 +---
 cpp/src/knn/knn.cu                          |  18 +-
 cpp/src/knn/knn_opg_common.cuh              |   2 +-
 cpp/src_prims/selection/knn.cuh             | 310 --------------------
 python/cuml/neighbors/ann.pxd               |   7 +-
 python/cuml/neighbors/nearest_neighbors.pyx |   4 -
 7 files changed, 23 insertions(+), 375 deletions(-)

diff --git a/cpp/cmake/thirdparty/get_raft.cmake b/cpp/cmake/thirdparty/get_raft.cmake
index c784df6d49..abef8830d4 100644
--- a/cpp/cmake/thirdparty/get_raft.cmake
+++ b/cpp/cmake/thirdparty/get_raft.cmake
@@ -30,7 +30,6 @@ function(find_and_configure_raft)
             SOURCE_SUBDIR  cpp
             OPTIONS
               "BUILD_TESTS OFF"
-
     )
 
     message(VERBOSE "CUML: Using RAFT located in ${raft_SOURCE_DIR}")
diff --git a/cpp/include/cuml/neighbors/knn.hpp b/cpp/include/cuml/neighbors/knn.hpp
index 4bdfa1a713..80e807ac80 100644
--- a/cpp/include/cuml/neighbors/knn.hpp
+++ b/cpp/include/cuml/neighbors/knn.hpp
@@ -16,59 +16,14 @@
 
 #pragma once
 
-#include <faiss/gpu/GpuIndex.h>
-#include <faiss/gpu/StandardGpuResources.h>
 #include <raft/linalg/distance_type.h>
+#include <raft/spatial/knn/ann_common.h>
 
 namespace raft {
 class handle_t;
 }
 
 namespace ML {
-struct knnIndex {
-  faiss::gpu::GpuIndex *index;
-  raft::distance::DistanceType metric;
-  float metricArg;
-
-  faiss::gpu::StandardGpuResources *gpu_res;
-  int device;
-  ~knnIndex() {
-    delete index;
-    delete gpu_res;
-  }
-};
-
-typedef enum {
-  QT_8bit,
-  QT_4bit,
-  QT_8bit_uniform,
-  QT_4bit_uniform,
-  QT_fp16,
-  QT_8bit_direct,
-  QT_6bit
-} QuantizerType;
-
-struct knnIndexParam {
-  virtual ~knnIndexParam() {}
-};
-
-struct IVFParam : knnIndexParam {
-  int nlist;
-  int nprobe;
-};
-
-struct IVFFlatParam : IVFParam {};
-
-struct IVFPQParam : IVFParam {
-  int M;
-  int n_bits;
-  bool usePrecomputedTables;
-};
-
-struct IVFSQParam : IVFParam {
-  QuantizerType qtype;
-  bool encodeResidual;
-};
 
 /**
  * @brief Flat C++ API function to perform a brute force knn on
@@ -112,8 +67,9 @@ void brute_force_knn(const raft::handle_t &handle, std::vector<float *> &input,
  * @param[in] n number of rows in the index array
  * @param[in] D the dimensionality of the index array
  */
-void approx_knn_build_index(raft::handle_t &handle, ML::knnIndex *index,
-                            ML::knnIndexParam *params,
+void approx_knn_build_index(raft::handle_t &handle,
+                            raft::spatial::knn::knnIndex *index,
+                            raft::spatial::knn::knnIndexParam *params,
                             raft::distance::DistanceType metric,
                             float metricArg, float *index_array, int n, int D);
 
@@ -131,8 +87,8 @@ void approx_knn_build_index(raft::handle_t &handle, ML::knnIndex *index,
  * @param[in] n number of rows in the query array
  */
 void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, ML::knnIndex *index, int k,
-                       float *query_array, int n);
+                       int64_t *indices, raft::spatial::knn::knnIndex *index,
+                       int k, float *query_array, int n);
 
 /**
  * @brief Flat C++ API function to perform a knn classification using a
diff --git a/cpp/src/knn/knn.cu b/cpp/src/knn/knn.cu
index adc0591518..15a5201d14 100644
--- a/cpp/src/knn/knn.cu
+++ b/cpp/src/knn/knn.cu
@@ -20,6 +20,7 @@
 #include <ml_mg_utils.cuh>
 
 #include <label/classlabels.cuh>
+#include <raft/spatial/knn/ann.hpp>
 #include <raft/spatial/knn/knn.hpp>
 #include <selection/knn.cuh>
 
@@ -44,19 +45,20 @@ void brute_force_knn(const raft::handle_t &handle, std::vector<float *> &input,
     rowMajorQuery, nullptr, metric, metric_arg);
 }
 
-void approx_knn_build_index(raft::handle_t &handle, ML::knnIndex *index,
-                            ML::knnIndexParam *params,
+void approx_knn_build_index(raft::handle_t &handle,
+                            raft::spatial::knn::knnIndex *index,
+                            raft::spatial::knn::knnIndexParam *params,
                             raft::distance::DistanceType metric,
                             float metricArg, float *index_array, int n, int D) {
-  MLCommon::Selection::approx_knn_build_index(handle, index, params, metric,
-                                              metricArg, index_array, n, D);
+  raft::spatial::knn::approx_knn_build_index(handle, index, params, metric,
+                                             metricArg, index_array, n, D);
 }
 
 void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, ML::knnIndex *index, int k,
-                       float *query_array, int n) {
-  MLCommon::Selection::approx_knn_search(handle, distances, indices, index, k,
-                                         query_array, n);
+                       int64_t *indices, raft::spatial::knn::knnIndex *index,
+                       int k, float *query_array, int n) {
+  raft::spatial::knn::approx_knn_search(handle, distances, indices, index, k,
+                                        query_array, n);
 }
 
 void knn_classify(raft::handle_t &handle, int *out, int64_t *knn_indices,
diff --git a/cpp/src/knn/knn_opg_common.cuh b/cpp/src/knn/knn_opg_common.cuh
index 80b64b0afa..78e3635540 100644
--- a/cpp/src/knn/knn_opg_common.cuh
+++ b/cpp/src/knn/knn_opg_common.cuh
@@ -692,7 +692,7 @@ void reduce(opg_knn_param<in_t, ind_t, dist_t, out_t> &params,
   }
 
   // Merge all KNN local results
-  MLCommon::Selection::knn_merge_parts(
+  raft::spatial::knn::knn_merge_parts(
     work.res_D.data(), work.res_I.data(), distances, indices, batch_size,
     work.idxRanks.size(), params.k, handle.get_stream(), trans.data());
   CUDA_CHECK(cudaStreamSynchronize(handle.get_stream()));
diff --git a/cpp/src_prims/selection/knn.cuh b/cpp/src_prims/selection/knn.cuh
index 4b5cd83539..0f08df8f36 100644
--- a/cpp/src_prims/selection/knn.cuh
+++ b/cpp/src_prims/selection/knn.cuh
@@ -63,316 +63,6 @@ inline __device__ T get_lbls(const T *labels, const int64_t *knn_indices,
   }
 }
 
-template <typename value_idx = int64_t, typename value_t = float, int warp_q,
-          int thread_q, int tpb>
-__global__ void knn_merge_parts_kernel(value_t *inK, value_idx *inV,
-                                       value_t *outK, value_idx *outV,
-                                       size_t n_samples, int n_parts,
-                                       value_t initK, value_idx initV, int k,
-                                       value_idx *translations) {
-  constexpr int kNumWarps = tpb / faiss::gpu::kWarpSize;
-
-  __shared__ value_t smemK[kNumWarps * warp_q];
-  __shared__ value_idx smemV[kNumWarps * warp_q];
-
-  /**
-   * Uses shared memory
-   */
-  faiss::gpu::BlockSelect<value_t, value_idx, false,
-                          faiss::gpu::Comparator<value_t>, warp_q, thread_q,
-                          tpb>
-    heap(initK, initV, smemK, smemV, k);
-
-  // Grid is exactly sized to rows available
-  int row = blockIdx.x;
-  int total_k = k * n_parts;
-
-  int i = threadIdx.x;
-
-  // Get starting pointers for cols in current thread
-  int part = i / k;
-  size_t row_idx = (row * k) + (part * n_samples * k);
-
-  int col = i % k;
-
-  value_t *inKStart = inK + (row_idx + col);
-  value_idx *inVStart = inV + (row_idx + col);
-
-  int limit = faiss::gpu::utils::roundDown(total_k, faiss::gpu::kWarpSize);
-  value_idx translation = 0;
-
-  for (; i < limit; i += tpb) {
-    translation = translations[part];
-    heap.add(*inKStart, (*inVStart) + translation);
-
-    part = (i + tpb) / k;
-    row_idx = (row * k) + (part * n_samples * k);
-
-    col = (i + tpb) % k;
-
-    inKStart = inK + (row_idx + col);
-    inVStart = inV + (row_idx + col);
-  }
-
-  // Handle last remainder fraction of a warp of elements
-  if (i < total_k) {
-    translation = translations[part];
-    heap.addThreadQ(*inKStart, (*inVStart) + translation);
-  }
-
-  heap.reduce();
-
-  for (int i = threadIdx.x; i < k; i += tpb) {
-    outK[row * k + i] = smemK[i];
-    outV[row * k + i] = smemV[i];
-  }
-}
-
-template <typename value_idx = int64_t, typename value_t = float, int warp_q,
-          int thread_q>
-inline void knn_merge_parts_impl(value_t *inK, value_idx *inV, value_t *outK,
-                                 value_idx *outV, size_t n_samples, int n_parts,
-                                 int k, cudaStream_t stream,
-                                 value_idx *translations) {
-  auto grid = dim3(n_samples);
-
-  constexpr int n_threads = (warp_q <= 1024) ? 128 : 64;
-  auto block = dim3(n_threads);
-
-  auto kInit = faiss::gpu::Limits<value_t>::getMax();
-  auto vInit = -1;
-  knn_merge_parts_kernel<value_idx, value_t, warp_q, thread_q, n_threads>
-    <<<grid, block, 0, stream>>>(inK, inV, outK, outV, n_samples, n_parts,
-                                 kInit, vInit, k, translations);
-  CUDA_CHECK(cudaPeekAtLastError());
-}
-
-/**
- * @brief Merge knn distances and index matrix, which have been partitioned
- * by row, into a single matrix with only the k-nearest neighbors.
- *
- * @param inK partitioned knn distance matrix
- * @param inV partitioned knn index matrix
- * @param outK merged knn distance matrix
- * @param outV merged knn index matrix
- * @param n_samples number of samples per partition
- * @param n_parts number of partitions
- * @param k number of neighbors per partition (also number of merged neighbors)
- * @param stream CUDA stream to use
- * @param translations mapping of index offsets for each partition
- */
-template <typename value_idx = int64_t, typename value_t = float>
-inline void knn_merge_parts(value_t *inK, value_idx *inV, value_t *outK,
-                            value_idx *outV, size_t n_samples, int n_parts,
-                            int k, cudaStream_t stream,
-                            value_idx *translations) {
-  if (k == 1)
-    knn_merge_parts_impl<value_idx, value_t, 1, 1>(
-      inK, inV, outK, outV, n_samples, n_parts, k, stream, translations);
-  else if (k <= 32)
-    knn_merge_parts_impl<value_idx, value_t, 32, 2>(
-      inK, inV, outK, outV, n_samples, n_parts, k, stream, translations);
-  else if (k <= 64)
-    knn_merge_parts_impl<value_idx, value_t, 64, 3>(
-      inK, inV, outK, outV, n_samples, n_parts, k, stream, translations);
-  else if (k <= 128)
-    knn_merge_parts_impl<value_idx, value_t, 128, 3>(
-      inK, inV, outK, outV, n_samples, n_parts, k, stream, translations);
-  else if (k <= 256)
-    knn_merge_parts_impl<value_idx, value_t, 256, 4>(
-      inK, inV, outK, outV, n_samples, n_parts, k, stream, translations);
-  else if (k <= 512)
-    knn_merge_parts_impl<value_idx, value_t, 512, 8>(
-      inK, inV, outK, outV, n_samples, n_parts, k, stream, translations);
-  else if (k <= 1024)
-    knn_merge_parts_impl<value_idx, value_t, 1024, 8>(
-      inK, inV, outK, outV, n_samples, n_parts, k, stream, translations);
-}
-
-inline faiss::MetricType build_faiss_metric(
-  raft::distance::DistanceType metric) {
-  switch (metric) {
-    case raft::distance::DistanceType::CosineExpanded:
-      return faiss::MetricType::METRIC_INNER_PRODUCT;
-    case raft::distance::DistanceType::CorrelationExpanded:
-      return faiss::MetricType::METRIC_INNER_PRODUCT;
-    case raft::distance::DistanceType::L2Expanded:
-      return faiss::MetricType::METRIC_L2;
-    case raft::distance::DistanceType::L2Unexpanded:
-      return faiss::MetricType::METRIC_L2;
-    case raft::distance::DistanceType::L2SqrtExpanded:
-      return faiss::MetricType::METRIC_L2;
-    case raft::distance::DistanceType::L2SqrtUnexpanded:
-      return faiss::MetricType::METRIC_L2;
-    case raft::distance::DistanceType::L1:
-      return faiss::MetricType::METRIC_L1;
-    case raft::distance::DistanceType::InnerProduct:
-      return faiss::MetricType::METRIC_INNER_PRODUCT;
-    case raft::distance::DistanceType::LpUnexpanded:
-      return faiss::MetricType::METRIC_Lp;
-    case raft::distance::DistanceType::Linf:
-      return faiss::MetricType::METRIC_Linf;
-    case raft::distance::DistanceType::Canberra:
-      return faiss::MetricType::METRIC_Canberra;
-    case raft::distance::DistanceType::BrayCurtis:
-      return faiss::MetricType::METRIC_BrayCurtis;
-    case raft::distance::DistanceType::JensenShannon:
-      return faiss::MetricType::METRIC_JensenShannon;
-    default:
-      THROW("MetricType not supported: %d", metric);
-  }
-}
-
-inline faiss::ScalarQuantizer::QuantizerType build_faiss_qtype(
-  ML::QuantizerType qtype) {
-  switch (qtype) {
-    case ML::QuantizerType::QT_8bit:
-      return faiss::ScalarQuantizer::QuantizerType::QT_8bit;
-    case ML::QuantizerType::QT_8bit_uniform:
-      return faiss::ScalarQuantizer::QuantizerType::QT_8bit_uniform;
-    case ML::QuantizerType::QT_4bit_uniform:
-      return faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform;
-    case ML::QuantizerType::QT_fp16:
-      return faiss::ScalarQuantizer::QuantizerType::QT_fp16;
-    case ML::QuantizerType::QT_8bit_direct:
-      return faiss::ScalarQuantizer::QuantizerType::QT_8bit_direct;
-    case ML::QuantizerType::QT_6bit:
-      return faiss::ScalarQuantizer::QuantizerType::QT_6bit;
-    default:
-      return (faiss::ScalarQuantizer::QuantizerType)qtype;
-  }
-}
-
-template <typename IntType = int>
-void approx_knn_ivfflat_build_index(ML::knnIndex *index, ML::IVFParam *params,
-                                    raft::distance::DistanceType metric,
-                                    IntType n, IntType D) {
-  faiss::gpu::GpuIndexIVFFlatConfig config;
-  config.device = index->device;
-  faiss::MetricType faiss_metric = build_faiss_metric(metric);
-  faiss::gpu::GpuIndexIVFFlat *faiss_index = new faiss::gpu::GpuIndexIVFFlat(
-    index->gpu_res, D, params->nlist, faiss_metric, config);
-  faiss_index->setNumProbes(params->nprobe);
-  index->index = faiss_index;
-}
-
-template <typename IntType = int>
-void approx_knn_ivfpq_build_index(ML::knnIndex *index, ML::IVFPQParam *params,
-                                  raft::distance::DistanceType metric,
-                                  IntType n, IntType D) {
-  faiss::gpu::GpuIndexIVFPQConfig config;
-  config.device = index->device;
-  config.usePrecomputedTables = params->usePrecomputedTables;
-  config.interleavedLayout = params->n_bits != 8;
-  faiss::MetricType faiss_metric = build_faiss_metric(metric);
-  faiss::gpu::GpuIndexIVFPQ *faiss_index =
-    new faiss::gpu::GpuIndexIVFPQ(index->gpu_res, D, params->nlist, params->M,
-                                  params->n_bits, faiss_metric, config);
-  faiss_index->setNumProbes(params->nprobe);
-  index->index = faiss_index;
-}
-
-template <typename IntType = int>
-void approx_knn_ivfsq_build_index(ML::knnIndex *index, ML::IVFSQParam *params,
-                                  raft::distance::DistanceType metric,
-                                  IntType n, IntType D) {
-  faiss::gpu::GpuIndexIVFScalarQuantizerConfig config;
-  config.device = index->device;
-  faiss::MetricType faiss_metric = build_faiss_metric(metric);
-  faiss::ScalarQuantizer::QuantizerType faiss_qtype =
-    build_faiss_qtype(params->qtype);
-  faiss::gpu::GpuIndexIVFScalarQuantizer *faiss_index =
-    new faiss::gpu::GpuIndexIVFScalarQuantizer(index->gpu_res, D, params->nlist,
-                                               faiss_qtype, faiss_metric,
-                                               params->encodeResidual);
-  faiss_index->setNumProbes(params->nprobe);
-  index->index = faiss_index;
-}
-
-template <typename IntType = int>
-void approx_knn_build_index(raft::handle_t &handle, ML::knnIndex *index,
-                            ML::knnIndexParam *params,
-                            raft::distance::DistanceType metric,
-                            float metricArg, float *index_array, IntType n,
-                            IntType D) {
-  int device;
-  CUDA_CHECK(cudaGetDevice(&device));
-
-  faiss::gpu::StandardGpuResources *gpu_res =
-    new faiss::gpu::StandardGpuResources();
-  gpu_res->noTempMemory();
-  gpu_res->setDefaultStream(device, handle.get_stream());
-  index->gpu_res = gpu_res;
-  index->device = device;
-  index->index = nullptr;
-  index->metric = metric;
-  index->metricArg = metricArg;
-
-  // perform preprocessing
-  // k set to 0 (unused during preprocessing / revertion)
-  std::unique_ptr<MetricProcessor<float>> query_metric_processor =
-    create_processor<float>(metric, n, D, 0, false, handle.get_stream(),
-                            handle.get_device_allocator());
-
-  query_metric_processor->preprocess(index_array);
-
-  if (dynamic_cast<ML::IVFFlatParam *>(params)) {
-    ML::IVFFlatParam *IVFFlat_param = dynamic_cast<ML::IVFFlatParam *>(params);
-    approx_knn_ivfflat_build_index(index, IVFFlat_param, metric, n, D);
-    std::vector<float> h_index_array(n * D);
-    raft::update_host(h_index_array.data(), index_array, h_index_array.size(),
-                      handle.get_stream());
-    query_metric_processor->revert(index_array);
-    index->index->train(n, h_index_array.data());
-    index->index->add(n, h_index_array.data());
-  } else {
-    if (dynamic_cast<ML::IVFPQParam *>(params)) {
-      ML::IVFPQParam *IVFPQ_param = dynamic_cast<ML::IVFPQParam *>(params);
-      approx_knn_ivfpq_build_index(index, IVFPQ_param, metric, n, D);
-    } else if (dynamic_cast<ML::IVFSQParam *>(params)) {
-      ML::IVFSQParam *IVFSQ_param = dynamic_cast<ML::IVFSQParam *>(params);
-      approx_knn_ivfsq_build_index(index, IVFSQ_param, metric, n, D);
-    } else {
-      ASSERT(index->index, "KNN index could not be initialized");
-    }
-
-    index->index->train(n, index_array);
-    index->index->add(n, index_array);
-    query_metric_processor->revert(index_array);
-  }
-}
-
-template <typename IntType = int>
-void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, ML::knnIndex *index, IntType k,
-                       float *query_array, IntType n) {
-  // perform preprocessing
-  std::unique_ptr<MetricProcessor<float>> query_metric_processor =
-    create_processor<float>(index->metric, n, index->index->d, k, false,
-                            handle.get_stream(), handle.get_device_allocator());
-
-  query_metric_processor->preprocess(query_array);
-  index->index->search(n, query_array, k, distances, indices);
-  query_metric_processor->revert(query_array);
-
-  // Perform necessary post-processing
-  if (index->metric == raft::distance::DistanceType::L2SqrtExpanded ||
-      index->metric == raft::distance::DistanceType::L2SqrtUnexpanded ||
-      index->metric == raft::distance::DistanceType::LpUnexpanded) {
-    /**
-  * post-processing
-  */
-    float p = 0.5;  // standard l2
-    if (index->metric == raft::distance::DistanceType::LpUnexpanded)
-      p = 1.0 / index->metricArg;
-    raft::linalg::unaryOp<float>(
-      distances, distances, n * k,
-      [p] __device__(float input) { return powf(input, p); },
-      handle.get_stream());
-  }
-  query_metric_processor->postprocess(distances);
-}
-
 template <typename OutType = float, bool precomp_lbls = false>
 __global__ void class_probs_kernel(OutType *out, const int64_t *knn_indices,
                                    const int *labels, int n_uniq_labels,
diff --git a/python/cuml/neighbors/ann.pxd b/python/cuml/neighbors/ann.pxd
index 3719485bf3..1d667fba84 100644
--- a/python/cuml/neighbors/ann.pxd
+++ b/python/cuml/neighbors/ann.pxd
@@ -21,7 +21,12 @@ from libc.stdint cimport uintptr_t
 from libcpp cimport bool
 
 
-cdef extern from "cuml/neighbors/knn.hpp" namespace "ML":
+cdef extern from "raft/spatial/knn/ann_common.h" \
+        namespace "raft::spatial::knn":
+
+    cdef cppclass knnIndex:
+        pass
+
     cdef cppclass knnIndexParam:
         pass
 
diff --git a/python/cuml/neighbors/nearest_neighbors.pyx b/python/cuml/neighbors/nearest_neighbors.pyx
index 4c7f82081a..01715d514b 100644
--- a/python/cuml/neighbors/nearest_neighbors.pyx
+++ b/python/cuml/neighbors/nearest_neighbors.pyx
@@ -64,10 +64,6 @@ if has_scipy():
     import scipy.sparse
 
 cdef extern from "cuml/neighbors/knn.hpp" namespace "ML":
-
-    cdef cppclass knnIndex:
-        pass
-
     void brute_force_knn(
         const handle_t &handle,
         vector[float*] &inputs,