Integrate RAFT FlatIndex / Distances

CMakeLists.txt

@@ -4,22 +4,44 @@
 # This source code is licensed under the BSD-style license found in the
 # LICENSE file in the root directory of this source tree.
 
-cmake_minimum_required(VERSION 3.17 FATAL_ERROR)
+cmake_minimum_required(VERSION 3.23.1 FATAL_ERROR)
+
+set(FAISS_LANGUAGES CXX)
+
+if(FAISS_ENABLE_GPU)
+  list(APPEND FAISS_LANGUAGES CUDA)
+endif()
+
+if(FAISS_ENABLE_RAFT)
+include(cmake/thirdparty/fetch_rapids.cmake)
+include(rapids-cmake)
+include(rapids-cpm)
+include(rapids-cuda)
+include(rapids-export)
+include(rapids-find)
+
+rapids_cuda_init_architectures(faiss)
+endif()
 
 project(faiss
   VERSION 1.7.3
   DESCRIPTION "A library for efficient similarity search and clustering of dense vectors."
   HOMEPAGE_URL "https://github.com/facebookresearch/faiss"
-  LANGUAGES CXX)
+  LANGUAGES ${FAISS_LANGUAGES})
 include(GNUInstallDirs)
 
+if(FAISS_ENABLE_RAFT)
+set(CMAKE_CXX_STANDARD 17)
+else()
 set(CMAKE_CXX_STANDARD 11)
+endif()
 
 list(APPEND CMAKE_MODULE_PATH "${PROJECT_SOURCE_DIR}/cmake")
 
 # Valid values are "generic", "avx2".
 option(FAISS_OPT_LEVEL "" "generic")
 option(FAISS_ENABLE_GPU "Enable support for GPU indexes." ON)
+option(FAISS_ENABLE_RAFT "Enable RAFT for GPU indexes." OFF)
 option(FAISS_ENABLE_PYTHON "Build Python extension." ON)
 option(FAISS_ENABLE_C_API "Build C API." OFF)
 
@@ -28,6 +50,11 @@ if(FAISS_ENABLE_GPU)
   enable_language(CUDA)
 endif()
 
+if(FAISS_ENABLE_RAFT)
+  rapids_cpm_init()
+  include(cmake/thirdparty/get_raft.cmake)
+endif()
+
 add_subdirectory(faiss)
 
 if(FAISS_ENABLE_GPU)

INSTALL.md

@@ -105,6 +105,9 @@ Several options can be passed to CMake, among which:
   values are `ON` and `OFF`),
   - `-DFAISS_ENABLE_PYTHON=OFF` in order to disable building python bindings
   (possible values are `ON` and `OFF`),
+  - `-DFAISS_ENABLE_RAFT=ON` in order to enable building the RAFT implementations
+    of the IVF-Flat and IVF-PQ GPU-accelerated indices (default is `OFF`, possible 
+    values are `ON` and `OFF`)
   - `-DBUILD_TESTING=OFF` in order to disable building C++ tests,
   - `-DBUILD_SHARED_LIBS=ON` in order to build a shared library (possible values
   are `ON` and `OFF`),

build.sh

@@ -0,0 +1,52 @@
+#!/bin/bash
+
+# NOTE: This file is temporary for the proof-of-concept branch and will be removed before this PR is merged
+
+BUILD_TYPE=Release
+BUILD_DIR=build/
+
+RAFT_REPO_REL=""
+EXTRA_CMAKE_ARGS=""
+set -e
+
+if [[ ${RAFT_REPO_REL} != "" ]]; then
+  RAFT_REPO_PATH="`readlink -f \"${RAFT_REPO_REL}\"`"
+  EXTRA_CMAKE_ARGS="${EXTRA_CMAKE_ARGS} -DCPM_raft_SOURCE=${RAFT_REPO_PATH}"
+fi
+
+if [ "$1" == "clean" ]; then
+  rm -rf build
+  exit 0
+fi
+
+if [ "$1" == "test" ]; then
+  make -C build -j test
+  exit 0
+fi
+
+if [ "$1" == "test-raft" ]; then
+  ./build/faiss/gpu/test/TestRaftIndexIVFFlat
+  exit 0
+fi
+
+mkdir -p $BUILD_DIR
+cd $BUILD_DIR
+
+cmake \
+ -DFAISS_ENABLE_GPU=ON \
+ -DFAISS_ENABLE_RAFT=ON \
+ -DFAISS_ENABLE_PYTHON=OFF \
+ -DBUILD_TESTING=ON \
+ -DBUILD_SHARED_LIBS=OFF \
+ -DCMAKE_BUILD_TYPE=${BUILD_TYPE} \
+ -DFAISS_OPT_LEVEL=avx2 \
+ -DRAFT_NVTX=OFF \
+ -DCMAKE_CUDA_ARCHITECTURES="NATIVE" \
+ -DCMAKE_EXPORT_COMPILE_COMMANDS=ON \
+ -DCMAKE_CUDA_COMPILER_LAUNCHER=ccache \
+ -DCMAKE_C_COMPILER_LAUNCHER=ccache \
+ -DCMAKE_CXX_COMPILER_LAUNCHER=ccache \
+ ${EXTRA_CMAKE_ARGS} \
+ ../
+
+cmake  --build . -j12

