From 79e04929cbd1c8b3b89aaadb08d4e3e44c563aa6 Mon Sep 17 00:00:00 2001 From: Nikolai Golub Date: Wed, 2 Sep 2026 14:23:17 +0200 Subject: [PATCH] perf(rsema1d): encode parity in cache-resident column stripes in parallel fill_parity encoded every row as one Leopard shard on the calling thread. For K=4096/N=12288 the high-rate work buffer is 16384 shards, so 32 KiB rows make it 512 MiB and every IFFT/FFT pass streams from DRAM: 770 ms per 128 MiB blob on one core, and running several encoders at once did not scale (4 processes: 1.3x total). Leopard applies the same transform independently to each 64-byte block position of a shard, so the rows are now split into column stripes and each stripe is encoded by its own HighRateEncoder on a rayon worker. The stripe width is chosen so that all workers' work buffers together stay within 32 MiB (cache-resident); encoders are kept in a thread-local and reused across calls. The parity region is faulted in sequentially before the striped scatter, since concurrent first-touch of interleaved pages serialised in the kernel (105 ms vs 6 ms for a 24 MiB region). Output is byte-identical to a single encoder over whole rows, which is what the Go verifier (klauspost Leopard GF16) recomputes: a unit test checks this for the production shape and for a row size that does not divide evenly, and the Go cross-implementation vectors (254 cases including K=4096/N=12288) pass. criterion, target-cpu=native, Ryzen 9 3950X (32 rayon threads): rsema1d encode_in_place/128MB_k4096_n12288 821 ms -> 207 ms rsema1d encode_in_place/8MB_k4096_n12288 52.6 ms -> 17.3 ms fibre blob_new/128MB 985 ms -> 408 ms fibre blob_new/8MB 64.1 ms -> 29.3 ms Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01VPtihoV1GJ9VcNkd7oMRUf --- rsema1d/src/codec/rs.rs | 219 +++++++++++++++++++++++++++++++++++----- 1 file changed, 196 insertions(+), 23 deletions(-) diff --git a/rsema1d/src/codec/rs.rs b/rsema1d/src/codec/rs.rs index 3d865b712..099639850 100644 --- a/rsema1d/src/codec/rs.rs +++ b/rsema1d/src/codec/rs.rs @@ -2,8 +2,142 @@ use crate::codec::rows::{OriginalRowsView, RowMatrix}; use crate::error::{Error, Result}; use crate::field::GF128; use crate::params::Parameters; +use rayon::prelude::*; use reed_solomon_simd::engine::DefaultEngine; use reed_solomon_simd::rate::{HighRateEncoder, RateEncoder}; +use std::cell::RefCell; + +/// Upper bound on the combined Leopard work set of all stripe encoders that +/// run concurrently, in bytes. +/// +/// Leopard's high-rate transform keeps `k.next_multiple_of(n.next_power_of_two())` +/// shards in a work buffer and sweeps it several times, so throughput is set by +/// whether that buffer sits in cache or in DRAM. Encoding every row as a single +/// shard (32 KiB rows at K=4096/N=12288) makes the buffer 512 MiB and the whole +/// encode DRAM-bound on one core. Splitting rows into column stripes gives +/// `row_size / stripe` independent encoders whose buffers are +/// `work_shards * stripe` bytes each; keeping all concurrently running buffers +/// within this budget keeps the transforms cache-resident. 32 MiB is half the +/// L3 of a 16-core Zen 2 part and was the best value measured there (16 threads +/// x 2 MiB or 32 threads x 1 MiB gave the same ~58 ms per 128 MiB blob). +const STRIPE_WORK_BUDGET: usize = 32 << 20; + +/// Smallest stripe worth encoding separately. Below one 64-byte Leopard block +/// per shard the per-shard overhead dominates. +const MIN_STRIPE: usize = 64; + +thread_local! { + /// One encoder per rayon worker, reused across calls so the Leopard work + /// buffer is allocated and page-faulted once per thread, not once per blob. + static ENCODER: RefCell>> = const { RefCell::new(None) }; +} + +/// Raw pointer to the parity region that can be shared across rayon tasks. +/// +/// Each stripe task writes only the byte ranges +/// `[row * row_size + offset, row * row_size + offset + len)` of its own stripe, +/// which are disjoint between tasks, so concurrent writes never alias. +#[derive(Clone, Copy)] +struct ParityOut(*mut u8); + +// SAFETY: see the type documentation; the pointer is only dereferenced for +// disjoint byte ranges, and the underlying `&mut [u8]` outlives the parallel +// section that uses it. +unsafe impl Send for ParityOut {} +unsafe impl Sync for ParityOut {} + +/// Page size assumed when faulting in freshly allocated parity memory. +const TOUCH_PAGE: usize = 4096; +/// Contiguous span each rayon worker faults in at a time. +const TOUCH_CHUNK: usize = 2 << 20; + +/// Write one byte per page of `buf`, contiguous chunks per worker. +/// +/// The striped scatter below has every worker writing 64-byte pieces all +/// over the parity region at once. When that region is freshly allocated, +/// the first touch of each page is a page fault, and many threads faulting +/// interleaved pages of one mapping serialize in the kernel: measured 105 ms +/// versus 6 ms for a 24 MiB parity region with 32 threads. Faulting the pages +/// in sequentially first takes the fast path; on already-resident memory this +/// pass costs one cached store per page. +fn touch_pages(buf: &mut [u8]) { + buf.par_chunks_mut(TOUCH_CHUNK).for_each(|chunk| { + for page in chunk.chunks_mut(TOUCH_PAGE) { + // SAFETY: `page` is a non-empty, valid, writable slice. + unsafe { std::ptr::write_volatile(page.as_mut_ptr(), page[0]) }; + } + }); +} + +/// Number of shards in the Leopard high-rate work buffer for `(k, n)`. +fn work_shards(k: usize, n: usize) -> usize { + let chunk = n.next_power_of_two(); + k.div_ceil(chunk) * chunk +} + +/// Stripe width in bytes for encoding `row_size`-byte rows with `(k, n)`. +/// +/// Returns the largest multiple of 64 such that all rayon workers' work +/// buffers fit in [`STRIPE_WORK_BUDGET`], clamped to `[MIN_STRIPE, row_size]`. +fn stripe_size(k: usize, n: usize, row_size: usize) -> usize { + let threads = rayon::current_num_threads().max(1); + let per_encoder = STRIPE_WORK_BUDGET / (threads * work_shards(k, n)); + let stripe = (per_encoder / MIN_STRIPE) * MIN_STRIPE; + stripe.clamp(MIN_STRIPE, row_size) +} + +/// Encode one column stripe of `original_rows` into the matching stripe of +/// the parity rows, using this thread's cached encoder. +fn encode_stripe( + original_rows: &[u8], + parity: ParityOut, + k: usize, + n: usize, + row_size: usize, + offset: usize, + len: usize, +) -> Result<()> { + ENCODER.with(|cell| { + let mut slot = cell.borrow_mut(); + let encoder = match slot.as_mut() { + Some(encoder) => { + encoder + .reset(k, n, len) + .map_err(|e| Error::ReedSolomon(e.to_string()))?; + encoder + } + None => slot.insert( + RateEncoder::new(k, n, len, DefaultEngine::new(), None) + .map_err(|e| Error::ReedSolomon(e.to_string()))?, + ), + }; + + for row in original_rows.chunks_exact(row_size) { + encoder + .add_original_shard(&row[offset..offset + len]) + .map_err(|e| Error::ReedSolomon(e.to_string()))?; + } + + let result = encoder + .encode() + .map_err(|e| Error::ReedSolomon(e.to_string()))?; + + for (i, recovery) in result.recovery_iter().enumerate() { + debug_assert_eq!(recovery.len(), len); + // SAFETY: `parity` points to `n * row_size` writable bytes that no + // other task touches in this stripe's byte ranges (see `ParityOut`), + // and `i < n`, `offset + len <= row_size`. + unsafe { + std::ptr::copy_nonoverlapping( + recovery.as_ptr(), + parity.0.add(i * row_size + offset), + len, + ); + } + } + Ok(()) + }) +} fn fill_parity( original_rows: &[u8], @@ -33,32 +167,28 @@ fn fill_parity( k, n, row_size ))); } - - let engine = DefaultEngine::new(); - let mut encoder: HighRateEncoder = - RateEncoder::new(k, n, row_size, engine, None) - .map_err(|e| Error::ReedSolomon(e.to_string()))?; - - // Add all original rows - for row in original_rows.chunks_exact(row_size) { - encoder - .add_original_shard(row) - .map_err(|e| Error::ReedSolomon(e.to_string()))?; + if !row_size.is_multiple_of(MIN_STRIPE) { + return Err(Error::InvalidParameters(format!( + "row_size {} is not a multiple of {}", + row_size, MIN_STRIPE + ))); } - // Generate parity rows - let result = encoder - .encode() - .map_err(|e| Error::ReedSolomon(e.to_string()))?; - - for (dst_row, src_row) in parity_rows - .chunks_exact_mut(row_size) - .zip(result.recovery_iter()) - { - dst_row.copy_from_slice(src_row); - } + // Leopard applies the same transform independently to every 64-byte + // block position of a shard, so encoding column stripes of the rows + // separately yields exactly the parity that one encoder over whole rows + // would, byte for byte. The last stripe may be shorter. + let stripe = stripe_size(k, n, row_size); + touch_pages(parity_rows); + let parity = ParityOut(parity_rows.as_mut_ptr()); - Ok(()) + (0..row_size.div_ceil(stripe)) + .into_par_iter() + .try_for_each(|s| { + let offset = s * stripe; + let len = stripe.min(row_size - offset); + encode_stripe(original_rows, parity, k, n, row_size, offset, len) + }) } /// Extend data using Reed-Solomon encoding. @@ -218,4 +348,47 @@ mod tests { } assert!(extended[k..(k + n)].iter().any(|rlc| *rlc != GF128::zero())); } + + /// Striped, parallel parity must be byte-identical to a single Leopard + /// encoder over whole rows (what validators recompute on the Go side). + #[test] + fn striped_parity_matches_single_encoder() { + use rand::{RngCore, SeedableRng}; + use rand_chacha::ChaCha8Rng; + + // Production shape (K=4096, N=12288) with a row size that splits into + // several stripes under any thread count, plus a row size that is not + // a multiple of the stripe (last stripe shorter) and a tiny case. + for (k, n, row_size, seed) in [ + (4096usize, 12288usize, 512usize, 1u64), + (4096, 12288, 64 * 21, 2), + (16, 48, 4096, 3), + (4, 4, 64, 4), + ] { + let mut rng = ChaCha8Rng::seed_from_u64(seed); + let mut original = vec![0u8; k * row_size]; + rng.fill_bytes(&mut original); + + let mut expected = vec![0u8; n * row_size]; + let mut encoder: HighRateEncoder = + RateEncoder::new(k, n, row_size, DefaultEngine::new(), None).unwrap(); + for row in original.chunks_exact(row_size) { + encoder.add_original_shard(row).unwrap(); + } + let result = encoder.encode().unwrap(); + for (dst, src) in expected + .chunks_exact_mut(row_size) + .zip(result.recovery_iter()) + { + dst.copy_from_slice(src); + } + + let mut parity = vec![0u8; n * row_size]; + fill_parity(&original, &mut parity, k, n, row_size).unwrap(); + assert!( + parity == expected, + "striped parity differs for k={k} n={n} row_size={row_size}" + ); + } + } }