Skip to content

Commit

Permalink
working
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevjue
kevjue committed Jul 30, 2024
1 parent 3a756ae commit f1b0b42
Show file tree
Hide file tree
Showing 4 changed files with 124 additions and 62 deletions.
4 changes: 4 additions & 0 deletions core/src/utils/prove.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -485,6 +485,10 @@ where
{
let start = Instant::now();
let prover = DefaultProver::new(machine);

let mut challenger = prover.config().challenger();
prover.debug_constraints(&pk, records.clone(), &mut challenger);

let mut challenger = prover.config().challenger();
let proof = prover
.prove(&pk, records, &mut challenger, SP1CoreOpts::default())
Expand Down
63 changes: 57 additions & 6 deletions recursion/core-v2/src/chips/poseidon2_wide/air.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,9 @@ use crate::builder::SP1RecursionAirBuilder;
use super::columns::permutation::Poseidon2;
use super::columns::preprocessed::Poseidon2PreprocessedCols;
use super::columns::{NUM_POSEIDON2_DEGREE3_COLS, NUM_POSEIDON2_DEGREE9_COLS};
use super::{external_linear_layer, Poseidon2WideChip, WIDTH};
use super::{
external_linear_layer, internal_linear_layer, Poseidon2WideChip, NUM_INTERNAL_ROUNDS, WIDTH,
};

impl<F, const DEGREE: usize> BaseAir<F> for Poseidon2WideChip<DEGREE> {
fn width(&self) -> usize {
Expand Down Expand Up @@ -67,10 +69,13 @@ where
)
});

// Apply the first half of the external rounds.
for r in 0..NUM_EXTERNAL_ROUNDS / 2 {
// Apply the external rounds.
for r in 0..NUM_EXTERNAL_ROUNDS {
self.eval_external_round(builder, local_row.as_ref(), r);
}

// Apply the internal rounds.
self.eval_internal_rounds(builder, local_row.as_ref());
}
}

Expand All @@ -92,8 +97,13 @@ impl<const DEGREE: usize> Poseidon2WideChip<DEGREE> {
}

// Add the round constants.
let round = if r < NUM_EXTERNAL_ROUNDS / 2 {
r
} else {
r + NUM_INTERNAL_ROUNDS
};
let add_rc: [AB::Expr; WIDTH] = array::from_fn(|i| {
local_state[i].clone() + AB::Expr::from_canonical_u32(RC_16_30_U32[r][i])
local_state[i].clone() + AB::F::from_wrapped_u32(RC_16_30_U32[round][i])
});

