merge branch-24.04

.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]

conda/environments/all_cuda-118_arch-aarch64.yaml

@@ -44,8 +44,8 @@ dependencies:
 - nvcc_linux-aarch64=11.8
 - pre-commit
 - pydata-sphinx-theme
-- pytest
 - pytest-cov
+- pytest==7.*
 - rapids-dask-dependency==24.4.*
 - recommonmark
 - rmm==24.4.*

conda/environments/all_cuda-118_arch-x86_64.yaml

@@ -44,8 +44,8 @@ dependencies:
 - nvcc_linux-64=11.8
 - pre-commit
 - pydata-sphinx-theme
-- pytest
 - pytest-cov
+- pytest==7.*
 - rapids-dask-dependency==24.4.*
 - recommonmark
 - rmm==24.4.*

conda/environments/all_cuda-120_arch-aarch64.yaml

@@ -40,8 +40,8 @@ dependencies:
 - numpydoc
 - pre-commit
 - pydata-sphinx-theme
-- pytest
 - pytest-cov
+- pytest==7.*
 - rapids-dask-dependency==24.4.*
 - recommonmark
 - rmm==24.4.*

conda/environments/all_cuda-120_arch-x86_64.yaml

@@ -40,8 +40,8 @@ dependencies:
 - numpydoc
 - pre-commit
 - pydata-sphinx-theme
-- pytest
 - pytest-cov
+- pytest==7.*
 - rapids-dask-dependency==24.4.*
 - recommonmark
 - rmm==24.4.*

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,6 +196,10 @@ 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)
@@ -203,6 +208,9 @@ add_library(raft::raft ALIAS raft)
 target_include_directories(
   raft INTERFACE "$<BUILD_INTERFACE:${RAFT_SOURCE_DIR}/include>" "$<INSTALL_INTERFACE:include>"
 )
+if(BUILD_CAGRA_HNSWLIB)
+  target_link_libraries(raft INTERFACE hnswlib::hnswlib)
+endif()
 
 if(NOT BUILD_CPU_ONLY)
   # Keep RAFT as lightweight as possible. Only CUDA libs and rmm should be used in global target.
@@ -427,6 +435,7 @@ if(RAFT_COMPILE_LIBRARY)
     src/raft_runtime/neighbors/cagra_search.cu
     src/raft_runtime/neighbors/cagra_serialize.cu
     src/raft_runtime/neighbors/eps_neighborhood.cu
+    $<$<BOOL:${BUILD_CAGRA_HNSWLIB}>:src/raft_runtime/neighbors/hnsw.cpp>
     src/raft_runtime/neighbors/ivf_flat_build.cu
     src/raft_runtime/neighbors/ivf_flat_search.cu
     src/raft_runtime/neighbors/ivf_flat_serialize.cu

cpp/bench/ann/src/common/cuda_huge_page_resource.hpp

