This repository has been archived by the owner on Aug 16, 2023. It is now read-only.
-
Notifications
You must be signed in to change notification settings - Fork 84
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
Provide support for RAFT-based indexes (#712)
* Add RAFT to thirdparty libraries Signed-off-by: William Hicks <[email protected]> * Add RAFT stubs Signed-off-by: William Hicks <[email protected]> * Update CMake config Signed-off-by: William Hicks <[email protected]> * Try different CMake configuration for RAFT Signed-off-by: William Hicks <[email protected]> * Remove RAFT and RMM submodules Signed-off-by: William Hicks <[email protected]> * Correct various RAFT invocation details Signed-off-by: William Hicks <[email protected]> * Correct build and metric options Signed-off-by: William Hicks <[email protected]> * Split compilation into two units for parallelism Also correct indices passed during extend Signed-off-by: William Hicks <[email protected]> * Add error-handling for RAFT-internal errors Signed-off-by: William Hicks <[email protected]> * Correct scope for dataset variables Signed-off-by: William Hicks <[email protected]> * Reindent for consistency with existing code Signed-off-by: William Hicks <[email protected]> * Update CMakeList and adapt to changes in main Signed-off-by: William Hicks <[email protected]> * Remove dependencies provided transitively through RAFT Signed-off-by: William Hicks <[email protected]> * Ensure that build does not require RAFT Signed-off-by: William Hicks <[email protected]> * Use pool allocator for RAFT indexes Signed-off-by: William Hicks <[email protected]> * Pass through all RAFT parameters from config Signed-off-by: William Hicks <[email protected]> * Fix style Signed-off-by: William Hicks <[email protected]> --------- Signed-off-by: William Hicks <[email protected]>
- Loading branch information
Showing
13 changed files
with
875 additions
and
2 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,20 @@ | ||
| # ============================================================================= | ||
| # Copyright (c) 2023, NVIDIA CORPORATION. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except | ||
| # in compliance with the License. You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software distributed under the License | ||
| # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express | ||
| # or implied. See the License for the specific language governing permissions and limitations under | ||
| # the License. | ||
|
|
||
| set(RAPIDS_VERSION "23.02") | ||
|
|
||
| if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) | ||
| file(DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-${RAPIDS_VERSION}/RAPIDS.cmake | ||
| ${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) | ||
| endif() | ||
| include(${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,103 @@ | ||
| # ============================================================================= | ||
| # Copyright (c) 2021-2022, NVIDIA CORPORATION. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except | ||
| # in compliance with the License. You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software distributed under the License | ||
| # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express | ||
| # or implied. See the License for the specific language governing permissions and limitations under | ||
| # the License. | ||
| # ============================================================================= | ||
|
|
||
| function(find_and_configure_cutlass) | ||
| set(oneValueArgs VERSION REPOSITORY PINNED_TAG) | ||
| cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) | ||
|
|
||
| # if(RAFT_ENABLE_DIST_DEPENDENCIES OR RAFT_COMPILE_LIBRARIES) | ||
| set(CUTLASS_ENABLE_HEADERS_ONLY | ||
| ON | ||
| CACHE BOOL "Enable only the header library" | ||
| ) | ||
| set(CUTLASS_NAMESPACE | ||
| "raft_cutlass" | ||
| CACHE STRING "Top level namespace of CUTLASS" | ||
| ) | ||
| set(CUTLASS_ENABLE_CUBLAS | ||
| OFF | ||
| CACHE BOOL "Disable CUTLASS to build with cuBLAS library." | ||
| ) | ||
|
|
||
| if (CUDA_STATIC_RUNTIME) | ||
| set(CUDART_LIBRARY "${CUDA_cudart_static_LIBRARY}" CACHE FILEPATH "fixing cutlass cmake code" FORCE) | ||
| endif() | ||
|
|
||
| rapids_cpm_find( | ||
| NvidiaCutlass ${PKG_VERSION} | ||
| GLOBAL_TARGETS nvidia::cutlass::cutlass | ||
| CPM_ARGS | ||
| GIT_REPOSITORY ${PKG_REPOSITORY} | ||
| GIT_TAG ${PKG_PINNED_TAG} | ||
| GIT_SHALLOW TRUE | ||
| OPTIONS "CUDAToolkit_ROOT ${CUDAToolkit_LIBRARY_DIR}" | ||
| ) | ||
|
|
||
| if(TARGET CUTLASS AND NOT TARGET nvidia::cutlass::cutlass) | ||
| add_library(nvidia::cutlass::cutlass ALIAS CUTLASS) | ||
| endif() | ||
|
|
||
| if(NvidiaCutlass_ADDED) | ||
| rapids_export( | ||
| BUILD NvidiaCutlass | ||
| EXPORT_SET NvidiaCutlass | ||
| GLOBAL_TARGETS nvidia::cutlass::cutlass | ||
| NAMESPACE nvidia::cutlass:: | ||
| ) | ||
| endif() | ||
| # endif() | ||
|
|
||
| # We generate the cutlass-config files when we built cutlass locally, so always do | ||
| # `find_dependency` | ||
| rapids_export_package( | ||
| BUILD NvidiaCutlass raft-distance-exports GLOBAL_TARGETS nvidia::cutlass::cutlass | ||
| ) | ||
| rapids_export_package( | ||
| INSTALL NvidiaCutlass raft-distance-exports GLOBAL_TARGETS nvidia::cutlass::cutlass | ||
| ) | ||
| rapids_export_package( | ||
| BUILD NvidiaCutlass raft-nn-exports GLOBAL_TARGETS nvidia::cutlass::cutlass | ||
| ) | ||
| rapids_export_package( | ||
| INSTALL NvidiaCutlass raft-nn-exports GLOBAL_TARGETS nvidia::cutlass::cutlass | ||
| ) | ||
|
|
||
| # Tell cmake where it can find the generated NvidiaCutlass-config.cmake we wrote. | ||
| include("${rapids-cmake-dir}/export/find_package_root.cmake") | ||
| rapids_export_find_package_root( | ||
| INSTALL NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}/../]=] raft-distance-exports | ||
| ) | ||
| rapids_export_find_package_root( | ||
| BUILD NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}]=] raft-distance-exports | ||
| ) | ||
| include("${rapids-cmake-dir}/export/find_package_root.cmake") | ||
| rapids_export_find_package_root( | ||
| INSTALL NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}/../]=] raft-nn-exports | ||
| ) | ||
| rapids_export_find_package_root( | ||
| BUILD NvidiaCutlass [=[${CMAKE_CURRENT_LIST_DIR}]=] raft-nn-exports | ||
| ) | ||
| endfunction() | ||
|
|
||
| if(NOT RAFT_CUTLASS_GIT_TAG) | ||
| set(RAFT_CUTLASS_GIT_TAG v2.9.1) | ||
| endif() | ||
|
|
||
| if(NOT RAFT_CUTLASS_GIT_REPOSITORY) | ||
| set(RAFT_CUTLASS_GIT_REPOSITORY https://github.com/NVIDIA/cutlass.git) | ||
| endif() | ||
|
|
||
| find_and_configure_cutlass( | ||
| VERSION 2.9.1 REPOSITORY ${RAFT_CUTLASS_GIT_REPOSITORY} PINNED_TAG ${RAFT_CUTLASS_GIT_TAG} | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,51 @@ | ||
| # ============================================================================= | ||
| # Copyright (c) 2023, NVIDIA CORPORATION. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except | ||
| # in compliance with the License. You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software distributed under the License | ||
| # is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express | ||
| # or implied. See the License for the specific language governing permissions and limitations under | ||
| # the License. | ||
|
|
||
| set(RAFT_VERSION "${RAPIDS_VERSION}") | ||
| set(RAFT_FORK "rapidsai") | ||
| set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION}") | ||
|
|
||
| function(find_and_configure_raft) | ||
| set(oneValueArgs VERSION FORK PINNED_TAG) | ||
| cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" | ||
| "${multiValueArgs}" ${ARGN} ) | ||
|
|
||
| #----------------------------------------------------- | ||
| # Invoke CPM find_package() | ||
| #----------------------------------------------------- | ||
| rapids_cpm_find(raft ${PKG_VERSION} | ||
| GLOBAL_TARGETS raft::raft | ||
| BUILD_EXPORT_SET faiss-exports | ||
| INSTALL_EXPORT_SET faiss-exports | ||
| COMPONENTS "distance nn" | ||
| CPM_ARGS | ||
| GIT_REPOSITORY https://github.com/${PKG_FORK}/raft.git | ||
| GIT_TAG ${PKG_PINNED_TAG} | ||
| SOURCE_SUBDIR cpp | ||
| OPTIONS | ||
| "BUILD_TESTS OFF" | ||
| "BUILD_BENCH OFF" | ||
| "RAFT_COMPILE_LIBRARIES OFF" | ||
| "RAFT_COMPILE_NN_LIBRARY OFF" | ||
| "RAFT_USE_FAISS_STATIC OFF" # Turn this on to build FAISS into your binary | ||
| "RAFT_ENABLE_NN_DEPENDENCIES OFF" | ||
| ) | ||
| endfunction() | ||
|
|
||
| # Change pinned tag here to test a commit in CI | ||
| # To use a different RAFT locally, set the CMake variable | ||
| # CPM_raft_SOURCE=/path/to/local/raft | ||
| find_and_configure_raft(VERSION ${RAFT_VERSION}.00 | ||
| FORK ${RAFT_FORK} | ||
| PINNED_TAG ${RAFT_PINNED_TAG} | ||
| ) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,23 @@ | ||
| #============================================================================= | ||
| # Copyright (c) 2021-2022, NVIDIA CORPORATION. | ||
| # | ||
| # Licensed under the Apache License, Version 2.0 (the "License"); | ||
| # you may not use this file except in compliance with the License. | ||
| # You may obtain a copy of the License at | ||
| # | ||
| # http://www.apache.org/licenses/LICENSE-2.0 | ||
| # | ||
| # Unless required by applicable law or agreed to in writing, software | ||
| # distributed under the License is distributed on an "AS IS" BASIS, | ||
| # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| # See the License for the specific language governing permissions and | ||
| # limitations under the License. | ||
| #============================================================================= | ||
|
|
||
| function(find_and_configure_rmm) | ||
| include(${rapids-cmake-dir}/cpm/rmm.cmake) | ||
| rapids_cpm_rmm(BUILD_EXPORT_SET raft-exports | ||
| INSTALL_EXPORT_SET raft-exports) | ||
| endfunction() | ||
|
|
||
| find_and_configure_rmm() |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,48 @@ | ||
| /* | ||
| * Copyright (c) 2023, NVIDIA CORPORATION. | ||
| * | ||
| * Licensed under the Apache License, Version 2.0 (the "License"); | ||
| * you may not use this file except in compliance with the License. | ||
| * You may obtain a copy of the License at | ||
| * | ||
| * http://www.apache.org/licenses/LICENSE-2.0 | ||
| * | ||
| * Unless required by applicable law or agreed to in writing, software | ||
| * distributed under the License is distributed on an "AS IS" BASIS, | ||
| * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. | ||
| * See the License for the specific language governing permissions and | ||
| * limitations under the License. | ||
| */ | ||
|
|
||
| #ifndef COMMON_RAFT_METRIC_H | ||
| #define COMMON_RAFT_METRIC_H | ||
|
|
||
| #include <algorithm> | ||
| #include <string> | ||
| #include <unordered_map> | ||
|
|
||
| #include "knowhere/comp/index_param.h" | ||
| #include "knowhere/expected.h" | ||
| #include "raft/distance/distance_types.hpp" | ||
|
|
||
| namespace knowhere { | ||
|
|
||
| inline expected<raft::distance::DistanceType, Status> | ||
| Str2RaftMetricType(std::string metric) { | ||
| static const std::unordered_map<std::string, raft::distance::DistanceType> metric_map = { | ||
| {metric::L2, raft::distance::DistanceType::L2Expanded}, | ||
| {metric::IP, raft::distance::DistanceType::InnerProduct}, | ||
| {metric::HAMMING, raft::distance::DistanceType::HammingUnexpanded}, | ||
| {metric::JACCARD, raft::distance::DistanceType::JaccardExpanded}, | ||
| }; | ||
|
|
||
| std::transform(metric.begin(), metric.end(), metric.begin(), toupper); | ||
| auto it = metric_map.find(metric); | ||
| if (it == metric_map.end()) | ||
| return unexpected(Status::invalid_metric_type); | ||
| return it->second; | ||
| } | ||
|
|
||
| } // namespace knowhere | ||
|
|
||
| #endif /* RAFT_METRIC_H */ |
Oops, something went wrong.