pub fn global_inclusive_cumsum<T: Equivalence + Zero + Copy, C: CommunicatorCollectives>(
arr: &[T],
comm: &C,
) -> Vec<T>
Expand description
Perform a global inclusive cumulative sum operation.
For the array [1, 3, 5, 7]
the output will be [1, 4, 9, 16]
.
Examples found in repository?
examples/mpi_cumsum.rs (line 32)
9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62
pub fn main() {
// Initialise MPI
let universe = mpi::initialize().unwrap();
// Get the world communicator
let comm = universe.world();
// Initialise a seeded Rng.
let mut rng = ChaCha8Rng::seed_from_u64(comm.rank() as u64);
// Create `npoints` per rank.
let nelems = 10;
// Generate random numbers
let mut elems = Vec::<usize>::with_capacity(nelems);
for _ in 0..nelems {
elems.push(rng.gen_range(0..100));
}
// Compute the cumulative sum.
let global_cum_sum = global_inclusive_cumsum(&elems, &comm);
// Copy array to root and compare with inclusive scan there.
if let (Some(cum_sum_root), Some(original_array)) = (
gather_to_root(&global_cum_sum, &comm),
gather_to_root(&elems, &comm),
) {
// Scan on root
let expected_cum_sum = original_array
.iter()
.scan(0, |state, x| {
*state += *x;
Some(*state)
})
.collect_vec();
// Check that the first element is not modified (inclusive cumsum)
assert_eq!(
original_array.first().unwrap(),
cum_sum_root.first().unwrap()
);
for (actual, expected) in izip!(cum_sum_root.iter(), expected_cum_sum.iter()) {
assert_eq!(*actual, *expected);
}
println!("Cumulative sum computed.");
}
}