Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Ecall demo with Keccak-f #717

Open
wants to merge 20 commits into
base: master
Choose a base branch
from
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Cargo.lock

Some generated files are not rendered by default. Learn more about how customized files appear on GitHub.

1 change: 1 addition & 0 deletions ceno_emul/Cargo.toml
Original file line number Diff line number Diff line change
Expand Up @@ -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 = { version = "2.0.2", features = ["keccak"] }
tracing.workspace = true

[dev-dependencies]
Expand Down
2 changes: 2 additions & 0 deletions ceno_emul/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -20,3 +20,5 @@ mod elf;
pub use elf::Program;

pub mod disassemble;

mod syscalls;
61 changes: 61 additions & 0 deletions ceno_emul/src/syscalls.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
use crate::{RegIdx, Tracer, VMState, Word, WordAddr, WriteOp};
use anyhow::Result;
use itertools::chain;

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;
naure marked this conversation as resolved.
Show resolved Hide resolved

/// 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)),
_ => 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_writes: Vec<WriteOp>,
pub reg_accesses: 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_accesses
.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_writes
.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);
}
self.witness
}
}
68 changes: 68 additions & 0 deletions ceno_emul/src/syscalls/keccak_permute.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,68 @@
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
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_accesses = 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_writes = 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_writes.len(), KECCAK_WORDS);
SyscallEffects {
witness: SyscallWitness {
mem_writes,
reg_accesses,
},
next_pc: None,
}
}
25 changes: 24 additions & 1 deletion ceno_emul/src/tracer.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand All @@ -29,6 +30,8 @@ pub struct StepRecord {
rd: Option<WriteOp>,

memory_op: Option<WriteOp>,

syscall: Option<SyscallWitness>,
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
Expand All @@ -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())
Expand Down Expand Up @@ -240,6 +251,7 @@ impl StepRecord {
}),
insn,
memory_op,
syscall: None,
}
}

Expand Down Expand Up @@ -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)]
Expand Down Expand Up @@ -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)
Expand Down
56 changes: 42 additions & 14 deletions ceno_emul/src/vm_state.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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};
Expand Down Expand Up @@ -104,30 +105,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_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)
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.
naure marked this conversation as resolved.
Show resolved Hide resolved
// 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_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)
}
Err(err) => {
tracing::error!("ecall error: {:?}", err);
self.trap(TrapCause::EcallError)
}
}
}
}

Expand Down
71 changes: 71 additions & 0 deletions ceno_emul/tests/test_elf.rs
Original file line number Diff line number Diff line change
@@ -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<()> {
Expand Down Expand Up @@ -72,6 +74,75 @@ 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.reg_accesses.len(), 1);
assert_eq!(
witness.reg_accesses[0].register_index(),
Platform::reg_arg0()
);

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)
naure marked this conversation as resolved.
Show resolved Hide resolved
})
.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());
Expand Down
1 change: 1 addition & 0 deletions examples-builder/build.rs
Original file line number Diff line number Diff line change
Expand Up @@ -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");

Expand Down
Loading
Loading