era_cudart 0.1.0

CUDA bindings for ZKsync
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

// Peer Device Memory Access
// https://docs.nvidia.com/cuda/cuda-runtime-api/group__CUDART__PEER.html

use std::mem::MaybeUninit;

use cudart_sys::*;

use crate::result::{CudaResult, CudaResultWrap};

pub fn device_can_access_peer(device_id: i32, device_peer_id: i32) -> CudaResult<bool> {
    let mut can_access_peer = MaybeUninit::<i32>::uninit();
    unsafe {
        cudaDeviceCanAccessPeer(can_access_peer.as_mut_ptr(), device_id, device_peer_id)
            .wrap_maybe_uninit(can_access_peer)
            .map(|value| value != 0)
    }
}

pub fn device_disable_peer_access(device_peer_id: i32) -> CudaResult<()> {
    unsafe { cudaDeviceDisablePeerAccess(device_peer_id).wrap() }
}

pub fn device_enable_peer_access(device_peer_id: i32) -> CudaResult<()> {
    unsafe { cudaDeviceEnablePeerAccess(device_peer_id, 0).wrap() }
}

#[cfg(test)]
mod tests {
    use serial_test::serial;

    use crate::device::device_reset;

    use super::*;

    #[test]
    #[serial]
    #[ignore = "needs multiple peer-enabled GPUs"]
    fn device_can_access_peer_is_true() {
        let result = device_can_access_peer(0, 1);
        assert_eq!(result, Ok(true));
    }

    #[test]
    #[serial]
    #[ignore = "needs multiple peer-enabled GPUs"]
    fn device_disable_peer_access_is_ok() {
        device_reset().unwrap();
        device_enable_peer_access(1).unwrap();
        let result = device_disable_peer_access(1);
        assert_eq!(result, Ok(()));
    }

    #[test]
    #[serial]
    #[ignore = "needs multiple peer-enabled GPUs"]
    fn device_enable_peer_access_is_ok() {
        device_reset().unwrap();
        let result = device_enable_peer_access(1);
        assert_eq!(result, Ok(()));
        device_disable_peer_access(1).unwrap();
    }
}