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" On)
+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,22 @@ include(cmake/san.cmake)
 # ----------------------------------------------------------------------------------------------
 project(VectorSimilarity)
 
+if (USE_CUDA)
+    # 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 "${CMAKE_CUDA_FLAGS} -O3 -std=c++17")
+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}")
 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,85 @@
+if(USE_CUDA)
+	# Set which version of RAPIDS to use
+	set(RAPIDS_VERSION 23.10)
+	# Set which version of RAFT to use (defined separately for testing
+	# minimal dependency changes if necessary)
+	set(RAFT_VERSION "${RAPIDS_VERSION}")
+    # TODO(wphicks): Reset to main fork and branch after
+    # https://github.com/rapidsai/raft/pull/1716 has been merged
+    # set(RAFT_FORK "rapidsai")
+    set(RAFT_FORK "wphicks")
+    # set(RAFT_PINNED_TAG "branch-${RAPIDS_VERSION}")
+    set(RAFT_PINNED_TAG "fea-resource_manager")
+
+	# 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

@@ -32,9 +32,15 @@ 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>
+)

src/VecSim/algorithms/brute_force/brute_force.h

@@ -35,6 +35,11 @@ class BruteForceIndex : public VecSimIndexAbstract<DistType> {
 public:
     BruteForceIndex(const BFParams *params, const AbstractIndexInitParams &abstractInitParams);
 
+    void clear() {
+        idToLabelMapping.clear();
+        vectorBlocks.clear();
+        count = idType{};
+    }
     size_t indexSize() const override;
     size_t indexCapacity() const override;
     vecsim_stl::vector<DistType> computeBlockScores(const DataBlock &block, const void *queryBlob,