From 9baadc8755c5c42d8d70bf70441a230b6ad3ec2a Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Wed, 7 Feb 2024 16:08:52 +0100
Subject: [PATCH] Expose IVF methods

---
 cpp/include/cuvs/neighbors/ivf_flat.hpp | 51 ++++++++++++++++++++++++-
 cpp/include/cuvs/neighbors/ivf_pq.hpp   | 41 ++++++++++++++++++++
 2 files changed, 91 insertions(+), 1 deletion(-)

diff --git a/cpp/include/cuvs/neighbors/ivf_flat.hpp b/cpp/include/cuvs/neighbors/ivf_flat.hpp
index cd74b9ff3..327490db0 100644
--- a/cpp/include/cuvs/neighbors/ivf_flat.hpp
+++ b/cpp/include/cuvs/neighbors/ivf_flat.hpp
@@ -27,4 +27,53 @@ using search_params = raft::neighbors::ivf_flat::search_params;
 template <typename T, typename IdxT>
 using index = raft::neighbors::ivf_flat::index<T, IdxT>;
 
-}  // namespace cuvs::neighbors::ivf_flat
+#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
\ No newline at end of file
diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp
index 3a1cfbff0..5dec958eb 100644
--- a/cpp/include/cuvs/neighbors/ivf_pq.hpp
+++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp
@@ -27,4 +27,45 @@ 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