Merge pull request #929 from powdr-labs/ff_arith

Finite field operations

pipeline/Cargo.toml

@@ -32,6 +32,8 @@ log = "0.4.17"
 mktemp = "0.5.0"
 serde = { version = "1.0", default-features = false, features = ["alloc", "derive", "rc"] }
 serde_cbor = "0.11.2"
+num-bigint = "0.4.3"
+num-traits = "0.2.15"
 
 [dev-dependencies]
 powdr-riscv = { path = "../riscv" }

pipeline/benches/executor_benchmark.rs

@@ -1,14 +1,16 @@
-use ::powdr_pipeline::inputs_to_query_callback;
-use ::powdr_pipeline::Pipeline;
-use criterion::{criterion_group, criterion_main, Criterion};
+use ::powdr_pipeline::{inputs_to_query_callback, Pipeline};
 use powdr_ast::analyzed::Analyzed;
-
-use mktemp::Temp;
 use powdr_number::{FieldElement, GoldilocksField};
-use powdr_riscv::continuations::bootloader::default_input;
-use powdr_riscv::{compile_rust_crate_to_riscv_asm, compile_rust_to_riscv_asm, compiler};
 
-use powdr_riscv::CoProcessors;
+use powdr_pipeline::test_util::{evaluate_integer_function, std_analyzed};
+use powdr_riscv::{
+    compile_rust_crate_to_riscv_asm, compile_rust_to_riscv_asm, compiler,
+    continuations::bootloader::default_input, CoProcessors,
+};
+
+use criterion::{criterion_group, criterion_main, Criterion};
+use mktemp::Temp;
+use num_traits::Num;
 
 type T = GoldilocksField;
 
@@ -23,7 +25,7 @@ fn run_witgen<T: FieldElement>(
         .generate();
 }
 
