From 90013ba85cb0d7d18e64b79152c4574bf17e6515 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:03:25 -0400
Subject: [PATCH 01/18] Adding pq ann

---
 .../raft/spatial/knn/detail/ivf_pq_ann.cuh    | 213 ++++++++++++++++++
 cpp/include/raft/spatial/knn/knn.hpp          |  86 +++++++
 2 files changed, 299 insertions(+)
 create mode 100644 cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
new file mode 100644
index 0000000000..935071d87d
--- /dev/null
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -0,0 +1,213 @@
+/*
+ * Copyright (c) 2020-2021, 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/cudart_utils.h>
+#include <raft/cuda_utils.cuh>
+
+#include <label/classlabels.cuh>
+#include <raft/distance/distance.cuh>
+
+#include <faiss/gpu/GpuDistance.h>
+#include <faiss/gpu/GpuIndexFlat.h>
+#include <faiss/gpu/GpuIndexIVFFlat.h>
+#include <faiss/gpu/GpuIndexIVFPQ.h>
+#include <faiss/gpu/GpuIndexIVFScalarQuantizer.h>
+#include <faiss/gpu/GpuResources.h>
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/utils/Heap.h>
+#include <faiss/gpu/utils/Limits.cuh>
+#include <faiss/gpu/utils/Select.cuh>
+#include <faiss/gpu/utils/Tensor.cuh>
+
+#include <thrust/device_vector.h>
+#include <thrust/iterator/transform_iterator.h>
+
+#include <raft/linalg/distance_type.h>
+#include "processing.cuh"
+
+#include "haversine_knn.cuh"
+
+#include <cuml/common/device_buffer.hpp>
+#include <cuml/neighbors/knn.hpp>
+#include <raft/mr/device/allocator.hpp>
+
+#include <iostream>
+#include <set>
+
+
+namespace raft {
+namespace spatial {
+namespace knn {
+namespace detail {
+
+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);
+}
+
+
+};
+};
+};
+};
\ No newline at end of file
diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index eae5a3c8cc..7049deb706 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -17,6 +17,10 @@
 #pragma once
 
 #include "detail/brute_force_knn.cuh"
+#include "detail/ivf_pq_ann.cuh"
+
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/GpuIndex.h>
 
 #include <raft/mr/device/allocator.hpp>
 #include <raft/mr/device/buffer.hpp>
@@ -27,6 +31,88 @@ namespace knn {
 
 using deviceAllocator = raft::mr::device::allocator;
 
+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 build an approximate nearest neighbors index
+ * from an index array and a set of parameters.
+ *
+ * @param[in] handle RAFT handle
+ * @param[out] index index to be built
+ * @param[in] params parametrization of the index to be built
+ * @param[in] metric distance metric to use. Euclidean (L2) is used by default
+ * @param[in] metricArg metric argument
+ * @param[in] index_array the index array to build the index with
+ * @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,
+                            raft::distance::DistanceType metric,
+                            float metricArg, float *index_array, IntType n,
+                            IntType D);
+
+/**
+ * @brief Flat C++ API function to perform an approximate nearest neighbors
+ * search from previously built index and a query array
+ *
+ * @param[in] handle RAFT handle
+ * @param[out] distances distances of the nearest neighbors toward
+ *                       their query point
+ * @param[out] indices indices of the nearest neighbors
+ * @param[in] index index to perform a search with
+ * @param[in] k the number of nearest neighbors to search for
+ * @param[in] query_array the query to perform a search with
+ * @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, IntType k,
+                       float *query_array, IntType n);
+
 /**
  * @brief Flat C++ API function to perform a brute force knn on
  * a series of input arrays and combine the results into a single

From 046944ff54dcb0e5101a3f647e629856148038d0 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:05:36 -0400
Subject: [PATCH 02/18] Removing unused inlcludes

---
 cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index 935071d87d..8d3f47f785 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -39,9 +39,6 @@
 #include <thrust/iterator/transform_iterator.h>
 
 #include <raft/linalg/distance_type.h>
-#include "processing.cuh"
-
-#include "haversine_knn.cuh"
 
 #include <cuml/common/device_buffer.hpp>
 #include <cuml/neighbors/knn.hpp>

From 08b5d601c9dd5eff08bb17afd06aa4d8be6b27d1 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:06:38 -0400
Subject: [PATCH 03/18] Using proper namespaces (to remove ML)

---
 .../raft/spatial/knn/detail/ivf_pq_ann.cuh    | 38 +++++++++----------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index 8d3f47f785..4494f0e81f 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -54,19 +54,19 @@ namespace knn {
 namespace detail {
 
 inline faiss::ScalarQuantizer::QuantizerType build_faiss_qtype(
-  ML::QuantizerType qtype) {
+  raft::spatial::knn::QuantizerType qtype) {
   switch (qtype) {
-    case ML::QuantizerType::QT_8bit:
+    case raft::spatial::knn::QuantizerType::QT_8bit:
       return faiss::ScalarQuantizer::QuantizerType::QT_8bit;
-    case ML::QuantizerType::QT_8bit_uniform:
+    case raft::spatial::knn::QuantizerType::QT_8bit_uniform:
       return faiss::ScalarQuantizer::QuantizerType::QT_8bit_uniform;
-    case ML::QuantizerType::QT_4bit_uniform:
+    case raft::spatial::knn::QuantizerType::QT_4bit_uniform:
       return faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform;
-    case ML::QuantizerType::QT_fp16:
+    case raft::spatial::knn::QuantizerType::QT_fp16:
       return faiss::ScalarQuantizer::QuantizerType::QT_fp16;
-    case ML::QuantizerType::QT_8bit_direct:
+    case raft::spatial::knn::QuantizerType::QT_8bit_direct:
       return faiss::ScalarQuantizer::QuantizerType::QT_8bit_direct;
-    case ML::QuantizerType::QT_6bit:
+    case raft::spatial::knn::QuantizerType::QT_6bit:
       return faiss::ScalarQuantizer::QuantizerType::QT_6bit;
     default:
       return (faiss::ScalarQuantizer::QuantizerType)qtype;
@@ -74,7 +74,7 @@ inline faiss::ScalarQuantizer::QuantizerType build_faiss_qtype(
 }
 
 template <typename IntType = int>
-void approx_knn_ivfflat_build_index(ML::knnIndex *index, ML::IVFParam *params,
+void approx_knn_ivfflat_build_index(raft::spatial::knn::knnIndex *index, raft::spatial::knn::IVFParam *params,
                                     raft::distance::DistanceType metric,
                                     IntType n, IntType D) {
   faiss::gpu::GpuIndexIVFFlatConfig config;
@@ -87,7 +87,7 @@ void approx_knn_ivfflat_build_index(ML::knnIndex *index, ML::IVFParam *params,
 }
 
 template <typename IntType = int>
-void approx_knn_ivfpq_build_index(ML::knnIndex *index, ML::IVFPQParam *params,
+void approx_knn_ivfpq_build_index(raft::spatial::knn::knnIndex *index, raft::spatial::knn::IVFPQParam *params,
                                   raft::distance::DistanceType metric,
                                   IntType n, IntType D) {
   faiss::gpu::GpuIndexIVFPQConfig config;
@@ -103,7 +103,7 @@ void approx_knn_ivfpq_build_index(ML::knnIndex *index, ML::IVFPQParam *params,
 }
 
 template <typename IntType = int>
-void approx_knn_ivfsq_build_index(ML::knnIndex *index, ML::IVFSQParam *params,
+void approx_knn_ivfsq_build_index(raft::spatial::knn::knnIndex *index, raft::spatial::knn::IVFSQParam *params,
                                   raft::distance::DistanceType metric,
                                   IntType n, IntType D) {
   faiss::gpu::GpuIndexIVFScalarQuantizerConfig config;
@@ -120,8 +120,8 @@ void approx_knn_ivfsq_build_index(ML::knnIndex *index, ML::IVFSQParam *params,
 }
 
 template <typename IntType = int>
-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, IntType n,
                             IntType D) {
@@ -146,8 +146,8 @@ void approx_knn_build_index(raft::handle_t &handle, ML::knnIndex *index,
 
   query_metric_processor->preprocess(index_array);
 
-  if (dynamic_cast<ML::IVFFlatParam *>(params)) {
-    ML::IVFFlatParam *IVFFlat_param = dynamic_cast<ML::IVFFlatParam *>(params);
+  if (dynamic_cast<raft::spatial::knn::IVFFlatParam *>(params)) {
+    raft::spatial::knn::IVFFlatParam *IVFFlat_param = dynamic_cast<raft::spatial::knn::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(),
@@ -156,11 +156,11 @@ void approx_knn_build_index(raft::handle_t &handle, ML::knnIndex *index,
     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);
+    if (dynamic_cast<raft::spatial::knn::IVFPQParam *>(params)) {
+      raft::spatial::knn::IVFPQParam *IVFPQ_param = dynamic_cast<raft::spatial::knn::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);
+    } else if (dynamic_cast<raft::spatial::knn::IVFSQParam *>(params)) {
+      raft::spatial::knn::IVFSQParam *IVFSQ_param = dynamic_cast<raft::spatial::knn::IVFSQParam *>(params);
       approx_knn_ivfsq_build_index(index, IVFSQ_param, metric, n, D);
     } else {
       ASSERT(index->index, "KNN index could not be initialized");
@@ -174,7 +174,7 @@ void approx_knn_build_index(raft::handle_t &handle, ML::knnIndex *index,
 
 template <typename IntType = int>
 void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, ML::knnIndex *index, IntType k,
+                       int64_t *indices, raft::spatial::knn::knnIndex *index, IntType k,
                        float *query_array, IntType n) {
   // perform preprocessing
   std::unique_ptr<MetricProcessor<float>> query_metric_processor =

From faac26b9d6c8331ec0b735daeb5ecc4076a72120 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:07:52 -0400
Subject: [PATCH 04/18] Linking in knn.hpp

---
 cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index 4494f0e81f..dc12a8d37a 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -17,6 +17,8 @@
 
 #pragma once
 
+#include "../knn.hpp"
+
 #include <raft/cudart_utils.h>
 #include <raft/cuda_utils.cuh>
 

From 8db42804764f64d5edfd8f8e05f364145d48e175 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:10:55 -0400
Subject: [PATCH 05/18] Removing unused absolute namespaces

---
 .../raft/spatial/knn/detail/ivf_pq_ann.cuh    | 38 +++++++++----------
 1 file changed, 19 insertions(+), 19 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index dc12a8d37a..ba14a808d4 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -56,19 +56,19 @@ namespace knn {
 namespace detail {
 
 inline faiss::ScalarQuantizer::QuantizerType build_faiss_qtype(
-  raft::spatial::knn::QuantizerType qtype) {
+  QuantizerType qtype) {
   switch (qtype) {
-    case raft::spatial::knn::QuantizerType::QT_8bit:
+    case QuantizerType::QT_8bit:
       return faiss::ScalarQuantizer::QuantizerType::QT_8bit;
-    case raft::spatial::knn::QuantizerType::QT_8bit_uniform:
+    case QuantizerType::QT_8bit_uniform:
       return faiss::ScalarQuantizer::QuantizerType::QT_8bit_uniform;
-    case raft::spatial::knn::QuantizerType::QT_4bit_uniform:
+    case QuantizerType::QT_4bit_uniform:
       return faiss::ScalarQuantizer::QuantizerType::QT_4bit_uniform;
-    case raft::spatial::knn::QuantizerType::QT_fp16:
+    case QuantizerType::QT_fp16:
       return faiss::ScalarQuantizer::QuantizerType::QT_fp16;
-    case raft::spatial::knn::QuantizerType::QT_8bit_direct:
+    case QuantizerType::QT_8bit_direct:
       return faiss::ScalarQuantizer::QuantizerType::QT_8bit_direct;
-    case raft::spatial::knn::QuantizerType::QT_6bit:
+    case QuantizerType::QT_6bit:
       return faiss::ScalarQuantizer::QuantizerType::QT_6bit;
     default:
       return (faiss::ScalarQuantizer::QuantizerType)qtype;
@@ -76,7 +76,7 @@ inline faiss::ScalarQuantizer::QuantizerType build_faiss_qtype(
 }
 
 template <typename IntType = int>
-void approx_knn_ivfflat_build_index(raft::spatial::knn::knnIndex *index, raft::spatial::knn::IVFParam *params,
+void approx_knn_ivfflat_build_index(knnIndex *index, IVFParam *params,
                                     raft::distance::DistanceType metric,
                                     IntType n, IntType D) {
   faiss::gpu::GpuIndexIVFFlatConfig config;
@@ -89,7 +89,7 @@ void approx_knn_ivfflat_build_index(raft::spatial::knn::knnIndex *index, raft::s
 }
 
 template <typename IntType = int>
-void approx_knn_ivfpq_build_index(raft::spatial::knn::knnIndex *index, raft::spatial::knn::IVFPQParam *params,
+void approx_knn_ivfpq_build_index(knnIndex *index, IVFPQParam *params,
                                   raft::distance::DistanceType metric,
                                   IntType n, IntType D) {
   faiss::gpu::GpuIndexIVFPQConfig config;
@@ -105,7 +105,7 @@ void approx_knn_ivfpq_build_index(raft::spatial::knn::knnIndex *index, raft::spa
 }
 
 template <typename IntType = int>
-void approx_knn_ivfsq_build_index(raft::spatial::knn::knnIndex *index, raft::spatial::knn::IVFSQParam *params,
+void approx_knn_ivfsq_build_index(knnIndex *index, IVFSQParam *params,
                                   raft::distance::DistanceType metric,
                                   IntType n, IntType D) {
   faiss::gpu::GpuIndexIVFScalarQuantizerConfig config;
@@ -122,8 +122,8 @@ void approx_knn_ivfsq_build_index(raft::spatial::knn::knnIndex *index, raft::spa
 }
 
 template <typename IntType = int>
-void approx_knn_build_index(raft::handle_t &handle, raft::spatial::knn::knnIndex *index,
-                            raft::spatial::knn::knnIndexParam *params,
+void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
+                            knnIndexParam *params,
                             raft::distance::DistanceType metric,
                             float metricArg, float *index_array, IntType n,
                             IntType D) {
@@ -148,8 +148,8 @@ void approx_knn_build_index(raft::handle_t &handle, raft::spatial::knn::knnIndex
 
   query_metric_processor->preprocess(index_array);
 
-  if (dynamic_cast<raft::spatial::knn::IVFFlatParam *>(params)) {
-    raft::spatial::knn::IVFFlatParam *IVFFlat_param = dynamic_cast<raft::spatial::knn::IVFFlatParam *>(params);
+  if (dynamic_cast<IVFFlatParam *>(params)) {
+    IVFFlatParam *IVFFlat_param = dynamic_cast<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(),
@@ -158,11 +158,11 @@ void approx_knn_build_index(raft::handle_t &handle, raft::spatial::knn::knnIndex
     index->index->train(n, h_index_array.data());
     index->index->add(n, h_index_array.data());
   } else {
-    if (dynamic_cast<raft::spatial::knn::IVFPQParam *>(params)) {
-      raft::spatial::knn::IVFPQParam *IVFPQ_param = dynamic_cast<raft::spatial::knn::IVFPQParam *>(params);
+    if (dynamic_cast<IVFPQParam *>(params)) {
+      IVFPQParam *IVFPQ_param = dynamic_cast<IVFPQParam *>(params);
       approx_knn_ivfpq_build_index(index, IVFPQ_param, metric, n, D);
-    } else if (dynamic_cast<raft::spatial::knn::IVFSQParam *>(params)) {
-      raft::spatial::knn::IVFSQParam *IVFSQ_param = dynamic_cast<raft::spatial::knn::IVFSQParam *>(params);
+    } else if (dynamic_cast<IVFSQParam *>(params)) {
+      IVFSQParam *IVFSQ_param = dynamic_cast<IVFSQParam *>(params);
       approx_knn_ivfsq_build_index(index, IVFSQ_param, metric, n, D);
     } else {
       ASSERT(index->index, "KNN index could not be initialized");
@@ -176,7 +176,7 @@ void approx_knn_build_index(raft::handle_t &handle, raft::spatial::knn::knnIndex
 
 template <typename IntType = int>
 void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, raft::spatial::knn::knnIndex *index, IntType k,
+                       int64_t *indices, knnIndex *index, IntType k,
                        float *query_array, IntType n) {
   // perform preprocessing
   std::unique_ptr<MetricProcessor<float>> query_metric_processor =

From 2673fb72718473baf8cfc3ac2e2ca40f71738b47 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:16:47 -0400
Subject: [PATCH 06/18] Calling functions

---
 cpp/include/raft/spatial/knn/knn.hpp | 22 +++++++++++++++-------
 1 file changed, 15 insertions(+), 7 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index 7049deb706..026c5e2173 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -44,7 +44,7 @@ struct knnIndex {
   }
 };
 
-typedef enum {
+enum QuantizerType {
   QT_8bit,
   QT_4bit,
   QT_8bit_uniform,
@@ -52,7 +52,7 @@ typedef enum {
   QT_fp16,
   QT_8bit_direct,
   QT_6bit
-} QuantizerType;
+};
 
 struct knnIndexParam {
   virtual ~knnIndexParam() {}
@@ -90,11 +90,14 @@ struct IVFSQParam : IVFParam {
  * @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, knnIndex *index,
+                            knnIndexParam *params,
                             raft::distance::DistanceType metric,
                             float metricArg, float *index_array, IntType n,
-                            IntType D);
+                            IntType D) {
+  detail::approx_knn_build_index(handle, index, params, metric,
+                                 metricArg, index_array, n, D);
+}
 
 /**
  * @brief Flat C++ API function to perform an approximate nearest neighbors
@@ -110,8 +113,13 @@ 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, IntType k,
-                       float *query_array, IntType n);
+                       int64_t *indices, knnIndex *index, IntType k,
+                       float *query_array, IntType n) {
+
+  detail::approx_knn_search(handle, distances, indices, index, k,
+                           query_array, n);
+
+}
 
 /**
  * @brief Flat C++ API function to perform a brute force knn on

From 30f389ef835b6d7bd68ea121f5c8c65c155a1dbd Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:19:36 -0400
Subject: [PATCH 07/18] the common

---
 cpp/include/raft/spatial/knn/common.h         | 81 +++++++++++++++++++
 .../raft/spatial/knn/detail/ivf_pq_ann.cuh    |  2 +-
 cpp/include/raft/spatial/knn/knn.hpp          | 46 +----------
 3 files changed, 83 insertions(+), 46 deletions(-)
 create mode 100644 cpp/include/raft/spatial/knn/common.h

diff --git a/cpp/include/raft/spatial/knn/common.h b/cpp/include/raft/spatial/knn/common.h
new file mode 100644
index 0000000000..d21cb4c80f
--- /dev/null
+++ b/cpp/include/raft/spatial/knn/common.h
@@ -0,0 +1,81 @@
+/*
+ * Copyright (c) 2020-2021, 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 "detail/brute_force_knn.cuh"
+#include "detail/ivf_pq_ann.cuh"
+
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/GpuIndex.h>
+
+#include <raft/mr/device/allocator.hpp>
+#include <raft/mr/device/buffer.hpp>
+
+namespace raft {
+namespace spatial {
+namespace knn {
+
+using deviceAllocator = raft::mr::device::allocator;
+
+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;
+  }
+};
+
+enum QuantizerType : unsigned int {
+  QT_8bit,
+  QT_4bit,
+  QT_8bit_uniform,
+  QT_4bit_uniform,
+  QT_fp16,
+  QT_8bit_direct,
+  QT_6bit
+};
+
+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;
+};
+
+};
+};
+};
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index ba14a808d4..36a98ce151 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -17,7 +17,7 @@
 
 #pragma once
 
-#include "../knn.hpp"
+#include "../common.h"
 
 #include <raft/cudart_utils.h>
 #include <raft/cuda_utils.cuh>
diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index 026c5e2173..7eed6a81c9 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -16,6 +16,7 @@
 
 #pragma once
 
+#include "common.h"
 #include "detail/brute_force_knn.cuh"
 #include "detail/ivf_pq_ann.cuh"
 
@@ -31,51 +32,6 @@ namespace knn {
 
 using deviceAllocator = raft::mr::device::allocator;
 
-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;
-  }
-};
-
-enum QuantizerType {
-  QT_8bit,
-  QT_4bit,
-  QT_8bit_uniform,
-  QT_4bit_uniform,
-  QT_fp16,
-  QT_8bit_direct,
-  QT_6bit
-};
-
-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 build an approximate nearest neighbors index

From e0e4b9b27898e35f38078d74769db7492eade199 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:31:35 -0400
Subject: [PATCH 08/18] Stuff

---
 cpp/include/raft/spatial/knn/common.h         | 20 +++++++++----------
 .../raft/spatial/knn/detail/ivf_pq_ann.cuh    |  2 +-
 cpp/include/raft/spatial/knn/knn.hpp          | 14 +++++++------
 3 files changed, 18 insertions(+), 18 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/common.h b/cpp/include/raft/spatial/knn/common.h
index d21cb4c80f..02f0ca9d28 100644
--- a/cpp/include/raft/spatial/knn/common.h
+++ b/cpp/include/raft/spatial/knn/common.h
@@ -29,8 +29,6 @@ namespace raft {
 namespace spatial {
 namespace knn {
 
-using deviceAllocator = raft::mr::device::allocator;
-
 struct knnIndex {
   faiss::gpu::GpuIndex *index;
   raft::distance::DistanceType metric;
@@ -42,7 +40,7 @@ struct knnIndex {
     delete index;
     delete gpu_res;
   }
-};
+}
 
 enum QuantizerType : unsigned int {
   QT_8bit,
@@ -52,16 +50,16 @@ enum QuantizerType : unsigned int {
   QT_fp16,
   QT_8bit_direct,
   QT_6bit
-};
+}
 
 struct knnIndexParam {
   virtual ~knnIndexParam() {}
-};
+}
 
 struct IVFParam : knnIndexParam {
   int nlist;
   int nprobe;
-};
+}
 
 struct IVFFlatParam : IVFParam {};
 
@@ -69,13 +67,13 @@ struct IVFPQParam : IVFParam {
   int M;
   int n_bits;
   bool usePrecomputedTables;
-};
+}
 
 struct IVFSQParam : IVFParam {
   QuantizerType qtype;
   bool encodeResidual;
-};
+}
 
-};
-};
-};
+}  // namespace knn
+}  // namespace spatial
+}  // namespace raft
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index 36a98ce151..436441993c 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -17,7 +17,7 @@
 
 #pragma once
 
-#include "../common.h"
+#include <raft/spatial/knn/common.h>
 
 #include <raft/cudart_utils.h>
 #include <raft/cuda_utils.cuh>
diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index 7eed6a81c9..8c35eb06d2 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -46,11 +46,12 @@ using deviceAllocator = raft::mr::device::allocator;
  * @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, knnIndex *index,
+template<typename value_idx = int>
+inline void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
                             knnIndexParam *params,
                             raft::distance::DistanceType metric,
-                            float metricArg, float *index_array, IntType n,
-                            IntType D) {
+                            float metricArg, float *index_array, value_idx n,
+                                   value_idx D) {
   detail::approx_knn_build_index(handle, index, params, metric,
                                  metricArg, index_array, n, D);
 }
@@ -68,9 +69,10 @@ void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
  * @param[in] query_array the query to perform a search with
  * @param[in] n number of rows in the query array
  */
