Merge branch 'branch-24.02' of https://github.com/rapidsai/raft into …

…faiss-ivf

ci/build_python.sh

@@ -15,11 +15,11 @@ CPP_CHANNEL=$(rapids-download-conda-from-s3 cpp)
 
 version=$(rapids-generate-version)
 git_commit=$(git rev-parse HEAD)
-export RAPIDS_PACKAGE_VERSION=${version} 
+export RAPIDS_PACKAGE_VERSION=${version}
 echo "${version}" > VERSION
 
 package_dir="python"
-for package_name in pylibraft raft-dask; do 
+for package_name in pylibraft raft-dask; do
   underscore_package_name=$(echo "${package_name}" | tr "-" "_")
   sed -i "/^__git_commit__/ s/= .*/= \"${git_commit}\"/g" "${package_dir}/${package_name}/${underscore_package_name}/_version.py"
 done
@@ -39,10 +39,10 @@ rapids-conda-retry mambabuild \
 
 # Build ann-bench for each cuda and python version
 rapids-conda-retry mambabuild \
---no-test \
---channel "${CPP_CHANNEL}" \
---channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
-conda/recipes/raft-ann-bench
+  --no-test \
+  --channel "${CPP_CHANNEL}" \
+  --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+  conda/recipes/raft-ann-bench
 
 # Build ann-bench-cpu only in CUDA 11 jobs since it only depends on python
 # version

conda/recipes/libraft/conda_build_config.yaml

@@ -73,7 +73,7 @@ cuda11_cuda_profiler_api_run_version:
   - ">=11.4.240,<12"
 
 spdlog_version:
-  - ">=1.11.0,<1.12"
+  - ">=1.12.0,<1.13"
 
 fmt_version:
-  - ">=9.1.0,<10"
+  - ">=10.1.1,<11"

conda/recipes/raft-ann-bench-cpu/conda_build_config.yaml

@@ -20,7 +20,7 @@ nlohmann_json_version:
   - ">=3.11.2"
 
 spdlog_version:
-  - ">=1.11.0,<1.12"
+  - ">=1.12.0,<1.13"
 
 fmt_version:
-  - ">=9.1.0,<10"
+  - ">=10.1.1,<11"

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

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)    \
-      }                                                                        \
     }                                                                          \
   }
 

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>

cpp/include/raft/neighbors/brute_force-inl.cuh

@@ -284,6 +284,44 @@ void fused_l2_knn(raft::resources const& handle,
 /**
  * @brief Build the index from the dataset for efficient search.
  *
+ * This function builds a brute force index for the given dataset. This lets you re-use
+ * precalculated norms for the dataset, leading to a speedup over calling
+ * raft::neighbors::brute_force::knn repeatedly.
+ *
+ * Example usage:
+ * @code{.cpp}
+ * #include <raft/neighbors/brute_force.cuh>
+ * #include <raft/core/device_mdarray.hpp>
+ * #include <raft/random/make_blobs.cuh>
+ *
+ * // create a random dataset
+ * int n_rows = 10000;
+ * int n_cols = 10000;
+ *
+ * raft::device_resources res;
+ * auto dataset = raft::make_device_matrix<float, int64_t>(res, n_rows, n_cols);
+ * auto labels = raft::make_device_vector<int64_t, int64_t>(res, n_rows);
+ *
+ * raft::random::make_blobs(res, dataset.view(), labels.view());
+ *
+ * // create a brute_force knn index from the dataset
+ * auto index = raft::neighbors::brute_force::build(res,
+ *                                                  raft::make_const_mdspan(dataset.view()));
+ *
+ * // Use the constructed index to search for the nearest 128 neighbors
+ * int k = 128;
+ * auto search = raft::make_const_mdspan(dataset.view());
+ *
+ * auto indices= raft::make_device_matrix<int, int64_t>(res, search.extent(0), k);
+ * auto distances = raft::make_device_matrix<float, int64_t>(res, search.extent(0), k);
+ *
+ * raft::neighbors::brute_force::search(res,
+ *                                      index,
+ *                                      search,
+ *                                      indices.view(),
+ *                                      distances.view());
+ * @endcode
+ *
  * @tparam T data element type
  *
  * @param[in] res
@@ -330,6 +368,8 @@ index<T> build(raft::resources const& res,
 /**
  * @brief Brute Force search using the constructed index.
  *
+ * See raft::neighbors::brute_force::build for a usage example
+ *
  * @tparam T data element type
  * @tparam IdxT type of the indices
  *

cpp/include/raft/neighbors/brute_force.cuh

@@ -44,10 +44,10 @@ namespace raft::neighbors::brute_force {
  * int n_cols = 10000;
 
  * raft::device_resources res;
- * auto dataset = raft::make_device_matrix<float, int>(res, n_rows, n_cols);
- * auto labels = raft::make_device_vector<float, int>(res, n_rows);
+ * auto dataset = raft::make_device_matrix<float, int64_t>(res, n_rows, n_cols);
+ * auto labels = raft::make_device_vector<int64_t, int64_t>(res, n_rows);
 
- * raft::make_blobs(res, dataset.view(), labels.view());
+ * raft::random::make_blobs(res, dataset.view(), labels.view());
  *
  * // create a brute_force knn index from the dataset
  * auto index = raft::neighbors::brute_force::build(res,