FEAT: parallelize witness assignment of bytecode circuit using assign_regions api

.github/workflows/ci.yml

@@ -134,8 +134,16 @@ jobs:
         with:
           command: test
           args: --verbose --release --features scroll --all --exclude integration-tests --exclude circuit-benchmarks serial_  -- --ignored
-
-
+      - name: Run parallel assignment tests(bytecode)
+        uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --release --package zkevm-circuits --lib bytecode_circuit::test --features scroll,parallel_syn -- --nocapture
+      - name: Run parallel assignment tests(state)
+        uses: actions-rs/cargo@v1
+        with:
+          command: test
+          args: --release --package zkevm-circuits --lib state_circuit::test --features scroll,parallel_syn -- --nocapture
   build:
     needs: [skip_check]
     if: |

zkevm-circuits/Cargo.toml

@@ -84,5 +84,6 @@ warn-unimplemented = ["eth-types/warn-unimplemented"]
 onephase = [] # debug only
 zktrie = []
 poseidon-codehash = []
+parallel_syn = ["halo2_proofs/parallel_syn"]
 
-debug-annotations = []
+debug-annotations = []

zkevm-circuits/src/bytecode_circuit/circuit.rs

@@ -545,6 +545,7 @@ impl<F: Field> BytecodeCircuitConfig<F> {
                         &push_data_left_is_zero_chip,
                         &index_length_diff_is_zero_chip,
                         empty_hash,
+                        0,
                         &mut offset,
                         last_row_offset,
                         fail_fast,
@@ -631,6 +632,7 @@ impl<F: Field> BytecodeCircuitConfig<F> {
         push_data_left_is_zero_chip: &IsZeroChip<F>,
         index_length_diff_is_zero_chip: &IsZeroChip<F>,
         empty_hash: Value<F>,
+        offset_begin: usize,
         offset: &mut usize,
         last_row_offset: usize,
         fail_fast: bool,
@@ -653,10 +655,13 @@ impl<F: Field> BytecodeCircuitConfig<F> {
         });
 
         for (idx, row) in bytecode.rows.iter().enumerate() {
-            if fail_fast && *offset > last_row_offset {
+            // for serial assignment api assign_region(), offset always begins with 0
+            // `offset` == global_offset
+            let global_offset = *offset + offset_begin;
+            if fail_fast && global_offset > last_row_offset {
                 log::error!(
-                    "Bytecode Circuit: offset={} > last_row_offset={}",
-                    offset,
+                    "Bytecode Circuit: global_offset={} > last_row_offset={}",
+                    global_offset,
                     last_row_offset
                 );
                 return Err(Error::Synthesis);
@@ -682,14 +687,14 @@ impl<F: Field> BytecodeCircuitConfig<F> {
             }
 
             // Set the data for this row
-            if *offset < last_row_offset {
+            if global_offset < last_row_offset {
                 self.set_row(
                     region,
                     push_data_left_is_zero_chip,
                     index_length_diff_is_zero_chip,
                     *offset,
                     true,
-                    *offset == last_row_offset,
+                    global_offset == last_row_offset,
                     code_hash,
                     row.tag,
                     row.index,
@@ -703,8 +708,8 @@ impl<F: Field> BytecodeCircuitConfig<F> {
                 /*
                 trace!(
                     "bytecode.set_row({}): last:{} h:{:?} t:{:?} i:{:?} c:{:?} v:{:?} pdl:{} rlc:{:?} l:{:?} pds:{:?}",
-                    offset,
-                    offset == last_row_offset,
+                    global_offset,
+                    global_offset == last_row_offset,
                     code_hash,
                     row.tag.get_lower_32(),
                     row.index.get_lower_32(),
@@ -720,14 +725,22 @@ impl<F: Field> BytecodeCircuitConfig<F> {
                 *offset += 1;
                 push_data_left = next_push_data_left
             }
-            if *offset == last_row_offset {
+
+            let global_offset = *offset + offset_begin;
+            if global_offset == last_row_offset {
+                // for serial assignment api assign_region(), offset always begins with 0
+                //     `offset` == global_offset,
+                //     so set `offset` here is equivalent to set `last_row_offset`
+                // if assign_bytecode() is called by assign_regions()
+                //     the rows are divided into multiple regions,
+                //     so the maximum value is the accumulated offset
                 self.set_padding_row(
                     region,
                     push_data_left_is_zero_chip,
                     index_length_diff_is_zero_chip,
                     empty_hash,
                     *offset,
-                    last_row_offset,
+                    *offset,
                 )?;
             }
         }

zkevm-circuits/src/bytecode_circuit/circuit/to_poseidon_hash.rs

@@ -24,6 +24,8 @@ use super::{
     BytecodeCircuitConfig, BytecodeCircuitConfigArgs,
 };
 
+use itertools::Itertools;
+
 #[derive(Clone, Debug)]
 /// Bytecode circuit (for hash block) configuration
 /// basically the BytecodeCircuit include two parts:
@@ -290,7 +292,6 @@ impl<F: Field, const BYTES_IN_FIELD: usize> ToHashBlockCircuitConfig<F, BYTES_IN
         #[cfg(feature = "scroll-trace")]
         {
             use hash_circuit::hash::HASHABLE_DOMAIN_SPEC;
-            use itertools::Itertools;
             let code_hash = bytecode_table.code_hash;
             let pick_hash_tbl_cols = |meta: &mut VirtualCells<F>, inp_i: usize| {
                 debug_assert_eq!(PoseidonTable::INPUT_WIDTH, 2);
@@ -417,6 +418,30 @@ impl<F: Field, const BYTES_IN_FIELD: usize> ToHashBlockCircuitConfig<F, BYTES_IN
         challenges: &Challenges<Value<F>>,
         fail_fast: bool,
     ) -> Result<(), Error> {
+        #[cfg(feature = "parallel_syn")]
+        {
+            // if feature "parallel_syn" is enabled,
+            // `parallel` assignment is turned on by default
+            // we can turn it off by set the environment variable
+            // `CIRCUIT_ASSIGNMENT_TYPE=serial`
+            let assignment_type = std::env::var("CIRCUIT_ASSIGNMENT_TYPE")
+                .ok()
+                .unwrap_or_default();
+            let is_parallel_assignment = match assignment_type.as_str() {
+                "serial" => false,
+                "parallel" => true,
+                &_ => true,
+            };
+            log::debug!("CIRCUIT_ASSIGNMENT_TYPE: {}", assignment_type);
+            log::debug!("is_parallel_assignment: {}", is_parallel_assignment);
+
+            if is_parallel_assignment {
+                return self.assign_internal_parallel(
+                    layouter, size, witness, overwrite, challenges, fail_fast,
+                );
+            }
+        }
+
         let base_conf = &self.base_conf;
         let push_data_left_is_zero_chip =
             IsZeroChip::construct(base_conf.push_data_left_is_zero.clone());
@@ -436,9 +461,23 @@ impl<F: Field, const BYTES_IN_FIELD: usize> ToHashBlockCircuitConfig<F, BYTES_IN
 
         let empty_hash = Value::known(POSEIDON_CODE_HASH_ZERO.to_word().to_scalar().unwrap());
 
+        let mut is_first_time = true;
         layouter.assign_region(
             || "assign bytecode with poseidon hash extension",
             |mut region| {
+                if is_first_time {
+                    is_first_time = false;
+                    base_conf.set_padding_row(
+                        &mut region,
+                        &push_data_left_is_zero_chip,
+                        &index_length_diff_is_zero_chip,
+                        empty_hash,
+                        last_row_offset,
+                        last_row_offset,
+                    )?;
+                    return Ok(());
+                }
+
                 let mut offset = 0;
                 let mut row_input = F::zero();
                 for bytecode in witness.iter() {
@@ -450,6 +489,7 @@ impl<F: Field, const BYTES_IN_FIELD: usize> ToHashBlockCircuitConfig<F, BYTES_IN
                         &push_data_left_is_zero_chip,
                         &index_length_diff_is_zero_chip,
                         empty_hash,
+                        0,
                         &mut offset,
                         last_row_offset,
                         fail_fast,
@@ -492,6 +532,140 @@ impl<F: Field, const BYTES_IN_FIELD: usize> ToHashBlockCircuitConfig<F, BYTES_IN
         )
     }
 
+    #[cfg(feature = "parallel_syn")]
+    pub(crate) fn assign_internal_parallel(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        size: usize,
+        witness: &[UnrolledBytecode<F>],
+        overwrite: &UnrolledBytecode<F>,
+        challenges: &Challenges<Value<F>>,
+        fail_fast: bool,
+    ) -> Result<(), Error> {
+        let base_conf = &self.base_conf;
+        let push_data_left_is_zero_chip =
+            IsZeroChip::construct(base_conf.push_data_left_is_zero.clone());
+        let index_length_diff_is_zero_chip =
+            IsZeroChip::construct(base_conf.index_length_diff_is_zero.clone());
+
+        // Subtract the unusable rows from the size
+        assert!(size > base_conf.minimum_rows);
+        let last_row_offset = size - base_conf.minimum_rows + 1;
+        let empty_hash = Value::known(POSEIDON_CODE_HASH_ZERO.to_word().to_scalar().unwrap());
+
+        let mut is_first_time_vec = vec![true; witness.len()];
+
+        let offsets_of_each_witness = witness
+            .iter()
+            .map(|bytecode| bytecode.rows.len())
+            .collect::<Vec<_>>();
+
+        //start from 0
+        let offsets_prefix_sum = [0]
+            .iter()
+            .chain(offsets_of_each_witness.iter())
+            .scan(0, |sum, &x| {
+                *sum += x;
+                Some(*sum)
+            })
+            .collect::<Vec<_>>();
+
+        layouter.assign_regions(
+            || "assign bytecode with poseidon hash extension(part1: regions)",
+            witness
+                .iter()
+                .zip(is_first_time_vec.iter_mut())
+                .enumerate()
+                .map(|(idx, (bytecode, is_first_time))| {
+                    let push_data_left_is_zero_chip = push_data_left_is_zero_chip.clone();
+                    let index_length_diff_is_zero_chip = index_length_diff_is_zero_chip.clone();
+                    let offsets_prefix_sum = offsets_prefix_sum.clone();
+                    move |region: Region<'_, F>| {
+                        let mut offset = 0;
+                        let mut region = region;
+
+                        if *is_first_time {
+                            *is_first_time = false;
+                            let mut dummy_offset = bytecode.rows.len() - 1;
+                            // last witness
+                            if idx == witness.len() - 1 {
+                                dummy_offset = last_row_offset - offsets_prefix_sum[idx];
+                            }
+                            base_conf.set_padding_row(
+                                &mut region,
+                                &push_data_left_is_zero_chip,
+                                &index_length_diff_is_zero_chip,
+                                empty_hash,
+                                dummy_offset,
+                                dummy_offset,
+                            )?;
+                            return Ok(());
+                        }
+
+                        base_conf.assign_bytecode(
+                            &mut region,
+                            bytecode,
+                            challenges,
+                            &push_data_left_is_zero_chip,
+                            &index_length_diff_is_zero_chip,
+                            empty_hash,
+                            offsets_prefix_sum[idx],
+                            &mut offset,
+                            last_row_offset,
+                            fail_fast,
+                        )?;
+
+                        let last_offset = last_row_offset - offsets_prefix_sum[idx];
+                        // last witness
+                        if idx == witness.len() - 1 {
+                            // Padding
+                            for idx in offset..=last_offset {
+                                base_conf.set_padding_row(
+                                    &mut region,
+                                    &push_data_left_is_zero_chip,
+                                    &index_length_diff_is_zero_chip,
+                                    empty_hash,
+                                    idx,
+                                    last_offset,
+                                )?;
+                                self.set_header_row(&mut region, 0, idx)?;
+                            }
+
+                            base_conf.assign_overwrite(&mut region, overwrite, challenges)?;
+                        }
+
+                        Ok(())
+                    }
+                })
+                .collect_vec(),
+        )?;
+
+        layouter.assign_region(
+            || "assign bytecode with poseidon hash extension(part2)",
+            |mut region| {
+                let mut row_input = F::zero();
+                for (witness_idx, bytecode) in witness.iter().enumerate() {
+                    for (idx, row) in bytecode.rows.iter().enumerate() {
+                        // if the base_conf's assignment not fail fast,
+                        // we also avoid the failure of "NotEnoughRowsAvailable"
+                        // in prover creation (so bytecode_incomplete test could pass)
+                        let offset = offsets_prefix_sum[witness_idx] + idx;
+                        if offset <= last_row_offset {
+                            row_input = self.assign_extended_row(
+                                &mut region,
+                                offset,
+                                row,
+                                row_input,
+                                bytecode.bytes.len(),
+                            )?;
+                        }
+                    }
+                }
+                Ok(())
+            },
+        )
+    }
+
     /// Assign a header row (at padding or start line of each bytecodes)
     fn set_header_row(
         &self,