cufile.hpp

/*
 * 版权所有 (c) 2022-2025, NVIDIA CORPORATION.
 *
 * 根据 Apache 许可证 2.0 版 (以下简称"许可证") 获得许可；
 * 您除非遵守本许可证的规定，否则不得使用此文件。
 * 您可以访问以下网址获取本许可证的副本：
 *
 * https://apache.ac.cn/licenses/LICENSE-2.0
 *
 * 除非适用法律要求或书面同意，根据本许可证分发的软件
 * 是按"原样"分发的，没有任何形式的明示或
 * 暗示的担保或条件。
 * 请参见许可证中规定的管理权限和
 * 限制条件。
 */
#pragma once

#include <kvikio/shim/cufile_h_wrapper.hpp>
#include <kvikio/shim/utils.hpp>

namespace kvikio {

/// @brief cuFile C-API 的 Shim 层。
class cuFileAPI {
 public
 decltype(cuFileHandleRegister)* HandleRegister{nullptr};
 decltype(cuFileHandleDeregister)* HandleDeregister{nullptr};
 decltype(cuFileRead)* Read{nullptr};
 decltype(cuFileWrite)* Write{nullptr};
 decltype(cuFileBufRegister)* BufRegister{nullptr};
 decltype(cuFileBufDeregister)* BufDeregister{nullptr};
 decltype(cuFileDriverGetProperties)* DriverGetProperties{nullptr};
 decltype(cuFileDriverSetPollMode)* DriverSetPollMode{nullptr};
 decltype(cuFileDriverSetMaxCacheSize)* DriverSetMaxCacheSize{nullptr};
 decltype(cuFileDriverSetMaxPinnedMemSize)* DriverSetMaxPinnedMemSize{nullptr};
 decltype(cuFileBatchIOSetUp)* BatchIOSetUp{nullptr};
 decltype(cuFileBatchIOSubmit)* BatchIOSubmit{nullptr};
 decltype(cuFileBatchIOGetStatus)* BatchIOGetStatus{nullptr};
 decltype(cuFileBatchIOCancel)* BatchIOCancel{nullptr};
 decltype(cuFileBatchIODestroy)* BatchIODestroy{nullptr};
 decltype(cuFileReadAsync)* ReadAsync{nullptr};
 decltype(cuFileWriteAsync)* WriteAsync{nullptr};
 decltype(cuFileStreamRegister)* StreamRegister{nullptr};
 decltype(cuFileStreamDeregister)* StreamDeregister{nullptr};

 private
 // 不要直接调用 driver open 和 close，使用 `.driver_open()` 和 `.driver_close()`。
 decltype(cuFileDriverOpen)* DriverOpen{nullptr};
 decltype(cuFileDriverClose)* DriverClose{nullptr};

 // 不要直接调用 `GetVersion`，使用 `cuFileAPI::instance().version`。
 decltype(cuFileGetVersion)* GetVersion{nullptr};

 public
 int version{0};

 private
 cuFileAPI();

#ifdef KVIKIO_CUFILE_FOUND
 // 注意，即使在 main 之后不允许调用 CUDA [1]，我们仍必须在程序退出时关闭驱动程序（如果打开了）。
 // 这是因为，如果驱动程序在程序退出时没有关闭，cuFile 会发生段错误，即我们进退两难，但
 // 这似乎是两害相权取其轻。
 // [1] <https://docs.nvda.net.cn/cuda/cuda-c-programming-guide/index.html#initialization>
 ~cuFileAPI();
#endif

 public
 cuFileAPI(cuFileAPI const&) = delete;
 void operator=(cuFileAPI const&) = delete;
 cuFileAPI(cuFileAPI const&&) = delete;
 void operator=(cuFileAPI const&&) = delete;

 KVIKIO_EXPORT static cuFileAPI& instance();

 /// @brief 打开 cuFile 驱动程序。
 void driver_open();
 /// @brief 关闭 cuFile 驱动程序。
 void driver_close();
};
/// @brief 检查 cuFile 库是否可用。
#ifdef KVIKIO_CUFILE_FOUND
bool is_cufile_library_available() noexcept;
#else
constexpr bool is_cufile_library_available() noexcept { return false; }
#endif
/// @brief 检查 cuFile 是否可用且预期能够工作。
bool is_cufile_available() noexcept;
/// @brief 获取 cuFile 版本（如果早于 v1.8 则为零）。
#ifdef KVIKIO_CUFILE_FOUND
int cufile_version() noexcept;
#else
constexpr int cufile_version() noexcept { return 0; }
#endif
/// @brief 检查 cuFile 的批处理 API 是否可用。
bool is_batch_api_available() noexcept;
/// @brief 检查 cuFile 的流 (异步) API 是否可用。
bool is_stream_api_available() noexcept;