Remove selection_faiss (#2027)

Remove the selection_faiss instantiations. Since #1985, we haven't been using the faiss select_k code and these aren't necessary anymore. This should lead to a 70MB improvement in libraft.so binary size. This also removes the raft::spatial::select_k code in favour of matrix:: select_k - the spatial version was marked deprecated, and didn't switch between the best selection algorithms for the input size. Authors: - Ben Frederickson (https://github.com/benfred) Approvers: - Tamas Bela Feher (https://github.com/tfeher) URL: #2027
rapidsai · Nov 30, 2023 · 4ba0139 · 4ba0139
1 parent ff60d3a
commit 4ba0139
Show file tree

Hide file tree

Showing 28 changed files with 25 additions and 1,495 deletions.
diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
@@ -372,16 +372,6 @@ if(RAFT_COMPILE_LIBRARY)
     src/neighbors/detail/refine_host_float_float.cpp
     src/neighbors/detail/refine_host_int8_t_float.cpp
     src/neighbors/detail/refine_host_uint8_t_float.cpp
-    src/neighbors/detail/selection_faiss_int32_t_float.cu
-    src/neighbors/detail/selection_faiss_int_double.cu
-    src/neighbors/detail/selection_faiss_long_float.cu
-    src/neighbors/detail/selection_faiss_size_t_double.cu
-    src/neighbors/detail/selection_faiss_size_t_float.cu
-    src/neighbors/detail/selection_faiss_uint32_t_float.cu
-    src/neighbors/detail/selection_faiss_int64_t_double.cu
-    src/neighbors/detail/selection_faiss_int64_t_half.cu
-    src/neighbors/detail/selection_faiss_uint32_t_double.cu
-    src/neighbors/detail/selection_faiss_uint32_t_half.cu
     src/neighbors/ivf_flat_build_float_int64_t.cu
     src/neighbors/ivf_flat_build_int8_t_int64_t.cu
     src/neighbors/ivf_flat_build_uint8_t_int64_t.cu

diff --git a/cpp/bench/prims/matrix/select_k.cu b/cpp/bench/prims/matrix/select_k.cu
@@ -279,9 +279,6 @@ const static size_t MAX_MEMORY = 16 * 1024 * 1024 * 1024ULL;
         SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kWarpDistributed, input)     \
         SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kWarpDistributedShm, input)  \
       }                                                                        \
-      if (input.k <= raft::neighbors::detail::kFaissMaxK<IdxT, KeyT>()) {      \
-        SELECTION_REGISTER_ALGO_INPUT(KeyT, IdxT, kFaissBlockSelect, input)    \
-      }                                                                        \
     }                                                                          \
   }
 

diff --git a/cpp/include/raft/matrix/detail/select_k-inl.cuh b/cpp/include/raft/matrix/detail/select_k-inl.cuh
@@ -25,7 +25,6 @@
 #include <raft/matrix/init.cuh>
 
 #include <raft/core/resource/thrust_policy.hpp>
-#include <raft/neighbors/detail/selection_faiss.cuh>
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/device/device_memory_resource.hpp>
 #include <thrust/scan.h>

diff --git a/cpp/include/raft/neighbors/detail/selection_faiss-ext.cuh b/cpp/include/raft/neighbors/detail/selection_faiss-ext.cuh
diff --git a/cpp/include/raft/neighbors/detail/selection_faiss-inl.cuh b/cpp/include/raft/neighbors/detail/selection_faiss-inl.cuh
diff --git a/cpp/include/raft/neighbors/detail/selection_faiss.cuh b/cpp/include/raft/neighbors/detail/selection_faiss.cuh
diff --git a/cpp/include/raft/neighbors/detail/selection_faiss_helpers.cuh b/cpp/include/raft/neighbors/detail/selection_faiss_helpers.cuh
diff --git a/cpp/include/raft/sparse/neighbors/detail/knn.cuh b/cpp/include/raft/sparse/neighbors/detail/knn.cuh
@@ -24,12 +24,12 @@
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/cudart_utils.hpp>
 
+#include <raft/matrix/select_k.cuh>
 #include <raft/sparse/coo.hpp>
 #include <raft/sparse/csr.hpp>
 #include <raft/sparse/detail/utils.h>
 #include <raft/sparse/distance/distance.cuh>
 #include <raft/sparse/op/slice.cuh>
-#include <raft/spatial/knn/knn.cuh>
 
 #include <algorithm>
 
@@ -365,15 +365,14 @@ class sparse_knn_t {
     bool ascending = raft::distance::is_min_close(metric);
 
     // kernel to slice first (min) k cols and copy into batched merge buffer
-    raft::spatial::knn::select_k(batch_dists,
-                                 batch_indices,
-                                 batch_rows,
-                                 batch_cols,
-                                 out_dists,
-                                 out_indices,
-                                 ascending,
-                                 n_neighbors,
-                                 resource::get_cuda_stream(handle));
+    raft::matrix::select_k<value_t, value_idx>(
+      handle,
+      make_device_matrix_view<const value_t, int64_t>(batch_dists, batch_rows, batch_cols),
+      make_device_matrix_view<const value_idx, int64_t>(batch_indices, batch_rows, batch_cols),
+      make_device_matrix_view<value_t, int64_t>(out_dists, batch_rows, n_neighbors),
+      make_device_matrix_view<value_idx, int64_t>(out_indices, batch_rows, n_neighbors),
+      ascending,
+      true);
   }
 
   void compute_distances(csr_batcher_t<value_idx, value_t>& idx_batcher,