openvm_deferral_circuit/call/
execution.rs

1use std::{
2    array::from_fn,
3    borrow::{Borrow, BorrowMut},
4    slice::from_raw_parts,
5};
6
7use itertools::Itertools;
8use openvm_circuit::{arch::*, system::memory::online::GuestMemory};
9use openvm_circuit_primitives::AlignedBytesBorrow;
10use openvm_deferral_transpiler::DeferralOpcode;
11use openvm_instructions::{
12    instruction::Instruction,
13    program::DEFAULT_PC_STEP,
14    riscv::{RV32_MEMORY_AS, RV32_REGISTER_AS},
15    LocalOpcode, DEFERRAL_AS,
16};
17use openvm_stark_sdk::config::baby_bear_poseidon2::DIGEST_SIZE;
18
19use super::DeferralCallExecutor;
20use crate::{
21    poseidon2::deferral_poseidon2_chip,
22    utils::{
23        byte_commit_to_f, combine_output, join_memory_ops, memory_op_chunk, COMMIT_MEMORY_OPS,
24        COMMIT_NUM_BYTES, DIGEST_MEMORY_OPS, OUTPUT_TOTAL_MEMORY_OPS,
25    },
26    DeferralFn, CALL_AIR_REL_IDX, POSEIDON2_AIR_REL_IDX,
27};
28
29#[derive(AlignedBytesBorrow, Clone)]
30#[repr(C)]
31struct DeferralCallPrecompute<'a> {
32    rd_ptr: u32,
33    rs_ptr: u32,
34    deferral_idx: u32,
35    input_acc_ptr: u32,
36    output_acc_ptr: u32,
37    deferral_fn: &'a DeferralFn,
38}
39
40impl DeferralCallExecutor {
41    #[inline(always)]
42    fn pre_compute_impl<'a, F: VmField>(
43        &'a self,
44        pc: u32,
45        inst: &Instruction<F>,
46        data: &mut DeferralCallPrecompute<'a>,
47    ) -> Result<(), StaticProgramError> {
48        let Instruction {
49            opcode,
50            a,
51            b,
52            c,
53            d,
54            e,
55            ..
56        } = inst;
57
58        if opcode.local_opcode_idx(DeferralOpcode::CLASS_OFFSET) != DeferralOpcode::CALL as usize
59            || d.as_canonical_u32() != RV32_REGISTER_AS
60            || e.as_canonical_u32() != RV32_MEMORY_AS
61        {
62            return Err(StaticProgramError::InvalidInstruction(pc));
63        }
64
65        let deferral_idx = c.as_canonical_u32();
66        let deferral_fn = self
67            .deferral_fns
68            .get(deferral_idx as usize)
69            .ok_or(StaticProgramError::InvalidInstruction(pc))?;
70
71        const DIGEST_SIZE_U32: u32 = DIGEST_SIZE as u32;
72        let input_acc_ptr = 2 * deferral_idx * DIGEST_SIZE_U32;
73        *data = DeferralCallPrecompute {
74            rd_ptr: a.as_canonical_u32(),
75            rs_ptr: b.as_canonical_u32(),
76            deferral_idx,
77            input_acc_ptr,
78            output_acc_ptr: input_acc_ptr + DIGEST_SIZE_U32,
79            deferral_fn: deferral_fn.as_ref(),
80        };
81
82        Ok(())
83    }
84}
85
86impl<F: VmField> InterpreterExecutor<F> for DeferralCallExecutor {
87    fn pre_compute_size(&self) -> usize {
88        size_of::<DeferralCallPrecompute>()
89    }
90
91    #[cfg(not(feature = "tco"))]
92    fn pre_compute<Ctx>(
93        &self,
94        pc: u32,
95        inst: &Instruction<F>,
96        data: &mut [u8],
97    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
98    where
99        Ctx: ExecutionCtxTrait,
100    {
101        let pre_compute: &mut DeferralCallPrecompute = data.borrow_mut();
102        self.pre_compute_impl(pc, inst, pre_compute)?;
103        Ok(execute_e1_impl::<_, _>)
104    }
105
106    #[cfg(feature = "tco")]
107    fn handler<Ctx>(
108        &self,
109        pc: u32,
110        inst: &Instruction<F>,
111        data: &mut [u8],
112    ) -> Result<Handler<F, Ctx>, StaticProgramError>
113    where
114        Ctx: ExecutionCtxTrait,
115    {
116        let pre_compute: &mut DeferralCallPrecompute = data.borrow_mut();
117        self.pre_compute_impl(pc, inst, pre_compute)?;
118        Ok(execute_e1_handler::<_, _>)
119    }
120}
121
122#[cfg(feature = "aot")]
123impl<F: VmField> AotExecutor<F> for DeferralCallExecutor {}
124
125impl<F: VmField> InterpreterMeteredExecutor<F> for DeferralCallExecutor {
126    fn metered_pre_compute_size(&self) -> usize {
127        size_of::<E2PreCompute<DeferralCallPrecompute>>()
128    }
129
130    #[cfg(not(feature = "tco"))]
131    fn metered_pre_compute<Ctx>(
132        &self,
133        air_idx: usize,
134        pc: u32,
135        inst: &Instruction<F>,
136        data: &mut [u8],
137    ) -> Result<ExecuteFunc<F, Ctx>, StaticProgramError>
138    where
139        Ctx: MeteredExecutionCtxTrait,
140    {
141        let pre_compute: &mut E2PreCompute<DeferralCallPrecompute> = data.borrow_mut();
142        pre_compute.chip_idx = air_idx as u32;
143        self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
144        Ok(execute_e2_impl::<_, _>)
145    }
146
147    #[cfg(feature = "tco")]
148    fn metered_handler<Ctx>(
149        &self,
150        air_idx: usize,
151        pc: u32,
152        inst: &Instruction<F>,
153        data: &mut [u8],
154    ) -> Result<Handler<F, Ctx>, StaticProgramError>
155    where
156        Ctx: MeteredExecutionCtxTrait,
157    {
158        let pre_compute: &mut E2PreCompute<DeferralCallPrecompute> = data.borrow_mut();
159        pre_compute.chip_idx = air_idx as u32;
160        self.pre_compute_impl(pc, inst, &mut pre_compute.data)?;
161        Ok(execute_e2_handler::<_, _>)
162    }
163}
164
165#[cfg(feature = "aot")]
166impl<F: VmField> AotMeteredExecutor<F> for DeferralCallExecutor {}
167
168#[inline(always)]
169unsafe fn execute_e12_impl<F: VmField, CTX: ExecutionCtxTrait>(
170    pre_compute: &DeferralCallPrecompute,
171    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
172) {
173    let output_ptr = u32::from_le_bytes(exec_state.vm_read(RV32_REGISTER_AS, pre_compute.rd_ptr));
174    let input_ptr = u32::from_le_bytes(exec_state.vm_read(RV32_REGISTER_AS, pre_compute.rs_ptr));
175
176    let input_commit_chunks: [[u8; DEFAULT_BLOCK_SIZE]; COMMIT_MEMORY_OPS] = from_fn(|i| {
177        exec_state.vm_read(RV32_MEMORY_AS, input_ptr + (i * DEFAULT_BLOCK_SIZE) as u32)
178    });
179    let input_commit_bytes: [_; COMMIT_NUM_BYTES] = join_memory_ops(input_commit_chunks);
180    let input_commit: [F; _] = byte_commit_to_f(&input_commit_bytes.map(F::from_u8));
181    let old_input_acc_chunks: [[F; DEFAULT_BLOCK_SIZE]; DIGEST_MEMORY_OPS] = from_fn(|i| {
182        exec_state.vm_read(
183            DEFERRAL_AS,
184            pre_compute.input_acc_ptr + (i * DEFAULT_BLOCK_SIZE) as u32,
185        )
186    });
187    let old_output_acc_chunks: [[F; DEFAULT_BLOCK_SIZE]; DIGEST_MEMORY_OPS] = from_fn(|i| {
188        exec_state.vm_read(
189            DEFERRAL_AS,
190            pre_compute.output_acc_ptr + (i * DEFAULT_BLOCK_SIZE) as u32,
191        )
192    });
193    let old_input_acc = join_memory_ops(old_input_acc_chunks);
194    let old_output_acc = join_memory_ops(old_output_acc_chunks);
195
196    let poseidon2_chip = deferral_poseidon2_chip();
197    let (output_commit, output_len) = pre_compute.deferral_fn.execute(
198        &input_commit_bytes.to_vec(),
199        &mut exec_state.streams.deferrals[pre_compute.deferral_idx as usize],
200        pre_compute.deferral_idx,
201        &poseidon2_chip,
202    );
203    let output_f_commit =
204        byte_commit_to_f(&output_commit.iter().map(|v| F::from_u8(*v)).collect_vec());
205
206    // (output_commit, output_len) pair, corresponds to guest struct OutputKey
207    let output_key = combine_output(output_commit, output_len.to_le_bytes());
208
209    let new_input_acc = poseidon2_chip.perm(&old_input_acc, &input_commit, true);
210    let new_output_acc = poseidon2_chip.perm(&old_output_acc, &output_f_commit, true);
211
212    for chunk_idx in 0..OUTPUT_TOTAL_MEMORY_OPS {
213        exec_state.vm_write::<u8, DEFAULT_BLOCK_SIZE>(
214            RV32_MEMORY_AS,
215            output_ptr + (chunk_idx * DEFAULT_BLOCK_SIZE) as u32,
216            &memory_op_chunk(&output_key, chunk_idx),
217        );
218    }
219    for chunk_idx in 0..DIGEST_MEMORY_OPS {
220        exec_state.vm_write::<F, DEFAULT_BLOCK_SIZE>(
221            DEFERRAL_AS,
222            pre_compute.input_acc_ptr + (chunk_idx * DEFAULT_BLOCK_SIZE) as u32,
223            &memory_op_chunk(&new_input_acc, chunk_idx),
224        );
225    }
226    for chunk_idx in 0..DIGEST_MEMORY_OPS {
227        exec_state.vm_write::<F, DEFAULT_BLOCK_SIZE>(
228            DEFERRAL_AS,
229            pre_compute.output_acc_ptr + (chunk_idx * DEFAULT_BLOCK_SIZE) as u32,
230            &memory_op_chunk(&new_output_acc, chunk_idx),
231        );
232    }
233
234    let pc = exec_state.pc();
235    exec_state.set_pc(pc.wrapping_add(DEFAULT_PC_STEP));
236}
237
238#[create_handler]
239#[inline(always)]
240unsafe fn execute_e1_impl<F: VmField, CTX: ExecutionCtxTrait>(
241    pre_compute: *const u8,
242    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
243) {
244    let pre_compute: &DeferralCallPrecompute =
245        from_raw_parts(pre_compute, size_of::<DeferralCallPrecompute>()).borrow();
246    execute_e12_impl(pre_compute, exec_state);
247}
248
249#[create_handler]
250#[inline(always)]
251unsafe fn execute_e2_impl<F: VmField, CTX: MeteredExecutionCtxTrait>(
252    pre_compute: *const u8,
253    exec_state: &mut VmExecState<F, GuestMemory, CTX>,
254) {
255    let pre_compute: &E2PreCompute<DeferralCallPrecompute> = from_raw_parts(
256        pre_compute,
257        size_of::<E2PreCompute<DeferralCallPrecompute>>(),
258    )
259    .borrow();
260
261    execute_e12_impl(&pre_compute.data, exec_state);
262    exec_state
263        .ctx
264        .on_height_change(pre_compute.chip_idx as usize, 1);
265
266    // The Poseidon2 peripheral chip's height also increases as a result of
267    // this opcode's execution. In DEFER_CALL, both the input and output
268    // hash accumulator for some deferral circuit are updated.
269    exec_state.ctx.on_height_change(
270        pre_compute.chip_idx as usize + (CALL_AIR_REL_IDX - POSEIDON2_AIR_REL_IDX),
271        2,
272    );
273}
openvm_deferral_circuit/call/execution.rs

openvm_deferral_circuit/call/
execution.rs
