ntt: two DRAM sweeps per encode, F192 on the F64 path, shift-XOR reduction

[tcoratger]

In one line

The WHIR encodes now make two passes over memory instead of three or four, so proving gets 10 to 22% faster across the four benchmarks, with bit-identical codewords.

The problem

A large NTT here is not limited by arithmetic. It is limited by memory bandwidth.

On a Ryzen 9 9950X3D, one read+write sweep of the 940 MB XMSS-900 codeword takes about 41 ms.
The base encode cost about 3.7 of those sweeps:

transpose            read message, write message       (plus a read before every write)
fused first pass     layers 1-3, gathered from the message
gathered pass        layers 4-8, radix 32 in L1
deep sub-blocks      layers 9-20

The F192 encodes of the deeper WHIR levels were worse off: a separate replicate pass, then three more sweeps.

The idea

Put as many layers as possible into each sweep, so each byte of the codeword travels to DRAM and back as few times as possible.

layers:  rate ............ deep_start ............ d
         |---- gathered pass ----|---- deep pass ----|

gathered pass:  one task per row group of 2^g rows, 2^(d - layer - g) apart,
                gathered into L2 scratch, all g layers run there, written back
deep pass:      one task per contiguous sub-block, its layers run in place

• A row group's rows are exactly the rows its g layers pair with each other, so in scratch they form a small transform of their own, with the same twiddles.
• For the base encode (56 lanes, 448 B a row), both units hold 2^10 rows, about 460 KB: one gathered pass and one deep pass cover all 20 layers.
• When every layer fits one L3-sized sub-block, the deep pass alone is the whole transform: one sweep.

Four changes

1. Two sweeps. The gathered first pass builds the replicas straight from the message: after r layers on the zero-padded message, every block is a copy of it. So no pass fills the replicas only to read them back.
2. F192 on the F64 path. A base-field twiddle scales the three coefficients of an F192 alike, and addition is coefficientwise. So a transform over n F192 lanes is exactly the F64 transform over 3n lanes. This deletes the dedicated F192 transform, its AVX-512 and NEON kernels, and the replicate pass in ligero_commit_ext.
3. Streaming stores where nothing reads the data back before it is evicted: the transpose tiles, the first pass's writes, and sub-blocks built in scratch. They skip the read an ordinary store makes before writing.
4. Shift-XOR reduction in the AVX-512 butterfly. Shifts and vpternlogq 0x96 replace two more carry-less multiplies per product, since Zen 5's carry-less multiplier is the scarce unit. On its own this takes the kernel from 0.277 to 0.165 ns per butterfly.

Tried and dropped: keeping all 8 rows of a radix-8 group in registers. It measured flat to slightly worse, because the kernel is bound by the multiplier and ALU, not by L1 loads.

Benchmarks

Zen 5 (Ryzen 9 9950X3D), rustc 1.98, target-cpu=native, 32 threads.
HEAD and this branch were run interleaved.

End to end, --repeat 3, two rounds each:

command	before	after	change
aggregate --xmss 900 --log-inv-rate 1	1.064 / 1.072 s	0.909 / 0.907 s	-15%, 846 -> 991 sig/s
aggregate --sphincs 220 --log-inv-rate 1	1.723 / 1.708 s	1.553 / 1.540 s	-10%, 128 -> 142 sig/s
recursion --n 2 --xmss-per-leaf 900 --log-inv-rate 2	0.570 / 0.593 s	0.454 / 0.455 s	-22%
aggregate --blobs 8 --log-inv-rate 1	1.282 / 1.257 s	1.095 / 1.094 s	-14%

The NTT stages alone, medians of 21 runs, three interleaved rounds, p10 to p90 within 2%:

stage	before	after	change
transpose + base encode, 2^21 x 56 F64	181 ms	103 ms	1.76x
(the transpose alone)	32.5 ms	19.7 ms	1.65x
F192 encode, 2^20 x 16 lanes, rate 1/16	75.3 ms	28.3 ms	2.66x
F192 encode, 2^19 x 16 lanes, rate 1/128	26.3 ms	6.2 ms	4.2x
butterfly kernel, L2-resident, 1 thread	0.277 ns	0.165 ns	1.68x

How close to the limit, measured pass by pass:

base, gathered pass   36.5 ms   reads 470 MB, streams 940 MB out   ~ the read ceiling
base, deep pass       44 ms     one read+write sweep of 1.88 GB     vs ~41 ms for a plain sweep
transpose             19.7 ms   940 MB moved                        ~ the read ceiling

Checks

• Proof sizes are unchanged in every benchmark, since the codewords are bit-identical.
• cargo testall passes, with and without ZK_ALLOC_POISON=1; clippy, rustdoc and fmt are clean.
• The pcs tests also pass on AVX2-only (x86-64-v3) and plain x86-64 builds.
• aarch64 type-checks, and a temporary compile_error! confirmed the NEON arm is compiled.
• The new tests force every plan the driver can pick. Wide lane counts shrink the cache budgets, which reaches two gathered passes and the streaming path cheaply. A deliberate mutation (visiting blocks in ascending order) fails four tests.

For the reviewer

• Apple silicon is unmeasured. The F192 encode now takes the F64 path on NEON too. It is the same butterfly kernel with fewer passes, so I expect a win, but please run the M4 numbers.
• Peak memory rises 40 to 80 MB, from per-thread scratch of up to 2 MB that lives as long as the thread.
• What remains is structural. Every pass is now near the bandwidth limit. The ~20 ms transpose could only go by changing the codeword layout across the whole PCS.

