KNN bruteforce and IVF methods C/C++ API

cpp/CMakeLists.txt

@@ -188,6 +188,8 @@ include(cmake/thirdparty/get_cutlass.cmake)
 
 add_library(
   cuvs SHARED
+  src/neighbors/brute_force.cpp
+
   src/neighbors/cagra_build_float.cpp
   src/neighbors/cagra_build_int8.cpp
   src/neighbors/cagra_build_uint8.cpp
@@ -198,6 +200,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_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(

cpp/include/cuvs/neighbors/brute_force.hpp

@@ -0,0 +1,39 @@
+/*
+ * 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 <raft/distance/distance_types.hpp>
+#include <raft_runtime/neighbors/brute_force.hpp>
+
+
+namespace cuvs::neighbors::brute_force {
+
+#define RAFT_INST_BFKNN(IDX_T, DATA_T, MATRIX_IDX_T, INDEX_LAYOUT, SEARCH_LAYOUT)               \
+  void knn(raft::resources const& handle,                                                       \
+           raft::device_matrix_view<const DATA_T, MATRIX_IDX_T, INDEX_LAYOUT> index,            \
+           raft::device_matrix_view<const DATA_T, MATRIX_IDX_T, SEARCH_LAYOUT> search,          \
+           raft::device_matrix_view<IDX_T, MATRIX_IDX_T, raft::row_major> indices,              \
+           raft::device_matrix_view<DATA_T, MATRIX_IDX_T, raft::row_major> distances,           \
+           raft::distance::DistanceType metric   = raft::distance::DistanceType::L2Unexpanded,  \
+           std::optional<float> metric_arg       = std::make_optional<float>(2.0f),             \
+           std::optional<IDX_T> global_id_offset = std::nullopt);
+
+RAFT_INST_BFKNN(int64_t, float, int64_t, raft::row_major, raft::row_major);
+
+#undef RAFT_INST_BFKNN
+
+}  // namespace cuvs::neighbors::brute_force

cpp/include/cuvs/neighbors/ivf_flat.hpp

@@ -0,0 +1,79 @@
+/*
+ * 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 <raft_runtime/neighbors/ivf_flat.hpp>
+
+
+namespace cuvs::neighbors::ivf_flat {
+
+using index_params = raft::neighbors::ivf_flat::index_params;
+using search_params = raft::neighbors::ivf_flat::search_params;
+
+template <typename T, typename IdxT>
+using index = raft::neighbors::ivf_flat::index<T, IdxT>;
+
+#define CUVS_IVF_FLAT(T, IdxT)                                                                \
+  auto build(raft::resources const& handle,                                                   \
+             const cuvs::neighbors::ivf_flat::index_params& params,                           \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset);               \
+                                                                                              \
+  void build(raft::resources const& handle,                                                   \
+             const cuvs::neighbors::ivf_flat::index_params& params,                           \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset,                \
+             cuvs::neighbors::ivf_flat::index<T, IdxT>& idx);                                 \
+                                                                                              \
+  auto extend(raft::resources const& handle,                                                  \
+              raft::device_matrix_view<const T, IdxT, raft::row_major> new_vectors,           \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,          \
+              const cuvs::neighbors::ivf_flat::index<T, IdxT>& orig_index);                   \
+                                                                                              \
+  void extend(raft::resources const& handle,                                                  \
+              raft::device_matrix_view<const T, IdxT, raft::row_major> new_vectors,           \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,          \
+              cuvs::neighbors::ivf_flat::index<T, IdxT>* idx);                                \
+                                                                                              \
+  void search(raft::resources const& handle,                                                  \
+              const cuvs::neighbors::ivf_flat::search_params& params,                         \
+              cuvs::neighbors::ivf_flat::index<T, IdxT>& index,                               \
+              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<float, IdxT, raft::row_major> distances);              \
+                                                                                              \
+  void serialize_file(raft::resources const& handle,                                          \
+                      const std::string& filename,                                            \
+                      const cuvs::neighbors::ivf_flat::index<T, IdxT>& index);                \
+                                                                                              \
+  void deserialize_file(raft::resources const& handle,                                        \
+                        const std::string& filename,                                          \
+                        cuvs::neighbors::ivf_flat::index<T, IdxT>* index);                    \
+                                                                                              \
+  void serialize(raft::resources const& handle,                                               \
+                 std::string& str,                                                            \
+                 const cuvs::neighbors::ivf_flat::index<T, IdxT>& index);                     \
+                                                                                              \
+  void deserialize(raft::resources const& handle,                                             \
+                   const std::string& str,                                                    \
+                   cuvs::neighbors::ivf_flat::index<T, IdxT>* index);
+
+CUVS_IVF_FLAT(float, uint64_t);
+CUVS_IVF_FLAT(int8_t, uint64_t);
+CUVS_IVF_FLAT(uint8_t, uint64_t);
+
+#undef CUVS_IVF_FLAT
+
+}  // namespace cuvs::neighbors::ivf_flat

cpp/include/cuvs/neighbors/ivf_pq.hpp

@@ -0,0 +1,71 @@
+/*
+ * 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 <raft_runtime/neighbors/ivf_pq.hpp>
+
+
+namespace cuvs::neighbors::ivf_pq {
+
+using index_params = raft::neighbors::ivf_pq::index_params;
+using search_params = raft::neighbors::ivf_pq::search_params;
+
+template <typename IdxT>
+using index = raft::neighbors::ivf_pq::index<IdxT>;
+
+#define CUVS_IVF_PQ(T, IdxT)                                                                    \
+  auto build(raft::resources const& handle,                                                     \
+             const cuvs::neighbors::ivf_pq::index_params& params,                               \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset);                 \
+                                                                                                \
+  void build(raft::resources const& handle,                                                     \
+             const cuvs::neighbors::ivf_pq::index_params& params,                               \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset,                  \
+             cuvs::neighbors::ivf_pq::index<IdxT>* idx);                                        \
+                                                                                                \
+  auto extend(raft::resources const& handle,                                                    \
+              raft::device_matrix_view<const T, IdxT, raft::row_major> new_vectors,             \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,            \
+              const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index);                          \
+                                                                                                \
+  void extend(raft::resources const& handle,                                                    \
+              raft::device_matrix_view<const T, IdxT, raft::row_major> new_vectors,             \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,            \
+              cuvs::neighbors::ivf_pq::index<IdxT>* idx);                                       \
+                                                                                                \
+  void search(raft::resources const& handle,                                                    \
+              const cuvs::neighbors::ivf_pq::search_params& params,                             \
+              cuvs::neighbors::ivf_pq::index<IdxT>& index,                                      \
+              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<float, IdxT, raft::row_major> distances);                \
+                                                                                                \
+  void serialize(raft::resources const& handle,                                                 \
+                 std::string& filename,                                                         \
+                 const cuvs::neighbors::ivf_pq::index<IdxT>& index);                            \
+                                                                                                \
+  void deserialize(raft::resources const& handle,                                               \
+                   const std::string& filename,                                                 \
+                   cuvs::neighbors::ivf_pq::index<IdxT>* index);
+
+CUVS_IVF_PQ(float, uint64_t);
+CUVS_IVF_PQ(int8_t, uint64_t);
+CUVS_IVF_PQ(uint8_t, uint64_t);
+
+#undef CUVS_IVF_PQ
+
+}  // namespace cuvs::neighbors::ivf_pq

cpp/src/neighbors/brute_force.cpp

@@ -0,0 +1,39 @@
+/*
+ * 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.
+ */
+
+#include <cuvs/neighbors/brute_force.hpp>
+
+namespace cuvs::neighbors::brute_force {
+
+#define RAFT_INST_BFKNN(IDX_T, DATA_T, MATRIX_IDX_T, INDEX_LAYOUT, SEARCH_LAYOUT)               \
+  void knn(raft::resources const& handle,                                                       \
+           raft::device_matrix_view<const DATA_T, MATRIX_IDX_T, INDEX_LAYOUT> index,            \
+           raft::device_matrix_view<const DATA_T, MATRIX_IDX_T, SEARCH_LAYOUT> search,          \
+           raft::device_matrix_view<IDX_T, MATRIX_IDX_T, raft::row_major> indices,              \
+           raft::device_matrix_view<DATA_T, MATRIX_IDX_T, raft::row_major> distances,           \
+           raft::distance::DistanceType metric,                                                 \
+           std::optional<float> metric_arg,                                                     \
+           std::optional<IDX_T> global_id_offset)                                               \
+  {                                                                                             \
+    raft::runtime::neighbors::brute_force::knn(handle, index, search, indices, distances,       \
+                                               metric, metric_arg, global_id_offset);           \
+  }
+
+RAFT_INST_BFKNN(int64_t, float, int64_t, raft::row_major, raft::row_major);
+
+#undef RAFT_INST_BFKNN
+
+}   // namespace cuvs::neighbors::brute_force

cpp/src/neighbors/ivf_flat_build_float.cpp

@@ -0,0 +1,41 @@
+/*
+ * 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.
+ */
+
+#include <cuvs/neighbors/ivf_flat.hpp>
+
+namespace cuvs::neighbors::ivf_flat {
+
+#define CUVS_INST_IVF_FLAT_BUILD(T, IdxT)                                         \
+  auto build(raft::resources const& handle,                                       \
+             const cuvs::neighbors::ivf_flat::index_params& params,               \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset)    \
+  {                                                                               \
+    return raft::runtime::neighbors::ivf_flat::build(handle, params, dataset);    \
+  }                                                                               \
+                                                                                  \
+  void build(raft::resources const& handle,                                       \
+             const cuvs::neighbors::ivf_flat::index_params& params,               \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset,    \
+             cuvs::neighbors::ivf_flat::index<T, IdxT>& idx)                      \
+  {                                                                               \
+    raft::runtime::neighbors::ivf_flat::build(handle, params, dataset, idx);      \
+  }
+
+CUVS_INST_IVF_FLAT_BUILD(float, int64_t);
+
+#undef CUVS_INST_IVF_FLAT_BUILD
+
+}   // namespace cuvs::neighbors::ivf_flat

cpp/src/neighbors/ivf_flat_build_int8.cpp

@@ -0,0 +1,41 @@
+/*
+ * 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.
+ */
+
+#include <cuvs/neighbors/ivf_flat.hpp>
+
+namespace cuvs::neighbors::ivf_flat {
+
+#define CUVS_INST_IVF_FLAT_BUILD(T, IdxT)                                         \
+  auto build(raft::resources const& handle,                                       \
+             const cuvs::neighbors::ivf_flat::index_params& params,               \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset)    \
+  {                                                                               \
+    return raft::runtime::neighbors::ivf_flat::build(handle, params, dataset);    \
+  }                                                                               \
+                                                                                  \
+  void build(raft::resources const& handle,                                       \
+             const cuvs::neighbors::ivf_flat::index_params& params,               \
+             raft::device_matrix_view<const T, IdxT, raft::row_major> dataset,    \
+             cuvs::neighbors::ivf_flat::index<T, IdxT>& idx)                      \
+  {                                                                               \
+    raft::runtime::neighbors::ivf_flat::build(handle, params, dataset, idx);      \
+  }
+
+CUVS_INST_IVF_FLAT_BUILD(int8_t, int64_t);
+
+#undef CUVS_INST_IVF_FLAT_BUILD
+
+}   // namespace cuvs::neighbors::ivf_flat