Implemented small Metal abstractions and make the FFT POC usable

.github/workflows/ci.yaml

@@ -65,7 +65,7 @@ jobs:
       - name: Install cargo-llvm-cov
         uses: taiki-e/install-action@cargo-llvm-cov
       - name: Run tests and generate code coverage
-        run: cargo llvm-cov --all-features --workspace --lcov --output-path lcov.info
+        run: cargo llvm-cov --lcov --output-path lcov.info
       - name: Upload coverage to Codecov
         uses: codecov/codecov-action@v3
         with:

Cargo.toml

@@ -1,10 +1,4 @@
 [workspace]
-default-members = [
-    "math",
-    "crypto",
-    "proving_system/stark",
-    "proving_system/plonk"
-]
 members = [
     "math",
     "crypto",

Makefile

@@ -4,7 +4,7 @@ test:
 	cargo test
 
 clippy:
-	cargo clippy --all-targets --all-features -- -D warnings
+	cargo clippy --all-targets -- -D warnings
 
 docker-shell:
 	docker build -t rust-curves .
@@ -20,5 +20,9 @@ benchmarks:
 benchmark:
 	cargo criterion --bench ${BENCH}
 
+METALPATH = gpu/src/metal/shaders
+build-metal:
+	xcrun -sdk macosx metal $(METALPATH)/all.metal -o $(METALPATH)/lib.metallib
+
 docs:
 	cd docs && mdbook serve --open

gpu/Cargo.toml

@@ -7,15 +7,28 @@ edition = "2021"
 
 [dependencies]
 lambdaworks-math = { path = "../math" }
-lambdaworks-crypto = { path = "../crypto"}
+thiserror = "1.0.38"
 rand = "0.8.5"
-metal = "0.24.0"
+metal = { version = "0.24.0", optional = true }
+objc = { version = "0.2.7", optional = true }
+
+[features]
+metal = ["dep:metal", "dep:objc"]
+cuda = []
+
+# Some features activate compilation of code which isn't
+# supported in all machines (e.g. metal, cuda), so we won't
+# use `--all-features` in any case, instead every feature
+# that should compile in all cases will require to be added
+# as default. If you don't want to compile with all of these
+# use `--no-default-features`.
+default = []
 
 [dev-dependencies]
 proptest = "1.1.0"
 criterion = "0.4.0"
-objc = "0.2.7"
 
 [[bench]]
 name = "metal_benchmarks"
 harness = false
+required-features = ["metal"]

gpu/benches/benchmarks/metal.rs

@@ -1,7 +1,6 @@
 use criterion::Criterion;
-use lambdaworks_gpu::fft::fft_metal::*;
+use lambdaworks_gpu::metal::{abstractions::state::MetalState, fft::ops::*};
 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},
 };
@@ -10,42 +9,35 @@ 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);
+fn gen_coeffs(order: usize) -> Vec<FE> {
+    let mut result = Vec::with_capacity(1 << order);
     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");
+    let mut group = c.benchmark_group("FFT");
 
-    for order in 20..=24 {
+    for order in 21..=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"),
+            format!("Parallel NR radix2 FFT for 2^{order} elements"),
             &coeffs,
             |bench, coeffs| {
                 bench.iter(|| {
                     // TODO: autoreleaspool hurts perf. by 2-3%. Search for an alternative
                     objc::rc::autoreleasepool(|| {
                         let coeffs = coeffs.clone();
+                        let metal_state = MetalState::new(None).unwrap();
                         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);
+                        fft(&coeffs, &twiddles, &metal_state).unwrap();
                     });
                 });
             },
@@ -56,22 +48,46 @@ pub fn metal_fft_benchmarks(c: &mut Criterion) {
 }
 
 pub fn metal_fft_twiddles_benchmarks(c: &mut Criterion) {
-    let mut group = c.benchmark_group("metal_fft");
+    let mut group = c.benchmark_group("FFT twiddles generation");
     group.sample_size(10); // it becomes too slow with the default of 100
 
-    for order in 2..=4 {
+    for order in 21..=24 {
         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"),
+            format!("Parallel twiddles generation for 2^({order}-1) elements"),
             &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();
+                        let metal_state = MetalState::new(None).unwrap();
+                        gen_twiddles::<F>(*order, RootsConfig::Natural, &metal_state).unwrap();
+                    });
+                });
+            },
+        );
+    }
+
+    group.finish();
+}
+
+pub fn metal_bitrev_permutation_benchmarks(c: &mut Criterion) {
+    let mut group = c.benchmark_group("Bit-reverse permutation");
+
+    for order in 21..=24 {
+        let coeffs = gen_coeffs(order);
+        group.throughput(criterion::Throughput::Elements(1 << order)); // info for criterion
+
+        group.bench_with_input(
+            format!("Parallel bitrev permutation for 2^{order} elements"),
+            &coeffs,
+            |bench, coeffs| {
+                bench.iter(|| {
+                    // TODO: autoreleaspool hurts perf. by 2-3%. Search for an alternative
+                    objc::rc::autoreleasepool(|| {
+                        let metal_state = MetalState::new(None).unwrap();
+                        bitrev_permutation(coeffs, &metal_state).unwrap();
                     });
                 });
             },

gpu/benches/metal_benchmarks.rs

@@ -3,6 +3,7 @@ use criterion::{criterion_group, criterion_main, Criterion};
 mod benchmarks;
 
 fn run_metal_benchmarks(c: &mut Criterion) {
+    benchmarks::metal::metal_bitrev_permutation_benchmarks(c);
     benchmarks::metal::metal_fft_twiddles_benchmarks(c);
     benchmarks::metal::metal_fft_benchmarks(c);
 }