From cb11327ab4b9f7d2735100f503af03e32791ed56 Mon Sep 17 00:00:00 2001
From: achirkin <achirkin@users.noreply.github.com>
Date: Tue, 12 Mar 2024 11:33:11 +0100
Subject: [PATCH] Further simplify deserialization code

---
 .../detail/cagra/cagra_serialize.cuh          |   4 +-
 .../neighbors/detail/dataset_serialize.hpp    | 103 ++++--------------
 2 files changed, 24 insertions(+), 83 deletions(-)

diff --git a/cpp/include/raft/neighbors/detail/cagra/cagra_serialize.cuh b/cpp/include/raft/neighbors/detail/cagra/cagra_serialize.cuh
index f5556b256a..cb8073d8f6 100644
--- a/cpp/include/raft/neighbors/detail/cagra/cagra_serialize.cuh
+++ b/cpp/include/raft/neighbors/detail/cagra/cagra_serialize.cuh
@@ -249,9 +249,7 @@ auto deserialize(raft::resources const& res, std::istream& is) -> index<T, IdxT>
   idx.update_graph(res, raft::make_const_mdspan(graph.view()));
   bool has_dataset = deserialize_scalar<bool>(res, is);
   if (has_dataset) {
-    std::unique_ptr<dataset<int64_t>> dataset;
-    neighbors::detail::deserialize(res, is, dataset);
-    idx.update_dataset(res, std::move(dataset));
+    idx.update_dataset(res, neighbors::detail::deserialize_dataset<int64_t>(res, is));
   }
   return idx;
 }
diff --git a/cpp/include/raft/neighbors/detail/dataset_serialize.hpp b/cpp/include/raft/neighbors/detail/dataset_serialize.hpp
index dc60a4782d..dc55d891be 100644
--- a/cpp/include/raft/neighbors/detail/dataset_serialize.hpp
+++ b/cpp/include/raft/neighbors/detail/dataset_serialize.hpp
@@ -122,51 +122,28 @@ void serialize(const raft::resources& res, std::ostream& os, const dataset<IdxT>
 }
 
 template <typename IdxT>
-void deserialize(raft::resources const& res,
-                 std::istream& is,
-                 std::unique_ptr<empty_dataset<IdxT>>& out)
+auto deserialize_empty(raft::resources const& res, std::istream& is)
+  -> std::unique_ptr<empty_dataset<IdxT>>
 {
   auto suggested_dim = deserialize_scalar<uint32_t>(res, is);
-  out                = std::make_unique<empty_dataset<IdxT>>(suggested_dim);
+  return std::make_unique<empty_dataset<IdxT>>(suggested_dim);
 }
 
 template <typename DataT, typename IdxT>
-void deserialize(raft::resources const& res,
-                 std::istream& is,
-                 std::unique_ptr<strided_dataset<DataT, IdxT>>& out)
+auto deserialize_strided(raft::resources const& res, std::istream& is)
+  -> std::unique_ptr<strided_dataset<DataT, IdxT>>
 {
-  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   = 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_matrix<DataT, IdxT>(n_rows, dim);
-  deserialize_mdspan(res, is, host_arrray.view());
-  RAFT_CUDA_TRY(cudaMemsetAsync(
-    out_array.data_handle(), 0, sizeof(DataT) * out_array.size(), resource::get_cuda_stream(res)));
-  RAFT_CUDA_TRY(cudaMemcpy2DAsync(out_array.data_handle(),
-                                  sizeof(DataT) * stride,
-                                  host_arrray.data_handle(),
-                                  sizeof(DataT) * dim,
-                                  sizeof(DataT) * dim,
-                                  n_rows,
-                                  cudaMemcpyDefault,
-                                  resource::get_cuda_stream(res)));
-
-  out = std::make_unique<out_owning_type>(std::move(out_array), out_layout);
+  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 host_array = make_host_matrix<DataT, IdxT>(n_rows, dim);
+  deserialize_mdspan(res, is, host_array.view());
+  return construct_strided_dataset(res, host_array, stride);
 }
 
 template <typename MathT, typename IdxT>
-void deserialize(raft::resources const& res,
-                 std::istream& is,
-                 std::unique_ptr<vpq_dataset<MathT, IdxT>>& out)
+auto deserialize_vpq(raft::resources const& res, std::istream& is)
+  -> std::unique_ptr<vpq_dataset<MathT, IdxT>>
 {
   auto n_rows             = deserialize_scalar<IdxT>(res, is);
   auto dim                = deserialize_scalar<uint32_t>(res, is);
@@ -183,62 +160,28 @@ void deserialize(raft::resources const& res,
   deserialize_mdspan(res, is, pq_code_book.view());
   deserialize_mdspan(res, is, data.view());
 
-  out = std::make_unique<vpq_dataset<MathT, IdxT>>(
+  return std::make_unique<vpq_dataset<MathT, IdxT>>(
     std::move(vq_code_book), std::move(pq_code_book), std::move(data));
 }
 
 template <typename IdxT>
-void deserialize(raft::resources const& res, std::istream& is, std::unique_ptr<dataset<IdxT>>& out)
+auto deserialize_dataset(raft::resources const& res, std::istream& is)
+  -> std::unique_ptr<dataset<IdxT>>
 {
   switch (deserialize_scalar<dataset_instance_tag>(res, is)) {
-    case kSerializeEmptyDataset: {
-      std::unique_ptr<empty_dataset<IdxT>> p;
-      deserialize(res, is, p);
-      out = std::move(p);
-      return;
-    }
+    case kSerializeEmptyDataset: return deserialize_empty<IdxT>(res, is);
     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);
-          out = std::move(p);
-          return;
-        }
-        case CUDA_R_16F: {
-          std::unique_ptr<strided_dataset<half, IdxT>> p;
-          deserialize(res, is, p);
-          out = std::move(p);
-          return;
-        }
-        case CUDA_R_8I: {
-          std::unique_ptr<strided_dataset<int8_t, IdxT>> p;
-          deserialize(res, is, p);
-          out = std::move(p);
-          return;
-        }
-        case CUDA_R_8U: {
-          std::unique_ptr<strided_dataset<uint8_t, IdxT>> p;
-          deserialize(res, is, p);
-          out = std::move(p);
-          return;
-        }
+        case CUDA_R_32F: return deserialize_strided<float, IdxT>(res, is);
+        case CUDA_R_16F: return deserialize_strided<half, IdxT>(res, is);
+        case CUDA_R_8I: return deserialize_strided<int8_t, IdxT>(res, is);
+        case CUDA_R_8U: return deserialize_strided<uint8_t, IdxT>(res, is);
         default: break;
       }
     case kSerializeVPQDataset:
       switch (deserialize_scalar<cudaDataType_t>(res, is)) {
-        case CUDA_R_32F: {
-          std::unique_ptr<vpq_dataset<float, IdxT>> p;
-          deserialize(res, is, p);
-          out = std::move(p);
-          return;
-        }
-        case CUDA_R_16F: {
-          std::unique_ptr<vpq_dataset<half, IdxT>> p;
-          deserialize(res, is, p);
-          out = std::move(p);
-          return;
-        }
+        case CUDA_R_32F: return deserialize_vpq<float, IdxT>(res, is);
+        case CUDA_R_16F: return deserialize_vpq<half, IdxT>(res, is);
         default: break;
       }
     default: break;