Ecall demo with Keccak-f (#717)

**Trace generation**

- Expand `StepRecord` to support multiple accesses to registers and
memory (beyond standard instructions).
- Introduce a module for syscall witness generation. See `fn
handle_syscall(…)`
- Witness generation for Keccak-f. Same format as sp1.
- Runtime wrapper function `syscall_keccak_permute`.
- Test program using Keccak-f and concrete values.

**Circuit part** ([standalone diff
here](https://github.com/Inversed-Tech/ceno/pull/1/files))

There is a placeholder for upcoming precompile circuits. The real
implementation must have the same effects as this one (same memory
writes, etc).

- `LargeEcallDummy` is an extension of the dummy circuit to support
precompiles.
- Support any number of registers and memory accesses.
- New testing technique using the real trace generator instead of
artificial StepRecord’s. See `keccak_step`.

---------

Co-authored-by: Aurélien Nicolas <[email protected]>

Cargo.toml

@@ -49,6 +49,7 @@ serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 strum = "0.26"
 strum_macros = "0.26"
+tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 tracing = { version = "0.1", features = [
   "attributes",
 ] }

ceno_emul/Cargo.toml

@@ -18,6 +18,7 @@ num-traits.workspace = true
 rrs_lib = { package = "rrs-succinct", version = "0.1.0" }
 strum.workspace = true
 strum_macros.workspace = true
+tiny-keccak.workspace = true
 tracing.workspace = true
 
 [features]

ceno_emul/src/host_utils.rs

@@ -5,7 +5,7 @@ use crate::{ByteAddr, EmuContext, VMState, WordAddr};
 const WORD_SIZE: usize = 4;
 const INFO_OUT_ADDR: WordAddr = ByteAddr(0xC000_0000).waddr();
 
-pub fn read_all_messages(state: &VMState) -> Vec<String> {
+pub fn read_all_messages(state: &VMState) -> Vec<Vec<u8>> {
     let mut offset: WordAddr = WordAddr::from(0);
     from_fn(move || match read_message(state, offset) {
         out if out.is_empty() => None,
@@ -17,16 +17,13 @@ pub fn read_all_messages(state: &VMState) -> Vec<String> {
     .collect()
 }
 
-fn read_message(state: &VMState, offset: WordAddr) -> String {
+fn read_message(state: &VMState, offset: WordAddr) -> Vec<u8> {
     let out_addr = INFO_OUT_ADDR + offset;
     let byte_len = state.peek_memory(out_addr) as usize;
 
-    String::from_utf8_lossy(
-        &(out_addr + 1_usize..)
-            .map(|address| state.peek_memory(address))
-            .flat_map(u32::to_le_bytes)
-            .take(byte_len)
-            .collect::<Vec<_>>(),
-    )
-    .to_string()
+    (out_addr + 1_usize..)
+        .map(|address| state.peek_memory(address))
+        .flat_map(u32::to_le_bytes)
+        .take(byte_len)
+        .collect::<Vec<_>>()
 }

ceno_emul/src/lib.rs

@@ -21,4 +21,10 @@ mod elf;
 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;

ceno_emul/src/syscalls.rs

@@ -0,0 +1,64 @@
+use crate::{RegIdx, Tracer, VMState, Word, WordAddr, WriteOp};
+use anyhow::Result;
+
+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
+
+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) -> Result<SyscallEffects> {
+    match function_code {
+        KECCAK_PERMUTE => Ok(keccak_permute::keccak_permute(vm)),
+        // TODO: introduce error types.
+        _ => Err(anyhow::anyhow!("Unknown syscall: {}", function_code)),
+    }
+}
+
+/// A syscall event, available to the circuit witness generators.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct SyscallWitness {
+    pub mem_ops: Vec<WriteOp>,
+    pub reg_ops: Vec<WriteOp>,
+}
+
+/// The effects of a syscall to apply on the VM.
+#[derive(Clone, Debug, Default, PartialEq, Eq)]
+pub struct SyscallEffects {
+    /// The witness being built. Get it with `finalize`.
+    witness: SyscallWitness,
+
+    /// The next PC after the syscall. Defaults to the next instruction.
+    pub next_pc: Option<u32>,
+}
+
+impl SyscallEffects {
+    /// Iterate over the register values after the syscall.
+    pub fn iter_reg_values(&self) -> impl Iterator<Item = (RegIdx, Word)> + '_ {
+        self.witness
+            .reg_ops
+            .iter()
+            .map(|op| (op.register_index(), op.value.after))
+    }
+
+    /// Iterate over the memory values after the syscall.
+    pub fn iter_mem_values(&self) -> impl Iterator<Item = (WordAddr, Word)> + '_ {
+        self.witness
+            .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 &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
+    }
+}

ceno_emul/src/syscalls/keccak_permute.rs

@@ -0,0 +1,65 @@
+use itertools::{Itertools, izip};
+use tiny_keccak::keccakf;
+
+use crate::{Change, EmuContext, Platform, VMState, WORD_SIZE, WordAddr, WriteOp};
+
+use super::{SyscallEffects, SyscallWitness};
+
+const KECCAK_CELLS: usize = 25; // u64 cells
+pub 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.
+pub fn keccak_permute(vm: &VMState) -> SyscallEffects {
+    let state_ptr = vm.peek_register(Platform::reg_arg0());
+
+    // Read the argument `state_ptr`.
+    let reg_ops = vec![WriteOp::new_register_op(
+        Platform::reg_arg0(),
+        Change::new(state_ptr, state_ptr),
+        0, // Cycle set later in finalize().
+    )];
+
+    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];
+        for (cell, (&lo, &hi)) in izip!(&mut state, input.iter().tuples()) {
+            *cell = lo as u64 | (hi as u64) << 32;
+        }
+
+        keccakf(&mut state);
+
+        state.into_iter().flat_map(|c| [c as u32, (c >> 32) as u32])
+    };
+
+    // Write permuted state.
+    let mem_ops = izip!(addrs, input, output)
+        .map(|(addr, before, after)| WriteOp {
+            addr,
+            value: Change { before, after },
+            previous_cycle: 0, // Cycle set later in finalize().
+        })
+        .collect_vec();
+
+    assert_eq!(mem_ops.len(), KECCAK_WORDS);
+    SyscallEffects {
+        witness: SyscallWitness { mem_ops, reg_ops },
+        next_pc: None,
+    }
+}

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<Instruction>) {
+    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::<Result<Vec<_>>>().unwrap();
+
+    (steps[2].clone(), instructions)
+}
+
+const fn load_immediate(rd: u32, imm: u32) -> Instruction {
+    encode_rv32u(InsnKind::ADDI, 0, 0, rd, imm)
+}

ceno_emul/src/tracer.rs

@@ -4,6 +4,7 @@ use crate::{
     CENO_PLATFORM, InsnKind, Instruction, PC_STEP_SIZE, Platform,
     addr::{ByteAddr, Cycle, RegIdx, Word, WordAddr},
     encode_rv32,
+    syscalls::{SyscallEffects, SyscallWitness},
 };
 
 /// An instruction and its context in an execution trace. That is concrete values of registers and memory.
@@ -29,6 +30,8 @@ pub struct StepRecord {
     rd: Option<WriteOp>,
 
     memory_op: Option<WriteOp>,
+
+    syscall: Option<SyscallWitness>,
 }
 
 #[derive(Clone, Debug, PartialEq, Eq, Hash)]
@@ -43,6 +46,14 @@ pub struct MemOp<T> {
 }
 
 impl<T> MemOp<T> {
+    pub fn new_register_op(idx: RegIdx, value: T, previous_cycle: Cycle) -> MemOp<T> {
+        MemOp {
+            addr: Platform::register_vma(idx).into(),
+            value,
+            previous_cycle,
+        }
+    }
+
     /// Get the register index of this operation.
     pub fn register_index(&self) -> RegIdx {
         Platform::register_index(self.addr.into())
@@ -240,6 +251,7 @@ impl StepRecord {
             }),
             insn,
             memory_op,
+            syscall: None,
         }
     }
 
