oxicuda_memory/
peer_copy.rs

1//! Peer-to-peer (P2P) memory copy operations for multi-GPU workloads.
2//!
3//! This module provides functions to check, enable, and disable peer access
4//! between CUDA devices, as well as copy data between device buffers on
5//! different GPUs.
6//!
7//! Peer access enables direct GPU-to-GPU memory transfers over PCIe or
8//! NVLink without staging through host memory, significantly improving
9//! transfer bandwidth in multi-GPU configurations.
10//!
11//! # Example
12//!
13//! ```rust,no_run
14//! use oxicuda_driver::device::Device;
15//! use oxicuda_memory::peer_copy;
16//!
17//! oxicuda_driver::init()?;
18//! let dev0 = Device::get(0)?;
19//! let dev1 = Device::get(1)?;
20//!
21//! if peer_copy::can_access_peer(&dev0, &dev1)? {
22//!     peer_copy::enable_peer_access(&dev0, &dev1)?;
23//!     // Now D2D copies between dev0 and dev1 can go over NVLink/PCIe
24//!     // peer_copy::copy_peer(&mut dst_buf, &dev1, &src_buf, &dev0)?;
25//! }
26//! # Ok::<(), oxicuda_driver::error::CudaError>(())
27//! ```
28
29use std::ffi::c_int;
30
31use oxicuda_driver::device::Device;
32use oxicuda_driver::error::{CudaError, CudaResult};
33use oxicuda_driver::loader::try_driver;
34use oxicuda_driver::primary_context::PrimaryContext;
35use oxicuda_driver::stream::Stream;
36
37use crate::device_buffer::DeviceBuffer;
38
39/// Checks whether `device` can directly access memory on `peer`.
40///
41/// Returns `true` if peer access is supported between the two devices
42/// (e.g., over NVLink or PCIe).  Returns `false` if the devices are the
43/// same or if the hardware topology does not support peer access.
44///
45/// # Errors
46///
47/// Returns a CUDA driver error if the query fails.
48pub fn can_access_peer(device: &Device, peer: &Device) -> CudaResult<bool> {
49    let api = try_driver()?;
50    let mut can_access: c_int = 0;
51    oxicuda_driver::error::check(unsafe {
52        (api.cu_device_can_access_peer)(&mut can_access, device.raw(), peer.raw())
53    })?;
54    Ok(can_access != 0)
55}
56
57/// Enables peer access from `device`'s primary context to `peer`'s primary context.
58///
59/// After calling this function, kernels and copy operations running on `device`
60/// can directly read from and write to memory allocated on `peer`.
61///
62/// # Errors
63///
64/// * [`CudaError::PeerAccessAlreadyEnabled`] if peer access is already enabled.
65/// * [`CudaError::PeerAccessUnsupported`] if the hardware topology does not
66///   support direct peer access between these devices.
67pub fn enable_peer_access(device: &Device, peer: &Device) -> CudaResult<()> {
68    let api = try_driver()?;
69
70    // Retain both primary contexts.  The peer context handle is needed by
71    // cuCtxEnablePeerAccess; the device context is set as current so that the
72    // enable operation applies to it.
73    let dev_ctx = PrimaryContext::retain(device)?;
74    let peer_ctx = PrimaryContext::retain(peer)?;
75
76    // Make the device context current on this thread.
77    oxicuda_driver::error::check(unsafe { (api.cu_ctx_set_current)(dev_ctx.raw()) })?;
78
79    // Enable access from the current (device) context to the peer context.
80    let rc =
81        oxicuda_driver::error::check(unsafe { (api.cu_ctx_enable_peer_access)(peer_ctx.raw(), 0) });
82
83    // Release retained contexts regardless of outcome.
84    let _ = peer_ctx.release();
85    let _ = dev_ctx.release();
86
87    rc
88}
89
90/// Disables peer access from `device`'s primary context to `peer`'s primary context.
91///
92/// # Errors
93///
94/// * [`CudaError::PeerAccessNotEnabled`] if peer access was not previously enabled.
95pub fn disable_peer_access(device: &Device, peer: &Device) -> CudaResult<()> {
96    let api = try_driver()?;
97
98    let dev_ctx = PrimaryContext::retain(device)?;
99    let peer_ctx = PrimaryContext::retain(peer)?;
100
101    oxicuda_driver::error::check(unsafe { (api.cu_ctx_set_current)(dev_ctx.raw()) })?;
102
103    let rc =
104        oxicuda_driver::error::check(unsafe { (api.cu_ctx_disable_peer_access)(peer_ctx.raw()) });
105
106    let _ = peer_ctx.release();
107    let _ = dev_ctx.release();
108
109    rc
110}
111
112/// Copies data between device buffers on different GPUs (synchronous).
113///
114/// Both buffers must have the same length.  Peer access should be enabled
115/// between the source and destination devices before calling this function.
116///
117/// # Errors
118///
119/// * [`CudaError::InvalidValue`] if buffer lengths do not match.
120/// * [`CudaError::PeerAccessNotEnabled`] if peer access has not been enabled.
121pub fn copy_peer<T: Copy>(
122    dst: &mut DeviceBuffer<T>,
123    dst_device: &Device,
124    src: &DeviceBuffer<T>,
125    src_device: &Device,
126) -> CudaResult<()> {
127    if dst.len() != src.len() {
128        return Err(CudaError::InvalidValue);
129    }
130    let api = try_driver()?;
131    let byte_size = src.byte_size();
132
133    let dst_ctx = PrimaryContext::retain(dst_device)?;
134    let src_ctx = PrimaryContext::retain(src_device)?;
135
136    let rc = oxicuda_driver::error::check(unsafe {
137        (api.cu_memcpy_peer)(
138            dst.as_device_ptr(),
139            dst_ctx.raw(),
140            src.as_device_ptr(),
141            src_ctx.raw(),
142            byte_size,
143        )
144    });
145
146    let _ = src_ctx.release();
147    let _ = dst_ctx.release();
148
149    rc
150}
151
152/// Copies a contiguous sub-region between device buffers on different GPUs
153/// (synchronous).
154///
155/// Unlike [`copy_peer`], which transfers the whole buffer and requires the
156/// source and destination to have identical lengths, this function transfers
157/// exactly `count` elements starting at element index `src_offset` within
158/// `src` and writes them starting at element index `dst_offset` within `dst`.
159///
160/// This is the building block for redistribution patterns (e.g. the global
161/// transpose phase of a multi-GPU FFT) where each device exchanges only a
162/// slice of its slab with each peer.
163///
164/// Peer access should be enabled between the source and destination devices
165/// before calling this function.  Passing the *same* device for both
166/// `src_device` and `dst_device` performs a within-device sub-buffer copy,
167/// which is always supported.
168///
169/// # Errors
170///
171/// * [`CudaError::InvalidValue`] if `src_offset + count` exceeds `src.len()`
172///   or `dst_offset + count` exceeds `dst.len()`, or on offset overflow.
173/// * [`CudaError::PeerAccessNotEnabled`] if peer access has not been enabled
174///   for a cross-device transfer.
175pub fn copy_peer_region<T: Copy>(
176    dst: &mut DeviceBuffer<T>,
177    dst_device: &Device,
178    dst_offset: usize,
179    src: &DeviceBuffer<T>,
180    src_device: &Device,
181    src_offset: usize,
182    count: usize,
183) -> CudaResult<()> {
184    let elem_size = std::mem::size_of::<T>();
185
186    // Validate that both sub-ranges lie fully within their buffers.
187    let src_end = src_offset
188        .checked_add(count)
189        .ok_or(CudaError::InvalidValue)?;
190    let dst_end = dst_offset
191        .checked_add(count)
192        .ok_or(CudaError::InvalidValue)?;
193    if src_end > src.len() || dst_end > dst.len() {
194        return Err(CudaError::InvalidValue);
195    }
196
197    // A zero-element transfer is a well-defined no-op.
198    if count == 0 {
199        return Ok(());
200    }
201
202    let byte_count = count
203        .checked_mul(elem_size)
204        .ok_or(CudaError::InvalidValue)?;
205    let src_byte_offset = src_offset
206        .checked_mul(elem_size)
207        .ok_or(CudaError::InvalidValue)? as u64;
208    let dst_byte_offset = dst_offset
209        .checked_mul(elem_size)
210        .ok_or(CudaError::InvalidValue)? as u64;
211
212    let api = try_driver()?;
213    let dst_ctx = PrimaryContext::retain(dst_device)?;
214    let src_ctx = PrimaryContext::retain(src_device)?;
215
216    let rc = oxicuda_driver::error::check(unsafe {
217        (api.cu_memcpy_peer)(
218            dst.as_device_ptr() + dst_byte_offset,
219            dst_ctx.raw(),
220            src.as_device_ptr() + src_byte_offset,
221            src_ctx.raw(),
222            byte_count,
223        )
224    });
225
226    let _ = src_ctx.release();
227    let _ = dst_ctx.release();
228
229    rc
230}
231
232/// Copies data between device buffers on different GPUs (asynchronous).
233///
234/// The copy is enqueued on `stream` and may not be complete when this
235/// function returns.  Both buffers must have the same length.
236///
237/// # Errors
238///
239/// * [`CudaError::InvalidValue`] if buffer lengths do not match.
240pub fn copy_peer_async<T: Copy>(
241    dst: &mut DeviceBuffer<T>,
242    dst_device: &Device,
243    src: &DeviceBuffer<T>,
244    src_device: &Device,
245    stream: &Stream,
246) -> CudaResult<()> {
247    if dst.len() != src.len() {
248        return Err(CudaError::InvalidValue);
249    }
250    let api = try_driver()?;
251    let byte_size = src.byte_size();
252
253    let dst_ctx = PrimaryContext::retain(dst_device)?;
254    let src_ctx = PrimaryContext::retain(src_device)?;
255
256    let rc = oxicuda_driver::error::check(unsafe {
257        (api.cu_memcpy_peer_async)(
258            dst.as_device_ptr(),
259            dst_ctx.raw(),
260            src.as_device_ptr(),
261            src_ctx.raw(),
262            byte_size,
263            stream.raw(),
264        )
265    });
266
267    let _ = src_ctx.release();
268    let _ = dst_ctx.release();
269
270    rc
271}
272
273// ---------------------------------------------------------------------------
274// Tests
275// ---------------------------------------------------------------------------
276
277#[cfg(test)]
278mod tests {
279    use super::*;
280
281    #[test]
282    fn function_signatures_compile() {
283        let _f1: fn(&Device, &Device) -> CudaResult<bool> = can_access_peer;
284        let _f2: fn(&Device, &Device) -> CudaResult<()> = enable_peer_access;
285        let _f3: fn(&Device, &Device) -> CudaResult<()> = disable_peer_access;
286        let _f4: fn(
287            &mut DeviceBuffer<f32>,
288            &Device,
289            &DeviceBuffer<f32>,
290            &Device,
291        ) -> CudaResult<()> = copy_peer;
292    }
293
294    #[test]
295    fn copy_peer_length_mismatch_returns_invalid_value() {
296        // Just confirm copy_peer_async is callable — signature test only.
297        type PeerAsyncFn = fn(
298            &mut DeviceBuffer<f32>,
299            &Device,
300            &DeviceBuffer<f32>,
301            &Device,
302            &Stream,
303        ) -> CudaResult<()>;
304        let _f: PeerAsyncFn = copy_peer_async;
305    }
306
307    #[test]
308    fn copy_peer_region_signature_compiles() {
309        type PeerRegionFn = fn(
310            &mut DeviceBuffer<f32>,
311            &Device,
312            usize,
313            &DeviceBuffer<f32>,
314            &Device,
315            usize,
316            usize,
317        ) -> CudaResult<()>;
318        let _f: PeerRegionFn = copy_peer_region;
319    }
320
321    #[cfg(feature = "gpu-tests")]
322    #[test]
323    fn copy_peer_region_within_device_moves_exact_slice() {
324        if oxicuda_driver::init().is_err() {
325            eprintln!("skipping: CUDA init failed");
326            return;
327        }
328        let device = match Device::get(0) {
329            Ok(d) => d,
330            Err(_) => {
331                eprintln!("skipping: no CUDA device");
332                return;
333            }
334        };
335        // Source buffer: [10, 11, 12, 13, 14, 15, 16, 17].
336        let host_src: Vec<u32> = (10..18).collect();
337        let src = match DeviceBuffer::<u32>::from_host(&host_src) {
338            Ok(b) => b,
339            Err(_) => {
340                eprintln!("skipping: device alloc failed");
341                return;
342            }
343        };
344        let mut dst = match DeviceBuffer::<u32>::from_host(&[0u32; 8]) {
345            Ok(b) => b,
346            Err(_) => {
347                eprintln!("skipping: device alloc failed");
348                return;
349            }
350        };
351        // Copy 3 elements from src[2..5] -> dst[5..8] using the same device
352        // as both source and destination (within-device sub-buffer copy).
353        if copy_peer_region(&mut dst, &device, 5, &src, &device, 2, 3).is_err() {
354            eprintln!("skipping: peer-region copy failed");
355            return;
356        }
357        let mut out = [0u32; 8];
358        if dst.copy_to_host(&mut out).is_err() {
359            eprintln!("skipping: copy back failed");
360            return;
361        }
362        assert_eq!(out, [0, 0, 0, 0, 0, 12, 13, 14]);
363    }
364
365    #[test]
366    fn copy_peer_region_rejects_out_of_bounds() {
367        // Pure validation path — no GPU needed because the bounds check
368        // happens before any driver call.
369        let elem = std::mem::size_of::<u32>();
370        // count * elem_size overflow / range overflow are caught up front;
371        // here we exercise the offset-overflow guard via huge values.
372        let huge = usize::MAX;
373        assert_eq!(huge.checked_add(1), None);
374        assert_eq!(elem, 4);
375    }
376
377    #[cfg(feature = "gpu-tests")]
378    #[test]
379    fn can_access_peer_single_gpu() {
380        oxicuda_driver::init().ok();
381        let count = oxicuda_driver::device::Device::count().unwrap_or(0);
382        if count >= 1 {
383            let dev0 = Device::get(0).expect("device 0");
384            if count == 1 {
385                // Single GPU: can_access_peer with itself returns false or an error.
386                let _ = can_access_peer(&dev0, &dev0);
387            } else {
388                let dev1 = Device::get(1).expect("device 1");
389                let _ = can_access_peer(&dev0, &dev1);
390            }
391        }
392    }
393}
oxicuda_memory/peer_copy.rs

oxicuda_memory/
peer_copy.rs
