CAGRA pad dataset for 128bit vectorized load

cpp/include/raft/neighbors/cagra.cuh

@@ -237,21 +237,28 @@ index<T, IdxT> build(raft::resources const& res,
                      const index_params& params,
                      mdspan<const T, matrix_extent<IdxT>, row_major, Accessor> dataset)
 {
-  size_t degree = params.intermediate_graph_degree;
-  if (degree >= static_cast<size_t>(dataset.extent(0))) {
+  size_t intermediate_degree = params.intermediate_graph_degree;
+  size_t graph_degree        = params.graph_degree;
+  if (intermediate_degree >= static_cast<size_t>(dataset.extent(0))) {
     RAFT_LOG_WARN(
       "Intermediate graph degree cannot be larger than dataset size, reducing it to %lu",
       dataset.extent(0));
-    degree = dataset.extent(0) - 1;
+    intermediate_degree = dataset.extent(0) - 1;
+  }
+  if (intermediate_degree < graph_degree) {
+    RAFT_LOG_WARN(
+      "Graph degree (%lu) cannot be larger than intermediate graph degree (%lu), reducing "
+      "graph_degree.",
+      graph_degree,
+      intermediate_degree);
+    graph_degree = intermediate_degree;
   }
-  RAFT_EXPECTS(degree >= params.graph_degree,
-               "Intermediate graph degree cannot be smaller than final graph degree");
 
-  auto knn_graph = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), degree);
+  auto knn_graph = raft::make_host_matrix<IdxT, IdxT>(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), params.graph_degree);
+  auto cagra_graph = raft::make_host_matrix<IdxT, IdxT>(dataset.extent(0), graph_degree);
 
   prune<IdxT>(res, knn_graph.view(), cagra_graph.view());
 
