Skip to content

Commit

Permalink
fix(recursion): poseidon2 external flag (#747)
Browse files Browse the repository at this point in the history
  • Loading branch information
@kevjue
kevjue authored May 15, 2024
1 parent 4f8f915 commit 00255c1
Show file tree
Hide file tree
Showing 2 changed files with 105 additions and 70 deletions.
93 changes: 57 additions & 36 deletions recursion/core/src/poseidon2/external.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -42,24 +42,26 @@ impl Poseidon2Chip {
receive_table: AB::Var,
memory_access: AB::Expr,
) {
let rounds_f = 8;
let rounds_p = 13;
let rounds_p_beginning = 2 + rounds_f / 2;
let rounds_p_end = rounds_p_beginning + rounds_p;
const NUM_ROUNDS_F: usize = 8;
const NUM_ROUNDS_P: usize = 13;
const ROUNDS_F_1_BEGINNING: usize = 2; // Previous rounds are memory read and initial.
const ROUNDS_P_BEGINNING: usize = ROUNDS_F_1_BEGINNING + NUM_ROUNDS_F / 2;
const ROUNDS_P_END: usize = ROUNDS_P_BEGINNING + NUM_ROUNDS_P;
const ROUND_F_2_END: usize = ROUNDS_P_END + NUM_ROUNDS_F / 2;

let is_memory_read = local.rounds[0];
let is_initial = local.rounds[1];

// First half of the external rounds.
let mut is_external_layer = (2..rounds_p_beginning)
let mut is_external_layer = (ROUNDS_F_1_BEGINNING..ROUNDS_P_BEGINNING)
.map(|i| local.rounds[i].into())
.sum::<AB::Expr>();

// Second half of the external rounds.
is_external_layer += (rounds_p_end..rounds_p + rounds_f)
is_external_layer += (ROUNDS_P_END..ROUND_F_2_END)
.map(|i| local.rounds[i].into())
.sum::<AB::Expr>();
let is_internal_layer = (rounds_p_beginning..rounds_p_end)
let is_internal_layer = (ROUNDS_P_BEGINNING..ROUNDS_P_END)
.map(|i| local.rounds[i].into())
.sum::<AB::Expr>();
let is_memory_write = local.rounds[local.rounds.len() - 1];
Expand All @@ -80,7 +82,7 @@ impl Poseidon2Chip {
is_initial.into(),
is_external_layer.clone(),
is_internal_layer.clone(),
rounds_f + rounds_p + 1,
NUM_ROUNDS_F + NUM_ROUNDS_P + 1,
);

self.eval_syscall(builder, local, receive_table);
Expand Down Expand Up @@ -413,39 +415,17 @@ mod tests {
}
}

#[test]
fn prove_babybear() {
let config = BabyBearPoseidon2::compressed();
let mut challenger = config.challenger();

let chip = Poseidon2Chip {
fixed_log2_rows: None,
};
let rng = &mut rand::thread_rng();

let test_inputs: Vec<[BabyBear; 16]> = (0..16)
.map(|_| core::array::from_fn(|_| BabyBear::rand(rng)))
.collect_vec();

let gt: Poseidon2<
BabyBear,
Poseidon2ExternalMatrixGeneral,
DiffusionMatrixBabyBear,
16,
7,
> = inner_perm();

let expected_outputs = test_inputs
.iter()
.map(|input| gt.permute(*input))
.collect::<Vec<_>>();

fn prove_babybear(inputs: Vec<[BabyBear; 16]>, outputs: Vec<[BabyBear; 16]>) {
let mut input_exec = ExecutionRecord::<BabyBear>::default();
for (input, output) in test_inputs.into_iter().zip_eq(expected_outputs) {
for (input, output) in inputs.into_iter().zip_eq(outputs) {
input_exec
.poseidon2_events
.push(Poseidon2Event::dummy_from_input(input, output));
}

let chip = Poseidon2Chip {
fixed_log2_rows: None,
};
let trace: RowMajorMatrix<BabyBear> =
chip.generate_trace(&input_exec, &mut ExecutionRecord::<BabyBear>::default());
println!(
Expand All @@ -455,6 +435,8 @@ mod tests {
);

let start = Instant::now();
let config = BabyBearPoseidon2::compressed();
let mut challenger = config.challenger();
let proof = uni_stark_prove(&config, &chip, &mut challenger, trace);
let duration = start.elapsed().as_secs_f64();
println!("proof duration = {:?}", duration);
Expand All @@ -471,4 +453,43 @@ mod tests {
let duration = start.elapsed().as_secs_f64();
println!("verify duration = {:?}", duration);
}

#[test]
fn prove_babybear_success() {
let rng = &mut rand::thread_rng();

let test_inputs: Vec<[BabyBear; 16]> = (0..16)
.map(|_| core::array::from_fn(|_| BabyBear::rand(rng)))
.collect_vec();

let gt: Poseidon2<
BabyBear,
Poseidon2ExternalMatrixGeneral,
DiffusionMatrixBabyBear,
16,
7,
> = inner_perm();

let expected_outputs = test_inputs
.iter()
.map(|input| gt.permute(*input))
.collect::<Vec<_>>();

prove_babybear(test_inputs, expected_outputs)
}

#[test]
#[should_panic]
fn prove_babybear_failure() {
let rng = &mut rand::thread_rng();
let test_inputs: Vec<[BabyBear; 16]> = (0..16)
.map(|_| core::array::from_fn(|_| BabyBear::rand(rng)))
.collect_vec();

let bad_outputs: Vec<[BabyBear; 16]> = (0..16)
.map(|_| core::array::from_fn(|_| BabyBear::rand(rng)))
.collect_vec();

prove_babybear(test_inputs, bad_outputs)
}
}
82 changes: 48 additions & 34 deletions recursion/core/src/poseidon2_wide/external.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -454,12 +454,12 @@ mod tests {
use p3_baby_bear::{BabyBear, DiffusionMatrixBabyBear};
use p3_field::AbstractField;
use p3_matrix::dense::RowMajorMatrix;
use p3_matrix::Matrix;
use p3_poseidon2::{Poseidon2, Poseidon2ExternalMatrixGeneral};
use p3_symmetric::Permutation;
use sp1_core::air::MachineAir;
use sp1_core::stark::StarkGenericConfig;
use sp1_core::utils::{inner_perm, uni_stark_prove, uni_stark_verify, BabyBearPoseidon2};
use zkhash::ark_ff::UniformRand;

fn generate_trace_degree<const DEGREE: usize>() {
let chip = Poseidon2WideChip::<DEGREE> {
Expand Down Expand Up @@ -504,46 +504,24 @@ mod tests {
generate_trace_degree::<7>();
}

fn poseidon2_wide_prove_babybear_degree<const DEGREE: usize>() {
const DEGREE: usize = 7;

let config = BabyBearPoseidon2::compressed();
let mut challenger = config.challenger();

fn poseidon2_wide_prove_babybear_degree<const DEGREE: usize>(
inputs: Vec<[BabyBear; 16]>,
outputs: Vec<[BabyBear; 16]>,
) {
let chip = Poseidon2WideChip::<DEGREE> {
fixed_log2_rows: None,
};

let test_inputs = (0..1000)
.map(|i| [BabyBear::from_canonical_u32(i); WIDTH])
.collect_vec();

let gt: Poseidon2<
BabyBear,
Poseidon2ExternalMatrixGeneral,
DiffusionMatrixBabyBear,
16,
7,
> = inner_perm();

let expected_outputs = test_inputs
.iter()
.map(|input| gt.permute(*input))
.collect::<Vec<_>>();

let mut input_exec = ExecutionRecord::<BabyBear>::default();
for (input, output) in test_inputs.into_iter().zip_eq(expected_outputs) {
for (input, output) in inputs.into_iter().zip_eq(outputs) {
input_exec
.poseidon2_events
.push(Poseidon2Event::dummy_from_input(input, output));
}
let trace: RowMajorMatrix<BabyBear> =
chip.generate_trace(&input_exec, &mut ExecutionRecord::<BabyBear>::default());
println!(
"trace dims is width: {:?}, height: {:?}",
trace.width(),
trace.height()
);

let config = BabyBearPoseidon2::compressed();
let mut challenger = config.challenger();

let start = Instant::now();
let proof = uni_stark_prove(&config, &chip, &mut challenger, trace);
Expand All @@ -560,8 +538,44 @@ mod tests {
}

#[test]
fn poseidon2_wide_prove_babybear() {
poseidon2_wide_prove_babybear_degree::<3>();
poseidon2_wide_prove_babybear_degree::<7>();
fn poseidon2_wide_prove_babybear_success() {
let rng = &mut rand::thread_rng();

let test_inputs: Vec<[BabyBear; 16]> = (0..1000)
.map(|_| core::array::from_fn(|_| BabyBear::rand(rng)))
.collect_vec();

let gt: Poseidon2<
BabyBear,
Poseidon2ExternalMatrixGeneral,
DiffusionMatrixBabyBear,
16,
7,
> = inner_perm();

let expected_outputs = test_inputs
.iter()
.map(|input| gt.permute(*input))
.collect::<Vec<_>>();

poseidon2_wide_prove_babybear_degree::<3>(test_inputs.clone(), expected_outputs.clone());
poseidon2_wide_prove_babybear_degree::<7>(test_inputs, expected_outputs);
}

#[test]
#[should_panic]
fn poseidon2_wide_prove_babybear_failure() {
let rng = &mut rand::thread_rng();

let test_inputs = (0..1000)
.map(|i| [BabyBear::from_canonical_u32(i); WIDTH])
.collect_vec();

let bad_outputs: Vec<[BabyBear; 16]> = (0..1000)
.map(|_| core::array::from_fn(|_| BabyBear::rand(rng)))
.collect_vec();

poseidon2_wide_prove_babybear_degree::<3>(test_inputs.clone(), bad_outputs.clone());
poseidon2_wide_prove_babybear_degree::<7>(test_inputs, bad_outputs);
}
}

0 comments on commit 00255c1

Please sign in to comment.