KNN bruteforce and IVF methods C/C++ API

.gitignore

@@ -65,3 +65,9 @@ _text
 # clang tooling
 compile_commands.json
 .clangd/
+
+# serialized ann indexes
+cagra_index
+ivf_flat_index
+ivf_pq_index
+

cpp/CMakeLists.txt

@@ -190,6 +190,7 @@ include(cmake/thirdparty/get_cutlass.cmake)
 
 add_library(
   cuvs SHARED
+  src/neighbors/brute_force.cu
   src/neighbors/cagra_build_float.cpp
   src/neighbors/cagra_build_int8.cpp
   src/neighbors/cagra_build_uint8.cpp
@@ -200,6 +201,28 @@ add_library(
   src/neighbors/cagra_serialize_float.cpp
   src/neighbors/cagra_serialize_int8.cpp
   src/neighbors/cagra_serialize_uint8.cpp
+  src/neighbors/ivf_flat_build_float.cpp
+  src/neighbors/ivf_flat_build_int8.cpp
+  src/neighbors/ivf_flat_build_uint8.cpp
+  src/neighbors/ivf_flat_extend_float.cpp
+  src/neighbors/ivf_flat_extend_int8.cpp
+  src/neighbors/ivf_flat_extend_uint8.cpp
+  src/neighbors/ivf_flat_search_float.cpp
+  src/neighbors/ivf_flat_search_int8.cpp
+  src/neighbors/ivf_flat_search_uint8.cpp
+  src/neighbors/ivf_flat_serialize_float.cpp
+  src/neighbors/ivf_flat_serialize_int8.cpp
+  src/neighbors/ivf_flat_serialize_uint8.cpp
+  src/neighbors/ivf_pq_build_float.cpp
+  src/neighbors/ivf_pq_build_int8.cpp
+  src/neighbors/ivf_pq_build_uint8.cpp
+  src/neighbors/ivf_pq_extend_float.cpp
+  src/neighbors/ivf_pq_extend_int8.cpp
+  src/neighbors/ivf_pq_extend_uint8.cpp
+  src/neighbors/ivf_pq_search_float.cpp
+  src/neighbors/ivf_pq_search_int8.cpp
+  src/neighbors/ivf_pq_search_uint8.cpp
+  src/neighbors/ivf_pq_serialize.cpp
 )
 
 target_compile_options(
@@ -300,7 +323,13 @@ 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)
 

cpp/include/cuvs/core/detail/interop.hpp

@@ -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;
 

cpp/include/cuvs/core/interop.hpp

@@ -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);
 }

cpp/include/cuvs/distance/distance_types.h

@@ -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

cpp/include/cuvs/neighbors/brute_force-inl.hpp

@@ -0,0 +1,97 @@
+/*
+ * 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 "ann_types.hpp"
+#include <raft/neighbors/brute_force-inl.cuh>
+
+namespace cuvs::neighbors::brute_force {
+
+/**
+ * @brief Brute Force index.
+ *
+ * The index stores the dataset and norms for the dataset in device memory.
+ *
+ * @tparam T data element type
+ */
+template <typename T>
+struct index : cuvs::neighbors::ann::index {
+ public:
+  // Don't allow copying the index for performance reasons (try avoiding copying data)
+  index(const index&)                    = delete;
+  index(index&&)                         = default;
+  auto operator=(const index&) -> index& = delete;
+  auto operator=(index&&) -> index&      = default;
+  ~index()                               = default;
+
+  /** Build a cuvs bruteforce index from an existing RAFT bruteforce index. */
+  index(raft::neighbors::brute_force::index<T>&& raft_idx)
+    : cuvs::neighbors::ann::index(),
+      raft_index_(std::make_unique<raft::neighbors::brute_force::index<T>>(std::move(raft_idx)))
+  {
+  }
+
+  /** Distance metric used for retrieval */
+  [[nodiscard]] cuvs::distance::DistanceType metric() const noexcept
+  {
+    return static_cast<cuvs::distance::DistanceType>((int)raft_index_->metric());
+  }
+
+  /** Total length of the index (number of vectors). */
+  [[nodiscard]] auto size() const noexcept { return raft_index_->size(); }
+
+  /** Dimensionality of the data. */
+  [[nodiscard]] auto dim() const noexcept { return raft_index_->dim(); }
+
+  /** Dataset [size, dim] */
+  [[nodiscard]] auto dataset() const noexcept
+    -> raft::device_matrix_view<const T, int64_t, raft::row_major>
+  {
+    return raft_index_->dataset();
+  }
+
+  /** Dataset norms */
+  [[nodiscard]] auto norms() const -> raft::device_vector_view<const T, int64_t, raft::row_major>
+  {
+    return raft_index_->norms();
+  }
+
+  /** Whether or not this index has dataset norms */
+  [[nodiscard]] bool has_norms() const noexcept { return raft_index_->has_norms(); }
+
+  [[nodiscard]] T metric_arg() const noexcept { return raft_index_->metric_arg(); }
+
+  /**
+   * Replace the dataset with a new dataset.
+   */
+  void update_dataset(raft::resources const& res,
+                      raft::device_matrix_view<const T, int64_t, raft::row_major> dataset)
+  {
+    raft_index_->update_dataset(res, dataset);
+  }
+
+  auto get_raft_index() const -> const raft::neighbors::brute_force::index<T>*
+  {
+    return raft_index_.get();
+  }
+  auto get_raft_index() -> raft::neighbors::brute_force::index<T>* { return raft_index_.get(); }
+
+ private:
+  std::unique_ptr<raft::neighbors::brute_force::index<T>> raft_index_;
+};
+
+}  // namespace cuvs::neighbors::brute_force