Reuse raft::mdarray/mdspan directly in dense dataset storage

c/src/neighbors/cagra.hpp

@@ -20,7 +20,7 @@ void convert_c_search_params(cuvsCagraSearchParams params,
 void* cagra_c_api_index_ptr(cuvsCagraIndex const* idx);
 
 namespace detail {
-template <typename T, typename IdxT, cuvs::neighbors::ann_dataset_view DatasetViewT>
+template <typename T, typename IdxT, cuvs::core::dataset_like DatasetViewT>
 int64_t merged_dataset_size(
   raft::resources const& res,
   std::vector<cuvs::neighbors::cagra::index<T, IdxT, DatasetViewT>*> const& indices,

c/src/neighbors/mg_cagra.cpp

@@ -75,8 +75,8 @@ static void with_mg_index_by_layout(mg_cagra_c_api_index_box* box,
 template <typename T, typename Fn>
 static void with_device_padded_dataset_view(cuvsDataset_t dataset, Fn&& fn)
 {
-  using owner_t = cuvs::neighbors::device_padded_dataset<T, int64_t>;
-  using view_t  = cuvs::neighbors::device_padded_dataset_view<T, int64_t>;
+  using owner_t = cuvs::core::device_padded_dataset<T, int64_t>;
+  using view_t  = cuvs::core::device_padded_dataset_view<T, int64_t>;
   if (dataset->is_owning) {
     auto* owner = reinterpret_cast<owner_t*>(dataset->addr);
     auto view   = owner->as_dataset_view();
@@ -204,13 +204,13 @@ void* _mg_build(cuvsResources_t res,
 
   if (layout == mg_cagra_dataset_layout::device_padded) {
     using padded_ann_t = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
-    auto padded_mds    = cuvs::neighbors::make_host_padded_dataset_view(mds);
+    auto padded_mds    = cuvs::core::make_host_padded_dataset_view(mds);
     auto* mg_index     = new mg_cagra_index_t<T, padded_ann_t>(
       cuvs::neighbors::cagra::build(*res_ptr, mg_params, padded_mds));
     return make_mg_cagra_box<T, padded_ann_t>(mg_index, mg_cagra_dataset_layout::device_padded);
   }
   using standard_ann_t = cuvs::neighbors::cagra::device_standard_index<T, uint32_t>;
-  auto standard_mds    = cuvs::neighbors::make_host_standard_dataset_view(mds);
+  auto standard_mds    = cuvs::core::make_host_standard_dataset_view(mds);
   auto* mg_index       = new mg_cagra_index_t<T, standard_ann_t>(
     cuvs::neighbors::cagra::build(*res_ptr, mg_params, standard_mds));
   return make_mg_cagra_box<T, standard_ann_t>(mg_index, mg_cagra_dataset_layout::device_standard);
@@ -296,13 +296,13 @@ void _mg_extend(cuvsResources_t res,
     using padded_ann_t = cuvs::neighbors::cagra::device_padded_index<T, uint32_t>;
     auto* mg_index_ptr =
       reinterpret_cast<mg_cagra_index_t<T, padded_ann_t>*>(box->index_ptr);
-    auto new_vectors = cuvs::neighbors::make_host_padded_dataset_view(new_vectors_mds);
+    auto new_vectors = cuvs::core::make_host_padded_dataset_view(new_vectors_mds);
     cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors, new_indices_mds);
   } else {
     using standard_ann_t = cuvs::neighbors::cagra::device_standard_index<T, uint32_t>;
     auto* mg_index_ptr =
       reinterpret_cast<mg_cagra_index_t<T, standard_ann_t>*>(box->index_ptr);
-    auto new_vectors = cuvs::neighbors::make_host_standard_dataset_view(new_vectors_mds);
+    auto new_vectors = cuvs::core::make_host_standard_dataset_view(new_vectors_mds);
     cuvs::neighbors::cagra::extend(*res_ptr, *mg_index_ptr, new_vectors, new_indices_mds);
   }
 }
@@ -372,28 +372,28 @@ extern "C" cuvsError_t cuvsMultiGpuCagraBuild(cuvsResources_t res,
     if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 32) {
       auto mds = cuvs::core::from_dlpack<raft::host_matrix_view<const float, int64_t, raft::row_major>>(
         dataset_tensor);
-      auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)
+      auto layout = cuvs::core::matrix_has_padded_row_width(mds)
                       ? mg_cagra_dataset_layout::device_padded
                       : mg_cagra_dataset_layout::device_standard;
       index->addr = reinterpret_cast<uintptr_t>(_mg_build<float>(res, *params, dataset_tensor, layout));
     } else if (dataset.dtype.code == kDLFloat && dataset.dtype.bits == 16) {
       auto mds = cuvs::core::from_dlpack<raft::host_matrix_view<const half, int64_t, raft::row_major>>(
         dataset_tensor);
-      auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)
+      auto layout = cuvs::core::matrix_has_padded_row_width(mds)
                       ? mg_cagra_dataset_layout::device_padded
                       : mg_cagra_dataset_layout::device_standard;
       index->addr = reinterpret_cast<uintptr_t>(_mg_build<half>(res, *params, dataset_tensor, layout));
     } else if (dataset.dtype.code == kDLInt && dataset.dtype.bits == 8) {
       auto mds = cuvs::core::from_dlpack<raft::host_matrix_view<const int8_t, int64_t, raft::row_major>>(
         dataset_tensor);
-      auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)
+      auto layout = cuvs::core::matrix_has_padded_row_width(mds)
                       ? mg_cagra_dataset_layout::device_padded
                       : mg_cagra_dataset_layout::device_standard;
       index->addr = reinterpret_cast<uintptr_t>(_mg_build<int8_t>(res, *params, dataset_tensor, layout));
     } else if (dataset.dtype.code == kDLUInt && dataset.dtype.bits == 8) {
       auto mds = cuvs::core::from_dlpack<raft::host_matrix_view<const uint8_t, int64_t, raft::row_major>>(
         dataset_tensor);
-      auto layout = cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)
+      auto layout = cuvs::core::matrix_has_padded_row_width(mds)
                       ? mg_cagra_dataset_layout::device_padded
                       : mg_cagra_dataset_layout::device_standard;
       index->addr = reinterpret_cast<uintptr_t>(_mg_build<uint8_t>(res, *params, dataset_tensor, layout));

c/src/neighbors/tiered_index.cpp

@@ -140,8 +140,8 @@ void* _build(cuvsResources_t res, cuvsTieredIndexParams params, DLManagedTensor*
     case CUVS_TIERED_INDEX_ALGO_CAGRA: {
       auto build_params = tiered_index::index_params<cagra::index_params>();
       convert_c_index_params(params, dataset.shape[0], dataset.shape[1], &build_params);
-      if (cuvs::neighbors::matrix_row_width_matches_cagra_required(mds)) {
-        auto padded_view = cuvs::neighbors::make_device_padded_dataset_view(*res_ptr, mds);
+      if (cuvs::core::matrix_has_padded_row_width(mds)) {
+        auto padded_view = cuvs::core::make_device_padded_dataset_view(*res_ptr, mds);
         auto* ptr = new tiered_index::index<cagra::device_padded_index<T, uint32_t>>(
           tiered_index::build(*res_ptr, build_params, padded_view));
         return make_tiered_index_box(

c/src/preprocessing/quantize/pq.cpp

@@ -1,5 +1,5 @@
 /*
- * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
+ * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
  * SPDX-License-Identifier: Apache-2.0
  */
 
@@ -242,9 +242,10 @@ extern "C" cuvsError_t cuvsProductQuantizerGetPqCodebook(cuvsProductQuantizer_t
     if (quantizer != nullptr) {
       auto quant_addr = quantizer->addr;
       if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) {
-        auto pq_mdspan =
+        auto const& vpq =
           (reinterpret_cast<cuvs::preprocessing::quantize::pq::quantizer<float>*>(quant_addr))
-            ->vpq_codebooks.pq_code_book.view();
+            ->vpq_codebooks.data();
+        auto pq_mdspan = vpq.pq_code_book.view();
         cuvs::core::to_dlpack(pq_mdspan, pq_codebook);
       } else {
         RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d",
@@ -264,9 +265,10 @@ extern "C" cuvsError_t cuvsProductQuantizerGetVqCodebook(cuvsProductQuantizer_t
     if (quantizer != nullptr) {
       auto quant_addr = quantizer->addr;
       if (quantizer->dtype.code == kDLFloat && quantizer->dtype.bits == 32) {
-        auto pq_mdspan =
+        auto const& vpq =
           (reinterpret_cast<cuvs::preprocessing::quantize::pq::quantizer<float>*>(quant_addr))
-            ->vpq_codebooks.vq_code_book.view();
+            ->vpq_codebooks.data();
+        auto pq_mdspan = vpq.vq_code_book.view();
         cuvs::core::to_dlpack(pq_mdspan, vq_codebook);
       } else {
         RAFT_FAIL("Unsupported quantizer dtype: %d and bits: %d",

cpp/bench/ann/src/cuvs/cuvs_ann_bench_param_parser.h

@@ -271,7 +271,8 @@ void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::nn_descent::
   }
 }
 
-inline void parse_build_param(const nlohmann::json& conf, cuvs::neighbors::vpq_params& param)
+inline void parse_build_param(const nlohmann::json& conf,
+                              cuvs::preprocessing::quantize::pq::vpq_params& param)
 {
   if (conf.contains("pq_bits")) { param.pq_bits = conf.at("pq_bits"); }
   if (conf.contains("pq_dim")) { param.pq_dim = conf.at("pq_dim"); }
@@ -445,7 +446,7 @@ void parse_build_param(const nlohmann::json& conf,
 
   nlohmann::json comp_search_conf = collect_conf_with_prefix(conf, "compression_");
   if (!comp_search_conf.empty()) {
-    auto vpq_pams = param.compression.value_or(cuvs::neighbors::vpq_params{});
+    auto vpq_pams = param.compression.value_or(cuvs::preprocessing::quantize::pq::vpq_params{});
     parse_build_param(comp_search_conf, vpq_pams);
     param.compression.emplace(vpq_pams);
   }

cpp/bench/ann/src/cuvs/cuvs_cagra_wrapper.h

@@ -85,11 +85,11 @@ template <typename T, typename SrcT>
 auto make_padded_view(const raft::resources& res,
                       SrcT src,
                       raft::device_matrix<T, int64_t, raft::row_major>& buffer)
-  -> cuvs::neighbors::device_padded_dataset_view<T, int64_t>
+  -> cuvs::core::device_padded_dataset_view<T, int64_t>
 {
   if constexpr (SrcT::accessor_type::is_device_accessible) {
-    if (cuvs::neighbors::matrix_row_width_matches_cagra_required(src)) {
-      return cuvs::neighbors::make_device_padded_dataset_view(res, src);
+    if (cuvs::core::matrix_has_padded_row_width(src)) {
+      return cuvs::core::make_device_padded_dataset_view(res, src);
     }
   }
   cuvs::neighbors::cagra::detail::copy_with_padding(res, buffer, src);
@@ -165,9 +165,9 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
     using dataset_dependent_params = std::function<cuvs::neighbors::cagra::index_params(
       raft::matrix_extent<int64_t>, cuvs::distance::DistanceType)>;
     dataset_dependent_params cagra_params;
-    std::optional<cuvs::neighbors::vpq_params> compression = std::nullopt;
-    size_t num_dataset_splits                              = 1;
-    CagraMergeType merge_type                              = CagraMergeType::kPhysical;
+    std::optional<cuvs::preprocessing::quantize::pq::vpq_params> compression = std::nullopt;
+    size_t num_dataset_splits                                                = 1;
+    CagraMergeType merge_type = CagraMergeType::kPhysical;
     cuvs::neighbors::cagra::merge_params merge_params;
   };
 
@@ -267,7 +267,8 @@ class cuvs_cagra : public algo<T>, public algo_gpu {
   std::shared_ptr<std::vector<raft::device_matrix<T, int64_t, raft::row_major>>>
     sub_dataset_buffers_ =
       std::make_shared<std::vector<raft::device_matrix<T, int64_t, raft::row_major>>>();
-  std::shared_ptr<cuvs::neighbors::device_vpq_dataset<half, int64_t>> vpq_dataset_;
+  std::shared_ptr<cuvs::preprocessing::quantize::pq::device_vpq_dataset<half, int64_t>>
+    vpq_dataset_;
   std::shared_ptr<cuvs::neighbors::cagra::device_pq_index<T, IdxT, half>> vpq_index_;
 
   inline rmm::device_async_resource_ref get_mr(AllocatorType mem_type)
@@ -302,7 +303,7 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
       // (and dispatches to ACE when it is configured), so the dataset is not uploaded here at all.
       // The single device copy needed for search is made later, by set_search_param.
       host_index_ = std::make_shared<host_index_type>(cuvs::neighbors::cagra::build(
-        handle_, host_params, cuvs::neighbors::make_host_standard_dataset_view(dataset_view_host)));
+        handle_, host_params, cuvs::core::make_host_standard_dataset_view(dataset_view_host)));
       index_ =
         std::make_shared<index_type>(detail::to_graph_only_index<T, IdxT>(handle_, *host_index_));
       // The graph moved into the index along with the file descriptors; nothing views the host
@@ -348,7 +349,7 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
           // As in the single-split case: graph only, the rows are uploaded by set_search_dataset.
           sub_host_indices_.push_back(
             std::make_shared<host_index_type>(cuvs::neighbors::cagra::build(
-              handle_, host_params, cuvs::neighbors::make_host_standard_dataset_view(sub_host))));
+              handle_, host_params, cuvs::core::make_host_standard_dataset_view(sub_host))));
           sub_index = detail::to_graph_only_index<T, IdxT>(handle_, *sub_host_indices_.back());
           if (sub_index.graph_fd().has_value()) { sub_host_indices_.pop_back(); }
         } else {
@@ -370,10 +371,10 @@ void cuvs_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
       for (auto* index : indices) {
         merged_rows += static_cast<int64_t>(index->size());
       }
-      auto const stride = static_cast<int64_t>(
-        cuvs::neighbors::cagra_required_row_width<T>(static_cast<uint32_t>(dim_)));
+      auto const stride =
+        static_cast<int64_t>(cuvs::core::padded_row_width<T>(static_cast<uint32_t>(dim_)));
       *dataset_                = raft::make_device_matrix<T, int64_t>(handle_, merged_rows, stride);
-      auto merged_dataset_view = cuvs::neighbors::device_padded_dataset_view<T, int64_t>(
+      auto merged_dataset_view = cuvs::core::device_padded_dataset_view<T, int64_t>(
         raft::make_const_mdspan(dataset_->view()), static_cast<uint32_t>(dim_));
       index_ =
         std::make_shared<index_type>(cuvs::neighbors::cagra::merge(handle_,
@@ -410,13 +411,15 @@ void cuvs_cagra<T, IdxT>::compress_dataset(const T* dataset, size_t nrow)
   // make_vpq_dataset() reads the rows wherever they are: host-resident ones are subsampled and
   // encoded in bounded batches instead of being staged on the device.
   auto src = raft::make_device_matrix_view<const T, int64_t, raft::row_major>(dataset, rows, dim_);
-  vpq_dataset_ = std::make_shared<cuvs::neighbors::device_vpq_dataset<half, int64_t>>(
-    cuvs::preprocessing::quantize::pq::make_vpq_dataset(handle_, *index_params_.compression, src));
+  vpq_dataset_ =
+    std::make_shared<cuvs::preprocessing::quantize::pq::device_vpq_dataset<half, int64_t>>(
+      cuvs::preprocessing::quantize::pq::make_vpq_dataset(
+        handle_, *index_params_.compression, src));
   vpq_index_ = std::make_shared<cuvs::neighbors::cagra::device_pq_index<T, IdxT, half>>(
     handle_, parse_metric_type(metric_), vpq_dataset_->as_dataset_view(), index_->graph());
 
   // Search runs on the compressed rows and the graph, so release the dense copy of the dataset.
-  cuvs::neighbors::device_padded_dataset_view<T, int64_t> empty_dv(
+  cuvs::core::device_padded_dataset_view<T, int64_t> empty_dv(
     raft::make_device_matrix_view(static_cast<T const*>(nullptr), 0, this->dim_), this->dim_);
   *index_   = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv);
   *dataset_ = raft::make_device_matrix<T, int64_t>(handle_, 0, 0);
@@ -483,7 +486,7 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
 
     // First free up existing memory
     *dataset_ = raft::make_device_matrix<T, int64_t>(handle_, 0, 0);
-    cuvs::neighbors::device_padded_dataset_view<T, int64_t> empty_dv(
+    cuvs::core::device_padded_dataset_view<T, int64_t> empty_dv(
       raft::make_device_matrix_view(static_cast<T const*>(nullptr), 0, this->dim_), this->dim_);
     *index_ = cuvs::neighbors::cagra::update_dataset(handle_, std::move(*index_), empty_dv);
 
@@ -494,7 +497,7 @@ void cuvs_cagra<T, IdxT>::set_search_param(const search_param_base& param,
     auto mr = get_mr(dataset_mem_);
     cuvs::neighbors::cagra::detail::copy_with_padding(handle_, *dataset_, *input_dataset_v_, mr);
 
-    cuvs::neighbors::device_padded_dataset_view<T, int64_t> dv(
+    cuvs::core::device_padded_dataset_view<T, int64_t> dv(
       raft::make_device_matrix_view(
         dataset_->data_handle(), dataset_->extent(0), dataset_->extent(1)),
       this->dim_);

cpp/bench/ann/src/cuvs/cuvs_mg_cagra_wrapper.h

@@ -95,8 +95,8 @@ void cuvs_mg_cagra<T, IdxT>::build(const T* dataset, size_t nrow)
     raft::make_host_matrix_view<const T, int64_t, raft::row_major>(dataset, nrow, dim_);
   // The row alignment of the host view is irrelevant: every per-rank device shard is padded
   // individually during the multi-GPU build.
-  cuvs::neighbors::host_padded_dataset_view<T, int64_t> dataset_view(dataset_mds,
-                                                                     static_cast<uint32_t>(dim_));
+  cuvs::core::host_padded_dataset_view<T, int64_t> dataset_view(dataset_mds,
+                                                                static_cast<uint32_t>(dim_));
   auto idx = cuvs::neighbors::cagra::build(clique_, build_params, dataset_view);
   index_   = std::make_shared<
       cuvs::neighbors::mg_index<cuvs::neighbors::cagra::device_padded_index<T, IdxT>, T, IdxT>>(