Skip to content

Commit 4f574d5

Browse files
cjnoletfacebook-github-bot
cjnolet
authored and
facebook-github-bot
committed
Integrate RAFT FlatIndex / Distances (facebookresearch#2707)
Summary: I've broken out the FlatIndex / Distances changes from facebookresearch#2521 into a separate PR to make things a litle easier to review. This does also include a couple other minor changes to the IVF Flat index which I used to make it easier to dispatch to the RAFT version. I can revert that change too if desired. Pull Request resolved: facebookresearch#2707 Reviewed By: wickedfoo Differential Revision: D44758912 Pulled By: algoriddle fbshipit-source-id: b2544990b4e941a2558f5004bceec4af4fa1ad09
1 parent cbff63e commit 4f574d5

24 files changed

+1293
-62
lines changed

‎CMakeLists.txt

+43-1
Original file line numberDiff line numberDiff line change
@@ -4,22 +4,58 @@
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
66

7+
# =============================================================================
8+
# Copyright (c) 2023, NVIDIA CORPORATION.
9+
#
10+
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
11+
# in compliance with the License. You may obtain a copy of the License at
12+
#
13+
# http://www.apache.org/licenses/LICENSE-2.0
14+
#
15+
# Unless required by applicable law or agreed to in writing, software distributed under the License
16+
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
17+
# or implied. See the License for the specific language governing permissions and limitations under
18+
# the License.
19+
# =============================================================================
20+
721
cmake_minimum_required(VERSION 3.23.1 FATAL_ERROR)
822

23+
set(FAISS_LANGUAGES CXX)
24+
25+
if(FAISS_ENABLE_GPU)
26+
list(APPEND FAISS_LANGUAGES CUDA)
27+
endif()
28+
29+
if(FAISS_ENABLE_RAFT)
30+
include(cmake/thirdparty/fetch_rapids.cmake)
31+
include(rapids-cmake)
32+
include(rapids-cpm)
33+
include(rapids-cuda)
34+
include(rapids-export)
35+
include(rapids-find)
36+
37+
rapids_cuda_init_architectures(faiss)
38+
endif()
39+
940
project(faiss
1041
VERSION 1.7.4
1142
DESCRIPTION "A library for efficient similarity search and clustering of dense vectors."
1243
HOMEPAGE_URL "https://github.com/facebookresearch/faiss"
13-
LANGUAGES CXX)
44+
LANGUAGES ${FAISS_LANGUAGES})
1445
include(GNUInstallDirs)
1546

47+
if(FAISS_ENABLE_RAFT)
48+
set(CMAKE_CXX_STANDARD 17)
49+
else()
1650
set(CMAKE_CXX_STANDARD 11)
51+
endif()
1752

1853
list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
1954

2055
# Valid values are "generic", "avx2".
2156
option(FAISS_OPT_LEVEL "" "generic")
2257
option(FAISS_ENABLE_GPU "Enable support for GPU indexes." ON)
58+
option(FAISS_ENABLE_RAFT "Enable RAFT for GPU indexes." OFF)
2359
option(FAISS_ENABLE_PYTHON "Build Python extension." ON)
2460
option(FAISS_ENABLE_C_API "Build C API." OFF)
2561

@@ -28,6 +64,12 @@ if(FAISS_ENABLE_GPU)
2864
enable_language(CUDA)
2965
endif()
3066

67+
if(FAISS_ENABLE_RAFT)
68+
rapids_cpm_init()
69+
include(cmake/thirdparty/get_raft.cmake)
70+
include(cmake/thirdparty/get_cutlass.cmake)
71+
endif()
72+
3173
add_subdirectory(faiss)
3274

3375
if(FAISS_ENABLE_GPU)

‎INSTALL.md

+3
Original file line numberDiff line numberDiff line change
@@ -105,6 +105,9 @@ Several options can be passed to CMake, among which:
105105
values are `ON` and `OFF`),
106106
- `-DFAISS_ENABLE_PYTHON=OFF` in order to disable building python bindings
107107
(possible values are `ON` and `OFF`),
108+
- `-DFAISS_ENABLE_RAFT=ON` in order to enable building the RAFT implementations
109+
of the IVF-Flat and IVF-PQ GPU-accelerated indices (default is `OFF`, possible
110+
values are `ON` and `OFF`)
108111
- `-DBUILD_TESTING=OFF` in order to disable building C++ tests,
109112
- `-DBUILD_SHARED_LIBS=ON` in order to build a shared library (possible values
110113
are `ON` and `OFF`),

‎cmake/thirdparty/fetch_rapids.cmake

