Parallel fft

[spapinistarkware]

This change is 

[spapinistarkware]

• Parallelize poseidon trace gen #822
• Parallelize domain eval #821
• Simd twiddles #820
• Parallel fft #819 👈
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @spapinistarkware and the rest of your teammates on Graphite

[spapinistarkware]

• Simd twiddles #820
• Parallel fft #819 👈
• dev

This stack of pull requests is managed by Graphite. Learn more about stacking.

Join @spapinistarkware and the rest of your teammates on Graphite

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.06%. Comparing base (3f93ed2) to head (f6810fa).
Report is 7 commits behind head on dev.

Additional details and impacted files

@@            Coverage Diff             @@
##              dev     #819      +/-   ##
==========================================
+ Coverage   92.01%   92.06%   +0.05%     
==========================================
  Files          90       90              
  Lines       12206    12254      +48     
  Branches    12206    12254      +48     
==========================================
+ Hits        11231    11282      +51     
+ Misses        869      865       -4     
- Partials      106      107       +1     

Flag	Coverage Δ
	92.06% <100.00%> (+0.05%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[andrewmilson]

Reviewable status: 0 of 3 files reviewed, 2 unresolved discussions (waiting on @Alon-Ti, @shaharsamocha7, and @spapinistarkware)

---

crates/prover/src/core/backend/simd/fft/ifft.rs line 157 at r1 (raw file):

    let values = UnsafeMutI32(values);
    iter.for_each(|index_h| {

Since this is so common might be worth creating macros for them like in https://github.com/arkworks-rs/std/blob/master/src/lib.rs.

Suggestion:

    iter!(0..1 << (log_size - fft_layers - LOG_N_LANES as usize)).for_each(|index_h| {

---

crates/prover/src/core/backend/simd/fft/mod.rs line 16 at r1 (raw file):

pub const MIN_FFT_LOG_SIZE: u32 = 5;

pub struct UnsafeMutI32(pub *mut u32);

I'm a little worried the quantity of unsafe code is getting out of hand in the SimdBackend and that there are as performant safe ways of doing the implementation. Haven't benchmarked with "safe" changes yet though just concerned these changes add more unsafe implementation

[alonh5]

Reviewable status: 0 of 5 files reviewed, 2 unresolved discussions (waiting on @Alon-Ti, @andrewmilson, @shaharsamocha7, and @spapinistarkware)

---

crates/prover/src/core/backend/simd/fft/ifft.rs line 157 at r1 (raw file):

Previously, andrewmilson (Andrew Milson) wrote…

Since this is so common might be worth creating macros for them like in https://github.com/arkworks-rs/std/blob/master/src/lib.rs.

Done.

[alonh5]

Reviewable status: 0 of 5 files reviewed, 2 unresolved discussions (waiting on @Alon-Ti, @andrewmilson, and @shaharsamocha7)

---

crates/prover/src/core/backend/simd/fft/mod.rs line 16 at r1 (raw file):

Previously, andrewmilson (Andrew Milson) wrote…

I'm a little worried the quantity of unsafe code is getting out of hand in the SimdBackend and that there are as performant safe ways of doing the implementation. Haven't benchmarked with "safe" changes yet though just concerned these changes add more unsafe implementation

I see the unsafe usage in this PR isn't new, just using old unsafe functions. Do you want to re-exmine those?

[andrewmilson]

Reviewed 1 of 5 files at r2, all commit messages.
Reviewable status: 1 of 5 files reviewed, 7 unresolved discussions (waiting on @Alon-Ti, @alonh5, @shaharsamocha7, and @spapinistarkware)

---

crates/prover/src/core/mod.rs line 65 at r2 (raw file):

    ($i: expr) => {{
        #[cfg(not(feature = "parallel"))]
        let iter = $i;

Suggestion:

let iter = $i.into_iter();

---

crates/prover/src/core/backend/simd/blake2s.rs line 57 at r2 (raw file):

            let iter = parallel_iter!(0..1 << log_size);

            return iter

Suggestion:

        return parallel_iter!(0..1 << log_size)

---

crates/prover/src/core/backend/simd/fft/ifft.rs line 157 at r1 (raw file):

Previously, alonh5 (Alon Haramati) wrote…

Done.

Allg since it's already been added but forgot to mention that I really like the way plonky3 handles this kind of functionality with their
https://github.com/Plonky3/Plonky3/tree/main/maybe-rayon crate. I think it actually has some benefits over the macro approach. For instance don't have to conditionally compile any preludes. Also cleaner syntax IMO.

---

crates/prover/src/core/backend/simd/fft/ifft.rs line 92 at r2 (raw file):

    let values = UnsafeMutI32(values);
    iter.for_each(|index_h| {

Nit:

Suggestion:

    let values = UnsafeMutI32(values);
    parallel_iter!(0..1 << (log_size - fft_layers)).for_each(|index_h| {

---

crates/prover/src/core/backend/simd/fft/mod.rs line 16 at r1 (raw file):

Previously, alonh5 (Alon Haramati) wrote…

I see the unsafe usage in this PR isn't new, just using old unsafe functions. Do you want to re-exmine those?

Yeah, maybe you can add this todo for me plz

// TODO(andrew): Examine usage of unsafe in SIMD FFT.

---

crates/prover/src/core/backend/simd/fft/rfft.rs line 97 at r2 (raw file):

    let src = UnsafeConstI32(src);
    let dst = UnsafeMutI32(dst);
    iter.for_each(|index_h| {

Same here

Code quote:

    let iter = parallel_iter!(0..1 << (log_size - fft_layers));

    let src = UnsafeConstI32(src);
    let dst = UnsafeMutI32(dst);
    iter.for_each(|index_h| {

---

crates/prover/src/core/backend/simd/fft/rfft.rs line 163 at r2 (raw file):

    let src = UnsafeConstI32(src);
    let dst = UnsafeMutI32(dst);
    iter.for_each(|index_h| {

Same here

Code quote:

    let iter = parallel_iter!(0..1 << (log_size - fft_layers - LOG_N_LANES as usize));

    let src = UnsafeConstI32(src);
    let dst = UnsafeMutI32(dst);
    iter.for_each(|index_h| {

[andrewmilson]

with comments

Reviewable status: 1 of 5 files reviewed, 7 unresolved discussions (waiting on @Alon-Ti, @alonh5, @shaharsamocha7, and @spapinistarkware)

[andrewmilson]

Reviewable status: 0 of 5 files reviewed, 1 unresolved discussion (waiting on @Alon-Ti, @shaharsamocha7, and @spapinistarkware)

[alonh5]

Merge activity

• Sep 24, 9:20 AM EDT: @alonh5 started a stack merge that includes this pull request via Graphite.
• Sep 24, 9:20 AM EDT: @alonh5 merged this pull request with Graphite.