use crate::numa::binding::NumaBinding;
use core::alloc::Layout;
pub trait NumaAllocator: Send + Sync {
unsafe fn alloc_on_node(&self, layout: Layout, node: u32) -> *mut u8;
unsafe fn dealloc_on_node(&self, ptr: *mut u8, layout: Layout, node: u32);
unsafe fn realloc_on_node(
&self,
ptr: *mut u8,
layout: Layout,
new_layout: Layout,
node: u32,
) -> *mut u8 {
let new_ptr = self.alloc_on_node(new_layout, node);
if !new_ptr.is_null() {
let copy_size = core::cmp::min(layout.size(), new_layout.size());
core::ptr::copy_nonoverlapping(ptr, new_ptr, copy_size);
self.dealloc_on_node(ptr, layout, node);
}
new_ptr
}
}
pub struct MnemosyneNumaAllocator;
impl NumaAllocator for MnemosyneNumaAllocator {
unsafe fn alloc_on_node(&self, layout: Layout, node: u32) -> *mut u8 {
#[cfg(feature = "mnemosyne-memory")]
{
let _binding = NumaBinding::bind(node);
unsafe { core::alloc::GlobalAlloc::alloc(&mnemosyne::Mnemosyne, layout) }
}
#[cfg(not(feature = "mnemosyne-memory"))]
{
let _binding = NumaBinding::bind(node);
alloc::alloc::alloc(layout)
}
}
unsafe fn dealloc_on_node(&self, ptr: *mut u8, layout: Layout, _node: u32) {
super::locality::bump_alloc_generation();
#[cfg(feature = "mnemosyne-memory")]
{
unsafe {
core::alloc::GlobalAlloc::dealloc(&mnemosyne::Mnemosyne, ptr, layout);
}
}
#[cfg(not(feature = "mnemosyne-memory"))]
{
let _ = _node;
alloc::alloc::dealloc(ptr, layout);
}
}
unsafe fn realloc_on_node(
&self,
ptr: *mut u8,
layout: Layout,
new_layout: Layout,
node: u32,
) -> *mut u8 {
#[cfg(feature = "mnemosyne-memory")]
{
let _binding = NumaBinding::bind(node);
let new_ptr = unsafe {
core::alloc::GlobalAlloc::realloc(
&mnemosyne::Mnemosyne,
ptr,
layout,
new_layout.size(),
)
};
if !new_ptr.is_null() && new_ptr != ptr {
super::locality::bump_alloc_generation();
}
new_ptr
}
#[cfg(not(feature = "mnemosyne-memory"))]
{
let new_ptr = self.alloc_on_node(new_layout, node);
if !new_ptr.is_null() {
let copy_size = core::cmp::min(layout.size(), new_layout.size());
core::ptr::copy_nonoverlapping(ptr, new_ptr, copy_size);
self.dealloc_on_node(ptr, layout, node);
}
new_ptr
}
}
}