-void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, knnIndex *index, IntType k,
-                       float *query_array, IntType n) {
+template<typename value_idx = int>
+inline void approx_knn_search(raft::handle_t &handle, float *distances,
+                       int64_t *indices, knnIndex *index, value_idx k,
+                       float *query_array, value_idx n) {
 
   detail::approx_knn_search(handle, distances, indices, index, k,
                            query_array, n);

From 0196db8c3b12469d44109959c4d53724d028658e Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:33:59 -0400
Subject: [PATCH 09/18] yup

---
 cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index 436441993c..9df0432f1b 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2021, NVIDIA CORPORATION.
+ * Copyright (c) 2021, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -206,7 +206,7 @@ void approx_knn_search(raft::handle_t &handle, float *distances,
 }
 
 
-};
-};
-};
-};
\ No newline at end of file
+}
+}
+}
+}
\ No newline at end of file

From 428928db535788a0ff4930b05116867f2a4ebdbd Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:35:52 -0400
Subject: [PATCH 10/18] It

---
 cpp/include/raft/spatial/knn/common.h | 6 ------
 1 file changed, 6 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/common.h b/cpp/include/raft/spatial/knn/common.h
index 02f0ca9d28..6378b4967a 100644
--- a/cpp/include/raft/spatial/knn/common.h
+++ b/cpp/include/raft/spatial/knn/common.h
@@ -16,15 +16,9 @@
 
 #pragma once
 
