libraft and pylibraft API for CAGRA build and HNSW search

.pre-commit-config.yaml

@@ -27,7 +27,7 @@ repos:
                 types_or: [python, cython]
                 additional_dependencies: ["flake8-force"]
       - repo: https://github.com/pre-commit/mirrors-mypy
-        rev: 'v0.971'
+        rev: 'v1.3.0'
         hooks:
               - id: mypy
                 additional_dependencies: [types-cachetools]

cpp/CMakeLists.txt

@@ -57,6 +57,7 @@ option(BUILD_SHARED_LIBS "Build raft shared libraries" ON)
 option(BUILD_TESTS "Build raft unit-tests" ON)
 option(BUILD_PRIMS_BENCH "Build raft C++ benchmark tests" OFF)
 option(BUILD_ANN_BENCH "Build raft ann benchmarks" OFF)
+option(BUILD_CAGRA_HNSWLIB "Build CAGRA+hnswlib interface" ON)
 option(CUDA_ENABLE_KERNELINFO "Enable kernel resource usage info" OFF)
 option(CUDA_ENABLE_LINEINFO
        "Enable the -lineinfo option for nvcc (useful for cuda-memcheck / profiler)" OFF
@@ -195,13 +196,18 @@ if(BUILD_PRIMS_BENCH OR BUILD_ANN_BENCH)
   rapids_cpm_gbench()
 endif()
 
+if (BUILD_CAGRA_HNSWLIB)
+  include(cmake/thirdparty/get_hnswlib.cmake)
+endif()
+
 # ##################################################################################################
 # * raft ---------------------------------------------------------------------
 add_library(raft INTERFACE)
 add_library(raft::raft ALIAS raft)
 
 target_include_directories(
   raft INTERFACE "$<BUILD_INTERFACE:${RAFT_SOURCE_DIR}/include>" "$<INSTALL_INTERFACE:include>"
+  "$<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:${CMAKE_CURRENT_BINARY_DIR}/_deps/hnswlib-src/hnswlib>"
 )
 
 if(NOT BUILD_CPU_ONLY)
@@ -408,6 +414,7 @@ if(RAFT_COMPILE_LIBRARY)
     src/raft_runtime/matrix/select_k_float_int64_t.cu
     src/raft_runtime/neighbors/brute_force_knn_int64_t_float.cu
     src/raft_runtime/neighbors/cagra_build.cu
+    src/raft_runtime/neighbors/cagra_hnswlib.cpp
     src/raft_runtime/neighbors/cagra_search.cu
     src/raft_runtime/neighbors/cagra_serialize.cu
     src/raft_runtime/neighbors/ivf_flat_build.cu

cpp/cmake/thirdparty/get_hnswlib.cmake

@@ -27,11 +27,20 @@ function(find_and_configure_hnswlib)
                 WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/_deps )
 
         message("SOURCE ${CMAKE_CURRENT_SOURCE_DIR}")
+        message("WORKING DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}")
         execute_process (
                 COMMAND git apply ${CMAKE_CURRENT_SOURCE_DIR}/cmake/patches/hnswlib.patch
                 WORKING_DIRECTORY ${CMAKE_CURRENT_BINARY_DIR}/_deps/hnswlib-src
         )
     endif ()
+    # rapids_cpm_find(hnswlib ${PKG_VERSION}
+    # GLOBAL_TARGETS      hnswlib
+    # BUILD_EXPORT_SET    raft-exports
+    # INSTALL_EXPORT_SET  raft-exports
+    # CPM_ARGS
+    # GIT_REPOSITORY         https://github.com/${PKG_FORK}/hnswlib.git
+    # GIT_TAG                ${PKG_PINNED_TAG}
+    # EXCLUDE_FROM_ALL       ${PKG_EXCLUDE_FROM_ALL})
 
     include(cmake/modules/FindAVX.cmake)
 

cpp/include/raft/neighbors/cagra_hnswlib.hpp