cmake/thirdparty/fetch_rapids.cmake

@@ -0,0 +1,12 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+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)

cmake/thirdparty/get_raft.cmake

@@ -0,0 +1,40 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+
+set(RAFT_VERSION "${RAPIDS_VERSION}")
+set(RAFT_FORK "benfred")
+set(RAFT_PINNED_TAG "bfknn")#"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
+            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_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}
+        )

faiss/gpu/CMakeLists.txt

@@ -164,6 +164,16 @@ set(FAISS_GPU_HEADERS
   utils/warpselect/WarpSelectImpl.cuh
 )
 
+if(FAISS_ENABLE_RAFT)
+  list(APPEND FAISS_GPU_HEADERS
+          impl/RaftFlatIndex.cuh)
+  list(APPEND FAISS_GPU_SRC
+          impl/RaftFlatIndex.cu)
+
+  target_compile_definitions(faiss PUBLIC USE_NVIDIA_RAFT=1)
+  target_compile_definitions(faiss_avx2 PUBLIC USE_NVIDIA_RAFT=1)
+endif()
+
 # Export FAISS_GPU_HEADERS variable to parent scope.
 set(FAISS_GPU_HEADERS ${FAISS_GPU_HEADERS} PARENT_SCOPE)
 
@@ -178,7 +188,7 @@ foreach(header ${FAISS_GPU_HEADERS})
 endforeach()
 
 find_package(CUDAToolkit REQUIRED)
-target_link_libraries(faiss PRIVATE CUDA::cudart CUDA::cublas)
-target_link_libraries(faiss_avx2 PRIVATE CUDA::cudart CUDA::cublas)
-target_compile_options(faiss PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all>)
-target_compile_options(faiss_avx2 PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all>)
+target_link_libraries(faiss PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft>)
+target_link_libraries(faiss_avx2 PRIVATE CUDA::cudart CUDA::cublas $<$<BOOL:${FAISS_ENABLE_RAFT}>:raft::raft>)
+target_compile_options(faiss PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)
+target_compile_options(faiss_avx2 PRIVATE $<$<COMPILE_LANGUAGE:CUDA>:-Xfatbin=-compress-all --expt-extended-lambda --expt-relaxed-constexpr>)

faiss/gpu/GpuDistance.cu

@@ -14,9 +14,24 @@
 #include <faiss/gpu/utils/CopyUtils.cuh>
 #include <faiss/gpu/utils/DeviceTensor.cuh>
 
