Minpoly quotient (Liam's trick) (Plonky3#177)

* ldt test

* something is buggy

* wow that took way too long to find

* little clean up

* nice quotient tests

* fix fri, tests pass

* simplify by computing a whole row at a time

* transposing saves 10%

* documentation

* fmt

* fmt, clippy

* fall back to scalar for misaligned inner_row

* fix bug

* spell Liam's name right

* no_std, rename xs->roots in binomial_expand

* assert F::Packing aligned in pack and unpack

* better quotient tests

* initial fast matrix reducer

* batched and packed reducer

* ExtensionField::from_base_fn

* messing with quotient

* i need to inline EF::D somehow

* add EF back into quotient so it can unroll

* cleanup

* fmt

* small cleanup

* rewrite comments for remainder poly packing

* use fancy new suffix function

* be more functional

* fix typo, make MatrixReducer stateless

field-testing/Cargo.toml

@@ -7,4 +7,7 @@ license = "MIT OR Apache-2.0"
 [dependencies]
 p3-field = { path = "../field" }
 rand = "0.8.5"
-criterion = "0.5.1"
+criterion = "0.5.1"
+
+[dev-dependencies]
+p3-baby-bear = { path = "../baby-bear" }

field-testing/src/lib.rs

@@ -165,3 +165,37 @@ macro_rules! test_two_adic_extension_field {
         }
     };
 }
+
+#[cfg(test)]
+mod tests {
+    use alloc::vec;
+    use alloc::vec::Vec;
+
+    use p3_baby_bear::BabyBear;
+    use p3_field::extension::{BinomialExtensionField, HasFrobenius};
+    use p3_field::{binomial_expand, eval_poly, AbstractExtensionField, AbstractField};
+    use rand::{thread_rng, Rng};
+
+    use super::*;
+
+    #[test]
+    fn test_minimal_poly() {
+        type F = BabyBear;
+        type EF = BinomialExtensionField<F, 4>;
+        for _ in 0..1024 {
+            let x: EF = thread_rng().gen();
+            let m: Vec<EF> = x.minimal_poly().into_iter().map(EF::from_base).collect();
+            assert!(eval_poly(&m, x).is_zero());
+        }
+    }
+
+    #[test]
+    fn test_binomial_expand() {
+        type F = BabyBear;
+        // (x - 1)(x - 2) = x^2 - 3x + 2
+        assert_eq!(
+            binomial_expand(&[F::one(), F::two()]),
+            vec![F::two(), -F::from_canonical_usize(3), F::one()]
+        );
+    }
+}

field/src/extension/binomial_extension.rs

@@ -9,7 +9,7 @@ use itertools::Itertools;
 use rand::distributions::Standard;
 use rand::prelude::Distribution;
 
-use super::{HasFrobenuis, HasTwoAdicBionmialExtension};
+use super::{HasFrobenius, HasTwoAdicBionmialExtension};
 use crate::extension::BinomiallyExtendable;
 use crate::field::Field;
 use crate::{field_to_array, AbstractExtensionField, AbstractField, ExtensionField, TwoAdicField};
@@ -40,7 +40,7 @@ impl<F: BinomiallyExtendable<D>, const D: usize> ExtensionField<F>
 {
 }
 
