use { crate::{align_down, align_up, error::MapError}, alloc::sync::Arc, core::{ convert::TryFrom as _, ptr::{copy_nonoverlapping, NonNull}, // sync::atomic::{AtomicU8, Ordering::*}, }, gpu_alloc_types::{MappedMemoryRange, MemoryDevice, MemoryPropertyFlags}, }; #[derive(Debug)] struct Relevant; impl Drop for Relevant { fn drop(&mut self) { report_error_on_drop!("Memory block wasn't deallocated"); } } /// Memory block allocated by `GpuAllocator`. #[derive(Debug)] pub struct MemoryBlock { memory_type: u32, props: MemoryPropertyFlags, offset: u64, size: u64, atom_mask: u64, mapped: bool, flavor: MemoryBlockFlavor, relevant: Relevant, } impl MemoryBlock { pub(crate) fn new( memory_type: u32, props: MemoryPropertyFlags, offset: u64, size: u64, atom_mask: u64, flavor: MemoryBlockFlavor, ) -> Self { isize::try_from(atom_mask).expect("`atom_mask` is too large"); MemoryBlock { memory_type, props, offset, size, atom_mask, flavor, mapped: false, relevant: Relevant, } } pub(crate) fn deallocate(self) -> MemoryBlockFlavor { core::mem::forget(self.relevant); self.flavor } } unsafe impl Sync for MemoryBlock where M: Sync {} unsafe impl Send for MemoryBlock where M: Send {} #[derive(Debug)] pub(crate) enum MemoryBlockFlavor { Dedicated { memory: M, }, Buddy { chunk: usize, index: usize, ptr: Option>, memory: Arc, }, FreeList { chunk: u64, ptr: Option>, memory: Arc, }, } impl MemoryBlock { /// Returns reference to parent memory object. #[inline(always)] pub fn memory(&self) -> &M { match &self.flavor { MemoryBlockFlavor::Dedicated { memory } => memory, MemoryBlockFlavor::Buddy { memory, .. } => memory, MemoryBlockFlavor::FreeList { memory, .. } => memory, } } /// Returns offset in bytes from start of memory object to start of this block. #[inline(always)] pub fn offset(&self) -> u64 { self.offset } /// Returns size of this memory block. #[inline(always)] pub fn size(&self) -> u64 { self.size } /// Returns memory property flags for parent memory object. #[inline(always)] pub fn props(&self) -> MemoryPropertyFlags { self.props } /// Returns index of type of parent memory object. #[inline(always)] pub fn memory_type(&self) -> u32 { self.memory_type } /// Returns pointer to mapped memory range of this block. /// This blocks becomes mapped. /// /// The user of returned pointer must guarantee that any previously submitted command that writes to this range has completed /// before the host reads from or writes to that range, /// and that any previously submitted command that reads from that range has completed /// before the host writes to that region. /// If the device memory was allocated without the `HOST_COHERENT` property flag set, /// these guarantees must be made for an extended range: /// the user must round down the start of the range to the nearest multiple of `non_coherent_atom_size`, /// and round the end of the range up to the nearest multiple of `non_coherent_atom_size`. /// /// # Panics /// /// This function panics if block is currently mapped. /// /// # Safety /// /// `block` must have been allocated from specified `device`. #[inline(always)] pub unsafe fn map( &mut self, device: &impl MemoryDevice, offset: u64, size: usize, ) -> Result, MapError> { let size_u64 = u64::try_from(size).expect("`size` doesn't fit device address space"); assert!(offset < self.size, "`offset` is out of memory block bounds"); assert!( size_u64 <= self.size - offset, "`offset + size` is out of memory block bounds" ); let ptr = match &mut self.flavor { MemoryBlockFlavor::Dedicated { memory } => { let end = align_up(offset + size_u64, self.atom_mask) .expect("mapping end doesn't fit device address space"); let aligned_offset = align_down(offset, self.atom_mask); if !acquire_mapping(&mut self.mapped) { return Err(MapError::AlreadyMapped); } let result = device.map_memory(memory, self.offset + aligned_offset, end - aligned_offset); match result { // the overflow is checked in `Self::new()` Ok(ptr) => { let ptr_offset = (offset - aligned_offset) as isize; ptr.as_ptr().offset(ptr_offset) } Err(err) => { release_mapping(&mut self.mapped); return Err(err.into()); } } } MemoryBlockFlavor::FreeList { ptr: Some(ptr), .. } | MemoryBlockFlavor::Buddy { ptr: Some(ptr), .. } => { if !acquire_mapping(&mut self.mapped) { return Err(MapError::AlreadyMapped); } let offset_isize = isize::try_from(offset) .expect("Buddy and linear block should fit host address space"); ptr.as_ptr().offset(offset_isize) } _ => return Err(MapError::NonHostVisible), }; Ok(NonNull::new_unchecked(ptr)) } /// Unmaps memory range of this block that was previously mapped with `Block::map`. /// This block becomes unmapped. /// /// # Panics /// /// This function panics if this block is not currently mapped. /// /// # Safety /// /// `block` must have been allocated from specified `device`. #[inline(always)] pub unsafe fn unmap(&mut self, device: &impl MemoryDevice) -> bool { if !release_mapping(&mut self.mapped) { return false; } match &mut self.flavor { MemoryBlockFlavor::Dedicated { memory } => { device.unmap_memory(memory); } MemoryBlockFlavor::Buddy { .. } => {} MemoryBlockFlavor::FreeList { .. } => {} } true } /// Transiently maps block memory range and copies specified data /// to the mapped memory range. /// /// # Panics /// /// This function panics if block is currently mapped. /// /// # Safety /// /// `block` must have been allocated from specified `device`. /// The caller must guarantee that any previously submitted command that reads or writes to this range has completed. #[inline(always)] pub unsafe fn write_bytes( &mut self, device: &impl MemoryDevice, offset: u64, data: &[u8], ) -> Result<(), MapError> { let size = data.len(); let ptr = self.map(device, offset, size)?; copy_nonoverlapping(data.as_ptr(), ptr.as_ptr(), size); let result = if !self.coherent() { let aligned_offset = align_down(offset, self.atom_mask); let end = align_up(offset + data.len() as u64, self.atom_mask).unwrap(); device.flush_memory_ranges(&[MappedMemoryRange { memory: self.memory(), offset: self.offset + aligned_offset, size: end - aligned_offset, }]) } else { Ok(()) }; self.unmap(device); result.map_err(Into::into) } /// Transiently maps block memory range and copies specified data /// from the mapped memory range. /// /// # Panics /// /// This function panics if block is currently mapped. /// /// # Safety /// /// `block` must have been allocated from specified `device`. /// The caller must guarantee that any previously submitted command that reads to this range has completed. #[inline(always)] pub unsafe fn read_bytes( &mut self, device: &impl MemoryDevice, offset: u64, data: &mut [u8], ) -> Result<(), MapError> { #[cfg(feature = "tracing")] { if !self.cached() { tracing::warn!("Reading from non-cached memory may be slow. Consider allocating HOST_CACHED memory block for host reads.") } } let size = data.len(); let ptr = self.map(device, offset, size)?; let result = if !self.coherent() { let aligned_offset = align_down(offset, self.atom_mask); let end = align_up(offset + data.len() as u64, self.atom_mask).unwrap(); device.invalidate_memory_ranges(&[MappedMemoryRange { memory: self.memory(), offset: self.offset + aligned_offset, size: end - aligned_offset, }]) } else { Ok(()) }; if result.is_ok() { copy_nonoverlapping(ptr.as_ptr(), data.as_mut_ptr(), size); } self.unmap(device); result.map_err(Into::into) } fn coherent(&self) -> bool { self.props.contains(MemoryPropertyFlags::HOST_COHERENT) } #[cfg(feature = "tracing")] fn cached(&self) -> bool { self.props.contains(MemoryPropertyFlags::HOST_CACHED) } } fn acquire_mapping(mapped: &mut bool) -> bool { if *mapped { false } else { *mapped = true; true } } fn release_mapping(mapped: &mut bool) -> bool { if *mapped { *mapped = false; true } else { false } }