// Apply the sboxes.
Expand All @@ -104,7 +114,7 @@ impl<const DEGREE: usize> Poseidon2WideChip<DEGREE> {
let calculated_sbox_deg_3 = add_rc[i].clone() * add_rc[i].clone() * add_rc[i].clone();

if let Some(external_sbox) = local_row.external_rounds_sbox() {
builder.assert_eq(external_sbox[r][i].into(), calculated_sbox_deg_3);
// builder.assert_eq(external_sbox[r][i].into(), calculated_sbox_deg_3);
sbox_deg_3[i] = external_sbox[r][i].into();
} else {
sbox_deg_3[i] = calculated_sbox_deg_3;
Expand All @@ -117,7 +127,7 @@ impl<const DEGREE: usize> Poseidon2WideChip<DEGREE> {
let mut state = sbox_deg_7;
external_linear_layer(&mut state);

let next_state = if r == NUM_EXTERNAL_ROUNDS / 2 {
let next_state = if r == (NUM_EXTERNAL_ROUNDS / 2) - 1 {
local_row.internal_rounds_state()
} else if r == NUM_EXTERNAL_ROUNDS - 1 {
local_row.perm_output()
Expand All @@ -129,4 +139,45 @@ impl<const DEGREE: usize> Poseidon2WideChip<DEGREE> {
builder.assert_eq(next_state[i], state[i].clone());
}
}

fn eval_internal_rounds<AB>(&self, builder: &mut AB, local_row: &dyn Poseidon2<AB::Var>)
where
AB: SP1RecursionAirBuilder + PairBuilder,
{
let state = &local_row.internal_rounds_state();
let s0 = local_row.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) = local_row.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 = local_row.external_rounds_state()[NUM_EXTERNAL_ROUNDS / 2];
for i in 0..WIDTH {
builder.assert_eq(external_state[i], state[i].clone())
}
}
}
2 changes: 0 additions & 2 deletions recursion/core-v2/src/chips/poseidon2_wide/columns/permutation.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -24,8 +24,6 @@ pub struct PermutationState<T: Copy> {
pub external_rounds_state: [[T; WIDTH]; NUM_EXTERNAL_ROUNDS],
pub internal_rounds_state: [T; WIDTH],
pub internal_rounds_s0: [T; NUM_INTERNAL_ROUNDS - 1],
pub external_rounds_sbox: [[T; WIDTH]; NUM_EXTERNAL_ROUNDS],
pub internal_rounds_sbox: [T; NUM_INTERNAL_ROUNDS],
pub output_state: [T; WIDTH],
}

Expand Down
117 changes: 63 additions & 54 deletions recursion/core-v2/src/chips/poseidon2_wide/trace.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -47,65 +47,18 @@ impl<F: PrimeField32, const DEGREE: usize> MachineAir<F> for Poseidon2WideChip<D
let num_columns = <Poseidon2WideChip<DEGREE> as BaseAir<F>>::width(self);

for event in &input.poseidon2_wide_events {
let mut input_row = vec![F::zero(); num_columns];

{
let permutation = permutation_mut::<F, DEGREE>(&mut input_row);

let (
external_rounds_state,
internal_rounds_state,
internal_rounds_s0,
mut external_sbox,
mut internal_sbox,
output_state,
) = permutation.get_cols_mut();

external_rounds_state[0] = event.input;

// Apply the first half of external rounds.
for r in 0..NUM_EXTERNAL_ROUNDS / 2 {
let next_state =
self.populate_external_round(external_rounds_state, &mut external_sbox, r);
if r == NUM_EXTERNAL_ROUNDS / 2 - 1 {
*internal_rounds_state = next_state;
} else {
external_rounds_state[r + 1] = next_state;
}
}

// Apply the internal rounds.
external_rounds_state[NUM_EXTERNAL_ROUNDS / 2] = self.populate_internal_rounds(
internal_rounds_state,
internal_rounds_s0,
&mut internal_sbox,
);

// Apply the second half of external rounds.
for r in NUM_EXTERNAL_ROUNDS / 2..NUM_EXTERNAL_ROUNDS {
let next_state =
self.populate_external_round(external_rounds_state, &mut external_sbox, r);
if r == NUM_EXTERNAL_ROUNDS - 1 {
for i in 0..WIDTH {
output_state[i] = next_state[i];
assert_eq!(event.output[i], next_state[i]);
}
} else {
external_rounds_state[r + 1] = next_state;
}
}
}
rows.push(input_row);
let mut row = vec![F::zero(); num_columns];
self.populate_perm(event.input, Some(event.output), row.as_mut_slice());
rows.push(row);
}

if self.pad {
// Pad the trace to a power of two.
// This will need to be adjusted when the AIR constraints are implemented.
pad_rows_fixed(
&mut rows,
|| vec![F::zero(); num_columns],
self.fixed_log2_rows,
);
let mut dummy_row = vec![F::zero(); num_columns];
self.populate_perm([F::zero(); WIDTH], None, &mut dummy_row);

pad_rows_fixed(&mut rows, || dummy_row.clone(), self.fixed_log2_rows);
}

// Convert the trace to a row major matrix.
Expand Down Expand Up @@ -186,6 +139,62 @@ impl<F: PrimeField32, const DEGREE: usize> MachineAir<F> for Poseidon2WideChip<D
}

impl<const DEGREE: usize> Poseidon2WideChip<DEGREE> {
fn populate_perm<F: PrimeField32>(
&self,
input: [F; WIDTH],
expected_output: Option<[F; WIDTH]>,
input_row: &mut [F],
) {
{
let permutation = permutation_mut::<F, DEGREE>(input_row);

let (
external_rounds_state,
internal_rounds_state,
internal_rounds_s0,
mut external_sbox,
mut internal_sbox,
output_state,
) = permutation.get_cols_mut();

external_rounds_state[0] = input;

// Apply the first half of external rounds.
for r in 0..NUM_EXTERNAL_ROUNDS / 2 {
let next_state =
self.populate_external_round(external_rounds_state, &mut external_sbox, r);
if r == NUM_EXTERNAL_ROUNDS / 2 - 1 {
*internal_rounds_state = next_state;
} else {
external_rounds_state[r + 1] = next_state;
}
}

// Apply the internal rounds.
external_rounds_state[NUM_EXTERNAL_ROUNDS / 2] = self.populate_internal_rounds(
internal_rounds_state,
internal_rounds_s0,
&mut internal_sbox,
);

// Apply the second half of external rounds.
for r in NUM_EXTERNAL_ROUNDS / 2..NUM_EXTERNAL_ROUNDS {
let next_state =
self.populate_external_round(external_rounds_state, &mut external_sbox, r);
if r == NUM_EXTERNAL_ROUNDS - 1 {
for i in 0..WIDTH {
output_state[i] = next_state[i];
if let Some(expected_output) = expected_output {
assert_eq!(expected_output[i], next_state[i]);
}
}
} else {
external_rounds_state[r + 1] = next_state;
}
}
}
}

fn populate_external_round<F: PrimeField32>(
&self,
external_rounds_state: &[[F; WIDTH]],
Expand Down

0 comments on commit f1b0b42

Please sign in to comment.