use lamellar::array::prelude::*;
#[lamellar::main]
fn main() {
let args: Vec<String> = std::env::args().collect();
let elem_per_pe = args
.get(1)
.and_then(|s| s.parse::<usize>().ok())
.unwrap_or_else(|| 2);
let world = lamellar::LamellarWorldBuilder::new().build();
let my_pe = world.my_pe();
let num_pes = world.num_pes();
let dim = elem_per_pe * num_pes;
let m = dim; let n = dim; let p = dim;
let a = LocalLockArray::<f32>::new(&world, m * n, Distribution::Block).block(); let b = LocalLockArray::<f32>::new(&world, n * p, Distribution::Block).block(); let c = AtomicArray::<f32>::new(&world, m * p, Distribution::Block).block();
a.dist_iter_mut()
.enumerate()
.for_each(|(i, x)| *x = i as f32)
.block();
b.dist_iter_mut()
.enumerate()
.for_each(move |(i, x)| {
let row = i / dim;
let col = i % dim;
if row == col {
*x = 1 as f32
} else {
*x = 0 as f32;
}
})
.block();
c.dist_iter_mut().for_each(|x| x.store(0.0)).block();
world.barrier();
let a = a.into_read_only().block();
let b = b.into_read_only().block();
let num_gops = ((2 * dim * dim * dim) - dim * dim) as f64 / 1_000_000_000.0;
if my_pe == 0 {
println!("starting");
}
let start = std::time::Instant::now();
if my_pe == 0 {
let a_c = a.clone();
let b_c = b.clone();
let c_c = c.clone();
world.block_on(async move {
for i in 0..m {
for j in 0..p {
let mut sum = 0.0;
for k in 0..n {
let a_val = a_c.get(k + i * m).spawn();
let b_val = b_c.get(j + k * n).spawn();
sum += a_val.await * b_val.await;
}
let _ = c_c.store(j + i * m, sum).spawn(); }
}
});
}
world.wait_all();
world.barrier();
let elapsed = start.elapsed().as_secs_f64();
c.wait_all();
c.barrier();
if my_pe == 0 {
println!("elapsed {:?} Gflops: {:?}", elapsed, num_gops / elapsed,);
}
}