From c6f8cc37292e93bc5ae22d5c6ef338e6acca8dd7 Mon Sep 17 00:00:00 2001 From: Anders Smedegaard Pedersen Date: Tue, 26 Nov 2024 17:59:23 +0100 Subject: [PATCH] closes 49. hipextmallocwithflags (#107) * memory methods on struct. Add malloc_with_flags() --- Cargo.lock | 1 + Cargo.toml | 1 + src/runtime/memory.rs | 20 +------ src/types/flags.rs | 11 ++++ src/types/memory.rs | 118 ++++++++++++++++++++++++++++++++++-------- src/types/mod.rs | 2 + 6 files changed, 111 insertions(+), 42 deletions(-) create mode 100644 src/types/flags.rs diff --git a/Cargo.lock b/Cargo.lock index c4f26aa..a776bb9 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -121,6 +121,7 @@ name = "hip-rs" version = "0.1.0" dependencies = [ "bindgen", + "bitflags", "cc", "env_logger", "log", diff --git a/Cargo.toml b/Cargo.toml index 01e206c..d3c391b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -18,6 +18,7 @@ semver = "1.0.23" uuid = "1.11.0" log = "0.4" env_logger = "0.10" +bitflags = "2.6.0" [build-dependencies] # For build script diff --git a/src/runtime/memory.rs b/src/runtime/memory.rs index bca60fd..d8645be 100644 --- a/src/runtime/memory.rs +++ b/src/runtime/memory.rs @@ -1,21 +1,3 @@ use super::sys; use crate::types::{Device, MemoryPointer, Result}; - -/// Allocates memory on a HIP device/accelerator. -/// -/// This function allocates a block of `size` bytes of device memory and returns a -/// MemoryPointer that safely manages the memory allocation. The memory will be -/// automatically freed when the MemoryPointer is dropped. -/// -/// If 0 is passed for `size`, `Ok(std::ptr::null_mut)` is returned. -/// -/// # Arguments -/// * `size` - Size of memory allocation in bytes -/// -/// # Returns -/// * `Ok(MemoryPointer)` - Handle to allocated device memory -/// * `Err(HipError)` - Error occurred during allocation -/// ``` -pub fn malloc(size: usize) -> Result> { - MemoryPointer::new(size) -} +use crate::DeviceMallocFlag; diff --git a/src/types/flags.rs b/src/types/flags.rs new file mode 100644 index 0000000..e5dfea2 --- /dev/null +++ b/src/types/flags.rs @@ -0,0 +1,11 @@ +use bitflags::bitflags; + +bitflags::bitflags! { + pub struct DeviceMallocFlag: u32 { + const DEFAULT = 0x0; + const FINEGRAINED = 0x1; + const SIGNAL_MEMORY = 0x2; + const UNCACHED = 0x3; + const CONTIGUOUS = 0x4; + } +} diff --git a/src/types/memory.rs b/src/types/memory.rs index b43f44c..97a7201 100644 --- a/src/types/memory.rs +++ b/src/types/memory.rs @@ -1,45 +1,90 @@ +use super::flags::DeviceMallocFlag; use super::{HipError, HipResult, Result}; use crate::sys; /// A wrapper for device memory allocated on the GPU. /// Automatically frees the memory when dropped. pub struct MemoryPointer { - ptr: *mut T, + pointer: *mut T, size: usize, } impl MemoryPointer { - pub fn new(size: usize) -> Result { + /// Private function that holds common logic for the + /// memory allocation functions. + /// + /// Takes the size to allocate and + fn allocate_with_fn(size: usize, alloc_fn: F) -> Result + where + F: FnOnce(*mut *mut std::ffi::c_void, usize) -> u32, + { // Handle zero size allocation according to spec if size == 0 { return Ok(MemoryPointer { - ptr: std::ptr::null_mut(), + pointer: std::ptr::null_mut(), size: 0, }); } let mut ptr = std::ptr::null_mut(); - - let code = unsafe { - sys::hipMalloc( - &mut ptr as *mut *mut T as *mut *mut std::ffi::c_void, - size * std::mem::size_of::(), - ) - }; + let code = alloc_fn( + &mut ptr as *mut *mut T as *mut *mut std::ffi::c_void, + size * std::mem::size_of::(), + ); let pointer = Self { - ptr: ptr as *mut T, + pointer: ptr as *mut T, size, }; (pointer, code).to_result() } + /// Allocates memory on a HIP device/accelerator. + /// + /// This function allocates a block of `size` bytes of device memory and returns a + /// MemoryPointer that safely manages the memory allocation. The memory will be + /// automatically freed when the MemoryPointer is dropped. + /// + /// If 0 is passed for `size`, `Ok(std::ptr::null_mut)` is returned. + /// + /// # Arguments + /// * `size` - Size of memory allocation in bytes + /// + /// # Returns + /// * `Ok(MemoryPointer)` - Handle to allocated device memory + /// * `Err(HipError)` - Error occurred during allocation + /// ``` + pub fn alloc(size: usize) -> Result { + Self::allocate_with_fn(size, |ptr, size| unsafe { sys::hipMalloc(ptr, size) }) + } + + /// Allocates memory on the default accelerator with specified allocation flags. + /// + /// # Arguments + /// * `size` - The requested memory size in bytes + /// * `flag` - The memory allocation flag. Must be one of: DeviceMallocDefault, + /// DeviceMallocFinegrained, DeviceMallocUncached, or MallocSignalMemory + /// + /// # Returns + /// * `Ok(MemoryPointer)` - Successfully allocated memory pointer + /// * `Err(_)` - If allocation fails due to out of memory or invalid flags + /// + /// # Notes + /// * If size is 0, returns null pointer with success status + /// * Invalid flags will result in hipErrorInvalidValue error + /// + pub fn alloc_with_flag(size: usize, flag: DeviceMallocFlag) -> Result { + Self::allocate_with_fn(size, |ptr, size| unsafe { + sys::hipExtMallocWithFlags(ptr, size, flag.bits()) + }) + } + /// Returns the raw memory pointer. /// /// Note: This pointer cannot be directly dereferenced from CPU code. pub fn as_ptr(&self) -> *mut T { - self.ptr + self.pointer } /// Returns the size in bytes of the allocated memory @@ -52,7 +97,7 @@ impl MemoryPointer { impl Drop for MemoryPointer { fn drop(&mut self) { unsafe { - let code = sys::hipFree(self.ptr as *mut std::ffi::c_void); + let code = sys::hipFree(self.pointer as *mut std::ffi::c_void); if code != 0 { let error = HipError::new(code); log::error!("MemoryPointer failed to free memory: {}", error); @@ -69,27 +114,27 @@ mod tests { #[test] fn test_new_zero_size() { - let result = MemoryPointer::::new(0).unwrap(); - assert!(result.ptr.is_null()); + let result = MemoryPointer::::alloc(0).unwrap(); + assert!(result.pointer.is_null()); assert_eq!(result.size, 0); } #[test] fn test_new_valid_size() { let size = 1024; - let result = MemoryPointer::::new(size).unwrap(); - assert!(!result.ptr.is_null()); + let result = MemoryPointer::::alloc(size).unwrap(); + assert!(!result.pointer.is_null()); assert_eq!(result.size, size); } #[test] fn test_new_different_types() { // Test with different sized types - let result = MemoryPointer::::new(100).unwrap(); - assert!(!result.ptr.is_null()); + let result = MemoryPointer::::alloc(100).unwrap(); + assert!(!result.pointer.is_null()); - let result = MemoryPointer::::new(100).unwrap(); - assert!(!result.ptr.is_null()); + let result = MemoryPointer::::alloc(100).unwrap(); + assert!(!result.pointer.is_null()); } #[test] @@ -97,9 +142,36 @@ mod tests { let mb = 1024 * 1024; let size = 3000 * mb; println!("Attempting to allocate {} bytes", size); - let result = MemoryPointer::::new(size); + let result = MemoryPointer::::alloc(size); sleep(Duration::from_secs(5)); - assert!(!result.unwrap().ptr.is_null()); + assert!(!result.unwrap().pointer.is_null()); + } + + #[test] + fn test_alloc_with_flag_success() { + let size = 1024; + let result = MemoryPointer::::alloc_with_flag(size, DeviceMallocFlag::DEFAULT); + assert!(result.is_ok()); + let ptr = result.unwrap(); + assert!(!ptr.pointer.is_null()); + } + + #[test] + fn test_alloc_with_flag_zero_size() { + let result = MemoryPointer::::alloc_with_flag(0, DeviceMallocFlag::DEFAULT); + assert!(result.is_ok()); + let ptr = result.unwrap(); + assert!(ptr.pointer.is_null()); + } + + #[test] + fn test_alloc_with_combined_flag() { + let size = 1024; + let flag = DeviceMallocFlag::DEFAULT | DeviceMallocFlag::FINEGRAINED; + let result = MemoryPointer::::alloc_with_flag(size, flag); + assert!(result.is_ok()); + let ptr = result.unwrap(); + assert!(!ptr.pointer.is_null()); } } diff --git a/src/types/mod.rs b/src/types/mod.rs index e1a7b82..d1d4adc 100644 --- a/src/types/mod.rs +++ b/src/types/mod.rs @@ -1,7 +1,9 @@ mod device; +mod flags; mod memory; mod result; pub use device::*; +pub use flags::*; pub use memory::*; pub use result::*;