torch_sys/
cuda.rs

1use libc::c_int;
2
3extern "C" {
4    /// Returns the number of CUDA devices available.
5    pub fn atc_cuda_device_count() -> c_int;
6
7    /// Returns true if at least one CUDA device is available.
8    pub fn atc_cuda_is_available() -> c_int;
9
10    /// Returns true if CUDA is available, and CuDNN is available.
11    pub fn atc_cudnn_is_available() -> c_int;
12
13    /// Sets the seed for the current GPU.
14    pub fn atc_manual_seed(seed: u64);
15
16    /// Sets the seed for all available GPUs.
17    pub fn atc_manual_seed_all(seed: u64);
18
19    /// Waits for all kernels in all streams on a CUDA device to complete.
20    pub fn atc_synchronize(device_index: i64);
21
22    /// Returns true if CUDNN is enabled by the user.
23    pub fn atc_user_enabled_cudnn() -> c_int;
24
25    /// Enable or disable CUDNN.
26    pub fn atc_set_user_enabled_cudnn(b: c_int);
27
28    /// Sets CUDNN benchmark mode.
29    pub fn atc_set_benchmark_cudnn(b: c_int);
30}