kryst 3.2.1

Krylov subspace and preconditioned iterative solvers for dense and sparse linear systems, with shared and distributed memory parallelism.
1
2
3
4
5
6
7
8
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
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118

use crate::algebra::prelude::*;

#[cfg(feature = "mpi")]
use super::MpiComm;
#[cfg(feature = "mpi")]
use mpi::traits::CommunicatorCollectives;

#[cfg(feature = "complex")]
#[inline]
pub fn pack_scalar_s(value: S) -> [R; 2] {
    [value.real(), value.imag()]
}

#[cfg(not(feature = "complex"))]
#[inline]
pub fn pack_scalar_s(value: S) -> [R; 1] {
    [value.real()]
}

#[cfg(feature = "complex")]
#[inline]
pub fn unpack_scalar_s(parts: [R; 2]) -> S {
    S::from_parts(parts[0], parts[1])
}

#[cfg(not(feature = "complex"))]
#[inline]
pub fn unpack_scalar_s(parts: [R; 1]) -> S {
    S::from_real(parts[0])
}

#[cfg(feature = "mpi")]
#[inline]
pub fn reduce_sum_real_rank_ordered(comm: &MpiComm, local: R) -> R {
    if comm.size <= 1 {
        return local;
    }

    let mut gathered = vec![0.0f64; comm.size];
    comm.world.all_gather_into(&local, gathered.as_mut_slice());
    let mut acc = 0.0f64;
    for value in gathered {
        acc += value;
    }
    acc
}

#[cfg(feature = "mpi")]
#[inline]
pub fn reduce_sum_scalar_rank_ordered(comm: &MpiComm, local: S) -> S {
    if comm.size <= 1 {
        return local;
    }

    #[cfg(feature = "complex")]
    const WIDTH: usize = 2;
    #[cfg(not(feature = "complex"))]
    const WIDTH: usize = 1;

    let packed = pack_scalar_s(local);
    let mut gathered = vec![0.0f64; WIDTH * comm.size];
    comm.world
        .all_gather_into(&packed[..], gathered.as_mut_slice());

    let mut acc = [0.0f64; WIDTH];
    for rank in 0..comm.size {
        for lane in 0..WIDTH {
            acc[lane] += gathered[WIDTH * rank + lane];
        }
    }

    #[cfg(feature = "complex")]
    {
        unpack_scalar_s([acc[0], acc[1]])
    }

    #[cfg(not(feature = "complex"))]
    {
        unpack_scalar_s([acc[0]])
    }
}

#[cfg(feature = "mpi")]
#[inline]
pub fn reduce_sum_scalars_rank_ordered(comm: &MpiComm, locals: &mut [S]) {
    if locals.is_empty() || comm.size <= 1 {
        return;
    }

    #[cfg(feature = "complex")]
    const WIDTH: usize = 2;
    #[cfg(not(feature = "complex"))]
    const WIDTH: usize = 1;

    let n = locals.len();
    let mut send = vec![0.0f64; WIDTH * n];
    for (idx, value) in locals.iter().enumerate() {
        let packed = pack_scalar_s(*value);
        for lane in 0..WIDTH {
            send[WIDTH * idx + lane] = packed[lane];
        }
    }

    let mut gathered = vec![0.0f64; WIDTH * n * comm.size];
    comm.world
        .all_gather_into(&send[..], gathered.as_mut_slice());

    for elem in 0..n {
        let mut acc = [0.0f64; WIDTH];
        for rank in 0..comm.size {
            for lane in 0..WIDTH {
                let offset = WIDTH * n * rank + WIDTH * elem + lane;
                acc[lane] += gathered[offset];
            }
        }
        locals[elem] = unpack_scalar_s(acc);
    }
}