Expose optimized Poseidon routines to the Poseidon gate (ARM) (#331)

* Expose optimized Poseidon routines to the Poseidon gate

* Daniel PR comments
This commit is contained in:
Jakub Nabaglo 2021-11-02 12:38:30 -07:00 committed by GitHub
parent bc57a561e6
commit 184f73c604
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 58 additions and 18 deletions

View File

@ -1006,8 +1006,8 @@ unsafe fn partial_rounds(
}
#[inline(always)]
pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
let state = [
fn unwrap_state(state: [GoldilocksField; 12]) -> [u64; 12] {
[
state[0].0,
state[1].0,
state[2].0,
@ -1020,22 +1020,11 @@ pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
state[9].0,
state[10].0,
state[11].0,
];
let state = const_layer_full(state, ALL_ROUND_CONSTANTS[0..WIDTH].try_into().unwrap());
let state = full_rounds(
state,
ALL_ROUND_CONSTANTS[WIDTH..WIDTH * (HALF_N_FULL_ROUNDS + 1)]
.try_into()
.unwrap(),
);
let state = partial_rounds(
state,
ALL_ROUND_CONSTANTS
[WIDTH * (HALF_N_FULL_ROUNDS + 1)..WIDTH * (HALF_N_FULL_ROUNDS + N_PARTIAL_ROUNDS + 1)]
.try_into()
.unwrap(),
);
let state = full_rounds(state, &FINAL_ROUND_CONSTANTS);
]
}
#[inline(always)]
fn wrap_state(state: [u64; 12]) -> [GoldilocksField; 12] {
[
GoldilocksField(state[0]),
GoldilocksField(state[1]),
@ -1051,3 +1040,38 @@ pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
GoldilocksField(state[11]),
]
}
#[inline(always)]
pub unsafe fn poseidon(state: [GoldilocksField; 12]) -> [GoldilocksField; 12] {
let state = unwrap_state(state);
let state = const_layer_full(state, ALL_ROUND_CONSTANTS[0..WIDTH].try_into().unwrap());
let state = full_rounds(
state,
ALL_ROUND_CONSTANTS[WIDTH..WIDTH * (HALF_N_FULL_ROUNDS + 1)]
.try_into()
.unwrap(),
);
let state = partial_rounds(
state,
ALL_ROUND_CONSTANTS
[WIDTH * (HALF_N_FULL_ROUNDS + 1)..WIDTH * (HALF_N_FULL_ROUNDS + N_PARTIAL_ROUNDS + 1)]
.try_into()
.unwrap(),
);
let state = full_rounds(state, &FINAL_ROUND_CONSTANTS);
wrap_state(state)
}
#[inline(always)]
pub unsafe fn sbox_layer(state: &mut [GoldilocksField; WIDTH]) {
*state = wrap_state(sbox_layer_full(unwrap_state(*state)));
}
#[inline(always)]
pub unsafe fn mds_layer(state: &[GoldilocksField; WIDTH]) -> [GoldilocksField; WIDTH] {
let state = unwrap_state(*state);
// We want to do an MDS layer without the constant layer.
let round_consts = [0u64; WIDTH];
let state = mds_const_layers_full(state, &round_consts);
wrap_state(state)
}

View File

@ -389,6 +389,22 @@ impl Poseidon<12> for GoldilocksField {
crate::hash::arch::aarch64::poseidon_goldilocks_neon::poseidon(input)
}
}
#[cfg(all(target_arch="aarch64", target_feature="neon"))]
#[inline(always)]
fn sbox_layer(state: &mut [Self; 12]) {
unsafe {
crate::hash::arch::aarch64::poseidon_goldilocks_neon::sbox_layer(state);
}
}
#[cfg(all(target_arch="aarch64", target_feature="neon"))]
#[inline(always)]
fn mds_layer(state: &[Self; 12]) -> [Self; 12] {
unsafe {
crate::hash::arch::aarch64::poseidon_goldilocks_neon::mds_layer(state)
}
}
}
#[cfg(test)]