use std::collections::HashMap;
use std::ops::Range;
use gfx_hal::adapter::{MemoryProperties, MemoryTypeId, PhysicalDevice};
use gfx_hal::memory::{Properties, Requirements};
use gfx_hal::{Backend, Device};
use bitvec::{bitvec, BitVec};
use slab::Slab;
use super::{pick_memory_type, Allocation, AllocationError, AllocationSpec, Allocator};
#[derive(Debug)]
pub struct Config {
pub chunk_size: u64,
pub min_block_size: u64,
}
#[derive(Debug)]
struct MemoryBlock<B: Backend> {
memory: B::Memory,
chunks: BitVec,
used_chunks: u64,
max_allocation_size: u64,
}
#[derive(Debug)]
struct MemoryPool<B: Backend> {
blocks: Slab<MemoryBlock<B>>,
flags: Properties,
id: MemoryTypeId,
block_map: HashMap<*const B::Memory, usize>,
}
impl<B: Backend> MemoryPool<B> {
fn new(id: MemoryTypeId, flags: Properties) -> MemoryPool<B> {
MemoryPool {
blocks: Slab::new(),
flags,
id,
block_map: Default::default(),
}
}
#[inline]
fn new_block(
&mut self,
device: &B::Device,
num_chunks: u64,
chunk_size: u64,
) -> Result<usize, AllocationError> {
let memory = unsafe { device.allocate_memory(self.id, num_chunks << chunk_size)? };
let block = MemoryBlock {
memory,
chunks: bitvec![0; num_chunks as usize],
used_chunks: 0,
max_allocation_size: num_chunks,
};
let id = self.blocks.insert(block);
self.block_map.insert(&self.blocks[id].memory as _, id);
Ok(id)
}
fn alloc(
&mut self,
device: &B::Device,
config: &Config,
num_chunks: u64,
) -> Result<ChunkAllocation<B>, AllocationError> {
let chunk_size = config.chunk_size;
let block = self
.blocks
.iter_mut()
.map(|(_, b)| b)
.find(|b| num_chunks <= b.max_allocation_size);
let block = if let Some(v) = block {
v
} else {
let block_size = std::cmp::max(num_chunks, config.min_block_size);
let block_id = self.new_block(device, block_size, chunk_size)?;
log::debug!(
"Creating new block #{} of size {} chunks",
block_id,
block_size
);
let block = &mut self.blocks[block_id];
block.max_allocation_size -= num_chunks;
block.used_chunks += num_chunks;
for i in 0..num_chunks as usize {
block.chunks.set(i, true);
}
let memory = &block.memory as _;
log::trace!("range: {:?}", 0..(num_chunks << chunk_size));
return Ok(ChunkAllocation {
memory,
range: 0..(num_chunks << chunk_size),
flags: self.flags,
pool_id: self.id,
relevant: relevant::Relevant,
});
};
block.used_chunks += num_chunks;
let mut count = 0;
let mut chunk_range: Range<usize> = 0..0;
for (i, chunk) in block.chunks.iter().enumerate() {
count = if chunk { 0 } else { count + 1 };
if count == num_chunks as usize {
chunk_range.start = i + 1 - count;
chunk_range.end = i + 1;
break;
}
}
for i in chunk_range.clone() {
block.chunks.set(i, true);
}
let mut max_count = 0;
let mut count = 0;
for chunk in block.chunks.iter() {
if chunk {
max_count = std::cmp::max(max_count, count);
count = 0;
} else {
count += 1;
}
}
max_count = std::cmp::max(max_count, count);
block.max_allocation_size = max_count;
let memory = &block.memory as _;
let range =
((chunk_range.start as u64) << chunk_size)..((chunk_range.end as u64) << chunk_size);
log::trace!("range: {:?}", range);
Ok(ChunkAllocation {
memory,
range,
flags: self.flags,
pool_id: self.id,
relevant: relevant::Relevant,
})
}
fn free(&mut self, device: &B::Device, config: &Config, allocation: ChunkAllocation<B>) {
let block_id = self.block_map[&allocation.memory];
let block = &mut self.blocks[block_id];
let num_chunks = (allocation.range.end - allocation.range.start) >> config.chunk_size;
let first_chunk = allocation.range.start >> config.chunk_size;
block.used_chunks -= num_chunks;
block.max_allocation_size = std::cmp::max(block.max_allocation_size, num_chunks);
for i in first_chunk..first_chunk + num_chunks {
block.chunks.set(i as _, false);
}
if block.used_chunks == 0 {
log::debug!("Freeing block #{}", block_id);
unsafe { device.free_memory(self.blocks.remove(block_id).memory) };
}
allocation.relevant.dispose();
}
}
#[derive(derivative::Derivative)]
#[derivative(Debug)]
pub struct ChunkAllocation<B: Backend> {
#[derivative(Debug = "ignore")]
memory: *const B::Memory,
range: Range<u64>,
flags: Properties,
pool_id: MemoryTypeId,
#[derivative(Debug = "ignore")]
relevant: relevant::Relevant,
}
impl<B: Backend> Allocation<B> for ChunkAllocation<B> {
fn memory(&self) -> &B::Memory {
unsafe { &*self.memory }
}
fn flags(&self) -> Properties {
self.flags
}
fn range(&self) -> Range<u64> {
self.range.clone()
}
}
#[derive(Debug)]
pub struct ChunkAllocator<B: Backend> {
pools: Vec<MemoryPool<B>>,
memory_props: MemoryProperties,
config: Config,
}
impl<B: Backend> ChunkAllocator<B> {
pub fn new(device: &B::PhysicalDevice) -> ChunkAllocator<B> {
let memory_props = device.memory_properties();
let pools: Vec<_> = memory_props
.memory_types
.iter()
.enumerate()
.map(|(id, ty)| MemoryPool::new(MemoryTypeId(id), ty.properties))
.collect();
log::debug!(
"Created a chunk allocator with {} pools 1K chunk size and 1M min block size",
pools.len()
);
ChunkAllocator {
pools,
memory_props,
config: Config {
chunk_size: 10,
min_block_size: 1024,
},
}
}
}
impl<B: Backend> Allocator<B> for ChunkAllocator<B> {
type Allocation = ChunkAllocation<B>;
fn alloc(
&mut self,
device: &B::Device,
req: Requirements,
spec: AllocationSpec,
) -> Result<Self::Allocation, AllocationError> {
debug_assert!(req.alignment <= 1 << self.config.chunk_size);
let pool_id = pick_memory_type(&self.memory_props, req, spec)?;
let pool = &mut self.pools[pool_id.0];
log::trace!(
"Allocating {} bytes in pool #{} ({:?})",
req.size,
pool_id.0,
pool.flags
);
let num_chunks = req.size
>> self.config.chunk_size
+ if req.size % (1 << self.config.chunk_size) == 0 {
0
} else {
1
};
pool.alloc(device, &self.config, num_chunks)
}
fn free(&mut self, device: &B::Device, allocation: Self::Allocation) {
let pool = &mut self.pools[allocation.pool_id.0];
log::trace!(
"Freeing {} bytes in pool #{} ({:?})",
allocation.range.end - allocation.range.start,
allocation.pool_id.0,
pool.flags
);
pool.free(device, &self.config, allocation);
}
}