From 81b3fe66fba04b8b60ba79b3641826df22fca67e Mon Sep 17 00:00:00 2001 From: ScottyThePilot Date: Sat, 2 Dec 2023 16:45:30 -0500 Subject: [PATCH] Add `rayon` parallel iterators (#2058) --- Cargo.toml | 1 + src/buffer.rs | 2 +- src/buffer_par.rs | 496 ++++++++++++++++++++++++++++++++++++++++++++++ src/lib.rs | 5 + 4 files changed, 503 insertions(+), 1 deletion(-) create mode 100644 src/buffer_par.rs diff --git a/Cargo.toml b/Cargo.toml index 8d3eb68d4f..9faa533b60 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -43,6 +43,7 @@ color_quant = "1.1" exr = { version = "1.5.0", optional = true } qoi = { version = "0.4", optional = true } libwebp = { package = "webp", version = "0.2.2", default-features = false, optional = true } +rayon = { version = "1.7.0", optional = true } [dev-dependencies] crc32fast = "1.2.0" diff --git a/src/buffer.rs b/src/buffer.rs index e3fd351e13..046bdebe76 100644 --- a/src/buffer.rs +++ b/src/buffer.rs @@ -888,7 +888,7 @@ where Container: Deref + DerefMut, { // TODO: choose name under which to expose. - fn inner_pixels_mut(&mut self) -> &mut [P::Subpixel] { + pub(crate) fn inner_pixels_mut(&mut self) -> &mut [P::Subpixel] { let len = Self::image_buffer_len(self.width, self.height).unwrap(); &mut self.data[..len] } diff --git a/src/buffer_par.rs b/src/buffer_par.rs new file mode 100644 index 0000000000..aa90a86664 --- /dev/null +++ b/src/buffer_par.rs @@ -0,0 +1,496 @@ +use rayon::iter::plumbing::*; +use rayon::iter::{IndexedParallelIterator, ParallelIterator}; +use rayon::slice::{ChunksExact, ChunksExactMut, ParallelSlice, ParallelSliceMut}; +use std::fmt; +use std::ops::{Deref, DerefMut}; + +use crate::traits::Pixel; +use crate::ImageBuffer; + +/// Parallel iterator over pixel refs. +#[derive(Clone)] +pub struct PixelsPar<'a, P> +where + P: Pixel + Sync + 'a, + P::Subpixel: Sync + 'a, +{ + chunks: ChunksExact<'a, P::Subpixel>, +} + +impl<'a, P> ParallelIterator for PixelsPar<'a, P> +where + P: Pixel + Sync + 'a, + P::Subpixel: Sync + 'a, +{ + type Item = &'a P; + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + self.chunks + .map(|v|

::from_slice(v)) + .drive_unindexed(consumer) + } + + fn opt_len(&self) -> Option { + Some(self.len()) + } +} + +impl<'a, P> IndexedParallelIterator for PixelsPar<'a, P> +where + P: Pixel + Sync + 'a, + P::Subpixel: Sync + 'a, +{ + fn drive>(self, consumer: C) -> C::Result { + self.chunks + .map(|v|

::from_slice(v)) + .drive(consumer) + } + + fn len(&self) -> usize { + self.chunks.len() + } + + fn with_producer>(self, callback: CB) -> CB::Output { + self.chunks + .map(|v|

::from_slice(v)) + .with_producer(callback) + } +} + +impl

fmt::Debug for PixelsPar<'_, P> +where + P: Pixel + Sync, + P::Subpixel: Sync + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("PixelsPar") + .field("chunks", &self.chunks) + .finish() + } +} + +/// Parallel iterator over mutable pixel refs. +pub struct PixelsMutPar<'a, P> +where + P: Pixel + Send + Sync + 'a, + P::Subpixel: Send + Sync + 'a, +{ + chunks: ChunksExactMut<'a, P::Subpixel>, +} + +impl<'a, P> ParallelIterator for PixelsMutPar<'a, P> +where + P: Pixel + Send + Sync + 'a, + P::Subpixel: Send + Sync + 'a, +{ + type Item = &'a mut P; + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + self.chunks + .map(|v|

::from_slice_mut(v)) + .drive_unindexed(consumer) + } + + fn opt_len(&self) -> Option { + Some(self.len()) + } +} + +impl<'a, P> IndexedParallelIterator for PixelsMutPar<'a, P> +where + P: Pixel + Send + Sync + 'a, + P::Subpixel: Send + Sync + 'a, +{ + fn drive>(self, consumer: C) -> C::Result { + self.chunks + .map(|v|

::from_slice_mut(v)) + .drive(consumer) + } + + fn len(&self) -> usize { + self.chunks.len() + } + + fn with_producer>(self, callback: CB) -> CB::Output { + self.chunks + .map(|v|

::from_slice_mut(v)) + .with_producer(callback) + } +} + +impl

