Adding FAISS cpu to raft-ann-bench

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

@@ -17,7 +17,6 @@ dependencies:
 - cudatoolkit
 - cxx-compiler
 - cython>=3.0.0
-- faiss-proc=*=cuda
 - gcc_linux-64=11.*
 - glog>=0.6.0
 - h5py>=3.8.0
@@ -30,11 +29,12 @@ dependencies:
 - libcusolver=11.4.1.48
 - libcusparse-dev=11.7.5.86
 - libcusparse=11.7.5.86
-- libfaiss>=1.7.1
+- libfaiss-avx2=1.7.4
 - nccl>=2.9.9
 - ninja
 - nlohmann_json>=3.11.2
 - nvcc_linux-64=11.8
+- rmm=23.10.*
 - scikit-build>=0.13.1
 - sysroot_linux-64==2.17
 name: bench_ann_cuda-118_arch-x86_64

conda/recipes/libraft/build_libraft.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-./build.sh libraft -v --allgpuarch --compile-lib --build-metrics=compile_lib --incl-cache-stats --no-nvtx
+./build.sh libraft --allgpuarch --compile-lib --build-metrics=compile_lib --incl-cache-stats --no-nvtx

conda/recipes/libraft/build_libraft_headers.sh

@@ -1,4 +1,4 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-./build.sh libraft -v --allgpuarch --no-nvtx
+./build.sh libraft --allgpuarch --no-nvtx

conda/recipes/libraft/build_libraft_template.sh

@@ -2,4 +2,4 @@
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 # Just building template so we verify it uses libraft.so and fail if it doesn't build
-./build.sh template -v
+./build.sh template

conda/recipes/libraft/build_libraft_tests.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 # Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
-./build.sh tests bench-prims -v --allgpuarch --no-nvtx --build-metrics=tests_bench_prims --incl-cache-stats
+./build.sh tests bench-prims --allgpuarch --no-nvtx --build-metrics=tests_bench_prims --incl-cache-stats
 cmake --install cpp/build --component testing

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

@@ -10,6 +10,9 @@ sysroot_version:
 cmake_version:
   - ">=3.26.4"
 
+faiss_version:
+  - "1.7.4"
+
 glog_version:
   - ">=0.6.0"
 

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

@@ -47,6 +47,7 @@ requirements:
   host:
     - glog {{ glog_version }}
     - matplotlib
+#    - faiss-cpu {{ faiss_version }}
     - nlohmann_json {{ nlohmann_json_version }}
     - python
     - pyyaml
@@ -55,13 +56,13 @@ requirements:
   run:
     - glog {{ glog_version }}
     - h5py {{ h5py_version }}
+#    - faiss-cpu {{ faiss_version }}
     - matplotlib
     - python
     - pyyaml
     - pandas
     - benchmark
-
 about:
   home: https://rapids.ai/
   license: Apache-2.0
-  summary: libraft ann bench
+  summary: RAFT ANN CPU benchmarks

conda/recipes/raft-ann-bench/build.sh

@@ -1,5 +1,5 @@
 #!/usr/bin/env bash
 # Copyright (c) 2023, NVIDIA CORPORATION.
 
-./build.sh bench-ann -v --allgpuarch --no-nvtx --build-metrics=bench_ann --incl-cache-stats
+./build.sh bench-ann --allgpuarch --no-nvtx --build-metrics=bench_ann --incl-cache-stats
 cmake --install cpp/build --component ann_bench

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

@@ -26,7 +26,7 @@ glog_version:
   - ">=0.6.0"
 
 faiss_version:
-  - ">=1.7.1"
+  - "1.7.4"
 
 h5py_version:
   - ">=3.8.0"

conda/recipes/raft-ann-bench/meta.yaml

@@ -71,10 +71,10 @@ requirements:
     - glog {{ glog_version }}
     - nlohmann_json {{ nlohmann_json_version }}
     # Temporarily ignore faiss benchmarks on CUDA 12 because packages do not exist yet
-    {% if cuda_major == "11" %}
-    - faiss-proc=*=cuda
-    - libfaiss {{ faiss_version }}
-    {% endif %}
+#    {% if cuda_major == "11" %}
+#    - faiss-proc=*=cuda
+#    - libfaiss {{ faiss_version }}
+#    {% endif %}
     - h5py {{ h5py_version }}
     - benchmark
     - matplotlib
@@ -93,10 +93,10 @@ requirements:
     {% endif %}
     - glog {{ glog_version }}
     # Temporarily ignore faiss benchmarks on CUDA 12 because packages do not exist yet
-    {% if cuda_major == "11" %}
-    - faiss-proc=*=cuda
-    - libfaiss {{ faiss_version }}
-    {% endif %}
+#    {% if cuda_major == "11" %}
+#    - faiss-proc=*=cuda
+#    - libfaiss {{ faiss_version }}
+#    {% endif %}
     - h5py {{ h5py_version }}
     - benchmark
     - glog {{ glog_version }}

cpp/bench/ann/CMakeLists.txt

@@ -15,9 +15,18 @@
 # ##################################################################################################
 # * benchmark options ------------------------------------------------------------------------------
 