@@ -0,0 +1,95 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "cagra_hnswlib_types.hpp"
+#include "detail/cagra_hnswlib.hpp"
+
+#include <cstddef>
+#include <raft/core/host_mdspan.hpp>
+#include <raft/core/resources.hpp>
+
+namespace raft::neighbors::cagra_hnswlib {
+
+/**
+ * @addtogroup cagra_hnswlib Build CAGRA index and search with hnswlib
+ * @{
+ */
+
+/**
+ * @brief Search hnswlib base layer only index constructed from a CAGRA index
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] res raft resources
+ * @param[in] params configure the search
+ * @param[in] idx cagra index
+ * @param[in] queries a host matrix view to a row-major matrix [n_queries, index->dim()]
+ * @param[out] neighbors a host matrix view to the indices of the neighbors in the source dataset
+ * [n_queries, k]
+ * @param[out] distances a host matrix view to the distances to the selected neighbors [n_queries,
+ * k]
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   // Build a CAGRA index
+ *   using namespace raft::neighbors;
+ *   // use default index parameters
+ *   cagra::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   auto index = cagra::build(res, index_params, dataset);
+ *
+ *   // Save CAGRA index as base layer only hnswlib index
+ *   cagra::serialize_to_hnswlib(res, "my_index.bin", index);
+ *
+ *   // Load CAGRA index as base layer only hnswlib index
+ *   cagra_hnswlib::index(D, "my_index.bin", raft::distance::L2Expanded);
+ *
+ *   // Search K nearest neighbors as an hnswlib index
+ *   // using host threads for concurrency
+ *   cagra_hnswlib::search_params search_params;
+ *   search_params.ef = 50 // ef >= K;
+ *   search_params.num_threads = 10;
+ *   auto neighbors = raft::make_host_matrix<uint32_t>(res, n_queries, k);
+ *   auto distances = raft::make_host_matrix<float>(res, n_queries, k);
+ *   cagra_hnswlib::search(res, search_params, index, queries, neighbors, distances);
+ * @endcode
+ */
+template <typename T>
+void search(raft::resources const& res,
+            const search_params& params,
+            const index<T>& idx,
+            raft::host_matrix_view<const T, int64_t, row_major> queries,
+            raft::host_matrix_view<uint64_t, int64_t, row_major> neighbors,
+            raft::host_matrix_view<float, int64_t, row_major> distances)
+{
+  RAFT_EXPECTS(
+    queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0),
+    "Number of rows in output neighbors and distances matrices must equal the number of queries.");
+
+  RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1),
+               "Number of columns in output neighbors and distances matrices must equal k");
+  RAFT_EXPECTS(queries.extent(1) == idx.dim(),
+               "Number of query dimensions should equal number of dimensions in the index.");
+
+  detail::search(res, params, idx, queries, neighbors, distances);
+}
+
+/**@}*/
+
+}  // namespace raft::neighbors::cagra_hnswlib

cpp/include/raft/neighbors/cagra_hnswlib_types.hpp

@@ -0,0 +1,70 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "ann_types.hpp"
+#include <memory>
+#include <raft/distance/distance_types.hpp>
+
+#include <cstdint>
+#include <sys/types.h>
+#include <type_traits>
+
+namespace raft::neighbors::cagra_hnswlib {
+
+/**
+ * @defgroup cagra_hnswlib Build CAGRA index and search with hnswlib
+ * @{
+ */
+
+struct search_params : ann::search_params {
+  int ef;               // size of the candidate list
+  int num_threads = 0;  // number of host threads to use for concurrent searches. Value of 0
+                        // automatically maximizes parallelism
+};
+
+template <typename T>
+struct index : ann::index {
+ public:
+  /**
+   * @brief load a base-layer-only hnswlib index originally saved from a built CAGRA index.
+   *  This is a virtual class and it cannot be used directly. To create an index, construct
+   *  an instance of `raft::neighbors::cagra_hnswlib::hnswlib_index` from the header
+   *  `raft/neighbores/hnswlib_types.hpp`
+   *
+   * @param[in] dim dimensions of the training dataset
+   * @param[in] metric distance metric to search. Supported metrics ("L2Expanded", "InnerProduct")
+   */
+  index(int dim, raft::distance::DistanceType metric) : dim_{dim}, metric_{metric} {}
+
+  /**
+  @brief Get underlying index
+  */
+  virtual auto get_index() const -> void const* = 0;
+
+  auto dim() const -> int const { return dim_; }
+
+  auto metric() const -> raft::distance::DistanceType { return metric_; }
+
+ private:
+  int dim_;
+  raft::distance::DistanceType metric_;
+};
+
+/**@}*/
+
+}  // namespace raft::neighbors::cagra_hnswlib

cpp/include/raft/neighbors/detail/cagra/cagra_serialize.cuh

