perf(xmss): remove dead work from the signing path#263
Open
wstran wants to merge 3 commits into
Open
Conversation
xmss_sign spends more time deriving a WOTS public key it never reads than it does producing the signature itself. gen_wots_secret_key goes through WotsSecretKey::new, which walks all V chains to their ends, but the signing path only calls sign_with_encoding, which reads pre_images. The public key is needed by leaf_layer alone, to hash it into a Merkle leaf. Root cause: WotsSecretKey cached `public_key` as a struct field, so every construction paid for the full chain walk whether the caller wanted it or not. Fix: drop the field and compute the key on demand in `public_key(public_param, slot)`. Key generation walks the same chains and produces the same leaves; signing no longer walks them at all. Measured (docker rust:latest, 6 cores, V=42, CHAIN_LENGTH=8, TARGET_SUM=184): chain hashes per signature: 184 useful + 294 dead -> 184 useful gen_wots_secret_key: 370.8 us -> negligible WotsSecretKey is private to the crate (only WotsSignature is re-exported), so this is not a public API change.
Signing grinds wots_encode until the codeword sums to TARGET_SUM, which takes on the order of 840 attempts per signature. Every attempt allocated three times over: to_little_endian_bits builds a Vec<bool> per encoding field element, flat_map collects a second Vec of 144 bools, and the chunk map collects the codeword Vec. Roughly 839 of 840 attempts discard all of it. Root cause: the codeword was assembled through a bit-vector pipeline instead of being read straight out of the field elements. Fix: read each W-bit chunk with a shift and a mask into a stack array, and accumulate the sum in the same pass. Masking to W bits already bounds every entry below CHAIN_LENGTH, so the target sum is the only remaining condition and is_valid_encoding is no longer needed. Reading chunks per field element rather than from one concatenated bit stream is only equivalent while W divides the bits taken from each element, so that invariant is now named (ENCODING_BITS_PER_FE, CHUNKS_PER_FE) and asserted at compile time, together with CHAIN_LENGTH being a power of two since CHAIN_LENGTH - 1 is used as the chunk mask. Measured (400k attempts per run, 5 interleaved rounds): wots_encode: 2750 ns -> 2534 ns per attempt (-7.9%) The same 502 codewords out of 420000 attempts are accepted before and after, so the encoding is unchanged.
Ignored tests, following the pattern of benchmark_poseidons.rs: signing and verification throughput, per-attempt cost of the grinding loop, and signature digests for fixed (seed, slot, message) triples so that a change meant to leave signatures untouched can be checked against them.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What
xmss_signdoes two pieces of work it does not need.V * (CHAIN_LENGTH - 1)= 294 chain hashes at the current parameters, that the signing path never reads. Signing walks each chain only part-way, 184 chain hashes in total, so the dead work was larger than the real work.Vec<bool>per field element insideto_little_endian_bits, then a 144-elementVec<bool>, then the codewordVec<u8>. About 839 of 840 attempts throw all of it away.This sits outside the prover but on a path that runs constantly, since
xmss_signis the signing entry point the lean clients build on.Root cause:
WotsSecretKeycachedpublic_keyas a struct field, so every construction paid for the full chain walk; andwots_encodeassembled the codeword through a bit-vector pipeline instead of reading the chunks out of the field elements.Fix: compute the WOTS public key on demand in
public_key(public_param, slot), so key generation still gets it and signing does not; and read eachW-bit chunk with a shift and a mask into a stack array while accumulating the sum in the same pass.WotsSecretKeyis private to the crate, so neither change touches the public API.Reading chunks per field element instead of from one concatenated bit stream is only equivalent while
Wdivides the bits taken from each element. That invariant was implicit inNUM_ENCODING_FE = V.div_ceil(24 / W), so it is now named and asserted:Setting
W = 5fails the build on the first of those rather than silently changing the encoding.Numbers
Machine: docker rust:latest, 6 cores, aarch64. Baseline and patched binaries are built first, then run interleaved on an idle machine, since building during a measurement was enough to swamp the effect.
xmss_sign, mean of 8 roundsxmss_sign, median of 8 roundsxmss_verifywots_encode, per attemptAll 8
xmss_signrounds improve, spread -8.1% to -12.1%. Thewots_encodemicrobenchmark reproduces within 0.6 percentage points across 5 rounds.repro:
Correctness
Signatures are byte identical. Signing five fixed (seed, slot, message) triples and hashing
as_ssz_bytesgives the same digests before and after:benchmark_sign::signature_digestsreproduces those. Thewots_encodemicrobenchmark accepts exactly the same 502 codewords out of 420000 attempts before and after, so the encoding is unchanged, andxmss --n-signatures 1550 --log-inv-rate 1reports the same 993,869 cycles, 3,911,273 poseidons and 327 KiB proof as on main.The uniformity rejection is untouched:
p - 1 = 127 * 2^24exactly, so dropping that single value is what makes the low 24 bits of each encoding element uniform, and the target sum remains the only acceptance condition.cargo test --release --workspacepasses,cargo clippy --all-targets --workspace -- -D warningsis clean, andcargo fmt --all --checkis clean. Each of the three commits builds on its own.Scope
Only the WOTS layer of the
xmsscrate. Key generation walks the same chains as before and produces the same leaves, so key generation timing does not change. The grinding loop still runs the same number of attempts under the same acceptance rule; only the per-attempt cost moves. The benchmarks are#[ignore]d tests behind#[cfg(test)], so they do not affect release builds.to_little_endian_bitsinpolynow has no callers in the workspace; it is left in place rather than removed from another crate here. Batching grinding attempts across SIMD lanes would attack the two Poseidon compressions that remain per attempt, and is left out of this PR.