keccak: switch to multi_pack (privacy-scaling-explorations#813)

* keccak: switch to multi_pack

* fix keccak bench

Makefile

@@ -34,9 +34,6 @@ evm_bench: ## Run Evm Circuit benchmarks
 state_bench: ## Run State Circuit benchmarks
 	@cargo test --profile bench bench_state_circuit_prover -p circuit-benchmarks --features benches  -- --nocapture
 
-keccak_round_bench: ## Run State Circuit benchmarks
-	@cargo test --profile bench bench_keccak_round -p circuit-benchmarks --features benches  -- --nocapture
-
 bit_keccak_bench: ## Run Bit Keccak Circuit benchmarks
 	@cargo test --profile bench bench_bit_keccak_circuit_prover -p circuit-benchmarks --features benches  -- --nocapture
 

README.md

@@ -11,13 +11,11 @@ To run the same tests as the CI, please use: `make test-all`.
 
 ## Running benchmarks
 
-There are currently two benchmarks to run in the workspace in regards to the circuits.
-Both use the `DEGREE` env var to specify the degree of the `K` parameter that you want 
+There are currently several benchmarks to run in the workspace in regards to the circuits.
+All use the `DEGREE` env var to specify the degree of the `K` parameter that you want 
 to use for your circuit in the bench process.
--   Keccak Permutation Round bench -> `DEGREE=16 make keccak_round_bench`
+-   Keccak Circuit prover benches. -> `DEGREE=16 make packed_multi_keccak_bench`
 -   EVM Circuit prover benches. -> `DEGREE=18 make evm_bench`.
--   State Circuit prover benches. -> `DEGREE=18, MEMORY_ADDRESS_MAX=3000, STACK_ADDRESS_MAX=1500 make state_bench`
-For state circuit, `MEMORY_ADDRESS_MAX` and `STACK_ADDRESS_MAX` do not need to be specified. 
-By default will take values: 2000 and 1300.
+-   State Circuit prover benches. -> `DEGREE=18 make state_bench`
 
-You can also run both benchmarks by running: `make circuit_benches DEGREE=18`.
+You can also run all benchmarks by running: `make circuit_benches DEGREE=18`.

circuit-benchmarks/src/bit_keccak.rs

@@ -28,8 +28,8 @@ mod tests {
             .parse()
             .expect("Cannot parse DEGREE env var as u32");
 
-        // Create the circuit
-        let mut circuit = KeccakBitCircuit::new(2usize.pow(degree));
+        // Create the circuit. Leave last dozens of rows for blinding.
+        let mut circuit = KeccakBitCircuit::new(2usize.pow(degree) - 64);
 
         // Use the complete circuit
         let inputs = vec![(0u8..135).collect::<Vec<_>>(); circuit.capacity()];
@@ -55,7 +55,7 @@ mod tests {
         let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]);
 
         // Bench proof generation time
-        let proof_message = format!("Bit Keccak Proof generation with {} rows", degree);
+        let proof_message = format!("Bit Keccak Proof generation with degree = {}", degree);
         let start2 = start_timer!(|| proof_message);
         create_proof::<
             KZGCommitmentScheme<Bn256>,

circuit-benchmarks/src/evm_circuit.rs

@@ -101,7 +101,7 @@ mod evm_circ_benches {
         let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]);
 
         // Bench proof generation time
-        let proof_message = format!("EVM circuit Proof generation with {} rows", degree);
+        let proof_message = format!("EVM circuit Proof generation with degree = {}", degree);
         let start2 = start_timer!(|| proof_message);
         create_proof::<
             KZGCommitmentScheme<Bn256>,

circuit-benchmarks/src/packed_keccak.rs

@@ -27,8 +27,8 @@ mod tests {
             .parse()
             .expect("Cannot parse DEGREE env var as u32");
 
-        // Create the circuit
-        let mut circuit = KeccakPackedCircuit::new(2usize.pow(degree));
+        // Create the circuit. Leave last dozens of rows for blinding.
+        let mut circuit = KeccakPackedCircuit::new(2usize.pow(degree) - 64);
 
         // Use the complete circuit
         let inputs = vec![(0u8..135).collect::<Vec<_>>(); circuit.capacity()];
@@ -54,7 +54,7 @@ mod tests {
         let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]);
 
         // Bench proof generation time
