KNN bruteforce and IVF methods C/C++ API

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(
@@ -216,16 +239,13 @@ target_include_directories(
 
 if(NOT BUILD_CPU_ONLY)
   # Keep cuVS as lightweight as possible. Only CUDA libs and rmm should be used in global target.
-  target_link_libraries(cuvs
-        PUBLIC
-            rmm::rmm
-            $<$<NOT:$<BOOL:${CUDA_STATIC_RUNTIME}>>:raft::raft>
-            $<$<NOT:$<BOOL:${CUDA_STATIC_RUNTIME}>>:raft::compiled>
-        PRIVATE
-            $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::raft>
-            $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::compiled_static>
-            nvidia::cutlass::cutlass
-    )
+  target_link_libraries(
+    cuvs
+    PUBLIC rmm::rmm $<$<NOT:$<BOOL:${CUDA_STATIC_RUNTIME}>>:raft::raft>
+           $<$<NOT:$<BOOL:${CUDA_STATIC_RUNTIME}>>:raft::compiled>
+    PRIVATE $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::raft>
+            $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::compiled_static> nvidia::cutlass::cutlass
+  )
 endif()
 
 # Endian detection
@@ -280,14 +300,14 @@ endif()
 
 set_target_properties(
   cuvs
-  PROPERTIES BUILD_RPATH                         "\$ORIGIN"
-             INSTALL_RPATH                       "\$ORIGIN"
-             CXX_STANDARD                        17
-             CXX_STANDARD_REQUIRED               ON
-             CUDA_STANDARD                       17
-             CUDA_STANDARD_REQUIRED              ON
+  PROPERTIES BUILD_RPATH "\$ORIGIN"
+             INSTALL_RPATH "\$ORIGIN"
+             CXX_STANDARD 17
+             CXX_STANDARD_REQUIRED ON
+             CUDA_STANDARD 17
+             CUDA_STANDARD_REQUIRED ON
              INTERFACE_POSITION_INDEPENDENT_CODE ON
-             POSITION_INDEPENDENT_CODE           ON
+             POSITION_INDEPENDENT_CODE ON
 )
 
 target_compile_options(
@@ -300,19 +320,22 @@ 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/ivf_flat_c.cpp src/neighbors/ivf_pq_c.cpp
+                  src/neighbors/cagra_c.cpp
+  )
 
   add_library(cuvs::c_api ALIAS cuvs_c)
 
   set_target_properties(
     cuvs_c
-    PROPERTIES BUILD_RPATH                         "\$ORIGIN"
-               INSTALL_RPATH                       "\$ORIGIN"
-               CXX_STANDARD                        17
-               CXX_STANDARD_REQUIRED               ON
-               POSITION_INDEPENDENT_CODE           ON
+    PROPERTIES BUILD_RPATH "\$ORIGIN"
+               INSTALL_RPATH "\$ORIGIN"
+               CXX_STANDARD 17
+               CXX_STANDARD_REQUIRED ON
+               POSITION_INDEPENDENT_CODE ON
                INTERFACE_POSITION_INDEPENDENT_CODE ON
-               EXPORT_NAME                         c_api
+               EXPORT_NAME c_api
   )
 
   target_compile_options(cuvs_c PRIVATE "$<$<COMPILE_LANGUAGE:CXX>:${CUVS_CXX_FLAGS}>")
@@ -323,12 +346,11 @@ if(BUILD_C_LIBRARY)
     INTERFACE "$<INSTALL_INTERFACE:include>"
   )
 
-  target_link_libraries(cuvs_c
-        PUBLIC
-            cuvs::cuvs
-        PRIVATE
-            $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::raft>
-    )
+  target_link_libraries(
+    cuvs_c
+    PUBLIC cuvs::cuvs
+    PRIVATE $<$<BOOL:${CUDA_STATIC_RUNTIME}>:raft::raft>
+  )
 
   # ensure CUDA symbols aren't relocated to the middle of the debug build binaries
   target_link_options(cuvs_c PRIVATE "${CMAKE_CURRENT_BINARY_DIR}/fatbin.ld")

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_c.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.hpp

@@ -0,0 +1,122 @@
+/*
+ * 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]] constexpr inline cuvs::distance::DistanceType metric() const noexcept
+  {
+    return raft_index_->metric_;
+  }
+
+  /** Total length of the index (number of vectors). */
+  [[nodiscard]] constexpr inline auto size() const noexcept
+  {
+    return raft_index_->dataset_view_.extent(0);
+  }
+
+  /** Dimensionality of the data. */
+  [[nodiscard]] constexpr inline auto dim() const noexcept
+  {
+    return raft_index_->dataset_view_.extent(1);
+  }
+
+  /** Dataset [size, dim] */
+  [[nodiscard]] inline auto dataset() const noexcept
+    -> raft::device_matrix_view<const T, int64_t, raft::row_major>
+  {
+    return raft_index_->dataset_view_;
+  }
+
+  /** Dataset norms */
+  [[nodiscard]] inline auto norms() const
+    -> raft::device_vector_view<const T, int64_t, raft::row_major>
+  {
+    return raft_index_->norms_view_.value();
+  }
+
+  /** Whether or not this index has dataset norms */
+  [[nodiscard]] inline bool has_norms() const noexcept
+  {
+    return raft_index_->norms_view_.has_value();
+  }
+
+  [[nodiscard]] inline 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_->dataset_view_ = 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_;
+};
+
+#define CUVS_INST_BFKNN(T, IdxT)                                                               \
+  auto build(raft::resources const& res,                                                       \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset,                 \
+             cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Unexpanded, \
+             float metric_arg                    = 0.0)                                        \
+    ->cuvs::neighbors::brute_force::index<T>;                                                  \
+                                                                                               \
+  void search(raft::resources const& res,                                                      \
+              const cuvs::neighbors::brute_force::index<T>& idx,                               \
+              raft::device_matrix_view<const T, IdxT, raft::row_major> queries,                \
+              raft::device_matrix_view<IdxT, IdxT, raft::row_major> neighbors,                 \
+              raft::device_matrix_view<T, IdxT, raft::row_major> distances);
+
+CUVS_INST_BFKNN(float, int64_t);
+
+}  // namespace cuvs::neighbors::brute_force