-impl<F: BinomiallyExtendable<D>, const D: usize> HasFrobenuis<F> for BinomialExtensionField<F, D> {
+impl<F: BinomiallyExtendable<D>, const D: usize> HasFrobenius<F> for BinomialExtensionField<F, D> {
     /// FrobeniusField automorphisms: x -> x^n, where n is the order of BaseField.
     fn frobenius(&self) -> Self {
         self.repeated_frobenius(1)
@@ -501,6 +501,13 @@ where
         }
     }
 
+    #[inline]
+    fn from_base_fn<F: FnMut(usize) -> AF>(f: F) -> Self {
+        Self {
+            value: array::from_fn(f),
+        }
+    }
+
     fn as_base_slice(&self) -> &[AF] {
         &self.value
     }

field/src/extension/mod.rs

@@ -1,8 +1,13 @@
+use core::{debug_assert, debug_assert_eq, iter};
+
 use crate::field::Field;
-use crate::ExtensionField;
+use crate::{naive_poly_mul, ExtensionField};
 
 mod binomial_extension;
 
+use alloc::vec;
+use alloc::vec::Vec;
+
 pub use binomial_extension::*;
 
 /// Binomial extension field trait.
@@ -19,10 +24,32 @@ pub trait BinomiallyExtendable<const D: usize>: Field {
     fn ext_generator() -> [Self; D];
 }
 
-pub trait HasFrobenuis<F: Field>: ExtensionField<F> {
+pub trait HasFrobenius<F: Field>: ExtensionField<F> {
     fn frobenius(&self) -> Self;
     fn repeated_frobenius(&self, count: usize) -> Self;
     fn frobenius_inv(&self) -> Self;
+
+    fn minimal_poly(mut self) -> Vec<F> {
+        let mut m = vec![Self::one()];
+        for _ in 0..Self::D {
+            m = naive_poly_mul(&m, &[-self, Self::one()]);
+            self = self.frobenius();
+        }
+        let mut m_iter = m
+            .into_iter()
+            .map(|c| c.as_base().expect("Extension is not algebraic?"));
+        let m: Vec<F> = m_iter.by_ref().take(Self::D + 1).collect();
+        debug_assert_eq!(m.len(), Self::D + 1);
+        debug_assert_eq!(m.last(), Some(&F::one()));
+        debug_assert!(m_iter.all(|c| c.is_zero()));
+        m
+    }
+
+    fn galois_group(self) -> Vec<Self> {
+        iter::successors(Some(self), |x| Some(x.frobenius()))
+            .take(Self::D)
+            .collect()
+    }
 }
 
 /// Optional trait for implementing Two Adic Binomial Extension Field.

field/src/field.rs

@@ -262,6 +262,10 @@ pub trait AbstractExtensionField<Base: AbstractField>:
     /// different f might have been used.
     fn from_base_slice(bs: &[Base]) -> Self;
 
+    /// Similar to `core:array::from_fn`, with the same caveats as
+    /// `from_base_slice`.
+    fn from_base_fn<F: FnMut(usize) -> Base>(f: F) -> Self;
+
     /// Suppose this field extension is represented by the quotient
     /// ring B[X]/(f(X)) where B is `Base` and f is an irreducible
     /// polynomial of degree `D`. This function takes a field element
@@ -288,6 +292,13 @@ pub trait ExtensionField<Base: Field>: Field + AbstractExtensionField<Base, F =
     fn is_in_basefield(&self) -> bool {
         self.as_base_slice()[1..].iter().all(Field::is_zero)
     }
+    fn as_base(&self) -> Option<Base> {
+        if self.is_in_basefield() {
+            Some(self.as_base_slice()[0])
+        } else {
+            None
+        }
+    }
 }
 
 impl<F: Field> ExtensionField<F> for F {}
@@ -304,6 +315,10 @@ impl<AF: AbstractField> AbstractExtensionField<AF> for AF {
         bs[0].clone()
     }
 
+    fn from_base_fn<F: FnMut(usize) -> AF>(mut f: F) -> Self {
+        f(0)
+    }
+
     fn as_base_slice(&self) -> &[AF] {
         slice::from_ref(self)
     }

field/src/helpers.rs

@@ -1,3 +1,4 @@
+use alloc::vec;
 use alloc::vec::Vec;
 use core::array;
 
@@ -64,3 +65,37 @@ pub fn field_to_array<AF: AbstractField, const D: usize>(x: AF) -> [AF; D] {
     arr[0] = x;
     arr
 }
+
+/// Naive polynomial multiplication.
+pub fn naive_poly_mul<AF: AbstractField>(a: &[AF], b: &[AF]) -> Vec<AF> {
+    // Grade school algorithm
+    let mut product = vec![AF::zero(); a.len() + b.len() - 1];
+    for (i, c1) in a.iter().enumerate() {
+        for (j, c2) in b.iter().enumerate() {
+            product[i + j] += c1.clone() * c2.clone();
+        }
+    }
+    product
+}
+
+/// Expand a product of binomials (x - roots[0])(x - roots[1]).. into polynomial coefficients.
+pub fn binomial_expand<AF: AbstractField>(roots: &[AF]) -> Vec<AF> {
+    let mut coeffs = vec![AF::zero(); roots.len() + 1];
+    coeffs[0] = AF::one();
+    for (i, x) in roots.iter().enumerate() {
+        for j in (1..i + 2).rev() {
+            coeffs[j] = coeffs[j - 1].clone() - x.clone() * coeffs[j].clone();
+        }
+        coeffs[0] *= -x.clone();
+    }
+    coeffs
+}
+
+pub fn eval_poly<AF: AbstractField>(poly: &[AF], x: AF) -> AF {
+    let mut acc = AF::zero();
+    for coeff in poly.iter().rev() {
+        acc *= x.clone();
+        acc += coeff.clone();
+    }
+    acc
+}

field/src/packed.rs

@@ -1,5 +1,5 @@
 use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Sub, SubAssign};
-use core::slice;
+use core::{mem, slice};
 
 use crate::field::Field;
 use crate::AbstractField;
