Convert IO tests' scratch rmm::device_uvector<T> to cuda::device_buffer<T> (#24963)

Use CCCL buffers instead of RMM vectors in IO tests.

Issue: rapidsai/build-planning#321

Authors:
  - Kyle Edwards (https://github.com/KyleFromNVIDIA)

Approvers:
  - Muhammad Haseeb (https://github.com/mhaseeb123)
  - Bradley Dice (https://github.com/bdice)

URL: #24963

Author: KyleFromNVIDIA

cpp/tests/io/comp/comp_test.cpp

@@ -13,13 +13,12 @@
 #include <cudf_test/memory_resource_utilities.hpp>
 #include <cudf_test/testing_main.hpp>
 
+#include <cudf/detail/utilities/buffer_factories.hpp>
+#include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/io/detail/codec.hpp>
 #include <cudf/io/types.hpp>
 #include <cudf/utilities/default_stream.hpp>
 
-#include <rmm/device_buffer.hpp>
-#include <rmm/device_uvector.hpp>
-
 #include <cuda/buffer>
 #include <cuda/iterator>
 
@@ -85,7 +84,8 @@ struct DecompressTest
                                               compressed.data(),
                                               compressed.data() + compressed.size()};
     }
-    rmm::device_uvector<uint8_t> dst{decompressed.size(), stream};
+    cuda::device_buffer<uint8_t> dst{
+      stream, cudf::get_current_device_resource_ref(), decompressed.size(), cuda::no_init};
 
     cudf::detail::hostdevice_vector<device_span<uint8_t const>> inf_in(
       1, stream, cudf::get_current_device_resource_ref());
@@ -480,10 +480,16 @@ void roundtrip_test(cudf::io::compression_type compression,
 
     auto const test_input = cudf::host_span<uint8_t const>{expected.data(), test_size};
 
-    auto d_comp = rmm::device_uvector<uint8_t>(
-      cudf::io::detail::max_compressed_size(compression, test_input.size()), stream, mr);
+    auto d_comp = cuda::device_buffer<uint8_t>(
+      stream,
+      mr,
+      cudf::io::detail::max_compressed_size(compression, test_input.size()),
+      cuda::no_init);
+    // Tracks the valid compressed bytes since device_buffer cannot be shrunk yet
+    // TODO: Remove this when https://github.com/NVIDIA/cccl/issues/12051 is resolved
+    device_span<uint8_t const> compressed = d_comp;
     {
-      auto const d_orig = cudf::detail::make_device_uvector_async(test_input, stream, mr);
+      auto const d_orig = cudf::detail::make_device_buffer_async(test_input, stream, mr);
       auto hd_srcs      = cudf::detail::hostdevice_vector<device_span<uint8_t const>>(
         1, stream, cudf::get_current_device_resource_ref());
       hd_srcs[0] = d_orig;
@@ -504,14 +510,15 @@ void roundtrip_test(cudf::io::compression_type compression,
       });
       hd_stats.device_to_host(stream);
       ASSERT_EQ(hd_stats[0].status, codec_status::SUCCESS);
-      d_comp.resize(hd_stats[0].bytes_written, stream);
+      // TODO: Resize the buffer once https://github.com/NVIDIA/cccl/issues/12051 is resolved.
+      compressed = compressed.first(hd_stats[0].bytes_written);
     }
 
-    auto d_got = rmm::device_uvector<uint8_t>(test_input.size(), stream, mr);
+    auto d_got = cuda::device_buffer<uint8_t>(stream, mr, test_input.size(), cuda::no_init);
     {
       auto hd_srcs = cudf::detail::hostdevice_vector<device_span<uint8_t const>>(
         1, stream, cudf::get_current_device_resource_ref());
-      hd_srcs[0] = d_comp;
+      hd_srcs[0] = compressed;
       hd_srcs.host_to_device_async(stream);
 
       auto hd_dsts = cudf::detail::hostdevice_vector<device_span<uint8_t>>(

cpp/tests/io/fst/logical_stack_test.cu

@@ -8,11 +8,10 @@
 #include <cudf_test/base_fixture.hpp>
 #include <cudf_test/testing_main.hpp>
 
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/types.hpp>
 #include <cudf/utilities/error.hpp>
 
-#include <rmm/device_uvector.hpp>
-
 #include <cuda/stream>
 #include <cuda_runtime_api.h>
 
@@ -199,34 +198,24 @@ TEST_F(LogicalStackTest, GroundTruth)
                           std::back_inserter(stack_symbols),
                           std::back_inserter(stack_op_indexes));
 
-  rmm::device_uvector<SymbolT> d_stack_ops{stack_symbols.size(), stream_view};
-  rmm::device_uvector<SymbolOffsetT> d_stack_op_indexes{stack_op_indexes.size(), stream_view};
+  auto d_stack_ops = cudf::detail::make_device_buffer_async(
+    stack_symbols, stream_view, cudf::get_current_device_resource_ref());
+  auto d_stack_op_indexes = cudf::detail::make_device_buffer_async(
+    stack_op_indexes, stream_view, cudf::get_current_device_resource_ref());
   cudf::detail::hostdevice_vector<SymbolT> top_of_stack_gpu{
     string_size, stream_view, cudf::get_current_device_resource_ref()};
-  cudf::device_span<SymbolOffsetT> d_stack_op_idx_span{d_stack_op_indexes};
-
-  CUDF_CUDA_TRY(cudaMemcpyAsync(d_stack_ops.data(),
-                                stack_symbols.data(),
-                                stack_symbols.size() * sizeof(SymbolT),
-                                cudaMemcpyDefault,
-                                stream.get()));
-
-  CUDF_CUDA_TRY(cudaMemcpyAsync(d_stack_op_indexes.data(),
-                                stack_op_indexes.data(),
-                                stack_op_indexes.size() * sizeof(SymbolOffsetT),
-                                cudaMemcpyDefault,
-                                stream.get()));
-
   // Run algorithm
-  fst::sparse_stack_op_to_top_of_stack<fst::stack_op_support::NO_RESET_SUPPORT, StackLevelT>(
-    d_stack_ops.data(),
-    d_stack_op_idx_span,
-    JSONToStackOp{},
-    top_of_stack_gpu.device_ptr(),
-    empty_stack_symbol,
-    read_symbol,
-    string_size,
-    stream.get());
+  fst::sparse_stack_op_to_top_of_stack<fst::stack_op_support::NO_RESET_SUPPORT,
+                                       StackLevelT,
+                                       SymbolT*,
+                                       SymbolOffsetT>(d_stack_ops.data(),
+                                                      d_stack_op_indexes,
+                                                      JSONToStackOp{},
+                                                      top_of_stack_gpu.device_ptr(),
+                                                      empty_stack_symbol,
+                                                      read_symbol,
+                                                      string_size,
+                                                      stream.get());
 
   // Async copy results from device to host
   top_of_stack_gpu.device_to_host_async(stream_view);

cpp/tests/io/json/json_tree_csr.cu

@@ -17,6 +17,7 @@
 #include <rmm/cuda_stream.hpp>
 #include <rmm/exec_policy.hpp>
 
+#include <cuda/buffer>
 #include <cuda/stream>
 #include <thrust/sequence.h>
 #include <thrust/sort.h>
@@ -166,12 +167,12 @@ void run_test(std::string const& input, bool enable_lines = true)
   auto& gpu_row_offsets = std::get<1>(tup);
 
   auto const num_nodes = gpu_col_id.size();
-  rmm::device_uvector<cudf::size_type> sorted_col_ids(gpu_col_id.size(), stream);  // make a copy
-  thrust::copy(
-    rmm::exec_policy_nosync(stream), gpu_col_id.begin(), gpu_col_id.end(), sorted_col_ids.begin());
+  cuda::device_buffer<cudf::size_type> sorted_col_ids(
+    stream, cudf::get_current_device_resource_ref(), gpu_col_id);  // make a copy
 
   // sort by {col_id} on {node_ids} stable
-  rmm::device_uvector<cudf::size_type> node_ids(gpu_col_id.size(), stream);
+  cuda::device_buffer<cudf::size_type> node_ids(
+    stream, cudf::get_current_device_resource_ref(), gpu_col_id.size(), cuda::no_init);
   thrust::sequence(rmm::exec_policy_nosync(stream), node_ids.begin(), node_ids.end());
   thrust::stable_sort_by_key(rmm::exec_policy_nosync(stream),
                              sorted_col_ids.begin(),

cpp/tests/io/json/nested_json_test.cpp

@@ -13,6 +13,8 @@
 #include <cudf_test/random.hpp>
 #include <cudf_test/testing_main.hpp>
 
+#include <cudf/detail/utilities/buffer_factories.hpp>
+#include <cudf/detail/utilities/vector_factories.hpp>
 #include <cudf/io/json.hpp>
 #include <cudf/scalar/scalar.hpp>
 #include <cudf/utilities/default_stream.hpp>
@@ -635,10 +637,10 @@ TEST_F(JsonParserTest, ExtractColumn)
   cudf::io::json_reader_options default_options{};
 
   std::string const input = R"( [{"a":0.0, "b":1.0}, {"a":0.1, "b":1.1}, {"a":0.2, "b":1.2}] )";
-  auto const d_input      = cudf::detail::make_device_uvector_async(
-    cudf::host_span<char const>{input.c_str(), input.size()},
-    stream,
-    cudf::get_current_device_resource_ref());
+  auto const d_input =
+    cudf::detail::make_device_buffer_async(cudf::host_span<char const>{input.c_str(), input.size()},
+                                           stream,
+                                           cudf::get_current_device_resource_ref());
   // Get the JSON's tree representation
   auto const cudf_table = json_parser(d_input, default_options, stream, mr);
 
@@ -843,12 +845,12 @@ TEST_F(JsonTest, PostProcessTokenStream)
   thrust::copy(input.cbegin(), input.cend(), token_tuples);
 
   // Initialize device-side test data
-  auto const d_offsets = cudf::detail::make_device_uvector_async(
+  auto const d_offsets = cudf::detail::make_device_buffer_async(
     cudf::host_span<token_index_t const>{offsets.data(), offsets.size()},
     stream,
     cudf::get_current_device_resource_ref());
-  auto const d_tokens = cudf::detail::make_device_uvector_async(
-    tokens, stream, cudf::get_current_device_resource_ref());
+  auto const d_tokens =
+    cudf::detail::make_device_buffer_async(tokens, stream, cudf::get_current_device_resource_ref());
 
   // Run system-under-test
   auto [d_filtered_tokens, d_filtered_indices] =
@@ -892,7 +894,7 @@ TEST_P(JsonDelimiterParamTest, UTF_JSON)
   {"a":1,"b":Infinity,"c":[null], "d": {"year":-600,"author": "Kaniyan"}}])";
   std::replace(ascii_pass.begin(), ascii_pass.end(), '\n', delimiter);
 
-  auto const d_ascii_pass = cudf::detail::make_device_uvector(
+  auto const d_ascii_pass = cudf::detail::make_device_buffer(
     cudf::host_span<char const>{ascii_pass.c_str(), ascii_pass.size()},
     stream,
     cudf::get_current_device_resource_ref());
@@ -909,7 +911,7 @@ TEST_P(JsonDelimiterParamTest, UTF_JSON)
   {"a":1,"b":Infinity,"c":[null], "d": {"year":-600,"author": "filip ʒakotɛ"}}])";
   std::replace(utf_failed.begin(), utf_failed.end(), '\n', delimiter);
 
-  auto const d_utf_failed = cudf::detail::make_device_uvector(
+  auto const d_utf_failed = cudf::detail::make_device_buffer(
     cudf::host_span<char const>{utf_failed.c_str(), utf_failed.size()},
     stream,
     cudf::get_current_device_resource_ref());
@@ -926,10 +928,10 @@ TEST_P(JsonDelimiterParamTest, UTF_JSON)
   {"a":1,"b":NaN,"c":[null, null], "d": {"year": 2, "author": "filip ʒakotɛ"}}])";
   std::replace(utf_pass.begin(), utf_pass.end(), '\n', delimiter);
 
