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#[cfg(feature = "aot")]
15use crate::common::*;
16
17#[derive(AlignedBytesBorrow, Clone)]
18#[repr(C)]
19struct AuiPcPreCompute {
20    imm: u32,
21    a: u8,
22}
23
24impl<A> Rv32AuipcExecutor<A> {
25    fn pre_compute_impl<F: PrimeField32>(
26        &self,
27        pc: u32,
28        inst: &Instruction<F>,
29        data: &mut AuiPcPreCompute,
30    ) -> Result<(), StaticProgramError> {
31        let Instruction { a, c: imm, d, .. } = inst;
32        if d.as_canonical_u32() != RV32_REGISTER_AS {
33            return Err(StaticProgramError::InvalidInstruction(pc));
34        }
35        let imm = imm.as_canonical_u32();
36        let data: &mut AuiPcPreCompute = data.borrow_mut();
37        *data = AuiPcPreCompute {
38            imm,
39            a: a.as_canonical_u32() as u8,
40        };
41        Ok(())
42    }
43}
44
45impl<F, A> InterpreterExecutor<F> for Rv32AuipcExecutor<A>
46where
47    F: PrimeField32,
48{
49    #[inline(always)]
50    fn pre_compute_size(&self) -> usize {
51        size_of::<AuiPcPreCompute>()
52    }
53
54    #[cfg(not(feature = "tco"))]
55    #[inline(always)]
56    fn pre_compute<Ctx: ExecutionCtxTrait>(
57        &self,
58        pc: u32,
59        inst: &Instruction<F>,
60        data: &mut [u8],
61    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError> {
62        let data: &mut AuiPcPreCompute = data.borrow_mut();
63        self.pre_compute_impl(pc, inst, data)?;
64        Ok(execute_e1_impl)
65    }
66
67    #[cfg(feature = "tco")]
68    fn handler<Ctx>(
69        &self,
70        pc: u32,
71        inst: &Instruction<F>,
72        data: &mut [u8],
73    ) -> Result<Handler<F, Ctx>, StaticProgramError>
74    where
75        Ctx: ExecutionCtxTrait,
76    {
77        let data: &mut AuiPcPreCompute = data.borrow_mut();
78        self.pre_compute_impl(pc, inst, data)?;
79        Ok(execute_e1_handler)
80    }
81}
82
83#[cfg(feature = "aot")]
84impl<F, A> AotExecutor<F> for Rv32AuipcExecutor<A>
85where
86    F: PrimeField32,
87{
88    fn generate_x86_asm(&self, inst: &Instruction<F>, pc: u32) -> Result<String, AotError> {
89        use openvm_instructions::riscv::RV32_CELL_BITS;
90
91        let mut asm_str = String::new();
92        let a = inst.a.as_canonical_u32() as u8;
93        let c = inst.c.as_canonical_u32();
94        let d = inst.d.as_canonical_u32();
95        let rd = pc.wrapping_add(c << RV32_CELL_BITS);
96
97        if d != RV32_REGISTER_AS {
98            return Err(AotError::InvalidInstruction);
99        }
100
101        let a_reg = a / 4;
102
103        if let Some(override_reg) = RISCV_TO_X86_OVERRIDE_MAP[a_reg as usize] {
104            asm_str += &format!("   mov {override_reg}, {rd}\n");
105        } else {
106            asm_str += &format!("   mov {REG_A_W}, {rd}\n");
107            asm_str += &gpr_to_xmm(REG_A_W, a_reg);
108        }
109
110        Ok(asm_str)
111    }
112
113    fn is_aot_supported(&self, _inst: &Instruction<F>) -> bool {
114        true
115    }
116}
117
118#[cfg(all(test, feature = "aot"))]
119mod tests {
120    use openvm_instructions::{riscv::RV32_REGISTER_AS, LocalOpcode};
121    use openvm_rv32im_transpiler::Rv32AuipcOpcode::AUIPC;
122    use openvm_stark_sdk::p3_baby_bear::BabyBear;
123
124    use super::*;
125
126    #[test]
127    fn aot_uses_full_24_bit_auipc_immediate() {
128        let executor = Rv32AuipcExecutor::new(());
129        let inst = Instruction::<BabyBear>::from_usize(
130            AUIPC.global_opcode(),
131            [4, 0, 0x10000, RV32_REGISTER_AS as usize, 0],
132        );
133
134        let asm = executor
135            .generate_x86_asm(&inst, 0)
136            .expect("valid AUIPC instruction should generate AOT assembly");
137
138        assert!(
139            asm.contains("16777216"),
140            "AUIPC AOT must shift the full raw 24-bit immediate"
141        );
142    }
143}
144
145impl<F, A> InterpreterMeteredExecutor<F> for Rv32AuipcExecutor<A>
146where
147    F: PrimeField32,
148{
149    fn metered_pre_compute_size(&self) -> usize {
150        size_of::<E2PreCompute<AuiPcPreCompute>>()
151    }
152
153    #[cfg(not(feature = "tco"))]
154    fn metered_pre_compute<Ctx>(
155        &self,
156        chip_idx: usize,
157        pc: u32,
158        inst: &Instruction<F>,
159        data: &mut [u8],
160    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
161    where
162        Ctx: MeteredExecutionCtxTrait,
163    {
164        let data: &mut E2PreCompute<AuiPcPreCompute> = data.borrow_mut();
165        data.chip_idx = chip_idx as u32;
166        self.pre_compute_impl(pc, inst, &mut data.data)?;
167        Ok(execute_e2_impl)
168    }
169
170    #[cfg(feature = "tco")]
171    fn metered_handler<Ctx>(
172        &self,
173        chip_idx: usize,
174        pc: u32,
175        inst: &Instruction<F>,
176        data: &mut [u8],
177    ) -> Result<Handler<F, Ctx>, StaticProgramError>
178    where
179        Ctx: MeteredExecutionCtxTrait,
180    {
181        let data: &mut E2PreCompute<AuiPcPreCompute> = data.borrow_mut();
182        data.chip_idx = chip_idx as u32;
183        self.pre_compute_impl(pc, inst, &mut data.data)?;
184        Ok(execute_e2_handler)
185    }
186}
187
188#[cfg(feature = "aot")]
189impl<F, A> AotMeteredExecutor<F> for Rv32AuipcExecutor<A>
190where
191    F: PrimeField32,
192{
193    fn is_aot_metered_supported(&self, _inst: &Instruction<F>) -> bool {
194        true
195    }
196    fn generate_x86_metered_asm(
197        &self,
198        inst: &Instruction<F>,
199        pc: u32,
200        chip_idx: usize,
201        _config: &SystemConfig,
202    ) -> Result<String, AotError> {
203        let mut asm_str = update_height_change_asm(chip_idx, 1)?;
204        asm_str += &self.generate_x86_asm(inst, pc)?;
205        Ok(asm_str)
206    }
207}
208
209#[inline(always)]
210unsafe fn execute_e12_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
211    pre_compute: &AuiPcPreCompute,
212    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
213) {
214    let pc = exec_state.pc();
215    let rd = run_auipc(pc, pre_compute.imm);
216    exec_state.vm_write(RV32_REGISTER_AS, pre_compute.a as u32, &rd);
217
218    exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP));
219}
220
221#[create_handler]
222#[inline(always)]
223unsafe fn execute_e1_impl<F: PrimeField32, CTX: ExecutionCtxTrait>(
224    pre_compute: *const u8,
225    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
226) {
227    let pre_compute: &AuiPcPreCompute =
228        std::slice::from_raw_parts(pre_compute, size_of::<AuiPcPreCompute>()).borrow();
229    execute_e12_impl(pre_compute, exec_state);
230}
231
232#[create_handler]
233#[inline(always)]
234unsafe fn execute_e2_impl<F: PrimeField32, CTX: MeteredExecutionCtxTrait>(
235    pre_compute: *const u8,
236    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
237) {
238    let pre_compute: &E2PreCompute<AuiPcPreCompute> =
239        std::slice::from_raw_parts(pre_compute, size_of::<E2PreCompute<AuiPcPreCompute>>())
240            .borrow();
241    exec_state
242        .ctx
243        .on_height_change(pre_compute.chip_idx as usize, 1);
244    execute_e12_impl(&pre_compute.data, exec_state);
245}
openvm_rv32im_circuit/auipc/execution.rs

openvm_rv32im_circuit/auipc/
execution.rs
