Use thread pool for assign_regions (#57)

* feat: use rayon threadpool * feat: add UT for many subregions * refact: move common struct out to module level * refact: reuse common configure code * fix ci errors --------- Co-authored-by: kunxian xia <[email protected]>
scroll-tech · Oct 17, 2023 · 92fe9b3 · 92fe9b3
1 parent e3fe25e
commit 92fe9b3
Show file tree

Hide file tree

Showing 3 changed files with 584 additions and 837 deletions.
diff --git a/.gitignore b/.gitignore
@@ -2,4 +2,4 @@
 **/*.rs.bk
 .vscode
 **/*.html
-.DS_Store
+.DS_Store
diff --git a/halo2_proofs/src/circuit/floor_planner/single_pass.rs b/halo2_proofs/src/circuit/floor_planner/single_pass.rs
@@ -6,6 +6,8 @@ use std::ops::Range;
 use std::sync::{Arc, Mutex};
 use std::time::Instant;
 
+use rayon::prelude::{IndexedParallelIterator, IntoParallelIterator, ParallelIterator};
+
 use ff::Field;
 
 use ark_std::{end_timer, start_timer};
@@ -258,39 +260,28 @@ impl<'a, F: Field, CS: Assignment<F> + 'a> Layouter<F> for SingleChipLayouter<'a
         let ref_sub_cs = sub_cs.iter_mut().collect();
         let sub_layouters = self.fork(ref_sub_cs)?;
         let regions_2nd_pass = Instant::now();
-        let ret = crossbeam::scope(|scope| {
-            let mut handles = vec![];
-            for (i, (mut assignment, mut sub_layouter)) in assignments
-                .into_iter()
-                .zip(sub_layouters.into_iter())
-                .enumerate()
-            {
+        let ret = assignments
+            .into_par_iter()
+            .zip(sub_layouters.into_par_iter())
+            .enumerate()
+            .map(|(i, (mut assignment, mut sub_layouter))| {
                 let region_name = format!("{}_{}", region_name, i);
-                handles.push(scope.spawn(move |_| {
-                    let sub_region_2nd_pass = Instant::now();
-                    sub_layouter.cs.enter_region(|| region_name.clone());
-                    let mut region =
-                        SingleChipLayouterRegion::new(&mut sub_layouter, (region_index + i).into());
-                    let region_ref: &mut dyn RegionLayouter<F> = &mut region;
-                    let result = assignment(region_ref.into());
-                    let constant = region.constants.clone();
-                    sub_layouter.cs.exit_region();
-                    log::debug!(
-                        "region {} 2nd pass synthesis took {:?}",
-                        region_name,
-                        sub_region_2nd_pass.elapsed()
-                    );
-
-                    (result, constant)
-                }));
-            }
-
-            handles
-                .into_iter()
-                .map(|handle| handle.join().expect("handle.join should never fail"))
-                .collect::<Vec<_>>()
-        })
-        .expect("scope should not fail");
+                let sub_region_2nd_pass = Instant::now();
+                sub_layouter.cs.enter_region(|| region_name.clone());
+                let mut region =
+                    SingleChipLayouterRegion::new(&mut sub_layouter, (region_index + i).into());
+                let region_ref: &mut dyn RegionLayouter<F> = &mut region;
+                let result = assignment(region_ref.into());
+                let constant = region.constants.clone();
+                sub_layouter.cs.exit_region();
+                log::debug!(
+                    "region {} 2nd pass synthesis took {:?}",
+                    region_name,
+                    sub_region_2nd_pass.elapsed()
+                );
+                (result, constant)
+            })
+            .collect::<Vec<_>>();
         let cs_merge_time = Instant::now();
         let num_sub_cs = sub_cs.len();
         self.cs.merge(sub_cs)?;