batch.hpp

/*
 * SPDX-FileCopyrightText: Copyright (c) 2023-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <cstddef>
#include <ctime>
#include <utility>
#include <vector>
 
#include <kvikio/error.hpp>
#include <kvikio/file_handle.hpp>
#include <kvikio/shim/cufile.hpp>
 
namespace kvikio {
 
struct BatchOp {
  // The file handle of the file to read or write
  FileHandle& file_handle;
  // Base address of buffer in device memory (host memory not supported).
  void* devPtr_base;
  // Offset in the file to read from or write to.
  off_t file_offset;
  // Offset relative to the `devPtr_base` pointer to write into or read from.
  off_t devPtr_offset;
  // Size in bytes to read or write.
  size_t size;
  // The operation type: CUFILE_READ or CUFILE_WRITE.
  CUfileOpcode_t opcode;
};
 
#ifdef KVIKIO_CUFILE_BATCH_API_FOUND
 
class BatchHandle {
 private:
  bool _initialized{false};
  int _max_num_events{};
  CUfileBatchHandle_t _handle{};
 
 public:
  BatchHandle() noexcept = default;
 
  BatchHandle(int max_num_events);
 
  BatchHandle(BatchHandle const&)            = delete;
  BatchHandle& operator=(BatchHandle const&) = delete;
  BatchHandle(BatchHandle&& o) noexcept;
  ~BatchHandle() noexcept;
 
  [[nodiscard]] bool closed() const noexcept;
 
  void close() noexcept;
 
  void submit(std::vector<BatchOp> const& operations);
 
  std::vector<CUfileIOEvents_t> status(unsigned min_nr,
                                       unsigned max_nr,
                                       struct timespec* timeout = nullptr);
 
  void cancel();
};
 
#else
 
class BatchHandle {
 public:
  BatchHandle() noexcept = default;
 
  BatchHandle(int max_num_events);
 
  [[nodiscard]] bool closed() const noexcept;
 
  void close() noexcept;
 
  void submit(std::vector<BatchOp> const& operations);
 
  std::vector<CUfileIOEvents_t> status(unsigned min_nr,
                                       unsigned max_nr,
                                       struct timespec* timeout = nullptr);
 
  void cancel();
};
 
#endif
 
}  // namespace kvikio