Python API for CAGRA+HNSW

cpp/CMakeLists.txt

@@ -582,6 +582,7 @@ if(BUILD_C_LIBRARY)
     src/neighbors/ivf_flat_c.cpp
     src/neighbors/ivf_pq_c.cpp
     src/neighbors/cagra_c.cpp
+    src/neighbors/hnsw_c.cpp
     src/neighbors/refine/refine_c.cpp
     src/distance/pairwise_distance_c.cpp
   )

cpp/include/cuvs/neighbors/cagra.h

@@ -267,6 +267,15 @@ cuvsError_t cuvsCagraIndexCreate(cuvsCagraIndex_t* index);
  */
 cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index);
 
+/**
+ * @brief Get dimension of the CAGRA index
+ *
+ * @param[in] index CAGRA index
+ * @param[out] dim return dimension of the index
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int* dim);
+
 /**
  * @}
  */
@@ -337,7 +346,10 @@ cuvsError_t cuvsCagraBuild(cuvsResources_t res,
  *        It is also important to note that the CAGRA Index must have been built
  *        with the same type of `queries`, such that `index.dtype.code ==
  * queries.dl_tensor.dtype.code` Types for input are:
- *        1. `queries`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *        1. `queries`:
+ *          a. kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *          b. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`
+ *          c. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`
  *        2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32`
  *        3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
  *
@@ -394,7 +406,7 @@ cuvsError_t cuvsCagraSearch(cuvsResources_t res,
  *
  * Experimental, both the API and the serialization format are subject to change.
  *
- * @code{.cpp}
+ * @code{.c}
  * #include <cuvs/neighbors/cagra.h>
  *
  * // Create cuvsResources_t
@@ -416,6 +428,32 @@ cuvsError_t cuvsCagraSerialize(cuvsResources_t res,
                                cuvsCagraIndex_t index,
                                bool include_dataset);
 
+/**
+ * Save the CAGRA index to file in hnswlib format.
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @code{.c}
+ * #include <cuvs/core/c_api.h>
+ * #include <cuvs/neighbors/cagra.h>
+ *
+ * // Create cuvsResources_t
+ * cuvsResources_t res;
+ * cuvsError_t res_create_status = cuvsResourcesCreate(&res);
+ *
+ * // create an index with `cuvsCagraBuild`
+ * cuvsCagraSerializeHnswlib(res, "/path/to/index", index);
+ * @endcode
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] filename the file name for saving the index
+ * @param[in] index CAGRA index
+ *
+ */
+cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res,
+                                        const char* filename,
+                                        cuvsCagraIndex_t index);
+
 /**
  * Load index from file.
  *

cpp/include/cuvs/neighbors/hnsw.h

@@ -0,0 +1,206 @@
+/*
+ * Copyright (c) 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.
+ * 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 <cuvs/core/c_api.h>
+#include <cuvs/distance/distance.h>
+#include <dlpack/dlpack.h>
+#include <stdbool.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @defgroup hnsw_c_search_params C API for hnswlib wrapper search params
+ * @{
+ */
+
+struct cuvsHnswSearchParams {
+  int32_t ef;
+  int32_t num_threads;
+};
+
+typedef struct cuvsHnswSearchParams* cuvsHnswSearchParams_t;
+
+/**
+ * @brief Allocate HNSW search params, and populate with default values
+ *
+ * @param[in] params cuvsHnswSearchParams_t to allocate
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsHnswSearchParamsCreate(cuvsHnswSearchParams_t* params);
+
+/**
+ * @brief De-allocate HNSW search params
+ *
+ * @param[in] params cuvsHnswSearchParams_t to de-allocate
+ * @return cuvsError_t
+ */
+cuvsError_t cuvsHnswSearchParamsDestroy(cuvsHnswSearchParams_t params);
+
+/**
+ * @}
+ */
+
+/**
+ * @defgroup hnsw_c_index C API for hnswlib wrapper index
+ * @{
+ */
+
+/**
+ * @brief Struct to hold address of cuvs::neighbors::Hnsw::index and its active trained dtype
+ *
+ */
+typedef struct {
+  uintptr_t addr;
+  DLDataType dtype;
+
+} cuvsHnswIndex;
+
+typedef cuvsHnswIndex* cuvsHnswIndex_t;
+
+/**
+ * @brief Allocate HNSW index
+ *
+ * @param[in] index cuvsHnswIndex_t to allocate
+ * @return HnswError_t
+ */
+cuvsError_t cuvsHnswIndexCreate(cuvsHnswIndex_t* index);
+
+/**
+ * @brief De-allocate HNSW index
+ *
+ * @param[in] index cuvsHnswIndex_t to de-allocate
+ */
+cuvsError_t cuvsHnswIndexDestroy(cuvsHnswIndex_t index);
+
+/**
+ * @}
+ */
+
+/**
+ * @defgroup hnsw_c_index_search C API for CUDA ANN Graph-based nearest neighbor search
+ * @{
+ */
+/**
+ * @brief Search a HNSW index with a `DLManagedTensor` which has underlying
+ *        `DLDeviceType` equal to `kDLCPU`, `kDLCUDAHost`, or `kDLCUDAManaged`.
+ *        It is also important to note that the HNSW Index must have been built
+ *        with the same type of `queries`, such that `index.dtype.code ==
+ *        queries.dl_tensor.dtype.code`
+ *        Supported types for input are:
+ *        1. `queries`:
+ *          a. kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *          b. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`
+ *          c. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`
+ *        2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 64`
+ *        3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *
+ * @code {.c}
+ * #include <cuvs/core/c_api.h>
+ * #include <cuvs/neighbors/hnsw.h>
+ *
+ * // Create cuvsResources_t
+ * cuvsResources_t res;
+ * cuvsError_t res_create_status = cuvsResourcesCreate(&res);
+ *
+ * // Assume a populated `DLManagedTensor` type here
+ * DLManagedTensor dataset;
+ * DLManagedTensor queries;
+ * DLManagedTensor neighbors;
+ *
+ * // Create default search params
+ * cuvsHnswSearchParams_t params;
+ * cuvsError_t params_create_status = cuvsHnswSearchParamsCreate(&params);
+ *
+ * // Search the `index` built using `cuvsHnswBuild`
+ * cuvsError_t search_status = cuvsHnswSearch(res, params, index, &queries, &neighbors,
+ * &distances);
+ *
+ * // de-allocate `params` and `res`
+ * cuvsError_t params_destroy_status = cuvsHnswSearchParamsDestroy(params);
+ * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);
+ * @endcode
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] params cuvsHnswSearchParams_t used to search Hnsw index
+ * @param[in] index cuvsHnswIndex which has been returned by `cuvsHnswBuild`
+ * @param[in] queries DLManagedTensor* queries dataset to search
+ * @param[out] neighbors DLManagedTensor* output `k` neighbors for queries
+ * @param[out] distances DLManagedTensor* output `k` distances for queries
+ */
+cuvsError_t cuvsHnswSearch(cuvsResources_t res,
+                           cuvsHnswSearchParams_t params,
+                           cuvsHnswIndex_t index,
+                           DLManagedTensor* queries,
+                           DLManagedTensor* neighbors,
+                           DLManagedTensor* distances);
+
+/**
+ * @}
+ */
+
+/**
+ * @defgroup hnsw_c_serialize HNSW C-API serialize functions
+ * @{
+ */
+
+/**
+ * Load hnswlib index from file which was serialized from a HNSW index.
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @code{.c}
+ * #include <cuvs/core/c_api.h>
+ * #include <cuvs/neighbors/cagra.h>
+ * #include <cuvs/neighbors/hnsw.h>
+ *
+ * // Create cuvsResources_t
+ * cuvsResources_t res;
+ * cuvsError_t res_create_status = cuvsResourcesCreate(&res);
+ *
+ * // create an index with `cuvsCagraBuild`
+ * cuvsCagraSerializeHnswlib(res, "/path/to/index", index);
+ *
+ * // Load the serialized CAGRA index from file as an hnswlib index
+ * // The index should have the same dtype as the one used to build CAGRA the index
+ * cuvsHnswIndex_t hnsw_index;
+ * cuvsHnswIndexCreate(&hnsw_index);
+ * hnsw_index->dtype = index->dtype;
+ * cuvsCagraDeserialize(res, "/path/to/index", hnsw_index);
+ * @endcode
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] filename the name of the file that stores the index
+ * @param[in] dim dimensions of the training dataset
+ * @param[in] metric distance metric to search. Supported metrics ("L2Expanded")
+ * @param[out] index HNSW index loaded disk
+ */
+cuvsError_t cuvsHnswDeserialize(cuvsResources_t res,
+                                const char* filename,
+                                int dim,
+                                cuvsDistanceType metric,
+                                cuvsHnswIndex_t index);
+/**
+ * @}
+ */
+
+#ifdef __cplusplus
+}
+#endif