@@ -103,6 +103,8 @@ void serialize_to_hnswlib(raft::resources const& res,
                           std::ostream& os,
                           const index<T, IdxT>& index_)
 {
+  // static_assert(std::is_same_v<IdxT, int> or std::is_same_v<IdxT, uint32_t>,
+  //               "An hnswlib index can only be trained with int32 or uint32 IdxT");
   common::nvtx::range<common::nvtx::domain::raft> fun_scope("cagra::serialize_to_hnswlib");
   RAFT_LOG_DEBUG("Saving CAGRA index to hnswlib format, size %zu, dim %u",
                  static_cast<size_t>(index_.size()),
@@ -120,14 +122,14 @@ void serialize_to_hnswlib(raft::resources const& res,
   // Example:M: 16, dim = 128, data_t = float, index_t = uint32_t, list_size_type = uint32_t,
   // labeltype: size_t size_data_per_element_ = M * 2 * sizeof(index_t) + sizeof(list_size_type) +
   // dim * sizeof(data_t) + sizeof(labeltype)
-  auto size_data_per_element =
-    static_cast<std::size_t>(index_.graph_degree() * 4 + 4 + index_.dim() * 4 + 8);
+  auto size_data_per_element = static_cast<std::size_t>(index_.graph_degree() * sizeof(IdxT) + 4 +
+                                                        index_.dim() * sizeof(T) + 8);
   os.write(reinterpret_cast<char*>(&size_data_per_element), sizeof(std::size_t));
   // label_offset
   std::size_t label_offset = size_data_per_element - 8;
   os.write(reinterpret_cast<char*>(&label_offset), sizeof(std::size_t));
   // offset_data
-  auto offset_data = static_cast<std::size_t>(index_.graph_degree() * 4 + 4);
+  auto offset_data = static_cast<std::size_t>(index_.graph_degree() * sizeof(IdxT) + 4);
   os.write(reinterpret_cast<char*>(&offset_data), sizeof(std::size_t));
   // max_level
   int max_level = 1;
@@ -184,17 +186,17 @@ void serialize_to_hnswlib(raft::resources const& res,
     }
 
     auto data_row = host_dataset.data_handle() + (index_.dim() * i);
-    if constexpr (std::is_same_v<T, float>) {
-      for (std::size_t j = 0; j < index_.dim(); ++j) {
-        auto data_elem = host_dataset(i, j);
-        os.write(reinterpret_cast<char*>(&data_elem), sizeof(T));
-      }
-    } else if constexpr (std::is_same_v<T, std::int8_t> or std::is_same_v<T, std::uint8_t>) {
-      for (std::size_t j = 0; j < index_.dim(); ++j) {
-        auto data_elem = static_cast<int>(host_dataset(i, j));
-        os.write(reinterpret_cast<char*>(&data_elem), sizeof(int));
-      }
+    // if constexpr (std::is_same_v<T, float>) {
+    for (std::size_t j = 0; j < index_.dim(); ++j) {
+      auto data_elem = host_dataset(i, j);
+      os.write(reinterpret_cast<char*>(&data_elem), sizeof(T));
     }
+    // } else if constexpr (std::is_same_v<T, std::int8_t> or std::is_same_v<T, std::uint8_t>) {
+    //   for (std::size_t j = 0; j < index_.dim(); ++j) {
+    //     auto data_elem = static_cast<int>(host_dataset(i, j));
+    //     os.write(reinterpret_cast<char*>(&data_elem), sizeof(int));
+    //   }
+    // }
 
     os.write(reinterpret_cast<char*>(&i), sizeof(std::size_t));
   }
@@ -204,7 +206,6 @@ void serialize_to_hnswlib(raft::resources const& res,
     auto zero = 0;
     os.write(reinterpret_cast<char*>(&zero), sizeof(int));
   }
-  // delete [] host_graph;
 }
 
 template <typename T, typename IdxT>

cpp/include/raft/neighbors/detail/cagra_hnswlib.hpp

@@ -0,0 +1,82 @@
+/*
+ * 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.
+ */
+
+#pragma once
+
+#include "../hnswlib_types.hpp"
+
+#include <cstdint>
+#include <raft/core/host_mdspan.hpp>
+#include <raft/core/resources.hpp>
+
+#include <omp.h>
+
+#include <hnswlib.h>
+
+namespace raft::neighbors::cagra_hnswlib::detail {
+
+template <typename T>
+void get_search_knn_results(hnswlib::HierarchicalNSW<typename hnsw_dist_t<T>::type> const* idx,
+                            const T* query,
+                            int k,
+                            uint64_t* indices,
+                            float* distances)
+{
+  auto result = idx->searchKnn(query, k);
+  assert(result.size() >= static_cast<size_t>(k));
+
+  for (int i = k - 1; i >= 0; --i) {
+    indices[i]   = result.top().second;
+    distances[i] = result.top().first;
+    result.pop();
+  }
+}
+
+template <typename T>
+void search(raft::resources const& res,
+            const search_params& params,
+            const index<T>& idx,
+            raft::host_matrix_view<const T, int64_t, row_major> queries,
+            raft::host_matrix_view<uint64_t, int64_t, row_major> neighbors,
+            raft::host_matrix_view<float, int64_t, row_major> distances)
+{
+  auto const* hnswlib_index =
+    reinterpret_cast<hnswlib::HierarchicalNSW<typename hnsw_dist_t<T>::type> const*>(
+      idx.get_index());
+
+  // when num_threads == 0, automatically maximize parallelism
+  if (params.num_threads) {
+#pragma omp parallel for num_threads(params.num_threads)
+    for (int64_t i = 0; i < queries.extent(0); ++i) {
+      get_search_knn_results(hnswlib_index,
+                             queries.data_handle() + i * queries.extent(1),
+                             neighbors.extent(1),
+                             neighbors.data_handle() + i * neighbors.extent(1),
+                             distances.data_handle() + i * distances.extent(1));
+    }
+  } else {
+#pragma omp parallel for
+    for (int64_t i = 0; i < queries.extent(0); ++i) {
+      get_search_knn_results(hnswlib_index,
+                             queries.data_handle() + i * queries.extent(1),
+                             neighbors.extent(1),
+                             neighbors.data_handle() + i * neighbors.extent(1),
+                             distances.data_handle() + i * distances.extent(1));
+    }
+  }
+}
+
+}  // namespace raft::neighbors::cagra_hnswlib::detail