feat: fptower and tests (#1239)

Co-authored-by: Tamir Hemo <[email protected]>

core/Cargo.toml

@@ -74,7 +74,7 @@ tiny-keccak = { version = "2.0.2", features = ["keccak"] }
 criterion = "0.5.1"
 num = { version = "0.4.3", features = ["rand"] }
 rand = "0.8.5"
-sp1-zkvm = { workspace = true }
+sp1-zkvm = { workspace = true, features = ["lib"] }
 
 [features]
 neon = ["p3-blake3/neon"]

core/src/operations/field/field_op.rs

@@ -1,8 +1,12 @@
 use std::fmt::Debug;
 
 use num::{BigUint, Zero};
+
+use serde::{Deserialize, Serialize};
+
 use p3_air::AirBuilder;
 use p3_field::PrimeField32;
+
 use sp1_derive::AlignedBorrow;
 
 use super::params::{FieldParameters, Limbs};
@@ -14,7 +18,7 @@ use crate::bytes::event::ByteRecord;
 use typenum::Unsigned;
 
 /// Airthmetic operation for emulating modular arithmetic.
-#[derive(PartialEq, Copy, Clone, Debug)]
+#[derive(PartialEq, Copy, Clone, Debug, Serialize, Deserialize)]
 pub enum FieldOperation {
     Add,
     Mul,
@@ -184,6 +188,54 @@ impl<F: PrimeField32, P: FieldParameters> FieldOpCols<F, P> {
 }
 
 impl<V: Copy, P: FieldParameters> FieldOpCols<V, P> {
+    /// Allows an evaluation over opetations specified by boolean flags.
+    #[allow(clippy::too_many_arguments)]
+    pub fn eval_variable<AB: SP1AirBuilder<Var = V>>(
+        &self,
+        builder: &mut AB,
+        a: &(impl Into<Polynomial<AB::Expr>> + Clone),
+        b: &(impl Into<Polynomial<AB::Expr>> + Clone),
+        modulus: &(impl Into<Polynomial<AB::Expr>> + Clone),
+        is_add: impl Into<AB::Expr> + Clone,
+        is_sub: impl Into<AB::Expr> + Clone,
+        is_mul: impl Into<AB::Expr> + Clone,
+        is_div: impl Into<AB::Expr> + Clone,
+        shard: impl Into<AB::Expr> + Clone,
+        channel: impl Into<AB::Expr> + Clone,
+        is_real: impl Into<AB::Expr> + Clone,
+    ) where
+        V: Into<AB::Expr>,
+        Limbs<V, P::Limbs>: Copy,
+    {
+        let p_a_param: Polynomial<AB::Expr> = (a).clone().into();
+        let p_b: Polynomial<AB::Expr> = (b).clone().into();
+        let p_res_param: Polynomial<AB::Expr> = self.result.into();
+
+        let is_add: AB::Expr = is_add.into();
+        let is_sub: AB::Expr = is_sub.into();
+        let is_mul: AB::Expr = is_mul.into();
+        let is_div: AB::Expr = is_div.into();
+
+        let p_result = p_res_param.clone() * (is_add.clone() + is_mul.clone())
+            + p_a_param.clone() * (is_sub.clone() + is_div.clone());
+
+        let p_add = p_a_param.clone() + p_b.clone();
+        let p_sub = p_res_param.clone() + p_b.clone();
+        let p_mul = p_a_param.clone() * p_b.clone();
+        let p_div = p_res_param * p_b.clone();
+        let p_op = p_add * is_add + p_sub * is_sub + p_mul * is_mul + p_div * is_div;
+
+        self.eval_with_polynomials(
+            builder,
+            p_op,
+            modulus.clone(),
+            p_result,
+            shard,
+            channel,
+            is_real,
+        );
+    }
+
     #[allow(clippy::too_many_arguments)]
     pub fn eval_with_modulus<AB: SP1AirBuilder<Var = V>>(
         &self,
@@ -201,17 +253,44 @@ impl<V: Copy, P: FieldParameters> FieldOpCols<V, P> {
     {
         let p_a_param: Polynomial<AB::Expr> = (a).clone().into();
         let p_b: Polynomial<AB::Expr> = (b).clone().into();
-        let p_modulus: Polynomial<AB::Expr> = (modulus).clone().into();
 
         let (p_a, p_result): (Polynomial<_>, Polynomial<_>) = match op {
             FieldOperation::Add | FieldOperation::Mul => (p_a_param, self.result.into()),
             FieldOperation::Sub | FieldOperation::Div => (self.result.into(), p_a_param),
         };
-        let p_carry: Polynomial<<AB as AirBuilder>::Expr> = self.carry.into();
-        let p_op = match op {
+        let p_op: Polynomial<<AB as AirBuilder>::Expr> = match op {
             FieldOperation::Add | FieldOperation::Sub => p_a + p_b,
             FieldOperation::Mul | FieldOperation::Div => p_a * p_b,
         };
+        self.eval_with_polynomials(
+            builder,
+            p_op,
+            modulus.clone(),
+            p_result,
+            shard,
+            channel,
+            is_real,
+        );
+    }
+
+    #[allow(clippy::too_many_arguments)]
+    pub fn eval_with_polynomials<AB: SP1AirBuilder<Var = V>>(
+        &self,
+        builder: &mut AB,
+        op: impl Into<Polynomial<AB::Expr>>,
+        modulus: impl Into<Polynomial<AB::Expr>>,
+        result: impl Into<Polynomial<AB::Expr>>,
+        shard: impl Into<AB::Expr> + Clone,
+        channel: impl Into<AB::Expr> + Clone,
+        is_real: impl Into<AB::Expr> + Clone,
+    ) where
+        V: Into<AB::Expr>,
+        Limbs<V, P::Limbs>: Copy,
+    {
+        let p_op: Polynomial<AB::Expr> = op.into();
+        let p_result: Polynomial<AB::Expr> = result.into();
+        let p_modulus: Polynomial<AB::Expr> = modulus.into();
+        let p_carry: Polynomial<<AB as AirBuilder>::Expr> = self.carry.into();
         let p_op_minus_result: Polynomial<AB::Expr> = p_op - &p_result;
         let p_vanishing = p_op_minus_result - &(&p_carry * &p_modulus);
         let p_witness_low = self.witness_low.0.iter().into();

core/src/runtime/mod.rs

@@ -882,8 +882,13 @@ impl<'a> Runtime<'a> {
                 exit_code = returned_exit_code;
 
                 // Update the syscall counts.
-                let syscall_count = self.state.syscall_counts.entry(syscall).or_insert(0);
-                let (threshold, multiplier) = match syscall {
+                let syscall_for_count = syscall.count_map();
+                let syscall_count = self
+                    .state
+                    .syscall_counts
+                    .entry(syscall_for_count)
+                    .or_insert(0);
+                let (threshold, multiplier) = match syscall_for_count {
                     SyscallCode::KECCAK_PERMUTE => {
                         (self.opts.split_opts.keccak_split_threshold, 24)
                     }

core/src/runtime/record.rs

@@ -18,6 +18,7 @@ use crate::runtime::MemoryInitializeFinalizeEvent;
 use crate::runtime::MemoryRecordEnum;
 use crate::stark::MachineRecord;
 use crate::syscall::precompiles::edwards::EdDecompressEvent;
+use crate::syscall::precompiles::fptower::{Fp2AddSubEvent, Fp2MulEvent, FpOpEvent};
 use crate::syscall::precompiles::keccak256::KeccakPermuteEvent;
 use crate::syscall::precompiles::sha256::{ShaCompressEvent, ShaExtendEvent};
 use crate::syscall::precompiles::uint256::Uint256MulEvent;
@@ -98,6 +99,18 @@ pub struct ExecutionRecord {
 
     pub bls12381_decompress_events: Vec<ECDecompressEvent>,
 
+    pub bls12381_fp_events: Vec<FpOpEvent>,
+
+    pub bls12381_fp2_addsub_events: Vec<Fp2AddSubEvent>,
+
+    pub bls12381_fp2_mul_events: Vec<Fp2MulEvent>,
+
+    pub bn254_fp_events: Vec<FpOpEvent>,
+
+    pub bn254_fp2_addsub_events: Vec<Fp2AddSubEvent>,
+
+    pub bn254_fp2_mul_events: Vec<Fp2MulEvent>,
+
     /// The public values.
     pub public_values: PublicValues<u32, u32>,
 
@@ -191,6 +204,27 @@ impl MachineRecord for ExecutionRecord {
             "uint256_mul_events".to_string(),
             self.uint256_mul_events.len(),
         );
+        stats.insert(
+            "bls12381_fp_event".to_string(),
+            self.bls12381_fp_events.len(),
+        );
+        stats.insert(
+            "bls12381_fp2_addsub_events".to_string(),
+            self.bls12381_fp2_addsub_events.len(),
+        );
+        stats.insert(
+            "bls12381_fp2_mul_events".to_string(),
+            self.bls12381_fp2_mul_events.len(),
+        );
+        stats.insert("bn254_fp_events".to_string(), self.bn254_fp_events.len());
+        stats.insert(
+            "bn254_fp2_addsub_events".to_string(),
+            self.bn254_fp2_addsub_events.len(),
+        );
+        stats.insert(
+            "bn254_fp2_mul_events".to_string(),
+            self.bn254_fp2_mul_events.len(),
+        );
         stats.insert(
             "bls12381_decompress_events".to_string(),
             self.bls12381_decompress_events.len(),
@@ -249,9 +283,19 @@ impl MachineRecord for ExecutionRecord {
             .append(&mut other.bls12381_double_events);
         self.uint256_mul_events
             .append(&mut other.uint256_mul_events);
+        self.bls12381_fp_events
+            .append(&mut other.bls12381_fp_events);
+        self.bls12381_fp2_addsub_events
+            .append(&mut other.bls12381_fp2_addsub_events);
+        self.bls12381_fp2_mul_events
+            .append(&mut other.bls12381_fp2_mul_events);
+        self.bn254_fp_events.append(&mut other.bn254_fp_events);
+        self.bn254_fp2_addsub_events
+            .append(&mut other.bn254_fp2_addsub_events);
+        self.bn254_fp2_mul_events
+            .append(&mut other.bn254_fp2_mul_events);
         self.bls12381_decompress_events
             .append(&mut other.bls12381_decompress_events);
-
         if self.byte_lookups.is_empty() {
             self.byte_lookups = std::mem::take(&mut other.byte_lookups);
         } else {
@@ -383,6 +427,8 @@ impl ExecutionRecord {
             ed_decompress_events: std::mem::take(&mut self.ed_decompress_events),
             k256_decompress_events: std::mem::take(&mut self.k256_decompress_events),
             uint256_mul_events: std::mem::take(&mut self.uint256_mul_events),
+            bls12381_fp_events: std::mem::take(&mut self.bls12381_fp_events),
+            bls12381_fp2_mul_events: std::mem::take(&mut self.bls12381_fp2_mul_events),
             bls12381_decompress_events: std::mem::take(&mut self.bls12381_decompress_events),
             memory_initialize_events: std::mem::take(&mut self.memory_initialize_events),
             memory_finalize_events: std::mem::take(&mut self.memory_finalize_events),
@@ -520,6 +566,48 @@ impl ExecutionRecord {
             opts.deferred_shift_threshold,
             last
         );
+        split_events!(
+            self,
+            bls12381_fp_events,
+            shards,
+            opts.deferred_shift_threshold,
+            last
+        );
+        split_events!(
+            self,
+            bls12381_fp2_addsub_events,
+            shards,
+            opts.deferred_shift_threshold,
+            last
+        );
+        split_events!(
+            self,
+            bls12381_fp2_mul_events,
+            shards,
+            opts.deferred_shift_threshold,
+            last
+        );
+        split_events!(
+            self,
+            bn254_fp_events,
+            shards,
+            opts.deferred_shift_threshold,
+            last
+        );
+        split_events!(
+            self,
+            bn254_fp2_addsub_events,
+            shards,
+            opts.deferred_shift_threshold,
+            last
+        );
+        split_events!(
+            self,
+            bn254_fp2_mul_events,
+            shards,
+            opts.deferred_shift_threshold,
+            last
+        );
         // _ = last_pct;
 
         if last {