Ecall demo with Keccak-f

ceno_emul/Cargo.toml

@@ -17,6 +17,7 @@ num-derive.workspace = true
 num-traits.workspace = true
 strum.workspace = true
 strum_macros.workspace = true
+tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 tracing.workspace = true
 
 [dev-dependencies]

ceno_emul/src/lib.rs

@@ -19,3 +19,5 @@ pub use elf::Program;
 
 mod rv32im_encode;
 pub use rv32im_encode::encode_rv32;
+
+mod syscalls;

ceno_emul/src/syscalls.rs

@@ -0,0 +1,79 @@
+use crate::{Change, EmuContext, VMState, WORD_SIZE, WordAddr, WriteOp};
+use anyhow::Result;
+use itertools::{Itertools, izip};
+use tiny_keccak::keccakf;
+
+/// A syscall event, available to the circuit witness generators.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct SyscallWitness {
+    pub mem_writes: Vec<WriteOp>,
+}
+
+/// The effects of a syscall to apply on the VM.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct SyscallEffects {
+    pub witness: SyscallWitness,
+    pub return_value: Option<u32>,
+    pub next_pc: Option<u32>,
+}
+
+pub const KECCAK_PERMUTE: u32 = 0x00_01_01_09;
+
+/// Trace the inputs and effects of a syscall.
+pub fn handle_syscall(vm: &VMState, function_code: u32, arg0: u32) -> Result<SyscallEffects> {
+    match function_code {
+        KECCAK_PERMUTE => Ok(keccak_permute(vm, arg0)),
+        _ => Err(anyhow::anyhow!("Unknown syscall: {}", function_code)),
+    }
+}
+
+const KECCAK_CELLS: usize = 25; // u64 cells
+const KECCAK_WORDS: usize = KECCAK_CELLS * 2; // u32 words
+
+/// Trace the execution of a Keccak permutation.
+///
+/// Compatible with:
+/// https://github.com/succinctlabs/sp1/blob/013c24ea2fa15a0e7ed94f7d11a7ada4baa39ab9/crates/core/executor/src/syscalls/precompiles/keccak256/permute.rs
+///
+/// TODO: test compatibility.
+fn keccak_permute(vm: &VMState, state_ptr: u32) -> SyscallEffects {
+    let addrs = (state_ptr..)
+        .step_by(WORD_SIZE)
+        .take(KECCAK_WORDS)
+        .map(WordAddr::from)
+        .collect_vec();
+
+    // Read Keccak state.
+    let input = addrs
+        .iter()
+        .map(|&addr| vm.peek_memory(addr))
+        .collect::<Vec<_>>();
+
+    // Compute Keccak permutation.
+    let output = {
+        let mut state = [0_u64; KECCAK_CELLS];
+        izip!(state.iter_mut(), input.chunks_exact(2)).for_each(|(cell, chunk)| {
+            let lo = chunk[0] as u64;
+            let hi = chunk[1] as u64;
+            *cell = lo | hi << 32;
+        });
+        keccakf(&mut state);
+        state.into_iter().flat_map(|c| [c as u32, (c >> 32) as u32])
+    };
+
+    // Write permuted state.
+    let mem_writes = izip!(addrs, input, output)
+        .map(|(addr, before, after)| WriteOp {
+            addr,
+            value: Change { before, after },
+            previous_cycle: 0, // Set later by Tracer.
+        })
+        .collect_vec();
+
+    assert_eq!(mem_writes.len(), KECCAK_WORDS);
+    SyscallEffects {
+        witness: SyscallWitness { mem_writes },
+        return_value: None,
+        next_pc: None,
+    }
+}

ceno_emul/src/tracer.rs

@@ -5,6 +5,7 @@ use crate::{
     addr::{ByteAddr, Cycle, RegIdx, Word, WordAddr},
     encode_rv32,
     rv32im::DecodedInstruction,
+    syscalls::{SyscallEffects, SyscallWitness},
 };
 
 /// An instruction and its context in an execution trace. That is concrete values of registers and memory.
@@ -30,6 +31,8 @@ pub struct StepRecord {
     rd: Option<WriteOp>,
 
     memory_op: Option<WriteOp>,
+
+    syscall: Option<SyscallWitness>,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
@@ -242,6 +245,7 @@ impl StepRecord {
                 previous_cycle,
             }),
             memory_op,
+            syscall: None,
         }
     }
 
@@ -282,6 +286,10 @@ impl StepRecord {
     pub fn is_busy_loop(&self) -> bool {
         self.pc.before == self.pc.after
     }
+
+    pub fn syscall(&self) -> Option<&SyscallWitness> {
+        self.syscall.as_ref()
+    }
 }
 
 #[derive(Debug)]
