Provide GPU-accelerated vector indexes with RAFT

CMakeLists.txt

@@ -1,4 +1,9 @@
-cmake_minimum_required(VERSION 3.10)
+option(USE_CUDA "Build Cuda code" OFF)
+if(USE_CUDA)
+  cmake_minimum_required(VERSION 3.26.4 FATAL_ERROR)
+else()
+  cmake_minimum_required(VERSION 3.10)
+endif()
 cmake_policy(SET CMP0077 NEW)
 
 set(CMAKE_CXX_STANDARD 20)
@@ -24,10 +29,30 @@ include(cmake/san.cmake)
 # ----------------------------------------------------------------------------------------------
 project(VectorSimilarity)
 
+if (USE_CUDA)
+	# List of architectures to generate device code
+	set(CMAKE_CUDA_ARCHITECTURES "native")
+	# Enable CUDA compilation for this project
+	enable_language(CUDA)
+	# Add definition for conditional compilation of CUDA components
+	add_definitions(-DUSE_CUDA)
+	# Perform all RAFT-specific CMake setup
+	include(cmake/raft.cmake)
+	# Required flags for compiling RAFT
+	set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} --expt-extended-lambda --expt-relaxed-constexpr")
+	set(CMAKE_CUDA_FLAGS_RELEASE "-O3")
+	set(CMAKE_CUDA_FLAGS_DEBUG "-g")
+	set(CMAKE_CUDA_STANDARD 17)
+	if(${CUDAToolkit_VERSION_MAJOR} GREATER 10)
+		# cuda11 support --threads for compile some large .cu more efficient
+		add_compile_options($<$<COMPILE_LANGUAGE:CUDA>:--threads=4>)
+	endif()
+endif()
+
 # Only do these if this is the main project, and not if it is included through add_subdirectory
 set_property(GLOBAL PROPERTY USE_FOLDERS ON)
 
-set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fexceptions -fPIC ${CLANG_SAN_FLAGS} ${LLVM_CXX_FLAGS} ${COV_CXX_FLAGS}")
+set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -fexceptions -fPIC -pthread ${CLANG_SAN_FLAGS} ${LLVM_CXX_FLAGS} ${COV_CXX_FLAGS} -lrt")
 set(CMAKE_SHARED_LINKER_FLAGS "${CMAKE_LINKER_FLAGS} ${LLVM_LD_FLAGS}")
 set(CMAKE_EXE_LINKER_FLAGS "${CMAKE_LINKER_FLAGS} ${LLVM_LD_FLAGS}")
 

check-format.sh

@@ -1,6 +1,6 @@
 #!/bin/bash
-CLANG_FMT_SRCS=$(find ./src/ \( -name '*.c' -o -name '*.cc' -o -name '*.cpp' -o -name '*.h' -o -name '*.hh' -o -name '*.hpp' \))
-CLANG_FMT_TESTS="$(find ./tests/ -type d \( -path ./tests/unit/build \) -prune -false -o  \( -name '*.c' -o -name '*.cc' -o -name '*.cpp' -o -name '*.h' -o -name '*.hh' -o -name '*.hpp' \))"
+CLANG_FMT_SRCS=$(find ./src/ \( -name '*.c' -o -name '*.cc' -o -name '*.cpp' -o -name '*.h' -o -name '*.hh' -o -name '*.hpp' -o -name '*.cuh' -o -name '*.cu' \))
+CLANG_FMT_TESTS="$(find ./tests/ -type d \( -path ./tests/unit/build \) -prune -false -o  \( -name '*.c' -o -name '*.cc' -o -name '*.cpp' -o -name '*.h' -o -name '*.hh' -o -name '*.hpp' -o -name '*.cuh' -o -name '*.cu' \))"
 
 E=0
 for filename in $CLANG_FMT_SRCS $CLANG_FMT_TESTS; do

cmake/raft.cmake

@@ -0,0 +1,81 @@
+if(USE_CUDA)
+	# Set which version of RAPIDS to use
+	set(RAPIDS_VERSION 23.12)
+	# Set which version of RAFT to use (defined separately for testing
+	# minimal dependency changes if necessary)
+	set(RAFT_VERSION "${RAPIDS_VERSION}")
+	set(RAFT_FORK "rapidsai")
+	set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION}")
+
+	# Download CMake file for bootstrapping RAPIDS-CMake, a utility that
+	# simplifies handling of complex RAPIDS dependencies
+	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}/RAPIDS.cmake)
+	endif()
+	include(${CMAKE_CURRENT_BINARY_DIR}/RAPIDS.cmake)
+
+	# General tool for orchestrating RAPIDS dependencies
+	include(rapids-cmake)
+	# CPM helper functions with dependency tracking
+	include(rapids-cpm)
+	rapids_cpm_init()
+	# Common CMake CUDA logic
+	include(rapids-cuda)
+	# Include required dependencies in Project-Config.cmake modules
+	# include(rapids-export)  TODO(wphicks)
+	# Functions to find system dependencies with dependency tracking
+	include(rapids-find)
+
+	# Correctly handle supported CUDA architectures
+	#    (From rapids-cuda)
+	rapids_cuda_init_architectures(VectorSimilarity)
+
+	# Find system CUDA toolkit
+	rapids_find_package(CUDAToolkit REQUIRED)
+
+	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 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()
+		#     (From rapids-cpm)
+		rapids_cpm_find(raft ${PKG_VERSION}
+			GLOBAL_TARGETS      raft::raft
+			BUILD_EXPORT_SET    VectorSimilarity-exports
+			INSTALL_EXPORT_SET  VectorSimilarity-exports
+			COMPONENTS          ${RAFT_COMPONENTS}
+			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_LIBRARY ${PKG_COMPILE_LIBRARY}"
+		)
+		if(raft_ADDED)
+			message(VERBOSE "VectorSimilarity: Using RAFT located in ${raft_SOURCE_DIR}")
+		else()
+			message(VERBOSE "VectorSimilarity: 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}
+		COMPILE_LIBRARY  OFF
+	)
+endif()

