Separate CAGRA index type from internal idx type (#1664)

This PR decouples CAGRA `IdxT` from the mdspan `IndexType` used in CAGRA API, by making the following change
`mdspan<Type, IdxT> -> mdspan<Type, int64_t>`.

In `cagra::index<T, IdxT>` the index type `IdxT` is the type that represents the neighbor indices returned during the ANN search, and the same type is also used as to store the neighborhood graph.

The build and search method take `mdspan` objects as input/output argument. These `mdspan`s have an `IndexType` that is used during offset calculation when indexing the matrix. Prior to this PR, we used `IdxT` as the mdspan index type.

There is a strong motivation to keep `IdxT` as `uint32_t` in order to minimize the memory footprint of the KNN graph. At the same time, the size and offset calculations of mdspans that represent dataset and knn graph require 64-bit indexing for large datasets. This PR changes the mdspan extent `IndexType` to `int64_t`.

This PR does not affect the performance of CAGRA: the search kernels already used correct index type. Only a few operations are affected by this change, and these are outside perf sensitive regions.

Authors:
  - Tamas Bela Feher (https://github.com/tfeher)

Approvers:
  - Corey J. Nolet (https://github.com/cjnolet)

URL: #1664

cpp/bench/ann/src/raft/raft_cagra_wrapper.h

@@ -71,7 +71,7 @@ class RaftCagra : public ANN<T> {
   AlgoProperty get_property() const override
   {
     AlgoProperty property;
-    property.dataset_memory_type      = MemoryType::Host;
+    property.dataset_memory_type      = MemoryType::HostMmap;
     property.query_memory_type        = MemoryType::Device;
     property.need_dataset_when_search = true;
     return property;
@@ -106,12 +106,13 @@ RaftCagra<T, IdxT>::RaftCagra(Metric metric, int dim, const BuildParam& param)
 template <typename T, typename IdxT>
 void RaftCagra<T, IdxT>::build(const T* dataset, size_t nrow, cudaStream_t)
 {
-  if (get_property().dataset_memory_type == MemoryType::Host) {
-    auto dataset_view = raft::make_host_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dimension_);
+  if (get_property().dataset_memory_type != MemoryType::Device) {
+    auto dataset_view =
+      raft::make_host_matrix_view<const T, int64_t>(dataset, IdxT(nrow), dimension_);
     index_.emplace(raft::neighbors::cagra::build(handle_, index_params_, dataset_view));
   } else {
     auto dataset_view =
-      raft::make_device_matrix_view<const T, IdxT>(dataset, IdxT(nrow), dimension_);
+      raft::make_device_matrix_view<const T, int64_t>(dataset, IdxT(nrow), dimension_);
     index_.emplace(raft::neighbors::cagra::build(handle_, index_params_, dataset_view));
   }
   return;
@@ -152,9 +153,10 @@ void RaftCagra<T, IdxT>::search(
     neighbors_IdxT = neighbors_storage.data();
   }
 
-  auto queries_view = raft::make_device_matrix_view<const T, IdxT>(queries, batch_size, dimension_);
-  auto neighbors_view = raft::make_device_matrix_view<IdxT, IdxT>(neighbors_IdxT, batch_size, k);
-  auto distances_view = raft::make_device_matrix_view<float, IdxT>(distances, batch_size, k);
+  auto queries_view =
+    raft::make_device_matrix_view<const T, int64_t>(queries, batch_size, dimension_);
+  auto neighbors_view = raft::make_device_matrix_view<IdxT, int64_t>(neighbors_IdxT, batch_size, k);
+  auto distances_view = raft::make_device_matrix_view<float, int64_t>(distances, batch_size, k);
 
   raft::neighbors::cagra::search(
     handle_, search_params_, *index_, queries_view, neighbors_view, distances_view);

cpp/bench/prims/neighbors/cagra_bench.cuh

@@ -47,9 +47,9 @@ struct CagraBench : public fixture {
   explicit CagraBench(const params& ps)
     : fixture(true),
       params_(ps),
-      queries_(make_device_matrix<T, IdxT>(handle, ps.n_queries, ps.n_dims)),
-      dataset_(make_device_matrix<T, IdxT>(handle, ps.n_samples, ps.n_dims)),
-      knn_graph_(make_device_matrix<IdxT, IdxT>(handle, ps.n_samples, ps.degree))
+      queries_(make_device_matrix<T, int64_t>(handle, ps.n_queries, ps.n_dims)),
+      dataset_(make_device_matrix<T, int64_t>(handle, ps.n_samples, ps.n_dims)),
+      knn_graph_(make_device_matrix<IdxT, int64_t>(handle, ps.n_samples, ps.degree))
   {
     // Generate random dataset and queriees
     raft::random::RngState state{42};
@@ -87,11 +87,11 @@ struct CagraBench : public fixture {
     search_params.thread_block_size = params_.block_size;
     search_params.num_parents       = params_.num_parents;
 
-    auto indices   = make_device_matrix<IdxT, IdxT>(handle, params_.n_queries, params_.k);
-    auto distances = make_device_matrix<float, IdxT>(handle, params_.n_queries, params_.k);
-    auto ind_v     = make_device_matrix_view<IdxT, IdxT, row_major>(
+    auto indices   = make_device_matrix<IdxT, int64_t>(handle, params_.n_queries, params_.k);
+    auto distances = make_device_matrix<float, int64_t>(handle, params_.n_queries, params_.k);
+    auto ind_v     = make_device_matrix_view<IdxT, int64_t, row_major>(
       indices.data_handle(), params_.n_queries, params_.k);
-    auto dist_v = make_device_matrix_view<float, IdxT, row_major>(
+    auto dist_v = make_device_matrix_view<float, int64_t, row_major>(
       distances.data_handle(), params_.n_queries, params_.k);
 
     auto queries_v = make_const_mdspan(queries_.view());
@@ -125,9 +125,9 @@ struct CagraBench : public fixture {
  private:
   const params params_;
   std::optional<const raft::neighbors::cagra::index<T, IdxT>> index_;
-  raft::device_matrix<T, IdxT, row_major> queries_;
-  raft::device_matrix<T, IdxT, row_major> dataset_;
-  raft::device_matrix<IdxT, IdxT, row_major> knn_graph_;
+  raft::device_matrix<T, int64_t, row_major> queries_;
+  raft::device_matrix<T, int64_t, row_major> dataset_;
+  raft::device_matrix<IdxT, int64_t, row_major> knn_graph_;
 };
 
 inline const std::vector<params> generate_inputs()

cpp/include/raft/neighbors/cagra.cuh

@@ -65,7 +65,7 @@ namespace raft::neighbors::cagra {
  * @endcode
  *
  * @tparam T data element type
- * @tparam IdxT type of the indices in the source dataset
+ * @tparam IdxT type of the dataset vector indices
  *
  * @param[in] res raft resources
  * @param[in] dataset a matrix view (host or device) to a row-major matrix [n_rows, dim]
@@ -76,19 +76,19 @@ namespace raft::neighbors::cagra {
  */
 template <typename DataT, typename IdxT, typename accessor>
 void build_knn_graph(raft::resources const& res,
-                     mdspan<const DataT, matrix_extent<IdxT>, row_major, accessor> dataset,
-                     raft::host_matrix_view<IdxT, IdxT, row_major> knn_graph,
+                     mdspan<const DataT, matrix_extent<int64_t>, row_major, accessor> dataset,
+                     raft::host_matrix_view<IdxT, int64_t, row_major> knn_graph,
                      std::optional<float> refine_rate                   = std::nullopt,
                      std::optional<ivf_pq::index_params> build_params   = std::nullopt,
                      std::optional<ivf_pq::search_params> search_params = std::nullopt)
 {
   using internal_IdxT = typename std::make_unsigned<IdxT>::type;
 
-  auto knn_graph_internal = make_host_matrix_view<internal_IdxT, internal_IdxT>(
+  auto knn_graph_internal = make_host_matrix_view<internal_IdxT, int64_t>(
     reinterpret_cast<internal_IdxT*>(knn_graph.data_handle()),
     knn_graph.extent(0),
     knn_graph.extent(1));
-  auto dataset_internal = mdspan<const DataT, matrix_extent<internal_IdxT>, row_major, accessor>(
+  auto dataset_internal = mdspan<const DataT, matrix_extent<int64_t>, row_major, accessor>(
     dataset.data_handle(), dataset.extent(0), dataset.extent(1));
 
   cagra::detail::build_knn_graph(
@@ -119,7 +119,7 @@ void build_knn_graph(raft::resources const& res,
  * @endcode
  *
  * @tparam DataT type of the data in the source dataset
- * @tparam IdxT type of the indices in the source dataset
+ * @tparam IdxT type of the dataset vector indices
  *
  * @param[in] res raft resources
  * @param[in] dataset a matrix view (host or device) to a row-major matrix [n_rows, dim]
@@ -133,20 +133,20 @@ template <typename DataT,
           typename g_accessor =
             host_device_accessor<std::experimental::default_accessor<IdxT>, memory_type::host>>
 void sort_knn_graph(raft::resources const& res,
-                    mdspan<const DataT, matrix_extent<IdxT>, row_major, d_accessor> dataset,
-                    mdspan<IdxT, matrix_extent<IdxT>, row_major, g_accessor> knn_graph)
+                    mdspan<const DataT, matrix_extent<int64_t>, row_major, d_accessor> dataset,
+                    mdspan<IdxT, matrix_extent<int64_t>, row_major, g_accessor> knn_graph)
 {
   using internal_IdxT = typename std::make_unsigned<IdxT>::type;
 
   using g_accessor_internal =
     host_device_accessor<std::experimental::default_accessor<internal_IdxT>, g_accessor::mem_type>;
   auto knn_graph_internal =
-    mdspan<internal_IdxT, matrix_extent<internal_IdxT>, row_major, g_accessor_internal>(
+    mdspan<internal_IdxT, matrix_extent<int64_t>, row_major, g_accessor_internal>(
       reinterpret_cast<internal_IdxT*>(knn_graph.data_handle()),
       knn_graph.extent(0),
       knn_graph.extent(1));
 
-  auto dataset_internal = mdspan<const DataT, matrix_extent<internal_IdxT>, row_major, d_accessor>(
+  auto dataset_internal = mdspan<const DataT, matrix_extent<int64_t>, row_major, d_accessor>(
     dataset.data_handle(), dataset.extent(0), dataset.extent(1));
 
   cagra::detail::graph::sort_knn_graph(res, dataset_internal, knn_graph_internal);
@@ -170,20 +170,20 @@ template <typename IdxT = uint32_t,
           typename g_accessor =
             host_device_accessor<std::experimental::default_accessor<IdxT>, memory_type::host>>
 void optimize(raft::resources const& res,
-              mdspan<IdxT, matrix_extent<IdxT>, row_major, g_accessor> knn_graph,
-              raft::host_matrix_view<IdxT, IdxT, row_major> new_graph)
+              mdspan<IdxT, matrix_extent<int64_t>, row_major, g_accessor> knn_graph,
+              raft::host_matrix_view<IdxT, int64_t, row_major> new_graph)
 {
   using internal_IdxT = typename std::make_unsigned<IdxT>::type;
 
-  auto new_graph_internal = raft::make_host_matrix_view<internal_IdxT, internal_IdxT>(
+  auto new_graph_internal = raft::make_host_matrix_view<internal_IdxT, int64_t>(
     reinterpret_cast<internal_IdxT*>(new_graph.data_handle()),
     new_graph.extent(0),
     new_graph.extent(1));
 
   using g_accessor_internal =
     host_device_accessor<std::experimental::default_accessor<internal_IdxT>, memory_type::host>;
   auto knn_graph_internal =
-    mdspan<internal_IdxT, matrix_extent<internal_IdxT>, row_major, g_accessor_internal>(
+    mdspan<internal_IdxT, matrix_extent<int64_t>, row_major, g_accessor_internal>(
       reinterpret_cast<internal_IdxT*>(knn_graph.data_handle()),
       knn_graph.extent(0),
       knn_graph.extent(1));
@@ -237,7 +237,7 @@ template <typename T,
             host_device_accessor<std::experimental::default_accessor<T>, memory_type::host>>
 index<T, IdxT> build(raft::resources const& res,
                      const index_params& params,
-                     mdspan<const T, matrix_extent<IdxT>, row_major, Accessor> dataset)
+                     mdspan<const T, matrix_extent<int64_t>, row_major, Accessor> dataset)
 {
   size_t intermediate_degree = params.intermediate_graph_degree;
   size_t graph_degree        = params.graph_degree;
@@ -256,11 +256,11 @@ index<T, IdxT> build(raft::resources const& res,
     graph_degree = intermediate_degree;
   }
 
-  auto knn_graph = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), intermediate_degree);
+  auto knn_graph = raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), intermediate_degree);
 
   build_knn_graph(res, dataset, knn_graph.view());
 
-  auto cagra_graph = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), graph_degree);
+  auto cagra_graph = raft::make_host_matrix<IdxT, int64_t>(dataset.extent(0), graph_degree);
 
   optimize<IdxT>(res, knn_graph.view(), cagra_graph.view());
 
@@ -289,9 +289,9 @@ template <typename T, typename IdxT>
 void search(raft::resources const& res,
             const search_params& params,
             const index<T, IdxT>& idx,
-            raft::device_matrix_view<const T, IdxT, row_major> queries,
-            raft::device_matrix_view<IdxT, IdxT, row_major> neighbors,
-            raft::device_matrix_view<float, IdxT, row_major> distances)
+            raft::device_matrix_view<const T, int64_t, row_major> queries,
+            raft::device_matrix_view<IdxT, int64_t, row_major> neighbors,
+            raft::device_matrix_view<float, int64_t, row_major> distances)
 {
   RAFT_EXPECTS(
     queries.extent(0) == neighbors.extent(0) && queries.extent(0) == distances.extent(0),
@@ -303,13 +303,13 @@ void search(raft::resources const& res,
                "Number of query dimensions should equal number of dimensions in the index.");
 
   using internal_IdxT   = typename std::make_unsigned<IdxT>::type;
-  auto queries_internal = raft::make_device_matrix_view<const T, internal_IdxT, row_major>(
+  auto queries_internal = raft::make_device_matrix_view<const T, int64_t, row_major>(
     queries.data_handle(), queries.extent(0), queries.extent(1));
-  auto neighbors_internal = raft::make_device_matrix_view<internal_IdxT, internal_IdxT, row_major>(
+  auto neighbors_internal = raft::make_device_matrix_view<internal_IdxT, int64_t, row_major>(
     reinterpret_cast<internal_IdxT*>(neighbors.data_handle()),
     neighbors.extent(0),
     neighbors.extent(1));
-  auto distances_internal = raft::make_device_matrix_view<float, internal_IdxT, row_major>(
+  auto distances_internal = raft::make_device_matrix_view<float, int64_t, row_major>(
     distances.data_handle(), distances.extent(0), distances.extent(1));
 
   cagra::detail::search_main<T, internal_IdxT, IdxT>(