rayonパッケージ
Revisión | 0aefd6178594f61ef340b2fbfaa16a9b43f05d6f (tree) |
---|---|
Tiempo | 2019-12-19 07:33:45 |
Autor | bors[bot] <26634292+bors[bot]@user...> |
Commiter | GitHub |
Merge #711
711: impl IntoParallelIterator for tuples => MultiZip r=nikomatsakis a=cuviper
This is implemented for tuples up to arity 12, much like the standard
library's trait implementations.
- For (a, b, ...), it calls into_par_iter() on each member.
- For &(a, b, ...), it calls par_iter() on each member.
- For &mut (a, b, ...), it calls par_iter_mut() on each member.
The resulting MultiZip iterator returns a tuple of the zipped items
from each input iterator. Internally, it is implemented with macros that
forward to a series of regular zips, mapping to a flattened tuple.
Closes #567.
Co-authored-by: Josh Stone <cuviper@gmail.com>
@@ -141,6 +141,8 @@ mod zip; | ||
141 | 141 | pub use self::zip::Zip; |
142 | 142 | mod zip_eq; |
143 | 143 | pub use self::zip_eq::ZipEq; |
144 | +mod multizip; | |
145 | +pub use self::multizip::MultiZip; | |
144 | 146 | mod interleave; |
145 | 147 | pub use self::interleave::Interleave; |
146 | 148 | mod interleave_shortest; |
@@ -0,0 +1,338 @@ | ||
1 | +use super::plumbing::*; | |
2 | +use super::*; | |
3 | + | |
4 | +/// `MultiZip` is an iterator that zips up a tuple of parallel iterators to | |
5 | +/// produce tuples of their items. | |
6 | +/// | |
7 | +/// It is created by calling `into_par_iter()` on a tuple of types that | |
8 | +/// implement `IntoParallelIterator`, or `par_iter()`/`par_iter_mut()` with | |
9 | +/// types that are iterable by reference. | |
10 | +/// | |
11 | +/// The implementation currently support tuples up to length 12. | |
12 | +/// | |
13 | +/// # Examples | |
14 | +/// | |
15 | +/// ``` | |
16 | +/// use rayon::prelude::*; | |
17 | +/// | |
18 | +/// // This will iterate `r` by mutable reference, like `par_iter_mut()`, while | |
19 | +/// // ranges are all iterated by value like `into_par_iter()`. | |
20 | +/// // Note that the zipped iterator is only as long as the shortest input. | |
21 | +/// let mut r = vec![0; 3]; | |
22 | +/// (&mut r, 1..10, 10..100, 100..1000).into_par_iter() | |
23 | +/// .for_each(|(r, x, y, z)| *r = x * y + z); | |
24 | +/// | |
25 | +/// assert_eq!(&r, &[1 * 10 + 100, 2 * 11 + 101, 3 * 12 + 102]); | |
26 | +/// ``` | |
27 | +/// | |
28 | +/// For a group that should all be iterated by reference, you can use a tuple reference. | |
29 | +/// | |
30 | +/// ``` | |
31 | +/// use rayon::prelude::*; | |
32 | +/// | |
33 | +/// let xs: Vec<_> = (1..10).collect(); | |
34 | +/// let ys: Vec<_> = (10..100).collect(); | |
35 | +/// let zs: Vec<_> = (100..1000).collect(); | |
36 | +/// | |
37 | +/// // Reference each input separately with `IntoParallelIterator`: | |
38 | +/// let r1: Vec<_> = (&xs, &ys, &zs).into_par_iter() | |
39 | +/// .map(|(x, y, z)| x * y + z) | |
40 | +/// .collect(); | |
41 | +/// | |
42 | +/// // Reference them all together with `IntoParallelRefIterator`: | |
43 | +/// let r2: Vec<_> = (xs, ys, zs).par_iter() | |
44 | +/// .map(|(x, y, z)| x * y + z) | |
45 | +/// .collect(); | |
46 | +/// | |
47 | +/// assert_eq!(r1, r2); | |
48 | +/// ``` | |
49 | +/// | |
50 | +/// Mutable references to a tuple will work similarly. | |
51 | +/// | |
52 | +/// ``` | |
53 | +/// use rayon::prelude::*; | |
54 | +/// | |
55 | +/// let mut xs: Vec<_> = (1..4).collect(); | |
56 | +/// let mut ys: Vec<_> = (-4..-1).collect(); | |
57 | +/// let mut zs = vec![0; 3]; | |
58 | +/// | |
59 | +/// // Mutably reference each input separately with `IntoParallelIterator`: | |
60 | +/// (&mut xs, &mut ys, &mut zs).into_par_iter().for_each(|(x, y, z)| { | |
61 | +/// *z += *x + *y; | |
62 | +/// std::mem::swap(x, y); | |
63 | +/// }); | |
64 | +/// | |
65 | +/// assert_eq!(xs, (vec![-4, -3, -2])); | |
66 | +/// assert_eq!(ys, (vec![1, 2, 3])); | |
67 | +/// assert_eq!(zs, (vec![-3, -1, 1])); | |
68 | +/// | |
69 | +/// // Mutably reference them all together with `IntoParallelRefMutIterator`: | |
70 | +/// let mut tuple = (xs, ys, zs); | |
71 | +/// tuple.par_iter_mut().for_each(|(x, y, z)| { | |
72 | +/// *z += *x + *y; | |
73 | +/// std::mem::swap(x, y); | |
74 | +/// }); | |
75 | +/// | |
76 | +/// assert_eq!(tuple, (vec![1, 2, 3], vec![-4, -3, -2], vec![-6, -2, 2])); | |
77 | +/// ``` | |
78 | +#[derive(Debug, Clone)] | |
79 | +pub struct MultiZip<T> { | |
80 | + tuple: T, | |
81 | +} | |
82 | + | |
83 | +// These macros greedily consume 4 or 2 items first to achieve log2 nesting depth. | |
84 | +// For example, 5 => 4,1 => (2,2),1. | |
85 | +// | |
86 | +// The tuples go up to 12, so we might want to greedily consume 8 too, but | |
87 | +// the depth works out the same if we let that expand on the right: | |
88 | +// 9 => 4,5 => (2,2),(4,1) => (2,2),((2,2),1) | |
89 | +// 12 => 4,8 => (2,2),(4,4) => (2,2),((2,2),(2,2)) | |
90 | +// | |
91 | +// But if we ever increase to 13, we would want to split 8,5 rather than 4,9. | |
92 | + | |
93 | +macro_rules! reduce { | |
94 | + ($a:expr, $b:expr, $c:expr, $d:expr, $( $x:expr ),+ => $fn:path) => { | |
95 | + reduce!(reduce!($a, $b, $c, $d => $fn), | |
96 | + reduce!($( $x ),+ => $fn) | |
97 | + => $fn) | |
98 | + }; | |
99 | + ($a:expr, $b:expr, $( $x:expr ),+ => $fn:path) => { | |
100 | + reduce!(reduce!($a, $b => $fn), | |
101 | + reduce!($( $x ),+ => $fn) | |
102 | + => $fn) | |
103 | + }; | |
104 | + ($a:expr, $b:expr => $fn:path) => { $fn($a, $b) }; | |
105 | + ($a:expr => $fn:path) => { $a }; | |
106 | +} | |
107 | + | |
108 | +macro_rules! nest { | |
109 | + ($A:tt, $B:tt, $C:tt, $D:tt, $( $X:tt ),+) => { | |
110 | + (nest!($A, $B, $C, $D), nest!($( $X ),+)) | |
111 | + }; | |
112 | + ($A:tt, $B:tt, $( $X:tt ),+) => { | |
113 | + (($A, $B), nest!($( $X ),+)) | |
114 | + }; | |
115 | + ($A:tt, $B:tt) => { ($A, $B) }; | |
116 | + ($A:tt) => { $A }; | |
117 | +} | |
118 | + | |
119 | +macro_rules! flatten { | |
120 | + ($( $T:ident ),+) => {{ | |
121 | + #[allow(non_snake_case)] | |
122 | + fn flatten<$( $T ),+>(nest!($( $T ),+) : nest!($( $T ),+)) -> ($( $T, )+) { | |
123 | + ($( $T, )+) | |
124 | + } | |
125 | + flatten | |
126 | + }}; | |
127 | +} | |
128 | + | |
129 | +macro_rules! multizip_impls { | |
130 | + ($( | |
131 | + $Tuple:ident { | |
132 | + $(($idx:tt) -> $T:ident)+ | |
133 | + } | |
134 | + )+) => { | |
135 | + $( | |
136 | + impl<$( $T, )+> IntoParallelIterator for ($( $T, )+) | |
137 | + where | |
138 | + $( | |
139 | + $T: IntoParallelIterator, | |
140 | + $T::Iter: IndexedParallelIterator, | |
141 | + )+ | |
142 | + { | |
143 | + type Item = ($( $T::Item, )+); | |
144 | + type Iter = MultiZip<($( $T::Iter, )+)>; | |
145 | + | |
146 | + fn into_par_iter(self) -> Self::Iter { | |
147 | + MultiZip { | |
148 | + tuple: ( $( self.$idx.into_par_iter(), )+ ), | |
149 | + } | |
150 | + } | |
151 | + } | |
152 | + | |
153 | + impl<'a, $( $T, )+> IntoParallelIterator for &'a ($( $T, )+) | |
154 | + where | |
155 | + $( | |
156 | + $T: IntoParallelRefIterator<'a>, | |
157 | + $T::Iter: IndexedParallelIterator, | |
158 | + )+ | |
159 | + { | |
160 | + type Item = ($( $T::Item, )+); | |
161 | + type Iter = MultiZip<($( $T::Iter, )+)>; | |
162 | + | |
163 | + fn into_par_iter(self) -> Self::Iter { | |
164 | + MultiZip { | |
165 | + tuple: ( $( self.$idx.par_iter(), )+ ), | |
166 | + } | |
167 | + } | |
168 | + } | |
169 | + | |
170 | + impl<'a, $( $T, )+> IntoParallelIterator for &'a mut ($( $T, )+) | |
171 | + where | |
172 | + $( | |
173 | + $T: IntoParallelRefMutIterator<'a>, | |
174 | + $T::Iter: IndexedParallelIterator, | |
175 | + )+ | |
176 | + { | |
177 | + type Item = ($( $T::Item, )+); | |
178 | + type Iter = MultiZip<($( $T::Iter, )+)>; | |
179 | + | |
180 | + fn into_par_iter(self) -> Self::Iter { | |
181 | + MultiZip { | |
182 | + tuple: ( $( self.$idx.par_iter_mut(), )+ ), | |
183 | + } | |
184 | + } | |
185 | + } | |
186 | + | |
187 | + impl<$( $T, )+> ParallelIterator for MultiZip<($( $T, )+)> | |
188 | + where | |
189 | + $( $T: IndexedParallelIterator, )+ | |
190 | + { | |
191 | + type Item = ($( $T::Item, )+); | |
192 | + | |
193 | + fn drive_unindexed<CONSUMER>(self, consumer: CONSUMER) -> CONSUMER::Result | |
194 | + where | |
195 | + CONSUMER: UnindexedConsumer<Self::Item>, | |
196 | + { | |
197 | + self.drive(consumer) | |
198 | + } | |
199 | + | |
200 | + fn opt_len(&self) -> Option<usize> { | |
201 | + Some(self.len()) | |
202 | + } | |
203 | + } | |
204 | + | |
205 | + impl<$( $T, )+> IndexedParallelIterator for MultiZip<($( $T, )+)> | |
206 | + where | |
207 | + $( $T: IndexedParallelIterator, )+ | |
208 | + { | |
209 | + fn drive<CONSUMER>(self, consumer: CONSUMER) -> CONSUMER::Result | |
210 | + where | |
211 | + CONSUMER: Consumer<Self::Item>, | |
212 | + { | |
213 | + reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip) | |
214 | + .map(flatten!($( $T ),+)) | |
215 | + .drive(consumer) | |
216 | + } | |
217 | + | |
218 | + fn len(&self) -> usize { | |
219 | + reduce!($( self.tuple.$idx.len() ),+ => Ord::min) | |
220 | + } | |
221 | + | |
222 | + fn with_producer<CB>(self, callback: CB) -> CB::Output | |
223 | + where | |
224 | + CB: ProducerCallback<Self::Item>, | |
225 | + { | |
226 | + reduce!($( self.tuple.$idx ),+ => IndexedParallelIterator::zip) | |
227 | + .map(flatten!($( $T ),+)) | |
228 | + .with_producer(callback) | |
229 | + } | |
230 | + } | |
231 | + )+ | |
232 | + } | |
233 | +} | |
234 | + | |
235 | +multizip_impls! { | |
236 | + Tuple1 { | |
237 | + (0) -> A | |
238 | + } | |
239 | + Tuple2 { | |
240 | + (0) -> A | |
241 | + (1) -> B | |
242 | + } | |
243 | + Tuple3 { | |
244 | + (0) -> A | |
245 | + (1) -> B | |
246 | + (2) -> C | |
247 | + } | |
248 | + Tuple4 { | |
249 | + (0) -> A | |
250 | + (1) -> B | |
251 | + (2) -> C | |
252 | + (3) -> D | |
253 | + } | |
254 | + Tuple5 { | |
255 | + (0) -> A | |
256 | + (1) -> B | |
257 | + (2) -> C | |
258 | + (3) -> D | |
259 | + (4) -> E | |
260 | + } | |
261 | + Tuple6 { | |
262 | + (0) -> A | |
263 | + (1) -> B | |
264 | + (2) -> C | |
265 | + (3) -> D | |
266 | + (4) -> E | |
267 | + (5) -> F | |
268 | + } | |
269 | + Tuple7 { | |
270 | + (0) -> A | |
271 | + (1) -> B | |
272 | + (2) -> C | |
273 | + (3) -> D | |
274 | + (4) -> E | |
275 | + (5) -> F | |
276 | + (6) -> G | |
277 | + } | |
278 | + Tuple8 { | |
279 | + (0) -> A | |
280 | + (1) -> B | |
281 | + (2) -> C | |
282 | + (3) -> D | |
283 | + (4) -> E | |
284 | + (5) -> F | |
285 | + (6) -> G | |
286 | + (7) -> H | |
287 | + } | |
288 | + Tuple9 { | |
289 | + (0) -> A | |
290 | + (1) -> B | |
291 | + (2) -> C | |
292 | + (3) -> D | |
293 | + (4) -> E | |
294 | + (5) -> F | |
295 | + (6) -> G | |
296 | + (7) -> H | |
297 | + (8) -> I | |
298 | + } | |
299 | + Tuple10 { | |
300 | + (0) -> A | |
301 | + (1) -> B | |
302 | + (2) -> C | |
303 | + (3) -> D | |
304 | + (4) -> E | |
305 | + (5) -> F | |
306 | + (6) -> G | |
307 | + (7) -> H | |
308 | + (8) -> I | |
309 | + (9) -> J | |
310 | + } | |
311 | + Tuple11 { | |
312 | + (0) -> A | |
313 | + (1) -> B | |
314 | + (2) -> C | |
315 | + (3) -> D | |
316 | + (4) -> E | |
317 | + (5) -> F | |
318 | + (6) -> G | |
319 | + (7) -> H | |
320 | + (8) -> I | |
321 | + (9) -> J | |
322 | + (10) -> K | |
323 | + } | |
324 | + Tuple12 { | |
325 | + (0) -> A | |
326 | + (1) -> B | |
327 | + (2) -> C | |
328 | + (3) -> D | |
329 | + (4) -> E | |
330 | + (5) -> F | |
331 | + (6) -> G | |
332 | + (7) -> H | |
333 | + (8) -> I | |
334 | + (9) -> J | |
335 | + (10) -> K | |
336 | + (11) -> L | |
337 | + } | |
338 | +} |
@@ -162,3 +162,20 @@ fn clone_repeat() { | ||
162 | 162 | fn clone_splitter() { |
163 | 163 | check(rayon::iter::split(0..1000, |x| (x, None))); |
164 | 164 | } |
165 | + | |
166 | +#[test] | |
167 | +fn clone_multizip() { | |
168 | + let v: &Vec<_> = &(0..1000).collect(); | |
169 | + check((v,).into_par_iter()); | |
170 | + check((v, v).into_par_iter()); | |
171 | + check((v, v, v).into_par_iter()); | |
172 | + check((v, v, v, v).into_par_iter()); | |
173 | + check((v, v, v, v, v).into_par_iter()); | |
174 | + check((v, v, v, v, v, v).into_par_iter()); | |
175 | + check((v, v, v, v, v, v, v).into_par_iter()); | |
176 | + check((v, v, v, v, v, v, v, v).into_par_iter()); | |
177 | + check((v, v, v, v, v, v, v, v, v).into_par_iter()); | |
178 | + check((v, v, v, v, v, v, v, v, v, v).into_par_iter()); | |
179 | + check((v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); | |
180 | + check((v, v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); | |
181 | +} |
@@ -173,3 +173,20 @@ fn debug_repeat() { | ||
173 | 173 | fn debug_splitter() { |
174 | 174 | check(rayon::iter::split(0..10, |x| (x, None))); |
175 | 175 | } |
176 | + | |
177 | +#[test] | |
178 | +fn debug_multizip() { | |
179 | + let v: &Vec<_> = &(0..10).collect(); | |
180 | + check((v,).into_par_iter()); | |
181 | + check((v, v).into_par_iter()); | |
182 | + check((v, v, v).into_par_iter()); | |
183 | + check((v, v, v, v).into_par_iter()); | |
184 | + check((v, v, v, v, v).into_par_iter()); | |
185 | + check((v, v, v, v, v, v).into_par_iter()); | |
186 | + check((v, v, v, v, v, v, v).into_par_iter()); | |
187 | + check((v, v, v, v, v, v, v, v).into_par_iter()); | |
188 | + check((v, v, v, v, v, v, v, v, v).into_par_iter()); | |
189 | + check((v, v, v, v, v, v, v, v, v, v).into_par_iter()); | |
190 | + check((v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); | |
191 | + check((v, v, v, v, v, v, v, v, v, v, v, v).into_par_iter()); | |
192 | +} |