-        let proof_message = format!("Packed Keccak Proof generation with {} rows", degree);
+        let proof_message = format!("Packed Keccak Proof generation with degree = {}", degree);
         let start2 = start_timer!(|| proof_message);
         create_proof::<
             KZGCommitmentScheme<Bn256>,

circuit-benchmarks/src/packed_multi_keccak.rs

@@ -27,8 +27,8 @@ mod tests {
             .parse()
             .expect("Cannot parse DEGREE env var as u32");
 
-        // Create the circuit
-        let mut circuit = KeccakPackedCircuit::new(2usize.pow(degree));
+        // Create the circuit. Leave last dozens of rows for blinding.
+        let mut circuit = KeccakPackedCircuit::new(2usize.pow(degree) - 64);
 
         // Use the complete circuit
         let inputs = vec![(0u8..135).collect::<Vec<_>>(); circuit.capacity()];
@@ -54,7 +54,10 @@ mod tests {
         let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]);
 
         // Bench proof generation time
-        let proof_message = format!("Packed Multi-Keccak Proof generation with {} rows", degree);
+        let proof_message = format!(
+            "Packed Multi-Keccak Proof generation with degree = {}",
+            degree
+        );
         let start2 = start_timer!(|| proof_message);
         create_proof::<
             KZGCommitmentScheme<Bn256>,

circuit-benchmarks/src/state_circuit.rs