-#include "detail/brute_force_knn.cuh"
-#include "detail/ivf_pq_ann.cuh"
-
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/GpuIndex.h>
 
-#include <raft/mr/device/allocator.hpp>
-#include <raft/mr/device/buffer.hpp>
-
 namespace raft {
 namespace spatial {
 namespace knn {

From 17d0543391ecff69a6f5849b78deef906bcef66e Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:36:40 -0400
Subject: [PATCH 11/18] The thing

---
 cpp/include/raft/spatial/knn/common.h | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/common.h b/cpp/include/raft/spatial/knn/common.h
index 6378b4967a..b9bf934167 100644
--- a/cpp/include/raft/spatial/knn/common.h
+++ b/cpp/include/raft/spatial/knn/common.h
@@ -34,7 +34,7 @@ struct knnIndex {
     delete index;
     delete gpu_res;
   }
-}
+};
 
 enum QuantizerType : unsigned int {
   QT_8bit,
@@ -44,16 +44,16 @@ enum QuantizerType : unsigned int {
   QT_fp16,
   QT_8bit_direct,
   QT_6bit
-}
+};
 
 struct knnIndexParam {
   virtual ~knnIndexParam() {}
-}
+};
 
 struct IVFParam : knnIndexParam {
   int nlist;
   int nprobe;
-}
+};
 
 struct IVFFlatParam : IVFParam {};
 