+24
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,24 @@
1+
# Copyright (c) Facebook, Inc. and its affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
# =============================================================================
6+
# Copyright (c) 2023, NVIDIA CORPORATION.
7+
#
8+
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
9+
# in compliance with the License. You may obtain a copy of the License at
10+
#
11+
# http://www.apache.org/licenses/LICENSE-2.0
12+
#
13+
# Unless required by applicable law or agreed to in writing, software distributed under the License
14+
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
15+
# or implied. See the License for the specific language governing permissions and limitations under
16+
# the License.
17+
# =============================================================================
18+
set(RAPIDS_VERSION "23.02")
19+
20+
if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake)
21+
file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-${RAPIDS_VERSION}/RAPIDS.cmake
22+
${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake)
23+
endif()
24+
include(${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake)

‎cmake/thirdparty/get_cutlass.cmake

+107
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,107 @@
1+
# Copyright (c) Facebook, Inc. and its affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
# =============================================================================
6+
# Copyright (c) 2021-2023, NVIDIA CORPORATION.
7+
#
8+
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
9+
# in compliance with the License. You may obtain a copy of the License at
10+
#
11+
# http://www.apache.org/licenses/LICENSE-2.0
12+
#
13+
# Unless required by applicable law or agreed to in writing, software distributed under the License
14+
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
15+
# or implied. See the License for the specific language governing permissions and limitations under
16+
# the License.
17+
# =============================================================================
18+
19+
function(find_and_configure_cutlass)
20+
set(oneValueArgs VERSION REPOSITORY PINNED_TAG)
21+
cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})
22+
23+
# if(RAFT_ENABLE_DIST_DEPENDENCIES OR RAFT_COMPILE_LIBRARIES)
24+
set(CUTLASS_ENABLE_HEADERS_ONLY
25+
ON
26+
CACHE BOOL "Enable only the header library"
27+
)
28+
set(CUTLASS_NAMESPACE
29+
"raft_cutlass"
30+
CACHE STRING "Top level namespace of CUTLASS"
31+
)
32+
set(CUTLASS_ENABLE_CUBLAS
33+
OFF
34+
CACHE BOOL "Disable CUTLASS to build with cuBLAS library."
35+
)
36+
37+
if (CUDA_STATIC_RUNTIME)
38+
set(CUDART_LIBRARY "${CUDA_cudart_static_LIBRARY}" CACHE FILEPATH "fixing cutlass cmake code" FORCE)
39+
endif()
40+
41+
rapids_cpm_find(
42+
NvidiaCutlass ${PKG_VERSION}
43+
GLOBAL_TARGETS nvidia::cutlass::cutlass
44+
CPM_ARGS
45+
GIT_REPOSITORY ${PKG_REPOSITORY}
46+
GIT_TAG ${PKG_PINNED_TAG}
47+
GIT_SHALLOW TRUE
48+
OPTIONS "CUDAToolkit_ROOT ${CUDAToolkit_LIBRARY_DIR}"
49+
)
50+
51+
if(TARGET CUTLASS AND NOT TARGET nvidia::cutlass::cutlass)
52+
add_library(nvidia::cutlass::cutlass ALIAS CUTLASS)
53+
endif()
54+
55+
if(NvidiaCutlass_ADDED)
56+
rapids_export(
57+
BUILD NvidiaCutlass
58+
EXPORT_SET NvidiaCutlass
59+
GLOBAL_TARGETS nvidia::cutlass::cutlass
60+
NAMESPACE nvidia::cutlass::
61+
)
62+
endif()
63+
# endif()
64+
65+
# We generate the cutlass-config files when we built cutlass locally, so always do
66+
# `find_dependency`
67+
rapids_export_package(
68+
BUILD NvidiaCutlass raft-distance-exports GLOBAL_TARGETS nvidia::cutlass::cutlass
69+
)
70+
rapids_export_package(
71+
INSTALL NvidiaCutlass raft-distance-exports GLOBAL_TARGETS nvidia::cutlass::cutlass
72+
)
73+
rapids_export_package(
74+
BUILD NvidiaCutlass raft-nn-exports GLOBAL_TARGETS nvidia::cutlass::cutlass
75+
)
76+
rapids_export_package(
77+
INSTALL NvidiaCutlass raft-nn-exports GLOBAL_TARGETS nvidia::cutlass::cutlass
78+
)
79+
80+
# Tell cmake where it can find the generated NvidiaCutlass-config.cmake we wrote.
81+
include("${rapids-cmake-dir}/export/find_package_root.cmake")
82+
rapids_export_find_package_root(
83+
INSTALL NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}/../]=] raft-distance-exports
84+
)
85+
rapids_export_find_package_root(
86+
BUILD NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}]=] raft-distance-exports
87+
)
88+
include("${rapids-cmake-dir}/export/find_package_root.cmake")
89+
rapids_export_find_package_root(
90+
INSTALL NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}/../]=] raft-nn-exports
91+
)
92+
rapids_export_find_package_root(
93+
BUILD NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}]=] raft-nn-exports
94+
)
95+
endfunction()
96+
97+
if(NOT RAFT_CUTLASS_GIT_TAG)
98+
set(RAFT_CUTLASS_GIT_TAG v2.9.1)
99+
endif()
100+
101+
if(NOT RAFT_CUTLASS_GIT_REPOSITORY)
102+
set(RAFT_CUTLASS_GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git)
103+
endif()
104+
105+
find_and_configure_cutlass(
106+
VERSION 2.9.1 REPOSITORY ${RAFT_CUTLASS_GIT_REPOSITORY} PINNED_TAG ${RAFT_CUTLASS_GIT_TAG}
107+
)

