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(|| 200);
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;
println!("m: {}, n: {}, p: {}", m, n, p);
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.wait_all();
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;
let rows_pe = m / num_pes;
let start = std::time::Instant::now();
b.onesided_iter() .chunks(p) .into_iter() .enumerate()
.for_each(|(j, col)| {
let col = col.clone();
let c = c.clone();
let _ = a
.local_chunks(n)
.enumerate()
.for_each(move |(i, row)| {
let sum = col.iter().zip(row).map(|(&i1, &i2)| i1 * i2).sum::<f32>(); c.mut_local_data().at(j + (i % rows_pe) * m).fetch_add(sum);
})
.spawn();
});
world.wait_all();
world.barrier();
let elapsed = start.elapsed().as_secs_f64();
let sum = c.sum().block();
println!("Elapsed: {:?}", elapsed);
if my_pe == 0 {
println!(
"elapsed {:?} Gflops: {:?} {:?}",
elapsed,
num_gops / elapsed,
sum
);
}
}