raft enabled cmake build (#2898)

Summary: Pull Request resolved: #2898 Reviewed By: mdouze Differential Revision: D46561295 Pulled By: algoriddle fbshipit-source-id: b9806c0c52acf82124c3b2e0095b1c1979318dcd
facebookresearch · Jun 13, 2023 · 6951466 · 6951466
1 parent 27ffd14
commit 6951466
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 3 deletions.
diff --git a/.circleci/config.yml b/.circleci/config.yml
@@ -150,6 +150,9 @@ jobs:
       gpu:
         type: string
         default: "OFF"
+      raft:
+        type: string
+        default: "OFF"
     executor: << parameters.exec >>
     environment:
       OMP_NUM_THREADS: 10
@@ -187,6 +190,7 @@ jobs:
                   -DBUILD_TESTING=ON \
                   -DBUILD_SHARED_LIBS=OFF \
                   -DFAISS_ENABLE_GPU=<< parameters.gpu >> \
+                  -DFAISS_ENABLE_RAFT=<< parameters.raft >> \
                   -DFAISS_OPT_LEVEL=<< parameters.opt_level >> \
                   -DFAISS_ENABLE_C_API=ON \
                   -DPYTHON_EXECUTABLE=$(which python) \
@@ -259,6 +263,13 @@ workflows:
           gpu: "ON"
           requires:
             - Linux x86_64 (cmake)
+      - build_cmake:
+          name: Linux x86_64 GPU w/ RAFT (cmake)
+          exec: linux-x86_64-gpu
+          gpu: "ON"
+          raft: "ON"
+          requires:
+            - Linux x86_64 (cmake)
       - build_conda:
           name: Linux x86_64 (conda)
           exec: linux-x86_64-cpu

diff --git a/CMakeLists.txt b/CMakeLists.txt
@@ -35,6 +35,8 @@ include(rapids-export)
 include(rapids-find)
 
 rapids_cuda_init_architectures(faiss)
+rapids_cuda_init_architectures(pyfaiss)
+rapids_cuda_init_architectures(faiss_c_library)
 endif()
 
 project(faiss

diff --git a/c_api/gpu/CMakeLists.txt b/c_api/gpu/CMakeLists.txt
@@ -16,7 +16,7 @@ file(GLOB FAISS_C_API_GPU_HEADERS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} "*.h")
 faiss_install_headers("${FAISS_C_API_GPU_HEADERS}" c_api/gpu)
 
 find_package(CUDAToolkit REQUIRED)
-target_link_libraries(faiss_c PUBLIC CUDA::cudart CUDA::cublas)
+target_link_libraries(faiss_c PUBLIC CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
 
 add_executable(example_gpu_c EXCLUDE_FROM_ALL example_gpu_c.c)
 target_link_libraries(example_gpu_c PRIVATE faiss_c)
diff --git a/faiss/python/CMakeLists.txt b/faiss/python/CMakeLists.txt
@@ -94,8 +94,8 @@ endif()
 
 if(FAISS_ENABLE_GPU)
   find_package(CUDAToolkit REQUIRED)
-  target_link_libraries(swigfaiss PRIVATE CUDA::cudart)
-  target_link_libraries(swigfaiss_avx2 PRIVATE CUDA::cudart)
+  target_link_libraries(swigfaiss PRIVATE CUDA::cudart $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
+  target_link_libraries(swigfaiss_avx2 PRIVATE CUDA::cudart $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
 endif()
 
 find_package(OpenMP REQUIRED)