+#if defined USE_NVIDIA_RAFT
+#include <faiss/gpu/impl/RaftUtils.h>
+#include <raft/core/device_resources.hpp>
+#include <raft/core/device_mdspan.hpp>
+#include <raft/neighbors/brute_force.cuh>
+#include <raft/core/error.hpp>
+#include <raft/core/mdspan_types.hpp>
+#define RAFT_NAME "raft"
+#endif
+
 namespace faiss {
 namespace gpu {
 
+#if defined USE_NVIDIA_RAFT
+using namespace raft::distance;
+using namespace raft::neighbors;
+#endif
+
 template <typename T>
 void bfKnnConvert(GpuResourcesProvider* prov, const GpuDistanceParams& args) {
     // Validate the input data
@@ -165,7 +180,6 @@ void bfKnnConvert(GpuResourcesProvider* prov, const GpuDistanceParams& args) {
                 tOutDistances,
                 tIntIndices,
                 args.ignoreOutDistances);
-
         // Convert and copy int indices out
         auto tOutIntIndices = toDeviceTemporary<int, 2>(
                 res,
@@ -193,13 +207,77 @@ void bfKnn(GpuResourcesProvider* res, const GpuDistanceParams& args) {
             "limitation: both vectorType and queryType must currently "
             "be the same (F32 or F16");
 
-    if (args.vectorType == DistanceDataType::F32) {
-        bfKnnConvert<float>(res, args);
-    } else if (args.vectorType == DistanceDataType::F16) {
-        bfKnnConvert<half>(res, args);
+#if defined USE_NVIDIA_RAFT
+    // Note: For now, RAFT bfknn requires queries and vectors to be same layout
+    if(args.should_use_raft() && args.queriesRowMajor == args.vectorsRowMajor) {
+
+        DistanceType distance = faiss_to_raft(args.metric, false);
+
+        auto resImpl = res->getResources();
+        auto res_impl = resImpl.get();
+        raft::handle_t& handle = res_impl->getRaftHandleCurrentDevice();
+
+        idx_t dims = args.dims;
+        idx_t num_vectors = args.numVectors;
+        const float *vectors = reinterpret_cast<const float*>(args.vectors);
+        const float *queries = reinterpret_cast<const float*>(args.queries);
+        idx_t num_queries = args.numQueries;
+        int k = args.k;
+        float metric_arg = args.metricArg;
+        idx_t *out_indices = reinterpret_cast<idx_t*>(args.outIndices);
+        float *out_distances = reinterpret_cast<float*>(args.outDistances);
+
+        //args.queriesRowMajor ? raft::row_major : raft::col_major;
+
+
+        auto inds = raft::make_device_matrix_view<idx_t, idx_t>(
+                out_indices, num_queries, k);
+        auto dists = raft::make_device_matrix_view<float, idx_t>(
+                out_distances,
+                num_queries,
+                k);
+
+
+        if(args.queriesRowMajor) {
+            auto index = raft::make_device_matrix_view<const float, idx_t, raft::row_major>(vectors, num_vectors, dims);
+            auto search = raft::make_device_matrix_view<const float, idx_t, raft::row_major>(queries, num_queries, dims);
+
+            std::vector<raft::device_matrix_view<const float, idx_t, raft::row_major>> index_vec = {index};
+            // For now, use RAFT's fused KNN when k <= 64 and L2 metric is used
+            if (args.k <= 64 && args.metric == MetricType::METRIC_L2 && args.numVectors > 0) {
+                RAFT_LOG_INFO("Invoking flat fused_l2_knn");
+                brute_force::fused_l2_knn(handle, index, search, inds, dists, distance);
+            } else {
+                RAFT_LOG_INFO("Invoking flat bfknn");
+                brute_force::knn(handle, index_vec, search, inds, dists, k, distance, metric_arg);
+            }
+        } else {
+            auto index = raft::make_device_matrix_view<const float, idx_t, raft::col_major>(vectors, num_vectors, dims);
+            auto search = raft::make_device_matrix_view<const float, idx_t, raft::col_major>(queries, num_queries, dims);
+
+            std::vector<raft::device_matrix_view<const float, idx_t, raft::col_major>> index_vec = {index};
+            // For now, use RAFT's fused KNN when k <= 64 and L2 metric is used
+            if (args.k <= 64 && args.metric == MetricType::METRIC_L2 && args.numVectors > 0) {
+                RAFT_LOG_INFO("Invoking flat fused_l2_knn");
+                brute_force::fused_l2_knn(handle, index, search, inds, dists, distance);
+            } else {
+                RAFT_LOG_INFO("Invoking flat bfknn");
+                brute_force::knn(handle, index_vec, search, inds, dists, k, distance, metric_arg);
+            }
+        }
     } else {
-        FAISS_THROW_MSG("unknown vectorType");
+#else
+        if (args.vectorType == DistanceDataType::F32) {
+            bfKnnConvert<float>(res, args);
+        } else if (args.vectorType == DistanceDataType::F16) {
+            bfKnnConvert<half>(res, args);
+        } else {
+            FAISS_THROW_MSG("unknown vectorType");
+        }
+#endif
+#if defined USE_NVIDIA_RAFT
     }
+#endif
 }
 
 // legacy version

faiss/gpu/GpuDistance.h

@@ -124,6 +124,18 @@ struct GpuDistanceParams {
     /// Otherwise, an integer 0 <= device < numDevices indicates the device for
     /// execution
     int device;
+
+    /// Should the index dispatch down to RAFT?
+    bool use_raft = false;
+
+    bool should_use_raft() const {
+#if defined USE_NVIDIA_RAFT
+        return use_raft;
+#else
+        FAISS_THROW_IF_NOT_MSG(!use_raft, "RAFT has not been compiled into the current version so it cannot be used.");
+#endif
+    }
+
 };
 
 /// A wrapper for gpu/impl/Distance.cuh to expose direct brute-force k-nearest

faiss/gpu/GpuIndex.h

@@ -23,6 +23,18 @@ struct GpuIndexConfig {
     /// On Pascal and above (CC 6+) architectures, allows GPUs to use
     /// more memory than is available on the GPU.
     MemorySpace memorySpace;
+
+    /// Should the index dispatch down to RAFT?
+    bool use_raft = false;
+
+    bool should_use_raft() {
+#if defined USE_NVIDIA_RAFT
+        return use_raft;
+#else
+        FAISS_THROW_IF_NOT_MSG(!use_raft, "RAFT has not been compiled into the current version so it cannot be used.");
+#endif
+    }
+
 };
 
 class GpuIndex : public faiss::Index {