cpp/include/cuvs/neighbors/hnsw.hpp

@@ -58,6 +58,8 @@ struct index : cuvs::neighbors::index {
    */
   index(int dim, cuvs::distance::DistanceType metric) : dim_{dim}, metric_{metric} {}
 
+  virtual ~index() {}
+
   /**
   @brief Get underlying index
   */
@@ -230,7 +232,7 @@ void search(raft::resources const& res,
 void search(raft::resources const& res,
             const search_params& params,
             const index<uint8_t>& idx,
-            raft::host_matrix_view<const int, int64_t, raft::row_major> queries,
+            raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> queries,
             raft::host_matrix_view<uint64_t, int64_t, raft::row_major> neighbors,
             raft::host_matrix_view<float, int64_t, raft::row_major> distances);
 
@@ -271,7 +273,7 @@ void search(raft::resources const& res,
 void search(raft::resources const& res,
             const search_params& params,
             const index<int8_t>& idx,
-            raft::host_matrix_view<const int, int64_t, raft::row_major> queries,
+            raft::host_matrix_view<const int8_t, int64_t, raft::row_major> queries,
             raft::host_matrix_view<uint64_t, int64_t, raft::row_major> neighbors,
             raft::host_matrix_view<float, int64_t, raft::row_major> distances);
 

cpp/src/neighbors/cagra_c.cpp

@@ -130,6 +130,14 @@ void _serialize(cuvsResources_t res,
   cuvs::neighbors::cagra::serialize(*res_ptr, std::string(filename), *index_ptr, include_dataset);
 }
 
+template <typename T>
+void _serialize_to_hnswlib(cuvsResources_t res, const char* filename, cuvsCagraIndex_t index)
+{
+  auto res_ptr   = reinterpret_cast<raft::resources*>(res);
+  auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<T, uint32_t>*>(index->addr);
+  cuvs::neighbors::cagra::serialize_to_hnswlib(*res_ptr, std::string(filename), *index_ptr);
+}
+
 template <typename T>
 void* _deserialize(cuvsResources_t res, const char* filename)
 {
@@ -168,6 +176,14 @@ extern "C" cuvsError_t cuvsCagraIndexDestroy(cuvsCagraIndex_t index_c_ptr)
   });
 }
 
+extern "C" cuvsError_t cuvsCagraIndexDim(cuvsCagraIndex_t index, int* dim)
+{
+  return cuvs::core::translate_exceptions([=] {
+    auto index_ptr = reinterpret_cast<cuvs::neighbors::cagra::index<float, uint32_t>*>(index->addr);
+    *dim           = index_ptr->dim();
+  });
+}
+
 extern "C" cuvsError_t cuvsCagraBuild(cuvsResources_t res,
                                       cuvsCagraIndexParams_t params,
                                       DLManagedTensor* dataset_tensor,
@@ -326,3 +342,20 @@ extern "C" cuvsError_t cuvsCagraSerialize(cuvsResources_t res,
     }
   });
 }
