Replace faiss bfKnn

build.sh

@@ -18,7 +18,7 @@ ARGS=$*
 # script, and that this script resides in the repo dir!
 REPODIR=$(cd $(dirname $0); pwd)
 
-VALIDARGS="clean libraft pylibraft raft-dask docs tests bench clean --uninstall  -v -g -n --compile-libs --compile-nn --compile-dist --allgpuarch --no-nvtx --show_depr_warn -h --buildfaiss --minimal-deps"
+VALIDARGS="clean libraft pylibraft raft-dask docs tests bench clean --uninstall  -v -g -n --compile-libs --compile-nn --compile-dist --allgpuarch --no-nvtx --show_depr_warn -h --minimal-deps"
 HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<tool>] [--limit-tests=<targets>] [--limit-bench=<targets>]
  where <target> is:
    clean            - remove all existing build artifacts and configuration (start over)
@@ -45,7 +45,6 @@ HELP="$0 [<target> ...] [<flag> ...] [--cmake-args=\"<args>\"] [--cache-tool=<to
    --limit-tests               - semicolon-separated list of test executables to compile (e.g. NEIGHBORS_TEST;CLUSTER_TEST)
    --limit-bench               - semicolon-separated list of benchmark executables to compute (e.g. NEIGHBORS_BENCH;CLUSTER_BENCH)
    --allgpuarch                - build for all supported GPU architectures
-   --buildfaiss                - build faiss statically into raft
    --no-nvtx                   - disable nvtx (profiling markers), but allow enabling it in downstream projects
    --show_depr_warn            - show cmake deprecation warnings
    --cmake-args=\\\"<args>\\\" - pass arbitrary list of CMake configuration options (escape all quotes in argument)
@@ -68,7 +67,6 @@ BUILD_ALL_GPU_ARCH=0
 BUILD_TESTS=OFF
 BUILD_TYPE=Release
 BUILD_BENCH=OFF
-BUILD_STATIC_FAISS=OFF
 COMPILE_LIBRARIES=OFF
 COMPILE_NN_LIBRARY=OFF
 COMPILE_DIST_LIBRARY=OFF
@@ -335,9 +333,6 @@ if hasArg bench || (( ${NUMARGS} == 0 )); then
 
 fi
 
-if hasArg --buildfaiss; then
-    BUILD_STATIC_FAISS=ON
-fi
 if hasArg --no-nvtx; then
     NVTX=OFF
 fi
@@ -407,7 +402,6 @@ if (( ${NUMARGS} == 0 )) || hasArg libraft || hasArg docs || hasArg tests || has
           -DCMAKE_MESSAGE_LOG_LEVEL=${CMAKE_LOG_LEVEL} \
           -DRAFT_COMPILE_NN_LIBRARY=${COMPILE_NN_LIBRARY} \
           -DRAFT_COMPILE_DIST_LIBRARY=${COMPILE_DIST_LIBRARY} \
-          -DRAFT_USE_FAISS_STATIC=${BUILD_STATIC_FAISS} \
           -DRAFT_ENABLE_thrust_DEPENDENCY=${ENABLE_thrust_DEPENDENCY} \
           ${CACHE_ARGS} \
           ${EXTRA_CMAKE_ARGS}

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -22,7 +22,6 @@ dependencies:
 - dask>=2022.12.0
 - distributed>=2022.12.0
 - doxygen>=1.8.20
-- faiss-proc=*=cuda
 - gcc_linux-64=9
 - libcublas-dev=11.11.3.6
 - libcublas=11.11.3.6
@@ -32,7 +31,6 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
-- libfaiss>=1.7.1=cuda*
 - ninja
 - pytest
 - pytest-cov

conda/recipes/libraft/conda_build_config.yaml

@@ -19,9 +19,6 @@ nccl_version:
 gtest_version:
   - "=1.10.0"
 
-libfaiss_version:
-  - "1.7.2 *_cuda"
-
 # The CTK libraries below are missing from the conda-forge::cudatoolkit
 # package. The "*_host_*" version specifiers correspond to `11.8` packages and the
 # "*_run_*" version specifiers correspond to `11.x` packages.

conda/recipes/libraft/meta.yaml

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 # Usage:
 #   conda build . -c conda-forge -c nvidia -c rapidsai
@@ -126,7 +126,6 @@ outputs:
       host:
         - {{ pin_subpackage('libraft-headers', exact=True) }}
         - cuda-profiler-api {{ cuda_profiler_api_host_version }}
-        - faiss-proc=*=cuda
         - lapack
         - libcublas {{ libcublas_host_version }}
         - libcublas-dev {{ libcublas_host_version }}
@@ -136,10 +135,7 @@ outputs:
         - libcusolver-dev {{ libcusolver_host_version }}
         - libcusparse {{ libcusparse_host_version }}
         - libcusparse-dev {{ libcusparse_host_version }}
-        - libfaiss {{ libfaiss_version }}
       run:
-        - faiss-proc=*=cuda
-        - libfaiss {{ libfaiss_version }}
         - {{ pin_subpackage('libraft-headers', exact=True) }}
     about:
       home: https://rapids.ai/

cpp/CMakeLists.txt

