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}