matrixmultiply/
util.rs

1// Copyright 2016 - 2018 Ulrik Sverdrup "bluss"
2//
3// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
4// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
5// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
6// option. This file may not be copied, modified, or distributed
7// except according to those terms.
8
9use core::cmp::min;
10
11#[derive(Copy, Clone)]
12pub struct RangeChunk { i: usize, n: usize, chunk: usize }
13
14/// Create an iterator that splits `n` in chunks of size `chunk`;
15/// the last item can be an uneven chunk.
16pub fn range_chunk(n: usize, chunk: usize) -> RangeChunk {
17    RangeChunk {
18        i: 0,
19        n: n,
20        chunk: chunk,
21    }
22}
23
24impl Iterator for RangeChunk {
25    type Item = (usize, usize);
26
27    #[inline]
28    fn next(&mut self) -> Option<Self::Item> {
29        if self.n == 0 {
30            None
31        } else {
32            let i = self.i;
33            let rem = min(self.n, self.chunk);
34            self.i += 1;
35            self.n -= rem;
36            Some((i, rem))
37        }
38    }
39}
40
41#[inline]
42pub fn round_up_to(x: usize, multiple_of: usize) -> usize {
43    let (mut d, r) = (x / multiple_of, x % multiple_of);
44    if r > 0 { d += 1; }
45    d * multiple_of
46}
47
48impl RangeChunk {
49    #[cfg(feature="threading")]
50    /// Split the iterator in `total` parts and only iterate the `index`th part of it.
51    /// The iterator must not have started when this is called.
52    pub(crate) fn part(self, index: usize, total: usize) -> Self {
53        debug_assert_eq!(self.i, 0, "range must be uniterated");
54        debug_assert_ne!(total, 0);
55        let (n, chunk) = (self.n, self.chunk);
56
57        // round up
58        let mut nchunks = n / chunk;
59        nchunks += (n % chunk != 0) as usize;
60
61        // chunks per thread
62        // round up
63        let mut chunks_per = nchunks / total;
64        chunks_per += (nchunks % total != 0) as usize;
65
66        let i = chunks_per * index;
67        let nn = min(n, (i + chunks_per) * chunk).saturating_sub(i * chunk);
68
69        RangeChunk { i, n: nn, chunk }
70    }
71}