1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116

use cpp::cpp;

use crate::device::DeviceId;
use crate::device::MemoryInfo;
use crate::ffi::result;

type Result<T> = std::result::Result<T, crate::error::Error>;

/// Synchronous implementation of [`crate::num_devices`].
///
/// Refer to [`crate::num_devices`] for documentation.
pub fn num_devices() -> Result<usize> {
    let mut num = 0_i32;
    let num_ptr = std::ptr::addr_of_mut!(num);
    let ret = cpp!(unsafe [
        num_ptr as "std::int32_t*"
    ] -> i32 as "std::int32_t" {
        return cudaGetDeviceCount(num_ptr);
    });

    result!(ret, num as usize)
}

/// Synchronous implementation of [`crate::Device`].
///
/// Refer to [`crate::Device`] for documentation.
pub struct Device;

impl Device {
    #[inline]
    pub fn get() -> Result<DeviceId> {
        let mut id: i32 = 0;
        let id_ptr = std::ptr::addr_of_mut!(id);
        let ret = cpp!(unsafe [
            id_ptr as "int*"
        ] -> i32 as "int" {
            return cudaGetDevice(id_ptr);
        });
        result!(ret, id)
    }

    #[inline(always)]
    pub fn get_or_panic() -> DeviceId {
        Device::get().unwrap_or_else(|err| panic!("failed to get device: {err}"))
    }

    #[inline]
    pub fn set(id: DeviceId) -> Result<()> {
        let ret = cpp!(unsafe [
            id as "int"
        ] -> i32 as "int" {
            return cudaSetDevice(id);
        });
        result!(ret)
    }

    #[inline(always)]
    pub fn set_or_panic(id: DeviceId) {
        Device::set(id).unwrap_or_else(|err| panic!("failed to set device {id}: {err}"));
    }

    pub fn synchronize() -> Result<()> {
        let ret = cpp!(unsafe [] -> i32 as "std::int32_t" {
            return cudaDeviceSynchronize();
        });
        result!(ret)
    }

    pub fn memory_info() -> Result<MemoryInfo> {
        let mut free: usize = 0;
        let free_ptr = std::ptr::addr_of_mut!(free);
        let mut total: usize = 0;
        let total_ptr = std::ptr::addr_of_mut!(total);

        let ret = cpp!(unsafe [
            free_ptr as "std::size_t*",
            total_ptr as "std::size_t*"
        ] -> i32 as "std::int32_t" {
            return cudaMemGetInfo(free_ptr, total_ptr);
        });
        result!(ret, MemoryInfo { free, total })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_num_devices() {
        assert!(matches!(num_devices(), Ok(num) if num > 0));
    }

    #[test]
    fn test_get_device() {
        assert!(matches!(Device::get(), Ok(0)));
    }

    #[test]
    fn test_set_device() {
        assert!(Device::set(0).is_ok());
        assert!(matches!(Device::get(), Ok(0)));
    }

    #[test]
    fn test_synchronize() {
        assert!(Device::synchronize().is_ok());
    }

    #[test]
    fn test_memory_info() {
        let memory_info = Device::memory_info().unwrap();
        assert!(memory_info.free > 0);
        assert!(memory_info.total > 0);
    }
}