From 1c9841158c962e7c96bbf378fde2f7a7b837df96 Mon Sep 17 00:00:00 2001 From: Georg Wiese Date: Tue, 14 May 2024 20:23:35 +0200 Subject: [PATCH] Halo2 backend: use fewer copy constraints (#1376) In Halo2, we need to simulate the "wrapping" behavior that we expect from the backend, by copying the first row of witness columns after the last. This was previously implemented by using copy constraints. Considering [how the argument works](https://zcash.github.io/halo2/design/proving-system/permutation.html), this leads to 1 additional fixed column during setup *per witness column*, and several additional witness columns (depending on the degree bound, but linear in the number of witness columns). But because Halo2 offers rotations of arbitrary offsets, we can instead constrain this as: `first_step * (witness_column - witness_column'degree)` On the `simple_sum` example, this improved the proof time by ~15% (353.01825ms -> 302.067125ms). The benchmarks on Keccak are still running... --- backend/src/halo2/circuit_builder.rs | 113 +++++++++++++++++---------- 1 file changed, 72 insertions(+), 41 deletions(-) diff --git a/backend/src/halo2/circuit_builder.rs b/backend/src/halo2/circuit_builder.rs index a0faad40d7..94c8fb9571 100644 --- a/backend/src/halo2/circuit_builder.rs +++ b/backend/src/halo2/circuit_builder.rs @@ -21,6 +21,7 @@ use powdr_ast::{ }; use powdr_number::FieldElement; +const FIRST_STEP_NAME: &str = "__first_step"; const ENABLE_NAME: &str = "__enable"; const INSTANCE_NAME: &str = "__instance"; @@ -84,25 +85,34 @@ pub(crate) struct PowdrCircuit<'a, T> { witgen_callback: Option>, } +fn get_publics(analyzed: &Analyzed) -> Vec<(String, usize)> { + let mut publics = analyzed + .public_declarations + .values() + .map(|public_declaration| { + let witness_name = public_declaration.referenced_poly_name(); + let witness_offset = public_declaration.index as usize; + (witness_name, witness_offset) + }) + .collect::>(); + + // Sort, so that the order is deterministic + publics.sort(); + publics +} + impl<'a, T: FieldElement> PowdrCircuit<'a, T> { pub(crate) fn new(analyzed: &'a Analyzed, fixed: &'a [(String, Vec)]) -> Self { - let mut publics = analyzed - .public_declarations - .values() - .map(|public_declaration| { - let witness_name = public_declaration.referenced_poly_name(); - let witness_offset = public_declaration.index as usize; - (witness_name, witness_offset) - }) - .collect::>(); - // Sort, so that the order is deterministic - publics.sort(); + for (fixed_name, _) in fixed { + assert!(fixed_name != FIRST_STEP_NAME); + assert!(fixed_name != ENABLE_NAME); + } Self { analyzed, fixed, witness: None, - publics, + publics: get_publics(analyzed), witgen_callback: None, } } @@ -187,11 +197,13 @@ impl<'a, T: FieldElement, F: PrimeField> Circuit for PowdrCi }) .collect(); + let first_step = meta.fixed_column(); let fixed = analyzed .constant_polys_in_source_order() .iter() .flat_map(|(symbol, _)| symbol.array_elements()) .map(|(name, _)| (name.clone(), meta.fixed_column())) + .chain(iter::once((FIRST_STEP_NAME.to_string(), first_step))) .collect(); let enable = meta.fixed_column(); @@ -222,11 +234,11 @@ impl<'a, T: FieldElement, F: PrimeField> Circuit for PowdrCi challenges, }; - // Enable equality for all witness columns & instance - for &column in config.advice.values() { - meta.enable_equality(column); - } + // Enable equality for instance column & all witness columns with public cells meta.enable_equality(config.instance); + for (column_name, _) in get_publics(&analyzed) { + meta.enable_equality(config.advice[&column_name]); + } // Add polynomial identities let identities = analyzed @@ -249,6 +261,25 @@ impl<'a, T: FieldElement, F: PrimeField> Circuit for PowdrCi }); } + // Add constraints that for all witness columns, cell 0 must equal cell . + // This forces the prover to simulate wrapping correctly. + meta.create_gate("enforce_wrapping", |meta| -> Vec<(String, Expression)> { + let first_step = meta.query_fixed(config.fixed[FIRST_STEP_NAME], Rotation::cur()); + config + .advice + .keys() + .map(|name| { + let first_row = meta.query_advice(config.advice[name], Rotation::cur()); + let last_row = meta.query_advice( + config.advice[name], + Rotation(analyzed.degree().try_into().unwrap()), + ); + let expr = first_step.clone() * (first_row - last_row); + (format!("enforce wrapping ({name})"), expr) + }) + .collect() + }); + // Challenge used to combine the lookup tuple with the selector let beta = Expression::Challenge(meta.challenge_usable_after(FirstPhase)); @@ -389,6 +420,12 @@ impl<'a, T: FieldElement, F: PrimeField> Circuit for PowdrCi i, || Value::known(value), )?; + region.assign_fixed( + || FIRST_STEP_NAME, + config.fixed[FIRST_STEP_NAME], + i, + || Value::known(F::from(if i == 0 { 1 } else { 0 })), + )?; } let publics = self @@ -406,31 +443,25 @@ impl<'a, T: FieldElement, F: PrimeField> Circuit for PowdrCi } else { Some(&new_witness) }; - for (name, &column) in config.advice.iter() { - // Note that we can't skip this loop if we don't have a witness, - // because the copy_advice() call below also adds copy constraints, - // which are needed to compute the verification key. - let values = witness.as_ref().and_then(|witness| { - witness.iter().find_map(|(witness_name, values)| { - (witness_name == name).then_some(values) - }) - }); - for i in 0..degree { - let value = values - .map(|values| Value::known(convert_field::(values[i]))) - .unwrap_or_default(); - - let assigned_cell = region.assign_advice(|| name, column, i, || value)?; - - // The first row needs to be copied to row - if i == 0 { - assigned_cell.copy_advice(|| name, &mut region, column, degree)?; - } - - // Collect public cells, which are later copy-constrained to equal - // a cell in the instance column. - if let Some(&instance_index) = publics.get(&(name.clone(), i)) { - public_cells.push((instance_index, assigned_cell)); + if let Some(witness) = witness { + for (name, values) in witness.iter() { + let column = *config.advice.get(name).unwrap(); + for (i, value) in values.iter().enumerate() { + let value = Value::known(convert_field::(*value)); + + let assigned_cell = + region.assign_advice(|| name, column, i, || value)?; + + // The first row needs to be copied to row + if i == 0 { + region.assign_advice(|| name, column, degree, || value)?; + } + + // Collect public cells, which are later copy-constrained to equal + // a cell in the instance column. + if let Some(&instance_index) = publics.get(&(name.clone(), i)) { + public_cells.push((instance_index, assigned_cell)); + } } } }