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}