Skip to main content

raphtory_api/
compute.rs

1use rayon::prelude::*;
2/// Compute cumulative sum in parallel over `num_chunks` chunks
3pub fn cum_sum(values: &mut [usize]) {
4    let mut sum = 0;
5    for v in values {
6        sum += *v;
7        *v = sum;
8    }
9}
10
11pub fn par_cum_sum(values: &mut [usize]) {
12    let num_chunks = rayon::current_num_threads();
13    let chunk_size = values.len().div_ceil(num_chunks);
14    let mut chunk_sums = Vec::with_capacity(num_chunks);
15    values
16        .par_chunks_mut(chunk_size)
17        .map(|chunk| {
18            let mut cum_sum = 0;
19            for v in chunk {
20                *v += cum_sum;
21                cum_sum = *v;
22            }
23            cum_sum
24        })
25        .collect_into_vec(&mut chunk_sums);
26
27    let mut cum_sum = 0;
28    for (partial_sum, next_chunk) in chunk_sums
29        .into_iter()
30        .zip(values.chunks_mut(chunk_size).skip(1))
31    {
32        cum_sum += partial_sum;
33        next_chunk.par_iter_mut().for_each(|v| *v += cum_sum);
34    }
35}
36
37#[cfg(test)]
38mod test {
39    use super::cum_sum;
40
41    #[test]
42    fn test_cum_sum() {
43        let mut values: Vec<_> = (0..100).collect();
44        cum_sum(&mut values);
45        let mut cum_sum = 0;
46        for (index, v) in values.into_iter().enumerate() {
47            cum_sum += index;
48            assert_eq!(v, cum_sum)
49        }
50    }
51}