The benchmark harness behind the NTT numbers (not committed)

Append it to the end of crates/pcs/src/whir_ntt_ext.rs, then run:

cargo test --release -p pcs --lib ntt_throughput -- --ignored --nocapture

/// Throughput of the WHIR encodes at the XMSS-900 shapes.
///
/// - Run: `cargo test --release -p pcs --lib ntt_throughput -- --ignored --nocapture`.
/// - `NTT_SAMPLES` sets the timed runs per line, 21 by default.
#[cfg(test)]
mod throughput {
    use super::*;
    use crate::ntt::transpose_lane_major;
    use primitives::test_rng::Rng;
    use std::time::Instant;

    /// Seconds per run, sorted, after two warm-up runs that fault in pages and fill caches.
    fn time(samples: usize, mut f: impl FnMut()) -> Vec<f64> {
        f();
        f();
        let mut secs: Vec<f64> = (0..samples)
            .map(|_| {
                let t = Instant::now();
                f();
                t.elapsed().as_secs_f64()
            })
            .collect();
        secs.sort_by(f64::total_cmp);
        secs
    }

    /// Print the median and the 10th and 90th percentiles, scaled to the given unit.
    fn report(name: &str, secs: &[f64], scale: f64, unit: &str) {
        // Percentile p of the sorted samples.
        let q = |p: f64| secs[((secs.len() - 1) as f64 * p) as usize] * scale;
        println!(
            "{name:<40} median {:>8.3} {unit}   p10 {:>8.3}   p90 {:>8.3}",
            q(0.5),
            q(0.1),
            q(0.9)
        );
    }

    #[test]
    #[ignore = "manual release benchmark"]
    fn ntt_throughput() {
        let samples = primitives::bench::env_usize("NTT_SAMPLES", 21);
        let mut rng = Rng::new(7);

        // The butterfly kernel alone.
        //
        // 2^9 rows of 56 words is 229 KiB: L2-resident, and below the size that splits across threads.
        // So this line measures arithmetic, not memory.
        let (log_d, lanes) = (9usize, 56usize);
        let ntt = AdditiveNttF64::standard(log_d);
        let mut data: Vec<F64> = (0..lanes << log_d).map(|_| F64(rng.next_u64())).collect();
        let secs = time(samples, || {
            ntt.forward_transform_interleaved_parallel_from_layer(&mut data, lanes, 0)
        });
        // 9 layers, each 2^8 row pairs of 56 lanes.
        let butterflies = (log_d * (lanes << (log_d - 1))) as f64;
        report(
            "kernel      2^9 x 56 F64, 1 thread",
            &secs,
            1e9 / butterflies,
            "ns/butterfly",
        );

        // Base encode: 2^20 message rows, 56 lanes, rate 1/2, so a 940 MB codeword.
        let (log_rows, lanes, rate) = (20usize, 56usize, 1usize);
        let msg: Vec<F64> = (0..lanes << log_rows).map(|_| F64(rng.next_u64())).collect();
        let ntt = AdditiveNttF64::standard(log_rows + rate);
        let mut cw = vec![F64::ZERO; msg.len() << rate];
        // Bandwidth ceiling: one in-place read and write of the whole codeword.
        //
        // A memory-bound pass costs at least this, so the encode lines read as a sweep count.
        let secs = time(samples, || {
            parallel::chunks_mut(&mut cw, 1 << 16, |_, c| {
                c.iter_mut().for_each(|x| x.0 = x.0.rotate_left(1))
            })
        });
        report("ceiling     read+write sweep, 940 MB", &secs, 1e3, "ms");
        // The lane-major to row-major transpose alone.
        let secs = time(samples, || {
            transpose_lane_major(&mut cw[..msg.len()], &msg, lanes, log_rows)
        });
        report("transpose   2^20 x 56 F64", &secs, 1e3, "ms");
        // Transpose then encode: the whole base commit NTT, as the prover runs it.
        let secs = time(samples, || {
            transpose_lane_major(&mut cw[..msg.len()], &msg, lanes, log_rows);
            ntt.encode_interleaved_in_place(&mut cw, lanes, rate);
        });
        report("+ encode    2^21 x 56 F64, rate 1/2", &secs, 1e3, "ms");
        drop(cw);

        // Extension encodes: the first two WHIR levels, 16 F192 lanes, rates 1/16 and 1/128.
        for (log_block, rate) in [(20usize, 4usize), (19, 7)] {
            let ntt = AdditiveNttF64::standard(log_block);
            let msg = rng.ext_vec((16 << log_block) >> rate);
            let mut mat = vec![F192::ZERO; 16 << log_block];
            let secs = time(samples, || encode_interleaved_ext(&ntt, &mut mat, &msg, 16, rate));
            report(
                &format!("ext encode  2^{log_block} x 16 F192, rate 1/2^{rate}"),
                &secs,
                1e3,
                "ms",
            );
        }
    }
}

🤖 Generated with Claude Code

[tcoratger]

@TomWambsgans Please double check for the Mac adaption, I expect again the speedup not to be identical on Mac device :)

[TomWambsgans]

nice
I observe 10% to 15% perf improvement on my linux server, and slight slowdown on M4 Max, fixed by Claude now (keeping the gains for x86, the fix only affects aarch64)