@@ -61,13 +61,13 @@ struct IVFPQParam : IVFParam {
   int M;
   int n_bits;
   bool usePrecomputedTables;
-}
+};
 
 struct IVFSQParam : IVFParam {
   QuantizerType qtype;
   bool encodeResidual;
-}
+};
 
-}  // namespace knn
-}  // namespace spatial
-}  // namespace raft
+};  // namespace knn
+};  // namespace spatial
+};  // namespace raft

From af21a27807854bbfc60f82cdc8bbbf3da8091f2c Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:42:03 -0400
Subject: [PATCH 12/18] Checking in

---
 cpp/include/raft/spatial/knn/knn.hpp | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index 8c35eb06d2..b035a6368e 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -33,6 +33,15 @@ namespace knn {
 using deviceAllocator = raft::mr::device::allocator;
 
 
+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) {
+  detail::knn_merge_parts(inK, inV, outK, outV, n_samples,
+                          n_parts, k, stream, translations);
+}
+
 /**
  * @brief Flat C++ API function to build an approximate nearest neighbors index
  * from an index array and a set of parameters.

From 71f3cf50aa1cfee5150bc931c81364c874a13781 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:50:22 -0400
Subject: [PATCH 13/18] Checking in

---
 cpp/include/raft/spatial/knn/common.h | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/cpp/include/raft/spatial/knn/common.h b/cpp/include/raft/spatial/knn/common.h
index b9bf934167..2f606c43ec 100644
--- a/cpp/include/raft/spatial/knn/common.h
+++ b/cpp/include/raft/spatial/knn/common.h
@@ -16,6 +16,8 @@
 
 #pragma once
 
+#include <raft/linalg/distance_type.h>
+
 #include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/GpuIndex.h>
 

From 446d25c54ca2155b2b2b825461da4c9c704c6c0b Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 16:55:12 -0400
Subject: [PATCH 14/18] Checking in

---
 cpp/include/raft/spatial/knn/ann.hpp          | 82 +++++++++++++++++++
 .../spatial/knn/{common.h => ann_common.h}    |  2 +-
 .../raft/spatial/knn/detail/ivf_pq_ann.cuh    |  2 +-
 cpp/include/raft/spatial/knn/knn.hpp          | 51 ------------
 4 files changed, 84 insertions(+), 53 deletions(-)
 create mode 100644 cpp/include/raft/spatial/knn/ann.hpp
 rename cpp/include/raft/spatial/knn/{common.h => ann_common.h} (100%)

diff --git a/cpp/include/raft/spatial/knn/ann.hpp b/cpp/include/raft/spatial/knn/ann.hpp
new file mode 100644
index 0000000000..71fdaaabd5
--- /dev/null
+++ b/cpp/include/raft/spatial/knn/ann.hpp
@@ -0,0 +1,82 @@
+/*
+ * Copyright (c) 2020-2021, 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 "ann_common.h"
+#include "detail/ivf_pq_ann.cuh"
+
+#include <faiss/gpu/StandardGpuResources.h>
+#include <faiss/gpu/GpuIndex.h>
+
+#include <raft/mr/device/allocator.hpp>
+#include <raft/mr/device/buffer.hpp>
+
+namespace raft {
+namespace spatial {
+namespace knn {
+
+using deviceAllocator = raft::mr::device::allocator;
+
+/**
+ * @brief Flat C++ API function to build an approximate nearest neighbors index
+ * from an index array and a set of parameters.
+ *
+ * @param[in] handle RAFT handle
+ * @param[out] index index to be built
+ * @param[in] params parametrization of the index to be built
+ * @param[in] metric distance metric to use. Euclidean (L2) is used by default
+ * @param[in] metricArg metric argument
+ * @param[in] index_array the index array to build the index with
+ * @param[in] n number of rows in the index array
+ * @param[in] D the dimensionality of the index array
+ */
+template<typename value_idx = int>
+inline void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
+                                   knnIndexParam *params,
+                                   raft::distance::DistanceType metric,
+                                   float metricArg, float *index_array, value_idx n,
+                                   value_idx D) {
+  detail::approx_knn_build_index(handle, index, params, metric,
+                                 metricArg, index_array, n, D);
+}
+
+/**
+ * @brief Flat C++ API function to perform an approximate nearest neighbors
+ * search from previously built index and a query array
+ *
+ * @param[in] handle RAFT handle
+ * @param[out] distances distances of the nearest neighbors toward
+ *                       their query point
+ * @param[out] indices indices of the nearest neighbors
+ * @param[in] index index to perform a search with
+ * @param[in] k the number of nearest neighbors to search for
+ * @param[in] query_array the query to perform a search with
+ * @param[in] n number of rows in the query array
+ */
+template<typename value_idx = int>
+inline void approx_knn_search(raft::handle_t &handle, float *distances,
+                              int64_t *indices, knnIndex *index, value_idx k,
+                              float *query_array, value_idx n) {
+
+  detail::approx_knn_search(handle, distances, indices, index, k,
+                            query_array, n);
+
+}
+
+}  // namespace knn
+}  // namespace spatial
+}  // namespace raft
diff --git a/cpp/include/raft/spatial/knn/common.h b/cpp/include/raft/spatial/knn/ann_common.h
similarity index 100%
rename from cpp/include/raft/spatial/knn/common.h
rename to cpp/include/raft/spatial/knn/ann_common.h
index 2f606c43ec..6a6c7751c2 100644
--- a/cpp/include/raft/spatial/knn/common.h
+++ b/cpp/include/raft/spatial/knn/ann_common.h
@@ -18,8 +18,8 @@
 
 #include <raft/linalg/distance_type.h>
 
