From 504fd62bdaa5aff2893496872e3701f20e452508 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Fri, 26 Jul 2024 16:50:08 -0700 Subject: [PATCH 01/18] external round constraints --- .../core-v2/src/chips/poseidon2_skinny/air.rs | 105 +++++++++++++++++- 1 file changed, 102 insertions(+), 3 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs index 34081801c6..71d1a3a1ce 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs @@ -3,14 +3,15 @@ use std::borrow::Borrow; -use p3_air::{Air, BaseAir, PairBuilder}; +use p3_air::{Air, AirBuilder, BaseAir, PairBuilder}; +use p3_field::AbstractField; use p3_matrix::Matrix; use crate::builder::SP1RecursionAirBuilder; use super::columns::preprocessed::Poseidon2PreprocessedCols; -use super::columns::{NUM_POSEIDON2_DEGREE3_COLS, NUM_POSEIDON2_DEGREE9_COLS}; -use super::{Poseidon2SkinnyChip, WIDTH}; +use super::columns::{Poseidon2, NUM_POSEIDON2_DEGREE3_COLS, NUM_POSEIDON2_DEGREE9_COLS}; +use super::{external_linear_layer, Poseidon2SkinnyChip, WIDTH}; impl BaseAir for Poseidon2SkinnyChip { fn width(&self) -> usize { @@ -33,6 +34,7 @@ where let main = builder.main(); let prepr = builder.preprocessed(); let local_row = Self::convert::(main.row_slice(0)); + let next_row = Self::convert::(main.row_slice(1)); let prep_local = prepr.row_slice(0); let prep_local: &Poseidon2PreprocessedCols<_> = (*prep_local).borrow(); @@ -61,5 +63,102 @@ where prep_local.memory_preprocessed[i].write_mult, ) }); + + self.eval_external_round( + builder, + local_row, + next_row, + prep_local.round_counters_preprocessed.round_constants, + prep_local.round_counters_preprocessed.is_external_round, + ); + } +} + +impl Poseidon2SkinnyChip { + fn eval_external_round( + &self, + builder: &mut AB, + local_row: Box>, + next_row: Box>, + round_constants: [AB::Var; WIDTH], + is_external_row: AB::Var, + ) where + AB::Var: 'static, + { + let local_state = local_row.state_var(); + + // Add the round constants. + let add_rc: [AB::Expr; WIDTH] = + core::array::from_fn(|i| local_state[i].into() + round_constants[i]); + + // Apply the sboxes. + // See `populate_external_round` for why we don't have columns for the sbox output here. + let mut sbox_deg_7: [AB::Expr; WIDTH] = core::array::from_fn(|_| AB::Expr::zero()); + let mut sbox_deg_3: [AB::Expr; WIDTH] = core::array::from_fn(|_| AB::Expr::zero()); + for i in 0..WIDTH { + let calculated_sbox_deg_3 = add_rc[i].clone() * add_rc[i].clone() * add_rc[i].clone(); + + if let Some(external_sbox) = local_row.s_box_state() { + builder.assert_eq(external_sbox[i].into(), calculated_sbox_deg_3); + sbox_deg_3[i] = external_sbox[i].into(); + } else { + sbox_deg_3[i] = calculated_sbox_deg_3; + } + + sbox_deg_7[i] = sbox_deg_3[i].clone() * sbox_deg_3[i].clone() * add_rc[i].clone(); + } + + // Apply the linear layer. + let mut state = sbox_deg_7; + external_linear_layer(&mut state); + + let next_state = next_row.state_var(); + for i in 0..WIDTH { + builder + .when(is_external_row) + .assert_eq(next_state[i], state[i].clone()); + } } + + // fn eval_internal_rounds( + // &self, + // builder: &mut AB, + // perm_cols: &dyn Permutation, + // ) { + // let state = &perm_cols.internal_rounds_state(); + // let s0 = perm_cols.internal_rounds_s0(); + // let mut state: [AB::Expr; WIDTH] = core::array::from_fn(|i| state[i].into()); + // for r in 0..NUM_INTERNAL_ROUNDS { + // // Add the round constant. + // let round = r + NUM_EXTERNAL_ROUNDS / 2; + // let add_rc = if r == 0 { + // state[0].clone() + // } else { + // s0[r - 1].into() + // } + AB::Expr::from_wrapped_u32(RC_16_30_U32[round][0]); + + // let mut sbox_deg_3 = add_rc.clone() * add_rc.clone() * add_rc.clone(); + // if let Some(internal_sbox) = perm_cols.internal_rounds_sbox() { + // builder.assert_eq(internal_sbox[r], sbox_deg_3); + // sbox_deg_3 = internal_sbox[r].into(); + // } + + // // See `populate_internal_rounds` for why we don't have columns for the sbox output here. + // let sbox_deg_7 = sbox_deg_3.clone() * sbox_deg_3.clone() * add_rc.clone(); + + // // Apply the linear layer. + // // See `populate_internal_rounds` for why we don't have columns for the new state here. + // state[0] = sbox_deg_7.clone(); + // internal_linear_layer(&mut state); + + // if r < NUM_INTERNAL_ROUNDS - 1 { + // builder.assert_eq(s0[r], state[0].clone()); + // } + // } + + // let external_state = perm_cols.external_rounds_state()[NUM_EXTERNAL_ROUNDS / 2]; + // for i in 0..WIDTH { + // builder.assert_eq(external_state[i], state[i].clone()) + // } + // } } From f7e94da30859cfe5106a29c7f85c6d21f07fd4c1 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Mon, 29 Jul 2024 13:42:22 -0700 Subject: [PATCH 02/18] trace gen works --- .../core-v2/src/chips/poseidon2_skinny/air.rs | 27 ++++--- .../poseidon2_skinny/columns/preprocessed.rs | 2 +- .../src/chips/poseidon2_skinny/trace.rs | 76 +++++++++++-------- 3 files changed, 60 insertions(+), 45 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs index 71d1a3a1ce..1e913e1b26 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs @@ -32,9 +32,9 @@ where { fn eval(&self, builder: &mut AB) { let main = builder.main(); - let prepr = builder.preprocessed(); let local_row = Self::convert::(main.row_slice(0)); let next_row = Self::convert::(main.row_slice(1)); + let prepr = builder.preprocessed(); let prep_local = prepr.row_slice(0); let prep_local: &Poseidon2PreprocessedCols<_> = (*prep_local).borrow(); @@ -64,13 +64,13 @@ where ) }); - self.eval_external_round( - builder, - local_row, - next_row, - prep_local.round_counters_preprocessed.round_constants, - prep_local.round_counters_preprocessed.is_external_round, - ); + // self.eval_external_round( + // builder, + // local_row, + // next_row, + // prep_local.round_counters_preprocessed.round_constants, + // prep_local.round_counters_preprocessed.is_external_round, + // ); } } @@ -99,7 +99,9 @@ impl Poseidon2SkinnyChip { let calculated_sbox_deg_3 = add_rc[i].clone() * add_rc[i].clone() * add_rc[i].clone(); if let Some(external_sbox) = local_row.s_box_state() { - builder.assert_eq(external_sbox[i].into(), calculated_sbox_deg_3); + builder + .when(is_external_row) + .assert_eq(external_sbox[i].into(), calculated_sbox_deg_3); sbox_deg_3[i] = external_sbox[i].into(); } else { sbox_deg_3[i] = calculated_sbox_deg_3; @@ -114,9 +116,10 @@ impl Poseidon2SkinnyChip { let next_state = next_row.state_var(); for i in 0..WIDTH { - builder - .when(is_external_row) - .assert_eq(next_state[i], state[i].clone()); + // builder + // .when_transition() + // .when(is_external_row) + // .assert_eq(next_state[i], state[i].clone()); } } diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs b/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs index 2d5cf5ebcd..02f5e041a8 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs @@ -5,9 +5,9 @@ use crate::chips::{mem::MemoryAccessCols, poseidon2_skinny::WIDTH}; #[derive(AlignedBorrow, Clone, Copy, Debug)] #[repr(C)] pub struct RoundCountersPreprocessedCols { + pub is_input_round: T, pub is_external_round: T, pub is_internal_round: T, - pub is_first_round: T, pub round_constants: [T; WIDTH], } diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index 79b4479735..21625975f7 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -26,6 +26,8 @@ use super::{ const PREPROCESSED_POSEIDON2_WIDTH: usize = size_of::>(); +const INTERNAL_ROUND_IDX: usize = NUM_EXTERNAL_ROUNDS / 2 + 1; + impl MachineAir for Poseidon2SkinnyChip { type Record = ExecutionRecord; @@ -46,45 +48,49 @@ impl MachineAir for Poseidon2SkinnyChip let num_columns = as BaseAir>::width(self); for event in &input.poseidon2_skinny_events { - let mut row_add = vec![vec![F::zero(); num_columns]; NUM_EXTERNAL_ROUNDS + 2]; + // We have one row for input, one row for output, NUM_EXTERNAL_ROUNDS rows for the external rounds, + // and one row for all internal rounds. + let mut row_add = vec![vec![F::zero(); num_columns]; NUM_EXTERNAL_ROUNDS + 3]; // The first row should have event.input and [event.input[0].clone(); NUM_INTERNAL_ROUNDS-1] // in its state columns. The sbox_state will be modified in the computation of the // first row. { - let mut cols = self.convert_mut(&mut row_add[0]); - *cols.get_cols_mut().0 = event.input; + let (first_row, second_row) = &mut row_add[0..2].split_at_mut(1); + let mut input_cols = self.convert_mut(&mut first_row[0]); + *input_cols.get_cols_mut().0 = event.input; + + let mut next_cols = self.convert_mut(&mut second_row[0]); + let (next_state, _, _) = next_cols.get_cols_mut(); + *next_state = event.input; + external_linear_layer(next_state); } // For each external round, and once for all the internal rounds at the same time, apply // the corresponding operation. This will change the sbox state in row i, and the state // and internal_rounds_s0 variable in row i+1. - for i in 0..NUM_EXTERNAL_ROUNDS + 1 { - let (next_state_var, next_internal_rounds_s0) = { + for i in 1..NUM_EXTERNAL_ROUNDS + 2 { + let next_state_var = { let mut cols = self.convert_mut(&mut row_add[i]); let (state, internal_state_s0, mut sbox_state) = cols.get_cols_mut(); - let mut state = *state; - if i == 0 { - external_linear_layer(&mut state); - } - let (next_state_var, next_internal_rounds_s0) = if i != NUM_EXTERNAL_ROUNDS / 2 - { - ( - self.populate_external_round(&state, &mut sbox_state, i), - [state[0]; NUM_INTERNAL_ROUNDS - 1], - ) + + if i != INTERNAL_ROUND_IDX { + self.populate_external_round(state, &mut sbox_state, i - 1) } else { - self.populate_internal_rounds(&state, internal_state_s0, &mut sbox_state) - }; - (next_state_var, next_internal_rounds_s0) + // Populate the internal rounds. + self.populate_internal_rounds(state, internal_state_s0, &mut sbox_state) + } }; let mut next_cols = self.convert_mut(&mut row_add[i + 1]); - let (next_state, internal_state_s0, _) = next_cols.get_cols_mut(); + let (next_state, _, _) = next_cols.get_cols_mut(); *next_state = next_state_var; - *internal_state_s0 = next_internal_rounds_s0; - if i == NUM_EXTERNAL_ROUNDS { - debug_assert_eq!(next_state_var, event.output); - } + } + + // Check that the permutation is computed correctly. + { + let last_row = &row_add[NUM_EXTERNAL_ROUNDS + 2]; + let output_cols = Poseidon2SkinnyChip::::convert(last_row.as_slice()); + debug_assert_eq!(*output_cols.state_var(), event.output); } rows.extend(row_add.into_iter()); } @@ -147,16 +153,22 @@ impl MachineAir for Poseidon2SkinnyChip (*row).as_mut_slice().borrow_mut(); // Set the round-counter columns. + cols.round_counters_preprocessed.is_input_round = F::from_bool(i == 0); + let is_external_round = i != INTERNAL_ROUND_IDX && i != 0; cols.round_counters_preprocessed.is_external_round = - F::from_bool((i != NUM_EXTERNAL_ROUNDS / 2 + 1) && i > 1); + F::from_bool(is_external_round); cols.round_counters_preprocessed.is_internal_round = - F::from_bool(i == NUM_EXTERNAL_ROUNDS / 2 + 1); - cols.round_counters_preprocessed.is_first_round = F::from_bool(i == 0); + F::from_bool(INTERNAL_ROUND_IDX == i); + (0..WIDTH).for_each(|j| { - cols.round_counters_preprocessed.round_constants[j] = if i <= 1 { - F::zero() + cols.round_counters_preprocessed.round_constants[j] = if is_external_round { + let external_round_num = + if i < INTERNAL_ROUND_IDX { i - 1 } else { i - 2 }; + F::from_wrapped_u32(RC_16_30_U32[external_round_num][j]) + } else if i == INTERNAL_ROUND_IDX { + F::from_wrapped_u32(RC_16_30_U32[NUM_EXTERNAL_ROUNDS / 2 + j][0]) } else { - F::from_wrapped_u32(RC_16_30_U32[i - 2][j]) + F::zero() }; }); @@ -169,7 +181,7 @@ impl MachineAir for Poseidon2SkinnyChip read_mult: F::one(), write_mult: F::zero(), }); - } else if i == NUM_EXTERNAL_ROUNDS + 1 { + } else if i == NUM_EXTERNAL_ROUNDS + 2 { cols.memory_preprocessed = instruction.addrs.output.map(|addr| MemoryAccessCols { addr, @@ -246,7 +258,7 @@ impl Poseidon2SkinnyChip { state: &[F; WIDTH], internal_rounds_s0: &mut [F; NUM_INTERNAL_ROUNDS - 1], sbox: &mut Option<&mut [F; WIDTH]>, - ) -> ([F; WIDTH], [F; NUM_INTERNAL_ROUNDS - 1]) { + ) -> [F; WIDTH] { let mut new_state = *state; let mut sbox_deg_3: [F; max(WIDTH, NUM_INTERNAL_ROUNDS)] = [F::zero(); max(WIDTH, NUM_INTERNAL_ROUNDS)]; @@ -283,7 +295,7 @@ impl Poseidon2SkinnyChip { *sbox = sbox_deg_3; } - (ret_state, *internal_rounds_s0) + ret_state } } From 43c9e28e031077b64e2813906e0aa3448698de06 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Mon, 29 Jul 2024 13:54:56 -0700 Subject: [PATCH 03/18] changes to preprocess trace --- recursion/core-v2/src/chips/poseidon2_skinny/trace.rs | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index 21625975f7..af02bf1d66 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -141,12 +141,12 @@ impl MachineAir for Poseidon2SkinnyChip let mut rows = vec![ [F::zero(); PREPROCESSED_POSEIDON2_WIDTH]; - num_instructions * (NUM_EXTERNAL_ROUNDS + 2) + num_instructions * (NUM_EXTERNAL_ROUNDS + 3) ]; - // Iterate over the instructions and take NUM_EXTERNAL_ROUNDS + 2 rows for each instruction. + // Iterate over the instructions and take NUM_EXTERNAL_ROUNDS + 3 rows for each instruction. instructions - .zip_eq(&rows.iter_mut().chunks(NUM_EXTERNAL_ROUNDS + 2)) + .zip_eq(&rows.iter_mut().chunks(NUM_EXTERNAL_ROUNDS + 3)) .for_each(|(instruction, row_add)| { row_add.into_iter().enumerate().for_each(|(i, row)| { let cols: &mut Poseidon2PreprocessedCols<_> = From 85a1e97ae8a2400fc790e679012052b3a5258ed4 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Mon, 29 Jul 2024 14:55:16 -0700 Subject: [PATCH 04/18] input round and external rounds working for degree 9 --- .../core-v2/src/chips/poseidon2_skinny/air.rs | 56 ++++++++++++++----- .../core-v2/src/chips/poseidon2_skinny/mod.rs | 16 +++--- .../src/chips/poseidon2_skinny/trace.rs | 16 ++++-- 3 files changed, 62 insertions(+), 26 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs index 1e913e1b26..4fcee9a7bf 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs @@ -1,6 +1,7 @@ //! The air module contains the AIR constraints for the poseidon2 chip. //! At the moment, we're only including memory constraints to test the new memory argument. +use std::array; use std::borrow::Borrow; use p3_air::{Air, AirBuilder, BaseAir, PairBuilder}; @@ -64,22 +65,51 @@ where ) }); - // self.eval_external_round( - // builder, - // local_row, - // next_row, - // prep_local.round_counters_preprocessed.round_constants, - // prep_local.round_counters_preprocessed.is_external_round, - // ); + self.eval_input_round( + builder, + local_row.as_ref(), + next_row.as_ref(), + prep_local.round_counters_preprocessed.is_input_round, + ); + + self.eval_external_round( + builder, + local_row.as_ref(), + next_row.as_ref(), + prep_local.round_counters_preprocessed.round_constants, + prep_local.round_counters_preprocessed.is_external_round, + ); } } impl Poseidon2SkinnyChip { + fn eval_input_round( + &self, + builder: &mut AB, + local_row: &dyn Poseidon2, + next_row: &dyn Poseidon2, + is_input_row: AB::Var, + ) where + AB::Var: 'static, + { + let mut local_state: [AB::Expr; WIDTH] = + array::from_fn(|i| local_row.state_var()[i].into()); + + external_linear_layer(&mut local_state); + + let next_state = next_row.state_var(); + for i in 0..WIDTH { + builder + .when(is_input_row) + .assert_eq(next_state[i], local_state[i].clone()); + } + } + fn eval_external_round( &self, builder: &mut AB, - local_row: Box>, - next_row: Box>, + local_row: &dyn Poseidon2, + next_row: &dyn Poseidon2, round_constants: [AB::Var; WIDTH], is_external_row: AB::Var, ) where @@ -116,10 +146,10 @@ impl Poseidon2SkinnyChip { let next_state = next_row.state_var(); for i in 0..WIDTH { - // builder - // .when_transition() - // .when(is_external_row) - // .assert_eq(next_state[i], state[i].clone()); + builder + .when_transition() + .when(is_external_row) + .assert_eq(next_state[i], state[i].clone()); } } diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs index 02f96cbe43..5bef1d8270 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs @@ -191,14 +191,14 @@ pub(crate) mod tests { Runtime::::new(&program, BabyBearPoseidon2::new().perm); runtime.run().unwrap(); - let config = SC::new(); - let machine_deg_3 = A::machine(config); - let (pk_3, vk_3) = machine_deg_3.setup(&program); - let result_deg_3 = - run_test_machine(vec![runtime.record.clone()], machine_deg_3, pk_3, vk_3); - if let Err(e) = result_deg_3 { - panic!("Verification failed: {:?}", e); - } + // let config = SC::new(); + // let machine_deg_3 = A::machine(config); + // let (pk_3, vk_3) = machine_deg_3.setup(&program); + // let result_deg_3 = + // run_test_machine(vec![runtime.record.clone()], machine_deg_3, pk_3, vk_3); + // if let Err(e) = result_deg_3 { + // panic!("Verification failed: {:?}", e); + // } let config = SC::new(); let machine_deg_9 = B::machine(config); diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index af02bf1d66..81ffda4ab8 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -154,17 +154,23 @@ impl MachineAir for Poseidon2SkinnyChip // Set the round-counter columns. cols.round_counters_preprocessed.is_input_round = F::from_bool(i == 0); - let is_external_round = i != INTERNAL_ROUND_IDX && i != 0; + let is_external_round = + i != 0 && i != INTERNAL_ROUND_IDX && i != NUM_EXTERNAL_ROUNDS + 2; cols.round_counters_preprocessed.is_external_round = F::from_bool(is_external_round); cols.round_counters_preprocessed.is_internal_round = - F::from_bool(INTERNAL_ROUND_IDX == i); + F::from_bool(i == INTERNAL_ROUND_IDX); (0..WIDTH).for_each(|j| { cols.round_counters_preprocessed.round_constants[j] = if is_external_round { - let external_round_num = - if i < INTERNAL_ROUND_IDX { i - 1 } else { i - 2 }; - F::from_wrapped_u32(RC_16_30_U32[external_round_num][j]) + let r = i - 1; + let round = if r < NUM_EXTERNAL_ROUNDS / 2 { + r + } else { + r + NUM_INTERNAL_ROUNDS - 1 + }; + + F::from_wrapped_u32(RC_16_30_U32[round][j]) } else if i == INTERNAL_ROUND_IDX { F::from_wrapped_u32(RC_16_30_U32[NUM_EXTERNAL_ROUNDS / 2 + j][0]) } else { From 1a6497ffc5b39003c559206c1d405182a8cf4fff Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Mon, 29 Jul 2024 15:30:28 -0700 Subject: [PATCH 05/18] internal rounds working for degree 9 --- .../core-v2/src/chips/poseidon2_skinny/air.rs | 105 +++++++++++------- 1 file changed, 63 insertions(+), 42 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs index 4fcee9a7bf..faf9e9a766 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs @@ -12,7 +12,9 @@ use crate::builder::SP1RecursionAirBuilder; use super::columns::preprocessed::Poseidon2PreprocessedCols; use super::columns::{Poseidon2, NUM_POSEIDON2_DEGREE3_COLS, NUM_POSEIDON2_DEGREE9_COLS}; -use super::{external_linear_layer, Poseidon2SkinnyChip, WIDTH}; +use super::{ + external_linear_layer, internal_linear_layer, Poseidon2SkinnyChip, NUM_INTERNAL_ROUNDS, WIDTH, +}; impl BaseAir for Poseidon2SkinnyChip { fn width(&self) -> usize { @@ -79,6 +81,14 @@ where prep_local.round_counters_preprocessed.round_constants, prep_local.round_counters_preprocessed.is_external_round, ); + + self.eval_internal_rounds( + builder, + local_row.as_ref(), + next_row.as_ref(), + prep_local.round_counters_preprocessed.round_constants, + prep_local.round_counters_preprocessed.is_internal_round, + ); } } @@ -153,45 +163,56 @@ impl Poseidon2SkinnyChip { } } - // fn eval_internal_rounds( - // &self, - // builder: &mut AB, - // perm_cols: &dyn Permutation, - // ) { - // let state = &perm_cols.internal_rounds_state(); - // let s0 = perm_cols.internal_rounds_s0(); - // let mut state: [AB::Expr; WIDTH] = core::array::from_fn(|i| state[i].into()); - // for r in 0..NUM_INTERNAL_ROUNDS { - // // Add the round constant. - // let round = r + NUM_EXTERNAL_ROUNDS / 2; - // let add_rc = if r == 0 { - // state[0].clone() - // } else { - // s0[r - 1].into() - // } + AB::Expr::from_wrapped_u32(RC_16_30_U32[round][0]); - - // let mut sbox_deg_3 = add_rc.clone() * add_rc.clone() * add_rc.clone(); - // if let Some(internal_sbox) = perm_cols.internal_rounds_sbox() { - // builder.assert_eq(internal_sbox[r], sbox_deg_3); - // sbox_deg_3 = internal_sbox[r].into(); - // } - - // // See `populate_internal_rounds` for why we don't have columns for the sbox output here. - // let sbox_deg_7 = sbox_deg_3.clone() * sbox_deg_3.clone() * add_rc.clone(); - - // // Apply the linear layer. - // // See `populate_internal_rounds` for why we don't have columns for the new state here. - // state[0] = sbox_deg_7.clone(); - // internal_linear_layer(&mut state); - - // if r < NUM_INTERNAL_ROUNDS - 1 { - // builder.assert_eq(s0[r], state[0].clone()); - // } - // } - - // let external_state = perm_cols.external_rounds_state()[NUM_EXTERNAL_ROUNDS / 2]; - // for i in 0..WIDTH { - // builder.assert_eq(external_state[i], state[i].clone()) - // } - // } + fn eval_internal_rounds( + &self, + builder: &mut AB, + local_row: &dyn Poseidon2, + next_row: &dyn Poseidon2, + round_constants: [AB::Var; WIDTH], + is_internal_row: AB::Var, + ) where + AB::Var: 'static, + { + let local_state = local_row.state_var(); + + let s0 = local_row.internal_rounds_s0(); + let mut state: [AB::Expr; WIDTH] = core::array::from_fn(|i| local_state[i].into()); + for r in 0..NUM_INTERNAL_ROUNDS { + // Add the round constant. + let add_rc = if r == 0 { + state[0].clone() + } else { + s0[r - 1].into() + } + round_constants[r]; + + let mut sbox_deg_3 = add_rc.clone() * add_rc.clone() * add_rc.clone(); + if let Some(internal_sbox) = local_row.s_box_state() { + builder + .when(is_internal_row) + .assert_eq(internal_sbox[r], sbox_deg_3); + sbox_deg_3 = internal_sbox[r].into(); + } + + // See `populate_internal_rounds` for why we don't have columns for the sbox output here. + let sbox_deg_7 = sbox_deg_3.clone() * sbox_deg_3.clone() * add_rc.clone(); + + // Apply the linear layer. + // See `populate_internal_rounds` for why we don't have columns for the new state here. + state[0] = sbox_deg_7.clone(); + internal_linear_layer(&mut state); + + if r < NUM_INTERNAL_ROUNDS - 1 { + builder + .when(is_internal_row) + .assert_eq(s0[r], state[0].clone()); + } + } + + let next_state = next_row.state_var(); + for i in 0..WIDTH { + builder + .when(is_internal_row) + .assert_eq(next_state[i], state[i].clone()) + } + } } From 652935b2b234b1edd163b2a4d8f93ce910c5388c Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Mon, 29 Jul 2024 15:32:35 -0700 Subject: [PATCH 06/18] uncomment degree 3 test --- .../core-v2/src/chips/poseidon2_skinny/mod.rs | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs index 5bef1d8270..02f96cbe43 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs @@ -191,14 +191,14 @@ pub(crate) mod tests { Runtime::::new(&program, BabyBearPoseidon2::new().perm); runtime.run().unwrap(); - // let config = SC::new(); - // let machine_deg_3 = A::machine(config); - // let (pk_3, vk_3) = machine_deg_3.setup(&program); - // let result_deg_3 = - // run_test_machine(vec![runtime.record.clone()], machine_deg_3, pk_3, vk_3); - // if let Err(e) = result_deg_3 { - // panic!("Verification failed: {:?}", e); - // } + let config = SC::new(); + let machine_deg_3 = A::machine(config); + let (pk_3, vk_3) = machine_deg_3.setup(&program); + let result_deg_3 = + run_test_machine(vec![runtime.record.clone()], machine_deg_3, pk_3, vk_3); + if let Err(e) = result_deg_3 { + panic!("Verification failed: {:?}", e); + } let config = SC::new(); let machine_deg_9 = B::machine(config); From 6e5dee6ffd89c7154da40439a854bf1e4337c92c Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 12:54:17 -0700 Subject: [PATCH 07/18] trace works, but the air doesn't yet --- .../core-v2/src/chips/poseidon2_skinny/air.rs | 131 ++++++---------- .../src/chips/poseidon2_skinny/columns/mod.rs | 117 ++------------ .../poseidon2_skinny/columns/permutation.rs | 148 ------------------ .../poseidon2_skinny/columns/preprocessed.rs | 2 +- .../core-v2/src/chips/poseidon2_skinny/mod.rs | 63 ++------ .../src/chips/poseidon2_skinny/trace.rs | 146 ++++++----------- 6 files changed, 115 insertions(+), 492 deletions(-) delete mode 100644 recursion/core-v2/src/chips/poseidon2_skinny/columns/permutation.rs diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs index faf9e9a766..7cf1ca19e5 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs @@ -9,22 +9,19 @@ use p3_field::AbstractField; use p3_matrix::Matrix; use crate::builder::SP1RecursionAirBuilder; +use crate::chips::poseidon2_skinny::columns::Poseidon2; use super::columns::preprocessed::Poseidon2PreprocessedCols; -use super::columns::{Poseidon2, NUM_POSEIDON2_DEGREE3_COLS, NUM_POSEIDON2_DEGREE9_COLS}; +use super::columns::NUM_POSEIDON2_COLS; use super::{ external_linear_layer, internal_linear_layer, Poseidon2SkinnyChip, NUM_INTERNAL_ROUNDS, WIDTH, }; impl BaseAir for Poseidon2SkinnyChip { fn width(&self) -> usize { - if DEGREE == 3 || DEGREE == 5 { - NUM_POSEIDON2_DEGREE3_COLS - } else if DEGREE == 9 || DEGREE == 17 { - NUM_POSEIDON2_DEGREE9_COLS - } else { - panic!("Unsupported degree: {}", DEGREE); - } + // We only support machines with degree 9. + assert!(DEGREE == 9); + NUM_POSEIDON2_COLS } } @@ -34,19 +31,23 @@ where AB::Var: 'static, { fn eval(&self, builder: &mut AB) { + // We only support machines with degree 9. + assert!(DEGREE == 9); + let main = builder.main(); - let local_row = Self::convert::(main.row_slice(0)); - let next_row = Self::convert::(main.row_slice(1)); + let (local_row, next_row) = (main.row_slice(0), main.row_slice(1)); + let local_row: &Poseidon2<_> = (*local_row).borrow(); + let next_row: &Poseidon2<_> = (*next_row).borrow(); let prepr = builder.preprocessed(); let prep_local = prepr.row_slice(0); let prep_local: &Poseidon2PreprocessedCols<_> = (*prep_local).borrow(); // Dummy constraints to normalize to DEGREE. let lhs = (0..DEGREE) - .map(|_| local_row.state_var()[0].into()) + .map(|_| local_row.state_var[0].into()) .product::(); let rhs = (0..DEGREE) - .map(|_| local_row.state_var()[0].into()) + .map(|_| local_row.state_var[0].into()) .product::(); builder.assert_eq(lhs, rhs); @@ -54,7 +55,7 @@ where (0..WIDTH).for_each(|i| { builder.receive_single( prep_local.memory_preprocessed[i].addr, - local_row.state_var()[i], + local_row.state_var[i], prep_local.memory_preprocessed[i].read_mult, ) }); @@ -62,30 +63,17 @@ where (0..WIDTH).for_each(|i| { builder.send_single( prep_local.memory_preprocessed[i].addr, - local_row.state_var()[i], + local_row.state_var[i], prep_local.memory_preprocessed[i].write_mult, ) }); - self.eval_input_round( - builder, - local_row.as_ref(), - next_row.as_ref(), - prep_local.round_counters_preprocessed.is_input_round, - ); - - self.eval_external_round( - builder, - local_row.as_ref(), - next_row.as_ref(), - prep_local.round_counters_preprocessed.round_constants, - prep_local.round_counters_preprocessed.is_external_round, - ); + self.eval_external_round(builder, local_row, prep_local, next_row); self.eval_internal_rounds( builder, - local_row.as_ref(), - next_row.as_ref(), + local_row, + next_row, prep_local.round_counters_preprocessed.round_constants, prep_local.round_counters_preprocessed.is_internal_round, ); @@ -93,72 +81,50 @@ where } impl Poseidon2SkinnyChip { - fn eval_input_round( + fn eval_external_round( &self, builder: &mut AB, - local_row: &dyn Poseidon2, - next_row: &dyn Poseidon2, - is_input_row: AB::Var, + local_row: &Poseidon2, + prep_local: &Poseidon2PreprocessedCols, + next_row: &Poseidon2, ) where AB::Var: 'static, { - let mut local_state: [AB::Expr; WIDTH] = - array::from_fn(|i| local_row.state_var()[i].into()); + let mut state_ell: [AB::Expr; WIDTH] = array::from_fn(|i| local_row.state_var[i].into()); - external_linear_layer(&mut local_state); + // For the first external round, we should apply the external linear layer. + external_linear_layer(&mut state_ell); - let next_state = next_row.state_var(); - for i in 0..WIDTH { - builder - .when(is_input_row) - .assert_eq(next_state[i], local_state[i].clone()); - } - } - - fn eval_external_round( - &self, - builder: &mut AB, - local_row: &dyn Poseidon2, - next_row: &dyn Poseidon2, - round_constants: [AB::Var; WIDTH], - is_external_row: AB::Var, - ) where - AB::Var: 'static, - { - let local_state = local_row.state_var(); + let state: [AB::Expr; WIDTH] = array::from_fn(|i| { + builder.if_else( + prep_local.round_counters_preprocessed.is_first_round, + state_ell[i].clone(), + local_row.state_var[i], + ) + }); // Add the round constants. - let add_rc: [AB::Expr; WIDTH] = - core::array::from_fn(|i| local_state[i].into() + round_constants[i]); + let add_rc: [AB::Expr; WIDTH] = core::array::from_fn(|i| { + state[i].clone() + prep_local.round_counters_preprocessed.round_constants[i] + }); // Apply the sboxes. // See `populate_external_round` for why we don't have columns for the sbox output here. let mut sbox_deg_7: [AB::Expr; WIDTH] = core::array::from_fn(|_| AB::Expr::zero()); - let mut sbox_deg_3: [AB::Expr; WIDTH] = core::array::from_fn(|_| AB::Expr::zero()); for i in 0..WIDTH { - let calculated_sbox_deg_3 = add_rc[i].clone() * add_rc[i].clone() * add_rc[i].clone(); - - if let Some(external_sbox) = local_row.s_box_state() { - builder - .when(is_external_row) - .assert_eq(external_sbox[i].into(), calculated_sbox_deg_3); - sbox_deg_3[i] = external_sbox[i].into(); - } else { - sbox_deg_3[i] = calculated_sbox_deg_3; - } - - sbox_deg_7[i] = sbox_deg_3[i].clone() * sbox_deg_3[i].clone() * add_rc[i].clone(); + let sbox_deg_3 = add_rc[i].clone() * add_rc[i].clone() * add_rc[i].clone(); + sbox_deg_7[i] = sbox_deg_3.clone() * sbox_deg_3.clone() * add_rc[i].clone(); } // Apply the linear layer. let mut state = sbox_deg_7; external_linear_layer(&mut state); - let next_state = next_row.state_var(); + let next_state = next_row.state_var; for i in 0..WIDTH { builder .when_transition() - .when(is_external_row) + .when(prep_local.round_counters_preprocessed.is_external_round) .assert_eq(next_state[i], state[i].clone()); } } @@ -166,16 +132,16 @@ impl Poseidon2SkinnyChip { fn eval_internal_rounds( &self, builder: &mut AB, - local_row: &dyn Poseidon2, - next_row: &dyn Poseidon2, + local_row: &Poseidon2, + next_row: &Poseidon2, round_constants: [AB::Var; WIDTH], is_internal_row: AB::Var, ) where AB::Var: 'static, { - let local_state = local_row.state_var(); + let local_state = local_row.state_var; - let s0 = local_row.internal_rounds_s0(); + let s0 = local_row.internal_rounds_s0; let mut state: [AB::Expr; WIDTH] = core::array::from_fn(|i| local_state[i].into()); for r in 0..NUM_INTERNAL_ROUNDS { // Add the round constant. @@ -185,14 +151,7 @@ impl Poseidon2SkinnyChip { s0[r - 1].into() } + round_constants[r]; - let mut sbox_deg_3 = add_rc.clone() * add_rc.clone() * add_rc.clone(); - if let Some(internal_sbox) = local_row.s_box_state() { - builder - .when(is_internal_row) - .assert_eq(internal_sbox[r], sbox_deg_3); - sbox_deg_3 = internal_sbox[r].into(); - } - + let sbox_deg_3 = add_rc.clone() * add_rc.clone() * add_rc.clone(); // See `populate_internal_rounds` for why we don't have columns for the sbox output here. let sbox_deg_7 = sbox_deg_3.clone() * sbox_deg_3.clone() * add_rc.clone(); @@ -208,7 +167,7 @@ impl Poseidon2SkinnyChip { } } - let next_state = next_row.state_var(); + let next_state = next_row.state_var; for i in 0..WIDTH { builder .when(is_internal_row) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/columns/mod.rs b/recursion/core-v2/src/chips/poseidon2_skinny/columns/mod.rs index 62c5bb4a3c..74aefe098b 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/columns/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/columns/mod.rs @@ -1,120 +1,23 @@ use std::mem::{size_of, transmute}; -use permutation::{PermutationNoSbox, PermutationSBox}; use sp1_core::utils::indices_arr; use sp1_derive::AlignedBorrow; -use super::{NUM_INTERNAL_ROUNDS, WIDTH}; +use crate::chips::poseidon2_skinny::{NUM_INTERNAL_ROUNDS, WIDTH}; -pub mod permutation; pub mod preprocessed; -pub const POSEIDON2_DEGREE3_COL_MAP: Poseidon2Degree3 = make_col_map_degree3(); -pub const NUM_POSEIDON2_DEGREE3_COLS: usize = size_of::>(); -/// Trait for getter methods for Poseidon2 columns. -pub trait Poseidon2<'a, T: Copy + 'a> { - fn state_var(&self) -> &[T; WIDTH]; - fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS - 1]; - fn s_box_state(&self) -> Option<&[T; WIDTH]>; +pub const NUM_POSEIDON2_COLS: usize = size_of::>(); +const fn make_col_map_degree9() -> Poseidon2 { + let indices_arr = indices_arr::(); + unsafe { transmute::<[usize; NUM_POSEIDON2_COLS], Poseidon2>(indices_arr) } } +pub const POSEIDON2_DEGREE9_COL_MAP: Poseidon2 = make_col_map_degree9(); -/// Trait for setter methods for Poseidon2 columns. Only need the memory columns are populated mutably. -pub trait Poseidon2Mut<'a, T: Copy + 'a> { - fn get_cols_mut( - &mut self, - ) -> ( - &mut [T; WIDTH], - &mut [T; NUM_INTERNAL_ROUNDS - 1], - Option<&mut [T; WIDTH]>, - ); -} - -const fn make_col_map_degree3() -> Poseidon2Degree3 { - let indices_arr = indices_arr::(); - unsafe { - transmute::<[usize; NUM_POSEIDON2_DEGREE3_COLS], Poseidon2Degree3>(indices_arr) - } -} - -/// Struct for the poseidon2 chip that contains sbox columns. -#[derive(AlignedBorrow, Clone, Copy)] -#[repr(C)] -pub struct Poseidon2Degree3 { - pub permutation_cols: PermutationSBox, -} - -impl<'a, T: Copy + 'a> Poseidon2<'a, T> for Poseidon2Degree3 { - fn state_var(&self) -> &[T; WIDTH] { - &self.permutation_cols.state.state_var - } - - fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS - 1] { - &self.permutation_cols.state.internal_rounds_s0 - } - - fn s_box_state(&self) -> Option<&[T; WIDTH]> { - Some(&self.permutation_cols.sbox_state.sbox_state) - } -} - -impl<'a, T: Copy + 'a> Poseidon2Mut<'a, T> for &'a mut Poseidon2Degree3 { - fn get_cols_mut( - &mut self, - ) -> ( - &mut [T; WIDTH], - &mut [T; NUM_INTERNAL_ROUNDS - 1], - Option<&mut [T; WIDTH]>, - ) { - ( - &mut self.permutation_cols.state.state_var, - &mut self.permutation_cols.state.internal_rounds_s0, - Some(&mut self.permutation_cols.sbox_state.sbox_state), - ) - } -} - -pub const NUM_POSEIDON2_DEGREE9_COLS: usize = size_of::>(); -const fn make_col_map_degree9() -> Poseidon2Degree9 { - let indices_arr = indices_arr::(); - unsafe { - transmute::<[usize; NUM_POSEIDON2_DEGREE9_COLS], Poseidon2Degree9>(indices_arr) - } -} -pub const POSEIDON2_DEGREE9_COL_MAP: Poseidon2Degree9 = make_col_map_degree9(); - -/// Struct for the poseidon2 chip that doesn't contain sbox columns. +/// Struct for the poseidon2 skinny non preprocessed column. #[derive(AlignedBorrow, Clone, Copy)] #[repr(C)] -pub struct Poseidon2Degree9 { - pub permutation_cols: PermutationNoSbox, -} - -impl<'a, T: Copy + 'a> Poseidon2<'a, T> for Poseidon2Degree9 { - fn state_var(&self) -> &[T; WIDTH] { - &self.permutation_cols.state.state_var - } - - fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS - 1] { - &self.permutation_cols.state.internal_rounds_s0 - } - - fn s_box_state(&self) -> Option<&[T; WIDTH]> { - None - } -} - -impl<'a, T: Copy + 'a> Poseidon2Mut<'a, T> for &'a mut Poseidon2Degree9 { - fn get_cols_mut( - &mut self, - ) -> ( - &mut [T; WIDTH], - &mut [T; NUM_INTERNAL_ROUNDS - 1], - Option<&mut [T; WIDTH]>, - ) { - ( - &mut self.permutation_cols.state.state_var, - &mut self.permutation_cols.state.internal_rounds_s0, - None, - ) - } +pub struct Poseidon2 { + pub state_var: [T; WIDTH], + pub internal_rounds_s0: [T; NUM_INTERNAL_ROUNDS - 1], } diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/columns/permutation.rs b/recursion/core-v2/src/chips/poseidon2_skinny/columns/permutation.rs deleted file mode 100644 index 0469707ece..0000000000 --- a/recursion/core-v2/src/chips/poseidon2_skinny/columns/permutation.rs +++ /dev/null @@ -1,148 +0,0 @@ -use std::{borrow::BorrowMut, mem::size_of}; - -use sp1_derive::AlignedBorrow; - -use crate::chips::poseidon2_skinny::{NUM_INTERNAL_ROUNDS, WIDTH}; - -use super::{POSEIDON2_DEGREE3_COL_MAP, POSEIDON2_DEGREE9_COL_MAP}; - -pub const fn max(a: usize, b: usize) -> usize { - if a > b { - a - } else { - b - } -} - -#[derive(AlignedBorrow, Clone, Copy)] -#[repr(C)] -pub struct PermutationState { - pub state_var: [T; WIDTH], - pub internal_rounds_s0: [T; NUM_INTERNAL_ROUNDS - 1], -} - -#[derive(AlignedBorrow, Clone, Copy)] -#[repr(C)] -pub struct PermutationSBoxState { - pub sbox_state: [T; max(WIDTH, NUM_INTERNAL_ROUNDS)], -} - -/// Trait that describes getter functions for the permutation columns. -pub trait Permutation { - fn state(&self) -> &[T; WIDTH]; - - fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS - 1]; - - fn sbox_state(&self) -> Option<&[T; max(WIDTH, NUM_INTERNAL_ROUNDS)]>; -} - -/// Trait that describes setter functions for the permutation columns. -pub trait PermutationMut { - #[allow(clippy::type_complexity)] - fn set_perm_state( - &mut self, - state: [T; WIDTH], - internal_rounds_s0: [T; NUM_INTERNAL_ROUNDS - 1], - ); - fn set_sbox_state(&mut self, sbox_state: Option<[T; max(WIDTH, NUM_INTERNAL_ROUNDS)]>); -} - -/// Permutation columns struct with S-boxes. -#[derive(AlignedBorrow, Clone, Copy)] -#[repr(C)] -pub struct PermutationSBox { - pub state: PermutationState, - pub sbox_state: PermutationSBoxState, -} - -impl Permutation for PermutationSBox { - fn state(&self) -> &[T; WIDTH] { - &self.state.state_var - } - - fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS - 1] { - &self.state.internal_rounds_s0 - } - - fn sbox_state(&self) -> Option<&[T; max(WIDTH, NUM_INTERNAL_ROUNDS)]> { - Some(&self.sbox_state.sbox_state) - } -} - -impl PermutationMut for PermutationSBox { - fn set_perm_state( - &mut self, - state: [T; WIDTH], - internal_rounds_s0: [T; NUM_INTERNAL_ROUNDS - 1], - ) { - self.state.state_var = state; - self.state.internal_rounds_s0 = internal_rounds_s0; - } - - fn set_sbox_state(&mut self, sbox_state: Option<[T; max(WIDTH, NUM_INTERNAL_ROUNDS)]>) { - if let Some(sbox_state) = sbox_state { - self.sbox_state.sbox_state = sbox_state; - } - } -} - -/// Permutation columns struct without S-boxes. -#[derive(AlignedBorrow, Clone, Copy)] -#[repr(C)] -pub struct PermutationNoSbox { - pub state: PermutationState, -} - -impl Permutation for PermutationNoSbox { - fn state(&self) -> &[T; WIDTH] { - &self.state.state_var - } - - fn internal_rounds_s0(&self) -> &[T; NUM_INTERNAL_ROUNDS - 1] { - &self.state.internal_rounds_s0 - } - - fn sbox_state(&self) -> Option<&[T; max(WIDTH, NUM_INTERNAL_ROUNDS)]> { - None - } -} - -impl PermutationMut for PermutationNoSbox { - fn set_perm_state( - &mut self, - state: [T; WIDTH], - internal_rounds_s0: [T; NUM_INTERNAL_ROUNDS - 1], - ) { - self.state.state_var = state; - self.state.internal_rounds_s0 = internal_rounds_s0; - } - - fn set_sbox_state(&mut self, _sbox_state: Option<[T; max(WIDTH, NUM_INTERNAL_ROUNDS)]>) {} -} - -/// Permutation columns struct without S-boxes and half of the external rounds. -/// In the past, all external rounds were stored in one row, so this was a distinct struct, but -/// now the structs don't track the number of external rounds. -pub type PermutationNoSboxHalfExternal = PermutationNoSbox; - -pub fn permutation_mut<'a, 'b: 'a, T, const DEGREE: usize>( - row: &'b mut [T], -) -> Box + 'a> -where - T: Copy, -{ - if DEGREE == 3 || DEGREE == 5 { - let start = POSEIDON2_DEGREE3_COL_MAP.permutation_cols.state.state_var[0]; - let end = start + size_of::>(); - let convert: PermutationSBox = *row[start..end].borrow_mut(); - Box::new(convert) - } else if DEGREE == 9 || DEGREE == 17 { - let start = POSEIDON2_DEGREE9_COL_MAP.permutation_cols.state.state_var[0]; - let end = start + size_of::>(); - - let convert: PermutationNoSbox = *row[start..end].borrow_mut(); - Box::new(convert) - } else { - panic!("Unsupported degree"); - } -} diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs b/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs index 02f5e041a8..466753def9 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs @@ -5,7 +5,7 @@ use crate::chips::{mem::MemoryAccessCols, poseidon2_skinny::WIDTH}; #[derive(AlignedBorrow, Clone, Copy, Debug)] #[repr(C)] pub struct RoundCountersPreprocessedCols { - pub is_input_round: T, + pub is_first_round: T, pub is_external_round: T, pub is_internal_round: T, pub round_constants: [T; WIDTH], diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs index 02f96cbe43..238d757639 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs @@ -1,9 +1,3 @@ -#![allow(clippy::needless_range_loop)] - -use std::borrow::Borrow; -use std::borrow::BorrowMut; -use std::ops::Deref; - use p3_baby_bear::{MONTY_INVERSE, POSEIDON2_INTERNAL_MATRIX_DIAG_16_BABYBEAR_MONTY}; use p3_field::AbstractField; use p3_field::PrimeField32; @@ -14,8 +8,6 @@ pub mod trace; use p3_poseidon2::matmul_internal; -use self::columns::{Poseidon2, Poseidon2Degree3, Poseidon2Degree9, Poseidon2Mut}; - /// The width of the permutation. pub const WIDTH: usize = 16; pub const RATE: usize = WIDTH / 2; @@ -24,6 +16,14 @@ pub const NUM_EXTERNAL_ROUNDS: usize = 8; pub const NUM_INTERNAL_ROUNDS: usize = 13; pub const NUM_ROUNDS: usize = NUM_EXTERNAL_ROUNDS + NUM_INTERNAL_ROUNDS; +pub const fn max(a: usize, b: usize) -> usize { + if a > b { + a + } else { + b + } +} + /// A chip that implements the Poseidon2 permutation in the skinny variant (one external round per /// row and one row for all internal rounds). pub struct Poseidon2SkinnyChip { @@ -33,47 +33,14 @@ pub struct Poseidon2SkinnyChip { impl Default for Poseidon2SkinnyChip { fn default() -> Self { + // We only support machines with degree 9. + assert!(DEGREE == 9); Self { fixed_log2_rows: None, pad: true, } } } - -impl<'a, const DEGREE: usize> Poseidon2SkinnyChip { - /// Transmute a row it to an immutable Poseidon2 instance. - pub(crate) fn convert(row: impl Deref) -> Box + 'a> - where - T: Copy + 'a, - { - if DEGREE == 3 || DEGREE == 5 { - let convert: &Poseidon2Degree3 = (*row).borrow(); - Box::new(*convert) - } else if DEGREE == 9 || DEGREE == 17 { - let convert: &Poseidon2Degree9 = (*row).borrow(); - Box::new(*convert) - } else { - panic!("Unsupported degree"); - } - } - - /// Transmute a row it to a mutable Poseidon2 instance. - pub(crate) fn convert_mut<'b: 'a, F: PrimeField32>( - &self, - row: &'b mut Vec, - ) -> Box + 'a> { - if DEGREE == 3 || DEGREE == 5 { - let convert: &mut Poseidon2Degree3 = row.as_mut_slice().borrow_mut(); - Box::new(convert) - } else if DEGREE == 9 || DEGREE == 17 { - let convert: &mut Poseidon2Degree9 = row.as_mut_slice().borrow_mut(); - Box::new(convert) - } else { - panic!("Unsupported degree"); - } - } -} - pub fn apply_m_4(x: &mut [AF]) where AF: AbstractField, @@ -143,7 +110,6 @@ pub(crate) mod tests { type SC = BabyBearPoseidon2Outer; type F = ::Val; type EF = ::Challenge; - type A = RecursionAir; type B = RecursionAir; let input = [1; WIDTH]; @@ -191,15 +157,6 @@ pub(crate) mod tests { Runtime::::new(&program, BabyBearPoseidon2::new().perm); runtime.run().unwrap(); - let config = SC::new(); - let machine_deg_3 = A::machine(config); - let (pk_3, vk_3) = machine_deg_3.setup(&program); - let result_deg_3 = - run_test_machine(vec![runtime.record.clone()], machine_deg_3, pk_3, vk_3); - if let Err(e) = result_deg_3 { - panic!("Verification failed: {:?}", e); - } - let config = SC::new(); let machine_deg_9 = B::machine(config); let (pk_9, vk_9) = machine_deg_9.setup(&program); diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index 81ffda4ab8..6bd115c8b1 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -1,4 +1,7 @@ -use std::{borrow::BorrowMut, mem::size_of}; +use std::{ + borrow::{Borrow, BorrowMut}, + mem::size_of, +}; use itertools::Itertools; use p3_air::BaseAir; @@ -12,21 +15,19 @@ use crate::{ chips::{ mem::MemoryAccessCols, poseidon2_skinny::{ - external_linear_layer, Poseidon2SkinnyChip, NUM_EXTERNAL_ROUNDS, NUM_INTERNAL_ROUNDS, + columns::Poseidon2, external_linear_layer, Poseidon2SkinnyChip, NUM_EXTERNAL_ROUNDS, + NUM_INTERNAL_ROUNDS, }, }, instruction::Instruction::Poseidon2Skinny, ExecutionRecord, RecursionProgram, }; -use super::{ - columns::{permutation::max, preprocessed::Poseidon2PreprocessedCols}, - internal_linear_layer, WIDTH, -}; +use super::{columns::preprocessed::Poseidon2PreprocessedCols, internal_linear_layer, max, WIDTH}; const PREPROCESSED_POSEIDON2_WIDTH: usize = size_of::>(); -const INTERNAL_ROUND_IDX: usize = NUM_EXTERNAL_ROUNDS / 2 + 1; +const INTERNAL_ROUND_IDX: usize = NUM_EXTERNAL_ROUNDS / 2; impl MachineAir for Poseidon2SkinnyChip { type Record = ExecutionRecord; @@ -50,47 +51,36 @@ impl MachineAir for Poseidon2SkinnyChip for event in &input.poseidon2_skinny_events { // We have one row for input, one row for output, NUM_EXTERNAL_ROUNDS rows for the external rounds, // and one row for all internal rounds. - let mut row_add = vec![vec![F::zero(); num_columns]; NUM_EXTERNAL_ROUNDS + 3]; + let mut row_add = vec![vec![F::zero(); num_columns]; NUM_EXTERNAL_ROUNDS + 2]; - // The first row should have event.input and [event.input[0].clone(); NUM_INTERNAL_ROUNDS-1] - // in its state columns. The sbox_state will be modified in the computation of the - // first row. - { - let (first_row, second_row) = &mut row_add[0..2].split_at_mut(1); - let mut input_cols = self.convert_mut(&mut first_row[0]); - *input_cols.get_cols_mut().0 = event.input; - - let mut next_cols = self.convert_mut(&mut second_row[0]); - let (next_state, _, _) = next_cols.get_cols_mut(); - *next_state = event.input; - external_linear_layer(next_state); - } + let first_row_cols: &mut Poseidon2 = row_add[0].as_mut_slice().borrow_mut(); + first_row_cols.state_var = event.input; // For each external round, and once for all the internal rounds at the same time, apply - // the corresponding operation. This will change the sbox state in row i, and the state - // and internal_rounds_s0 variable in row i+1. - for i in 1..NUM_EXTERNAL_ROUNDS + 2 { + // the corresponding operation. This will change the state and internal_rounds_s0 variable + // in row r+1. + for r in 0..NUM_EXTERNAL_ROUNDS + 1 { let next_state_var = { - let mut cols = self.convert_mut(&mut row_add[i]); - let (state, internal_state_s0, mut sbox_state) = cols.get_cols_mut(); + let cols: &mut Poseidon2 = row_add[r].as_mut_slice().borrow_mut(); + let state = cols.state_var; - if i != INTERNAL_ROUND_IDX { - self.populate_external_round(state, &mut sbox_state, i - 1) + if r != INTERNAL_ROUND_IDX { + self.populate_external_round(&state, r) } else { // Populate the internal rounds. - self.populate_internal_rounds(state, internal_state_s0, &mut sbox_state) + self.populate_internal_rounds(&state, &mut cols.internal_rounds_s0) } }; - let mut next_cols = self.convert_mut(&mut row_add[i + 1]); - let (next_state, _, _) = next_cols.get_cols_mut(); - *next_state = next_state_var; + let next_row_cols: &mut Poseidon2 = row_add[r + 1].as_mut_slice().borrow_mut(); + next_row_cols.state_var = next_state_var; } // Check that the permutation is computed correctly. { - let last_row = &row_add[NUM_EXTERNAL_ROUNDS + 2]; - let output_cols = Poseidon2SkinnyChip::::convert(last_row.as_slice()); - debug_assert_eq!(*output_cols.state_var(), event.output); + let last_row_cols: &Poseidon2 = + row_add[NUM_EXTERNAL_ROUNDS + 1].as_slice().borrow(); + + debug_assert_eq!(last_row_cols.state_var, event.output); } rows.extend(row_add.into_iter()); } @@ -141,37 +131,36 @@ impl MachineAir for Poseidon2SkinnyChip let mut rows = vec![ [F::zero(); PREPROCESSED_POSEIDON2_WIDTH]; - num_instructions * (NUM_EXTERNAL_ROUNDS + 3) + num_instructions * (NUM_EXTERNAL_ROUNDS + 2) ]; - // Iterate over the instructions and take NUM_EXTERNAL_ROUNDS + 3 rows for each instruction. + // Iterate over the instructions and take NUM_EXTERNAL_ROUNDS + 2 rows for each instruction. + // We have one extra round for the internal rounds and one extra round for the output. instructions - .zip_eq(&rows.iter_mut().chunks(NUM_EXTERNAL_ROUNDS + 3)) + .zip_eq(&rows.iter_mut().chunks(NUM_EXTERNAL_ROUNDS + 2)) .for_each(|(instruction, row_add)| { - row_add.into_iter().enumerate().for_each(|(i, row)| { + row_add.into_iter().enumerate().for_each(|(r, row)| { let cols: &mut Poseidon2PreprocessedCols<_> = (*row).as_mut_slice().borrow_mut(); // Set the round-counter columns. - cols.round_counters_preprocessed.is_input_round = F::from_bool(i == 0); - let is_external_round = - i != 0 && i != INTERNAL_ROUND_IDX && i != NUM_EXTERNAL_ROUNDS + 2; + cols.round_counters_preprocessed.is_first_round = F::from_bool(r == 0); + let is_external_round = r != INTERNAL_ROUND_IDX && r != NUM_EXTERNAL_ROUNDS + 1; cols.round_counters_preprocessed.is_external_round = F::from_bool(is_external_round); cols.round_counters_preprocessed.is_internal_round = - F::from_bool(i == INTERNAL_ROUND_IDX); + F::from_bool(r == INTERNAL_ROUND_IDX); (0..WIDTH).for_each(|j| { cols.round_counters_preprocessed.round_constants[j] = if is_external_round { - let r = i - 1; let round = if r < NUM_EXTERNAL_ROUNDS / 2 { r } else { - r + NUM_INTERNAL_ROUNDS - 1 + r + NUM_INTERNAL_ROUNDS }; F::from_wrapped_u32(RC_16_30_U32[round][j]) - } else if i == INTERNAL_ROUND_IDX { + } else if r == INTERNAL_ROUND_IDX { F::from_wrapped_u32(RC_16_30_U32[NUM_EXTERNAL_ROUNDS / 2 + j][0]) } else { F::zero() @@ -180,19 +169,19 @@ impl MachineAir for Poseidon2SkinnyChip // Set the memory columns. We read once, at the first iteration, // and write once, at the last iteration. - if i == 0 { + if r == 0 { cols.memory_preprocessed = instruction.addrs.input.map(|addr| MemoryAccessCols { addr, read_mult: F::one(), write_mult: F::zero(), }); - } else if i == NUM_EXTERNAL_ROUNDS + 2 { + } else if r == NUM_EXTERNAL_ROUNDS + 1 { cols.memory_preprocessed = instruction.addrs.output.map(|addr| MemoryAccessCols { addr, read_mult: F::zero(), - write_mult: instruction.mults[i], + write_mult: instruction.mults[r], }); } }); @@ -218,10 +207,15 @@ impl Poseidon2SkinnyChip { fn populate_external_round( &self, round_state: &[F; WIDTH], - sbox: &mut Option<&mut [F; WIDTH]>, r: usize, ) -> [F; WIDTH] { let mut state = { + // For the first round, apply the external linear layer. + let mut state = *round_state; + if r == 0 { + external_linear_layer(&mut state); + } + // Add round constants. // Optimization: Since adding a constant is a degree 1 operation, we can avoid adding @@ -231,10 +225,7 @@ impl Poseidon2SkinnyChip { } else { r + NUM_INTERNAL_ROUNDS - 1 }; - let mut add_rc = *round_state; - for i in 0..WIDTH { - add_rc[i] += F::from_wrapped_u32(RC_16_30_U32[round][i]); - } + (0..WIDTH).for_each(|i| state[i] += F::from_wrapped_u32(RC_16_30_U32[round][i])); // Apply the sboxes. // Optimization: since the linear layer that comes after the sbox is degree 1, we can @@ -243,12 +234,8 @@ impl Poseidon2SkinnyChip { let mut sbox_deg_7: [F; 16] = [F::zero(); WIDTH]; let mut sbox_deg_3: [F; 16] = [F::zero(); WIDTH]; for i in 0..WIDTH { - sbox_deg_3[i] = add_rc[i] * add_rc[i] * add_rc[i]; - sbox_deg_7[i] = sbox_deg_3[i] * sbox_deg_3[i] * add_rc[i]; - } - - if let Some(sbox) = sbox.as_deref_mut() { - *sbox = sbox_deg_3; + sbox_deg_3[i] = state[i] * state[i] * state[i]; + sbox_deg_7[i] = sbox_deg_3[i] * sbox_deg_3[i] * state[i]; } sbox_deg_7 @@ -263,7 +250,6 @@ impl Poseidon2SkinnyChip { &self, state: &[F; WIDTH], internal_rounds_s0: &mut [F; NUM_INTERNAL_ROUNDS - 1], - sbox: &mut Option<&mut [F; WIDTH]>, ) -> [F; WIDTH] { let mut new_state = *state; let mut sbox_deg_3: [F; max(WIDTH, NUM_INTERNAL_ROUNDS)] = @@ -295,13 +281,7 @@ impl Poseidon2SkinnyChip { } } - let ret_state = new_state; - - if let Some(sbox) = sbox.as_deref_mut() { - *sbox = sbox_deg_3; - } - - ret_state + new_state } } @@ -321,35 +301,7 @@ mod tests { }; #[test] - fn generate_trace_deg_3() { - type F = BabyBear; - let input_0 = [F::one(); WIDTH]; - let permuter = inner_perm(); - let output_0 = permuter.permute(input_0); - let mut rng = rand::thread_rng(); - - let input_1 = [F::rand(&mut rng); WIDTH]; - let output_1 = permuter.permute(input_1); - - let shard = ExecutionRecord { - poseidon2_skinny_events: vec![ - Poseidon2SkinnyEvent { - input: input_0, - output: output_0, - }, - Poseidon2SkinnyEvent { - input: input_1, - output: output_1, - }, - ], - ..Default::default() - }; - let chip_3 = Poseidon2SkinnyChip::<3>::default(); - let _: RowMajorMatrix = chip_3.generate_trace(&shard, &mut ExecutionRecord::default()); - } - - #[test] - fn generate_trace_deg_9() { + fn generate_trace() { type F = BabyBear; let input_0 = [F::one(); WIDTH]; let permuter = inner_perm(); From 140b7a6fe025a25faefdd66ec1e55e9c561b61f7 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 13:54:00 -0700 Subject: [PATCH 08/18] fixed bug --- recursion/core-v2/src/chips/poseidon2_skinny/trace.rs | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index 6bd115c8b1..07611af173 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -153,15 +153,15 @@ impl MachineAir for Poseidon2SkinnyChip (0..WIDTH).for_each(|j| { cols.round_counters_preprocessed.round_constants[j] = if is_external_round { - let round = if r < NUM_EXTERNAL_ROUNDS / 2 { + let round = if r < INTERNAL_ROUND_IDX { r } else { - r + NUM_INTERNAL_ROUNDS + r + NUM_INTERNAL_ROUNDS - 1 }; F::from_wrapped_u32(RC_16_30_U32[round][j]) } else if r == INTERNAL_ROUND_IDX { - F::from_wrapped_u32(RC_16_30_U32[NUM_EXTERNAL_ROUNDS / 2 + j][0]) + F::from_wrapped_u32(RC_16_30_U32[INTERNAL_ROUND_IDX + j][0]) } else { F::zero() }; @@ -220,7 +220,7 @@ impl Poseidon2SkinnyChip { // Optimization: Since adding a constant is a degree 1 operation, we can avoid adding // columns for it, and instead include it in the constraint for the x^3 part of the sbox. - let round = if r < NUM_EXTERNAL_ROUNDS / 2 { + let round = if r < INTERNAL_ROUND_IDX { r } else { r + NUM_INTERNAL_ROUNDS - 1 @@ -258,7 +258,7 @@ impl Poseidon2SkinnyChip { // Add the round constant to the 0th state element. // Optimization: Since adding a constant is a degree 1 operation, we can avoid adding // columns for it, just like for external rounds. - let round = r + NUM_EXTERNAL_ROUNDS / 2; + let round = r + INTERNAL_ROUND_IDX; let add_rc = new_state[0] + F::from_wrapped_u32(RC_16_30_U32[round][0]); // Apply the sboxes. From a09b566265643c0f2451d27953c97993334f3276 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 15:58:33 -0700 Subject: [PATCH 09/18] working air --- recursion/circuit-v2/src/challenger.rs | 15 ++- .../core-v2/src/chips/poseidon2_skinny/air.rs | 46 ++++--- .../poseidon2_skinny/columns/preprocessed.rs | 2 +- .../core-v2/src/chips/poseidon2_skinny/mod.rs | 8 -- .../src/chips/poseidon2_skinny/trace.rs | 112 +++++++++--------- 5 files changed, 98 insertions(+), 85 deletions(-) diff --git a/recursion/circuit-v2/src/challenger.rs b/recursion/circuit-v2/src/challenger.rs index 35f469c5a7..fe627a570a 100644 --- a/recursion/circuit-v2/src/challenger.rs +++ b/recursion/circuit-v2/src/challenger.rs @@ -255,9 +255,18 @@ mod tests { let records = vec![runtime.record]; - let machine = RecursionAir::<_, 3, 0>::machine_with_all_chips(BabyBearPoseidon2::default()); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(records.clone(), machine, pk, vk); + // Run with the poseidon2 wide chip. + // let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(BabyBearPoseidon2::default()); + // let (pk, vk) = wide_machine.setup(&program); + // let result = run_test_machine(records.clone(), wide_machine, pk, vk); + // if let Err(e) = result { + // panic!("Verification failed: {:?}", e); + // } + + // Run with the poseidon2 skinny chip. + let skinny_machine = RecursionAir::<_, 9, 0>::machine(BabyBearPoseidon2::compressed()); + let (pk, vk) = skinny_machine.setup(&program); + let result = run_test_machine(records.clone(), skinny_machine, pk, vk); if let Err(e) = result { panic!("Verification failed: {:?}", e); } diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs index 7cf1ca19e5..882662220d 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs @@ -68,6 +68,8 @@ where ) }); + self.eval_input_round(builder, local_row, prep_local, next_row); + self.eval_external_round(builder, local_row, prep_local, next_row); self.eval_internal_rounds( @@ -81,31 +83,39 @@ where } impl Poseidon2SkinnyChip { + fn eval_input_round( + &self, + builder: &mut AB, + local_row: &Poseidon2, + prep_local: &Poseidon2PreprocessedCols, + next_row: &Poseidon2, + ) { + let mut state: [AB::Expr; WIDTH] = array::from_fn(|i| local_row.state_var[i].into()); + + // Apply the linear layer. + external_linear_layer(&mut state); + + let next_state = next_row.state_var; + for i in 0..WIDTH { + builder + .when_transition() + .when(prep_local.round_counters_preprocessed.is_input_round) + .assert_eq(next_state[i], state[i].clone()); + } + } + fn eval_external_round( &self, builder: &mut AB, local_row: &Poseidon2, prep_local: &Poseidon2PreprocessedCols, next_row: &Poseidon2, - ) where - AB::Var: 'static, - { - let mut state_ell: [AB::Expr; WIDTH] = array::from_fn(|i| local_row.state_var[i].into()); - - // For the first external round, we should apply the external linear layer. - external_linear_layer(&mut state_ell); - - let state: [AB::Expr; WIDTH] = array::from_fn(|i| { - builder.if_else( - prep_local.round_counters_preprocessed.is_first_round, - state_ell[i].clone(), - local_row.state_var[i], - ) - }); + ) { + let local_state = local_row.state_var; // Add the round constants. let add_rc: [AB::Expr; WIDTH] = core::array::from_fn(|i| { - state[i].clone() + prep_local.round_counters_preprocessed.round_constants[i] + local_state[i].into() + prep_local.round_counters_preprocessed.round_constants[i] }); // Apply the sboxes. @@ -136,9 +146,7 @@ impl Poseidon2SkinnyChip { next_row: &Poseidon2, round_constants: [AB::Var; WIDTH], is_internal_row: AB::Var, - ) where - AB::Var: 'static, - { + ) { let local_state = local_row.state_var; let s0 = local_row.internal_rounds_s0; diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs b/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs index 466753def9..02f5e041a8 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/columns/preprocessed.rs @@ -5,7 +5,7 @@ use crate::chips::{mem::MemoryAccessCols, poseidon2_skinny::WIDTH}; #[derive(AlignedBorrow, Clone, Copy, Debug)] #[repr(C)] pub struct RoundCountersPreprocessedCols { - pub is_first_round: T, + pub is_input_round: T, pub is_external_round: T, pub is_internal_round: T, pub round_constants: [T; WIDTH], diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs index 238d757639..0d08be6431 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs @@ -16,14 +16,6 @@ pub const NUM_EXTERNAL_ROUNDS: usize = 8; pub const NUM_INTERNAL_ROUNDS: usize = 13; pub const NUM_ROUNDS: usize = NUM_EXTERNAL_ROUNDS + NUM_INTERNAL_ROUNDS; -pub const fn max(a: usize, b: usize) -> usize { - if a > b { - a - } else { - b - } -} - /// A chip that implements the Poseidon2 permutation in the skinny variant (one external round per /// row and one row for all internal rounds). pub struct Poseidon2SkinnyChip { diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index 07611af173..fcf0aa9aad 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -4,7 +4,6 @@ use std::{ }; use itertools::Itertools; -use p3_air::BaseAir; use p3_field::PrimeField32; use p3_matrix::{dense::RowMajorMatrix, Matrix}; use sp1_core::{air::MachineAir, utils::pad_rows_fixed}; @@ -15,19 +14,21 @@ use crate::{ chips::{ mem::MemoryAccessCols, poseidon2_skinny::{ - columns::Poseidon2, external_linear_layer, Poseidon2SkinnyChip, NUM_EXTERNAL_ROUNDS, - NUM_INTERNAL_ROUNDS, + columns::{Poseidon2, NUM_POSEIDON2_COLS}, + external_linear_layer, Poseidon2SkinnyChip, NUM_EXTERNAL_ROUNDS, NUM_INTERNAL_ROUNDS, }, }, instruction::Instruction::Poseidon2Skinny, ExecutionRecord, RecursionProgram, }; -use super::{columns::preprocessed::Poseidon2PreprocessedCols, internal_linear_layer, max, WIDTH}; +use super::{columns::preprocessed::Poseidon2PreprocessedCols, internal_linear_layer, WIDTH}; const PREPROCESSED_POSEIDON2_WIDTH: usize = size_of::>(); -const INTERNAL_ROUND_IDX: usize = NUM_EXTERNAL_ROUNDS / 2; +const INTERNAL_ROUND_IDX: usize = NUM_EXTERNAL_ROUNDS / 2 + 1; +const INPUT_ROUND_IDX: usize = 0; +const OUTPUT_ROUND_IDX: usize = NUM_EXTERNAL_ROUNDS + 2; impl MachineAir for Poseidon2SkinnyChip { type Record = ExecutionRecord; @@ -46,40 +47,46 @@ impl MachineAir for Poseidon2SkinnyChip ) -> RowMajorMatrix { let mut rows = Vec::new(); - let num_columns = as BaseAir>::width(self); - for event in &input.poseidon2_skinny_events { // We have one row for input, one row for output, NUM_EXTERNAL_ROUNDS rows for the external rounds, // and one row for all internal rounds. - let mut row_add = vec![vec![F::zero(); num_columns]; NUM_EXTERNAL_ROUNDS + 2]; + let mut row_add = [[F::zero(); NUM_POSEIDON2_COLS]; NUM_EXTERNAL_ROUNDS + 3]; - let first_row_cols: &mut Poseidon2 = row_add[0].as_mut_slice().borrow_mut(); - first_row_cols.state_var = event.input; + // The first row should have event.input and [event.input[0].clone(); NUM_INTERNAL_ROUNDS-1] + // in its state columns. The sbox_state will be modified in the computation of the + // first row. + { + let (first_row, second_row) = &mut row_add[0..2].split_at_mut(1); + let input_cols: &mut Poseidon2 = first_row[0].as_mut_slice().borrow_mut(); + input_cols.state_var = event.input; + + let next_cols: &mut Poseidon2 = second_row[0].as_mut_slice().borrow_mut(); + next_cols.state_var = event.input; + external_linear_layer(&mut next_cols.state_var); + } // For each external round, and once for all the internal rounds at the same time, apply // the corresponding operation. This will change the state and internal_rounds_s0 variable // in row r+1. - for r in 0..NUM_EXTERNAL_ROUNDS + 1 { + for i in 1..OUTPUT_ROUND_IDX { let next_state_var = { - let cols: &mut Poseidon2 = row_add[r].as_mut_slice().borrow_mut(); + let cols: &mut Poseidon2 = row_add[i].as_mut_slice().borrow_mut(); let state = cols.state_var; - if r != INTERNAL_ROUND_IDX { - self.populate_external_round(&state, r) + if i != INTERNAL_ROUND_IDX { + self.populate_external_round(&state, i - 1) } else { // Populate the internal rounds. self.populate_internal_rounds(&state, &mut cols.internal_rounds_s0) } }; - let next_row_cols: &mut Poseidon2 = row_add[r + 1].as_mut_slice().borrow_mut(); + let next_row_cols: &mut Poseidon2 = row_add[i + 1].as_mut_slice().borrow_mut(); next_row_cols.state_var = next_state_var; } // Check that the permutation is computed correctly. { - let last_row_cols: &Poseidon2 = - row_add[NUM_EXTERNAL_ROUNDS + 1].as_slice().borrow(); - + let last_row_cols: &Poseidon2 = row_add[OUTPUT_ROUND_IDX].as_slice().borrow(); debug_assert_eq!(last_row_cols.state_var, event.output); } rows.extend(row_add.into_iter()); @@ -90,14 +97,16 @@ impl MachineAir for Poseidon2SkinnyChip // This will need to be adjusted when the AIR constraints are implemented. pad_rows_fixed( &mut rows, - || vec![F::zero(); num_columns], + || [F::zero(); NUM_POSEIDON2_COLS], self.fixed_log2_rows, ); } // Convert the trace to a row major matrix. - let trace = - RowMajorMatrix::new(rows.into_iter().flatten().collect::>(), num_columns); + let trace = RowMajorMatrix::new( + rows.into_iter().flatten().collect::>(), + NUM_POSEIDON2_COLS, + ); #[cfg(debug_assertions)] println!( @@ -131,37 +140,41 @@ impl MachineAir for Poseidon2SkinnyChip let mut rows = vec![ [F::zero(); PREPROCESSED_POSEIDON2_WIDTH]; - num_instructions * (NUM_EXTERNAL_ROUNDS + 2) + num_instructions * (NUM_EXTERNAL_ROUNDS + 3) ]; - // Iterate over the instructions and take NUM_EXTERNAL_ROUNDS + 2 rows for each instruction. - // We have one extra round for the internal rounds and one extra round for the output. + // Iterate over the instructions and take NUM_EXTERNAL_ROUNDS + 3 rows for each instruction. + // We have one extra round for the internal rounds, one extra round for the input, + // and one extra round for the output. instructions - .zip_eq(&rows.iter_mut().chunks(NUM_EXTERNAL_ROUNDS + 2)) + .zip_eq(&rows.iter_mut().chunks(NUM_EXTERNAL_ROUNDS + 3)) .for_each(|(instruction, row_add)| { - row_add.into_iter().enumerate().for_each(|(r, row)| { + row_add.into_iter().enumerate().for_each(|(i, row)| { let cols: &mut Poseidon2PreprocessedCols<_> = (*row).as_mut_slice().borrow_mut(); // Set the round-counter columns. - cols.round_counters_preprocessed.is_first_round = F::from_bool(r == 0); - let is_external_round = r != INTERNAL_ROUND_IDX && r != NUM_EXTERNAL_ROUNDS + 1; + cols.round_counters_preprocessed.is_input_round = + F::from_bool(i == INPUT_ROUND_IDX); + let is_external_round = + i != INPUT_ROUND_IDX && i != INTERNAL_ROUND_IDX && i != OUTPUT_ROUND_IDX; cols.round_counters_preprocessed.is_external_round = F::from_bool(is_external_round); cols.round_counters_preprocessed.is_internal_round = - F::from_bool(r == INTERNAL_ROUND_IDX); + F::from_bool(i == INTERNAL_ROUND_IDX); (0..WIDTH).for_each(|j| { cols.round_counters_preprocessed.round_constants[j] = if is_external_round { - let round = if r < INTERNAL_ROUND_IDX { + let r = i - 1; + let round = if i < INTERNAL_ROUND_IDX { r } else { r + NUM_INTERNAL_ROUNDS - 1 }; F::from_wrapped_u32(RC_16_30_U32[round][j]) - } else if r == INTERNAL_ROUND_IDX { - F::from_wrapped_u32(RC_16_30_U32[INTERNAL_ROUND_IDX + j][0]) + } else if i == INTERNAL_ROUND_IDX { + F::from_wrapped_u32(RC_16_30_U32[NUM_EXTERNAL_ROUNDS / 2 + j][0]) } else { F::zero() }; @@ -169,19 +182,19 @@ impl MachineAir for Poseidon2SkinnyChip // Set the memory columns. We read once, at the first iteration, // and write once, at the last iteration. - if r == 0 { + if i == INPUT_ROUND_IDX { cols.memory_preprocessed = instruction.addrs.input.map(|addr| MemoryAccessCols { addr, read_mult: F::one(), write_mult: F::zero(), }); - } else if r == NUM_EXTERNAL_ROUNDS + 1 { + } else if i == OUTPUT_ROUND_IDX { cols.memory_preprocessed = instruction.addrs.output.map(|addr| MemoryAccessCols { addr, read_mult: F::zero(), - write_mult: instruction.mults[r], + write_mult: instruction.mults[i], }); } }); @@ -210,37 +223,30 @@ impl Poseidon2SkinnyChip { r: usize, ) -> [F; WIDTH] { let mut state = { - // For the first round, apply the external linear layer. - let mut state = *round_state; - if r == 0 { - external_linear_layer(&mut state); - } - // Add round constants. // Optimization: Since adding a constant is a degree 1 operation, we can avoid adding // columns for it, and instead include it in the constraint for the x^3 part of the sbox. - let round = if r < INTERNAL_ROUND_IDX { + let round = if r < NUM_EXTERNAL_ROUNDS / 2 { r } else { r + NUM_INTERNAL_ROUNDS - 1 }; - (0..WIDTH).for_each(|i| state[i] += F::from_wrapped_u32(RC_16_30_U32[round][i])); + let mut add_rc = *round_state; + (0..WIDTH).for_each(|i| add_rc[i] += F::from_wrapped_u32(RC_16_30_U32[round][i])); // Apply the sboxes. // Optimization: since the linear layer that comes after the sbox is degree 1, we can // avoid adding columns for the result of the sbox, and instead include the x^3 -> x^7 // part of the sbox in the constraint for the linear layer let mut sbox_deg_7: [F; 16] = [F::zero(); WIDTH]; - let mut sbox_deg_3: [F; 16] = [F::zero(); WIDTH]; for i in 0..WIDTH { - sbox_deg_3[i] = state[i] * state[i] * state[i]; - sbox_deg_7[i] = sbox_deg_3[i] * sbox_deg_3[i] * state[i]; + let sbox_deg_3 = add_rc[i] * add_rc[i] * add_rc[i]; + sbox_deg_7[i] = sbox_deg_3 * sbox_deg_3 * add_rc[i]; } sbox_deg_7 }; - // Apply the linear layer. external_linear_layer(&mut state); state @@ -252,20 +258,18 @@ impl Poseidon2SkinnyChip { internal_rounds_s0: &mut [F; NUM_INTERNAL_ROUNDS - 1], ) -> [F; WIDTH] { let mut new_state = *state; - let mut sbox_deg_3: [F; max(WIDTH, NUM_INTERNAL_ROUNDS)] = - [F::zero(); max(WIDTH, NUM_INTERNAL_ROUNDS)]; - for r in 0..NUM_INTERNAL_ROUNDS { + (0..NUM_INTERNAL_ROUNDS).for_each(|r| { // Add the round constant to the 0th state element. // Optimization: Since adding a constant is a degree 1 operation, we can avoid adding // columns for it, just like for external rounds. - let round = r + INTERNAL_ROUND_IDX; + let round = r + NUM_EXTERNAL_ROUNDS / 2; let add_rc = new_state[0] + F::from_wrapped_u32(RC_16_30_U32[round][0]); // Apply the sboxes. // Optimization: since the linear layer that comes after the sbox is degree 1, we can // avoid adding columns for the result of the sbox, just like for external rounds. - sbox_deg_3[r] = add_rc * add_rc * add_rc; - let sbox_deg_7 = sbox_deg_3[r] * sbox_deg_3[r] * add_rc; + let sbox_deg_3 = add_rc * add_rc * add_rc; + let sbox_deg_7 = sbox_deg_3 * sbox_deg_3 * add_rc; // Apply the linear layer. new_state[0] = sbox_deg_7; @@ -279,7 +283,7 @@ impl Poseidon2SkinnyChip { if r < NUM_INTERNAL_ROUNDS - 1 { internal_rounds_s0[r] = new_state[0]; } - } + }); new_state } From 072aeaf749c4c87cd39b96f0db63a07d3db3bc28 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 16:54:53 -0700 Subject: [PATCH 10/18] merged the p2 skinny and wide non air components --- recursion/circuit-v2/src/challenger.rs | 2 +- recursion/compiler/src/circuit/builder.rs | 21 ++----- recursion/compiler/src/circuit/compiler.rs | 55 ++----------------- recursion/compiler/src/ir/instructions.rs | 4 +- .../core-v2/src/chips/poseidon2_skinny/mod.rs | 4 +- .../src/chips/poseidon2_skinny/trace.rs | 30 +++++----- .../core-v2/src/chips/poseidon2_wide/mod.rs | 4 +- .../core-v2/src/chips/poseidon2_wide/trace.rs | 22 ++++---- recursion/core-v2/src/lib.rs | 4 +- recursion/core-v2/src/runtime/instruction.rs | 21 +------ recursion/core-v2/src/runtime/mod.rs | 27 +-------- recursion/core-v2/src/runtime/record.rs | 9 +-- 12 files changed, 55 insertions(+), 148 deletions(-) diff --git a/recursion/circuit-v2/src/challenger.rs b/recursion/circuit-v2/src/challenger.rs index fe627a570a..448c5fca03 100644 --- a/recursion/circuit-v2/src/challenger.rs +++ b/recursion/circuit-v2/src/challenger.rs @@ -96,7 +96,7 @@ impl DuplexChallengerVariable { self.sponge_state[0..self.input_buffer.len()].copy_from_slice(self.input_buffer.as_slice()); self.input_buffer.clear(); - self.sponge_state = builder.poseidon2_permute_v2_wide(self.sponge_state); + self.sponge_state = builder.poseidon2_permute_v2(self.sponge_state); self.output_buffer.clear(); self.output_buffer.extend_from_slice(&self.sponge_state); diff --git a/recursion/compiler/src/circuit/builder.rs b/recursion/compiler/src/circuit/builder.rs index 7affd03979..240ae34482 100644 --- a/recursion/compiler/src/circuit/builder.rs +++ b/recursion/compiler/src/circuit/builder.rs @@ -11,8 +11,7 @@ pub trait CircuitV2Builder { fn num2bits_v2_f(&mut self, num: Felt) -> Vec>; fn exp_reverse_bits_v2(&mut self, input: Felt, power_bits: Vec>) -> Felt; - fn poseidon2_permute_v2_skinny(&mut self, state: [Felt; WIDTH]) -> [Felt; WIDTH]; - fn poseidon2_permute_v2_wide(&mut self, state: [Felt; WIDTH]) -> [Felt; WIDTH]; + fn poseidon2_permute_v2(&mut self, state: [Felt; WIDTH]) -> [Felt; WIDTH]; fn poseidon2_hash_v2(&mut self, array: &[Felt]) -> [Felt; DIGEST_SIZE]; fn poseidon2_compress_v2( &mut self, @@ -55,21 +54,13 @@ impl CircuitV2Builder for Builder { output } /// Applies the Poseidon2 permutation to the given array. - fn poseidon2_permute_v2_skinny(&mut self, array: [Felt; WIDTH]) -> [Felt; WIDTH] { + fn poseidon2_permute_v2(&mut self, array: [Felt; WIDTH]) -> [Felt; WIDTH] { let output: [Felt; WIDTH] = core::array::from_fn(|_| self.uninit()); self.operations - .push(DslIr::CircuitV2Poseidon2PermuteBabyBearSkinny( - output, array, - )); - output - } - /// Applies the Poseidon2 permutation to the given array using the wide precompile. - fn poseidon2_permute_v2_wide(&mut self, array: [Felt; WIDTH]) -> [Felt; WIDTH] { - let output: [Felt; WIDTH] = core::array::from_fn(|_| self.uninit()); - self.operations - .push(DslIr::CircuitV2Poseidon2PermuteBabyBearWide(output, array)); + .push(DslIr::CircuitV2Poseidon2PermuteBabyBear(output, array)); output } + /// Applies the Poseidon2 permutation to the given array. /// /// Reference: [p3_symmetric::PaddingFreeSponge] @@ -78,7 +69,7 @@ impl CircuitV2Builder for Builder { let mut state = core::array::from_fn(|_| self.eval(C::F::zero())); for input_chunk in input.chunks(HASH_RATE) { state[..input_chunk.len()].copy_from_slice(input_chunk); - state = self.poseidon2_permute_v2_skinny(state); + state = self.poseidon2_permute_v2(state); } let state: [Felt; DIGEST_SIZE] = state[..DIGEST_SIZE].try_into().unwrap(); state @@ -93,7 +84,7 @@ impl CircuitV2Builder for Builder { // debug_assert!(DIGEST_SIZE * N <= WIDTH); let mut pre_iter = input.into_iter().chain(repeat(self.eval(C::F::default()))); let pre = core::array::from_fn(move |_| pre_iter.next().unwrap()); - let post = self.poseidon2_permute_v2_skinny(pre); + let post = self.poseidon2_permute_v2(pre); let post: [Felt; DIGEST_SIZE] = post[..DIGEST_SIZE].try_into().unwrap(); post } diff --git a/recursion/compiler/src/circuit/compiler.rs b/recursion/compiler/src/circuit/compiler.rs index cfaa0a9651..c8fccf6d35 100644 --- a/recursion/compiler/src/circuit/compiler.rs +++ b/recursion/compiler/src/circuit/compiler.rs @@ -227,26 +227,12 @@ impl AsmCompiler { ] } - fn poseidon2_permute_skinny( + fn poseidon2_permute( &mut self, dst: [impl Reg; WIDTH], src: [impl Reg; WIDTH], ) -> Instruction { - Instruction::Poseidon2Skinny(Poseidon2WideInstr { - addrs: Poseidon2Io { - input: src.map(|r| r.read(self)), - output: dst.map(|r| r.write(self)), - }, - mults: [C::F::zero(); WIDTH], - }) - } - - fn poseidon2_permute_wide( - &mut self, - dst: [impl Reg; WIDTH], - src: [impl Reg; WIDTH], - ) -> Instruction { - Instruction::Poseidon2Wide(Poseidon2WideInstr { + Instruction::Poseidon2(Poseidon2Instr { addrs: Poseidon2Io { input: src.map(|r| r.read(self)), output: dst.map(|r| r.write(self)), @@ -422,11 +408,8 @@ impl AsmCompiler { DslIr::AssertNeFI(lhs, rhs) => self.base_assert_ne(lhs, Imm::F(rhs)), DslIr::AssertNeEI(lhs, rhs) => self.ext_assert_ne(lhs, Imm::EF(rhs)), - DslIr::CircuitV2Poseidon2PermuteBabyBearSkinny(dst, src) => { - vec![self.poseidon2_permute_skinny(dst, src)] - } - DslIr::CircuitV2Poseidon2PermuteBabyBearWide(dst, src) => { - vec![self.poseidon2_permute_wide(dst, src)] + DslIr::CircuitV2Poseidon2PermuteBabyBear(dst, src) => { + vec![self.poseidon2_permute(dst, src)] } DslIr::CircuitV2ExpReverseBits(dst, base, exp) => { vec![self.exp_reverse_bits(dst, base, exp)] @@ -521,11 +504,7 @@ impl AsmCompiler { kind: MemAccessKind::Write, .. }) => vec![(mult, inner)], - Instruction::Poseidon2Skinny(Poseidon2SkinnyInstr { - addrs: Poseidon2Io { ref output, .. }, - mults, - }) => mults.iter_mut().zip(output).collect(), - Instruction::Poseidon2Wide(Poseidon2WideInstr { + Instruction::Poseidon2(Poseidon2SkinnyInstr { addrs: Poseidon2Io { ref output, .. }, mults, }) => mults.iter_mut().zip(output).collect(), @@ -779,29 +758,7 @@ mod tests { let output_1 = inner_perm().permute(input_1); let input_1_felts = input_1.map(|x| builder.eval(x)); - let output_1_felts = builder.poseidon2_permute_v2_skinny(input_1_felts); - let expected: [Felt<_>; WIDTH] = output_1.map(|x| builder.eval(x)); - for (lhs, rhs) in output_1_felts.into_iter().zip(expected) { - builder.assert_felt_eq(lhs, rhs); - } - } - - test_operations(builder.operations); - } - - #[test] - fn test_poseidon2_wide() { - setup_logger(); - - let mut builder = AsmBuilder::::default(); - let mut rng = StdRng::seed_from_u64(0xCAFEDA7E) - .sample_iter::<[F; WIDTH], _>(rand::distributions::Standard); - for _ in 0..100 { - let input_1: [F; WIDTH] = rng.next().unwrap(); - let output_1 = inner_perm().permute(input_1); - - let input_1_felts = input_1.map(|x| builder.eval(x)); - let output_1_felts = builder.poseidon2_permute_v2_wide(input_1_felts); + let output_1_felts = builder.poseidon2_permute_v2(input_1_felts); let expected: [Felt<_>; WIDTH] = output_1.map(|x| builder.eval(x)); for (lhs, rhs) in output_1_felts.into_iter().zip(expected) { builder.assert_felt_eq(lhs, rhs); diff --git a/recursion/compiler/src/ir/instructions.rs b/recursion/compiler/src/ir/instructions.rs index a135a3b8eb..6f110880a9 100644 --- a/recursion/compiler/src/ir/instructions.rs +++ b/recursion/compiler/src/ir/instructions.rs @@ -210,9 +210,7 @@ pub enum DslIr { /// Permutates an array of BabyBear elements in the circuit. CircuitPoseidon2PermuteBabyBear([Felt; 16]), /// Permutates an array of BabyBear elements in the circuit using the skinny precompile. - CircuitV2Poseidon2PermuteBabyBearSkinny([Felt; 16], [Felt; 16]), - /// Permutates an array of BabyBear elements in the circuit using the wide precompile. - CircuitV2Poseidon2PermuteBabyBearWide([Felt; 16], [Felt; 16]), + CircuitV2Poseidon2PermuteBabyBear([Felt; 16], [Felt; 16]), // Miscellaneous instructions. /// Decompose hint operation of a usize into an array. (output = num2bits(usize)). diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs index 0d08be6431..d6fee33832 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs @@ -119,7 +119,7 @@ pub(crate) mod tests { let instructions = (0..WIDTH) .map(|i| instr::mem(MemAccessKind::Write, 1, i as u32, input[i])) - .chain(once(instr::poseidon2_skinny( + .chain(once(instr::poseidon2( [1; WIDTH], std::array::from_fn(|i| (i + WIDTH) as u32), std::array::from_fn(|i| i as u32), @@ -131,7 +131,7 @@ pub(crate) mod tests { .chain((0..WIDTH).map(|i| { instr::mem(MemAccessKind::Write, 1, (2 * WIDTH + i) as u32, input_1[i]) })) - .chain(once(instr::poseidon2_skinny( + .chain(once(instr::poseidon2( [1; WIDTH], std::array::from_fn(|i| (i + 3 * WIDTH) as u32), std::array::from_fn(|i| (i + 2 * WIDTH) as u32), diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index fcf0aa9aad..64913d0dc1 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -14,11 +14,11 @@ use crate::{ chips::{ mem::MemoryAccessCols, poseidon2_skinny::{ - columns::{Poseidon2, NUM_POSEIDON2_COLS}, + columns::{Poseidon2 as Poseidon2Cols, NUM_POSEIDON2_COLS}, external_linear_layer, Poseidon2SkinnyChip, NUM_EXTERNAL_ROUNDS, NUM_INTERNAL_ROUNDS, }, }, - instruction::Instruction::Poseidon2Skinny, + instruction::Instruction::Poseidon2, ExecutionRecord, RecursionProgram, }; @@ -39,7 +39,7 @@ impl MachineAir for Poseidon2SkinnyChip format!("Poseidon2Skinny {}", DEGREE) } - #[instrument(name = "generate poseidon2 skinny trace", level = "debug", skip_all, fields(rows = input.poseidon2_skinny_events.len()))] + #[instrument(name = "generate poseidon2 skinny trace", level = "debug", skip_all, fields(rows = input.poseidon2_events.len()))] fn generate_trace( &self, input: &ExecutionRecord, @@ -47,7 +47,7 @@ impl MachineAir for Poseidon2SkinnyChip ) -> RowMajorMatrix { let mut rows = Vec::new(); - for event in &input.poseidon2_skinny_events { + for event in &input.poseidon2_events { // We have one row for input, one row for output, NUM_EXTERNAL_ROUNDS rows for the external rounds, // and one row for all internal rounds. let mut row_add = [[F::zero(); NUM_POSEIDON2_COLS]; NUM_EXTERNAL_ROUNDS + 3]; @@ -57,10 +57,10 @@ impl MachineAir for Poseidon2SkinnyChip // first row. { let (first_row, second_row) = &mut row_add[0..2].split_at_mut(1); - let input_cols: &mut Poseidon2 = first_row[0].as_mut_slice().borrow_mut(); + let input_cols: &mut Poseidon2Cols = first_row[0].as_mut_slice().borrow_mut(); input_cols.state_var = event.input; - let next_cols: &mut Poseidon2 = second_row[0].as_mut_slice().borrow_mut(); + let next_cols: &mut Poseidon2Cols = second_row[0].as_mut_slice().borrow_mut(); next_cols.state_var = event.input; external_linear_layer(&mut next_cols.state_var); } @@ -70,7 +70,7 @@ impl MachineAir for Poseidon2SkinnyChip // in row r+1. for i in 1..OUTPUT_ROUND_IDX { let next_state_var = { - let cols: &mut Poseidon2 = row_add[i].as_mut_slice().borrow_mut(); + let cols: &mut Poseidon2Cols = row_add[i].as_mut_slice().borrow_mut(); let state = cols.state_var; if i != INTERNAL_ROUND_IDX { @@ -80,13 +80,15 @@ impl MachineAir for Poseidon2SkinnyChip self.populate_internal_rounds(&state, &mut cols.internal_rounds_s0) } }; - let next_row_cols: &mut Poseidon2 = row_add[i + 1].as_mut_slice().borrow_mut(); + let next_row_cols: &mut Poseidon2Cols = + row_add[i + 1].as_mut_slice().borrow_mut(); next_row_cols.state_var = next_state_var; } // Check that the permutation is computed correctly. { - let last_row_cols: &Poseidon2 = row_add[OUTPUT_ROUND_IDX].as_slice().borrow(); + let last_row_cols: &Poseidon2Cols = + row_add[OUTPUT_ROUND_IDX].as_slice().borrow(); debug_assert_eq!(last_row_cols.state_var, event.output); } rows.extend(row_add.into_iter()); @@ -132,7 +134,7 @@ impl MachineAir for Poseidon2SkinnyChip .instructions .iter() .filter_map(|instruction| match instruction { - Poseidon2Skinny(instr) => Some(instr), + Poseidon2(instr) => Some(instr), _ => None, }); @@ -301,7 +303,7 @@ mod tests { use crate::{ chips::poseidon2_skinny::{Poseidon2SkinnyChip, WIDTH}, - ExecutionRecord, Poseidon2SkinnyEvent, + ExecutionRecord, Poseidon2Event, }; #[test] @@ -315,12 +317,12 @@ mod tests { let input_1 = [F::rand(&mut rng); WIDTH]; let output_1 = permuter.permute(input_1); let shard = ExecutionRecord { - poseidon2_skinny_events: vec![ - Poseidon2SkinnyEvent { + poseidon2_events: vec![ + Poseidon2Event { input: input_0, output: output_0, }, - Poseidon2SkinnyEvent { + Poseidon2Event { input: input_1, output: output_1, }, diff --git a/recursion/core-v2/src/chips/poseidon2_wide/mod.rs b/recursion/core-v2/src/chips/poseidon2_wide/mod.rs index d41b928a56..c269177e99 100644 --- a/recursion/core-v2/src/chips/poseidon2_wide/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_wide/mod.rs @@ -151,7 +151,7 @@ pub(crate) mod tests { let instructions = (0..WIDTH) .map(|i| instr::mem(MemAccessKind::Write, 1, i as u32, input[i])) - .chain(once(instr::poseidon2_wide( + .chain(once(instr::poseidon2( [1; WIDTH], std::array::from_fn(|i| (i + WIDTH) as u32), std::array::from_fn(|i| i as u32), @@ -163,7 +163,7 @@ pub(crate) mod tests { .chain((0..WIDTH).map(|i| { instr::mem(MemAccessKind::Write, 1, (2 * WIDTH + i) as u32, input_1[i]) })) - .chain(once(instr::poseidon2_wide( + .chain(once(instr::poseidon2( [1; WIDTH], std::array::from_fn(|i| (i + 3 * WIDTH) as u32), std::array::from_fn(|i| (i + 2 * WIDTH) as u32), diff --git a/recursion/core-v2/src/chips/poseidon2_wide/trace.rs b/recursion/core-v2/src/chips/poseidon2_wide/trace.rs index 0941cb579f..fa5de2e7b2 100644 --- a/recursion/core-v2/src/chips/poseidon2_wide/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_wide/trace.rs @@ -16,7 +16,7 @@ use crate::{ NUM_EXTERNAL_ROUNDS, WIDTH, }, }, - instruction::Instruction::Poseidon2Wide, + instruction::Instruction::Poseidon2, ExecutionRecord, RecursionProgram, }; @@ -36,7 +36,7 @@ impl MachineAir for Poseidon2WideChip, @@ -46,7 +46,7 @@ impl MachineAir for Poseidon2WideChip as BaseAir>::width(self); - for event in &input.poseidon2_wide_events { + for event in &input.poseidon2_events { let mut input_row = vec![F::zero(); num_columns]; { @@ -138,7 +138,7 @@ impl MachineAir for Poseidon2WideChip Some(instr), + Poseidon2(instr) => Some(instr), _ => None, }); @@ -291,7 +291,7 @@ mod tests { use crate::{ chips::poseidon2_wide::{Poseidon2WideChip, WIDTH}, - ExecutionRecord, Poseidon2SkinnyEvent, + ExecutionRecord, Poseidon2Event, }; #[test] @@ -306,12 +306,12 @@ mod tests { let output_1 = permuter.permute(input_1); let shard = ExecutionRecord { - poseidon2_wide_events: vec![ - Poseidon2SkinnyEvent { + poseidon2_events: vec![ + Poseidon2Event { input: input_0, output: output_0, }, - Poseidon2SkinnyEvent { + Poseidon2Event { input: input_1, output: output_1, }, @@ -334,12 +334,12 @@ mod tests { let output_1 = permuter.permute(input_1); let shard = ExecutionRecord { - poseidon2_wide_events: vec![ - Poseidon2SkinnyEvent { + poseidon2_events: vec![ + Poseidon2Event { input: input_0, output: output_0, }, - Poseidon2SkinnyEvent { + Poseidon2Event { input: input_1, output: output_1, }, diff --git a/recursion/core-v2/src/lib.rs b/recursion/core-v2/src/lib.rs index 05afb74e5b..3c134d779a 100644 --- a/recursion/core-v2/src/lib.rs +++ b/recursion/core-v2/src/lib.rs @@ -97,7 +97,7 @@ pub struct Poseidon2SkinnyInstr { pub mults: [F; WIDTH], } -pub type Poseidon2SkinnyEvent = Poseidon2Io; +pub type Poseidon2Event = Poseidon2Io; /// The inputs and outputs to an exp-reverse-bits operation. #[derive(Clone, Debug, PartialEq, Eq, Serialize, Deserialize)] @@ -109,7 +109,7 @@ pub struct ExpReverseBitsIo { } pub type Poseidon2WideEvent = Poseidon2Io; -pub type Poseidon2WideInstr = Poseidon2SkinnyInstr; +pub type Poseidon2Instr = Poseidon2SkinnyInstr; /// An instruction invoking the exp-reverse-bits operation. #[derive(Clone, Debug, Serialize, Deserialize)] diff --git a/recursion/core-v2/src/runtime/instruction.rs b/recursion/core-v2/src/runtime/instruction.rs index 757ae8e94f..af8576e2e4 100644 --- a/recursion/core-v2/src/runtime/instruction.rs +++ b/recursion/core-v2/src/runtime/instruction.rs @@ -8,8 +8,7 @@ pub enum Instruction { BaseAlu(BaseAluInstr), ExtAlu(ExtAluInstr), Mem(MemInstr), - Poseidon2Skinny(Poseidon2SkinnyInstr), - Poseidon2Wide(Poseidon2WideInstr), + Poseidon2(Poseidon2Instr), ExpReverseBitsLen(ExpReverseBitsInstr), HintBits(HintBitsInstr), FriFold(FriFoldInstr), @@ -124,26 +123,12 @@ pub fn mem_block( }) } -pub fn poseidon2_skinny( +pub fn poseidon2( mults: [u32; WIDTH], output: [u32; WIDTH], input: [u32; WIDTH], ) -> Instruction { - Instruction::Poseidon2Skinny(Poseidon2SkinnyInstr { - mults: mults.map(F::from_canonical_u32), - addrs: Poseidon2Io { - output: output.map(F::from_canonical_u32).map(Address), - input: input.map(F::from_canonical_u32).map(Address), - }, - }) -} - -pub fn poseidon2_wide( - mults: [u32; WIDTH], - output: [u32; WIDTH], - input: [u32; WIDTH], -) -> Instruction { - Instruction::Poseidon2Wide(Poseidon2WideInstr { + Instruction::Poseidon2(Poseidon2Instr { mults: mults.map(F::from_canonical_u32), addrs: Poseidon2Io { output: output.map(F::from_canonical_u32).map(Address), diff --git a/recursion/core-v2/src/runtime/mod.rs b/recursion/core-v2/src/runtime/mod.rs index e8e8384585..3d9ee683d4 100644 --- a/recursion/core-v2/src/runtime/mod.rs +++ b/recursion/core-v2/src/runtime/mod.rs @@ -369,7 +369,7 @@ where } self.record.mem_events.push(MemEvent { inner: val }); } - Instruction::Poseidon2Skinny(Poseidon2SkinnyInstr { + Instruction::Poseidon2(Poseidon2Instr { addrs: Poseidon2Io { input, output }, mults, }) => { @@ -384,30 +384,7 @@ where .for_each(|((&val, addr), mult)| { self.mw(addr, Block::from(val), mult); }); - self.record - .poseidon2_skinny_events - .push(Poseidon2SkinnyEvent { - input: in_vals, - output: perm_output, - }); - } - - Instruction::Poseidon2Wide(Poseidon2WideInstr { - addrs: Poseidon2Io { input, output }, - mults, - }) => { - self.nb_wide_poseidons += 1; - let in_vals = std::array::from_fn(|i| self.mr(input[i]).val[0]); - let perm_output = self.perm.as_ref().unwrap().permute(in_vals); - - perm_output - .iter() - .zip(output) - .zip(mults) - .for_each(|((&val, addr), mult)| { - self.mw(addr, Block::from(val), mult); - }); - self.record.poseidon2_wide_events.push(Poseidon2WideEvent { + self.record.poseidon2_events.push(Poseidon2Event { input: in_vals, output: perm_output, }); diff --git a/recursion/core-v2/src/runtime/record.rs b/recursion/core-v2/src/runtime/record.rs index d76e95b94c..0e24a9fa99 100644 --- a/recursion/core-v2/src/runtime/record.rs +++ b/recursion/core-v2/src/runtime/record.rs @@ -18,8 +18,7 @@ pub struct ExecutionRecord { /// The public values. pub public_values: PublicValues, - pub poseidon2_skinny_events: Vec>, - pub poseidon2_wide_events: Vec>, + pub poseidon2_events: Vec>, pub exp_reverse_bits_len_events: Vec>, pub fri_fold_events: Vec>, } @@ -40,16 +39,14 @@ impl MachineRecord for ExecutionRecord { ext_alu_events, mem_events, public_values: _, - poseidon2_wide_events, - poseidon2_skinny_events, + poseidon2_events, exp_reverse_bits_len_events, fri_fold_events, } = self; base_alu_events.append(&mut other.base_alu_events); ext_alu_events.append(&mut other.ext_alu_events); mem_events.append(&mut other.mem_events); - poseidon2_wide_events.append(&mut other.poseidon2_wide_events); - poseidon2_skinny_events.append(&mut other.poseidon2_skinny_events); + poseidon2_events.append(&mut other.poseidon2_events); exp_reverse_bits_len_events.append(&mut other.exp_reverse_bits_len_events); fri_fold_events.append(&mut other.fri_fold_events); } From 274888135c019777ab84146f62e6e9b09b1abc14 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 19:07:00 -0700 Subject: [PATCH 11/18] bug fix --- .../core-v2/src/chips/poseidon2_skinny/trace.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs index 64913d0dc1..6037876622 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/trace.rs @@ -1,4 +1,5 @@ use std::{ + array, borrow::{Borrow, BorrowMut}, mem::size_of, }; @@ -192,12 +193,11 @@ impl MachineAir for Poseidon2SkinnyChip write_mult: F::zero(), }); } else if i == OUTPUT_ROUND_IDX { - cols.memory_preprocessed = - instruction.addrs.output.map(|addr| MemoryAccessCols { - addr, - read_mult: F::zero(), - write_mult: instruction.mults[i], - }); + cols.memory_preprocessed = array::from_fn(|i| MemoryAccessCols { + addr: instruction.addrs.output[i], + read_mult: F::zero(), + write_mult: instruction.mults[i], + }); } }); }); From 92506cd098874bb5efba356c85f4bb398fbc986f Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 19:10:37 -0700 Subject: [PATCH 12/18] uncomment a machine test --- recursion/circuit-v2/src/challenger.rs | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/recursion/circuit-v2/src/challenger.rs b/recursion/circuit-v2/src/challenger.rs index 448c5fca03..305600ae10 100644 --- a/recursion/circuit-v2/src/challenger.rs +++ b/recursion/circuit-v2/src/challenger.rs @@ -256,12 +256,12 @@ mod tests { let records = vec![runtime.record]; // Run with the poseidon2 wide chip. - // let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(BabyBearPoseidon2::default()); - // let (pk, vk) = wide_machine.setup(&program); - // let result = run_test_machine(records.clone(), wide_machine, pk, vk); - // if let Err(e) = result { - // panic!("Verification failed: {:?}", e); - // } + let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(BabyBearPoseidon2::default()); + let (pk, vk) = wide_machine.setup(&program); + let result = run_test_machine(records.clone(), wide_machine, pk, vk); + if let Err(e) = result { + panic!("Verification failed: {:?}", e); + } // Run with the poseidon2 skinny chip. let skinny_machine = RecursionAir::<_, 9, 0>::machine(BabyBearPoseidon2::compressed()); From 39e507f256a75bdd648b84e01198da35fc72b702 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 19:14:30 -0700 Subject: [PATCH 13/18] allow for higher degrees --- recursion/core-v2/src/chips/poseidon2_skinny/air.rs | 4 ++-- recursion/core-v2/src/chips/poseidon2_skinny/mod.rs | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs index 882662220d..17db9b24b2 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/air.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/air.rs @@ -20,7 +20,7 @@ use super::{ impl BaseAir for Poseidon2SkinnyChip { fn width(&self) -> usize { // We only support machines with degree 9. - assert!(DEGREE == 9); + assert!(DEGREE >= 9); NUM_POSEIDON2_COLS } } @@ -32,7 +32,7 @@ where { fn eval(&self, builder: &mut AB) { // We only support machines with degree 9. - assert!(DEGREE == 9); + assert!(DEGREE >= 9); let main = builder.main(); let (local_row, next_row) = (main.row_slice(0), main.row_slice(1)); diff --git a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs index d6fee33832..6cd418eea2 100644 --- a/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs +++ b/recursion/core-v2/src/chips/poseidon2_skinny/mod.rs @@ -26,7 +26,7 @@ pub struct Poseidon2SkinnyChip { impl Default for Poseidon2SkinnyChip { fn default() -> Self { // We only support machines with degree 9. - assert!(DEGREE == 9); + assert!(DEGREE >= 9); Self { fixed_log2_rows: None, pad: true, From 7d376659a81cffa1f781831ad2a0ec3f5e1cd586 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Wed, 31 Jul 2024 19:24:52 -0700 Subject: [PATCH 14/18] change config to get unit tests to work --- recursion/core-v2/src/chips/alu_base.rs | 2 +- recursion/core-v2/src/chips/alu_ext.rs | 2 +- recursion/core-v2/src/chips/exp_reverse_bits.rs | 2 +- recursion/core-v2/src/chips/fri_fold.rs | 2 +- recursion/core-v2/src/chips/mem.rs | 2 +- recursion/core-v2/src/machine.rs | 2 +- 6 files changed, 6 insertions(+), 6 deletions(-) diff --git a/recursion/core-v2/src/chips/alu_base.rs b/recursion/core-v2/src/chips/alu_base.rs index 9583f8b3fe..910d097196 100644 --- a/recursion/core-v2/src/chips/alu_base.rs +++ b/recursion/core-v2/src/chips/alu_base.rs @@ -279,7 +279,7 @@ mod tests { let mut runtime = Runtime::::new(&program, BabyBearPoseidon2::new().perm); runtime.run().unwrap(); - let machine = A::machine(config); + let machine = A::machine_wide(config); let (pk, vk) = machine.setup(&program); let result = run_test_machine(vec![runtime.record], machine, pk, vk); if let Err(e) = result { diff --git a/recursion/core-v2/src/chips/alu_ext.rs b/recursion/core-v2/src/chips/alu_ext.rs index f318ce9042..f6a3b06cb4 100644 --- a/recursion/core-v2/src/chips/alu_ext.rs +++ b/recursion/core-v2/src/chips/alu_ext.rs @@ -265,7 +265,7 @@ mod tests { runtime.run().unwrap(); let config = SC::new(); - let machine = A::machine(config); + let machine = A::machine_wide(config); let (pk, vk) = machine.setup(&program); let result = run_test_machine(vec![runtime.record], machine, pk, vk); if let Err(e) = result { diff --git a/recursion/core-v2/src/chips/exp_reverse_bits.rs b/recursion/core-v2/src/chips/exp_reverse_bits.rs index 5a16fdd3a6..68e567e0f3 100644 --- a/recursion/core-v2/src/chips/exp_reverse_bits.rs +++ b/recursion/core-v2/src/chips/exp_reverse_bits.rs @@ -439,7 +439,7 @@ mod tests { let mut runtime = Runtime::::new(&program, BabyBearPoseidon2::new().perm); runtime.run().unwrap(); - let machine = A::machine(config); + let machine = A::machine_wide(config); let (pk, vk) = machine.setup(&program); let result = run_test_machine(vec![runtime.record], machine, pk, vk); if let Err(e) = result { diff --git a/recursion/core-v2/src/chips/fri_fold.rs b/recursion/core-v2/src/chips/fri_fold.rs index eef9bd2c6d..c4a77e932d 100644 --- a/recursion/core-v2/src/chips/fri_fold.rs +++ b/recursion/core-v2/src/chips/fri_fold.rs @@ -577,7 +577,7 @@ mod tests { let mut runtime = Runtime::::new(&program, BabyBearPoseidon2::new().perm); runtime.run().unwrap(); - let machine = A::machine(config); + let machine = A::machine_wide(config); let (pk, vk) = machine.setup(&program); let result = run_test_machine(vec![runtime.record], machine, pk, vk); if let Err(e) = result { diff --git a/recursion/core-v2/src/chips/mem.rs b/recursion/core-v2/src/chips/mem.rs index 5d883b6d0c..72f5f6e43a 100644 --- a/recursion/core-v2/src/chips/mem.rs +++ b/recursion/core-v2/src/chips/mem.rs @@ -222,7 +222,7 @@ mod tests { type SC = BabyBearPoseidon2Outer; type F = ::Val; type EF = ::Challenge; - type A = RecursionAir; + type A = RecursionAir; pub fn prove_program(program: RecursionProgram) { let mut runtime = Runtime::::new( diff --git a/recursion/core-v2/src/machine.rs b/recursion/core-v2/src/machine.rs index 4100a68ca2..ec853fc43f 100644 --- a/recursion/core-v2/src/machine.rs +++ b/recursion/core-v2/src/machine.rs @@ -238,7 +238,7 @@ mod tests { type SC = BabyBearPoseidon2Outer; type F = ::Val; type EF = ::Challenge; - type A = RecursionAir; + type A = RecursionAir; fn test_instructions(instructions: Vec>) { let program = RecursionProgram { From b1c84516d7573b09a4001ea11ba0657ef94d915b Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Thu, 1 Aug 2024 10:15:59 -0700 Subject: [PATCH 15/18] some cleanup of the recursion core-v2 tests --- recursion/compiler/src/circuit/compiler.rs | 19 ++++-- recursion/core-v2/src/chips/alu_base.rs | 27 ++------ recursion/core-v2/src/chips/alu_ext.rs | 27 ++------ .../core-v2/src/chips/exp_reverse_bits.rs | 20 +----- recursion/core-v2/src/chips/fri_fold.rs | 20 +----- recursion/core-v2/src/chips/mem.rs | 56 ++++++--------- recursion/core-v2/src/machine.rs | 68 ++++++++----------- 7 files changed, 75 insertions(+), 162 deletions(-) diff --git a/recursion/compiler/src/circuit/compiler.rs b/recursion/compiler/src/circuit/compiler.rs index c8fccf6d35..7b53f5686a 100644 --- a/recursion/compiler/src/circuit/compiler.rs +++ b/recursion/compiler/src/circuit/compiler.rs @@ -717,7 +717,6 @@ mod tests { type SC = BabyBearPoseidon2; type F = ::Val; type EF = ::Challenge; - type A = RecursionAir; fn test_operations(operations: TracedVec>>) { test_operations_with_runner(operations, |program| { let mut runtime = Runtime::::new( @@ -737,17 +736,25 @@ mod tests { let program = compiler.compile(operations); let record = run(&program); - let config = SC::new(); - let machine = A::machine_with_all_chips(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![record], machine, pk, vk); + // Run with the poseidon2 wide chip. + let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(BabyBearPoseidon2::default()); + let (pk, vk) = wide_machine.setup(&program); + let result = run_test_machine(vec![record.clone()], wide_machine, pk, vk); + if let Err(e) = result { + panic!("Verification failed: {:?}", e); + } + + // Run with the poseidon2 skinny chip. + let skinny_machine = RecursionAir::<_, 9, 0>::machine(BabyBearPoseidon2::compressed()); + let (pk, vk) = skinny_machine.setup(&program); + let result = run_test_machine(vec![record.clone()], skinny_machine, pk, vk); if let Err(e) = result { panic!("Verification failed: {:?}", e); } } #[test] - fn test_poseidon2_skinny() { + fn test_poseidon2() { setup_logger(); let mut builder = AsmBuilder::::default(); diff --git a/recursion/core-v2/src/chips/alu_base.rs b/recursion/core-v2/src/chips/alu_base.rs index 910d097196..946d325069 100644 --- a/recursion/core-v2/src/chips/alu_base.rs +++ b/recursion/core-v2/src/chips/alu_base.rs @@ -201,18 +201,13 @@ where #[cfg(test)] mod tests { - use machine::RecursionAir; - use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear}; + use machine::tests::run_recursion_test_machines; + use p3_baby_bear::BabyBear; use p3_field::AbstractField; use p3_matrix::dense::RowMajorMatrix; use rand::{rngs::StdRng, Rng, SeedableRng}; - use sp1_core::{ - air::MachineAir, - stark::StarkGenericConfig, - utils::{run_test_machine, BabyBearPoseidon2}, - }; - use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer; + use sp1_core::{air::MachineAir, stark::StarkGenericConfig, utils::BabyBearPoseidon2}; use super::*; @@ -237,10 +232,8 @@ mod tests { #[test] pub fn four_ops() { - type SC = BabyBearPoseidon2Outer; + type SC = BabyBearPoseidon2; type F = ::Val; - type EF = ::Challenge; - type A = RecursionAir; let mut rng = StdRng::seed_from_u64(0xDEADBEEF); let mut random_felt = move || -> F { rng.sample(rand::distributions::Standard) }; @@ -274,16 +267,6 @@ mod tests { traces: Default::default(), }; - let config = SC::new(); - - let mut runtime = - Runtime::::new(&program, BabyBearPoseidon2::new().perm); - runtime.run().unwrap(); - let machine = A::machine_wide(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![runtime.record], machine, pk, vk); - if let Err(e) = result { - panic!("Verification failed: {:?}", e); - } + run_recursion_test_machines(program); } } diff --git a/recursion/core-v2/src/chips/alu_ext.rs b/recursion/core-v2/src/chips/alu_ext.rs index f6a3b06cb4..4221f53565 100644 --- a/recursion/core-v2/src/chips/alu_ext.rs +++ b/recursion/core-v2/src/chips/alu_ext.rs @@ -183,17 +183,13 @@ where #[cfg(test)] mod tests { - use machine::RecursionAir; - use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear}; + use machine::tests::run_recursion_test_machines; + use p3_baby_bear::BabyBear; use p3_field::{extension::BinomialExtensionField, AbstractExtensionField, AbstractField}; use p3_matrix::dense::RowMajorMatrix; use rand::{rngs::StdRng, Rng, SeedableRng}; - use sp1_core::{ - air::MachineAir, - stark::StarkGenericConfig, - utils::{run_test_machine, BabyBearPoseidon2Inner}, - }; + use sp1_core::{air::MachineAir, stark::StarkGenericConfig}; use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer; use super::*; @@ -221,8 +217,6 @@ mod tests { pub fn four_ops() { type SC = BabyBearPoseidon2Outer; type F = ::Val; - type EF = ::Challenge; - type A = RecursionAir; let mut rng = StdRng::seed_from_u64(0xDEADBEEF); let mut random_extfelt = move || { @@ -258,18 +252,7 @@ mod tests { instructions, traces: Default::default(), }; - let mut runtime = Runtime::::new( - &program, - BabyBearPoseidon2Inner::new().perm, - ); - runtime.run().unwrap(); - - let config = SC::new(); - let machine = A::machine_wide(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![runtime.record], machine, pk, vk); - if let Err(e) = result { - panic!("Verification failed: {:?}", e); - } + + run_recursion_test_machines(program); } } diff --git a/recursion/core-v2/src/chips/exp_reverse_bits.rs b/recursion/core-v2/src/chips/exp_reverse_bits.rs index 68e567e0f3..00141c8c81 100644 --- a/recursion/core-v2/src/chips/exp_reverse_bits.rs +++ b/recursion/core-v2/src/chips/exp_reverse_bits.rs @@ -346,35 +346,29 @@ mod tests { use rand::Rng; use rand::SeedableRng; use sp1_core::air::MachineAir; - use sp1_core::utils::run_test_machine; use sp1_core::utils::setup_logger; - use sp1_core::utils::BabyBearPoseidon2; use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer; use std::iter::once; use p3_baby_bear::BabyBear; - use p3_baby_bear::DiffusionMatrixBabyBear; use p3_field::{AbstractField, PrimeField32}; use p3_matrix::dense::RowMajorMatrix; use sp1_core::stark::StarkGenericConfig; use crate::chips::exp_reverse_bits::ExpReverseBitsLenChip; - use crate::machine::RecursionAir; + use crate::machine::tests::run_recursion_test_machines; use crate::runtime::instruction as instr; use crate::runtime::ExecutionRecord; use crate::ExpReverseBitsEvent; use crate::Instruction; use crate::MemAccessKind; use crate::RecursionProgram; - use crate::Runtime; #[test] fn prove_babybear_circuit_erbl() { setup_logger(); type SC = BabyBearPoseidon2Outer; type F = ::Val; - type EF = ::Challenge; - type A = RecursionAir; let mut rng = StdRng::seed_from_u64(0xDEADBEEF); let mut random_felt = move || -> F { F::from_canonical_u32(rng.gen_range(0..1 << 16)) }; @@ -434,17 +428,7 @@ mod tests { traces: Default::default(), }; - let config = SC::new(); - - let mut runtime = - Runtime::::new(&program, BabyBearPoseidon2::new().perm); - runtime.run().unwrap(); - let machine = A::machine_wide(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![runtime.record], machine, pk, vk); - if let Err(e) = result { - panic!("Verification failed: {:?}", e); - } + run_recursion_test_machines(program); } #[test] diff --git a/recursion/core-v2/src/chips/fri_fold.rs b/recursion/core-v2/src/chips/fri_fold.rs index c4a77e932d..99908724f7 100644 --- a/recursion/core-v2/src/chips/fri_fold.rs +++ b/recursion/core-v2/src/chips/fri_fold.rs @@ -409,25 +409,22 @@ mod tests { use rand::Rng; use rand::SeedableRng; use sp1_core::air::MachineAir; - use sp1_core::utils::run_test_machine; use sp1_core::utils::setup_logger; - use sp1_core::utils::BabyBearPoseidon2; use sp1_recursion_core::air::Block; use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer; use std::mem::size_of; use p3_baby_bear::BabyBear; - use p3_baby_bear::DiffusionMatrixBabyBear; use p3_field::AbstractField; use p3_matrix::dense::RowMajorMatrix; use sp1_core::stark::StarkGenericConfig; use crate::chips::fri_fold::FriFoldChip; + use crate::machine::tests::run_recursion_test_machines; use crate::{ - machine::RecursionAir, runtime::{instruction as instr, ExecutionRecord}, FriFoldBaseIo, FriFoldEvent, FriFoldExtSingleIo, FriFoldExtVecIo, Instruction, - MemAccessKind, RecursionProgram, Runtime, + MemAccessKind, RecursionProgram, }; #[test] @@ -436,7 +433,6 @@ mod tests { type SC = BabyBearPoseidon2Outer; type F = ::Val; type EF = ::Challenge; - type A = RecursionAir; let mut rng = StdRng::seed_from_u64(0xDEADBEEF); let mut random_felt = move || -> F { F::from_canonical_u32(rng.gen_range(0..1 << 16)) }; @@ -572,17 +568,7 @@ mod tests { traces: Default::default(), }; - let config = SC::new(); - - let mut runtime = - Runtime::::new(&program, BabyBearPoseidon2::new().perm); - runtime.run().unwrap(); - let machine = A::machine_wide(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![runtime.record], machine, pk, vk); - if let Err(e) = result { - panic!("Verification failed: {:?}", e); - } + run_recursion_test_machines(program); } #[test] diff --git a/recursion/core-v2/src/chips/mem.rs b/recursion/core-v2/src/chips/mem.rs index 72f5f6e43a..97eafdb02f 100644 --- a/recursion/core-v2/src/chips/mem.rs +++ b/recursion/core-v2/src/chips/mem.rs @@ -203,43 +203,17 @@ you will also need to write your own execution record struct but look at recursi #[cfg(test)] mod tests { - use machine::RecursionAir; - use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear}; + use machine::tests::run_recursion_test_machines; + use p3_baby_bear::BabyBear; use p3_field::AbstractField; use p3_matrix::dense::RowMajorMatrix; - use sp1_core::{ - air::MachineAir, - stark::StarkGenericConfig, - utils::{run_test_machine, BabyBearPoseidon2Inner}, - }; - use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer; + use sp1_core::air::MachineAir; use super::*; use crate::runtime::instruction as instr; - type SC = BabyBearPoseidon2Outer; - type F = ::Val; - type EF = ::Challenge; - type A = RecursionAir; - - pub fn prove_program(program: RecursionProgram) { - let mut runtime = Runtime::::new( - &program, - BabyBearPoseidon2Inner::new().perm, - ); - runtime.run().unwrap(); - - let config = SC::new(); - let machine = A::machine(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![runtime.record], machine, pk, vk); - if let Err(e) = result { - panic!("Verification failed: {:?}", e); - } - } - #[test] pub fn generate_trace() { let shard = ExecutionRecord:: { @@ -261,48 +235,56 @@ mod tests { #[test] pub fn prove_basic_mem() { - prove_program(RecursionProgram { + let program = RecursionProgram { instructions: vec![ instr::mem(MemAccessKind::Write, 1, 1, 2), instr::mem(MemAccessKind::Read, 1, 1, 2), ], traces: Default::default(), - }); + }; + + run_recursion_test_machines(program); } #[test] #[should_panic] pub fn basic_mem_bad_mult() { - prove_program(RecursionProgram { + let program = RecursionProgram { instructions: vec![ instr::mem(MemAccessKind::Write, 1, 1, 2), instr::mem(MemAccessKind::Read, 999, 1, 2), ], traces: Default::default(), - }); + }; + + run_recursion_test_machines(program); } #[test] #[should_panic] pub fn basic_mem_bad_address() { - prove_program(RecursionProgram { + let program = RecursionProgram { instructions: vec![ instr::mem(MemAccessKind::Write, 1, 1, 2), instr::mem(MemAccessKind::Read, 1, 999, 2), ], traces: Default::default(), - }); + }; + + run_recursion_test_machines(program); } #[test] #[should_panic] pub fn basic_mem_bad_value() { - prove_program(RecursionProgram { + let program = RecursionProgram { instructions: vec![ instr::mem(MemAccessKind::Write, 1, 1, 2), instr::mem(MemAccessKind::Read, 1, 1, 999), ], traces: Default::default(), - }); + }; + + run_recursion_test_machines(program); } } diff --git a/recursion/core-v2/src/machine.rs b/recursion/core-v2/src/machine.rs index ec853fc43f..9d1e5b863a 100644 --- a/recursion/core-v2/src/machine.rs +++ b/recursion/core-v2/src/machine.rs @@ -55,15 +55,6 @@ impl, const DEGREE: usize, const COL_P .collect::>(); StarkMachine::new(config, chips, PROOF_MAX_NUM_PVS) } - pub fn machine_with_all_chips>( - config: SC, - ) -> StarkMachine { - let chips = Self::get_skinny_and_wide() - .into_iter() - .map(Chip::new) - .collect::>(); - StarkMachine::new(config, chips, PROOF_MAX_NUM_PVS) - } pub fn machine_with_padding>( config: SC, @@ -132,19 +123,6 @@ impl, const DEGREE: usize, const COL_P ] } - pub fn get_skinny_and_wide() -> Vec { - vec![ - // RecursionAir::Program(ProgramChip::default()), - RecursionAir::Memory(MemoryChip::default()), - RecursionAir::BaseAlu(BaseAluChip::default()), - RecursionAir::ExtAlu(ExtAluChip::default()), - RecursionAir::Poseidon2Skinny(Poseidon2SkinnyChip::::default()), - RecursionAir::Poseidon2Wide(Poseidon2WideChip::::default()), - RecursionAir::ExpReverseBitsLen(ExpReverseBitsLenChip::::default()), - RecursionAir::FriFold(FriFoldChip::::default()), - ] - } - pub fn get_all_with_padding( fri_fold_padding: usize, poseidon2_padding: usize, @@ -217,7 +195,7 @@ impl, const DEGREE: usize, const COL_P } #[cfg(test)] -mod tests { +pub mod tests { use machine::RecursionAir; use p3_baby_bear::DiffusionMatrixBabyBear; @@ -228,36 +206,46 @@ mod tests { use rand::prelude::*; use sp1_core::{ stark::StarkGenericConfig, - utils::{run_test_machine, BabyBearPoseidon2Inner}, + utils::{run_test_machine, BabyBearPoseidon2}, }; - use sp1_recursion_core::stark::config::BabyBearPoseidon2Outer; // TODO expand glob import use crate::{runtime::instruction as instr, *}; - type SC = BabyBearPoseidon2Outer; + type SC = BabyBearPoseidon2; type F = ::Val; type EF = ::Challenge; - type A = RecursionAir; + type A = RecursionAir; + type B = RecursionAir; - fn test_instructions(instructions: Vec>) { - let program = RecursionProgram { - instructions, - traces: Default::default(), - }; - let mut runtime = Runtime::::new( - &program, - BabyBearPoseidon2Inner::new().perm, - ); + /// Runs the given program on machines that use the wide and skinny Poseidon2 chips. + pub fn run_recursion_test_machines(program: RecursionProgram) { + let mut runtime = Runtime::::new(&program, SC::new().perm); runtime.run().unwrap(); - let config = SC::new(); - let machine = A::machine(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![runtime.record], machine, pk, vk); + // Run with the poseidon2 wide chip. + let wide_machine = A::machine_wide(BabyBearPoseidon2::default()); + let (pk, vk) = wide_machine.setup(&program); + let result = run_test_machine(vec![runtime.record.clone()], wide_machine, pk, vk); if let Err(e) = result { panic!("Verification failed: {:?}", e); } + + // Run with the poseidon2 skinny chip. + let skinny_machine = B::machine(BabyBearPoseidon2::compressed()); + let (pk, vk) = skinny_machine.setup(&program); + let result = run_test_machine(vec![runtime.record], skinny_machine, pk, vk); + if let Err(e) = result { + panic!("Verification failed: {:?}", e); + } + } + + fn test_instructions(instructions: Vec>) { + let program = RecursionProgram { + instructions, + traces: Default::default(), + }; + run_recursion_test_machines(program); } #[test] From d1b16ffe9891eb0399cc7e9e2b2535e9319b9f58 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Thu, 1 Aug 2024 10:19:16 -0700 Subject: [PATCH 16/18] fix --- recursion/circuit-v2/src/challenger.rs | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/recursion/circuit-v2/src/challenger.rs b/recursion/circuit-v2/src/challenger.rs index 0520d9daa7..0b05f60b2d 100644 --- a/recursion/circuit-v2/src/challenger.rs +++ b/recursion/circuit-v2/src/challenger.rs @@ -224,7 +224,7 @@ pub(crate) mod tests { use p3_field::AbstractField; use sp1_core::stark::StarkGenericConfig; use sp1_core::utils::setup_logger; - use sp1_core::utils::BabyBearPoseidon2Inner; + use sp1_core::utils::BabyBearPoseidon2; use sp1_recursion_compiler::asm::AsmBuilder; use sp1_recursion_compiler::asm::AsmConfig; use sp1_recursion_compiler::circuit::AsmCompiler; @@ -241,7 +241,7 @@ pub(crate) mod tests { use sp1_core::utils::run_test_machine; - type SC = BabyBearPoseidon2Inner; + type SC = BabyBearPoseidon2; type F = ::Val; type EF = ::Challenge; @@ -262,7 +262,7 @@ pub(crate) mod tests { let records = vec![runtime.record]; // Run with the poseidon2 wide chip. - let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(BabyBearPoseidon2::default()); + let wide_machine = RecursionAir::<_, 3, 0>::machine_wide(SC::default()); let (pk, vk) = wide_machine.setup(&program); let result = run_test_machine(records.clone(), wide_machine, pk, vk); if let Err(e) = result { @@ -270,7 +270,7 @@ pub(crate) mod tests { } // Run with the poseidon2 skinny chip. - let skinny_machine = RecursionAir::<_, 9, 0>::machine(BabyBearPoseidon2::compressed()); + let skinny_machine = RecursionAir::<_, 9, 0>::machine(SC::compressed()); let (pk, vk) = skinny_machine.setup(&program); let result = run_test_machine(records.clone(), skinny_machine, pk, vk); if let Err(e) = result { From f1ee9ff04a4626e03918edb5b48b49bd9100b412 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Thu, 8 Aug 2024 10:37:30 -0700 Subject: [PATCH 17/18] fixed tests --- recursion/core-v2/src/chips/mem/constant.rs | 6 +++--- recursion/core-v2/src/chips/public_values.rs | 21 +++----------------- 2 files changed, 6 insertions(+), 21 deletions(-) diff --git a/recursion/core-v2/src/chips/mem/constant.rs b/recursion/core-v2/src/chips/mem/constant.rs index feef160cde..b477afa0bd 100644 --- a/recursion/core-v2/src/chips/mem/constant.rs +++ b/recursion/core-v2/src/chips/mem/constant.rs @@ -155,7 +155,7 @@ where #[cfg(test)] mod tests { - use machine::RecursionAir; + use machine::{tests::run_recursion_test_machines, RecursionAir}; use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear}; use p3_field::AbstractField; use p3_matrix::dense::RowMajorMatrix; @@ -184,7 +184,7 @@ mod tests { runtime.run().unwrap(); let config = SC::new(); - let machine = A::machine(config); + let machine = A::machine_wide(config); let (pk, vk) = machine.setup(&program); let result = run_test_machine(vec![runtime.record], machine, pk, vk); if let Err(e) = result { @@ -213,7 +213,7 @@ mod tests { #[test] pub fn prove_basic_mem() { - prove_program(RecursionProgram { + run_recursion_test_machines(RecursionProgram { instructions: vec![ instr::mem(MemAccessKind::Write, 1, 1, 2), instr::mem(MemAccessKind::Read, 1, 1, 2), diff --git a/recursion/core-v2/src/chips/public_values.rs b/recursion/core-v2/src/chips/public_values.rs index 2e4d78a136..58bed3a4c6 100644 --- a/recursion/core-v2/src/chips/public_values.rs +++ b/recursion/core-v2/src/chips/public_values.rs @@ -180,9 +180,7 @@ mod tests { use rand::Rng; use rand::SeedableRng; use sp1_core::air::MachineAir; - use sp1_core::utils::run_test_machine; use sp1_core::utils::setup_logger; - use sp1_core::utils::BabyBearPoseidon2; use sp1_core::utils::DIGEST_SIZE; use sp1_recursion_core::air::RecursionPublicValues; use sp1_recursion_core::air::NUM_PV_ELMS_TO_HASH; @@ -192,17 +190,16 @@ mod tests { use std::borrow::Borrow; use p3_baby_bear::BabyBear; - use p3_baby_bear::DiffusionMatrixBabyBear; use p3_field::AbstractField; use p3_matrix::dense::RowMajorMatrix; use sp1_core::stark::StarkGenericConfig; use crate::chips::public_values::PublicValuesChip; + use crate::machine::tests::run_recursion_test_machines; use crate::CommitPublicValuesEvent; use crate::{ - machine::RecursionAir, runtime::{instruction as instr, ExecutionRecord}, - MemAccessKind, RecursionProgram, Runtime, + MemAccessKind, RecursionProgram, }; #[test] @@ -210,8 +207,6 @@ mod tests { setup_logger(); type SC = BabyBearPoseidon2Outer; type F = ::Val; - type EF = ::Challenge; - type A = RecursionAir; let mut rng = StdRng::seed_from_u64(0xDEADBEEF); let mut random_felt = move || -> F { F::from_canonical_u32(rng.gen_range(0..1 << 16)) }; @@ -240,17 +235,7 @@ mod tests { traces: Default::default(), }; - let config = SC::new(); - - let mut runtime = - Runtime::::new(&program, BabyBearPoseidon2::new().perm); - runtime.run().unwrap(); - let machine = A::machine(config); - let (pk, vk) = machine.setup(&program); - let result = run_test_machine(vec![runtime.record], machine, pk, vk); - if let Err(e) = result { - panic!("Verification failed: {:?}", e); - } + run_recursion_test_machines(program); } #[test] From e1676f68cbc7c23706f3a00a46587c303ab50ea2 Mon Sep 17 00:00:00 2001 From: Kevin Jue Date: Thu, 8 Aug 2024 11:23:48 -0700 Subject: [PATCH 18/18] fixed unit test --- recursion/compiler/src/circuit/mod.rs | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/recursion/compiler/src/circuit/mod.rs b/recursion/compiler/src/circuit/mod.rs index eb1a350d86..14036c5929 100644 --- a/recursion/compiler/src/circuit/mod.rs +++ b/recursion/compiler/src/circuit/mod.rs @@ -158,7 +158,7 @@ mod tests { .into(); runtime.run().unwrap(); - let machine = A::machine(SC::new()); + let machine = A::machine_wide(SC::new()); let (pk, vk) = machine.setup(&program); let result =