+
+extern "C" cuvsError_t cuvsCagraSerializeToHnswlib(cuvsResources_t res,
+                                                   const char* filename,
+                                                   cuvsCagraIndex_t index)
+{
+  return cuvs::core::translate_exceptions([=] {
+    if (index->dtype.code == kDLFloat && index->dtype.bits == 32) {
+      _serialize_to_hnswlib<float>(res, filename, index);
+    } else if (index->dtype.code == kDLInt && index->dtype.bits == 8) {
+      _serialize_to_hnswlib<int8_t>(res, filename, index);
+    } else if (index->dtype.code == kDLUInt && index->dtype.bits == 8) {
+      _serialize_to_hnswlib<uint8_t>(res, filename, index);
+    } else {
+      RAFT_FAIL("Unsupported index dtype: %d and bits: %d", index->dtype.code, index->dtype.bits);
+    }
+  });
+}

cpp/src/neighbors/detail/cagra/cagra_serialize.cuh

@@ -119,9 +119,9 @@ void serialize_to_hnswlib(raft::resources const& res,
   os.write(reinterpret_cast<char*>(&curr_element_count), sizeof(std::size_t));
   // 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 * 4 + sizeof(labeltype)
-  auto size_data_per_element =
-    static_cast<std::size_t>(index_.graph_degree() * sizeof(IdxT) + 4 + index_.dim() * 4 + 8);
+  // dim * sizeof(T) + sizeof(labeltype)
+  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;
@@ -184,16 +184,9 @@ 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 = static_cast<float>(host_dataset(i, j));
-        os.write(reinterpret_cast<char*>(&data_elem), sizeof(float));
-      }
-    } 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));
-      }
+    for (std::size_t j = 0; j < index_.dim(); ++j) {
+      auto data_elem = static_cast<T>(host_dataset(i, j));
+      os.write(reinterpret_cast<char*>(&data_elem), sizeof(T));
     }
 
     os.write(reinterpret_cast<char*>(&i), sizeof(std::size_t));