Merge pull request #1234 from rapidsai/branch-23.02

Forward-merge branch-23.02 to branch-23.04
rapidsai · Feb 3, 2023 · f8d3f5b · f8d3f5b
2 parents d15fbe5 + 4ad7dab
commit f8d3f5b
Show file tree

Hide file tree

Showing 42 changed files with 565 additions and 2,201 deletions.
diff --git a/conda/recipes/libraft/build_libraft_distance.sh b/conda/recipes/libraft/build_libraft_distance.sh
@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-PARALLEL_LEVEL=8 ./build.sh libraft -v --allgpuarch --compile-dist --no-nvtx
+./build.sh libraft -v --allgpuarch --compile-dist --no-nvtx
diff --git a/conda/recipes/libraft/build_libraft_nn.sh b/conda/recipes/libraft/build_libraft_nn.sh
@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-PARALLEL_LEVEL=8 ./build.sh libraft -v --allgpuarch --compile-nn --no-nvtx
+./build.sh libraft -v --allgpuarch --compile-nn --no-nvtx
diff --git a/conda/recipes/libraft/build_libraft_tests.sh b/conda/recipes/libraft/build_libraft_tests.sh
@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-PARALLEL_LEVEL=8 ./build.sh tests bench -v --allgpuarch --no-nvtx
+./build.sh tests bench -v --allgpuarch --no-nvtx
 cmake --install cpp/build --component testing
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
@@ -305,7 +305,6 @@ if(RAFT_COMPILE_DIST_LIBRARY)
     src/distance/neighbors/specializations/refine_h_uint64_t_float.cu
     src/distance/neighbors/specializations/refine_h_uint64_t_int8_t.cu
     src/distance/neighbors/specializations/refine_h_uint64_t_uint8_t.cu
-    src/distance/neighbors/ivfpq_search.cu
     src/distance/cluster/kmeans_fit_float.cu
     src/distance/cluster/kmeans_fit_double.cu
     src/distance/distance/specializations/detail/canberra.cu
@@ -356,6 +355,11 @@ if(RAFT_COMPILE_DIST_LIBRARY)
     src/distance/distance/specializations/fused_l2_nn_float_int.cu
     src/distance/distance/specializations/fused_l2_nn_float_int64.cu
     src/distance/neighbors/ivfpq_build.cu
+    src/distance/neighbors/ivfpq_deserialize.cu
+    src/distance/neighbors/ivfpq_serialize.cu
+    src/distance/neighbors/ivfpq_search_float_uint64_t.cu
+    src/distance/neighbors/ivfpq_search_int8_t_uint64_t.cu
+    src/distance/neighbors/ivfpq_search_uint8_t_uint64_t.cu
     src/distance/neighbors/specializations/detail/ivfpq_compute_similarity_float_fast.cu
     src/distance/neighbors/specializations/detail/ivfpq_compute_similarity_float_no_basediff.cu
     src/distance/neighbors/specializations/detail/ivfpq_compute_similarity_float_no_smem_lut.cu
@@ -427,16 +431,21 @@ set_target_properties(raft_nn PROPERTIES EXPORT_NAME nn)
 if(RAFT_COMPILE_NN_LIBRARY)
   add_library(
     raft_nn_lib
-    src/nn/specializations/ball_cover.cu
     src/nn/specializations/detail/ball_cover_lowdim_pass_one_2d.cu
     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/ball_cover_all_knn_query.cu
+    src/nn/specializations/ball_cover_build_index.cu
+    src/nn/specializations/ball_cover_knn_query.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
     src/nn/specializations/fused_l2_knn_int_float_false.cu
-    src/nn/specializations/knn.cu
+    src/nn/specializations/brute_force_knn_long_float_int.cu
+    src/nn/specializations/brute_force_knn_long_float_uint.cu
+    src/nn/specializations/brute_force_knn_uint32_t_float_int.cu
+    src/nn/specializations/brute_force_knn_uint32_t_float_uint.cu
   )
   set_target_properties(
     raft_nn_lib

diff --git a/cpp/bench/CMakeLists.txt b/cpp/bench/CMakeLists.txt
@@ -82,7 +82,6 @@ if(BUILD_BENCH)
     bench/distance/distance_l1.cu
     bench/distance/distance_unexp_l2.cu
     bench/distance/fused_l2_nn.cu
-    bench/distance/masked_nn.cu
     bench/distance/kernels.cu
     bench/main.cpp
     OPTIONAL

diff --git a/cpp/bench/distance/distance_common.cuh b/cpp/bench/distance/distance_common.cuh
@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -24,14 +24,14 @@
 
 namespace raft::bench::distance {
 
-struct distance_params {
+struct distance_inputs {
   int m, n, k;
   bool isRowMajor;
-};  // struct distance_params
+};  // struct distance_inputs
 
 template <typename T, raft::distance::DistanceType DType>
 struct distance : public fixture {
-  distance(const distance_params& p)
+  distance(const distance_inputs& p)
     : params(p),
       x(p.m * p.k, stream),
       y(p.n * p.k, stream),
@@ -63,13 +63,13 @@ struct distance : public fixture {
   }
 
  private:
-  distance_params params;
+  distance_inputs params;
   rmm::device_uvector<T> x, y, out;
   rmm::device_uvector<char> workspace;
   size_t worksize;
 };  // struct Distance
 
-const std::vector<distance_params> dist_input_vecs{
+const std::vector<distance_inputs> dist_input_vecs{
   {32, 16384, 16384, true},    {64, 16384, 16384, true},    {128, 16384, 16384, true},
   {256, 16384, 16384, true},   {512, 16384, 16384, true},   {1024, 16384, 16384, true},
   {16384, 32, 16384, true},    {16384, 64, 16384, true},    {16384, 128, 16384, true},

diff --git a/cpp/bench/distance/masked_nn.cu b/cpp/bench/distance/masked_nn.cu
diff --git a/cpp/bench/neighbors/knn.cuh b/cpp/bench/neighbors/knn.cuh
@@ -33,8 +33,6 @@
 #if defined RAFT_DISTANCE_COMPILED
 #include <raft/cluster/specializations.cuh>
 #include <raft/neighbors/specializations.cuh>
-#else
-#pragma message("NN / Distance specializations are not enabled; expect very long building times.")
 #endif
 #endif