-#include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/GpuIndex.h>
+#include <faiss/gpu/StandardGpuResources.h>
 
 namespace raft {
 namespace spatial {
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index 9df0432f1b..1af1521d6f 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -17,7 +17,7 @@
 
 #pragma once
 
-#include <raft/spatial/knn/common.h>
+#include "../ann_common.h"
 
 #include <raft/cudart_utils.h>
 #include <raft/cuda_utils.cuh>
diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index b035a6368e..024af0607e 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -16,12 +16,7 @@
 
 #pragma once
 
-#include "common.h"
 #include "detail/brute_force_knn.cuh"
-#include "detail/ivf_pq_ann.cuh"
-
-#include <faiss/gpu/StandardGpuResources.h>
-#include <faiss/gpu/GpuIndex.h>
 
 #include <raft/mr/device/allocator.hpp>
 #include <raft/mr/device/buffer.hpp>
@@ -42,52 +37,6 @@ inline void knn_merge_parts(value_t *inK, value_idx *inV, value_t *outK,
                           n_parts, k, stream, translations);
 }
 
-/**
- * @brief Flat C++ API function to build an approximate nearest neighbors index
- * from an index array and a set of parameters.
- *
- * @param[in] handle RAFT handle
- * @param[out] index index to be built
- * @param[in] params parametrization of the index to be built
- * @param[in] metric distance metric to use. Euclidean (L2) is used by default
- * @param[in] metricArg metric argument
- * @param[in] index_array the index array to build the index with
- * @param[in] n number of rows in the index array
- * @param[in] D the dimensionality of the index array
- */
-template<typename value_idx = int>
-inline void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
-                            knnIndexParam *params,
-                            raft::distance::DistanceType metric,
-                            float metricArg, float *index_array, value_idx n,
-                                   value_idx D) {
-  detail::approx_knn_build_index(handle, index, params, metric,
-                                 metricArg, index_array, n, D);
-}
-
-/**
- * @brief Flat C++ API function to perform an approximate nearest neighbors
- * search from previously built index and a query array
- *
- * @param[in] handle RAFT handle
- * @param[out] distances distances of the nearest neighbors toward
- *                       their query point
- * @param[out] indices indices of the nearest neighbors
- * @param[in] index index to perform a search with
- * @param[in] k the number of nearest neighbors to search for
- * @param[in] query_array the query to perform a search with
- * @param[in] n number of rows in the query array
- */
-template<typename value_idx = int>
-inline void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, knnIndex *index, value_idx k,
-                       float *query_array, value_idx n) {
-
-  detail::approx_knn_search(handle, distances, indices, index, k,
-                           query_array, n);
-
-}
-
 /**
  * @brief Flat C++ API function to perform a brute force knn on
  * a series of input arrays and combine the results into a single

From fc8d48d2022139c7f80cd12a301dd59562bfc042 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Tue, 8 Jun 2021 17:50:42 -0400
Subject: [PATCH 15/18] Fixing style

---
 cpp/include/raft/spatial/knn/ann.hpp          | 20 +++++++++----------
 .../raft/spatial/knn/detail/ivf_pq_ann.cuh    | 15 ++++++--------
 cpp/include/raft/spatial/knn/knn.hpp          |  5 ++---
 3 files changed, 17 insertions(+), 23 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/ann.hpp b/cpp/include/raft/spatial/knn/ann.hpp
index 71fdaaabd5..6b0b4a64e3 100644
--- a/cpp/include/raft/spatial/knn/ann.hpp
+++ b/cpp/include/raft/spatial/knn/ann.hpp
@@ -19,8 +19,8 @@
 #include "ann_common.h"
 #include "detail/ivf_pq_ann.cuh"
 
-#include <faiss/gpu/StandardGpuResources.h>
 #include <faiss/gpu/GpuIndex.h>
+#include <faiss/gpu/StandardGpuResources.h>
 
 #include <raft/mr/device/allocator.hpp>
 #include <raft/mr/device/buffer.hpp>
@@ -44,14 +44,14 @@ using deviceAllocator = raft::mr::device::allocator;
  * @param[in] n number of rows in the index array
  * @param[in] D the dimensionality of the index array
  */
