perf(RAM): Drop cached polynomials

The cached polynomials represent the same information as the original randomized trace, except in a different basis. If they are needed, the relevant randomized trace column is interpolated on the fly. Through various techniques, like batched barycentric evaluation, the drawback of losing the cached polynomials is offset almost completely. In the process, the code becomes notably simpler and, more importantly, requires considerably less RAM. Co-authored-by: Thorkil Schmidiger <[email protected]>
TritonVM · Oct 1, 2024 · 9bf637d · 9bf637d
1 parent ce592a9
commit 9bf637d
Show file tree

Hide file tree

Showing 5 changed files with 207 additions and 177 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -84,7 +84,7 @@ strum = { version = "0.26", features = ["derive"] }
 syn = "2.0"
 test-strategy = "0.4.0"
 thiserror = "1.0"
-twenty-first = "0.42.0-alpha.9"
+twenty-first = "0.42.0-alpha.10"
 unicode-width = "0.1"
 divan = "0.1.14"
 

diff --git a/triton-vm/benches/barycentric_eval.rs b/triton-vm/benches/barycentric_eval.rs
@@ -6,7 +6,8 @@ use itertools::Itertools;
 use rand::prelude::StdRng;
 use rand::Rng;
 use rand_core::SeedableRng;
-use triton_vm::fri::barycentric_evaluate;
+use twenty_first::math::polynomial::barycentric_evaluate;
+use twenty_first::prelude::XFieldElement;
 
 criterion_main!(benches);
 criterion_group!(
@@ -25,7 +26,7 @@ fn barycentric_eval<const LOG2_N: usize>(c: &mut Criterion) {
     c.bench_function(&format!("barycentric_evaluation_(1<<{LOG2_N})"), |b| {
         b.iter_batched(
             || ((0..1 << LOG2_N).map(|_| rng.gen()).collect_vec(), rng.gen()),
-            |(codeword, indeterminate)| barycentric_evaluate(&codeword, indeterminate),
+            |(cw, ind): (Vec<XFieldElement>, XFieldElement)| barycentric_evaluate(&cw, ind),
             BatchSize::SmallInput,
         )
     });

diff --git a/triton-vm/src/fri.rs b/triton-vm/src/fri.rs
@@ -1,8 +1,8 @@
 use itertools::Itertools;
 use num_traits::Zero;
 use rayon::prelude::*;
+use std::ops::Mul;
 use twenty_first::math::traits::FiniteField;
-use twenty_first::math::traits::PrimitiveRootOfUnity;
 use twenty_first::prelude::*;
 
 use crate::arithmetic_domain::ArithmeticDomain;
@@ -471,6 +471,9 @@ impl<'stream> FriVerifier<'stream> {
     fn assert_last_round_codeword_corresponds_to_low_degree_polynomial(
         &mut self,
     ) -> VerifierResult<()> {
+        // todo: remove once deprecated local function `barycentric_evaluate` is removed
+        use twenty_first::math::polynomial::barycentric_evaluate;
+
         if self.last_round_polynomial.degree() > self.last_round_max_degree.try_into().unwrap() {
             return Err(FriValidationError::LastRoundPolynomialHasTooHighDegree);
         }
@@ -624,33 +627,15 @@ fn codeword_as_digests(codeword: &[XFieldElement]) -> Vec<Digest> {
 /// Panics if the codeword is some length that is not a power of 2 or greater than (1 << 32).
 ///
 /// [Credit]: https://github.com/0xPolygonMiden/miden-vm/issues/568
-pub fn barycentric_evaluate(
-    codeword: &[XFieldElement],
+#[deprecated(
+    since = "0.42.2",
+    note = "use `twenty_first::math::polynomial::barycentric_evaluate` instead"
+)]
+pub fn barycentric_evaluate<FF: FiniteField + Mul<XFieldElement, Output = XFieldElement>>(
+    codeword: &[FF],
     indeterminate: XFieldElement,
 ) -> XFieldElement {
-    let root_order = codeword.len().try_into().unwrap();
-    let generator = BFieldElement::primitive_root_of_unity(root_order).unwrap();
-    let domain_iter = (0..root_order)
-        .scan(bfe!(1), |acc, _| {
-            let to_yield = Some(*acc);
-            *acc *= generator;
-            to_yield
-        })
-        .collect_vec();
-
-    let domain_shift = domain_iter.iter().map(|&d| indeterminate - d).collect();
-    let domain_shift_inverses = XFieldElement::batch_inversion(domain_shift);
-    let domain_over_domain_shift = domain_iter
-        .into_iter()
-        .zip(domain_shift_inverses)
-        .map(|(d, inv)| d * inv);
-    let numerator = domain_over_domain_shift
-        .clone()
-        .zip(codeword)
-        .map(|(dsi, &abscis)| dsi * abscis)
-        .sum::<XFieldElement>();
-    let denominator = domain_over_domain_shift.sum::<XFieldElement>();
-    numerator / denominator
+    twenty_first::math::polynomial::barycentric_evaluate(codeword, indeterminate)
 }
 
 #[cfg(test)]
@@ -661,7 +646,6 @@ mod tests {
     use assert2::assert;
     use assert2::let_assert;
     use itertools::Itertools;
-    use proptest::collection::vec;
     use proptest::prelude::*;
     use proptest_arbitrary_interop::arb;
     use rand::prelude::*;
@@ -1066,21 +1050,4 @@ mod tests {
     ) {
         let _verdict = fri.verify(&mut proof_stream);
     }
-
-    #[proptest]
-    fn polynomial_evaluation_and_barycentric_evaluation_are_equivalent(
-        #[strategy(1_usize..13)] _log_num_coefficients: usize,
-        #[strategy(1_usize..6)] log_expansion_factor: usize,
-        #[strategy(vec(arb(), 1 << #_log_num_coefficients))] coefficients: Vec<XFieldElement>,
-        #[strategy(arb())] indeterminate: XFieldElement,
-    ) {
-        let domain_len = coefficients.len() * (1 << log_expansion_factor);
-        let domain = ArithmeticDomain::of_length(domain_len).unwrap();
-        let polynomial = Polynomial::from(&coefficients);
-        let codeword = domain.evaluate(&polynomial);
-        prop_assert_eq!(
-            polynomial.evaluate(indeterminate),
-            barycentric_evaluate(&codeword, indeterminate)
-        );
-    }
 }