From 1cc38bb032c2e3d11f0faa1b798462e2630e3d6e Mon Sep 17 00:00:00 2001 From: wborgeaud Date: Thu, 9 Jun 2022 22:31:33 +0200 Subject: [PATCH 1/4] Add Keccak input registers --- evm/src/keccak/keccak_stark.rs | 45 +++++++++++++++++++++++++--------- evm/src/keccak/registers.rs | 28 +++++++++++++++++++-- 2 files changed, 60 insertions(+), 13 deletions(-) diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index cbc2408b..1af02d7f 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -17,7 +17,7 @@ use crate::keccak::logic::{ }; use crate::keccak::registers::{ reg_a, reg_a_prime, reg_a_prime_prime, reg_a_prime_prime_0_0_bit, reg_a_prime_prime_prime, - reg_b, reg_c, reg_c_partial, reg_step, NUM_REGISTERS, + reg_b, reg_c, reg_c_partial, reg_input_limb, reg_step, NUM_REGISTERS, }; use crate::keccak::round_flags::{eval_round_flags, eval_round_flags_recursively}; use crate::stark::Stark; @@ -65,6 +65,7 @@ impl, const D: usize> KeccakStark { fn generate_trace_rows_for_perm(&self, input: [u64; INPUT_LIMBS]) -> Vec<[F; NUM_REGISTERS]> { let mut rows = vec![[F::ZERO; NUM_REGISTERS]; NUM_ROUNDS]; + self.copy_input(input, &mut rows[0]); for x in 0..5 { for y in 0..5 { let input_xy = input[x * 5 + y]; @@ -76,6 +77,7 @@ impl, const D: usize> KeccakStark { self.generate_trace_row_for_round(&mut rows[0], 0); for round in 1..24 { + self.copy_input(input, &mut rows[round]); self.copy_output_to_input(rows[round - 1], &mut rows[round]); self.generate_trace_row_for_round(&mut rows[round], round); } @@ -188,6 +190,14 @@ impl, const D: usize> KeccakStark { row[out_reg_hi] = F::from_canonical_u64(row[in_reg_hi].to_canonical_u64() ^ rc_hi); } + fn copy_input(&self, input: [u64; INPUT_LIMBS], row: &mut [F; NUM_REGISTERS]) { + for i in 0..INPUT_LIMBS { + let (low, high) = (input[i] as u32, input[i] >> 32); + row[reg_input_limb(2 * i)] = F::from_canonical_u32(low); + row[reg_input_limb(2 * i + 1)] = F::from_canonical_u64(high); + } + } + pub fn generate_trace(&self, inputs: Vec<[u64; INPUT_LIMBS]>) -> Vec> { let mut timing = TimingTree::new("generate trace", log::Level::Debug); @@ -223,6 +233,16 @@ impl, const D: usize> Stark for KeccakStark, const D: usize> Stark for KeccakStark>(); let mut keccak_input: [[u64; 5]; 5] = [ input[0..5].try_into().unwrap(), diff --git a/evm/src/keccak/registers.rs b/evm/src/keccak/registers.rs index 3a891828..9eb265f2 100644 --- a/evm/src/keccak/registers.rs +++ b/evm/src/keccak/registers.rs @@ -1,4 +1,4 @@ -use crate::keccak::keccak_stark::NUM_ROUNDS; +use crate::keccak::keccak_stark::{INPUT_LIMBS, NUM_ROUNDS}; /// A register which is set to 1 if we are in the `i`th round, otherwise 0. pub(crate) const fn reg_step(i: usize) -> usize { @@ -6,6 +6,30 @@ pub(crate) const fn reg_step(i: usize) -> usize { i } +/// Registers to hold permutation inputs. +/// `reg_input_limb(2*i) -> input[i] as u32` +/// `reg_input_limb(2*i+1) -> input[i] >> 32` +pub(crate) const fn reg_input_limb(i: usize) -> usize { + debug_assert!(i < 2 * INPUT_LIMBS); + NUM_ROUNDS + i +} + +/// Registers to hold permutation outputs. +/// `reg_output_limb(2*i) -> output[i] as u32` +/// `reg_output_limb(2*i+1) -> output[i] >> 32` +#[allow(dead_code)] // TODO: Remove once it is used. +pub(crate) const fn reg_output_limb(i: usize) -> usize { + debug_assert!(i < 2 * INPUT_LIMBS); + let ii = i / 2; + let x = ii / 5; + let y = ii % 5; + if i % 2 == 0 { + reg_a_prime_prime_prime(x, y) + } else { + reg_a_prime_prime_prime(x, y) + 1 + } +} + const R: [[u8; 5]; 5] = [ [0, 36, 3, 41, 18], [1, 44, 10, 45, 2], @@ -14,7 +38,7 @@ const R: [[u8; 5]; 5] = [ [27, 20, 39, 8, 14], ]; -const START_A: usize = NUM_ROUNDS; +const START_A: usize = NUM_ROUNDS + 2 * INPUT_LIMBS; pub(crate) const fn reg_a(x: usize, y: usize, z: usize) -> usize { debug_assert!(x < 5); debug_assert!(y < 5); From 8af99cba5d5dffd0f06039f423f58f060c84c97e Mon Sep 17 00:00:00 2001 From: wborgeaud Date: Thu, 9 Jun 2022 22:42:48 +0200 Subject: [PATCH 2/4] Progress --- evm/src/keccak/keccak_stark.rs | 9 ++++++++- evm/src/prover.rs | 8 +++++++- 2 files changed, 15 insertions(+), 2 deletions(-) diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index 1af02d7f..56388626 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -233,14 +233,21 @@ impl, const D: usize> Stark for KeccakStark( &ctl_vars, &mut consumer, ); - consumer.accumulators() + let accs = consumer.accumulators(); + for a in &accs { + if !a.is_zero() { + dbg!(i); + } + } + accs }) .collect::>(); From 30abe19e4264a6a7cb95fa643fc8765876dc73ec Mon Sep 17 00:00:00 2001 From: wborgeaud Date: Fri, 10 Jun 2022 00:49:22 +0200 Subject: [PATCH 3/4] Fix --- evm/src/all_stark.rs | 4 ++-- evm/src/keccak/keccak_stark.rs | 21 ++++++++++++++------- evm/src/prover.rs | 8 +------- 3 files changed, 17 insertions(+), 16 deletions(-) diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs index eb1d7087..5b61e663 100644 --- a/evm/src/all_stark.rs +++ b/evm/src/all_stark.rs @@ -87,11 +87,11 @@ mod tests { let keccak_stark = KeccakStark:: { f: Default::default(), }; - let keccak_rows = (2 * NUM_ROUNDS + 1).next_power_of_two(); + let keccak_rows = (6 * NUM_ROUNDS + 1).next_power_of_two(); let keccak_looked_col = 3; let mut rng = ChaCha8Rng::seed_from_u64(0x6feb51b7ec230f25); - let num_inputs = 2; + let num_inputs = 6; let keccak_inputs = (0..num_inputs) .map(|_| [0u64; INPUT_LIMBS].map(|_| rng.gen())) .collect_vec(); diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index ac56462d..a834ad85 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -51,14 +51,11 @@ impl, const D: usize> KeccakStark { rows.extend(self.generate_trace_rows_for_perm(*input)); } - // Pad rows to power of two. - for i in rows.len()..num_rows { - let mut row = [F::ZERO; NUM_REGISTERS]; - self.copy_output_to_input(rows[i - 1], &mut row); - self.generate_trace_row_for_round(&mut row, i % NUM_ROUNDS); - rows.push(row); + let pad_rows = self.generate_trace_rows_for_perm([0; INPUT_LIMBS]); + while rows.len() < num_rows { + rows.extend(&pad_rows); } - + rows.drain(num_rows..); rows } @@ -400,6 +397,16 @@ impl, const D: usize> Stark for KeccakStark>(); + let expected_input_limb = reduce_with_powers_ext_circuit(builder, &bits, two); + let is_first_round = vars.local_values[reg_step(0)]; + let diff = builder.sub_extension(local_input_limb, expected_input_limb); + let constraint = builder.mul_extension(is_first_round, diff); + yield_constr.constraint(builder, constraint); } // C_partial[x] = xor(A[x, 0], A[x, 1], A[x, 2]) diff --git a/evm/src/prover.rs b/evm/src/prover.rs index b83e083b..c0cc387a 100644 --- a/evm/src/prover.rs +++ b/evm/src/prover.rs @@ -523,13 +523,7 @@ fn check_constraints<'a, F, C, S, const D: usize>( &ctl_vars, &mut consumer, ); - let accs = consumer.accumulators(); - for a in &accs { - if !a.is_zero() { - dbg!(i); - } - } - accs + consumer.accumulators() }) .collect::>(); From 8bd6bebdc7f7df38f06289947a54ecdba5f18651 Mon Sep 17 00:00:00 2001 From: wborgeaud Date: Fri, 10 Jun 2022 19:30:51 +0200 Subject: [PATCH 4/4] INPUT_LIMBS -> NUM_INPUTS --- evm/src/all_stark.rs | 8 ++++---- evm/src/keccak/keccak_stark.rs | 26 +++++++++++++------------- evm/src/keccak/registers.rs | 8 ++++---- 3 files changed, 21 insertions(+), 21 deletions(-) diff --git a/evm/src/all_stark.rs b/evm/src/all_stark.rs index 5b61e663..3e9cc58c 100644 --- a/evm/src/all_stark.rs +++ b/evm/src/all_stark.rs @@ -64,7 +64,7 @@ mod tests { use crate::cpu; use crate::cpu::cpu_stark::CpuStark; use crate::cross_table_lookup::CrossTableLookup; - use crate::keccak::keccak_stark::{KeccakStark, INPUT_LIMBS, NUM_ROUNDS}; + use crate::keccak::keccak_stark::{KeccakStark, NUM_INPUTS, NUM_ROUNDS}; use crate::proof::AllProof; use crate::prover::prove; use crate::recursive_verifier::{ @@ -87,13 +87,13 @@ mod tests { let keccak_stark = KeccakStark:: { f: Default::default(), }; - let keccak_rows = (6 * NUM_ROUNDS + 1).next_power_of_two(); + let keccak_rows = (2 * NUM_ROUNDS + 1).next_power_of_two(); let keccak_looked_col = 3; let mut rng = ChaCha8Rng::seed_from_u64(0x6feb51b7ec230f25); - let num_inputs = 6; + let num_inputs = 2; let keccak_inputs = (0..num_inputs) - .map(|_| [0u64; INPUT_LIMBS].map(|_| rng.gen())) + .map(|_| [0u64; NUM_INPUTS].map(|_| rng.gen())) .collect_vec(); let keccak_trace = keccak_stark.generate_trace(keccak_inputs); let column_to_copy: Vec<_> = keccak_trace[keccak_looked_col].values[..].into(); diff --git a/evm/src/keccak/keccak_stark.rs b/evm/src/keccak/keccak_stark.rs index a834ad85..d9fdd06d 100644 --- a/evm/src/keccak/keccak_stark.rs +++ b/evm/src/keccak/keccak_stark.rs @@ -27,8 +27,8 @@ use crate::vars::{StarkEvaluationTargets, StarkEvaluationVars}; /// Number of rounds in a Keccak permutation. pub(crate) const NUM_ROUNDS: usize = 24; -/// Number of 64-bit limbs in a preimage of the Keccak permutation. -pub(crate) const INPUT_LIMBS: usize = 25; +/// Number of 64-bit elements in the Keccak permutation input. +pub(crate) const NUM_INPUTS: usize = 25; pub(crate) const NUM_PUBLIC_INPUTS: usize = 0; @@ -42,7 +42,7 @@ impl, const D: usize> KeccakStark { /// in our lookup arguments, as those are computed after transposing to column-wise form. pub(crate) fn generate_trace_rows( &self, - inputs: Vec<[u64; INPUT_LIMBS]>, + inputs: Vec<[u64; NUM_INPUTS]>, ) -> Vec<[F; NUM_REGISTERS]> { let num_rows = (inputs.len() * NUM_ROUNDS).next_power_of_two(); info!("{} rows", num_rows); @@ -51,7 +51,7 @@ impl, const D: usize> KeccakStark { rows.extend(self.generate_trace_rows_for_perm(*input)); } - let pad_rows = self.generate_trace_rows_for_perm([0; INPUT_LIMBS]); + let pad_rows = self.generate_trace_rows_for_perm([0; NUM_INPUTS]); while rows.len() < num_rows { rows.extend(&pad_rows); } @@ -59,7 +59,7 @@ impl, const D: usize> KeccakStark { rows } - fn generate_trace_rows_for_perm(&self, input: [u64; INPUT_LIMBS]) -> Vec<[F; NUM_REGISTERS]> { + fn generate_trace_rows_for_perm(&self, input: [u64; NUM_INPUTS]) -> Vec<[F; NUM_REGISTERS]> { let mut rows = vec![[F::ZERO; NUM_REGISTERS]; NUM_ROUNDS]; self.copy_input(input, &mut rows[0]); @@ -187,15 +187,15 @@ impl, const D: usize> KeccakStark { row[out_reg_hi] = F::from_canonical_u64(row[in_reg_hi].to_canonical_u64() ^ rc_hi); } - fn copy_input(&self, input: [u64; INPUT_LIMBS], row: &mut [F; NUM_REGISTERS]) { - for i in 0..INPUT_LIMBS { + fn copy_input(&self, input: [u64; NUM_INPUTS], row: &mut [F; NUM_REGISTERS]) { + for i in 0..NUM_INPUTS { let (low, high) = (input[i] as u32, input[i] >> 32); row[reg_input_limb(2 * i)] = F::from_canonical_u32(low); row[reg_input_limb(2 * i + 1)] = F::from_canonical_u64(high); } } - pub fn generate_trace(&self, inputs: Vec<[u64; INPUT_LIMBS]>) -> Vec> { + pub fn generate_trace(&self, inputs: Vec<[u64; NUM_INPUTS]>) -> Vec> { let mut timing = TimingTree::new("generate trace", log::Level::Debug); // Generate the witness, except for permuted columns in the lookup argument. @@ -230,7 +230,7 @@ impl, const D: usize> Stark for KeccakStark, const D: usize> Stark for KeccakStark Result<()> { - let input: [u64; INPUT_LIMBS] = rand::random(); + let input: [u64; NUM_INPUTS] = rand::random(); const D: usize = 2; type C = PoseidonGoldilocksConfig; @@ -609,7 +609,7 @@ mod tests { let rows = stark.generate_trace_rows(vec![input.try_into().unwrap()]); let last_row = rows[NUM_ROUNDS - 1]; let base = F::from_canonical_u64(1 << 32); - let output = (0..INPUT_LIMBS) + let output = (0..NUM_INPUTS) .map(|i| last_row[reg_output_limb(2 * i)] + base * last_row[reg_output_limb(2 * i + 1)]) .collect::>(); diff --git a/evm/src/keccak/registers.rs b/evm/src/keccak/registers.rs index 9eb265f2..2dc019e3 100644 --- a/evm/src/keccak/registers.rs +++ b/evm/src/keccak/registers.rs @@ -1,4 +1,4 @@ -use crate::keccak::keccak_stark::{INPUT_LIMBS, NUM_ROUNDS}; +use crate::keccak::keccak_stark::{NUM_INPUTS, NUM_ROUNDS}; /// A register which is set to 1 if we are in the `i`th round, otherwise 0. pub(crate) const fn reg_step(i: usize) -> usize { @@ -10,7 +10,7 @@ pub(crate) const fn reg_step(i: usize) -> usize { /// `reg_input_limb(2*i) -> input[i] as u32` /// `reg_input_limb(2*i+1) -> input[i] >> 32` pub(crate) const fn reg_input_limb(i: usize) -> usize { - debug_assert!(i < 2 * INPUT_LIMBS); + debug_assert!(i < 2 * NUM_INPUTS); NUM_ROUNDS + i } @@ -19,7 +19,7 @@ pub(crate) const fn reg_input_limb(i: usize) -> usize { /// `reg_output_limb(2*i+1) -> output[i] >> 32` #[allow(dead_code)] // TODO: Remove once it is used. pub(crate) const fn reg_output_limb(i: usize) -> usize { - debug_assert!(i < 2 * INPUT_LIMBS); + debug_assert!(i < 2 * NUM_INPUTS); let ii = i / 2; let x = ii / 5; let y = ii % 5; @@ -38,7 +38,7 @@ const R: [[u8; 5]; 5] = [ [27, 20, 39, 8, 14], ]; -const START_A: usize = NUM_ROUNDS + 2 * INPUT_LIMBS; +const START_A: usize = NUM_ROUNDS + 2 * NUM_INPUTS; pub(crate) const fn reg_a(x: usize, y: usize, z: usize) -> usize { debug_assert!(x < 5); debug_assert!(y < 5);