added threads option

SalOne22 · Oct 8, 2023 · 704eea2 · 704eea2
1 parent aa9d3bd
commit 704eea2
Show file tree

Hide file tree

Showing 3 changed files with 25 additions and 0 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -4,6 +4,10 @@ All notable changes to the Rimage library will be documented in this file.
 
 ## **[Unreleased]** v0.9.1
 
+### New features
+
+- Added previously removed `--threads` option when parallel feature is enabled
+
 ## [v0.9.0](https://github.com/SalOne22/rimage/releases/tag/v0.9.0)
 
 ### Breaking Changes

diff --git a/Cargo.lock b/Cargo.lock
diff --git a/src/bin/rimage/main.rs b/src/bin/rimage/main.rs
@@ -34,6 +34,9 @@ fn main() -> Result<(), Box<dyn Error>> {
             arg!(-s --suffix [SUFFIX] "Appends suffix to output file(s) names"),
             arg!(-b --backup "Appends '.backup' to input file(s) names")
                 .action(ArgAction::SetTrue),
+            #[cfg(feature = "parallel")]
+            arg!(-t --threads "Number of threads to use [default: number of cores]")
+                .value_parser(value_parser!(usize)),
         ])
         .next_help_heading("Quantization")
         .args([
@@ -59,6 +62,14 @@ fn main() -> Result<(), Box<dyn Error>> {
     let codec = matches.get_one::<Codec>("codec").unwrap();
     let quality = matches.get_one::<f32>("quality").unwrap();
 
+    #[cfg(feature = "parallel")]
+    if let Some(threads) = matches.get_one::<usize>("threads") {
+        rayon::ThreadPoolBuilder::new()
+            .num_threads(*threads)
+            .build_global()
+            .unwrap();
+    }
+
     let mut quantization_config = QuantizationConfig::new();
 
     if let Some(quality) = matches.get_one::<u8>("quantization") {