CAGRA: decouple source idx type from graph idx type and add a mapping between them (#1251)

Add a mapping from internal graph index type to the source (user) index type for CAGRA index.

This PR contains a minimal set of changes required to decouple the graph index type from the public API index type:
  - Add the mapping - a device vector of `n_rows` elements storing the original indices (which are provided by the user)
  - Add corresponding setters/getters and update the serialization code
  - Modify all search kernels accordingly: (1) using the map before coercing the output to `OutputIdxT` and (2) using the map every time the filter operator is used (so that the filtering happens on "user-space" indices rather than on internal indices).

In the detail namespace, I add a new template parameter `SourceIdxT` , which doesn't currently multiply the template instantiations. `SourceIdxT` is bound directly to IdxT in `index<T, IdxT>`, whereas the `IndexT` in the detail namespace becomes hard-coded to `graph_idx_t`, which is `uint32_t`.

***What's not here***
This PR does not modify the order of the graph or the dataset during construction. Hence to test this feature now, a user needs to reorder both the dataset and the graph manually using the `update_xxx` methods of the CAGRA index.

***IVF-XX Feature parity***
With the addition of the index mapping, it will be possible to add the indices argument to the `extend` function similar to the other algorithms.

***Multi-vector index***
(related issue #1191)
This PR can be seen as a prerequisite for <single_id - multiple_vectors> mapping: since the ids in the source_indicies vector are not necessarily unique. To actually allow multi-vector usage in a meaningful way, we also need to:
  - De-duplicate the indices in the search results
  - Extend the public API - how do add new vectors and indices alongside them, and where/how to keep the dataset.

Authors:
  - Artem M. Chirkin (https://github.com/achirkin)
  - Corey J. Nolet (https://github.com/cjnolet)
  - Ben Frederickson (https://github.com/benfred)

Approvers:
  - Divye Gala (https://github.com/divyegala)
  - Ben Frederickson (https://github.com/benfred)

URL: #1251

Author: achirkin

cpp/include/cuvs/neighbors/cagra.hpp

@@ -291,6 +291,7 @@ struct index : cuvs::neighbors::index {
   using index_type         = IdxT;
   using value_type         = T;
   using dataset_index_type = int64_t;
+  using graph_index_type   = uint32_t;
 
   static_assert(!raft::is_narrowing_v<uint32_t, IdxT>,
                 "IdxT must be able to represent all values of uint32_t");
@@ -334,11 +335,21 @@ struct index : cuvs::neighbors::index {
 
   /** neighborhood graph [size, graph-degree] */
   [[nodiscard]] inline auto graph() const noexcept
-    -> raft::device_matrix_view<const IdxT, int64_t, raft::row_major>
+    -> raft::device_matrix_view<const graph_index_type, int64_t, raft::row_major>
   {
     return graph_view_;
   }
 
+  /** Mapping from internal graph node indices to the original user-provided indices. */
+  [[nodiscard]] inline auto source_indices() const noexcept
+    -> std::optional<raft::device_vector_view<const index_type, int64_t>>
+  {
+    return source_indices_.has_value()
+             ? std::optional<raft::device_vector_view<const index_type, int64_t>>(
+                 source_indices_->view())
+             : std::nullopt;
+  }
+
   /** Dataset norms for cosine distance [size] */
   [[nodiscard]] inline auto dataset_norms() const noexcept
     -> std::optional<raft::device_vector_view<const float, int64_t>>
@@ -361,7 +372,7 @@ struct index : cuvs::neighbors::index {
         cuvs::distance::DistanceType metric = cuvs::distance::DistanceType::L2Expanded)
     : cuvs::neighbors::index(),
       metric_(metric),
-      graph_(raft::make_device_matrix<IdxT, int64_t>(res, 0, 0)),
+      graph_(raft::make_device_matrix<graph_index_type, int64_t>(res, 0, 0)),
       dataset_(new cuvs::neighbors::empty_dataset<int64_t>(0)),
       dataset_norms_(std::nullopt)
   {
@@ -405,7 +416,7 @@ struct index : cuvs::neighbors::index {
    *   using namespace raft::neighbors::experimental;
    *
    *   auto dataset = raft::make_device_matrix<float, int64_t>(res, n_rows, n_cols);
-   *   auto knn_graph = raft::make_device_matrix<uint32_n, int64_t>(res, n_rows, graph_degree);
+   *   auto knn_graph = raft::make_device_matrix<uint32_t, int64_t>(res, n_rows, graph_degree);
    *
    *   // custom loading and graph creation
    *   // load_dataset(dataset.view());
@@ -424,11 +435,13 @@ struct index : cuvs::neighbors::index {
   index(raft::resources const& res,
         cuvs::distance::DistanceType metric,
         raft::mdspan<const T, raft::matrix_extent<int64_t>, raft::row_major, data_accessor> dataset,
-        raft::mdspan<const IdxT, raft::matrix_extent<int64_t>, raft::row_major, graph_accessor>
-          knn_graph)
+        raft::mdspan<const graph_index_type,
+                     raft::matrix_extent<int64_t>,
+                     raft::row_major,
+                     graph_accessor> knn_graph)
     : cuvs::neighbors::index(),
       metric_(metric),
-      graph_(raft::make_device_matrix<IdxT, int64_t>(res, 0, 0)),
+      graph_(raft::make_device_matrix<graph_index_type, int64_t>(res, 0, 0)),
       dataset_(make_aligned_dataset(res, dataset, 16)),
       dataset_norms_(std::nullopt)
   {
@@ -536,8 +549,9 @@ struct index : cuvs::neighbors::index {
    * Since the new graph is a device array, we store a reference to that, and it is
    * the caller's responsibility to ensure that knn_graph stays alive as long as the index.
    */
-  void update_graph(raft::resources const& res,
-                    raft::device_matrix_view<const IdxT, int64_t, raft::row_major> knn_graph)
+  void update_graph(
+    raft::resources const& res,
+    raft::device_matrix_view<const graph_index_type, int64_t, raft::row_major> knn_graph)
   {
     graph_view_ = knn_graph;
   }
@@ -547,16 +561,19 @@ struct index : cuvs::neighbors::index {
    *
    * We create a copy of the graph on the device. The index manages the lifetime of this copy.
    */
-  void update_graph(raft::resources const& res,
-                    raft::host_matrix_view<const IdxT, int64_t, raft::row_major> knn_graph)
+  void update_graph(
+    raft::resources const& res,
+    raft::host_matrix_view<const graph_index_type, int64_t, raft::row_major> knn_graph)
   {
     RAFT_LOG_DEBUG("Copying CAGRA knn graph from host to device");
 
     if ((graph_.extent(0) != knn_graph.extent(0)) || (graph_.extent(1) != knn_graph.extent(1))) {
       // clear existing memory before allocating to prevent OOM errors on large graphs
-      if (graph_.size()) { graph_ = raft::make_device_matrix<IdxT, int64_t>(res, 0, 0); }
-      graph_ =
-        raft::make_device_matrix<IdxT, int64_t>(res, knn_graph.extent(0), knn_graph.extent(1));
+      if (graph_.size()) {
+        graph_ = raft::make_device_matrix<graph_index_type, int64_t>(res, 0, 0);
+      }
+      graph_ = raft::make_device_matrix<graph_index_type, int64_t>(
+        res, knn_graph.extent(0), knn_graph.extent(1));
     }
     raft::copy(graph_.data_handle(),
                knn_graph.data_handle(),
@@ -565,11 +582,52 @@ struct index : cuvs::neighbors::index {
     graph_view_ = graph_.view();
   }
 
+  /**
+   * Replace the source indices with a new source indices taking the ownership of the passed vector.
+   */
+  void update_source_indices(raft::device_vector<index_type, int64_t>&& source_indices)
+  {
+    RAFT_EXPECTS(source_indices.extent(0) == size(),
+                 "Source indices must have the same number of rows as the index");
+    source_indices_.emplace(std::move(source_indices));
+  }
+
+  /**
+   * Copy the provided source indices into the index.
+   */
+  template <typename Accessor>
+  void update_source_indices(
+    raft::resources const& res,
+    raft::mdspan<const index_type, raft::vector_extent<int64_t>, raft::row_major, Accessor>
+      source_indices)
+  {
+    RAFT_EXPECTS(source_indices.extent(0) == size(),
+                 "Source indices must have the same number of rows as the index");
+    // Reset the array if it's not compatible to avoid using more memory than necessary.
+    // NB: this likely is never triggered because we check the invariant above (but it doesn't
+    // hurt).
+    if (source_indices_.has_value()) {
+      if (source_indices_->extent(0) != source_indices.extent(0)) { source_indices_.reset(); }
+    }
+    // Allocate the new array if needed.
+    if (!source_indices_.has_value()) {
+      source_indices_.emplace(
+        raft::make_device_vector<index_type, int64_t>(res, source_indices.extent(0)));
+    }
+    // Copy the data.
+    raft::copy(source_indices_->data_handle(),
+               source_indices.data_handle(),
+               source_indices.extent(0),
+               raft::resource::get_cuda_stream(res));
+  }
+
  private:
   cuvs::distance::DistanceType metric_;
-  raft::device_matrix<IdxT, int64_t, raft::row_major> graph_;
-  raft::device_matrix_view<const IdxT, int64_t, raft::row_major> graph_view_;
+  raft::device_matrix<graph_index_type, int64_t, raft::row_major> graph_;
+  raft::device_matrix_view<const graph_index_type, int64_t, raft::row_major> graph_view_;
   std::unique_ptr<neighbors::dataset<dataset_index_type>> dataset_;
+  // Mapping from internal graph node indices to the original user-provided indices.
+  std::optional<raft::device_vector<IdxT, int64_t>> source_indices_;
   // only float distances supported at the moment
   std::optional<raft::device_vector<float, int64_t>> dataset_norms_;
 

cpp/src/neighbors/detail/cagra/cagra_search.cuh

@@ -50,16 +50,21 @@ template <typename DataT,
           typename IndexT,
           typename DistanceT,
           typename CagraSampleFilterT,
-          typename OutputIdxT = IndexT>
-void search_main_core(raft::resources const& res,
-                      search_params params,
-                      const dataset_descriptor_host<DataT, IndexT, DistanceT>& dataset_desc,
-                      raft::device_matrix_view<const IndexT, int64_t, raft::row_major> graph,
-                      raft::device_matrix_view<const DataT, int64_t, raft::row_major> queries,
-                      raft::device_matrix_view<OutputIdxT, int64_t, raft::row_major> neighbors,
-                      raft::device_matrix_view<DistanceT, int64_t, raft::row_major> distances,
-                      CagraSampleFilterT sample_filter = CagraSampleFilterT())
+          typename SourceIdxT = IndexT,
+          typename OutputIdxT = SourceIdxT>
+void search_main_core(
+  raft::resources const& res,
+  search_params params,
+  const dataset_descriptor_host<DataT, IndexT, DistanceT>& dataset_desc,
+  raft::device_matrix_view<const IndexT, int64_t, raft::row_major> graph,
+  std::optional<raft::device_vector_view<const SourceIdxT, int64_t>> source_indices,
+  raft::device_matrix_view<const DataT, int64_t, raft::row_major> queries,
+  raft::device_matrix_view<OutputIdxT, int64_t, raft::row_major> neighbors,
+  raft::device_matrix_view<DistanceT, int64_t, raft::row_major> distances,
+  CagraSampleFilterT sample_filter = CagraSampleFilterT())
 {
+  static_assert(std::is_same_v<IndexT, uint32_t>,
+                "Only uint32_t is supported as the graph element type (internal index type)");
   RAFT_LOG_DEBUG("# dataset size = %lu, dim = %lu\n",
                  static_cast<size_t>(graph.extent(0)),
                  static_cast<size_t>(queries.extent(1)));
@@ -80,8 +85,9 @@ void search_main_core(raft::resources const& res,
     queries.extent(1));
 
   using CagraSampleFilterT_s = typename CagraSampleFilterT_Selector<CagraSampleFilterT>::type;
-  std::unique_ptr<search_plan_impl<DataT, IndexT, DistanceT, CagraSampleFilterT_s, OutputIdxT>>
-    plan = factory<DataT, IndexT, DistanceT, CagraSampleFilterT_s, OutputIdxT>::create(
+  std::unique_ptr<
+    search_plan_impl<DataT, IndexT, DistanceT, CagraSampleFilterT_s, SourceIdxT, OutputIdxT>>
+    plan = factory<DataT, IndexT, DistanceT, CagraSampleFilterT_s, SourceIdxT, OutputIdxT>::create(
       res, params, dataset_desc, queries.extent(1), graph.extent(0), graph.extent(1), topk);
 
   plan->check(topk);
@@ -104,6 +110,7 @@ void search_main_core(raft::resources const& res,
 
     (*plan)(res,
             graph,
+            source_indices,
             _topk_indices_ptr,
             _topk_distances_ptr,
             _query_ptr,
@@ -147,10 +154,12 @@ void search_main(raft::resources const& res,
                  CagraSampleFilterT sample_filter = CagraSampleFilterT())
 {
   // n_rows has the same type as the dataset index (the array extents type)
-  using ds_idx_type = decltype(index.data().n_rows());
+  using ds_idx_type    = decltype(index.data().n_rows());
+  using graph_idx_type = uint32_t;
   // Dispatch search parameters based on the dataset kind.
   if (auto* strided_dset = dynamic_cast<const strided_dataset<T, ds_idx_type>*>(&index.data());
       strided_dset != nullptr) {
+    // Search using a plain (strided) row-major dataset
     RAFT_EXPECTS(index.metric() != cuvs::distance::DistanceType::CosineExpanded ||
                    index.dataset_norms().has_value(),
                  "Dataset norms must be provided for CosineExpanded metric");
@@ -159,20 +168,36 @@ void search_main(raft::resources const& res,
     if (index.metric() == cuvs::distance::DistanceType::CosineExpanded) {
       dataset_norms_ptr = index.dataset_norms().value().data_handle();
     }
-    auto desc = dataset_descriptor_init_with_cache<T, IdxT, DistanceT>(
+    auto desc = dataset_descriptor_init_with_cache<T, graph_idx_type, DistanceT>(
       res, params, *strided_dset, index.metric(), dataset_norms_ptr);
-    search_main_core<T, IdxT, DistanceT, CagraSampleFilterT, OutputIdxT>(
-      res, params, desc, index.graph(), queries, neighbors, distances, sample_filter);
+    search_main_core<T, graph_idx_type, DistanceT, CagraSampleFilterT, IdxT, OutputIdxT>(
+      res,
+      params,
+      desc,
+      index.graph(),
+      index.source_indices(),
+      queries,
+      neighbors,
+      distances,
+      sample_filter);
   } else if (auto* vpq_dset = dynamic_cast<const vpq_dataset<float, ds_idx_type>*>(&index.data());
              vpq_dset != nullptr) {
     // Search using a compressed dataset
     RAFT_FAIL("FP32 VPQ dataset support is coming soon");
   } else if (auto* vpq_dset = dynamic_cast<const vpq_dataset<half, ds_idx_type>*>(&index.data());
              vpq_dset != nullptr) {
-    auto desc = dataset_descriptor_init_with_cache<T, IdxT, DistanceT>(
+    auto desc = dataset_descriptor_init_with_cache<T, graph_idx_type, DistanceT>(
       res, params, *vpq_dset, index.metric(), nullptr);
-    search_main_core<T, IdxT, DistanceT, CagraSampleFilterT, OutputIdxT>(
-      res, params, desc, index.graph(), queries, neighbors, distances, sample_filter);
+    search_main_core<T, graph_idx_type, DistanceT, CagraSampleFilterT, IdxT, OutputIdxT>(
+      res,
+      params,
+      desc,
+      index.graph(),
+      index.source_indices(),
+      queries,
+      neighbors,
+      distances,
+      sample_filter);
   } else if (auto* empty_dset = dynamic_cast<const empty_dataset<ds_idx_type>*>(&index.data());
              empty_dset != nullptr) {
     // Forgot to add a dataset.

cpp/src/neighbors/detail/cagra/cagra_serialize.cuh

@@ -35,7 +35,7 @@
 
 namespace cuvs::neighbors::cagra::detail {
 
-constexpr int serialization_version = 4;
+constexpr int serialization_version = 5;
 
 /**
  * Save the index to file.
@@ -71,14 +71,18 @@ void serialize(raft::resources const& res,
   raft::serialize_mdspan(res, os, index_.graph());
 
   include_dataset &= (index_.data().n_rows() > 0);
+  bool has_source_indices = index_.source_indices().has_value();
+  uint32_t content_map    = 0x1u * include_dataset + 0x2u * has_source_indices;
 
-  raft::serialize_scalar(res, os, include_dataset);
+  raft::serialize_scalar(res, os, content_map);
   if (include_dataset) {
     RAFT_LOG_DEBUG("Saving CAGRA index with dataset");
     neighbors::detail::serialize(res, os, index_.data());
   } else {
     RAFT_LOG_DEBUG("Saving CAGRA index WITHOUT dataset");
   }
+
+  if (has_source_indices) { raft::serialize_mdspan(res, os, index_.source_indices().value()); }
 }
 
 template <typename T, typename IdxT>
@@ -284,10 +288,21 @@ void deserialize(raft::resources const& res, std::istream& is, index<T, IdxT>* i
 
   *index_ = index<T, IdxT>(res, metric);
   index_->update_graph(res, raft::make_const_mdspan(graph.view()));
-  bool has_dataset = raft::deserialize_scalar<bool>(res, is);
+
+  auto content_map = raft::deserialize_scalar<uint32_t>(res, is);
+  bool has_dataset = content_map & 0x1u;
   if (has_dataset) {
     index_->update_dataset(res, cuvs::neighbors::detail::deserialize_dataset<int64_t>(res, is));
   }
+
+  bool has_source_indices = content_map & 0x2u;
+  if (has_source_indices) {
+    auto source_indices = raft::make_host_vector<IdxT, int64_t>(n_rows);
+    deserialize_mdspan(res, is, source_indices.view());
+    index_->update_source_indices(res, raft::make_const_mdspan(source_indices.view()));
+    raft::resource::sync_stream(
+      res);  // Don't let the vector out of the scope before the copy is finished
+  }
 }
 
 template <typename T, typename IdxT>

cpp/src/neighbors/detail/cagra/factory.cuh

@@ -30,14 +30,15 @@ template <typename DataT,
           typename IndexT,
           typename DistanceT,
           typename CagraSampleFilterT = cuvs::neighbors::filtering::none_sample_filter,
-          typename OutputIndexT       = IndexT>
+          typename SourceIndexT       = IndexT,
+          typename OutputIndexT       = SourceIndexT>
 class factory {
  public:
   /**
    * Create a search structure for dataset with dim features.
    */
   static std::unique_ptr<
-    search_plan_impl<DataT, IndexT, DistanceT, CagraSampleFilterT, OutputIndexT>>
+    search_plan_impl<DataT, IndexT, DistanceT, CagraSampleFilterT, SourceIndexT, OutputIndexT>>
   create(raft::resources const& res,
          search_params const& params,
          const dataset_descriptor_host<DataT, IndexT, DistanceT>& dataset_desc,
@@ -52,22 +53,25 @@ class factory {
 
  private:
   static std::unique_ptr<
-    search_plan_impl<DataT, IndexT, DistanceT, CagraSampleFilterT, OutputIndexT>>
+    search_plan_impl<DataT, IndexT, DistanceT, CagraSampleFilterT, SourceIndexT, OutputIndexT>>
   dispatch_kernel(raft::resources const& res,
                   search_plan_impl_base& plan,
                   const dataset_descriptor_host<DataT, IndexT, DistanceT>& dataset_desc)
   {
     if (plan.algo == search_algo::SINGLE_CTA) {
       return std::make_unique<
-        single_cta_search::search<DataT, IndexT, DistanceT, CagraSampleFilterT, OutputIndexT>>(
+        single_cta_search::
+          search<DataT, IndexT, DistanceT, CagraSampleFilterT, SourceIndexT, OutputIndexT>>(
         res, plan, dataset_desc, plan.dim, plan.dataset_size, plan.graph_degree, plan.topk);
     } else if (plan.algo == search_algo::MULTI_CTA) {
       return std::make_unique<
-        multi_cta_search::search<DataT, IndexT, DistanceT, CagraSampleFilterT, OutputIndexT>>(
+        multi_cta_search::
+          search<DataT, IndexT, DistanceT, CagraSampleFilterT, SourceIndexT, OutputIndexT>>(
         res, plan, dataset_desc, plan.dim, plan.dataset_size, plan.graph_degree, plan.topk);
     } else {
       return std::make_unique<
-        multi_kernel_search::search<DataT, IndexT, DistanceT, CagraSampleFilterT, OutputIndexT>>(
+        multi_kernel_search::
+          search<DataT, IndexT, DistanceT, CagraSampleFilterT, SourceIndexT, OutputIndexT>>(
         res, plan, dataset_desc, plan.dim, plan.dataset_size, plan.graph_degree, plan.topk);
     }
   }