From 53caff9f6aeb7eaf184dafb9d0448545b598bb62 Mon Sep 17 00:00:00 2001
From: Tamas Bela Feher <tfeher@nvidia.com>
Date: Thu, 23 May 2024 18:15:27 +0200
Subject: [PATCH 1/4] Accept host_mdspan for IVF-PQ build and extend

---
 cpp/include/cuvs/neighbors/ivf_pq.hpp         | 140 ++++++++++++++++++
 .../neighbors/detail/cagra/cagra_build.cuh    |   3 +-
 .../ivf_pq/detail/generate_ivf_pq.py          |  39 +----
 .../ivf_pq_build_extend_float_int64_t.cu      |  35 +----
 .../detail/ivf_pq_build_extend_inst.cuh       |  92 ++++++++++++
 .../ivf_pq_build_extend_int8_t_int64_t.cu     |  35 +----
 .../ivf_pq_build_extend_uint8_t_int64_t.cu    |  35 +----
 cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh     |  47 +++---
 8 files changed, 256 insertions(+), 170 deletions(-)
 create mode 100644 cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh

diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp
index d2d6ca9bb..d6dda82ac 100644
--- a/cpp/include/cuvs/neighbors/ivf_pq.hpp
+++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp
@@ -594,6 +594,146 @@ void build(raft::resources const& handle,
            const cuvs::neighbors::ivf_pq::index_params& index_params,
            raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
            cuvs::neighbors::ivf_pq::index<int64_t>* idx);
+/**
+ * @brief Build the index from the dataset for efficient search.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   // use default index parameters
+ *   ivf_pq::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   auto index = ivf_pq::build(handle, index_params, dataset);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] index_params configure the index building
+ * @param[in] dataset a host_matrix_view to a row-major matrix [n_rows, dim]
+ *
+ * @return the constructed ivf-pq index
+ */
+auto build(raft::resources const& handle,
+           const cuvs::neighbors::ivf_pq::index_params& index_params,
+           raft::host_matrix_view<const float, int64_t, raft::row_major> dataset)
+  -> cuvs::neighbors::ivf_pq::index<int64_t>;
+
+/**
+ * @brief Build the index from the dataset for efficient search.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   // use default index parameters
+ *   ivf_pq::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
+ *   ivf_pq::build(handle, index_params, dataset, index);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] index_params configure the index building
+ * @param[in] dataset raft::host_matrix_view to a row-major matrix [n_rows, dim]
+ * @param[out] idx reference to ivf_pq::index
+ *
+ */
+void build(raft::resources const& handle,
+           const cuvs::neighbors::ivf_pq::index_params& index_params,
+           raft::host_matrix_view<const float, int64_t, raft::row_major> dataset,
+           cuvs::neighbors::ivf_pq::index<int64_t>* idx);
+
+/**
+ * @brief Build the index from the dataset for efficient search.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   // use default index parameters
+ *   ivf_pq::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   auto index = ivf_pq::build(handle, index_params, dataset);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] index_params configure the index building
+ * @param[in] dataset a host_matrix_view to a row-major matrix [n_rows, dim]
+ *
+ * @return the constructed ivf-pq index
+ */
+auto build(raft::resources const& handle,
+           const cuvs::neighbors::ivf_pq::index_params& index_params,
+           raft::host_matrix_view<const int8_t, int64_t, raft::row_major> dataset)
+  -> cuvs::neighbors::ivf_pq::index<int64_t>;
+
+/**
+ * @brief Build the index from the dataset for efficient search.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   // use default index parameters
+ *   ivf_pq::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
+ *   ivf_pq::build(handle, index_params, dataset, index);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] index_params configure the index building
+ * @param[in] dataset raft::host_matrix_view to a row-major matrix [n_rows, dim]
+ * @param[out] idx reference to ivf_pq::index
+ *
+ */
+void build(raft::resources const& handle,
+           const cuvs::neighbors::ivf_pq::index_params& index_params,
+           raft::host_matrix_view<const int8_t, int64_t, raft::row_major> dataset,
+           cuvs::neighbors::ivf_pq::index<int64_t>* idx);
+
+/**
+ * @brief Build the index from the dataset for efficient search.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   // use default index parameters
+ *   ivf_pq::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   auto index = ivf_pq::build(handle, index_params, dataset);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] index_params configure the index building
+ * @param[in] dataset a host_matrix_view to a row-major matrix [n_rows, dim]
+ *
+ * @return the constructed ivf-pq index
+ */
+auto build(raft::resources const& handle,
+           const cuvs::neighbors::ivf_pq::index_params& index_params,
+           raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset)
+  -> cuvs::neighbors::ivf_pq::index<int64_t>;
+
+/**
+ * @brief Build the index from the dataset for efficient search.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   // use default index parameters
+ *   ivf_pq::index_params index_params;
+ *   // create and fill the index from a [N, D] dataset
+ *   ivf_pq::index<decltype(dataset::value_type), decltype(dataset::index_type)> index;
+ *   ivf_pq::build(handle, index_params, dataset, index);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] index_params configure the index building
+ * @param[in] dataset raft::host_matrix_view to a row-major matrix [n_rows, dim]
+ * @param[out] idx reference to ivf_pq::index
+ *
+ */
+void build(raft::resources const& handle,
+           const cuvs::neighbors::ivf_pq::index_params& index_params,
+           raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> dataset,
+           cuvs::neighbors::ivf_pq::index<int64_t>* idx);
 /**
  * @}
  */
