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}