Skip to content

Commit

Permalink
Use thread pool for assign_regions (#57)
Browse files Browse the repository at this point in the history
* 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 <xiakunxian130@gmail.com>
  • Loading branch information
alannotnerd and kunxian-xia authored Oct 17, 2023
1 parent e3fe25e commit 92fe9b3
Show file tree
Hide file tree
Showing 3 changed files with 584 additions and 837 deletions.
2 changes: 1 addition & 1 deletion .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,4 @@
**/*.rs.bk
.vscode
**/*.html
.DS_Store
.DS_Store
55 changes: 23 additions & 32 deletions halo2_proofs/src/circuit/floor_planner/single_pass.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -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)?;
Expand Down
Loading

0 comments on commit 92fe9b3

Please sign in to comment.