-  auto const d_utf_pass = cudf::detail::make_device_uvector(
-    cudf::host_span<char const>{utf_pass.c_str(), utf_pass.size()},
-    stream,
-    cudf::get_current_device_resource_ref());
+  auto const d_utf_pass =
+    cudf::detail::make_device_buffer(cudf::host_span<char const>{utf_pass.c_str(), utf_pass.size()},
+                                     stream,
+                                     cudf::get_current_device_resource_ref());
   CUDF_EXPECT_NO_THROW(json_parser(d_utf_pass, default_options, stream, mr));
 }
 
@@ -947,10 +949,10 @@ TEST_F(JsonParserTest, ExtractColumnWithQuotes)
   options.enable_keep_quotes(true);
 
   std::string const input = R"( [{"a":"0.0", "b":1.0}, {"b":1.1}, {"b":2.1, "a":"2.0"}] )";
-  auto const d_input      = cudf::detail::make_device_uvector_async(
-    cudf::host_span<char const>{input.c_str(), input.size()},
-    stream,
-    cudf::get_current_device_resource_ref());
+  auto const d_input =
+    cudf::detail::make_device_buffer_async(cudf::host_span<char const>{input.c_str(), input.size()},
+                                           stream,
+                                           cudf::get_current_device_resource_ref());
   // Get the JSON's tree representation
   auto const cudf_table = json_parser(d_input, options, stream, mr);
 
