rlst 0.8.0

A Rust native linear algebra library.
Documentation 
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
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140

//! An implementation of distributed sample sort.

use std::cmp::Ordering;

use itertools::Itertools;
use mpi::traits::{CommunicatorCollectives, Equivalence};
use paste::paste;
use rand::Rng;

use crate::{
    RlstResult, TotalCmp,
    distributed_tools::{
        all_to_allv,
        array_tools::{gather_to_all, global_min, global_size},
        sort_to_bins,
    },
};

/// Maximum number of weighted stratified representatives per nonempty rank.
pub const OVERSAMPLING: usize = 4;

/// A helper trait for parallel sorts to convert items into unique items.
pub trait AsUnique: TotalCmp + Clone + Copy {
    /// The unique output type.
    type Unique: Default + Copy + Clone + TotalCmp + Equivalence;
    /// Convert `self` into the unique type.
    fn into_unique(value: Self, rank: usize, index: usize) -> Self::Unique;
    /// Convert a unique item back into the original type.
    fn from_unique(elem: Self::Unique) -> Self;
}

/// Globally sort `arr` while retaining every input record and stable tie order.
///
/// Empty local inputs and empty output buckets are valid. The output balance is
/// approximate; no rank is guaranteed to receive an element.
pub fn parsort<Item: AsUnique, C: CommunicatorCollectives, R: Rng + ?Sized>(
    arr: &[Item],
    comm: &C,
    rng: &mut R,
) -> RlstResult<Vec<Item>> {
    let rank = comm.rank() as usize;
    let mut local = arr.to_vec();
    if comm.size() == 1 {
        local.sort_by(|left, right| left.total_cmp(*right));
        return Ok(local);
    }
    if global_size(arr, comm) == 0 {
        return Ok(Vec::new());
    }
    let mut unique = local
        .iter()
        .enumerate()
        .map(|(index, item)| Item::into_unique(*item, rank, index))
        .collect_vec();
    unique.sort_by(|left, right| left.total_cmp(*right));
    let bins = get_bins::<Item, _, _>(&unique, comm, rng);
    let counts = sort_to_bins(&unique, &bins);
    let mut received = all_to_allv(comm, &counts, &unique).1;
    received.sort_by(|left, right| left.total_cmp(*right));
    Ok(received.into_iter().map(Item::from_unique).collect())
}

/// Create exactly one lower bound per rank from weighted local representatives.
fn get_bins<Item, C, R>(arr: &[Item::Unique], comm: &C, _rng: &mut R) -> Vec<Item::Unique>
where
    Item: AsUnique,
    C: CommunicatorCollectives,
    R: Rng + ?Sized,
{
    let size = comm.size() as usize;
    let samples = arr.len().min(OVERSAMPLING);
    let mut local_values = Vec::with_capacity(samples);
    let mut local_weights = Vec::with_capacity(samples);
    for sample in 0..samples {
        let start = sample * arr.len() / samples;
        let end = (sample + 1) * arr.len() / samples;
        local_values.push(arr[start]);
        local_weights.push(end - start);
    }
    let values = gather_to_all(&local_values, comm);
    let weights = gather_to_all(&local_weights, comm);
    let mut weighted = values.into_iter().zip(weights).collect_vec();
    weighted.sort_by(|(left, _), (right, _)| left.total_cmp(*right));
    let minimum = match global_min(arr, comm) {
        Some(value) => value,
        None => unreachable!("global input was checked nonempty"),
    };
    let total_weight: usize = weighted.iter().map(|(_, weight)| weight).sum();
    debug_assert!(total_weight > 0);
    let mut bounds = Vec::with_capacity(size);
    bounds.push(minimum);
    for bucket in 1..size {
        let target = bucket * total_weight / size;
        let mut accumulated = 0;
        let mut bound = weighted.last().expect("nonempty global samples").0;
        for &(value, weight) in &weighted {
            accumulated += weight;
            if accumulated >= target {
                bound = value;
                break;
            }
        }
        bounds.push(bound);
    }
    bounds
}

macro_rules! impl_as_unique_for {
    ($dtype:ty) => { paste! {
        #[derive(Copy, Clone, PartialOrd, PartialEq, Equivalence, Default)]
        #[allow(non_camel_case_types, missing_docs)]
        pub struct [<UniqueItem_ $dtype>] { pub value: $dtype, pub rank: usize, pub index: usize }
        impl TotalCmp for [<UniqueItem_$dtype>] {
            fn total_cmp(self, other: Self) -> Ordering {
                match self.value.total_cmp(other.value) {
                    Ordering::Equal => match self.rank.cmp(&other.rank) {
                        Ordering::Equal => self.index.cmp(&other.index), order => order,
                    }, order => order,
                }
            }
        }
        impl AsUnique for $dtype {
            type Unique = [<UniqueItem_$dtype>];
            fn into_unique(value: $dtype, rank: usize, index: usize) -> Self::Unique { Self::Unique { value, rank, index } }
            fn from_unique(elem: Self::Unique) -> Self { elem.value }
        }
    }};
}
impl_as_unique_for!(f64);
impl_as_unique_for!(f32);
impl_as_unique_for!(i8);
impl_as_unique_for!(i16);
impl_as_unique_for!(i32);
impl_as_unique_for!(i64);
impl_as_unique_for!(isize);
impl_as_unique_for!(u8);
impl_as_unique_for!(u16);
impl_as_unique_for!(u32);
impl_as_unique_for!(u64);
impl_as_unique_for!(usize);