use RAFT for select_k on the gpu (#656)

This changes to use RAFT https://github.com/rapidsai/raft for GPU top-k code instead of faiss. The RAFT version is quite a bit faster, and also doesn't have the same performance issues with small batch sizes that faiss has (meaning we can delete a bunch of code that was trying to work around that). RAFT also doesn't have limitations on the size of K, where faiss is limited to k less than 2048.
benfred · May 24, 2023 · 3cb4532 · 3cb4532
1 parent f3f1a1f
commit 3cb4532
Show file tree

Hide file tree

Showing 3 changed files with 77 additions and 194 deletions.
diff --git a/implicit/gpu/CMakeLists.txt b/implicit/gpu/CMakeLists.txt
@@ -14,7 +14,7 @@ else()
     add_cython_target(_cuda CXX)
 
     # use rapids-cmake to install dependencies
-    file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-22.02/RAPIDS.cmake
+    file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-23.04/RAPIDS.cmake
         ${CMAKE_BINARY_DIR}/RAPIDS.cmake)
     include(${CMAKE_BINARY_DIR}/RAPIDS.cmake)
     include(rapids-cmake)
@@ -29,30 +29,25 @@ else()
     include(${rapids-cmake-dir}/cpm/rmm.cmake)
     rapids_cpm_rmm(BUILD_EXPORT_SET implicit-exports INSTALL_EXPORT_SET implicit-exports)
 
-    # get faiss from git
-    rapids_cpm_find(FAISS 1.7.2
+    # get raft
+    # note: we're using RAFT in header only mode right now - mainly to reduce binary
+    # size of the compiled wheels
+    rapids_cpm_find(raft 23.06
         CPM_ARGS
-          GIT_REPOSITORY  https://github.com/facebookresearch/faiss.git
-          VERSION         1.7.2
+          GIT_REPOSITORY  https://github.com/rapidsai/raft.git
+          GIT_TAG         branch-23.06
           DOWNLOAD_ONLY   YES
     )
-    FILE(GLOB FAISS_BLOCK_SELECT ${FAISS_SOURCE_DIR}/faiss/gpu/utils/blockselect/*.cu)
-    include_directories(${FAISS_SOURCE_DIR})
+    include_directories(${raft_SOURCE_DIR}/cpp/include)
 
-    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --extended-lambda -Wno-deprecated-gpu-targets -Xfatbin=-compress-all")
+    set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --extended-lambda -Wno-deprecated-gpu-targets -Xfatbin=-compress-all --expt-relaxed-constexpr")
 
     add_library(_cuda MODULE ${_cuda}
         als.cu
         bpr.cu
         matrix.cu
         random.cu
         knn.cu
-        # we need a limited subset of faiss for the blockselect code. rather than compile all of
-        # faiss, just include the minimal subset here
-        ${FAISS_BLOCK_SELECT}
-        ${FAISS_SOURCE_DIR}/faiss/gpu/utils/BlockSelectFloat.cu
-        ${FAISS_SOURCE_DIR}/faiss/gpu/utils/DeviceUtils.cu
-        ${FAISS_SOURCE_DIR}/faiss/gpu/GpuResources.cpp
     )
 
     python_extension_module(_cuda)
@@ -62,9 +57,9 @@ else()
         set_target_properties(_cuda PROPERTIES CUDA_ARCHITECTURES $ENV{IMPLICIT_CUDA_ARCH})
     else()
         if (${CUDAToolkit_VERSION} VERSION_LESS "11.1.0")
-            set_target_properties(_cuda PROPERTIES CUDA_ARCHITECTURES "37;50;70;60;80")
+            set_target_properties(_cuda PROPERTIES CUDA_ARCHITECTURES "60;70;80")
         else()
-            set_target_properties(_cuda PROPERTIES CUDA_ARCHITECTURES "37;50;70;60;80;86")
+            set_target_properties(_cuda PROPERTIES CUDA_ARCHITECTURES "60;70;80;86")
         endif()
     endif()
     target_link_libraries(_cuda CUDA::cublas CUDA::curand rmm::rmm)