Skip to main content

oxicuda_driver/
memory_info.rs

1//! Safe wrappers for GPU memory information and bulk memory operations.
2//!
3//! This module provides convenient functions for querying device memory
4//! usage and performing device-to-device copies and memsets.
5//!
6//! # Example
7//!
8//! ```rust,no_run
9//! # use oxicuda_driver::memory_info;
10//! # fn main() -> Result<(), oxicuda_driver::CudaError> {
11//! let (free, total) = memory_info::device_memory_info()?;
12//! println!("GPU memory: {free} / {total} bytes free");
13//! # Ok(())
14//! # }
15//! ```
16
17use crate::error::{CudaError, CudaResult};
18use crate::ffi::CUdeviceptr;
19use crate::loader::try_driver;
20use crate::stream::Stream;
21
22/// Returns the amount of free and total device memory in bytes.
23///
24/// This queries the current context's device for its memory utilisation.
25///
26/// # Returns
27///
28/// A tuple `(free_bytes, total_bytes)`.
29///
30/// # Errors
31///
32/// Returns a [`CudaError`] if no context is current or the driver call fails.
33pub fn device_memory_info() -> CudaResult<(usize, usize)> {
34    let api = try_driver()?;
35    let mut free: usize = 0;
36    let mut total: usize = 0;
37    crate::cuda_call!((api.cu_mem_get_info_v2)(&mut free, &mut total))?;
38    Ok((free, total))
39}
40
41/// Copies `bytes` bytes from one device pointer to another.
42///
43/// Both `src` and `dst` must point to valid device memory allocations of
44/// at least `bytes` bytes.
45///
46/// # Errors
47///
48/// Returns a [`CudaError`] if the copy fails.
49pub fn memcpy_device_to_device(dst: CUdeviceptr, src: CUdeviceptr, bytes: usize) -> CudaResult<()> {
50    let api = try_driver()?;
51    crate::cuda_call!((api.cu_memcpy_dtod_v2)(dst, src, bytes))
52}
53
54/// Asynchronously copies `bytes` bytes from one device pointer to another.
55///
56/// The copy is enqueued on the given `stream` and returns immediately.
57///
58/// # Errors
59///
60/// Returns [`CudaError::NotSupported`] if the driver does not export the
61/// async DtoD copy entry point, or another [`CudaError`] on failure.
62pub fn memcpy_device_to_device_async(
63    dst: CUdeviceptr,
64    src: CUdeviceptr,
65    bytes: usize,
66    stream: &Stream,
67) -> CudaResult<()> {
68    let api = try_driver()?;
69    let f = api.cu_memcpy_dtod_async_v2.ok_or(CudaError::NotSupported)?;
70    crate::cuda_call!(f(dst, src, bytes, stream.raw()))
71}
72
73/// Sets `count` 32-bit elements starting at `ptr` to `value`.
74///
75/// # Errors
76///
77/// Returns a [`CudaError`] if the memset fails.
78pub fn memset_d32(ptr: CUdeviceptr, value: u32, count: usize) -> CudaResult<()> {
79    let api = try_driver()?;
80    crate::cuda_call!((api.cu_memset_d32_v2)(ptr, value, count))
81}
82
83/// Asynchronously sets `count` 32-bit elements starting at `ptr` to `value`.
84///
85/// The operation is enqueued on the given `stream`.
86///
87/// # Errors
88///
89/// Returns [`CudaError::NotSupported`] if the driver does not export the
90/// async memset entry point, or another [`CudaError`] on failure.
91pub fn memset_d32_async(
92    ptr: CUdeviceptr,
93    value: u32,
94    count: usize,
95    stream: &Stream,
96) -> CudaResult<()> {
97    let api = try_driver()?;
98    let f = api.cu_memset_d32_async.ok_or(CudaError::NotSupported)?;
99    crate::cuda_call!(f(ptr, value, count, stream.raw()))
100}
101
102// ---------------------------------------------------------------------------
103// Tests
104// ---------------------------------------------------------------------------
105
106#[cfg(test)]
107mod tests {
108    use super::*;
109
110    #[test]
111    fn device_memory_info_returns_error_without_gpu() {
112        // On macOS or systems without a GPU, this should return an error
113        // rather than panicking.
114        let result = device_memory_info();
115        // We just verify it does not panic; on CI without a GPU it will be Err.
116        let _ = result;
117    }
118
119    #[test]
120    fn memcpy_dtod_returns_error_without_gpu() {
121        let result = memcpy_device_to_device(0, 0, 0);
122        let _ = result;
123    }
124
125    #[test]
126    fn memset_d32_returns_error_without_gpu() {
127        let result = memset_d32(0, 0, 0);
128        let _ = result;
129    }
130}