use crate::{
lamellae::{
comm::atomic::atomic_type_supported,
comm::{AtomicOp, CommInfo, CommProgress, CommShutdown},
AllocError, AllocResult, CollectiveOpKind, CommAlloc, CommAllocAddr, CommAllocInner,
},
Backend,
};
use parking_lot::Mutex;
use std::collections::HashMap;
use std::sync::{
atomic::{AtomicUsize, Ordering},
Arc,
};
#[derive(Debug)]
pub(crate) struct LocalAlloc {
pub(crate) ptr: *mut u8,
pub(crate) layout: std::alloc::Layout,
}
unsafe impl Send for LocalAlloc {}
unsafe impl Sync for LocalAlloc {}
impl LocalAlloc {
pub(crate) fn start(&self) -> usize {
self.ptr as usize
}
pub(crate) unsafe fn as_mut_ptr<T>(&self) -> *mut T {
self.ptr as *mut T
}
pub(crate) unsafe fn as_mut_slice<T>(&self) -> &mut [T] {
unsafe {
std::slice::from_raw_parts_mut(
self.as_mut_ptr::<T>(),
self.num_bytes() / std::mem::size_of::<T>(),
)
}
}
pub(crate) fn num_bytes(&self) -> usize {
self.layout.size()
}
pub(crate) fn sub_alloc(&self, offset: usize, len: usize) -> AllocResult<Arc<LocalAlloc>> {
if offset + len > self.layout.size() {
return Err(AllocError::InvalidSubAlloc(offset, len));
}
let new_data = unsafe { self.ptr.add(offset) };
Ok(Arc::new(LocalAlloc {
ptr: new_data,
layout: std::alloc::Layout::from_size_align(len, self.layout.align()).unwrap(),
}))
}
pub(crate) fn leak(self: Arc<Self>) -> Option<CommAllocAddr> {
None
}
pub(crate) fn wait(&self) {
}
}
impl From<Arc<LocalAlloc>> for CommAlloc {
fn from(alloc: Arc<LocalAlloc>) -> Self {
CommAlloc {
inner_alloc: Arc::new(CommAllocInner::LocalAlloc(alloc)),
}
}
}
#[derive(Debug)]
pub(crate) struct LocalComm {
pub(crate) num_pes: usize,
pub(crate) my_pe: usize,
pub(crate) allocs: Arc<Mutex<HashMap<usize, Arc<LocalAlloc>>>>,
pub(crate) heap_allocs: Arc<Mutex<HashMap<usize, Arc<LocalAlloc>>>>,
pub(crate) put_amt: Arc<AtomicUsize>,
pub(crate) get_amt: Arc<AtomicUsize>,
}
impl LocalComm {
pub(crate) fn new() -> LocalComm {
LocalComm {
num_pes: 1,
my_pe: 0,
allocs: Arc::new(Mutex::new(HashMap::new())),
heap_allocs: Arc::new(Mutex::new(HashMap::new())),
put_amt: Arc::new(AtomicUsize::new(0)),
get_amt: Arc::new(AtomicUsize::new(0)),
}
}
}
impl CommShutdown for LocalComm {
fn force_shutdown(&self) {}
}
impl CommProgress for LocalComm {
fn flush_all(&self) {}
fn wait_all(&self) {}
fn barrier(&self) {}
}
impl CommInfo for LocalComm {
fn my_pe(&self) -> usize {
self.my_pe
}
fn num_pes(&self) -> usize {
self.num_pes
}
fn backend(&self) -> Backend {
Backend::Local
}
fn atomic_avail<T: 'static>(&self) -> bool {
atomic_type_supported::<T>()
}
fn atomic_op_avail<T: 'static>(&self, _op: AtomicOp<T>) -> bool {
atomic_type_supported::<T>()
}
fn MB_sent(&self) -> f64 {
(self.put_amt.load(Ordering::SeqCst) + self.get_amt.load(Ordering::SeqCst)) as f64
/ 1_000_000.0
}
fn collective_avail<T: 'static>(&self, _op: CollectiveOpKind) -> bool {
true }
}