optim-invert: cache the inversion of step height and avoid zero inver…

…sions (scroll-tech#987)

Co-authored-by: Aurélien Nicolas <[email protected]>

zkevm-circuits/src/evm_circuit/execution.rs

@@ -5,7 +5,7 @@ use super::{
         N_BYTE_LOOKUPS, N_COPY_COLUMNS, N_PHASE1_COLUMNS, POW_OF_RAND_TABLE_LOOKUPS,
         RW_TABLE_LOOKUPS, SIG_TABLE_LOOKUPS, TX_TABLE_LOOKUPS,
     },
-    util::{instrumentation::Instrument, CachedRegion, CellManager, StoredExpression},
+    util::{instrumentation::Instrument, CachedRegion, CellManager, Inverter, StoredExpression},
     EvmCircuitExports,
 };
 use crate::{
@@ -987,6 +987,7 @@ impl<F: Field> ExecutionConfig<F> {
     fn assign_q_step(
         &self,
         region: &mut Region<'_, F>,
+        inverter: &Inverter<F>,
         offset: usize,
         height: usize,
     ) -> Result<(), Error> {
@@ -1000,22 +1001,18 @@ impl<F: Field> ExecutionConfig<F> {
                 offset,
                 || Value::known(if idx == 0 { F::one() } else { F::zero() }),
             )?;
-            let value = if idx == 0 {
-                F::zero()
-            } else {
-                F::from((height - idx) as u64)
-            };
+            let value = if idx == 0 { 0 } else { (height - idx) as u64 };
             region.assign_advice(
                 || "step height",
                 self.num_rows_until_next_step,
                 offset,
-                || Value::known(value),
+                || Value::known(F::from(value)),
             )?;
             region.assign_advice(
                 || "step height inv",
                 self.num_rows_inv,
                 offset,
-                || Value::known(value.invert().unwrap_or(F::zero())),
+                || Value::known(inverter.get(value)),
             )?;
         }
         Ok(())
@@ -1047,6 +1044,8 @@ impl<F: Field> ExecutionConfig<F> {
                 }
                 let mut offset = 0;
 
+                let inverter = Inverter::new(MAX_STEP_HEIGHT as u64);
+
                 // Annotate the EVMCircuit columns within it's single region.
                 self.annotate_circuit(&mut region);
 
@@ -1131,7 +1130,7 @@ impl<F: Field> ExecutionConfig<F> {
                     )?;
 
                     // q_step logic
-                    self.assign_q_step(&mut region, offset, height)?;
+                    self.assign_q_step(&mut region, &inverter, offset, height)?;
 
                     offset += height;
                 }
@@ -1168,7 +1167,7 @@ impl<F: Field> ExecutionConfig<F> {
                     )?;
 
                     for row_idx in offset..last_row {
-                        self.assign_q_step(&mut region, row_idx, height)?;
+                        self.assign_q_step(&mut region, &inverter, row_idx, height)?;
                     }
                     offset = last_row;
                 }
@@ -1188,7 +1187,7 @@ impl<F: Field> ExecutionConfig<F> {
                     None,
                     challenges,
                 )?;
-                self.assign_q_step(&mut region, offset, height)?;
+                self.assign_q_step(&mut region, &inverter, offset, height)?;
                 // enable q_step_last
                 self.q_step_last.enable(&mut region, offset)?;
                 offset += height;

zkevm-circuits/src/evm_circuit/util.rs

@@ -15,6 +15,7 @@ use eth_types::{Address, ToLittleEndian, ToWord, U256};
 use halo2_proofs::{
     arithmetic::FieldExt,
     circuit::{AssignedCell, Region, Value},
+    halo2curves::group::ff::BatchInvert,
     plonk::{Advice, Assigned, Column, ConstraintSystem, Error, Expression, VirtualCells},
     poly::Rotation,
 };
@@ -736,3 +737,27 @@ impl<'a> StepRws<'a> {
         rw
     }
 }
+
+/// A struct to cache field inversions.
+pub struct Inverter<F> {
+    inverses: Vec<F>,
+}
+
+impl<F: FieldExt> Inverter<F> {
+    /// Create a new Inverter with preloaded inverses up to `preload_up_to` inclusive.
+    pub fn new(preload_up_to: u64) -> Self {
+        let mut inverses = (0..=preload_up_to).map(F::from).collect::<Vec<F>>();
+
+        inverses.iter_mut().skip(1).batch_invert();
+
+        Self { inverses }
+    }
+
+    /// Return the inverse of `value`, from cache or calculated.
+    pub fn get(&self, value: u64) -> F {
+        match self.inverses.get(value as usize) {
+            Some(i) => *i,
+            None => F::from(value).invert().unwrap(),
+        }
+    }
+}

zkevm-circuits/src/evm_circuit/util/math_gadget/is_zero.rs

@@ -49,13 +49,14 @@ impl<F: Field> IsZeroGadget<F> {
         offset: usize,
         value: F,
     ) -> Result<F, Error> {
-        let inverse = value.invert().unwrap_or(F::zero());
-        self.inverse.assign(region, offset, Value::known(inverse))?;
-        Ok(if value.is_zero().into() {
-            F::one()
-        } else {
+        let is_zero = value.is_zero_vartime();
+        let inverse = if is_zero {
             F::zero()
-        })
+        } else {
+            value.invert().unwrap()
+        };
+        self.inverse.assign(region, offset, Value::known(inverse))?;
+        Ok(F::from(is_zero))
     }
 
     pub(crate) fn assign_value(