fmt::Debug for PixelsMutPar<'_, P> +where + P: Pixel + Send + Sync, + P::Subpixel: Send + Sync + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("PixelsMutPar") + .field("chunks", &self.chunks) + .finish() + } +} + +/// Parallel iterator over pixel refs and their coordinates. +#[derive(Clone)] +pub struct EnumeratePixelsPar<'a, P> +where + P: Pixel + Sync + 'a, + P::Subpixel: Sync + 'a, +{ + pixels: PixelsPar<'a, P>, + width: u32, +} + +impl<'a, P> ParallelIterator for EnumeratePixelsPar<'a, P> +where + P: Pixel + Sync + 'a, + P::Subpixel: Sync + 'a, +{ + type Item = (u32, u32, &'a P); + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + self.pixels + .enumerate() + .map(|(i, p)| { + ( + (i % self.width as usize) as u32, + (i / self.width as usize) as u32, + p, + ) + }) + .drive_unindexed(consumer) + } + + fn opt_len(&self) -> Option { + Some(self.len()) + } +} + +impl<'a, P> IndexedParallelIterator for EnumeratePixelsPar<'a, P> +where + P: Pixel + Sync + 'a, + P::Subpixel: Sync + 'a, +{ + fn drive>(self, consumer: C) -> C::Result { + self.pixels + .enumerate() + .map(|(i, p)| { + ( + (i % self.width as usize) as u32, + (i / self.width as usize) as u32, + p, + ) + }) + .drive(consumer) + } + + fn len(&self) -> usize { + self.pixels.len() + } + + fn with_producer>(self, callback: CB) -> CB::Output { + self.pixels + .enumerate() + .map(|(i, p)| { + ( + (i % self.width as usize) as u32, + (i / self.width as usize) as u32, + p, + ) + }) + .with_producer(callback) + } +} + +impl

fmt::Debug for EnumeratePixelsPar<'_, P> +where + P: Pixel + Sync, + P::Subpixel: Sync + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("EnumeratePixelsPar") + .field("pixels", &self.pixels) + .field("width", &self.width) + .finish() + } +} + +/// Parallel iterator over mutable pixel refs and their coordinates. +pub struct EnumeratePixelsMutPar<'a, P> +where + P: Pixel + Send + Sync + 'a, + P::Subpixel: Send + Sync + 'a, +{ + pixels: PixelsMutPar<'a, P>, + width: u32, +} + +impl<'a, P> ParallelIterator for EnumeratePixelsMutPar<'a, P> +where + P: Pixel + Send + Sync + 'a, + P::Subpixel: Send + Sync + 'a, +{ + type Item = (u32, u32, &'a mut P); + + fn drive_unindexed(self, consumer: C) -> C::Result + where + C: UnindexedConsumer, + { + self.pixels + .enumerate() + .map(|(i, p)| { + ( + (i % self.width as usize) as u32, + (i / self.width as usize) as u32, + p, + ) + }) + .drive_unindexed(consumer) + } + + fn opt_len(&self) -> Option { + Some(self.len()) + } +} + +impl<'a, P> IndexedParallelIterator for EnumeratePixelsMutPar<'a, P> +where + P: Pixel + Send + Sync + 'a, + P::Subpixel: Send + Sync + 'a, +{ + fn drive>(self, consumer: C) -> C::Result { + self.pixels + .enumerate() + .map(|(i, p)| { + ( + (i % self.width as usize) as u32, + (i / self.width as usize) as u32, + p, + ) + }) + .drive(consumer) + } + + fn len(&self) -> usize { + self.pixels.len() + } + + fn with_producer>(self, callback: CB) -> CB::Output { + self.pixels + .enumerate() + .map(|(i, p)| { + ( + (i % self.width as usize) as u32, + (i / self.width as usize) as u32, + p, + ) + }) + .with_producer(callback) + } +} + +impl

fmt::Debug for EnumeratePixelsMutPar<'_, P> +where + P: Pixel + Send + Sync, + P::Subpixel: Send + Sync + fmt::Debug, +{ + fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result { + f.debug_struct("EnumeratePixelsMutPar") + .field("pixels", &self.pixels) + .field("width", &self.width) + .finish() + } +} + +impl ImageBuffer +where + P: Pixel + Sync, + P::Subpixel: Sync, + Container: Deref, +{ + /// Returns a parallel iterator over the pixels of this image, usable with `rayon`. + /// See [`pixels`] for more information. + /// + /// [`pixels`]: #method.pixels + pub fn par_pixels(&self) -> PixelsPar

{ + PixelsPar { + chunks: self + .inner_pixels() + .par_chunks_exact(

::CHANNEL_COUNT as usize), + } + } + + /// Returns a parallel iterator over the pixels of this image and their coordinates, usable with `rayon`. + /// See [`enumerate_pixels`] for more information. + /// + /// [`enumerate_pixels`]: #method.enumerate_pixels + pub fn par_enumerate_pixels(&self) -> EnumeratePixelsPar

