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

openvm_algebra_circuit/extension/
mod.rs