src/VecSim/CMakeLists.txt

@@ -19,6 +19,8 @@ add_library(VectorSimilarity ${VECSIM_LIBTYPE}
     index_factories/hnsw_factory.cpp
     index_factories/tiered_factory.cpp
     index_factories/index_factory.cpp
+    $<$<BOOL:${USE_CUDA}>:index_factories/raft_ivf_factory.cu>
+    $<$<BOOL:${USE_CUDA}>:index_factories/raft_ivf_tiered_factory.cpp>
     algorithms/hnsw/visited_nodes_handler.cpp
     vec_sim.cpp
     vec_sim_interface.cpp
@@ -31,9 +33,19 @@ add_library(VectorSimilarity ${VECSIM_LIBTYPE}
     ${HEADER_LIST}
 )
 
-target_link_libraries(VectorSimilarity VectorSimilaritySpaces)
 
 if(VECSIM_BUILD_TESTS)
     add_library(VectorSimilaritySerializer utils/serializer.cpp)
-    target_link_libraries(VectorSimilarity VectorSimilaritySerializer)
 endif()
+
+target_link_libraries(VectorSimilarity
+PUBLIC
+	VectorSimilaritySpaces
+	$<$<BOOL:${VECSIM_BUILD_TESTS}>:VectorSimilaritySerializer>
+PRIVATE
+	$<$<BOOL:${USE_CUDA}>:raft::raft>
+    $<$<BOOL:${USE_CUDA}>:CUDA::cusolver>
+    $<$<BOOL:${USE_CUDA}>:CUDA::cublas>
+    $<$<BOOL:${USE_CUDA}>:CUDA::curand>
+    $<$<BOOL:${USE_CUDA}>:CUDA::cusparse>
+)

src/VecSim/algorithms/brute_force/brute_force.h

@@ -35,6 +35,7 @@ class BruteForceIndex : public VecSimIndexAbstract<DistType> {
 public:
     BruteForceIndex(const BFParams *params, const AbstractIndexInitParams &abstractInitParams);
 
+    virtual void clear() = 0;
     size_t indexSize() const override;
     size_t indexCapacity() const override;
     vecsim_stl::vector<DistType> computeBlockScores(const DataBlock &block, const void *queryBlob,
@@ -54,6 +55,7 @@ class BruteForceIndex : public VecSimIndexAbstract<DistType> {
                                                   VecSimQueryParams *queryParams) const override;
     bool preferAdHocSearch(size_t subsetSize, size_t k, bool initial_check) const override;
     inline labelType getVectorLabel(idType id) const { return idToLabelMapping.at(id); }
+    inline vecsim_stl::vector<labelType> getLabels() const { return idToLabelMapping; }
 
     inline const vecsim_stl::vector<DataBlock> &getVectorBlocks() const { return vectorBlocks; }
     inline const labelType getLabelByInternalId(idType internal_id) const {

src/VecSim/algorithms/brute_force/brute_force_multi.h

@@ -23,6 +23,14 @@ class BruteForceIndex_Multi : public BruteForceIndex<DataType, DistType> {
 
     ~BruteForceIndex_Multi() {}
 
+    void clear() override {
+        this->labelToIdsLookup.clear();
+        this->idToLabelMapping.clear();
+        this->idToLabelMapping.shrink_to_fit();
+        this->vectorBlocks.clear();
+        this->vectorBlocks.shrink_to_fit();
+        this->count = idType{};
+    }
     int addVector(const void *vector_data, labelType label, void *auxiliaryCtx = nullptr) override;
     int deleteVector(labelType labelType) override;
     int deleteVectorById(labelType label, idType id) override;

src/VecSim/algorithms/brute_force/brute_force_single.h

@@ -21,6 +21,14 @@ class BruteForceIndex_Single : public BruteForceIndex<DataType, DistType> {
                            const AbstractIndexInitParams &abstractInitParams);
     ~BruteForceIndex_Single();
 
+    void clear() override {
+        this->labelToIdLookup.clear();
+        this->idToLabelMapping.clear();
+        this->idToLabelMapping.shrink_to_fit();
+        this->vectorBlocks.clear();
+        this->vectorBlocks.shrink_to_fit();
+        this->count = idType{};
+    }
     int addVector(const void *vector_data, labelType label, void *auxiliaryCtx = nullptr) override;
     int deleteVector(labelType label) override;
     int deleteVectorById(labelType label, idType id) override;