Merge pull request #230 from image-rs/rayon-blocking

Restructure multithreading

Author: HeroicKatora

.github/workflows/rust.yml

@@ -49,6 +49,7 @@ jobs:
     steps:
       - name: Installing emulator and linker
         run: |
+          sudo apt-get update
           sudo apt-get install qemu binfmt-support qemu-user-static gcc-aarch64-linux-gnu binutils-aarch64-linux-gnu
 
       - name: Installing Rust toolchain

Cargo.toml

@@ -19,6 +19,13 @@ png = "0.16"
 walkdir = "2.0"
 criterion = "0.3"
 
+[features]
+default = ["rayon"]
+platform_independent = []
+nightly_aarch64_neon = []
+
+## Internal development configuration: testing and benchmarking
+
 [[bench]]
 name = "decoding_benchmark"
 harness = false
@@ -27,8 +34,18 @@ harness = false
 name = "large_image"
 harness = false
 
-[features]
-default = ["rayon"]
-platform_independent = []
-nightly_aarch64_neon = []
+[[test]]
+name = "rayon"
+required-features = ["rayon"]
+
+[[test]]
+name = "rayon-0"
+required-features = ["rayon"]
+
+[[test]]
+name = "rayon-1"
+required-features = ["rayon"]
 
+[[test]]
+name = "rayon-2"
+required-features = ["rayon"]

src/decoder.rs

@@ -17,7 +17,7 @@ use crate::parser::{
     IccChunk, ScanInfo,
 };
 use crate::upsampler::Upsampler;
-use crate::worker::{PlatformWorker, RowData, Worker};
+use crate::worker::{PreferWorkerKind, RowData, Worker, with_worker};
 
 pub const MAX_COMPONENTS: usize = 4;
 