@@ -53,7 +53,7 @@ mod tests {
         let instances: Vec<&[Fr]> = instance.iter().map(|v| v.as_slice()).collect();
 
         // Bench proof generation time
-        let proof_message = format!("State Circuit Proof generation with {} rows", degree);
+        let proof_message = format!("State Circuit Proof generation with degree = {}", degree);
         let start2 = start_timer!(|| proof_message);
         create_proof::<
             KZGCommitmentScheme<Bn256>,

circuit-benchmarks/src/super_circuit.rs

@@ -89,7 +89,7 @@ mod tests {
         let mut transcript = Blake2bWrite::<_, G1Affine, Challenge255<_>>::init(vec![]);
 
         // Bench proof generation time
-        let proof_message = format!("SuperCircuit Proof generation with {} rows", degree);
+        let proof_message = format!("SuperCircuit Proof generation with degree = {}", degree);
         let start2 = start_timer!(|| proof_message);
         create_proof::<
             KZGCommitmentScheme<Bn256>,

circuit-benchmarks/src/tx_circuit.rs

@@ -78,7 +78,10 @@ mod tests {
         let instance_slices: Vec<&[Fr]> = instance.iter().map(|v| &v[..]).collect();
 
         // Bench proof generation time
-        let proof_message = format!("Packed Multi-Keccak Proof generation with {} rows", DEGREE);
+        let proof_message = format!(
+            "Packed Multi-Keccak Proof generation with degree = {}",
+            DEGREE
+        );
         let start2 = start_timer!(|| proof_message);
         create_proof::<
             KZGCommitmentScheme<Bn256>,

zkevm-circuits/src/keccak_circuit/keccak_packed.rs

@@ -9,6 +9,7 @@ use crate::keccak_circuit::util::{
     NUM_BITS_PER_BYTE, NUM_BITS_PER_WORD, NUM_WORDS_TO_ABSORB, NUM_WORDS_TO_SQUEEZE, RATE,
     RATE_IN_BITS, RHO_MATRIX,
 };
+use crate::table::KeccakTable;
 use crate::{evm_circuit::util::constraint_builder::BaseConstraintBuilder, util::Expr};
 use eth_types::Field;
 use gadgets::util::{and, select, sum};
@@ -126,10 +127,8 @@ pub struct KeccakPackedConfig<F> {
     q_round_last: Column<Fixed>,
     q_padding: Column<Fixed>,
     q_padding_last: Column<Fixed>,
-    is_final: Column<Advice>,
-    length: Column<Advice>,
-    data_rlc: Column<Advice>,
-    hash_rlc: Column<Advice>,
+    /// The columns for other circuits to lookup Keccak hash results
+    pub keccak_table: KeccakTable,
     state: [Column<Advice>; KECCAK_WIDTH],
     cell_values: Vec<Column<Advice>>,
     absorb_from: Column<Advice>,
@@ -178,7 +177,7 @@ impl<F: Field> Circuit<F> for KeccakPackedCircuit<F> {
         mut layouter: impl Layouter<F>,
     ) -> Result<(), Error> {
         config.load(&mut layouter)?;
-        config.assign(layouter, self.size, &self.witness)?;
+        config.assign(&mut layouter, &self.witness)?;
         Ok(())
     }
 }
@@ -203,11 +202,6 @@ impl<F: Field> KeccakPackedCircuit<F> {
     pub fn generate_witness(&mut self, inputs: &[Vec<u8>]) {
         self.witness = multi_keccak(inputs, KeccakPackedCircuit::r());
     }
-
-    /// Sets the witness using the witness data directly
-    fn set_witness(&mut self, witness: &[KeccakRow<F>]) {
-        self.witness = witness.to_vec();
-    }
 }
 
 /// Splits a word into parts
@@ -477,10 +471,13 @@ impl<F: Field> KeccakPackedConfig<F> {
         let q_round_last = meta.fixed_column();
         let q_padding = meta.fixed_column();
         let q_padding_last = meta.fixed_column();
-        let is_final = meta.advice_column();
-        let length = meta.advice_column();
-        let data_rlc = meta.advice_column();
-        let hash_rlc = meta.advice_column();
+
+        let keccak_table = KeccakTable::construct(meta);
+        let is_final = keccak_table.is_enabled;
+        let length = keccak_table.input_len;
+        let data_rlc = keccak_table.input_rlc;
+        let hash_rlc = keccak_table.output_rlc;
+
         let state = array_init::array_init(|_| meta.advice_column());
         let absorb_from = meta.advice_column();
         let absorb_data = meta.advice_column();
@@ -1201,10 +1198,7 @@ impl<F: Field> KeccakPackedConfig<F> {
             q_round_last,
             q_padding,
             q_padding_last,
-            is_final,
-            length,
-            data_rlc,
-            hash_rlc,
+            keccak_table,
             state,
             cell_values,
             absorb_from,
@@ -1222,10 +1216,20 @@ impl<F: Field> KeccakPackedConfig<F> {
         }
     }
 
+    /// Sets the witness using the data to be hashed
+    pub fn assign_from_witness(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        inputs: &[Vec<u8>],
+        r: F,
+    ) -> Result<(), Error> {
+        let witness = multi_keccak(inputs, r);
+        self.assign(layouter, &witness)
+    }
+
     pub(crate) fn assign(
         &self,
-        mut layouter: impl Layouter<F>,
-        _size: usize,
+        layouter: &mut impl Layouter<F>,
         witness: &[KeccakRow<F>],
     ) -> Result<(), Error> {
         layouter.assign_region(
@@ -1283,10 +1287,6 @@ impl<F: Field> KeccakPackedConfig<F> {
                 self.squeeze_packed,
                 row.squeeze_data.packed,
             ),
-            ("is_final", self.is_final, F::from(row.is_final)),
-            ("length", self.length, F::from(row.length as u64)),
-            ("data_rlc", self.data_rlc, row.data_rlc),
-            ("hash_rlc", self.hash_rlc, row.hash_rlc),
         ] {
             region.assign_advice(
                 || format!("assign {} {}", name, offset),
@@ -1296,6 +1296,17 @@ impl<F: Field> KeccakPackedConfig<F> {
             )?;
         }
 
+        self.keccak_table.assign_row(
+            region,
+            offset,
+            [
+                F::from(row.is_final),
+                row.data_rlc,
+                F::from(row.length as u64),
+                row.hash_rlc,
+            ],
+        )?;
+
         // State words
         for (idx, (word, column)) in row.state.iter().zip(self.state.iter()).enumerate() {
             region.assign_advice(

zkevm-circuits/src/keccak_circuit/keccak_packed_multi.rs

@@ -8,6 +8,7 @@ use crate::keccak_circuit::util::{
     rotate, scatter, target_part_sizes, to_bytes, unpack, BIT_SIZE, NUM_WORDS_TO_ABSORB,
     NUM_WORDS_TO_SQUEEZE, RATE, RATE_IN_BITS, RHO_MATRIX, ROUND_CST,
 };
+use crate::table::KeccakTable;
 use crate::{evm_circuit::util::constraint_builder::BaseConstraintBuilder, util::Expr};
 use eth_types::Field;
 use gadgets::util::{and, select, sum};
@@ -321,10 +322,8 @@ pub struct KeccakPackedConfig<F> {
     q_round_last: Column<Fixed>,
     q_padding: Column<Fixed>,
     q_padding_last: Column<Fixed>,
-    is_final: Column<Advice>,
-    length: Column<Advice>,
-    data_rlc: Column<Advice>,
-    hash_rlc: Column<Advice>,
+    /// The columns for other circuits to lookup Keccak hash results
+    pub keccak_table: KeccakTable,
     cell_manager: CellManager<F>,
     round_cst: Column<Fixed>,
     normalize_3: [TableColumn; 2],
@@ -367,7 +366,7 @@ impl<F: Field> Circuit<F> for KeccakPackedCircuit<F> {
         mut layouter: impl Layouter<F>,
     ) -> Result<(), Error> {
         config.load(&mut layouter)?;
-        config.assign(layouter, self.size, &self.witness)?;
+        config.assign(&mut layouter, &self.witness)?;
         Ok(())
     }
 }
@@ -392,11 +391,6 @@ impl<F: Field> KeccakPackedCircuit<F> {
     pub fn generate_witness(&mut self, inputs: &[Vec<u8>]) {
         self.witness = multi_keccak(inputs, KeccakPackedCircuit::r());
     }
-
-    /// Sets the witness using the witness data directly
-    fn set_witness(&mut self, witness: &[KeccakRow<F>]) {
-        self.witness = witness.to_vec();
-    }
 }
 
 /// Recombines parts back together
@@ -792,11 +786,14 @@ impl<F: Field> KeccakPackedConfig<F> {
         let q_round_last = meta.fixed_column();
         let q_padding = meta.fixed_column();
         let q_padding_last = meta.fixed_column();
-        let is_final = meta.advice_column();
         let round_cst = meta.fixed_column();
-        let length = meta.advice_column();
-        let data_rlc = meta.advice_column();
-        let hash_rlc = meta.advice_column();
+
+        let keccak_table = KeccakTable::construct(meta);
+        let is_final = keccak_table.is_enabled;
+        let length = keccak_table.input_len;
+        let data_rlc = keccak_table.input_rlc;
+        let hash_rlc = keccak_table.output_rlc;
+
         let normalize_3 = array_init::array_init(|_| meta.lookup_table_column());
         let normalize_4 = array_init::array_init(|_| meta.lookup_table_column());
         let normalize_6 = array_init::array_init(|_| meta.lookup_table_column());
@@ -1529,10 +1526,7 @@ impl<F: Field> KeccakPackedConfig<F> {
             q_round_last,
             q_padding,
             q_padding_last,
-            is_final,
-            length,
-            data_rlc,
-            hash_rlc,
+            keccak_table,
             cell_manager,
             round_cst,
             normalize_3,
@@ -1544,10 +1538,20 @@ impl<F: Field> KeccakPackedConfig<F> {
         }
     }
 
+    /// Sets the witness using the data to be hashed
+    pub fn assign_from_witness(
+        &self,
+        layouter: &mut impl Layouter<F>,
+        inputs: &[Vec<u8>],
+        r: F,
+    ) -> Result<(), Error> {
+        let witness = multi_keccak(inputs, r);
+        self.assign(layouter, &witness)
+    }
+
     pub(crate) fn assign(
         &self,
-        mut layouter: impl Layouter<F>,
-        _size: usize,
+        layouter: &mut impl Layouter<F>,
         witness: &[KeccakRow<F>],
     ) -> Result<(), Error> {
         layouter.assign_region(
@@ -1589,20 +1593,16 @@ impl<F: Field> KeccakPackedConfig<F> {
             )?;
         }
 
-        // Keccak data
-        for (name, column, value) in &[
-            ("is_final", self.is_final, F::from(row.is_final)),
-            ("length", self.length, F::from(row.length as u64)),
-            ("data_rlc", self.data_rlc, row.data_rlc),
-            ("hash_rlc", self.hash_rlc, row.hash_rlc),
-        ] {
-            region.assign_advice(
-                || format!("assign {} {}", name, offset),
-                *column,
-                offset,
-                || Value::known(*value),
-            )?;
-        }
+        self.keccak_table.assign_row(
+            region,
+            offset,
+            [
+                F::from(row.is_final),
+                row.data_rlc,
+                F::from(row.length as u64),
+                row.hash_rlc,
+            ],
+        )?;
 
         // Cell values
         for (idx, (bit, column)) in row