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

openvm_algebra_circuit/extension/
mod.rs
