From db1a50c2122d081969d192f696e2c864438f7007 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Thu, 27 Aug 2026 10:17:08 +0200 Subject: [PATCH 1/2] SME2 streaming-mode BLAKE2s, on one worker per cluster CREDIT: https://github.com/zooko/blake3-sme2 M4 and later expose 512-bit streaming SVE2 behind SMSTART, so sixteen BLAKE2s lanes fit one register and `xar` fuses each xor with its rotation: a G function is ten instructions instead of sixteen. ZA carries both transposes, the message in as rows and out as columns, the digests out the same way. NEON is illegal inside streaming mode, so everything from the loads to the digest stores is hand-written. The width itself buys nothing. A streaming vector is four times a NEON one and issues at a quarter of the rate, and the two cancel: 251 GB/s of operand width per thread against 276 for one NEON core. What is left is the instruction saving, about 1.7x per thread. The block is shared by a whole core cluster and one thread saturates it, so extra streaming workers only take cores away from NEON: fmopa measures 3.94 G/s on one thread and 7.98 G/s on twelve, exactly the two performance clusters. Hence one slot per block, workers 0 and 1 plus the first efficiency worker, and NEON everywhere else. LEANVM_SME_WORKERS tunes it, 0 disabling the backend. A batch under sixteen inputs, and every remainder, still goes to NEON. Measured on an M4 Max, 2026-08-27, medians of interleaved runs: hash_bench multithreaded_throughput 441 -> 492 Mhash/s (+11.6%) aggregate --xmss 900 --log-inv-rate 1 1062 -> 1083 sig/s (+2.0%) The end-to-end gain is small because Merkle hashing is 14.5% of proving time and that stage is memory-bandwidth bound: it takes 0.198 s on four threads against 0.128 s on eleven, so it was never short of arithmetic. Same runs confirmed the pool defaults are already at their optimum, 11 performance workers against 1036 sig/s for 12 and 985 for 13, and all four efficiency workers against 995 with none. Also corrects the AGENTS.md claim that the M4 has no SVE. It has, in streaming mode, but with no 64-bit polynomial multiply: `pmullb z.q` faults there, which is what rules the field arithmetic out. Co-Authored-By: Claude Opus 5 (1M context) --- AGENTS.md | 2 +- crates/parallel/src/lib.rs | 9 + crates/primitives/src/hash.rs | 34 +++ crates/primitives/src/hash/blake2s_sme2.s | 311 ++++++++++++++++++++++ crates/primitives/src/hash/sme2.rs | 151 +++++++++++ 5 files changed, 506 insertions(+), 1 deletion(-) create mode 100644 crates/primitives/src/hash/blake2s_sme2.s create mode 100644 crates/primitives/src/hash/sme2.rs diff --git a/AGENTS.md b/AGENTS.md index 9a2c5ada..0dac0906 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -105,7 +105,7 @@ Understand the third before changing the verifier. `guests/aggregate.py` is zkDS ## Conventions that bite - **The prover is memory-bandwidth bound above four cores.** Doubling four cores to eight buys well under two, and `Commit` is slower on sixteen threads than on eight. What pays there is deleting traffic, not instructions. `primitives::stream::Stream` publishes a buffer without the read-for-ownership an ordinary store pays, but ONLY where nothing reads the destination again before it is evicted. Where a consumer follows in the same pass, the fetch it avoids becomes that consumer's miss: fold kernels earn it by building their round message from registers, or by folding into an L1 stage first (`whir::fold_and_msg_blocks`). That fetch is an x86 cost only: on Apple silicon a store-only fill already sustains what a read-only pass does and `STNP` measures identical to `STP`, so `Stream` is a plain copy there and the L1 stage earns its keep for the read locality alone, which is still better than writing through. -- **NEON is the width ceiling on Apple silicon**, so an AVX-512 win that is purely width has no counterpart: the M4 has no SVE, and its SME2 is streaming-mode matrix work with no polynomial multiply. What does port is *shape*. A fused NTT pass wants a butterfly at a time over whole rows, not the register-resident tile the AVX-512 arms use: they transpose anyway and want to pay for it once per pass, while NEON transposes nothing and a tile leaves only its own width of independent work to cover the reduction's dependent PMULL folds, where a row leaves the whole lane count. Measured both directions: the tile costs the extension NTT, and costs the base encode's `Commit` again. +- **NEON is the width ceiling on Apple silicon**, so an AVX-512 win that is purely width has no counterpart. `SMSTART` does give 512-bit streaming SVE2 on M4 and later, but the SME block is shared by a whole core cluster and one thread saturates it, so it is faster per thread and slower per machine, and the streaming subset has no 64-bit polynomial multiply (`pmullb z.q` faults there), which rules out the field arithmetic. `hash::sme2` uses it anyway for BLAKE2s, on one worker per cluster and NEON everywhere else, which is the only shape that wins; the gain is `xar` fusing each xor with its rotation, not the width. What does port is *shape*. A fused NTT pass wants a butterfly at a time over whole rows, not the register-resident tile the AVX-512 arms use: they transpose anyway and want to pay for it once per pass, while NEON transposes nothing and a tile leaves only its own width of independent work to cover the reduction's dependent PMULL folds, where a row leaves the whole lane count. Measured both directions: the tile costs the extension NTT, and costs the base encode's `Commit` again. - **A `[F192; N]` in a NEON kernel is a memory object, where on AVX-512 it is the register.** Four tower products are four independent PMULL chains wanting most of the 32 vector registers, so an array of them spills and the spill costs more than batching the products saves; the same array is free on AVX-512, where the quad IS one register. Keep the quad as a tuple or as named values and let arrays exist only inside the batched-product helper, on the target that wants them (`flock::zerocheck::multilinear`'s `mul_quad`). The symptom is indirect, so suspect the shape rather than the arithmetic: the products measure the same either way, destructuring the results changes nothing, and forcing the helper to inline recovers almost none of it. - **On Zen 4, 512-bit cross-lane data movement is half-rate** (every 512-bit shuffle is two 256-bit uops), so packing scalars into vector lanes with `vpermi2q`/`vpermq` and extracting with `vextracti64x4` loses to the scalar moves it replaces. Widening the arithmetic still pays: `mul4` beats the same products issued one at a time. Prefer kernels where both qwords of every 128-bit lane carry a product and nothing crosses lanes. - Use comments only when necessary: uncommented but readable and simple code is better than commented slop. And when you use comments, be concise. diff --git a/crates/parallel/src/lib.rs b/crates/parallel/src/lib.rs index fb712065..e5bec73b 100644 --- a/crates/parallel/src/lib.rs +++ b/crates/parallel/src/lib.rs @@ -150,6 +150,15 @@ fn pool() -> &'static Pool { }) } +/// Which worker the calling thread is: `0` for the dispatcher, `1..perf` for the +/// performance workers, and `perf..` for the efficiency ones. Any thread +/// outside the pool reads as `0`. +#[must_use] +#[inline] +pub fn worker_id() -> usize { + WORKER_ID.get() +} + fn worker_main(pool: &'static Pool, id: usize, qos: Qos) { WORKER_ID.set(id); set_qos(qos); diff --git a/crates/primitives/src/hash.rs b/crates/primitives/src/hash.rs index e8ea49d4..7f2acfcd 100644 --- a/crates/primitives/src/hash.rs +++ b/crates/primitives/src/hash.rs @@ -639,6 +639,10 @@ mod x86 { } } +/// The streaming-mode backend, for the workers that can reach an SME block. +#[cfg(all(target_arch = "aarch64", target_os = "macos"))] +mod sme2; + #[cfg(target_arch = "aarch64")] mod arm { use super::{Lanes32, OUT_LEN}; @@ -1156,6 +1160,13 @@ pub fn hash_many_dyn_from_state(data: &[u8], len: usize, state: &[u32; 8], t_off unsafe { hash_many_with::(data, len, state, t_offset, out) } + #[cfg(all(target_arch = "aarch64", target_os = "macos"))] + if sme2::enabled() { + // SAFETY: the asserts above, and `sme2::enabled` has already checked + // that this machine runs the kernel at sixteen lanes. + unsafe { sme2::hash_many(data, len, state, t_offset, out) }; + return; + } #[cfg(target_arch = "aarch64")] unsafe { hash_many_with::(data, len, state, t_offset, out) @@ -1285,6 +1296,29 @@ mod tests { } } + /// The streaming backend is the one place where the whole loop, transposes + /// and stores included, is hand-written, so it gets its own comparison + /// against the scalar reference across lane counts and block counts. + #[cfg(all(target_arch = "aarch64", target_os = "macos"))] + #[test] + fn streaming_backend_matches_scalar() { + if !sme2::available() { + return; + } + for len in [64usize, 128, 320] { + for n in [16usize, 17, 33, 48] { + let data: Vec = (0..n * len).map(|i| (i * 31 + 7) as u8).collect(); + let mut got = vec![0u8; n * OUT_LEN]; + // SAFETY: `data` is `n * len` bytes and `len` is a multiple of 64. + unsafe { sme2::hash_many(&data, len, &PARAM_IV, 0, &mut got) }; + for i in 0..n { + let want = hash(&data[i * len..(i + 1) * len]); + assert_eq!(&got[i * OUT_LEN..(i + 1) * OUT_LEN], &want[..], "len {len}, input {i}"); + } + } + } + } + /// Every SIMD backend compiled into this build agrees with the scalar /// hash, not just the one the dispatch picks. Both the round arithmetic and /// the transpose network are per-backend, so this is what keeps an untaken diff --git a/crates/primitives/src/hash/blake2s_sme2.s b/crates/primitives/src/hash/blake2s_sme2.s new file mode 100644 index 00000000..5f3651b9 --- /dev/null +++ b/crates/primitives/src/hash/blake2s_sme2.s @@ -0,0 +1,311 @@ +// BLAKE2s-256 over sixteen inputs at once, in SME2 streaming mode. +// +// Sixteen 32-bit lanes fill one 512-bit streaming vector, so a whole batch is +// one register per state word. Against the NEON backend this trades three of +// the four issue slots for four times the width, which cancels; the win is +// `xar`, which fuses each xor with its rotation, so a G function is ten +// instructions instead of sixteen. ZA carries the two transposes: the message +// goes in as sixteen rows and comes out as sixteen columns, and the digests +// leave the same way. +// +// NEON is illegal inside streaming mode, so everything from the loads to the +// digest stores lives here. +// +// uint64_t blake2s_hash16_sme2( +// const uint8_t *const *inputs, // x0, sixteen pointers +// const uint32_t *state, // x1, the chaining value, splatted +// uint64_t t_offset, // x2 +// uint64_t len, // x3, bytes per input, a multiple of 64 +// uint8_t *out, // x4, sixteen 32-byte digests +// const uint32_t *iv); // x5, the eight BLAKE2s IV words +// +// Returns the streaming vector length in 32-bit lanes. The caller must treat +// anything other than 16 as "nothing was written". +// +// ZA use: +// ZA0.S the message block, rows in, columns out +// ZA1.S the digest transpose +// ZA2.S the chaining value across blocks, rows 0..7 + +.text +.arch armv9-a+sme2 +.p2align 6 + +// One input's 64-byte block into a ZA row. Rows 0..15 are addressed as a +// selector holding 0, 4, 8 or 12 plus a slice of 0..3. +.macro LOAD_ROW ptr_off, selector, slice + ldr x8, [x0, #\ptr_off] + add x8, x8, x10 + ld1w { za0h.s[\selector, \slice] }, p0/z, [x8] +.endm + +.macro COLUMN_G4 mx0, my0, mx1, my1, mx2, my2, mx3, my3 + add z0.s, z0.s, z4.s + add z1.s, z1.s, z5.s + add z2.s, z2.s, z6.s + add z3.s, z3.s, z7.s + + add z0.s, z0.s, \mx0 + add z1.s, z1.s, \mx1 + add z2.s, z2.s, \mx2 + add z3.s, z3.s, \mx3 + + xar z12.s, z12.s, z0.s, #16 + xar z13.s, z13.s, z1.s, #16 + xar z14.s, z14.s, z2.s, #16 + xar z15.s, z15.s, z3.s, #16 + + add z8.s, z8.s, z12.s + add z9.s, z9.s, z13.s + add z10.s, z10.s, z14.s + add z11.s, z11.s, z15.s + + xar z4.s, z4.s, z8.s, #12 + xar z5.s, z5.s, z9.s, #12 + xar z6.s, z6.s, z10.s, #12 + xar z7.s, z7.s, z11.s, #12 + + add z0.s, z0.s, z4.s + add z1.s, z1.s, z5.s + add z2.s, z2.s, z6.s + add z3.s, z3.s, z7.s + + add z0.s, z0.s, \my0 + add z1.s, z1.s, \my1 + add z2.s, z2.s, \my2 + add z3.s, z3.s, \my3 + + xar z12.s, z12.s, z0.s, #8 + xar z13.s, z13.s, z1.s, #8 + xar z14.s, z14.s, z2.s, #8 + xar z15.s, z15.s, z3.s, #8 + + add z8.s, z8.s, z12.s + add z9.s, z9.s, z13.s + add z10.s, z10.s, z14.s + add z11.s, z11.s, z15.s + + xar z4.s, z4.s, z8.s, #7 + xar z5.s, z5.s, z9.s, #7 + xar z6.s, z6.s, z10.s, #7 + xar z7.s, z7.s, z11.s, #7 + +.endm + +.macro DIAG_G4 mx0, my0, mx1, my1, mx2, my2, mx3, my3 + add z0.s, z0.s, z5.s + add z1.s, z1.s, z6.s + add z2.s, z2.s, z7.s + add z3.s, z3.s, z4.s + + add z0.s, z0.s, \mx0 + add z1.s, z1.s, \mx1 + add z2.s, z2.s, \mx2 + add z3.s, z3.s, \mx3 + + xar z15.s, z15.s, z0.s, #16 + xar z12.s, z12.s, z1.s, #16 + xar z13.s, z13.s, z2.s, #16 + xar z14.s, z14.s, z3.s, #16 + + add z10.s, z10.s, z15.s + add z11.s, z11.s, z12.s + add z8.s, z8.s, z13.s + add z9.s, z9.s, z14.s + + xar z5.s, z5.s, z10.s, #12 + xar z6.s, z6.s, z11.s, #12 + xar z7.s, z7.s, z8.s, #12 + xar z4.s, z4.s, z9.s, #12 + + add z0.s, z0.s, z5.s + add z1.s, z1.s, z6.s + add z2.s, z2.s, z7.s + add z3.s, z3.s, z4.s + + add z0.s, z0.s, \my0 + add z1.s, z1.s, \my1 + add z2.s, z2.s, \my2 + add z3.s, z3.s, \my3 + + xar z15.s, z15.s, z0.s, #8 + xar z12.s, z12.s, z1.s, #8 + xar z13.s, z13.s, z2.s, #8 + xar z14.s, z14.s, z3.s, #8 + + add z10.s, z10.s, z15.s + add z11.s, z11.s, z12.s + add z8.s, z8.s, z13.s + add z9.s, z9.s, z14.s + + xar z5.s, z5.s, z10.s, #7 + xar z6.s, z6.s, z11.s, #7 + xar z7.s, z7.s, z8.s, #7 + xar z4.s, z4.s, z9.s, #7 + +.endm + +.macro ROUND s0, s1, s2, s3, s4, s5, s6, s7, s8, s9, s10, s11, s12, s13, s14, s15 + COLUMN_G4 \s0, \s1, \s2, \s3, \s4, \s5, \s6, \s7 + DIAG_G4 \s8, \s9, \s10, \s11, \s12, \s13, \s14, \s15 +.endm + +.globl _blake2s_hash16_sme2 +_blake2s_hash16_sme2: + // Entering streaming mode zeroes the vector registers, d8-d15 included. + stp d8, d9, [sp, #-64]! + stp d10, d11, [sp, #16] + stp d12, d13, [sp, #32] + stp d14, d15, [sp, #48] + + smstart + cntw x9 + cmp x9, #16 + b.ne Lwrong_vl + + ptrue p0.s // all sixteen lanes + ptrue p1.s, vl8 // the eight words of one digest + mov w12, #0 + mov w13, #4 + mov w14, #8 + mov w15, #12 + + // h starts as the caller's chaining value, one splat per word, in ZA2. + ld1rw { z0.s }, p0/z, [x1] + ld1rw { z1.s }, p0/z, [x1, #4] + ld1rw { z2.s }, p0/z, [x1, #8] + ld1rw { z3.s }, p0/z, [x1, #12] + ld1rw { z4.s }, p0/z, [x1, #16] + ld1rw { z5.s }, p0/z, [x1, #20] + ld1rw { z6.s }, p0/z, [x1, #24] + ld1rw { z7.s }, p0/z, [x1, #28] + mov za2h.s[w12, 0:3], { z0.s - z3.s } + mov za2h.s[w13, 0:3], { z4.s - z7.s } + + lsr x6, x3, #6 // block count + mov x7, #0 // block index + mov x10, #0 // byte offset into every input + mov x11, x2 // running byte counter t + +Lblock: + LOAD_ROW 0, w12, 0 + LOAD_ROW 8, w12, 1 + LOAD_ROW 16, w12, 2 + LOAD_ROW 24, w12, 3 + LOAD_ROW 32, w13, 0 + LOAD_ROW 40, w13, 1 + LOAD_ROW 48, w13, 2 + LOAD_ROW 56, w13, 3 + LOAD_ROW 64, w14, 0 + LOAD_ROW 72, w14, 1 + LOAD_ROW 80, w14, 2 + LOAD_ROW 88, w14, 3 + LOAD_ROW 96, w15, 0 + LOAD_ROW 104, w15, 1 + LOAD_ROW 112, w15, 2 + LOAD_ROW 120, w15, 3 + + // Vertical reads transpose: z(16+w) holds word w of every lane. + mov { z16.s - z19.s }, za0v.s[w12, 0:3] + mov { z20.s - z23.s }, za0v.s[w13, 0:3] + mov { z24.s - z27.s }, za0v.s[w14, 0:3] + mov { z28.s - z31.s }, za0v.s[w15, 0:3] + + mov { z0.s - z3.s }, za2h.s[w12, 0:3] + mov { z4.s - z7.s }, za2h.s[w13, 0:3] + ld1rw { z8.s }, p0/z, [x5] + ld1rw { z9.s }, p0/z, [x5, #4] + ld1rw { z10.s }, p0/z, [x5, #8] + ld1rw { z11.s }, p0/z, [x5, #12] + + add x11, x11, #64 // t counts bytes fed so far + ldr w8, [x5, #16] + eor w8, w8, w11 + dup z12.s, w8 + lsr x16, x11, #32 + ldr w17, [x5, #20] + eor w8, w17, w16 + dup z13.s, w8 + add x16, x7, #1 + cmp x16, x6 + csetm w17, eq // the last block inverts IV[6] + ldr w8, [x5, #24] + eor w8, w8, w17 + dup z14.s, w8 + ld1rw { z15.s }, p0/z, [x5, #28] + + ROUND z16.s, z17.s, z18.s, z19.s, z20.s, z21.s, z22.s, z23.s, z24.s, z25.s, z26.s, z27.s, z28.s, z29.s, z30.s, z31.s + ROUND z30.s, z26.s, z20.s, z24.s, z25.s, z31.s, z29.s, z22.s, z17.s, z28.s, z16.s, z18.s, z27.s, z23.s, z21.s, z19.s + ROUND z27.s, z24.s, z28.s, z16.s, z21.s, z18.s, z31.s, z29.s, z26.s, z30.s, z19.s, z22.s, z23.s, z17.s, z25.s, z20.s + ROUND z23.s, z25.s, z19.s, z17.s, z29.s, z28.s, z27.s, z30.s, z18.s, z22.s, z21.s, z26.s, z20.s, z16.s, z31.s, z24.s + ROUND z25.s, z16.s, z21.s, z23.s, z18.s, z20.s, z26.s, z31.s, z30.s, z17.s, z27.s, z28.s, z22.s, z24.s, z19.s, z29.s + ROUND z18.s, z28.s, z22.s, z26.s, z16.s, z27.s, z24.s, z19.s, z20.s, z29.s, z23.s, z21.s, z31.s, z30.s, z17.s, z25.s + ROUND z28.s, z21.s, z17.s, z31.s, z30.s, z29.s, z20.s, z26.s, z16.s, z23.s, z22.s, z19.s, z25.s, z18.s, z24.s, z27.s + ROUND z29.s, z27.s, z23.s, z30.s, z28.s, z17.s, z19.s, z25.s, z21.s, z16.s, z31.s, z20.s, z24.s, z22.s, z18.s, z26.s + ROUND z22.s, z31.s, z30.s, z25.s, z27.s, z19.s, z16.s, z24.s, z28.s, z18.s, z29.s, z23.s, z17.s, z20.s, z26.s, z21.s + ROUND z26.s, z18.s, z24.s, z20.s, z23.s, z22.s, z17.s, z21.s, z31.s, z27.s, z25.s, z30.s, z19.s, z28.s, z29.s, z16.s + + // h ^= v[i] ^ v[i + 8] + mov { z16.s - z19.s }, za2h.s[w12, 0:3] + mov { z20.s - z23.s }, za2h.s[w13, 0:3] + eor3 z16.d, z16.d, z0.d, z8.d + eor3 z17.d, z17.d, z1.d, z9.d + eor3 z18.d, z18.d, z2.d, z10.d + eor3 z19.d, z19.d, z3.d, z11.d + eor3 z20.d, z20.d, z4.d, z12.d + eor3 z21.d, z21.d, z5.d, z13.d + eor3 z22.d, z22.d, z6.d, z14.d + eor3 z23.d, z23.d, z7.d, z15.d + mov za2h.s[w12, 0:3], { z16.s - z19.s } + mov za2h.s[w13, 0:3], { z20.s - z23.s } + + add x10, x10, #64 + add x7, x7, #1 + cmp x7, x6 + b.lo Lblock + + // Writing h as columns of ZA1 makes row l lane l's eight digest words. + mov { z0.s - z3.s }, za2h.s[w12, 0:3] + mov { z4.s - z7.s }, za2h.s[w13, 0:3] + mov za1v.s[w12, 0:3], { z0.s - z3.s } + mov za1v.s[w13, 0:3], { z4.s - z7.s } + st1w { za1h.s[w12, 0] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w12, 1] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w12, 2] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w12, 3] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w13, 0] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w13, 1] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w13, 2] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w13, 3] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w14, 0] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w14, 1] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w14, 2] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w14, 3] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w15, 0] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w15, 1] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w15, 2] }, p1, [x4] + add x4, x4, #32 + st1w { za1h.s[w15, 3] }, p1, [x4] + +Lwrong_vl: + mov x0, x9 + smstop + ldp d10, d11, [sp, #16] + ldp d12, d13, [sp, #32] + ldp d14, d15, [sp, #48] + ldp d8, d9, [sp], #64 + ret diff --git a/crates/primitives/src/hash/sme2.rs b/crates/primitives/src/hash/sme2.rs new file mode 100644 index 00000000..f4505714 --- /dev/null +++ b/crates/primitives/src/hash/sme2.rs @@ -0,0 +1,151 @@ +//! BLAKE2s over sixteen inputs at once, in SME2 streaming mode. +//! +//! A streaming vector is 512 bits, so one register holds one state word across +//! sixteen lanes, and `xar` fuses each xor with its rotation. The catch is that +//! the SME block is shared by a whole core cluster and one thread saturates it, +//! so this is faster per thread and slower per machine: [`enabled`] hands it to +//! the few workers that can reach a block of their own and leaves everyone else +//! on NEON. + +use std::sync::OnceLock; + +use super::{BLOCK_LEN, IV, OUT_LEN, PARAM_IV, arm::Neon, hash_many_with}; + +std::arch::global_asm!(include_str!("blake2s_sme2.s"), options(raw)); + +unsafe extern "C" { + /// Returns the streaming vector length in 32-bit lanes; anything but 16 + /// means it wrote nothing. + fn blake2s_hash16_sme2( + inputs: *const *const u8, + state: *const u32, + t_offset: u64, + len: u64, + out: *mut u8, + iv: *const u32, + ) -> u64; +} + +/// Inputs per call: the 32-bit lanes of one streaming vector. +pub(super) const LANES: usize = 16; + +/// Read an integer `sysctl`, `None` if it does not exist. +fn sysctl(name: &core::ffi::CStr) -> Option { + let mut value: i32 = 0; + let mut len = core::mem::size_of::(); + // SAFETY: a read-only sysctl with a correctly sized destination and a null + // new-value pointer. + let rc = unsafe { + libc::sysctlbyname( + name.as_ptr(), + (&raw mut value).cast(), + &raw mut len, + std::ptr::null_mut(), + 0, + ) + }; + (rc == 0).then_some(value) +} + +/// Whether this machine has the kernel's vector length and agrees with the +/// scalar backend on one block. +fn probe() -> bool { + if sysctl(c"hw.optional.arm.FEAT_SME2") != Some(1) { + return false; + } + let block = [0u8; BLOCK_LEN]; + let inputs = [block.as_ptr(); LANES]; + let mut out = [0u8; LANES * OUT_LEN]; + // SAFETY: sixteen pointers to a whole readable block, and room for sixteen + // digests. + let lanes = unsafe { + blake2s_hash16_sme2( + inputs.as_ptr(), + PARAM_IV.as_ptr(), + 0, + BLOCK_LEN as u64, + out.as_mut_ptr(), + IV.as_ptr(), + ) + }; + lanes == LANES as u64 && out.chunks_exact(OUT_LEN).all(|d| d == super::hash(&block)) +} + +/// How many workers take the streaming path, at most one per cluster. +/// +/// A second thread on a cluster splits one block rather than finding another, +/// so past three this only takes cores away from NEON. `LEANVM_SME_WORKERS` +/// overrides it, `0` disabling the backend. +fn streaming_workers() -> usize { + static N: OnceLock = OnceLock::new(); + *N.get_or_init(|| { + std::env::var("LEANVM_SME_WORKERS") + .ok() + .and_then(|v| v.parse().ok()) + .unwrap_or(SLOTS) + .min(SLOTS) + }) +} + +/// One streaming slot per SME block: two performance clusters and the +/// efficiency one. +const SLOTS: usize = 3; + +/// Whether the calling worker should take the streaming path. +/// +/// Workers 0 and 1 are dispatcher and first performance worker, which the +/// scheduler tends to place on different performance clusters, and worker +/// `perf` is the first efficiency one, which reaches the third block. A bad +/// landing costs little: guided self-scheduling leaves a streaming worker that +/// shares a block claiming fewer chunks. +pub(super) fn enabled() -> bool { + if !available() { + return false; + } + let slots = [0, 1, parallel::topology().perf]; + slots[..streaming_workers()].contains(¶llel::worker_id()) +} + +/// Whether this host runs the kernel at all, probed once. +pub(super) fn available() -> bool { + static AVAILABLE: OnceLock = OnceLock::new(); + *AVAILABLE.get_or_init(probe) +} + +/// [`super::hash_many_dyn_from_state`] on the streaming backend. +/// +/// # Safety +/// `data` must hold `n * len` bytes for `n = out.len() / OUT_LEN`, and `len` +/// must be a nonzero multiple of [`BLOCK_LEN`]. +pub(super) unsafe fn hash_many(data: &[u8], len: usize, state: &[u32; 8], t_offset: u64, out: &mut [u8]) { + let n = out.len() / OUT_LEN; + let mut inputs = [std::ptr::null::(); LANES]; + for g in 0..n / LANES { + let base = g * LANES; + for (l, slot) in inputs.iter_mut().enumerate() { + *slot = data[(base + l) * len..].as_ptr(); + } + // SAFETY: every pointer has `len` readable bytes, and the sixteen + // digests at `base` are inside `out`. + unsafe { + blake2s_hash16_sme2( + inputs.as_ptr(), + state.as_ptr(), + t_offset, + len as u64, + out.as_mut_ptr().add(base * OUT_LEN), + IV.as_ptr(), + ); + } + } + // Fewer than sixteen inputs left, and a batch under sixteen never entered + // the loop at all: NEON finishes those, being the faster of the two below + // one full vector. + let done = n - n % LANES; + if done < n { + // SAFETY: `n - done` inputs of `len` bytes remain, with as many digests. + unsafe { + hash_many_with::(&data[done * len..], len, state, t_offset, &mut out[done * OUT_LEN..]); + } + } +} From 8d9bf6502cd176df086ab536bc30073e93b56289 Mon Sep 17 00:00:00 2001 From: Tom Wambsgans Date: Thu, 27 Aug 2026 11:12:24 +0200 Subject: [PATCH 2/2] Never leave streaming mode mid-batch, and guard the slot list The batch tail was the defect. Finishing a partial group on NEON meant leaving streaming mode mid-call, and that transition costs several hundred nanoseconds, far more than the handful of inputs it was there to hash, so any batch whose size was not a multiple of sixteen could be slower than plain NEON: at 31 inputs 1263 ns against NEON's 946, at 63 inputs 2149 against 1707. The remainder now rides a padded group instead, repeating the last input to fill the spare lanes, which is free because the kernel drives sixteen lanes either way. 31 inputs go to 589 ns and 63 to 1072, both now 1.6x ahead of NEON. A batch under one full vector still goes to NEON and never enters streaming mode at all. Between 17 and 24 inputs the two are within 0.8x to 1.0x of each other; that band is not worth a threshold fitted to noise, and the Merkle tiler does not land in it. `LEANVM_SME_WORKERS=0` did not do what its documentation said. `enabled` consulted `available` first, which probes by running the kernel, so the documented kill switch could not rescue a host that reports the feature but faults on SMSTART. The count is now checked before anything touches the hardware. Slot 0 needed a guard. `worker_id` returns 0 for the dispatcher and for every thread that is not a pool worker alike, so an off-pool thread hashing concurrently would have split the dispatcher's block, which is the one thing the slot list exists to prevent. `parallel::in_task` separates them, being set only inside a dispatch. A strictly sequential pool dispatches nothing, so that case is admitted explicitly, and it is where the kernel's own gain is visible: 40 -> 67 Mhash/s on one thread. The probe hashed sixteen copies of the same zero block, so it could not have caught a lane-routing error despite claiming to agree with the scalar backend. It now hashes sixteen distinct blocks. Two hazards are now recorded in the module documentation, both latent rather than live: a signal handler using Advanced SIMD takes SIGILL if it fires inside the kernel, because the handler runs with PSTATE.SM still set, and bare `smstart` enables ZA without committing a pending lazy save, which is safe only while nothing else in the process uses ZA. The assembly's own contract comment omitted that `len` must be nonzero, which the do-while loop requires and the Rust side already asserts. Measured again on the M4 Max, 2026-08-27, medians of interleaved runs: hash_bench multithreaded_throughput 452 -> 512 Mhash/s (+13%) aggregate --xmss 900 --log-inv-rate 1 1055 -> 1070 sig/s (+1.4%) That end-to-end figure corrects the +2.0% claimed in the previous commit, which sat at the optimistic end of the spread; six interleaved rounds put it at +1.4%, with five of six pairs favouring the backend. One thing is left alone deliberately. `pcs::merkle::BATCH_LEAVES` is `hash::LANES * 2`, which is 8 on aarch64 while both backends there consume 16 per whole batch (NEON interleaves four 4-lane groups), so some leaf shapes shed 8 leaves per tile to the slower path. It predates this work, it does not fire in any benchmarked configuration, and changing the tiler deserves its own measured commit. Co-Authored-By: Claude Opus 5 (1M context) --- AGENTS.md | 1 + crates/parallel/src/lib.rs | 11 ++++ crates/primitives/src/hash/blake2s_sme2.s | 2 +- crates/primitives/src/hash/sme2.rs | 71 +++++++++++++++++------ 4 files changed, 67 insertions(+), 18 deletions(-) diff --git a/AGENTS.md b/AGENTS.md index 0dac0906..18c2d6b5 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -133,6 +133,7 @@ Understand the third before changing the verifier. `guests/aggregate.py` is zkDS | `LEANVM_NUM_THREADS` / `RAYON_NUM_THREADS` | performance-worker count; `1` = sequential | | `LEANVM_PROFILE` | per-stage prover timings | | `LEANVM_NO_ARENA` | disable the arena (less memory, slower) | +| `LEANVM_SME_WORKERS` | workers on the SME2 BLAKE2s backend; `0` = NEON only | | `ZK_ALLOC_STATS` | arena bytes/phase, high water, overflow | | `BENCH_REPEAT`, `BENCH_COOLDOWN` | `--repeat`/`--cooldown` for `#[ignore]`d benches | | `LEANVM_XMSS_N`, `LEANVM_HASH_N`, `LEANVM_HASH_UNROLL` | workload sizes in tests | diff --git a/crates/parallel/src/lib.rs b/crates/parallel/src/lib.rs index e5bec73b..fcab1d26 100644 --- a/crates/parallel/src/lib.rs +++ b/crates/parallel/src/lib.rs @@ -150,6 +150,17 @@ fn pool() -> &'static Pool { }) } +/// Whether the calling thread is running a pool task right now. +/// +/// This is what separates the dispatcher from every other thread in the +/// process: both read as worker 0, but only the dispatcher inside a dispatch +/// reads as being in a task. +#[must_use] +#[inline] +pub fn in_task() -> bool { + IN_TASK.get() +} + /// Which worker the calling thread is: `0` for the dispatcher, `1..perf` for the /// performance workers, and `perf..` for the efficiency ones. Any thread /// outside the pool reads as `0`. diff --git a/crates/primitives/src/hash/blake2s_sme2.s b/crates/primitives/src/hash/blake2s_sme2.s index 5f3651b9..78c951db 100644 --- a/crates/primitives/src/hash/blake2s_sme2.s +++ b/crates/primitives/src/hash/blake2s_sme2.s @@ -15,7 +15,7 @@ // const uint8_t *const *inputs, // x0, sixteen pointers // const uint32_t *state, // x1, the chaining value, splatted // uint64_t t_offset, // x2 -// uint64_t len, // x3, bytes per input, a multiple of 64 +// uint64_t len, // x3, bytes per input, a nonzero multiple of 64 // uint8_t *out, // x4, sixteen 32-byte digests // const uint32_t *iv); // x5, the eight BLAKE2s IV words // diff --git a/crates/primitives/src/hash/sme2.rs b/crates/primitives/src/hash/sme2.rs index f4505714..fa0d8c9a 100644 --- a/crates/primitives/src/hash/sme2.rs +++ b/crates/primitives/src/hash/sme2.rs @@ -6,6 +6,13 @@ //! so this is faster per thread and slower per machine: [`enabled`] hands it to //! the few workers that can reach a block of their own and leaves everyone else //! on NEON. +//! +//! Two hazards come with the mode. A signal delivered inside the kernel runs +//! its handler with `PSTATE.SM` still set, where Advanced SIMD is illegal, so a +//! handler using vector code takes `SIGILL`; nothing in this workspace installs +//! one. And bare `smstart` enables ZA without committing a pending lazy save, +//! which is safe only because nothing else here uses ZA, so `TPIDR2_EL0` is +//! always zero. use std::sync::OnceLock; @@ -53,8 +60,10 @@ fn probe() -> bool { if sysctl(c"hw.optional.arm.FEAT_SME2") != Some(1) { return false; } - let block = [0u8; BLOCK_LEN]; - let inputs = [block.as_ptr(); LANES]; + // Sixteen different blocks, so a lane that reads the wrong input fails here + // rather than in a proof. + let blocks: [[u8; BLOCK_LEN]; LANES] = std::array::from_fn(|l| std::array::from_fn(|i| (l * 61 + i * 7) as u8)); + let inputs: [*const u8; LANES] = std::array::from_fn(|l| blocks[l].as_ptr()); let mut out = [0u8; LANES * OUT_LEN]; // SAFETY: sixteen pointers to a whole readable block, and room for sixteen // digests. @@ -68,7 +77,7 @@ fn probe() -> bool { IV.as_ptr(), ) }; - lanes == LANES as u64 && out.chunks_exact(OUT_LEN).all(|d| d == super::hash(&block)) + lanes == LANES as u64 && (0..LANES).all(|l| out[l * OUT_LEN..(l + 1) * OUT_LEN] == super::hash(&blocks[l])[..]) } /// How many workers take the streaming path, at most one per cluster. @@ -93,17 +102,24 @@ const SLOTS: usize = 3; /// Whether the calling worker should take the streaming path. /// -/// Workers 0 and 1 are dispatcher and first performance worker, which the -/// scheduler tends to place on different performance clusters, and worker -/// `perf` is the first efficiency one, which reaches the third block. A bad -/// landing costs little: guided self-scheduling leaves a streaming worker that -/// shares a block claiming fewer chunks. +/// Worker 0 is the dispatcher and worker 1 the first performance worker, which +/// the scheduler does place on different performance clusters; worker `perf` is +/// the first efficiency one, which reaches the third block. The +/// [`parallel::in_task`] guard is what makes slot 0 safe: every thread outside +/// the pool reads as worker 0 too, and without it one of those could split the +/// dispatcher's block, which is the single thing this list exists to prevent. A +/// bad landing costs little anyway, since guided self-scheduling leaves a +/// worker that shares a block claiming fewer chunks. pub(super) fn enabled() -> bool { - if !available() { - return false; - } + // Before `available`, which probes by running the kernel: on a host that + // claims the feature but faults on `SMSTART`, setting the count to zero has + // to be enough to stay away from it. let slots = [0, 1, parallel::topology().perf]; - slots[..streaming_workers()].contains(¶llel::worker_id()) + let slots = &slots[..streaming_workers()]; + // A sequential pool dispatches nothing, so nothing is ever in a task, and + // with one thread there is no block to contend for either. + let in_pool_work = parallel::in_task() || parallel::num_threads() <= 1; + in_pool_work && slots.contains(¶llel::worker_id()) && available() } /// Whether this host runs the kernel at all, probed once. @@ -119,6 +135,13 @@ pub(super) fn available() -> bool { /// must be a nonzero multiple of [`BLOCK_LEN`]. pub(super) unsafe fn hash_many(data: &[u8], len: usize, state: &[u32; 8], t_offset: u64, out: &mut [u8]) { let n = out.len() / OUT_LEN; + if n < LANES { + // Under one vector there is nothing for the wide lanes to carry, and + // NEON does not pay to enter streaming mode. + // SAFETY: the caller's sizes, unchanged. + unsafe { hash_many_with::(data, len, state, t_offset, out) }; + return; + } let mut inputs = [std::ptr::null::(); LANES]; for g in 0..n / LANES { let base = g * LANES; @@ -138,14 +161,28 @@ pub(super) unsafe fn hash_many(data: &[u8], len: usize, state: &[u32; 8], t_offs ); } } - // Fewer than sixteen inputs left, and a batch under sixteen never entered - // the loop at all: NEON finishes those, being the faster of the two below - // one full vector. + // The remainder rides a padded group rather than a NEON call. Leaving + // streaming mode to finish a handful of inputs costs several times what + // the spare lanes do, those being free: the kernel drives sixteen either + // way. Repeating the last input fills them. let done = n - n % LANES; if done < n { - // SAFETY: `n - done` inputs of `len` bytes remain, with as many digests. + for (l, slot) in inputs.iter_mut().enumerate() { + *slot = data[(done + l).min(n - 1) * len..].as_ptr(); + } + let mut padded = [0u8; LANES * OUT_LEN]; + // SAFETY: every pointer is one of this batch's own inputs, so each has + // `len` readable bytes, and `padded` holds all sixteen digests. unsafe { - hash_many_with::(&data[done * len..], len, state, t_offset, &mut out[done * OUT_LEN..]); + blake2s_hash16_sme2( + inputs.as_ptr(), + state.as_ptr(), + t_offset, + len as u64, + padded.as_mut_ptr(), + IV.as_ptr(), + ); } + out[done * OUT_LEN..].copy_from_slice(&padded[..(n - done) * OUT_LEN]); } }