-template<typename value_idx = int>
+template <typename value_idx = int>
 inline void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
                                    knnIndexParam *params,
                                    raft::distance::DistanceType metric,
-                                   float metricArg, float *index_array, value_idx n,
-                                   value_idx D) {
-  detail::approx_knn_build_index(handle, index, params, metric,
-                                 metricArg, index_array, n, D);
+                                   float metricArg, float *index_array,
+                                   value_idx n, value_idx D) {
+  detail::approx_knn_build_index(handle, index, params, metric, metricArg,
+                                 index_array, n, D);
 }
 
 /**
@@ -67,14 +67,12 @@ inline void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
  * @param[in] query_array the query to perform a search with
  * @param[in] n number of rows in the query array
  */
-template<typename value_idx = int>
+template <typename value_idx = int>
 inline void approx_knn_search(raft::handle_t &handle, float *distances,
                               int64_t *indices, knnIndex *index, value_idx k,
                               float *query_array, value_idx n) {
-
-  detail::approx_knn_search(handle, distances, indices, index, k,
-                            query_array, n);
-
+  detail::approx_knn_search(handle, distances, indices, index, k, query_array,
+                            n);
 }
 
 }  // namespace knn
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
index 1af1521d6f..042936a7e8 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
@@ -14,7 +14,6 @@
  * limitations under the License.
  */
 
-
 #pragma once
 
 #include "../ann_common.h"
@@ -49,7 +48,6 @@
 #include <iostream>
 #include <set>
 
-
 namespace raft {
 namespace spatial {
 namespace knn {
@@ -199,14 +197,13 @@ void approx_knn_search(raft::handle_t &handle, float *distances,
       p = 1.0 / index->metricArg;
     raft::linalg::unaryOp<float>(
       distances, distances, n * k,
-    [p] __device__(float input) { return powf(input, p); },
-    handle.get_stream());
+      [p] __device__(float input) { return powf(input, p); },
+      handle.get_stream());
   }
   query_metric_processor->postprocess(distances);
 }
 
-
-}
-}
-}
-}
\ No newline at end of file
+}  // namespace detail
+}  // namespace knn
+}  // namespace spatial
+}  // namespace raft
\ No newline at end of file
diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index 024af0607e..ebbb5061cd 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -27,14 +27,13 @@ namespace knn {
 
 using deviceAllocator = raft::mr::device::allocator;
 
-
 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) {
-  detail::knn_merge_parts(inK, inV, outK, outV, n_samples,
-                          n_parts, k, stream, translations);
+  detail::knn_merge_parts(inK, inV, outK, outV, n_samples, n_parts, k, stream,
+                          translations);
 }
 
 /**

From fbb665c8bbb90849e64bf4b59b536500f5a81b9a Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Wed, 9 Jun 2021 13:37:23 -0400
Subject: [PATCH 16/18] Renaming headers to reflect category and faiss.

---
 cpp/include/raft/spatial/knn/ann.hpp                      | 8 +++++---
 .../detail/{ivf_pq_ann.cuh => ann_quantized_faiss.cuh}    | 6 +++---
 .../{brute_force_knn.cuh => knn_brute_force_faiss.cuh}    | 0
 cpp/include/raft/spatial/knn/knn.hpp                      | 2 +-
 4 files changed, 9 insertions(+), 7 deletions(-)
 rename cpp/include/raft/spatial/knn/detail/{ivf_pq_ann.cuh => ann_quantized_faiss.cuh} (96%)
 rename cpp/include/raft/spatial/knn/detail/{brute_force_knn.cuh => knn_brute_force_faiss.cuh} (100%)

diff --git a/cpp/include/raft/spatial/knn/ann.hpp b/cpp/include/raft/spatial/knn/ann.hpp
index 6b0b4a64e3..77d7831b4a 100644
--- a/cpp/include/raft/spatial/knn/ann.hpp
+++ b/cpp/include/raft/spatial/knn/ann.hpp
@@ -17,7 +17,7 @@
 #pragma once
 
 #include "ann_common.h"
-#include "detail/ivf_pq_ann.cuh"
+#include "detail/ann_quantized_faiss.cuh"
 
 #include <faiss/gpu/GpuIndex.h>
 #include <faiss/gpu/StandardGpuResources.h>
@@ -45,7 +45,8 @@ using deviceAllocator = raft::mr::device::allocator;
  * @param[in] D the dimensionality of the index array
  */
 template <typename value_idx = int>