@@ -290,7 +297,6 @@ void search(raft::resources const& res,
 
   RAFT_EXPECTS(neighbors.extent(1) == distances.extent(1),
                "Number of columns in output neighbors and distances matrices must equal k");
-
   RAFT_EXPECTS(queries.extent(1) == idx.dim(),
                "Number of query dimensions should equal number of dimensions in the index.");
 

cpp/include/raft/neighbors/cagra_types.hpp

@@ -26,6 +26,7 @@
 #include <raft/core/resources.hpp>
 #include <raft/distance/distance_types.hpp>
 #include <raft/util/integer_utils.hpp>
+#include <raft/util/pow2_utils.cuh>
 
 #include <memory>
 #include <optional>
@@ -82,8 +83,6 @@ struct search_params : ann::search_params {
   /** Lower limit of search iterations. */
   size_t min_iterations = 0;
 
-  /** Bit length for reading the dataset vectors. 0, 64 or 128. Auto selection when 0. */
-  size_t load_bit_length = 0;
   /** Thread block size. 0, 64, 128, 256, 512, 1024. Auto selection when 0. */
   size_t thread_block_size = 0;
   /** Hashmap type. Auto selection when AUTO. */
@@ -113,6 +112,7 @@ static_assert(std::is_aggregate_v<search_params>);
  */
 template <typename T, typename IdxT>
 struct index : ann::index {
+  using AlignDim = raft::Pow2<16 / sizeof(T)>;
   static_assert(!raft::is_narrowing_v<uint32_t, IdxT>,
                 "IdxT must be able to represent all values of uint32_t");
 
@@ -124,12 +124,15 @@ struct index : ann::index {
   }
 
   // /** Total length of the index. */
-  [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT { return dataset_.extent(0); }
+  [[nodiscard]] constexpr inline auto size() const noexcept -> IdxT
+  {
+    return dataset_view_.extent(0);
+  }
 
   /** Dimensionality of the data. */
   [[nodiscard]] constexpr inline auto dim() const noexcept -> uint32_t
   {
-    return dataset_.extent(1);
+    return dataset_view_.extent(1);
   }
   /** Graph degree */
   [[nodiscard]] constexpr inline auto graph_degree() const noexcept -> uint32_t
@@ -138,9 +141,10 @@ struct index : ann::index {
   }
 
   /** Dataset [size, dim] */
-  [[nodiscard]] inline auto dataset() const noexcept -> device_matrix_view<const T, IdxT, row_major>
+  [[nodiscard]] inline auto dataset() const noexcept
+    -> device_matrix_view<const T, IdxT, layout_stride>
   {
-    return dataset_.view();
+    return dataset_view_;
   }
 
   /** neighborhood graph [size, graph-degree] */
@@ -179,15 +183,36 @@ struct index : ann::index {
         mdspan<IdxT, matrix_extent<IdxT>, row_major, graph_accessor> knn_graph)
     : ann::index(),
       metric_(metric),
-      dataset_(make_device_matrix<T, IdxT>(res, dataset.extent(0), dataset.extent(1))),
+      dataset_(
+        make_device_matrix<T, IdxT>(res, dataset.extent(0), AlignDim::roundUp(dataset.extent(1)))),
       graph_(make_device_matrix<IdxT, IdxT>(res, knn_graph.extent(0), knn_graph.extent(1)))
   {
     RAFT_EXPECTS(dataset.extent(0) == knn_graph.extent(0),
                  "Dataset and knn_graph must have equal number of rows");
-    raft::copy(dataset_.data_handle(),
-               dataset.data_handle(),
-               dataset.size(),
-               resource::get_cuda_stream(res));
+    if (dataset_.extent(1) == dataset.extent(1)) {
+      raft::copy(dataset_.data_handle(),
+                 dataset.data_handle(),
+                 dataset.size(),
+                 resource::get_cuda_stream(res));
+    } else {
+      // copy with padding
+      RAFT_CUDA_TRY(cudaMemsetAsync(
+        dataset_.data_handle(), 0, dataset_.size() * sizeof(T), resource::get_cuda_stream(res)));
+      RAFT_CUDA_TRY(cudaMemcpy2DAsync(dataset_.data_handle(),
+                                      sizeof(T) * dataset_.extent(1),
+                                      dataset.data_handle(),
+                                      sizeof(T) * dataset.extent(1),
+                                      sizeof(T) * dataset.extent(1),
+                                      dataset.extent(0),
+                                      cudaMemcpyDefault,
+                                      resource::get_cuda_stream(res)));
+    }
+    dataset_view_ = make_device_strided_matrix_view<T, IdxT>(
+      dataset_.data_handle(), dataset_.extent(0), dataset.extent(1), dataset_.extent(1));
+    RAFT_LOG_DEBUG("CAGRA dataset strided matrix view %zux%zu, stride %zu",
+                   static_cast<size_t>(dataset_view_.extent(0)),
+                   static_cast<size_t>(dataset_view_.extent(1)),
+                   static_cast<size_t>(dataset_view_.stride(0)));
     raft::copy(graph_.data_handle(),
                knn_graph.data_handle(),
                knn_graph.size(),
@@ -199,6 +224,7 @@ struct index : ann::index {
   raft::distance::DistanceType metric_;
   raft::device_matrix<T, IdxT, row_major> dataset_;
   raft::device_matrix<IdxT, IdxT, row_major> graph_;
+  raft::device_matrix_view<T, IdxT, layout_stride> dataset_view_;
 };
 
 /** @} */

cpp/include/raft/neighbors/detail/cagra/cagra_build.cuh

@@ -46,10 +46,6 @@ void build_knn_graph(raft::resources const& res,
                      std::optional<ivf_pq::index_params> build_params   = std::nullopt,
                      std::optional<ivf_pq::search_params> search_params = std::nullopt)
 {
-  RAFT_EXPECTS(
-    dataset.extent(1) * sizeof(DataT) % 8 == 0,
-    "Dataset rows are expected to have at least 8 bytes alignment. Try padding feature dims.");
-
   RAFT_EXPECTS(!build_params || build_params->metric == distance::DistanceType::L2Expanded,
                "Currently only L2Expanded metric is supported");
 

cpp/include/raft/neighbors/detail/cagra/cagra_search.cuh

@@ -27,8 +27,6 @@
 #include <rmm/cuda_stream_view.hpp>
 
 #include "factory.cuh"
-#include "search_multi_cta.cuh"
-#include "search_multi_kernel.cuh"
 #include "search_plan.cuh"
 #include "search_single_cta.cuh"
 
@@ -92,8 +90,11 @@ void search_main(raft::resources const& res,
         : nullptr;
     uint32_t* _num_executed_iterations = nullptr;
 
-    auto dataset_internal = raft::make_device_matrix_view<const T, internal_IdxT, row_major>(
-      index.dataset().data_handle(), index.dataset().extent(0), index.dataset().extent(1));
+    auto dataset_internal = make_device_strided_matrix_view<const T, internal_IdxT, row_major>(
+      index.dataset().data_handle(),
+      index.dataset().extent(0),
+      index.dataset().extent(1),
+      index.dataset().stride(0));
     auto graph_internal =
       raft::make_device_matrix_view<const internal_IdxT, internal_IdxT, row_major>(
         reinterpret_cast<const internal_IdxT*>(index.graph().data_handle()),

cpp/include/raft/neighbors/detail/cagra/cagra_serialize.cuh

@@ -25,7 +25,7 @@
 namespace raft::neighbors::experimental::cagra::detail {
 
 // Serialization version 1.
-constexpr int serialization_version = 1;
+constexpr int serialization_version = 2;
 
 // NB: we wrap this check in a struct, so that the updated RealSize is easy to see in the error
 // message.
@@ -36,7 +36,8 @@ 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>), 136>;
+constexpr size_t expected_size = 176;
+template struct check_index_layout<sizeof(index<double, std::uint64_t>), expected_size>;
 
 /**
  * Save the index to file.
@@ -59,7 +60,19 @@ void serialize(raft::resources const& res, std::ostream& os, const index<T, IdxT
   serialize_scalar(res, os, index_.dim());
   serialize_scalar(res, os, index_.graph_degree());
   serialize_scalar(res, os, index_.metric());
-  serialize_mdspan(res, os, index_.dataset());
+  auto dataset = index_.dataset();
+  // Remove padding before saving the dataset
+  auto host_dataset = make_host_matrix<T, IdxT>(dataset.extent(0), dataset.extent(1));
+  RAFT_CUDA_TRY(cudaMemcpy2DAsync(host_dataset.data_handle(),
+                                  sizeof(T) * host_dataset.extent(1),
+                                  dataset.data_handle(),
+                                  sizeof(T) * dataset.stride(0),
+                                  sizeof(T) * host_dataset.extent(1),
+                                  dataset.extent(0),
+                                  cudaMemcpyDefault,
+                                  resource::get_cuda_stream(res)));
+  resource::sync_stream(res);
+  serialize_mdspan(res, os, host_dataset.view());
   serialize_mdspan(res, os, index_.graph());
 }
 
@@ -100,7 +113,6 @@ auto deserialize(raft::resources const& res, std::istream& is) -> index<T, IdxT>
 
   auto dataset = raft::make_host_matrix<T, IdxT>(n_rows, dim);
   auto graph   = raft::make_host_matrix<IdxT, IdxT>(n_rows, graph_degree);
-
   deserialize_mdspan(res, is, dataset.view());
   deserialize_mdspan(res, is, graph.view());
 

cpp/include/raft/neighbors/detail/cagra/compute_distance.hpp

@@ -56,6 +56,7 @@ _RAFT_DEVICE void compute_distance_to_random_nodes(
   const DATA_T* const dataset_ptr,         // [dataset_size, dataset_dim]
   const std::size_t dataset_dim,
   const std::size_t dataset_size,
+  const std::size_t dataset_ld,
   const std::size_t num_pickup,
   const unsigned num_distilation,
   const uint64_t rand_xor_mask,
@@ -93,7 +94,7 @@ _RAFT_DEVICE void compute_distance_to_random_nodes(
         for (uint32_t e = 0; e < nelem; e++) {
           const uint32_t k = (lane_id + (TEAM_SIZE * e)) * vlen;
           if (k >= dataset_dim) break;
-          dl_buff[e].load = ((LOAD_T*)(dataset_ptr + k + (dataset_dim * seed_index)))[0];
+          dl_buff[e].load = ((LOAD_T*)(dataset_ptr + k + (dataset_ld * seed_index)))[0];
         }
 #pragma unroll
         for (uint32_t e = 0; e < nelem; e++) {
@@ -146,6 +147,7 @@ _RAFT_DEVICE void compute_distance_to_child_nodes(INDEX_T* const result_child_in
                                                   // [dataset_dim, dataset_size]
                                                   const DATA_T* const dataset_ptr,
                                                   const std::size_t dataset_dim,
+                                                  const std::size_t dataset_ld,
                                                   // [knn_k, dataset_size]
                                                   const INDEX_T* const knn_graph,
                                                   const std::uint32_t knn_k,
@@ -215,7 +217,7 @@ _RAFT_DEVICE void compute_distance_to_child_nodes(INDEX_T* const result_child_in
       for (unsigned e = 0; e < nelem; e++) {
         const unsigned k = (lane_id + (TEAM_SIZE * e)) * vlen;
         if (k >= dataset_dim) break;
-        dl_buff[e].load = ((LOAD_T*)(dataset_ptr + k + (dataset_dim * child_id)))[0];
+        dl_buff[e].load = ((LOAD_T*)(dataset_ptr + k + (dataset_ld * child_id)))[0];
       }
 #pragma unroll
       for (unsigned e = 0; e < nelem; e++) {