|
| 1 | +use std::{ |
| 2 | + ptr::NonNull, |
| 3 | + sync::{Arc, Mutex}, |
| 4 | +}; |
| 5 | + |
| 6 | +use super::*; |
| 7 | + |
| 8 | +#[derive(Clone)] |
| 9 | +pub(crate) struct UnsafeNonNullPtr(pub(crate) Arc<NonNull<[u8]>>); |
| 10 | +unsafe impl Send for UnsafeNonNullPtr {} |
| 11 | +unsafe impl Sync for UnsafeNonNullPtr {} |
| 12 | + |
| 13 | +impl UnsafeNonNullPtr { |
| 14 | + pub(crate) fn new(ptr: NonNull<[u8]>) -> Self { |
| 15 | + Self(Arc::new(ptr)) |
| 16 | + } |
| 17 | + |
| 18 | + pub(crate) fn as_ptr(&self) -> *const u8 { |
| 19 | + self.0.as_ptr().cast() |
| 20 | + } |
| 21 | + pub(crate) fn as_mut_ptr(&mut self) -> *mut u8 { |
| 22 | + self.0.as_ptr().cast() |
| 23 | + } |
| 24 | +} |
| 25 | + |
| 26 | +#[derive(Clone)] |
| 27 | +pub(crate) struct StaticBitmapAllocator { |
| 28 | + pub(crate) memory: UnsafeNonNullPtr, |
| 29 | + pub(crate) memory_size: usize, |
| 30 | + pub(crate) block_size_in_bytes: usize, |
| 31 | + pub(crate) bitmap: Arc<Mutex<Vec<bool>>>, |
| 32 | +} |
| 33 | + |
| 34 | +impl StaticBitmapAllocator { |
| 35 | + pub(crate) fn init( |
| 36 | + memory: NonNull<[u8]>, |
| 37 | + num_blocks: usize, |
| 38 | + block_size_in_bytes: usize, |
| 39 | + ) -> Self { |
| 40 | + let memory_size_in_bytes = num_blocks * block_size_in_bytes; |
| 41 | + Self { |
| 42 | + memory: UnsafeNonNullPtr::new(memory), |
| 43 | + memory_size: memory_size_in_bytes, |
| 44 | + block_size_in_bytes, |
| 45 | + bitmap: Arc::new(Mutex::new(vec![false; num_blocks])), |
| 46 | + } |
| 47 | + } |
| 48 | + |
| 49 | + pub(crate) fn as_ptr(&self) -> *const u8 { |
| 50 | + self.memory.as_ptr().cast() |
| 51 | + } |
| 52 | + |
| 53 | + pub(crate) fn find_free_block(&self) -> Option<usize> { |
| 54 | + for (idx, entry) in self.bitmap.lock().unwrap().iter_mut().enumerate() { |
| 55 | + if !*entry { |
| 56 | + *entry = true; |
| 57 | + return Some(idx); |
| 58 | + } |
| 59 | + } |
| 60 | + None |
| 61 | + } |
| 62 | + |
| 63 | + #[allow(unreachable_code)] |
| 64 | + pub(crate) fn find_adjacent_free_blocks( |
| 65 | + &self, |
| 66 | + requested_num_blocks: usize, |
| 67 | + ) -> Option<std::ops::Range<usize>> { |
| 68 | + let mut bitmap = self.bitmap.lock().unwrap(); |
| 69 | + if requested_num_blocks > bitmap.len() { |
| 70 | + return None; |
| 71 | + } |
| 72 | + let _range_of_blocks_found = false; |
| 73 | + let _found_range = 0..0; |
| 74 | + |
| 75 | + let mut start = 0; |
| 76 | + let mut end = requested_num_blocks; |
| 77 | + let mut busy_block_idx = 0; |
| 78 | + loop { |
| 79 | + let mut has_busy_block = false; |
| 80 | + for (idx, sub_entry) in bitmap[start..end].iter().copied().enumerate() { |
| 81 | + if sub_entry { |
| 82 | + has_busy_block = true; |
| 83 | + busy_block_idx = start + idx; |
| 84 | + } |
| 85 | + } |
| 86 | + if !has_busy_block { |
| 87 | + for entry in bitmap[start..end].iter_mut() { |
| 88 | + *entry = true; |
| 89 | + } |
| 90 | + return Some(start..end); |
| 91 | + } else { |
| 92 | + start = busy_block_idx + 1; |
| 93 | + end = start + requested_num_blocks; |
| 94 | + if end > bitmap.len() { |
| 95 | + break; |
| 96 | + } |
| 97 | + } |
| 98 | + } |
| 99 | + // panic!("not found block {} {} {}", start, end, self.bitmap.len()); |
| 100 | + None |
| 101 | + } |
| 102 | + |
| 103 | + pub(crate) fn free_blocks(&self, index: usize, num_blocks: usize) { |
| 104 | + assert!(num_blocks > 0); |
| 105 | + let mut guard = self.bitmap.lock().unwrap(); |
| 106 | + for i in index..index + num_blocks { |
| 107 | + guard[i] = false; |
| 108 | + } |
| 109 | + } |
| 110 | + |
| 111 | + pub(crate) fn allocate( |
| 112 | + &self, |
| 113 | + layout: std::alloc::Layout, |
| 114 | + ) -> CudaResult<std::ptr::NonNull<[u8]>> { |
| 115 | + let size = layout.size(); |
| 116 | + assert!(size > 0); |
| 117 | + assert_eq!(size % self.block_size_in_bytes, 0); |
| 118 | + let num_blocks = size / self.block_size_in_bytes; |
| 119 | + |
| 120 | + if size > self.block_size_in_bytes { |
| 121 | + if let Some(range) = self.find_adjacent_free_blocks(num_blocks) { |
| 122 | + let index = range.start; |
| 123 | + let offset = index * self.block_size_in_bytes; |
| 124 | + let ptr = unsafe { self.as_ptr().add(offset) }; |
| 125 | + let ptr = unsafe { NonNull::new_unchecked(ptr as _) }; |
| 126 | + return Ok(NonNull::slice_from_raw_parts(ptr, size)); |
| 127 | + } |
| 128 | + panic!("allocation of {} blocks has failed", num_blocks); |
| 129 | + // return Err(CudaError::AllocationError(format!( |
| 130 | + // "allocation of {} blocks has failed", |
| 131 | + // num_blocks |
| 132 | + // ))); |
| 133 | + } |
| 134 | + |
| 135 | + if let Some(index) = self.find_free_block() { |
| 136 | + let offset = index * self.block_size_in_bytes; |
| 137 | + let ptr = unsafe { self.as_ptr().add(offset) }; |
| 138 | + let ptr = unsafe { NonNull::new_unchecked(ptr as _) }; |
| 139 | + Ok(NonNull::slice_from_raw_parts(ptr, size)) |
| 140 | + } else { |
| 141 | + panic!("allocation of 1 block has failed"); |
| 142 | + // return Err(CudaError::AllocationError(format!( |
| 143 | + // "allocation of 1 block has failed", |
| 144 | + // ))); |
| 145 | + } |
| 146 | + } |
| 147 | + |
| 148 | + pub(crate) fn deallocate(&self, ptr: std::ptr::NonNull<u8>, layout: std::alloc::Layout) { |
| 149 | + let size = layout.size(); |
| 150 | + assert!(size > 0); |
| 151 | + assert_eq!(size % self.block_size_in_bytes, 0); |
| 152 | + let offset = unsafe { ptr.as_ptr().offset_from(self.as_ptr()) } as usize; |
| 153 | + if offset >= self.memory_size { |
| 154 | + return; |
| 155 | + } |
| 156 | + assert_eq!(offset % self.block_size_in_bytes, 0); |
| 157 | + let index = offset / self.block_size_in_bytes; |
| 158 | + let num_blocks = size / self.block_size_in_bytes; |
| 159 | + self.free_blocks(index, num_blocks); |
| 160 | + } |
| 161 | +} |
0 commit comments