-inline void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
+inline void approx_knn_build_index(raft::handle_t &handle,
+                                   raft::spatial::knn::knnIndex *index,
                                    knnIndexParam *params,
                                    raft::distance::DistanceType metric,
                                    float metricArg, float *index_array,
@@ -69,7 +70,8 @@ inline void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
  */
 template <typename value_idx = int>
 inline void approx_knn_search(raft::handle_t &handle, float *distances,
-                              int64_t *indices, knnIndex *index, value_idx k,
+                              int64_t *indices,
+                              raft::spatial::knn::knnIndex *index, value_idx k,
                               float *query_array, value_idx n) {
   detail::approx_knn_search(handle, distances, indices, index, k, query_array,
                             n);
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh b/cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh
similarity index 96%
rename from cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
rename to cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh
index 042936a7e8..8ba8cbf9c8 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_pq_ann.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh
@@ -120,8 +120,8 @@ void approx_knn_ivfsq_build_index(knnIndex *index, IVFSQParam *params,
 }
 
 template <typename IntType = int>
-void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
-                            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, IntType n,
                             IntType D) {
@@ -174,7 +174,7 @@ void approx_knn_build_index(raft::handle_t &handle, knnIndex *index,
 
 template <typename IntType = int>
 void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, knnIndex *index, IntType k,
+                       int64_t *indices, raft::spatial::knn::knnIndex *index, IntType k,
                        float *query_array, IntType n) {
   // perform preprocessing
   std::unique_ptr<MetricProcessor<float>> query_metric_processor =
diff --git a/cpp/include/raft/spatial/knn/detail/brute_force_knn.cuh b/cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh
similarity index 100%
rename from cpp/include/raft/spatial/knn/detail/brute_force_knn.cuh
rename to cpp/include/raft/spatial/knn/detail/knn_brute_force_faiss.cuh
diff --git a/cpp/include/raft/spatial/knn/knn.hpp b/cpp/include/raft/spatial/knn/knn.hpp
index ebbb5061cd..a3a1972c13 100644
--- a/cpp/include/raft/spatial/knn/knn.hpp
+++ b/cpp/include/raft/spatial/knn/knn.hpp
@@ -16,7 +16,7 @@
 
 #pragma once
 
-#include "detail/brute_force_knn.cuh"
+#include "detail/knn_brute_force_faiss.cuh"
 
 #include <raft/mr/device/allocator.hpp>
 #include <raft/mr/device/buffer.hpp>

From 4b478e03bb28b98d03a10a1def3d911836d96253 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Wed, 9 Jun 2021 14:42:13 -0400
Subject: [PATCH 17/18] Fixing style

---
 .../raft/spatial/knn/detail/ann_quantized_faiss.cuh        | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh b/cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh
index 8ba8cbf9c8..bb37fd03db 100644
--- a/cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ann_quantized_faiss.cuh
@@ -120,7 +120,8 @@ void approx_knn_ivfsq_build_index(knnIndex *index, IVFSQParam *params,
 }
 
 template <typename IntType = int>
-void approx_knn_build_index(raft::handle_t &handle, raft::spatial::knn::knnIndex *index,
+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, IntType n,
@@ -174,8 +175,8 @@ void approx_knn_build_index(raft::handle_t &handle, raft::spatial::knn::knnIndex
 
 template <typename IntType = int>
 void approx_knn_search(raft::handle_t &handle, float *distances,
-                       int64_t *indices, raft::spatial::knn::knnIndex *index, IntType k,
-                       float *query_array, IntType n) {
+                       int64_t *indices, raft::spatial::knn::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,

From 67cb4a2bb90510d15366768a0282b64ca78c8e47 Mon Sep 17 00:00:00 2001
From: "Corey J. Nolet" <cjnolet@gmail.com>
Date: Wed, 9 Jun 2021 15:57:40 -0400
Subject: [PATCH 18/18] Using public API for merge_parts

---
 cpp/include/raft/sparse/selection/knn.cuh | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/cpp/include/raft/sparse/selection/knn.cuh b/cpp/include/raft/sparse/selection/knn.cuh
index 93b2a83bbd..e327386d13 100644
--- a/cpp/include/raft/sparse/selection/knn.cuh
+++ b/cpp/include/raft/sparse/selection/knn.cuh
@@ -33,7 +33,7 @@
 #include <raft/sparse/csr.cuh>
 #include <raft/sparse/distance/distance.cuh>
 #include <raft/sparse/selection/selection.cuh>
-#include <raft/spatial/knn/detail/brute_force_knn.cuh>
+
 #include <raft/spatial/knn/knn.hpp>
 
 #include <raft/cudart_utils.h>
@@ -310,7 +310,7 @@ class sparse_knn_t {
                         stream);
 
     // combine merge buffers only if there's more than 1 partition to combine
-    raft::spatial::knn::detail::knn_merge_parts(
+    raft::spatial::knn::knn_merge_parts(
       merge_buffer_dists, merge_buffer_indices, out_dists, out_indices,
       query_batcher.batch_rows(), 2, k, stream, trans.data());
   }