all_gather/
all_gather.rs

1#![deny(warnings)]
2extern crate mpi_fork_fnsp as mpi;
3
4use mpi::datatype::{MutView, UserDatatype, View};
5use mpi::traits::*;
6use mpi::Count;
7
8fn main() {
9    let universe = mpi::initialize().unwrap();
10    let world = universe.world();
11    let root_rank = 0;
12
13    let count = world.size() as usize;
14    let i = 2_u64.pow(world.rank() as u32 + 1);
15    let mut a = vec![0u64; count];
16
17    world.all_gather_into(&i, &mut a[..]);
18
19    if world.rank() == root_rank {
20        println!("Root gathered sequence: {:?}.", a);
21    }
22    assert!(a
23        .iter()
24        .enumerate()
25        .all(|(a, &b)| b == 2u64.pow(a as u32 + 1)));
26
27    let factor = world.rank() as u64 + 1;
28    let a = (1_u64..)
29        .take(count)
30        .map(|x| x * factor)
31        .collect::<Vec<_>>();
32    let mut t = vec![0u64; count * count];
33
34    world.all_gather_into(&a[..], &mut t[..]);
35
36    if world.rank() == root_rank {
37        println!("Root gathered table:");
38        for r in t.chunks(count) {
39            println!("{:?}", r);
40        }
41    }
42    assert!((0_u64..)
43        .zip(t.iter())
44        .all(|(a, &b)| b == (a / count as u64 + 1) * (a % count as u64 + 1)));
45
46    let d = UserDatatype::contiguous(count as Count, &u64::equivalent_datatype());
47    t = vec![0u64; count * count];
48
49    {
50        let sv = unsafe { View::with_count_and_datatype(&a[..], 1, &d) };
51        let mut rv = unsafe { MutView::with_count_and_datatype(&mut t[..], count as Count, &d) };
52
53        world.all_gather_into(&sv, &mut rv);
54    }
55
56    if world.rank() == root_rank {
57        println!("Root gathered table:");
58        for r in t.chunks(count) {
59            println!("{:?}", r);
60        }
61    }
62    assert!((0_u64..)
63        .zip(t.iter())
64        .all(|(a, &b)| b == (a / count as u64 + 1) * (a % count as u64 + 1)));
65}