From 292406c390ff16ea2bdd2ba0e0df1bbae470c0d2 Mon Sep 17 00:00:00 2001
From: achirkin <achirkin@users.noreply.github.com>
Date: Mon, 11 Mar 2024 18:21:36 +0100
Subject: [PATCH] Simplify unique_ptr arithmetics

---
 cpp/include/raft/neighbors/cagra_types.hpp    | 12 ++---
 cpp/include/raft/neighbors/dataset.hpp        | 20 ++------
 .../neighbors/detail/dataset_serialize.hpp    | 46 +++++++++++--------
 3 files changed, 36 insertions(+), 42 deletions(-)

diff --git a/cpp/include/raft/neighbors/cagra_types.hpp b/cpp/include/raft/neighbors/cagra_types.hpp
index 868f214320..9b86ca29f2 100644
--- a/cpp/include/raft/neighbors/cagra_types.hpp
+++ b/cpp/include/raft/neighbors/cagra_types.hpp
@@ -261,7 +261,7 @@ struct index : ann::index {
         mdspan<const IdxT, matrix_extent<int64_t>, row_major, graph_accessor> knn_graph)
     : ann::index(),
       metric_(metric),
-      dataset_(upcast_dataset_ptr(construct_aligned_dataset(res, dataset, 16))),
+      dataset_(construct_aligned_dataset(res, dataset, 16)),
       graph_(make_device_matrix<IdxT, int64_t>(res, 0, 0))
   {
     RAFT_EXPECTS(dataset.extent(0) == knn_graph.extent(0),
@@ -280,14 +280,14 @@ struct index : ann::index {
   void update_dataset(raft::resources const& res,
                       raft::device_matrix_view<const T, int64_t, row_major> dataset)
   {
-    upcast_dataset_ptr(construct_aligned_dataset(res, dataset, 16)).swap(dataset_);
+    dataset_ = construct_aligned_dataset(res, dataset, 16);
   }
 
   /** Set the dataset reference explicitly to a device matrix view with padding. */
   void update_dataset(raft::resources const& res,
                       raft::device_matrix_view<const T, int64_t, layout_stride> dataset)
   {
-    upcast_dataset_ptr(construct_aligned_dataset(res, dataset, 16)).swap(dataset_);
+    dataset_ = construct_aligned_dataset(res, dataset, 16);
   }
 
   /**
@@ -298,7 +298,7 @@ struct index : ann::index {
   void update_dataset(raft::resources const& res,
                       raft::host_matrix_view<const T, int64_t, row_major> dataset)
   {
-    upcast_dataset_ptr(construct_aligned_dataset(res, dataset, 16)).swap(dataset_);
+    dataset_ = construct_aligned_dataset(res, dataset, 16);
   }
 
   /** Replace the dataset with a new dataset. */
@@ -306,14 +306,14 @@ struct index : ann::index {
   auto update_dataset(raft::resources const& res, DatasetT&& dataset)
     -> std::enable_if_t<std::is_base_of_v<neighbors::dataset<int64_t>, DatasetT>>
   {
-    upcast_dataset_ptr(std::make_unique<DatasetT>(std::move(dataset))).swap(dataset_);
+    dataset_ = std::make_unique<DatasetT>(std::move(dataset));
   }
 
   template <typename DatasetT>
   auto update_dataset(raft::resources const& res, std::unique_ptr<DatasetT>&& dataset)
     -> std::enable_if_t<std::is_base_of_v<neighbors::dataset<int64_t>, DatasetT>>
   {
-    upcast_dataset_ptr(std::move(dataset)).swap(dataset_);
+    dataset_ = std::move(dataset);
   }
 
   /**
diff --git a/cpp/include/raft/neighbors/dataset.hpp b/cpp/include/raft/neighbors/dataset.hpp
index dd346e6d9d..23ea6054bd 100644
--- a/cpp/include/raft/neighbors/dataset.hpp
+++ b/cpp/include/raft/neighbors/dataset.hpp
@@ -118,7 +118,6 @@ auto construct_strided_dataset(const raft::resources& res,
   using value_type   = typename SrcT::value_type;
   using index_type   = typename SrcT::index_type;
   using layout_type  = typename SrcT::layout_type;
-  using out_type     = strided_dataset<value_type, index_type>;
   static_assert(extents_type::rank() == 2, "The input must be a matrix.");
   static_assert(std::is_same_v<layout_type, layout_right> ||
                   std::is_same_v<layout_type, layout_right_padded<value_type>> ||
@@ -133,9 +132,9 @@ auto construct_strided_dataset(const raft::resources& res,
 
   if (device_accessible && row_major && stride_matches) {
     // Everything matches: make a non-owning dataset
-    return std::unique_ptr<out_type>{new non_owning_dataset<value_type, index_type>{
+    return std::make_unique<non_owning_dataset<value_type, index_type>>(
       make_device_strided_matrix_view<const value_type, index_type>(
-        src.data_handle(), src.extent(0), src.extent(1), required_stride)}};
+        src.data_handle(), src.extent(0), src.extent(1), required_stride));
   }
   // Something is wrong: have to make a copy and produce an owning dataset
   auto out_layout =
@@ -161,7 +160,7 @@ auto construct_strided_dataset(const raft::resources& res,
                                   cudaMemcpyDefault,
                                   resource::get_cuda_stream(res)));
 
-  return std::unique_ptr<out_type>{new out_owning_type{std::move(out_array), out_layout}};
+  return std::make_unique<out_owning_type>(std::move(out_array), out_layout);
 }
 
 template <typename SrcT>
@@ -169,21 +168,10 @@ auto construct_aligned_dataset(const raft::resources& res, const SrcT& src, uint
   -> std::unique_ptr<strided_dataset<typename SrcT::value_type, typename SrcT::index_type>>
 {
   using value_type       = typename SrcT::value_type;
-  using index_type       = typename SrcT::index_type;
-  using out_type         = strided_dataset<value_type, index_type>;
   constexpr size_t kSize = sizeof(value_type);
   uint32_t required_stride =
     raft::round_up_safe<size_t>(src.extent(1) * kSize, std::lcm(align_bytes, kSize)) / kSize;
-  return std::unique_ptr<out_type>{construct_strided_dataset(res, src, required_stride).release()};
-}
-
-template <typename DatasetT>
-auto upcast_dataset_ptr(std::unique_ptr<DatasetT>&& src)
-  -> std::unique_ptr<dataset<typename DatasetT::index_type>>
-{
-  using out_type = dataset<typename DatasetT::index_type>;
-  static_assert(std::is_base_of_v<out_type, DatasetT>, "The source must be a child of `dataset`");
-  return std::unique_ptr<out_type>{src.release()};
+  return construct_strided_dataset(res, src, required_stride);
 }
 
 /** Parameters for VPQ compression. */
diff --git a/cpp/include/raft/neighbors/detail/dataset_serialize.hpp b/cpp/include/raft/neighbors/detail/dataset_serialize.hpp
index 2864d260c1..529569865b 100644
--- a/cpp/include/raft/neighbors/detail/dataset_serialize.hpp
+++ b/cpp/include/raft/neighbors/detail/dataset_serialize.hpp
@@ -124,7 +124,7 @@ void deserialize(raft::resources const& res,
                  std::unique_ptr<empty_dataset<IdxT>>& out)
 {
   auto suggested_dim = deserialize_scalar<uint32_t>(res, is);
-  return std::make_unique<empty_dataset<IdxT>>(suggested_dim).swap(out);
+  out                = std::make_unique<empty_dataset<IdxT>>(suggested_dim);
 }
 
 template <typename DataT, typename IdxT>
@@ -132,17 +132,18 @@ void deserialize(raft::resources const& res,
                  std::istream& is,
                  std::unique_ptr<strided_dataset<DataT, IdxT>>& out)
 {
-  using out_mdarray_type          = device_mdarray<DataT, matrix_extent<IdxT>, layout_stride>;
-  using out_layout_type           = typename out_mdarray_type::layout_type;
-  using out_container_policy_type = typename out_mdarray_type::container_policy_type;
-  using out_owning_type = owning_dataset<DataT, IdxT, out_layout_type, out_container_policy_type>;
-
   auto n_rows      = deserialize_scalar<IdxT>(res, is);
   auto dim         = deserialize_scalar<uint32_t>(res, is);
   auto stride      = deserialize_scalar<uint32_t>(res, is);
   auto out_extents = make_extents<IdxT>(n_rows, dim);
   auto out_layout  = make_strided_layout(out_extents, std::array<IdxT, 2>{stride, 1});
-  auto out_array   = out_mdarray_type{res, out_layout, out_container_policy_type{}};
+  auto out_array   = make_device_matrix<DataT, IdxT>(res, n_rows, stride);
+
+  using out_mdarray_type          = decltype(out_array);
+  using out_layout_type           = typename out_mdarray_type::layout_type;
+  using out_container_policy_type = typename out_mdarray_type::container_policy_type;
+  using out_owning_type = owning_dataset<DataT, IdxT, out_layout_type, out_container_policy_type>;
+
   auto host_arrray = make_host_mdarray<DataT, IdxT>(out_extents);
   deserialize_mdspan(res, is, host_arrray.view());
   RAFT_CUDA_TRY(cudaMemcpy2DAsync(out_array.data_handle(),
@@ -153,9 +154,8 @@ void deserialize(raft::resources const& res,
                                   n_rows,
                                   cudaMemcpyDefault,
                                   resource::get_cuda_stream(res)));
-  return std::unique_ptr<strided_dataset<DataT, IdxT>>{
-    new out_owning_type{std::move(out_array), out_layout}}
-    .swap(out);
+
+  out = std::make_unique<out_owning_type>(std::move(out_array), out_layout);
 }
 
 template <typename MathT, typename IdxT>
@@ -178,9 +178,8 @@ void deserialize(raft::resources const& res,
   deserialize_mdspan(res, is, pq_code_book.view());
   deserialize_mdspan(res, is, data.view());
 
-  return std::unique_ptr<vpq_dataset<MathT, IdxT>>{
-    new vpq_dataset{std::move(vq_code_book), std::move(pq_code_book), std::move(data)}}
-    .swap(out);
+  out = std::make_unique<vpq_dataset<MathT, IdxT>>(
+    std::move(vq_code_book), std::move(pq_code_book), std::move(data));
 }
 
 template <typename IdxT>
@@ -190,29 +189,34 @@ void deserialize(raft::resources const& res, std::istream& is, std::unique_ptr<d
     case kSerializeEmptyDataset: {
       std::unique_ptr<empty_dataset<IdxT>> p;
       deserialize(res, is, p);
-      return upcast_dataset_ptr(std::move(p)).swap(out);
+      out = std::move(p);
+      return;
     }
     case kSerializeStridedDataset:
       switch (deserialize_scalar<cudaDataType_t>(res, is)) {
         case CUDA_R_32F: {
           std::unique_ptr<strided_dataset<float, IdxT>> p;
           deserialize(res, is, p);
-          return upcast_dataset_ptr(std::move(p)).swap(out);
+          out = std::move(p);
+          return;
         }
         case CUDA_R_16F: {
           std::unique_ptr<strided_dataset<half, IdxT>> p;
           deserialize(res, is, p);
-          return upcast_dataset_ptr(std::move(p)).swap(out);
+          out = std::move(p);
+          return;
         }
         case CUDA_R_8I: {
           std::unique_ptr<strided_dataset<int8_t, IdxT>> p;
           deserialize(res, is, p);
-          return upcast_dataset_ptr(std::move(p)).swap(out);
+          out = std::move(p);
+          return;
         }
         case CUDA_R_8U: {
           std::unique_ptr<strided_dataset<uint8_t, IdxT>> p;
           deserialize(res, is, p);
-          return upcast_dataset_ptr(std::move(p)).swap(out);
+          out = std::move(p);
+          return;
         }
         default: break;
       }
@@ -221,12 +225,14 @@ void deserialize(raft::resources const& res, std::istream& is, std::unique_ptr<d
         case CUDA_R_32F: {
           std::unique_ptr<vpq_dataset<float, IdxT>> p;
           deserialize(res, is, p);
-          return upcast_dataset_ptr(std::move(p)).swap(out);
+          out = std::move(p);
+          return;
         }
         case CUDA_R_16F: {
           std::unique_ptr<vpq_dataset<half, IdxT>> p;
           deserialize(res, is, p);
-          return upcast_dataset_ptr(std::move(p)).swap(out);
+          out = std::move(p);
+          return;
         }
         default: break;
       }