perf: Avoid heap allocations in state, input buffers and byte hash path - #61
Draft
vadorovsky wants to merge 2 commits into
Draft
vadorovsky wants to merge 2 commits into
vadorovsky wants to merge 2 commits into
Conversation
The Poseidon state holds at most `MAX_X5_LEN` (13) field elements and the number of inputs is at most 12, so neither buffer needs dynamic allocation. `Poseidon.state` is now an `ArrayVec<F, MAX_X5_LEN>`. The per-round `collect()` in `apply_mds` no longer mallocs and frees a fresh `Vec` (one allocation pair per round, 64-74 per hash) - the replacement state is built on the stack and moved into place. `hash_bytes_be`/`hash_bytes_le` deserialize inputs into a single stack-resident `ArrayVec<F, MAX_INPUTS>` instead of building two intermediate `Vec`s (validation results and converted elements). This adds the `arrayvec` dependency and a public `MAX_INPUTS` constant. Two incidental refactors: the input-length check is extracted into `validate_inputs_length`, and `validate_bytes_length` now returns `Result<(), PoseidonError>` instead of `Result<&[u8], PoseidonError>` - the identity return was only used by the old map/collect pattern. Note that this last part is a public API change.
The byte-input path still performed small heap allocations per call. `bytes_to_prime_field_element_be/le` built a `num_bigint::BigUint` (a heap `Vec<u32>`) for every input. The value is now assembled directly in `F::BigInt` (a fixed-size stack array of limbs) from the bytes, most significant byte first, and checked with `F::from_bigint`, which performs the modulus check in ark-ff 0.5. `hash_bytes_be/le` converted the result via `to_bytes_be/le`, which return a `Vec<u8>`. A new `IntoHashBytes` extension trait for `PrimeField` serializes the result into the fixed-size output array directly from its limbs instead. This removes the `num-bigint` dependency. Edge-case behavior is preserved: short inputs and oversized inputs with leading zero bytes are still accepted, `BytesToBigInt` is still returned for values that do not fit in `F::BigInt`, and `InputLargerThanModulus` for values >= the modulus. Also drops a needless borrow in `hash` flagged by clippy. Benchmarked on AMD EPYC 7H12 64-core with the following results. Current `main`: ``` poseidon_bn254_x5_1 time: [24.407 µs 24.626 µs 24.894 µs] poseidon_bn254_x5_2 time: [34.412 µs 34.441 µs 34.485 µs] poseidon_bn254_x5_3 time: [47.040 µs 47.303 µs 47.542 µs] poseidon_bn254_x5_4 time: [65.993 µs 66.204 µs 66.439 µs] poseidon_bn254_x5_5 time: [85.580 µs 85.980 µs 86.484 µs] poseidon_bn254_x5_6 time: [112.77 µs 112.91 µs 113.08 µs] poseidon_bn254_x5_7 time: [143.64 µs 143.69 µs 143.76 µs] poseidon_bn254_x5_8 time: [175.33 µs 175.41 µs 175.48 µs] poseidon_bn254_x5_9 time: [215.06 µs 223.51 µs 231.91 µs] poseidon_bn254_x5_10 time: [274.29 µs 274.95 µs 275.64 µs] poseidon_bn254_x5_11 time: [297.08 µs 297.61 µs 298.29 µs] poseidon_bn254_x5_12 time: [374.82 µs 375.18 µs 375.67 µs] ``` This change + arrayvec migration: ``` poseidon_bn254_x5_1 time: [22.651 µs 22.673 µs 22.700 µs] poseidon_bn254_x5_2 time: [33.034 µs 33.140 µs 33.285 µs] poseidon_bn254_x5_3 time: [44.899 µs 44.931 µs 44.976 µs] poseidon_bn254_x5_4 time: [64.334 µs 64.487 µs 64.700 µs] poseidon_bn254_x5_5 time: [83.545 µs 83.695 µs 83.888 µs] poseidon_bn254_x5_6 time: [113.27 µs 113.38 µs 113.54 µs] poseidon_bn254_x5_7 time: [143.47 µs 144.17 µs 145.06 µs] poseidon_bn254_x5_8 time: [175.51 µs 175.62 µs 175.73 µs] poseidon_bn254_x5_9 time: [206.25 µs 206.52 µs 206.83 µs] poseidon_bn254_x5_10 time: [272.00 µs 272.19 µs 272.54 µs] poseidon_bn254_x5_11 time: [296.73 µs 297.05 µs 297.51 µs] poseidon_bn254_x5_12 time: [375.78 µs 376.22 µs 376.66 µs] ```
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
1 participant
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.
See individual commits for details.
Current
main:This change + arrayvec migration: