diff --git a/ceno_emul/src/lib.rs b/ceno_emul/src/lib.rs index 63da9823c..5e9d871ef 100644 --- a/ceno_emul/src/lib.rs +++ b/ceno_emul/src/lib.rs @@ -23,5 +23,8 @@ pub use elf::Program; pub mod disassemble; mod syscalls; +pub use syscalls::{KECCAK_PERMUTE, keccak_permute::KECCAK_WORDS}; + +pub mod test_utils; pub mod host_utils; diff --git a/ceno_emul/src/syscalls.rs b/ceno_emul/src/syscalls.rs index 5980506cb..56419016e 100644 --- a/ceno_emul/src/syscalls.rs +++ b/ceno_emul/src/syscalls.rs @@ -1,8 +1,7 @@ use crate::{RegIdx, Tracer, VMState, Word, WordAddr, WriteOp}; use anyhow::Result; -use itertools::chain; -mod keccak_permute; +pub mod keccak_permute; // Using the same function codes as sp1: // https://github.com/succinctlabs/sp1/blob/013c24ea2fa15a0e7ed94f7d11a7ada4baa39ab9/crates/core/executor/src/syscalls/code.rs @@ -20,8 +19,8 @@ pub fn handle_syscall(vm: &VMState, function_code: u32) -> Result, - pub reg_accesses: Vec, + pub mem_ops: Vec, + pub reg_ops: Vec, } /// The effects of a syscall to apply on the VM. @@ -38,7 +37,7 @@ impl SyscallEffects { /// Iterate over the register values after the syscall. pub fn iter_reg_values(&self) -> impl Iterator + '_ { self.witness - .reg_accesses + .reg_ops .iter() .map(|op| (op.register_index(), op.value.after)) } @@ -46,15 +45,18 @@ impl SyscallEffects { /// Iterate over the memory values after the syscall. pub fn iter_mem_values(&self) -> impl Iterator + '_ { self.witness - .mem_writes + .mem_ops .iter() .map(|op| (op.addr, op.value.after)) } /// Keep track of the cycles of registers and memory accesses. pub fn finalize(mut self, tracer: &mut Tracer) -> SyscallWitness { - for op in chain(&mut self.witness.reg_accesses, &mut self.witness.mem_writes) { - op.previous_cycle = tracer.track_access(op.addr, 0); + for op in &mut self.witness.reg_ops { + op.previous_cycle = tracer.track_access(op.addr, Tracer::SUBCYCLE_RD); + } + for op in &mut self.witness.mem_ops { + op.previous_cycle = tracer.track_access(op.addr, Tracer::SUBCYCLE_MEM); } self.witness } diff --git a/ceno_emul/src/syscalls/keccak_permute.rs b/ceno_emul/src/syscalls/keccak_permute.rs index 63decd3eb..091725f9c 100644 --- a/ceno_emul/src/syscalls/keccak_permute.rs +++ b/ceno_emul/src/syscalls/keccak_permute.rs @@ -6,7 +6,7 @@ use crate::{Change, EmuContext, Platform, VMState, WORD_SIZE, WordAddr, WriteOp} use super::{SyscallEffects, SyscallWitness}; const KECCAK_CELLS: usize = 25; // u64 cells -const KECCAK_WORDS: usize = KECCAK_CELLS * 2; // u32 words +pub const KECCAK_WORDS: usize = KECCAK_CELLS * 2; // u32 words /// Trace the execution of a Keccak permutation. /// @@ -18,7 +18,7 @@ pub fn keccak_permute(vm: &VMState) -> SyscallEffects { let state_ptr = vm.peek_register(Platform::reg_arg0()); // Read the argument `state_ptr`. - let reg_accesses = vec![WriteOp::new_register_op( + let reg_ops = vec![WriteOp::new_register_op( Platform::reg_arg0(), Change::new(state_ptr, state_ptr), 0, // Cycle set later in finalize(). @@ -49,7 +49,7 @@ pub fn keccak_permute(vm: &VMState) -> SyscallEffects { }; // Write permuted state. - let mem_writes = izip!(addrs, input, output) + let mem_ops = izip!(addrs, input, output) .map(|(addr, before, after)| WriteOp { addr, value: Change { before, after }, @@ -57,12 +57,9 @@ pub fn keccak_permute(vm: &VMState) -> SyscallEffects { }) .collect_vec(); - assert_eq!(mem_writes.len(), KECCAK_WORDS); + assert_eq!(mem_ops.len(), KECCAK_WORDS); SyscallEffects { - witness: SyscallWitness { - mem_writes, - reg_accesses, - }, + witness: SyscallWitness { mem_ops, reg_ops }, next_pc: None, } } diff --git a/ceno_emul/src/test_utils.rs b/ceno_emul/src/test_utils.rs new file mode 100644 index 000000000..07a5a817c --- /dev/null +++ b/ceno_emul/src/test_utils.rs @@ -0,0 +1,28 @@ +use crate::{ + CENO_PLATFORM, InsnKind, Instruction, Platform, Program, StepRecord, VMState, encode_rv32, + encode_rv32u, syscalls::KECCAK_PERMUTE, +}; +use anyhow::Result; + +pub fn keccak_step() -> (StepRecord, Vec) { + let instructions = vec![ + // Call Keccak-f. + load_immediate(Platform::reg_arg0() as u32, CENO_PLATFORM.heap.start), + load_immediate(Platform::reg_ecall() as u32, KECCAK_PERMUTE), + encode_rv32(InsnKind::ECALL, 0, 0, 0, 0), + // Halt. + load_immediate(Platform::reg_ecall() as u32, Platform::ecall_halt()), + encode_rv32(InsnKind::ECALL, 0, 0, 0, 0), + ]; + + let pc = CENO_PLATFORM.pc_base(); + let program = Program::new(pc, pc, instructions.clone(), Default::default()); + let mut vm = VMState::new(CENO_PLATFORM, program.into()); + let steps = vm.iter_until_halt().collect::>>().unwrap(); + + (steps[2].clone(), instructions) +} + +const fn load_immediate(rd: u32, imm: u32) -> Instruction { + encode_rv32u(InsnKind::ADDI, 0, 0, rd, imm) +} diff --git a/ceno_host/tests/test_elf.rs b/ceno_host/tests/test_elf.rs index 7b665ddff..0fb6c0f46 100644 --- a/ceno_host/tests/test_elf.rs +++ b/ceno_host/tests/test_elf.rs @@ -127,15 +127,12 @@ fn test_ceno_rt_keccak() -> Result<()> { // Check the syscall effects. for (witness, expect) in izip!(syscalls, keccak_outs) { - assert_eq!(witness.reg_accesses.len(), 1); - assert_eq!( - witness.reg_accesses[0].register_index(), - Platform::reg_arg0() - ); + assert_eq!(witness.reg_ops.len(), 1); + assert_eq!(witness.reg_ops[0].register_index(), Platform::reg_arg0()); - assert_eq!(witness.mem_writes.len(), expect.len() * 2); + assert_eq!(witness.mem_ops.len(), expect.len() * 2); let got = witness - .mem_writes + .mem_ops .chunks_exact(2) .map(|write_ops| { assert_eq!( diff --git a/ceno_zkvm/src/instructions/riscv/dummy/dummy_circuit.rs b/ceno_zkvm/src/instructions/riscv/dummy/dummy_circuit.rs index b73c932a1..0ac8bfaa6 100644 --- a/ceno_zkvm/src/instructions/riscv/dummy/dummy_circuit.rs +++ b/ceno_zkvm/src/instructions/riscv/dummy/dummy_circuit.rs @@ -95,7 +95,7 @@ pub struct DummyConfig { impl DummyConfig { #[allow(clippy::too_many_arguments)] - fn construct_circuit( + pub(super) fn construct_circuit( circuit_builder: &mut CircuitBuilder, kind: InsnKind, with_rs1: bool, @@ -213,7 +213,7 @@ impl DummyConfig { }) } - fn assign_instance( + pub(super) fn assign_instance( &self, instance: &mut [::BaseField], lk_multiplicity: &mut LkMultiplicity, @@ -264,4 +264,8 @@ impl DummyConfig { Ok(()) } + + pub(super) fn ts(&self) -> WitIn { + self.vm_state.ts + } } diff --git a/ceno_zkvm/src/instructions/riscv/dummy/dummy_ecall.rs b/ceno_zkvm/src/instructions/riscv/dummy/dummy_ecall.rs new file mode 100644 index 000000000..cd3156387 --- /dev/null +++ b/ceno_zkvm/src/instructions/riscv/dummy/dummy_ecall.rs @@ -0,0 +1,137 @@ +use std::marker::PhantomData; + +use ceno_emul::{Change, InsnKind, KECCAK_WORDS, StepRecord, WORD_SIZE}; +use ff_ext::ExtensionField; +use itertools::Itertools; + +use super::{super::insn_base::WriteMEM, dummy_circuit::DummyConfig}; +use crate::{ + Value, + circuit_builder::CircuitBuilder, + error::ZKVMError, + expression::{ToExpr, WitIn}, + instructions::{ + Instruction, + riscv::{constants::UInt, insn_base::WriteRD}, + }, + set_val, + witness::LkMultiplicity, +}; + +trait EcallSpec { + const NAME: &'static str; + + const REG_OPS_COUNT: usize; + const MEM_OPS_COUNT: usize; +} + +pub struct KeccakSpec; + +impl EcallSpec for KeccakSpec { + const NAME: &'static str = "KECCAK"; + + const REG_OPS_COUNT: usize = 1; + const MEM_OPS_COUNT: usize = KECCAK_WORDS; +} + +/// LargeEcallDummy can handle any instruction and produce its effects, +/// including multiple memory operations. +/// +/// Unsafe: The content is not constrained. +pub struct LargeEcallDummy(PhantomData<(E, S)>); + +impl Instruction for LargeEcallDummy { + type InstructionConfig = LargeEcallConfig; + + fn name() -> String { + format!("{}_DUMMY", S::NAME) + } + + fn construct_circuit(cb: &mut CircuitBuilder) -> Result { + let dummy_insn = DummyConfig::construct_circuit( + cb, + InsnKind::ECALL, + true, // Read the ecall function code. + false, + false, + false, + false, + false, + )?; + + let start_addr = cb.create_witin(|| "mem_addr"); + + let reg_writes = (0..S::REG_OPS_COUNT) + .map(|i| { + let val_after = UInt::new_unchecked(|| format!("reg_after_{}", i), cb)?; + + WriteRD::construct_circuit(cb, val_after.register_expr(), dummy_insn.ts()) + .map(|writer| (val_after, writer)) + }) + .collect::, _>>()?; + + let mem_writes = (0..S::MEM_OPS_COUNT) + .map(|i| { + let val_before = cb.create_witin(|| format!("mem_before_{}", i)); + let val_after = cb.create_witin(|| format!("mem_after_{}", i)); + + WriteMEM::construct_circuit( + cb, + start_addr.expr() + (i * WORD_SIZE) as u64, + val_before.expr(), + val_after.expr(), + dummy_insn.ts(), + ) + .map(|writer| (Change::new(val_before, val_after), writer)) + }) + .collect::, _>>()?; + + Ok(LargeEcallConfig { + dummy_insn, + start_addr, + reg_writes, + mem_writes, + }) + } + + fn assign_instance( + config: &Self::InstructionConfig, + instance: &mut [E::BaseField], + lk_multiplicity: &mut LkMultiplicity, + step: &StepRecord, + ) -> Result<(), ZKVMError> { + let ops = &step.syscall().expect("syscall step"); + + // Assign instruction. + config + .dummy_insn + .assign_instance(instance, lk_multiplicity, step)?; + + set_val!(instance, config.start_addr, u64::from(ops.mem_ops[0].addr)); + + // Assign registers. + for ((value, writer), op) in config.reg_writes.iter().zip_eq(&ops.reg_ops) { + value.assign_value(instance, Value::new_unchecked(op.value.after)); + writer.assign_op(instance, lk_multiplicity, step.cycle(), op)?; + } + + // Assign memory. + for ((value, writer), op) in config.mem_writes.iter().zip_eq(&ops.mem_ops) { + set_val!(instance, value.before, op.value.before as u64); + set_val!(instance, value.after, op.value.after as u64); + writer.assign_op(instance, lk_multiplicity, step.cycle(), op)?; + } + + Ok(()) + } +} + +#[derive(Debug)] +pub struct LargeEcallConfig { + dummy_insn: DummyConfig, + + reg_writes: Vec<(UInt, WriteRD)>, + + start_addr: WitIn, + mem_writes: Vec<(Change, WriteMEM)>, +} diff --git a/ceno_zkvm/src/instructions/riscv/dummy/mod.rs b/ceno_zkvm/src/instructions/riscv/dummy/mod.rs index 9c912f422..02114e737 100644 --- a/ceno_zkvm/src/instructions/riscv/dummy/mod.rs +++ b/ceno_zkvm/src/instructions/riscv/dummy/mod.rs @@ -12,5 +12,8 @@ mod dummy_circuit; pub use dummy_circuit::DummyInstruction; +mod dummy_ecall; +pub use dummy_ecall::LargeEcallDummy; + #[cfg(test)] mod test; diff --git a/ceno_zkvm/src/instructions/riscv/dummy/test.rs b/ceno_zkvm/src/instructions/riscv/dummy/test.rs index 5f3a03bb9..199d8a3a8 100644 --- a/ceno_zkvm/src/instructions/riscv/dummy/test.rs +++ b/ceno_zkvm/src/instructions/riscv/dummy/test.rs @@ -1,4 +1,5 @@ use ceno_emul::{Change, InsnKind, StepRecord, encode_rv32}; +use dummy_ecall::KeccakSpec; use goldilocks::GoldilocksExt2; use super::*; @@ -37,6 +38,30 @@ fn test_dummy_ecall() { MockProver::assert_satisfied_raw(&cb, raw_witin, &[insn_code], None, Some(lkm)); } +#[test] +fn test_dummy_keccak() { + type KeccakDummy = LargeEcallDummy; + + let mut cs = ConstraintSystem::::new(|| "riscv"); + let mut cb = CircuitBuilder::new(&mut cs); + let config = cb + .namespace( + || "keccak_dummy", + |cb| { + let config = KeccakDummy::construct_circuit(cb); + Ok(config) + }, + ) + .unwrap() + .unwrap(); + + let (step, program) = ceno_emul::test_utils::keccak_step(); + let (raw_witin, lkm) = + KeccakDummy::assign_instances(&config, cb.cs.num_witin as usize, vec![step]).unwrap(); + + MockProver::assert_satisfied_raw(&cb, raw_witin, &program, None, Some(lkm)); +} + #[test] fn test_dummy_r() { let mut cs = ConstraintSystem::::new(|| "riscv"); diff --git a/ceno_zkvm/src/instructions/riscv/insn_base.rs b/ceno_zkvm/src/instructions/riscv/insn_base.rs index 468eb3623..057a27bb5 100644 --- a/ceno_zkvm/src/instructions/riscv/insn_base.rs +++ b/ceno_zkvm/src/instructions/riscv/insn_base.rs @@ -1,6 +1,7 @@ -use ceno_emul::{StepRecord, Word}; +use ceno_emul::{Cycle, StepRecord, Word, WriteOp}; use ff::Field; use ff_ext::ExtensionField; +use goldilocks::SmallField; use itertools::Itertools; use super::constants::{PC_STEP_SIZE, UINT_LIMBS, UInt}; @@ -221,6 +222,16 @@ impl WriteRD { step: &StepRecord, ) -> Result<(), ZKVMError> { let op = step.rd().expect("rd op"); + self.assign_op(instance, lk_multiplicity, step.cycle(), &op) + } + + pub fn assign_op( + &self, + instance: &mut [E::BaseField], + lk_multiplicity: &mut LkMultiplicity, + cycle: Cycle, + op: &WriteOp, + ) -> Result<(), ZKVMError> { set_val!(instance, self.id, op.register_index() as u64); set_val!(instance, self.prev_ts, op.previous_cycle); @@ -235,7 +246,7 @@ impl WriteRD { instance, lk_multiplicity, op.previous_cycle, - step.cycle() + Tracer::SUBCYCLE_RD, + cycle + Tracer::SUBCYCLE_RD, )?; Ok(()) @@ -331,17 +342,24 @@ impl WriteMEM { lk_multiplicity: &mut LkMultiplicity, step: &StepRecord, ) -> Result<(), ZKVMError> { - set_val!( - instance, - self.prev_ts, - step.memory_op().unwrap().previous_cycle - ); + let op = step.memory_op().unwrap(); + self.assign_op(instance, lk_multiplicity, step.cycle(), &op) + } + + pub fn assign_op( + &self, + instance: &mut [F], + lk_multiplicity: &mut LkMultiplicity, + cycle: Cycle, + op: &WriteOp, + ) -> Result<(), ZKVMError> { + set_val!(instance, self.prev_ts, op.previous_cycle); self.lt_cfg.assign_instance( instance, lk_multiplicity, - step.memory_op().unwrap().previous_cycle, - step.cycle() + Tracer::SUBCYCLE_MEM, + op.previous_cycle, + cycle + Tracer::SUBCYCLE_MEM, )?; Ok(())