Complex M31 field extension and NTT

[weijiekoh]

This draft PR adds a cm31_ntt member to ProveKit.

cm31_ntt is an implementation of the complex Mersenne-31 field extension (CM31) as well as the number-theoretic transform NTT algorithm for polynomials with CM31 coefficients. It is optimised for the ARM platform, and has been benchmarked on a Raspberry Pi 5.

RF and CF

CM31 builds upon Mersenne-31 field elements whose arithmetic is implemented using redundant representation. We repersent these redundant M31 field elements as 32-bit unsigned integers wrapped in the RF type defined in src/rm31.rs. Please refer to this note by Solbert and Domb to learn more about the algorithms used.

Each CM31 field element consists of a real RF part, and an imaginary RF part. The CF type and complex arithmetic operations are defined in src/cm31.rs.

NTT

ntt.rs contains an optimised implementation of the NTT algorithm. Use ntt() as such:

// Assume that f is a Vec of CM31 coefficients 
// and is at least 8 elements long.

// precomp contains precomputed twiddle factors.
// It can be computed in advance. It can be
// easily serialised and deserialised for easy
// storage and retrieval.
let precomp = precompute_twiddles(n).unwrap();
let res = ntt(&f, &precomp).unwrap();

Optimisations

The fist optimisation is simply to precompute twiddle factors.

The most important optimisation is to combine an approach that allocates memory with each recursive iteration with one that performs the NTT in-place.

The traditional divide-and-conquer approach allocates memory with each recursive iteration. It is very fast, but leaves some room for improvement. Take a look at ntt_r8_vec() and ntt_r8_vec_p() in ntt.rs. They are straightforward implementations of the divide-and-conquer algorithm (the former does not precompute twiddle factors, and the latter does), but they allocate new Vecs with each recursive implementation, resulting in some performance overhead.

An in-place algorithm (ntt_r8_ip and ntt_r8_ip_p) which avoids memory allocation altogether, however, is extremely slow, especially past 32768 elements. This is likely due to the small CPU cache space leading to costly cache misses.

We had a breakthrough when we found that for NTTs over sizes lower than 262144, ntt_r8_vec_p() is slower than its in-place counterpart ntt_r8_ip_p. This led us to develop our most efficient algorithm by combining the two approaches. It is implemented in ntt_r8_hybrid_p() which uses the divide-and-conquer approach via recursion, but when the NTT size is NTT_BLOCK_SIZE_FOR_CACHE (hardcoded to 32768), it uses the in-place algorithm.

Benchmark results can be found in README.md.

[recmo]

LGTM, but could use some help making it more canonically Rust (which also helps performance).
@Dzejkop Can you help here?

[recmo]

I'm surprised we are not exploiting the special structure here. Does the compiler turn this into bitshifts?

[weijiekoh]

Resolved in commit 496dbce, though benchmarks don't show a significant improvement.

[recmo]

@xrvdg We will want to aggressively optimize this function.

[weijiekoh]

Thanks for the feedback and suggestions! I'll incorporate them into a new commit.

[weijiekoh]

Thank you @Dzejkop ! Your suggestions are super helpful. Things look much cleaner and consistent now!

[recmo]

Thank you @Dzejkop and @weijiekoh. This looks good to merge (we can always follow up with further PRs).