@@ -990,7 +992,7 @@ TEST_F(JsonParserTest, ExpectFailMixStructAndList)
 
   // libcudf does not currently support a mix of lists and structs.
   for (auto const& input : inputs_fail) {
-    auto const d_input = cudf::detail::make_device_uvector_async(
+    auto const d_input = cudf::detail::make_device_buffer_async(
       cudf::host_span<char const>{input.c_str(), input.size()},
       stream,
       cudf::get_current_device_resource_ref());
@@ -999,7 +1001,7 @@ TEST_F(JsonParserTest, ExpectFailMixStructAndList)
   }
 
   for (auto const& input : inputs_succeed) {
-    auto const d_input = cudf::detail::make_device_uvector_async(
+    auto const d_input = cudf::detail::make_device_buffer_async(
       cudf::host_span<char const>{input.c_str(), input.size()},
       stream,
       cudf::get_current_device_resource_ref());
@@ -1021,9 +1023,9 @@ TEST_F(JsonParserTest, EmptyString)
 
   std::string const input = R"([])";
   auto const d_input =
-    cudf::detail::make_device_uvector(cudf::host_span<char const>{input.c_str(), input.size()},
-                                      stream,
-                                      cudf::get_current_device_resource_ref());
+    cudf::detail::make_device_buffer(cudf::host_span<char const>{input.c_str(), input.size()},
+                                     stream,
+                                     cudf::get_current_device_resource_ref());
   // Get the JSON's tree representation
   auto const cudf_table = json_parser(d_input, default_options, stream, mr);
 
@@ -1410,9 +1412,9 @@ TEST_F(JsonTest, RejectsUnquotedValuesWithInvalidLeadingChar)
                                              token_t::StructEnd,
                                              token_t::LineEnd};
     std::vector<cuio_json::SymbolOffsetT> token_indices{0, 0, 1, 3, 5, 6, 6, 6, 7};
-    auto d_tokens = cudf::detail::make_device_uvector_async(
+    auto d_tokens = cudf::detail::make_device_buffer_async(
       tokens, stream, cudf::get_current_device_resource_ref());
-    auto d_token_indices = cudf::detail::make_device_uvector_async(
+    auto d_token_indices = cudf::detail::make_device_buffer_async(
       token_indices, stream, cudf::get_current_device_resource_ref());
 
     cuio_json::detail::validate_token_stream(d_input, d_tokens, d_token_indices, opts, stream);

cpp/tests/io/parquet_writer_test.cpp

@@ -15,6 +15,7 @@
 
 #include <cudf/column/column_factories.hpp>
 #include <cudf/copying.hpp>
+#include <cudf/detail/utilities/buffer_factories.hpp>
 #include <cudf/filling.hpp>
 #include <cudf/io/data_sink.hpp>
 #include <cudf/io/detail/codec.hpp>
@@ -242,7 +243,7 @@ TEST_F(ParquetWriterTest, BufferSource)
 
   // device buffer
   {
-    auto const d_input = cudf::detail::make_device_uvector(
+    auto const d_input = cudf::detail::make_device_buffer(
       cudf::host_span<uint8_t const>{reinterpret_cast<uint8_t const*>(out_buffer.data()),
                                      out_buffer.size()},
       cudf::get_default_stream(),