airbender_host/runner/
transpiler_runner.rs

1use super::{resolve_cycles, ExecutionResult, FlamegraphConfig, Runner};
2use crate::error::{HostError, Result};
3use crate::receipt::Receipt;
4use riscv_transpiler::abstractions::non_determinism::QuasiUARTSource;
5use riscv_transpiler::common_constants::{
6    rom::ROM_SECOND_WORD_BITS, INITIAL_TIMESTAMP, TIMESTAMP_STEP,
7};
8use riscv_transpiler::cycle::CycleMarkerHooks;
9use riscv_transpiler::ir::{preprocess_bytecode, FullUnsignedMachineDecoderConfig};
10#[cfg(target_arch = "x86_64")]
11use riscv_transpiler::jit::JittedCode;
12use riscv_transpiler::jit::RAM_SIZE;
13use riscv_transpiler::vm::{
14    DelegationsCounters, FlamegraphConfig as VmFlamegraphConfig, RamWithRomRegion, SimpleTape,
15    State, VmFlamegraphProfiler, VM,
16};
17use std::io::Read;
18use std::path::{Path, PathBuf};
19
20/// Builder for creating a configured transpiler runner.
21pub struct TranspilerRunnerBuilder {
22    app_bin_path: PathBuf,
23    cycles: Option<usize>,
24    text_path: Option<PathBuf>,
25    flamegraph: Option<FlamegraphConfig>,
26    use_jit: bool,
27}
28
29impl TranspilerRunnerBuilder {
30    pub fn new(app_bin_path: impl AsRef<Path>) -> Self {
31        Self {
32            app_bin_path: app_bin_path.as_ref().to_path_buf(),
33            cycles: None,
34            text_path: None,
35            flamegraph: None,
36            use_jit: false,
37        }
38    }
39
40    pub fn with_cycles(mut self, cycles: usize) -> Self {
41        self.cycles = Some(cycles);
42        self
43    }
44
45    pub fn maybe_cycles(self, cycles: Option<usize>) -> Self {
46        match cycles {
47            Some(v) => self.with_cycles(v),
48            None => self,
49        }
50    }
51
52    pub fn with_text_path(mut self, text_path: impl AsRef<Path>) -> Self {
53        self.text_path = Some(text_path.as_ref().to_path_buf());
54        self
55    }
56
57    pub fn maybe_text_path(self, text_path: Option<impl AsRef<Path>>) -> Self {
58        match text_path {
59            Some(v) => self.with_text_path(v),
60            None => self,
61        }
62    }
63
64    pub fn with_flamegraph(mut self, flamegraph: FlamegraphConfig) -> Self {
65        self.flamegraph = Some(flamegraph);
66        self
67    }
68
69    pub fn with_jit(mut self) -> Self {
70        self.use_jit = true;
71        self
72    }
73
74    pub fn build(self) -> Result<TranspilerRunner> {
75        if self.use_jit && cfg!(not(target_arch = "x86_64")) {
76            return Err(HostError::Transpiler(
77                "JIT execution is only available on x86_64 targets".to_string(),
78            ));
79        }
80
81        let app_bin_path = resolve_app_bin_path(&self.app_bin_path)?;
82        let app_text_path = self
83            .text_path
84            .as_deref()
85            .map(resolve_text_path)
86            .unwrap_or_else(|| resolve_text_path(&derive_text_path(&app_bin_path)))?;
87        let cycles = resolve_cycles(self.cycles)?;
88
89        Ok(TranspilerRunner {
90            app_bin_path,
91            app_text_path,
92            cycles,
93            flamegraph: self.flamegraph,
94            use_jit: self.use_jit,
95        })
96    }
97}
98
99/// Transpiler based execution runner.
100pub struct TranspilerRunner {
101    app_bin_path: PathBuf,
102    app_text_path: PathBuf,
103    cycles: usize,
104    flamegraph: Option<FlamegraphConfig>,
105    use_jit: bool,
106}
107
108impl Runner for TranspilerRunner {
109    fn run(&self, input_words: &[u32]) -> Result<ExecutionResult> {
110        if self.flamegraph.is_some() {
111            return self.run_without_jit_with_flamegraph(input_words);
112        }
113
114        if self.use_jit {
115            return self.run_with_jit(input_words);
116        }
117
118        self.run_without_jit(input_words)
119    }
120}
121
122impl TranspilerRunner {
123    #[cfg(target_arch = "x86_64")]
124    fn run_with_jit(&self, input_words: &[u32]) -> Result<ExecutionResult> {
125        let bin_words = read_u32_words(&self.app_bin_path)?;
126        let text_words = read_u32_words(&self.app_text_path)?;
127        let mut non_determinism_source = QuasiUARTSource::new_with_reads(input_words.to_vec());
128
129        let cycles_bound = match u32::try_from(self.cycles) {
130            Ok(value) => Some(value),
131            Err(_) => {
132                tracing::warn!(
133                    "cycles limit {} exceeds u32::MAX; running transpiler without a cycle bound",
134                    self.cycles
135                );
136                None
137            }
138        };
139
140        let (state, _memory) = JittedCode::run_alternative_simulator(
141            &text_words,
142            &mut non_determinism_source,
143            &bin_words,
144            cycles_bound,
145        );
146        let cycles_executed = ((state.timestamp - INITIAL_TIMESTAMP) / TIMESTAMP_STEP) as usize;
147
148        Ok(ExecutionResult {
149            receipt: Receipt::from_registers(state.registers),
150            cycles_executed,
151            reached_end: true,
152            cycle_markers: None,
153        })
154    }
155
156    #[cfg(not(target_arch = "x86_64"))]
157    fn run_with_jit(&self, _input_words: &[u32]) -> Result<ExecutionResult> {
158        Err(HostError::Transpiler(
159            "JIT execution is only available on x86_64 targets".to_string(),
160        ))
161    }
162
163    fn run_without_jit(&self, input_words: &[u32]) -> Result<ExecutionResult> {
164        self.run_without_jit_internal(input_words, None)
165    }
166
167    fn run_without_jit_with_flamegraph(&self, input_words: &[u32]) -> Result<ExecutionResult> {
168        let flamegraph = self
169            .flamegraph
170            .as_ref()
171            .ok_or_else(|| HostError::Transpiler("flamegraph options are missing".to_string()))?;
172
173        let symbols_path = flamegraph
174            .elf_path
175            .clone()
176            .unwrap_or_else(|| derive_elf_path(&self.app_bin_path));
177        let mut profiler_config = VmFlamegraphConfig::new(symbols_path, flamegraph.output.clone());
178        profiler_config.frequency_recip = flamegraph.sampling_rate;
179        profiler_config.reverse_graph = flamegraph.inverse;
180        let mut profiler = VmFlamegraphProfiler::new(profiler_config).map_err(|err| {
181            HostError::Transpiler(format!("failed to initialize flamegraph profiler: {err}"))
182        })?;
183
184        self.run_without_jit_internal(input_words, Some(&mut profiler))
185    }
186
187    fn run_without_jit_internal(
188        &self,
189        input_words: &[u32],
190        profiler: Option<&mut VmFlamegraphProfiler>,
191    ) -> Result<ExecutionResult> {
192        let bin_words = read_u32_words(&self.app_bin_path)?;
193        let text_words = read_u32_words(&self.app_text_path)?;
194        let instructions = preprocess_bytecode::<FullUnsignedMachineDecoderConfig>(&text_words);
195        let instruction_tape = SimpleTape::new(&instructions);
196        let mut ram =
197            RamWithRomRegion::<{ ROM_SECOND_WORD_BITS }>::from_rom_content(&bin_words, RAM_SIZE);
198        let mut state = State::initial_with_counters(DelegationsCounters::default());
199        let mut non_determinism_source = QuasiUARTSource::new_with_reads(input_words.to_vec());
200
201        let (reached_end, cycle_markers) = CycleMarkerHooks::with(|| match profiler {
202            Some(profiler) => {
203                VM::<DelegationsCounters, CycleMarkerHooks>::run_basic_unrolled_with_flamegraph::<
204                    _,
205                    _,
206                    _,
207                >(
208                    &mut state,
209                    &mut ram,
210                    &mut (),
211                    &instruction_tape,
212                    self.cycles,
213                    &mut non_determinism_source,
214                    profiler,
215                )
216                .map_err(|err| {
217                    HostError::Transpiler(format!("failed to generate flamegraph: {err}"))
218                })
219            }
220            None => Ok(
221                VM::<DelegationsCounters, CycleMarkerHooks>::run_basic_unrolled::<_, _, _>(
222                    &mut state,
223                    &mut ram,
224                    &mut (),
225                    &instruction_tape,
226                    self.cycles,
227                    &mut non_determinism_source,
228                ),
229            ),
230        });
231        let reached_end = reached_end?;
232
233        let cycles_executed = ((state.timestamp - INITIAL_TIMESTAMP) / TIMESTAMP_STEP) as usize;
234        let registers = state.registers.map(|register| register.value);
235
236        Ok(ExecutionResult {
237            receipt: Receipt::from_registers(registers),
238            cycles_executed,
239            reached_end,
240            cycle_markers: Some(cycle_markers.into()),
241        })
242    }
243}
244
245fn resolve_app_bin_path(path: &Path) -> Result<PathBuf> {
246    if !path.exists() {
247        return Err(HostError::Transpiler(format!(
248            "binary not found: {}",
249            path.display()
250        )));
251    }
252
253    path.canonicalize().map_err(|err| {
254        HostError::Transpiler(format!(
255            "failed to canonicalize binary path {}: {err}",
256            path.display()
257        ))
258    })
259}
260
261fn resolve_text_path(path: &Path) -> Result<PathBuf> {
262    if !path.exists() {
263        return Err(HostError::Transpiler(format!(
264            "text file not found: {}",
265            path.display()
266        )));
267    }
268
269    path.canonicalize().map_err(|err| {
270        HostError::Transpiler(format!(
271            "failed to canonicalize text path {}: {err}",
272            path.display()
273        ))
274    })
275}
276
277fn derive_text_path(bin_path: &Path) -> PathBuf {
278    let mut text_path = bin_path.to_path_buf();
279    text_path.set_extension("text");
280    text_path
281}
282
283fn derive_elf_path(bin_path: &Path) -> PathBuf {
284    let mut elf_path = bin_path.to_path_buf();
285    elf_path.set_extension("elf");
286    elf_path
287}
288
289fn read_u32_words(path: &Path) -> Result<Vec<u32>> {
290    let mut file = std::fs::File::open(path).map_err(|err| {
291        HostError::Transpiler(format!("failed to open {}: {err}", path.display()))
292    })?;
293    let mut bytes = Vec::new();
294    file.read_to_end(&mut bytes).map_err(|err| {
295        HostError::Transpiler(format!("failed to read {}: {err}", path.display()))
296    })?;
297
298    if bytes.len() % 4 != 0 {
299        return Err(HostError::Transpiler(format!(
300            "file length is not a multiple of 4: {}",
301            path.display()
302        )));
303    }
304
305    let mut words = Vec::with_capacity(bytes.len() / 4);
306    for chunk in bytes.as_chunks::<4>().0 {
307        words.push(u32::from_le_bytes(*chunk));
308    }
309    Ok(words)
310}
311
312#[cfg(test)]
313mod tests {
314    use super::TranspilerRunnerBuilder;
315    use crate::runner::Runner;
316    use std::path::Path;
317
318    const MARKER_OPCODE: u32 = 0x7ff01073; // csrrw x0, 2047, x0
319    const ADDI_OPCODE: u32 = 0x00100093; // addi x1, x0, 1
320    const LOOP_OPCODE: u32 = 0x0000006f; // jal x0, 0
321
322    // TODO: Evaluate how low-level do we want tests to be
323    #[test]
324    fn collects_cycle_markers_for_interpreter_runs() {
325        let dir = tempfile::tempdir().expect("create temp dir");
326        let bin_path = dir.path().join("app.bin");
327        let text_path = dir.path().join("app.text");
328        let program = [MARKER_OPCODE, ADDI_OPCODE, MARKER_OPCODE, LOOP_OPCODE];
329        write_program(&bin_path, &program);
330        write_program(&text_path, &program);
331
332        let runner = TranspilerRunnerBuilder::new(&bin_path)
333            .with_text_path(&text_path)
334            .with_cycles(program.len())
335            .build()
336            .expect("build runner");
337        let execution = runner.run(&[]).expect("run program");
338        let markers = execution.cycle_markers.expect("cycle markers");
339
340        assert!(execution.reached_end);
341        assert_eq!(execution.receipt.registers[1], 1);
342        assert_eq!(markers.markers.len(), 2);
343        assert!(markers.delegation_counter.is_empty());
344        let diff = markers.markers[1].diff(&markers.markers[0]);
345        assert_eq!(diff.cycles, 1);
346        assert!(diff.delegations.is_empty());
347    }
348
349    #[cfg(target_arch = "x86_64")]
350    #[test]
351    fn jit_runs_do_not_collect_cycle_markers() {
352        let dir = tempfile::tempdir().expect("create temp dir");
353        let bin_path = dir.path().join("app.bin");
354        let text_path = dir.path().join("app.text");
355        let program = [ADDI_OPCODE, LOOP_OPCODE];
356        write_program(&bin_path, &program);
357        write_program(&text_path, &program);
358
359        let runner = TranspilerRunnerBuilder::new(&bin_path)
360            .with_text_path(&text_path)
361            .with_cycles(program.len())
362            .with_jit()
363            .build()
364            .expect("build runner");
365        let execution = runner.run(&[]).expect("run program");
366
367        assert_eq!(execution.receipt.registers[1], 1);
368        assert!(execution.cycle_markers.is_none());
369    }
370
371    fn write_program(path: &Path, program: &[u32]) {
372        let bytes: Vec<u8> = program.iter().flat_map(|word| word.to_le_bytes()).collect();
373        std::fs::write(path, bytes).expect("write test program");
374    }
375}
airbender_host/runner/transpiler_runner.rs

airbender_host/runner/
transpiler_runner.rs
