diff --git a/recursion/core/src/cpu/air.rs b/recursion/core/src/cpu/air.rs index c08f39ced1..dde09e0b08 100644 --- a/recursion/core/src/cpu/air.rs +++ b/recursion/core/src/cpu/air.rs @@ -5,6 +5,8 @@ use p3_air::BaseAir; use p3_field::PrimeField32; use p3_matrix::dense::RowMajorMatrix; use p3_matrix::MatrixRowSlices; +use sp1_core::air::AirInteraction; +use sp1_core::lookup::InteractionKind; use sp1_core::stark::SP1AirBuilder; use sp1_core::{air::MachineAir, utils::pad_to_power_of_two}; use std::borrow::Borrow; @@ -42,6 +44,7 @@ impl MachineAir for CpuChip { cols.instruction.op_c = event.instruction.op_c; cols.instruction.imm_b = F::from_canonical_u32(event.instruction.imm_b as u32); cols.instruction.imm_c = F::from_canonical_u32(event.instruction.imm_c as u32); + cols.is_real = F::one(); row }) .collect::>(); @@ -72,7 +75,19 @@ where { fn eval(&self, builder: &mut AB) { let main = builder.main(); - let _: &CpuCols = main.row_slice(0).borrow(); + let local: &CpuCols = main.row_slice(0).borrow(); let _: &CpuCols = main.row_slice(1).borrow(); + builder.send(AirInteraction::new( + vec![ + local.instruction.opcode.into(), + local.instruction.op_a.into(), + local.instruction.op_b.into(), + local.instruction.op_c.into(), + local.instruction.imm_b.into(), + local.instruction.imm_c.into(), + ], + local.is_real.into(), + InteractionKind::Program, + )); } } diff --git a/recursion/core/src/cpu/columns.rs b/recursion/core/src/cpu/columns.rs index c7536951b8..00e39a3384 100644 --- a/recursion/core/src/cpu/columns.rs +++ b/recursion/core/src/cpu/columns.rs @@ -4,7 +4,7 @@ use sp1_core::operations::IsZeroOperation; use sp1_derive::AlignedBorrow; /// The column layout for the chip. -#[derive(AlignedBorrow, Default, Clone, Copy)] +#[derive(AlignedBorrow, Default, Clone, Copy, Debug)] #[repr(C)] pub struct CpuCols { pub clk: T, @@ -54,6 +54,8 @@ pub struct CpuCols { // c = a == b; pub a_eq_b: IsZeroOperation, + + pub is_real: T, } #[derive(AlignedBorrow, Default, Debug, Clone, Copy)] diff --git a/recursion/core/src/lib.rs b/recursion/core/src/lib.rs index b80e4a62ed..dbee40520d 100644 --- a/recursion/core/src/lib.rs +++ b/recursion/core/src/lib.rs @@ -1,12 +1,13 @@ pub mod air; pub mod cpu; pub mod memory; +pub mod program; pub mod runtime; pub mod stark; #[cfg(test)] pub mod tests { - use crate::runtime::{ExecutionRecord, Instruction, Opcode, Program, Runtime}; + use crate::runtime::{Instruction, Opcode, Program, Runtime}; use crate::stark::RecursionAir; use p3_baby_bear::BabyBear; @@ -18,9 +19,9 @@ pub mod tests { pub fn fibonacci_program() -> Program { // .main - // imm 0(fp) 1 <-- a = 1 - // imm 1(fp) 1 <-- b = 1 - // imm 2(fp) 10 <-- iterations = 10 + // imm 0(fp) 1 <-- a = 1 + // imm 1(fp) 1 <-- b = 1 + // imm 2(fp) 10 <-- iterations = 10 // .body: // add 3(fp) 0(fp) 1(fp) <-- tmp = a + b // sw 0(fp) 1(fp) <-- a = b @@ -45,18 +46,11 @@ pub mod tests { #[test] fn test_fibonacci_execute() { - let program = fibonacci_program(); - let mut runtime = Runtime:: { - clk: BabyBear::zero(), - program, - fp: BabyBear::zero(), - pc: BabyBear::zero(), - memory: vec![BabyBear::zero(); 1024 * 1024], - record: ExecutionRecord::::default(), - }; + let program = fibonacci_program::(); + let mut runtime = Runtime::new(&program); runtime.run(); println!("{:#?}", runtime.record.cpu_events); - assert_eq!(runtime.memory[1], BabyBear::from_canonical_u32(144)); + assert_eq!(runtime.memory[1].value, BabyBear::from_canonical_u32(144)); } #[test] diff --git a/recursion/core/src/memory/global.rs b/recursion/core/src/memory/global.rs new file mode 100644 index 0000000000..4d3245e8d2 --- /dev/null +++ b/recursion/core/src/memory/global.rs @@ -0,0 +1,119 @@ +use core::mem::size_of; +use p3_air::{Air, BaseAir}; +use p3_field::PrimeField32; +use p3_matrix::dense::RowMajorMatrix; +use p3_matrix::MatrixRowSlices; +use sp1_core::air::SP1AirBuilder; +use sp1_core::utils::indices_arr; +use sp1_core::{air::MachineAir, utils::pad_to_power_of_two}; +use sp1_derive::AlignedBorrow; +use std::borrow::Borrow; +use std::mem::transmute; + +use crate::memory::Word; +use crate::runtime::ExecutionRecord; + +#[allow(dead_code)] +#[derive(PartialEq)] +pub enum MemoryChipKind { + Init, + Finalize, + Program, +} + +pub struct MemoryGlobalChip { + pub kind: MemoryChipKind, +} + +#[allow(dead_code)] +impl MemoryGlobalChip { + pub fn new(kind: MemoryChipKind) -> Self { + Self { kind } + } +} + +impl BaseAir for MemoryGlobalChip { + fn width(&self) -> usize { + NUM_MEMORY_INIT_COLS + } +} + +impl MachineAir for MemoryGlobalChip { + type Record = ExecutionRecord; + + fn name(&self) -> String { + match self.kind { + MemoryChipKind::Init => "MemoryInit".to_string(), + MemoryChipKind::Finalize => "MemoryFinalize".to_string(), + MemoryChipKind::Program => "MemoryProgram".to_string(), + } + } + + #[allow(unused_variables)] + fn generate_trace( + &self, + input: &ExecutionRecord, + _output: &mut ExecutionRecord, + ) -> RowMajorMatrix { + let memory_record = match self.kind { + MemoryChipKind::Init => &input.first_memory_record, + MemoryChipKind::Finalize => &input.last_memory_record, + MemoryChipKind::Program => &input.program_memory_record, + }; + let rows: Vec<[F; 8]> = (0..memory_record.len()) // TODO: change this back to par_iter + .map(|i| [F::zero(); NUM_MEMORY_INIT_COLS]) + .collect::>(); + + let mut trace = RowMajorMatrix::new( + rows.into_iter().flatten().collect::>(), + NUM_MEMORY_INIT_COLS, + ); + + pad_to_power_of_two::(&mut trace.values); + + trace + } + + fn included(&self, shard: &Self::Record) -> bool { + match self.kind { + MemoryChipKind::Init => !shard.first_memory_record.is_empty(), + MemoryChipKind::Finalize => !shard.last_memory_record.is_empty(), + MemoryChipKind::Program => !shard.program_memory_record.is_empty(), + } + } +} + +#[derive(AlignedBorrow, Default, Debug, Clone, Copy)] +#[repr(C)] +pub struct MemoryInitCols { + pub shard: T, + pub timestamp: T, + pub addr: T, + pub value: Word, + pub is_real: T, +} + +pub(crate) const NUM_MEMORY_INIT_COLS: usize = size_of::>(); +#[allow(dead_code)] +pub(crate) const MEMORY_INIT_COL_MAP: MemoryInitCols = make_col_map(); + +const fn make_col_map() -> MemoryInitCols { + let indices_arr = indices_arr::(); + unsafe { transmute::<[usize; NUM_MEMORY_INIT_COLS], MemoryInitCols>(indices_arr) } +} + +impl Air for MemoryGlobalChip +where + AB: SP1AirBuilder, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local: &MemoryInitCols = main.row_slice(0).borrow(); + + // Dummy constraint of degree 3. + builder.assert_eq( + local.is_real * local.is_real * local.is_real, + local.is_real * local.is_real * local.is_real, + ); + } +} diff --git a/recursion/core/src/memory/mod.rs b/recursion/core/src/memory/mod.rs index 5d0ac6c411..a2bb95982e 100644 --- a/recursion/core/src/memory/mod.rs +++ b/recursion/core/src/memory/mod.rs @@ -1,9 +1,12 @@ +mod global; + use crate::air::Word; use sp1_derive::AlignedBorrow; use std::mem::size_of; #[derive(Debug, Clone)] pub struct MemoryRecord { + pub addr: F, pub value: F, pub timestamp: F, pub prev_value: F, diff --git a/recursion/core/src/program/mod.rs b/recursion/core/src/program/mod.rs new file mode 100644 index 0000000000..b54b1eb1a1 --- /dev/null +++ b/recursion/core/src/program/mod.rs @@ -0,0 +1,123 @@ +use crate::{cpu::columns::InstructionCols, runtime::ExecutionRecord}; +use core::mem::size_of; +use hashbrown::HashMap; +use p3_air::{Air, BaseAir}; +use p3_field::PrimeField32; +use p3_matrix::dense::RowMajorMatrix; +use p3_matrix::MatrixRowSlices; +use sp1_core::lookup::InteractionKind; +use sp1_core::{ + air::{AirInteraction, MachineAir, SP1AirBuilder}, + utils::pad_to_power_of_two, +}; +use sp1_derive::AlignedBorrow; +use std::borrow::Borrow; +use std::borrow::BorrowMut; + +pub const NUM_PROGRAM_COLS: usize = size_of::>(); + +#[derive(Default)] +pub struct ProgramChip; + +#[derive(AlignedBorrow, Clone, Copy, Debug, Default)] +#[repr(C)] +pub struct ProgramCols { + pub pc: T, + pub instruction: InstructionCols, + pub multiplicity: T, +} + +impl MachineAir for ProgramChip { + type Record = ExecutionRecord; + + fn name(&self) -> String { + "Program".to_string() + } + + fn generate_trace( + &self, + input: &ExecutionRecord, + _output: &mut ExecutionRecord, + ) -> RowMajorMatrix { + let mut instruction_counts = HashMap::new(); + input.cpu_events.iter().for_each(|event| { + let pc = event.pc; + instruction_counts + .entry(pc) + .and_modify(|count| *count += 1) + .or_insert(1); + }); + let rows = input + .program + .instructions + .clone() + .into_iter() + .enumerate() + .map(|(i, instruction)| { + let pc = F::from_canonical_u32(i as u32); + let mut row = [F::zero(); NUM_PROGRAM_COLS]; + let cols: &mut ProgramCols = row.as_mut_slice().borrow_mut(); + cols.pc = pc; + cols.instruction.opcode = F::from_canonical_u32(instruction.opcode as u32); + cols.instruction.op_a = instruction.op_a; + cols.instruction.op_b = instruction.op_b; + cols.instruction.op_c = instruction.op_c; + cols.instruction.imm_b = F::from_bool(instruction.imm_b); + cols.instruction.imm_c = F::from_bool(instruction.imm_c); + cols.multiplicity = + F::from_canonical_usize(*instruction_counts.get(&cols.pc).unwrap_or(&0)); + row + }) + .collect::>(); + + // Convert the trace to a row major matrix. + let mut trace = RowMajorMatrix::new( + rows.into_iter().flatten().collect::>(), + NUM_PROGRAM_COLS, + ); + + // Pad the trace to a power of two. + pad_to_power_of_two::(&mut trace.values); + + trace + } + + fn included(&self, _: &Self::Record) -> bool { + true + } +} + +impl BaseAir for ProgramChip { + fn width(&self) -> usize { + NUM_PROGRAM_COLS + } +} + +impl Air for ProgramChip +where + AB: SP1AirBuilder, +{ + fn eval(&self, builder: &mut AB) { + let main = builder.main(); + let local: &ProgramCols = main.row_slice(0).borrow(); + + // Dummy constraint of degree 3. + builder.assert_eq( + local.pc * local.pc * local.pc, + local.pc * local.pc * local.pc, + ); + + builder.receive(AirInteraction::new( + vec![ + local.instruction.opcode.into(), + local.instruction.op_a.into(), + local.instruction.op_b.into(), + local.instruction.op_c.into(), + local.instruction.imm_b.into(), + local.instruction.imm_c.into(), + ], + local.multiplicity.into(), + InteractionKind::Program, + )); + } +} diff --git a/recursion/core/src/runtime/mod.rs b/recursion/core/src/runtime/mod.rs index 8d5cf8ba8c..32ef8e319c 100644 --- a/recursion/core/src/runtime/mod.rs +++ b/recursion/core/src/runtime/mod.rs @@ -3,14 +3,31 @@ mod opcode; mod program; mod record; +use crate::memory::MemoryRecord; +use std::sync::Arc; + pub use instruction::*; pub use opcode::*; pub use program::*; pub use record::*; +use sp1_core::runtime::AccessPosition; use crate::cpu::CpuEvent; use p3_field::PrimeField32; +#[derive(Debug, Clone, Default)] +pub struct CpuRecord { + pub a: Option>, + pub b: Option>, + pub c: Option>, +} + +#[derive(Debug, Clone, Default)] +pub struct MemoryEntry { + pub value: F, + pub timestamp: F, +} + pub struct Runtime { /// The current clock. pub clk: F, @@ -25,34 +42,84 @@ pub struct Runtime { pub program: Program, /// Memory. - pub memory: Vec, + pub memory: Vec>, /// The execution record. pub record: ExecutionRecord, + + pub access: CpuRecord, } impl Runtime { pub fn new(program: &Program) -> Self { + let record = ExecutionRecord:: { + program: Arc::new(program.clone()), + ..Default::default() + }; Self { clk: F::zero(), program: program.clone(), fp: F::zero(), pc: F::zero(), - memory: vec![F::zero(); 1024 * 1024], - record: ExecutionRecord::::default(), + memory: vec![MemoryEntry::default(); 1024 * 1024], + record, + access: CpuRecord::default(), } } + fn mr(&mut self, addr: F, position: AccessPosition) -> F { + let entry = self.memory[addr.as_canonical_u32() as usize].clone(); + let (prev_value, prev_timestamp) = (entry.value, entry.timestamp); + let record = MemoryRecord { + addr, + value: prev_value, + timestamp: self.timestamp(&position), + prev_value, + prev_timestamp, + }; + match position { + AccessPosition::A => self.access.a = Some(record), + AccessPosition::B => self.access.b = Some(record), + AccessPosition::C => self.access.c = Some(record), + _ => unreachable!(), + }; + prev_value + } + + fn mw(&mut self, addr: F, value: F, position: AccessPosition) { + let entry = self.memory[addr.as_canonical_u32() as usize].clone(); + let (prev_value, prev_timestamp) = (entry.value, entry.timestamp); + let timestamp = self.timestamp(&position); + let record = MemoryRecord { + addr, + value: prev_value, + timestamp, + prev_value, + prev_timestamp, + }; + self.memory[addr.as_canonical_u32() as usize] = MemoryEntry { value, timestamp }; + match position { + AccessPosition::A => self.access.a = Some(record), + AccessPosition::B => self.access.b = Some(record), + AccessPosition::C => self.access.c = Some(record), + _ => unreachable!(), + }; + } + + fn timestamp(&self, position: &AccessPosition) -> F { + self.clk + F::from_canonical_u32(*position as u32) + } + /// Fetch the destination address and input operand values for an ALU instruction. fn alu_rr(&mut self, instruction: &Instruction) -> (F, F, F) { if !instruction.imm_c { let a_ptr = self.fp + instruction.op_a; - let b_val = self.memory[(self.fp + instruction.op_b).as_canonical_u32() as usize]; - let c_val = self.memory[(self.fp + instruction.op_c).as_canonical_u32() as usize]; + let b_val = self.mr(self.fp + instruction.op_b, AccessPosition::B); + let c_val = self.mr(self.fp + instruction.op_c, AccessPosition::C); (a_ptr, b_val, c_val) } else { let a_ptr = self.fp + instruction.op_a; - let b_val = self.memory[(self.fp + instruction.op_b).as_canonical_u32() as usize]; + let b_val = self.mr(self.fp + instruction.op_b, AccessPosition::B); let c_val = instruction.op_c; (a_ptr, b_val, c_val) } @@ -62,8 +129,7 @@ impl Runtime { fn load_rr(&mut self, instruction: &Instruction) -> (F, F) { if !instruction.imm_b { let a_ptr = self.fp + instruction.op_a; - let b_ptr = self.memory[(self.fp + instruction.op_b).as_canonical_u32() as usize]; - let b = self.memory[(b_ptr).as_canonical_u32() as usize]; + let b = self.mr(self.fp + instruction.op_b, AccessPosition::B); (a_ptr, b) } else { let a_ptr = self.fp + instruction.op_a; @@ -75,7 +141,7 @@ impl Runtime { fn store_rr(&mut self, instruction: &Instruction) -> (F, F) { if !instruction.imm_b { let a_ptr = self.fp + instruction.op_a; - let b = self.memory[(self.fp + instruction.op_b).as_canonical_u32() as usize]; + let b = self.mr(self.fp + instruction.op_b, AccessPosition::B); (a_ptr, b) } else { let a_ptr = self.fp + instruction.op_a; @@ -85,9 +151,9 @@ impl Runtime { /// Fetch the input operand values for a branch instruction. fn branch_rr(&mut self, instruction: &Instruction) -> (F, F, F) { - let a = self.memory[(self.fp + instruction.op_a).as_canonical_u32() as usize]; + let a = self.mr(self.fp + instruction.op_a, AccessPosition::A); let b = if !instruction.imm_b { - self.memory[(self.fp + instruction.op_b).as_canonical_u32() as usize] + self.mr(self.fp + instruction.op_b, AccessPosition::B) } else { instruction.op_b }; @@ -105,37 +171,37 @@ impl Runtime { Opcode::ADD => { let (a_ptr, b_val, c_val) = self.alu_rr(&instruction); let a_val = b_val + c_val; - self.memory[(a_ptr).as_canonical_u32() as usize] = a_val; + self.mw(a_ptr, a_val, AccessPosition::A); (a, b, c) = (a_val, b_val, c_val); } Opcode::SUB => { let (a_ptr, b_val, c_val) = self.alu_rr(&instruction); let a_val = b_val - c_val; - self.memory[(a_ptr).as_canonical_u32() as usize] = a_val; + self.mw(a_ptr, a_val, AccessPosition::A); (a, b, c) = (a_val, b_val, c_val); } Opcode::MUL => { let (a_ptr, b_val, c_val) = self.alu_rr(&instruction); let a_val = b_val * c_val; - self.memory[(a_ptr).as_canonical_u32() as usize] = a_val; + self.mw(a_ptr, a_val, AccessPosition::A); (a, b, c) = (a_val, b_val, c_val); } Opcode::DIV => { let (a_ptr, b_val, c_val) = self.alu_rr(&instruction); let a_val = b_val / c_val; - self.memory[(a_ptr).as_canonical_u32() as usize] = a_val; + self.mw(a_ptr, a_val, AccessPosition::A); (a, b, c) = (a_val, b_val, c_val); } Opcode::LW => { let (a_ptr, b_val) = self.load_rr(&instruction); let a_val = b_val; - self.memory[(a_ptr).as_canonical_u32() as usize] = a_val; + self.mw(a_ptr, a_val, AccessPosition::A); (a, b, c) = (a_val, b_val, F::zero()); } Opcode::SW => { let (a_ptr, b_val) = self.store_rr(&instruction); let a_val = b_val; - self.memory[(a_ptr).as_canonical_u32() as usize] = a_val; + self.mw(a_ptr, a_val, AccessPosition::A); (a, b, c) = (a_val, b_val, F::zero()); } Opcode::BEQ => { @@ -153,7 +219,7 @@ impl Runtime { Opcode::JAL => { let imm = instruction.op_b; let a_ptr = instruction.op_a + self.fp; - self.memory[(a_ptr).as_canonical_u32() as usize] = self.pc; + self.mw(a_ptr, self.pc, AccessPosition::A); next_pc = self.pc + imm; (a, b, c) = (a_ptr, F::zero(), F::zero()); } @@ -161,16 +227,15 @@ impl Runtime { let imm = instruction.op_c; let b_ptr = instruction.op_b + self.fp; let a_ptr = instruction.op_a + self.fp; - let b_val = self.memory[(b_ptr).as_canonical_u32() as usize]; + let b_val = self.mr(b_ptr, AccessPosition::B); let c_val = imm; let a_val = self.pc + F::one(); - self.memory[(a_ptr).as_canonical_u32() as usize] = a_val; + self.mw(a_ptr, a_val, AccessPosition::A); next_pc = b_val + c_val; (a, b, c) = (a_val, b_val, c_val); } }; - self.pc = next_pc; let event = CpuEvent { clk: self.clk, pc: self.pc, @@ -183,6 +248,7 @@ impl Runtime { c, c_record: None, }; + self.pc = next_pc; self.record.cpu_events.push(event); self.clk += F::one(); } diff --git a/recursion/core/src/runtime/program.rs b/recursion/core/src/runtime/program.rs index d2abce11f8..0384e91480 100644 --- a/recursion/core/src/runtime/program.rs +++ b/recursion/core/src/runtime/program.rs @@ -1,6 +1,6 @@ use super::Instruction; -#[derive(Debug, Clone)] +#[derive(Debug, Clone, Default)] pub struct Program { /// The instructions of the program. pub instructions: Vec>, diff --git a/recursion/core/src/runtime/record.rs b/recursion/core/src/runtime/record.rs index 5aeaa54a4a..e04dd3501a 100644 --- a/recursion/core/src/runtime/record.rs +++ b/recursion/core/src/runtime/record.rs @@ -1,12 +1,21 @@ +use std::sync::Arc; + use hashbrown::HashMap; use p3_field::PrimeField32; use sp1_core::stark::MachineRecord; use crate::cpu::CpuEvent; +use crate::memory::MemoryRecord; + +use super::Program; #[derive(Default, Debug, Clone)] pub struct ExecutionRecord { + pub program: Arc>, pub cpu_events: Vec>, + pub first_memory_record: Vec<(u32, MemoryRecord, u32)>, + pub last_memory_record: Vec<(u32, MemoryRecord, u32)>, + pub program_memory_record: Vec<(u32, MemoryRecord, u32)>, } impl MachineRecord for ExecutionRecord { diff --git a/recursion/core/src/stark/mod.rs b/recursion/core/src/stark/mod.rs index da2b829fc5..47ebf5949a 100644 --- a/recursion/core/src/stark/mod.rs +++ b/recursion/core/src/stark/mod.rs @@ -1,4 +1,4 @@ -use crate::cpu::CpuChip; +use crate::{cpu::CpuChip, program::ProgramChip}; use p3_field::PrimeField32; use sp1_core::stark::{Chip, MachineStark, StarkGenericConfig}; use sp1_derive::MachineAir; @@ -8,6 +8,7 @@ use sp1_derive::MachineAir; #[execution_record_path = "crate::runtime::ExecutionRecord"] pub enum RecursionAir { Cpu(CpuChip), + Program(ProgramChip), } #[allow(dead_code)] @@ -24,6 +25,8 @@ impl RecursionAir { let mut chips = vec![]; let cpu = CpuChip::default(); chips.push(RecursionAir::Cpu(cpu)); + let program = ProgramChip; + chips.push(RecursionAir::Program(program)); chips } }