risc0_zkvm/host/server/exec/
executor.rs

1// Copyright 2024 RISC Zero, Inc.
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::{cell::RefCell, io::Write, rc::Rc, sync::Arc, time::Instant};
16
17use anyhow::{Context as _, Result};
18use risc0_binfmt::{MemoryImage, Program};
19use risc0_circuit_rv32im::prove::emu::{
20    addr::ByteAddr,
21    exec::{
22        Executor, Syscall as NewSyscall, SyscallContext as NewSyscallContext,
23        DEFAULT_SEGMENT_LIMIT_PO2,
24    },
25};
26use risc0_core::scope;
27use risc0_zkp::core::digest::Digest;
28use risc0_zkvm_platform::{fileno, memory::GUEST_MAX_MEM, PAGE_SIZE};
29use tempfile::tempdir;
30
31use crate::{
32    host::client::env::SegmentPath, Assumptions, ExecutorEnv, FileSegmentRef, Output, Segment,
33    SegmentRef, Session,
34};
35
36use super::{
37    profiler::Profiler,
38    syscall::{SyscallContext, SyscallTable},
39};
40
41// The Executor provides an implementation for the execution phase.
42///
43/// The proving phase uses an execution trace generated by the Executor.
44pub struct ExecutorImpl<'a> {
45    env: ExecutorEnv<'a>,
46    image: MemoryImage,
47    pub(crate) syscall_table: SyscallTable<'a>,
48    profiler: Option<Rc<RefCell<Profiler>>>,
49}
50
51impl<'a> ExecutorImpl<'a> {
52    /// Construct a new [ExecutorImpl] from a [MemoryImage] and entry point.
53    ///
54    /// Before a guest program is proven, the [ExecutorImpl] is responsible for
55    /// deciding where a zkVM program should be split into [Segment]s and what
56    /// work will be done in each segment. This is the execution phase:
57    /// the guest program is executed to determine how its proof should be
58    /// divided into subparts.
59    pub fn new(env: ExecutorEnv<'a>, image: MemoryImage) -> Result<Self> {
60        Self::with_details(env, image, None)
61    }
62
63    /// Construct a new [ExecutorImpl] from the ELF binary of the guest program
64    /// you want to run and an [ExecutorEnv] containing relevant
65    /// environmental configuration details.
66    ///
67    /// # Example
68    /// ```
69    /// use risc0_zkvm::{ExecutorImpl, ExecutorEnv, Session};
70    /// use risc0_zkvm_methods::{BENCH_ELF, bench::BenchmarkSpec};
71    ///
72    /// let env = ExecutorEnv::builder()
73    ///     .write(&BenchmarkSpec::SimpleLoop { iters: 1 })
74    ///     .unwrap()
75    ///     .build()
76    ///     .unwrap();
77    /// let mut exec = ExecutorImpl::from_elf(env, BENCH_ELF).unwrap();
78    /// ```
79    pub fn from_elf(mut env: ExecutorEnv<'a>, elf: &[u8]) -> Result<Self> {
80        let program = Program::load_elf(elf, GUEST_MAX_MEM as u32)?;
81        let image = MemoryImage::new(&program, PAGE_SIZE as u32)?;
82
83        let profiler = if env.pprof_out.is_some() {
84            let profiler = Rc::new(RefCell::new(Profiler::new(elf, None)?));
85            env.trace.push(profiler.clone());
86            Some(profiler)
87        } else {
88            None
89        };
90
91        Self::with_details(env, image, profiler)
92    }
93
94    fn with_details(
95        env: ExecutorEnv<'a>,
96        image: MemoryImage,
97        profiler: Option<Rc<RefCell<Profiler>>>,
98    ) -> Result<Self> {
99        let syscall_table = SyscallTable::from_env(&env);
100        Ok(Self {
101            env,
102            image,
103            syscall_table,
104            profiler,
105        })
106    }
107
108    /// This will run the executor to get a [Session] which contain the results
109    /// of the execution.
110    pub fn run(&mut self) -> Result<Session> {
111        if self.env.segment_path.is_none() {
112            self.env.segment_path = Some(SegmentPath::TempDir(Arc::new(tempdir()?)));
113        }
114
115        let path = self.env.segment_path.clone().unwrap();
116        self.run_with_callback(|segment| Ok(Box::new(FileSegmentRef::new(&segment, &path)?)))
117    }
118
119    /// Run the executor until [crate::ExitCode::Halted] or
120    /// [crate::ExitCode::Paused] is reached, producing a [Session] as a result.
121    pub fn run_with_callback<F>(&mut self, mut callback: F) -> Result<Session>
122    where
123        F: FnMut(Segment) -> Result<Box<dyn SegmentRef>>,
124    {
125        scope!("execute");
126
127        let journal = Journal::default();
128        self.env
129            .posix_io
130            .borrow_mut()
131            .with_write_fd(fileno::JOURNAL, journal.clone());
132
133        let segment_limit_po2 = self
134            .env
135            .segment_limit_po2
136            .unwrap_or(DEFAULT_SEGMENT_LIMIT_PO2 as u32) as usize;
137
138        let mut refs = Vec::new();
139        let mut exec = Executor::new(
140            self.image.clone(),
141            self,
142            self.env.input_digest,
143            self.env.trace.clone(),
144        );
145
146        let start_time = Instant::now();
147        let result = exec.run(segment_limit_po2, self.env.session_limit, |inner| {
148            let output = inner
149                .exit_code
150                .expects_output()
151                .then(|| -> Option<Result<_>> {
152                    inner
153                        .output_digest
154                        .and_then(|digest| {
155                            (digest != Digest::ZERO).then(|| journal.buf.borrow().clone())
156                        })
157                        .map(|journal| {
158                            Ok(Output {
159                                journal: journal.into(),
160                                assumptions: Assumptions(
161                                    self.syscall_table
162                                        .assumptions_used
163                                        .borrow()
164                                        .iter()
165                                        .map(|(a, _)| a.clone().into())
166                                        .collect::<Vec<_>>(),
167                                )
168                                .into(),
169                            })
170                        })
171                })
172                .flatten()
173                .transpose()?;
174
175            let segment = Segment {
176                index: inner.index as u32,
177                inner,
178                output,
179            };
180            let segment_ref = callback(segment)?;
181            refs.push(segment_ref);
182            Ok(())
183        })?;
184        let elapsed = start_time.elapsed();
185
186        // Set the session_journal to the committed data iff the guest set a non-zero output.
187        let session_journal = result
188            .output_digest
189            .and_then(|digest| (digest != Digest::ZERO).then(|| journal.buf.take()));
190        if !result.exit_code.expects_output() && session_journal.is_some() {
191            tracing::debug!(
192                "dropping non-empty journal due to exit code {:?}: 0x{}",
193                result.exit_code,
194                hex::encode(journal.buf.borrow().as_slice())
195            );
196        };
197
198        // Take (clear out) the list of accessed assumptions.
199        // Leave the assumptions cache so it can be used if execution is resumed from pause.
200        let assumptions = self.syscall_table.assumptions_used.take();
201        let pending_zkrs = self.syscall_table.pending_zkrs.take();
202        let pending_keccaks = self.syscall_table.pending_keccaks.take();
203
204        if let Some(profiler) = self.profiler.take() {
205            let report = profiler.borrow_mut().finalize_to_vec();
206            std::fs::write(self.env.pprof_out.as_ref().unwrap(), report)?;
207        }
208
209        self.image = result.post_image.clone();
210        let syscall_metrics = self.syscall_table.metrics.borrow().clone();
211
212        let session = Session::new(
213            refs,
214            self.env.input_digest.unwrap_or_default(),
215            session_journal,
216            result.exit_code,
217            result.post_image,
218            assumptions,
219            result.user_cycles,
220            result.paging_cycles,
221            result.reserved_cycles,
222            result.total_cycles,
223            result.pre_state,
224            result.post_state,
225            pending_zkrs,
226            pending_keccaks,
227            result.ecall_metrics,
228            syscall_metrics,
229        );
230
231        tracing::info!("execution time: {elapsed:?}");
232        session.log();
233
234        Ok(session)
235    }
236}
237
238struct ContextAdapter<'a, 'b> {
239    ctx: &'b mut dyn NewSyscallContext,
240    syscall_table: SyscallTable<'a>,
241}
242
243impl<'a, 'b> SyscallContext<'a> for ContextAdapter<'a, 'b> {
244    fn get_pc(&self) -> u32 {
245        self.ctx.get_pc()
246    }
247
248    fn get_cycle(&self) -> u64 {
249        self.ctx.get_cycle()
250    }
251
252    fn load_register(&mut self, idx: usize) -> u32 {
253        self.ctx.peek_register(idx).unwrap()
254    }
255
256    fn load_page(&mut self, page_idx: u32) -> Result<Vec<u8>> {
257        self.ctx.peek_page(page_idx)
258    }
259
260    fn load_u8(&mut self, addr: ByteAddr) -> Result<u8> {
261        self.ctx.peek_u8(addr)
262    }
263
264    fn load_u32(&mut self, addr: ByteAddr) -> Result<u32> {
265        self.ctx.peek_u32(addr)
266    }
267
268    fn syscall_table(&self) -> &SyscallTable<'a> {
269        &self.syscall_table
270    }
271}
272
273impl<'a> NewSyscall for ExecutorImpl<'a> {
274    fn syscall(
275        &self,
276        syscall: &str,
277        ctx: &mut dyn NewSyscallContext,
278        into_guest: &mut [u32],
279    ) -> Result<(u32, u32)> {
280        let mut ctx = ContextAdapter {
281            ctx,
282            syscall_table: self.syscall_table.clone(),
283        };
284        self.syscall_table
285            .get_syscall(syscall)
286            .context(format!("Unknown syscall: {syscall:?}"))?
287            .borrow_mut()
288            .syscall(syscall, &mut ctx, into_guest)
289    }
290}
291
292// Capture the journal output in a buffer that we can access afterwards.
293#[derive(Clone, Default)]
294struct Journal {
295    buf: Rc<RefCell<Vec<u8>>>,
296}
297
298impl Write for Journal {
299    fn write(&mut self, bytes: &[u8]) -> std::io::Result<usize> {
300        self.buf.borrow_mut().write(bytes)
301    }
302
303    fn flush(&mut self) -> std::io::Result<()> {
304        self.buf.borrow_mut().flush()
305    }
306}
risc0_zkvm/host/server/exec/executor.rs

risc0_zkvm/host/server/exec/
executor.rs
