From 8ed3e20f3d800adde15c94799ac743ae52d251a6 Mon Sep 17 00:00:00 2001
From: Nghia Truong <7416935+ttnghia@users.noreply.github.com>
Date: Mon, 11 Mar 2024 16:01:20 -0600
Subject: [PATCH] Cleanup `hostdevice_vector` and add more APIs (#15252)

This work includes:
 * Fix a bug in `hostdevice_vector` when the host buffer does not change its size when appending new elements. Instead, the new elements are written directly into raw memory (which is out of bounds). Previously, this did not trigger any issue since the host buffer has reserved plenty of memory upon its construction, until I attempted to access the `front()` and `back()` elements of it.
 * Add `front()` and `back()` accessors which return the first and last elements in the host buffer.

Authors:
  - Nghia Truong (https://github.com/ttnghia)

Approvers:
  - https://github.com/nvdbaranec
  - Vukasin Milovanovic (https://github.com/vuule)

URL: https://github.com/rapidsai/cudf/pull/15252
---
 cpp/src/io/utilities/hostdevice_vector.hpp | 22 ++++++++++------------
 1 file changed, 10 insertions(+), 12 deletions(-)

diff --git a/cpp/src/io/utilities/hostdevice_vector.hpp b/cpp/src/io/utilities/hostdevice_vector.hpp
index a1e8af51858..0883ac3609f 100644
--- a/cpp/src/io/utilities/hostdevice_vector.hpp
+++ b/cpp/src/io/utilities/hostdevice_vector.hpp
@@ -26,13 +26,9 @@
 #include <cudf/utilities/span.hpp>
 
 #include <rmm/cuda_stream_view.hpp>
-#include <rmm/device_buffer.hpp>
+#include <rmm/device_uvector.hpp>
 #include <rmm/mr/host/host_memory_resource.hpp>
 
-#include <thrust/host_vector.h>
-
-#include <variant>
-
 namespace cudf::detail {
 
 /**
@@ -57,26 +53,23 @@ class hostdevice_vector {
   }
 
   explicit hostdevice_vector(size_t initial_size, size_t max_size, rmm::cuda_stream_view stream)
-    : h_data({cudf::io::get_host_memory_resource(), stream}), d_data(0, stream)
+    : h_data({cudf::io::get_host_memory_resource(), stream}), d_data(max_size, stream)
   {
     CUDF_EXPECTS(initial_size <= max_size, "initial_size cannot be larger than max_size");
 
     h_data.reserve(max_size);
     h_data.resize(initial_size);
-
-    current_size = initial_size;
-    d_data.resize(max_size, stream);
   }
 
   void push_back(T const& data)
   {
     CUDF_EXPECTS(size() < capacity(),
                  "Cannot insert data into hostdevice_vector because capacity has been exceeded.");
-    h_data[current_size++] = data;
+    h_data.push_back(data);
   }
 
   [[nodiscard]] size_t capacity() const noexcept { return d_data.size(); }
-  [[nodiscard]] size_t size() const noexcept { return current_size; }
+  [[nodiscard]] size_t size() const noexcept { return h_data.size(); }
   [[nodiscard]] size_t size_bytes() const noexcept { return sizeof(T) * size(); }
   [[nodiscard]] bool empty() const noexcept { return size() == 0; }
 
@@ -92,6 +85,12 @@ class hostdevice_vector {
   [[nodiscard]] T* end() { return host_ptr(size()); }
   [[nodiscard]] T const* end() const { return host_ptr(size()); }
 
+  [[nodiscard]] T& front() { return h_data.front(); }
+  [[nodiscard]] T const& front() const { return front(); }
+
+  [[nodiscard]] T& back() { return h_data.back(); }
+  [[nodiscard]] T const& back() const { return back(); }
+
   [[nodiscard]] T* device_ptr(size_t offset = 0) { return d_data.data() + offset; }
   [[nodiscard]] T const* device_ptr(size_t offset = 0) const { return d_data.data() + offset; }
 
@@ -175,7 +174,6 @@ class hostdevice_vector {
 
  private:
   cudf::detail::rmm_host_vector<T> h_data;
-  size_t current_size = 0;
   rmm::device_uvector<T> d_data;
 };