Skip to main content

reduce/
reduce.rs

1//! Reduction examples: reduce-to-root and all-reduce with built-in operations.
2
3use mpi::collective::SystemOperation;
4use mpi::traits::*;
5
6fn main() {
7    let universe = mpi::initialize().unwrap();
8    let world = universe.world();
9    let rank = world.rank();
10    let size = world.size();
11    let root_rank = 0;
12    let root = world.process_at_rank(root_rank);
13
14    // Sum of all ranks lands on the root.
15    if rank == root_rank {
16        let mut sum = 0i32;
17        root.reduce_into_root(&rank, &mut sum, SystemOperation::sum());
18        let expected = (0..size).sum::<i32>();
19        println!("Sum of ranks = {sum} (expected {expected})");
20        assert_eq!(sum, expected);
21    } else {
22        root.reduce_into(&rank, SystemOperation::sum());
23    }
24
25    // Every rank learns the maximum rank via all-reduce.
26    let mut max = 0i32;
27    world.all_reduce_into(&rank, &mut max, SystemOperation::max());
28    assert_eq!(max, size - 1);
29
30    // Element-wise all-reduce over a vector.
31    let local = [rank as f64, (rank * rank) as f64];
32    let mut total = vec![0.0f64; 2];
33    world.all_reduce_into(&local[..], &mut total[..], SystemOperation::sum());
34    let expected0: f64 = (0..size).map(|r| r as f64).sum();
35    let expected1: f64 = (0..size).map(|r| (r * r) as f64).sum();
36    assert!((total[0] - expected0).abs() < 1e-9);
37    assert!((total[1] - expected1).abs() < 1e-9);
38
39    world.barrier();
40    if rank == root_rank {
41        println!("Reductions verified. vector all-reduce = {total:?}");
42    }
43}