@@ -65,9 +65,7 @@ option(
 option(RAFT_COMPILE_DIST_LIBRARY "Enable building raft distant shared library instantiations"
        ${RAFT_COMPILE_LIBRARIES}
 )
-option(RAFT_ENABLE_NN_DEPENDENCIES "Search for raft::nn dependencies like faiss"
-       ${RAFT_COMPILE_LIBRARIES}
-)
+option(RAFT_ENABLE_NN_DEPENDENCIES "Search for raft::nn dependencies" ${RAFT_COMPILE_LIBRARIES})
 
 option(RAFT_ENABLE_thrust_DEPENDENCY "Enable Thrust dependency" ON)
 
@@ -83,16 +81,7 @@ if(BUILD_TESTS AND NOT RAFT_ENABLE_thrust_DEPENDENCY)
   set(RAFT_ENABLE_thrust_DEPENDENCY ON)
 endif()
 
-option(RAFT_EXCLUDE_FAISS_FROM_ALL "Exclude FAISS targets from RAFT's 'all' target" ON)
-
 include(CMakeDependentOption)
-cmake_dependent_option(
-  RAFT_USE_FAISS_STATIC
-  "Build and statically link the FAISS library for nearest neighbors search on GPU"
-  ON
-  RAFT_COMPILE_LIBRARIES
-  OFF
-)
 
 message(VERBOSE "RAFT: Building optional components: ${raft_FIND_COMPONENTS}")
 message(VERBOSE "RAFT: Build RAFT unit-tests: ${BUILD_TESTS}")
@@ -177,7 +166,6 @@ rapids_cpm_init()
 # thrust before rmm/cuco so we get the right version of thrust/cub
 include(cmake/thirdparty/get_thrust.cmake)
 include(cmake/thirdparty/get_rmm.cmake)
-include(cmake/thirdparty/get_faiss.cmake)
 include(cmake/thirdparty/get_cutlass.cmake)
 
 if(RAFT_ENABLE_cuco_DEPENDENCY)
@@ -303,7 +291,6 @@ if(RAFT_COMPILE_DIST_LIBRARY)
     src/distance/distance/specializations/detail/chebyshev.cu
     src/distance/distance/specializations/detail/correlation.cu
     src/distance/distance/specializations/detail/cosine.cu
-    src/distance/distance/specializations/detail/cosine.cu
     src/distance/distance/specializations/detail/hamming_unexpanded.cu
     src/distance/distance/specializations/detail/hellinger_expanded.cu
     src/distance/distance/specializations/detail/jensen_shannon_float_float_float_int.cu
@@ -427,7 +414,10 @@ if(RAFT_COMPILE_NN_LIBRARY)
     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/knn_long_float_int.cu
+    src/nn/specializations/knn_long_float_uint.cu
+    src/nn/specializations/knn_uint_float_int.cu
+    src/nn/specializations/knn_uint_float_uint.cu
   )
   set_target_properties(
     raft_nn_lib
@@ -444,7 +434,7 @@ if(RAFT_COMPILE_NN_LIBRARY)
 
   target_link_libraries(
     raft_nn_lib
-    PUBLIC faiss::faiss raft::raft
+    PUBLIC raft::raft
     PRIVATE nvidia::cutlass::cutlass
   )
   target_compile_options(
@@ -652,12 +642,6 @@ endif()
 
 if(nn IN_LIST raft_FIND_COMPONENTS)
   enable_language(CUDA)
-
-  if(TARGET faiss AND (NOT TARGET faiss::faiss))
-      add_library(faiss::faiss ALIAS faiss)
-  elseif(TARGET faiss::faiss AND (NOT TARGET faiss))
-      add_library(faiss ALIAS faiss::faiss)
-  endif()
 endif()
 ]=]
 )

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.

cpp/include/raft/matrix/detail/select_radix.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.

cpp/include/raft/spatial/knn/detail/faiss_select/DistanceUtils.h

@@ -0,0 +1,52 @@
+/**
+ * Copyright (c) Facebook, Inc. and its affiliates.
+ *
+ * This source code is licensed under the MIT license found in the
+ * LICENSE file thirdparty/LICENSES/LICENSE.faiss
+ */
+
+#pragma once
+
+namespace raft::spatial::knn::detail::faiss_select {
+// If the inner size (dim) of the vectors is small, we want a larger query tile
+// size, like 1024
+inline void chooseTileSize(size_t numQueries,
+                           size_t numCentroids,
+                           size_t dim,
+                           size_t elementSize,
+                           size_t totalMem,
+                           size_t& tileRows,
+                           size_t& tileCols)
+{
+  // The matrix multiplication should be large enough to be efficient, but if
+  // it is too large, we seem to lose efficiency as opposed to
+  // double-streaming. Each tile size here defines 1/2 of the memory use due
+  // to double streaming. We ignore available temporary memory, as that is
+  // adjusted independently by the user and can thus meet these requirements
+  // (or not). For <= 4 GB GPUs, prefer 512 MB of usage. For <= 8 GB GPUs,
+  // prefer 768 MB of usage. Otherwise, prefer 1 GB of usage.
+  size_t targetUsage = 0;
+
+  if (totalMem <= ((size_t)4) * 1024 * 1024 * 1024) {
+    targetUsage = 512 * 1024 * 1024;
+  } else if (totalMem <= ((size_t)8) * 1024 * 1024 * 1024) {
+    targetUsage = 768 * 1024 * 1024;
+  } else {
+    targetUsage = 1024 * 1024 * 1024;
+  }
+
+  targetUsage /= 2 * elementSize;
+
+  // 512 seems to be a batch size sweetspot for float32.
+  // If we are on float16, increase to 512.
+  // If the k size (vec dim) of the matrix multiplication is small (<= 32),
+  // increase to 1024.
+  size_t preferredTileRows = 512;
+  if (dim <= 32) { preferredTileRows = 1024; }
+
+  tileRows = std::min(preferredTileRows, numQueries);
+
+  // tileCols is the remainder size
+  tileCols = std::min(targetUsage / preferredTileRows, numCentroids);
+}
+}  // namespace raft::spatial::knn::detail::faiss_select