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 <fcntl.h>
#include <sys/stat.h>
#include <sys/types.h>
#include <unistd.h>
 
#include <cstddef>
#include <cstdlib>
#include <iostream>
#include <numeric>
#include <optional>
#include <system_error>
#include <utility>
 
#include <kvikio/buffer.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 {
namespace detail {
 
inline int open_fd_parse_flags(const std::string& flags, bool o_direct)
{
  int file_flags = -1;
  if (flags.empty()) { throw std::invalid_argument("Unknown file open flag"); }
  switch (flags[0]) {
    case 'r':
      file_flags = O_RDONLY;
      if (flags[1] == '+') { file_flags = O_RDWR; }
      break;
    case 'w':
      file_flags = O_WRONLY;
      if (flags[1] == '+') { file_flags = O_RDWR; }
      file_flags |= O_CREAT | O_TRUNC;
      break;
    case 'a': throw std::invalid_argument("Open flag 'a' isn't supported");
    default: throw std::invalid_argument("Unknown file open flag");
  }
  file_flags |= O_CLOEXEC;
  if (o_direct) { file_flags |= O_DIRECT; }
  return file_flags;
}
 
inline int open_fd(const std::string& file_path,
                   const std::string& flags,
                   bool o_direct,
                   mode_t mode)
{
  // NOLINTNEXTLINE(cppcoreguidelines-pro-type-vararg)
  int fd = ::open(file_path.c_str(), open_fd_parse_flags(flags, o_direct), mode);
  if (fd == -1) { throw std::system_error(errno, std::generic_category(), "Unable to open file"); }
  return fd;
}
 
[[nodiscard]] inline int open_flags(int fd)
{
  int ret = fcntl(fd, F_GETFL);  // NOLINT(cppcoreguidelines-pro-type-vararg)
  if (ret == -1) {
    throw std::system_error(errno, std::generic_category(), "Unable to retrieve open flags");
  }
  return ret;
}
 
[[nodiscard]] inline std::size_t get_file_size(int file_descriptor)
{
  struct stat st {};
  int ret = fstat(file_descriptor, &st);
  if (ret == -1) {
    throw std::system_error(errno, std::generic_category(), "Unable to query file size");
  }
  return static_cast<std::size_t>(st.st_size);
}
 
}  // namespace detail
 
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};
  bool _compat_mode{false};
  mutable std::size_t _nbytes{0};  // The size of the underlying file, zero means unknown.
  CUfileHandle_t _handle{};
 
 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,
             bool compat_mode         = defaults::compat_mode())
    : _fd_direct_off{detail::open_fd(file_path, flags, false, mode)},
      _initialized{true},
      _compat_mode{compat_mode}
  {
    try {
      _fd_direct_on = detail::open_fd(file_path, flags, true, mode);
    } catch (const std::system_error&) {
      _compat_mode = true;  // Fall back to compat mode if we cannot open the file with O_DIRECT
    }
 
    if (!_compat_mode) {
      CUfileDescr_t desc{};  // It is important to set to zero!
      desc.type = CU_FILE_HANDLE_TYPE_OPAQUE_FD;
      // NOLINTNEXTLINE(cppcoreguidelines-pro-type-union-access)
      desc.handle.fd = _fd_direct_on;
      CUFILE_TRY(cuFileAPI::instance().HandleRegister(&_handle, &desc));
    }
  }
 
  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, false)},
      _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, false);
    _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 (!_compat_mode) { cuFileAPI::instance().HandleDeregister(_handle); }
    ::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 (_compat_mode) {
      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 { return detail::open_flags(_fd_direct_off); }
 
  [[nodiscard]] inline std::size_t nbytes() const
  {
    if (closed()) { return 0; }
    if (_nbytes == 0) { _nbytes = detail::get_file_size(_fd_direct_off); }
    return _nbytes;
  }
 
  std::size_t read(void* devPtr_base,
                   std::size_t size,
                   std::size_t file_offset,
                   std::size_t devPtr_offset)
  {
    if (_compat_mode) {
      return posix_device_read(_fd_direct_off, devPtr_base, size, file_offset, devPtr_offset);
    }
    ssize_t ret = cuFileAPI::instance().Read(
      _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 read file");
    }
    if (ret < -1) {
      throw CUfileException(std::string{"cuFile error at: "} + __FILE__ + ":" +
                            KVIKIO_STRINGIFY(__LINE__) + ": " + CUFILE_ERRSTR(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)
  {
    _nbytes = 0;  // Invalidate the computed file size
 
    if (_compat_mode) {
      return posix_device_write(_fd_direct_off, devPtr_base, size, file_offset, devPtr_offset);
    }
    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())
  {
    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 posix_host_read(_fd_direct_off, buf, size, file_offset, false);
      };
 
      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 posix_device_read(_fd_direct_off, buf, size, file_offset, 0);
      };
      return std::async(std::launch::deferred, task);
    }
 
    // 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);
    };
    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())
  {
    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 posix_host_write(_fd_direct_off, buf, size, file_offset, false);
      };
 
      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 posix_device_write(_fd_direct_off, buf, size, file_offset, 0);
      };
      return std::async(std::launch::deferred, task);
    }
 
    // 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);
    };
    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 (kvikio::is_batch_and_stream_available() && !_compat_mode) {
      CUFILE_TRY(cuFileAPI::instance().ReadAsync(
        _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_read_p, stream));
      return;
    }
    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));
  }
 
  [[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 (kvikio::is_batch_and_stream_available() && !_compat_mode) {
      CUFILE_TRY(cuFileAPI::instance().WriteAsync(
        _handle, devPtr_base, size_p, file_offset_p, devPtr_offset_p, bytes_written_p, stream));
      return;
    }
    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));
  }
 
  [[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_on() const noexcept { return _compat_mode; }
};
 
}  // namespace kvikio