rapidsai · rapids-bot · Mar 20, 2024 · Feb 7, 2024 · Feb 7, 2024 · Feb 7, 2024
diff --git a/.gitignore b/.gitignore
@@ -66,3 +66,9 @@ docs/source/_static/rust
 # clang tooling
 compile_commands.json
 .clangd/
+
+# serialized ann indexes
+cagra_index
+ivf_flat_index
+ivf_pq_index
+
@@ -187,6 +187,8 @@ include(cmake/thirdparty/get_cutlass.cmake)
 
 add_library(
   cuvs SHARED
+  src/neighbors/brute_force_index.cu
+  src/neighbors/brute_force.cu
   src/neighbors/cagra_build_float.cpp
   src/neighbors/cagra_build_int8.cpp
   src/neighbors/cagra_build_uint8.cpp
@@ -197,6 +199,30 @@ add_library(
   src/neighbors/cagra_serialize_float.cpp
   src/neighbors/cagra_serialize_int8.cpp
   src/neighbors/cagra_serialize_uint8.cpp
+  src/neighbors/ivf_flat_index.cpp
+  src/neighbors/ivf_flat/ivf_flat_build_float_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_build_int8_t_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_build_uint8_t_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_extend_float_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_extend_int8_t_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_extend_uint8_t_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_search_float_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_search_int8_t_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_search_uint8_t_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_serialize_float_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_serialize_int8_t_int64_t.cpp
+  src/neighbors/ivf_flat/ivf_flat_serialize_uint8_t_int64_t.cpp
+  src/neighbors/ivf_pq_index.cpp
+  src/neighbors/ivf_pq/ivf_pq_build_float_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_build_int8_t_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_build_uint8_t_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_extend_float_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_extend_int8_t_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_extend_uint8_t_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_search_float_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_search_int8_t_int64_t.cpp
+  src/neighbors/ivf_pq/ivf_pq_search_uint8_t_int64_t.cpp
+  src/neighbors/ivf_pq_serialize.cpp
 )
 
 target_compile_options(
@@ -297,7 +323,14 @@ target_link_options(cuvs PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld")
 # ##################################################################################################
 # * cuvs_c -------------------------------------------------------------------------------
 if(BUILD_C_LIBRARY)
-  add_library(cuvs_c SHARED src/core/c_api.cpp src/neighbors/cagra_c.cpp)
+  add_library(
+    cuvs_c SHARED
+    src/core/c_api.cpp
+    src/neighbors/brute_force_c.cpp
+    src/neighbors/ivf_flat_c.cpp
+    src/neighbors/ivf_pq_c.cpp
+    src/neighbors/cagra_c.cpp
+  )
 
   add_library(cuvs::c_api ALIAS cuvs_c)
 

@@ -53,20 +53,20 @@ DLDataType data_type_to_DLDataType()
   }
 }
 
-bool is_dlpack_device_compatible(DLTensor tensor)
+inline bool is_dlpack_device_compatible(DLTensor tensor)
 {
   return tensor.device.device_type == kDLCUDAManaged || tensor.device.device_type == kDLCUDAHost ||
          tensor.device.device_type == kDLCUDA;
 }
 
-bool is_dlpack_host_compatible(DLTensor tensor)
+inline bool is_dlpack_host_compatible(DLTensor tensor)
 {
   return tensor.device.device_type == kDLCUDAManaged || tensor.device.device_type == kDLCUDAHost ||
          tensor.device.device_type == kDLCPU;
 }
 
 template <typename MdspanType, typename = raft::is_mdspan_t<MdspanType>>
-MdspanType from_dlpack(DLManagedTensor* managed_tensor)
+inline MdspanType from_dlpack(DLManagedTensor* managed_tensor)
 {
   auto tensor = managed_tensor->dl_tensor;
 

@@ -33,7 +33,7 @@ namespace cuvs::core {
  * @param[in] tensor DLTensor object to check underlying memory type
  * @return bool
  */
-bool is_dlpack_device_compatible(DLTensor tensor)
+inline bool is_dlpack_device_compatible(DLTensor tensor)
 {
   return detail::is_dlpack_device_compatible(tensor);
 }
@@ -46,7 +46,7 @@ bool is_dlpack_device_compatible(DLTensor tensor)
  * @param tensor DLTensor object to check underlying memory type
  * @return bool
  */
-bool is_dlpack_host_compatible(DLTensor tensor)
+inline bool is_dlpack_host_compatible(DLTensor tensor)
 {
   return detail::is_dlpack_host_compatible(tensor);
 }
@@ -72,7 +72,7 @@ bool is_dlpack_host_compatible(DLTensor tensor)
  * @return MdspanType
  */
 template <typename MdspanType, typename = raft::is_mdspan_t<MdspanType>>
-MdspanType from_dlpack(DLManagedTensor* managed_tensor)
+inline MdspanType from_dlpack(DLManagedTensor* managed_tensor)
 {
   return detail::from_dlpack<MdspanType>(managed_tensor);
 }

@@ -0,0 +1,70 @@
+/*
+ * 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.
+ */
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/** enum to tell how to compute distance */
+enum DistanceType {
+
+  /** evaluate as dist_ij = sum(x_ik^2) + sum(y_ij)^2 - 2*sum(x_ik * y_jk) */
+  L2Expanded = 0,
+  /** same as above, but inside the epilogue, perform square root operation */
+  L2SqrtExpanded = 1,
+  /** cosine distance */
+  CosineExpanded = 2,
+  /** L1 distance */
+  L1 = 3,
+  /** evaluate as dist_ij += (x_ik - y-jk)^2 */
+  L2Unexpanded = 4,
+  /** same as above, but inside the epilogue, perform square root operation */
+  L2SqrtUnexpanded = 5,
+  /** basic inner product **/
+  InnerProduct = 6,
+  /** Chebyshev (Linf) distance **/
+  Linf = 7,
+  /** Canberra distance **/
+  Canberra = 8,
+  /** Generalized Minkowski distance **/
+  LpUnexpanded = 9,
+  /** Correlation distance **/
+  CorrelationExpanded = 10,
+  /** Jaccard distance **/
+  JaccardExpanded = 11,
+  /** Hellinger distance **/
+  HellingerExpanded = 12,
+  /** Haversine distance **/
+  Haversine = 13,
+  /** Bray-Curtis distance **/
+  BrayCurtis = 14,
+  /** Jensen-Shannon distance**/
+  JensenShannon = 15,
+  /** Hamming distance **/
+  HammingUnexpanded = 16,
+  /** KLDivergence **/
+  KLDivergence = 17,
+  /** RusselRao **/
+  RusselRaoExpanded = 18,
+  /** Dice-Sorensen distance **/
+  DiceExpanded = 19,
+  /** Precomputed (special value) **/
+  Precomputed = 100
+};
+
+#ifdef __cplusplus
+}
+#endif
@@ -0,0 +1,163 @@
+/*
+ * 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_types.h>
+#include <dlpack/dlpack.h>
+#include <stdint.h>
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+/**
+ * @defgroup bruteforce_c_index Bruteforce index
+ * @{
+ */
+/**
+ * @brief Struct to hold address of cuvs::neighbors::brute_force::index and its active trained dtype
+ *
+ */
+typedef struct {
+  uintptr_t addr;
+  DLDataType dtype;
+} bruteForceIndex;
+
+typedef bruteForceIndex* cuvsBruteForceIndex_t;
+
+/**
+ * @brief Allocate BRUTEFORCE index
+ *
+ * @param[in] index cuvsBruteForceIndex_t to allocate
+ * @return cuvsError_t
+ */
+cuvsError_t bruteForceIndexCreate(cuvsBruteForceIndex_t* index);
+
+/**
+ * @brief De-allocate BRUTEFORCE index
+ *
+ * @param[in] index cuvsBruteForceIndex_t to de-allocate
+ */
+cuvsError_t bruteForceIndexDestroy(cuvsBruteForceIndex_t index);
+/**
+ * @}
+ */
+
+/**
+ * @defgroup bruteforce_c_index_build Bruteforce index build
+ * @{
+ */
+/**
+ * @brief Build a BRUTEFORCE index with a `DLManagedTensor` which has underlying
+ *        `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`,
+ *        or `kDLCPU`. Also, acceptable underlying types are:
+ *        1. `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *        2. `kDLDataType.code == kDLInt` and `kDLDataType.bits = 8`
+ *        3. `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 8`
+ *
+ * @code {.c}
+ * #include <cuvs/core/c_api.h>
+ * #include <cuvs/neighbors/brute_force.h>
+ *
+ * // Create cuvsResources_t
+ * cuvsResources_t res;
+ * cuvsError_t res_create_status = cuvsResourcesCreate(&res);
+ *
+ * // Assume a populated `DLManagedTensor` type here
+ * DLManagedTensor dataset;
+ *
+ * // Create BRUTEFORCE index
+ * cuvsBruteForceIndex_t index;
+ * cuvsError_t index_create_status = bruteForceIndexCreate(&index);
+ *
+ * // Build the BRUTEFORCE Index
+ * cuvsError_t build_status = bruteForceBuild(res, &dataset_tensor, L2Expanded, 0.f, index);
+ *
+ * // de-allocate `index` and `res`
+ * cuvsError_t index_destroy_status = bruteForceIndexDestroy(index);
+ * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);
+ * @endcode
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] dataset DLManagedTensor* training dataset
+ * @param[in] metric metric
+ * @param[in] metric_arg metric_arg
+ * @param[out] index cuvsBruteForceIndex_t Newly built BRUTEFORCE index
+ * @return cuvsError_t
+ */
+cuvsError_t bruteForceBuild(cuvsResources_t res,
+                            DLManagedTensor* dataset,
+                            enum DistanceType metric,
+                            float metric_arg,
+                            cuvsBruteForceIndex_t index);
+/**
+ * @}
+ */
+
+/**
+ * @defgroup bruteforce_c_index_search Bruteforce index search
+ * @{
+ */
+/**
+ * @brief Search a BRUTEFORCE index with a `DLManagedTensor` which has underlying
+ *        `DLDeviceType` equal to `kDLCUDA`, `kDLCUDAHost`, `kDLCUDAManaged`.
+ *        It is also important to note that the BRUTEFORCE 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`
+ *        2. `neighbors`: `kDLDataType.code == kDLUInt` and `kDLDataType.bits = 32`
+ *        3. `distances`: `kDLDataType.code == kDLFloat` and `kDLDataType.bits = 32`
+ *
+ * @code {.c}
+ * #include <cuvs/core/c_api.h>
+ * #include <cuvs/neighbors/brute_force.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;
+ *
+ * // Search the `index` built using `bruteForceBuild`
+ * cuvsError_t search_status = bruteForceSearch(res, index, &queries, &neighbors, &distances);
+ *
+ * // de-allocate `res`
+ * cuvsError_t res_destroy_status = cuvsResourcesDestroy(res);
+ * @endcode
+ *
+ * @param[in] res cuvsResources_t opaque C handle
+ * @param[in] index bruteForceIndex which has been returned by `bruteForceBuild`
+ * @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 bruteForceSearch(cuvsResources_t res,
+                             cuvsBruteForceIndex_t index,
+                             DLManagedTensor* queries,
+                             DLManagedTensor* neighbors,
+                             DLManagedTensor* distances);
+/**
+ * @}
+ */
+
+#ifdef __cplusplus
+}
+#endif