file_handle.hpp

/*
 * Copyright (c) 2021-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.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */
#pragma once
 
#include <sys/stat.h>
#include <sys/types.h>
 
#include <cstddef>
#include <cstdlib>
#include <stdexcept>
#include <system_error>
#include <utility>
 
#include <kvikio/buffer.hpp>
#include <kvikio/cufile/config.hpp>
#include <kvikio/defaults.hpp>
#include <kvikio/error.hpp>
#include <kvikio/parallel_operation.hpp>
#include <kvikio/posix_io.hpp>
#include <kvikio/shim/cufile.hpp>
#include <kvikio/stream.hpp>
#include <kvikio/utils.hpp>
 
namespace kvikio {
 
class FileHandle {
 private:
  // We use two file descriptors, one opened with the O_DIRECT flag and one without.
  int _fd_direct_on{-1};
  int _fd_direct_off{-1};
  bool _initialized{false};
  CompatMode _compat_mode{CompatMode::AUTO};
  mutable std::size_t _nbytes{0};  // The size of the underlying file, zero means unknown.
  CUfileHandle_t _handle{};
 
  bool is_compat_mode_preferred_for_async(CompatMode requested_compat_mode)
  {
    if (defaults::is_compat_mode_preferred(requested_compat_mode)) { return true; }
 
    if (!is_stream_api_available()) {
      if (requested_compat_mode == CompatMode::AUTO) { return true; }
      throw std::runtime_error("Missing the cuFile stream api.");
    }
 
    // When checking for availability, we also check if cuFile's config file exists. This is
    // because even when the stream API is available, it doesn't work if no config file exists.
    if (config_path().empty()) {
      if (requested_compat_mode == CompatMode::AUTO) { return true; }
      throw std::runtime_error("Missing cuFile configuration file.");
    }
    return false;
  }
 
 public:
  static constexpr mode_t m644 = S_IRUSR | S_IWUSR | S_IRGRP | S_IWGRP | S_IROTH;
  FileHandle() noexcept        = default;
 
  FileHandle(const std::string& file_path,
             const std::string& flags = "r",
             mode_t mode              = m644,
             CompatMode compat_mode   = defaults::compat_mode());
 
  FileHandle(const FileHandle&)            = delete;
  FileHandle& operator=(FileHandle const&) = delete;
  FileHandle(FileHandle&& o) noexcept
    : _fd_direct_on{std::exchange(o._fd_direct_on, -1)},
      _fd_direct_off{std::exchange(o._fd_direct_off, -1)},
      _initialized{std::exchange(o._initialized, false)},
      _compat_mode{std::exchange(o._compat_mode, CompatMode::AUTO)},
      _nbytes{std::exchange(o._nbytes, 0)},
      _handle{std::exchange(o._handle, CUfileHandle_t{})}
  {
  }
  FileHandle& operator=(FileHandle&& o) noexcept
  {
    _fd_direct_on  = std::exchange(o._fd_direct_on, -1);
    _fd_direct_off = std::exchange(o._fd_direct_off, -1);
    _initialized   = std::exchange(o._initialized, false);
    _compat_mode   = std::exchange(o._compat_mode, CompatMode::AUTO);
    _nbytes        = std::exchange(o._nbytes, 0);
    _handle        = std::exchange(o._handle, CUfileHandle_t{});
    return *this;
  }
  ~FileHandle() noexcept { close(); }
 
  [[nodiscard]] bool closed() const noexcept { return !_initialized; }
 
  void close() noexcept
  {
    if (closed()) { return; }
 
    if (!is_compat_mode_preferred()) { cuFileAPI::instance().HandleDeregister(_handle); }
    _compat_mode = CompatMode::AUTO;
    ::close(_fd_direct_off);
    if (_fd_direct_on != -1) { ::close(_fd_direct_on); }
    _fd_direct_on  = -1;
    _fd_direct_off = -1;
    _initialized   = false;
  }
 
  [[nodiscard]] CUfileHandle_t handle()
  {
    if (closed()) { throw CUfileException("File handle is closed"); }
    if (is_compat_mode_preferred()) {
      throw CUfileException("The underlying cuFile handle isn't available in compatibility mode");
    }
    return _handle;
  }
 
  [[nodiscard]] int fd() const noexcept { return _fd_direct_off; }
 
  [[nodiscard]] int fd_open_flags() const;
 
  [[nodiscard]] std::size_t nbytes() const;
 
