nexar-rdma 0.1.2

RDMA and GPUDirect transport extensions for nexar
Documentation 
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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139

//! CUDA device adapter for GPU memory staging.
//!
//! Uses cudarc 0.19 to copy GPU memory to/from host for network transfer.
//! Without GPUDirect RDMA, the path is GPU → Host → NIC → Host → GPU.

use nexar::device::DeviceAdapter;
use nexar::error::{NexarError, Result};
use nexar::types::{DataType, ReduceOp};
use std::sync::Arc;

/// Device adapter for NVIDIA GPUs using cudarc.
///
/// Handles device↔host staging for network I/O and CPU-side reduction.
pub struct CudaAdapter {
    ctx: Arc<cudarc::driver::CudaContext>,
}

impl CudaAdapter {
    /// Create a new CUDA adapter for the given GPU ordinal.
    pub fn new(device_ordinal: u32) -> Result<Self> {
        let ctx = cudarc::driver::CudaContext::new(device_ordinal as usize).map_err(|e| {
            NexarError::device(format!(
                "failed to create CUDA context for GPU {device_ordinal}: {e}"
            ))
        })?;
        Ok(Self { ctx })
    }
}

impl DeviceAdapter for CudaAdapter {
    unsafe fn stage_for_send_on_stream(
        &self,
        ptr: u64,
        size_bytes: usize,
        stream: u64,
    ) -> Result<Vec<u8>> {
        self.ctx
            .bind_to_thread()
            .map_err(|e| NexarError::device(format!("CUDA bind_to_thread failed: {e}")))?;
        let mut buf = vec![0u8; size_bytes];
        unsafe {
            cudarc::driver::result::memcpy_dtoh_async(
                &mut buf,
                ptr as cudarc::driver::sys::CUdeviceptr,
                stream as cudarc::driver::sys::CUstream,
            )
        }
        .map_err(|e| NexarError::device(format!("cuMemcpyDtoHAsync failed: {e}")))?;
        Ok(buf)
    }

    unsafe fn receive_to_device_on_stream(
        &self,
        data: &[u8],
        dst_ptr: u64,
        stream: u64,
    ) -> Result<()> {
        self.ctx
            .bind_to_thread()
            .map_err(|e| NexarError::device(format!("CUDA bind_to_thread failed: {e}")))?;
        unsafe {
            cudarc::driver::result::memcpy_htod_async(
                dst_ptr as cudarc::driver::sys::CUdeviceptr,
                data,
                stream as cudarc::driver::sys::CUstream,
            )
        }
        .map_err(|e| NexarError::device(format!("cuMemcpyHtoDAsync failed: {e}")))?;
        Ok(())
    }

    unsafe fn stage_for_send(&self, ptr: u64, size_bytes: usize) -> Result<Vec<u8>> {
        self.ctx
            .bind_to_thread()
            .map_err(|e| NexarError::device(format!("CUDA bind_to_thread failed: {e}")))?;
        let mut buf = vec![0u8; size_bytes];
        unsafe {
            cudarc::driver::result::memcpy_dtoh_sync(
                &mut buf,
                ptr as cudarc::driver::sys::CUdeviceptr,
            )
        }
        .map_err(|e| NexarError::device(format!("cuMemcpyDtoH failed: {e}")))?;
        Ok(buf)
    }

    unsafe fn receive_to_device(&self, data: &[u8], dst_ptr: u64) -> Result<()> {
        self.ctx
            .bind_to_thread()
            .map_err(|e| NexarError::device(format!("CUDA bind_to_thread failed: {e}")))?;
        unsafe {
            cudarc::driver::result::memcpy_htod_sync(
                dst_ptr as cudarc::driver::sys::CUdeviceptr,
                data,
            )
        }
        .map_err(|e| NexarError::device(format!("cuMemcpyHtoD failed: {e}")))?;
        Ok(())
    }

    unsafe fn reduce_inplace(
        &self,
        dst_ptr: u64,
        src: &[u8],
        count: usize,
        dtype: DataType,
        op: ReduceOp,
    ) -> Result<()> {
        let elem_size = dtype.size_in_bytes();
        let total_bytes = count * elem_size;

        self.ctx
            .bind_to_thread()
            .map_err(|e| NexarError::device(format!("CUDA bind_to_thread failed: {e}")))?;

        let mut dst_host = vec![0u8; total_bytes];
        unsafe {
            cudarc::driver::result::memcpy_dtoh_sync(
                &mut dst_host,
                dst_ptr as cudarc::driver::sys::CUdeviceptr,
            )
        }
        .map_err(|e| NexarError::device(format!("cuMemcpyDtoH (reduce) failed: {e}")))?;

        // Perform CPU-side reduction into dst_host.
        let dst_host_ptr = dst_host.as_mut_ptr() as u64;
        unsafe { nexar::reduce::reduce_ptr(dst_host_ptr, src, count, dtype, op)? };

        unsafe {
            cudarc::driver::result::memcpy_htod_sync(
                dst_ptr as cudarc::driver::sys::CUdeviceptr,
                &dst_host,
            )
        }
        .map_err(|e| NexarError::device(format!("cuMemcpyHtoD (reduce) failed: {e}")))?;

        Ok(())
    }
}