From 328a1794569c0b1986c4dec99ad5c42917f8360f Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Tue, 7 Feb 2023 02:38:58 +0100
Subject: [PATCH 01/33] Initial index splitting

---
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 103 +++++++++++++++++-
 1 file changed, 98 insertions(+), 5 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index d234822a23..b67cf38139 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -65,6 +65,35 @@ struct search_params : ann::search_params {
 static_assert(std::is_aggregate_v<index_params>);
 static_assert(std::is_aggregate_v<search_params>);
 
+/** The data for a single list (cluster). */
+template <typename T, typename IdxT = uint32_t>
+struct list_data {
+  /** Cluster data. */
+  device_mdarray<T, extent_2d<IdxT>, row_major> data;
+  /** Source indices. */
+  device_mdarray<IdxT, extent_1d<IdxT>, row_major> indices;
+  /** The actual size of the content. */
+  std::atomic<IdxT> size;
+
+  list_data(raft::device_resources const& handle, IdxT n_rows, uint32_t dim)
+    : size{n_rows}
+  {
+    auto capacity = round_up_safe<IdxT>(bound_by_power_of_two<IdxT>(size), kIndexGroupSize);
+    try {
+      data = make_device_mdarray<T>(handle, make_extents<IdxT>(capacity, dim));
+      indices = make_device_mdarray<IdxT>(handle, make_extents<IdxT>(capacity));
+    } catch (std::bad_alloc& e) {
+      RAFT_FAIL(
+        "ivf-flat: failed to allocate a big enough index list to hold all data "
+        "(requested size: %zu records, selected capacity: %zu records). "
+        "Allocator exception: %s",
+        size_t(size),
+        size_t(capacity),
+        e.what());
+    }
+  }
+};
+
 /**
  * @brief IVF-flat index.
  *
@@ -78,6 +107,12 @@ struct index : ann::index {
                 "IdxT must be able to represent all values of uint32_t");
 
  public:
+  /**
+   * Default value filled in the `indices()` array.
+   * One may encounter it trying to access a record within a cluster that is outside of the
+   * `list_sizes()` bound (due to the record alignment `kIndexGroupSize`).
+   */
+  constexpr static IdxT kInvalidRecord = std::numeric_limits<IdxT>::max() - 1;
   /**
    * Vectorized load/store size in elements, determines the size of interleaved data chunks.
    *
@@ -252,8 +287,60 @@ struct index : ann::index {
    * Replace the content of the index with new uninitialized mdarrays to hold the indicated amount
    * of data.
    */
-  void allocate(raft::device_resources const& handle, IdxT index_size)
+  void resize_list(raft::device_resources const& handle, uint32_t label, uint32_t list_size)
   {
+    uint32_t prev_size = 0;
+    auto& list         = lists()(label);
+    bool skip_resize   = false;
+    if (list) {
+      copy(&prev_size, &list_sizes()(label), 1, handle.get_stream());
+      handle.sync_stream();
+      if (list_size <= list->indices.extent(0)) {
+        auto shared_list_size = prev_size;
+        if (list_size <= prev_size ||
+            list->size.compare_exchange_strong(shared_list_size, list_size)) {
+          // We don't need to resize the list if:
+          //  1. The list exists
+          //  2. The new size fits in the list
+          //  3. The list doesn't grow or no-one else has grown it yet
+          skip_resize = true;
+        }
+      }
+    }
+    // Sic! We're writing the min(list_size, prev_size)
+    //       to keep the number of _valid_ records after update
+    const auto new_list_size = std::min(list_size, prev_size);
+    raft::copy(&list_sizes()(label), &new_list_size, 1, handle.get_stream());
+    if (skip_resize) { return; }
+    auto new_list = new list_data<T, IdxT>(handle, list_size, dim());
+    if (prev_size > 0) {
+      auto copied_data_extents = make_extents<size_t>(prev_size, dim());
+      auto copied_view         = make_mdspan<T, size_t, row_major, false, true>(
+        new_list->data.data_handle(), copied_data_extents);
+      copy(copied_view.data_handle(),
+           list->data.data_handle(),
+           copied_view.size(),
+           handle.get_stream());
+      copy(new_list->indices.data_handle(),
+           list->indices.data_handle(),
+           prev_size,
+           handle.get_stream());
+    }
+    // swap the shared pointer content with the new list
+    list.reset(new_list);
+    // fill unused index spaces with placeholder values for easier debugging
+    thrust::fill_n(handle.get_thrust_policy(),
+                   list->indices.data_handle() + prev_size,
+                   list->indices.size() - prev_size,
+                   kInvalidRecord);
+    // keep the device pointers updated
+    const auto new_data_ptr = list->data.data_handle();
+    copy(&(data_ptrs()(label)), &new_data_ptr, 1, handle.get_stream());
+    const auto new_inds_ptr = list->indices.data_handle();
+    copy(&(inds_ptrs()(label)), &new_inds_ptr, 1, handle.get_stream());
+
+
+
     data_    = make_device_mdarray<T>(handle, make_extents<IdxT>(index_size, dim()));
     indices_ = make_device_mdarray<IdxT>(handle, make_extents<IdxT>(index_size));
 
@@ -278,13 +365,19 @@ struct index : ann::index {
   uint32_t veclen_;
   raft::distance::DistanceType metric_;
   bool adaptive_centers_;
-  device_mdarray<T, extent_2d<IdxT>, row_major> data_;
-  device_mdarray<IdxT, extent_1d<IdxT>, row_major> indices_;
-  device_mdarray<uint32_t, extent_1d<uint32_t>, row_major> list_sizes_;
-  device_mdarray<IdxT, extent_1d<uint32_t>, row_major> list_offsets_;
+  host_mdarray<std::shared_ptr<list_data<T, IdxT>, extent_1d<uint32_t>, row_major> lists_;
+  //device_mdarray<T, extent_2d<IdxT>, row_major> data_;
+  //device_mdarray<IdxT, extent_1d<IdxT>, row_major> indices_;
+  //device_mdarray<uint32_t, extent_1d<uint32_t>, row_major> list_sizes_;
+  //device_mdarray<IdxT, extent_1d<uint32_t>, row_major> list_offsets_;
   device_mdarray<float, extent_2d<uint32_t>, row_major> centers_;
   std::optional<device_mdarray<float, extent_1d<uint32_t>, row_major>> center_norms_;
 
+  // Computed members
+  device_mdarray<T*, extent_1d<uint32_t>, row_major> data_ptrs_;
+  device_mdarray<IdxT*, extent_1d<uint32_t>, row_major> inds_ptrs_;
+  host_mdarray<IdxT, extent_1d<uint32_t>, row_major> accum_sorted_sizes_;
+
   /** Throw an error if the index content is inconsistent. */
   void check_consistency()
   {

From a83aca4d312e6c2299e89abb4f4034f705605113 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Tue, 7 Feb 2023 18:59:14 +0100
Subject: [PATCH 02/33] Adapt `extend`

---
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 205 +++++++-----------
 .../spatial/knn/detail/ivf_flat_build.cuh     | 122 ++++++++---
 2 files changed, 175 insertions(+), 152 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index b67cf38139..b384cbcbb3 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -66,22 +66,22 @@ static_assert(std::is_aggregate_v<index_params>);
 static_assert(std::is_aggregate_v<search_params>);
 
 /** The data for a single list (cluster). */
-template <typename T, typename IdxT = uint32_t>
+template <typename T, typename IdxT, typename SizeT = uint32_t>
 struct list_data {
   /** Cluster data. */
-  device_mdarray<T, extent_2d<IdxT>, row_major> data;
+  device_mdarray<T, extent_2d<SizeT>, row_major> data;
   /** Source indices. */
-  device_mdarray<IdxT, extent_1d<IdxT>, row_major> indices;
+  device_mdarray<IdxT, extent_1d<SizeT>, row_major> indices;
   /** The actual size of the content. */
-  std::atomic<IdxT> size;
+  std::atomic<SizeT> size;
 
-  list_data(raft::device_resources const& handle, IdxT n_rows, uint32_t dim)
+  list_data(raft::device_resources const& handle, SizeT n_rows, uint32_t dim)
     : size{n_rows}
   {
-    auto capacity = round_up_safe<IdxT>(bound_by_power_of_two<IdxT>(size), kIndexGroupSize);
+    auto capacity = round_up_safe<SizeT>(bound_by_power_of_two<SizeT>(size), kIndexGroupSize);
     try {
-      data = make_device_mdarray<T>(handle, make_extents<IdxT>(capacity, dim));
-      indices = make_device_mdarray<IdxT>(handle, make_extents<IdxT>(capacity));
+      data = make_device_mdarray<T>(handle, make_extents<SizeT>(capacity, dim));
+      indices = make_device_mdarray<IdxT>(handle, make_extents<SizeT>(capacity));
     } catch (std::bad_alloc& e) {
       RAFT_FAIL(
         "ivf-flat: failed to allocate a big enough index list to hold all data "
@@ -153,27 +153,6 @@ struct index : ann::index {
    *     x[16, 4], x[16, 5], x[17, 4], x[17, 5], ... x[30, 4], x[30, 5],    -    ,    -    ,
    *
    */
-  inline auto data() noexcept -> device_mdspan<T, extent_2d<IdxT>, row_major>
-  {
-    return data_.view();
-  }
-  [[nodiscard]] inline auto data() const noexcept
-    -> device_mdspan<const T, extent_2d<size_t>, row_major>
-  {
-    return data_.view();
-  }
-
-  /** Inverted list indices: ids of items in the source data [size] */
-  inline auto indices() noexcept -> device_mdspan<IdxT, extent_1d<IdxT>, row_major>
-  {
-    return indices_.view();
-  }
-  [[nodiscard]] inline auto indices() const noexcept
-    -> device_mdspan<const IdxT, extent_1d<IdxT>, row_major>
-  {
-    return indices_.view();
-  }
-
   /** Sizes of the lists (clusters) [n_lists] */
   inline auto list_sizes() noexcept -> device_mdspan<uint32_t, extent_1d<uint32_t>, row_major>
   {
@@ -185,20 +164,6 @@ struct index : ann::index {
     return list_sizes_.view();
   }
 
-  /**
-   * Offsets into the lists [n_lists + 1].
-   * The last value contains the total length of the index.
-   */
-  inline auto list_offsets() noexcept -> device_mdspan<IdxT, extent_1d<uint32_t>, row_major>
-  {
-    return list_offsets_.view();
-  }
-  [[nodiscard]] inline auto list_offsets() const noexcept
-    -> device_mdspan<const IdxT, extent_1d<uint32_t>, row_major>
-  {
-    return list_offsets_.view();
-  }
-
   /** k-means cluster centers corresponding to the lists [n_lists, dim] */
   inline auto centers() noexcept -> device_mdspan<float, extent_2d<uint32_t>, row_major>
   {
@@ -238,7 +203,10 @@ struct index : ann::index {
   }
 
   /** Total length of the index. */
-  [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT { return indices_.extent(0); }
+  [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT
+  {
+    return total_size_;
+  }
   /** Dimensionality of the data. */
   [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t
   {
@@ -247,7 +215,7 @@ struct index : ann::index {
   /** Number of clusters/inverted lists. */
   [[nodiscard]] constexpr inline auto n_lists() const noexcept -> uint32_t
   {
-    return centers_.extent(0);
+    return lists_.extent(0);
   }
 
   // Don't allow copying the index for performance reasons (try avoiding copying data)
@@ -267,14 +235,17 @@ struct index : ann::index {
       veclen_(calculate_veclen(dim)),
       metric_(metric),
       adaptive_centers_(adaptive_centers),
-      data_(make_device_mdarray<T>(handle, make_extents<IdxT>(0, dim))),
-      indices_(make_device_mdarray<IdxT>(handle, make_extents<IdxT>(0))),
-      list_sizes_(make_device_mdarray<uint32_t>(handle, make_extents<uint32_t>(n_lists))),
-      list_offsets_(make_device_mdarray<IdxT>(handle, make_extents<uint32_t>(n_lists + 1))),
       centers_(make_device_mdarray<float>(handle, make_extents<uint32_t>(n_lists, dim))),
-      center_norms_(std::nullopt)
+      center_norms_(std::nullopt),
+      lists_{make_host_mdarray<std::shared_ptr<list_data<T, IdxT>>>(make_extents<uint32_t>(n_lists))},
+      list_sizes_{make_device_mdarray<uint32_t>(handle, make_extents<uint32_t>(n_lists))},
+      data_ptrs_{make_device_mdarray<T*>(handle, make_extents<uint32_t>(n_lists))},
+      inds_ptrs_{make_device_mdarray<IdxT*>(handle, make_extents<uint32_t>(n_lists))}
   {
     check_consistency();
+    for (uint32_t i = 0; i < n_lists; i++) {
+      lists_(i) = std::shared_ptr<list_data<T, IdxT>>();
+    }
   }
 
   /** Construct an empty index. It needs to be trained and then populated. */
@@ -283,78 +254,61 @@ struct index : ann::index {
   {
   }
 
-  /**
-   * Replace the content of the index with new uninitialized mdarrays to hold the indicated amount
-   * of data.
-   */
-  void resize_list(raft::device_resources const& handle, uint32_t label, uint32_t list_size)
+  /** Pointers to the inverted lists (clusters) data  [n_lists]. */
+  inline auto data_ptrs() noexcept -> device_mdspan<T*, extent_1d<uint32_t>, row_major>
   {
-    uint32_t prev_size = 0;
-    auto& list         = lists()(label);
-    bool skip_resize   = false;
-    if (list) {
-      copy(&prev_size, &list_sizes()(label), 1, handle.get_stream());
-      handle.sync_stream();
-      if (list_size <= list->indices.extent(0)) {
-        auto shared_list_size = prev_size;
-        if (list_size <= prev_size ||
-            list->size.compare_exchange_strong(shared_list_size, list_size)) {
-          // We don't need to resize the list if:
-          //  1. The list exists
-          //  2. The new size fits in the list
-          //  3. The list doesn't grow or no-one else has grown it yet
-          skip_resize = true;
-        }
-      }
-    }
-    // Sic! We're writing the min(list_size, prev_size)
-    //       to keep the number of _valid_ records after update
-    const auto new_list_size = std::min(list_size, prev_size);
-    raft::copy(&list_sizes()(label), &new_list_size, 1, handle.get_stream());
-    if (skip_resize) { return; }
-    auto new_list = new list_data<T, IdxT>(handle, list_size, dim());
-    if (prev_size > 0) {
-      auto copied_data_extents = make_extents<size_t>(prev_size, dim());
-      auto copied_view         = make_mdspan<T, size_t, row_major, false, true>(
-        new_list->data.data_handle(), copied_data_extents);
-      copy(copied_view.data_handle(),
-           list->data.data_handle(),
-           copied_view.size(),
-           handle.get_stream());
-      copy(new_list->indices.data_handle(),
-           list->indices.data_handle(),
-           prev_size,
-           handle.get_stream());
-    }
-    // swap the shared pointer content with the new list
-    list.reset(new_list);
-    // fill unused index spaces with placeholder values for easier debugging
-    thrust::fill_n(handle.get_thrust_policy(),
-                   list->indices.data_handle() + prev_size,
-                   list->indices.size() - prev_size,
-                   kInvalidRecord);
-    // keep the device pointers updated
-    const auto new_data_ptr = list->data.data_handle();
-    copy(&(data_ptrs()(label)), &new_data_ptr, 1, handle.get_stream());
-    const auto new_inds_ptr = list->indices.data_handle();
-    copy(&(inds_ptrs()(label)), &new_inds_ptr, 1, handle.get_stream());
-
-
+    return data_ptrs_.view();
+  }
+  [[nodiscard]] inline auto data_ptrs() const noexcept
+    -> device_mdspan<T* const, extent_1d<uint32_t>, row_major>
+  {
+    return data_ptrs_.view();
+  }
 
-    data_    = make_device_mdarray<T>(handle, make_extents<IdxT>(index_size, dim()));
-    indices_ = make_device_mdarray<IdxT>(handle, make_extents<IdxT>(index_size));
+  /** Pointers to the inverted lists (clusters) indices  [n_lists]. */
+  inline auto inds_ptrs() noexcept -> device_mdspan<IdxT*, extent_1d<uint32_t>, row_major>
+  {
+    return inds_ptrs_.view();
+  }
+  [[nodiscard]] inline auto inds_ptrs() const noexcept
+    -> device_mdspan<IdxT* const, extent_1d<uint32_t>, row_major>
+  {
+    return inds_ptrs_.view();
+  }
 
-    switch (metric_) {
-      case raft::distance::DistanceType::L2Expanded:
-      case raft::distance::DistanceType::L2SqrtExpanded:
-      case raft::distance::DistanceType::L2Unexpanded:
-      case raft::distance::DistanceType::L2SqrtUnexpanded:
-        center_norms_ = make_device_mdarray<float>(handle, make_extents<uint32_t>(n_lists()));
-        break;
-      default: center_norms_ = std::nullopt;
+  /**
+   * Update the state of the dependent index members.
+   */
+  void recompute_internal_state(raft::device_resources const& res)
+  {
+    auto stream  = res.get_stream();
+
+    // Actualize the list pointers
+    auto lists     = lists();
+    auto data_ptrs = data_ptrs();
+    auto inds_ptrs = inds_ptrs();
+    IdxT recompute_total_size = 0;
+    for (uint32_t label = 0; label < lists.size(); label++) {
+      const auto data_ptr = lists(label) ? lists(label)->data.data_handle() : nullptr;
+      const auto inds_ptr = lists(label) ? lists(label)->indices.data_handle() : nullptr;
+      const auto list_size = lists(label) ? lists(label)->size() : 0;
+      copy(&data_ptrs(label), &data_ptr, 1, stream);
+      copy(&inds_ptrs(label), &inds_ptr, 1, stream);
+      recompute_total_size += list_size;
     }
+    total_size_ = recompute_total_size;
+  }
 
-    check_consistency();
+  /** Lists' data and indices. */
+  inline auto lists() noexcept
+    -> host_mdspan<std::shared_ptr<list_data<T, IdxT>>, extent_1d<uint32_t>, row_major>
+  {
+    return lists_.view();
+  }
+  [[nodiscard]] inline auto lists() const noexcept
+    -> host_mdspan<const std::shared_ptr<list_data<T, IdxT>>, extent_1d<uint32_t>, row_major>
+  {
+    return lists_.view();
   }
 
  private:
@@ -365,10 +319,10 @@ struct index : ann::index {
   uint32_t veclen_;
   raft::distance::DistanceType metric_;
   bool adaptive_centers_;
-  host_mdarray<std::shared_ptr<list_data<T, IdxT>, extent_1d<uint32_t>, row_major> lists_;
+  host_mdarray<std::shared_ptr<list_data<T, IdxT>>, extent_1d<uint32_t>, row_major> lists_;
   //device_mdarray<T, extent_2d<IdxT>, row_major> data_;
   //device_mdarray<IdxT, extent_1d<IdxT>, row_major> indices_;
-  //device_mdarray<uint32_t, extent_1d<uint32_t>, row_major> list_sizes_;
+  device_mdarray<uint32_t, extent_1d<uint32_t>, row_major> list_sizes_;
   //device_mdarray<IdxT, extent_1d<uint32_t>, row_major> list_offsets_;
   device_mdarray<float, extent_2d<uint32_t>, row_major> centers_;
   std::optional<device_mdarray<float, extent_1d<uint32_t>, row_major>> center_norms_;
@@ -376,21 +330,22 @@ struct index : ann::index {
   // Computed members
   device_mdarray<T*, extent_1d<uint32_t>, row_major> data_ptrs_;
   device_mdarray<IdxT*, extent_1d<uint32_t>, row_major> inds_ptrs_;
-  host_mdarray<IdxT, extent_1d<uint32_t>, row_major> accum_sorted_sizes_;
+  IdxT total_size_;
 
   /** Throw an error if the index content is inconsistent. */
   void check_consistency()
   {
+    auto n_lists = lists_.extent(0);
     RAFT_EXPECTS(dim() % veclen_ == 0, "dimensionality is not a multiple of the veclen");
-    RAFT_EXPECTS(data_.extent(0) == indices_.extent(0), "inconsistent index size");
-    RAFT_EXPECTS(data_.extent(1) == IdxT(centers_.extent(1)), "inconsistent data dimensionality");
+    RAFT_EXPECTS(list_sizes_.extent(0) == n_lists, "inconsistent list size");
+    RAFT_EXPECTS(data_ptrs_.extent(0) == n_lists, "inconsistent list size");
+    RAFT_EXPECTS(inds_ptrs_.extent(0) == n_lists, "inconsistent list size");
     RAFT_EXPECTS(                                               //
       (centers_.extent(0) == list_sizes_.extent(0)) &&          //
-        (centers_.extent(0) + 1 == list_offsets_.extent(0)) &&  //
         (!center_norms_.has_value() || centers_.extent(0) == center_norms_->extent(0)),
       "inconsistent number of lists (clusters)");
-    RAFT_EXPECTS(reinterpret_cast<size_t>(data_.data_handle()) % (veclen_ * sizeof(T)) == 0,
-                 "The data storage pointer is not aligned to the vector length");
+    //RAFT_EXPECTS(reinterpret_cast<size_t>(data_.data_handle()) % (veclen_ * sizeof(T)) == 0,
+    //             "The data storage pointer is not aligned to the vector length");
   }
 
   static auto calculate_veclen(uint32_t dim) -> uint32_t
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
index c417a97531..1b8bc307b6 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
@@ -41,6 +41,63 @@ namespace raft::spatial::knn::ivf_flat::detail {
 
 using namespace raft::spatial::knn::detail;  // NOLINT
 
+/**
+  * Resize a list by the given id, so that it can contain the given number of records;
+  * possibly, copy the data.
+  *
+  * Besides resizing the corresponding list_data, this function updates the device pointers
+  *   data_ptrs, inds_ptrs, and the list_sizes if necessary.
+  *
+  * The new `list_sizes(label)` represents the number of valid records in the index;
+  * it can be `list_size` if the previous size was not smaller; otherwise it's not updated.
+  *
+  * @param[in] handle
+  * @param[in] label list id
+  * @param[in] list_size the minimum size the list should grow.
+  */
+template <typename T, typename IdxT, typename SizeT = uint32_t>
+void resize_list(raft::device_resources const& handle, 
+                 std::shared_ptr<list_data<T, IdxT, SizeT>>& orig_list,
+                 SizeT new_used_size,
+                 SizeT old_used_size,
+                 uint32_t dim)
+{
+  bool skip_resize = false;
+  // TODO update total_size
+  if (orig_list) {
+    if (new_used_size <= orig_list->indices.extent(0)) {
+      auto shared_list_size = old_used_size;
+      if (new_used_size <= old_used_size ||
+          orig_list->size.compare_exchange_strong(shared_list_size, new_used_size)) {
+        // We don't need to resize the list if:
+        //  1. The list exists
+        //  2. The new size fits in the list
+        //  3. The list doesn't grow or no-one else has grown it yet
+        skip_resize = true;
+      }
+    }
+  } else {
+    old_used_size = 0;
+  }
+  if (skip_resize) { return; }
+  auto new_list = std::make_shared<list_data<T, IdxT>>(res, new_used_size, dim);
+  if (old_used_size > 0) {
+    auto copied_data_extents = make_extents<SizeT>(old_used_size, dim);
+    auto copied_view         = make_mdspan<T, SizeT, row_major, false, true>(
+      new_list->data.data_handle(), copied_data_extents);
+    copy(copied_view.data_handle(),
+         orig_list->data.data_handle(),
+         copied_view.size(),
+         res.get_stream());
+    copy(new_list->indices.data_handle(),
+         orig_list->indices.data_handle(),
+         old_used_size,
+         res.get_stream());
+  }
+  // swap the shared pointer content with the new list
+  new_list.swap(orig_list);
+}
+
 /**
  * @brief Record the dataset into the index, one source row at a time.
  *
@@ -119,29 +176,30 @@ __global__ void build_index_kernel(const LabelT* labels,
 
 /** See raft::spatial::knn::ivf_flat::extend docs */
 template <typename T, typename IdxT>
-inline auto extend(raft::device_resources const& handle,
-                   const index<T, IdxT>& orig_index,
-                   const T* new_vectors,
-                   const IdxT* new_indices,
-                   IdxT n_rows) -> index<T, IdxT>
+void extend(raft::device_resources const& handle,
+            index<T, IdxT>* index,
+            const T* new_vectors,
+            const IdxT* new_indices,
+            IdxT n_rows)
 {
   using LabelT = uint32_t;
+  RAFT_EXPECTS(orig_index != nullptr, "index cannot be empty.");
 
   auto stream  = handle.get_stream();
-  auto n_lists = orig_index.n_lists();
-  auto dim     = orig_index.dim();
+  auto n_lists = index->n_lists();
+  auto dim     = index->dim();
   common::nvtx::range<common::nvtx::domain::raft> fun_scope(
     "ivf_flat::extend(%zu, %u)", size_t(n_rows), dim);
 
-  RAFT_EXPECTS(new_indices != nullptr || orig_index.size() == 0,
+  RAFT_EXPECTS(new_indices != nullptr || index->size() == 0,
                "You must pass data indices when the index is non-empty.");
 
   rmm::device_uvector<LabelT> new_labels(n_rows, stream);
   raft::cluster::kmeans_balanced_params kmeans_params;
-  kmeans_params.metric     = orig_index.metric();
+  kmeans_params.metric     = index->metric();
   auto new_vectors_view    = raft::make_device_matrix_view<const T, IdxT>(new_vectors, n_rows, dim);
   auto orig_centroids_view = raft::make_device_matrix_view<const float, IdxT>(
-    orig_index.centers().data_handle(), n_lists, dim);
+    index->centers().data_handle(), n_lists, dim);
   auto labels_view = raft::make_device_vector_view<LabelT, IdxT>(new_labels.data(), n_rows);
   raft::cluster::kmeans_balanced::predict(handle,
                                           kmeans_params,
@@ -150,20 +208,13 @@ inline auto extend(raft::device_resources const& handle,
                                           labels_view,
                                           utils::mapping<float>{});
 
-  index<T, IdxT> ext_index(
-    handle, orig_index.metric(), n_lists, orig_index.adaptive_centers(), dim);
-
-  auto list_sizes_ptr   = ext_index.list_sizes().data_handle();
-  auto list_offsets_ptr = ext_index.list_offsets().data_handle();
-  auto centers_ptr      = ext_index.centers().data_handle();
+  auto list_sizes_ptr   = index->list_sizes().data_handle();
+  //auto list_offsets_ptr = ext_index.list_offsets().data_handle();
 
   // Calculate the centers and sizes on the new data, starting from the original values
-  raft::copy(centers_ptr, orig_index.centers().data_handle(), ext_index.centers().size(), stream);
 
-  if (ext_index.adaptive_centers()) {
-    raft::copy(
-      list_sizes_ptr, orig_index.list_sizes().data_handle(), ext_index.list_sizes().size(), stream);
-    auto centroids_view = raft::make_device_matrix_view<float, IdxT>(centers_ptr, n_lists, dim);
+  if (index->adaptive_centers()) {
+    auto centroids_view = index->centers().view();
     auto list_sizes_view =
       raft::make_device_vector_view<std::remove_pointer_t<decltype(list_sizes_ptr)>, IdxT>(
         list_sizes_ptr, n_lists);
@@ -177,6 +228,7 @@ inline auto extend(raft::device_resources const& handle,
                                                                     false,
                                                                     utils::mapping<float>{});
   } else {
+    auto new_list_sizes   = raft::make_device_vector<uint32_t, IdxT>(handle, n_lists);
     raft::stats::histogram<uint32_t, IdxT>(raft::stats::HistTypeAuto,
                                            reinterpret_cast<int32_t*>(list_sizes_ptr),
                                            IdxT(n_lists),
@@ -188,8 +240,24 @@ inline auto extend(raft::device_resources const& handle,
       list_sizes_ptr, list_sizes_ptr, orig_index.list_sizes().data_handle(), n_lists, stream);
   }
 
-  // Calculate new offsets
-  IdxT index_size = 0;
+  // Calculate and allocate new list data
+  {
+    std::vector<uint32_t> new_cluster_sizes(n_lists);
+    std::vector<uint32_t> old_cluster_sizes(n_lists);
+    copy(new_cluster_sizes.data(), list_sizes, n_lists, stream);
+    copy(old_cluster_sizes.data(), orig_list_sizes.data(), n_lists, stream);
+    handle.sync_stream();
+    auto lists = index->lists();
+    for (uint32_t label = 0; label < n_lists; label++) {
+      resize_list(handle,
+                  lists(label),
+                  new_cluster_sizes[label],
+                  old_cluster_sizes[label],
+                  index->dim());
+  }
+  // Update the pointers and the sizes
+  index->recompute_internal_state(handle);
+  /*IdxT index_size = 0;
   update_device(list_offsets_ptr, &index_size, 1, stream);
   thrust::inclusive_scan(
     rmm::exec_policy(stream),
@@ -219,7 +287,7 @@ inline auto extend(raft::device_resources const& handle,
                       ext_index.indices().data_handle(),
                       IdxT(1),
                       stream);
-  }
+  }*/
 
   // Copy the old sizes, so we can start from the current state of the index;
   // we'll rebuild the `list_sizes_ptr` in the following kernel, using it as an atomic counter.
@@ -357,13 +425,13 @@ inline void fill_refinement_index(raft::device_resources const& handle,
     raft::compose_op(raft::cast_op<LabelT>(), raft::div_const_op<IdxT>(n_candidates)));
 
   auto list_sizes_ptr   = refinement_index->list_sizes().data_handle();
-  auto list_offsets_ptr = refinement_index->list_offsets().data_handle();
+  //auto list_offsets_ptr = refinement_index->list_offsets().data_handle();
   // We do not fill centers and center norms, since we will not run coarse search.
 
   // Calculate new offsets
   uint32_t n_roundup     = Pow2<kIndexGroupSize>::roundUp(n_candidates);
-  auto list_offsets_view = raft::make_device_vector_view<IdxT, IdxT>(
-    list_offsets_ptr, refinement_index->list_offsets().size());
+  //auto list_offsets_view = raft::make_device_vector_view<IdxT, IdxT>(
+  //  list_offsets_ptr, refinement_index->list_offsets().size());
   linalg::map_offset(handle,
                      list_offsets_view,
                      raft::compose_op(raft::cast_op<IdxT>(), raft::mul_const_op<IdxT>(n_roundup)));

From 843904a46307b6cb0b179631d836909f1f38168f Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Thu, 9 Feb 2023 23:09:30 +0100
Subject: [PATCH 03/33] Refactoring: build and extend fix

---
 cpp/include/raft/neighbors/ivf_flat.cuh       |  12 +-
 cpp/include/raft/neighbors/ivf_flat_types.hpp |  21 +-
 .../spatial/knn/detail/ivf_flat_build.cuh     | 189 ++++++++----------
 cpp/test/neighbors/ann_ivf_flat.cuh           |  28 +--
 4 files changed, 115 insertions(+), 135 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh
index f18611b9f1..d210d2c74b 100644
--- a/cpp/include/raft/neighbors/ivf_flat.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat.cuh
@@ -252,7 +252,7 @@ void extend(raft::device_resources const& handle,
             const IdxT* new_indices,
             IdxT n_rows)
 {
-  *index = extend(handle, *index, new_vectors, new_indices, n_rows);
+  raft::spatial::knn::ivf_flat::detail::extend(handle, index, new_vectors, new_indices, n_rows);
 }
 
 /**
@@ -293,11 +293,11 @@ void extend(raft::device_resources const& handle,
             raft::device_matrix_view<const value_t, idx_t, row_major> new_vectors,
             std::optional<raft::device_vector_view<const idx_t, idx_t>> new_indices = std::nullopt)
 {
-  *index = extend(handle,
-                  *index,
-                  new_vectors.data_handle(),
-                  new_indices.has_value() ? new_indices.value().data_handle() : nullptr,
-                  static_cast<idx_t>(new_vectors.extent(0)));
+  extend(handle,
+         index,
+         new_vectors.data_handle(),
+         new_indices.has_value() ? new_indices.value().data_handle() : nullptr,
+         static_cast<idx_t>(new_vectors.extent(0)));
 }
 
 /** @} */
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index b384cbcbb3..fcca6cd858 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -19,10 +19,13 @@
 #include "ann_types.hpp"
 
 #include <raft/core/device_mdarray.hpp>
+#include <raft/core/host_mdarray.hpp>
+#include <raft/core/mdspan_types.hpp>
 #include <raft/core/error.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <raft/util/integer_utils.hpp>
 
+#include <memory>
 #include <optional>
 #include <type_traits>
 
@@ -284,16 +287,16 @@ struct index : ann::index {
     auto stream  = res.get_stream();
 
     // Actualize the list pointers
-    auto lists     = lists();
-    auto data_ptrs = data_ptrs();
-    auto inds_ptrs = inds_ptrs();
+    auto this_lists     = lists();
+    auto this_data_ptrs = data_ptrs();
+    auto this_inds_ptrs = inds_ptrs();
     IdxT recompute_total_size = 0;
-    for (uint32_t label = 0; label < lists.size(); label++) {
-      const auto data_ptr = lists(label) ? lists(label)->data.data_handle() : nullptr;
-      const auto inds_ptr = lists(label) ? lists(label)->indices.data_handle() : nullptr;
-      const auto list_size = lists(label) ? lists(label)->size() : 0;
-      copy(&data_ptrs(label), &data_ptr, 1, stream);
-      copy(&inds_ptrs(label), &inds_ptr, 1, stream);
+    for (uint32_t label = 0; label < this_lists.size(); label++) {
+      const auto data_ptr = this_lists(label) ? this_lists(label)->data.data_handle() : nullptr;
+      const auto inds_ptr = this_lists(label) ? this_lists(label)->indices.data_handle() : nullptr;
+      const auto list_size = this_lists(label) ? IdxT(this_lists(label)->size) : 0;
+      copy(&this_data_ptrs(label), &data_ptr, 1, stream);
+      copy(&this_inds_ptrs(label), &inds_ptr, 1, stream);
       recompute_total_size += list_size;
     }
     total_size_ = recompute_total_size;
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
index 1b8bc307b6..ca595f7303 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
@@ -16,7 +16,6 @@
 
 #pragma once
 
-#include "../ivf_flat_types.hpp"
 #include "ann_utils.cuh"
 
 #include <raft/cluster/kmeans_balanced.cuh>
@@ -29,6 +28,7 @@
 #include <raft/linalg/add.cuh>
 #include <raft/linalg/map.cuh>
 #include <raft/linalg/norm.cuh>
+#include <raft/neighbors/ivf_flat_types.hpp>
 #include <raft/stats/histogram.cuh>
 #include <raft/util/pow2_utils.cuh>
 
@@ -40,7 +40,13 @@
 namespace raft::spatial::knn::ivf_flat::detail {
 
 using namespace raft::spatial::knn::detail;  // NOLINT
+using namespace raft::neighbors::ivf_flat;   // NOLINT
 
+using raft::neighbors::ivf_flat::index;
+using raft::neighbors::ivf_flat::index_params;
+using raft::neighbors::ivf_flat::kIndexGroupSize;
+using raft::neighbors::ivf_flat::search_params;
+using raft::neighbors::ivf_flat::list_data;
 /**
   * Resize a list by the given id, so that it can contain the given number of records;
   * possibly, copy the data.
@@ -80,7 +86,7 @@ void resize_list(raft::device_resources const& handle,
     old_used_size = 0;
   }
   if (skip_resize) { return; }
-  auto new_list = std::make_shared<list_data<T, IdxT>>(res, new_used_size, dim);
+  auto new_list = std::make_shared<list_data<T, IdxT>>(handle, new_used_size, dim);
   if (old_used_size > 0) {
     auto copied_data_extents = make_extents<SizeT>(old_used_size, dim);
     auto copied_view         = make_mdspan<T, SizeT, row_major, false, true>(
@@ -88,11 +94,11 @@ void resize_list(raft::device_resources const& handle,
     copy(copied_view.data_handle(),
          orig_list->data.data_handle(),
          copied_view.size(),
-         res.get_stream());
+         handle.get_stream());
     copy(new_list->indices.data_handle(),
          orig_list->indices.data_handle(),
          old_used_size,
-         res.get_stream());
+         handle.get_stream());
   }
   // swap the shared pointer content with the new list
   new_list.swap(orig_list);
@@ -131,11 +137,11 @@ void resize_list(raft::device_resources const& handle,
  */
 template <typename T, typename IdxT, typename LabelT, bool gather_src = false>
 __global__ void build_index_kernel(const LabelT* labels,
-                                   const IdxT* list_offsets,
+                                   //const IdxT* list_offsets,
                                    const T* source_vecs,
                                    const IdxT* source_ixs,
-                                   T* list_data,
-                                   IdxT* list_index,
+                                   T** list_data_ptrs,
+                                   IdxT** list_index_ptrs,
                                    uint32_t* list_sizes_ptr,
                                    IdxT n_rows,
                                    uint32_t dim,
@@ -146,10 +152,11 @@ __global__ void build_index_kernel(const LabelT* labels,
 
   auto list_id     = labels[i];
   auto inlist_id   = atomicAdd(list_sizes_ptr + list_id, 1);
-  auto list_offset = list_offsets[list_id];
+  auto* list_index = list_index_ptrs[list_id];
+  auto* list_data  = list_data_ptrs[list_id];
 
   // Record the source vector id in the index
-  list_index[list_offset + inlist_id] = source_ixs == nullptr ? i : source_ixs[i];
+  list_index[inlist_id] = source_ixs == nullptr ? i : source_ixs[i];
 
   // The data is written in interleaved groups of `index::kGroupSize` vectors
   using interleaved_group = Pow2<kIndexGroupSize>;
@@ -157,7 +164,7 @@ __global__ void build_index_kernel(const LabelT* labels,
   auto ingroup_id         = interleaved_group::mod(inlist_id) * veclen;
 
   // Point to the location of the interleaved group of vectors
-  list_data += (list_offset + group_offset) * dim;
+  list_data += group_offset * dim;
 
   // Point to the source vector
   if constexpr (gather_src) {
@@ -183,7 +190,7 @@ void extend(raft::device_resources const& handle,
             IdxT n_rows)
 {
   using LabelT = uint32_t;
-  RAFT_EXPECTS(orig_index != nullptr, "index cannot be empty.");
+  RAFT_EXPECTS(index != nullptr, "index cannot be empty.");
 
   auto stream  = handle.get_stream();
   auto n_lists = index->n_lists();
@@ -194,32 +201,30 @@ void extend(raft::device_resources const& handle,
   RAFT_EXPECTS(new_indices != nullptr || index->size() == 0,
                "You must pass data indices when the index is non-empty.");
 
-  rmm::device_uvector<LabelT> new_labels(n_rows, stream);
+  auto new_labels = raft::make_device_vector<LabelT, IdxT>(handle, n_rows);
   raft::cluster::kmeans_balanced_params kmeans_params;
   kmeans_params.metric     = index->metric();
   auto new_vectors_view    = raft::make_device_matrix_view<const T, IdxT>(new_vectors, n_rows, dim);
   auto orig_centroids_view = raft::make_device_matrix_view<const float, IdxT>(
     index->centers().data_handle(), n_lists, dim);
-  auto labels_view = raft::make_device_vector_view<LabelT, IdxT>(new_labels.data(), n_rows);
   raft::cluster::kmeans_balanced::predict(handle,
                                           kmeans_params,
                                           new_vectors_view,
                                           orig_centroids_view,
-                                          labels_view,
+                                          new_labels.view(),
                                           utils::mapping<float>{});
 
-  auto list_sizes_ptr   = index->list_sizes().data_handle();
-  //auto list_offsets_ptr = ext_index.list_offsets().data_handle();
+  auto* list_sizes_ptr   = index->list_sizes().data_handle();
+  auto old_list_sizes_dev =  raft::make_device_vector<uint32_t, IdxT>(handle, n_lists);
+  copy(old_list_sizes_dev.data_handle(), list_sizes_ptr, n_lists, stream);
 
   // Calculate the centers and sizes on the new data, starting from the original values
-
   if (index->adaptive_centers()) {
-    auto centroids_view = index->centers().view();
+    auto centroids_view = raft::make_device_matrix_view<float, IdxT>(index->centers().data_handle(), index->centers().extent(0), index->centers().extent(1));
     auto list_sizes_view =
       raft::make_device_vector_view<std::remove_pointer_t<decltype(list_sizes_ptr)>, IdxT>(
         list_sizes_ptr, n_lists);
-    auto const_labels_view =
-      raft::make_device_vector_view<const LabelT, IdxT>(new_labels.data(), n_rows);
+    auto const_labels_view = make_const_mdspan(new_labels.view());
     raft::cluster::kmeans_balanced::helpers::calc_centers_and_sizes(handle,
                                                                     new_vectors_view,
                                                                     const_labels_view,
@@ -228,109 +233,82 @@ void extend(raft::device_resources const& handle,
                                                                     false,
                                                                     utils::mapping<float>{});
   } else {
-    auto new_list_sizes   = raft::make_device_vector<uint32_t, IdxT>(handle, n_lists);
     raft::stats::histogram<uint32_t, IdxT>(raft::stats::HistTypeAuto,
                                            reinterpret_cast<int32_t*>(list_sizes_ptr),
                                            IdxT(n_lists),
-                                           new_labels.data(),
+                                           new_labels.data_handle(),
                                            n_rows,
                                            1,
                                            stream);
     raft::linalg::add(
-      list_sizes_ptr, list_sizes_ptr, orig_index.list_sizes().data_handle(), n_lists, stream);
+      list_sizes_ptr, list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream);
   }
 
   // Calculate and allocate new list data
   {
-    std::vector<uint32_t> new_cluster_sizes(n_lists);
-    std::vector<uint32_t> old_cluster_sizes(n_lists);
-    copy(new_cluster_sizes.data(), list_sizes, n_lists, stream);
-    copy(old_cluster_sizes.data(), orig_list_sizes.data(), n_lists, stream);
+    std::vector<uint32_t> new_list_sizes(n_lists);
+    std::vector<uint32_t> old_list_sizes(n_lists);
+    copy(old_list_sizes.data(), old_list_sizes_dev.data_handle(), n_lists, stream);
+    copy(new_list_sizes.data(), list_sizes_ptr, n_lists, stream);
     handle.sync_stream();
     auto lists = index->lists();
     for (uint32_t label = 0; label < n_lists; label++) {
       resize_list(handle,
                   lists(label),
-                  new_cluster_sizes[label],
-                  old_cluster_sizes[label],
+                  new_list_sizes[label],
+                  old_list_sizes[label],
                   index->dim());
+      }
   }
   // Update the pointers and the sizes
   index->recompute_internal_state(handle);
-  /*IdxT index_size = 0;
-  update_device(list_offsets_ptr, &index_size, 1, stream);
-  thrust::inclusive_scan(
-    rmm::exec_policy(stream),
-    list_sizes_ptr,
-    list_sizes_ptr + n_lists,
-    list_offsets_ptr + 1,
-    [] __device__(IdxT s, uint32_t l) { return s + Pow2<kIndexGroupSize>::roundUp(l); });
-  update_host(&index_size, list_offsets_ptr + n_lists, 1, stream);
-  handle.sync_stream(stream);
-
-  ext_index.allocate(handle, index_size);
-
-  // Populate index with the old data
-  if (orig_index.size() > 0) {
-    utils::block_copy(orig_index.list_offsets().data_handle(),
-                      list_offsets_ptr,
-                      IdxT(n_lists),
-                      orig_index.data().data_handle(),
-                      ext_index.data().data_handle(),
-                      IdxT(dim),
-                      stream);
-
-    utils::block_copy(orig_index.list_offsets().data_handle(),
-                      list_offsets_ptr,
-                      IdxT(n_lists),
-                      orig_index.indices().data_handle(),
-                      ext_index.indices().data_handle(),
-                      IdxT(1),
-                      stream);
-  }*/
 
   // Copy the old sizes, so we can start from the current state of the index;
   // we'll rebuild the `list_sizes_ptr` in the following kernel, using it as an atomic counter.
   raft::copy(
-    list_sizes_ptr, orig_index.list_sizes().data_handle(), ext_index.list_sizes().size(), stream);
+    list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream);
 
+  // Kernel to insert the new vectors
   const dim3 block_dim(256);
   const dim3 grid_dim(raft::ceildiv<IdxT>(n_rows, block_dim.x));
-  build_index_kernel<<<grid_dim, block_dim, 0, stream>>>(new_labels.data(),
-                                                         list_offsets_ptr,
+  build_index_kernel<<<grid_dim, block_dim, 0, stream>>>(new_labels.data_handle(),
                                                          new_vectors,
                                                          new_indices,
-                                                         ext_index.data().data_handle(),
-                                                         ext_index.indices().data_handle(),
+                                                         index->data_ptrs().data_handle(),//ext_index.data().data_handle(),
+                                                         index->inds_ptrs().data_handle(),//ext_index.indices().data_handle(),
                                                          list_sizes_ptr,
                                                          n_rows,
                                                          dim,
-                                                         ext_index.veclen());
+                                                         index->veclen());
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 
   // Precompute the centers vector norms for L2Expanded distance
-  if (ext_index.center_norms().has_value()) {
-    if (!ext_index.adaptive_centers() && orig_index.center_norms().has_value()) {
-      raft::copy(ext_index.center_norms()->data_handle(),
-                 orig_index.center_norms()->data_handle(),
-                 orig_index.center_norms()->size(),
-                 stream);
-    } else {
-      raft::linalg::rowNorm(ext_index.center_norms()->data_handle(),
-                            ext_index.centers().data_handle(),
-                            dim,
-                            n_lists,
-                            raft::linalg::L2Norm,
-                            true,
-                            stream);
-      RAFT_LOG_TRACE_VEC(ext_index.center_norms()->data_handle(), std::min<uint32_t>(dim, 20));
-    }
+  if (index->center_norms().has_value() && index->adaptive_centers()) {
+    raft::linalg::rowNorm(index->center_norms()->data_handle(),
+                          index->centers().data_handle(),
+                          dim,
+                          n_lists,
+                          raft::linalg::L2Norm,
+                          true,
+                          stream);
+    RAFT_LOG_TRACE_VEC(index->center_norms()->data_handle(), std::min<uint32_t>(dim, 20));
   }
+}
 
-  // assemble the index
+/** See raft::spatial::knn::ivf_flat::extend docs */
+template <typename T, typename IdxT>
+auto extend(raft::device_resources const& handle,
+            const index<T, IdxT>& orig_index,
+            const T* new_vectors,
+            const IdxT* new_indices,
+            IdxT n_rows) -> index<T, IdxT>
+{
+  auto ext_index = clone(handle, &orig_index);
+  extend(handle, &ext_index, new_vectors, new_indices, n_rows);
   return ext_index;
 }
 
+
 /** See raft::spatial::knn::ivf_flat::build docs */
 template <typename T, typename IdxT>
 inline auto build(raft::device_resources const& handle,
@@ -348,7 +326,8 @@ inline auto build(raft::device_resources const& handle,
 
   index<T, IdxT> index(handle, params, dim);
   utils::memzero(index.list_sizes().data_handle(), index.list_sizes().size(), stream);
-  utils::memzero(index.list_offsets().data_handle(), index.list_offsets().size(), stream);
+  utils::memzero(index.data_ptrs().data_handle(), index.data_ptrs().size(), stream);
+  utils::memzero(index.inds_ptrs().data_handle(), index.inds_ptrs().size(), stream);
 
   // Train the kmeans clustering
   {
@@ -378,10 +357,9 @@ inline auto build(raft::device_resources const& handle,
 
   // add the data if necessary
   if (params.add_data_on_build) {
-    return detail::extend<T, IdxT>(handle, index, dataset, nullptr, n_rows);
-  } else {
-    return index;
+    detail::extend<T, IdxT>(handle, &index, dataset, nullptr, n_rows);
   }
+  return index;
 }
 
 /**
@@ -425,19 +403,17 @@ inline void fill_refinement_index(raft::device_resources const& handle,
     raft::compose_op(raft::cast_op<LabelT>(), raft::div_const_op<IdxT>(n_candidates)));
 
   auto list_sizes_ptr   = refinement_index->list_sizes().data_handle();
-  //auto list_offsets_ptr = refinement_index->list_offsets().data_handle();
   // We do not fill centers and center norms, since we will not run coarse search.
 
-  // Calculate new offsets
-  uint32_t n_roundup     = Pow2<kIndexGroupSize>::roundUp(n_candidates);
-  //auto list_offsets_view = raft::make_device_vector_view<IdxT, IdxT>(
-  //  list_offsets_ptr, refinement_index->list_offsets().size());
-  linalg::map_offset(handle,
-                     list_offsets_view,
-                     raft::compose_op(raft::cast_op<IdxT>(), raft::mul_const_op<IdxT>(n_roundup)));
-
-  IdxT index_size = n_roundup * n_lists;
-  refinement_index->allocate(handle, index_size);
+  // Allocate new memory
+  auto lists = refinement_index->lists();
+  for (uint32_t label = 0; label < n_lists; label++) {
+    resize_list(handle,
+                lists(label),
+                n_candidates,
+                uint32_t(0),
+                refinement_index->dim());
+  }
 
   RAFT_CUDA_TRY(cudaMemsetAsync(list_sizes_ptr, 0, n_lists * sizeof(uint32_t), stream));
 
@@ -445,11 +421,10 @@ inline void fill_refinement_index(raft::device_resources const& handle,
   const dim3 grid_dim(raft::ceildiv<IdxT>(n_queries * n_candidates, block_dim.x));
   build_index_kernel<T, IdxT, LabelT, true>
     <<<grid_dim, block_dim, 0, stream>>>(new_labels.data(),
-                                         list_offsets_ptr,
                                          dataset,
                                          candidate_idx,
-                                         refinement_index->data().data_handle(),
-                                         refinement_index->indices().data_handle(),
+                                         refinement_index->data_ptrs().data_handle(),
+                                         refinement_index->inds_ptrs().data_handle(),
                                          list_sizes_ptr,
                                          n_queries * n_candidates,
                                          refinement_index->dim(),
@@ -464,9 +439,10 @@ inline void fill_refinement_index(raft::device_resources const& handle,
 //             compatible fashion.
 constexpr int serialization_version = 2;
 
+/* TODO
 static_assert(sizeof(index<double, std::uint64_t>) == 408,
               "The size of the index struct has changed since the last update; "
-              "paste in the new size and consider updating the save/load logic");
+              "paste in the new size and consider updating the save/load logic");*/
 
 /**
  * Save the index to file.
@@ -488,6 +464,7 @@ void serialize(raft::device_resources const& handle,
 
   RAFT_LOG_DEBUG(
     "Saving IVF-PQ index, size %zu, dim %u", static_cast<size_t>(index_.size()), index_.dim());
+  /* TODO
   serialize_scalar(handle, of, serialization_version);
   serialize_scalar(handle, of, index_.size());
   serialize_scalar(handle, of, index_.dim());
@@ -507,7 +484,7 @@ void serialize(raft::device_resources const& handle,
   } else {
     bool has_norms = false;
     serialize_scalar(handle, of, has_norms);
-  }
+  }*/
   of.close();
   if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); }
 }
@@ -541,8 +518,8 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   bool adaptive_centers = deserialize_scalar<bool>(handle, infile);
 
   index<T, IdxT> index_ =
-    raft::spatial::knn::ivf_flat::index<T, IdxT>(handle, metric, n_lists, adaptive_centers, dim);
-
+    index<T, IdxT>(handle, metric, n_lists, adaptive_centers, dim);
+  /* TODO
   index_.allocate(handle, n_rows);
   auto data = index_.data();
   deserialize_mdspan(handle, infile, data);
@@ -558,7 +535,7 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
       auto center_norms = *index_.center_norms();
       deserialize_mdspan(handle, infile, center_norms);
     }
-  }
+  }*/
   infile.close();
   return index_;
 }
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index c100afb2c4..3807e414c0 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -169,7 +169,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         auto half_of_data_view = raft::make_device_matrix_view<const DataT, IdxT>(
           (const DataT*)database.data(), half_of_data, ps.dim);
 
-        auto index_2 = ivf_flat::extend(handle_, index, half_of_data_view);
+        ivf_flat::extend(handle_, &index, half_of_data_view);
 
         auto new_half_of_data_view = raft::make_device_matrix_view<const DataT, IdxT>(
           database.data() + half_of_data * ps.dim, IdxT(ps.num_db_vecs) - half_of_data, ps.dim);
@@ -178,7 +178,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
           vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data);
 
         ivf_flat::extend(handle_,
-                         &index_2,
+                         &index,
                          new_half_of_data_view,
                          std::make_optional<raft::device_vector_view<const IdxT, IdxT>>(
                            new_half_of_data_indices_view));
@@ -189,7 +189,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
           indices_ivfflat_dev.data(), ps.num_queries, ps.k);
         auto dists_out_view = raft::make_device_matrix_view<T, IdxT>(
           distances_ivfflat_dev.data(), ps.num_queries, ps.k);
-        raft::spatial::knn::ivf_flat::detail::serialize(handle_, "ivf_flat_index", index_2);
+        raft::spatial::knn::ivf_flat::detail::serialize(handle_, "ivf_flat_index", index);
 
         auto index_loaded =
           raft::spatial::knn::ivf_flat::detail::deserialize<DataT, IdxT>(handle_, "ivf_flat_index");
@@ -207,30 +207,30 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         handle_.sync_stream(stream_);
 
         // Test the centroid invariants
-        if (index_2.adaptive_centers()) {
+        if (index.adaptive_centers()) {
           // The centers must be up-to-date with the corresponding data
-          std::vector<uint32_t> list_sizes(index_2.n_lists());
-          std::vector<IdxT> list_offsets(index_2.n_lists());
+          std::vector<uint32_t> list_sizes(index.n_lists());
+          //std::vector<IdxT> list_offsets(index.n_lists());
           rmm::device_uvector<float> centroid(ps.dim, stream_);
           raft::copy(
-            list_sizes.data(), index_2.list_sizes().data_handle(), index_2.n_lists(), stream_);
-          raft::copy(
-            list_offsets.data(), index_2.list_offsets().data_handle(), index_2.n_lists(), stream_);
+            list_sizes.data(), index.list_sizes().data_handle(), index.n_lists(), stream_);
+          //raft::copy(
+          //  list_offsets.data(), index.list_offsets().data_handle(), index.n_lists(), stream_);
           handle_.sync_stream(stream_);
-          for (uint32_t l = 0; l < index_2.n_lists(); l++) {
+          for (uint32_t l = 0; l < index.n_lists(); l++) {
             rmm::device_uvector<float> cluster_data(list_sizes[l] * ps.dim, stream_);
             raft::spatial::knn::detail::utils::copy_selected<float>(
               (IdxT)list_sizes[l],
               (IdxT)ps.dim,
               database.data(),
-              index_2.indices().data_handle() + list_offsets[l],
+              index.inds_ptrs()(l),
               (IdxT)ps.dim,
               cluster_data.data(),
               (IdxT)ps.dim,
               stream_);
             raft::stats::mean<float, uint32_t>(
               centroid.data(), cluster_data.data(), ps.dim, list_sizes[l], false, true, stream_);
-            ASSERT_TRUE(raft::devArrMatch(index_2.centers().data_handle() + ps.dim * l,
+            ASSERT_TRUE(raft::devArrMatch(index.centers().data_handle() + ps.dim * l,
                                           centroid.data(),
                                           ps.dim,
                                           raft::CompareApprox<float>(0.001),
@@ -238,9 +238,9 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
           }
         } else {
           // The centers must be immutable
-          ASSERT_TRUE(raft::devArrMatch(index_2.centers().data_handle(),
+          ASSERT_TRUE(raft::devArrMatch(index.centers().data_handle(),
                                         index.centers().data_handle(),
-                                        index_2.centers().size(),
+                                        index.centers().size(),
                                         raft::Compare<float>(),
                                         stream_));
         }

From f09b3a003949c3b7838b1c197c0b64ee61789858 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Fri, 10 Feb 2023 15:49:39 +0100
Subject: [PATCH 04/33] Refactor ivf flat search for index splitting

---
 .../spatial/knn/detail/ivf_flat_search.cuh    | 33 +++++++++++--------
 1 file changed, 20 insertions(+), 13 deletions(-)

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
index 7f70d4b8a5..77ecae745a 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
@@ -16,7 +16,6 @@
 
 #pragma once
 
-#include "../ivf_flat_types.hpp"
 #include "ann_utils.cuh"
 
 #include <raft/core/cudart_utils.hpp>
@@ -30,6 +29,7 @@
 #include <raft/linalg/unary_op.cuh>
 #include <raft/matrix/detail/select_k.cuh>
 #include <raft/matrix/detail/select_warpsort.cuh>
+#include <raft/neighbors/ivf_flat_types.hpp>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/device_loads_stores.cuh>
 #include <raft/util/integer_utils.hpp>
@@ -42,6 +42,12 @@
 namespace raft::spatial::knn::ivf_flat::detail {
 
 using namespace raft::spatial::knn::detail;  // NOLINT
+using namespace raft::neighbors::ivf_flat;   // NOLINT
+
+using raft::neighbors::ivf_flat::index;
+using raft::neighbors::ivf_flat::index_params;
+using raft::neighbors::ivf_flat::kIndexGroupSize;
+using raft::neighbors::ivf_flat::search_params;
 
 constexpr int kThreadsPerBlock = 128;
 
@@ -673,10 +679,10 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
                           const uint32_t query_smem_elems,
                           const T* query,
                           const uint32_t* coarse_index,
-                          const IdxT* list_indices,
-                          const T* list_data,
+                          const IdxT* const* list_indices_ptrs, // const IdxT* list_indices
+                          const T* const* list_data_ptrs, // const T* list_data
                           const uint32_t* list_sizes,
-                          const IdxT* list_offsets,
+                          //const IdxT* list_offsets,
                           const uint32_t n_probes,
                           const uint32_t k,
                           const uint32_t dim,
@@ -723,7 +729,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
     // Every CUDA block scans one cluster at a time.
     for (int probe_id = blockIdx.x; probe_id < n_probes; probe_id += gridDim.x) {
       const uint32_t list_id   = coarse_index[probe_id];  // The id of cluster(list)
-      const size_t list_offset = list_offsets[list_id];
+      //const size_t list_offset = list_offsets[list_id];
 
       // The number of vectors in each cluster(list); [nlist]
       const uint32_t list_length = list_sizes[list_id];
@@ -741,7 +747,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
            group_id += kNumWarps) {
         AccT dist = 0;
         // This is where this warp begins reading data (start position of an interleaved group)
-        const T* data = list_data + (list_offset + group_id * kIndexGroupSize) * dim;
+        //const T* data = list_data + (list_offset + group_id * kIndexGroupSize) * dim;
+        const T* data = list_data_ptrs[list_id] + (group_id * kIndexGroupSize) * dim;
 
         // This is the vector a given lane/thread handles
         const uint32_t vec_id = group_id * WarpSize + lane_id;
@@ -778,7 +785,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
 
         // Enqueue one element per thread
         const float val  = valid ? static_cast<float>(dist) : block_sort_t::queue_t::kDummy;
-        const size_t idx = valid ? static_cast<size_t>(list_indices[list_offset + vec_id]) : 0;
+        const size_t idx = valid ? static_cast<size_t>(list_indices_ptrs[list_id][vec_id]) : 0;
         queue.add(val, idx);
       }
     }
@@ -819,7 +826,7 @@ template <int Capacity,
           typename PostLambda>
 void launch_kernel(Lambda lambda,
                    PostLambda post_process,
-                   const ivf_flat::index<T, IdxT>& index,
+                   const index<T, IdxT>& index,
                    const T* queries,
                    const uint32_t* coarse_index,
                    const uint32_t num_queries,
@@ -869,10 +876,10 @@ void launch_kernel(Lambda lambda,
                                                         query_smem_elems,
                                                         queries,
                                                         coarse_index,
-                                                        index.indices().data_handle(),
-                                                        index.data().data_handle(),
-                                                        index.list_sizes().data_handle(),
-                                                        index.list_offsets().data_handle(),
+                                                        index.inds_ptrs().data_handle(),//indices().data_handle(),
+                                                        index.data_ptrs().data_handle(),//data().data_handle(),
+                                                        index.list_sizes().data_handle(),// TODO
+                                                        //index.list_offsets().data_handle(),
                                                         n_probes,
                                                         k,
                                                         index.dim(),
@@ -1056,7 +1063,7 @@ struct select_interleaved_scan_kernel {
  * @param stream
  */
 template <typename T, typename AccT, typename IdxT>
-void ivfflat_interleaved_scan(const ivf_flat::index<T, IdxT>& index,
+void ivfflat_interleaved_scan(const index<T, IdxT>& index,
                               const T* queries,
                               const uint32_t* coarse_query_results,
                               const uint32_t n_queries,

From 93d5b35863b1d1b940fbedd624be3ea3176d45ad Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Sun, 12 Feb 2023 22:57:51 +0100
Subject: [PATCH 05/33] Use mdpsan/mdarray aliases

---
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 88 +++++++++----------
 .../spatial/knn/detail/ivf_flat_build.cuh     | 14 ++-
 .../spatial/knn/detail/ivf_flat_search.cuh    |  7 +-
 3 files changed, 57 insertions(+), 52 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index fcca6cd858..c57e5e5c6e 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -23,6 +23,7 @@
 #include <raft/core/mdspan_types.hpp>
 #include <raft/core/error.hpp>
 #include <raft/distance/distance_types.hpp>
+#include <raft/matrix/init.cuh>
 #include <raft/util/integer_utils.hpp>
 
 #include <memory>
@@ -38,6 +39,14 @@ namespace raft::neighbors::ivf_flat {
 /** Size of the interleaved group (see `index::data` description). */
 constexpr static uint32_t kIndexGroupSize = 32;
 
+/**
+ * Default value filled in the `indices()` array.
+ * One may encounter it trying to access a record within a cluster that is outside of the
+ * `list_sizes()` bound (due to the record alignment `kIndexGroupSize`).
+ */
+template <typename IdxT>
+constexpr static IdxT kInvalidRecord = std::numeric_limits<IdxT>::max() - 1;
+
 struct index_params : ann::index_params {
   /** The number of inverted lists (clusters) */
   uint32_t n_lists = 1024;
@@ -72,19 +81,19 @@ static_assert(std::is_aggregate_v<search_params>);
 template <typename T, typename IdxT, typename SizeT = uint32_t>
 struct list_data {
   /** Cluster data. */
-  device_mdarray<T, extent_2d<SizeT>, row_major> data;
+  device_matrix<T, SizeT, row_major> data;
   /** Source indices. */
-  device_mdarray<IdxT, extent_1d<SizeT>, row_major> indices;
+  device_vector<IdxT, SizeT> indices;
   /** The actual size of the content. */
   std::atomic<SizeT> size;
 
-  list_data(raft::device_resources const& handle, SizeT n_rows, uint32_t dim)
+  list_data(raft::device_resources const& res, SizeT n_rows, uint32_t dim)
     : size{n_rows}
   {
     auto capacity = round_up_safe<SizeT>(bound_by_power_of_two<SizeT>(size), kIndexGroupSize);
     try {
-      data = make_device_mdarray<T>(handle, make_extents<SizeT>(capacity, dim));
-      indices = make_device_mdarray<IdxT>(handle, make_extents<SizeT>(capacity));
+      data = make_device_matrix<T, SizeT, row_major>(res, capacity, dim);
+      indices = make_device_vector<IdxT, SizeT>(res, capacity);
     } catch (std::bad_alloc& e) {
       RAFT_FAIL(
         "ivf-flat: failed to allocate a big enough index list to hold all data "
@@ -94,6 +103,8 @@ struct list_data {
         size_t(capacity),
         e.what());
     }
+    // Fill the index buffer with a pre-defined marker for easier debugging
+    matrix::fill(res, indices.view(), ivf_flat::kInvalidRecord<IdxT>);
   }
 };
 
@@ -110,12 +121,6 @@ struct index : ann::index {
                 "IdxT must be able to represent all values of uint32_t");
 
  public:
-  /**
-   * Default value filled in the `indices()` array.
-   * One may encounter it trying to access a record within a cluster that is outside of the
-   * `list_sizes()` bound (due to the record alignment `kIndexGroupSize`).
-   */
-  constexpr static IdxT kInvalidRecord = std::numeric_limits<IdxT>::max() - 1;
   /**
    * Vectorized load/store size in elements, determines the size of interleaved data chunks.
    *
@@ -157,23 +162,23 @@ struct index : ann::index {
    *
    */
   /** Sizes of the lists (clusters) [n_lists] */
-  inline auto list_sizes() noexcept -> device_mdspan<uint32_t, extent_1d<uint32_t>, row_major>
+  inline auto list_sizes() noexcept -> device_vector_view<uint32_t, uint32_t>
   {
     return list_sizes_.view();
   }
   [[nodiscard]] inline auto list_sizes() const noexcept
-    -> device_mdspan<const uint32_t, extent_1d<uint32_t>, row_major>
+    -> device_vector_view<const uint32_t, uint32_t>
   {
     return list_sizes_.view();
   }
 
   /** k-means cluster centers corresponding to the lists [n_lists, dim] */
-  inline auto centers() noexcept -> device_mdspan<float, extent_2d<uint32_t>, row_major>
+  inline auto centers() noexcept -> device_matrix_view<float, uint32_t, row_major>
   {
     return centers_.view();
   }
   [[nodiscard]] inline auto centers() const noexcept
-    -> device_mdspan<const float, extent_2d<uint32_t>, row_major>
+    -> device_matrix_view<const float, uint32_t, row_major>
   {
     return centers_.view();
   }
@@ -185,20 +190,20 @@ struct index : ann::index {
    * calculation.
    */
   inline auto center_norms() noexcept
-    -> std::optional<device_mdspan<float, extent_1d<uint32_t>, row_major>>
+    -> std::optional<device_vector_view<float, uint32_t>>
   {
     if (center_norms_.has_value()) {
-      return std::make_optional<device_mdspan<float, extent_1d<uint32_t>, row_major>>(
+      return std::make_optional<device_vector_view<float, uint32_t>>(
         center_norms_->view());
     } else {
       return std::nullopt;
     }
   }
   [[nodiscard]] inline auto center_norms() const noexcept
-    -> std::optional<device_mdspan<const float, extent_1d<uint32_t>, row_major>>
+    -> std::optional<device_vector_view<const float, uint32_t>>
   {
     if (center_norms_.has_value()) {
-      return std::make_optional<device_mdspan<const float, extent_1d<uint32_t>, row_major>>(
+      return std::make_optional<device_vector_view<const float, uint32_t>>(
         center_norms_->view());
     } else {
       return std::nullopt;
@@ -229,7 +234,7 @@ struct index : ann::index {
   ~index()                          = default;
 
   /** Construct an empty index. It needs to be trained and then populated. */
-  index(raft::device_resources const& handle,
+  index(raft::device_resources const& res,
         raft::distance::DistanceType metric,
         uint32_t n_lists,
         bool adaptive_centers,
@@ -238,12 +243,12 @@ struct index : ann::index {
       veclen_(calculate_veclen(dim)),
       metric_(metric),
       adaptive_centers_(adaptive_centers),
-      centers_(make_device_mdarray<float>(handle, make_extents<uint32_t>(n_lists, dim))),
+      centers_(make_device_matrix<float, uint32_t>(res, n_lists, dim)),
       center_norms_(std::nullopt),
-      lists_{make_host_mdarray<std::shared_ptr<list_data<T, IdxT>>>(make_extents<uint32_t>(n_lists))},
-      list_sizes_{make_device_mdarray<uint32_t>(handle, make_extents<uint32_t>(n_lists))},
-      data_ptrs_{make_device_mdarray<T*>(handle, make_extents<uint32_t>(n_lists))},
-      inds_ptrs_{make_device_mdarray<IdxT*>(handle, make_extents<uint32_t>(n_lists))}
+      lists_{make_host_vector<std::shared_ptr<list_data<T, IdxT>>, uint32_t>(n_lists)},
+      list_sizes_{make_device_vector<uint32_t, uint32_t>(res, n_lists)},
+      data_ptrs_{make_device_vector<T*, uint32_t>(res, n_lists)},
+      inds_ptrs_{make_device_vector<IdxT*, uint32_t>(res, n_lists)}
   {
     check_consistency();
     for (uint32_t i = 0; i < n_lists; i++) {
@@ -252,29 +257,29 @@ struct index : ann::index {
   }
 
   /** Construct an empty index. It needs to be trained and then populated. */
-  index(raft::device_resources const& handle, const index_params& params, uint32_t dim)
-    : index(handle, params.metric, params.n_lists, params.adaptive_centers, dim)
+  index(raft::device_resources const& res, const index_params& params, uint32_t dim)
+    : index(res, params.metric, params.n_lists, params.adaptive_centers, dim)
   {
   }
 
   /** Pointers to the inverted lists (clusters) data  [n_lists]. */
-  inline auto data_ptrs() noexcept -> device_mdspan<T*, extent_1d<uint32_t>, row_major>
+  inline auto data_ptrs() noexcept -> device_vector_view<T*, uint32_t>
   {
     return data_ptrs_.view();
   }
   [[nodiscard]] inline auto data_ptrs() const noexcept
-    -> device_mdspan<T* const, extent_1d<uint32_t>, row_major>
+    -> device_vector_view<T* const, uint32_t>
   {
     return data_ptrs_.view();
   }
 
   /** Pointers to the inverted lists (clusters) indices  [n_lists]. */
-  inline auto inds_ptrs() noexcept -> device_mdspan<IdxT*, extent_1d<uint32_t>, row_major>
+  inline auto inds_ptrs() noexcept -> device_vector_view<IdxT*, uint32_t>
   {
     return inds_ptrs_.view();
   }
   [[nodiscard]] inline auto inds_ptrs() const noexcept
-    -> device_mdspan<IdxT* const, extent_1d<uint32_t>, row_major>
+    -> device_vector_view<IdxT* const, uint32_t>
   {
     return inds_ptrs_.view();
   }
@@ -304,12 +309,12 @@ struct index : ann::index {
 
   /** Lists' data and indices. */
   inline auto lists() noexcept
-    -> host_mdspan<std::shared_ptr<list_data<T, IdxT>>, extent_1d<uint32_t>, row_major>
+    -> host_vector_view<std::shared_ptr<list_data<T, IdxT>>, uint32_t>
   {
     return lists_.view();
   }
   [[nodiscard]] inline auto lists() const noexcept
-    -> host_mdspan<const std::shared_ptr<list_data<T, IdxT>>, extent_1d<uint32_t>, row_major>
+    -> host_vector_view<const std::shared_ptr<list_data<T, IdxT>>, uint32_t>
   {
     return lists_.view();
   }
@@ -322,17 +327,14 @@ struct index : ann::index {
   uint32_t veclen_;
   raft::distance::DistanceType metric_;
   bool adaptive_centers_;
-  host_mdarray<std::shared_ptr<list_data<T, IdxT>>, extent_1d<uint32_t>, row_major> lists_;
-  //device_mdarray<T, extent_2d<IdxT>, row_major> data_;
-  //device_mdarray<IdxT, extent_1d<IdxT>, row_major> indices_;
-  device_mdarray<uint32_t, extent_1d<uint32_t>, row_major> list_sizes_;
-  //device_mdarray<IdxT, extent_1d<uint32_t>, row_major> list_offsets_;
-  device_mdarray<float, extent_2d<uint32_t>, row_major> centers_;
-  std::optional<device_mdarray<float, extent_1d<uint32_t>, row_major>> center_norms_;
+  host_vector<std::shared_ptr<list_data<T, IdxT>>, uint32_t> lists_;
+  device_vector<uint32_t, uint32_t> list_sizes_;
+  device_matrix<float, uint32_t, row_major> centers_;
+  std::optional<device_vector<float, uint32_t>> center_norms_;
 
   // Computed members
-  device_mdarray<T*, extent_1d<uint32_t>, row_major> data_ptrs_;
-  device_mdarray<IdxT*, extent_1d<uint32_t>, row_major> inds_ptrs_;
+  device_vector<T*, uint32_t> data_ptrs_;
+  device_vector<IdxT*, uint32_t> inds_ptrs_;
   IdxT total_size_;
 
   /** Throw an error if the index content is inconsistent. */
@@ -347,8 +349,6 @@ struct index : ann::index {
       (centers_.extent(0) == list_sizes_.extent(0)) &&          //
         (!center_norms_.has_value() || centers_.extent(0) == center_norms_->extent(0)),
       "inconsistent number of lists (clusters)");
-    //RAFT_EXPECTS(reinterpret_cast<size_t>(data_.data_handle()) % (veclen_ * sizeof(T)) == 0,
-    //             "The data storage pointer is not aligned to the vector length");
   }
 
   static auto calculate_veclen(uint32_t dim) -> uint32_t
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
index ca595f7303..1f31a54dce 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
@@ -439,10 +439,16 @@ inline void fill_refinement_index(raft::device_resources const& handle,
 //             compatible fashion.
 constexpr int serialization_version = 2;
 
-/* TODO
-static_assert(sizeof(index<double, std::uint64_t>) == 408,
-              "The size of the index struct has changed since the last update; "
-              "paste in the new size and consider updating the save/load logic");*/
+// NB: we wrap this check in a struct, so that the updated RealSize is easy to see in the error
+// message.
+template <size_t RealSize, size_t ExpectedSize>
+struct check_index_layout {
+  static_assert(RealSize == ExpectedSize,
+                "The size of the index struct has changed since the last update; "
+                "paste in the new size and consider updating the serialization logic");
+};
+
+template struct check_index_layout<sizeof(index<double, std::uint64_t>), 376>;
 
 /**
  * Save the index to file.
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
index 77ecae745a..8ab3d06efc 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
@@ -876,10 +876,9 @@ void launch_kernel(Lambda lambda,
                                                         query_smem_elems,
                                                         queries,
                                                         coarse_index,
-                                                        index.inds_ptrs().data_handle(),//indices().data_handle(),
-                                                        index.data_ptrs().data_handle(),//data().data_handle(),
-                                                        index.list_sizes().data_handle(),// TODO
-                                                        //index.list_offsets().data_handle(),
+                                                        index.inds_ptrs().data_handle(),
+                                                        index.data_ptrs().data_handle(),
+                                                        index.list_sizes().data_handle(),
                                                         n_probes,
                                                         k,
                                                         index.dim(),

From 5fcf5647077b7da6cf3ec392990140aa083d833c Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Mon, 13 Feb 2023 16:47:36 +0100
Subject: [PATCH 06/33] Add serialization

---
 cpp/include/raft/neighbors/ivf_flat_types.hpp |  9 +-
 .../spatial/knn/detail/ivf_flat_build.cuh     | 84 ++++++++++++++++---
 2 files changed, 78 insertions(+), 15 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index c57e5e5c6e..476fa755e0 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -23,11 +23,11 @@
 #include <raft/core/mdspan_types.hpp>
 #include <raft/core/error.hpp>
 #include <raft/distance/distance_types.hpp>
-#include <raft/matrix/init.cuh>
 #include <raft/util/integer_utils.hpp>
 
 #include <memory>
 #include <optional>
+#include <thrust/fill.h>
 #include <type_traits>
 
 namespace raft::neighbors::ivf_flat {
@@ -90,7 +90,7 @@ struct list_data {
   list_data(raft::device_resources const& res, SizeT n_rows, uint32_t dim)
     : size{n_rows}
   {
-    auto capacity = round_up_safe<SizeT>(bound_by_power_of_two<SizeT>(size), kIndexGroupSize);
+    auto capacity = round_up_safe<SizeT>(n_rows, kIndexGroupSize);
     try {
       data = make_device_matrix<T, SizeT, row_major>(res, capacity, dim);
       indices = make_device_vector<IdxT, SizeT>(res, capacity);
@@ -99,12 +99,13 @@ struct list_data {
         "ivf-flat: failed to allocate a big enough index list to hold all data "
         "(requested size: %zu records, selected capacity: %zu records). "
         "Allocator exception: %s",
-        size_t(size),
+        size_t(n_rows),
         size_t(capacity),
         e.what());
     }
     // Fill the index buffer with a pre-defined marker for easier debugging
-    matrix::fill(res, indices.view(), ivf_flat::kInvalidRecord<IdxT>);
+    thrust::fill_n(
+      res.get_thrust_policy(), indices.data_handle(), indices.size(), kInvalidRecord<IdxT>);
   }
 };
 
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
index 1f31a54dce..5a90be93ae 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
@@ -274,8 +274,8 @@ void extend(raft::device_resources const& handle,
   build_index_kernel<<<grid_dim, block_dim, 0, stream>>>(new_labels.data_handle(),
                                                          new_vectors,
                                                          new_indices,
-                                                         index->data_ptrs().data_handle(),//ext_index.data().data_handle(),
-                                                         index->inds_ptrs().data_handle(),//ext_index.indices().data_handle(),
+                                                         index->data_ptrs().data_handle(),
+                                                         index->inds_ptrs().data_handle(),
                                                          list_sizes_ptr,
                                                          n_rows,
                                                          dim,
@@ -432,12 +432,12 @@ inline void fill_refinement_index(raft::device_resources const& handle,
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }
 
-// Serialization version 2
+// Serialization version 3
 // No backward compatibility yet; that is, can't add additional fields without breaking
 // backward compatibility.
 // TODO(hcho3) Implement next-gen serializer for IVF that allows for expansion in a backward
 //             compatible fashion.
-constexpr int serialization_version = 2;
+constexpr int serialization_version = 3;
 
 // NB: we wrap this check in a struct, so that the updated RealSize is easy to see in the error
 // message.
@@ -450,6 +450,59 @@ struct check_index_layout {
 
 template struct check_index_layout<sizeof(index<double, std::uint64_t>), 376>;
 
+
+template <typename T, typename IdxT, typename SizeT>
+void serialize_list(const raft::device_resources& handle,
+                    std::ostream& os,
+                    const list_data<T, IdxT, SizeT>& ld,
+                    std::optional<SizeT> size_override = std::nullopt)
+{
+  auto size = size_override.value_or(ld.size.load());
+  serialize_scalar(handle, os, size);
+  if (size == 0) { return; }
+
+  auto data_extents = make_extents<SizeT>(size, ld.data.extent(1));
+  auto data_array   = make_host_mdarray<T, SizeT, row_major>(data_extents);
+  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+  copy(data_array.data_handle(), ld.data.data_handle(), data_array.size(), handle.get_stream());
+  copy(inds_array.data_handle(), ld.indices.data_handle(), inds_array.size(), handle.get_stream());
+  handle.sync_stream();
+  serialize_mdspan(handle, os, data_array.view());
+  serialize_mdspan(handle, os, inds_array.view());
+}
+
+template <typename T, typename IdxT, typename SizeT>
+void serialize_list(const raft::device_resources& handle,
+                    std::ostream& os,
+                    const std::shared_ptr<const list_data<T, IdxT, SizeT>>& ld,
+                    std::optional<SizeT> size_override = std::nullopt)
+{
+  if (ld) {
+    return serialize_list(handle, os, *ld, size_override);
+  } else {
+    return serialize_scalar(handle, os, SizeT{0});
+  }
+}
+
+template <typename T, typename IdxT, typename SizeT>
+void deserialize_list(const raft::device_resources& handle,
+                      std::istream& is,
+                      std::shared_ptr<list_data<T, IdxT, SizeT>>& ld,
+                      uint32_t dim)
+{
+  auto size = deserialize_scalar<SizeT>(handle, is);
+  if (size == 0) { return ld.reset(); }
+  std::make_shared<list_data<IdxT, SizeT>>(handle, size, dim).swap(ld);
+  auto data_extents = make_extents<SizeT>(size, ld->data.extent(1));
+  auto data_array   = make_host_mdarray<T, SizeT, row_major>(data_extents);
+  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+  deserialize_mdspan(handle, is, data_array.view());
+  deserialize_mdspan(handle, is, inds_array.view());
+  copy(ld->data.data_handle(), data_array.data_handle(), data_array.size(), handle.get_stream());
+  // NB: copying exactly 'size' indices to leave the rest 'kInvalidRecord' intact.
+  copy(ld->indices.data_handle(), inds_array.data_handle(), size, handle.get_stream());
+}
+
 /**
  * Save the index to file.
  *
@@ -469,8 +522,8 @@ void serialize(raft::device_resources const& handle,
   if (!of) { RAFT_FAIL("Cannot open %s", filename.c_str()); }
 
   RAFT_LOG_DEBUG(
-    "Saving IVF-PQ index, size %zu, dim %u", static_cast<size_t>(index_.size()), index_.dim());
-  /* TODO
+    "Saving IVF-Flat index, size %zu, dim %u", static_cast<size_t>(index_.size()), index_.dim());
+
   serialize_scalar(handle, of, serialization_version);
   serialize_scalar(handle, of, index_.size());
   serialize_scalar(handle, of, index_.dim());
@@ -478,10 +531,7 @@ void serialize(raft::device_resources const& handle,
   serialize_scalar(handle, of, index_.metric());
   serialize_scalar(handle, of, index_.veclen());
   serialize_scalar(handle, of, index_.adaptive_centers());
-  serialize_mdspan(handle, of, index_.data());
-  serialize_mdspan(handle, of, index_.indices());
   serialize_mdspan(handle, of, index_.list_sizes());
-  serialize_mdspan(handle, of, index_.list_offsets());
   serialize_mdspan(handle, of, index_.centers());
   if (index_.center_norms()) {
     bool has_norms = true;
@@ -490,7 +540,17 @@ void serialize(raft::device_resources const& handle,
   } else {
     bool has_norms = false;
     serialize_scalar(handle, of, has_norms);
-  }*/
+  }
+  auto sizes_host = make_host_vector<uint32_t, uint32_t>(index_.list_sizes().extent(0));
+  copy(sizes_host.data_handle(),
+       index_.list_sizes().data_handle(),
+       sizes_host.size(),
+       handle.get_stream());
+  handle.sync_stream();
+  serialize_mdspan(handle, of, sizes_host.view());
+  for (uint32_t label = 0; label < index_.n_lists(); label++) {
+    serialize_list<T, IdxT, uint32_t>(handle, of, index_.lists()(label), sizes_host(label));
+  }
   of.close();
   if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); }
 }
@@ -541,7 +601,9 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
       auto center_norms = *index_.center_norms();
       deserialize_mdspan(handle, infile, center_norms);
     }
-  }*/
+  }
+  index_.recompute_internal_state(handle);
+  */
   infile.close();
   return index_;
 }

From c49bf6728d1dbeb2c9b256cb009bcd72e279fe7f Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Thu, 16 Feb 2023 18:44:49 +0100
Subject: [PATCH 07/33] Deserialize ivf_flat and style fix

---
 cpp/include/raft/neighbors/ivf_flat_types.hpp |  66 +++++----
 .../spatial/knn/detail/ivf_flat_build.cuh     | 134 +++++++++++-------
 .../spatial/knn/detail/ivf_flat_search.cuh    |   9 +-
 cpp/test/neighbors/ann_ivf_flat.cuh           |  32 +++--
 4 files changed, 133 insertions(+), 108 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index 476fa755e0..53c73ae6fe 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -19,9 +19,9 @@
 #include "ann_types.hpp"
 
 #include <raft/core/device_mdarray.hpp>
+#include <raft/core/error.hpp>
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/mdspan_types.hpp>
-#include <raft/core/error.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <raft/util/integer_utils.hpp>
 
@@ -87,12 +87,11 @@ struct list_data {
   /** The actual size of the content. */
   std::atomic<SizeT> size;
 
-  list_data(raft::device_resources const& res, SizeT n_rows, uint32_t dim)
-    : size{n_rows}
+  list_data(raft::device_resources const& res, SizeT n_rows, uint32_t dim) : size{n_rows}
   {
     auto capacity = round_up_safe<SizeT>(n_rows, kIndexGroupSize);
     try {
-      data = make_device_matrix<T, SizeT, row_major>(res, capacity, dim);
+      data    = make_device_matrix<T, SizeT, row_major>(res, capacity, dim);
       indices = make_device_vector<IdxT, SizeT>(res, capacity);
     } catch (std::bad_alloc& e) {
       RAFT_FAIL(
@@ -190,12 +189,10 @@ struct index : ann::index {
    * NB: this may be empty if the index is empty or if the metric does not require the center norms
    * calculation.
    */
-  inline auto center_norms() noexcept
-    -> std::optional<device_vector_view<float, uint32_t>>
+  inline auto center_norms() noexcept -> std::optional<device_vector_view<float, uint32_t>>
   {
     if (center_norms_.has_value()) {
-      return std::make_optional<device_vector_view<float, uint32_t>>(
-        center_norms_->view());
+      return std::make_optional<device_vector_view<float, uint32_t>>(center_norms_->view());
     } else {
       return std::nullopt;
     }
@@ -204,18 +201,14 @@ struct index : ann::index {
     -> std::optional<device_vector_view<const float, uint32_t>>
   {
     if (center_norms_.has_value()) {
-      return std::make_optional<device_vector_view<const float, uint32_t>>(
-        center_norms_->view());
+      return std::make_optional<device_vector_view<const float, uint32_t>>(center_norms_->view());
     } else {
       return std::nullopt;
     }
   }
 
   /** Total length of the index. */
-  [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT
-  {
-    return total_size_;
-  }
+  [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT { return total_size_; }
   /** Dimensionality of the data. */
   [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t
   {
@@ -249,12 +242,22 @@ struct index : ann::index {
       lists_{make_host_vector<std::shared_ptr<list_data<T, IdxT>>, uint32_t>(n_lists)},
       list_sizes_{make_device_vector<uint32_t, uint32_t>(res, n_lists)},
       data_ptrs_{make_device_vector<T*, uint32_t>(res, n_lists)},
-      inds_ptrs_{make_device_vector<IdxT*, uint32_t>(res, n_lists)}
+      inds_ptrs_{make_device_vector<IdxT*, uint32_t>(res, n_lists)},
+      total_size_{0}
   {
-    check_consistency();
     for (uint32_t i = 0; i < n_lists; i++) {
       lists_(i) = std::shared_ptr<list_data<T, IdxT>>();
     }
+    switch (metric_) {
+      case raft::distance::DistanceType::L2Expanded:
+      case raft::distance::DistanceType::L2SqrtExpanded:
+      case raft::distance::DistanceType::L2Unexpanded:
+      case raft::distance::DistanceType::L2SqrtUnexpanded:
+        center_norms_ = make_device_vector<float, uint32_t>(res, n_lists);
+        break;
+      default: center_norms_ = std::nullopt;
+    }
+    check_consistency();
   }
 
   /** Construct an empty index. It needs to be trained and then populated. */
@@ -264,12 +267,8 @@ struct index : ann::index {
   }
 
   /** Pointers to the inverted lists (clusters) data  [n_lists]. */
-  inline auto data_ptrs() noexcept -> device_vector_view<T*, uint32_t>
-  {
-    return data_ptrs_.view();
-  }
-  [[nodiscard]] inline auto data_ptrs() const noexcept
-    -> device_vector_view<T* const, uint32_t>
+  inline auto data_ptrs() noexcept -> device_vector_view<T*, uint32_t> { return data_ptrs_.view(); }
+  [[nodiscard]] inline auto data_ptrs() const noexcept -> device_vector_view<T* const, uint32_t>
   {
     return data_ptrs_.view();
   }
@@ -279,8 +278,7 @@ struct index : ann::index {
   {
     return inds_ptrs_.view();
   }
-  [[nodiscard]] inline auto inds_ptrs() const noexcept
-    -> device_vector_view<IdxT* const, uint32_t>
+  [[nodiscard]] inline auto inds_ptrs() const noexcept -> device_vector_view<IdxT* const, uint32_t>
   {
     return inds_ptrs_.view();
   }
@@ -290,27 +288,27 @@ struct index : ann::index {
    */
   void recompute_internal_state(raft::device_resources const& res)
   {
-    auto stream  = res.get_stream();
+    auto stream = res.get_stream();
 
     // Actualize the list pointers
-    auto this_lists     = lists();
-    auto this_data_ptrs = data_ptrs();
-    auto this_inds_ptrs = inds_ptrs();
+    auto this_lists           = lists();
+    auto this_data_ptrs       = data_ptrs();
+    auto this_inds_ptrs       = inds_ptrs();
     IdxT recompute_total_size = 0;
     for (uint32_t label = 0; label < this_lists.size(); label++) {
-      const auto data_ptr = this_lists(label) ? this_lists(label)->data.data_handle() : nullptr;
-      const auto inds_ptr = this_lists(label) ? this_lists(label)->indices.data_handle() : nullptr;
+      const auto data_ptr  = this_lists(label) ? this_lists(label)->data.data_handle() : nullptr;
+      const auto inds_ptr  = this_lists(label) ? this_lists(label)->indices.data_handle() : nullptr;
       const auto list_size = this_lists(label) ? IdxT(this_lists(label)->size) : 0;
       copy(&this_data_ptrs(label), &data_ptr, 1, stream);
       copy(&this_inds_ptrs(label), &inds_ptr, 1, stream);
       recompute_total_size += list_size;
     }
     total_size_ = recompute_total_size;
+    check_consistency();
   }
 
   /** Lists' data and indices. */
-  inline auto lists() noexcept
-    -> host_vector_view<std::shared_ptr<list_data<T, IdxT>>, uint32_t>
+  inline auto lists() noexcept -> host_vector_view<std::shared_ptr<list_data<T, IdxT>>, uint32_t>
   {
     return lists_.view();
   }
@@ -346,8 +344,8 @@ struct index : ann::index {
     RAFT_EXPECTS(list_sizes_.extent(0) == n_lists, "inconsistent list size");
     RAFT_EXPECTS(data_ptrs_.extent(0) == n_lists, "inconsistent list size");
     RAFT_EXPECTS(inds_ptrs_.extent(0) == n_lists, "inconsistent list size");
-    RAFT_EXPECTS(                                               //
-      (centers_.extent(0) == list_sizes_.extent(0)) &&          //
+    RAFT_EXPECTS(                                       //
+      (centers_.extent(0) == list_sizes_.extent(0)) &&  //
         (!center_norms_.has_value() || centers_.extent(0) == center_norms_->extent(0)),
       "inconsistent number of lists (clusters)");
   }
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
index 5a90be93ae..86724c8cea 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
@@ -45,24 +45,24 @@ using namespace raft::neighbors::ivf_flat;   // NOLINT
 using raft::neighbors::ivf_flat::index;
 using raft::neighbors::ivf_flat::index_params;
 using raft::neighbors::ivf_flat::kIndexGroupSize;
-using raft::neighbors::ivf_flat::search_params;
 using raft::neighbors::ivf_flat::list_data;
+using raft::neighbors::ivf_flat::search_params;
 /**
-  * Resize a list by the given id, so that it can contain the given number of records;
-  * possibly, copy the data.
-  *
-  * Besides resizing the corresponding list_data, this function updates the device pointers
-  *   data_ptrs, inds_ptrs, and the list_sizes if necessary.
-  *
-  * The new `list_sizes(label)` represents the number of valid records in the index;
-  * it can be `list_size` if the previous size was not smaller; otherwise it's not updated.
-  *
-  * @param[in] handle
-  * @param[in] label list id
-  * @param[in] list_size the minimum size the list should grow.
-  */
+ * Resize a list by the given id, so that it can contain the given number of records;
+ * possibly, copy the data.
+ *
+ * Besides resizing the corresponding list_data, this function updates the device pointers
+ *   data_ptrs, inds_ptrs, and the list_sizes if necessary.
+ *
+ * The new `list_sizes(label)` represents the number of valid records in the index;
+ * it can be `list_size` if the previous size was not smaller; otherwise it's not updated.
+ *
+ * @param[in] handle
+ * @param[in] label list id
+ * @param[in] list_size the minimum size the list should grow.
+ */
 template <typename T, typename IdxT, typename SizeT = uint32_t>
-void resize_list(raft::device_resources const& handle, 
+void resize_list(raft::device_resources const& handle,
                  std::shared_ptr<list_data<T, IdxT, SizeT>>& orig_list,
                  SizeT new_used_size,
                  SizeT old_used_size,
@@ -89,8 +89,8 @@ void resize_list(raft::device_resources const& handle,
   auto new_list = std::make_shared<list_data<T, IdxT>>(handle, new_used_size, dim);
   if (old_used_size > 0) {
     auto copied_data_extents = make_extents<SizeT>(old_used_size, dim);
-    auto copied_view         = make_mdspan<T, SizeT, row_major, false, true>(
-      new_list->data.data_handle(), copied_data_extents);
+    auto copied_view = make_mdspan<T, SizeT, row_major, false, true>(new_list->data.data_handle(),
+                                                                     copied_data_extents);
     copy(copied_view.data_handle(),
          orig_list->data.data_handle(),
          copied_view.size(),
@@ -104,6 +104,44 @@ void resize_list(raft::device_resources const& handle,
   new_list.swap(orig_list);
 }
 
+template <typename T, typename IdxT>
+auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> index<T, IdxT>
+{
+  auto stream = res.get_stream();
+
+  // Allocate the new index
+  index<T, IdxT> target(
+    res, source.metric(), source.n_lists(), source.adaptive_centers(), source.dim());
+
+  // Copy the independent parts
+  copy(target.list_sizes().data_handle(),
+       source.list_sizes().data_handle(),
+       source.list_sizes().size(),
+       stream);
+  copy(target.centers().data_handle(),
+       source.centers().data_handle(),
+       source.centers().size(),
+       stream);
+  if (source.center_norms().has_value())
+    copy(target.center_norms().value.data_handle(),
+         source.center_norms().value.data_handle(),
+         source.center_norms().value.size(),
+         stream);
+  // Copy shared pointers
+  {
+    auto source_lists = source.lists();
+    auto target_lists = target.lists();
+    for (uint32_t label = 0; label < source_lists.size(); label++) {
+      target_lists(label) = source_lists(label);
+    }
+  }
+
+  // Make sure the device pointers point to the new lists
+  target.recompute_internal_state(res);
+
+  return target;
+}
+
 /**
  * @brief Record the dataset into the index, one source row at a time.
  *
@@ -137,7 +175,7 @@ void resize_list(raft::device_resources const& handle,
  */
 template <typename T, typename IdxT, typename LabelT, bool gather_src = false>
 __global__ void build_index_kernel(const LabelT* labels,
-                                   //const IdxT* list_offsets,
+                                   // const IdxT* list_offsets,
                                    const T* source_vecs,
                                    const IdxT* source_ixs,
                                    T** list_data_ptrs,
@@ -203,10 +241,10 @@ void extend(raft::device_resources const& handle,
 
   auto new_labels = raft::make_device_vector<LabelT, IdxT>(handle, n_rows);
   raft::cluster::kmeans_balanced_params kmeans_params;
-  kmeans_params.metric     = index->metric();
-  auto new_vectors_view    = raft::make_device_matrix_view<const T, IdxT>(new_vectors, n_rows, dim);
-  auto orig_centroids_view = raft::make_device_matrix_view<const float, IdxT>(
-    index->centers().data_handle(), n_lists, dim);
+  kmeans_params.metric  = index->metric();
+  auto new_vectors_view = raft::make_device_matrix_view<const T, IdxT>(new_vectors, n_rows, dim);
+  auto orig_centroids_view =
+    raft::make_device_matrix_view<const float, IdxT>(index->centers().data_handle(), n_lists, dim);
   raft::cluster::kmeans_balanced::predict(handle,
                                           kmeans_params,
                                           new_vectors_view,
@@ -214,13 +252,14 @@ void extend(raft::device_resources const& handle,
                                           new_labels.view(),
                                           utils::mapping<float>{});
 
-  auto* list_sizes_ptr   = index->list_sizes().data_handle();
-  auto old_list_sizes_dev =  raft::make_device_vector<uint32_t, IdxT>(handle, n_lists);
+  auto* list_sizes_ptr    = index->list_sizes().data_handle();
+  auto old_list_sizes_dev = raft::make_device_vector<uint32_t, IdxT>(handle, n_lists);
   copy(old_list_sizes_dev.data_handle(), list_sizes_ptr, n_lists, stream);
 
   // Calculate the centers and sizes on the new data, starting from the original values
   if (index->adaptive_centers()) {
-    auto centroids_view = raft::make_device_matrix_view<float, IdxT>(index->centers().data_handle(), index->centers().extent(0), index->centers().extent(1));
+    auto centroids_view = raft::make_device_matrix_view<float, IdxT>(
+      index->centers().data_handle(), index->centers().extent(0), index->centers().extent(1));
     auto list_sizes_view =
       raft::make_device_vector_view<std::remove_pointer_t<decltype(list_sizes_ptr)>, IdxT>(
         list_sizes_ptr, n_lists);
@@ -253,20 +292,15 @@ void extend(raft::device_resources const& handle,
     handle.sync_stream();
     auto lists = index->lists();
     for (uint32_t label = 0; label < n_lists; label++) {
-      resize_list(handle,
-                  lists(label),
-                  new_list_sizes[label],
-                  old_list_sizes[label],
-                  index->dim());
-      }
+      resize_list(handle, lists(label), new_list_sizes[label], old_list_sizes[label], index->dim());
+    }
   }
   // Update the pointers and the sizes
   index->recompute_internal_state(handle);
 
   // Copy the old sizes, so we can start from the current state of the index;
   // we'll rebuild the `list_sizes_ptr` in the following kernel, using it as an atomic counter.
-  raft::copy(
-    list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream);
+  raft::copy(list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream);
 
   // Kernel to insert the new vectors
   const dim3 block_dim(256);
@@ -308,7 +342,6 @@ auto extend(raft::device_resources const& handle,
   return ext_index;
 }
 
-
 /** See raft::spatial::knn::ivf_flat::build docs */
 template <typename T, typename IdxT>
 inline auto build(raft::device_resources const& handle,
@@ -402,17 +435,13 @@ inline void fill_refinement_index(raft::device_resources const& handle,
     new_labels_view,
     raft::compose_op(raft::cast_op<LabelT>(), raft::div_const_op<IdxT>(n_candidates)));
 
-  auto list_sizes_ptr   = refinement_index->list_sizes().data_handle();
+  auto list_sizes_ptr = refinement_index->list_sizes().data_handle();
   // We do not fill centers and center norms, since we will not run coarse search.
 
   // Allocate new memory
   auto lists = refinement_index->lists();
   for (uint32_t label = 0; label < n_lists; label++) {
-    resize_list(handle,
-                lists(label),
-                n_candidates,
-                uint32_t(0),
-                refinement_index->dim());
+    resize_list(handle, lists(label), n_candidates, uint32_t(0), refinement_index->dim());
   }
 
   RAFT_CUDA_TRY(cudaMemsetAsync(list_sizes_ptr, 0, n_lists * sizeof(uint32_t), stream));
@@ -450,7 +479,6 @@ struct check_index_layout {
 
 template struct check_index_layout<sizeof(index<double, std::uint64_t>), 376>;
 
-
 template <typename T, typename IdxT, typename SizeT>
 void serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
@@ -492,7 +520,7 @@ void deserialize_list(const raft::device_resources& handle,
 {
   auto size = deserialize_scalar<SizeT>(handle, is);
   if (size == 0) { return ld.reset(); }
-  std::make_shared<list_data<IdxT, SizeT>>(handle, size, dim).swap(ld);
+  std::make_shared<list_data<T, IdxT, SizeT>>(handle, size, dim).swap(ld);
   auto data_extents = make_extents<SizeT>(size, ld->data.extent(1));
   auto data_array   = make_host_mdarray<T, SizeT, row_major>(data_extents);
   auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
@@ -531,7 +559,6 @@ void serialize(raft::device_resources const& handle,
   serialize_scalar(handle, of, index_.metric());
   serialize_scalar(handle, of, index_.veclen());
   serialize_scalar(handle, of, index_.adaptive_centers());
-  serialize_mdspan(handle, of, index_.list_sizes());
   serialize_mdspan(handle, of, index_.centers());
   if (index_.center_norms()) {
     bool has_norms = true;
@@ -583,15 +610,8 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   auto veclen           = deserialize_scalar<std::uint32_t>(handle, infile);
   bool adaptive_centers = deserialize_scalar<bool>(handle, infile);
 
-  index<T, IdxT> index_ =
-    index<T, IdxT>(handle, metric, n_lists, adaptive_centers, dim);
-  /* TODO
-  index_.allocate(handle, n_rows);
-  auto data = index_.data();
-  deserialize_mdspan(handle, infile, data);
-  deserialize_mdspan(handle, infile, index_.indices());
-  deserialize_mdspan(handle, infile, index_.list_sizes());
-  deserialize_mdspan(handle, infile, index_.list_offsets());
+  index<T, IdxT> index_ = index<T, IdxT>(handle, metric, n_lists, adaptive_centers, dim);
+
   deserialize_mdspan(handle, infile, index_.centers());
   bool has_norms = deserialize_scalar<bool>(handle, infile);
   if (has_norms) {
@@ -602,9 +622,15 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
       deserialize_mdspan(handle, infile, center_norms);
     }
   }
-  index_.recompute_internal_state(handle);
-  */
+  deserialize_mdspan(handle, infile, index_.list_sizes());
+  for (uint32_t label = 0; label < index_.n_lists(); label++) {
+    deserialize_list<T, IdxT, uint32_t>(handle, infile, index_.lists()(label), index_.dim());
+  }
+  handle.sync_stream();
   infile.close();
+
+  index_.recompute_internal_state(handle);
+
   return index_;
 }
 }  // namespace raft::spatial::knn::ivf_flat::detail
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh b/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
index 8ab3d06efc..905d7abc07 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
+++ b/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
@@ -679,10 +679,9 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
                           const uint32_t query_smem_elems,
                           const T* query,
                           const uint32_t* coarse_index,
-                          const IdxT* const* list_indices_ptrs, // const IdxT* list_indices
-                          const T* const* list_data_ptrs, // const T* list_data
+                          const IdxT* const* list_indices_ptrs,  // const IdxT* list_indices
+                          const T* const* list_data_ptrs,        // const T* list_data
                           const uint32_t* list_sizes,
-                          //const IdxT* list_offsets,
                           const uint32_t n_probes,
                           const uint32_t k,
                           const uint32_t dim,
@@ -728,8 +727,7 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
 
     // Every CUDA block scans one cluster at a time.
     for (int probe_id = blockIdx.x; probe_id < n_probes; probe_id += gridDim.x) {
-      const uint32_t list_id   = coarse_index[probe_id];  // The id of cluster(list)
-      //const size_t list_offset = list_offsets[list_id];
+      const uint32_t list_id = coarse_index[probe_id];  // The id of cluster(list)
 
       // The number of vectors in each cluster(list); [nlist]
       const uint32_t list_length = list_sizes[list_id];
@@ -747,7 +745,6 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
            group_id += kNumWarps) {
         AccT dist = 0;
         // This is where this warp begins reading data (start position of an interleaved group)
-        //const T* data = list_data + (list_offset + group_id * kIndexGroupSize) * dim;
         const T* data = list_data_ptrs[list_id] + (group_id * kIndexGroupSize) * dim;
 
         // This is the vector a given lane/thread handles
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index 3807e414c0..fb0dbb130b 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -58,6 +58,15 @@ struct AnnIvfFlatInputs {
   bool adaptive_centers;
 };
 
+template <typename IdxT>
+::std::ostream& operator<<(::std::ostream& os, const AnnIvfFlatInputs<IdxT>& p)
+{
+  os << "{ " << p.num_queries << ", " << p.num_db_vecs << ", " << p.dim << ", " << p.k << ", "
+     << p.nprobe << ", " << p.nlist << ", " << static_cast<int>(p.metric) << ", "
+     << p.adaptive_centers << '}' << std::endl;
+  return os;
+}
+
 template <typename T, typename DataT, typename IdxT>
 class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
  public:
@@ -210,24 +219,19 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         if (index.adaptive_centers()) {
           // The centers must be up-to-date with the corresponding data
           std::vector<uint32_t> list_sizes(index.n_lists());
-          //std::vector<IdxT> list_offsets(index.n_lists());
           rmm::device_uvector<float> centroid(ps.dim, stream_);
-          raft::copy(
-            list_sizes.data(), index.list_sizes().data_handle(), index.n_lists(), stream_);
-          //raft::copy(
-          //  list_offsets.data(), index.list_offsets().data_handle(), index.n_lists(), stream_);
+          raft::copy(list_sizes.data(), index.list_sizes().data_handle(), index.n_lists(), stream_);
           handle_.sync_stream(stream_);
           for (uint32_t l = 0; l < index.n_lists(); l++) {
             rmm::device_uvector<float> cluster_data(list_sizes[l] * ps.dim, stream_);
-            raft::spatial::knn::detail::utils::copy_selected<float>(
-              (IdxT)list_sizes[l],
-              (IdxT)ps.dim,
-              database.data(),
-              index.inds_ptrs()(l),
-              (IdxT)ps.dim,
-              cluster_data.data(),
-              (IdxT)ps.dim,
-              stream_);
+            raft::spatial::knn::detail::utils::copy_selected<float>((IdxT)list_sizes[l],
+                                                                    (IdxT)ps.dim,
+                                                                    database.data(),
+                                                                    index.inds_ptrs()(l),
+                                                                    (IdxT)ps.dim,
+                                                                    cluster_data.data(),
+                                                                    (IdxT)ps.dim,
+                                                                    stream_);
             raft::stats::mean<float, uint32_t>(
               centroid.data(), cluster_data.data(), ps.dim, list_sizes[l], false, true, stream_);
             ASSERT_TRUE(raft::devArrMatch(index.centers().data_handle() + ps.dim * l,

From 7e2d80bae4bf3cd7f527c7a545eaab4a6c87e3d9 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Tue, 21 Feb 2023 15:29:14 +0100
Subject: [PATCH 08/33] Integrate ivf::list to ivf_flat index splitting

---
 .../detail/ivf_flat_build.cuh                 | 266 ++----------------
 .../detail/ivf_flat_search.cuh                |  10 +-
 .../neighbors/detail/ivf_flat_serialize.cuh   | 158 +++++++++++
 .../neighbors/detail/ivf_pq_serialize.cuh     |   4 +-
 cpp/include/raft/neighbors/detail/refine.cuh  |  22 +-
 cpp/include/raft/neighbors/ivf_flat.cuh       |  23 +-
 cpp/include/raft/neighbors/ivf_flat_types.hpp |  92 +++---
 cpp/include/raft/neighbors/ivf_list.hpp       |  45 ++-
 cpp/include/raft/neighbors/ivf_list_types.hpp |   7 +-
 cpp/include/raft/neighbors/ivf_pq_types.hpp   |   2 +-
 cpp/test/neighbors/ann_ivf_flat.cuh           |  10 +-
 11 files changed, 295 insertions(+), 344 deletions(-)
 rename cpp/include/raft/{spatial/knn => neighbors}/detail/ivf_flat_build.cuh (63%)
 rename cpp/include/raft/{spatial/knn => neighbors}/detail/ivf_flat_search.cuh (99%)
 create mode 100644 cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh

diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
similarity index 63%
rename from cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
rename to cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 86724c8cea..9dafb63149 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -16,93 +16,35 @@
 
 #pragma once
 
-#include "ann_utils.cuh"
-
 #include <raft/cluster/kmeans_balanced.cuh>
 #include <raft/core/device_resources.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/mdarray.hpp>
 #include <raft/core/nvtx.hpp>
 #include <raft/core/operators.hpp>
-#include <raft/core/serialize.hpp>
 #include <raft/linalg/add.cuh>
 #include <raft/linalg/map.cuh>
 #include <raft/linalg/norm.cuh>
 #include <raft/neighbors/ivf_flat_types.hpp>
+#include <raft/neighbors/ivf_list.hpp>
+#include <raft/neighbors/ivf_list_types.hpp>
+#include <raft/spatial/knn/detail/ann_utils.cuh>
 #include <raft/stats/histogram.cuh>
 #include <raft/util/pow2_utils.cuh>
 
 #include <rmm/cuda_stream_view.hpp>
 
 #include <cstdint>
-#include <fstream>
 
-namespace raft::spatial::knn::ivf_flat::detail {
+namespace raft::neighbors::ivf_flat::detail {
 
 using namespace raft::spatial::knn::detail;  // NOLINT
-using namespace raft::neighbors::ivf_flat;   // NOLINT
 
 using raft::neighbors::ivf_flat::index;
 using raft::neighbors::ivf_flat::index_params;
 using raft::neighbors::ivf_flat::kIndexGroupSize;
 using raft::neighbors::ivf_flat::list_data;
 using raft::neighbors::ivf_flat::search_params;
-/**
- * Resize a list by the given id, so that it can contain the given number of records;
- * possibly, copy the data.
- *
- * Besides resizing the corresponding list_data, this function updates the device pointers
- *   data_ptrs, inds_ptrs, and the list_sizes if necessary.
- *
- * The new `list_sizes(label)` represents the number of valid records in the index;
- * it can be `list_size` if the previous size was not smaller; otherwise it's not updated.
- *
- * @param[in] handle
- * @param[in] label list id
- * @param[in] list_size the minimum size the list should grow.
- */
-template <typename T, typename IdxT, typename SizeT = uint32_t>
-void resize_list(raft::device_resources const& handle,
-                 std::shared_ptr<list_data<T, IdxT, SizeT>>& orig_list,
-                 SizeT new_used_size,
-                 SizeT old_used_size,
-                 uint32_t dim)
-{
-  bool skip_resize = false;
-  // TODO update total_size
-  if (orig_list) {
-    if (new_used_size <= orig_list->indices.extent(0)) {
-      auto shared_list_size = old_used_size;
-      if (new_used_size <= old_used_size ||
-          orig_list->size.compare_exchange_strong(shared_list_size, new_used_size)) {
-        // We don't need to resize the list if:
-        //  1. The list exists
-        //  2. The new size fits in the list
-        //  3. The list doesn't grow or no-one else has grown it yet
-        skip_resize = true;
-      }
-    }
-  } else {
-    old_used_size = 0;
-  }
-  if (skip_resize) { return; }
-  auto new_list = std::make_shared<list_data<T, IdxT>>(handle, new_used_size, dim);
-  if (old_used_size > 0) {
-    auto copied_data_extents = make_extents<SizeT>(old_used_size, dim);
-    auto copied_view = make_mdspan<T, SizeT, row_major, false, true>(new_list->data.data_handle(),
-                                                                     copied_data_extents);
-    copy(copied_view.data_handle(),
-         orig_list->data.data_handle(),
-         copied_view.size(),
-         handle.get_stream());
-    copy(new_list->indices.data_handle(),
-         orig_list->indices.data_handle(),
-         old_used_size,
-         handle.get_stream());
-  }
-  // swap the shared pointer content with the new list
-  new_list.swap(orig_list);
-}
 
 template <typename T, typename IdxT>
 auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> index<T, IdxT>
@@ -110,8 +52,12 @@ auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> i
   auto stream = res.get_stream();
 
   // Allocate the new index
-  index<T, IdxT> target(
-    res, source.metric(), source.n_lists(), source.adaptive_centers(), source.dim());
+  index<T, IdxT> target(res,
+                        source.metric(),
+                        source.n_lists(),
+                        source.adaptive_centers(),
+                        source.conservative_memory_allocation(),
+                        source.dim());
 
   // Copy the independent parts
   copy(target.list_sizes().data_handle(),
@@ -219,7 +165,7 @@ __global__ void build_index_kernel(const LabelT* labels,
   }
 }
 
-/** See raft::spatial::knn::ivf_flat::extend docs */
+/** See raft::neighbors::ivf_flat::extend docs */
 template <typename T, typename IdxT>
 void extend(raft::device_resources const& handle,
             index<T, IdxT>* index,
@@ -290,9 +236,11 @@ void extend(raft::device_resources const& handle,
     copy(old_list_sizes.data(), old_list_sizes_dev.data_handle(), n_lists, stream);
     copy(new_list_sizes.data(), list_sizes_ptr, n_lists, stream);
     handle.sync_stream();
-    auto lists = index->lists();
+    auto lists            = index->lists();
+    auto list_device_spec = list_spec<uint32_t>{index->dim(), false};
     for (uint32_t label = 0; label < n_lists; label++) {
-      resize_list(handle, lists(label), new_list_sizes[label], old_list_sizes[label], index->dim());
+      ivf::resize_list(
+        handle, lists(label), list_device_spec, new_list_sizes[label], old_list_sizes[label]);
     }
   }
   // Update the pointers and the sizes
@@ -329,7 +277,7 @@ void extend(raft::device_resources const& handle,
   }
 }
 
-/** See raft::spatial::knn::ivf_flat::extend docs */
+/** See raft::neighbors::ivf_flat::extend docs */
 template <typename T, typename IdxT>
 auto extend(raft::device_resources const& handle,
             const index<T, IdxT>& orig_index,
@@ -342,7 +290,7 @@ auto extend(raft::device_resources const& handle,
   return ext_index;
 }
 
-/** See raft::spatial::knn::ivf_flat::build docs */
+/** See raft::neighbors::ivf_flat::build docs */
 template <typename T, typename IdxT>
 inline auto build(raft::device_resources const& handle,
                   const index_params& params,
@@ -439,9 +387,10 @@ inline void fill_refinement_index(raft::device_resources const& handle,
   // We do not fill centers and center norms, since we will not run coarse search.
 
   // Allocate new memory
-  auto lists = refinement_index->lists();
+  auto lists            = refinement_index->lists();
+  auto list_device_spec = list_spec<uint32_t>{refinement_index->dim(), false};
   for (uint32_t label = 0; label < n_lists; label++) {
-    resize_list(handle, lists(label), n_candidates, uint32_t(0), refinement_index->dim());
+    ivf::resize_list(handle, lists(label), list_device_spec, n_candidates, uint32_t(0));
   }
 
   RAFT_CUDA_TRY(cudaMemsetAsync(list_sizes_ptr, 0, n_lists * sizeof(uint32_t), stream));
@@ -460,177 +409,4 @@ inline void fill_refinement_index(raft::device_resources const& handle,
                                          refinement_index->veclen());
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }
-
-// Serialization version 3
-// No backward compatibility yet; that is, can't add additional fields without breaking
-// backward compatibility.
-// TODO(hcho3) Implement next-gen serializer for IVF that allows for expansion in a backward
-//             compatible fashion.
-constexpr int serialization_version = 3;
-
-// NB: we wrap this check in a struct, so that the updated RealSize is easy to see in the error
-// message.
-template <size_t RealSize, size_t ExpectedSize>
-struct check_index_layout {
-  static_assert(RealSize == ExpectedSize,
-                "The size of the index struct has changed since the last update; "
-                "paste in the new size and consider updating the serialization logic");
-};
-
-template struct check_index_layout<sizeof(index<double, std::uint64_t>), 376>;
-
-template <typename T, typename IdxT, typename SizeT>
-void serialize_list(const raft::device_resources& handle,
-                    std::ostream& os,
-                    const list_data<T, IdxT, SizeT>& ld,
-                    std::optional<SizeT> size_override = std::nullopt)
-{
-  auto size = size_override.value_or(ld.size.load());
-  serialize_scalar(handle, os, size);
-  if (size == 0) { return; }
-
-  auto data_extents = make_extents<SizeT>(size, ld.data.extent(1));
-  auto data_array   = make_host_mdarray<T, SizeT, row_major>(data_extents);
-  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
-  copy(data_array.data_handle(), ld.data.data_handle(), data_array.size(), handle.get_stream());
-  copy(inds_array.data_handle(), ld.indices.data_handle(), inds_array.size(), handle.get_stream());
-  handle.sync_stream();
-  serialize_mdspan(handle, os, data_array.view());
-  serialize_mdspan(handle, os, inds_array.view());
-}
-
-template <typename T, typename IdxT, typename SizeT>
-void serialize_list(const raft::device_resources& handle,
-                    std::ostream& os,
-                    const std::shared_ptr<const list_data<T, IdxT, SizeT>>& ld,
-                    std::optional<SizeT> size_override = std::nullopt)
-{
-  if (ld) {
-    return serialize_list(handle, os, *ld, size_override);
-  } else {
-    return serialize_scalar(handle, os, SizeT{0});
-  }
-}
-
-template <typename T, typename IdxT, typename SizeT>
-void deserialize_list(const raft::device_resources& handle,
-                      std::istream& is,
-                      std::shared_ptr<list_data<T, IdxT, SizeT>>& ld,
-                      uint32_t dim)
-{
-  auto size = deserialize_scalar<SizeT>(handle, is);
-  if (size == 0) { return ld.reset(); }
-  std::make_shared<list_data<T, IdxT, SizeT>>(handle, size, dim).swap(ld);
-  auto data_extents = make_extents<SizeT>(size, ld->data.extent(1));
-  auto data_array   = make_host_mdarray<T, SizeT, row_major>(data_extents);
-  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
-  deserialize_mdspan(handle, is, data_array.view());
-  deserialize_mdspan(handle, is, inds_array.view());
-  copy(ld->data.data_handle(), data_array.data_handle(), data_array.size(), handle.get_stream());
-  // NB: copying exactly 'size' indices to leave the rest 'kInvalidRecord' intact.
-  copy(ld->indices.data_handle(), inds_array.data_handle(), size, handle.get_stream());
-}
-
-/**
- * Save the index to file.
- *
- * Experimental, both the API and the serialization format are subject to change.
- *
- * @param[in] handle the raft handle
- * @param[in] filename the file name for saving the index
- * @param[in] index_ IVF-Flat index
- *
- */
-template <typename T, typename IdxT>
-void serialize(raft::device_resources const& handle,
-               const std::string& filename,
-               const index<T, IdxT>& index_)
-{
-  std::ofstream of(filename, std::ios::out | std::ios::binary);
-  if (!of) { RAFT_FAIL("Cannot open %s", filename.c_str()); }
-
-  RAFT_LOG_DEBUG(
-    "Saving IVF-Flat index, size %zu, dim %u", static_cast<size_t>(index_.size()), index_.dim());
-
-  serialize_scalar(handle, of, serialization_version);
-  serialize_scalar(handle, of, index_.size());
-  serialize_scalar(handle, of, index_.dim());
-  serialize_scalar(handle, of, index_.n_lists());
-  serialize_scalar(handle, of, index_.metric());
-  serialize_scalar(handle, of, index_.veclen());
-  serialize_scalar(handle, of, index_.adaptive_centers());
-  serialize_mdspan(handle, of, index_.centers());
-  if (index_.center_norms()) {
-    bool has_norms = true;
-    serialize_scalar(handle, of, has_norms);
-    serialize_mdspan(handle, of, *index_.center_norms());
-  } else {
-    bool has_norms = false;
-    serialize_scalar(handle, of, has_norms);
-  }
-  auto sizes_host = make_host_vector<uint32_t, uint32_t>(index_.list_sizes().extent(0));
-  copy(sizes_host.data_handle(),
-       index_.list_sizes().data_handle(),
-       sizes_host.size(),
-       handle.get_stream());
-  handle.sync_stream();
-  serialize_mdspan(handle, of, sizes_host.view());
-  for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    serialize_list<T, IdxT, uint32_t>(handle, of, index_.lists()(label), sizes_host(label));
-  }
-  of.close();
-  if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); }
-}
-
-/** Load an index from file.
- *
- * Experimental, both the API and the serialization format are subject to change.
- *
- * @param[in] handle the raft handle
- * @param[in] filename the name of the file that stores the index
- * @param[in] index_ IVF-Flat index
- *
- */
-template <typename T, typename IdxT>
-auto deserialize(raft::device_resources const& handle, const std::string& filename)
-  -> index<T, IdxT>
-{
-  std::ifstream infile(filename, std::ios::in | std::ios::binary);
-
-  if (!infile) { RAFT_FAIL("Cannot open %s", filename.c_str()); }
-
-  auto ver = deserialize_scalar<int>(handle, infile);
-  if (ver != serialization_version) {
-    RAFT_FAIL("serialization version mismatch, expected %d, got %d ", serialization_version, ver);
-  }
-  auto n_rows           = deserialize_scalar<IdxT>(handle, infile);
-  auto dim              = deserialize_scalar<std::uint32_t>(handle, infile);
-  auto n_lists          = deserialize_scalar<std::uint32_t>(handle, infile);
-  auto metric           = deserialize_scalar<raft::distance::DistanceType>(handle, infile);
-  auto veclen           = deserialize_scalar<std::uint32_t>(handle, infile);
-  bool adaptive_centers = deserialize_scalar<bool>(handle, infile);
-
-  index<T, IdxT> index_ = index<T, IdxT>(handle, metric, n_lists, adaptive_centers, dim);
-
-  deserialize_mdspan(handle, infile, index_.centers());
-  bool has_norms = deserialize_scalar<bool>(handle, infile);
-  if (has_norms) {
-    if (!index_.center_norms()) {
-      RAFT_FAIL("Error inconsistent center norms");
-    } else {
-      auto center_norms = *index_.center_norms();
-      deserialize_mdspan(handle, infile, center_norms);
-    }
-  }
-  deserialize_mdspan(handle, infile, index_.list_sizes());
-  for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    deserialize_list<T, IdxT, uint32_t>(handle, infile, index_.lists()(label), index_.dim());
-  }
-  handle.sync_stream();
-  infile.close();
-
-  index_.recompute_internal_state(handle);
-
-  return index_;
-}
-}  // namespace raft::spatial::knn::ivf_flat::detail
+}  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_search.cuh
similarity index 99%
rename from cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
rename to cpp/include/raft/neighbors/detail/ivf_flat_search.cuh
index 905d7abc07..0003b02fc5 100644
--- a/cpp/include/raft/spatial/knn/detail/ivf_flat_search.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_search.cuh
@@ -16,8 +16,6 @@
 
 #pragma once
 
-#include "ann_utils.cuh"
-
 #include <raft/core/cudart_utils.hpp>
 #include <raft/core/device_resources.hpp>
 #include <raft/core/logger.hpp>
@@ -30,6 +28,7 @@
 #include <raft/matrix/detail/select_k.cuh>
 #include <raft/matrix/detail/select_warpsort.cuh>
 #include <raft/neighbors/ivf_flat_types.hpp>
+#include <raft/spatial/knn/detail/ann_utils.cuh>
 #include <raft/util/cuda_utils.cuh>
 #include <raft/util/device_loads_stores.cuh>
 #include <raft/util/integer_utils.hpp>
@@ -39,10 +38,9 @@
 #include <rmm/cuda_stream_view.hpp>
 #include <rmm/mr/device/per_device_resource.hpp>
 
-namespace raft::spatial::knn::ivf_flat::detail {
+namespace raft::neighbors::ivf_flat::detail {
 
 using namespace raft::spatial::knn::detail;  // NOLINT
-using namespace raft::neighbors::ivf_flat;   // NOLINT
 
 using raft::neighbors::ivf_flat::index;
 using raft::neighbors::ivf_flat::index_params;
@@ -1271,7 +1269,7 @@ inline bool is_min_close(distance::DistanceType metric)
   return select_min;
 }
 
-/** See raft::spatial::knn::ivf_flat::search docs */
+/** See raft::neighbors::ivf_flat::search docs */
 template <typename T, typename IdxT>
 inline void search(raft::device_resources const& handle,
                    const search_params& params,
@@ -1308,4 +1306,4 @@ inline void search(raft::device_resources const& handle,
                                      mr);
 }
 
-}  // namespace raft::spatial::knn::ivf_flat::detail
+}  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
new file mode 100644
index 0000000000..405bc3252f
--- /dev/null
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -0,0 +1,158 @@
+/*
+ * Copyright (c) 2023, 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/core/mdarray.hpp>
+#include <raft/core/serialize.hpp>
+#include <raft/neighbors/ivf_flat_types.hpp>
+#include <raft/neighbors/ivf_list.hpp>
+#include <raft/neighbors/ivf_list_types.hpp>
+
+#include <fstream>
+
+namespace raft::neighbors::ivf_flat::detail {
+
+// Serialization version 3
+// No backward compatibility yet; that is, can't add additional fields without breaking
+// backward compatibility.
+// TODO(hcho3) Implement next-gen serializer for IVF that allows for expansion in a backward
+//             compatible fashion.
+constexpr int serialization_version = 3;
+
+// NB: we wrap this check in a struct, so that the updated RealSize is easy to see in the error
+// message.
+template <size_t RealSize, size_t ExpectedSize>
+struct check_index_layout {
+  static_assert(RealSize == ExpectedSize,
+                "The size of the index struct has changed since the last update; "
+                "paste in the new size and consider updating the serialization logic");
+};
+
+template struct check_index_layout<sizeof(index<double, std::uint64_t>), 376>;
+
+/**
+ * Save the index to file.
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @param[in] handle the raft handle
+ * @param[in] filename the file name for saving the index
+ * @param[in] index_ IVF-Flat index
+ *
+ */
+template <typename T, typename IdxT>
+void serialize(raft::device_resources const& handle,
+               const std::string& filename,
+               const index<T, IdxT>& index_)
+{
+  std::ofstream of(filename, std::ios::out | std::ios::binary);
+  if (!of) { RAFT_FAIL("Cannot open %s", filename.c_str()); }
+
+  RAFT_LOG_DEBUG(
+    "Saving IVF-Flat index, size %zu, dim %u", static_cast<size_t>(index_.size()), index_.dim());
+
+  serialize_scalar(handle, of, serialization_version);
+  serialize_scalar(handle, of, index_.size());
+  serialize_scalar(handle, of, index_.dim());
+  serialize_scalar(handle, of, index_.n_lists());
+  serialize_scalar(handle, of, index_.metric());
+  serialize_scalar(handle, of, index_.veclen());
+  serialize_scalar(handle, of, index_.adaptive_centers());
+  serialize_scalar(handle, of, index_.conservative_memory_allocation());
+  serialize_mdspan(handle, of, index_.centers());
+  if (index_.center_norms()) {
+    bool has_norms = true;
+    serialize_scalar(handle, of, has_norms);
+    serialize_mdspan(handle, of, *index_.center_norms());
+  } else {
+    bool has_norms = false;
+    serialize_scalar(handle, of, has_norms);
+  }
+  auto sizes_host = make_host_vector<uint32_t, uint32_t>(index_.list_sizes().extent(0));
+  copy(sizes_host.data_handle(),
+       index_.list_sizes().data_handle(),
+       sizes_host.size(),
+       handle.get_stream());
+  handle.sync_stream();
+  serialize_mdspan(handle, of, sizes_host.view());
+
+  auto list_store_spec = list_spec<uint32_t>{index_.dim(), true};
+  for (uint32_t label = 0; label < index_.n_lists(); label++) {
+    ivf::serialize_list<list_spec, T, IdxT, uint32_t>(
+      handle, of, index_.lists()(label), list_store_spec, sizes_host(label));
+  }
+  of.close();
+  if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); }
+}
+
+/** Load an index from file.
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @param[in] handle the raft handle
+ * @param[in] filename the name of the file that stores the index
+ * @param[in] index_ IVF-Flat index
+ *
+ */
+template <typename T, typename IdxT>
+auto deserialize(raft::device_resources const& handle, const std::string& filename)
+  -> index<T, IdxT>
+{
+  std::ifstream infile(filename, std::ios::in | std::ios::binary);
+
+  if (!infile) { RAFT_FAIL("Cannot open %s", filename.c_str()); }
+
+  auto ver = deserialize_scalar<int>(handle, infile);
+  if (ver != serialization_version) {
+    RAFT_FAIL("serialization version mismatch, expected %d, got %d ", serialization_version, ver);
+  }
+  auto n_rows           = deserialize_scalar<IdxT>(handle, infile);
+  auto dim              = deserialize_scalar<std::uint32_t>(handle, infile);
+  auto n_lists          = deserialize_scalar<std::uint32_t>(handle, infile);
+  auto metric           = deserialize_scalar<raft::distance::DistanceType>(handle, infile);
+  auto veclen           = deserialize_scalar<std::uint32_t>(handle, infile);
+  bool adaptive_centers = deserialize_scalar<bool>(handle, infile);
+  bool cma              = deserialize_scalar<bool>(handle, infile);
+
+  index<T, IdxT> index_ = index<T, IdxT>(handle, metric, n_lists, adaptive_centers, cma, dim);
+
+  deserialize_mdspan(handle, infile, index_.centers());
+  bool has_norms = deserialize_scalar<bool>(handle, infile);
+  if (has_norms) {
+    if (!index_.center_norms()) {
+      RAFT_FAIL("Error inconsistent center norms");
+    } else {
+      auto center_norms = *index_.center_norms();
+      deserialize_mdspan(handle, infile, center_norms);
+    }
+  }
+  deserialize_mdspan(handle, infile, index_.list_sizes());
+
+  auto list_device_spec = list_spec<uint32_t>{index_.dim(), cma};
+  auto list_store_spec  = list_spec<uint32_t>{index_.dim(), true};
+  for (uint32_t label = 0; label < index_.n_lists(); label++) {
+    ivf::deserialize_list<list_spec, T, IdxT, uint32_t>(
+      handle, infile, index_.lists()(label), list_store_spec, list_device_spec);
+  }
+  handle.sync_stream();
+  infile.close();
+
+  index_.recompute_internal_state(handle);
+
+  return index_;
+}
+}  // namespace raft::neighbors::ivf_flat::detail
\ No newline at end of file
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
index 33d9b363ba..bab9a52ed1 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
@@ -96,7 +96,7 @@ void serialize(raft::device_resources const& handle_,
   serialize_mdspan(handle_, of, sizes_host.view());
   auto list_store_spec = list_spec<uint32_t>{index.pq_bits(), index.pq_dim(), true};
   for (uint32_t label = 0; label < index.n_lists(); label++) {
-    ivf::serialize_list<list_spec, IdxT, uint32_t>(
+    ivf::serialize_list<list_spec, uint8_t, IdxT, uint32_t>(
       handle_, of, index.lists()[label], list_store_spec, sizes_host(label));
   }
 
@@ -154,7 +154,7 @@ auto deserialize(raft::device_resources const& handle_, const std::string& filen
   auto list_device_spec = list_spec<uint32_t>{pq_bits, pq_dim, cma};
   auto list_store_spec  = list_spec<uint32_t>{pq_bits, pq_dim, true};
   for (auto& list : index.lists()) {
-    ivf::deserialize_list<list_spec, IdxT, uint32_t>(
+    ivf::deserialize_list<list_spec, uint8_t, IdxT, uint32_t>(
       handle_, infile, list, list_store_spec, list_device_spec);
   }
 
diff --git a/cpp/include/raft/neighbors/detail/refine.cuh b/cpp/include/raft/neighbors/detail/refine.cuh
index b264643584..5d226132a4 100644
--- a/cpp/include/raft/neighbors/detail/refine.cuh
+++ b/cpp/include/raft/neighbors/detail/refine.cuh
@@ -20,9 +20,9 @@
 #include <raft/core/device_resources.hpp>
 #include <raft/core/host_mdspan.hpp>
 #include <raft/core/nvtx.hpp>
+#include <raft/neighbors/detail/ivf_flat_build.cuh>
+#include <raft/neighbors/detail/ivf_flat_search.cuh>
 #include <raft/spatial/knn/detail/ann_utils.cuh>
-#include <raft/spatial/knn/detail/ivf_flat_build.cuh>
-#include <raft/spatial/knn/detail/ivf_flat_search.cuh>
 
 #include <cstdlib>
 #include <omp.h>
@@ -108,17 +108,17 @@ void refine_device(raft::device_resources const& handle,
     handle.get_thrust_policy(), fake_coarse_idx.data(), fake_coarse_idx.data() + n_queries);
 
   raft::neighbors::ivf_flat::index<data_t, idx_t> refinement_index(
-    handle, metric, n_queries, false, dim);
+    handle, metric, n_queries, false, false, dim);
 
-  raft::spatial::knn::ivf_flat::detail::fill_refinement_index(handle,
-                                                              &refinement_index,
-                                                              dataset.data_handle(),
-                                                              neighbor_candidates.data_handle(),
-                                                              n_queries,
-                                                              n_candidates);
+  raft::neighbors::ivf_flat::detail::fill_refinement_index(handle,
+                                                           &refinement_index,
+                                                           dataset.data_handle(),
+                                                           neighbor_candidates.data_handle(),
+                                                           n_queries,
+                                                           n_candidates);
 
   uint32_t grid_dim_x = 1;
-  raft::spatial::knn::ivf_flat::detail::ivfflat_interleaved_scan<
+  raft::neighbors::ivf_flat::detail::ivfflat_interleaved_scan<
     data_t,
     typename raft::spatial::knn::detail::utils::config<data_t>::value_t,
     idx_t>(refinement_index,
@@ -128,7 +128,7 @@ void refine_device(raft::device_resources const& handle,
            refinement_index.metric(),
            1,
            k,
-           raft::spatial::knn::ivf_flat::detail::is_min_close(metric),
+           raft::neighbors::ivf_flat::detail::is_min_close(metric),
            indices.data_handle(),
            distances.data_handle(),
            grid_dim_x,
diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh
index d210d2c74b..4a5e2a3641 100644
--- a/cpp/include/raft/neighbors/ivf_flat.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat.cuh
@@ -16,9 +16,10 @@
 
 #pragma once
 
+#include <raft/neighbors/detail/ivf_flat_build.cuh>
+#include <raft/neighbors/detail/ivf_flat_search.cuh>
+#include <raft/neighbors/detail/ivf_flat_serialize.cuh>
 #include <raft/neighbors/ivf_flat_types.hpp>
-#include <raft/spatial/knn/detail/ivf_flat_build.cuh>
-#include <raft/spatial/knn/detail/ivf_flat_search.cuh>
 
 #include <raft/core/device_resources.hpp>
 
@@ -67,7 +68,7 @@ auto build(raft::device_resources const& handle,
            IdxT n_rows,
            uint32_t dim) -> index<T, IdxT>
 {
-  return raft::spatial::knn::ivf_flat::detail::build(handle, params, dataset, n_rows, dim);
+  return raft::neighbors::ivf_flat::detail::build(handle, params, dataset, n_rows, dim);
 }
 
 /**
@@ -112,11 +113,11 @@ auto build(raft::device_resources const& handle,
            raft::device_matrix_view<const value_t, idx_t, row_major> dataset,
            const index_params& params) -> index<value_t, idx_t>
 {
-  return raft::spatial::knn::ivf_flat::detail::build(handle,
-                                                     params,
-                                                     dataset.data_handle(),
-                                                     static_cast<idx_t>(dataset.extent(0)),
-                                                     static_cast<idx_t>(dataset.extent(1)));
+  return raft::neighbors::ivf_flat::detail::build(handle,
+                                                  params,
+                                                  dataset.data_handle(),
+                                                  static_cast<idx_t>(dataset.extent(0)),
+                                                  static_cast<idx_t>(dataset.extent(1)));
 }
 
 /** @} */
@@ -160,7 +161,7 @@ auto extend(raft::device_resources const& handle,
             const IdxT* new_indices,
             IdxT n_rows) -> index<T, IdxT>
 {
-  return raft::spatial::knn::ivf_flat::detail::extend(
+  return raft::neighbors::ivf_flat::detail::extend(
     handle, orig_index, new_vectors, new_indices, n_rows);
 }
 
@@ -252,7 +253,7 @@ void extend(raft::device_resources const& handle,
             const IdxT* new_indices,
             IdxT n_rows)
 {
-  raft::spatial::knn::ivf_flat::detail::extend(handle, index, new_vectors, new_indices, n_rows);
+  raft::neighbors::ivf_flat::detail::extend(handle, index, new_vectors, new_indices, n_rows);
 }
 
 /**
@@ -355,7 +356,7 @@ void search(raft::device_resources const& handle,
             float* distances,
             rmm::mr::device_memory_resource* mr = nullptr)
 {
-  return raft::spatial::knn::ivf_flat::detail::search(
+  return raft::neighbors::ivf_flat::detail::search(
     handle, params, index, queries, n_queries, k, neighbors, distances, mr);
 }
 
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index 53c73ae6fe..b60d5b3aa7 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -23,6 +23,7 @@
 #include <raft/core/host_mdarray.hpp>
 #include <raft/core/mdspan_types.hpp>
 #include <raft/distance/distance_types.hpp>
+#include <raft/neighbors/ivf_list_types.hpp>
 #include <raft/util/integer_utils.hpp>
 
 #include <memory>
@@ -39,14 +40,6 @@ namespace raft::neighbors::ivf_flat {
 /** Size of the interleaved group (see `index::data` description). */
 constexpr static uint32_t kIndexGroupSize = 32;
 
-/**
- * Default value filled in the `indices()` array.
- * One may encounter it trying to access a record within a cluster that is outside of the
- * `list_sizes()` bound (due to the record alignment `kIndexGroupSize`).
- */
-template <typename IdxT>
-constexpr static IdxT kInvalidRecord = std::numeric_limits<IdxT>::max() - 1;
-
 struct index_params : ann::index_params {
   /** The number of inverted lists (clusters) */
   uint32_t n_lists = 1024;
@@ -67,6 +60,16 @@ struct index_params : ann::index_params {
    * `index.centers()` "drift" together with the changing distribution of the newly added data.
    */
   bool adaptive_centers = false;
+  /**
+   * By default, the algorithm allocates more space than necessary for individual clusters
+   * (`list_data`). This allows to amortize the cost of memory allocation and reduce the number of
+   * data copies during repeated calls to `extend` (extending the database).
+   *
+   * The alternative is the conservative allocation behavior; when enabled, the algorithm always
+   * allocates the minimum amount of memory required to store the given number of records. Set this
+   * flag to `true` if you prefer to use as little GPU memory for the database as possible.
+   */
+  bool conservative_memory_allocation = false;
 };
 
 struct search_params : ann::search_params {
@@ -77,37 +80,38 @@ struct search_params : ann::search_params {
 static_assert(std::is_aggregate_v<index_params>);
 static_assert(std::is_aggregate_v<search_params>);
 
-/** The data for a single list (cluster). */
-template <typename T, typename IdxT, typename SizeT = uint32_t>
-struct list_data {
-  /** Cluster data. */
-  device_matrix<T, SizeT, row_major> data;
-  /** Source indices. */
-  device_vector<IdxT, SizeT> indices;
-  /** The actual size of the content. */
-  std::atomic<SizeT> size;
-
-  list_data(raft::device_resources const& res, SizeT n_rows, uint32_t dim) : size{n_rows}
+template <typename SizeT = uint32_t>
+struct list_spec {
+  using list_extents = matrix_extent<SizeT>;
+
+  SizeT align_max;
+  SizeT align_min;
+  uint32_t dim;
+
+  constexpr list_spec(uint32_t dim, bool conservative_memory_allocation)
+    : dim(dim),
+      align_min(kIndexGroupSize),
+      align_max(conservative_memory_allocation ? kIndexGroupSize : 1024)
   {
-    auto capacity = round_up_safe<SizeT>(n_rows, kIndexGroupSize);
-    try {
-      data    = make_device_matrix<T, SizeT, row_major>(res, capacity, dim);
-      indices = make_device_vector<IdxT, SizeT>(res, capacity);
-    } catch (std::bad_alloc& e) {
-      RAFT_FAIL(
-        "ivf-flat: failed to allocate a big enough index list to hold all data "
-        "(requested size: %zu records, selected capacity: %zu records). "
-        "Allocator exception: %s",
-        size_t(n_rows),
-        size_t(capacity),
-        e.what());
-    }
-    // Fill the index buffer with a pre-defined marker for easier debugging
-    thrust::fill_n(
-      res.get_thrust_policy(), indices.data_handle(), indices.size(), kInvalidRecord<IdxT>);
+  }
+
+  // Allow casting between different size-types (for safer size and offset calculations)
+  template <typename OtherSizeT>
+  constexpr explicit list_spec(const list_spec<OtherSizeT>& other_spec)
+    : dim{other_spec.dim}, align_min{other_spec.align_min}, align_max{other_spec.align_max}
+  {
+  }
+
+  /** Determine the extents of an array enough to hold a given amount of data. */
+  constexpr auto make_list_extents(SizeT n_rows) const -> list_extents
+  {
+    return make_extents<SizeT>(n_rows, dim);
   }
 };
 
+template <typename ValueT, typename IdxT, typename SizeT = uint32_t>
+using list_data = ivf::list<list_spec, ValueT, IdxT, SizeT>;
+
 /**
  * @brief IVF-flat index.
  *
@@ -232,11 +236,13 @@ struct index : ann::index {
         raft::distance::DistanceType metric,
         uint32_t n_lists,
         bool adaptive_centers,
+        bool conservative_memory_allocation,
         uint32_t dim)
     : ann::index(),
       veclen_(calculate_veclen(dim)),
       metric_(metric),
       adaptive_centers_(adaptive_centers),
+      conservative_memory_allocation_{conservative_memory_allocation},
       centers_(make_device_matrix<float, uint32_t>(res, n_lists, dim)),
       center_norms_(std::nullopt),
       lists_{make_host_vector<std::shared_ptr<list_data<T, IdxT>>, uint32_t>(n_lists)},
@@ -262,7 +268,12 @@ struct index : ann::index {
 
   /** Construct an empty index. It needs to be trained and then populated. */
   index(raft::device_resources const& res, const index_params& params, uint32_t dim)
-    : index(res, params.metric, params.n_lists, params.adaptive_centers, dim)
+    : index(res,
+            params.metric,
+            params.n_lists,
+            params.adaptive_centers,
+            params.conservative_memory_allocation,
+            dim)
   {
   }
 
@@ -282,6 +293,14 @@ struct index : ann::index {
   {
     return inds_ptrs_.view();
   }
+  /**
+   * Whether to use convervative memory allocation when extending the list (cluster) data
+   * (see index_params.conservative_memory_allocation).
+   */
+  [[nodiscard]] constexpr inline auto conservative_memory_allocation() const noexcept -> bool
+  {
+    return conservative_memory_allocation_;
+  }
 
   /**
    * Update the state of the dependent index members.
@@ -326,6 +345,7 @@ struct index : ann::index {
   uint32_t veclen_;
   raft::distance::DistanceType metric_;
   bool adaptive_centers_;
+  bool conservative_memory_allocation_;
   host_vector<std::shared_ptr<list_data<T, IdxT>>, uint32_t> lists_;
   device_vector<uint32_t, uint32_t> list_sizes_;
   device_matrix<float, uint32_t, row_major> centers_;
diff --git a/cpp/include/raft/neighbors/ivf_list.hpp b/cpp/include/raft/neighbors/ivf_list.hpp
index 4644143057..2fe3d2cf3d 100644
--- a/cpp/include/raft/neighbors/ivf_list.hpp
+++ b/cpp/include/raft/neighbors/ivf_list.hpp
@@ -35,10 +35,10 @@
 namespace raft::neighbors::ivf {
 
 /** The data for a single IVF list. */
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
-list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res,
-                               const SpecT<SizeT>& spec,
-                               SizeT n_rows)
+template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
+list<SpecT, ValueT, IdxT, SizeT>::list(raft::device_resources const& res,
+                                       const SpecT<SizeT>& spec,
+                                       SizeT n_rows)
   : size{n_rows}
 {
   auto capacity = round_up_safe<SizeT>(n_rows, spec.align_max);
@@ -47,7 +47,7 @@ list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res,
     capacity = std::min<SizeT>(capacity, spec.align_max);
   }
   try {
-    data    = make_device_mdarray<value_type>(res, spec.make_list_extents(capacity));
+    data    = make_device_mdarray<ValueT>(res, spec.make_list_extents(capacity));
     indices = make_device_vector<IdxT, SizeT>(res, capacity);
   } catch (std::bad_alloc& e) {
     RAFT_FAIL(
@@ -67,9 +67,9 @@ list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res,
  * Resize a list by the given id, so that it can contain the given number of records;
  * copy the data if necessary.
  */
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
+template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
 void resize_list(raft::device_resources const& res,
-                 std::shared_ptr<list<SpecT, IdxT, SizeT>>& orig_list,  // NOLINT
+                 std::shared_ptr<list<SpecT, ValueT, IdxT, SizeT>>& orig_list,  // NOLINT
                  const SpecT<SizeT>& spec,
                  SizeT new_used_size,
                  SizeT old_used_size)
@@ -91,12 +91,11 @@ void resize_list(raft::device_resources const& res,
     old_used_size = 0;
   }
   if (skip_resize) { return; }
-  auto new_list = std::make_shared<list<SpecT, IdxT, SizeT>>(res, spec, new_used_size);
+  auto new_list = std::make_shared<list<SpecT, ValueT, IdxT, SizeT>>(res, spec, new_used_size);
   if (old_used_size > 0) {
     auto copied_data_extents = SpecT<size_t>{spec}.make_list_extents(old_used_size);
-    auto copied_view =
-      make_mdspan<typename SpecT<SizeT>::value_type, size_t, row_major, false, true>(
-        new_list->data.data_handle(), copied_data_extents);
+    auto copied_view         = make_mdspan<ValueT, size_t, row_major, false, true>(
+      new_list->data.data_handle(), copied_data_extents);
     copy(copied_view.data_handle(),
          orig_list->data.data_handle(),
          copied_view.size(),
@@ -110,10 +109,10 @@ void resize_list(raft::device_resources const& res,
   new_list.swap(orig_list);
 }
 
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
+template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
 void serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
-                    const list<SpecT, IdxT, SizeT>& ld,
+                    const list<SpecT, ValueT, IdxT, SizeT>& ld,
                     const SpecT<SizeT>& store_spec,
                     std::optional<SizeT> size_override = std::nullopt)
 {
@@ -122,9 +121,8 @@ void serialize_list(const raft::device_resources& handle,
   if (size == 0) { return; }
 
   auto data_extents = store_spec.make_list_extents(size);
-  auto data_array =
-    make_host_mdarray<typename SpecT<SizeT>::value_type, SizeT, row_major>(data_extents);
-  auto inds_array = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+  auto data_array   = make_host_mdarray<ValueT, SizeT, row_major>(data_extents);
+  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
   copy(data_array.data_handle(), ld.data.data_handle(), data_array.size(), handle.get_stream());
   copy(inds_array.data_handle(), ld.indices.data_handle(), inds_array.size(), handle.get_stream());
   handle.sync_stream();
@@ -132,10 +130,10 @@ void serialize_list(const raft::device_resources& handle,
   serialize_mdspan(handle, os, inds_array.view());
 }
 
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
+template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
 void serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
-                    const std::shared_ptr<const list<SpecT, IdxT, SizeT>>& ld,
+                    const std::shared_ptr<const list<SpecT, ValueT, IdxT, SizeT>>& ld,
                     const SpecT<SizeT>& store_spec,
                     std::optional<SizeT> size_override = std::nullopt)
 {
@@ -146,20 +144,19 @@ void serialize_list(const raft::device_resources& handle,
   }
 }
 
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
+template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
 void deserialize_list(const raft::device_resources& handle,
                       std::istream& is,
-                      std::shared_ptr<list<SpecT, IdxT, SizeT>>& ld,
+                      std::shared_ptr<list<SpecT, ValueT, IdxT, SizeT>>& ld,
                       const SpecT<SizeT>& store_spec,
                       const SpecT<SizeT>& device_spec)
 {
   auto size = deserialize_scalar<SizeT>(handle, is);
   if (size == 0) { return ld.reset(); }
-  std::make_shared<list<SpecT, IdxT, SizeT>>(handle, device_spec, size).swap(ld);
+  std::make_shared<list<SpecT, ValueT, IdxT, SizeT>>(handle, device_spec, size).swap(ld);
   auto data_extents = store_spec.make_list_extents(size);
-  auto data_array =
-    make_host_mdarray<typename SpecT<SizeT>::value_type, SizeT, row_major>(data_extents);
-  auto inds_array = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+  auto data_array   = make_host_mdarray<ValueT, SizeT, row_major>(data_extents);
+  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
   deserialize_mdspan(handle, is, data_array.view());
   deserialize_mdspan(handle, is, inds_array.view());
   copy(ld->data.data_handle(), data_array.data_handle(), data_array.size(), handle.get_stream());
diff --git a/cpp/include/raft/neighbors/ivf_list_types.hpp b/cpp/include/raft/neighbors/ivf_list_types.hpp
index aae6a63794..2a12ab04e6 100644
--- a/cpp/include/raft/neighbors/ivf_list_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_list_types.hpp
@@ -36,9 +36,12 @@ constexpr static IdxT kInvalidRecord = (std::is_signed_v<IdxT> ? IdxT{0}
                                        1;
 
 /** The data for a single IVF list. */
-template <template <typename> typename SpecT, typename IdxT, typename SizeT = uint32_t>
+template <template <typename> typename SpecT,
+          typename ValueT,
+          typename IdxT,
+          typename SizeT = uint32_t>
 struct list {
-  using value_type   = typename SpecT<SizeT>::value_type;
+  using value_type   = ValueT;
   using list_extents = typename SpecT<SizeT>::list_extents;
 
   /** Possibly encoded data; it's layout is defined by `SpecT`. */
diff --git a/cpp/include/raft/neighbors/ivf_pq_types.hpp b/cpp/include/raft/neighbors/ivf_pq_types.hpp
index 8c00b08452..79add9ab16 100644
--- a/cpp/include/raft/neighbors/ivf_pq_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_pq_types.hpp
@@ -211,7 +211,7 @@ struct list_spec {
 };
 
 template <typename IdxT, typename SizeT = uint32_t>
-using list_data = ivf::list<list_spec, IdxT, SizeT>;
+using list_data = ivf::list<list_spec, uint8_t, IdxT, SizeT>;
 
 /**
  * @brief IVF-PQ index.
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index c7007cdf5b..192b553bea 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -155,9 +155,8 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
                                   0.001,
                                   min_recall));
       {
-        // new interface
-        raft::spatial::knn::ivf_flat::index_params index_params;
-        raft::spatial::knn::ivf_flat::search_params search_params;
+        ivf_flat::index_params index_params;
+        ivf_flat::search_params search_params;
         index_params.n_lists   = ps.nlist;
         index_params.metric    = ps.metric;
         search_params.n_probes = ps.nprobe;
@@ -201,10 +200,9 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
           indices_ivfflat_dev.data(), ps.num_queries, ps.k);
         auto dists_out_view = raft::make_device_matrix_view<T, IdxT>(
           distances_ivfflat_dev.data(), ps.num_queries, ps.k);
-        raft::spatial::knn::ivf_flat::detail::serialize(handle_, "ivf_flat_index", index);
+        ivf_flat::detail::serialize(handle_, "ivf_flat_index", index);
 
-        auto index_loaded =
-          raft::spatial::knn::ivf_flat::detail::deserialize<DataT, IdxT>(handle_, "ivf_flat_index");
+        auto index_loaded = ivf_flat::detail::deserialize<DataT, IdxT>(handle_, "ivf_flat_index");
 
         ivf_flat::search(handle_,
                          index_loaded,

From 74e0a8c57d0ac2d4310114651b93d48b2489cddb Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Tue, 21 Feb 2023 22:56:21 +0100
Subject: [PATCH 09/33] Update refine

---
 cpp/include/raft/neighbors/detail/ivf_flat_build.cuh     | 9 +++++----
 cpp/include/raft/neighbors/detail/ivf_flat_search.cuh    | 9 ++-------
 cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh | 1 +
 3 files changed, 8 insertions(+), 11 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 9dafb63149..ee4612bf50 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -107,11 +107,11 @@ auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> i
  *     we use source_vecs[source_ixs[i],:]. In both cases i=0..n_rows-1.
  *
  * @param[in] labels device pointer to the cluster ids for each row [n_rows]
- * @param[in] list_offsets device pointer to the cluster offsets in the output (index) [n_lists]
  * @param[in] source_vecs device pointer to the input data [n_rows, dim]
  * @param[in] source_ixs device pointer to the input indices [n_rows]
- * @param[out] list_data device pointer to the output [index_size, dim]
- * @param[out] list_index device pointer to the source ids corr. to the output [index_size]
+ * @param[out] list_data_ptrs device pointer to the index data of size [n_lists][index_size, dim]
+ * @param[out] list_index_ptrs device pointer to the source ids corr. to the output [n_lists]
+ * [index_size]
  * @param[out] list_sizes_ptr device pointer to the cluster sizes [n_lists];
  *                          it's used as an atomic counter, and must be initialized with zeros.
  * @param n_rows source length
@@ -121,7 +121,6 @@ auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> i
  */
 template <typename T, typename IdxT, typename LabelT, bool gather_src = false>
 __global__ void build_index_kernel(const LabelT* labels,
-                                   // const IdxT* list_offsets,
                                    const T* source_vecs,
                                    const IdxT* source_ixs,
                                    T** list_data_ptrs,
@@ -392,6 +391,8 @@ inline void fill_refinement_index(raft::device_resources const& handle,
   for (uint32_t label = 0; label < n_lists; label++) {
     ivf::resize_list(handle, lists(label), list_device_spec, n_candidates, uint32_t(0));
   }
+  // Update the pointers and the sizes
+  refinement_index->recompute_internal_state(handle);
 
   RAFT_CUDA_TRY(cudaMemsetAsync(list_sizes_ptr, 0, n_lists * sizeof(uint32_t), stream));
 
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_search.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_search.cuh
index 0003b02fc5..b2bfd18610 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_search.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_search.cuh
@@ -42,11 +42,6 @@ namespace raft::neighbors::ivf_flat::detail {
 
 using namespace raft::spatial::knn::detail;  // NOLINT
 
-using raft::neighbors::ivf_flat::index;
-using raft::neighbors::ivf_flat::index_params;
-using raft::neighbors::ivf_flat::kIndexGroupSize;
-using raft::neighbors::ivf_flat::search_params;
-
 constexpr int kThreadsPerBlock = 128;
 
 /**
@@ -677,8 +672,8 @@ __global__ void __launch_bounds__(kThreadsPerBlock)
                           const uint32_t query_smem_elems,
                           const T* query,
                           const uint32_t* coarse_index,
-                          const IdxT* const* list_indices_ptrs,  // const IdxT* list_indices
-                          const T* const* list_data_ptrs,        // const T* list_data
+                          const IdxT* const* list_indices_ptrs,
+                          const T* const* list_data_ptrs,
                           const uint32_t* list_sizes,
                           const uint32_t n_probes,
                           const uint32_t k,
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 405bc3252f..3fa99fe5cb 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -95,6 +95,7 @@ void serialize(raft::device_resources const& handle,
     ivf::serialize_list<list_spec, T, IdxT, uint32_t>(
       handle, of, index_.lists()(label), list_store_spec, sizes_host(label));
   }
+  handle.sync_stream();
   of.close();
   if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); }
 }

From 135a9b683f3d733397721c2ae2526bb4c9728768 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Fri, 24 Feb 2023 09:56:13 +0100
Subject: [PATCH 10/33] Use std vector for ivf flat index list

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  | 29 ++++++-------------
 .../neighbors/detail/ivf_flat_serialize.cuh   | 10 +++----
 cpp/include/raft/neighbors/detail/refine.cuh  |  2 +-
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 29 ++++++++-----------
 cpp/test/neighbors/ann_ivf_flat.cuh           |  1 +
 5 files changed, 27 insertions(+), 44 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index ee4612bf50..c1ed796899 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -40,12 +40,6 @@ namespace raft::neighbors::ivf_flat::detail {
 
 using namespace raft::spatial::knn::detail;  // NOLINT
 
-using raft::neighbors::ivf_flat::index;
-using raft::neighbors::ivf_flat::index_params;
-using raft::neighbors::ivf_flat::kIndexGroupSize;
-using raft::neighbors::ivf_flat::list_data;
-using raft::neighbors::ivf_flat::search_params;
-
 template <typename T, typename IdxT>
 auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> index<T, IdxT>
 {
@@ -69,18 +63,12 @@ auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> i
        source.centers().size(),
        stream);
   if (source.center_norms().has_value())
-    copy(target.center_norms().value.data_handle(),
-         source.center_norms().value.data_handle(),
-         source.center_norms().value.size(),
+    copy(target.center_norms().value().data_handle(),
+         source.center_norms().value().data_handle(),
+         source.center_norms().value().size(),
          stream);
   // Copy shared pointers
-  {
-    auto source_lists = source.lists();
-    auto target_lists = target.lists();
-    for (uint32_t label = 0; label < source_lists.size(); label++) {
-      target_lists(label) = source_lists(label);
-    }
-  }
+  target.lists() = source.lists();
 
   // Make sure the device pointers point to the new lists
   target.recompute_internal_state(res);
@@ -178,6 +166,8 @@ void extend(raft::device_resources const& handle,
   auto stream  = handle.get_stream();
   auto n_lists = index->n_lists();
   auto dim     = index->dim();
+  auto list_device_spec =
+    list_spec<uint32_t>{index->dim(), index->conservative_memory_allocation()};
   common::nvtx::range<common::nvtx::domain::raft> fun_scope(
     "ivf_flat::extend(%zu, %u)", size_t(n_rows), dim);
 
@@ -235,11 +225,10 @@ void extend(raft::device_resources const& handle,
     copy(old_list_sizes.data(), old_list_sizes_dev.data_handle(), n_lists, stream);
     copy(new_list_sizes.data(), list_sizes_ptr, n_lists, stream);
     handle.sync_stream();
-    auto lists            = index->lists();
-    auto list_device_spec = list_spec<uint32_t>{index->dim(), false};
+    auto& lists = index->lists();
     for (uint32_t label = 0; label < n_lists; label++) {
       ivf::resize_list(
-        handle, lists(label), list_device_spec, new_list_sizes[label], old_list_sizes[label]);
+        handle, lists[label], list_device_spec, new_list_sizes[label], old_list_sizes[label]);
     }
   }
   // Update the pointers and the sizes
@@ -389,7 +378,7 @@ inline void fill_refinement_index(raft::device_resources const& handle,
   auto lists            = refinement_index->lists();
   auto list_device_spec = list_spec<uint32_t>{refinement_index->dim(), false};
   for (uint32_t label = 0; label < n_lists; label++) {
-    ivf::resize_list(handle, lists(label), list_device_spec, n_candidates, uint32_t(0));
+    ivf::resize_list(handle, lists[label], list_device_spec, n_candidates, uint32_t(0));
   }
   // Update the pointers and the sizes
   refinement_index->recompute_internal_state(handle);
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 3fa99fe5cb..85688ae561 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -42,7 +42,7 @@ struct check_index_layout {
                 "paste in the new size and consider updating the serialization logic");
 };
 
-template struct check_index_layout<sizeof(index<double, std::uint64_t>), 376>;
+template struct check_index_layout<sizeof(index<double, std::uint64_t>), 368>;
 
 /**
  * Save the index to file.
@@ -70,7 +70,6 @@ void serialize(raft::device_resources const& handle,
   serialize_scalar(handle, of, index_.dim());
   serialize_scalar(handle, of, index_.n_lists());
   serialize_scalar(handle, of, index_.metric());
-  serialize_scalar(handle, of, index_.veclen());
   serialize_scalar(handle, of, index_.adaptive_centers());
   serialize_scalar(handle, of, index_.conservative_memory_allocation());
   serialize_mdspan(handle, of, index_.centers());
@@ -93,7 +92,7 @@ void serialize(raft::device_resources const& handle,
   auto list_store_spec = list_spec<uint32_t>{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
     ivf::serialize_list<list_spec, T, IdxT, uint32_t>(
-      handle, of, index_.lists()(label), list_store_spec, sizes_host(label));
+      handle, of, index_.lists()[label], list_store_spec, sizes_host(label));
   }
   handle.sync_stream();
   of.close();
@@ -125,7 +124,6 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   auto dim              = deserialize_scalar<std::uint32_t>(handle, infile);
   auto n_lists          = deserialize_scalar<std::uint32_t>(handle, infile);
   auto metric           = deserialize_scalar<raft::distance::DistanceType>(handle, infile);
-  auto veclen           = deserialize_scalar<std::uint32_t>(handle, infile);
   bool adaptive_centers = deserialize_scalar<bool>(handle, infile);
   bool cma              = deserialize_scalar<bool>(handle, infile);
 
@@ -137,7 +135,7 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
     if (!index_.center_norms()) {
       RAFT_FAIL("Error inconsistent center norms");
     } else {
-      auto center_norms = *index_.center_norms();
+      auto center_norms = index_.center_norms().value();
       deserialize_mdspan(handle, infile, center_norms);
     }
   }
@@ -147,7 +145,7 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   auto list_store_spec  = list_spec<uint32_t>{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
     ivf::deserialize_list<list_spec, T, IdxT, uint32_t>(
-      handle, infile, index_.lists()(label), list_store_spec, list_device_spec);
+      handle, infile, index_.lists()[label], list_store_spec, list_device_spec);
   }
   handle.sync_stream();
   infile.close();
diff --git a/cpp/include/raft/neighbors/detail/refine.cuh b/cpp/include/raft/neighbors/detail/refine.cuh
index 5d226132a4..b0aebe28b6 100644
--- a/cpp/include/raft/neighbors/detail/refine.cuh
+++ b/cpp/include/raft/neighbors/detail/refine.cuh
@@ -108,7 +108,7 @@ void refine_device(raft::device_resources const& handle,
     handle.get_thrust_policy(), fake_coarse_idx.data(), fake_coarse_idx.data() + n_queries);
 
   raft::neighbors::ivf_flat::index<data_t, idx_t> refinement_index(
-    handle, metric, n_queries, false, false, dim);
+    handle, metric, n_queries, false, true, dim);
 
   raft::neighbors::ivf_flat::detail::fill_refinement_index(handle,
                                                            &refinement_index,
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index b60d5b3aa7..4c79fe8cf7 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -219,10 +219,7 @@ struct index : ann::index {
     return centers_.extent(1);
   }
   /** Number of clusters/inverted lists. */
-  [[nodiscard]] constexpr inline auto n_lists() const noexcept -> uint32_t
-  {
-    return lists_.extent(0);
-  }
+  [[nodiscard]] constexpr inline auto n_lists() const noexcept -> uint32_t { return lists_.size(); }
 
   // Don't allow copying the index for performance reasons (try avoiding copying data)
   index(const index&) = delete;
@@ -245,15 +242,12 @@ struct index : ann::index {
       conservative_memory_allocation_{conservative_memory_allocation},
       centers_(make_device_matrix<float, uint32_t>(res, n_lists, dim)),
       center_norms_(std::nullopt),
-      lists_{make_host_vector<std::shared_ptr<list_data<T, IdxT>>, uint32_t>(n_lists)},
+      lists_{n_lists},
       list_sizes_{make_device_vector<uint32_t, uint32_t>(res, n_lists)},
       data_ptrs_{make_device_vector<T*, uint32_t>(res, n_lists)},
       inds_ptrs_{make_device_vector<IdxT*, uint32_t>(res, n_lists)},
       total_size_{0}
   {
-    for (uint32_t i = 0; i < n_lists; i++) {
-      lists_(i) = std::shared_ptr<list_data<T, IdxT>>();
-    }
     switch (metric_) {
       case raft::distance::DistanceType::L2Expanded:
       case raft::distance::DistanceType::L2SqrtExpanded:
@@ -315,9 +309,10 @@ struct index : ann::index {
     auto this_inds_ptrs       = inds_ptrs();
     IdxT recompute_total_size = 0;
     for (uint32_t label = 0; label < this_lists.size(); label++) {
-      const auto data_ptr  = this_lists(label) ? this_lists(label)->data.data_handle() : nullptr;
-      const auto inds_ptr  = this_lists(label) ? this_lists(label)->indices.data_handle() : nullptr;
-      const auto list_size = this_lists(label) ? IdxT(this_lists(label)->size) : 0;
+      auto& list           = this_lists[label];
+      const auto data_ptr  = list ? list->data.data_handle() : nullptr;
+      const auto inds_ptr  = list ? list->indices.data_handle() : nullptr;
+      const auto list_size = list ? IdxT(list->size) : 0;
       copy(&this_data_ptrs(label), &data_ptr, 1, stream);
       copy(&this_inds_ptrs(label), &inds_ptr, 1, stream);
       recompute_total_size += list_size;
@@ -327,14 +322,14 @@ struct index : ann::index {
   }
 
   /** Lists' data and indices. */
-  inline auto lists() noexcept -> host_vector_view<std::shared_ptr<list_data<T, IdxT>>, uint32_t>
+  inline auto lists() noexcept -> std::vector<std::shared_ptr<list_data<T, IdxT>>>&
   {
-    return lists_.view();
+    return lists_;
   }
   [[nodiscard]] inline auto lists() const noexcept
-    -> host_vector_view<const std::shared_ptr<list_data<T, IdxT>>, uint32_t>
+    -> const std::vector<std::shared_ptr<list_data<T, IdxT>>>&
   {
-    return lists_.view();
+    return lists_;
   }
 
  private:
@@ -346,7 +341,7 @@ struct index : ann::index {
   raft::distance::DistanceType metric_;
   bool adaptive_centers_;
   bool conservative_memory_allocation_;
-  host_vector<std::shared_ptr<list_data<T, IdxT>>, uint32_t> lists_;
+  std::vector<std::shared_ptr<list_data<T, IdxT>>> lists_;
   device_vector<uint32_t, uint32_t> list_sizes_;
   device_matrix<float, uint32_t, row_major> centers_;
   std::optional<device_vector<float, uint32_t>> center_norms_;
@@ -359,7 +354,7 @@ struct index : ann::index {
   /** Throw an error if the index content is inconsistent. */
   void check_consistency()
   {
-    auto n_lists = lists_.extent(0);
+    auto n_lists = lists_.size();
     RAFT_EXPECTS(dim() % veclen_ == 0, "dimensionality is not a multiple of the veclen");
     RAFT_EXPECTS(list_sizes_.extent(0) == n_lists, "inconsistent list size");
     RAFT_EXPECTS(data_ptrs_.extent(0) == n_lists, "inconsistent list size");
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index 192b553bea..a616e92681 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -163,6 +163,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
 
         index_params.add_data_on_build        = false;
         index_params.kmeans_trainset_fraction = 0.5;
+        index_params.metric_arg               = 0;
 
         auto database_view = raft::make_device_matrix_view<const DataT, IdxT>(
           (const DataT*)database.data(), ps.num_db_vecs, ps.dim);

From 81b2cbfd63110da1ccd1248c21d5fa1dd77272a6 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Fri, 24 Feb 2023 10:07:29 +0100
Subject: [PATCH 11/33] Test second variant of `ivf_flat::extend`

---
 cpp/include/raft/neighbors/detail/ivf_flat_build.cuh | 4 ++--
 cpp/test/neighbors/ann_ivf_flat.cuh                  | 4 ++--
 2 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index c1ed796899..35ed31932f 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -273,8 +273,8 @@ auto extend(raft::device_resources const& handle,
             const IdxT* new_indices,
             IdxT n_rows) -> index<T, IdxT>
 {
-  auto ext_index = clone(handle, &orig_index);
-  extend(handle, &ext_index, new_vectors, new_indices, n_rows);
+  auto ext_index = clone(handle, orig_index);
+  detail::extend(handle, &ext_index, new_vectors, new_indices, n_rows);
   return ext_index;
 }
 
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index a616e92681..7c31b59cd2 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -181,7 +181,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         auto half_of_data_view = raft::make_device_matrix_view<const DataT, IdxT>(
           (const DataT*)database.data(), half_of_data, ps.dim);
 
-        ivf_flat::extend(handle_, &index, half_of_data_view);
+        auto index_2 = ivf_flat::extend(handle_, index, half_of_data_view);
 
         auto new_half_of_data_view = raft::make_device_matrix_view<const DataT, IdxT>(
           database.data() + half_of_data * ps.dim, IdxT(ps.num_db_vecs) - half_of_data, ps.dim);
@@ -201,7 +201,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
           indices_ivfflat_dev.data(), ps.num_queries, ps.k);
         auto dists_out_view = raft::make_device_matrix_view<T, IdxT>(
           distances_ivfflat_dev.data(), ps.num_queries, ps.k);
-        ivf_flat::detail::serialize(handle_, "ivf_flat_index", index);
+        ivf_flat::detail::serialize(handle_, "ivf_flat_index", index_2);
 
         auto index_loaded = ivf_flat::detail::deserialize<DataT, IdxT>(handle_, "ivf_flat_index");
 

From fd33dbcb3caa7d0d4ce80a72a821916794784e2a Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Sat, 25 Feb 2023 20:46:14 +0100
Subject: [PATCH 12/33] Use ValueT template on spec

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  |  4 +-
 .../neighbors/detail/ivf_flat_serialize.cuh   | 10 ++--
 .../raft/neighbors/detail/ivf_pq_build.cuh    |  5 +-
 .../neighbors/detail/ivf_pq_serialize.cuh     | 10 ++--
 cpp/include/raft/neighbors/ivf_flat_types.hpp |  5 +-
 cpp/include/raft/neighbors/ivf_list.hpp       | 48 ++++++++++---------
 cpp/include/raft/neighbors/ivf_list_types.hpp | 11 ++---
 cpp/include/raft/neighbors/ivf_pq_types.hpp   | 10 ++--
 8 files changed, 52 insertions(+), 51 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 35ed31932f..f5baffb0d3 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -167,7 +167,7 @@ void extend(raft::device_resources const& handle,
   auto n_lists = index->n_lists();
   auto dim     = index->dim();
   auto list_device_spec =
-    list_spec<uint32_t>{index->dim(), index->conservative_memory_allocation()};
+    list_spec<T, uint32_t>{index->dim(), index->conservative_memory_allocation()};
   common::nvtx::range<common::nvtx::domain::raft> fun_scope(
     "ivf_flat::extend(%zu, %u)", size_t(n_rows), dim);
 
@@ -376,7 +376,7 @@ inline void fill_refinement_index(raft::device_resources const& handle,
 
   // Allocate new memory
   auto lists            = refinement_index->lists();
-  auto list_device_spec = list_spec<uint32_t>{refinement_index->dim(), false};
+  auto list_device_spec = list_spec<T, uint32_t>{refinement_index->dim(), false};
   for (uint32_t label = 0; label < n_lists; label++) {
     ivf::resize_list(handle, lists[label], list_device_spec, n_candidates, uint32_t(0));
   }
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 85688ae561..b7e6955616 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -89,9 +89,9 @@ void serialize(raft::device_resources const& handle,
   handle.sync_stream();
   serialize_mdspan(handle, of, sizes_host.view());
 
-  auto list_store_spec = list_spec<uint32_t>{index_.dim(), true};
+  auto list_store_spec = list_spec<T, uint32_t>{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::serialize_list<list_spec, T, IdxT, uint32_t>(
+    ivf::serialize_list<list_spec<T, uint32_t>, IdxT, uint32_t>(
       handle, of, index_.lists()[label], list_store_spec, sizes_host(label));
   }
   handle.sync_stream();
@@ -141,10 +141,10 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   }
   deserialize_mdspan(handle, infile, index_.list_sizes());
 
-  auto list_device_spec = list_spec<uint32_t>{index_.dim(), cma};
-  auto list_store_spec  = list_spec<uint32_t>{index_.dim(), true};
+  auto list_device_spec = list_spec<T, uint32_t>{index_.dim(), cma};
+  auto list_store_spec  = list_spec<T, uint32_t>{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::deserialize_list<list_spec, T, IdxT, uint32_t>(
+    ivf::deserialize_list<list_spec<T, uint32_t>, IdxT, uint32_t>(
       handle, infile, index_.lists()[label], list_store_spec, list_device_spec);
   }
   handle.sync_stream();
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
index bf3014568a..6ba43f4c3e 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
@@ -664,7 +664,8 @@ __launch_bounds__(BlockSize) __global__ void process_and_fill_codes_kernel(
   const uint32_t pq_len     = pq_centers.extent(1);
   const uint32_t pq_dim     = new_vectors.extent(1) / pq_len;
 
-  auto pq_extents = list_spec<uint32_t>{PqBits, pq_dim, true}.make_list_extents(out_ix + 1);
+  auto pq_extents =
+    list_spec<uint8_t, uint32_t>{PqBits, pq_dim, true}.make_list_extents(out_ix + 1);
   auto pq_extents_vectorized =
     make_extents<uint32_t>(pq_extents.extent(0), pq_extents.extent(1), pq_extents.extent(2));
   auto pq_dataset = make_mdspan<pq_vec_t, uint32_t, row_major, false, true>(
@@ -906,7 +907,7 @@ void extend(raft::device_resources const& handle,
   std::optional<list_data<IdxT, size_t>> placeholder_list(
     std::in_place_t{},
     handle,
-    list_spec<size_t>(spec),
+    list_spec<uint8_t, size_t>(spec),
     n_rows + (kIndexGroupSize - 1) * std::min<IdxT>(n_clusters, n_rows));
 
   // Available device memory
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
index bab9a52ed1..1b90a21d54 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
@@ -94,9 +94,9 @@ void serialize(raft::device_resources const& handle_,
        handle_.get_stream());
   handle_.sync_stream();
   serialize_mdspan(handle_, of, sizes_host.view());
-  auto list_store_spec = list_spec<uint32_t>{index.pq_bits(), index.pq_dim(), true};
+  auto list_store_spec = list_spec<uint8_t, uint32_t>{index.pq_bits(), index.pq_dim(), true};
   for (uint32_t label = 0; label < index.n_lists(); label++) {
-    ivf::serialize_list<list_spec, uint8_t, IdxT, uint32_t>(
+    ivf::serialize_list<list_spec<uint8_t, uint32_t>, IdxT, uint32_t>(
       handle_, of, index.lists()[label], list_store_spec, sizes_host(label));
   }
 
@@ -151,10 +151,10 @@ auto deserialize(raft::device_resources const& handle_, const std::string& filen
   deserialize_mdspan(handle_, infile, index.centers_rot());
   deserialize_mdspan(handle_, infile, index.rotation_matrix());
   deserialize_mdspan(handle_, infile, index.list_sizes());
-  auto list_device_spec = list_spec<uint32_t>{pq_bits, pq_dim, cma};
-  auto list_store_spec  = list_spec<uint32_t>{pq_bits, pq_dim, true};
+  auto list_device_spec = list_spec<uint8_t, uint32_t>{pq_bits, pq_dim, cma};
+  auto list_store_spec  = list_spec<uint8_t, uint32_t>{pq_bits, pq_dim, true};
   for (auto& list : index.lists()) {
-    ivf::deserialize_list<list_spec, uint8_t, IdxT, uint32_t>(
+    ivf::deserialize_list<list_spec<uint8_t, uint32_t>, IdxT, uint32_t>(
       handle_, infile, list, list_store_spec, list_device_spec);
   }
 
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index 4c79fe8cf7..9faa794f6d 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -80,8 +80,9 @@ struct search_params : ann::search_params {
 static_assert(std::is_aggregate_v<index_params>);
 static_assert(std::is_aggregate_v<search_params>);
 
-template <typename SizeT = uint32_t>
+template <typename ValueT, typename SizeT = uint32_t>
 struct list_spec {
+  using value_type   = ValueT;
   using list_extents = matrix_extent<SizeT>;
 
   SizeT align_max;
@@ -110,7 +111,7 @@ struct list_spec {
 };
 
 template <typename ValueT, typename IdxT, typename SizeT = uint32_t>
-using list_data = ivf::list<list_spec, ValueT, IdxT, SizeT>;
+using list_data = ivf::list<list_spec<ValueT, SizeT>, IdxT, SizeT>;
 
 /**
  * @brief IVF-flat index.
diff --git a/cpp/include/raft/neighbors/ivf_list.hpp b/cpp/include/raft/neighbors/ivf_list.hpp
index 2fe3d2cf3d..93c855139c 100644
--- a/cpp/include/raft/neighbors/ivf_list.hpp
+++ b/cpp/include/raft/neighbors/ivf_list.hpp
@@ -35,10 +35,8 @@
 namespace raft::neighbors::ivf {
 
 /** The data for a single IVF list. */
-template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
-list<SpecT, ValueT, IdxT, SizeT>::list(raft::device_resources const& res,
-                                       const SpecT<SizeT>& spec,
-                                       SizeT n_rows)
+template <typename SpecT, typename IdxT, typename SizeT>
+list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res, const SpecT& spec, SizeT n_rows)
   : size{n_rows}
 {
   auto capacity = round_up_safe<SizeT>(n_rows, spec.align_max);
@@ -47,7 +45,7 @@ list<SpecT, ValueT, IdxT, SizeT>::list(raft::device_resources const& res,
     capacity = std::min<SizeT>(capacity, spec.align_max);
   }
   try {
-    data    = make_device_mdarray<ValueT>(res, spec.make_list_extents(capacity));
+    data = make_device_mdarray<typename SpecT::value_type>(res, spec.make_list_extents(capacity));
     indices = make_device_vector<IdxT, SizeT>(res, capacity);
   } catch (std::bad_alloc& e) {
     RAFT_FAIL(
@@ -67,10 +65,13 @@ list<SpecT, ValueT, IdxT, SizeT>::list(raft::device_resources const& res,
  * Resize a list by the given id, so that it can contain the given number of records;
  * copy the data if necessary.
  */
-template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
+template <template <typename, typename> typename SpecT,
+          typename ValueT,
+          typename IdxT,
+          typename SizeT>
 void resize_list(raft::device_resources const& res,
-                 std::shared_ptr<list<SpecT, ValueT, IdxT, SizeT>>& orig_list,  // NOLINT
-                 const SpecT<SizeT>& spec,
+                 std::shared_ptr<list<SpecT<ValueT, SizeT>, IdxT, SizeT>>& orig_list,  // NOLINT
+                 const SpecT<ValueT, SizeT>& spec,
                  SizeT new_used_size,
                  SizeT old_used_size)
 {
@@ -91,9 +92,10 @@ void resize_list(raft::device_resources const& res,
     old_used_size = 0;
   }
   if (skip_resize) { return; }
-  auto new_list = std::make_shared<list<SpecT, ValueT, IdxT, SizeT>>(res, spec, new_used_size);
+  auto new_list =
+    std::make_shared<list<SpecT<ValueT, SizeT>, IdxT, SizeT>>(res, spec, new_used_size);
   if (old_used_size > 0) {
-    auto copied_data_extents = SpecT<size_t>{spec}.make_list_extents(old_used_size);
+    auto copied_data_extents = SpecT<ValueT, size_t>{spec}.make_list_extents(old_used_size);
     auto copied_view         = make_mdspan<ValueT, size_t, row_major, false, true>(
       new_list->data.data_handle(), copied_data_extents);
     copy(copied_view.data_handle(),
@@ -109,11 +111,11 @@ void resize_list(raft::device_resources const& res,
   new_list.swap(orig_list);
 }
 
-template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
+template <typename SpecT, typename IdxT, typename SizeT>
 void serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
-                    const list<SpecT, ValueT, IdxT, SizeT>& ld,
-                    const SpecT<SizeT>& store_spec,
+                    const list<SpecT, IdxT, SizeT>& ld,
+                    const SpecT& store_spec,
                     std::optional<SizeT> size_override = std::nullopt)
 {
   auto size = size_override.value_or(ld.size.load());
@@ -121,7 +123,7 @@ void serialize_list(const raft::device_resources& handle,
   if (size == 0) { return; }
 
   auto data_extents = store_spec.make_list_extents(size);
-  auto data_array   = make_host_mdarray<ValueT, SizeT, row_major>(data_extents);
+  auto data_array   = make_host_mdarray<typename SpecT::value_type, SizeT, row_major>(data_extents);
   auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
   copy(data_array.data_handle(), ld.data.data_handle(), data_array.size(), handle.get_stream());
   copy(inds_array.data_handle(), ld.indices.data_handle(), inds_array.size(), handle.get_stream());
@@ -130,11 +132,11 @@ void serialize_list(const raft::device_resources& handle,
   serialize_mdspan(handle, os, inds_array.view());
 }
 
-template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
+template <typename SpecT, typename IdxT, typename SizeT>
 void serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
-                    const std::shared_ptr<const list<SpecT, ValueT, IdxT, SizeT>>& ld,
-                    const SpecT<SizeT>& store_spec,
+                    const std::shared_ptr<const list<SpecT, IdxT, SizeT>>& ld,
+                    const SpecT& store_spec,
                     std::optional<SizeT> size_override = std::nullopt)
 {
   if (ld) {
@@ -144,18 +146,18 @@ void serialize_list(const raft::device_resources& handle,
   }
 }
 
-template <template <typename> typename SpecT, typename ValueT, typename IdxT, typename SizeT>
+template <typename SpecT, typename IdxT, typename SizeT>
 void deserialize_list(const raft::device_resources& handle,
                       std::istream& is,
-                      std::shared_ptr<list<SpecT, ValueT, IdxT, SizeT>>& ld,
-                      const SpecT<SizeT>& store_spec,
-                      const SpecT<SizeT>& device_spec)
+                      std::shared_ptr<list<SpecT, IdxT, SizeT>>& ld,
+                      const SpecT& store_spec,
+                      const SpecT& device_spec)
 {
   auto size = deserialize_scalar<SizeT>(handle, is);
   if (size == 0) { return ld.reset(); }
-  std::make_shared<list<SpecT, ValueT, IdxT, SizeT>>(handle, device_spec, size).swap(ld);
+  std::make_shared<list<SpecT, IdxT, SizeT>>(handle, device_spec, size).swap(ld);
   auto data_extents = store_spec.make_list_extents(size);
-  auto data_array   = make_host_mdarray<ValueT, SizeT, row_major>(data_extents);
+  auto data_array   = make_host_mdarray<typename SpecT::value_type, SizeT, row_major>(data_extents);
   auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
   deserialize_mdspan(handle, is, data_array.view());
   deserialize_mdspan(handle, is, inds_array.view());
diff --git a/cpp/include/raft/neighbors/ivf_list_types.hpp b/cpp/include/raft/neighbors/ivf_list_types.hpp
index 2a12ab04e6..cc191b12ea 100644
--- a/cpp/include/raft/neighbors/ivf_list_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_list_types.hpp
@@ -36,13 +36,10 @@ constexpr static IdxT kInvalidRecord = (std::is_signed_v<IdxT> ? IdxT{0}
                                        1;
 
 /** The data for a single IVF list. */
-template <template <typename> typename SpecT,
-          typename ValueT,
-          typename IdxT,
-          typename SizeT = uint32_t>
+template <typename SpecT, typename IdxT, typename SizeT = uint32_t>
 struct list {
-  using value_type   = ValueT;
-  using list_extents = typename SpecT<SizeT>::list_extents;
+  using value_type   = typename SpecT::value_type;
+  using list_extents = typename SpecT::list_extents;
 
   /** Possibly encoded data; it's layout is defined by `SpecT`. */
   device_mdarray<value_type, list_extents, row_major> data;
@@ -52,7 +49,7 @@ struct list {
   std::atomic<SizeT> size;
 
   /** Allocate a new list capable of holding at least `n_rows` data records and indices. */
-  list(raft::device_resources const& res, const SpecT<SizeT>& spec, SizeT n_rows);
+  list(raft::device_resources const& res, const SpecT& spec, SizeT n_rows);
 };
 
 }  // namespace raft::neighbors::ivf
diff --git a/cpp/include/raft/neighbors/ivf_pq_types.hpp b/cpp/include/raft/neighbors/ivf_pq_types.hpp
index 79add9ab16..df90a25de8 100644
--- a/cpp/include/raft/neighbors/ivf_pq_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_pq_types.hpp
@@ -161,9 +161,9 @@ constexpr static uint32_t kIndexGroupVecLen = 16;
 template <typename IdxT>
 constexpr static IdxT kOutOfBoundsRecord = std::numeric_limits<IdxT>::max();
 
-template <typename SizeT = uint32_t>
+template <typename ValueT = uint8_t, typename SizeT = uint32_t>
 struct list_spec {
-  using value_type = uint8_t;
+  using value_type = ValueT;
   /** PQ-encoded data stored in the interleaved format:
    *
    *    [ ceildiv(list_size, kIndexGroupSize)
@@ -189,8 +189,8 @@ struct list_spec {
   }
 
   // Allow casting between different size-types (for safer size and offset calculations)
-  template <typename OtherSizeT>
-  constexpr explicit list_spec(const list_spec<OtherSizeT>& other_spec)
+  template <typename OtherValueT, typename OtherSizeT>
+  constexpr explicit list_spec(const list_spec<OtherValueT, OtherSizeT>& other_spec)
     : pq_bits{other_spec.pq_bits},
       pq_dim{other_spec.pq_dim},
       align_min{other_spec.align_min},
@@ -211,7 +211,7 @@ struct list_spec {
 };
 
 template <typename IdxT, typename SizeT = uint32_t>
-using list_data = ivf::list<list_spec, uint8_t, IdxT, SizeT>;
+using list_data = ivf::list<list_spec<uint8_t, SizeT>, IdxT, SizeT>;
 
 /**
  * @brief IVF-PQ index.

From 539fbc567963eb925638509a4f30f51150ec6048 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Sun, 26 Feb 2023 13:00:05 +0100
Subject: [PATCH 13/33] Use second variant of ivf_extend

---
 cpp/test/neighbors/ann_ivf_flat.cuh | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index 7c31b59cd2..b8b86e26b4 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -190,7 +190,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
           vector_indices.data() + half_of_data, IdxT(ps.num_db_vecs) - half_of_data);
 
         ivf_flat::extend(handle_,
-                         &index,
+                         &index_2,
                          new_half_of_data_view,
                          std::make_optional<raft::device_vector_view<const IdxT, IdxT>>(
                            new_half_of_data_indices_view));
@@ -218,25 +218,26 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         handle_.sync_stream(stream_);
 
         // Test the centroid invariants
-        if (index.adaptive_centers()) {
+        if (index_2.adaptive_centers()) {
           // The centers must be up-to-date with the corresponding data
-          std::vector<uint32_t> list_sizes(index.n_lists());
+          std::vector<uint32_t> list_sizes(index_2.n_lists());
           rmm::device_uvector<float> centroid(ps.dim, stream_);
-          raft::copy(list_sizes.data(), index.list_sizes().data_handle(), index.n_lists(), stream_);
+          raft::copy(
+            list_sizes.data(), index_2.list_sizes().data_handle(), index_2.n_lists(), stream_);
           handle_.sync_stream(stream_);
-          for (uint32_t l = 0; l < index.n_lists(); l++) {
+          for (uint32_t l = 0; l < index_2.n_lists(); l++) {
             rmm::device_uvector<float> cluster_data(list_sizes[l] * ps.dim, stream_);
             raft::spatial::knn::detail::utils::copy_selected<float>((IdxT)list_sizes[l],
                                                                     (IdxT)ps.dim,
                                                                     database.data(),
-                                                                    index.inds_ptrs()(l),
+                                                                    index_2.inds_ptrs()(l),
                                                                     (IdxT)ps.dim,
                                                                     cluster_data.data(),
                                                                     (IdxT)ps.dim,
                                                                     stream_);
             raft::stats::mean<float, uint32_t>(
               centroid.data(), cluster_data.data(), ps.dim, list_sizes[l], false, true, stream_);
-            ASSERT_TRUE(raft::devArrMatch(index.centers().data_handle() + ps.dim * l,
+            ASSERT_TRUE(raft::devArrMatch(index_2.centers().data_handle() + ps.dim * l,
                                           centroid.data(),
                                           ps.dim,
                                           raft::CompareApprox<float>(0.001),
@@ -244,9 +245,9 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
           }
         } else {
           // The centers must be immutable
-          ASSERT_TRUE(raft::devArrMatch(index.centers().data_handle(),
+          ASSERT_TRUE(raft::devArrMatch(index_2.centers().data_handle(),
                                         index.centers().data_handle(),
-                                        index.centers().size(),
+                                        index_2.centers().size(),
                                         raft::Compare<float>(),
                                         stream_));
         }

From 33bfb827e5c191521e0fc8fe0652cf1eba40adf5 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Tue, 28 Feb 2023 02:17:16 +0100
Subject: [PATCH 14/33] Fix spec template

---
 ci/docs/build.sh                              | 2 +-
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 9 ++++++---
 2 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/ci/docs/build.sh b/ci/docs/build.sh
index 7ac7059ef2..e3062107c0 100644
--- a/ci/docs/build.sh
+++ b/ci/docs/build.sh
@@ -1,5 +1,5 @@
 #!/bin/bash
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 #################################
 # RAFT docs build script for CI #
 #################################
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index 9faa794f6d..e5751ff90a 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -97,8 +97,8 @@ struct list_spec {
   }
 
   // Allow casting between different size-types (for safer size and offset calculations)
-  template <typename OtherSizeT>
-  constexpr explicit list_spec(const list_spec<OtherSizeT>& other_spec)
+  template <typename OtherValueT, typename OtherSizeT>
+  constexpr explicit list_spec(const list_spec<OtherValueT, OtherSizeT>& other_spec)
     : dim{other_spec.dim}, align_min{other_spec.align_min}, align_max{other_spec.align_max}
   {
   }
@@ -166,7 +166,10 @@ struct index : ann::index {
    *     x[16, 4], x[16, 5], x[17, 4], x[17, 5], ... x[30, 4], x[30, 5],    -    ,    -    ,
    *
    */
-  /** Sizes of the lists (clusters) [n_lists] */
+  /** Sizes of the lists (clusters) [n_lists]
+   * NB: This may differ from the actual list size if the shared lists have been extended by another
+   * index
+   */
   inline auto list_sizes() noexcept -> device_vector_view<uint32_t, uint32_t>
   {
     return list_sizes_.view();

From adb96e407d0894c7c6d8f14d8d25b430d006a4b0 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Mon, 6 Mar 2023 20:58:25 +0100
Subject: [PATCH 15/33] Revert ValueT on ivfpq

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  |  9 +--
 .../neighbors/detail/ivf_flat_serialize.cuh   | 12 ++--
 .../raft/neighbors/detail/ivf_pq_build.cuh    |  5 +-
 .../neighbors/detail/ivf_pq_serialize.cuh     | 10 ++--
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 57 ++++++++++---------
 cpp/include/raft/neighbors/ivf_list.hpp       | 54 +++++++++---------
 cpp/include/raft/neighbors/ivf_list_types.hpp |  8 +--
 cpp/include/raft/neighbors/ivf_pq_types.hpp   | 10 ++--
 8 files changed, 84 insertions(+), 81 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index f5baffb0d3..5bdcce6053 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -166,8 +166,8 @@ void extend(raft::device_resources const& handle,
   auto stream  = handle.get_stream();
   auto n_lists = index->n_lists();
   auto dim     = index->dim();
-  auto list_device_spec =
-    list_spec<T, uint32_t>{index->dim(), index->conservative_memory_allocation()};
+  typename list_spec_wrapper<T>::list_spec<uint32_t> list_device_spec{
+    index->dim(), index->conservative_memory_allocation()};
   common::nvtx::range<common::nvtx::domain::raft> fun_scope(
     "ivf_flat::extend(%zu, %u)", size_t(n_rows), dim);
 
@@ -375,8 +375,9 @@ inline void fill_refinement_index(raft::device_resources const& handle,
   // We do not fill centers and center norms, since we will not run coarse search.
 
   // Allocate new memory
-  auto lists            = refinement_index->lists();
-  auto list_device_spec = list_spec<T, uint32_t>{refinement_index->dim(), false};
+  auto lists = refinement_index->lists();
+  typename list_spec_wrapper<T>::list_spec<uint32_t> list_device_spec{refinement_index->dim(),
+                                                                      false};
   for (uint32_t label = 0; label < n_lists; label++) {
     ivf::resize_list(handle, lists[label], list_device_spec, n_candidates, uint32_t(0));
   }
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index b7e6955616..18f52e48c8 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -89,10 +89,9 @@ void serialize(raft::device_resources const& handle,
   handle.sync_stream();
   serialize_mdspan(handle, of, sizes_host.view());
 
-  auto list_store_spec = list_spec<T, uint32_t>{index_.dim(), true};
+  typename list_spec_wrapper<T>::list_spec<uint32_t> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::serialize_list<list_spec<T, uint32_t>, IdxT, uint32_t>(
-      handle, of, index_.lists()[label], list_store_spec, sizes_host(label));
+    ivf::serialize_list(handle, of, index_.lists()[label], list_store_spec /*, sizes_host(label)*/);
   }
   handle.sync_stream();
   of.close();
@@ -141,11 +140,10 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   }
   deserialize_mdspan(handle, infile, index_.list_sizes());
 
-  auto list_device_spec = list_spec<T, uint32_t>{index_.dim(), cma};
-  auto list_store_spec  = list_spec<T, uint32_t>{index_.dim(), true};
+  typename list_spec_wrapper<T>::list_spec<uint32_t> list_device_spec{index_.dim(), cma};
+  typename list_spec_wrapper<T>::list_spec<uint32_t> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::deserialize_list<list_spec<T, uint32_t>, IdxT, uint32_t>(
-      handle, infile, index_.lists()[label], list_store_spec, list_device_spec);
+    ivf::deserialize_list(handle, infile, index_.lists()[label], list_store_spec, list_device_spec);
   }
   handle.sync_stream();
   infile.close();
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
index 6ba43f4c3e..bf3014568a 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
@@ -664,8 +664,7 @@ __launch_bounds__(BlockSize) __global__ void process_and_fill_codes_kernel(
   const uint32_t pq_len     = pq_centers.extent(1);
   const uint32_t pq_dim     = new_vectors.extent(1) / pq_len;
 
-  auto pq_extents =
-    list_spec<uint8_t, uint32_t>{PqBits, pq_dim, true}.make_list_extents(out_ix + 1);
+  auto pq_extents = list_spec<uint32_t>{PqBits, pq_dim, true}.make_list_extents(out_ix + 1);
   auto pq_extents_vectorized =
     make_extents<uint32_t>(pq_extents.extent(0), pq_extents.extent(1), pq_extents.extent(2));
   auto pq_dataset = make_mdspan<pq_vec_t, uint32_t, row_major, false, true>(
@@ -907,7 +906,7 @@ void extend(raft::device_resources const& handle,
   std::optional<list_data<IdxT, size_t>> placeholder_list(
     std::in_place_t{},
     handle,
-    list_spec<uint8_t, size_t>(spec),
+    list_spec<size_t>(spec),
     n_rows + (kIndexGroupSize - 1) * std::min<IdxT>(n_clusters, n_rows));
 
   // Available device memory
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
index 1b90a21d54..33d9b363ba 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
@@ -94,9 +94,9 @@ void serialize(raft::device_resources const& handle_,
        handle_.get_stream());
   handle_.sync_stream();
   serialize_mdspan(handle_, of, sizes_host.view());
-  auto list_store_spec = list_spec<uint8_t, uint32_t>{index.pq_bits(), index.pq_dim(), true};
+  auto list_store_spec = list_spec<uint32_t>{index.pq_bits(), index.pq_dim(), true};
   for (uint32_t label = 0; label < index.n_lists(); label++) {
-    ivf::serialize_list<list_spec<uint8_t, uint32_t>, IdxT, uint32_t>(
+    ivf::serialize_list<list_spec, IdxT, uint32_t>(
       handle_, of, index.lists()[label], list_store_spec, sizes_host(label));
   }
 
@@ -151,10 +151,10 @@ auto deserialize(raft::device_resources const& handle_, const std::string& filen
   deserialize_mdspan(handle_, infile, index.centers_rot());
   deserialize_mdspan(handle_, infile, index.rotation_matrix());
   deserialize_mdspan(handle_, infile, index.list_sizes());
-  auto list_device_spec = list_spec<uint8_t, uint32_t>{pq_bits, pq_dim, cma};
-  auto list_store_spec  = list_spec<uint8_t, uint32_t>{pq_bits, pq_dim, true};
+  auto list_device_spec = list_spec<uint32_t>{pq_bits, pq_dim, cma};
+  auto list_store_spec  = list_spec<uint32_t>{pq_bits, pq_dim, true};
   for (auto& list : index.lists()) {
-    ivf::deserialize_list<list_spec<uint8_t, uint32_t>, IdxT, uint32_t>(
+    ivf::deserialize_list<list_spec, IdxT, uint32_t>(
       handle_, infile, list, list_store_spec, list_device_spec);
   }
 
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index e5751ff90a..cf66240b12 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -80,38 +80,41 @@ struct search_params : ann::search_params {
 static_assert(std::is_aggregate_v<index_params>);
 static_assert(std::is_aggregate_v<search_params>);
 
-template <typename ValueT, typename SizeT = uint32_t>
-struct list_spec {
-  using value_type   = ValueT;
-  using list_extents = matrix_extent<SizeT>;
-
-  SizeT align_max;
-  SizeT align_min;
-  uint32_t dim;
-
-  constexpr list_spec(uint32_t dim, bool conservative_memory_allocation)
-    : dim(dim),
-      align_min(kIndexGroupSize),
-      align_max(conservative_memory_allocation ? kIndexGroupSize : 1024)
-  {
-  }
+template <typename ValueT>
+struct list_spec_wrapper {
+  template <typename SizeT = uint32_t>
+  struct list_spec {
+    using value_type   = ValueT;
+    using list_extents = matrix_extent<SizeT>;
+
+    SizeT align_max;
+    SizeT align_min;
+    uint32_t dim;
+
+    constexpr list_spec(uint32_t dim, bool conservative_memory_allocation)
+      : dim(dim),
+        align_min(kIndexGroupSize),
+        align_max(conservative_memory_allocation ? kIndexGroupSize : 1024)
+    {
+    }
 
-  // Allow casting between different size-types (for safer size and offset calculations)
-  template <typename OtherValueT, typename OtherSizeT>
-  constexpr explicit list_spec(const list_spec<OtherValueT, OtherSizeT>& other_spec)
-    : dim{other_spec.dim}, align_min{other_spec.align_min}, align_max{other_spec.align_max}
-  {
-  }
+    // Allow casting between different size-types (for safer size and offset calculations)
+    template <typename OtherSizeT>
+    constexpr explicit list_spec(const list_spec<OtherSizeT>& other_spec)
+      : dim{other_spec.dim}, align_min{other_spec.align_min}, align_max{other_spec.align_max}
+    {
+    }
 
-  /** Determine the extents of an array enough to hold a given amount of data. */
-  constexpr auto make_list_extents(SizeT n_rows) const -> list_extents
-  {
-    return make_extents<SizeT>(n_rows, dim);
-  }
+    /** Determine the extents of an array enough to hold a given amount of data. */
+    constexpr auto make_list_extents(SizeT n_rows) const -> list_extents
+    {
+      return make_extents<SizeT>(n_rows, dim);
+    }
+  };
 };
 
 template <typename ValueT, typename IdxT, typename SizeT = uint32_t>
-using list_data = ivf::list<list_spec<ValueT, SizeT>, IdxT, SizeT>;
+using list_data = ivf::list<list_spec_wrapper<ValueT>::list_spec, IdxT, SizeT>;
 
 /**
  * @brief IVF-flat index.
diff --git a/cpp/include/raft/neighbors/ivf_list.hpp b/cpp/include/raft/neighbors/ivf_list.hpp
index 93c855139c..4839ac90f6 100644
--- a/cpp/include/raft/neighbors/ivf_list.hpp
+++ b/cpp/include/raft/neighbors/ivf_list.hpp
@@ -35,8 +35,10 @@
 namespace raft::neighbors::ivf {
 
 /** The data for a single IVF list. */
-template <typename SpecT, typename IdxT, typename SizeT>
-list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res, const SpecT& spec, SizeT n_rows)
+template <template <typename> typename SpecT, typename IdxT, typename SizeT>
+list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res,
+                               const SpecT<SizeT>& spec,
+                               SizeT n_rows)
   : size{n_rows}
 {
   auto capacity = round_up_safe<SizeT>(n_rows, spec.align_max);
@@ -45,7 +47,8 @@ list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res, const SpecT& s
     capacity = std::min<SizeT>(capacity, spec.align_max);
   }
   try {
-    data = make_device_mdarray<typename SpecT::value_type>(res, spec.make_list_extents(capacity));
+    data =
+      make_device_mdarray<typename SpecT<SizeT>::value_type>(res, spec.make_list_extents(capacity));
     indices = make_device_vector<IdxT, SizeT>(res, capacity);
   } catch (std::bad_alloc& e) {
     RAFT_FAIL(
@@ -65,13 +68,10 @@ list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res, const SpecT& s
  * Resize a list by the given id, so that it can contain the given number of records;
  * copy the data if necessary.
  */
-template <template <typename, typename> typename SpecT,
-          typename ValueT,
-          typename IdxT,
-          typename SizeT>
+template <template <typename> typename SpecT, typename IdxT, typename SizeT>
 void resize_list(raft::device_resources const& res,
-                 std::shared_ptr<list<SpecT<ValueT, SizeT>, IdxT, SizeT>>& orig_list,  // NOLINT
-                 const SpecT<ValueT, SizeT>& spec,
+                 std::shared_ptr<list<SpecT, IdxT, SizeT>>& orig_list,  // NOLINT
+                 const SpecT<SizeT>& spec,
                  SizeT new_used_size,
                  SizeT old_used_size)
 {
@@ -92,12 +92,12 @@ void resize_list(raft::device_resources const& res,
     old_used_size = 0;
   }
   if (skip_resize) { return; }
-  auto new_list =
-    std::make_shared<list<SpecT<ValueT, SizeT>, IdxT, SizeT>>(res, spec, new_used_size);
+  auto new_list = std::make_shared<list<SpecT, IdxT, SizeT>>(res, spec, new_used_size);
   if (old_used_size > 0) {
-    auto copied_data_extents = SpecT<ValueT, size_t>{spec}.make_list_extents(old_used_size);
-    auto copied_view         = make_mdspan<ValueT, size_t, row_major, false, true>(
-      new_list->data.data_handle(), copied_data_extents);
+    auto copied_data_extents = SpecT<size_t>{spec}.make_list_extents(old_used_size);
+    auto copied_view =
+      make_mdspan<typename SpecT<SizeT>::value_type, size_t, row_major, false, true>(
+        new_list->data.data_handle(), copied_data_extents);
     copy(copied_view.data_handle(),
          orig_list->data.data_handle(),
          copied_view.size(),
@@ -111,11 +111,11 @@ void resize_list(raft::device_resources const& res,
   new_list.swap(orig_list);
 }
 
-template <typename SpecT, typename IdxT, typename SizeT>
+template <template <typename> typename SpecT, typename IdxT, typename SizeT>
 void serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
                     const list<SpecT, IdxT, SizeT>& ld,
-                    const SpecT& store_spec,
+                    const SpecT<SizeT>& store_spec,
                     std::optional<SizeT> size_override = std::nullopt)
 {
   auto size = size_override.value_or(ld.size.load());
@@ -123,8 +123,9 @@ void serialize_list(const raft::device_resources& handle,
   if (size == 0) { return; }
 
   auto data_extents = store_spec.make_list_extents(size);
-  auto data_array   = make_host_mdarray<typename SpecT::value_type, SizeT, row_major>(data_extents);
-  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+  auto data_array =
+    make_host_mdarray<typename SpecT<SizeT>::value_type, SizeT, row_major>(data_extents);
+  auto inds_array = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
   copy(data_array.data_handle(), ld.data.data_handle(), data_array.size(), handle.get_stream());
   copy(inds_array.data_handle(), ld.indices.data_handle(), inds_array.size(), handle.get_stream());
   handle.sync_stream();
@@ -132,11 +133,11 @@ void serialize_list(const raft::device_resources& handle,
   serialize_mdspan(handle, os, inds_array.view());
 }
 
-template <typename SpecT, typename IdxT, typename SizeT>
+template <template <typename> typename SpecT, typename IdxT, typename SizeT>
 void serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
-                    const std::shared_ptr<const list<SpecT, IdxT, SizeT>>& ld,
-                    const SpecT& store_spec,
+                    const std::shared_ptr<list<SpecT, IdxT, SizeT>>& ld,
+                    const SpecT<SizeT>& store_spec,
                     std::optional<SizeT> size_override = std::nullopt)
 {
   if (ld) {
@@ -146,19 +147,20 @@ void serialize_list(const raft::device_resources& handle,
   }
 }
 
-template <typename SpecT, typename IdxT, typename SizeT>
+template <template <typename> typename SpecT, typename IdxT, typename SizeT>
 void deserialize_list(const raft::device_resources& handle,
                       std::istream& is,
                       std::shared_ptr<list<SpecT, IdxT, SizeT>>& ld,
-                      const SpecT& store_spec,
-                      const SpecT& device_spec)
+                      const SpecT<SizeT>& store_spec,
+                      const SpecT<SizeT>& device_spec)
 {
   auto size = deserialize_scalar<SizeT>(handle, is);
   if (size == 0) { return ld.reset(); }
   std::make_shared<list<SpecT, IdxT, SizeT>>(handle, device_spec, size).swap(ld);
   auto data_extents = store_spec.make_list_extents(size);
-  auto data_array   = make_host_mdarray<typename SpecT::value_type, SizeT, row_major>(data_extents);
-  auto inds_array   = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+  auto data_array =
+    make_host_mdarray<typename SpecT<SizeT>::value_type, SizeT, row_major>(data_extents);
+  auto inds_array = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
   deserialize_mdspan(handle, is, data_array.view());
   deserialize_mdspan(handle, is, inds_array.view());
   copy(ld->data.data_handle(), data_array.data_handle(), data_array.size(), handle.get_stream());
diff --git a/cpp/include/raft/neighbors/ivf_list_types.hpp b/cpp/include/raft/neighbors/ivf_list_types.hpp
index cc191b12ea..aae6a63794 100644
--- a/cpp/include/raft/neighbors/ivf_list_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_list_types.hpp
@@ -36,10 +36,10 @@ constexpr static IdxT kInvalidRecord = (std::is_signed_v<IdxT> ? IdxT{0}
                                        1;
 
 /** The data for a single IVF list. */
-template <typename SpecT, typename IdxT, typename SizeT = uint32_t>
+template <template <typename> typename SpecT, typename IdxT, typename SizeT = uint32_t>
 struct list {
-  using value_type   = typename SpecT::value_type;
-  using list_extents = typename SpecT::list_extents;
+  using value_type   = typename SpecT<SizeT>::value_type;
+  using list_extents = typename SpecT<SizeT>::list_extents;
 
   /** Possibly encoded data; it's layout is defined by `SpecT`. */
   device_mdarray<value_type, list_extents, row_major> data;
@@ -49,7 +49,7 @@ struct list {
   std::atomic<SizeT> size;
 
   /** Allocate a new list capable of holding at least `n_rows` data records and indices. */
-  list(raft::device_resources const& res, const SpecT& spec, SizeT n_rows);
+  list(raft::device_resources const& res, const SpecT<SizeT>& spec, SizeT n_rows);
 };
 
 }  // namespace raft::neighbors::ivf
diff --git a/cpp/include/raft/neighbors/ivf_pq_types.hpp b/cpp/include/raft/neighbors/ivf_pq_types.hpp
index df90a25de8..8c00b08452 100644
--- a/cpp/include/raft/neighbors/ivf_pq_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_pq_types.hpp
@@ -161,9 +161,9 @@ constexpr static uint32_t kIndexGroupVecLen = 16;
 template <typename IdxT>
 constexpr static IdxT kOutOfBoundsRecord = std::numeric_limits<IdxT>::max();
 
-template <typename ValueT = uint8_t, typename SizeT = uint32_t>
+template <typename SizeT = uint32_t>
 struct list_spec {
-  using value_type = ValueT;
+  using value_type = uint8_t;
   /** PQ-encoded data stored in the interleaved format:
    *
    *    [ ceildiv(list_size, kIndexGroupSize)
@@ -189,8 +189,8 @@ struct list_spec {
   }
 
   // Allow casting between different size-types (for safer size and offset calculations)
-  template <typename OtherValueT, typename OtherSizeT>
-  constexpr explicit list_spec(const list_spec<OtherValueT, OtherSizeT>& other_spec)
+  template <typename OtherSizeT>
+  constexpr explicit list_spec(const list_spec<OtherSizeT>& other_spec)
     : pq_bits{other_spec.pq_bits},
       pq_dim{other_spec.pq_dim},
       align_min{other_spec.align_min},
@@ -211,7 +211,7 @@ struct list_spec {
 };
 
 template <typename IdxT, typename SizeT = uint32_t>
-using list_data = ivf::list<list_spec<uint8_t, SizeT>, IdxT, SizeT>;
+using list_data = ivf::list<list_spec, IdxT, SizeT>;
 
 /**
  * @brief IVF-PQ index.

From 32936c97117c145e7103c39c17cd46fa148ca97a Mon Sep 17 00:00:00 2001
From: achirkin <achirkin@users.noreply.github.com>
Date: Tue, 7 Mar 2023 12:53:20 +0100
Subject: [PATCH 16/33] Make ivf::list more flexible

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  |  7 +-
 .../neighbors/detail/ivf_flat_serialize.cuh   |  8 +-
 .../raft/neighbors/detail/ivf_pq_build.cuh    |  7 +-
 .../neighbors/detail/ivf_pq_serialize.cuh     | 12 +--
 cpp/include/raft/neighbors/ivf_list.hpp       | 91 ++++++++++---------
 cpp/include/raft/neighbors/ivf_list_types.hpp | 38 ++++++--
 cpp/include/raft/neighbors/ivf_pq_types.hpp   |  7 +-
 7 files changed, 102 insertions(+), 68 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 5bdcce6053..7e745d2f9f 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -166,8 +166,8 @@ void extend(raft::device_resources const& handle,
   auto stream  = handle.get_stream();
   auto n_lists = index->n_lists();
   auto dim     = index->dim();
-  typename list_spec_wrapper<T>::list_spec<uint32_t> list_device_spec{
-    index->dim(), index->conservative_memory_allocation()};
+  list_spec<uint32_t, T, IdxT> list_device_spec{index->dim(),
+                                                index->conservative_memory_allocation()};
   common::nvtx::range<common::nvtx::domain::raft> fun_scope(
     "ivf_flat::extend(%zu, %u)", size_t(n_rows), dim);
 
@@ -376,8 +376,7 @@ inline void fill_refinement_index(raft::device_resources const& handle,
 
   // Allocate new memory
   auto lists = refinement_index->lists();
-  typename list_spec_wrapper<T>::list_spec<uint32_t> list_device_spec{refinement_index->dim(),
-                                                                      false};
+  list_spec<uint32_t, T, IdxT> list_device_spec{refinement_index->dim(), false};
   for (uint32_t label = 0; label < n_lists; label++) {
     ivf::resize_list(handle, lists[label], list_device_spec, n_candidates, uint32_t(0));
   }
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 18f52e48c8..880a196206 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -89,7 +89,7 @@ void serialize(raft::device_resources const& handle,
   handle.sync_stream();
   serialize_mdspan(handle, of, sizes_host.view());
 
-  typename list_spec_wrapper<T>::list_spec<uint32_t> list_store_spec{index_.dim(), true};
+  list_spec<uint32_t, T, IdxT> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
     ivf::serialize_list(handle, of, index_.lists()[label], list_store_spec /*, sizes_host(label)*/);
   }
@@ -140,8 +140,8 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   }
   deserialize_mdspan(handle, infile, index_.list_sizes());
 
-  typename list_spec_wrapper<T>::list_spec<uint32_t> list_device_spec{index_.dim(), cma};
-  typename list_spec_wrapper<T>::list_spec<uint32_t> list_store_spec{index_.dim(), true};
+  list_spec<uint32_t, T, IdxT> list_device_spec{index_.dim(), cma};
+  list_spec<uint32_t, T, IdxT> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
     ivf::deserialize_list(handle, infile, index_.lists()[label], list_store_spec, list_device_spec);
   }
@@ -152,4 +152,4 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
 
   return index_;
 }
-}  // namespace raft::neighbors::ivf_flat::detail
\ No newline at end of file
+}  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
index bf3014568a..1a563d213e 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_build.cuh
@@ -664,7 +664,7 @@ __launch_bounds__(BlockSize) __global__ void process_and_fill_codes_kernel(
   const uint32_t pq_len     = pq_centers.extent(1);
   const uint32_t pq_dim     = new_vectors.extent(1) / pq_len;
 
-  auto pq_extents = list_spec<uint32_t>{PqBits, pq_dim, true}.make_list_extents(out_ix + 1);
+  auto pq_extents = list_spec<uint32_t, IdxT>{PqBits, pq_dim, true}.make_list_extents(out_ix + 1);
   auto pq_extents_vectorized =
     make_extents<uint32_t>(pq_extents.extent(0), pq_extents.extent(1), pq_extents.extent(2));
   auto pq_dataset = make_mdspan<pq_vec_t, uint32_t, row_major, false, true>(
@@ -899,14 +899,15 @@ void extend(raft::device_resources const& handle,
     &managed_memory_upstream, 1024 * 1024);
 
   // The spec defines how the clusters look like
-  auto spec = list_spec{index->pq_bits(), index->pq_dim(), index->conservative_memory_allocation()};
+  auto spec = list_spec<uint32_t, IdxT>{
+    index->pq_bits(), index->pq_dim(), index->conservative_memory_allocation()};
   // Try to allocate an index with the same parameters and the projected new size
   // (which can be slightly larger than index->size() + n_rows, due to padding).
   // If this fails, the index would be too big to fit in the device anyway.
   std::optional<list_data<IdxT, size_t>> placeholder_list(
     std::in_place_t{},
     handle,
-    list_spec<size_t>(spec),
+    list_spec<size_t, IdxT>{spec},
     n_rows + (kIndexGroupSize - 1) * std::min<IdxT>(n_clusters, n_rows));
 
   // Available device memory
diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
index 33d9b363ba..7770ce53e9 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
@@ -94,10 +94,9 @@ void serialize(raft::device_resources const& handle_,
        handle_.get_stream());
   handle_.sync_stream();
   serialize_mdspan(handle_, of, sizes_host.view());
-  auto list_store_spec = list_spec<uint32_t>{index.pq_bits(), index.pq_dim(), true};
+  auto list_store_spec = list_spec<uint32_t, IdxT>{index.pq_bits(), index.pq_dim(), true};
   for (uint32_t label = 0; label < index.n_lists(); label++) {
-    ivf::serialize_list<list_spec, IdxT, uint32_t>(
-      handle_, of, index.lists()[label], list_store_spec, sizes_host(label));
+    ivf::serialize_list(handle_, of, index.lists()[label], list_store_spec, sizes_host(label));
   }
 
   of.close();
@@ -151,11 +150,10 @@ auto deserialize(raft::device_resources const& handle_, const std::string& filen
   deserialize_mdspan(handle_, infile, index.centers_rot());
   deserialize_mdspan(handle_, infile, index.rotation_matrix());
   deserialize_mdspan(handle_, infile, index.list_sizes());
-  auto list_device_spec = list_spec<uint32_t>{pq_bits, pq_dim, cma};
-  auto list_store_spec  = list_spec<uint32_t>{pq_bits, pq_dim, true};
+  auto list_device_spec = list_spec<uint32_t, IdxT>{pq_bits, pq_dim, cma};
+  auto list_store_spec  = list_spec<uint32_t, IdxT>{pq_bits, pq_dim, true};
   for (auto& list : index.lists()) {
-    ivf::deserialize_list<list_spec, IdxT, uint32_t>(
-      handle_, infile, list, list_store_spec, list_device_spec);
+    ivf::deserialize_list(handle_, infile, list, list_store_spec, list_device_spec);
   }
 
   handle_.sync_stream();
diff --git a/cpp/include/raft/neighbors/ivf_list.hpp b/cpp/include/raft/neighbors/ivf_list.hpp
index 4839ac90f6..e2aa661d12 100644
--- a/cpp/include/raft/neighbors/ivf_list.hpp
+++ b/cpp/include/raft/neighbors/ivf_list.hpp
@@ -35,10 +35,12 @@
 namespace raft::neighbors::ivf {
 
 /** The data for a single IVF list. */
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
-list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res,
-                               const SpecT<SizeT>& spec,
-                               SizeT n_rows)
+template <template <typename, typename...> typename SpecT,
+          typename SizeT,
+          typename... SpecExtraArgs>
+list<SpecT, SizeT, SpecExtraArgs...>::list(raft::device_resources const& res,
+                                           const spec_type& spec,
+                                           size_type n_rows)
   : size{n_rows}
 {
   auto capacity = round_up_safe<SizeT>(n_rows, spec.align_max);
@@ -47,9 +49,8 @@ list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res,
     capacity = std::min<SizeT>(capacity, spec.align_max);
   }
   try {
-    data =
-      make_device_mdarray<typename SpecT<SizeT>::value_type>(res, spec.make_list_extents(capacity));
-    indices = make_device_vector<IdxT, SizeT>(res, capacity);
+    data    = make_device_mdarray<value_type>(res, spec.make_list_extents(capacity));
+    indices = make_device_vector<index_type, SizeT>(res, capacity);
   } catch (std::bad_alloc& e) {
     RAFT_FAIL(
       "ivf::list: failed to allocate a big enough list to hold all data "
@@ -60,20 +61,22 @@ list<SpecT, IdxT, SizeT>::list(raft::device_resources const& res,
       e.what());
   }
   // Fill the index buffer with a pre-defined marker for easier debugging
-  thrust::fill_n(
-    res.get_thrust_policy(), indices.data_handle(), indices.size(), ivf::kInvalidRecord<IdxT>);
+  thrust::fill_n(res.get_thrust_policy(),
+                 indices.data_handle(),
+                 indices.size(),
+                 ivf::kInvalidRecord<index_type>);
 }
 
 /**
  * Resize a list by the given id, so that it can contain the given number of records;
  * copy the data if necessary.
  */
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
+template <typename ListT>
 void resize_list(raft::device_resources const& res,
-                 std::shared_ptr<list<SpecT, IdxT, SizeT>>& orig_list,  // NOLINT
-                 const SpecT<SizeT>& spec,
-                 SizeT new_used_size,
-                 SizeT old_used_size)
+                 std::shared_ptr<ListT>& orig_list,  // NOLINT
+                 const typename ListT::spec_type& spec,
+                 typename ListT::size_type new_used_size,
+                 typename ListT::size_type old_used_size)
 {
   bool skip_resize = false;
   if (orig_list) {
@@ -92,11 +95,11 @@ void resize_list(raft::device_resources const& res,
     old_used_size = 0;
   }
   if (skip_resize) { return; }
-  auto new_list = std::make_shared<list<SpecT, IdxT, SizeT>>(res, spec, new_used_size);
+  auto new_list = std::make_shared<ListT>(res, spec, new_used_size);
   if (old_used_size > 0) {
-    auto copied_data_extents = SpecT<size_t>{spec}.make_list_extents(old_used_size);
+    auto copied_data_extents = spec.make_list_extents(old_used_size);
     auto copied_view =
-      make_mdspan<typename SpecT<SizeT>::value_type, size_t, row_major, false, true>(
+      make_mdspan<typename ListT::value_type, typename ListT::size_type, row_major, false, true>(
         new_list->data.data_handle(), copied_data_extents);
     copy(copied_view.data_handle(),
          orig_list->data.data_handle(),
@@ -111,21 +114,24 @@ void resize_list(raft::device_resources const& res,
   new_list.swap(orig_list);
 }
 
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
-void serialize_list(const raft::device_resources& handle,
+template <typename ListT>
+auto serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
-                    const list<SpecT, IdxT, SizeT>& ld,
-                    const SpecT<SizeT>& store_spec,
-                    std::optional<SizeT> size_override = std::nullopt)
+                    const ListT& ld,
+                    const typename ListT::spec_type& store_spec,
+                    std::optional<typename ListT::size_type> size_override = std::nullopt)
+  -> enable_if_valid_list_t<ListT>
 {
-  auto size = size_override.value_or(ld.size.load());
+  using size_type = typename ListT::size_type;
+  auto size       = size_override.value_or(ld.size.load());
   serialize_scalar(handle, os, size);
   if (size == 0) { return; }
 
   auto data_extents = store_spec.make_list_extents(size);
   auto data_array =
-    make_host_mdarray<typename SpecT<SizeT>::value_type, SizeT, row_major>(data_extents);
-  auto inds_array = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+    make_host_mdarray<typename ListT::value_type, size_type, row_major>(data_extents);
+  auto inds_array = make_host_mdarray<typename ListT::index_type, size_type, row_major>(
+    make_extents<size_type>(size));
   copy(data_array.data_handle(), ld.data.data_handle(), data_array.size(), handle.get_stream());
   copy(inds_array.data_handle(), ld.indices.data_handle(), inds_array.size(), handle.get_stream());
   handle.sync_stream();
@@ -133,34 +139,37 @@ void serialize_list(const raft::device_resources& handle,
   serialize_mdspan(handle, os, inds_array.view());
 }
 
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
-void serialize_list(const raft::device_resources& handle,
+template <typename ListT>
+auto serialize_list(const raft::device_resources& handle,
                     std::ostream& os,
-                    const std::shared_ptr<list<SpecT, IdxT, SizeT>>& ld,
-                    const SpecT<SizeT>& store_spec,
-                    std::optional<SizeT> size_override = std::nullopt)
+                    const std::shared_ptr<ListT>& ld,
+                    const typename ListT::spec_type& store_spec,
+                    std::optional<typename ListT::size_type> size_override = std::nullopt)
+  -> enable_if_valid_list_t<ListT>
 {
   if (ld) {
-    return serialize_list(handle, os, *ld, store_spec, size_override);
+    return serialize_list<ListT>(handle, os, *ld, store_spec, size_override);
   } else {
-    return serialize_scalar(handle, os, SizeT{0});
+    return serialize_scalar(handle, os, typename ListT::size_type{0});
   }
 }
 
-template <template <typename> typename SpecT, typename IdxT, typename SizeT>
-void deserialize_list(const raft::device_resources& handle,
+template <typename ListT>
+auto deserialize_list(const raft::device_resources& handle,
                       std::istream& is,
-                      std::shared_ptr<list<SpecT, IdxT, SizeT>>& ld,
-                      const SpecT<SizeT>& store_spec,
-                      const SpecT<SizeT>& device_spec)
+                      std::shared_ptr<ListT>& ld,
+                      const typename ListT::spec_type& store_spec,
+                      const typename ListT::spec_type& device_spec) -> enable_if_valid_list_t<ListT>
 {
-  auto size = deserialize_scalar<SizeT>(handle, is);
+  using size_type = typename ListT::size_type;
+  auto size       = deserialize_scalar<size_type>(handle, is);
   if (size == 0) { return ld.reset(); }
-  std::make_shared<list<SpecT, IdxT, SizeT>>(handle, device_spec, size).swap(ld);
+  std::make_shared<ListT>(handle, device_spec, size).swap(ld);
   auto data_extents = store_spec.make_list_extents(size);
   auto data_array =
-    make_host_mdarray<typename SpecT<SizeT>::value_type, SizeT, row_major>(data_extents);
-  auto inds_array = make_host_mdarray<IdxT, SizeT, row_major>(make_extents<SizeT>(size));
+    make_host_mdarray<typename ListT::value_type, size_type, row_major>(data_extents);
+  auto inds_array = make_host_mdarray<typename ListT::index_type, size_type, row_major>(
+    make_extents<size_type>(size));
   deserialize_mdspan(handle, is, data_array.view());
   deserialize_mdspan(handle, is, inds_array.view());
   copy(ld->data.data_handle(), data_array.data_handle(), data_array.size(), handle.get_stream());
diff --git a/cpp/include/raft/neighbors/ivf_list_types.hpp b/cpp/include/raft/neighbors/ivf_list_types.hpp
index aae6a63794..233775ea39 100644
--- a/cpp/include/raft/neighbors/ivf_list_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_list_types.hpp
@@ -36,20 +36,46 @@ constexpr static IdxT kInvalidRecord = (std::is_signed_v<IdxT> ? IdxT{0}
                                        1;
 
 /** The data for a single IVF list. */
-template <template <typename> typename SpecT, typename IdxT, typename SizeT = uint32_t>
+template <template <typename, typename...> typename SpecT,
+          typename SizeT,
+          typename... SpecExtraArgs>
 struct list {
-  using value_type   = typename SpecT<SizeT>::value_type;
-  using list_extents = typename SpecT<SizeT>::list_extents;
+  using size_type    = SizeT;
+  using spec_type    = SpecT<size_type, SpecExtraArgs...>;
+  using value_type   = typename spec_type::value_type;
+  using index_type   = typename spec_type::index_type;
+  using list_extents = typename spec_type::list_extents;
 
   /** Possibly encoded data; it's layout is defined by `SpecT`. */
   device_mdarray<value_type, list_extents, row_major> data;
   /** Source indices. */
-  device_mdarray<IdxT, extent_1d<SizeT>, row_major> indices;
+  device_mdarray<index_type, extent_1d<size_type>, row_major> indices;
   /** The actual size of the content. */
-  std::atomic<SizeT> size;
+  std::atomic<size_type> size;
 
   /** Allocate a new list capable of holding at least `n_rows` data records and indices. */
-  list(raft::device_resources const& res, const SpecT<SizeT>& spec, SizeT n_rows);
+  list(raft::device_resources const& res, const spec_type& spec, size_type n_rows);
 };
 
+template <typename ListT, class T = void>
+struct enable_if_valid_list {
+};
+
+template <class T,
+          template <typename, typename...>
+          typename SpecT,
+          typename SizeT,
+          typename... SpecExtraArgs>
+struct enable_if_valid_list<list<SpecT, SizeT, SpecExtraArgs...>, T> {
+  using type = T;
+};
+
+/**
+ * Designed after `std::enable_if_t`, this trait is helpful in the instance resolution;
+ * plug this in the return type of a function that has an instance of `ivf::list` as
+ * a template parameter.
+ */
+template <typename ListT, class T = void>
+using enable_if_valid_list_t = typename enable_if_valid_list<ListT, T>::type;
+
 }  // namespace raft::neighbors::ivf
diff --git a/cpp/include/raft/neighbors/ivf_pq_types.hpp b/cpp/include/raft/neighbors/ivf_pq_types.hpp
index 8c00b08452..78570fbd1a 100644
--- a/cpp/include/raft/neighbors/ivf_pq_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_pq_types.hpp
@@ -161,9 +161,10 @@ constexpr static uint32_t kIndexGroupVecLen = 16;
 template <typename IdxT>
 constexpr static IdxT kOutOfBoundsRecord = std::numeric_limits<IdxT>::max();
 
-template <typename SizeT = uint32_t>
+template <typename SizeT, typename IdxT>
 struct list_spec {
   using value_type = uint8_t;
+  using index_type = IdxT;
   /** PQ-encoded data stored in the interleaved format:
    *
    *    [ ceildiv(list_size, kIndexGroupSize)
@@ -190,7 +191,7 @@ struct list_spec {
 
   // Allow casting between different size-types (for safer size and offset calculations)
   template <typename OtherSizeT>
-  constexpr explicit list_spec(const list_spec<OtherSizeT>& other_spec)
+  constexpr explicit list_spec(const list_spec<OtherSizeT, IdxT>& other_spec)
     : pq_bits{other_spec.pq_bits},
       pq_dim{other_spec.pq_dim},
       align_min{other_spec.align_min},
@@ -211,7 +212,7 @@ struct list_spec {
 };
 
 template <typename IdxT, typename SizeT = uint32_t>
-using list_data = ivf::list<list_spec, IdxT, SizeT>;
+using list_data = ivf::list<list_spec, SizeT, IdxT>;
 
 /**
  * @brief IVF-PQ index.

From acf1888bbc6b425b1156a71ced25e53819263820 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Tue, 7 Mar 2023 14:00:22 +0100
Subject: [PATCH 17/33] Use new list API for ivf flat

---
 .../neighbors/detail/ivf_flat_serialize.cuh   |  2 +-
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 58 +++++++++----------
 2 files changed, 29 insertions(+), 31 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 880a196206..8449a2370c 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -91,7 +91,7 @@ void serialize(raft::device_resources const& handle,
 
   list_spec<uint32_t, T, IdxT> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::serialize_list(handle, of, index_.lists()[label], list_store_spec /*, sizes_host(label)*/);
+    ivf::serialize_list(handle, of, index_.lists()[label], list_store_spec, sizes_host(label));
   }
   handle.sync_stream();
   of.close();
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index cf66240b12..d350801cb8 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -80,41 +80,39 @@ struct search_params : ann::search_params {
 static_assert(std::is_aggregate_v<index_params>);
 static_assert(std::is_aggregate_v<search_params>);
 
-template <typename ValueT>
-struct list_spec_wrapper {
-  template <typename SizeT = uint32_t>
-  struct list_spec {
-    using value_type   = ValueT;
-    using list_extents = matrix_extent<SizeT>;
-
-    SizeT align_max;
-    SizeT align_min;
-    uint32_t dim;
-
-    constexpr list_spec(uint32_t dim, bool conservative_memory_allocation)
-      : dim(dim),
-        align_min(kIndexGroupSize),
-        align_max(conservative_memory_allocation ? kIndexGroupSize : 1024)
-    {
-    }
+template <typename SizeT, typename ValueT, typename IdxT>
+struct list_spec {
+  using value_type   = ValueT;
+  using list_extents = matrix_extent<SizeT>;
+  using index_type   = IdxT;
+
+  SizeT align_max;
+  SizeT align_min;
+  uint32_t dim;
+
+  constexpr list_spec(uint32_t dim, bool conservative_memory_allocation)
+    : dim(dim),
+      align_min(kIndexGroupSize),
+      align_max(conservative_memory_allocation ? kIndexGroupSize : 1024)
+  {
+  }
 
-    // Allow casting between different size-types (for safer size and offset calculations)
-    template <typename OtherSizeT>
-    constexpr explicit list_spec(const list_spec<OtherSizeT>& other_spec)
-      : dim{other_spec.dim}, align_min{other_spec.align_min}, align_max{other_spec.align_max}
-    {
-    }
+  // Allow casting between different size-types (for safer size and offset calculations)
+  template <typename OtherSizeT>
+  constexpr explicit list_spec(const list_spec<OtherSizeT, ValueT, IdxT>& other_spec)
+    : dim{other_spec.dim}, align_min{other_spec.align_min}, align_max{other_spec.align_max}
+  {
+  }
 
-    /** Determine the extents of an array enough to hold a given amount of data. */
-    constexpr auto make_list_extents(SizeT n_rows) const -> list_extents
-    {
-      return make_extents<SizeT>(n_rows, dim);
-    }
-  };
+  /** Determine the extents of an array enough to hold a given amount of data. */
+  constexpr auto make_list_extents(SizeT n_rows) const -> list_extents
+  {
+    return make_extents<SizeT>(n_rows, dim);
+  }
 };
 
 template <typename ValueT, typename IdxT, typename SizeT = uint32_t>
-using list_data = ivf::list<list_spec_wrapper<ValueT>::list_spec, IdxT, SizeT>;
+using list_data = ivf::list<list_spec, SizeT, ValueT, IdxT>;
 
 /**
  * @brief IVF-flat index.

From f790abf4e77089da4f4098661f38be11616b0d18 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Thu, 9 Mar 2023 02:13:29 +0100
Subject: [PATCH 18/33] Fix adaptive centers

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  | 29 ++++++++++++++-----
 .../neighbors/detail/ivf_flat_serialize.cuh   |  1 +
 cpp/include/raft/neighbors/ivf_flat_types.hpp | 22 ++++++++------
 3 files changed, 35 insertions(+), 17 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 7e745d2f9f..7c6ca7cb2d 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -62,11 +62,13 @@ auto clone(const raft::device_resources& res, const index<T, IdxT>& source) -> i
        source.centers().data_handle(),
        source.centers().size(),
        stream);
-  if (source.center_norms().has_value())
-    copy(target.center_norms().value().data_handle(),
-         source.center_norms().value().data_handle(),
-         source.center_norms().value().size(),
+  if (source.center_norms().has_value()) {
+    target.allocate_center_norms(res);
+    copy(target.center_norms()->data_handle(),
+         source.center_norms()->data_handle(),
+         source.center_norms()->size(),
          stream);
+  }
   // Copy shared pointers
   target.lists() = source.lists();
 
@@ -219,9 +221,9 @@ void extend(raft::device_resources const& handle,
   }
 
   // Calculate and allocate new list data
+  std::vector<uint32_t> new_list_sizes(n_lists);
+  std::vector<uint32_t> old_list_sizes(n_lists);
   {
-    std::vector<uint32_t> new_list_sizes(n_lists);
-    std::vector<uint32_t> old_list_sizes(n_lists);
     copy(old_list_sizes.data(), old_list_sizes_dev.data_handle(), n_lists, stream);
     copy(new_list_sizes.data(), list_sizes_ptr, n_lists, stream);
     handle.sync_stream();
@@ -233,7 +235,6 @@ void extend(raft::device_resources const& handle,
   }
   // Update the pointers and the sizes
   index->recompute_internal_state(handle);
-
   // Copy the old sizes, so we can start from the current state of the index;
   // we'll rebuild the `list_sizes_ptr` in the following kernel, using it as an atomic counter.
   raft::copy(list_sizes_ptr, old_list_sizes_dev.data_handle(), n_lists, stream);
@@ -253,7 +254,19 @@ void extend(raft::device_resources const& handle,
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 
   // Precompute the centers vector norms for L2Expanded distance
-  if (index->center_norms().has_value() && index->adaptive_centers()) {
+  if (!index->center_norms().has_value()) {
+    index->allocate_center_norms(handle);
+    if (index->center_norms().has_value()) {
+      raft::linalg::rowNorm(index->center_norms()->data_handle(),
+                            index->centers().data_handle(),
+                            dim,
+                            n_lists,
+                            raft::linalg::L2Norm,
+                            true,
+                            stream);
+      RAFT_LOG_TRACE_VEC(index->center_norms()->data_handle(), std::min<uint32_t>(dim, 20));
+    }
+  } else if (index->center_norms().has_value() && index->adaptive_centers()) {
     raft::linalg::rowNorm(index->center_norms()->data_handle(),
                           index->centers().data_handle(),
                           dim,
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 8449a2370c..1195197fa3 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -131,6 +131,7 @@ auto deserialize(raft::device_resources const& handle, const std::string& filena
   deserialize_mdspan(handle, infile, index_.centers());
   bool has_norms = deserialize_scalar<bool>(handle, infile);
   if (has_norms) {
+    index_.allocate_center_norms(handle);
     if (!index_.center_norms()) {
       RAFT_FAIL("Error inconsistent center norms");
     } else {
diff --git a/cpp/include/raft/neighbors/ivf_flat_types.hpp b/cpp/include/raft/neighbors/ivf_flat_types.hpp
index d350801cb8..20bca6e3e6 100644
--- a/cpp/include/raft/neighbors/ivf_flat_types.hpp
+++ b/cpp/include/raft/neighbors/ivf_flat_types.hpp
@@ -253,15 +253,6 @@ struct index : ann::index {
       inds_ptrs_{make_device_vector<IdxT*, uint32_t>(res, n_lists)},
       total_size_{0}
   {
-    switch (metric_) {
-      case raft::distance::DistanceType::L2Expanded:
-      case raft::distance::DistanceType::L2SqrtExpanded:
-      case raft::distance::DistanceType::L2Unexpanded:
-      case raft::distance::DistanceType::L2SqrtUnexpanded:
-        center_norms_ = make_device_vector<float, uint32_t>(res, n_lists);
-        break;
-      default: center_norms_ = std::nullopt;
-    }
     check_consistency();
   }
 
@@ -326,6 +317,19 @@ struct index : ann::index {
     check_consistency();
   }
 
+  void allocate_center_norms(raft::device_resources const& res)
+  {
+    switch (metric_) {
+      case raft::distance::DistanceType::L2Expanded:
+      case raft::distance::DistanceType::L2SqrtExpanded:
+      case raft::distance::DistanceType::L2Unexpanded:
+      case raft::distance::DistanceType::L2SqrtUnexpanded:
+        center_norms_ = make_device_vector<float, uint32_t>(res, n_lists());
+        break;
+      default: center_norms_ = std::nullopt;
+    }
+  }
+
   /** Lists' data and indices. */
   inline auto lists() noexcept -> std::vector<std::shared_ptr<list_data<T, IdxT>>>&
   {

From 6a051960999462f07af162e0ab26e8d823736a59 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Thu, 9 Mar 2023 18:19:31 +0100
Subject: [PATCH 19/33] Fix adaptive center norms testing

---
 cpp/test/neighbors/ann_ivf_flat.cuh | 12 ++++++++----
 1 file changed, 8 insertions(+), 4 deletions(-)

diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index b8b86e26b4..34531ddb89 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -157,9 +157,10 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
       {
         ivf_flat::index_params index_params;
         ivf_flat::search_params search_params;
-        index_params.n_lists   = ps.nlist;
-        index_params.metric    = ps.metric;
-        search_params.n_probes = ps.nprobe;
+        index_params.n_lists          = ps.nlist;
+        index_params.metric           = ps.metric;
+        index_params.adaptive_centers = ps.adaptive_centers;
+        search_params.n_probes        = ps.nprobe;
 
         index_params.add_data_on_build        = false;
         index_params.kmeans_trainset_fraction = 0.5;
@@ -221,16 +222,19 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         if (index_2.adaptive_centers()) {
           // The centers must be up-to-date with the corresponding data
           std::vector<uint32_t> list_sizes(index_2.n_lists());
+          std::vector<IdxT*> list_indices(index_2.n_lists());
           rmm::device_uvector<float> centroid(ps.dim, stream_);
           raft::copy(
             list_sizes.data(), index_2.list_sizes().data_handle(), index_2.n_lists(), stream_);
+          raft::copy(
+            list_indices.data(), index_2.inds_ptrs().data_handle(), index_2.n_lists(), stream_);
           handle_.sync_stream(stream_);
           for (uint32_t l = 0; l < index_2.n_lists(); l++) {
             rmm::device_uvector<float> cluster_data(list_sizes[l] * ps.dim, stream_);
             raft::spatial::knn::detail::utils::copy_selected<float>((IdxT)list_sizes[l],
                                                                     (IdxT)ps.dim,
                                                                     database.data(),
-                                                                    index_2.inds_ptrs()(l),
+                                                                    list_indices[l],
                                                                     (IdxT)ps.dim,
                                                                     cluster_data.data(),
                                                                     (IdxT)ps.dim,

From f0a7031d0d349a4217fac18e429fdd39505efc84 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Thu, 9 Mar 2023 18:42:08 +0100
Subject: [PATCH 20/33] Fix list spec template after merge

---
 cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
index a675cb94fb..826ed90db1 100644
--- a/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_pq_serialize.cuh
@@ -89,10 +89,9 @@ void serialize(raft::device_resources const& handle_, std::ostream& os, const in
        handle_.get_stream());
   handle_.sync_stream();
   serialize_mdspan(handle_, os, sizes_host.view());
-  auto list_store_spec = list_spec<uint32_t>{index.pq_bits(), index.pq_dim(), true};
+  auto list_store_spec = list_spec<uint32_t, IdxT>{index.pq_bits(), index.pq_dim(), true};
   for (uint32_t label = 0; label < index.n_lists(); label++) {
-    ivf::serialize_list<list_spec, IdxT, uint32_t>(
-      handle_, os, index.lists()[label], list_store_spec, sizes_host(label));
+    ivf::serialize_list(handle_, os, index.lists()[label], list_store_spec, sizes_host(label));
   }
 }
 

From a6c54a40e1a3a528f5a3ca405f2c4da6830ab7f6 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Fri, 10 Mar 2023 00:17:30 +0100
Subject: [PATCH 21/33] Add serialization overload

---
 .../neighbors/detail/ivf_flat_serialize.cuh   |  92 ++++++-----
 cpp/include/raft/neighbors/ivf_flat.cuh       |   8 +-
 .../raft/neighbors/ivf_flat_serialize.cuh     | 156 ++++++++++++++++++
 cpp/test/neighbors/ann_ivf_flat.cuh           |   1 +
 4 files changed, 212 insertions(+), 45 deletions(-)
 create mode 100644 cpp/include/raft/neighbors/ivf_flat_serialize.cuh

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 1195197fa3..18e8c2862e 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -55,31 +55,26 @@ template struct check_index_layout<sizeof(index<double, std::uint64_t>), 368>;
  *
  */
 template <typename T, typename IdxT>
-void serialize(raft::device_resources const& handle,
-               const std::string& filename,
-               const index<T, IdxT>& index_)
+void serialize(raft::device_resources const& handle, std::ostream& os, const index<T, IdxT>& index_)
 {
-  std::ofstream of(filename, std::ios::out | std::ios::binary);
-  if (!of) { RAFT_FAIL("Cannot open %s", filename.c_str()); }
-
   RAFT_LOG_DEBUG(
     "Saving IVF-Flat index, size %zu, dim %u", static_cast<size_t>(index_.size()), index_.dim());
 
-  serialize_scalar(handle, of, serialization_version);
-  serialize_scalar(handle, of, index_.size());
-  serialize_scalar(handle, of, index_.dim());
-  serialize_scalar(handle, of, index_.n_lists());
-  serialize_scalar(handle, of, index_.metric());
-  serialize_scalar(handle, of, index_.adaptive_centers());
-  serialize_scalar(handle, of, index_.conservative_memory_allocation());
-  serialize_mdspan(handle, of, index_.centers());
+  serialize_scalar(handle, os, serialization_version);
+  serialize_scalar(handle, os, index_.size());
+  serialize_scalar(handle, os, index_.dim());
+  serialize_scalar(handle, os, index_.n_lists());
+  serialize_scalar(handle, os, index_.metric());
+  serialize_scalar(handle, os, index_.adaptive_centers());
+  serialize_scalar(handle, os, index_.conservative_memory_allocation());
+  serialize_mdspan(handle, os, index_.centers());
   if (index_.center_norms()) {
     bool has_norms = true;
-    serialize_scalar(handle, of, has_norms);
-    serialize_mdspan(handle, of, *index_.center_norms());
+    serialize_scalar(handle, os, has_norms);
+    serialize_mdspan(handle, os, *index_.center_norms());
   } else {
     bool has_norms = false;
-    serialize_scalar(handle, of, has_norms);
+    serialize_scalar(handle, os, has_norms);
   }
   auto sizes_host = make_host_vector<uint32_t, uint32_t>(index_.list_sizes().extent(0));
   copy(sizes_host.data_handle(),
@@ -87,13 +82,25 @@ void serialize(raft::device_resources const& handle,
        sizes_host.size(),
        handle.get_stream());
   handle.sync_stream();
-  serialize_mdspan(handle, of, sizes_host.view());
+  serialize_mdspan(handle, os, sizes_host.view());
 
   list_spec<uint32_t, T, IdxT> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::serialize_list(handle, of, index_.lists()[label], list_store_spec, sizes_host(label));
+    ivf::serialize_list(handle, os, index_.lists()[label], list_store_spec, sizes_host(label));
   }
   handle.sync_stream();
+}
+
+template <typename T, typename IdxT>
+void serialize(raft::device_resources const& handle,
+               const std::string& filename,
+               const index<T, IdxT>& index_)
+{
+  std::ofstream of(filename, std::ios::out | std::ios::binary);
+  if (!of) { RAFT_FAIL("Cannot open file %s", filename.c_str()); }
+
+  detail::serialize(handle, of, index_);
+
   of.close();
   if (!of) { RAFT_FAIL("Error writing output %s", filename.c_str()); }
 }
@@ -108,49 +115,58 @@ void serialize(raft::device_resources const& handle,
  *
  */
 template <typename T, typename IdxT>
-auto deserialize(raft::device_resources const& handle, const std::string& filename)
-  -> index<T, IdxT>
+auto deserialize(raft::device_resources const& handle, std::istream& is) -> index<T, IdxT>
 {
-  std::ifstream infile(filename, std::ios::in | std::ios::binary);
-
-  if (!infile) { RAFT_FAIL("Cannot open %s", filename.c_str()); }
-
-  auto ver = deserialize_scalar<int>(handle, infile);
+  auto ver = deserialize_scalar<int>(handle, is);
   if (ver != serialization_version) {
     RAFT_FAIL("serialization version mismatch, expected %d, got %d ", serialization_version, ver);
   }
-  auto n_rows           = deserialize_scalar<IdxT>(handle, infile);
-  auto dim              = deserialize_scalar<std::uint32_t>(handle, infile);
-  auto n_lists          = deserialize_scalar<std::uint32_t>(handle, infile);
-  auto metric           = deserialize_scalar<raft::distance::DistanceType>(handle, infile);
-  bool adaptive_centers = deserialize_scalar<bool>(handle, infile);
-  bool cma              = deserialize_scalar<bool>(handle, infile);
+  auto n_rows           = deserialize_scalar<IdxT>(handle, is);
+  auto dim              = deserialize_scalar<std::uint32_t>(handle, is);
+  auto n_lists          = deserialize_scalar<std::uint32_t>(handle, is);
+  auto metric           = deserialize_scalar<raft::distance::DistanceType>(handle, is);
+  bool adaptive_centers = deserialize_scalar<bool>(handle, is);
+  bool cma              = deserialize_scalar<bool>(handle, is);
 
   index<T, IdxT> index_ = index<T, IdxT>(handle, metric, n_lists, adaptive_centers, cma, dim);
 
-  deserialize_mdspan(handle, infile, index_.centers());
-  bool has_norms = deserialize_scalar<bool>(handle, infile);
+  deserialize_mdspan(handle, is, index_.centers());
+  bool has_norms = deserialize_scalar<bool>(handle, is);
   if (has_norms) {
     index_.allocate_center_norms(handle);
     if (!index_.center_norms()) {
       RAFT_FAIL("Error inconsistent center norms");
     } else {
       auto center_norms = index_.center_norms().value();
-      deserialize_mdspan(handle, infile, center_norms);
+      deserialize_mdspan(handle, is, center_norms);
     }
   }
-  deserialize_mdspan(handle, infile, index_.list_sizes());
+  deserialize_mdspan(handle, is, index_.list_sizes());
 
   list_spec<uint32_t, T, IdxT> list_device_spec{index_.dim(), cma};
   list_spec<uint32_t, T, IdxT> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::deserialize_list(handle, infile, index_.lists()[label], list_store_spec, list_device_spec);
+    ivf::deserialize_list(handle, is, index_.lists()[label], list_store_spec, list_device_spec);
   }
   handle.sync_stream();
-  infile.close();
 
   index_.recompute_internal_state(handle);
 
   return index_;
 }
+
+template <typename T, typename IdxT>
+auto deserialize(raft::device_resources const& handle, const std::string& filename)
+  -> index<T, IdxT>
+{
+  std::ifstream is(filename, std::ios::in | std::ios::binary);
+
+  if (!is) { RAFT_FAIL("Cannot open file %s", filename.c_str()); }
+
+  auto index = detail::deserialize<T, IdxT>(handle, is);
+
+  is.close();
+
+  return index;
+}
 }  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh
index 4a5e2a3641..f42bfe66c7 100644
--- a/cpp/include/raft/neighbors/ivf_flat.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat.cuh
@@ -18,7 +18,7 @@
 
 #include <raft/neighbors/detail/ivf_flat_build.cuh>
 #include <raft/neighbors/detail/ivf_flat_search.cuh>
-#include <raft/neighbors/detail/ivf_flat_serialize.cuh>
+#include <raft/neighbors/ivf_flat_serialize.cuh>
 #include <raft/neighbors/ivf_flat_types.hpp>
 
 #include <raft/core/device_resources.hpp>
@@ -100,7 +100,6 @@ auto build(raft::device_resources const& handle,
  * @tparam value_t data element type
  * @tparam idx_t type of the indices in the source dataset
  * @tparam int_t precision / type of integral arguments
- * @tparam matrix_idx_t matrix indexing type
  *
  * @param[in] handle
  * @param[in] params configure the index building
@@ -191,8 +190,6 @@ auto extend(raft::device_resources const& handle,
  *
  * @tparam value_t data element type
  * @tparam idx_t type of the indices in the source dataset
- * @tparam int_t precision / type of integral arguments
- * @tparam matrix_idx_t matrix indexing type
  *
  * @param[in] handle
  * @param[in] orig_index original index
@@ -278,8 +275,6 @@ void extend(raft::device_resources const& handle,
  *
  * @tparam value_t data element type
  * @tparam idx_t type of the indices in the source dataset
- * @tparam int_t precision / type of integral arguments
- * @tparam matrix_idx_t matrix indexing type
  *
  * @param[in] handle
  * @param[inout] index
@@ -389,7 +384,6 @@ void search(raft::device_resources const& handle,
  * @tparam value_t data element type
  * @tparam idx_t type of the indices
  * @tparam int_t precision / type of integral arguments
- * @tparam matrix_idx_t matrix indexing type
  *
  * @param[in] handle
  * @param[in] index ivf-flat constructed index
diff --git a/cpp/include/raft/neighbors/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/ivf_flat_serialize.cuh
new file mode 100644
index 0000000000..d2ec9a39bd
--- /dev/null
+++ b/cpp/include/raft/neighbors/ivf_flat_serialize.cuh
@@ -0,0 +1,156 @@
+/*
+ * Copyright (c) 2023, 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 "detail/ivf_flat_serialize.cuh"
+
+namespace raft::neighbors::ivf_flat {
+
+/**
+ * \defgroup ivf_flat_serialize IVF-Flat Serialize
+ * @{
+ */
+
+/**
+ * Write the index to an output stream
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @code{.cpp}
+ * #include <raft/core/device_resources.hpp>
+ *
+ * raft::device_resources handle;
+ *
+ * // create an output stream
+ * std::ostream os(std::cout.rdbuf());
+ * // create an index with `auto index = ivf_flat::build(...);`
+ * raft::serialize(handle, os, index);
+ * @endcode
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle the raft handle
+ * @param[in] os output stream
+ * @param[in] index IVF-Flat index
+ *
+ * @return raft::neighbors::ivf_flat::index<T, IdxT>
+ */
+template <typename T, typename IdxT>
+void serialize(raft::device_resources const& handle, std::ostream& os, const index<T, IdxT>& index)
+{
+  detail::serialize(handle, os, index);
+}
+
+/**
+ * Save the index to file.
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @code{.cpp}
+ * #include <raft/core/device_resources.hpp>
+ *
+ * raft::device_resources handle;
+ *
+ * // create a string with a filepath
+ * std::string filename("/path/to/index");
+ * // create an index with `auto index = ivf_flat::build(...);`
+ * raft::serialize(handle, filename, index);
+ * @endcode
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle the raft handle
+ * @param[in] filename the file name for saving the index
+ * @param[in] index IVF-Flat index
+ *
+ * @return raft::neighbors::ivf_flat::index<T, IdxT>
+ */
+template <typename T, typename IdxT>
+void serialize(raft::device_resources const& handle,
+               const std::string& filename,
+               const index<T, IdxT>& index)
+{
+  detail::serialize(handle, filename, index);
+}
+
+/**
+ * Load index from input stream
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @code{.cpp}
+ * #include <raft/core/device_resources.hpp>
+ *
+ * raft::device_resources handle;
+ *
+ * // create an input stream
+ * std::istream is(std::cin.rdbuf());
+ * using T    = float; // data element type
+ * using IdxT = int; // type of the index
+ * auto index = raft::deserialize<T, IdxT>(handle, is);
+ * @endcode
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle the raft handle
+ * @param[in] is input stream
+ *
+ * @return raft::neighbors::ivf_flat::index<T, IdxT>
+ */
+template <typename T, typename IdxT>
+index<T, IdxT> deserialize(raft::device_resources const& handle, std::istream& is)
+{
+  return detail::deserialize<T, IdxT>(handle, is);
+}
+
+/**
+ * Load index from file.
+ *
+ * Experimental, both the API and the serialization format are subject to change.
+ *
+ * @code{.cpp}
+ * #include <raft/core/device_resources.hpp>
+ *
+ * raft::device_resources handle;
+ *
+ * // create a string with a filepath
+ * std::string filename("/path/to/index");
+ * using T    = float; // data element type
+ * using IdxT = int; // type of the index
+ * auto index = raft::deserialize<T, IdxT>(handle, filename);
+ * @endcode
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle the raft handle
+ * @param[in] filename the name of the file that stores the index
+ *
+ * @return raft::neighbors::ivf_flat::index<T, IdxT>
+ */
+template <typename T, typename IdxT>
+index<T, IdxT> deserialize(raft::device_resources const& handle, const std::string& filename)
+{
+  return detail::deserialize<T, IdxT>(handle, filename);
+}
+
+/**@}*/
+
+}  // namespace raft::neighbors::ivf_flat
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index 34531ddb89..bc4de89319 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -230,6 +230,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
             list_indices.data(), index_2.inds_ptrs().data_handle(), index_2.n_lists(), stream_);
           handle_.sync_stream(stream_);
           for (uint32_t l = 0; l < index_2.n_lists(); l++) {
+            if (list_sizes[l] == 0) continue;
             rmm::device_uvector<float> cluster_data(list_sizes[l] * ps.dim, stream_);
             raft::spatial::knn::detail::utils::copy_selected<float>((IdxT)list_sizes[l],
                                                                     (IdxT)ps.dim,

From 85b4aa11b694e09430e31f9d62dc9cb6f3eff7c2 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Fri, 10 Mar 2023 03:05:45 +0100
Subject: [PATCH 22/33] Fix list size override

---
 cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
index 18e8c2862e..dabd06fa89 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_serialize.cuh
@@ -86,7 +86,11 @@ void serialize(raft::device_resources const& handle, std::ostream& os, const ind
 
   list_spec<uint32_t, T, IdxT> list_store_spec{index_.dim(), true};
   for (uint32_t label = 0; label < index_.n_lists(); label++) {
-    ivf::serialize_list(handle, os, index_.lists()[label], list_store_spec, sizes_host(label));
+    ivf::serialize_list(handle,
+                        os,
+                        index_.lists()[label],
+                        list_store_spec,
+                        Pow2<kIndexGroupSize>::roundUp(sizes_host(label)));
   }
   handle.sync_stream();
 }

From 692af0ddd685059d1390442c8ec4cc9401ad783c Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Fri, 10 Mar 2023 16:20:38 +0100
Subject: [PATCH 23/33] Fix resizelist with interleaved format

---
 cpp/include/raft/neighbors/detail/ivf_flat_build.cuh | 7 +++++--
 1 file changed, 5 insertions(+), 2 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 7c6ca7cb2d..c9b26847b6 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -229,8 +229,11 @@ void extend(raft::device_resources const& handle,
     handle.sync_stream();
     auto& lists = index->lists();
     for (uint32_t label = 0; label < n_lists; label++) {
-      ivf::resize_list(
-        handle, lists[label], list_device_spec, new_list_sizes[label], old_list_sizes[label]);
+      ivf::resize_list(handle,
+                       lists[label],
+                       list_device_spec,
+                       new_list_sizes[label],
+                       Pow2<kIndexGroupSize>::roundUp(old_list_sizes[label]));
     }
   }
   // Update the pointers and the sizes

From a8b96a7d51d543005e3505b4dd9f352a1eb4a974 Mon Sep 17 00:00:00 2001
From: Mickael Ide <mide@nvidia.com>
Date: Sun, 12 Mar 2023 20:35:14 +0100
Subject: [PATCH 24/33] Fix refine list resize operation

---
 cpp/include/raft/neighbors/detail/ivf_flat_build.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index c9b26847b6..bf7248b983 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -391,7 +391,7 @@ inline void fill_refinement_index(raft::device_resources const& handle,
   // We do not fill centers and center norms, since we will not run coarse search.
 
   // Allocate new memory
-  auto lists = refinement_index->lists();
+  auto& lists = refinement_index->lists();
   list_spec<uint32_t, T, IdxT> list_device_spec{refinement_index->dim(), false};
   for (uint32_t label = 0; label < n_lists; label++) {
     ivf::resize_list(handle, lists[label], list_device_spec, n_candidates, uint32_t(0));

From beb9264f2c4fbe6c5c69278e67b89be1985cb899 Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Tue, 14 Mar 2023 19:18:29 +0100
Subject: [PATCH 25/33] IVF-Flat reconstruction

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  | 101 ++++++++++++++++++
 cpp/include/raft/neighbors/ivf_flat.cuh       |  25 +++++
 cpp/test/neighbors/ann_ivf_flat.cuh           |  32 ++++++
 3 files changed, 158 insertions(+)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index bf7248b983..82b6c55bbb 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -34,6 +34,8 @@
 
 #include <rmm/cuda_stream_view.hpp>
 
+#include <thrust/extrema.h>
+
 #include <cstdint>
 
 namespace raft::neighbors::ivf_flat::detail {
@@ -415,4 +417,103 @@ inline void fill_refinement_index(raft::device_resources const& handle,
                                          refinement_index->veclen());
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }
+
+template <typename T, typename IdxT>
+__global__ void get_data_ptr_kernel(const uint32_t* list_sizes,
+                                    const T* const* list_data_ptrs,
+                                    const IdxT* const* list_indices_ptrs,
+                                    uint32_t dim,
+                                    uint32_t veclen,
+                                    uint32_t n_list,
+                                    IdxT max_indice,
+                                    T** ptrs_to_data)
+{
+  const IdxT list_id = IdxT(blockDim.x) * IdxT(blockIdx.x) + threadIdx.x;
+  if (list_id >= n_list) { return; }
+  const IdxT inlist_id     = IdxT(blockDim.y) * IdxT(blockIdx.y) + threadIdx.y;
+  const uint32_t list_size = list_sizes[list_id];
+  if (inlist_id >= list_size) { return; }
+
+  auto* list_indices = list_indices_ptrs[list_id];
+  IdxT id            = list_indices[inlist_id];
+  if (id > max_indice) { return; }
+
+  using interleaved_group = Pow2<kIndexGroupSize>;
+  auto group_offset       = interleaved_group::roundDown(inlist_id);
+  auto ingroup_id         = interleaved_group::mod(inlist_id) * veclen;
+
+  auto* list_data  = list_data_ptrs[list_id];
+  const T* ptr     = list_data + (group_offset * dim) + ingroup_id;
+  ptrs_to_data[id] = (T*)ptr;
+}
+
+template <typename T, typename IdxT>
+__global__ void reconstruct_batch_kernel(const IdxT* vector_ids,
+                                         const T** ptrs_to_data,
+                                         uint32_t dim,
+                                         uint32_t veclen,
+                                         IdxT n_rows,
+                                         T* reconstr)
+{
+  const IdxT i = IdxT(blockDim.x) * IdxT(blockIdx.x) + threadIdx.x;
+  if (i >= n_rows) { return; }
+
+  const IdxT vector_id = vector_ids[i];
+  const T* src         = ptrs_to_data[vector_id];
+  if (!src) { return; }
+
+  reconstr += i * dim;
+  for (uint32_t l = 0; l < dim; l += veclen) {
+    for (uint32_t j = 0; j < veclen; j++) {
+      reconstr[l + j] = src[l * kIndexGroupSize + j];
+    }
+  }
+}
+
+template <typename T, typename IdxT>
+void reconstruct_batch(raft::device_resources const& handle,
+                       const index<T, IdxT>& index,
+                       device_mdspan<const IdxT, extent_1d<IdxT>, row_major> vector_ids,
+                       device_mdspan<T, extent_2d<IdxT>, row_major> vector_out)
+{
+  thrust::device_ptr<const IdxT> vector_ids_ptr =
+    thrust::device_pointer_cast(vector_ids.data_handle());
+  IdxT max_indice = *thrust::max_element(
+    handle.get_thrust_policy(), vector_ids_ptr, vector_ids_ptr + vector_ids.extent(0));
+
+  rmm::device_uvector<T*> ptrs_to_data(max_indice + 1, handle.get_stream());
+  utils::memzero(ptrs_to_data.data(), ptrs_to_data.size(), handle.get_stream());
+
+  thrust::device_ptr<const uint32_t> list_sizes_ptr =
+    thrust::device_pointer_cast(index.list_sizes().data_handle());
+  uint32_t max_list_size = *thrust::max_element(
+    handle.get_thrust_policy(), list_sizes_ptr, list_sizes_ptr + index.list_sizes().extent(0));
+
+  const dim3 block_dim1(16, 16);
+  const dim3 grid_dim1(raft::ceildiv<size_t>(index.n_lists(), block_dim1.x),
+                       raft::ceildiv<size_t>(max_list_size, block_dim1.y));
+  get_data_ptr_kernel<<<grid_dim1, block_dim1, 0, handle.get_stream()>>>(
+    index.list_sizes().data_handle(),
+    index.data_ptrs().data_handle(),
+    index.inds_ptrs().data_handle(),
+    index.dim(),
+    index.veclen(),
+    index.n_lists(),
+    max_indice,
+    ptrs_to_data.data());
+  RAFT_CUDA_TRY(cudaPeekAtLastError());
+
+  auto n_reconstruction = vector_ids.extent(0);
+  const dim3 block_dim2(256);
+  const dim3 grid_dim2(raft::ceildiv<size_t>(n_reconstruction, block_dim2.x));
+  reconstruct_batch_kernel<<<grid_dim2, block_dim2, 0, handle.get_stream()>>>(
+    vector_ids.data_handle(),
+    (const T**)ptrs_to_data.data(),
+    index.dim(),
+    index.veclen(),
+    n_reconstruction,
+    vector_out.data_handle());
+  RAFT_CUDA_TRY(cudaPeekAtLastError());
+}
+
 }  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/include/raft/neighbors/ivf_flat.cuh b/cpp/include/raft/neighbors/ivf_flat.cuh
index f42bfe66c7..40080e060d 100644
--- a/cpp/include/raft/neighbors/ivf_flat.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat.cuh
@@ -425,6 +425,31 @@ void search(raft::device_resources const& handle,
                 nullptr);
 }
 
+/**
+ * @brief Reconstruction of the index
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices in the source dataset
+ *
+ * @param[in] handle
+ * @param[in] idx
+ * @param[in] vector_ids a mdspan array containing the ids of vectors to be reconstructed [n_rows]
+ * @param[out] vector_out a mdspan array to hold the vectors to be reconstructed [n_rows,
+ * index.dim()]
+ */
+template <typename T, typename IdxT>
+void reconstruct_batch(raft::device_resources const& handle,
+                       const index<T, IdxT>& idx,
+                       device_mdspan<const IdxT, extent_1d<IdxT>, row_major> vector_ids,
+                       device_mdspan<T, extent_2d<IdxT>, row_major> vector_out)
+{
+  RAFT_EXPECTS(vector_ids.extent(0) == vector_out.extent(0),
+               "row index inputs and vectors outputs should have the same length");
+  RAFT_EXPECTS(vector_out.extent(1) == idx.dim(),
+               "vectors outputs should have the right dimension");
+  return raft::neighbors::ivf_flat::detail::reconstruct_batch(handle, idx, vector_ids, vector_out);
+}
+
 /** @} */
 
 }  // namespace raft::neighbors::ivf_flat
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index eb0f84d104..78c1b409b6 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -263,6 +263,38 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
                                   0.001,
                                   min_recall));
     }
+
+    {
+      raft::spatial::knn::ivf_flat::index_params index_params;
+      index_params.n_lists                  = ps.nlist;
+      index_params.metric                   = ps.metric;
+      index_params.adaptive_centers         = ps.adaptive_centers;
+      index_params.add_data_on_build        = true;
+      index_params.kmeans_trainset_fraction = 0.5;
+
+      auto vectors_data =
+        raft::make_device_matrix_view<const DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
+      auto index = ivf_flat::build(handle_, vectors_data, index_params);
+
+      rmm::device_uvector<IdxT> vecs_ids(ps.num_db_vecs, stream_);
+      thrust::sequence(handle_.get_thrust_policy(),
+                       thrust::device_pointer_cast(vecs_ids.data()),
+                       thrust::device_pointer_cast(vecs_ids.data() + ps.num_db_vecs));
+      handle_.sync_stream(stream_);
+
+      auto vectors_ids =
+        raft::make_device_vector_view<const IdxT, IdxT>(vecs_ids.data(), ps.num_db_vecs);
+      auto vectors_out =
+        raft::make_device_matrix<DataT, IdxT, row_major>(handle_, ps.num_db_vecs, ps.dim);
+      ivf_flat::reconstruct_batch(handle_, index, vectors_ids, vectors_out.view());
+      handle_.sync_stream(stream_);
+
+      ASSERT_TRUE(raft::devArrMatch(vectors_data.data_handle(),
+                                    vectors_out.data_handle(),
+                                    ps.num_db_vecs * ps.dim,
+                                    raft::Compare<DataT>(),
+                                    stream_));
+    }
   }
 
   void SetUp() override

From f7986575e7e8d5068b28adf6b398c2cf386a1fed Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Tue, 23 May 2023 15:17:37 +0200
Subject: [PATCH 26/33] addressing review + expose function

---
 cpp/CMakeLists.txt                            | 11 ++-
 cpp/bench/ann/src/faiss/faiss_benchmark.cu    |  2 +-
 .../raft/neighbors/detail/ivf_flat_build.cuh  | 51 ++++++-------
 cpp/include/raft/neighbors/ivf_flat-ext.cuh   | 19 +++++
 cpp/include/raft/neighbors/ivf_flat-inl.cuh   | 22 ++++++
 .../raft_runtime/neighbors/ivf_flat.hpp       |  5 ++
 .../ivf_flat_reconstruct_float_int64_t.cu     | 36 ++++++++++
 .../ivf_flat_reconstruct_int8_t_int64_t.cu    | 36 ++++++++++
 .../ivf_flat_reconstruct_uint8_t_int64_t.cu   | 36 ++++++++++
 .../raft_runtime/neighbors/ivf_flat_build.cu  | 72 ++++++++++---------
 cpp/test/neighbors/ann_ivf_flat.cuh           | 10 +--
 11 files changed, 237 insertions(+), 63 deletions(-)
 create mode 100644 cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
 create mode 100644 cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
 create mode 100644 cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu

diff --git a/cpp/CMakeLists.txt b/cpp/CMakeLists.txt
index eb35554768..b849eafd53 100644
--- a/cpp/CMakeLists.txt
+++ b/cpp/CMakeLists.txt
@@ -337,6 +337,9 @@ if(RAFT_COMPILE_LIBRARY)
     src/neighbors/ivf_flat_search_float_int64_t.cu
     src/neighbors/ivf_flat_search_int8_t_int64_t.cu
     src/neighbors/ivf_flat_search_uint8_t_int64_t.cu
+    src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
+    src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
+    src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
     src/neighbors/ivfpq_build_float_int64_t.cu
     src/neighbors/ivfpq_build_int8_t_int64_t.cu
     src/neighbors/ivfpq_build_uint8_t_int64_t.cu
@@ -602,7 +605,9 @@ target_link_libraries(raft::raft INTERFACE
 # Use `rapids_export` for 22.04 as it will have COMPONENT support
 rapids_export(
   INSTALL raft
-  EXPORT_SET raft-exports COMPONENTS ${raft_components} COMPONENTS_EXPORT_SET ${raft_export_sets}
+  EXPORT_SET raft-exports
+  COMPONENTS ${raft_components}
+  COMPONENTS_EXPORT_SET ${raft_export_sets}
   GLOBAL_TARGETS raft compiled distributed
   NAMESPACE raft::
   DOCUMENTATION doc_string
@@ -613,7 +618,9 @@ rapids_export(
 # * build export -------------------------------------------------------------
 rapids_export(
   BUILD raft
-  EXPORT_SET raft-exports COMPONENTS ${raft_components} COMPONENTS_EXPORT_SET ${raft_export_sets}
+  EXPORT_SET raft-exports
+  COMPONENTS ${raft_components}
+  COMPONENTS_EXPORT_SET ${raft_export_sets}
   GLOBAL_TARGETS raft compiled distributed
   DOCUMENTATION doc_string
   NAMESPACE raft::
diff --git a/cpp/bench/ann/src/faiss/faiss_benchmark.cu b/cpp/bench/ann/src/faiss/faiss_benchmark.cu
index 294da9a14f..60269f72a5 100644
--- a/cpp/bench/ann/src/faiss/faiss_benchmark.cu
+++ b/cpp/bench/ann/src/faiss/faiss_benchmark.cu
@@ -104,10 +104,10 @@ std::unique_ptr<raft::bench::ann::ANN<T>> create_algo(const std::string& algo,
   // stop compiler warning; not all algorithms support multi-GPU so it may not be used
   (void)dev_list;
 
-  raft::bench::ann::Metric metric = parse_metric(distance);
   std::unique_ptr<raft::bench::ann::ANN<T>> ann;
 
   if constexpr (std::is_same_v<T, float>) {
+    raft::bench::ann::Metric metric = parse_metric(distance);
     if (algo == "faiss_gpu_ivf_flat") {
       ann = make_algo<T, raft::bench::ann::FaissGpuIVFFlat>(metric, dim, conf, dev_list);
     } else if (algo == "faiss_gpu_ivf_pq") {
diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index e1e8aecc8d..506a2e81a5 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -17,11 +17,13 @@
 #pragma once
 
 #include <raft/cluster/kmeans_balanced.cuh>
+#include <raft/core/device_mdspan.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/core/mdarray.hpp>
 #include <raft/core/nvtx.hpp>
 #include <raft/core/operators.hpp>
 #include <raft/core/resource/cuda_stream.hpp>
+#include <raft/core/resource/thrust_policy.hpp>
 #include <raft/core/resources.hpp>
 #include <raft/linalg/add.cuh>
 #include <raft/linalg/map.cuh>
@@ -472,48 +474,49 @@ __global__ void reconstruct_batch_kernel(const IdxT* vector_ids,
 }
 
 template <typename T, typename IdxT>
-void reconstruct_batch(raft::device_resources const& handle,
+void reconstruct_batch(raft::resources const& handle,
                        const index<T, IdxT>& index,
-                       device_mdspan<const IdxT, extent_1d<IdxT>, row_major> vector_ids,
-                       device_mdspan<T, extent_2d<IdxT>, row_major> vector_out)
+                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
+                       raft::device_matrix_view<T, IdxT, row_major> vector_out)
 {
+  auto stream      = raft::resource::get_cuda_stream(handle);
+  auto exec_policy = raft::resource::get_thrust_policy(handle);
+
   thrust::device_ptr<const IdxT> vector_ids_ptr =
     thrust::device_pointer_cast(vector_ids.data_handle());
-  IdxT max_indice = *thrust::max_element(
-    handle.get_thrust_policy(), vector_ids_ptr, vector_ids_ptr + vector_ids.extent(0));
+  IdxT max_indice =
+    *thrust::max_element(exec_policy, vector_ids_ptr, vector_ids_ptr + vector_ids.extent(0));
 
-  rmm::device_uvector<T*> ptrs_to_data(max_indice + 1, handle.get_stream());
-  utils::memzero(ptrs_to_data.data(), ptrs_to_data.size(), handle.get_stream());
+  rmm::device_uvector<T*> ptrs_to_data(max_indice + 1, stream);
+  utils::memzero(ptrs_to_data.data(), ptrs_to_data.size(), stream);
 
   thrust::device_ptr<const uint32_t> list_sizes_ptr =
     thrust::device_pointer_cast(index.list_sizes().data_handle());
   uint32_t max_list_size = *thrust::max_element(
-    handle.get_thrust_policy(), list_sizes_ptr, list_sizes_ptr + index.list_sizes().extent(0));
+    exec_policy, list_sizes_ptr, list_sizes_ptr + index.list_sizes().extent(0));
 
   const dim3 block_dim1(16, 16);
   const dim3 grid_dim1(raft::ceildiv<size_t>(index.n_lists(), block_dim1.x),
                        raft::ceildiv<size_t>(max_list_size, block_dim1.y));
-  get_data_ptr_kernel<<<grid_dim1, block_dim1, 0, handle.get_stream()>>>(
-    index.list_sizes().data_handle(),
-    index.data_ptrs().data_handle(),
-    index.inds_ptrs().data_handle(),
-    index.dim(),
-    index.veclen(),
-    index.n_lists(),
-    max_indice,
-    ptrs_to_data.data());
+  get_data_ptr_kernel<<<grid_dim1, block_dim1, 0, stream>>>(index.list_sizes().data_handle(),
+                                                            index.data_ptrs().data_handle(),
+                                                            index.inds_ptrs().data_handle(),
+                                                            index.dim(),
+                                                            index.veclen(),
+                                                            index.n_lists(),
+                                                            max_indice,
+                                                            ptrs_to_data.data());
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 
   auto n_reconstruction = vector_ids.extent(0);
   const dim3 block_dim2(256);
   const dim3 grid_dim2(raft::ceildiv<size_t>(n_reconstruction, block_dim2.x));
-  reconstruct_batch_kernel<<<grid_dim2, block_dim2, 0, handle.get_stream()>>>(
-    vector_ids.data_handle(),
-    (const T**)ptrs_to_data.data(),
-    index.dim(),
-    index.veclen(),
-    n_reconstruction,
-    vector_out.data_handle());
+  reconstruct_batch_kernel<<<grid_dim2, block_dim2, 0, stream>>>(vector_ids.data_handle(),
+                                                                 (const T**)ptrs_to_data.data(),
+                                                                 index.dim(),
+                                                                 index.veclen(),
+                                                                 n_reconstruction,
+                                                                 vector_out.data_handle());
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }
 
diff --git a/cpp/include/raft/neighbors/ivf_flat-ext.cuh b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
index dff7b6b2ab..acead6f7dc 100644
--- a/cpp/include/raft/neighbors/ivf_flat-ext.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
@@ -93,6 +93,12 @@ void search(raft::resources const& handle,
             raft::device_matrix_view<IdxT, IdxT, row_major> neighbors,
             raft::device_matrix_view<float, IdxT, row_major> distances) RAFT_EXPLICIT;
 
+template <typename T, typename IdxT>
+void reconstruct_batch(raft::resources const& handle,
+                       const index<T, IdxT>& index,
+                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
+                       raft::device_matrix_view<T, IdxT, row_major> vector_out) RAFT_EXPLICIT;
+
 }  // namespace raft::neighbors::ivf_flat
 
 #endif  // RAFT_EXPLICIT_INSTANTIATE_ONLY
@@ -183,3 +189,16 @@ instantiate_raft_neighbors_ivf_flat_search(int8_t, int64_t);
 instantiate_raft_neighbors_ivf_flat_search(uint8_t, int64_t);
 
 #undef instantiate_raft_neighbors_ivf_flat_search
+
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)              \
+  extern template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
+    raft::resources const& handle,                                            \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& index,                   \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,                    \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+
+instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
+instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
+instantiate_raft_neighbors_ivf_flat_reconstruct(uint8_t, int64_t);
+
+#undef instantiate_raft_neighbors_ivf_flat_reconstruct
diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
index 739e012e08..5aa829eadd 100644
--- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
@@ -464,6 +464,28 @@ void search(raft::resources const& handle,
                 nullptr);
 }
 
+/**
+ * @brief Reconstruct vector database
+ *
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle
+ * @param[in] index ivf-flat constructed index
+ * @param[in] vector_ids vector with the ids of vectors to look for
+ * @param[out] vector_out matrix with the database of vectors
+ */
+template <typename T, typename IdxT>
+void reconstruct_batch(raft::resources const& handle,
+                       const index<T, IdxT>& index,
+                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
+                       raft::device_matrix_view<T, IdxT, row_major> vector_out)
+{
+  return raft::neighbors::ivf_flat::detail::reconstruct_batch(
+    handle, index, vector_ids, vector_out);
+}
+
 /** @} */
 
 }  // namespace raft::neighbors::ivf_flat
diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
index 5b8918ec7f..e61e962a1d 100644
--- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
+++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
@@ -46,6 +46,11 @@ namespace raft::runtime::neighbors::ivf_flat {
               std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, \
               raft::neighbors::ivf_flat::index<T, IdxT>* idx);                       \
                                                                                      \
+  void reconstruct_batch(raft::resources const& handle,                              \
+                         raft::neighbors::ivf_flat::index<T, IdxT>& idx,             \
+                         raft::device_vector_view<const IdxT, IdxT> vector_ids,      \
+                         raft::device_matrix_view<T, IdxT, row_major> vector_out);   \
+                                                                                     \
   void serialize_file(raft::resources const& handle,                                 \
                       const std::string& filename,                                   \
                       const raft::neighbors::ivf_flat::index<T, IdxT>& index);       \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
new file mode 100644
index 0000000000..14b4003bc7
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2023, 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 generated by ivf_flat_00_generate.py
+ *
+ * Make changes there and run in this directory:
+ *
+ * > python ivf_flat_00_generate.py
+ *
+ */
+
+#include <raft/neighbors/ivf_flat-inl.cuh>
+
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)       \
+  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
+    raft::resources const& handle,                                     \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,             \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
+
+#undef instantiate_raft_neighbors_ivf_flat_reconstruct
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
new file mode 100644
index 0000000000..3b5093dcd0
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2023, 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 generated by ivf_flat_00_generate.py
+ *
+ * Make changes there and run in this directory:
+ *
+ * > python ivf_flat_00_generate.py
+ *
+ */
+
+#include <raft/neighbors/ivf_flat-inl.cuh>
+
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)       \
+  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
+    raft::resources const& handle,                                     \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,             \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
+
+#undef instantiate_raft_neighbors_ivf_flat_reconstruct
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
new file mode 100644
index 0000000000..2b4ebb1aa8
--- /dev/null
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
@@ -0,0 +1,36 @@
+/*
+ * Copyright (c) 2023, 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 generated by ivf_flat_00_generate.py
+ *
+ * Make changes there and run in this directory:
+ *
+ * > python ivf_flat_00_generate.py
+ *
+ */
+
+#include <raft/neighbors/ivf_flat-inl.cuh>
+
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)       \
+  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
+    raft::resources const& handle,                                     \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,             \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+instantiate_raft_neighbors_ivf_flat_reconstruct(uint8_t, int64_t);
+
+#undef instantiate_raft_neighbors_ivf_flat_reconstruct
diff --git a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
index 7fccb95411..26bab4f1e7 100644
--- a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
+++ b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
@@ -19,38 +19,46 @@
 
 namespace raft::runtime::neighbors::ivf_flat {
 
-#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                                \
-  auto build(raft::resources const& handle,                                            \
-             const raft::neighbors::ivf_flat::index_params& params,                    \
-             raft::device_matrix_view<const T, IdxT, row_major> dataset)               \
-    ->raft::neighbors::ivf_flat::index<T, IdxT>                                        \
-  {                                                                                    \
-    return raft::neighbors::ivf_flat::build<T, IdxT>(handle, params, dataset);         \
-  }                                                                                    \
-  auto extend(raft::resources const& handle,                                           \
-              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,          \
-              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,   \
-              const raft::neighbors::ivf_flat::index<T, IdxT>& orig_index)             \
-    ->raft::neighbors::ivf_flat::index<T, IdxT>                                        \
-  {                                                                                    \
-    return raft::neighbors::ivf_flat::extend<T, IdxT>(                                 \
-      handle, new_vectors, new_indices, orig_index);                                   \
-  }                                                                                    \
-                                                                                       \
-  void build(raft::resources const& handle,                                            \
-             const raft::neighbors::ivf_flat::index_params& params,                    \
-             raft::device_matrix_view<const T, IdxT, row_major> dataset,               \
-             raft::neighbors::ivf_flat::index<T, IdxT>& idx)                           \
-  {                                                                                    \
-    idx = build(handle, params, dataset);                                              \
-  }                                                                                    \
-                                                                                       \
-  void extend(raft::resources const& handle,                                           \
-              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,          \
-              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,   \
-              raft::neighbors::ivf_flat::index<T, IdxT>* idx)                          \
-  {                                                                                    \
-    raft::neighbors::ivf_flat::extend<T, IdxT>(handle, new_vectors, new_indices, idx); \
+#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                                         \
+  auto build(raft::resources const& handle,                                                     \
+             const raft::neighbors::ivf_flat::index_params& params,                             \
+             raft::device_matrix_view<const T, IdxT, row_major> dataset)                        \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>                                                 \
+  {                                                                                             \
+    return raft::neighbors::ivf_flat::build<T, IdxT>(handle, params, dataset);                  \
+  }                                                                                             \
+  auto extend(raft::resources const& handle,                                                    \
+              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,                   \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,            \
+              const raft::neighbors::ivf_flat::index<T, IdxT>& orig_index)                      \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>                                                 \
+  {                                                                                             \
+    return raft::neighbors::ivf_flat::extend<T, IdxT>(                                          \
+      handle, new_vectors, new_indices, orig_index);                                            \
+  }                                                                                             \
+                                                                                                \
+  void build(raft::resources const& handle,                                                     \
+             const raft::neighbors::ivf_flat::index_params& params,                             \
+             raft::device_matrix_view<const T, IdxT, row_major> dataset,                        \
+             raft::neighbors::ivf_flat::index<T, IdxT>& idx)                                    \
+  {                                                                                             \
+    idx = build(handle, params, dataset);                                                       \
+  }                                                                                             \
+                                                                                                \
+  void extend(raft::resources const& handle,                                                    \
+              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,                   \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,            \
+              raft::neighbors::ivf_flat::index<T, IdxT>* idx)                                   \
+  {                                                                                             \
+    raft::neighbors::ivf_flat::extend<T, IdxT>(handle, new_vectors, new_indices, idx);          \
+  }                                                                                             \
+                                                                                                \
+  void reconstruct_batch(raft::resources const& handle,                                         \
+                         const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+                         raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
+                         raft::device_matrix_view<T, IdxT, row_major> vector_out)               \
+  {                                                                                             \
+    raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(handle, idx, vector_ids, vector_out); \
   }
 
 RAFT_INST_BUILD_EXTEND(float, int64_t);
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index 7a909081e4..f368d4e368 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -19,6 +19,7 @@
 #include "ann_utils.cuh"
 #include <raft/core/resource/cuda_stream.hpp>
 #include <raft/core/resource/thrust_policy.hpp>
+#include <raft/core/resources.hpp>
 
 #include <raft_internal/neighbors/naive_knn.cuh>
 
@@ -273,20 +274,21 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
 
       auto vectors_data =
         raft::make_device_matrix_view<const DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
-      auto index = ivf_flat::build(handle_, vectors_data, index_params);
+      auto index = ivf_flat::build(handle_, index_params, vectors_data);
 
       rmm::device_uvector<IdxT> vecs_ids(ps.num_db_vecs, stream_);
-      thrust::sequence(handle_.get_thrust_policy(),
+      thrust::sequence(resource::get_thrust_policy(handle_),
                        thrust::device_pointer_cast(vecs_ids.data()),
                        thrust::device_pointer_cast(vecs_ids.data() + ps.num_db_vecs));
-      handle_.sync_stream(stream_);
+      resource::sync_stream(handle_);
 
       auto vectors_ids =
         raft::make_device_vector_view<const IdxT, IdxT>(vecs_ids.data(), ps.num_db_vecs);
       auto vectors_out =
         raft::make_device_matrix<DataT, IdxT, row_major>(handle_, ps.num_db_vecs, ps.dim);
       ivf_flat::reconstruct_batch(handle_, index, vectors_ids, vectors_out.view());
-      handle_.sync_stream(stream_);
+
+      resource::sync_stream(handle_);
 
       ASSERT_TRUE(raft::devArrMatch(vectors_data.data_handle(),
                                     vectors_out.data_handle(),

From 080c7842b8c4854ed8c841059b2ecd48bbae308e Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Mon, 29 May 2023 15:48:37 +0200
Subject: [PATCH 27/33] Cluster reconstruction

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  | 48 +++++++++++++++++++
 cpp/include/raft/neighbors/ivf_flat-ext.cuh   | 26 +++++++---
 cpp/include/raft/neighbors/ivf_flat-inl.cuh   | 24 ++++++++++
 .../raft_runtime/neighbors/ivf_flat.hpp       |  6 +++
 .../ivf_flat_reconstruct_float_int64_t.cu     | 20 +++++---
 .../ivf_flat_reconstruct_int8_t_int64_t.cu    | 20 +++++---
 .../ivf_flat_reconstruct_uint8_t_int64_t.cu   | 20 +++++---
 .../raft_runtime/neighbors/ivf_flat_build.cu  | 10 ++++
 cpp/test/neighbors/ann_ivf_flat.cuh           | 34 +++++++++++++
 9 files changed, 184 insertions(+), 24 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 506a2e81a5..02175248f3 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -520,4 +520,52 @@ void reconstruct_batch(raft::resources const& handle,
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }
 
+template <typename T>
+__global__ void reconstruct_list_data_kernel(
+  T* out_vectors,
+  T* in_list_data,
+  std::variant<uint32_t, const uint32_t*> offset_or_indices,
+  uint32_t len,
+  size_t veclen,
+  uint32_t dim)
+{
+  for (uint32_t ix = threadIdx.x + blockDim.x * blockIdx.x; ix < len; ix += blockDim.x) {
+    const uint32_t src_ix = std::holds_alternative<uint32_t>(offset_or_indices)
+                              ? std::get<uint32_t>(offset_or_indices) + ix
+                              : std::get<const uint32_t*>(offset_or_indices)[ix];
+
+    using group_align         = Pow2<kIndexGroupSize>;
+    const uint32_t group_ix   = group_align::div(src_ix);
+    const uint32_t ingroup_ix = group_align::mod(src_ix) * veclen;
+
+    for (uint32_t l = 0; l < dim; l += veclen) {
+      for (uint32_t j = 0; j < veclen; j++) {
+        out_vectors[ix * dim + l + j] = in_list_data[l * kIndexGroupSize + ingroup_ix + j];
+      }
+    }
+  }
+}
+
+/** Decode the list data; see the public interface for the api and usage. */
+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::resources const& handle,
+                           const index<T, IdxT>& index,
+                           device_matrix_view<T, IdxT, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset)
+{
+  auto stream = raft::resource::get_cuda_stream(handle);
+
+  uint32_t len = out_vectors.extent(0);
+  const dim3 block_dim(256);
+  const dim3 grid_dim(raft::ceildiv<size_t>(len, block_dim.x));
+  reconstruct_list_data_kernel<<<grid_dim, block_dim, 0, stream>>>(
+    out_vectors.data_handle(),
+    index.lists()[label]->data.data_handle(),
+    offset,
+    len,
+    (size_t)index.veclen(),
+    index.dim());
+}
+
 }  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/include/raft/neighbors/ivf_flat-ext.cuh b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
index acead6f7dc..aad2d2f5e5 100644
--- a/cpp/include/raft/neighbors/ivf_flat-ext.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
@@ -99,6 +99,13 @@ void reconstruct_batch(raft::resources const& handle,
                        raft::device_vector_view<const IdxT, IdxT> vector_ids,
                        raft::device_matrix_view<T, IdxT, row_major> vector_out) RAFT_EXPLICIT;
 
+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::resources const& handle,
+                           const index<T, IdxT>& index,
+                           device_matrix_view<T, IdxT, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset) RAFT_EXPLICIT;
+
 }  // namespace raft::neighbors::ivf_flat
 
 #endif  // RAFT_EXPLICIT_INSTANTIATE_ONLY
@@ -190,12 +197,19 @@ instantiate_raft_neighbors_ivf_flat_search(uint8_t, int64_t);
 
 #undef instantiate_raft_neighbors_ivf_flat_search
 
-#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)              \
-  extern template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
-    raft::resources const& handle,                                            \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& index,                   \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,                    \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)                  \
+  extern template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
+    raft::resources const& handle,                                                \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,                        \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);                     \
+                                                                                  \
+  extern template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
+    raft::resources const& handle,                                                \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
+    raft::device_matrix_view<T, IdxT, row_major> out_vectors,                     \
+    uint32_t label,                                                               \
+    uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
 instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
index 5aa829eadd..fd322f2965 100644
--- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
@@ -486,6 +486,30 @@ void reconstruct_batch(raft::resources const& handle,
     handle, index, vector_ids, vector_out);
 }
 
+/**
+ * @brief Reconstruct vectors of a given cluster
+ *
+ *
+ * @tparam T data element type
+ * @tparam IdxT type of the indices
+ *
+ * @param[in] handle
+ * @param[in] index ivf-flat constructed index
+ * @param[out] vector_out matrix with the vectors contained in the cluster
+ * @param[in] label cluster index
+ * @param[in] offset offset
+ */
+template <typename T, typename IdxT>
+void reconstruct_list_data(raft::resources const& handle,
+                           const index<T, IdxT>& index,
+                           device_matrix_view<T, IdxT, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset)
+{
+  return raft::neighbors::ivf_flat::detail::reconstruct_list_data(
+    handle, index, out_vectors, label, offset);
+}
+
 /** @} */
 
 }  // namespace raft::neighbors::ivf_flat
diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
index e61e962a1d..74e8bad00a 100644
--- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
+++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
@@ -51,6 +51,12 @@ namespace raft::runtime::neighbors::ivf_flat {
                          raft::device_vector_view<const IdxT, IdxT> vector_ids,      \
                          raft::device_matrix_view<T, IdxT, row_major> vector_out);   \
                                                                                      \
+  void reconstruct_list_data(raft::resources const& handle,                          \
+                             const raft::neighbors::ivf_flat::index<T, IdxT>& idx,   \
+                             device_matrix_view<T, IdxT, row_major> out_vectors,     \
+                             uint32_t label,                                         \
+                             uint32_t offset);                                       \
+                                                                                     \
   void serialize_file(raft::resources const& handle,                                 \
                       const std::string& filename,                                   \
                       const raft::neighbors::ivf_flat::index<T, IdxT>& index);       \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
index 14b4003bc7..91ed15c41a 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
@@ -25,12 +25,20 @@
 
 #include <raft/neighbors/ivf_flat-inl.cuh>
 
-#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)       \
-  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
-    raft::resources const& handle,                                     \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,             \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)           \
+  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);              \
+                                                                           \
+  template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
+    uint32_t label,                                                        \
+    uint32_t offset);
+
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
 
 #undef instantiate_raft_neighbors_ivf_flat_reconstruct
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
index 3b5093dcd0..0f214e66ca 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
@@ -25,12 +25,20 @@
 
 #include <raft/neighbors/ivf_flat-inl.cuh>
 
-#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)       \
-  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
-    raft::resources const& handle,                                     \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,             \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)           \
+  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);              \
+                                                                           \
+  template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
+    uint32_t label,                                                        \
+    uint32_t offset);
+
 instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
 
 #undef instantiate_raft_neighbors_ivf_flat_reconstruct
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
index 2b4ebb1aa8..b693a1415f 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
@@ -25,12 +25,20 @@
 
 #include <raft/neighbors/ivf_flat-inl.cuh>
 
-#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)       \
-  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>( \
-    raft::resources const& handle,                                     \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,             \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);
+#define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)           \
+  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
+    raft::device_matrix_view<T, IdxT, row_major> vector_out);              \
+                                                                           \
+  template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
+    raft::resources const& handle,                                         \
+    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
+    raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
+    uint32_t label,                                                        \
+    uint32_t offset);
+
 instantiate_raft_neighbors_ivf_flat_reconstruct(uint8_t, int64_t);
 
 #undef instantiate_raft_neighbors_ivf_flat_reconstruct
diff --git a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
index 26bab4f1e7..39b2962782 100644
--- a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
+++ b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
@@ -59,6 +59,16 @@ namespace raft::runtime::neighbors::ivf_flat {
                          raft::device_matrix_view<T, IdxT, row_major> vector_out)               \
   {                                                                                             \
     raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(handle, idx, vector_ids, vector_out); \
+  }                                                                                             \
+                                                                                                \
+  void reconstruct_list_data(raft::resources const& handle,                                     \
+                             const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
+                             device_matrix_view<T, IdxT, row_major> out_vectors,                \
+                             uint32_t label,                                                    \
+                             uint32_t offset)                                                   \
+  {                                                                                             \
+    raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>(                                  \
+      handle, idx, out_vectors, label, offset);                                                 \
   }
 
 RAFT_INST_BUILD_EXTEND(float, int64_t);
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index f368d4e368..8b0bd6bdc2 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -296,6 +296,40 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
                                     raft::Compare<DataT>(),
                                     stream_));
     }
+
+    {
+      raft::spatial::knn::ivf_flat::index_params index_params;
+      index_params.n_lists                  = ps.nlist;
+      index_params.metric                   = ps.metric;
+      index_params.adaptive_centers         = ps.adaptive_centers;
+      index_params.add_data_on_build        = true;
+      index_params.kmeans_trainset_fraction = 0.5;
+
+      auto vectors_data =
+        raft::make_device_matrix_view<const DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
+      auto index = ivf_flat::build(handle_, index_params, vectors_data);
+
+      uint32_t cluster = 0;
+      uint32_t offset  = 0;
+      uint32_t n_rows  = index.lists()[cluster]->size;
+      if (n_rows > 0 && n_rows <= 30) {
+        auto vectors_out =
+          raft::make_device_matrix<DataT, IdxT, row_major>(handle_, n_rows, ps.dim);
+        ivf_flat::reconstruct_list_data(handle_, index, vectors_out.view(), cluster, offset);
+
+        std::vector<IdxT> h_indices(n_rows);
+        raft::update_host(
+          h_indices.data(), index.lists()[cluster]->indices.data_handle(), n_rows, stream_);
+        for (IdxT i = 0; i < n_rows; ++i) {
+          IdxT idx = h_indices[i];
+          ASSERT_TRUE(raft::devArrMatch(&(vectors_data.data_handle()[idx * ps.dim]),
+                                        &(vectors_out.data_handle()[i * ps.dim]),
+                                        ps.dim,
+                                        raft::Compare<DataT>(),
+                                        stream_));
+        }
+      }
+    }
   }
 
   void SetUp() override

From 0f6cdbe495cd690d05dfad0d525c4406b23907aa Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Wed, 12 Jul 2023 16:13:24 +0200
Subject: [PATCH 28/33] addressing review 1/2

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  | 51 +++++++++----------
 cpp/include/raft/neighbors/ivf_flat-ext.cuh   |  8 +--
 cpp/include/raft/neighbors/ivf_flat-inl.cuh   |  4 +-
 .../raft_runtime/neighbors/ivf_flat.hpp       |  4 +-
 .../ivf_flat_reconstruct_float_int64_t.cu     |  4 +-
 .../ivf_flat_reconstruct_int8_t_int64_t.cu    |  4 +-
 .../ivf_flat_reconstruct_uint8_t_int64_t.cu   |  4 +-
 .../raft_runtime/neighbors/ivf_flat_build.cu  |  4 +-
 cpp/test/neighbors/ann_ivf_flat.cuh           |  6 +--
 9 files changed, 44 insertions(+), 45 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index 02175248f3..d938e07fd1 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -520,26 +520,25 @@ void reconstruct_batch(raft::resources const& handle,
   RAFT_CUDA_TRY(cudaPeekAtLastError());
 }
 
-template <typename T>
-__global__ void reconstruct_list_data_kernel(
-  T* out_vectors,
-  T* in_list_data,
-  std::variant<uint32_t, const uint32_t*> offset_or_indices,
-  uint32_t len,
-  size_t veclen,
-  uint32_t dim)
+template <typename T, typename IdxT>
+__global__ void reconstruct_list_data_kernel(T* out_vectors,
+                                             T* in_list_data,
+                                             std::variant<IdxT, const IdxT*> offset_or_indices,
+                                             IdxT len,
+                                             size_t veclen,
+                                             IdxT dim)
 {
-  for (uint32_t ix = threadIdx.x + blockDim.x * blockIdx.x; ix < len; ix += blockDim.x) {
-    const uint32_t src_ix = std::holds_alternative<uint32_t>(offset_or_indices)
-                              ? std::get<uint32_t>(offset_or_indices) + ix
-                              : std::get<const uint32_t*>(offset_or_indices)[ix];
+  for (IdxT ix = threadIdx.x + blockDim.x * blockIdx.x; ix < len; ix += blockDim.x) {
+    const IdxT src_ix = std::holds_alternative<IdxT>(offset_or_indices)
+                              ? std::get<IdxT>(offset_or_indices) + ix
+                              : std::get<const IdxT*>(offset_or_indices)[ix];
 
     using group_align         = Pow2<kIndexGroupSize>;
-    const uint32_t group_ix   = group_align::div(src_ix);
-    const uint32_t ingroup_ix = group_align::mod(src_ix) * veclen;
+    const IdxT group_ix   = group_align::div(src_ix);
+    const IdxT ingroup_ix = group_align::mod(src_ix) * veclen;
 
-    for (uint32_t l = 0; l < dim; l += veclen) {
-      for (uint32_t j = 0; j < veclen; j++) {
+    for (IdxT l = 0; l < dim; l += veclen) {
+      for (IdxT j = 0; j < veclen; j++) {
         out_vectors[ix * dim + l + j] = in_list_data[l * kIndexGroupSize + ingroup_ix + j];
       }
     }
@@ -551,21 +550,21 @@ template <typename T, typename IdxT>
 void reconstruct_list_data(raft::resources const& handle,
                            const index<T, IdxT>& index,
                            device_matrix_view<T, IdxT, row_major> out_vectors,
-                           uint32_t label,
-                           uint32_t offset)
+                           IdxT label,
+                           IdxT offset)
 {
   auto stream = raft::resource::get_cuda_stream(handle);
 
-  uint32_t len = out_vectors.extent(0);
+  IdxT len = out_vectors.extent(0);
   const dim3 block_dim(256);
-  const dim3 grid_dim(raft::ceildiv<size_t>(len, block_dim.x));
-  reconstruct_list_data_kernel<<<grid_dim, block_dim, 0, stream>>>(
-    out_vectors.data_handle(),
-    index.lists()[label]->data.data_handle(),
-    offset,
-    len,
+  const dim3 grid_dim(raft::div_rounding_up_safe<size_t>(len, block_dim.x));
+  reconstruct_list_data_kernel<T, IdxT><<<grid_dim, block_dim, 0, stream>>>(
+    (T*)out_vectors.data_handle(),
+    (T*)index.lists()[label]->data.data_handle(),
+    (IdxT)offset,
+    (IdxT)len,
     (size_t)index.veclen(),
-    index.dim());
+    (IdxT)index.dim());
 }
 
 }  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/include/raft/neighbors/ivf_flat-ext.cuh b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
index aad2d2f5e5..384446d1db 100644
--- a/cpp/include/raft/neighbors/ivf_flat-ext.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
@@ -103,8 +103,8 @@ template <typename T, typename IdxT>
 void reconstruct_list_data(raft::resources const& handle,
                            const index<T, IdxT>& index,
                            device_matrix_view<T, IdxT, row_major> out_vectors,
-                           uint32_t label,
-                           uint32_t offset) RAFT_EXPLICIT;
+                           IdxT label,
+                           IdxT offset) RAFT_EXPLICIT;
 
 }  // namespace raft::neighbors::ivf_flat
 
@@ -208,8 +208,8 @@ instantiate_raft_neighbors_ivf_flat_search(uint8_t, int64_t);
     raft::resources const& handle,                                                \
     const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
     raft::device_matrix_view<T, IdxT, row_major> out_vectors,                     \
-    uint32_t label,                                                               \
-    uint32_t offset);
+    IdxT label,                                                                   \
+    IdxT offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
 instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
index fd322f2965..65ce62e9cd 100644
--- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
@@ -503,8 +503,8 @@ template <typename T, typename IdxT>
 void reconstruct_list_data(raft::resources const& handle,
                            const index<T, IdxT>& index,
                            device_matrix_view<T, IdxT, row_major> out_vectors,
-                           uint32_t label,
-                           uint32_t offset)
+                           IdxT label,
+                           IdxT offset)
 {
   return raft::neighbors::ivf_flat::detail::reconstruct_list_data(
     handle, index, out_vectors, label, offset);
diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
index 74e8bad00a..efe40f3a54 100644
--- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
+++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
@@ -54,8 +54,8 @@ namespace raft::runtime::neighbors::ivf_flat {
   void reconstruct_list_data(raft::resources const& handle,                          \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,   \
                              device_matrix_view<T, IdxT, row_major> out_vectors,     \
-                             uint32_t label,                                         \
-                             uint32_t offset);                                       \
+                             IdxT label,                                             \
+                             IdxT offset);                                           \
                                                                                      \
   void serialize_file(raft::resources const& handle,                                 \
                       const std::string& filename,                                   \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
index 91ed15c41a..fffe9dc66b 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
@@ -36,8 +36,8 @@
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
     raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
-    uint32_t label,                                                        \
-    uint32_t offset);
+    IdxT label,                                                            \
+    IdxT offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
 
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
index 0f214e66ca..933a9a6b3a 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
@@ -36,8 +36,8 @@
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
     raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
-    uint32_t label,                                                        \
-    uint32_t offset);
+    IdxT label,                                                            \
+    IdxT offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
 
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
index b693a1415f..492e1caee7 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
@@ -36,8 +36,8 @@
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
     raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
-    uint32_t label,                                                        \
-    uint32_t offset);
+    IdxT label,                                                            \
+    IdxT offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(uint8_t, int64_t);
 
diff --git a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
index 39b2962782..b62b09e32c 100644
--- a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
+++ b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
@@ -64,8 +64,8 @@ namespace raft::runtime::neighbors::ivf_flat {
   void reconstruct_list_data(raft::resources const& handle,                                     \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
                              device_matrix_view<T, IdxT, row_major> out_vectors,                \
-                             uint32_t label,                                                    \
-                             uint32_t offset)                                                   \
+                             IdxT label,                                                        \
+                             IdxT offset)                                                       \
   {                                                                                             \
     raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>(                                  \
       handle, idx, out_vectors, label, offset);                                                 \
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index 8b0bd6bdc2..1bae97b096 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -309,9 +309,9 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         raft::make_device_matrix_view<const DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
       auto index = ivf_flat::build(handle_, index_params, vectors_data);
 
-      uint32_t cluster = 0;
-      uint32_t offset  = 0;
-      uint32_t n_rows  = index.lists()[cluster]->size;
+      IdxT cluster = 0;
+      IdxT offset  = 0;
+      IdxT n_rows  = index.lists()[cluster]->size;
       if (n_rows > 0 && n_rows <= 30) {
         auto vectors_out =
           raft::make_device_matrix<DataT, IdxT, row_major>(handle_, n_rows, ps.dim);

From cad88b38a1de065fed933e8a910002d543b8cfa2 Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Wed, 12 Jul 2023 16:22:25 +0200
Subject: [PATCH 29/33] addressing review 2/2

---
 cpp/include/raft/neighbors/ivf_flat-ext.cuh   | 12 ----------
 cpp/include/raft/neighbors/ivf_flat-inl.cuh   | 22 -------------------
 .../raft_runtime/neighbors/ivf_flat.hpp       |  5 -----
 .../ivf_flat_reconstruct_float_int64_t.cu     |  6 -----
 .../ivf_flat_reconstruct_int8_t_int64_t.cu    |  6 -----
 .../ivf_flat_reconstruct_uint8_t_int64_t.cu   |  6 -----
 .../raft_runtime/neighbors/ivf_flat_build.cu  |  8 -------
 cpp/test/neighbors/ann_ivf_flat.cuh           |  3 ++-
 8 files changed, 2 insertions(+), 66 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat-ext.cuh b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
index 384446d1db..2273e1e354 100644
--- a/cpp/include/raft/neighbors/ivf_flat-ext.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
@@ -93,12 +93,6 @@ void search(raft::resources const& handle,
             raft::device_matrix_view<IdxT, IdxT, row_major> neighbors,
             raft::device_matrix_view<float, IdxT, row_major> distances) RAFT_EXPLICIT;
 
-template <typename T, typename IdxT>
-void reconstruct_batch(raft::resources const& handle,
-                       const index<T, IdxT>& index,
-                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
-                       raft::device_matrix_view<T, IdxT, row_major> vector_out) RAFT_EXPLICIT;
-
 template <typename T, typename IdxT>
 void reconstruct_list_data(raft::resources const& handle,
                            const index<T, IdxT>& index,
@@ -198,12 +192,6 @@ instantiate_raft_neighbors_ivf_flat_search(uint8_t, int64_t);
 #undef instantiate_raft_neighbors_ivf_flat_search
 
 #define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)                  \
-  extern template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
-    raft::resources const& handle,                                                \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,                        \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);                     \
-                                                                                  \
   extern template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                                \
     const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
index 65ce62e9cd..40c9a7a8a2 100644
--- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
@@ -464,28 +464,6 @@ void search(raft::resources const& handle,
                 nullptr);
 }
 
-/**
- * @brief Reconstruct vector database
- *
- *
- * @tparam T data element type
- * @tparam IdxT type of the indices
- *
- * @param[in] handle
- * @param[in] index ivf-flat constructed index
- * @param[in] vector_ids vector with the ids of vectors to look for
- * @param[out] vector_out matrix with the database of vectors
- */
-template <typename T, typename IdxT>
-void reconstruct_batch(raft::resources const& handle,
-                       const index<T, IdxT>& index,
-                       raft::device_vector_view<const IdxT, IdxT> vector_ids,
-                       raft::device_matrix_view<T, IdxT, row_major> vector_out)
-{
-  return raft::neighbors::ivf_flat::detail::reconstruct_batch(
-    handle, index, vector_ids, vector_out);
-}
-
 /**
  * @brief Reconstruct vectors of a given cluster
  *
diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
index efe40f3a54..a076e870b9 100644
--- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
+++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
@@ -46,11 +46,6 @@ namespace raft::runtime::neighbors::ivf_flat {
               std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices, \
               raft::neighbors::ivf_flat::index<T, IdxT>* idx);                       \
                                                                                      \
-  void reconstruct_batch(raft::resources const& handle,                              \
-                         raft::neighbors::ivf_flat::index<T, IdxT>& idx,             \
-                         raft::device_vector_view<const IdxT, IdxT> vector_ids,      \
-                         raft::device_matrix_view<T, IdxT, row_major> vector_out);   \
-                                                                                     \
   void reconstruct_list_data(raft::resources const& handle,                          \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,   \
                              device_matrix_view<T, IdxT, row_major> out_vectors,     \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
index fffe9dc66b..1ae8cfdcf9 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
@@ -26,12 +26,6 @@
 #include <raft/neighbors/ivf_flat-inl.cuh>
 
 #define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)           \
-  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
-    raft::resources const& handle,                                         \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);              \
-                                                                           \
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
index 933a9a6b3a..1d5250487d 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
@@ -26,12 +26,6 @@
 #include <raft/neighbors/ivf_flat-inl.cuh>
 
 #define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)           \
-  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
-    raft::resources const& handle,                                         \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);              \
-                                                                           \
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
index 492e1caee7..2d521d9703 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
@@ -26,12 +26,6 @@
 #include <raft/neighbors/ivf_flat-inl.cuh>
 
 #define instantiate_raft_neighbors_ivf_flat_reconstruct(T, IdxT)           \
-  template void raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(     \
-    raft::resources const& handle,                                         \
-    const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
-    raft::device_matrix_view<T, IdxT, row_major> vector_out);              \
-                                                                           \
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
diff --git a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
index b62b09e32c..c5a608555b 100644
--- a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
+++ b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
@@ -53,14 +53,6 @@ namespace raft::runtime::neighbors::ivf_flat {
     raft::neighbors::ivf_flat::extend<T, IdxT>(handle, new_vectors, new_indices, idx);          \
   }                                                                                             \
                                                                                                 \
-  void reconstruct_batch(raft::resources const& handle,                                         \
-                         const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-                         raft::device_vector_view<const IdxT, IdxT> vector_ids,                 \
-                         raft::device_matrix_view<T, IdxT, row_major> vector_out)               \
-  {                                                                                             \
-    raft::neighbors::ivf_flat::reconstruct_batch<T, IdxT>(handle, idx, vector_ids, vector_out); \
-  }                                                                                             \
-                                                                                                \
   void reconstruct_list_data(raft::resources const& handle,                                     \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
                              device_matrix_view<T, IdxT, row_major> out_vectors,                \
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index 1bae97b096..a557fd0ea6 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -27,6 +27,7 @@
 #include <raft/core/logger.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <raft/neighbors/ivf_flat.cuh>
+#include <raft/neighbors/detail/ivf_flat_build.cuh>
 #include <raft/random/rng.cuh>
 #include <raft/spatial/knn/ann.cuh>
 #include <raft/spatial/knn/knn.cuh>
@@ -286,7 +287,7 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         raft::make_device_vector_view<const IdxT, IdxT>(vecs_ids.data(), ps.num_db_vecs);
       auto vectors_out =
         raft::make_device_matrix<DataT, IdxT, row_major>(handle_, ps.num_db_vecs, ps.dim);
-      ivf_flat::reconstruct_batch(handle_, index, vectors_ids, vectors_out.view());
+      ivf_flat::detail::reconstruct_batch(handle_, index, vectors_ids, vectors_out.view());
 
       resource::sync_stream(handle_);
 

From 6cd1abc74ea0206e62e6a573d949e16ee417d3e2 Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Wed, 12 Jul 2023 18:10:00 +0200
Subject: [PATCH 30/33] Fix style

---
 .../raft/neighbors/detail/ivf_flat_build.cuh  | 20 ++---
 .../raft_runtime/neighbors/ivf_flat_build.cu  | 84 +++++++++----------
 cpp/test/neighbors/ann_ivf_flat.cuh           |  2 +-
 3 files changed, 53 insertions(+), 53 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index d938e07fd1..d370f77f6c 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -530,10 +530,10 @@ __global__ void reconstruct_list_data_kernel(T* out_vectors,
 {
   for (IdxT ix = threadIdx.x + blockDim.x * blockIdx.x; ix < len; ix += blockDim.x) {
     const IdxT src_ix = std::holds_alternative<IdxT>(offset_or_indices)
-                              ? std::get<IdxT>(offset_or_indices) + ix
-                              : std::get<const IdxT*>(offset_or_indices)[ix];
+                          ? std::get<IdxT>(offset_or_indices) + ix
+                          : std::get<const IdxT*>(offset_or_indices)[ix];
 
-    using group_align         = Pow2<kIndexGroupSize>;
+    using group_align     = Pow2<kIndexGroupSize>;
     const IdxT group_ix   = group_align::div(src_ix);
     const IdxT ingroup_ix = group_align::mod(src_ix) * veclen;
 
@@ -558,13 +558,13 @@ void reconstruct_list_data(raft::resources const& handle,
   IdxT len = out_vectors.extent(0);
   const dim3 block_dim(256);
   const dim3 grid_dim(raft::div_rounding_up_safe<size_t>(len, block_dim.x));
-  reconstruct_list_data_kernel<T, IdxT><<<grid_dim, block_dim, 0, stream>>>(
-    (T*)out_vectors.data_handle(),
-    (T*)index.lists()[label]->data.data_handle(),
-    (IdxT)offset,
-    (IdxT)len,
-    (size_t)index.veclen(),
-    (IdxT)index.dim());
+  reconstruct_list_data_kernel<T, IdxT>
+    <<<grid_dim, block_dim, 0, stream>>>((T*)out_vectors.data_handle(),
+                                         (T*)index.lists()[label]->data.data_handle(),
+                                         (IdxT)offset,
+                                         (IdxT)len,
+                                         (size_t)index.veclen(),
+                                         (IdxT)index.dim());
 }
 
 }  // namespace raft::neighbors::ivf_flat::detail
diff --git a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
index c5a608555b..bbeda2e3aa 100644
--- a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
+++ b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
@@ -19,48 +19,48 @@
 
 namespace raft::runtime::neighbors::ivf_flat {
 
-#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                                         \
-  auto build(raft::resources const& handle,                                                     \
-             const raft::neighbors::ivf_flat::index_params& params,                             \
-             raft::device_matrix_view<const T, IdxT, row_major> dataset)                        \
-    ->raft::neighbors::ivf_flat::index<T, IdxT>                                                 \
-  {                                                                                             \
-    return raft::neighbors::ivf_flat::build<T, IdxT>(handle, params, dataset);                  \
-  }                                                                                             \
-  auto extend(raft::resources const& handle,                                                    \
-              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,                   \
-              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,            \
-              const raft::neighbors::ivf_flat::index<T, IdxT>& orig_index)                      \
-    ->raft::neighbors::ivf_flat::index<T, IdxT>                                                 \
-  {                                                                                             \
-    return raft::neighbors::ivf_flat::extend<T, IdxT>(                                          \
-      handle, new_vectors, new_indices, orig_index);                                            \
-  }                                                                                             \
-                                                                                                \
-  void build(raft::resources const& handle,                                                     \
-             const raft::neighbors::ivf_flat::index_params& params,                             \
-             raft::device_matrix_view<const T, IdxT, row_major> dataset,                        \
-             raft::neighbors::ivf_flat::index<T, IdxT>& idx)                                    \
-  {                                                                                             \
-    idx = build(handle, params, dataset);                                                       \
-  }                                                                                             \
-                                                                                                \
-  void extend(raft::resources const& handle,                                                    \
-              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,                   \
-              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,            \
-              raft::neighbors::ivf_flat::index<T, IdxT>* idx)                                   \
-  {                                                                                             \
-    raft::neighbors::ivf_flat::extend<T, IdxT>(handle, new_vectors, new_indices, idx);          \
-  }                                                                                             \
-                                                                                                \
-  void reconstruct_list_data(raft::resources const& handle,                                     \
-                             const raft::neighbors::ivf_flat::index<T, IdxT>& idx,              \
-                             device_matrix_view<T, IdxT, row_major> out_vectors,                \
-                             IdxT label,                                                        \
-                             IdxT offset)                                                       \
-  {                                                                                             \
-    raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>(                                  \
-      handle, idx, out_vectors, label, offset);                                                 \
+#define RAFT_INST_BUILD_EXTEND(T, IdxT)                                                \
+  auto build(raft::resources const& handle,                                            \
+             const raft::neighbors::ivf_flat::index_params& params,                    \
+             raft::device_matrix_view<const T, IdxT, row_major> dataset)               \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>                                        \
+  {                                                                                    \
+    return raft::neighbors::ivf_flat::build<T, IdxT>(handle, params, dataset);         \
+  }                                                                                    \
+  auto extend(raft::resources const& handle,                                           \
+              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,          \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,   \
+              const raft::neighbors::ivf_flat::index<T, IdxT>& orig_index)             \
+    ->raft::neighbors::ivf_flat::index<T, IdxT>                                        \
+  {                                                                                    \
+    return raft::neighbors::ivf_flat::extend<T, IdxT>(                                 \
+      handle, new_vectors, new_indices, orig_index);                                   \
+  }                                                                                    \
+                                                                                       \
+  void build(raft::resources const& handle,                                            \
+             const raft::neighbors::ivf_flat::index_params& params,                    \
+             raft::device_matrix_view<const T, IdxT, row_major> dataset,               \
+             raft::neighbors::ivf_flat::index<T, IdxT>& idx)                           \
+  {                                                                                    \
+    idx = build(handle, params, dataset);                                              \
+  }                                                                                    \
+                                                                                       \
+  void extend(raft::resources const& handle,                                           \
+              raft::device_matrix_view<const T, IdxT, row_major> new_vectors,          \
+              std::optional<raft::device_vector_view<const IdxT, IdxT>> new_indices,   \
+              raft::neighbors::ivf_flat::index<T, IdxT>* idx)                          \
+  {                                                                                    \
+    raft::neighbors::ivf_flat::extend<T, IdxT>(handle, new_vectors, new_indices, idx); \
+  }                                                                                    \
+                                                                                       \
+  void reconstruct_list_data(raft::resources const& handle,                            \
+                             const raft::neighbors::ivf_flat::index<T, IdxT>& idx,     \
+                             device_matrix_view<T, IdxT, row_major> out_vectors,       \
+                             IdxT label,                                               \
+                             IdxT offset)                                              \
+  {                                                                                    \
+    raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>(                         \
+      handle, idx, out_vectors, label, offset);                                        \
   }
 
 RAFT_INST_BUILD_EXTEND(float, int64_t);
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index a557fd0ea6..f4db78a7b7 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -26,8 +26,8 @@
 #include <raft/core/device_mdspan.hpp>
 #include <raft/core/logger.hpp>
 #include <raft/distance/distance_types.hpp>
-#include <raft/neighbors/ivf_flat.cuh>
 #include <raft/neighbors/detail/ivf_flat_build.cuh>
+#include <raft/neighbors/ivf_flat.cuh>
 #include <raft/random/rng.cuh>
 #include <raft/spatial/knn/ann.cuh>
 #include <raft/spatial/knn/knn.cuh>

From de6c13dc92fb456b4fd9198efa45137202c8373c Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Fri, 14 Jul 2023 15:21:17 +0200
Subject: [PATCH 31/33] addressing review

---
 cpp/include/raft/neighbors/detail/ivf_flat_build.cuh |  6 +++---
 cpp/include/raft/neighbors/ivf_flat-ext.cuh          | 12 ++++++------
 cpp/include/raft/neighbors/ivf_flat-inl.cuh          |  6 +++---
 cpp/include/raft_runtime/neighbors/ivf_flat.hpp      |  6 +++---
 .../neighbors/ivf_flat_reconstruct_float_int64_t.cu  |  6 +++---
 .../neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu |  6 +++---
 .../ivf_flat_reconstruct_uint8_t_int64_t.cu          |  6 +++---
 cpp/src/raft_runtime/neighbors/ivf_flat_build.cu     |  6 +++---
 cpp/test/neighbors/ann_ivf_flat.cuh                  |  8 ++++----
 9 files changed, 31 insertions(+), 31 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
index d370f77f6c..92c8c8ed81 100644
--- a/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
+++ b/cpp/include/raft/neighbors/detail/ivf_flat_build.cuh
@@ -549,9 +549,9 @@ __global__ void reconstruct_list_data_kernel(T* out_vectors,
 template <typename T, typename IdxT>
 void reconstruct_list_data(raft::resources const& handle,
                            const index<T, IdxT>& index,
-                           device_matrix_view<T, IdxT, row_major> out_vectors,
-                           IdxT label,
-                           IdxT offset)
+                           device_matrix_view<T, uint32_t, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset)
 {
   auto stream = raft::resource::get_cuda_stream(handle);
 
diff --git a/cpp/include/raft/neighbors/ivf_flat-ext.cuh b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
index 2273e1e354..0198dff10d 100644
--- a/cpp/include/raft/neighbors/ivf_flat-ext.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
@@ -96,9 +96,9 @@ void search(raft::resources const& handle,
 template <typename T, typename IdxT>
 void reconstruct_list_data(raft::resources const& handle,
                            const index<T, IdxT>& index,
-                           device_matrix_view<T, IdxT, row_major> out_vectors,
-                           IdxT label,
-                           IdxT offset) RAFT_EXPLICIT;
+                           device_matrix_view<T, uint32_t, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset) RAFT_EXPLICIT;
 
 }  // namespace raft::neighbors::ivf_flat
 
@@ -195,9 +195,9 @@ instantiate_raft_neighbors_ivf_flat_search(uint8_t, int64_t);
   extern template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                                \
     const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
-    raft::device_matrix_view<T, IdxT, row_major> out_vectors,                     \
-    IdxT label,                                                                   \
-    IdxT offset);
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,                     \
+    uint32_t label,                                                                   \
+    uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
 instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
index 40c9a7a8a2..888020a813 100644
--- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
@@ -480,9 +480,9 @@ void search(raft::resources const& handle,
 template <typename T, typename IdxT>
 void reconstruct_list_data(raft::resources const& handle,
                            const index<T, IdxT>& index,
-                           device_matrix_view<T, IdxT, row_major> out_vectors,
-                           IdxT label,
-                           IdxT offset)
+                           device_matrix_view<T, uint32_t, row_major> out_vectors,
+                           uint32_t label,
+                           uint32_t offset)
 {
   return raft::neighbors::ivf_flat::detail::reconstruct_list_data(
     handle, index, out_vectors, label, offset);
diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
index a076e870b9..b678b1a515 100644
--- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
+++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
@@ -48,9 +48,9 @@ namespace raft::runtime::neighbors::ivf_flat {
                                                                                      \
   void reconstruct_list_data(raft::resources const& handle,                          \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,   \
-                             device_matrix_view<T, IdxT, row_major> out_vectors,     \
-                             IdxT label,                                             \
-                             IdxT offset);                                           \
+                             device_matrix_view<T, uint32_t, row_major> out_vectors,     \
+                             uint32_t label,                                             \
+                             uint32_t offset);                                           \
                                                                                      \
   void serialize_file(raft::resources const& handle,                                 \
                       const std::string& filename,                                   \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
index 1ae8cfdcf9..d4333ed8aa 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
@@ -29,9 +29,9 @@
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
-    IdxT label,                                                            \
-    IdxT offset);
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,              \
+    uint32_t label,                                                            \
+    uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
 
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
index 1d5250487d..256f54042c 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
@@ -29,9 +29,9 @@
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
-    IdxT label,                                                            \
-    IdxT offset);
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,              \
+    uint32_t label,                                                            \
+    uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
 
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
index 2d521d9703..bb5028d3e0 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
@@ -29,9 +29,9 @@
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_matrix_view<T, IdxT, row_major> out_vectors,              \
-    IdxT label,                                                            \
-    IdxT offset);
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,              \
+    uint32_t label,                                                            \
+    uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(uint8_t, int64_t);
 
diff --git a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
index bbeda2e3aa..0338e177ff 100644
--- a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
+++ b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
@@ -55,9 +55,9 @@ namespace raft::runtime::neighbors::ivf_flat {
                                                                                        \
   void reconstruct_list_data(raft::resources const& handle,                            \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,     \
-                             device_matrix_view<T, IdxT, row_major> out_vectors,       \
-                             IdxT label,                                               \
-                             IdxT offset)                                              \
+                             device_matrix_view<T, uint32_t, row_major> out_vectors,       \
+                             uint32_t label,                                               \
+                             uint32_t offset)                                              \
   {                                                                                    \
     raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>(                         \
       handle, idx, out_vectors, label, offset);                                        \
diff --git a/cpp/test/neighbors/ann_ivf_flat.cuh b/cpp/test/neighbors/ann_ivf_flat.cuh
index f4db78a7b7..2795c649d6 100644
--- a/cpp/test/neighbors/ann_ivf_flat.cuh
+++ b/cpp/test/neighbors/ann_ivf_flat.cuh
@@ -310,12 +310,12 @@ class AnnIVFFlatTest : public ::testing::TestWithParam<AnnIvfFlatInputs<IdxT>> {
         raft::make_device_matrix_view<const DataT, IdxT>(database.data(), ps.num_db_vecs, ps.dim);
       auto index = ivf_flat::build(handle_, index_params, vectors_data);
 
-      IdxT cluster = 0;
-      IdxT offset  = 0;
-      IdxT n_rows  = index.lists()[cluster]->size;
+      uint32_t cluster = 0;
+      uint32_t offset  = 0;
+      uint32_t n_rows  = index.lists()[cluster]->size;
       if (n_rows > 0 && n_rows <= 30) {
         auto vectors_out =
-          raft::make_device_matrix<DataT, IdxT, row_major>(handle_, n_rows, ps.dim);
+          raft::make_device_matrix<DataT, uint32_t, row_major>(handle_, n_rows, ps.dim);
         ivf_flat::reconstruct_list_data(handle_, index, vectors_out.view(), cluster, offset);
 
         std::vector<IdxT> h_indices(n_rows);

From 5e02d7fdfd8f0bb2e16d25f6ffe735acf5b42081 Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Fri, 14 Jul 2023 16:47:48 +0200
Subject: [PATCH 32/33] Fix style

---
 cpp/include/raft/neighbors/ivf_flat-ext.cuh               | 4 ++--
 cpp/include/raft_runtime/neighbors/ivf_flat.hpp           | 6 +++---
 cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu   | 4 ++--
 cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu  | 4 ++--
 cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu | 4 ++--
 cpp/src/raft_runtime/neighbors/ivf_flat_build.cu          | 6 +++---
 6 files changed, 14 insertions(+), 14 deletions(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat-ext.cuh b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
index 1b8fed86d6..2bcb361242 100644
--- a/cpp/include/raft/neighbors/ivf_flat-ext.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-ext.cuh
@@ -216,8 +216,8 @@ instantiate_raft_neighbors_ivf_flat_search(uint8_t, int64_t);
   extern template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                                \
     const raft::neighbors::ivf_flat::index<T, IdxT>& index,                       \
-    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,                     \
-    uint32_t label,                                                                   \
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,                 \
+    uint32_t label,                                                               \
     uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
diff --git a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
index b678b1a515..7b6daad427 100644
--- a/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
+++ b/cpp/include/raft_runtime/neighbors/ivf_flat.hpp
@@ -48,9 +48,9 @@ namespace raft::runtime::neighbors::ivf_flat {
                                                                                      \
   void reconstruct_list_data(raft::resources const& handle,                          \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,   \
-                             device_matrix_view<T, uint32_t, row_major> out_vectors,     \
-                             uint32_t label,                                             \
-                             uint32_t offset);                                           \
+                             device_matrix_view<T, uint32_t, row_major> out_vectors, \
+                             uint32_t label,                                         \
+                             uint32_t offset);                                       \
                                                                                      \
   void serialize_file(raft::resources const& handle,                                 \
                       const std::string& filename,                                   \
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
index d4333ed8aa..d13d2b5829 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_float_int64_t.cu
@@ -29,8 +29,8 @@
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,              \
-    uint32_t label,                                                            \
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,          \
+    uint32_t label,                                                        \
     uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(float, int64_t);
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
index 256f54042c..8940c18209 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_int8_t_int64_t.cu
@@ -29,8 +29,8 @@
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,              \
-    uint32_t label,                                                            \
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,          \
+    uint32_t label,                                                        \
     uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(int8_t, int64_t);
diff --git a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
index bb5028d3e0..8329fe15ff 100644
--- a/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
+++ b/cpp/src/neighbors/ivf_flat_reconstruct_uint8_t_int64_t.cu
@@ -29,8 +29,8 @@
   template void raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>( \
     raft::resources const& handle,                                         \
     const raft::neighbors::ivf_flat::index<T, IdxT>& idx,                  \
-    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,              \
-    uint32_t label,                                                            \
+    raft::device_matrix_view<T, uint32_t, row_major> out_vectors,          \
+    uint32_t label,                                                        \
     uint32_t offset);
 
 instantiate_raft_neighbors_ivf_flat_reconstruct(uint8_t, int64_t);
diff --git a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
index 0338e177ff..68cf13e704 100644
--- a/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
+++ b/cpp/src/raft_runtime/neighbors/ivf_flat_build.cu
@@ -55,9 +55,9 @@ namespace raft::runtime::neighbors::ivf_flat {
                                                                                        \
   void reconstruct_list_data(raft::resources const& handle,                            \
                              const raft::neighbors::ivf_flat::index<T, IdxT>& idx,     \
-                             device_matrix_view<T, uint32_t, row_major> out_vectors,       \
-                             uint32_t label,                                               \
-                             uint32_t offset)                                              \
+                             device_matrix_view<T, uint32_t, row_major> out_vectors,   \
+                             uint32_t label,                                           \
+                             uint32_t offset)                                          \
   {                                                                                    \
     raft::neighbors::ivf_flat::reconstruct_list_data<T, IdxT>(                         \
       handle, idx, out_vectors, label, offset);                                        \

From 87edda0e358d0fe0be95a1e50d58babe7709a812 Mon Sep 17 00:00:00 2001
From: viclafargue <viclafargue@nvidia.com>
Date: Tue, 25 Jul 2023 17:02:58 +0200
Subject: [PATCH 33/33] Fix doc

---
 cpp/include/raft/neighbors/ivf_flat-inl.cuh | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/cpp/include/raft/neighbors/ivf_flat-inl.cuh b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
index 78d6a14f2e..c28e8e1f46 100644
--- a/cpp/include/raft/neighbors/ivf_flat-inl.cuh
+++ b/cpp/include/raft/neighbors/ivf_flat-inl.cuh
@@ -600,7 +600,7 @@ void search(raft::resources const& handle,
  *
  * @param[in] handle
  * @param[in] index ivf-flat constructed index
- * @param[out] vector_out matrix with the vectors contained in the cluster
+ * @param[out] out_vectors matrix with the vectors contained in the cluster
  * @param[in] label cluster index
  * @param[in] offset offset
  */