openvm_deferral_circuit/extension/
mod.rs

1use std::sync::Arc;
2
3use derive_more::derive::From;
4use openvm_circuit::{
5    arch::{
6        AirInventory, AirInventoryError, ChipInventory, ChipInventoryError, ExecutionBridge,
7        ExecutorInventoryBuilder, ExecutorInventoryError, InitFileGenerator, MatrixRecordArena,
8        RowMajorMatrixArena, SystemConfig, VmBuilder, VmChipComplex, VmCircuitExtension,
9        VmExecutionExtension, VmField, VmProverExtension,
10    },
11    system::{memory::SharedMemoryHelper, SystemChipInventory, SystemCpuBuilder, SystemExecutor},
12};
13use openvm_circuit_derive::{AnyEnum, Executor, MeteredExecutor, PreflightExecutor, VmConfig};
14use openvm_circuit_primitives::bitwise_op_lookup::{
15    BitwiseOperationLookupAir, BitwiseOperationLookupBus, BitwiseOperationLookupChip,
16    SharedBitwiseOperationLookupChip,
17};
18use openvm_cpu_backend::{CpuBackend, CpuDevice};
19use openvm_deferral_transpiler::DeferralOpcode;
20use openvm_instructions::LocalOpcode;
21use openvm_rv32im_circuit::{
22    Rv32I, Rv32IExecutor, Rv32ImCpuProverExt, Rv32Io, Rv32IoExecutor, Rv32M, Rv32MExecutor,
23};
24use openvm_stark_backend::{StarkEngine, StarkProtocolConfig, Val};
25use serde::{Deserialize, Serialize};
26
27use crate::{
28    call::{
29        DeferralCallAdapterAir, DeferralCallAdapterExecutor, DeferralCallAdapterFiller,
30        DeferralCallAir, DeferralCallChip, DeferralCallCoreAir, DeferralCallCoreFiller,
31        DeferralCallExecutor,
32    },
33    count::{DeferralCircuitCountAir, DeferralCircuitCountBus, DeferralCircuitCountChip},
34    output::{DeferralOutputAir, DeferralOutputChip, DeferralOutputExecutor, DeferralOutputFiller},
35    poseidon2::{
36        deferral_poseidon2_air, deferral_poseidon2_chip, DeferralPoseidon2Air, DeferralPoseidon2Bus,
37    },
38    utils::COMMIT_NUM_BYTES,
39    DeferralFn,
40};
41
42cfg_if::cfg_if! {
43    if #[cfg(feature = "cuda")] {
44        mod cuda;
45        pub use self::cuda::DeferralGpuProverExt as DeferralProverExt;
46        pub use self::cuda::Rv32DeferralGpuBuilder as Rv32DeferralBuilder;
47
48    } else {
49        pub use self::DeferralCpuProverExt as DeferralProverExt;
50        pub use self::Rv32DeferralCpuBuilder as Rv32DeferralBuilder;
51    }
52}
53
54// SAFETY: These deferral AIRs must be at these indices within the extension
55pub(crate) const POSEIDON2_AIR_REL_IDX: usize = 1;
56pub(crate) const CALL_AIR_REL_IDX: usize = 2;
57pub(crate) const OUTPUT_AIR_REL_IDX: usize = 3;
58
59// =================================== VM Extension Implementation =================================
60
61#[derive(Clone, Debug, Default, Serialize, Deserialize, derive_new::new)]
62pub struct DeferralExtension {
63    #[serde(skip)]
64    pub fns: Vec<Arc<DeferralFn>>,
65    pub def_circuit_commits: Vec<[u8; COMMIT_NUM_BYTES]>,
66}
67
68#[derive(Clone, From, AnyEnum, Executor, MeteredExecutor, PreflightExecutor)]
69#[cfg_attr(
70    feature = "aot",
71    derive(
72        openvm_circuit_derive::AotExecutor,
73        openvm_circuit_derive::AotMeteredExecutor
74    )
75)]
76pub enum DeferralExecutor {
77    Call(DeferralCallExecutor),
78    Output(DeferralOutputExecutor),
79}
80
81impl<F: VmField> VmExecutionExtension<F> for DeferralExtension {
82    type Executor = DeferralExecutor;
83
84    fn extend_execution(
85        &self,
86        inventory: &mut ExecutorInventoryBuilder<F, DeferralExecutor>,
87    ) -> Result<(), ExecutorInventoryError> {
88        let call = DeferralCallExecutor::new(DeferralCallAdapterExecutor, self.fns.clone());
89        inventory.add_executor(call, [DeferralOpcode::CALL.global_opcode()])?;
90
91        inventory.add_executor(
92            DeferralOutputExecutor::new(),
93            [DeferralOpcode::OUTPUT.global_opcode()],
94        )?;
95
96        Ok(())
97    }
98}
99
100impl<SC: StarkProtocolConfig> VmCircuitExtension<SC> for DeferralExtension
101where
102    Val<SC>: VmField,
103{
104    fn extend_circuit(&self, inventory: &mut AirInventory<SC>) -> Result<(), AirInventoryError> {
105        let memory_bridge = inventory.system().port().memory_bridge;
106        let execution_bridge = ExecutionBridge::new(
107            inventory.system().port().execution_bus,
108            inventory.system().port().program_bus,
109        );
110
111        let count_bus = DeferralCircuitCountBus::new(inventory.new_bus_idx());
112        let poseidon2_bus = DeferralPoseidon2Bus::new(inventory.new_bus_idx());
113        let bitwise_bus = {
114            let existing_air = inventory.find_air::<BitwiseOperationLookupAir<8>>().next();
115            if let Some(air) = existing_air {
116                air.bus
117            } else {
118                let bus = BitwiseOperationLookupBus::new(inventory.new_bus_idx());
119                let air = BitwiseOperationLookupAir::<8>::new(bus);
120                inventory.add_air(air);
121                air.bus
122            }
123        };
124
125        let base_num_airs = inventory.num_airs();
126        let address_bits = inventory.pointer_max_bits();
127
128        inventory.add_air(DeferralCircuitCountAir::new(count_bus, self.fns.len()));
129
130        assert_eq!(inventory.num_airs() - base_num_airs, POSEIDON2_AIR_REL_IDX);
131        inventory.add_air_ref(Arc::new(deferral_poseidon2_air(poseidon2_bus.0)));
132
133        assert_eq!(inventory.num_airs() - base_num_airs, CALL_AIR_REL_IDX);
134        inventory.add_air(DeferralCallAir::new(
135            DeferralCallAdapterAir::new(execution_bridge, memory_bridge, bitwise_bus, address_bits),
136            DeferralCallCoreAir::new(count_bus, poseidon2_bus, bitwise_bus),
137        ));
138
139        assert_eq!(inventory.num_airs() - base_num_airs, OUTPUT_AIR_REL_IDX);
140        inventory.add_air(DeferralOutputAir::new(
141            execution_bridge,
142            memory_bridge,
143            count_bus,
144            poseidon2_bus,
145            bitwise_bus,
146            address_bits,
147        ));
148
149        Ok(())
150    }
151}
152
153pub struct DeferralCpuProverExt;
154
155impl<SC, E, RA> VmProverExtension<E, RA, DeferralExtension> for DeferralCpuProverExt
156where
157    SC: StarkProtocolConfig,
158    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
159    RA: RowMajorMatrixArena<Val<SC>>,
160    Val<SC>: VmField,
161    SC::EF: Ord,
162{
163    fn extend_prover(
164        &self,
165        extension: &DeferralExtension,
166        inventory: &mut ChipInventory<SC, RA, CpuBackend<SC>>,
167    ) -> Result<(), ChipInventoryError> {
168        let range_checker = inventory.range_checker()?.clone();
169        let timestamp_max_bits = inventory.timestamp_max_bits();
170        let address_bits = inventory.airs().pointer_max_bits();
171        let mem_helper = SharedMemoryHelper::new(range_checker.clone(), timestamp_max_bits);
172        let bitwise_lu = {
173            let existing_chip = inventory
174                .find_chip::<SharedBitwiseOperationLookupChip<8>>()
175                .next();
176            if let Some(chip) = existing_chip {
177                chip.clone()
178            } else {
179                let air: &BitwiseOperationLookupAir<8> = inventory.next_air()?;
180                let chip = Arc::new(BitwiseOperationLookupChip::new(air.bus));
181                inventory.add_periphery_chip(chip.clone());
182                chip
183            }
184        };
185        let count_chip = Arc::new(DeferralCircuitCountChip::new(extension.fns.len()));
186        let poseidon2_chip = Arc::new(deferral_poseidon2_chip());
187
188        inventory.next_air::<DeferralCircuitCountAir>()?;
189        inventory.add_periphery_chip(count_chip.clone());
190
191        inventory.next_air::<DeferralPoseidon2Air<Val<SC>>>()?;
192        inventory.add_periphery_chip(poseidon2_chip.clone());
193
194        inventory.next_air::<DeferralCallAir>()?;
195        inventory.add_executor_chip(DeferralCallChip::new(
196            DeferralCallCoreFiller::new(
197                DeferralCallAdapterFiller::new(bitwise_lu.clone(), address_bits),
198                count_chip.clone(),
199                poseidon2_chip.clone(),
200                bitwise_lu.clone(),
201                address_bits,
202            ),
203            mem_helper.clone(),
204        ));
205
206        inventory.next_air::<DeferralOutputAir>()?;
207        inventory.add_executor_chip(DeferralOutputChip::new(
208            DeferralOutputFiller::new(
209                count_chip.clone(),
210                poseidon2_chip.clone(),
211                bitwise_lu,
212                address_bits,
213            ),
214            mem_helper,
215        ));
216
217        Ok(())
218    }
219}
220
221// =================================== VM Rv32 Config and Builder =================================
222
223#[derive(Clone, VmConfig, Serialize, Deserialize)]
224pub struct Rv32DeferralConfig {
225    #[config(executor = "SystemExecutor<F>")]
226    pub system: SystemConfig,
227    #[extension]
228    pub rv32i: Rv32I,
229    #[extension]
230    pub rv32m: Rv32M,
231    #[extension]
232    pub io: Rv32Io,
233    #[serde(skip)]
234    #[extension(executor = "DeferralExecutor")]
235    pub deferral: DeferralExtension,
236}
237
238impl InitFileGenerator for Rv32DeferralConfig {}
239
240#[derive(Clone)]
241pub struct Rv32DeferralCpuBuilder;
242
243impl<SC, E> VmBuilder<E> for Rv32DeferralCpuBuilder
244where
245    SC: StarkProtocolConfig,
246    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
247    Val<SC>: VmField,
248    SC::EF: Ord,
249{
250    type VmConfig = Rv32DeferralConfig;
251    type SystemChipInventory = SystemChipInventory<SC>;
252    type RecordArena = MatrixRecordArena<Val<SC>>;
253
254    fn create_chip_complex(
255        &self,
256        config: &Self::VmConfig,
257        circuit: AirInventory<SC>,
258        device_ctx: &openvm_stark_backend::EngineDeviceCtx<E>,
259    ) -> Result<
260        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
261        ChipInventoryError,
262    > {
263        let mut chip_complex = VmBuilder::<E>::create_chip_complex(
264            &SystemCpuBuilder,
265            &config.system,
266            circuit,
267            device_ctx,
268        )?;
269        let inventory = &mut chip_complex.inventory;
270        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32i, inventory)?;
271        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.rv32m, inventory)?;
272        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.io, inventory)?;
273        VmProverExtension::<E, _, _>::extend_prover(
274            &DeferralCpuProverExt,
275            &config.deferral,
276            inventory,
277        )?;
278        Ok(chip_complex)
279    }
280}
openvm_deferral_circuit/extension/mod.rs

openvm_deferral_circuit/extension/
mod.rs