diff --git a/cpp/src/neighbors/detail/cagra/cagra_build.cuh b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
index 9a92a5350..c5611e586 100644
--- a/cpp/src/neighbors/detail/cagra/cagra_build.cuh
+++ b/cpp/src/neighbors/detail/cagra/cagra_build.cuh
@@ -91,8 +91,7 @@ void build_knn_graph(
   }();
 
   RAFT_LOG_DEBUG("# Building IVF-PQ index %s", model_name.c_str());
-  auto index = cuvs::neighbors::ivf_pq::detail::build<DataT, int64_t>(
-    res, *build_params, dataset.data_handle(), dataset.extent(0), dataset.extent(1));
+  auto index = cuvs::neighbors::ivf_pq::detail::build<DataT, int64_t>(res, *build_params, dataset);
 
   //
   // search top (k + 1) neighbors
diff --git a/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py b/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py
index 2ba513579..1d3704fac 100644
--- a/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py
+++ b/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py
@@ -41,7 +41,7 @@
 """
 
 build_include_macro = """
-#include "../ivf_pq_build.cuh"
+#include "ivf_pq_build_extent_inst.cuh"
 """
 search_include_macro = """
 #include "../ivf_pq_search.cuh"
@@ -61,42 +61,7 @@
     uint8_t_int64_t=("uint8_t", "int64_t"),
 )
 
