From 6951466b437baedae2c3d738e561de3228275a2e Mon Sep 17 00:00:00 2001
From: Gergely Szilvasy <gsz@meta.com>
Date: Tue, 13 Jun 2023 08:43:18 -0700
Subject: [PATCH] raft enabled cmake build (#2898)

Summary: Pull Request resolved: https://github.com/facebookresearch/faiss/pull/2898

Reviewed By: mdouze

Differential Revision: D46561295

Pulled By: algoriddle

fbshipit-source-id: b9806c0c52acf82124c3b2e0095b1c1979318dcd
---
 .circleci/config.yml        | 11 +++++++++++
 CMakeLists.txt              |  2 ++
 c_api/gpu/CMakeLists.txt    |  2 +-
 faiss/python/CMakeLists.txt |  4 ++--
 4 files changed, 16 insertions(+), 3 deletions(-)

diff --git a/.circleci/config.yml b/.circleci/config.yml
index c7d8d03230..f5eb52cdaf 100644
--- 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
index e6aced198a..eabbe066c0 100644
--- 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
index 397462c775..4ec926439d 100644
--- 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
index fef27d9b2a..d7da50ceb7 100644
--- 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)