-fn criterion_benchmark(c: &mut Criterion) {
+fn executor_benchmark(c: &mut Criterion) {
     let mut group = c.benchmark_group("executor-benchmark");
     group.sample_size(10);
 
@@ -71,5 +73,65 @@ fn criterion_benchmark(c: &mut Criterion) {
     group.finish();
 }
 
-criterion_group!(benches, criterion_benchmark);
+fn evaluator_benchmark(c: &mut Criterion) {
+    let mut group = c.benchmark_group("evaluator-benchmark");
+
+    let analyzed = std_analyzed::<GoldilocksField>();
+
+    group.bench_function("std::math::ff::inverse", |b| {
+        b.iter(|| {
+            let modulus = num_bigint::BigInt::from_str_radix(
+                "fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f",
+                16,
+            )
+            .unwrap();
+            let x = modulus.clone() - num_bigint::BigInt::from(17);
+
+            evaluate_integer_function(
+                &analyzed,
+                "std::math::ff::inverse",
+                vec![x.clone(), modulus.clone()],
+            );
+        })
+    });
+
+    group.bench_function("std::math::ff::reduce", |b| {
+        b.iter(|| {
+            let modulus = num_bigint::BigInt::from_str_radix(
+                "fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f",
+                16,
+            )
+            .unwrap();
+            let x = modulus.clone() + num_bigint::BigInt::from(17);
+
+            evaluate_integer_function(
+                &analyzed,
+                "std::math::ff::reduce",
+                vec![x.clone(), modulus.clone()],
+            );
+        })
+    });
+
+    group.bench_function("std::math::ff::mul", |b| {
+        b.iter(|| {
+            let modulus = num_bigint::BigInt::from_str_radix(
+                "fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f",
+                16,
+            )
+            .unwrap();
+            let x = modulus.clone() - num_bigint::BigInt::from(17);
+            let y = modulus.clone() - num_bigint::BigInt::from(11);
+
+            evaluate_integer_function(
+                &analyzed,
+                "std::math::ff::mul",
+                vec![x.clone(), y.clone(), modulus.clone()],
+            );
+        })
+    });
+
+    group.finish();
+}
+
+criterion_group!(benches, evaluator_benchmark, executor_benchmark);
 criterion_main!(benches);

pipeline/src/pipeline.rs

@@ -129,6 +129,13 @@ impl<T: FieldElement> Artifact<T> {
         }
     }
 
+    pub fn to_analyzed_pil(&self) -> Option<&Analyzed<T>> {
+        match self {
+            Artifact::AnalyzedPil(analyzed) => Some(analyzed),
+            _ => None,
+        }
+    }
+
     pub fn to_optimized_pil(&self) -> Option<&Analyzed<T>> {
         match self {
             Artifact::OptimzedPil(optimized_pil) => Some(optimized_pil),
@@ -745,6 +752,14 @@ impl<T: FieldElement> Pipeline<T> {
         }
     }
 
+    pub fn analyzed_pil(mut self) -> Result<Analyzed<T>, Vec<String>> {
+        self.advance_to(Stage::AnalyzedPil)?;
+        let Artifact::AnalyzedPil(analyzed) = self.artifact.unwrap() else {
+            panic!()
+        };
+        Ok(analyzed)
+    }
+
     pub fn optimized_pil(mut self) -> Result<Analyzed<T>, Vec<String>> {
         self.advance_to(Stage::OptimizedPil)?;
         let Artifact::OptimzedPil(optimized_pil) = self.artifact.unwrap() else {

pipeline/src/test_util.rs

@@ -1,7 +1,9 @@
+use powdr_ast::analyzed::Analyzed;
 use powdr_backend::BackendType;
-use powdr_number::FieldElement;
-use powdr_number::{Bn254Field, GoldilocksField};
+use powdr_number::{Bn254Field, FieldElement, GoldilocksField};
+use powdr_pil_analyzer::evaluator::{self, SymbolLookup};
 use std::path::PathBuf;
+use std::rc::Rc;
 
 use crate::pipeline::{Pipeline, Stage};
 use crate::verify::verify;
@@ -80,3 +82,40 @@ pub fn gen_halo2_proof(file_name: &str, inputs: Vec<Bn254Field>) {
 
 #[cfg(not(feature = "halo2"))]
 pub fn gen_halo2_proof(_file_name: &str, _inputs: Vec<Bn254Field>) {}
+
+/// Returns the analyzed PIL containing only the std library.
+pub fn std_analyzed<T: FieldElement>() -> Analyzed<T> {
+    // airgen needs a main machine.
+    let code = "machine Main { }".to_string();
+    let mut pipeline = Pipeline::default().from_asm_string(code, None);
+    pipeline.advance_to(Stage::AnalyzedPil).unwrap();
+    pipeline.analyzed_pil().unwrap()
+}
+
+/// Evaluates a function call.
+pub fn evaluate_function<'a, T: FieldElement>(
+    analyzed: &'a Analyzed<T>,
+    function: &'a str,
+    arguments: Vec<Rc<evaluator::Value<'a, T, evaluator::NoCustom>>>,
+) -> evaluator::Value<'a, T, evaluator::NoCustom> {
+    let symbols = evaluator::Definitions(&analyzed.definitions);
+    let function = symbols.lookup(function).unwrap();
+    evaluator::evaluate_function_call(function, arguments, &symbols).unwrap()
+}
+
+/// Evaluates a function call assuming inputs and outputs are integers.
+pub fn evaluate_integer_function<T: FieldElement>(
+    analyzed: &Analyzed<T>,
+    function: &str,
+    arguments: Vec<num_bigint::BigInt>,
+) -> num_bigint::BigInt {
+    let arguments = arguments
+        .into_iter()
+        .map(|x| Rc::new(evaluator::Value::Integer(x)))
+        .collect();
+    if let evaluator::Value::Integer(x) = evaluate_function(analyzed, function, arguments) {
+        x
+    } else {
+        panic!("Expected integer.");
+    }
+}

pipeline/tests/powdr_std.rs

@@ -1,7 +1,12 @@
-use powdr_number::GoldilocksField;
-use powdr_pipeline::test_util::{gen_estark_proof, gen_halo2_proof, verify_test_file};
+use powdr_number::{FieldElement, GoldilocksField};
+
+use powdr_pipeline::test_util::{
+    evaluate_integer_function, gen_estark_proof, gen_halo2_proof, std_analyzed, verify_test_file,
+};
 use test_log::test;
 
+use num_traits::Num;
+
 #[test]
 fn poseidon_bn254_test() {
     let f = "std/poseidon_bn254_test.asm";
@@ -36,3 +41,126 @@ fn arith_test() {
     // Halo2 test runs out of memory on CI
     // gen_halo2_proof(f, Default::default());
 }
+
+#[test]
+fn ff_reduce_mod_7() {
+    let test_inputs = vec![
+        -22, -21, -20, -8, -7, -6, -2, -1, -0, -1, -2, -3, 4, 5, 6, 7, 8, 9, 13, 14, 15, 20, 21, 22,
+    ];
+    let analyzed = std_analyzed::<GoldilocksField>();
+    for x in test_inputs {
+        let x = num_bigint::BigInt::from(x);
+        let modulus = num_bigint::BigInt::from(7);
+        let result = evaluate_integer_function(
+            &analyzed,
+            "std::math::ff::reduce",
+            vec![x.clone(), modulus.clone()],
+        );
+        assert!(num_bigint::BigInt::from(0) <= result && result < modulus);
+        if x < result {
+            assert_eq!((result - x) % modulus, 0.into());
+        } else {
+            assert_eq!((x - result) % modulus, 0.into());
+        }
+    }
+}
+
+#[test]
+fn ff_inverse() {
+    let test_inputs = vec![
+        (1, 11),
+        (1, 7),
+        (2, 7),
+        (3, 7),
+        (4, 7),
+        (5, 7),
+        (6, 7),
+        (2, 17),
+        (3, 17),
+        (9, 17),
+        (15, 17),
+        (16, 17),
+    ];
+    let analyzed = std_analyzed::<GoldilocksField>();
+    for (x, modulus) in test_inputs {
+        let x = num_bigint::BigInt::from(x);
+        let modulus = num_bigint::BigInt::from(modulus);
+        let result = evaluate_integer_function(
+            &analyzed,
+            "std::math::ff::inverse",
+            vec![x.clone(), modulus.clone()],
+        );
+        assert_eq!((result * x) % modulus, 1.into());
+    }
+}
+
+#[test]
+fn ff_add_sub_mul_div() {
+    let inputs = vec![
+        (1, 0, 11),
+        (1, 6, 7),
+        (6, 6, 7),
+        (0, 0, 17),
+        (0, 16, 17),
+        (16, 16, 17),
+        (3, 8, 17),
+        (16, 1, 17),
+        (5, 9, 17),
+        (3, 14, 17),
+    ];
+    let analyzed = std_analyzed::<GoldilocksField>();
+    for (x, y, modulus) in inputs {
+        let x = num_bigint::BigInt::from(x);
+        let y = num_bigint::BigInt::from(y);
+        let modulus = num_bigint::BigInt::from(modulus);
+        let result = evaluate_integer_function(
+            &analyzed,
+            "std::math::ff::add",
+            vec![x.clone(), y.clone(), modulus.clone()],
+        );
+        assert_eq!((x.clone() + y.clone()) % modulus.clone(), result);
+
+        let result = evaluate_integer_function(
+            &analyzed,
+            "std::math::ff::sub",
+            vec![x.clone(), y.clone(), modulus.clone()],
+        );
+        assert_eq!(
+            (x.clone() - y.clone() + modulus.clone()) % modulus.clone(),
+            result
+        );
+
+        let result = evaluate_integer_function(
+            &analyzed,
+            "std::math::ff::mul",
+            vec![x.clone(), y.clone(), modulus.clone()],
+        );
+        assert_eq!((x.clone() * y.clone()) % modulus.clone(), result);
+
+        if y != 0.into() {
+            let result = evaluate_integer_function(
+                &analyzed,
+                "std::math::ff::div",
+                vec![x.clone(), y.clone(), modulus.clone()],
+            );
+            assert_eq!(x, (result * y) % modulus);
+        }
+    }
+}
+
+#[test]
+fn ff_inv_big() {
+    let analyzed = std_analyzed::<GoldilocksField>();
+    // modulus of the secp256k1 base field
+    let modulus = num_bigint::BigInt::from_str_radix(
+        "fffffffffffffffffffffffffffffffffffffffffffffffffffffffefffffc2f",
+        16,
+    )
+    .unwrap();
+    let x = modulus.clone() - num_bigint::BigInt::from(17);
+    let result = evaluate_integer_function(
+        &analyzed,
+        "std::math::ff::inverse",
+        vec![x.clone(), modulus.clone()],
+    );
+}

std/math/ff.asm

@@ -0,0 +1,43 @@
+/// Inverts `x` in the finite field with modulus `modulus`.
+/// Assumes that `modulus` is prime, but does not check it.
+let inverse = |x, modulus|
+    if x <= 0 || x >= modulus {
+        std::check::panic("Tried to compute the inverse of zero, of a negative number or a number outside the field.")
+    } else {
+        reduce(extended_gcd(x, modulus)[0], modulus)
+    };
+
+/// Computes `x + y` modulo the modulus.
+let add = |x, y, modulus| reduce(x + y, modulus);
+
+/// Computes `x - y` modulo the modulus.
+let sub = |x, y, modulus| reduce(x - y, modulus);
+
+/// Computes `x * y` modulo the modulus.
+let mul = |x, y, modulus| reduce(x * y, modulus);
+
+/// Computes `x / y` modulo the modulus.
+let div = |x, y, modulus| mul(x, inverse(y, modulus), modulus);
+
+/// Reduces `x` modulo `modulus`, so that it is in the range
+/// between `0` and `modulus`. Works on negative `x`.
+let reduce = |x, modulus|
+    if x < 0 {
+        (modulus - ((-x) % modulus)) % modulus
+    } else {
+        x % modulus
+    };
+
+let extended_gcd = |a, b|
+    if b == 0 {
+        if a == 1 {
+            [1, 0]
+        } else {
+            // a is the gcd, but we do not really want to compute it.
+            std::check::panic("Inputs are not co-prime, inverse does not exist.")
+        }
+    } else {
+        // TODO this is written in a complicated way
+        // because we do not have tuple destructuring assignment
+        (|r| [r[1], r[0] - (a / b) * r[1]])(extended_gcd(b, a % b))
+    };

std/math/mod.asm

@@ -0,0 +1 @@
+mod ff;

std/mod.asm

@@ -5,6 +5,7 @@ mod convert;
 mod debug;
 mod field;
 mod hash;
+mod math;
 mod shift;
 mod split;
 mod utils;