Merge branch 'main' into java/int-types-tests

Author: mythrocks

conda/environments/bench_ann_cuda-129_arch-aarch64.yaml

@@ -35,7 +35,7 @@ dependencies:
 - matplotlib-base>=3.9
 - nccl>=2.19
 - ninja
-- nlohmann_json>=3.11.2
+- nlohmann_json>=3.12.0
 - openblas
 - pandas
 - pylibraft==25.12.*,>=0.0.0a0

conda/environments/bench_ann_cuda-129_arch-x86_64.yaml

@@ -38,7 +38,7 @@ dependencies:
 - mkl-devel=2023
 - nccl>=2.19
 - ninja
-- nlohmann_json>=3.11.2
+- nlohmann_json>=3.12.0
 - openblas
 - pandas
 - pylibraft==25.12.*,>=0.0.0a0

conda/environments/bench_ann_cuda-130_arch-aarch64.yaml

@@ -35,7 +35,7 @@ dependencies:
 - matplotlib-base>=3.9
 - nccl>=2.19
 - ninja
-- nlohmann_json>=3.11.2
+- nlohmann_json>=3.12.0
 - openblas
 - pandas
 - pylibraft==25.12.*,>=0.0.0a0

conda/environments/bench_ann_cuda-130_arch-x86_64.yaml

@@ -38,7 +38,7 @@ dependencies:
 - mkl-devel=2023
 - nccl>=2.19
 - ninja
-- nlohmann_json>=3.11.2
+- nlohmann_json>=3.12.0
 - openblas
 - pandas
 - pylibraft==25.12.*,>=0.0.0a0

conda/recipes/cuvs-bench-cpu/conda_build_config.yaml

@@ -20,4 +20,4 @@ h5py_version:
   - ">=3.8.0"
 
 nlohmann_json_version:
-  - ">=3.11.2"
+  - ">=3.12.0"

cpp/cmake/thirdparty/get_nlohmann_json.cmake

@@ -1,5 +1,5 @@
 #=============================================================================
-# Copyright (c) 2024, NVIDIA CORPORATION.
+# Copyright (c) 2024-2025, NVIDIA CORPORATION.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -35,8 +35,8 @@ function(find_and_configure_nlohmann_json)
 
 endfunction()
 
-find_and_configure_nlohmann_json(VERSION  3.11.2
+find_and_configure_nlohmann_json(VERSION  3.12.0
         FORK             nlohmann
-        PINNED_TAG       v3.11.2
+        PINNED_TAG       v3.12.0
         EXCLUDE_FROM_ALL ON
         )

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_build.cuh

@@ -133,7 +133,7 @@ void build_knn_graph(
 
   uint32_t node_degree = knn_graph.extent(1);
   raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope(
-    "cagra::build_graph(%zu, %zu, %u)",
+    "cagra::build_knn_graph<IVF-PQ>(%zu, %zu, %u)",
     size_t(dataset.extent(0)),
     size_t(dataset.extent(1)),
     node_degree);
@@ -391,6 +391,12 @@ void build_knn_graph(
   raft::host_matrix_view<IdxT, int64_t, raft::row_major> knn_graph,
   cuvs::neighbors::nn_descent::index_params build_params)
 {
+  raft::common::nvtx::range<cuvs::common::nvtx::domain::cuvs> fun_scope(
+    "cagra::build_knn_graph<NN-DESCENT>(%zu, %zu, %u)",
+    size_t(dataset.extent(0)),
+    size_t(dataset.extent(1)),
+    size_t(knn_graph.extent(1)));
+
   std::optional<raft::host_matrix_view<IdxT, int64_t, row_major>> graph_view = knn_graph;
   auto nn_descent_idx = cuvs::neighbors::nn_descent::build(res, build_params, dataset, graph_view);
 

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>