Integrate KNN implementation: ivf-pq (#789)

Integrate a new KNN implementation.

Authors:
  - Artem M. Chirkin (https://github.com/achirkin)

Approvers:
  - Tamas Bela Feher (https://github.com/tfeher)
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #789

cpp/CMakeLists.txt

@@ -340,6 +340,21 @@ if(RAFT_COMPILE_NN_LIBRARY)
           src/nn/specializations/detail/ball_cover_lowdim_pass_two_2d.cu
           src/nn/specializations/detail/ball_cover_lowdim_pass_one_3d.cu
           src/nn/specializations/detail/ball_cover_lowdim_pass_two_3d.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_float_fast.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_float_no_basediff.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_float_no_smem_lut.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_fast.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_basediff.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_fp8s_no_smem_lut.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_fast.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_basediff.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_fp8u_no_smem_lut.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_half_fast.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_half_no_basediff.cu
+          src/nn/specializations/detail/ivfpq_compute_similarity_half_no_smem_lut.cu
+          src/nn/specializations/detail/ivfpq_search_float_int64_t.cu
+          src/nn/specializations/detail/ivfpq_search_float_uint32_t.cu
+          src/nn/specializations/detail/ivfpq_search_float_uint64_t.cu
           src/nn/specializations/fused_l2_knn_long_float_true.cu
           src/nn/specializations/fused_l2_knn_long_float_false.cu
           src/nn/specializations/fused_l2_knn_int_float_true.cu

cpp/bench/CMakeLists.txt

@@ -33,7 +33,16 @@ add_executable(${RAFT_CPP_BENCH_TARGET}
   bench/random/rng.cu
   bench/sparse/convert_csr.cu
   bench/spatial/fused_l2_nn.cu
-  bench/spatial/knn.cu
+  bench/spatial/knn/brute_force_float_int64_t.cu
+  bench/spatial/knn/brute_force_float_uint32_t.cu
+  bench/spatial/knn/ivf_flat_float_int64_t.cu
+  bench/spatial/knn/ivf_flat_float_uint32_t.cu
+  bench/spatial/knn/ivf_flat_int8_t_int64_t.cu
+  bench/spatial/knn/ivf_flat_uint8_t_uint32_t.cu
+  bench/spatial/knn/ivf_pq_float_int64_t.cu
+  bench/spatial/knn/ivf_pq_float_uint32_t.cu
+  bench/spatial/knn/ivf_pq_int8_t_int64_t.cu
+  bench/spatial/knn/ivf_pq_uint8_t_uint32_t.cu
   bench/spatial/selection.cu
   bench/main.cpp
 )

cpp/bench/spatial/knn.cu → cpp/bench/spatial/knn.cuh

@@ -14,11 +14,14 @@
  * limitations under the License.
  */
 
+#pragma once
+
 #include <common/benchmark.hpp>
 
 #include <raft/random/rng.cuh>
 
 #include <raft/spatial/knn/ivf_flat.cuh>
+#include <raft/spatial/knn/ivf_pq.cuh>
 #include <raft/spatial/knn/knn.cuh>
 #if defined RAFT_NN_COMPILED
 #include <raft/spatial/knn/specializations.cuh>
@@ -45,16 +48,16 @@ struct params {
   size_t k;
 };
 
-auto operator<<(std::ostream& os, const params& p) -> std::ostream&
+inline auto operator<<(std::ostream& os, const params& p) -> std::ostream&
 {
   os << p.n_samples << "#" << p.n_dims << "#" << p.n_queries << "#" << p.k;
   return os;
 }
 
-enum class TransferStrategy { NO_COPY, COPY_PLAIN, COPY_PINNED, MAP_PINNED, MANAGED };
-enum class Scope { BUILD, SEARCH, BUILD_SEARCH };
+enum class TransferStrategy { NO_COPY, COPY_PLAIN, COPY_PINNED, MAP_PINNED, MANAGED };  // NOLINT
+enum class Scope { BUILD, SEARCH, BUILD_SEARCH };                                       // NOLINT
 
-auto operator<<(std::ostream& os, const TransferStrategy& ts) -> std::ostream&
+inline auto operator<<(std::ostream& os, const TransferStrategy& ts) -> std::ostream&
 {
   switch (ts) {
     case TransferStrategy::NO_COPY: os << "NO_COPY"; break;
@@ -67,7 +70,7 @@ auto operator<<(std::ostream& os, const TransferStrategy& ts) -> std::ostream&
   return os;
 }
 
-auto operator<<(std::ostream& os, const Scope& s) -> std::ostream&
+inline auto operator<<(std::ostream& os, const Scope& s) -> std::ostream&
 {
   switch (s) {
     case Scope::BUILD: os << "BUILD"; break;
@@ -156,6 +159,34 @@ struct ivf_flat_knn {
   }
 };
 
+template <typename ValT, typename IdxT>
+struct ivf_pq_knn {
+  using dist_t = float;
+
+  std::optional<const raft::spatial::knn::ivf_pq::index<IdxT>> index;
+  raft::spatial::knn::ivf_pq::index_params index_params;
+  raft::spatial::knn::ivf_pq::search_params search_params;
+  params ps;
+
+  ivf_pq_knn(const raft::handle_t& handle, const params& ps, const ValT* data) : ps(ps)
+  {
+    index_params.n_lists = 4096;
+    index_params.metric  = raft::distance::DistanceType::L2Expanded;
+    index.emplace(raft::spatial::knn::ivf_pq::build(
+      handle, index_params, data, IdxT(ps.n_samples), uint32_t(ps.n_dims)));
+  }
+
+  void search(const raft::handle_t& handle,
+              const ValT* search_items,
+              dist_t* out_dists,
+              IdxT* out_idxs)
+  {
+    search_params.n_probes = 20;
+    raft::spatial::knn::ivf_pq::search(
+      handle, search_params, *index, search_items, ps.n_queries, ps.k, out_idxs, out_dists);
+  }
+};
+
 template <typename ValT, typename IdxT>
 struct brute_force_knn {
   using dist_t = ValT;
@@ -217,7 +248,7 @@ struct knn : public fixture {
   }
 
   template <typename T>
-  void gen_data(raft::random::RngState& state,
+  void gen_data(raft::random::RngState& state,  // NOLINT
                 rmm::device_uvector<T>& vec,
                 size_t n,
                 rmm::cuda_stream_view stream)
@@ -338,15 +369,15 @@ struct knn : public fixture {
   rmm::device_uvector<IdxT> out_idxs_;
 };
 
-const std::vector<params> kInputs{
+inline const std::vector<params> kInputs{
   {2000000, 128, 1000, 32}, {10000000, 128, 1000, 32}, {10000, 8192, 1000, 32}};
 
-const std::vector<TransferStrategy> kAllStrategies{
+inline const std::vector<TransferStrategy> kAllStrategies{
   TransferStrategy::NO_COPY, TransferStrategy::MAP_PINNED, TransferStrategy::MANAGED};
-const std::vector<TransferStrategy> kNoCopyOnly{TransferStrategy::NO_COPY};
+inline const std::vector<TransferStrategy> kNoCopyOnly{TransferStrategy::NO_COPY};
 
-const std::vector<Scope> kScopeFull{Scope::BUILD_SEARCH};
-const std::vector<Scope> kAllScopes{Scope::BUILD_SEARCH, Scope::SEARCH, Scope::BUILD};
+inline const std::vector<Scope> kScopeFull{Scope::BUILD_SEARCH};
+inline const std::vector<Scope> kAllScopes{Scope::BUILD_SEARCH, Scope::SEARCH, Scope::BUILD};
 
 #define KNN_REGISTER(ValT, IdxT, ImplT, inputs, strats, scope)                   \
   namespace BENCHMARK_PRIVATE_NAME(knn)                                          \
@@ -355,14 +386,4 @@ const std::vector<Scope> kAllScopes{Scope::BUILD_SEARCH, Scope::SEARCH, Scope::B
     RAFT_BENCH_REGISTER(KNN, #ValT "/" #IdxT "/" #ImplT, inputs, strats, scope); \
   }
 
-KNN_REGISTER(float, int64_t, brute_force_knn, kInputs, kAllStrategies, kScopeFull);
-KNN_REGISTER(float, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-KNN_REGISTER(int8_t, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-KNN_REGISTER(uint8_t, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-
-KNN_REGISTER(float, uint32_t, brute_force_knn, kInputs, kNoCopyOnly, kScopeFull);
-KNN_REGISTER(float, uint32_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-KNN_REGISTER(int8_t, uint32_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-KNN_REGISTER(uint8_t, uint32_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
-
 }  // namespace raft::bench::spatial

cpp/bench/spatial/knn/brute_force_float_int64_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(float, int64_t, brute_force_knn, kInputs, kAllStrategies, kScopeFull);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/brute_force_float_uint32_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(float, uint32_t, brute_force_knn, kInputs, kAllStrategies, kScopeFull);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/ivf_flat_float_int64_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(float, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/ivf_flat_float_uint32_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(float, uint32_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/ivf_flat_int8_t_int64_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(int8_t, int64_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/ivf_flat_uint8_t_uint32_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(uint8_t, uint32_t, ivf_flat_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/ivf_pq_float_int64_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(float, int64_t, ivf_pq_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/ivf_pq_float_uint32_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(float, uint32_t, ivf_pq_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+}  // namespace raft::bench::spatial

cpp/bench/spatial/knn/ivf_pq_int8_t_int64_t.cu

@@ -0,0 +1,23 @@
+/*
+ * 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 "../knn.cuh"
+
+namespace raft::bench::spatial {
+
+KNN_REGISTER(int8_t, int64_t, ivf_pq_knn, kInputs, kNoCopyOnly, kAllScopes);
+
+}  // namespace raft::bench::spatial