openvm_algebra_circuit/
fp2_extension.rs

1use derive_more::derive::From;
2use num_bigint::BigUint;
3use openvm_algebra_transpiler::Fp2Opcode;
4use openvm_circuit::{
5    arch::{SystemPort, VmExtension, VmInventory, VmInventoryBuilder, VmInventoryError},
6    system::phantom::PhantomChip,
7};
8use openvm_circuit_derive::{AnyEnum, InstructionExecutor};
9use openvm_circuit_primitives::bitwise_op_lookup::{
10    BitwiseOperationLookupBus, SharedBitwiseOperationLookupChip,
11};
12use openvm_circuit_primitives_derive::{Chip, ChipUsageGetter};
13use openvm_instructions::{LocalOpcode, VmOpcode};
14use openvm_mod_circuit_builder::ExprBuilderConfig;
15use openvm_rv32_adapters::Rv32VecHeapAdapterChip;
16use openvm_stark_backend::p3_field::PrimeField32;
17use serde::{Deserialize, Serialize};
18use serde_with::{serde_as, DisplayFromStr};
19use strum::EnumCount;
20
21use crate::fp2_chip::{Fp2AddSubChip, Fp2MulDivChip};
22
23#[serde_as]
24#[derive(Clone, Debug, derive_new::new, Serialize, Deserialize)]
25pub struct Fp2Extension {
26    #[serde_as(as = "Vec<DisplayFromStr>")]
27    pub supported_modulus: Vec<BigUint>,
28}
29
30#[derive(ChipUsageGetter, Chip, InstructionExecutor, AnyEnum, From)]
31pub enum Fp2ExtensionExecutor<F: PrimeField32> {
32    // 32 limbs prime
33    Fp2AddSubRv32_32(Fp2AddSubChip<F, 2, 32>),
34    Fp2MulDivRv32_32(Fp2MulDivChip<F, 2, 32>),
35    // 48 limbs prime
36    Fp2AddSubRv32_48(Fp2AddSubChip<F, 6, 16>),
37    Fp2MulDivRv32_48(Fp2MulDivChip<F, 6, 16>),
38}
39
40#[derive(ChipUsageGetter, Chip, AnyEnum, From)]
41pub enum Fp2ExtensionPeriphery<F: PrimeField32> {
42    BitwiseOperationLookup(SharedBitwiseOperationLookupChip<8>),
43    // We put this only to get the <F> generic to work
44    Phantom(PhantomChip<F>),
45}
46
47impl<F: PrimeField32> VmExtension<F> for Fp2Extension {
48    type Executor = Fp2ExtensionExecutor<F>;
49    type Periphery = Fp2ExtensionPeriphery<F>;
50
51    fn build(
52        &self,
53        builder: &mut VmInventoryBuilder<F>,
54    ) -> Result<VmInventory<Self::Executor, Self::Periphery>, VmInventoryError> {
55        let mut inventory = VmInventory::new();
56        let SystemPort {
57            execution_bus,
58            program_bus,
59            memory_bridge,
60        } = builder.system_port();
61        let bitwise_lu_chip = if let Some(&chip) = builder
62            .find_chip::<SharedBitwiseOperationLookupChip<8>>()
63            .first()
64        {
65            chip.clone()
66        } else {
67            let bitwise_lu_bus = BitwiseOperationLookupBus::new(builder.new_bus_idx());
68            let chip = SharedBitwiseOperationLookupChip::new(bitwise_lu_bus);
69            inventory.add_periphery_chip(chip.clone());
70            chip
71        };
72        let offline_memory = builder.system_base().offline_memory();
73        let range_checker = builder.system_base().range_checker_chip.clone();
74        let address_bits = builder.system_config().memory_config.pointer_max_bits;
75
76        let addsub_opcodes = (Fp2Opcode::ADD as usize)..=(Fp2Opcode::SETUP_ADDSUB as usize);
77        let muldiv_opcodes = (Fp2Opcode::MUL as usize)..=(Fp2Opcode::SETUP_MULDIV as usize);
78
79        for (i, modulus) in self.supported_modulus.iter().enumerate() {
80            // determine the number of bytes needed to represent a prime field element
81            let bytes = modulus.bits().div_ceil(8);
82            let start_offset = Fp2Opcode::CLASS_OFFSET + i * Fp2Opcode::COUNT;
83
84            let config32 = ExprBuilderConfig {
85                modulus: modulus.clone(),
86                num_limbs: 32,
87                limb_bits: 8,
88            };
89            let config48 = ExprBuilderConfig {
90                modulus: modulus.clone(),
91                num_limbs: 48,
92                limb_bits: 8,
93            };
94            let adapter_chip_32 = Rv32VecHeapAdapterChip::new(
95                execution_bus,
96                program_bus,
97                memory_bridge,
98                address_bits,
99                bitwise_lu_chip.clone(),
100            );
101            let adapter_chip_48 = Rv32VecHeapAdapterChip::new(
102                execution_bus,
103                program_bus,
104                memory_bridge,
105                address_bits,
106                bitwise_lu_chip.clone(),
107            );
108
109            if bytes <= 32 {
110                let addsub_chip = Fp2AddSubChip::new(
111                    adapter_chip_32.clone(),
112                    config32.clone(),
113                    start_offset,
114                    range_checker.clone(),
115                    offline_memory.clone(),
116                );
117                inventory.add_executor(
118                    Fp2ExtensionExecutor::Fp2AddSubRv32_32(addsub_chip),
119                    addsub_opcodes
120                        .clone()
121                        .map(|x| VmOpcode::from_usize(x + start_offset)),
122                )?;
123                let muldiv_chip = Fp2MulDivChip::new(
124                    adapter_chip_32.clone(),
125                    config32.clone(),
126                    start_offset,
127                    range_checker.clone(),
128                    offline_memory.clone(),
129                );
130                inventory.add_executor(
131                    Fp2ExtensionExecutor::Fp2MulDivRv32_32(muldiv_chip),
132                    muldiv_opcodes
133                        .clone()
134                        .map(|x| VmOpcode::from_usize(x + start_offset)),
135                )?;
136            } else if bytes <= 48 {
137                let addsub_chip = Fp2AddSubChip::new(
138                    adapter_chip_48.clone(),
139                    config48.clone(),
140                    start_offset,
141                    range_checker.clone(),
142                    offline_memory.clone(),
143                );
144                inventory.add_executor(
145                    Fp2ExtensionExecutor::Fp2AddSubRv32_48(addsub_chip),
146                    addsub_opcodes
147                        .clone()
148                        .map(|x| VmOpcode::from_usize(x + start_offset)),
149                )?;
150                let muldiv_chip = Fp2MulDivChip::new(
151                    adapter_chip_48.clone(),
152                    config48.clone(),
153                    start_offset,
154                    range_checker.clone(),
155                    offline_memory.clone(),
156                );
157                inventory.add_executor(
158                    Fp2ExtensionExecutor::Fp2MulDivRv32_48(muldiv_chip),
159                    muldiv_opcodes
160                        .clone()
161                        .map(|x| VmOpcode::from_usize(x + start_offset)),
162                )?;
163            } else {
164                panic!("Modulus too large");
165            }
166        }
167
168        Ok(inventory)
169    }
170}