@@ -386,6 +394,21 @@ impl Tracer {
         });
     }
 
+    pub fn track_syscall(&mut self, mut effects: SyscallEffects) {
+        let cycle = self.record.cycle + Self::SUBCYCLE_MEM;
+        for op in &mut effects.witness.mem_writes {
+            op.previous_cycle = self.track_access(op.addr, Self::SUBCYCLE_MEM);
+            assert_ne!(
+                op.previous_cycle, cycle,
+                "Memory address {:?} was accessed twice in the same cycle",
+                op.addr
+            );
+        }
+
+        assert!(self.record.syscall.is_none(), "Only one syscall per step");
+        self.record.syscall = Some(effects.witness);
+    }
+
     /// - Return the cycle when an address was last accessed.
     /// - Return 0 if this is the first access.
     /// - Record the current instruction as the origin of the latest access.

ceno_emul/src/vm_state.rs

@@ -6,6 +6,7 @@ use crate::{
     addr::{ByteAddr, RegIdx, Word, WordAddr},
     platform::Platform,
     rv32im::{DecodedInstruction, Emulator, TrapCause},
+    syscalls::{KECCAK_PERMUTE, SyscallEffects, handle_syscall},
     tracer::{Change, StepRecord, Tracer},
 };
 use anyhow::{Result, anyhow};
@@ -105,6 +106,22 @@ impl VMState {
         self.set_pc(0.into());
         self.halted = true;
     }