{ + EnumeratePixelsPar { + pixels: self.par_pixels(), + width: self.width(), + } + } +} + +impl ImageBuffer +where + P: Pixel + Send + Sync, + P::Subpixel: Send + Sync, + Container: Deref + DerefMut, +{ + /// Returns a parallel iterator over the mutable pixels of this image, usable with `rayon`. + /// See [`pixels_mut`] for more information. + /// + /// [`pixels_mut`]: #method.pixels_mut + pub fn par_pixels_mut(&mut self) -> PixelsMutPar

{ + PixelsMutPar { + chunks: self + .inner_pixels_mut() + .par_chunks_exact_mut(

::CHANNEL_COUNT as usize), + } + } + + /// Returns a parallel iterator over the mutable pixels of this image and their coordinates, usable with `rayon`. + /// See [`enumerate_pixels_mut`] for more information. + /// + /// [`enumerate_pixels_mut`]: #method.enumerate_pixels_mut + pub fn par_enumerate_pixels_mut(&mut self) -> EnumeratePixelsMutPar

{ + let width = self.width(); + EnumeratePixelsMutPar { + pixels: self.par_pixels_mut(), + width, + } + } +} + +impl

ImageBuffer> +where + P: Pixel + Send + Sync, + P::Subpixel: Send + Sync, +{ + /// Constructs a new ImageBuffer by repeated application of the supplied function, + /// utilizing multi-threading via `rayon`. + /// + /// The arguments to the function are the pixel's x and y coordinates. + /// + /// # Panics + /// + /// Panics when the resulting image is larger the the maximum size of a vector. + pub fn from_par_fn(width: u32, height: u32, f: F) -> ImageBuffer> + where + F: Fn(u32, u32) -> P + Send + Sync, + { + let mut buf = ImageBuffer::new(width, height); + buf.par_enumerate_pixels_mut().for_each(|(x, y, p)| { + *p = f(x, y); + }); + + buf + } +} + +#[cfg(test)] +mod test { + use crate::{Rgb, RgbImage}; + use rayon::iter::{IndexedParallelIterator, ParallelIterator}; + + fn test_width_height(width: u32, height: u32, len: usize) { + let mut image = RgbImage::new(width, height); + + assert_eq!(image.par_enumerate_pixels_mut().len(), len); + assert_eq!(image.par_enumerate_pixels().len(), len); + assert_eq!(image.par_pixels_mut().len(), len); + assert_eq!(image.par_pixels().len(), len); + } + + #[test] + fn zero_width_zero_height() { + test_width_height(0, 0, 0); + } + + #[test] + fn zero_width_nonzero_height() { + test_width_height(0, 2, 0); + } + + #[test] + fn nonzero_width_zero_height() { + test_width_height(2, 0, 0); + } + + #[test] + fn iter_parity() { + let mut image1 = RgbImage::from_fn(17, 29, |x, y| { + Rgb(std::array::from_fn(|i| { + ((x + y * 98 + i as u32 * 27) % 255) as u8 + })) + }); + let mut image2 = image1.clone(); + + assert_eq!( + image1.enumerate_pixels_mut().collect::>(), + image2.par_enumerate_pixels_mut().collect::>() + ); + assert_eq!( + image1.enumerate_pixels().collect::>(), + image2.par_enumerate_pixels().collect::>() + ); + assert_eq!( + image1.pixels_mut().collect::>(), + image2.par_pixels_mut().collect::>() + ); + assert_eq!( + image1.pixels().collect::>(), + image2.par_pixels().collect::>() + ); + } +} + +#[cfg(test)] +#[cfg(feature = "benchmarks")] +mod benchmarks { + use crate::{Rgb, RgbImage}; + + const S: u32 = 1024; + + #[bench] + fn creation(b: &mut test::Bencher) { + let mut bytes = 0; + b.iter(|| { + let img = RgbImage::from_fn(S, S, |_, _| test::black_box(pixel_func())); + + bytes += img.as_raw().len() as u64; + }); + + b.bytes = bytes; + } + + #[bench] + fn creation_par(b: &mut test::Bencher) { + let mut bytes = 0; + b.iter(|| { + let img = RgbImage::from_par_fn(S, S, |_, _| test::black_box(pixel_func())); + + bytes += img.as_raw().len() as u64; + }); + + b.bytes = bytes; + } + + fn pixel_func() -> Rgb { + use std::collections::hash_map::RandomState; + use std::hash::{BuildHasher, Hasher}; + Rgb(std::array::from_fn(|_| { + RandomState::new().build_hasher().finish() as u8 + })) + } +} diff --git a/src/lib.rs b/src/lib.rs index 304ee72cfd..8122d05abb 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -185,6 +185,9 @@ pub mod buffer { ConvertBuffer, EnumeratePixels, EnumeratePixelsMut, EnumerateRows, EnumerateRowsMut, Pixels, PixelsMut, Rows, RowsMut, }; + + #[cfg(feature = "rayon")] + pub use crate::buffer_par::*; } // Math utils @@ -282,6 +285,8 @@ pub mod codecs { mod animation; #[path = "buffer.rs"] mod buffer_; +#[cfg(feature = "rayon")] +mod buffer_par; mod color; mod dynimage; mod image;