-option(RAFT_ANN_BENCH_USE_FAISS_BFKNN "Include faiss' brute-force knn algorithm in benchmark" ON)
-option(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT "Include faiss' ivf flat algorithm in benchmark" ON)
-option(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ "Include faiss' ivf pq algorithm in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT "Include faiss' brute-force knn algorithm in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT "Include faiss' ivf flat algorithm in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ "Include faiss' ivf pq algorithm in benchmark" ON)
+option(RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT
+       "Include faiss' cpu brute-force knn algorithm in benchmark" ON
+)
+option(RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT "Include faiss' cpu brute-force algorithm in benchmark" ON)
+
+option(RAFT_ANN_BENCH_USE_FAISS_CPU_IVF_FLAT "Include faiss' cpu ivf flat algorithm in benchmark"
+       ON
+)
+option(RAFT_ANN_BENCH_USE_FAISS_CPU_IVF_PQ "Include faiss' cpu ivf pq algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT "Include raft's ivf flat algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ "Include raft's ivf pq algorithm in benchmark" ON)
 option(RAFT_ANN_BENCH_USE_RAFT_CAGRA "Include raft's CAGRA in benchmark" ON)
@@ -33,9 +42,15 @@ option(RAFT_ANN_BENCH_SINGLE_EXE
 find_package(Threads REQUIRED)
 
 if(BUILD_CPU_ONLY)
-  set(RAFT_ANN_BENCH_USE_FAISS_BFKNN OFF)
-  set(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT OFF)
-  set(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ OFF)
+
+  # Include necessary logging dependencies
+  include(cmake/thirdparty/get_fmt.cmake)
+  include(cmake/thirdparty/get_spdlog.cmake)
+
+  set(RAFT_FAISS_ENABLE_GPU OFF)
+  set(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT OFF)
+  set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT OFF)
+  set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ OFF)
   set(RAFT_ANN_BENCH_USE_RAFT_CAGRA OFF)
@@ -44,22 +59,30 @@ else()
   # Disable faiss benchmarks on CUDA 12 since faiss is not yet CUDA 12-enabled.
   # https://github.com/rapidsai/raft/issues/1627
   if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0.0)
-    set(RAFT_ANN_BENCH_USE_FAISS_BFKNN OFF)
-    set(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT OFF)
-    set(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ OFF)
+    set(RAFT_FAISS_ENABLE_GPU OFF)
+    set(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT OFF)
+    set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT OFF)
+    set(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ OFF)
+  else()
+    set(RAFT_FAISS_ENABLE_GPU ON)
   endif()
 endif()
 
 set(RAFT_ANN_BENCH_USE_FAISS OFF)
-if(RAFT_ANN_BENCH_USE_FAISS_BFKNN
-   OR RAFT_ANN_BENCH_USE_FAISS_IVFPQ
+if(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT
+   OR RAFT_ANN_BENCH_USE_FAISS_GPU_IVFPQ
    OR RAFT_ANN_BENCH_USE_FAISS_IFFLAT
+   OR RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT
+   OR RAFT_ANN_BENCH_USE_FAISS_CPU_IVFPQ
+   OR RAFT_ANN_BENCH_USE_FAISS_CPU_IFFLAT
 )
   set(RAFT_ANN_BENCH_USE_FAISS ON)
+  set(RAFT_USE_FAISS_STATIC ON)
 endif()
 
 set(RAFT_ANN_BENCH_USE_RAFT OFF)
 if(RAFT_ANN_BENCH_USE_RAFT_IVF_PQ
+   OR RAFT_ANN_BENCH_USE_RAFT_BRUTE_FORCE
    OR RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT
    OR RAFT_ANN_BENCH_USE_RAFT_CAGRA
 )
@@ -80,6 +103,12 @@ if(RAFT_ANN_BENCH_USE_GGNN)
 endif()
 
 if(RAFT_ANN_BENCH_USE_FAISS)
+  # We need to ensure that faiss has all the conda
+  # information. So we currently use the very ugly
+  # hammer of `link_libraries` to ensure that all
+  # targets in this directory and the faiss directory
+  # will have the conda includes/link dirs
+  link_libraries($<TARGET_NAME_IF_EXISTS:conda_env>)
   include(cmake/thirdparty/get_faiss.cmake)
 endif()
 
@@ -116,14 +145,15 @@ function(ConfigureAnnBench)
     ${BENCH_NAME}
     PRIVATE raft::raft
             nlohmann_json::nlohmann_json
-            $<$<BOOL:${GPU_BUILD}>:$<$<BOOL:${RAFT_ANN_BENCH_USE_MULTIGPU}>:NCCL::NCCL>>
             ${ConfigureAnnBench_LINKS}
             Threads::Threads
             $<$<BOOL:${GPU_BUILD}>:${RAFT_CTK_MATH_DEPENDENCIES}>
             $<TARGET_NAME_IF_EXISTS:OpenMP::OpenMP_CXX>
             $<TARGET_NAME_IF_EXISTS:conda_env>
             -static-libgcc
             -static-libstdc++
+            $<$<BOOL:${BUILD_CPU_ONLY}>:fmt::fmt-header-only>
+            $<$<BOOL:${BUILD_CPU_ONLY}>:spdlog::spdlog_header_only>
   )
 
   set_target_properties(
@@ -201,6 +231,12 @@ if(RAFT_ANN_BENCH_USE_RAFT_IVF_FLAT)
   )
 endif()
 