+
+    fn apply_syscall(&mut self, effects: SyscallEffects) -> Result<()> {
+        for write_op in &effects.witness.mem_writes {
+            self.memory.insert(write_op.addr, write_op.value.after);
+        }
+
+        if let Some(return_value) = effects.return_value {
+            self.store_register(Platform::reg_arg0(), return_value)?;
+        }
+
+        let next_pc = effects.next_pc.unwrap_or(self.pc + PC_STEP_SIZE as u32);
+        self.set_pc(next_pc.into());
+
+        self.tracer.track_syscall(effects);
+        Ok(())
+    }
 }
 
 impl EmuContext for VMState {
@@ -118,15 +135,22 @@ impl EmuContext for VMState {
             self.halt();
             Ok(true)
         } else if self.platform.unsafe_ecall_nop {
-            // Treat unknown ecalls as all powerful instructions:
-            // Read two registers, write one register, write one memory word, and branch.
-            tracing::warn!("ecall ignored: syscall_id={}", function);
-            self.store_register(DecodedInstruction::RD_NULL as RegIdx, 0)?;
-            // Example ecall effect - any writable address will do.
-            let addr = (self.platform.stack_top - WORD_SIZE as u32).into();
-            self.store_memory(addr, self.peek_memory(addr))?;
-            self.set_pc(ByteAddr(self.pc) + PC_STEP_SIZE);
-            Ok(true)
+            if function == KECCAK_PERMUTE {
+                let effects = handle_syscall(self, function, arg0)?;
+                self.apply_syscall(effects)?;
+                Ok(true)
+            } else {
+                // TODO: remove this example.
+                // Treat unknown ecalls as all powerful instructions:
+                // Read two registers, write one register, write one memory word, and branch.
+                tracing::warn!("ecall ignored: syscall_id={}", function);
+                self.store_register(DecodedInstruction::RD_NULL as RegIdx, 0)?;
+                // Example ecall effect - any writable address will do.
+                let addr = (self.platform.stack_top - WORD_SIZE as u32).into();
+                self.store_memory(addr, self.peek_memory(addr))?;
+                self.set_pc(ByteAddr(self.pc) + PC_STEP_SIZE);
+                Ok(true)
+            }
         } else {
             self.trap(TrapCause::EcallError)
         }

ceno_emul/tests/test_elf.rs

@@ -1,5 +1,7 @@
 use anyhow::Result;
 use ceno_emul::{ByteAddr, CENO_PLATFORM, EmuContext, InsnKind, Platform, StepRecord, VMState};
+use itertools::{Itertools, izip};
+use tiny_keccak::keccakf;
 
 #[test]
 fn test_ceno_rt_mini() -> Result<()> {
@@ -72,6 +74,70 @@ fn test_ceno_rt_io() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_ceno_rt_keccak() -> Result<()> {
+    let program_elf = ceno_examples::ceno_rt_keccak;
+    let mut state = VMState::new_from_elf(unsafe_platform(), program_elf)?;
+    let steps = run(&mut state)?;
+
+    // Expect the program to have written successive states between Keccak permutations.
+    const ITERATIONS: usize = 3;
+    let keccak_outs = sample_keccak_f(ITERATIONS);
+
+    let all_messages = read_all_messages(&state);
+    assert_eq!(all_messages.len(), ITERATIONS);
+    for (got, expect) in izip!(&all_messages, &keccak_outs) {
+        let got = got
+            .chunks_exact(8)
+            .map(|chunk| u64::from_le_bytes(chunk.try_into().unwrap()))
+            .collect_vec();
+        assert_eq!(&got, expect);
+    }
+
+    // Find the syscall records.
+    let syscalls = steps.iter().filter_map(|step| step.syscall()).collect_vec();
+    assert_eq!(syscalls.len(), ITERATIONS);
+
+    // Check the syscall effects.
+    for (witness, expect) in izip!(syscalls, keccak_outs) {
+        assert_eq!(witness.mem_writes.len(), expect.len() * 2);
+
+        let got = witness
+            .mem_writes
+            .chunks_exact(2)
+            .map(|write_ops| {
+                assert_eq!(
+                    write_ops[1].addr.baddr(),
+                    write_ops[0].addr.baddr() + WORD_SIZE as u32
+                );
+                let lo = write_ops[0].value.after as u64;
+                let hi = write_ops[1].value.after as u64;
+                lo | (hi << 32)
+            })
+            .collect_vec();
+        assert_eq!(got, expect);
+    }
+
+    Ok(())
+}
+
+fn unsafe_platform() -> Platform {
+    let mut platform = CENO_PLATFORM;
+    platform.unsafe_ecall_nop = true;
+    platform
+}
+
+fn sample_keccak_f(count: usize) -> Vec<Vec<u64>> {
+    let mut state = [0_u64; 25];
+
+    (0..count)
+        .map(|_| {
+            keccakf(&mut state);
+            state.into()
+        })
+        .collect_vec()
+}
+
 fn run(state: &mut VMState) -> Result<Vec<StepRecord>> {
     let steps = state.iter_until_halt().collect::<Result<Vec<_>>>()?;
     eprintln!("Emulator ran for {} steps.", steps.len());

ceno_rt/src/lib.rs

@@ -12,6 +12,9 @@ pub use io::info_out;
 mod params;
 pub use params::*;
 
+mod syscalls;
+pub use syscalls::*;
+
 #[cfg(not(test))]
 mod panic_handler {
     use core::panic::PanicInfo;

ceno_rt/src/syscalls.rs

@@ -0,0 +1,24 @@
+// Based on https://github.com/succinctlabs/sp1/blob/013c24ea2fa15a0e7ed94f7d11a7ada4baa39ab9/crates/zkvm/entrypoint/src/syscalls/keccak_permute.rs
+
+const KECCAK_PERMUTE: u32 = 0x00_01_01_09;
+
+use core::arch::asm;
+
+/// Executes the Keccak256 permutation on the given state.
+///
+/// ### Safety
+///
+/// The caller must ensure that `state` is valid pointer to data that is aligned along a four
+/// byte boundary.
+#[allow(unused_variables)]
+#[no_mangle]
+pub extern "C" fn syscall_keccak_permute(state: &mut [u64; 25]) {
+    unsafe {
+        asm!(
+            "ecall",
+            in("t0") KECCAK_PERMUTE,
+            in("a0") state as *mut [u64; 25],
+            in("a1") 0
+        );
+    }
+}

examples-builder/build.rs

@@ -14,6 +14,7 @@ const EXAMPLES: &[&str] = &[
     "ceno_rt_mem",
     "ceno_rt_mini",
     "ceno_rt_panic",
+    "ceno_rt_keccak",
 ];
 const CARGO_MANIFEST_DIR: &str = env!("CARGO_MANIFEST_DIR");
 

examples/examples/ceno_rt_keccak.rs

@@ -0,0 +1,28 @@
+//! Compute the Keccak permutation using a syscall.
+//!
+//! Iterate multiple times and log the state after each iteration.
+
+#![no_main]
+#![no_std]
+extern crate ceno_rt;
+use ceno_rt::{info_out, syscall_keccak_permute};
+use core::slice;
+
+const ITERATIONS: usize = 3;
+
+ceno_rt::entry!(main);
+fn main() {
+    let mut state = [0_u64; 25];
+
+    for _ in 0..ITERATIONS {
+        syscall_keccak_permute(&mut state);
+        log_state(&state);
+    }
+}
+
+fn log_state(state: &[u64; 25]) {
+    let out = unsafe {
+        slice::from_raw_parts(state.as_ptr() as *const u8, state.len() * size_of::<u64>())
+    };
+    info_out().write_frame(out);
+}