Files
xtrain/crates/xtrain-cuda/src/device.rs

17 lines
407 B
Rust

use crate::error::{self, Result};
use crate::ffi;
pub fn device_count() -> Result<i32> {
let mut count = 0;
error::check(unsafe { ffi::cudaGetDeviceCount(&mut count) })?;
Ok(count)
}
pub fn set_device(device: u32) -> Result<()> {
error::check(unsafe { ffi::cudaSetDevice(device as i32) })
}
pub fn synchronize() -> Result<()> {
error::check(unsafe { ffi::cudaDeviceSynchronize() })
}