‎cmake/thirdparty/get_raft.cmake

+51
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,51 @@
1+
# Copyright (c) Facebook, Inc. and its affiliates.
2+
#
3+
# This source code is licensed under the MIT license found in the
4+
# LICENSE file in the root directory of this source tree.
5+
# =============================================================================
6+
# Copyright (c) 2021-2023, NVIDIA CORPORATION.
7+
#
8+
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
9+
# in compliance with the License. You may obtain a copy of the License at
10+
#
11+
# http://www.apache.org/licenses/LICENSE-2.0
12+
#
13+
# Unless required by applicable law or agreed to in writing, software distributed under the License
14+
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
15+
# or implied. See the License for the specific language governing permissions and limitations under
16+
# the License.
17+
# =============================================================================
18+
19+
set(RAFT_VERSION "${RAPIDS_VERSION}")
20+
set(RAFT_FORK "rapidsai")
21+
set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION}")
22+
23+
function(find_and_configure_raft)
24+
set(oneValueArgs VERSION FORK PINNED_TAG)
25+
cmake_parse_arguments(PKG "${options}" "${oneValueArgs}"
26+
"${multiValueArgs}" ${ARGN} )
27+
28+
#-----------------------------------------------------
29+
# Invoke CPM find_package()
30+
#-----------------------------------------------------
31+
rapids_cpm_find(raft ${PKG_VERSION}
32+
GLOBAL_TARGETS raft::raft
33+
BUILD_EXPORT_SET faiss-exports
34+
INSTALL_EXPORT_SET faiss-exports
35+
CPM_ARGS
36+
GIT_REPOSITORY https://github.com/${PKG_FORK}/raft.git
37+
GIT_TAG ${PKG_PINNED_TAG}
38+
SOURCE_SUBDIR cpp
39+
OPTIONS
40+
"BUILD_TESTS OFF"
41+
"RAFT_COMPILE_LIBRARY OFF"
42+
)
43+
endfunction()
44+
45+
# Change pinned tag here to test a commit in CI
46+
# To use a different RAFT locally, set the CMake variable
47+
# CPM_raft_SOURCE=/path/to/local/raft
48+
find_and_configure_raft(VERSION ${RAFT_VERSION}.00
49+
FORK ${RAFT_FORK}
50+
PINNED_TAG ${RAFT_PINNED_TAG}
51+
)

‎faiss/gpu/CMakeLists.txt

+27-4
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,19 @@
33
#
44
# This source code is licensed under the BSD-style license found in the
55
# LICENSE file in the root directory of this source tree.
6+
# =============================================================================
7+
# Copyright (c) 2023, NVIDIA CORPORATION.
8+
#
9+
# Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except
10+
# in compliance with the License. You may obtain a copy of the License at
11+
#
12+
# http://www.apache.org/licenses/LICENSE-2.0
13+
#
14+
# Unless required by applicable law or agreed to in writing, software distributed under the License
15+
# is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express
16+
# or implied. See the License for the specific language governing permissions and limitations under
17+
# the License.
18+
# =============================================================================
619

720
set(FAISS_GPU_SRC
821
GpuAutoTune.cpp
@@ -163,6 +176,16 @@ set(FAISS_GPU_HEADERS
163176
utils/warpselect/WarpSelectImpl.cuh
164177
)
165178

179+
if(FAISS_ENABLE_RAFT)
180+
list(APPEND FAISS_GPU_HEADERS
181+
impl/RaftFlatIndex.cuh)
182+
list(APPEND FAISS_GPU_SRC
183+
impl/RaftFlatIndex.cu)
184+
185+
target_compile_definitions(faiss PUBLIC USE_NVIDIA_RAFT=1)
186+
target_compile_definitions(faiss_avx2 PUBLIC USE_NVIDIA_RAFT=1)
187+
endif()
188+
166189
# Export FAISS_GPU_HEADERS variable to parent scope.
167190
set(FAISS_GPU_HEADERS ${FAISS_GPU_HEADERS} PARENT_SCOPE)
168191

@@ -177,7 +200,7 @@ foreach(header ${FAISS_GPU_HEADERS})
177200
endforeach()
178201

179202
find_package(CUDAToolkit REQUIRED)
180-
target_link_libraries(faiss PRIVATE CUDA::cudart CUDA::cublas)
181-
target_link_libraries(faiss_avx2 PRIVATE CUDA::cudart CUDA::cublas)
182-
target_compile_options(faiss PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all>)
183-
target_compile_options(faiss_avx2 PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all>)
203+
target_link_libraries(faiss PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
204+
target_link_libraries(faiss_avx2 PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft> $<$<BOOL:${FAISS_ENABLE_RAFT}>:nvidia::cutlass::cutlass>)
205+
target_compile_options(faiss PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)
206+
target_compile_options(faiss_avx2 PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)

0 commit comments

Comments
 (0)