batch.hpp

/*
 * Copyright (c) 2023-2025, 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 <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(const BatchHandle&)            = delete;
  BatchHandle& operator=(BatchHandle const&) = delete;
  BatchHandle(BatchHandle&& o) noexcept;
  ~BatchHandle() noexcept;
 
  [[nodiscard]] bool closed() const noexcept;
 
  void close() noexcept;
 
  void submit(const std::vector<BatchOp>& 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(const std::vector<BatchOp>& operations);
 
  std::vector<CUfileIOEvents_t> status(unsigned min_nr,
                                       unsigned max_nr,
                                       struct timespec* timeout = nullptr);
 
  void cancel();
};
 
#endif
 
}  // namespace kvikio