fil_rustacuda/memory/
malloc.rs

1use super::DeviceCopy;
2use crate::error::*;
3use crate::memory::DevicePointer;
4use crate::memory::UnifiedPointer;
5use std::mem;
6use std::os::raw::c_void;
7use std::ptr;
8
9/// Unsafe wrapper around the `cuMemAlloc` function, which allocates some device memory and
10/// returns a [`DevicePointer`](struct.DevicePointer.html) pointing to it. The memory is not cleared.
11///
12/// Note that `count` is in units of T; thus a `count` of 3 will allocate `3 * size_of::<T>()` bytes
13/// of memory.
14///
15/// Memory buffers allocated using `cuda_malloc` must be freed using [`cuda_free`](fn.cuda_free.html).
16///
17/// # Errors
18///
19/// If allocating memory fails, returns the CUDA error value.
20/// If the number of bytes to allocate is zero (either because count is zero or because T is a
21/// zero-sized type), or if the size of the allocation would overflow a usize, returns InvalidValue.
22///
23/// # Safety
24///
25/// Since the allocated memory is not initialized, the caller must ensure that it is initialized
26/// before copying it to the host in any way. Additionally, the caller must ensure that the memory
27/// allocated is freed using cuda_free, or the memory will be leaked.
28///
29/// # Examples
30///
31/// ```
32/// # let _context = rustacuda::quick_init().unwrap();
33/// use rustacuda::memory::*;
34/// unsafe {
35///     // Allocate space for 5 u64s
36///     let device_buffer = cuda_malloc::<u64>(5).unwrap();
37///     cuda_free(device_buffer).unwrap();
38/// }
39/// ```
40pub unsafe fn cuda_malloc<T>(count: usize) -> CudaResult<DevicePointer<T>> {
41    let size = count.checked_mul(mem::size_of::<T>()).unwrap_or(0);
42    if size == 0 {
43        return Err(CudaError::InvalidMemoryAllocation);
44    }
45
46    let mut ptr: *mut c_void = ptr::null_mut();
47    cuda_driver_sys::cuMemAlloc_v2(&mut ptr as *mut *mut c_void as *mut u64, size).to_result()?;
48    let ptr = ptr as *mut T;
49    Ok(DevicePointer::wrap(ptr as *mut T))
50}
51
52/// Unsafe wrapper around the `cuMemAllocManaged` function, which allocates some unified memory and
53/// returns a [`UnifiedPointer`](struct.UnifiedPointer.html) pointing to it. The memory is not cleared.
54///
55/// Note that `count` is in units of T; thus a `count` of 3 will allocate `3 * size_of::<T>()` bytes
56/// of memory.
57///
58/// Memory buffers allocated using `cuda_malloc_unified` must be freed using [`cuda_free`](fn.cuda_free.html).
59///
60/// # Errors
61///
62/// If allocating memory fails, returns the CUDA error value.
63/// If the number of bytes to allocate is zero (either because count is zero or because T is a
64/// zero-sized type), or if the size of the allocation would overflow a usize, returns InvalidValue.
65///
66/// # Safety
67///
68/// Since the allocated memory is not initialized, the caller must ensure that it is initialized
69/// before reading from it in any way. Additionally, the caller must ensure that the memory
70/// allocated is freed using cuda_free, or the memory will be leaked.
71///
72/// # Examples
73///
74/// ```
75/// # let _context = rustacuda::quick_init().unwrap();
76/// use rustacuda::memory::*;
77/// unsafe {
78///     // Allocate space for a u64
79///     let mut unified_buffer = cuda_malloc_unified::<u64>(1).unwrap();
80///     // Write to it
81///     *unified_buffer.as_raw_mut() = 5u64;
82///     cuda_free_unified(unified_buffer).unwrap();
83/// }
84/// ```
85pub unsafe fn cuda_malloc_unified<T: DeviceCopy>(count: usize) -> CudaResult<UnifiedPointer<T>> {
86    let size = count.checked_mul(mem::size_of::<T>()).unwrap_or(0);
87    if size == 0 {
88        return Err(CudaError::InvalidMemoryAllocation);
89    }
90
91    let mut ptr: *mut c_void = ptr::null_mut();
92    cuda_driver_sys::cuMemAllocManaged(
93        &mut ptr as *mut *mut c_void as *mut u64,
94        size,
95        cuda_driver_sys::CUmemAttach_flags_enum::CU_MEM_ATTACH_GLOBAL as u32,
96    )
97    .to_result()?;
98    let ptr = ptr as *mut T;
99    Ok(UnifiedPointer::wrap(ptr as *mut T))
100}
101
102/// Free memory allocated with [`cuda_malloc`](fn.cuda_malloc.html).
103///
104/// # Errors
105///
106/// If freeing memory fails, returns the CUDA error value. If the given pointer is null, returns
107/// InvalidValue.
108///
109/// # Safety
110///
111/// The given pointer must have been allocated with `cuda_malloc`, or null.
112/// The caller is responsible for ensuring that no other pointers to the deallocated buffer exist.
113///
114/// # Examples
115///
116/// ```
117/// # let _context = rustacuda::quick_init().unwrap();
118/// use rustacuda::memory::*;
119/// unsafe {
120///     let device_buffer = cuda_malloc::<u64>(5).unwrap();
121///     // Free allocated memory.
122///     cuda_free(device_buffer).unwrap();
123/// }
124/// ```
125pub unsafe fn cuda_free<T>(mut p: DevicePointer<T>) -> CudaResult<()> {
126    let ptr = p.as_raw_mut();
127    if ptr.is_null() {
128        return Err(CudaError::InvalidMemoryAllocation);
129    }
130
131    cuda_driver_sys::cuMemFree_v2(ptr as u64).to_result()?;
132    Ok(())
133}
134
135/// Free memory allocated with [`cuda_malloc_unified`](fn.cuda_malloc_unified.html).
136///
137/// # Errors
138///
139/// If freeing memory fails, returns the CUDA error value. If the given pointer is null, returns
140/// InvalidValue.
141///
142/// # Safety
143///
144/// The given pointer must have been allocated with `cuda_malloc_unified`, or null.
145/// The caller is responsible for ensuring that no other pointers to the deallocated buffer exist.
146///
147/// # Examples
148///
149/// ```
150/// # let _context = rustacuda::quick_init().unwrap();
151/// use rustacuda::memory::*;
152/// unsafe {
153///     let unified_buffer = cuda_malloc_unified::<u64>(5).unwrap();
154///     // Free allocated memory.
155///     cuda_free_unified(unified_buffer).unwrap();
156/// }
157/// ```
158pub unsafe fn cuda_free_unified<T: DeviceCopy>(mut p: UnifiedPointer<T>) -> CudaResult<()> {
159    let ptr = p.as_raw_mut();
160    if ptr.is_null() {
161        return Err(CudaError::InvalidMemoryAllocation);
162    }
163
164    cuda_driver_sys::cuMemFree_v2(ptr as u64).to_result()?;
165    Ok(())
166}
167
168/// Unsafe wrapper around the `cuMemAllocHost` function, which allocates some page-locked host memory
169/// and returns a raw pointer pointing to it. The memory is not cleared.
170///
171/// Note that `count` is in units of T; thus a `count` of 3 will allocate `3 * size_of::<T>()` bytes
172/// of memory.
173///
174/// Memory buffers allocated using `cuda_malloc_locked` must be freed using [`cuda_free_locked`](fn.cuda_free_locked.html).
175///
176/// # Errors
177///
178/// If allocating memory fails, returns the CUDA error value.
179/// If the number of bytes to allocate is zero (either because count is zero or because T is a
180/// zero-sized type), or if the size of the allocation would overflow a usize, returns InvalidValue.
181///
182/// # Safety
183///
184/// Since the allocated memory is not initialized, the caller must ensure that it is initialized
185/// before reading from it in any way. Additionally, the caller must ensure that the memory
186/// allocated is freed using `cuda_free_locked`, or the memory will be leaked.
187///
188/// # Examples
189///
190/// ```
191/// # let _context = rustacuda::quick_init().unwrap();
192/// use rustacuda::memory::*;
193/// unsafe {
194///     // Allocate space for 5 u64s
195///     let locked_buffer = cuda_malloc_locked::<u64>(5).unwrap();
196///     cuda_free_locked(locked_buffer).unwrap();
197/// }
198/// ```
199pub unsafe fn cuda_malloc_locked<T>(count: usize) -> CudaResult<*mut T> {
200    let size = count.checked_mul(mem::size_of::<T>()).unwrap_or(0);
201    if size == 0 {
202        return Err(CudaError::InvalidMemoryAllocation);
203    }
204
205    let mut ptr: *mut c_void = ptr::null_mut();
206    cuda_driver_sys::cuMemAllocHost_v2(&mut ptr as *mut *mut c_void, size).to_result()?;
207    let ptr = ptr as *mut T;
208    Ok(ptr as *mut T)
209}
210
211/// Free page-locked memory allocated with [`cuda_malloc_host`](fn.cuda_malloc_host.html).
212///
213/// # Errors
214///
215/// If freeing memory fails, returns the CUDA error value. If the given pointer is null, returns
216/// InvalidValue.
217///
218/// # Safety
219///
220/// The given pointer must have been allocated with `cuda_malloc_locked`, or null.
221/// The caller is responsible for ensuring that no other pointers to the deallocated buffer exist.
222///
223/// # Examples
224///
225/// ```
226/// # let _context = rustacuda::quick_init().unwrap();
227/// use rustacuda::memory::*;
228/// unsafe {
229///     let locked_buffer = cuda_malloc_locked::<u64>(5).unwrap();
230///     // Free allocated memory
231///     cuda_free_locked(locked_buffer).unwrap();
232/// }
233/// ```
234pub unsafe fn cuda_free_locked<T>(ptr: *mut T) -> CudaResult<()> {
235    if ptr.is_null() {
236        return Err(CudaError::InvalidMemoryAllocation);
237    }
238
239    cuda_driver_sys::cuMemFreeHost(ptr as *mut c_void).to_result()?;
240    Ok(())
241}
242
243#[cfg(test)]
244mod test {
245    use super::*;
246
247    #[derive(Clone, Debug)]
248    struct ZeroSizedType;
249    unsafe impl DeviceCopy for ZeroSizedType {}
250
251    #[test]
252    fn test_cuda_malloc() {
253        let _context = crate::quick_init().unwrap();
254        unsafe {
255            let device_mem = cuda_malloc::<u64>(1).unwrap();
256            assert!(!device_mem.is_null());
257            cuda_free(device_mem).unwrap();
258        }
259    }
260
261    #[test]
262    fn test_cuda_malloc_zero_bytes() {
263        let _context = crate::quick_init().unwrap();
264        unsafe {
265            assert_eq!(
266                CudaError::InvalidMemoryAllocation,
267                cuda_malloc::<u64>(0).unwrap_err()
268            );
269        }
270    }
271
272    #[test]
273    fn test_cuda_malloc_zero_sized() {
274        let _context = crate::quick_init().unwrap();
275        unsafe {
276            assert_eq!(
277                CudaError::InvalidMemoryAllocation,
278                cuda_malloc::<ZeroSizedType>(10).unwrap_err()
279            );
280        }
281    }
282
283    #[test]
284    fn test_cuda_alloc_overflow() {
285        let _context = crate::quick_init().unwrap();
286        unsafe {
287            assert_eq!(
288                CudaError::InvalidMemoryAllocation,
289                cuda_malloc::<u64>(::std::usize::MAX - 1).unwrap_err()
290            );
291        }
292    }
293
294    #[test]
295    fn test_cuda_malloc_unified() {
296        let _context = crate::quick_init().unwrap();
297        unsafe {
298            let mut unified = cuda_malloc_unified::<u64>(1).unwrap();
299            assert!(!unified.is_null());
300
301            // Write to the allocated memory
302            *unified.as_raw_mut() = 64;
303
304            cuda_free_unified(unified).unwrap();
305        }
306    }
307
308    #[test]
309    fn test_cuda_malloc_unified_zero_bytes() {
310        let _context = crate::quick_init().unwrap();
311        unsafe {
312            assert_eq!(
313                CudaError::InvalidMemoryAllocation,
314                cuda_malloc_unified::<u64>(0).unwrap_err()
315            );
316        }
317    }
318
319    #[test]
320    fn test_cuda_malloc_unified_zero_sized() {
321        let _context = crate::quick_init().unwrap();
322        unsafe {
323            assert_eq!(
324                CudaError::InvalidMemoryAllocation,
325                cuda_malloc_unified::<ZeroSizedType>(10).unwrap_err()
326            );
327        }
328    }
329
330    #[test]
331    fn test_cuda_malloc_unified_overflow() {
332        let _context = crate::quick_init().unwrap();
333        unsafe {
334            assert_eq!(
335                CudaError::InvalidMemoryAllocation,
336                cuda_malloc_unified::<u64>(::std::usize::MAX - 1).unwrap_err()
337            );
338        }
339    }
340
341    #[test]
342    fn test_cuda_free_null() {
343        let _context = crate::quick_init().unwrap();
344        let null = ::std::ptr::null_mut::<u64>();
345        unsafe {
346            assert_eq!(
347                CudaError::InvalidMemoryAllocation,
348                cuda_free(DevicePointer::wrap(null)).unwrap_err()
349            );
350        }
351    }
352
353    #[test]
354    fn test_cuda_malloc_locked() {
355        let _context = crate::quick_init().unwrap();
356        unsafe {
357            let locked = cuda_malloc_locked::<u64>(1).unwrap();
358            assert!(!locked.is_null());
359
360            // Write to the allocated memory
361            *locked = 64;
362
363            cuda_free_locked(locked).unwrap();
364        }
365    }
366
367    #[test]
368    fn test_cuda_malloc_locked_zero_bytes() {
369        let _context = crate::quick_init().unwrap();
370        unsafe {
371            assert_eq!(
372                CudaError::InvalidMemoryAllocation,
373                cuda_malloc_locked::<u64>(0).unwrap_err()
374            );
375        }
376    }
377
378    #[test]
379    fn test_cuda_malloc_locked_zero_sized() {
380        let _context = crate::quick_init().unwrap();
381        unsafe {
382            assert_eq!(
383                CudaError::InvalidMemoryAllocation,
384                cuda_malloc_locked::<ZeroSizedType>(10).unwrap_err()
385            );
386        }
387    }
388
389    #[test]
390    fn test_cuda_malloc_locked_overflow() {
391        let _context = crate::quick_init().unwrap();
392        unsafe {
393            assert_eq!(
394                CudaError::InvalidMemoryAllocation,
395                cuda_malloc_locked::<u64>(::std::usize::MAX - 1).unwrap_err()
396            );
397        }
398    }
399
400    #[test]
401    fn test_cuda_free_locked_null() {
402        let _context = crate::quick_init().unwrap();
403        unsafe {
404            assert_eq!(
405                CudaError::InvalidMemoryAllocation,
406                cuda_free_locked(::std::ptr::null_mut::<u64>()).unwrap_err()
407            );
408        }
409    }
410}