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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169

//! An interface for dealing with the kinds of parallel computations involved.
use std::env;

use crossbeam_channel::{bounded, Receiver, SendError};
use log::trace;
use once_cell::sync::Lazy;
use yastl::Pool;

/// The number of threads the thread pool should use.
///
/// By default it's equal to the number of CPUs, but it can be changed with the
/// `EC_GPU_NUM_THREADS` environment variable.
static NUM_THREADS: Lazy<usize> = Lazy::new(read_num_threads);

/// The thread pool that is used for the computations.
///
/// By default, it's size is equal to the number of CPUs. It can be set to a different value with
/// the `EC_GPU_NUM_THREADS` environment variable.
pub static THREAD_POOL: Lazy<Pool> = Lazy::new(|| Pool::new(*NUM_THREADS));

/// Returns the number of threads.
///
/// The number can be set with the `EC_GPU_NUM_THREADS` environment variable. If it isn't set, it
/// defaults to the number of CPUs the system has.
fn read_num_threads() -> usize {
    env::var("EC_GPU_NUM_THREADS")
        .ok()
        .and_then(|num| num.parse::<usize>().ok())
        .unwrap_or_else(num_cpus::get)
}

/// A worker operates on a pool of threads.
#[derive(Clone, Default)]
pub struct Worker {}

impl Worker {
    /// Returns a new worker.
    pub fn new() -> Worker {
        Worker {}
    }

    /// Returns binary logarithm (floored) of the number of threads.
    ///
    /// This means, the number of threads is `2^log_num_threads()`.
    pub fn log_num_threads(&self) -> u32 {
        log2_floor(*NUM_THREADS)
    }

    /// Executes a function in a thread and returns a [`Waiter`] immediately.
    pub fn compute<F, R>(&self, f: F) -> Waiter<R>
    where
        F: FnOnce() -> R + Send + 'static,
        R: Send + 'static,
    {
        let (sender, receiver) = bounded(1);

        THREAD_POOL.spawn(move || {
            let res = f();
            // Best effort. We run it in a separate thread, so the receiver might not exist
            // anymore, but that's OK. It only means that we are not interested in the result.
            // A message is logged though, as concurrency issues are hard to debug and this might
            // help in such cases.
            if let Err(SendError(_)) = sender.send(res) {
                trace!("Cannot send result");
            }
        });

        Waiter { receiver }
    }

    /// Executes a function and returns the result once it is finished.
    ///
    /// The function gets the [`yastl::Scope`] as well as the `chunk_size` as parameters. THe
    /// `chunk_size` is number of elements per thread.
    pub fn scope<'a, F, R>(&self, elements: usize, f: F) -> R
    where
        F: FnOnce(&yastl::Scope<'a>, usize) -> R,
    {
        let chunk_size = if elements < *NUM_THREADS {
            1
        } else {
            elements / *NUM_THREADS
        };

        THREAD_POOL.scoped(|scope| f(scope, chunk_size))
    }

    /// Executes the passed in function, and returns the result once it is finished.
    pub fn scoped<'a, F, R>(&self, f: F) -> R
    where
        F: FnOnce(&yastl::Scope<'a>) -> R,
    {
        let (sender, receiver) = bounded(1);
        THREAD_POOL.scoped(|s| {
            let res = f(s);
            sender.send(res).unwrap();
        });

        receiver.recv().unwrap()
    }
}

/// A future that is waiting for a result.
pub struct Waiter<T> {
    receiver: Receiver<T>,
}

impl<T> Waiter<T> {
    /// Wait for the result.
    pub fn wait(&self) -> T {
        self.receiver.recv().unwrap()
    }

    /// One off sending.
    pub fn done(val: T) -> Self {
        let (sender, receiver) = bounded(1);
        sender.send(val).unwrap();

        Waiter { receiver }
    }
}

fn log2_floor(num: usize) -> u32 {
    assert!(num > 0);

    let mut pow = 0;

    while (1 << (pow + 1)) <= num {
        pow += 1;
    }

    pow
}

#[cfg(test)]
pub mod tests {
    use super::*;

    #[test]
    fn test_log2_floor() {
        assert_eq!(log2_floor(1), 0);
        assert_eq!(log2_floor(3), 1);
        assert_eq!(log2_floor(4), 2);
        assert_eq!(log2_floor(5), 2);
        assert_eq!(log2_floor(6), 2);
        assert_eq!(log2_floor(7), 2);
        assert_eq!(log2_floor(8), 3);
    }

    #[test]
    fn test_read_num_threads() {
        let num_cpus = num_cpus::get();
        temp_env::with_var("EC_GPU_NUM_THREADS", None::<&str>, || {
            assert_eq!(
                read_num_threads(),
                num_cpus,
                "By default the number of threads matches the number of CPUs."
            );
        });

        temp_env::with_var("EC_GPU_NUM_THREADS", Some("1234"), || {
            assert_eq!(
                read_num_threads(),
                1234,
                "Number of threads matches the environment variable."
            );
        });
    }
}