stream/
kernel.rs

1use std::fmt;
2
3use num_traits::{cast, Num, NumCast};
4#[cfg(feature = "rayon")]
5use rayon::prelude::*;
6use size::Size;
7
8pub trait Kernel: fmt::Display {
9    fn size(&self) -> Size;
10
11    fn init(&mut self);
12
13    fn run(&mut self);
14}
15
16/// Filling kernel
17///
18/// ```
19/// a(i) = x
20/// ```
21///
22/// Bytes/iter read = sizeof(T)
23///
24/// Bytes/iter write = sizeof(T)
25///
26/// FLOPS/iter = 0
27pub struct KernelFill<T> {
28    a: Vec<T>,
29    size: usize,
30    offset: usize,
31}
32
33impl<T> KernelFill<T> {
34    pub fn new(size: usize, offset: usize) -> Self {
35        Self {
36            a: Vec::new(),
37            size, offset,
38        }
39    }
40}
41
42impl<T> Kernel for KernelFill<T>
43where
44    T: Num + NumCast + Copy + Send
45{
46    fn size(&self) -> Size {
47        Size::from_bytes(size_of::<T>() * self.size)
48    }
49
50    fn init(&mut self) {
51        self.a = vec![cast(0).unwrap(); self.size + self.offset];
52    }
53
54    fn run(&mut self) {
55        #[cfg(not(feature = "rayon"))]
56        self.a.iter_mut().for_each(|a| *a = cast(3).unwrap());
57        #[cfg(feature = "rayon")]
58        self.a.par_iter_mut().for_each(|a| *a = cast(3).unwrap());
59    }
60}
61
62impl<T> fmt::Display for KernelFill<T> {
63    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
64        write!(f, "Fill {}", Size::from_bytes(size_of::<T>() * self.size))
65    }
66}
67
68/// Copying kernel
69///
70/// ```
71/// a(i) = b(i)
72/// ```
73///
74/// Bytes/iter read = 2 * sizeof(T)
75///
76/// Bytes/iter write = sizeof(T)
77///
78/// FLOPS/iter = 1
79pub struct KernelCopy<T> {
80    a: Vec<T>,
81    b: Vec<T>,
82    size: usize,
83    offset: usize,
84}
85
86impl<T> KernelCopy<T> {
87    pub fn new(size: usize, offset: usize) -> Self {
88        Self {
89            a: Vec::new(),
90            b: Vec::new(),
91            size, offset,
92        }
93    }
94}
95
96impl<T> Kernel for KernelCopy<T>
97where
98    T: Num + NumCast + Copy + Send + Sync
99{
100    fn size(&self) -> Size {
101        Size::from_bytes(2 * size_of::<T>() * self.size)
102    }
103
104    fn init(&mut self) {
105        self.a = vec![cast(0).unwrap(); self.size + self.offset];
106        self.b = vec![cast(1).unwrap(); self.size + self.offset];
107    }
108
109    fn run(&mut self) {
110        #[cfg(not(feature = "rayon"))]
111        self.a.iter_mut().zip(self.b.iter())
112            .for_each(|(a, &b)| *a = b);
113        #[cfg(feature = "rayon")]
114        self.a.par_iter_mut().zip(self.b.par_iter())
115            .for_each(|(a, &b)| *a = b);
116    }
117}
118
119impl<T> fmt::Display for KernelCopy<T> {
120    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
121        write!(f, "Copy {}", Size::from_bytes(2 * size_of::<T>() * self.size))
122    }
123}
124
125/// Scaling kernel
126///
127/// ```
128/// a(i) = b(i) * x
129/// ```
130///
131/// Bytes/iter read = 2 * sizeof(T)
132///
133/// Bytes/iter write = sizeof(T)
134///
135/// FLOPS/iter = 1
136pub struct KernelScale<T> {
137    a: Vec<T>,
138    b: Vec<T>,
139    size: usize,
140    offset: usize,
141}
142
143impl<T> KernelScale<T> {
144    pub fn new(size: usize, offset: usize) -> Self {
145        Self {
146            a: Vec::new(),
147            b: Vec::new(),
148            size, offset,
149        }
150    }
151}
152
153impl<T> Kernel for KernelScale<T>
154where
155    T: Num + NumCast + Copy + Send + Sync
156{
157    fn size(&self) -> Size {
158        Size::from_bytes(2 * size_of::<T>() * self.size)
159    }
160
161    fn init(&mut self) {
162        self.a = vec![cast(0).unwrap(); self.size + self.offset];
163        self.b = vec![cast(1).unwrap(); self.size + self.offset];
164    }
165
166    fn run(&mut self) {
167        #[cfg(not(feature = "rayon"))]
168        self.a.iter_mut().zip(self.b.iter())
169            .for_each(|(a, &b)| *a = b * cast(3).unwrap());
170        #[cfg(feature = "rayon")]
171        self.a.par_iter_mut().zip(self.b.par_iter())
172            .for_each(|(a, &b)| *a = b * cast(3).unwrap());
173    }
174}
175
176impl<T> fmt::Display for KernelScale<T> {
177    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
178        write!(f, "Scale {}", Size::from_bytes(2 * size_of::<T>() * self.size))
179    }
180}
181
182/// Adding kernel
183///
184/// ```
185/// a(i) = b(i) + c(i)
186/// ```
187///
188/// Bytes/iter read = 3 * sizeof(T)
189///
190/// Bytes/iter write = sizeof(T)
191///
192/// FLOPS/iter = 1
193pub struct KernelAdd<T> {
194    a: Vec<T>,
195    b: Vec<T>,
196    c: Vec<T>,
197    size: usize,
198    offset: usize,
199}
200
201impl<T> KernelAdd<T> {
202    pub fn new(size: usize, offset: usize) -> Self {
203        Self {
204            a: Vec::new(),
205            b: Vec::new(),
206            c: Vec::new(),
207            size, offset,
208        }
209    }
210}
211
212impl<T> Kernel for KernelAdd<T>
213where
214    T: Num + NumCast + Copy + Send + Sync
215{
216    fn size(&self) -> Size {
217        Size::from_bytes(3 * size_of::<T>() * self.size)
218    }
219
220    fn init(&mut self) {
221        self.a = vec![cast(0).unwrap(); self.size + self.offset];
222        self.b = vec![cast(1).unwrap(); self.size + self.offset];
223        self.c = vec![cast(2).unwrap(); self.size + self.offset];
224    }
225
226    fn run(&mut self) {
227        #[cfg(not(feature = "rayon"))]
228        self.a.iter_mut().zip(self.b.iter().zip(self.c.iter()))
229            .for_each(|(a, (&b, &c))| *a = b + c);
230        #[cfg(feature = "rayon")]
231        self.a.par_iter_mut().zip(self.b.par_iter().zip(self.c.par_iter()))
232            .for_each(|(a, (&b, &c))| *a = b + c);
233    }
234}
235
236impl<T> fmt::Display for KernelAdd<T> {
237    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
238        write!(f, "Add {}", Size::from_bytes(3 * size_of::<T>() * self.size))
239    }
240}
241
242/// Triad kernel
243///
244/// ```
245/// a(i) = b(i) + c(i) * x
246/// ```
247///
248/// Bytes/iter read = 3 * sizeof(T)
249///
250/// Bytes/iter write = sizeof(T)
251///
252/// FLOPS/iter = 2
253pub struct KernelTriad<T> {
254    a: Vec<T>,
255    b: Vec<T>,
256    c: Vec<T>,
257    size: usize,
258    offset: usize,
259}
260
261impl<T> KernelTriad<T> {
262    pub fn new(size: usize, offset: usize) -> Self {
263        Self {
264            a: Vec::new(),
265            b: Vec::new(),
266            c: Vec::new(),
267            size, offset,
268        }
269    }
270}
271
272impl<T> Kernel for KernelTriad<T>
273where
274    T: Num + NumCast + Copy + Send + Sync
275{
276    fn size(&self) -> Size {
277        Size::from_bytes(3 * size_of::<T>() * self.size)
278    }
279
280    fn init(&mut self) {
281        self.a = vec![cast(0).unwrap(); self.size + self.offset];
282        self.b = vec![cast(1).unwrap(); self.size + self.offset];
283        self.c = vec![cast(2).unwrap(); self.size + self.offset];
284    }
285
286    fn run(&mut self) {
287        #[cfg(not(feature = "rayon"))]
288        self.a.iter_mut().zip(self.b.iter().zip(self.c.iter()))
289            .for_each(|(a, (&b, &c))| *a = b + c * cast(3).unwrap());
290        #[cfg(feature = "rayon")]
291        self.a.par_iter_mut().zip(self.b.par_iter().zip(self.c.par_iter()))
292            .for_each(|(a, (&b, &c))| *a = b + c * cast(3).unwrap());
293    }
294}
295
296impl<T> fmt::Display for KernelTriad<T> {
297    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
298        write!(f, "Triad {}", Size::from_bytes(3 * size_of::<T>() * self.size))
299    }
300}
301
302/// Summing kernel
303///
304/// ```
305/// x += a(i)
306/// ```
307///
308/// Bytes/iter read = sizeof(T)
309///
310/// Bytes/iter write = 0
311///
312/// FLOPS/iter = 1
313pub struct KernelSum<T> {
314    a: Vec<T>,
315    size: usize,
316    offset: usize,
317}
318
319impl<T> KernelSum<T> {
320    pub fn new(size: usize, offset: usize) -> Self {
321        Self {
322            a: Vec::new(),
323            size, offset,
324        }
325    }
326}
327
328impl<T> Kernel for KernelSum<T>
329where
330    T: Num + NumCast + Copy + Send + Sync,
331{
332    fn size(&self) -> Size {
333        Size::from_bytes(size_of::<T>() * self.size)
334    }
335
336    fn init(&mut self) {
337        self.a = vec![cast(1).unwrap(); self.size + self.offset];
338    }
339
340    fn run(&mut self) {
341        #[cfg(not(feature = "rayon"))]
342        self.a.iter().fold(T::zero(), |acc, &x| acc + x);
343        #[cfg(feature = "rayon")]
344        self.a.par_iter().cloned().reduce(|| T::zero(), |x, y| x + y);
345    }
346}
347
348impl<T> fmt::Display for KernelSum<T> {
349    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
350        write!(f, "Sum {}", Size::from_bytes(size_of::<T>() * self.size))
351    }
352}