use lamellar::array::prelude::*;
use lamellar::memregion::prelude::*;
use tracing_subscriber::prelude::*;
const ARRAY_LEN: usize = 100;
#[lamellar::AmData(Clone)]
struct RdmaAM {
array: UnsafeArray<u8>,
orig_pe: usize,
index: usize,
}
#[lamellar::am]
impl LamellarAM for RdmaAM {
async fn exec(&self) {
let num_pes = lamellar::num_pes;
let max_i =
unsafe { std::cmp::min(self.array.local_as_slice().len(), self.index + num_pes) };
println!(
"\t({},{}) in RdmaAM on pe {:?}, originating from pe {:?} with index {} and max_i {}",
self.orig_pe,
self.index,
lamellar::current_pe,
self.orig_pe,
self.index,
max_i
);
println!("\tlocal segement of array: {:?}..", unsafe {
&self.array.local_as_slice()[self.index..max_i]
});
let local = unsafe { self.array.get_buffer(0, ARRAY_LEN).await };
let my_index = self.index * num_pes + lamellar::current_pe;
println!("\t({},{}) current view of remote segment on pe {:?}: {:?}..{:?}\n\tpe: {:?} updating index {:?} on pe {:?}", self.orig_pe,
self.index,self.orig_pe, &local[self.index..max_i], &local[local.len()-max_i..],lamellar::current_pe, my_index, self.orig_pe);
unsafe {
self.array.put(my_index, lamellar::current_pe as u8).await;
}
}
}
#[lamellar::main]
fn main() {
let _subscriber = tracing_subscriber::registry()
.with(tracing_subscriber::EnvFilter::from_default_env())
.with(
tracing_subscriber::fmt::layer()
.with_thread_ids(true)
.with_file(true)
.with_line_number(true)
.with_level(true),
)
.init();
let world = lamellar::LamellarWorldBuilder::new().build();
let my_pe = world.my_pe();
let num_pes = world.num_pes();
println!("creating array");
let array = UnsafeArray::<u8>::new(world.team(), ARRAY_LEN, Distribution::Block).block();
println!("creating memregion");
let local_mem_region = world.alloc_one_sided_mem_region::<u8>(ARRAY_LEN);
println!("about to initialize array");
array.print();
if my_pe == 0 {
unsafe {
for i in local_mem_region.as_mut_slice() {
*i = 255_u8;
}
}
unsafe { array.put_buffer(0, &local_mem_region).block() };
}
array.print();
for i in unsafe { array.local_as_slice() } {
while *i != 255_u8 {
std::thread::yield_now();
}
}
if my_pe == 0 {
println!("------------------------------------------------------------");
}
world.barrier();
println!("about to free mem region");
drop(local_mem_region);
println!("freed mem region");
println!("[{:?}] Before {:?}", my_pe, unsafe {
array.local_as_slice()
});
world.barrier();
if my_pe == 0 {
println!("------------------------------------------------------------");
}
world.barrier();
let index = 0;
if my_pe == 0 {
let _ = world.spawn_am_pe(
num_pes - 1,
RdmaAM {
array: array.clone(),
orig_pe: my_pe,
index: index,
},
);
}
world.wait_all();
world.barrier();
println!("[{:?}] after {:?}", my_pe, unsafe {
array.local_as_slice()
});
world.barrier();
array.print();
if my_pe == 0 {
println!("------------------------------------------------------------");
}
world.barrier();
if my_pe == 0 {
let sum = unsafe { array.sum().block() };
println!("sum: {:?}", sum);
println!("------------------------------------------------------------");
}
}