diff --git a/cpp/benchmarks/io/cuio_common.cpp b/cpp/benchmarks/io/cuio_common.cpp index 943b329a364..b5318b45eb4 100644 --- a/cpp/benchmarks/io/cuio_common.cpp +++ b/cpp/benchmarks/io/cuio_common.cpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -41,7 +41,8 @@ std::string random_file_in_dir(std::string const& dir_path) cuio_source_sink_pair::cuio_source_sink_pair(io_type type) : type{type}, d_buffer{0, cudf::get_default_stream()}, - file_name{random_file_in_dir(tmpdir.path())} + file_name{random_file_in_dir(tmpdir.path())}, + void_sink{cudf::io::data_sink::create()} { } @@ -67,7 +68,7 @@ cudf::io::source_info cuio_source_sink_pair::make_source_info() cudf::io::sink_info cuio_source_sink_pair::make_sink_info() { switch (type) { - case io_type::VOID: return cudf::io::sink_info(&void_sink); + case io_type::VOID: return cudf::io::sink_info(void_sink.get()); case io_type::FILEPATH: return cudf::io::sink_info(file_name); case io_type::HOST_BUFFER: [[fallthrough]]; case io_type::DEVICE_BUFFER: return cudf::io::sink_info(&h_buffer); @@ -78,7 +79,7 @@ cudf::io::sink_info cuio_source_sink_pair::make_sink_info() size_t cuio_source_sink_pair::size() { switch (type) { - case io_type::VOID: return void_sink.bytes_written(); + case io_type::VOID: return void_sink->bytes_written(); case io_type::FILEPATH: return static_cast( std::ifstream(file_name, std::ifstream::ate | std::ifstream::binary).tellg()); diff --git a/cpp/benchmarks/io/cuio_common.hpp b/cpp/benchmarks/io/cuio_common.hpp index fe509f196be..3d5be41e25f 100644 --- a/cpp/benchmarks/io/cuio_common.hpp +++ b/cpp/benchmarks/io/cuio_common.hpp @@ -1,5 +1,5 @@ /* - * Copyright (c) 2020-2023, NVIDIA CORPORATION. + * Copyright (c) 2020-2024, NVIDIA CORPORATION. * * Licensed under the Apache License, Version 2.0 (the "License"); * you may not use this file except in compliance with the License. @@ -32,15 +32,6 @@ std::string random_file_in_dir(std::string const& dir_path); * @brief Class to create a coupled `source_info` and `sink_info` of given type. */ class cuio_source_sink_pair { - class bytes_written_only_sink : public cudf::io::data_sink { - size_t _bytes_written = 0; - - public: - void host_write(void const* data, size_t size) override { _bytes_written += size; } - void flush() override {} - size_t bytes_written() override { return _bytes_written; } - }; - public: cuio_source_sink_pair(io_type type); ~cuio_source_sink_pair() @@ -79,7 +70,7 @@ class cuio_source_sink_pair { std::vector h_buffer; rmm::device_uvector d_buffer; std::string const file_name; - bytes_written_only_sink void_sink; + std::unique_ptr void_sink; }; /** diff --git a/cpp/src/io/utilities/data_sink.cpp b/cpp/src/io/utilities/data_sink.cpp index 0b14d060b05..5786e9dd6d1 100644 --- a/cpp/src/io/utilities/data_sink.cpp +++ b/cpp/src/io/utilities/data_sink.cpp @@ -139,6 +139,8 @@ class void_sink : public data_sink { [[nodiscard]] bool supports_device_write() const override { return true; } + [[nodiscard]] bool is_device_write_preferred(size_t size) const override { return true; } + void device_write(void const* gpu_data, size_t size, rmm::cuda_stream_view stream) override { _bytes_written += size; @@ -189,6 +191,11 @@ class user_sink_wrapper : public data_sink { return user_sink->device_write_async(gpu_data, size, stream); } + [[nodiscard]] bool is_device_write_preferred(size_t size) const override + { + return user_sink->is_device_write_preferred(size); + } + void flush() override { user_sink->flush(); } size_t bytes_written() override { return user_sink->bytes_written(); }