Rewrite SIMD implementation to use Iterator

Author: GabrielMajeri

examples/mandelbrot/src/scalar_par.rs

@@ -6,7 +6,7 @@ use *;
 /// Complex number
 #[repr(align(16))]
 #[derive(Copy, Clone)]
-pub struct Complex {
+struct Complex {
     real: f64,
     imag: f64,
 }
@@ -26,7 +26,7 @@ impl Complex {
 }
 
 /// An iterator yielding the infinite Mandelbrot sequence
-pub struct MandelbrotIter {
+struct MandelbrotIter {
     /// Initial value which generated this sequence
     start: Complex,
     /// Current iteration value

examples/mandelbrot/src/simd_par.rs

@@ -5,45 +5,96 @@ use packed_simd::*;
 use rayon::prelude::*;
 use *;
 
-pub type u64s = u64x8;
-pub type u32s = u32x8;
-pub type f64s = f64x8;
-
-/// This function will operate on N complex numbers at once, where N is the
-/// number of lanes in a SIMD vector of doubles.
-fn mandelbrot(c_x: f64s, c_y: f64s) -> u32s {
-    let mut x = c_x;
-    let mut y = c_y;
+type u64s = u64x8;
+type u32s = u32x8;
+type f64s = f64x8;
+type m64s = m64x8;
+
+/// Storage for complex numbers in SIMD format.
+/// The real and imaginary parts are kept in separate registers.
+#[derive(Copy, Clone)]
+struct Complex {
+    real: f64s,
+    imag: f64s,
+}
 
-    let mut count = u64s::splat(0);
+impl Complex {
+    /// Returns a mask describing which members of the Mandelbrot sequence
+    /// haven't diverged yet
+    #[inline]
+    fn undiverged(&self) -> m64s {
+        let Self { real: x, imag: y } = *self;
 
-    for i in 0..ITER_LIMIT {
         let xx = x * x;
         let yy = y * y;
-        let xy = x * y;
-        let new_x = c_x + xx - yy;
-        let new_y = c_y + xy + xy;
+        let sum = xx + yy;
 
-        let sum = x * x + y * y;
+        sum.le(f64s::splat(THRESHOLD))
+    }
+}
 
-        // Keep track of those lanes which haven't diverged yet. The other ones
-        // will be masked off.
-        let undiverged = sum.le(f64s::splat(4.));
+/// Mandelbrot sequence iterator using SIMD.
+struct MandelbrotIter {
+    /// Initial value which generated this sequence
+    start: Complex,
+    /// Current iteration value
+    current: Complex,
+}
+
+impl MandelbrotIter {
+    /// Creates a new Mandelbrot sequence iterator for a given starting point
+    fn new(start: Complex) -> Self {
+        Self { start, current: start }
+    }
 
-        // Stop the iteration if they all diverged. Note that we don't do this
-        // check every iteration, since a branch misprediction can hurt more
-        // than doing some extra calculations.
-        if i % 5 == 0 && undiverged.none() {
-            break;
+    /// Returns the number of iterations it takes for each member of the Mandelbrot
+    /// sequence to diverge at this point, or `ITER_LIMIT` if they don't diverge.
+    ///
+    /// This function will operate on N complex numbers at once, where N is the
+    /// number of lanes in a SIMD vector of doubles.
+    fn count(mut self) -> u32s {
+        let mut z = self.start;
+        let mut count = u64s::splat(0);
+        for _ in 0..ITER_LIMIT {
+            // Keep track of those lanes which haven't diverged yet. The other ones
+            // will be masked off.
+            let undiverged = z.undiverged();
+
+            // Stop the iteration if they all diverged. Note that we don't do this
+            // check every iteration, since a branch misprediction can hurt more
+            // than doing some extra calculations.
+            if undiverged.none() {
+                break;
+            }
+
+            count += undiverged.select(u64s::splat(1), u64s::splat(0));
+
+            z = self.next().unwrap();
         }
+        count.cast()
+    }
+}
 
-        count += undiverged.select(u64s::splat(1), u64s::splat(0));
+impl Iterator for MandelbrotIter {
+    type Item = Complex;
 
-        x = new_x;
-        y = new_y;
-    }
+    /// Generates the next values in the sequence
+    #[inline]
+    fn next(&mut self) -> Option<Complex> {
+        let Complex { real: c_x, imag: c_y } = self.start;
+        let Complex { real: x, imag: y } = self.current;
 
-    count.cast()
+        let xx = x * x;
+        let yy = y * y;
+        let xy = x * y;
+
+        let new_x = c_x + (xx - yy);
+        let new_y = c_y + (xy + xy);
+
+        self.current = Complex { real: new_x, imag: new_y };
+
+        Some(self.current)
+    }
 }
 
 pub fn generate(dims: Dimensions, xr: Range, yr: Range) -> Vec<u32> {
@@ -89,7 +140,8 @@ pub fn generate(dims: Dimensions, xr: Range, yr: Range) -> Vec<u32> {
         let y = f64s::splat(yr.start + dy * (i as f64));
         row.iter_mut().enumerate().for_each(|(j, count)| {
             let x = xs[j];
-            *count = mandelbrot(x, y);
+            let z = Complex { real: x, imag: y };
+            *count = MandelbrotIter::new(z).count();
         });
     });