+if(RAFT_ANN_BENCH_USE_RAFT_BRUTE_FORCE)
+  ConfigureAnnBench(
+    NAME RAFT_BRUTE_FORCE PATH bench/ann/src/raft/raft_benchmark.cu LINKS raft::compiled
+  )
+endif()
+
 if(RAFT_ANN_BENCH_USE_RAFT_CAGRA)
   ConfigureAnnBench(
     NAME
@@ -213,20 +249,52 @@ if(RAFT_ANN_BENCH_USE_RAFT_CAGRA)
   )
 endif()
 
-if(RAFT_ANN_BENCH_USE_FAISS_IVF_FLAT)
+set(RAFT_FAISS_TARGETS faiss::faiss)
+if(TARGET faiss::faiss_avx2)
+  set(RAFT_FAISS_TARGETS faiss::faiss_avx2)
+endif()
+
+message("RAFT_FAISS_TARGETS: ${RAFT_FAISS_TARGETS}")
+message("CUDAToolkit_LIBRARY_DIR: ${CUDAToolkit_LIBRARY_DIR}")
+if(RAFT_ANN_BENCH_USE_FAISS_CPU_FLAT)
+  ConfigureAnnBench(
+    NAME FAISS_CPU_FLAT PATH bench/ann/src/faiss/faiss_cpu_benchmark.cpp LINKS
+    ${RAFT_FAISS_TARGETS}
+  )
+endif()
+
+if(RAFT_ANN_BENCH_USE_FAISS_CPU_IVF_FLAT)
+  ConfigureAnnBench(
+    NAME FAISS_CPU_IVF_FLAT PATH bench/ann/src/faiss/faiss_cpu_benchmark.cpp LINKS
+    ${RAFT_FAISS_TARGETS}
+  )
+endif()
+
+if(RAFT_ANN_BENCH_USE_FAISS_CPU_IVF_PQ)
   ConfigureAnnBench(
-    NAME FAISS_IVF_FLAT PATH bench/ann/src/faiss/faiss_benchmark.cu LINKS faiss::faiss
+    NAME FAISS_CPU_IVF_PQ PATH bench/ann/src/faiss/faiss_cpu_benchmark.cpp LINKS
+    ${RAFT_FAISS_TARGETS}
   )
 endif()
 
-if(RAFT_ANN_BENCH_USE_FAISS_IVF_PQ)
+if(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_FLAT)
   ConfigureAnnBench(
-    NAME FAISS_IVF_PQ PATH bench/ann/src/faiss/faiss_benchmark.cu LINKS faiss::faiss
+    NAME FAISS_GPU_IVF_FLAT PATH bench/ann/src/faiss/faiss_gpu_benchmark.cu LINKS
+    ${RAFT_FAISS_TARGETS}
   )
 endif()
 
-if(RAFT_ANN_BENCH_USE_FAISS_BFKNN)
-  ConfigureAnnBench(NAME FAISS_BFKNN PATH bench/ann/src/faiss/faiss_benchmark.cu LINKS faiss::faiss)
+if(RAFT_ANN_BENCH_USE_FAISS_GPU_IVF_PQ)
+  ConfigureAnnBench(
+    NAME FAISS_GPU_IVF_PQ PATH bench/ann/src/faiss/faiss_gpu_benchmark.cu LINKS
+    ${RAFT_FAISS_TARGETS}
+  )
+endif()
+
+if(RAFT_ANN_BENCH_USE_FAISS_GPU_FLAT)
+  ConfigureAnnBench(
+    NAME FAISS_GPU_FLAT PATH bench/ann/src/faiss/faiss_gpu_benchmark.cu LINKS ${RAFT_FAISS_TARGETS}
+  )
 endif()
 
 if(RAFT_ANN_BENCH_USE_GGNN)
@@ -277,7 +345,8 @@ if(RAFT_ANN_BENCH_SINGLE_EXE)
   target_compile_definitions(
     ANN_BENCH
     PRIVATE
-      $<$<BOOL:${CUDAToolkit_FOUND}>:ANN_BENCH_LINK_CUDART="libcudart.so.${CUDAToolkit_VERSION_MAJOR}.${CUDAToolkit_VERSION_MINOR}.${CUDAToolkit_VERSION_PATCH}">
+      $<$<BOOL:${CUDAToolkit_FOUND}>:ANN_BENCH_LINK_CUDART="libcudart.so.${CUDAToolkit_VERSION_MAJOR}.${CUDAToolkit_VERSION_MINOR}.${CUDAToolkit_VERSION_PATCH}
+  ">
       $<$<BOOL:${NVTX3_HEADERS_FOUND}>:ANN_BENCH_NVTX3_HEADERS_FOUND>
   )