async_cuda_core/ffi/
device.rs

1use cpp::cpp;
2
3use crate::device::DeviceId;
4use crate::device::MemoryInfo;
5use crate::ffi::result;
6
7type Result<T> = std::result::Result<T, crate::error::Error>;
8
9/// Synchronous implementation of [`crate::num_devices`].
10///
11/// Refer to [`crate::num_devices`] for documentation.
12pub fn num_devices() -> Result<usize> {
13    let mut num = 0_i32;
14    let num_ptr = std::ptr::addr_of_mut!(num);
15    let ret = cpp!(unsafe [
16        num_ptr as "std::int32_t*"
17    ] -> i32 as "std::int32_t" {
18        return cudaGetDeviceCount(num_ptr);
19    });
20
21    result!(ret, num as usize)
22}
23
24/// Synchronous implementation of [`crate::Device`].
25///
26/// Refer to [`crate::Device`] for documentation.
27pub struct Device;
28
29impl Device {
30    pub fn get() -> Result<DeviceId> {
31        let mut id: i32 = 0;
32        let id_ptr = std::ptr::addr_of_mut!(id);
33        let ret = cpp!(unsafe [
34            id_ptr as "std::int32_t*"
35        ] -> i32 as "std::int32_t" {
36            return cudaGetDevice(id_ptr);
37        });
38        result!(ret, id as DeviceId)
39    }
40
41    pub fn set(id: DeviceId) -> Result<()> {
42        let id = id as i32;
43        let ret = cpp!(unsafe [
44            id as "std::int32_t"
45        ] -> i32 as "std::int32_t" {
46            return cudaSetDevice(id);
47        });
48        result!(ret)
49    }
50
51    pub fn synchronize() -> Result<()> {
52        let ret = cpp!(unsafe [] -> i32 as "std::int32_t" {
53            return cudaDeviceSynchronize();
54        });
55        result!(ret)
56    }
57
58    pub fn memory_info() -> Result<MemoryInfo> {
59        let mut free: usize = 0;
60        let free_ptr = std::ptr::addr_of_mut!(free);
61        let mut total: usize = 0;
62        let total_ptr = std::ptr::addr_of_mut!(total);
63
64        let ret = cpp!(unsafe [
65            free_ptr as "std::size_t*",
66            total_ptr as "std::size_t*"
67        ] -> i32 as "std::int32_t" {
68            return cudaMemGetInfo(free_ptr, total_ptr);
69        });
70        result!(ret, MemoryInfo { free, total })
71    }
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77
78    #[test]
79    fn test_num_devices() {
80        assert!(matches!(num_devices(), Ok(num) if num > 0));
81    }
82
83    #[test]
84    fn test_get_device() {
85        assert!(matches!(Device::get(), Ok(0)));
86    }
87
88    #[test]
89    fn test_set_device() {
90        assert!(Device::set(0).is_ok());
91        assert!(matches!(Device::get(), Ok(0)));
92    }
93
94    #[test]
95    fn test_synchronize() {
96        assert!(Device::synchronize().is_ok());
97    }
98
99    #[test]
100    fn test_memory_info() {
101        let memory_info = Device::memory_info().unwrap();
102        assert!(memory_info.free > 0);
103        assert!(memory_info.total > 0);
104    }
105}
async_cuda_core/ffi/device.rs

async_cuda_core/ffi/
device.rs
