Simd twiddles

[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

[codecov-commenter]

Codecov Report

All modified and coverable lines are covered by tests ✅

Project coverage is 92.15%. Comparing base (39763f5) to head (80f3bbb).

Additional details and impacted files

@@            Coverage Diff             @@
##              dev     #820      +/-   ##
==========================================
+ Coverage   91.97%   92.15%   +0.17%     
==========================================
  Files          90       90              
  Lines       12230    12306      +76     
  Branches    12230    12306      +76     
==========================================
+ Hits        11249    11340      +91     
+ Misses        875      860      -15     
  Partials      106      106              

Flag	Coverage Δ
	92.15% <100.00%> (+0.17%)	⬆️

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.

[alonh5]

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

---

crates/prover/src/core/backend/simd/circle.rs line 282 at r1 (raw file):

    #[allow(clippy::int_plus_one)]
    fn precompute_twiddles(mut coset: Coset) -> TwiddleTree<Self> {

Please add a unit test to compare with the CPU implementation.

---

crates/prover/src/core/backend/simd/circle.rs line 317 at r1 (raw file):

        }

        xs.push(PackedM31::from_array(extra.try_into().unwrap()));

Does this try into work because it's initialized with capacity N_LANES?

Code quote:

xs.push(PackedM31::from_array(extra.try_into().unwrap()));

---

crates/prover/src/core/backend/simd/circle.rs line 340 at r1 (raw file):

#[allow(clippy::int_plus_one)]
fn gen_coset_xs(coset: Coset, res: &mut Vec<PackedM31>) {

Can you add documentation? Maybe a reference to the algorithm

---

crates/prover/src/core/backend/simd/circle.rs line 344 at r1 (raw file):

    assert!(log_size >= LOG_N_LANES);

    let initial_points = std::array::from_fn(|i| coset.at(bit_reverse_index(i, log_size)));

Can this bit reverse be done for the whole coset at once?

Code quote:

let initial_points = std::array::from_fn(|i| coset.at(bit_reverse_index(i, log_size)));

[alonh5]

Reviewable status: 0 of 1 files reviewed, all discussions resolved (waiting on @Alon-Ti and @shaharsamocha7)

---

crates/prover/src/core/backend/simd/circle.rs line 282 at r1 (raw file):

Previously, alonh5 (Alon Haramati) wrote…

Please add a unit test to compare with the CPU implementation.

Done.

---

crates/prover/src/core/backend/simd/circle.rs line 340 at r1 (raw file):

Previously, alonh5 (Alon Haramati) wrote…

Can you add documentation? Maybe a reference to the algorithm

Done.

[shaharsamocha7]

Reviewed 2 of 3 files at r3, all commit messages.
Reviewable status: 2 of 3 files reviewed, 3 unresolved discussions (waiting on @Alon-Ti and @spapinistarkware)

---

crates/prover/src/core/backend/cpu/circle.rs line 199 at r3 (raw file):

        coset = coset.double();
    }
    twiddles.push(1.into());

Add a comment:
// Pad twiddles with Arbitrary value to a power of 2 vector.

Code quote:

twiddles.push(1.into());

---

crates/prover/src/core/backend/simd/circle.rs line 303 at r3 (raw file):

        twiddles.push(PackedM31::from_array(
            remaining_twiddles.try_into().unwrap(),
        ));

Doesn't this code fail if the coset is smaller than 16?

Code quote:

        let remaining_twiddles = slow_precompute_twiddles(coset);

        twiddles.push(PackedM31::from_array(
            remaining_twiddles.try_into().unwrap(),
        ));

---

crates/prover/src/core/backend/simd/circle.rs line 310 at r3 (raw file):

        let dbl_twiddles = twiddles
            .into_iter()
            .flat_map(|x| x.to_array().map(|x| x.0 * 2))

Mul by PackedM31::broadcast(2) instead of unpacking.
I think this unpacking is slow

Code quote:

 .flat_map(|x| x.to_array().map(|x| x.0 * 2))

[alonh5]

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

---

crates/prover/src/core/backend/cpu/circle.rs line 199 at r3 (raw file):

Previously, shaharsamocha7 wrote…

Add a comment:
// Pad twiddles with Arbitrary value to a power of 2 vector.

Done.

---

crates/prover/src/core/backend/simd/circle.rs line 303 at r3 (raw file):

Previously, shaharsamocha7 wrote…

Doesn't this code fail if the coset is smaller than 16?

It won't reach this if it's smaller than 16. It will go into the early return.

---

crates/prover/src/core/backend/simd/circle.rs line 310 at r3 (raw file):

Previously, shaharsamocha7 wrote…

Mul by PackedM31::broadcast(2) instead of unpacking.
I think this unpacking is slow

That's a problem because of the reduction, we want to multiply by 2 as a u32 not a m31.

[shaharsamocha7]

Reviewed 1 of 1 files at r4.
Reviewable status: 2 of 3 files reviewed, 1 unresolved discussion (waiting on @Alon-Ti and @alonh5)

---

crates/prover/src/core/backend/simd/circle.rs line 310 at r3 (raw file):

Previously, alonh5 (Alon Haramati) wrote…

That's a problem because of the reduction, we want to multiply by 2 as a u32 not a m31.

you have the into_simd() func, can you try using that?

[alonh5]

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

---

crates/prover/src/core/backend/simd/circle.rs line 310 at r3 (raw file):

Previously, shaharsamocha7 wrote…

you have the into_simd() func, can you try using that?

Done.

[shaharsamocha7]

Reviewed 1 of 1 files at r6, all commit messages.
Reviewable status: complete! all files reviewed, all discussions resolved (waiting on @Alon-Ti)

[alonh5]

Reviewable status: complete! all files reviewed, all discussions resolved (waiting on @Alon-Ti)

[alonh5]

Merge activity

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