@@ -275,6 +287,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)]
@@ -376,11 +392,18 @@ impl Tracer {
         });
     }
 
+    pub fn track_syscall(&mut self, effects: SyscallEffects) {
+        let witness = effects.finalize(self);
+
+        assert!(self.record.syscall.is_none(), "Only one syscall per step");
+        self.record.syscall = Some(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.
     /// - Accesses within the same instruction are distinguished by `subcycle ∈ [0, 3]`.
-    fn track_access(&mut self, addr: WordAddr, subcycle: Cycle) -> Cycle {
+    pub fn track_access(&mut self, addr: WordAddr, subcycle: Cycle) -> Cycle {
         self.latest_accesses
             .insert(addr, self.record.cycle + subcycle)
             .unwrap_or(0)

ceno_emul/src/vm_state.rs

@@ -6,6 +6,7 @@ use crate::{
     addr::{ByteAddr, RegIdx, Word, WordAddr},
     platform::Platform,
     rv32im::{Instruction, TrapCause},
+    syscalls::{SyscallEffects, handle_syscall},
     tracer::{Change, StepRecord, Tracer},
 };
 use anyhow::{Result, anyhow};
@@ -52,6 +53,10 @@ impl VMState {
 
     pub fn new_from_elf(platform: Platform, elf: &[u8]) -> Result<Self> {
         let program = Arc::new(Program::load_elf(elf, u32::MAX)?);
+        let platform = Platform {
+            prog_data: program.image.keys().copied().collect(),
+            ..platform
+        };
         Ok(Self::new(platform, program))
     }
 
@@ -104,30 +109,57 @@ impl VMState {
         self.set_pc(0.into());
         self.halted = true;
     }
+
+    fn apply_syscall(&mut self, effects: SyscallEffects) -> Result<()> {
+        for (addr, value) in effects.iter_mem_values() {
+            self.memory.insert(addr, value);
+        }
+
+        for (idx, value) in effects.iter_reg_values() {
+            self.registers[idx] = 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 {
     // Expect an ecall to terminate the program: function HALT with argument exit_code.
     fn ecall(&mut self) -> Result<bool> {
         let function = self.load_register(Platform::reg_ecall())?;
-        let arg0 = self.load_register(Platform::reg_arg0())?;
         if function == Platform::ecall_halt() {
-            tracing::debug!("halt with exit_code={}", arg0);
-
+            let exit_code = self.load_register(Platform::reg_arg0())?;
+            tracing::debug!("halt with exit_code={}", exit_code);
             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(Instruction::RD_NULL as RegIdx, 0)?;
-            // Example ecall effect - any writable address will do.
-            let addr = (self.platform.stack.end - 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)
+            match handle_syscall(self, function) {
+                Ok(effects) => {
+                    self.apply_syscall(effects)?;
+                    Ok(true)
+                }
+                Err(err) if self.platform.unsafe_ecall_nop => {
+                    tracing::warn!("ecall ignored with unsafe_ecall_nop: {:?}", err);
+                    // TODO: remove this example.
+                    // Treat unknown ecalls as all powerful instructions:
+                    // Read two registers, write one register, write one memory word, and branch.
+                    let _arg0 = self.load_register(Platform::reg_arg0())?;
+                    self.store_register(Instruction::RD_NULL as RegIdx, 0)?;
+                    // Example ecall effect - any writable address will do.
+                    let addr = (self.platform.stack.end - WORD_SIZE as u32).into();
+                    self.store_memory(addr, self.peek_memory(addr))?;
+                    self.set_pc(ByteAddr(self.pc) + PC_STEP_SIZE);
+                    Ok(true)
+                }
+                Err(err) => {
+                    tracing::error!("ecall error: {:?}", err);
+                    self.trap(TrapCause::EcallError)
+                }
+            }
         }
     }