Extract FFT code into its own crate (#252)

* Remove fft from `math` module.

* Make a separate fft crate to avoid circular dependencies.

* Bring back benches.

* Move tests to gpu.

* Fix error handling.

* Remove more repeated code.

* Rename abstractions to ops.

* Move metal tests back to gpu.

* Fix linting error.

* Remove duplicate tests.

* nit.

Author: mpaulucci

.rusty-hook.toml

@@ -1,5 +1,5 @@
 [hooks]
-pre-commit = "cargo test && cargo clippy --all-targets --all-features -- -D warnings && cargo fmt --all -- --check"
+pre-commit = "cargo test && cargo clippy --all-targets -- -D warnings && cargo fmt --all -- --check"
 
 [logging]
 verbose = true

Cargo.toml

@@ -1,6 +1,7 @@
 [workspace]
 members = [
     "math",
+    "fft",
     "crypto",
     "proving_system/stark",
     "proving_system/plonk",

fft/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "lambdaworks-fft"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+rand = "0.8.5"
+thiserror = "1.0.38"
+lambdaworks-math = { path = "../math" }
+lambdaworks-gpu = { path = "../gpu", optional = true }
+
+[dev-dependencies]
+proptest = "1.1.0"
+criterion = "0.4.0"
+objc = "0.2.7"
+
+[features]
+metal = ["dep:lambdaworks-gpu"]
+cuda = ["dep:lambdaworks-gpu"]
+
+[[bench]]
+name = "fft_benchmarks"
+harness = false

math/benches/benchmarks/fft.rs renamed to fft/benches/benchmarks/fft.rs

@@ -1,8 +1,15 @@
 use criterion::Criterion;
+use lambdaworks_fft::{
+    bit_reversing::in_place_bit_reverse_permute,
+    fft_iterative::{in_place_nr_2radix_fft, in_place_rn_2radix_fft},
+    polynomial::FFTPoly,
+    roots_of_unity::get_twiddles,
+};
 use lambdaworks_math::{
-    fft::{bit_reversing::*, fft_iterative::*},
-    field::{element::FieldElement, traits::IsTwoAdicField},
-    field::{fields::fft_friendly::stark_252_prime_field::Stark252PrimeField, traits::RootsConfig},
+    field::{
+        element::FieldElement, fields::fft_friendly::stark_252_prime_field::Stark252PrimeField,
+        traits::RootsConfig,
+    },
     polynomial::Polynomial,
     unsigned_integer::element::UnsignedInteger,
 };
@@ -29,8 +36,8 @@ pub fn fft_benchmarks(c: &mut Criterion) {
         group.throughput(criterion::Throughput::Elements(1 << order));
 
         let input = rand_field_elements(order);
-        let twiddles_bitrev = F::get_twiddles(order, RootsConfig::BitReverse).unwrap();
-        let twiddles_nat = F::get_twiddles(order, RootsConfig::Natural).unwrap();
+        let twiddles_bitrev = get_twiddles(order, RootsConfig::BitReverse).unwrap();
+        let twiddles_nat = get_twiddles(order, RootsConfig::Natural).unwrap();
 
         // the objective is to bench ordered FFT, that's why a bitrev permutation is added.
         group.bench_with_input("Sequential from NR radix2", &input, |bench, input| {
@@ -61,7 +68,7 @@ pub fn twiddles_benchmarks(c: &mut Criterion) {
 
         group.bench_with_input("Sequential", &order, |bench, order| {
             bench.iter(|| {
-                F::get_twiddles(*order, RootsConfig::Natural).unwrap();
+                get_twiddles::<F>(*order, RootsConfig::Natural).unwrap();
             });
         });
     }

fft/benches/benchmarks/mod.rs

@@ -0,0 +1 @@
+pub mod fft;

math/benches/fft_benchmarks.rs renamed to fft/benches/fft_benchmarks.rs

@@ -1,3 +1,4 @@
+use lambdaworks_math::field::errors::FieldError;
 use thiserror::Error;
 
 #[derive(Debug, Error)]
@@ -6,7 +7,6 @@ pub enum FFTError {
     InvalidOrder(String),
     #[error("Could not calculate {1} root of unity")]
     RootOfUnityError(String, u64),
-
     #[error("Couldn't find a system default device for Metal")]
     MetalDeviceNotFound(),
     #[error("Couldn't create a new Metal library: {0}")]
@@ -16,3 +16,14 @@ pub enum FFTError {
     #[error("Couldn't create a new Metal compute pipeline: {0}")]
     MetalPipelineError(String),
 }
+
+impl From<FieldError> for FFTError {
+    fn from(error: FieldError) -> Self {
+        match error {
+            FieldError::DivisionByZero => {
+                panic!("Can't divide by zero during FFT");
+            }
+            FieldError::RootOfUnityError(error, order) => FFTError::RootOfUnityError(error, order),
+        }
+    }
+}

math/src/fft/bit_reversing.rs renamed to fft/src/bit_reversing.rs

@@ -1,4 +1,4 @@
-use crate::field::{element::FieldElement, traits::IsTwoAdicField};
+use lambdaworks_math::field::{element::FieldElement, traits::IsTwoAdicField};
 
 /// In-Place Radix-2 NR DIT FFT algorithm over a slice of two-adic field elements.
 /// It's required that the twiddle factors are in bit-reverse order. Else this function will not
@@ -93,10 +93,10 @@ where
 
 #[cfg(test)]
 mod tests {
-    use crate::fft::helpers::log2;
-    use crate::fft::test_helpers::naive_matrix_dft_test;
-    use crate::field::test_fields::u64_test_field::U64TestField;
-    use crate::{fft::bit_reversing::in_place_bit_reverse_permute, field::traits::RootsConfig};
+    use crate::helpers::log2;
+    use crate::test_helpers::naive_matrix_dft_test;
+    use crate::{bit_reversing::in_place_bit_reverse_permute, roots_of_unity::get_twiddles};
+    use lambdaworks_math::field::{test_fields::u64_test_field::U64TestField, traits::RootsConfig};
     use proptest::{collection, prelude::*};
 
     use super::*;
@@ -127,7 +127,7 @@ mod tests {
             let expected = naive_matrix_dft_test(&coeffs);
 
             let order = log2(coeffs.len()).unwrap();
-            let twiddles = F::get_twiddles(order, RootsConfig::BitReverse).unwrap();
+            let twiddles = get_twiddles(order, RootsConfig::BitReverse).unwrap();
 
             let mut result = coeffs;
             in_place_nr_2radix_fft(&mut result, &twiddles);
@@ -142,7 +142,7 @@ mod tests {
             let expected = naive_matrix_dft_test(&coeffs);
 
             let order = log2(coeffs.len()).unwrap();
-            let twiddles = F::get_twiddles(order, RootsConfig::Natural).unwrap();
+            let twiddles = get_twiddles(order, RootsConfig::Natural).unwrap();
 
             let mut result = coeffs;
             in_place_bit_reverse_permute(&mut result);

math/src/fft/errors.rs renamed to fft/src/errors.rs

@@ -1,8 +1,10 @@
-pub mod abstractions;
 pub mod bit_reversing;
 pub mod errors;
 pub mod fft_iterative;
 pub(crate) mod helpers;
+pub mod ops;
+pub mod polynomial;
+pub mod roots_of_unity;
 
 #[cfg(test)]
 pub(crate) mod test_helpers;

math/src/fft/fft_iterative.rs renamed to fft/src/fft_iterative.rs

@@ -1,8 +1,10 @@
-use crate::field::{
+use lambdaworks_math::field::{
     element::FieldElement,
     traits::{IsTwoAdicField, RootsConfig},
 };
 
+use crate::roots_of_unity::get_twiddles;
+
 use super::{
     bit_reversing::in_place_bit_reverse_permute, errors::FFTError,
     fft_iterative::in_place_nr_2radix_fft, helpers::log2,
@@ -16,7 +18,7 @@ pub fn fft_with_blowup<F: IsTwoAdicField>(
 ) -> Result<Vec<FieldElement<F>>, FFTError> {
     let domain_size = coeffs.len() * blowup_factor;
     let order = log2(domain_size)?;
-    let twiddles = F::get_twiddles(order, RootsConfig::BitReverse)?;
+    let twiddles = get_twiddles(order, RootsConfig::BitReverse)?;
 
     let mut results = coeffs.to_vec();
     results.resize(domain_size, FieldElement::zero());
@@ -33,7 +35,7 @@ pub fn fft<F: IsTwoAdicField>(
     coeffs: &[FieldElement<F>],
 ) -> Result<Vec<FieldElement<F>>, FFTError> {
     let order = log2(coeffs.len())?;
-    let twiddles = F::get_twiddles(order, RootsConfig::BitReverse)?;
+    let twiddles = get_twiddles(order, RootsConfig::BitReverse)?;
 
     let mut results = coeffs.to_vec();
     in_place_nr_2radix_fft(&mut results, &twiddles);
@@ -48,7 +50,7 @@ pub fn inverse_fft<F: IsTwoAdicField>(
     coeffs: &[FieldElement<F>],
 ) -> Result<Vec<FieldElement<F>>, FFTError> {
     let order = log2(coeffs.len())?;
-    let twiddles = F::get_twiddles(order, RootsConfig::BitReverseInversed)?;
+    let twiddles = get_twiddles(order, RootsConfig::BitReverseInversed)?;
 
     let mut results = coeffs.to_vec();
     in_place_nr_2radix_fft(&mut results, &twiddles);