-build_extend_macro = """
-#define CUVS_INST_IVF_PQ_BUILD_EXTEND(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)               \\
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                   \\
-  {                                                                                          \\
-    return cuvs::neighbors::ivf_pq::detail::build(handle, params, 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)                                      \\
-  {                                                                                          \\
-    cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, 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)                        \\
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                   \\
-  {                                                                                          \\
-    return cuvs::neighbors::ivf_pq::detail::extend(                                          \\
-      handle, new_vectors, new_indices, 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)                                     \\
-  {                                                                                          \\
-    cuvs::neighbors::ivf_pq::detail::extend(                                                 \\
-      handle, new_vectors, new_indices, idx);                                                \\
-  }             
-"""
+build_extend_macro = ""  # moved to header ivf_pq_build_extend_inst.cuh
 
 search_macro = """
 #define CUVS_INST_IVF_PQ_SEARCH(T, IdxT)                                        \\
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu
index ee7f26381..0269481fb 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu
@@ -25,42 +25,9 @@
 
 #include <cuvs/neighbors/ivf_pq.hpp>
 
-#include "../ivf_pq_build.cuh"
+#include "ivf_pq_build_extent_inst.cuh"
 
 namespace cuvs::neighbors::ivf_pq {
-
-#define CUVS_INST_IVF_PQ_BUILD_EXTEND(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)                    \
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
-  {                                                                                               \
-    return cuvs::neighbors::ivf_pq::detail::build(handle, params, 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)                                           \
-  {                                                                                               \
-    cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, 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)                             \
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
-  {                                                                                               \
-    return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, 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)                                          \
-  {                                                                                               \
-    cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, idx);               \
-  }
 CUVS_INST_IVF_PQ_BUILD_EXTEND(float, int64_t);
 
 #undef CUVS_INST_IVF_PQ_BUILD_EXTEND
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
new file mode 100644
index 000000000..e79779e59
--- /dev/null
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
@@ -0,0 +1,92 @@
+/*
+ * 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.
+ */
+
+/*
+ * NOTE: this file is used by generate_ivf_pq.py
+ *
+ */
+
+#include <cuvs/neighbors/ivf_pq.hpp>
+
+#include "../ivf_pq_build.cuh"
+
+namespace cuvs::neighbors::ivf_pq {
+
+#define CUVS_INST_IVF_PQ_BUILD_EXTEND(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)                    \
+    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+  {                                                                                               \
+    return cuvs::neighbors::ivf_pq::detail::build(handle, params, 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)                                           \
+  {                                                                                               \
+    cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, idx);                         \
+  }                                                                                               \
+                                                                                                  \
+  auto build(raft::resources const& handle,                                                       \
+             const cuvs::neighbors::ivf_pq::index_params& params,                                 \
+             raft::host_matrix_view<const T, IdxT, raft::row_major> dataset)                      \
+    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+  {                                                                                               \
+    return cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset);                       \
+  }                                                                                               \
+                                                                                                  \
+  void build(raft::resources const& handle,                                                       \
+             const cuvs::neighbors::ivf_pq::index_params& params,                                 \
+             raft::host_matrix_view<const T, IdxT, raft::row_major> dataset,                      \
+             cuvs::neighbors::ivf_pq::index<IdxT>* idx)                                           \
+  {                                                                                               \
+    cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, 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)                             \
+    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+  {                                                                                               \
+    return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, 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)                                          \
+  {                                                                                               \
+    cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, idx);               \
+  }                                                                                               \
+  auto extend(raft::resources const& handle,                                                      \
+              raft::host_matrix_view<const T, IdxT, raft::row_major> new_vectors,                 \
+              std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,                \
+              const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index)                             \
+    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+  {                                                                                               \
+    return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, orig_index); \
+  }                                                                                               \
+                                                                                                  \
+  void extend(raft::resources const& handle,                                                      \
+              raft::host_matrix_view<const T, IdxT, raft::row_major> new_vectors,                 \
+              std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,                \
+              cuvs::neighbors::ivf_pq::index<IdxT>* idx)                                          \
+  {                                                                                               \
+    cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, idx);               \
+  }
+}  // namespace cuvs::neighbors::ivf_pq
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu
index 8b612439d..943fffaf2 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu
@@ -25,42 +25,9 @@
 
 #include <cuvs/neighbors/ivf_pq.hpp>
 
-#include "../ivf_pq_build.cuh"
+#include "ivf_pq_build_extent_inst.cuh"
 
 namespace cuvs::neighbors::ivf_pq {
-
-#define CUVS_INST_IVF_PQ_BUILD_EXTEND(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)                    \
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
-  {                                                                                               \
-    return cuvs::neighbors::ivf_pq::detail::build(handle, params, 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)                                           \
-  {                                                                                               \
-    cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, 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)                             \
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
-  {                                                                                               \
-    return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, 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)                                          \
-  {                                                                                               \
-    cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, idx);               \
-  }
 CUVS_INST_IVF_PQ_BUILD_EXTEND(int8_t, int64_t);
 
 #undef CUVS_INST_IVF_PQ_BUILD_EXTEND
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu
index df6e12b7e..b25330ca3 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu
@@ -25,42 +25,9 @@
 
 #include <cuvs/neighbors/ivf_pq.hpp>
 
-#include "../ivf_pq_build.cuh"
+#include "ivf_pq_build_extent_inst.cuh"
 
 namespace cuvs::neighbors::ivf_pq {
-
-#define CUVS_INST_IVF_PQ_BUILD_EXTEND(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)                    \
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
-  {                                                                                               \
-    return cuvs::neighbors::ivf_pq::detail::build(handle, params, 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)                                           \
-  {                                                                                               \
-    cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, 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)                             \
-    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
-  {                                                                                               \
-    return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, 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)                                          \
-  {                                                                                               \
-    cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, idx);               \
-  }
 CUVS_INST_IVF_PQ_BUILD_EXTEND(uint8_t, int64_t);
 
 #undef CUVS_INST_IVF_PQ_BUILD_EXTEND
diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
index 5536d6edb..42c741d97 100644
--- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
+++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
@@ -1659,13 +1659,14 @@ auto extend(raft::resources const& handle,
   return ext_index;
 }
 
-template <typename T, typename IdxT>
+template <typename T, typename IdxT, typename accessor>
 auto build(raft::resources const& handle,
            const index_params& params,
-           const T* dataset,
-           IdxT n_rows,
-           uint32_t dim) -> index<IdxT>
+           raft::mdspan<const T, raft::matrix_extent<IdxT>, raft::row_major, accessor> dataset)
+  -> index<IdxT>
 {
+  IdxT n_rows = dataset.extent(0);
+  IdxT dim    = dataset.extent(1);
   raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope(
     "ivf_pq::build(%zu, %u)", size_t(n_rows), dim);
   static_assert(std::is_same_v<T, float> || std::is_same_v<T, half> || std::is_same_v<T, uint8_t> ||
@@ -1700,7 +1701,7 @@ auto build(raft::resources const& handle,
     if constexpr (std::is_same_v<T, float>) {
       RAFT_CUDA_TRY(cudaMemcpy2DAsync(trainset.data(),
                                       sizeof(T) * index.dim(),
-                                      dataset,
+                                      dataset.data_handle(),
                                       sizeof(T) * index.dim() * trainset_ratio,
                                       sizeof(T) * index.dim(),
                                       n_rows_train,
@@ -1709,7 +1710,7 @@ auto build(raft::resources const& handle,
     } else {
       size_t dim = index.dim();
       cudaPointerAttributes dataset_attr;
-      RAFT_CUDA_TRY(cudaPointerGetAttributes(&dataset_attr, dataset));
+      RAFT_CUDA_TRY(cudaPointerGetAttributes(&dataset_attr, dataset.data_handle()));
       if (dataset_attr.devicePointer != nullptr) {
         // data is available on device: just run the kernel to raft::copy and map the data
         auto p = reinterpret_cast<T*>(dataset_attr.devicePointer);
@@ -1728,7 +1729,7 @@ auto build(raft::resources const& handle,
         // T at the end of float rows.
         RAFT_CUDA_TRY(cudaMemcpy2DAsync(trainset_tmp,
                                         sizeof(float) * index.dim(),
-                                        dataset,
+                                        dataset.data_handle(),
                                         sizeof(T) * index.dim() * trainset_ratio,
                                         sizeof(T) * index.dim(),
                                         n_rows_train,
@@ -1809,36 +1810,24 @@ auto build(raft::resources const& handle,
 
   // add the data if necessary
   if (params.add_data_on_build) {
-    detail::extend<T, IdxT>(handle, &index, dataset, nullptr, n_rows);
+    detail::extend<T, IdxT>(handle, &index, dataset.data_handle(), nullptr, n_rows);
   }
   return index;
 }
 
-template <typename T, typename IdxT>
-auto build(raft::resources const& handle,
-           const index_params& params,
-           raft::device_matrix_view<const T, IdxT, raft::row_major> dataset) -> index<IdxT>
-{
-  IdxT n_rows = dataset.extent(0);
-  IdxT dim    = dataset.extent(1);
-  return build(handle, params, dataset.data_handle(), n_rows, dim);
-}
-
-template <typename T, typename IdxT>
+template <typename T, typename IdxT, typename accessor>
 void build(raft::resources const& handle,
            const index_params& params,
-           raft::device_matrix_view<const T, IdxT, raft::row_major> dataset,
+           raft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, accessor> dataset,
            index<IdxT>* index)
 {
-  IdxT n_rows = dataset.extent(0);
-  IdxT dim    = dataset.extent(1);
-  *index      = build(handle, params, dataset.data_handle(), n_rows, dim);
+  *index = build(handle, params, dataset);
 }
 
-template <typename T, typename IdxT>
+template <typename T, typename IdxT, typename accessor>
 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,
+            raft::mdspan<const T, IdxT, raft::row_major, accessor> new_vectors,
+            raft::mdspan<const IdxT, IdxT, raft::row_major, accessor> new_indices,
             const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index) -> index<IdxT>
 {
   ASSERT(new_vectors.extent(1) == orig_index.dim(),
@@ -1857,10 +1846,10 @@ auto extend(raft::resources const& handle,
                 n_rows);
 }
 
-template <typename T, typename IdxT>
+template <typename T, typename IdxT, typename accessor>
 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,
+            raft::mdspan<const T, IdxT, raft::row_major, accessor> new_vectors,
+            std::optional<raft::mdspan<const IdxT, IdxT, raft::row_major, accessor>> new_indices,
             index<IdxT>* index)
 {
   ASSERT(new_vectors.extent(1) == index->dim(),

From 9f3118240abaec4f228d7d1e17f52e5e4db81b36 Mon Sep 17 00:00:00 2001
From: Tamas Bela Feher <tfeher@nvidia.com>
Date: Fri, 24 May 2024 15:54:57 +0200
Subject: [PATCH 2/4] Fix extend API

---
 cpp/include/cuvs/neighbors/ivf_pq.hpp         | 171 ++++++++++++++++++
 .../ivf_pq/detail/generate_ivf_pq.py          |   2 +-
 .../ivf_pq_build_extend_float_int64_t.cu      |   2 +-
 .../detail/ivf_pq_build_extend_inst.cuh       |  31 ++--
 .../ivf_pq_build_extend_int8_t_int64_t.cu     |   2 +-
 .../ivf_pq_build_extend_uint8_t_int64_t.cu    |   2 +-
 cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh     |  24 ++-
 7 files changed, 205 insertions(+), 29 deletions(-)

diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp
index d6dda82ac..1965fe15f 100644
--- a/cpp/include/cuvs/neighbors/ivf_pq.hpp
+++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp
@@ -912,6 +912,177 @@ void extend(raft::resources const& handle,
             raft::device_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors,
             std::optional<raft::device_vector_view<const int64_t, int64_t>> new_indices,
             cuvs::neighbors::ivf_pq::index<int64_t>* idx);
+
+/**
+ * @brief Extend the index with the new data.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   ivf_pq::index_params index_params;
+ *   index_params.add_data_on_build = false;      // don't populate index on build
+ *   index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
+ *   // train the index from a [N, D] dataset
+ *   auto index_empty = ivf_pq::build(handle, index_params, dataset);
+ *   // fill the index with the data
+ *   std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
+ *   auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()]
+ * @param[in] new_indices a device vector view to a vector of indices [n_rows].
+ *    If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt`
+ *    here to imply a continuous range `[0...n_rows)`.
+ * @param[inout] idx
+ */
+auto extend(raft::resources const& handle,
+            raft::host_matrix_view<const float, int64_t, raft::row_major> new_vectors,
+            std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
+            const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
+  -> cuvs::neighbors::ivf_pq::index<int64_t>;
+
+/**
+ * @brief Extend the index with the new data.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   ivf_pq::index_params index_params;
+ *   index_params.add_data_on_build = false;      // don't populate index on build
+ *   index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
+ *   // train the index from a [N, D] dataset
+ *   auto index_empty = ivf_pq::build(handle, index_params, dataset);
+ *   // fill the index with the data
+ *   std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
+ *   ivf_pq::extend(handle, new_vectors, no_op, &index_empty);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()]
+ * @param[in] new_indices a device vector view to a vector of indices [n_rows].
+ *    If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt`
+ *    here to imply a continuous range `[0...n_rows)`.
+ * @param[inout] idx
+ */
+void extend(raft::resources const& handle,
+            raft::host_matrix_view<const float, int64_t, raft::row_major> new_vectors,
+            std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
+            cuvs::neighbors::ivf_pq::index<int64_t>* idx);
+
+/**
+ * @brief Extend the index with the new data.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   ivf_pq::index_params index_params;
+ *   index_params.add_data_on_build = false;      // don't populate index on build
+ *   index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
+ *   // train the index from a [N, D] dataset
+ *   auto index_empty = ivf_pq::build(handle, index_params, dataset);
+ *   // fill the index with the data
+ *   std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
+ *   auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()]
+ * @param[in] new_indices a device vector view to a vector of indices [n_rows].
+ *    If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt`
+ *    here to imply a continuous range `[0...n_rows)`.
+ * @param[inout] idx
+ */
+auto extend(raft::resources const& handle,
+            raft::host_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors,
+            std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
+            const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
+  -> cuvs::neighbors::ivf_pq::index<int64_t>;
+
+/**
+ * @brief Extend the index with the new data.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   ivf_pq::index_params index_params;
+ *   index_params.add_data_on_build = false;      // don't populate index on build
+ *   index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
+ *   // train the index from a [N, D] dataset
+ *   auto index_empty = ivf_pq::build(handle, index_params, dataset);
+ *   // fill the index with the data
+ *   std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
+ *   ivf_pq::extend(handle, new_vectors, no_op, &index_empty);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()]
+ * @param[in] new_indices a device vector view to a vector of indices [n_rows].
+ *    If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt`
+ *    here to imply a continuous range `[0...n_rows)`.
+ * @param[inout] idx
+ */
+void extend(raft::resources const& handle,
+            raft::host_matrix_view<const int8_t, int64_t, raft::row_major> new_vectors,
+            std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
+            cuvs::neighbors::ivf_pq::index<int64_t>* idx);
+
+/**
+ * @brief Extend the index with the new data.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   ivf_pq::index_params index_params;
+ *   index_params.add_data_on_build = false;      // don't populate index on build
+ *   index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
+ *   // train the index from a [N, D] dataset
+ *   auto index_empty = ivf_pq::build(handle, index_params, dataset);
+ *   // fill the index with the data
+ *   std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
+ *   auto index = ivf_pq::extend(handle, new_vectors, no_op, index_empty);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()]
+ * @param[in] new_indices a device vector view to a vector of indices [n_rows].
+ *    If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt`
+ *    here to imply a continuous range `[0...n_rows)`.
+ * @param[inout] idx
+ */
+auto extend(raft::resources const& handle,
+            raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors,
+            std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
+            const cuvs::neighbors::ivf_pq::index<int64_t>& idx)
+  -> cuvs::neighbors::ivf_pq::index<int64_t>;
+
+/**
+ * @brief Extend the index with the new data.
+ *
+ * Usage example:
+ * @code{.cpp}
+ *   using namespace cuvs::neighbors;
+ *   ivf_pq::index_params index_params;
+ *   index_params.add_data_on_build = false;      // don't populate index on build
+ *   index_params.kmeans_trainset_fraction = 1.0; // use whole dataset for kmeans training
+ *   // train the index from a [N, D] dataset
+ *   auto index_empty = ivf_pq::build(handle, index_params, dataset);
+ *   // fill the index with the data
+ *   std::optional<raft::host_vector_view<const IdxT, IdxT>> no_op = std::nullopt;
+ *   ivf_pq::extend(handle, new_vectors, no_op, &index_empty);
+ * @endcode
+ *
+ * @param[in] handle
+ * @param[in] new_vectors a device matrix view to a row-major matrix [n_rows, idx.dim()]
+ * @param[in] new_indices a device vector view to a vector of indices [n_rows].
+ *    If the original index is empty (`idx.size() == 0`), you can pass `std::nullopt`
+ *    here to imply a continuous range `[0...n_rows)`.
+ * @param[inout] idx
+ */
+void extend(raft::resources const& handle,
+            raft::host_matrix_view<const uint8_t, int64_t, raft::row_major> new_vectors,
+            std::optional<raft::host_vector_view<const int64_t, int64_t>> new_indices,
+            cuvs::neighbors::ivf_pq::index<int64_t>* idx);
 /**
  * @}
  */
diff --git a/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py b/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py
index 1d3704fac..878c7ee21 100644
--- a/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py
+++ b/cpp/src/neighbors/ivf_pq/detail/generate_ivf_pq.py
@@ -41,7 +41,7 @@
 """
 
 build_include_macro = """
-#include "ivf_pq_build_extent_inst.cuh"
+#include "ivf_pq_build_extend_inst.cuh"
 """
 search_include_macro = """
 #include "../ivf_pq_search.cuh"
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu
index 0269481fb..65c443d88 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_float_int64_t.cu
@@ -25,7 +25,7 @@
 
 #include <cuvs/neighbors/ivf_pq.hpp>
 
-#include "ivf_pq_build_extent_inst.cuh"
+#include "ivf_pq_build_extend_inst.cuh"
 
 namespace cuvs::neighbors::ivf_pq {
 CUVS_INST_IVF_PQ_BUILD_EXTEND(float, int64_t);
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
index e79779e59..6060d2b7a 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
@@ -28,7 +28,7 @@ namespace cuvs::neighbors::ivf_pq {
 #define CUVS_INST_IVF_PQ_BUILD_EXTEND(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)                    \
+             raft::device_matrix_view<const T, int64_t, raft::row_major> dataset)                 \
     -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
   {                                                                                               \
     return cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset);                       \
@@ -36,7 +36,7 @@ namespace cuvs::neighbors::ivf_pq {
                                                                                                   \
   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,                    \
+             raft::device_matrix_view<const T, int64_t, raft::row_major> dataset,                 \
              cuvs::neighbors::ivf_pq::index<IdxT>* idx)                                           \
   {                                                                                               \
     cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, idx);                         \
@@ -44,7 +44,7 @@ namespace cuvs::neighbors::ivf_pq {
                                                                                                   \
   auto build(raft::resources const& handle,                                                       \
              const cuvs::neighbors::ivf_pq::index_params& params,                                 \
-             raft::host_matrix_view<const T, IdxT, raft::row_major> dataset)                      \
+             raft::host_matrix_view<const T, int64_t, raft::row_major> dataset)                   \
     -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
   {                                                                                               \
     return cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset);                       \
@@ -52,30 +52,30 @@ namespace cuvs::neighbors::ivf_pq {
                                                                                                   \
   void build(raft::resources const& handle,                                                       \
              const cuvs::neighbors::ivf_pq::index_params& params,                                 \
-             raft::host_matrix_view<const T, IdxT, raft::row_major> dataset,                      \
+             raft::host_matrix_view<const T, int64_t, raft::row_major> dataset,                   \
              cuvs::neighbors::ivf_pq::index<IdxT>* idx)                                           \
   {                                                                                               \
     cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset, 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)                             \
+  auto extend(                                                                                    \
+    raft::resources const& handle,                                                                \
+    raft::device_matrix_view<const T, int64_t, raft::row_major> new_vectors,                      \
+    std::optional<raft::device_vector_view<const IdxT, int64_t, raft::row_major>> new_indices,    \
+    const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index)                                       \
     -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
   {                                                                                               \
     return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, 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,              \
+              raft::device_matrix_view<const T, int64_t, raft::row_major> new_vectors,            \
+              std::optional<raft::device_vector_view<const IdxT, int64_t>> new_indices,           \
               cuvs::neighbors::ivf_pq::index<IdxT>* idx)                                          \
   {                                                                                               \
     cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, idx);               \
   }                                                                                               \
   auto extend(raft::resources const& handle,                                                      \
-              raft::host_matrix_view<const T, IdxT, raft::row_major> new_vectors,                 \
-              std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,                \
+              raft::host_matrix_view<const T, int64_t, raft::row_major> new_vectors,              \
+              std::optional<raft::host_vector_view<const IdxT, int64_t>> new_indices,             \
               const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index)                             \
     -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
   {                                                                                               \
@@ -83,10 +83,11 @@ namespace cuvs::neighbors::ivf_pq {
   }                                                                                               \
                                                                                                   \
   void extend(raft::resources const& handle,                                                      \
-              raft::host_matrix_view<const T, IdxT, raft::row_major> new_vectors,                 \
-              std::optional<raft::host_vector_view<const IdxT, IdxT>> new_indices,                \
+              raft::host_matrix_view<const T, int64_t, raft::row_major> new_vectors,              \
+              std::optional<raft::host_vector_view<const IdxT, int64_t>> new_indices,             \
               cuvs::neighbors::ivf_pq::index<IdxT>* idx)                                          \
   {                                                                                               \
     cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, idx);               \
   }
+
 }  // namespace cuvs::neighbors::ivf_pq
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu
index 943fffaf2..048560b5a 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_int8_t_int64_t.cu
@@ -25,7 +25,7 @@
 
 #include <cuvs/neighbors/ivf_pq.hpp>
 
-#include "ivf_pq_build_extent_inst.cuh"
+#include "ivf_pq_build_extend_inst.cuh"
 
 namespace cuvs::neighbors::ivf_pq {
 CUVS_INST_IVF_PQ_BUILD_EXTEND(int8_t, int64_t);
diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu
index b25330ca3..755fb9f8d 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_uint8_t_int64_t.cu
@@ -25,7 +25,7 @@
 
 #include <cuvs/neighbors/ivf_pq.hpp>
 
-#include "ivf_pq_build_extent_inst.cuh"
+#include "ivf_pq_build_extend_inst.cuh"
 
 namespace cuvs::neighbors::ivf_pq {
 CUVS_INST_IVF_PQ_BUILD_EXTEND(uint8_t, int64_t);
diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
index 42c741d97..3b0e475f4 100644
--- a/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
+++ b/cpp/src/neighbors/ivf_pq/ivf_pq_build.cuh
@@ -1824,11 +1824,13 @@ void build(raft::resources const& handle,
   *index = build(handle, params, dataset);
 }
 
-template <typename T, typename IdxT, typename accessor>
-auto extend(raft::resources const& handle,
-            raft::mdspan<const T, IdxT, raft::row_major, accessor> new_vectors,
-            raft::mdspan<const IdxT, IdxT, raft::row_major, accessor> new_indices,
-            const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index) -> index<IdxT>
+template <typename T, typename IdxT, typename accessor, typename accessor2>
+auto extend(
+  raft::resources const& handle,
+  raft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, accessor> new_vectors,
+  std::optional<raft::mdspan<const IdxT, raft::vector_extent<int64_t>, raft::row_major, accessor2>>
+    new_indices,
+  const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index) -> index<IdxT>
 {
   ASSERT(new_vectors.extent(1) == orig_index.dim(),
          "new_vectors should have the same dimension as the index");
@@ -1846,11 +1848,13 @@ auto extend(raft::resources const& handle,
                 n_rows);
 }
 
-template <typename T, typename IdxT, typename accessor>
-void extend(raft::resources const& handle,
-            raft::mdspan<const T, IdxT, raft::row_major, accessor> new_vectors,
-            std::optional<raft::mdspan<const IdxT, IdxT, raft::row_major, accessor>> new_indices,
-            index<IdxT>* index)
+template <typename T, typename IdxT, typename accessor, typename accessor2>
+void extend(
+  raft::resources const& handle,
+  raft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, accessor> new_vectors,
+  std::optional<raft::mdspan<const IdxT, raft::vector_extent<int64_t>, raft::row_major, accessor2>>
+    new_indices,
+  index<IdxT>* index)
 {
   ASSERT(new_vectors.extent(1) == index->dim(),
          "new_vectors should have the same dimension as the index");

From 0d214c045d87185f179ea4e2962d6458695f6c34 Mon Sep 17 00:00:00 2001
From: Tamas Bela Feher <tfeher@nvidia.com>
Date: Fri, 24 May 2024 19:51:10 +0200
Subject: [PATCH 3/4] fix style

---
 .../neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh  | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
index 6060d2b7a..4b963b0b3 100644
--- a/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
+++ b/cpp/src/neighbors/ivf_pq/detail/ivf_pq_build_extend_inst.cuh
@@ -29,7 +29,7 @@ namespace cuvs::neighbors::ivf_pq {
   auto build(raft::resources const& handle,                                                       \
              const cuvs::neighbors::ivf_pq::index_params& params,                                 \
              raft::device_matrix_view<const T, int64_t, raft::row_major> dataset)                 \
-    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
   {                                                                                               \
     return cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset);                       \
   }                                                                                               \
@@ -45,7 +45,7 @@ namespace cuvs::neighbors::ivf_pq {
   auto build(raft::resources const& handle,                                                       \
              const cuvs::neighbors::ivf_pq::index_params& params,                                 \
              raft::host_matrix_view<const T, int64_t, raft::row_major> dataset)                   \
-    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
   {                                                                                               \
     return cuvs::neighbors::ivf_pq::detail::build(handle, params, dataset);                       \
   }                                                                                               \
@@ -62,7 +62,7 @@ namespace cuvs::neighbors::ivf_pq {
     raft::device_matrix_view<const T, int64_t, raft::row_major> new_vectors,                      \
     std::optional<raft::device_vector_view<const IdxT, int64_t, raft::row_major>> new_indices,    \
     const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index)                                       \
-    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
   {                                                                                               \
     return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, orig_index); \
   }                                                                                               \
@@ -77,7 +77,7 @@ namespace cuvs::neighbors::ivf_pq {
               raft::host_matrix_view<const T, int64_t, raft::row_major> new_vectors,              \
               std::optional<raft::host_vector_view<const IdxT, int64_t>> new_indices,             \
               const cuvs::neighbors::ivf_pq::index<IdxT>& orig_index)                             \
-    -> cuvs::neighbors::ivf_pq::index<IdxT>                                                       \
+    ->cuvs::neighbors::ivf_pq::index<IdxT>                                                        \
   {                                                                                               \
     return cuvs::neighbors::ivf_pq::detail::extend(handle, new_vectors, new_indices, orig_index); \
   }                                                                                               \

From 55924cf96ae085766f42512168a05def5e442d93 Mon Sep 17 00:00:00 2001
From: Tamas Bela Feher <tfeher@nvidia.com>
Date: Mon, 27 May 2024 01:01:30 +0200
Subject: [PATCH 4/4] Fix test compilation issue

---
 cpp/include/cuvs/neighbors/ivf_pq.hpp         | 69 ++++++++++++++++---
 .../neighbors/ivf_pq/ivf_pq_deserialize.cu    | 15 +++-
 cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cu  | 15 +++-
 cpp/test/neighbors/ann_ivf_pq.cuh             | 18 ++---
 4 files changed, 91 insertions(+), 26 deletions(-)

diff --git a/cpp/include/cuvs/neighbors/ivf_pq.hpp b/cpp/include/cuvs/neighbors/ivf_pq.hpp
index 1965fe15f..61a0de275 100644
--- a/cpp/include/cuvs/neighbors/ivf_pq.hpp
+++ b/cpp/include/cuvs/neighbors/ivf_pq.hpp
@@ -1321,9 +1321,30 @@ void search_with_filtering(
  * @{
  */
 /**
- * Save the index to file.
+ * Serialize the index to an output string.
+ *
+ * @code{.cpp}
+ * #include <raft/core/resources.hpp>
+ *
+ * raft::resources handle;
+ *
+ * // create an input string
+ * std::string str
+ * // create an index with `auto index = ivf_pq::build(...);`
+ * cuvs::serialize(handle, str, index);
+ * @endcode
+ *
+ * @param[in] handle the raft handle
+ * @param[out] str output string
+ * @param[in] index IVF-PQ index
  *
- * Experimental, both the API and the serialization format are subject to change.
+ */
+void serialize(raft::resources const& handle,
+               std::string& str,
+               const cuvs::neighbors::ivf_pq::index<int64_t>& index);
+
+/**
+ * Save the index to file.
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
@@ -1341,14 +1362,38 @@ void search_with_filtering(
  * @param[in] index IVF-PQ index
  *
  */
-void serialize(raft::resources const& handle,
-               std::string& filename,
-               const cuvs::neighbors::ivf_pq::index<int64_t>& index);
+void serialize_file(raft::resources const& handle,
+                    const std::string& filename,
+                    const cuvs::neighbors::ivf_pq::index<int64_t>& index);
 
 /**
- * Load index from file.
+ * Load index from input string.
+ *
+ * @code{.cpp}
+ * #include <raft/core/resources.hpp>
+ *
+ * raft::resources handle;
+ *
+ * std::string str = ...
+ *
+ * using IdxT = int64_t; // type of the index
+ * // create an empty index
+ * cuvs::neighbors::ivf_pq::index<IdxT> index(handl, index_params, dim);
+ *
+ * cuvs::deserialize(handle, filename, &index);
+ * @endcode
+ *
+ * @param[in] handle the raft handle
+ * @param[in] str the name of the file that stores the index
+ * @param[out] index IVF-PQ index
  *
- * Experimental, both the API and the serialization format are subject to change.
+ */
+
+void deserialize(raft::resources const& handle,
+                 const std::string& str,
+                 cuvs::neighbors::ivf_pq::index<int64_t>* index);
+/**
+ * Load index from file.
  *
  * @code{.cpp}
  * #include <raft/core/resources.hpp>
@@ -1358,7 +1403,9 @@ void serialize(raft::resources const& handle,
  * // create a string with a filepath
  * std::string filename("/path/to/index");
  * using IdxT = int64_t; // type of the index
- * // create an empty index with `ivf_pq::index<IdxT> index(handle, index_params, dim);`
+ * // create an empty index with
+ * ivf_pq::index<IdxT> index(handle, index_params, dim);
+ *
  * cuvs::deserialize(handle, filename, &index);
  * @endcode
  *
@@ -1367,9 +1414,9 @@ void serialize(raft::resources const& handle,
  * @param[out] index IVF-PQ index
  *
  */
-void deserialize(raft::resources const& handle,
-                 const std::string& filename,
-                 cuvs::neighbors::ivf_pq::index<int64_t>* index);
+void deserialize_file(raft::resources const& handle,
+                      const std::string& filename,
+                      cuvs::neighbors::ivf_pq::index<int64_t>* index);
 /**
  * @}
  */
diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_deserialize.cu b/cpp/src/neighbors/ivf_pq/ivf_pq_deserialize.cu
index 9515cd2d5..7827e7892 100644
--- a/cpp/src/neighbors/ivf_pq/ivf_pq_deserialize.cu
+++ b/cpp/src/neighbors/ivf_pq/ivf_pq_deserialize.cu
@@ -16,15 +16,24 @@
 
 #include "ivf_pq_serialize.cuh"
 #include <cuvs/neighbors/ivf_pq.hpp>
+#include <sstream>
 
 namespace cuvs::neighbors::ivf_pq {
 
-void deserialize(raft::resources const& handle,
-                 const std::string& filename,
-                 cuvs::neighbors::ivf_pq::index<int64_t>* index)
+void deserialize_file(raft::resources const& handle,
+                      const std::string& filename,
+                      cuvs::neighbors::ivf_pq::index<int64_t>* index)
 {
   if (!index) { RAFT_FAIL("Invalid index pointer"); }
   *index = cuvs::neighbors::ivf_pq::detail::deserialize<int64_t>(handle, filename);
 }
 
+void deserialize(raft::resources const& handle,
+                 const std::string& str,
+                 cuvs::neighbors::ivf_pq::index<int64_t>* index)
+{
+  if (!index) { RAFT_FAIL("Invalid index pointer"); }
+  std::istringstream is(str);
+  *index = cuvs::neighbors::ivf_pq::detail::deserialize<int64_t>(handle, is);
+}
 }  // namespace cuvs::neighbors::ivf_pq
\ No newline at end of file
diff --git a/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cu b/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cu
index 0fb71d031..f0214f4bb 100644
--- a/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cu
+++ b/cpp/src/neighbors/ivf_pq/ivf_pq_serialize.cu
@@ -16,14 +16,23 @@
 
 #include "ivf_pq_serialize.cuh"
 #include <cuvs/neighbors/ivf_pq.hpp>
+#include <sstream>
 
 namespace cuvs::neighbors::ivf_pq {
 
-void serialize(raft::resources const& handle,
-               std::string& filename,
-               const cuvs::neighbors::ivf_pq::index<int64_t>& index)
+void serialize_file(raft::resources const& handle,
+                    const std::string& filename,
+                    const cuvs::neighbors::ivf_pq::index<int64_t>& index)
 {
   cuvs::neighbors::ivf_pq::detail::serialize(handle, filename, index);
 }
 
+void serialize(raft::resources const& handle,
+               std::string& str,
+               const cuvs::neighbors::ivf_pq::index<int64_t>& index)
+{
+  std::ostringstream os;
+  cuvs::neighbors::ivf_pq::detail::serialize(handle, os, index);
+  str = os.str();
+}
 }  // namespace cuvs::neighbors::ivf_pq
\ No newline at end of file
diff --git a/cpp/test/neighbors/ann_ivf_pq.cuh b/cpp/test/neighbors/ann_ivf_pq.cuh
index fda32c6df..f716a8efe 100644
--- a/cpp/test/neighbors/ann_ivf_pq.cuh
+++ b/cpp/test/neighbors/ann_ivf_pq.cuh
@@ -204,7 +204,7 @@ class ivf_pq_test : public ::testing::TestWithParam<ivf_pq_inputs> {
     ipams.add_data_on_build = true;
 
     auto index_view =
-      raft::make_device_matrix_view<DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
+      raft::make_device_matrix_view<const DataT, int64_t>(database.data(), ps.num_db_vecs, ps.dim);
     return cuvs::neighbors::ivf_pq::build(handle_, ipams, index_view);
   }
 
@@ -224,16 +224,16 @@ class ivf_pq_test : public ::testing::TestWithParam<ivf_pq_inputs> {
     ipams.add_data_on_build = false;
 
     auto database_view =
-      raft::make_device_matrix_view<DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
+      raft::make_device_matrix_view<const DataT, int64_t>(database.data(), ps.num_db_vecs, ps.dim);
     auto idx = cuvs::neighbors::ivf_pq::build(handle_, ipams, database_view);
 
-    auto vecs_2_view = raft::make_device_matrix_view<DataT, IdxT>(vecs_2, size_2, ps.dim);
-    auto inds_2_view = raft::make_device_vector_view<IdxT, IdxT>(inds_2, size_2);
+    auto vecs_2_view = raft::make_device_matrix_view<const DataT, int64_t>(vecs_2, size_2, ps.dim);
+    auto inds_2_view = raft::make_device_vector_view<const IdxT, int64_t>(inds_2, size_2);
     cuvs::neighbors::ivf_pq::extend(handle_, vecs_2_view, inds_2_view, &idx);
 
     auto vecs_1_view =
-      raft::make_device_matrix_view<DataT, IdxT, raft::row_major>(vecs_1, size_1, ps.dim);
-    auto inds_1_view = raft::make_device_vector_view<const IdxT, IdxT>(inds_1, size_1);
+      raft::make_device_matrix_view<const DataT, int64_t, raft::row_major>(vecs_1, size_1, ps.dim);
+    auto inds_1_view = raft::make_device_vector_view<const IdxT, int64_t>(inds_1, size_1);
     cuvs::neighbors::ivf_pq::extend(handle_, vecs_1_view, inds_1_view, &idx);
     return idx;
   }
@@ -241,9 +241,9 @@ class ivf_pq_test : public ::testing::TestWithParam<ivf_pq_inputs> {
   auto build_serialize()
   {
     std::string filename = "ivf_pq_index";
-    cuvs::neighbors::ivf_pq::serialize(handle_, filename, build_only());
+    cuvs::neighbors::ivf_pq::serialize_file(handle_, filename, build_only());
     cuvs::neighbors::ivf_pq::index<IdxT> index(handle_, ps.index_params, ps.dim);
-    cuvs::neighbors::ivf_pq::deserialize(handle_, filename, &index);
+    cuvs::neighbors::ivf_pq::deserialize_file(handle_, filename, &index);
     return index;
   }
 
@@ -564,7 +564,7 @@ class ivf_pq_filter_test : public ::testing::TestWithParam<ivf_pq_inputs> {
     ipams.add_data_on_build = true;
 
     auto index_view =
-      raft::make_device_matrix_view<DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
+      raft::make_device_matrix_view<const DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
     return cuvs::neighbors::ivf_pq::build(handle_, ipams, index_view);
   }