diff --git a/Cargo.lock b/Cargo.lock index 53615191e8..fdbfea5a2f 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -4661,6 +4661,7 @@ dependencies = [ "p3-merkle-tree", "p3-poseidon2", "p3-symmetric", + "rand", "serde", "serde_with", "sp1-core", diff --git a/prover/src/lib.rs b/prover/src/lib.rs index fd772c84e9..44ad4b7602 100644 --- a/prover/src/lib.rs +++ b/prover/src/lib.rs @@ -557,6 +557,7 @@ impl SP1Prover { runtime.witness_stream = witness_stream.into(); runtime.run(); let mut checkpoint = runtime.memory.clone(); + let checkpoint_uninit = runtime.uninitialized_memory.clone(); // Execute runtime. let machine = RecursionAirWideDeg3::machine(InnerSC::default()); @@ -568,6 +569,7 @@ impl SP1Prover { e.1.timestamp = BabyBear::zero(); }); runtime.memory = checkpoint; + runtime.uninitialized_memory = checkpoint_uninit; runtime.run(); runtime.print_stats(); tracing::info!( diff --git a/recursion/core/Cargo.toml b/recursion/core/Cargo.toml index 385b337b3c..006db2a1c9 100644 --- a/recursion/core/Cargo.toml +++ b/recursion/core/Cargo.toml @@ -30,3 +30,6 @@ serde_with = "3.6.1" backtrace = { version = "0.3.71", features = ["serde"] } arrayref = "0.3.6" static_assertions = "1.1.0" + +[dev-dependencies] +rand = "0.8.5" \ No newline at end of file diff --git a/recursion/core/src/multi/mod.rs b/recursion/core/src/multi/mod.rs index fec939e011..5904de276c 100644 --- a/recursion/core/src/multi/mod.rs +++ b/recursion/core/src/multi/mod.rs @@ -190,6 +190,7 @@ where poseidon2_chip.eval_poseidon2( &mut sub_builder, local.poseidon2(), + next.poseidon2(), local.poseidon2_receive_table, local.poseidon2_memory_access.into(), ); diff --git a/recursion/core/src/poseidon2/columns.rs b/recursion/core/src/poseidon2/columns.rs index 4b8a2169a1..9322c5d809 100644 --- a/recursion/core/src/poseidon2/columns.rs +++ b/recursion/core/src/poseidon2/columns.rs @@ -11,8 +11,6 @@ pub struct Poseidon2Cols { pub left_input: T, pub right_input: T, pub rounds: [T; 24], // 1 round for memory input; 1 round for initialize; 8 rounds for external; 13 rounds for internal; 1 round for memory output - pub is_computation: T, - pub is_memory_access: T, pub round_specific_cols: RoundSpecificCols, } diff --git a/recursion/core/src/poseidon2/external.rs b/recursion/core/src/poseidon2/external.rs index cacffd2323..a6124d798d 100644 --- a/recursion/core/src/poseidon2/external.rs +++ b/recursion/core/src/poseidon2/external.rs @@ -9,6 +9,7 @@ use sp1_primitives::RC_16_30_U32; use std::ops::Add; use crate::air::{RecursionInteractionAirBuilder, RecursionMemoryAirBuilder}; +use crate::memory::MemoryCols; use crate::poseidon2_wide::{apply_m_4, internal_linear_layer}; use crate::runtime::Opcode; @@ -37,6 +38,7 @@ impl Poseidon2Chip { &self, builder: &mut AB, local: &Poseidon2Cols, + next: &Poseidon2Cols, receive_table: AB::Var, memory_access: AB::Expr, ) { @@ -65,6 +67,7 @@ impl Poseidon2Chip { self.eval_mem( builder, local, + next, is_memory_read, is_memory_write, memory_access, @@ -73,6 +76,7 @@ impl Poseidon2Chip { self.eval_computation( builder, local, + next, is_initial.into(), is_external_layer.clone(), is_internal_layer.clone(), @@ -94,12 +98,12 @@ impl Poseidon2Chip { &self, builder: &mut AB, local: &Poseidon2Cols, + next: &Poseidon2Cols, is_memory_read: AB::Var, is_memory_write: AB::Var, memory_access: AB::Expr, ) { let memory_access_cols = local.round_specific_cols.memory_access(); - builder .when(is_memory_read) .assert_eq(local.left_input, memory_access_cols.addr_first_half); @@ -128,12 +132,24 @@ impl Poseidon2Chip { memory_access.clone(), ); } + + // For the memory read round, need to connect the memory val to the input of the next + // computation round. + let next_computation_col = next.round_specific_cols.computation(); + for i in 0..WIDTH { + builder.when_transition().when(is_memory_read).assert_eq( + *memory_access_cols.mem_access[i].value(), + next_computation_col.input[i], + ); + } } + #[allow(clippy::too_many_arguments)] fn eval_computation( &self, builder: &mut AB, local: &Poseidon2Cols, + next: &Poseidon2Cols, is_initial: AB::Expr, is_external_layer: AB::Expr, is_internal_layer: AB::Expr, @@ -158,11 +174,11 @@ impl Poseidon2Chip { let mut result: AB::Expr = computation_cols.input[i].into(); for r in 0..rounds { if i == 0 { - result += local.rounds[r + 1] + result += local.rounds[r + 2] * constants[r][i] * (is_external_layer.clone() + is_internal_layer.clone()); } else { - result += local.rounds[r + 1] * constants[r][i] * is_external_layer.clone(); + result += local.rounds[r + 2] * constants[r][i] * is_external_layer.clone(); } } builder @@ -251,9 +267,26 @@ impl Poseidon2Chip { let mut state: [AB::Expr; WIDTH] = sbox_result.clone(); internal_linear_layer(&mut state); builder - .when(is_internal_layer) + .when(is_internal_layer.clone()) .assert_all_eq(state.clone(), computation_cols.output); } + + // Assert that the round's output values are equal the the next round's input values. For the + // last computation round, assert athat the output values are equal to the output memory values. + let next_row_computation = next.round_specific_cols.computation(); + let next_row_memory_access = next.round_specific_cols.memory_access(); + for i in 0..WIDTH { + let next_round_value = builder.if_else( + local.rounds[22], + *next_row_memory_access.mem_access[i].value(), + next_row_computation.input[i], + ); + + builder + .when_transition() + .when(is_initial.clone() + is_external_layer.clone() + is_internal_layer.clone()) + .assert_eq(computation_cols.output[i], next_round_value); + } } fn eval_syscall( @@ -295,9 +328,13 @@ where let main = builder.main(); let local = main.row_slice(0); let local: &Poseidon2Cols = (*local).borrow(); + let next = main.row_slice(1); + let next: &Poseidon2Cols = (*next).borrow(); + self.eval_poseidon2::( builder, local, + next, Self::do_receive_table::(local), Self::do_memory_access::(local), ); @@ -309,10 +346,10 @@ mod tests { use itertools::Itertools; use std::borrow::Borrow; use std::time::Instant; + use zkhash::ark_ff::UniformRand; use p3_baby_bear::BabyBear; use p3_baby_bear::DiffusionMatrixBabyBear; - use p3_field::AbstractField; use p3_matrix::{dense::RowMajorMatrix, Matrix}; use p3_poseidon2::Poseidon2; use p3_poseidon2::Poseidon2ExternalMatrixGeneral; @@ -324,7 +361,7 @@ mod tests { }; use crate::{ - poseidon2::{Poseidon2Chip, Poseidon2Event, WIDTH}, + poseidon2::{Poseidon2Chip, Poseidon2Event}, runtime::ExecutionRecord, }; use p3_symmetric::Permutation; @@ -338,12 +375,12 @@ mod tests { let chip = Poseidon2Chip { fixed_log2_rows: None, }; - let test_inputs = vec![ - [BabyBear::from_canonical_u32(1); WIDTH], - [BabyBear::from_canonical_u32(2); WIDTH], - [BabyBear::from_canonical_u32(3); WIDTH], - [BabyBear::from_canonical_u32(4); WIDTH], - ]; + + 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, @@ -384,9 +421,10 @@ mod tests { let chip = Poseidon2Chip { fixed_log2_rows: None, }; + let rng = &mut rand::thread_rng(); - let test_inputs = (0..16) - .map(|i| [BabyBear::from_canonical_u32(i); WIDTH]) + let test_inputs: Vec<[BabyBear; 16]> = (0..16) + .map(|_| core::array::from_fn(|_| BabyBear::rand(rng))) .collect_vec(); let gt: Poseidon2< diff --git a/recursion/core/src/poseidon2/mod.rs b/recursion/core/src/poseidon2/mod.rs index 93283b97c9..2c42ce7219 100644 --- a/recursion/core/src/poseidon2/mod.rs +++ b/recursion/core/src/poseidon2/mod.rs @@ -30,7 +30,7 @@ impl Poseidon2Event { MemoryRecord::new_read(F::zero(), Block::from(input[i]), F::one(), F::zero()) }); let output_records: [MemoryRecord; WIDTH] = core::array::from_fn(|i| { - MemoryRecord::new_read(F::zero(), Block::from(output[i]), F::one(), F::zero()) + MemoryRecord::new_read(F::zero(), Block::from(output[i]), F::two(), F::zero()) }); Self { diff --git a/recursion/core/src/poseidon2/trace.rs b/recursion/core/src/poseidon2/trace.rs index ccbb85baaf..619c8c7743 100644 --- a/recursion/core/src/poseidon2/trace.rs +++ b/recursion/core/src/poseidon2/trace.rs @@ -115,7 +115,7 @@ impl MachineAir for Poseidon2Chip { // Apply the round constants. for j in 0..WIDTH { computation_cols.add_rc[j] = computation_cols.input[j] - + F::from_wrapped_u32(RC_16_30_U32[r - 1][j]); + + F::from_wrapped_u32(RC_16_30_U32[r - 2][j]); } } else { // Apply the round constants only on the first element. @@ -123,7 +123,7 @@ impl MachineAir for Poseidon2Chip { .add_rc .copy_from_slice(&computation_cols.input); computation_cols.add_rc[0] = - computation_cols.input[0] + F::from_wrapped_u32(RC_16_30_U32[r - 1][0]); + computation_cols.input[0] + F::from_wrapped_u32(RC_16_30_U32[r - 2][0]); }; // Apply the sbox. diff --git a/recursion/core/src/poseidon2_wide/external.rs b/recursion/core/src/poseidon2_wide/external.rs index bf5f74b766..4525f23048 100644 --- a/recursion/core/src/poseidon2_wide/external.rs +++ b/recursion/core/src/poseidon2_wide/external.rs @@ -493,10 +493,8 @@ mod tests { .push(Poseidon2Event::dummy_from_input(input, output)); } - let trace: RowMajorMatrix = - chip.generate_trace(&input_exec, &mut ExecutionRecord::::default()); - - assert_eq!(trace.height(), test_inputs.len()); + // Generate trace will assert for the expected outputs. + chip.generate_trace(&input_exec, &mut ExecutionRecord::::default()); } /// A test generating a trace for a single permutation that checks that the output is correct