driver.hpp

/*
 * SPDX-FileCopyrightText: Copyright (c) 2021-2025, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */
#pragma once
 
#include <vector>
 
#include <kvikio/shim/cufile.hpp>
#include <kvikio/shim/cufile_h_wrapper.hpp>
 
namespace kvikio {
 
#ifdef KVIKIO_CUFILE_FOUND
 
class DriverInitializer {
  // Optional, if not used cuFiles opens the driver automatically
 public:
  DriverInitializer();
 
  DriverInitializer(DriverInitializer const&)                = delete;
  DriverInitializer& operator=(DriverInitializer const&)     = delete;
  DriverInitializer(DriverInitializer&&) noexcept            = delete;
  DriverInitializer& operator=(DriverInitializer&&) noexcept = delete;
 
  ~DriverInitializer() noexcept;
};
 
class DriverProperties {
 private:
  CUfileDrvProps_t _props{};
  bool _initialized{false};
 
  // Because Cython does not handle exceptions in the default
  // constructor, we initialize `_props` lazily.
  void lazy_init();
 
 public:
  DriverProperties() = default;
 
  bool is_gds_available();
 
  [[nodiscard]] unsigned int get_nvfs_major_version();
 
  [[nodiscard]] unsigned int get_nvfs_minor_version();
 
  [[nodiscard]] bool get_nvfs_allow_compat_mode();
 
  [[nodiscard]] bool get_nvfs_poll_mode();
 
  [[nodiscard]] std::size_t get_nvfs_poll_thresh_size();
 
  void set_nvfs_poll_mode(bool enable);
 
  void set_nvfs_poll_thresh_size(std::size_t size_in_kb);
 
  [[nodiscard]] std::vector<CUfileDriverControlFlags> get_nvfs_statusflags();
 
  [[nodiscard]] std::size_t get_max_device_cache_size();
 
  void set_max_device_cache_size(std::size_t size_in_kb);
 
  [[nodiscard]] std::size_t get_per_buffer_cache_size();
 
  [[nodiscard]] std::size_t get_max_pinned_memory_size();
 
  void set_max_pinned_memory_size(std::size_t size_in_kb);
 
  [[nodiscard]] std::size_t get_max_batch_io_size();
};
 
#else
struct DriverInitializer {
  // Implement a non-default constructor to avoid `unused variable` warnings downstream
  DriverInitializer();
};
 
struct DriverProperties {
  // Implement a non-default constructor to avoid `unused variable` warnings downstream
  DriverProperties();
 
  static bool is_gds_available();
 
  [[nodiscard]] static unsigned int get_nvfs_major_version();
 
  [[nodiscard]] static unsigned int get_nvfs_minor_version();
 
  [[nodiscard]] static bool get_nvfs_allow_compat_mode();
 
  [[nodiscard]] static bool get_nvfs_poll_mode();
 
  [[nodiscard]] static std::size_t get_nvfs_poll_thresh_size();
 
  static void set_nvfs_poll_mode(bool enable);
 
  static void set_nvfs_poll_thresh_size(std::size_t size_in_kb);
 
  [[nodiscard]] static std::vector<CUfileDriverControlFlags> get_nvfs_statusflags();
 
  [[nodiscard]] static std::size_t get_max_device_cache_size();
 
  static void set_max_device_cache_size(std::size_t size_in_kb);
 
  [[nodiscard]] static std::size_t get_per_buffer_cache_size();
 
  [[nodiscard]] static std::size_t get_max_pinned_memory_size();
 
  static void set_max_pinned_memory_size(std::size_t size_in_kb);
 
  [[nodiscard]] std::size_t get_max_batch_io_size();
};
#endif
 
}  // namespace kvikio