multithreading via Rayon

[Shnatsel]

• Extend this to f32
• Multi-thread the part that spans the two halves Doesn't seem to be beneficial, difficult to measure due to noise
• Feature-gate rayon dependency
• Do something about the regression for smaller sizes in the planner

[Shnatsel]

The initial numbers are promising, with over 3x faster DiT at size 16777216. The gains for smaller sizes are smaller, and the smallest sizes that use this codepath regress.

[codecov-commenter]

Codecov Report

✅ All modified and coverable lines are covered by tests.
✅ Project coverage is 99.82%. Comparing base (e3d75e5) to head (2d4130e).

Additional details and impacted files

@@            Coverage Diff             @@
##             main      #50      +/-   ##
==========================================
+ Coverage   99.29%   99.82%   +0.52%     
==========================================
  Files          12       13       +1     
  Lines        2277     2281       +4     
==========================================
+ Hits         2261     2277      +16     
+ Misses         16        4      -12     

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

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[Shnatsel]

With rayon thread pool also being used for COBRA instead of spawning brand new threads, sizes 131072 and above are 2x faster compared to main, and our largest size 16777216 is 2.5x faster. Admittedly my CPU has a lot of threads so the gains may be less pronounced on lower core counts.

@smu160 I'd appreciate benchmarks on ARM

[Shnatsel]

So, fun fact: disabling rayon feature and statically knowing that multi-threading for COBRA will not be used improves benchmarks for the smallest size considerably:

Forward f32/PhastFT DIT/64
                        time:   [156.70 ns 161.32 ns 164.25 ns]
                        thrpt:  [389.66 Melem/s 396.73 Melem/s 408.42 Melem/s]
                        thrpt:  [1.4516 GiB/s 1.4779 GiB/s 1.5215 GiB/s]
                 change:
                        time:   [−29.253% −26.600% −23.937%] (p = 0.00 < 0.05)
                        thrpt:  [+31.469% +36.240% +41.348%]
                        Performance has improved.

Not sure if it's a code layout artifact or the result of more inlining or what. We never had the chance to see these gains before because std::thread::scope codepath was compiled unconditionally.

I'm not going to do anything about it just yet, but it is something to keep in mind for the future.

[smu160]

Results on apple m2