Setting a threshold for KvikIO IO

cpp/src/io/utilities/data_sink.cpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2020-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2020-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -63,8 +63,8 @@ class file_sink : public data_sink {
 
   [[nodiscard]] bool is_device_write_preferred(size_t size) const override
   {
-    return !_kvikio_file.closed() ||
-           (_cufile_out != nullptr && _cufile_out->is_cufile_io_preferred(size));
+    if (size < _gds_io_preferred_threshold) { return false; }
+    return supports_device_write();
   }
 
   std::future<void> device_write_async(void const* gpu_data,
@@ -96,6 +96,8 @@ class file_sink : public data_sink {
   size_t _bytes_written = 0;
   std::unique_ptr<detail::cufile_output_impl> _cufile_out;
   kvikio::FileHandle _kvikio_file;
+  // The write size above which GDS is faster then d2h-copy + posix-write
+  static constexpr size_t _gds_io_preferred_threshold = 128 << 10;
 };
 
 /**

cpp/src/io/utilities/datasource.cpp

@@ -56,8 +56,8 @@ class file_source : public datasource {
 
   [[nodiscard]] bool is_device_read_preferred(size_t size) const override
   {
-    return !_kvikio_file.closed() ||
-           (_cufile_in != nullptr && _cufile_in->is_cufile_io_preferred(size));
+    if (size < _gds_io_preferred_threshold) { return false; }
+    return supports_device_read();
   }
 
   std::future<size_t> device_read_async(size_t offset,
@@ -98,6 +98,8 @@ class file_source : public datasource {
  private:
   std::unique_ptr<detail::cufile_input_impl> _cufile_in;
   kvikio::FileHandle _kvikio_file;
+  // The read size above which GDS is faster then posix-read + h2d-copy
+  static constexpr size_t _gds_io_preferred_threshold = 128 << 10;
 };
 
 /**

cpp/src/io/utilities/file_io_utilities.hpp

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2021-2022, NVIDIA CORPORATION.
+ * Copyright (c) 2021-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -49,36 +49,10 @@ class file_wrapper {
   [[nodiscard]] auto desc() const { return fd; }
 };
 
-/**
- * @brief Base class for cuFile input/output.
- *
- * Contains the common API for cuFile input and output classes.
- */
-class cufile_io_base {
- public:
-  /**
-   * @brief Returns an estimate of whether the cuFile operation is the optimal option.
-   *
-   * @param size Read/write operation size, in bytes.
-   * @return Whether a cuFile operation with the given size is expected to be faster than a host
-   * read + H2D copy
-   */
-  static bool is_cufile_io_preferred(size_t size) { return size > op_size_threshold; }
-
- protected:
-  /**
-   * @brief The read/write size above which cuFile is faster then host read + copy
-   *
-   * This may not be the optimal threshold for all systems. Derived `is_cufile_io_preferred`
-   * implementations can use a different logic.
-   */
-  static constexpr size_t op_size_threshold = 128 << 10;
-};
-
 /**
  * @brief Interface class for cufile input.
  */
-class cufile_input : public cufile_io_base {
+class cufile_input {
  public:
   /**
    * @brief Asynchronously reads into existing device memory.
@@ -101,7 +75,7 @@ class cufile_input : public cufile_io_base {
 /**
  * @brief Interface class for cufile output.
  */
-class cufile_output : public cufile_io_base {
+class cufile_output {
  public:
   /**
    * @brief Asynchronously writes the data from a device buffer into a file.