openvm_circuit/arch/testing/program/
mod.rs

1use std::{borrow::BorrowMut, mem::size_of, sync::Arc};
2
3use openvm_instructions::instruction::Instruction;
4use openvm_stark_backend::{
5    config::{StarkGenericConfig, Val},
6    p3_field::{Field, FieldAlgebra, PrimeField32},
7    p3_matrix::dense::RowMajorMatrix,
8    prover::{cpu::CpuBackend, types::AirProvingContext},
9    Chip, ChipUsageGetter,
10};
11
12use crate::{
13    arch::ExecutionState,
14    system::program::{ProgramBus, ProgramExecutionCols},
15};
16
17pub mod air;
18#[cfg(feature = "cuda")]
19mod cuda;
20#[cfg(feature = "cuda")]
21pub use cuda::*;
22
23#[derive(Debug)]
24pub struct ProgramTester<F: Field> {
25    pub bus: ProgramBus,
26    pub records: Vec<ProgramExecutionCols<F>>,
27}
28
29impl<F: PrimeField32> ProgramTester<F> {
30    pub fn new(bus: ProgramBus) -> Self {
31        Self {
32            bus,
33            records: vec![],
34        }
35    }
36
37    pub fn execute(&mut self, instruction: &Instruction<F>, initial_state: &ExecutionState<u32>) {
38        self.records.push(ProgramExecutionCols {
39            pc: F::from_canonical_u32(initial_state.pc),
40            opcode: instruction.opcode.to_field(),
41            a: instruction.a,
42            b: instruction.b,
43            c: instruction.c,
44            d: instruction.d,
45            e: instruction.e,
46            f: instruction.f,
47            g: instruction.g,
48        });
49    }
50}
51
52impl<F: Field> ProgramTester<F> {
53    fn width() -> usize {
54        size_of::<ProgramExecutionCols<u8>>() + 1
55    }
56}
57
58impl<SC: StarkGenericConfig, RA> Chip<RA, CpuBackend<SC>> for ProgramTester<Val<SC>> {
59    fn generate_proving_ctx(&self, _: RA) -> AirProvingContext<CpuBackend<SC>> {
60        let height = self.records.len().next_power_of_two();
61        let width = self.trace_width();
62        let mut values = Val::<SC>::zero_vec(height * width);
63        // This zip only goes through records. The padding rows between records.len()..height
64        // are filled with zeros - in particular count = 0 so nothing is added to bus.
65        for (row, record) in values.chunks_mut(width).zip(&self.records) {
66            *(row[..width - 1]).borrow_mut() = *record;
67            row[width - 1] = Val::<SC>::ONE;
68        }
69        AirProvingContext::simple_no_pis(Arc::new(RowMajorMatrix::new(values, width)))
70    }
71}
72
73impl<F: Field> ChipUsageGetter for ProgramTester<F> {
74    fn air_name(&self) -> String {
75        "ProgramDummyAir".to_string()
76    }
77    fn current_trace_height(&self) -> usize {
78        self.records.len()
79    }
80    fn trace_width(&self) -> usize {
81        Self::width()
82    }
83}