  std::size_t read(void* devPtr_base,
                   std::size_t size,
                   std::size_t file_offset,
                   std::size_t devPtr_offset,
                   bool sync_default_stream = true)
  {
    if (is_compat_mode_preferred()) {
      return detail::posix_device_read(
        _fd_direct_off, devPtr_base, size, file_offset, devPtr_offset);
    }
    if (sync_default_stream) { CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(nullptr)); }
 
    KVIKIO_NVTX_SCOPED_RANGE("cufileRead()", size);
    ssize_t ret = cuFileAPI::instance().Read(
      _handle, devPtr_base, size, convert_size2off(file_offset), convert_size2off(devPtr_offset));
    CUFILE_CHECK_BYTES_DONE(ret);
    return ret;
  }
 
  std::size_t write(const void* devPtr_base,
                    std::size_t size,
                    std::size_t file_offset,
                    std::size_t devPtr_offset,
                    bool sync_default_stream = true)
  {
    _nbytes = 0;  // Invalidate the computed file size
 
    if (is_compat_mode_preferred()) {
      return detail::posix_device_write(
        _fd_direct_off, devPtr_base, size, file_offset, devPtr_offset);
    }
    if (sync_default_stream) { CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(nullptr)); }
 
    KVIKIO_NVTX_SCOPED_RANGE("cufileWrite()", size);
    ssize_t ret = cuFileAPI::instance().Write(
      _handle, devPtr_base, size, convert_size2off(file_offset), convert_size2off(devPtr_offset));
    if (ret == -1) {
      throw std::system_error(errno, std::generic_category(), "Unable to write file");
    }
    if (ret < -1) {
      throw CUfileException(std::string{"cuFile error at: "} + __FILE__ + ":" +
                            KVIKIO_STRINGIFY(__LINE__) + ": " + CUFILE_ERRSTR(ret));
    }
    return ret;
  }
 
  std::future<std::size_t> pread(void* buf,
                                 std::size_t size,
                                 std::size_t file_offset   = 0,
                                 std::size_t task_size     = defaults::task_size(),
                                 std::size_t gds_threshold = defaults::gds_threshold(),
                                 bool sync_default_stream  = true)
  {
    KVIKIO_NVTX_MARKER("FileHandle::pread()", size);
    if (is_host_memory(buf)) {
      auto op = [this](void* hostPtr_base,
                       std::size_t size,
                       std::size_t file_offset,
                       std::size_t hostPtr_offset) -> std::size_t {
        char* buf = static_cast<char*>(hostPtr_base) + hostPtr_offset;
        return detail::posix_host_read<detail::PartialIO::NO>(
          _fd_direct_off, buf, size, file_offset);
      };
 
      return parallel_io(op, buf, size, file_offset, task_size, 0);
    }
 
    CUcontext ctx = get_context_from_pointer(buf);
 
    // Shortcut that circumvent the threadpool and use the POSIX backend directly.
    if (size < gds_threshold) {
      auto task = [this, ctx, buf, size, file_offset]() -> std::size_t {
        PushAndPopContext c(ctx);
        return detail::posix_device_read(_fd_direct_off, buf, size, file_offset, 0);
      };
      return std::async(std::launch::deferred, task);
    }
 
    // Let's synchronize once instead of in each task.
    if (sync_default_stream && !is_compat_mode_preferred()) {
      PushAndPopContext c(ctx);
      CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(nullptr));
    }
 
    // Regular case that use the threadpool and run the tasks in parallel
    auto task = [this, ctx](void* devPtr_base,
                            std::size_t size,
                            std::size_t file_offset,
                            std::size_t devPtr_offset) -> std::size_t {
      PushAndPopContext c(ctx);
      return read(devPtr_base, size, file_offset, devPtr_offset, /* sync_default_stream = */ false);
    };
    auto [devPtr_base, base_size, devPtr_offset] = get_alloc_info(buf, &ctx);
    return parallel_io(task, devPtr_base, size, file_offset, task_size, devPtr_offset);
  }
 
  std::future<std::size_t> pwrite(const void* buf,
                                  std::size_t size,
                                  std::size_t file_offset   = 0,
                                  std::size_t task_size     = defaults::task_size(),
                                  std::size_t gds_threshold = defaults::gds_threshold(),
                                  bool sync_default_stream  = true)
  {
    KVIKIO_NVTX_MARKER("FileHandle::pwrite()", size);
    if (is_host_memory(buf)) {
      auto op = [this](const void* hostPtr_base,
                       std::size_t size,
                       std::size_t file_offset,
                       std::size_t hostPtr_offset) -> std::size_t {
        const char* buf = static_cast<const char*>(hostPtr_base) + hostPtr_offset;
        return detail::posix_host_write<detail::PartialIO::NO>(
          _fd_direct_off, buf, size, file_offset);
      };
 
      return parallel_io(op, buf, size, file_offset, task_size, 0);
    }
 
    CUcontext ctx = get_context_from_pointer(buf);
 
    // Shortcut that circumvent the threadpool and use the POSIX backend directly.
    if (size < gds_threshold) {
      auto task = [this, ctx, buf, size, file_offset]() -> std::size_t {
        PushAndPopContext c(ctx);
        return detail::posix_device_write(_fd_direct_off, buf, size, file_offset, 0);
      };
      return std::async(std::launch::deferred, task);
    }
 
    // Let's synchronize once instead of in each task.
    if (sync_default_stream && !is_compat_mode_preferred()) {
      PushAndPopContext c(ctx);
      CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(nullptr));
    }
 
    // Regular case that use the threadpool and run the tasks in parallel
    auto op = [this, ctx](const void* devPtr_base,
                          std::size_t size,
                          std::size_t file_offset,
                          std::size_t devPtr_offset) -> std::size_t {
      PushAndPopContext c(ctx);
      return write(
        devPtr_base, size, file_offset, devPtr_offset, /* sync_default_stream = */ false);
    };
    auto [devPtr_base, base_size, devPtr_offset] = get_alloc_info(buf, &ctx);
    return parallel_io(op, devPtr_base, size, file_offset, task_size, devPtr_offset);
  }
 
  void read_async(void* devPtr_base,
                  std::size_t* size_p,
                  off_t* file_offset_p,
                  off_t* devPtr_offset_p,
                  ssize_t* bytes_read_p,
                  CUstream stream)
  {
    if (is_compat_mode_preferred_for_async(_compat_mode)) {
      CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(stream));
      *bytes_read_p =
        static_cast<ssize_t>(read(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p));
    } else {
      CUFILE_TRY(cuFileAPI::instance().ReadAsync(
        _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_read_p, stream));
    }
  }
 
  [[nodiscard]] StreamFuture read_async(void* devPtr_base,
                                        std::size_t size,
                                        off_t file_offset   = 0,
                                        off_t devPtr_offset = 0,
                                        CUstream stream     = nullptr)
  {
    StreamFuture ret(devPtr_base, size, file_offset, devPtr_offset, stream);
    auto [devPtr_base_, size_p, file_offset_p, devPtr_offset_p, bytes_read_p, stream_] =
      ret.get_args();
    read_async(devPtr_base_, size_p, file_offset_p, devPtr_offset_p, bytes_read_p, stream_);
    return ret;
  }
 
  void write_async(void* devPtr_base,
                   std::size_t* size_p,
                   off_t* file_offset_p,
                   off_t* devPtr_offset_p,
                   ssize_t* bytes_written_p,
                   CUstream stream)
  {
    if (is_compat_mode_preferred_for_async(_compat_mode)) {
      CUDA_DRIVER_TRY(cudaAPI::instance().StreamSynchronize(stream));
      *bytes_written_p =
        static_cast<ssize_t>(write(devPtr_base, *size_p, *file_offset_p, *devPtr_offset_p));
    } else {
      CUFILE_TRY(cuFileAPI::instance().WriteAsync(
        _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_written_p, stream));
    }
  }
 
  [[nodiscard]] StreamFuture write_async(void* devPtr_base,
                                         std::size_t size,
                                         off_t file_offset   = 0,
                                         off_t devPtr_offset = 0,
                                         CUstream stream     = nullptr)
  {
    StreamFuture ret(devPtr_base, size, file_offset, devPtr_offset, stream);
    auto [devPtr_base_, size_p, file_offset_p, devPtr_offset_p, bytes_written_p, stream_] =
      ret.get_args();
    write_async(devPtr_base_, size_p, file_offset_p, devPtr_offset_p, bytes_written_p, stream_);
    return ret;
  }
 
  [[nodiscard]] bool is_compat_mode_preferred() const noexcept
  {
    return defaults::is_compat_mode_preferred(_compat_mode);
  }
 
  [[nodiscard]] bool is_compat_mode_preferred_for_async() const noexcept
  {
    static bool is_extra_symbol_available = is_stream_api_available();
    static bool is_config_path_empty      = config_path().empty();
    return is_compat_mode_preferred() || !is_extra_symbol_available || is_config_path_empty;
  }
};
 
}  // namespace kvikio