From 7a67021c415ebdfe1b7fecb900d1397c50a28c6d Mon Sep 17 00:00:00 2001 From: Mickael Ide Date: Tue, 11 Apr 2023 18:32:49 +0200 Subject: [PATCH 1/4] Adding mdspan calls to ivf raft Signed-off-by: Mickael Ide --- CMakeLists.txt | 2 +- cmake/libs/libraft.cmake | 26 +++++++++------- cmake/utils/fetch_rapids.cmake | 8 ++--- src/index/ivf_raft/ivf_raft.cuh | 54 ++++++++++++++++++--------------- 4 files changed, 50 insertions(+), 40 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 78dbc2866..64c0ffec4 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -146,7 +146,7 @@ list(APPEND KNOWHERE_LINKER_LIBS prometheus-cpp::core prometheus-cpp::pull prome add_library(knowhere SHARED ${KNOWHERE_SRCS}) add_dependencies(knowhere ${KNOWHERE_LINKER_LIBS}) if(WITH_RAFT) - list(APPEND KNOWHERE_LINKER_LIBS raft::raft) + list(APPEND KNOWHERE_LINKER_LIBS raft::raft raft::compiled) endif() target_link_libraries(knowhere PUBLIC ${KNOWHERE_LINKER_LIBS}) target_include_directories( diff --git a/cmake/libs/libraft.cmake b/cmake/libs/libraft.cmake index 75e11b5b0..e0f858ffe 100644 --- a/cmake/libs/libraft.cmake +++ b/cmake/libs/libraft.cmake @@ -32,10 +32,14 @@ set(RAFT_FORK "rapidsai") set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION}") function(find_and_configure_raft) - set(oneValueArgs VERSION FORK PINNED_TAG) + set(oneValueArgs VERSION FORK PINNED_TAG COMPILE_LIBRARY) cmake_parse_arguments(PKG "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN}) + set(RAFT_COMPONENTS "") + if(PKG_COMPILE_LIBRARY) + string(APPEND RAFT_COMPONENTS " compiled") + endif() # ----------------------------------------------------- # Invoke CPM find_package() # ----------------------------------------------------- @@ -44,12 +48,8 @@ function(find_and_configure_raft) ${PKG_VERSION} GLOBAL_TARGETS raft::raft - BUILD_EXPORT_SET - faiss-exports - INSTALL_EXPORT_SET - faiss-exports COMPONENTS - "distance nn" + ${RAFT_COMPONENTS} CPM_ARGS GIT_REPOSITORY https://github.com/${PKG_FORK}/raft.git @@ -60,13 +60,17 @@ function(find_and_configure_raft) 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") + "RAFT_COMPILE_LIBRARY ${PKG_COMPILE_LIBRARY}" + "RAFT_USE_FAISS_STATIC OFF") # Turn this on to build FAISS into your binary + + if(raft_ADDED) + message(VERBOSE "KNOWHERE: Using RAFT located in ${raft_SOURCE_DIR}") + else() + message(VERBOSE "KNOWHERE: Using RAFT located in ${raft_DIR}") + endif() 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}) + ${RAFT_PINNED_TAG} COMPILE_LIBRARY OFF) diff --git a/cmake/utils/fetch_rapids.cmake b/cmake/utils/fetch_rapids.cmake index e6f76f6e0..56899f2c5 100644 --- a/cmake/utils/fetch_rapids.cmake +++ b/cmake/utils/fetch_rapids.cmake @@ -13,12 +13,12 @@ # License for the specific language governing permissions and limitations under # the License. -set(RAPIDS_VERSION "23.02") +set(RAPIDS_VERSION "23.04") -if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) +if(NOT EXISTS ${CMAKE_CURRENT_BINARY_DIR}/RAPIDS.cmake) file( DOWNLOAD https://raw.githubusercontent.com/rapidsai/rapids-cmake/branch-${RAPIDS_VERSION}/RAPIDS.cmake - ${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) + ${CMAKE_CURRENT_BINARY_DIR}/RAPIDS.cmake) endif() -include(${CMAKE_CURRENT_BINARY_DIR}/FAISS_RAPIDS.cmake) +include(${CMAKE_CURRENT_BINARY_DIR}/RAPIDS.cmake) diff --git a/src/index/ivf_raft/ivf_raft.cuh b/src/index/ivf_raft/ivf_raft.cuh index 96712396e..e91bac917 100644 --- a/src/index/ivf_raft/ivf_raft.cuh +++ b/src/index/ivf_raft/ivf_raft.cuh @@ -36,6 +36,10 @@ #include "thrust/execution_policy.h" #include "thrust/sequence.h" +#ifdef RAFT_COMPILED +#include +#endif + namespace knowhere { namespace raft_res_pool { @@ -248,9 +252,9 @@ class RaftIvfIndexNode : public IndexNode { auto* data = reinterpret_cast(dataset.GetTensor()); auto stream = res_->get_stream(); - auto data_gpu = rmm::device_uvector(rows * dim, stream); - RAFT_CUDA_TRY(cudaMemcpyAsync(data_gpu.data(), data, data_gpu.size() * sizeof(float), cudaMemcpyDefault, - stream.value())); + auto data_gpu = raft::make_device_matrix(*res_, rows, dim); + RAFT_CUDA_TRY(cudaMemcpyAsync(data_gpu.data_handle(), data, data_gpu.size() * sizeof(float), + cudaMemcpyDefault, stream.value())); if constexpr (std::is_same_v) { auto build_params = raft::neighbors::ivf_flat::index_params{}; build_params.metric = metric.value(); @@ -259,7 +263,7 @@ class RaftIvfIndexNode : public IndexNode { build_params.kmeans_trainset_fraction = ivf_raft_cfg.kmeans_trainset_fraction; build_params.adaptive_centers = ivf_raft_cfg.adaptive_centers; gpu_index_ = raft::neighbors::ivf_flat::build(*res_, build_params, - data_gpu.data(), rows, dim); + data_gpu.view()); } else if constexpr (std::is_same_v) { auto build_params = raft::neighbors::ivf_pq::index_params{}; build_params.metric = metric.value(); @@ -276,7 +280,7 @@ class RaftIvfIndexNode : public IndexNode { build_params.codebook_kind = codebook_kind.value(); build_params.force_random_rotation = ivf_raft_cfg.force_random_rotation; gpu_index_ = raft::neighbors::ivf_pq::build(*res_, build_params, - data_gpu.data(), rows, dim); + data_gpu.view()); } else { static_assert(std::is_same_v); } @@ -312,19 +316,21 @@ class RaftIvfIndexNode : public IndexNode { auto stream = res_->get_stream(); // TODO(wphicks): Clean up transfer with raft // buffer objects when available - auto data_gpu = rmm::device_uvector(rows * dim, stream); - RAFT_CUDA_TRY(cudaMemcpyAsync(data_gpu.data(), data, data_gpu.size() * sizeof(float), cudaMemcpyDefault, - stream.value())); + auto data_gpu = raft::make_device_matrix(*res_, rows, dim); + RAFT_CUDA_TRY(cudaMemcpyAsync(data_gpu.data_handle(), data, data_gpu.size() * sizeof(float), + cudaMemcpyDefault, stream.value())); auto indices = rmm::device_uvector(rows, stream); thrust::sequence(thrust::device, indices.begin(), indices.end(), gpu_index_->size()); if constexpr (std::is_same_v) { - raft::neighbors::ivf_flat::extend(*res_, *gpu_index_, data_gpu.data(), - indices.data(), rows); + raft::neighbors::ivf_flat::extend(*res_, raft::make_const_mdspan(data_gpu.view()), + std::make_optional(raft::make_device_vector_view(indices.data(), rows)), + gpu_index_.value()); } else if constexpr (std::is_same_v) { - raft::neighbors::ivf_pq::extend(*res_, *gpu_index_, data_gpu.data(), - indices.data(), rows); + raft::neighbors::ivf_pq::extend(*res_, raft::make_const_mdspan(data_gpu.view()), + std::make_optional(raft::make_device_matrix_view(indices.data(), rows, 1)), + gpu_index_.value()); } else { static_assert(std::is_same_v); } @@ -356,19 +362,19 @@ class RaftIvfIndexNode : public IndexNode { auto stream = res_->get_stream(); // TODO(wphicks): Clean up transfer with raft // buffer objects when available - auto data_gpu = rmm::device_uvector(rows * dim, stream); - RAFT_CUDA_TRY(cudaMemcpyAsync(data_gpu.data(), data, data_gpu.size() * sizeof(float), cudaMemcpyDefault, - stream.value())); + auto data_gpu = raft::make_device_matrix(*res_, rows, dim); + RAFT_CUDA_TRY(cudaMemcpyAsync(data_gpu.data_handle(), data, data_gpu.size() * sizeof(float), + cudaMemcpyDefault, stream.value())); - auto ids_gpu = rmm::device_uvector(output_size, stream); - auto dis_gpu = rmm::device_uvector(output_size, stream); + auto ids_gpu = raft::make_device_matrix(*res_, rows, ivf_raft_cfg.k); + auto dis_gpu = raft::make_device_matrix(*res_, rows, ivf_raft_cfg.k); if constexpr (std::is_same_v) { auto search_params = raft::neighbors::ivf_flat::search_params{}; search_params.n_probes = ivf_raft_cfg.nprobe; raft::neighbors::ivf_flat::search(*res_, search_params, *gpu_index_, - data_gpu.data(), rows, ivf_raft_cfg.k, - ids_gpu.data(), dis_gpu.data()); + raft::make_const_mdspan(data_gpu.view()), + ids_gpu.view(), dis_gpu.view()); } else if constexpr (std::is_same_v) { auto search_params = raft::neighbors::ivf_pq::search_params{}; search_params.n_probes = ivf_raft_cfg.nprobe; @@ -396,15 +402,15 @@ class RaftIvfIndexNode : public IndexNode { } search_params.internal_distance_dtype = internal_distance_dtype.value(); search_params.preferred_shmem_carveout = search_params.preferred_shmem_carveout; - raft::neighbors::ivf_pq::search(*res_, search_params, *gpu_index_, data_gpu.data(), - rows, ivf_raft_cfg.k, ids_gpu.data(), - dis_gpu.data()); + raft::neighbors::ivf_pq::search(*res_, search_params, *gpu_index_, + raft::make_const_mdspan(data_gpu.view()), + ids_gpu.view(), dis_gpu.view()); } else { static_assert(std::is_same_v); } - RAFT_CUDA_TRY(cudaMemcpyAsync(ids.get(), ids_gpu.data(), ids_gpu.size() * sizeof(std::int64_t), + RAFT_CUDA_TRY(cudaMemcpyAsync(ids.get(), ids_gpu.data_handle(), ids_gpu.size() * sizeof(std::int64_t), cudaMemcpyDefault, stream.value())); - RAFT_CUDA_TRY(cudaMemcpyAsync(dis.get(), dis_gpu.data(), dis_gpu.size() * sizeof(float), cudaMemcpyDefault, + RAFT_CUDA_TRY(cudaMemcpyAsync(dis.get(), dis_gpu.data_handle(), dis_gpu.size() * sizeof(float), cudaMemcpyDefault, stream.value())); stream.synchronize(); } catch (std::exception& e) { From 439aee24b8ee87f9477207a25d76bf40dd763e60 Mon Sep 17 00:00:00 2001 From: Mickael Ide Date: Fri, 14 Apr 2023 15:45:41 +0200 Subject: [PATCH 2/4] Add pylibraft Signed-off-by: Mickael Ide --- CMakeLists.txt | 9 +++++++++ README.md | 4 ++-- conanfile.py | 6 ++++++ 3 files changed, 17 insertions(+), 2 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 64c0ffec4..4aa9acd92 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -29,6 +29,7 @@ knowhere_option(WITH_BENCHMARK "Build with benchmark" OFF) knowhere_option(WITH_COVERAGE "Build with coverage" OFF) knowhere_option(WITH_CCACHE "Build with ccache" ON) knowhere_option(WITH_PROFILER "Build with profiler" OFF) +knowhere_option(HINT_LIBRAFT "Hint directory for libraft" "") if(KNOWHERE_VERSION) message(STATUS "Building KNOWHERE version: ${KNOWHERE_VERSION}") @@ -147,6 +148,14 @@ add_library(knowhere SHARED ${KNOWHERE_SRCS}) add_dependencies(knowhere ${KNOWHERE_LINKER_LIBS}) if(WITH_RAFT) list(APPEND KNOWHERE_LINKER_LIBS raft::raft raft::compiled) + find_library(LIBRAFT_FOUND raft HINTS ${HINT_LIBRAFT}) + if (NOT LIBRAFT_FOUND) + message(WARNING "libraft not found") + else() + message(STATUS "libraft found") + list(APPEND KNOWHERE_LINKER_LIBS ${LIBRAFT_FOUND}) + add_definitions(-DRAFT_COMPILED) + endif() endif() target_link_libraries(knowhere PUBLIC ${KNOWHERE_LINKER_LIBS}) target_include_directories( diff --git a/README.md b/README.md index 51992f292..8b93486ef 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ Here's a list of verified OS types where Knowhere can successfully build and run ```bash $ sudo apt install build-essential libopenblas-dev libaio-dev python3-dev python3-pip -$ pip3 install conan==1.59.0 --user +$ pip3 install conan==1.59.0 pylibraft-cu11 --user --extra-index=https://pypi.nvidia.com $ export PATH=$PATH:$HOME/.local/bin ``` @@ -48,7 +48,7 @@ $ conan install .. --build=missing -o with_ut=True -o with_raft=True -s compiler #DISKANN SUPPORT $ conan install .. --build=missing -o with_ut=True -o with_diskann=True -s compiler.libcxx=libstdc++11 -s build_type=Debug/Release #build with conan -$conan build .. +$ conan build .. #verbose export VERBOSE=1 ``` diff --git a/conanfile.py b/conanfile.py index d66ecac11..d927df28a 100644 --- a/conanfile.py +++ b/conanfile.py @@ -132,6 +132,12 @@ def generate(self): if cxx_std_flag else "c++{}".format(self._minimum_cpp_standard) ) + if self.options.with_raft: + try: + import pylibraft + tc.variables["HINT_LIBRAFT"] = pylibraft.__path__[0] + except: + pass tc.variables["CXX_STD"] = cxx_std_value if is_msvc(self): tc.variables["MSVC_LANGUAGE_VERSION"] = cxx_std_value From 27c2fe13c21703a3f826477d9ade7be011ec0f11 Mon Sep 17 00:00:00 2001 From: Mickael Ide Date: Fri, 21 Apr 2023 16:14:07 +0200 Subject: [PATCH 3/4] Remove pylibraft Signed-off-by: Mickael Ide --- CMakeLists.txt | 3 +-- README.md | 2 +- conanfile.py | 6 ------ 3 files changed, 2 insertions(+), 9 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 05ba4254d..ab50f61e3 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -29,7 +29,6 @@ knowhere_option(WITH_BENCHMARK "Build with benchmark" OFF) knowhere_option(WITH_COVERAGE "Build with coverage" OFF) knowhere_option(WITH_CCACHE "Build with ccache" ON) knowhere_option(WITH_PROFILER "Build with profiler" OFF) -knowhere_option(HINT_LIBRAFT "Hint directory for libraft" "") if(KNOWHERE_VERSION) message(STATUS "Building KNOWHERE version: ${KNOWHERE_VERSION}") @@ -149,7 +148,7 @@ add_library(knowhere SHARED ${KNOWHERE_SRCS}) add_dependencies(knowhere ${KNOWHERE_LINKER_LIBS}) if(WITH_RAFT) list(APPEND KNOWHERE_LINKER_LIBS raft::raft raft::compiled) - find_library(LIBRAFT_FOUND raft HINTS ${HINT_LIBRAFT}) + find_library(LIBRAFT_FOUND raft) if (NOT LIBRAFT_FOUND) message(WARNING "libraft not found") else() diff --git a/README.md b/README.md index 8b93486ef..c6bd686c3 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ Here's a list of verified OS types where Knowhere can successfully build and run ```bash $ sudo apt install build-essential libopenblas-dev libaio-dev python3-dev python3-pip -$ pip3 install conan==1.59.0 pylibraft-cu11 --user --extra-index=https://pypi.nvidia.com +$ pip3 install conan==1.59.0 --user $ export PATH=$PATH:$HOME/.local/bin ``` diff --git a/conanfile.py b/conanfile.py index fb01356ba..67366c2ec 100644 --- a/conanfile.py +++ b/conanfile.py @@ -133,12 +133,6 @@ def generate(self): if cxx_std_flag else "c++{}".format(self._minimum_cpp_standard) ) - if self.options.with_raft: - try: - import pylibraft - tc.variables["HINT_LIBRAFT"] = pylibraft.__path__[0] - except: - pass tc.variables["CXX_STD"] = cxx_std_value if is_msvc(self): tc.variables["MSVC_LANGUAGE_VERSION"] = cxx_std_value From 47e020b5c9535213b2076bada627af0617c5dd80 Mon Sep 17 00:00:00 2001 From: Mickael Ide Date: Thu, 4 May 2023 14:48:14 +0200 Subject: [PATCH 4/4] Fix style Signed-off-by: Mickael Ide --- src/index/ivf_raft/ivf_raft.cuh | 28 ++++++++++++++++------------ 1 file changed, 16 insertions(+), 12 deletions(-) diff --git a/src/index/ivf_raft/ivf_raft.cuh b/src/index/ivf_raft/ivf_raft.cuh index da20d2f71..2213f5562 100644 --- a/src/index/ivf_raft/ivf_raft.cuh +++ b/src/index/ivf_raft/ivf_raft.cuh @@ -277,8 +277,8 @@ class RaftIvfIndexNode : public IndexNode { build_params.kmeans_n_iters = ivf_raft_cfg.kmeans_n_iters; build_params.kmeans_trainset_fraction = ivf_raft_cfg.kmeans_trainset_fraction; build_params.adaptive_centers = ivf_raft_cfg.adaptive_centers; - gpu_index_ = raft::neighbors::ivf_flat::build(*res_, build_params, - data_gpu.view()); + gpu_index_ = + raft::neighbors::ivf_flat::build(*res_, build_params, data_gpu.view()); } else if constexpr (std::is_same_v) { auto build_params = raft::neighbors::ivf_pq::index_params{}; build_params.metric = metric.value(); @@ -294,8 +294,8 @@ class RaftIvfIndexNode : public IndexNode { } build_params.codebook_kind = codebook_kind.value(); build_params.force_random_rotation = ivf_raft_cfg.force_random_rotation; - gpu_index_ = raft::neighbors::ivf_pq::build(*res_, build_params, - data_gpu.view()); + gpu_index_ = + raft::neighbors::ivf_pq::build(*res_, build_params, data_gpu.view()); } else { static_assert(std::is_same_v); } @@ -338,13 +338,17 @@ class RaftIvfIndexNode : public IndexNode { thrust::sequence(thrust::device, indices.begin(), indices.end(), gpu_index_->size()); if constexpr (std::is_same_v) { - raft::neighbors::ivf_flat::extend(*res_, raft::make_const_mdspan(data_gpu.view()), - std::make_optional(raft::make_device_vector_view(indices.data(), rows)), - gpu_index_.value()); + raft::neighbors::ivf_flat::extend( + *res_, raft::make_const_mdspan(data_gpu.view()), + std::make_optional( + raft::make_device_vector_view(indices.data(), rows)), + gpu_index_.value()); } else if constexpr (std::is_same_v) { - raft::neighbors::ivf_pq::extend(*res_, raft::make_const_mdspan(data_gpu.view()), - std::make_optional(raft::make_device_matrix_view(indices.data(), rows, 1)), - gpu_index_.value()); + raft::neighbors::ivf_pq::extend( + *res_, raft::make_const_mdspan(data_gpu.view()), + std::make_optional( + raft::make_device_matrix_view(indices.data(), rows, 1)), + gpu_index_.value()); } else { static_assert(std::is_same_v); } @@ -424,8 +428,8 @@ class RaftIvfIndexNode : public IndexNode { } RAFT_CUDA_TRY(cudaMemcpyAsync(ids.get(), ids_gpu.data_handle(), ids_gpu.size() * sizeof(std::int64_t), cudaMemcpyDefault, stream.value())); - RAFT_CUDA_TRY(cudaMemcpyAsync(dis.get(), dis_gpu.data_handle(), dis_gpu.size() * sizeof(float), cudaMemcpyDefault, - stream.value())); + RAFT_CUDA_TRY(cudaMemcpyAsync(dis.get(), dis_gpu.data_handle(), dis_gpu.size() * sizeof(float), + cudaMemcpyDefault, stream.value())); stream.synchronize(); } catch (std::exception& e) { LOG_KNOWHERE_WARNING_ << "RAFT inner error, " << e.what();