openvm_rv32im_circuit/auipc/
execution.rs

1use std::{
2    borrow::{Borrow, BorrowMut},
3    mem::size_of,
4};
5
6use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
7use openvm_circuit_primitives_derive::AlignedBytesBorrow;
8use openvm_instructions::{
9    instruction::Instruction, program::DEFAULT_PC_STEP, riscv::RV32_REGISTER_AS,
10};
11use openvm_stark_backend::p3_field::PrimeField32;
12
13use super::{run_auipc, Rv32AuipcExecutor};
14
15#[derive(AlignedBytesBorrow, Clone)]
16#[repr(C)]
17struct AuiPcPreCompute {
18    imm: u32,
19    a: u8,
20}
21
22impl<A> Rv32AuipcExecutor<A> {
23    fn pre_compute_impl<F: PrimeField32>(
24        &self,
25        pc: u32,
26        inst: &Instruction<F>,
27        data: &mut AuiPcPreCompute,
28    ) -> Result<(), StaticProgramError> {
29        let Instruction { a, c: imm, d, .. } = inst;
30        if d.as_canonical_u32() != RV32_REGISTER_AS {
31            return Err(StaticProgramError::InvalidInstruction(pc));
32        }
33        let imm = imm.as_canonical_u32();
34        let data: &mut AuiPcPreCompute = data.borrow_mut();
35        *data = AuiPcPreCompute {
36            imm,
37            a: a.as_canonical_u32() as u8,
38        };
39        Ok(())
40    }
41}
42
43impl<F, A> Executor<F> for Rv32AuipcExecutor<A>
44where
45    F: PrimeField32,
46{
47    #[inline(always)]
48    fn pre_compute_size(&self) -> usize {
49        size_of::<AuiPcPreCompute>()
50    }
51
52    #[inline(always)]
53    fn pre_compute<Ctx: ExecutionCtxTrait>(
54        &self,
55        pc: u32,
56        inst: &Instruction<F>,
57        data: &mut [u8],
58    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError> {
59        let data: &mut AuiPcPreCompute = data.borrow_mut();
60        self.pre_compute_impl(pc, inst, data)?;
61        Ok(execute_e1_impl)
62    }
63
64    #[cfg(feature = "tco")]
65    fn handler<Ctx>(
66        &self,
67        pc: u32,
68        inst: &Instruction<F>,
69        data: &mut [u8],
70    ) -> Result<Handler<F, Ctx>, StaticProgramError>
71    where
72        Ctx: ExecutionCtxTrait,
73    {
74        let data: &mut AuiPcPreCompute = data.borrow_mut();
75        self.pre_compute_impl(pc, inst, data)?;
76        Ok(execute_e1_tco_handler)
77    }
78}
79
80impl<F, A> MeteredExecutor<F> for Rv32AuipcExecutor<A>
81where
82    F: PrimeField32,
83{
84    fn metered_pre_compute_size(&self) -> usize {
85        size_of::<E2PreCompute<AuiPcPreCompute>>()
86    }
87
88    fn metered_pre_compute<Ctx>(
89        &self,
90        chip_idx: usize,
91        pc: u32,
92        inst: &Instruction<F>,
93        data: &mut [u8],
94    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
95    where
96        Ctx: MeteredExecutionCtxTrait,
97    {
98        let data: &mut E2PreCompute<AuiPcPreCompute> = data.borrow_mut();
99        data.chip_idx = chip_idx as u32;
100        self.pre_compute_impl(pc, inst, &mut data.data)?;
101        Ok(execute_e2_impl)
102    }
103
104    #[cfg(feature = "tco")]
105    fn metered_handler<Ctx>(
106        &self,
107        chip_idx: usize,
108        pc: u32,
109        inst: &Instruction<F>,
110        data: &mut [u8],
111    ) -> Result<Handler<F, Ctx>, StaticProgramError>
112    where
113        Ctx: MeteredExecutionCtxTrait,
114    {
115        let data: &mut E2PreCompute<AuiPcPreCompute> = data.borrow_mut();
116        data.chip_idx = chip_idx as u32;
117        self.pre_compute_impl(pc, inst, &mut data.data)?;
118        Ok(execute_e2_tco_handler)
119    }
120}
121
122#[inline(always)]
123unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
124    pre_compute: &AuiPcPreCompute,
125    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
126) {
127    let rd = run_auipc(vm_state.pc, pre_compute.imm);
128    vm_state.vm_write(RV32_REGISTER_AS, pre_compute.a as u32, &rd);
129
130    vm_state.pc = vm_state.pc.wrapping_add(DEFAULT_PC_STEP);
131    vm_state.instret += 1;
132}
133
134#[create_tco_handler]
135unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
136    pre_compute: &[u8],
137    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
138) {
139    let pre_compute: &AuiPcPreCompute = pre_compute.borrow();
140    execute_e12_impl(pre_compute, vm_state);
141}
142
143#[create_tco_handler]
144unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait>(
145    pre_compute: &[u8],
146    vm_state: &mut VmExecState<F, GuestMemory, CTX>,
147) {
148    let pre_compute: &E2PreCompute<AuiPcPreCompute> = pre_compute.borrow();
149    vm_state
150        .ctx
151        .on_height_change(pre_compute.chip_idx as usize, 1);
152    execute_e12_impl(&pre_compute.data, vm_state);
153}