openvm_algebra_circuit/extension/
mod.rs

1use std::result::Result;
2
3use num_bigint::BigUint;
4use openvm_circuit::{
5    arch::{
6        AirInventory, ChipInventoryError, InitFileGenerator, MatrixRecordArena, SystemConfig,
7        VmBuilder, VmChipComplex, VmField, VmProverExtension,
8    },
9    system::{SystemChipInventory, SystemCpuBuilder, SystemExecutor},
10};
11use openvm_circuit_derive::VmConfig;
12use openvm_rv32im_circuit::{
13    Rv32I, Rv32IExecutor, Rv32ImCpuProverExt, Rv32Io, Rv32IoExecutor, Rv32M, Rv32MExecutor,
14};
15use openvm_stark_backend::{
16    config::{StarkGenericConfig, Val},
17    prover::cpu::{CpuBackend, CpuDevice},
18};
19use openvm_stark_sdk::engine::StarkEngine;
20use serde::{Deserialize, Serialize};
21
22mod modular;
23pub use modular::*;
24mod fp2;
25pub use fp2::*;
26
27cfg_if::cfg_if! {
28    if #[cfg(feature = "cuda")] {
29        mod cuda;
30        mod hybrid;
31        pub use cuda::*;
32        pub use hybrid::*;
33        pub use {
34            AlgebraHybridProverExt as AlgebraProverExt,
35            Rv32ModularHybridBuilder as Rv32ModularBuilder,
36            Rv32ModularWithFp2HybridBuilder as Rv32ModularWithFp2Builder,
37        };
38    } else {
39        pub use self::{
40            AlgebraCpuProverExt as AlgebraProverExt,
41            Rv32ModularCpuBuilder as Rv32ModularBuilder,
42            Rv32ModularWithFp2CpuBuilder as Rv32ModularWithFp2Builder,
43        };
44    }
45}
46
47pub struct AlgebraCpuProverExt;
48
49#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)]
50pub struct Rv32ModularConfig {
51    #[config(executor = "SystemExecutor<F>")]
52    pub system: SystemConfig,
53    #[extension]
54    pub base: Rv32I,
55    #[extension]
56    pub mul: Rv32M,
57    #[extension]
58    pub io: Rv32Io,
59    #[extension]
60    pub modular: ModularExtension,
61}
62
63impl InitFileGenerator for Rv32ModularConfig {
64    fn generate_init_file_contents(&self) -> Option<String> {
65        Some(format!(
66            "// This file is automatically generated by cargo openvm. Do not rename or edit.\n{}\n",
67            self.modular.generate_moduli_init()
68        ))
69    }
70}
71
72impl Rv32ModularConfig {
73    pub fn new(moduli: Vec<BigUint>) -> Self {
74        Self {
75            system: SystemConfig::default(),
76            base: Default::default(),
77            mul: Default::default(),
78            io: Default::default(),
79            modular: ModularExtension::new(moduli),
80        }
81    }
82}
83
84#[derive(Clone, Debug, VmConfig, Serialize, Deserialize)]
85pub struct Rv32ModularWithFp2Config {
86    #[config(generics = true)]
87    pub modular: Rv32ModularConfig,
88    #[extension]
89    pub fp2: Fp2Extension,
90}
91
92impl Rv32ModularWithFp2Config {
93    pub fn new(moduli_with_names: Vec<(String, BigUint)>) -> Self {
94        let moduli = moduli_with_names
95            .iter()
96            .map(|(_, modulus)| modulus.clone())
97            .collect();
98        Self {
99            modular: Rv32ModularConfig::new(moduli),
100            fp2: Fp2Extension::new(moduli_with_names),
101        }
102    }
103}
104
105impl InitFileGenerator for Rv32ModularWithFp2Config {
106    fn generate_init_file_contents(&self) -> Option<String> {
107        Some(format!(
108            "// This file is automatically generated by cargo openvm. Do not rename or edit.\n{}\n{}\n",
109            self.modular.modular.generate_moduli_init(),
110            self.fp2.generate_complex_init(&self.modular.modular)
111        ))
112    }
113}
114
115#[derive(Clone)]
116pub struct Rv32ModularCpuBuilder;
117
118impl<E, SC> VmBuilder<E> for Rv32ModularCpuBuilder
119where
120    SC: StarkGenericConfig,
121    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
122    Val<SC>: VmField,
123{
124    type VmConfig = Rv32ModularConfig;
125    type SystemChipInventory = SystemChipInventory<SC>;
126    type RecordArena = MatrixRecordArena<Val<SC>>;
127
128    fn create_chip_complex(
129        &self,
130        config: &Rv32ModularConfig,
131        circuit: AirInventory<SC>,
132    ) -> Result<
133        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
134        ChipInventoryError,
135    > {
136        let mut chip_complex =
137            VmBuilder::<E>::create_chip_complex(&SystemCpuBuilder, &config.system, circuit)?;
138        let inventory = &mut chip_complex.inventory;
139        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.base, inventory)?;
140        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.mul, inventory)?;
141        VmProverExtension::<E, _, _>::extend_prover(&Rv32ImCpuProverExt, &config.io, inventory)?;
142        VmProverExtension::<E, _, _>::extend_prover(
143            &AlgebraCpuProverExt,
144            &config.modular,
145            inventory,
146        )?;
147        Ok(chip_complex)
148    }
149}
150
151#[derive(Clone)]
152pub struct Rv32ModularWithFp2CpuBuilder;
153
154impl<E, SC> VmBuilder<E> for Rv32ModularWithFp2CpuBuilder
155where
156    SC: StarkGenericConfig,
157    E: StarkEngine<SC = SC, PB = CpuBackend<SC>, PD = CpuDevice<SC>>,
158    Val<SC>: VmField,
159{
160    type VmConfig = Rv32ModularWithFp2Config;
161    type SystemChipInventory = SystemChipInventory<SC>;
162    type RecordArena = MatrixRecordArena<Val<SC>>;
163
164    fn create_chip_complex(
165        &self,
166        config: &Rv32ModularWithFp2Config,
167        circuit: AirInventory<SC>,
168    ) -> Result<
169        VmChipComplex<SC, Self::RecordArena, E::PB, Self::SystemChipInventory>,
170        ChipInventoryError,
171    > {
172        let mut chip_complex =
173            VmBuilder::<E>::create_chip_complex(&Rv32ModularCpuBuilder, &config.modular, circuit)?;
174        let inventory = &mut chip_complex.inventory;
175        VmProverExtension::<E, _, _>::extend_prover(&AlgebraCpuProverExt, &config.fp2, inventory)?;
176        Ok(chip_complex)
177    }
178}
openvm_algebra_circuit/extension/mod.rs

openvm_algebra_circuit/extension/
mod.rs
