[bn254, poseidon] Add new poseidon implementation - #108
Conversation
6dc9cd3 to
be6839b
Compare
samkim-crypto
changed the title
zz-sol
left a comment
zz-sol
left a comment
There was a problem hiding this comment.
I think we need to add a constant.rs file to store a set of default constant.
Similar to, for instance,
- https://github.com/Plonky3/Plonky3/blob/main/mersenne-31/src/poseidon1.rs#L135
- https://github.com/Lightprotocol/light-poseidon/blob/main/light-poseidon/src/parameters/bn254_x5.rs
then we can have a separate function that allows users to derive their own parameters.
we can use the default parameters for consistency checks against other libraries
| }; | ||
|
|
||
| /// Generates a randomized field element safely constrained below the BN254 Fr modulus. | ||
| fn random_fr() -> U256 { |
There was a problem hiding this comment.
nit: maybe sample ark_bn254::Fr then convert to U256? or rename the function non_uniform_random_fr_for_bench()
|
|
||
| macro_rules! bench_width { | ||
| ($c:expr, $t:literal, $partial_rounds:literal) => { | ||
| let constants = make_dummy_constants::<$t>($partial_rounds); |
There was a problem hiding this comment.
is it possible to use real constants for benchmark?
| let a_limbs = [a.l0, a.l1, a.l2, a.l3, a.l4]; | ||
|
|
||
| // CIOS Algorithm: Loop is fully unrolled by the LLVM compiler. | ||
| for i in 0..5 { |
There was a problem hiding this comment.
Here
- five radix-2^52 CIOS iterations use the Montgomery radix
R=2^260, - the serial backend and
Fr::R2useR=2^256. pack_8xonly rearranges limbs and does not convert between the two Montgomery domains.
So the SIMD path computes a*b*2^-260 instead of a*b*2^-256.
i think it may be helpful to have a test for consistency with light_poseidon test vectors and used across platforms
| authors = { workspace = true } | ||
| repository = { workspace = true } | ||
| license = { workspace = true } | ||
| readme = "README.md" |
| } | ||
| new_state[i] = sum; | ||
| } | ||
| *state = new_state; |
There was a problem hiding this comment.
It seems that the CIOS result may still be in [r, 2r) after carry propagation. apply_dense_matrix_simd subsequently passes this value to the scalar add implementation, which requires inputs to be strictly below r. please document here and in apply_dense_matrix_simd on the input criteria, if this is fine; or maybe we need to do a mod reduction first?
| /// by 4 bits), followed by a multiplication with the precomputed | ||
| /// window value that matches the next 4 bits of the exponent. | ||
| #[inline(always)] | ||
| pub fn invert(a: &U256) -> U256 { |
There was a problem hiding this comment.
short circuit for a \in {0, 1}.
also should the return type be Option<U256>?
not a strong preference as this only captures 0 case.
|
|
||
| /// Computes `a^((p+1)/4) mod p` using a static 4-bit window addition chain. | ||
| #[inline(always)] | ||
| pub fn sqrt(a: &U256) -> U256 { |
There was a problem hiding this comment.
seems that we don't need sqrt.
For sqrt, i think after the computation we need to check t^2 == a and return None if it fails -- there exist quadratic non-residues.
Also we can add some short circuit for 0 and 1
| use super::*; | ||
|
|
||
| #[test] | ||
| fn test_fq_invert() { |
There was a problem hiding this comment.
i think we don't use invert and sqrt in poseidon. but if we do want to include those functions, it may be better to have more tests for random inputs and for edge cases.
| /// Like the Fq implementation, this strictly avoids dynamic branching or | ||
| /// loops in favor of a static 4-bit window addition chain. | ||
| #[inline(always)] | ||
| pub fn invert(a: &U256) -> U256 { |
There was a problem hiding this comment.
same comments as on Fq::invert()
| //! to prioritize cycle efficiency and lowest possible Compute Units. | ||
|
|
||
| pub mod backend; | ||
| pub mod poseidon; |
There was a problem hiding this comment.
can we please have a test and check against light_poseidon (and perhaps arkworks?)
2 participants
Adding poseidon implementation that has zero dependencies and has avx-512 optimization.
I was able to get good speedups over the reference
light-poseidoncrate:[U256; T]). This completely eliminates the heap allocation (Vec) and dynamic sizing overhead.O(T^2)dense matrix-vector multiplications during the partial rounds down to anO(T)sparse matrix computation. This is probably the source of most improvement overlight-poseidon.x^5) and Montgomery addition chains are explicitly unrolled and inlined to remove unnecessary loops and branching overhead.I wrote code for architectures that have and don't have avx acceleration:
For CPUs lacking AVX-512 (e.g., Apple Silicon, older x86), I used a 4-limb (64-bit) pure-Rust backend.
On supported modern hardware, the dense matrix and full rounds are dynamically routed to a SIMD engine using
vpmadd52instructions. The 256-bit state is repacked into five 52-bit limbs across 64-bit AVX lanes. I left 12 bits of "headroom" per limb. This allows us to accumulate up to 4,096 matrix cross-products simultaneously without needing to propagate a carry or performing a Montgomery reduction until the matrix multiplication is finished. At peak state widths (T=13), the engine calculates 8 S-boxes perfectly in parallel, and collapses 169 scalar Montgomery multiplications into just 26 SIMD multiplications.Currently, I named this crate to be
solana-bn254. I am planning on adding a native bn254 crate for the alt-bn128 syscalls as well. Once these are merged and audited, then I think we can rename the crate to besolana-bn254-syscalland replace the existingsolana-bn254-syscallcrate that depends onark-bn254.