update fork

succinctlabs · Mar 17, 2024 · a5da07c · a5da07c
2 parents 7d11ba5 + 6144d64
commit a5da07c
Show file tree

Hide file tree

Showing 102 changed files with 7,172 additions and 2,232 deletions.
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -122,12 +122,10 @@ jobs:
         cargo check --verbose --package p3-merkle-tree
         cargo check --verbose --package p3-mersenne-31
         cargo check --verbose --package p3-monolith
-        cargo check --verbose --package p3-multi-stark
         cargo check --verbose --package p3-poseidon
         cargo check --verbose --package p3-poseidon2
         cargo check --verbose --package p3-reed-solomon
         cargo check --verbose --package p3-rescue
         cargo check --verbose --package p3-symmetric
-        cargo check --verbose --package p3-tensor-pcs
         cargo check --verbose --package p3-uni-stark
         cargo check --verbose --package p3-util
diff --git a/Cargo.toml b/Cargo.toml
@@ -6,6 +6,7 @@ members = [
     "blake3",
     "brakedown",
     "challenger",
+    "circle",
     "code",
     "commit",
     "dft",
@@ -22,13 +23,11 @@ members = [
     "maybe-rayon",
     "mersenne-31",
     "monolith",
-    "multi-stark",
     "poseidon",
     "poseidon2",
     "reed-solomon",
     "rescue",
     "symmetric",
-    "tensor-pcs",
     "util",
     "uni-stark",
 ]
diff --git a/README.md b/README.md
@@ -12,12 +12,12 @@ Fields:
   - [x] "complex" extension field
   - [x] ~128 bit extension field
   - [x] AVX2
-  - [ ] AVX-512
+  - [x] AVX-512
   - [x] NEON
 - [x] BabyBear
   - [x] ~128 bit extension field
-  - [ ] AVX2
-  - [ ] AVX-512
+  - [x] AVX2
+  - [x] AVX-512
   - [x] NEON
 - [x] Goldilocks
   - [x] ~128 bit extension field
@@ -64,6 +64,22 @@ We sometimes use a Keccak AIR to compare Plonky3's performance to other librarie
 RUST_LOG=info cargo run --example prove_baby_bear_keccak --release --features parallel
 ```
 
+## CPU features
+
+Plonky3 contains optimizations that rely on newer CPU instructions that are not available in older processors. These instruction sets include x86's [BMI1 and 2](https://en.wikipedia.org/wiki/X86_Bit_manipulation_instruction_set), [AVX2](https://en.wikipedia.org/wiki/Advanced_Vector_Extensions#Advanced_Vector_Extensions_2), and [AVX-512](https://en.wikipedia.org/wiki/AVX-512). Rustc does not emit those instructions by default; they must be explicitly enabled through the `target-feature` compiler option (or implicitly by setting `target-cpu`). To enable all features that are supported on your machine, you can set `target-cpu` to `native`. For example, to run the tests:
+```
+RUSTFLAGS="-Ctarget-cpu=native" cargo test
+```
+
+Support for some instructions, such as AVX-512, is still experimental. They are only available in the nightly build of Rustc and are enabled by the [`nightly-features` feature flag](#nightly-only-optimizations). To use them, you must enable the flag in Rustc (e.g. by setting `target-feature`) and you must also enable the `nightly-features` feature.
+
+## Nightly-only optimizations
+
+Some optimizations (in particular, AVX-512-optimized math) rely on features that are currently available only in the nightly build of Rustc. To use them, you need to enable the `nightly-features` feature. For example, to run the tests:
+```
+cargo test --features nightly-features
+```
+
 
 ## License
 

diff --git a/air/src/virtual_column.rs b/air/src/virtual_column.rs
@@ -1,3 +1,4 @@
+use alloc::borrow::Cow;
 use alloc::vec;
 use alloc::vec::Vec;
 use core::ops::Mul;
@@ -6,8 +7,8 @@ use p3_field::{AbstractField, Field};
 
 /// An affine function over columns in a PAIR.
 #[derive(Clone, Debug)]
-pub struct VirtualPairCol<F: Field> {
-    column_weights: Vec<(PairCol, F)>,
+pub struct VirtualPairCol<'a, F: Field> {
+    column_weights: Cow<'a, [(PairCol, F)]>,
     constant: F,
 }
 
@@ -19,22 +20,36 @@ pub enum PairCol {
 }
 
 impl PairCol {
-    fn get<T: Copy>(&self, preprocessed: &[T], main: &[T]) -> T {
+    pub const fn get<T: Copy>(&self, preprocessed: &[T], main: &[T]) -> T {
         match self {
             PairCol::Preprocessed(i) => preprocessed[*i],
             PairCol::Main(i) => main[*i],
         }
     }
 }
 
-impl<F: Field> VirtualPairCol<F> {
-    pub fn new(column_weights: Vec<(PairCol, F)>, constant: F) -> Self {
+impl<'a, F: Field> VirtualPairCol<'a, F> {
+    pub const fn new(column_weights: Cow<'a, [(PairCol, F)]>, constant: F) -> Self {
         Self {
             column_weights,
             constant,
         }
     }
 
+    pub const fn new_owned(column_weights: Vec<(PairCol, F)>, constant: F) -> Self {
+        Self {
+            column_weights: Cow::Owned(column_weights),
+            constant,
+        }
+    }
+
+    pub const fn new_borrowed(column_weights: &'a [(PairCol, F)], constant: F) -> Self {
+        Self {
+            column_weights: Cow::Borrowed(column_weights),
+            constant,
+        }
+    }
+
     pub fn new_preprocessed(column_weights: Vec<(usize, F)>, constant: F) -> Self {
         Self::new(
             column_weights
@@ -63,15 +78,15 @@ impl<F: Field> VirtualPairCol<F> {
     #[must_use]
     pub fn constant(x: F) -> Self {
         Self {
-            column_weights: vec![],
+            column_weights: Cow::Owned(vec![]),
             constant: x,
         }
     }
 
     #[must_use]
     pub fn single(column: PairCol) -> Self {
         Self {
-            column_weights: vec![(column, F::one())],
+            column_weights: Cow::Owned(vec![(column, F::one())]),
             constant: F::zero(),
         }
     }
@@ -117,7 +132,7 @@ impl<F: Field> VirtualPairCol<F> {
         Var: Into<Expr> + Copy,
     {
         let mut result = self.constant.into();
-        for (column, weight) in &self.column_weights {
+        for (column, weight) in self.column_weights.iter() {
             result += column.get(preprocessed, main).into() * *weight;
         }
         result

diff --git a/baby-bear/Cargo.toml b/baby-bear/Cargo.toml
@@ -4,13 +4,22 @@ version = "0.1.0"
 edition = "2021"
 license = "MIT OR Apache-2.0"
 
+[features]
+nightly-features = []
+
 [dependencies]
 p3-field = { path = "../field" }
+p3-mds = { path = "../mds" }
+p3-poseidon2 = { path = "../poseidon2" }
+p3-symmetric = { path = "../symmetric" }
 rand = "0.8.5"
 serde = { version = "1.0", default-features = false, features = ["derive"] }
 
 [dev-dependencies]
 p3-field-testing = { path = "../field-testing" }
+ark-ff = { version = "^0.4.0", default-features = false }
+zkhash = { git = "https://github.com/HorizenLabs/poseidon2" }
+rand = { version = "0.8.5", features = ["min_const_gen"] }
 criterion = "0.5.1"
 rand_chacha = "0.3.1"
 serde_json = "1.0.113"

diff --git a/baby-bear/src/baby_bear.rs b/baby-bear/src/baby_bear.rs
@@ -3,8 +3,8 @@ use core::iter::{Product, Sum};
 use core::ops::{Add, AddAssign, Div, Mul, MulAssign, Neg, Sub, SubAssign};
 
 use p3_field::{
-    exp_1725656503, exp_u64_by_squaring, AbstractField, Field, Packable, PrimeField, PrimeField32,
-    PrimeField64, TwoAdicField,
+    exp_1725656503, exp_u64_by_squaring, halve_u32, AbstractField, Field, Packable, PrimeField,
+    PrimeField32, PrimeField64, TwoAdicField,
 };
 use rand::distributions::{Distribution, Standard};
 use rand::Rng;
@@ -186,11 +186,30 @@ impl AbstractField for BabyBear {
 impl Field for BabyBear {
     #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
     type Packing = crate::PackedBabyBearNeon;
-    #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
+    #[cfg(all(
+        target_arch = "x86_64",
+        target_feature = "avx2",
+        not(all(feature = "nightly-features", target_feature = "avx512f"))
+    ))]
     type Packing = crate::PackedBabyBearAVX2;
+    #[cfg(all(
+        feature = "nightly-features",
+        target_arch = "x86_64",
+        target_feature = "avx512f"
+    ))]
+    type Packing = crate::PackedBabyBearAVX512;
     #[cfg(not(any(
         all(target_arch = "aarch64", target_feature = "neon"),
-        all(target_arch = "x86_64", target_feature = "avx2"),
+        all(
+            target_arch = "x86_64",
+            target_feature = "avx2",
+            not(all(feature = "nightly-features", target_feature = "avx512f"))
+        ),
+        all(
+            feature = "nightly-features",
+            target_arch = "x86_64",
+            target_feature = "avx512f"
+        ),
     )))]
     type Packing = Self;
 
@@ -238,6 +257,13 @@ impl Field for BabyBear {
 
         Some(p1110111111111111111111111111111)
     }
+
+    #[inline]
+    fn halve(&self) -> Self {
+        BabyBear {
+            value: halve_u32::<P>(self.value),
+        }
+    }
 }
 
 impl PrimeField for BabyBear {}
@@ -249,20 +275,6 @@ impl PrimeField64 for BabyBear {
     fn as_canonical_u64(&self) -> u64 {
         u64::from(self.as_canonical_u32())
     }
-
-    #[inline]
-    fn linear_combination_u64<const N: usize>(u: [u64; N], v: &[Self; N]) -> Self {
-        // In order not to overflow a u64, we must have sum(u) <= 2^32.
-        debug_assert!(u.iter().sum::<u64>() <= (1u64 << 32));
-
-        let mut dot = u[0] * v[0].value as u64;
-        for i in 1..N {
-            dot += u[i] * v[i].value as u64;
-        }
-        Self {
-            value: (dot % (P as u64)) as u32,
-        }
-    }
 }
 
 impl PrimeField32 for BabyBear {
@@ -418,6 +430,23 @@ const fn to_monty(x: u32) -> u32 {
     (((x as u64) << MONTY_BITS) % P as u64) as u32
 }
 
+/// Convert a constant u32 array into a constant Babybear array.
+/// Saves every element in Monty Form
+#[inline]
+#[must_use]
+pub(crate) const fn to_babybear_array<const N: usize>(input: [u32; N]) -> [BabyBear; N] {
+    let mut output = [BabyBear { value: 0 }; N];
+    let mut i = 0;
+    loop {
+        if i == N {
+            break;
+        }
+        output[i].value = to_monty(input[i]);
+        i += 1;
+    }
+    output
+}
+
 #[inline]
 #[must_use]
 fn to_monty_64(x: u64) -> u32 {

diff --git a/baby-bear/src/lib.rs b/baby-bear/src/lib.rs
@@ -1,11 +1,23 @@
 #![no_std]
+#![cfg_attr(
+    all(
+        feature = "nightly-features",
+        target_arch = "x86_64",
+        target_feature = "avx512f"
+    ),
+    feature(stdarch_x86_avx512)
+)]
 
 extern crate alloc;
 
 mod baby_bear;
 mod extension;
+mod mds;
+mod poseidon2;
 
 pub use baby_bear::*;
+pub use mds::*;
+pub use poseidon2::DiffusionMatrixBabybear;
 
 #[cfg(all(target_arch = "aarch64", target_feature = "neon"))]
 mod aarch64_neon;
@@ -16,3 +28,16 @@ pub use aarch64_neon::*;
 mod x86_64_avx2;
 #[cfg(all(target_arch = "x86_64", target_feature = "avx2"))]
 pub use x86_64_avx2::*;
+
+#[cfg(all(
+    feature = "nightly-features",
+    target_arch = "x86_64",
+    target_feature = "avx512f"
+))]
+mod x86_64_avx512;
+#[cfg(all(
+    feature = "nightly-features",
+    target_arch = "x86_64",
+    target_feature = "avx512f"
+))]
+pub use x86_64_avx512::*;