@@ -73,6 +73,10 @@ pub unsafe trait PackedField: AbstractField<F = Self::Scalar>
     fn interleave(&self, other: Self, block_len: usize) -> (Self, Self);
 
     fn pack_slice(buf: &[Self::Scalar]) -> &[Self] {
+        // Sources vary, but this should be true on all platforms we care about.
+        // This should be a const assert, but trait methods can't access `Self` in a const context,
+        // even with inner struct instantiation. So we will trust LLVM to optimize this out.
+        assert!(mem::align_of::<Self>() <= mem::align_of::<Self::Scalar>());
         assert!(
             buf.len() % Self::WIDTH == 0,
             "Slice length (got {}) must be a multiple of packed field width ({}).",
@@ -84,7 +88,13 @@ pub unsafe trait PackedField: AbstractField<F = Self::Scalar>
         unsafe { slice::from_raw_parts(buf_ptr, n) }
     }
 
+    fn pack_slice_with_suffix(buf: &[Self::Scalar]) -> (&[Self], &[Self::Scalar]) {
+        let (packed, suffix) = buf.split_at(buf.len() - buf.len() % Self::WIDTH);
+        (Self::pack_slice(packed), suffix)
+    }
+
     fn pack_slice_mut(buf: &mut [Self::Scalar]) -> &mut [Self] {
+        assert!(mem::align_of::<Self>() <= mem::align_of::<Self::Scalar>());
         assert!(
             buf.len() % Self::WIDTH == 0,
             "Slice length (got {}) must be a multiple of packed field width ({}).",
@@ -95,6 +105,13 @@ pub unsafe trait PackedField: AbstractField<F = Self::Scalar>
         let n = buf.len() / Self::WIDTH;
         unsafe { slice::from_raw_parts_mut(buf_ptr, n) }
     }
+
+    fn unpack_slice(buf: &[Self]) -> &[Self::Scalar] {
+        assert!(mem::align_of::<Self>() >= mem::align_of::<Self::Scalar>());
+        let buf_ptr = buf.as_ptr().cast::<Self::Scalar>();
+        let n = buf.len() * Self::WIDTH;
+        unsafe { slice::from_raw_parts(buf_ptr, n) }
+    }
 }
 
 unsafe impl<F: Field> PackedField for F {

fri/src/config.rs

@@ -8,7 +8,7 @@ pub trait FriConfig {
     type Val: PrimeField64;
     type Challenge: ExtensionField<Self::Val> + TwoAdicField;
 
-    type InputMmcs: Mmcs<Self::Challenge>;
+    type InputMmcs: Mmcs<Self::Val>;
     type CommitPhaseMmcs: DirectMmcs<Self::Challenge>;
 
     type Challenger: FieldChallenger<Self::Val>
@@ -50,7 +50,7 @@ impl<Val, Challenge, InputMmcs, CommitPhaseMmcs, Challenger> FriConfig
 where
     Val: PrimeField64,
     Challenge: ExtensionField<Val> + TwoAdicField,
-    InputMmcs: Mmcs<Challenge>,
+    InputMmcs: Mmcs<Val>,
     CommitPhaseMmcs: DirectMmcs<Challenge>,
     Challenger: FieldChallenger<Val> + CanObserve<<CommitPhaseMmcs as Mmcs<Challenge>>::Commitment>,
 {

fri/src/lib.rs

@@ -12,6 +12,7 @@ use crate::verifier::verify;
 
 mod config;
 mod fold_even_odd;
+mod matrix_reducer;
 mod proof;
 mod prover;
 mod verifier;
@@ -24,7 +25,7 @@ pub struct FriLdt<FC: FriConfig> {
     pub config: FC,
 }
 
-impl<FC: FriConfig> Ldt<FC::Val, FC::Challenge, FC::InputMmcs, FC::Challenger> for FriLdt<FC> {
+impl<FC: FriConfig> Ldt<FC::Val, FC::InputMmcs, FC::Challenger> for FriLdt<FC> {
     type Proof = FriProof<FC>;
     type Error = ();
 
@@ -35,7 +36,7 @@ impl<FC: FriConfig> Ldt<FC::Val, FC::Challenge, FC::InputMmcs, FC::Challenger> f
     fn prove(
         &self,
         input_mmcs: &[FC::InputMmcs],
-        input_data: &[&<FC::InputMmcs as Mmcs<FC::Challenge>>::ProverData],
+        input_data: &[&<FC::InputMmcs as Mmcs<FC::Val>>::ProverData],
         challenger: &mut FC::Challenger,
     ) -> Self::Proof {
         prove::<FC>(&self.config, input_mmcs, input_data, challenger)
@@ -44,7 +45,7 @@ impl<FC: FriConfig> Ldt<FC::Val, FC::Challenge, FC::InputMmcs, FC::Challenger> f
     fn verify(
         &self,
         input_mmcs: &[FC::InputMmcs],
-        _input_commits: &[<FC::InputMmcs as Mmcs<FC::Challenge>>::Commitment],
+        _input_commits: &[<FC::InputMmcs as Mmcs<FC::Val>>::Commitment],
         proof: &Self::Proof,
         challenger: &mut FC::Challenger,
     ) -> Result<(), Self::Error> {