Implemented POC of parallel FFT and twiddle generation in Metal (#190)

* Added FFTMetalState

* Added device field to FFTMetalState

* Added WIP twiddle factors generator in GPU

* Small fixes in WIP twiddle GPU generator

* Changed twiddle generator function constants to buffers, formatting and small fixes

* Removed unused helper, formatting

* [WIP] Workaround for result buffer, completed gen_twiddles, tests

* Update compiled metal lib

* Small fixes in gen_twiddles

* Changed gen_twiddles_cpu loop

* Fix twiddle generation on gpu.

* Fix artifacts from broken main

* Fixed threadgroup dispatching from twiddles

* [WIP] POC Parallel FFT

* Move all gpu related features to a new crate

* [WIP] Fix metal kernel

* [WIP] Various changes related to Metal FFT

* [WIP] For testing purposes

* Use real U32 field.

* [WIP] Fixed undefined behaviours

* Finished FFT in Metal POC

* Removed unused line

* Added comments for unsafe lines

* Changed storage mode for metal buffers

* Replaced static array with unsafe vector from pointer

* Proptest for Metal FFT

* Simplified stage loop

* Split Metal setup from execution

Will be useful for benchmarks

* Fixed basetype size

* Changed FFT Metal test name

* Small fix

* Fixed test memory leaks

* Added benchmarks

* Fixed leaks on benchmarks

* Moved and revamped parallel twiddle generation

* Formatting

* Add build-essential to try to fix dependency problem

* Add sudo to apt install

* Swap apt install to g++

* Try other solutions to dependency problem

* Try reinstall g++

* Try installing gobjc++-mingw

* Try gobjc++

* Try gobjc

* Try all

* Add gnustep-base-common

* Try cargo check as command

* Try CFLAGS

* Update CFLAGS

* Export CFLAGS

* Update CFLAGS

* Try cross

* Fix syntax

* Try CFLAGS

* Fix syntax

* Try macos runners

* Swap other CI steps to macos runner

* Delete unused function

* Add benchmark for twiddles to the general benchmark function

* Add conditional compilation for objc crate

* Swap runners back to ubuntu

* Add conditional compilation for metal crate to compile only in macos

* Revert conditional compilation for metal and objc crates

* Deleted operations.rs

* Added newline at end of fp.h.metal

* Reverted name change of U64TestField

* Change runners in CI to run in macos

* Ignore metal tests

* Removed some unsafe blocks, added comments

* Removed metal dependency from math crate

* Removed GPU from default members

* Reverted to use ubuntu on CI

* Changed coverage host to macos

---------

Co-authored-by: Martin Paulucci <martin.paulucci@lambdaclass.com>
Co-authored-by: IAvecilla <iavecilla@fi.uba.ar>
Co-authored-by: klaus993 <18153834+klaus993@users.noreply.github.com>
Co-authored-by: MauroFab <maurotoscano2@gmail.com>

Author: 5 people

.github/workflows/tests.yaml

@@ -30,7 +30,7 @@ jobs:
           command: check
 
   coverage:
-    runs-on: ubuntu-latest
+    runs-on: macos-latest
     env:
       CARGO_TERM_COLOR: always
     steps:

Cargo.toml

@@ -1,7 +1,12 @@
 [workspace]
-
+default-members = [
+    "math", 
+    "crypto",
+    "proving_system/stark"
+]
 members = [
     "math",
     "crypto",
     "proving_system/stark",
+    "gpu",
 ]

gpu/Cargo.toml

@@ -0,0 +1,21 @@
+[package]
+name = "lambdaworks-gpu"
+version = "0.1.0"
+edition = "2021"
+
+# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
+
+[dependencies]
+lambdaworks-math = { path = "../math" }
+lambdaworks-crypto = { path = "../crypto"}
+rand = "0.8.5"
+metal = "0.24.0"
+
+[dev-dependencies]
+proptest = "1.1.0"
+criterion = "0.4.0"
+objc = "0.2.7"
+
+[[bench]]
+name = "metal_benchmarks"
+harness = false

gpu/benches/all_benchmarks.rs

@@ -0,0 +1,11 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+
+mod benchmarks;
+
+fn run_all_benchmarks(c: &mut Criterion) {
+    benchmarks::metal::metal_fft_twiddles_benchmarks(c);
+    benchmarks::metal::metal_fft_benchmarks(c);
+}
+
+criterion_group!(benches, run_all_benchmarks);
+criterion_main!(benches);

gpu/benches/benchmarks/metal.rs

@@ -0,0 +1,82 @@
+use criterion::Criterion;
+use lambdaworks_gpu::fft::fft_metal::*;
+use lambdaworks_math::{
+    fft::bit_reversing::in_place_bit_reverse_permute,
+    field::{element::FieldElement, traits::IsTwoAdicField},
+    field::{test_fields::u32_test_field::U32TestField, traits::RootsConfig},
+};
+use rand::random;
+
+type F = U32TestField;
+type FE = FieldElement<F>;
+
+fn gen_coeffs(pow: usize) -> Vec<FE> {
+    let mut result = Vec::with_capacity(1 << pow);
+    for _ in 0..result.capacity() {
+        result.push(FE::new(random()));
+    }
+    result
+}
+
+pub fn metal_fft_benchmarks(c: &mut Criterion) {
+    let mut group = c.benchmark_group("metal_fft");
+
+    for order in 20..=24 {
+        let coeffs = gen_coeffs(order);
+        group.throughput(criterion::Throughput::Elements(1 << order)); // info for criterion
+
+        // the objective is to bench ordered FFT, including twiddles generation and Metal setup
+        group.bench_with_input(
+            format!("parallel_nr_2radix_2^{order}_coeffs"),
+            &coeffs,
+            |bench, coeffs| {
+                bench.iter(|| {
+                    // TODO: autoreleaspool hurts perf. by 2-3%. Search for an alternative
+                    objc::rc::autoreleasepool(|| {
+                        let coeffs = coeffs.clone();
+                        let twiddles =
+                            F::get_twiddles(order as u64, RootsConfig::BitReverse).unwrap();
+                        let fft_metal = FFTMetalState::new(None).unwrap();
+                        let command_buff_encoder = fft_metal
+                            .setup_fft("radix2_dit_butterfly", &twiddles)
+                            .unwrap();
+
+                        let mut result = fft_metal
+                            .execute_fft(&coeffs, command_buff_encoder)
+                            .unwrap();
+
+                        in_place_bit_reverse_permute(&mut result);
+                    });
+                });
+            },
+        );
+    }
+
+    group.finish();
+}
+
+pub fn metal_fft_twiddles_benchmarks(c: &mut Criterion) {
+    let mut group = c.benchmark_group("metal_fft");
+    group.sample_size(10); // it becomes too slow with the default of 100
+
+    for order in 2..=4 {
+        group.throughput(criterion::Throughput::Elements(1 << order)); // info for criterion
+
+        // the objective is to bench ordered FFT, including twiddles generation and Metal setup
+        group.bench_with_input(
+            format!("parallel_twiddle_factors_2^({order}-1)_elems"),
+            &order,
+            |bench, order| {
+                bench.iter(|| {
+                    // TODO: autoreleaspool hurts perf. by 2-3%. Search for an alternative
+                    objc::rc::autoreleasepool(|| {
+                        let metal_state = FFTMetalState::new(None).unwrap();
+                        let _gpu_twiddles = metal_state.gen_twiddles::<F>(*order).unwrap();
+                    });
+                });
+            },
+        );
+    }
+
+    group.finish();
+}

gpu/benches/benchmarks/mod.rs

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

gpu/benches/metal_benchmarks.rs

@@ -0,0 +1,11 @@
+use criterion::{criterion_group, criterion_main, Criterion};
+
+mod benchmarks;
+
+fn run_metal_benchmarks(c: &mut Criterion) {
+    benchmarks::metal::metal_fft_twiddles_benchmarks(c);
+    benchmarks::metal::metal_fft_benchmarks(c);
+}
+
+criterion_group!(benches, run_metal_benchmarks);
+criterion_main!(benches);