SME2 streaming-mode BLAKE2s

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.
@@ -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                          |

crates/parallel/src/lib.rs

@@ -150,6 +150,26 @@ 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`.
+#[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);

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::<x86::Avx2>(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::<arm::Neon>(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<u8> = (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