@@ -191,7 +191,9 @@ impl<R: Read> Decoder<R> {
     ///
     /// If successful, the metadata can be obtained using the `info` method.
     pub fn read_info(&mut self) -> Result<()> {
-        self.decode_internal(true).map(|_| ())
+        with_worker(PreferWorkerKind::Multithreaded, |worker| {
+            self.decode_internal(true, worker)
+        }).map(|_| ())
     }
 
     /// Configure the decoder to scale the image during decoding.
@@ -219,10 +221,16 @@ impl<R: Read> Decoder<R> {
 
     /// Decodes the image and returns the decoded pixels if successful.
     pub fn decode(&mut self) -> Result<Vec<u8>> {
-        self.decode_internal(false)
+        with_worker(PreferWorkerKind::Multithreaded, |worker| {
+            self.decode_internal(false, worker)
+        })
     }
 
-    fn decode_internal(&mut self, stop_after_metadata: bool) -> Result<Vec<u8>> {
+    fn decode_internal(
+        &mut self,
+        stop_after_metadata: bool,
+        worker: &mut dyn Worker,
+    ) -> Result<Vec<u8>> {
         if stop_after_metadata && self.frame.is_some() {
             // The metadata has already been read.
             return Ok(Vec::new());
@@ -237,7 +245,6 @@ impl<R: Read> Decoder<R> {
 
         let mut previous_marker = Marker::SOI;
         let mut pending_marker = None;
-        let mut worker = None;
         let mut scans_processed = 0;
         let mut planes = vec![
             Vec::<u8>::new();
@@ -318,9 +325,6 @@ impl<R: Read> Decoder<R> {
                     if self.frame.is_none() {
                         return Err(Error::Format("scan encountered before frame".to_owned()));
                     }
-                    if worker.is_none() {
-                        worker = Some(PlatformWorker::new()?);
-                    }
 
                     let frame = self.frame.clone().unwrap();
                     let scan = parse_sos(&mut self.reader, &frame)?;
@@ -383,7 +387,7 @@ impl<R: Read> Decoder<R> {
                         }
 
                         let (marker, data) =
-                            self.decode_scan(&frame, &scan, worker.as_mut().unwrap(), &finished)?;
+                            self.decode_scan(&frame, &scan, worker, &finished)?;
 
                         if let Some(data) = data {
                             for (i, plane) in data
@@ -545,10 +549,6 @@ impl<R: Read> Decoder<R> {
                     };
 
                 // Get the worker prepared
-                if worker.is_none() {
-                    worker = Some(PlatformWorker::new()?);
-                }
-                let worker = worker.as_mut().unwrap();
                 let row_data = RowData {
                     index: i,
                     component: component.clone(),
@@ -560,14 +560,17 @@ impl<R: Read> Decoder<R> {
                 let coefficients_per_mcu_row = usize::from(component.block_size.width)
                     * usize::from(component.vertical_sampling_factor)
                     * 64;
-                for mcu_y in 0..frame.mcu_size.height {
-                    let row_coefficients = {
+
+                let mut tasks = (0..frame.mcu_size.height)
+                    .map(|mcu_y| {
                         let offset = usize::from(mcu_y) * coefficients_per_mcu_row;
-                        self.coefficients[i][offset..offset + coefficients_per_mcu_row].to_vec()
-                    };
+                        let row_coefficients = self.coefficients[i][offset..offset + coefficients_per_mcu_row].to_vec();
+                        (i, row_coefficients)
+                    });
 
-                    worker.append_row((i, row_coefficients))?;
-                }
+                // FIXME: additional potential work stealing opportunities for rayon case if we
+                // also internally can parallelize over components.
+                worker.append_rows(&mut tasks)?;
                 planes[i] = worker.get_result(i)?;
             }
         }
@@ -616,7 +619,7 @@ impl<R: Read> Decoder<R> {
         &mut self,
         frame: &FrameInfo,
         scan: &ScanInfo,
-        worker: &mut PlatformWorker,
+        worker: &mut dyn Worker,
         finished: &[bool; MAX_COMPONENTS],
     ) -> Result<(Option<Marker>, Option<Vec<Vec<u8>>>)> {
         assert!(scan.component_indices.len() <= MAX_COMPONENTS);
@@ -871,6 +874,8 @@ impl<R: Read> Decoder<R> {
                         )
                     };
 
+                    // FIXME: additional potential work stealing opportunities for rayon case if we
+                    // also internally can parallelize over components.
                     worker.append_row((i, row_coefficients))?;
                 }
             }

src/worker/immediate.rs

@@ -16,6 +16,11 @@ pub struct ImmediateWorker {
     quantization_tables: Vec<Option<Arc<[u16; 64]>>>,
 }
 
+pub fn with_immediate<T>(f: impl FnOnce(&mut dyn Worker) -> T) -> T {
+    let mut worker = ImmediateWorker::new_immediate();
+    f(&mut worker)
+}
+
 impl ImmediateWorker {
     pub fn new_immediate() -> ImmediateWorker {
         ImmediateWorker {
@@ -25,6 +30,7 @@ impl ImmediateWorker {
             quantization_tables: vec![None; MAX_COMPONENTS],
         }
     }
+
     pub fn start_immediate(&mut self, data: RowData) {
         assert!(self.results[data.index].is_empty());
 
@@ -33,6 +39,7 @@ impl ImmediateWorker {
         self.components[data.index] = Some(data.component);
         self.quantization_tables[data.index] = Some(data.quantization_table);
     }
+
     pub fn append_row_immediate(&mut self, (index, data): (usize, Vec<i16>)) {
         // Convert coefficients from a MCU row to samples.
 
@@ -55,15 +62,13 @@ impl ImmediateWorker {
 
         self.offsets[index] += block_count * component.dct_scale * component.dct_scale;
     }
+
     pub fn get_result_immediate(&mut self, index: usize) -> Vec<u8> {
         mem::take(&mut self.results[index])
     }
 }
 
 impl Worker for ImmediateWorker {
-    fn new() -> Result<Self> {
-        Ok(ImmediateWorker::new_immediate())
-    }
     fn start(&mut self, data: RowData) -> Result<()> {
         self.start_immediate(data);
         Ok(())

src/worker/mod.rs

@@ -1,10 +1,7 @@
 mod immediate;
 mod multithreaded;
-
-#[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
-pub use self::multithreaded::MultiThreadedWorker as PlatformWorker;
-#[cfg(any(target_arch = "asmjs", target_arch = "wasm32"))]
-pub use self::immediate::ImmediateWorker as PlatformWorker;
+#[cfg(feature = "rayon")]
+mod rayon;
 
 use alloc::sync::Arc;
 use alloc::vec::Vec;
@@ -17,9 +14,34 @@ pub struct RowData {
     pub quantization_table: Arc<[u16; 64]>,
 }
 
-pub trait Worker: Sized {
-    fn new() -> Result<Self>;
+pub trait Worker {
     fn start(&mut self, row_data: RowData) -> Result<()>;
     fn append_row(&mut self, row: (usize, Vec<i16>)) -> Result<()>;
     fn get_result(&mut self, index: usize) -> Result<Vec<u8>>;
+    /// Default implementation for spawning multiple tasks.
+    fn append_rows(&mut self, row: &mut dyn Iterator<Item=(usize, Vec<i16>)>)
+        -> Result<()>
+    {
+        for item in row {
+            self.append_row(item)?;
+        }
+        Ok(())
+    }
+}
+
+pub enum PreferWorkerKind {
+    Immediate,
+    Multithreaded,
+}
+
+/// Execute something with a worker system.
+pub fn with_worker<T>(prefer: PreferWorkerKind, f: impl FnOnce(&mut dyn Worker) -> T) -> T {
+    match prefer {
+        #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
+        #[cfg(feature = "rayon")]
+        PreferWorkerKind::Multithreaded => self::rayon::with_rayon(f),
+        #[cfg(not(any(target_arch = "asmjs", target_arch = "wasm32")))]
+        PreferWorkerKind::Multithreaded => self::multithreaded::with_multithreading(f),
+        _ => self::immediate::with_immediate(f),
+    }
 }