@@ -41,14 +41,6 @@ class cuda_huge_page_resource final : public rmm::mr::device_memory_resource {
   cuda_huge_page_resource& operator=(cuda_huge_page_resource const&) = default;
   cuda_huge_page_resource& operator=(cuda_huge_page_resource&&)      = default;
 
-  /**
-   * @brief Query whether the resource supports use of non-null CUDA streams for
-   * allocation/deallocation. `cuda_huge_page_resource` does not support streams.
-   *
-   * @returns bool false
-   */
-  [[nodiscard]] bool supports_streams() const noexcept override { return false; }
-
  private:
   /**
    * @brief Allocates memory of size at least `bytes` using cudaMalloc.

cpp/bench/ann/src/common/cuda_pinned_resource.hpp

@@ -45,14 +45,6 @@ class cuda_pinned_resource final : public rmm::mr::device_memory_resource {
   cuda_pinned_resource& operator=(cuda_pinned_resource const&) = default;
   cuda_pinned_resource& operator=(cuda_pinned_resource&&)      = default;
 
-  /**
-   * @brief Query whether the resource supports use of non-null CUDA streams for
-   * allocation/deallocation. `cuda_pinned_resource` does not support streams.
-   *
-   * @returns bool false
-   */
-  [[nodiscard]] bool supports_streams() const noexcept override { return false; }
-
  private:
   /**
    * @brief Allocates memory of size at least `bytes` using cudaMalloc.

cpp/bench/ann/src/hnswlib/hnswlib_wrapper.h

@@ -52,6 +52,11 @@ struct hnsw_dist_t<uint8_t> {
   using type = int;
 };
 
+template <>
+struct hnsw_dist_t<int8_t> {
+  using type = int;
+};
+
 template <typename T>
 class HnswLib : public ANN<T> {
  public:
@@ -135,7 +140,7 @@ void HnswLib<T>::build(const T* dataset, size_t nrow, cudaStream_t)
       space_ = std::make_shared<hnswlib::L2Space>(dim_);
     }
   } else if constexpr (std::is_same_v<T, uint8_t>) {
-    space_ = std::make_shared<hnswlib::L2SpaceI>(dim_);
+    space_ = std::make_shared<hnswlib::L2SpaceI<T>>(dim_);
   }
 
   appr_alg_ = std::make_shared<hnswlib::HierarchicalNSW<typename hnsw_dist_t<T>::type>>(
@@ -205,7 +210,7 @@ void HnswLib<T>::load(const std::string& path_to_index)
       space_ = std::make_shared<hnswlib::L2Space>(dim_);
     }
   } else if constexpr (std::is_same_v<T, uint8_t>) {
-    space_ = std::make_shared<hnswlib::L2SpaceI>(dim_);
+    space_ = std::make_shared<hnswlib::L2SpaceI<T>>(dim_);
   }
 
   appr_alg_ = std::make_shared<hnswlib::HierarchicalNSW<typename hnsw_dist_t<T>::type>>(

cpp/cmake/patches/hnswlib.diff

@@ -105,6 +105,63 @@
                          }
                      }
                  }
+diff --git a/hnswlib/space_l2.h b/hnswlib/space_l2.h
+index 4413537..c3240f3 100644
+--- a/hnswlib/space_l2.h
++++ b/hnswlib/space_l2.h
+@@ -252,13 +252,14 @@ namespace hnswlib {
+         ~L2Space() {}
+     };
+
++    template <typename T>
+     static int
+     L2SqrI4x(const void *__restrict pVect1, const void *__restrict pVect2, const void *__restrict qty_ptr) {
+
+         size_t qty = *((size_t *) qty_ptr);
+         int res = 0;
+-        unsigned char *a = (unsigned char *) pVect1;
+-        unsigned char *b = (unsigned char *) pVect2;
++        T *a = (T *) pVect1;
++        T *b = (T *) pVect2;
+
+         qty = qty >> 2;
+         for (size_t i = 0; i < qty; i++) {
+@@ -279,11 +280,12 @@ namespace hnswlib {
+         return (res);
+     }
+
++    template <typename T>
+     static int L2SqrI(const void* __restrict pVect1, const void* __restrict pVect2, const void* __restrict qty_ptr) {
+         size_t qty = *((size_t*)qty_ptr);
+         int res = 0;
+-        unsigned char* a = (unsigned char*)pVect1;
+-        unsigned char* b = (unsigned char*)pVect2;
++        T* a = (T*)pVect1;
++        T* b = (T*)pVect2;
+
+         for(size_t i = 0; i < qty; i++)
+         {
+@@ -294,6 +296,7 @@ namespace hnswlib {
+         return (res);
+     }
+
++    template <typename T>
+     class L2SpaceI : public SpaceInterface<int> {
+
+         DISTFUNC<int> fstdistfunc_;
+@@ -302,10 +305,10 @@ namespace hnswlib {
+     public:
+         L2SpaceI(size_t dim) {
+             if(dim % 4 == 0) {
+-                fstdistfunc_ = L2SqrI4x;
++                fstdistfunc_ = L2SqrI4x<T>;
+             }
+             else {
+-                fstdistfunc_ = L2SqrI;
++                fstdistfunc_ = L2SqrI<T>;
+             }
+             dim_ = dim;
+             data_size_ = dim * sizeof(unsigned char);
 diff --git a/hnswlib/visited_list_pool.h b/hnswlib/visited_list_pool.h
 index 5e1a4a5..4195ebd 100644
 --- a/hnswlib/visited_list_pool.h

cpp/cmake/thirdparty/get_hnswlib.cmake

@@ -30,6 +30,8 @@ function(find_and_configure_hnswlib)
   rapids_cpm_find(
     hnswlib ${PKG_VERSION}
     GLOBAL_TARGETS hnswlib::hnswlib
+    BUILD_EXPORT_SET raft-exports
+    INSTALL_EXPORT_SET raft-exports
     CPM_ARGS
     GIT_REPOSITORY ${PKG_REPOSITORY}
     GIT_TAG ${PKG_PINNED_TAG}
@@ -51,11 +53,13 @@ function(find_and_configure_hnswlib)
       # write export rules
       rapids_export(
         BUILD hnswlib
+        VERSION ${PKG_VERSION}
         EXPORT_SET hnswlib-exports
         GLOBAL_TARGETS hnswlib
         NAMESPACE hnswlib::)
       rapids_export(
         INSTALL hnswlib
+        VERSION ${PKG_VERSION}
         EXPORT_SET hnswlib-exports
         GLOBAL_TARGETS hnswlib
         NAMESPACE hnswlib::)
@@ -74,5 +78,5 @@ endif()
 find_and_configure_hnswlib(VERSION 0.6.2
         REPOSITORY       ${RAFT_HNSWLIB_GIT_REPOSITORY}
         PINNED_TAG       ${RAFT_HNSWLIB_GIT_TAG}
-        EXCLUDE_FROM_ALL ON
+        EXCLUDE_FROM_ALL OFF
         )

cpp/include/raft/neighbors/cagra_serialize.cuh

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2023, NVIDIA CORPORATION.
+ * Copyright (c) 2023-2024, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -32,13 +32,14 @@ namespace raft::neighbors::cagra {
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
+ * #include <raft/neighbors/cagra_serialize.hpp>
  *
  * raft::resources handle;
  *
  * // create an output stream
  * std::ostream os(std::cout.rdbuf());
- * // create an index with `auto index = cagra::build(...);`
- * raft::serialize(handle, os, index);
+ * // create an index with `auto index = raft::cagra::build(...);`
+ * raft::cagra::serialize(handle, os, index);
  * @endcode
  *
  * @tparam T data element type
@@ -66,13 +67,14 @@ void serialize(raft::resources const& handle,
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
+ * #include <raft/neighbors/cagra_serialize.hpp>
  *
  * raft::resources handle;
  *
  * // create a string with a filepath
  * std::string filename("/path/to/index");
- * // create an index with `auto index = cagra::build(...);`
- * raft::serialize(handle, filename, index);
+ * // create an index with `auto index = raft::cagra::build(...);`
+ * raft::cagra::serialize(handle, filename, index);
  * @endcode
  *
  * @tparam T data element type
@@ -100,13 +102,14 @@ void serialize(raft::resources const& handle,
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
+ * #include <raft/neighbors/cagra_serialize.hpp>
  *
  * raft::resources handle;
  *
  * // create an output stream
  * std::ostream os(std::cout.rdbuf());
- * // create an index with `auto index = cagra::build(...);`
- * raft::serialize_to_hnswlib(handle, os, index);
+ * // create an index with `auto index = raft::cagra::build(...);`
+ * raft::cagra::serialize_to_hnswlib(handle, os, index);
  * @endcode
  *
  * @tparam T data element type
@@ -120,25 +123,26 @@ void serialize(raft::resources const& handle,
 template <typename T, typename IdxT>
 void serialize_to_hnswlib(raft::resources const& handle,
                           std::ostream& os,
-                          const index<T, IdxT>& index)
+                          const raft::neighbors::cagra::index<T, IdxT>& index)
 {
   detail::serialize_to_hnswlib<T, IdxT>(handle, os, index);
 }
 
 /**
- * Write the CAGRA built index as a base layer HNSW index to file
+ * Save a CAGRA build index in hnswlib base-layer-only serialized format
  *
  * Experimental, both the API and the serialization format are subject to change.
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
+ * #include <raft/neighbors/cagra_serialize.hpp>
  *
  * raft::resources handle;
  *
  * // create a string with a filepath
  * std::string filename("/path/to/index");
- * // create an index with `auto index = cagra::build(...);`
- * raft::serialize_to_hnswlib(handle, filename, index);
+ * // create an index with `auto index = raft::cagra::build(...);`
+ * raft::cagra::serialize_to_hnswlib(handle, filename, index);
  * @endcode
  *
  * @tparam T data element type
@@ -152,7 +156,7 @@ void serialize_to_hnswlib(raft::resources const& handle,
 template <typename T, typename IdxT>
 void serialize_to_hnswlib(raft::resources const& handle,
                           const std::string& filename,
-                          const index<T, IdxT>& index)
+                          const raft::neighbors::cagra::index<T, IdxT>& index)
 {
   detail::serialize_to_hnswlib<T, IdxT>(handle, filename, index);
 }
@@ -164,14 +168,15 @@ void serialize_to_hnswlib(raft::resources const& handle,
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
+ * #include <raft/neighbors/cagra_serialize.hpp>
  *
  * raft::resources handle;
  *
  * // create an input stream
  * std::istream is(std::cin.rdbuf());
  * using T    = float; // data element type
  * using IdxT = int; // type of the index
- * auto index = raft::deserialize<T, IdxT>(handle, is);
+ * auto index = raft::cagra::deserialize<T, IdxT>(handle, is);
  * @endcode
  *
  * @tparam T data element type
@@ -195,14 +200,15 @@ index<T, IdxT> deserialize(raft::resources const& handle, std::istream& is)
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
+ * #include <raft/neighbors/cagra_serialize.hpp>
  *
  * raft::resources handle;
  *
  * // create a string with a filepath
  * std::string filename("/path/to/index");
  * using T    = float; // data element type
  * using IdxT = int; // type of the index
- * auto index = raft::deserialize<T, IdxT>(handle, filename);
+ * auto index = raft::cagra::deserialize<T, IdxT>(handle, filename);
  * @endcode
  *
  * @tparam T data element type