logos-storage/plonky2
1
0
Fork 0
mirror of https://github.com/logos-storage/plonky2.git synced 2026-01-10 09:43:09 +00:00
Code Issues Packages Projects Releases Wiki Activity

Add additional methods for Poseidon with PackedField (#1526)

Browse Source 
* Add additional methods for Poseidon with PackedField

* Apply comments

* Comment

---------

Co-authored-by: Linda Guiga <lindaguiga3@gmail.com>
This commit is contained in:
Robin Salen 2024-02-15 08:52:34 -05:00 committed by GitHub
parent 710225c9e0
commit 6f42ed5ca3
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
1 changed files with 170 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

172
plonky2/src/hash/poseidon.rs
View File

@ -5,6 +5,7 @@
use alloc::{vec, vec::Vec};
use core::fmt::Debug;
use plonky2_field::packed::PackedField;
use unroll::unroll_for_loops;
use crate::field::extension::{Extendable, FieldExtension};
@ -160,8 +161,9 @@ pub trait Poseidon: PrimeField64 {
// times number of rounds.
const N_ROUND_CONSTANTS: usize = SPONGE_WIDTH * N_ROUNDS;
// The MDS matrix we use is C + D, where C is the circulant matrix whose first row is given by
// `MDS_MATRIX_CIRC`, and D is the diagonal matrix whose diagonal is given by `MDS_MATRIX_DIAG`.
// The MDS matrix we use is C + D, where C is the circulant matrix whose first
// row is given by `MDS_MATRIX_CIRC`, and D is the diagonal matrix whose
// diagonal is given by `MDS_MATRIX_DIAG`.
const MDS_MATRIX_CIRC: [u64; SPONGE_WIDTH];
const MDS_MATRIX_DIAG: [u64; SPONGE_WIDTH];
@ -213,6 +215,33 @@ pub trait Poseidon: PrimeField64 {
res
}
/// Same as `mds_row_shf` for `PackedField`.
fn mds_row_shf_packed_field<
F: RichField + Extendable<D>,
const D: usize,
FE,
P,
const D2: usize,
>(
r: usize,
v: &[P; SPONGE_WIDTH],
) -> P
where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>,
{
debug_assert!(r < SPONGE_WIDTH);
let mut res = P::ZEROS;
for i in 0..SPONGE_WIDTH {
res +=
v[(i + r) % SPONGE_WIDTH] * P::Scalar::from_canonical_u64(Self::MDS_MATRIX_CIRC[i]);
}
res += v[r] * P::Scalar::from_canonical_u64(Self::MDS_MATRIX_DIAG[r]);
res
}
/// Recursive version of `mds_row_shf`.
fn mds_row_shf_circuit<const D: usize>(
builder: &mut CircuitBuilder<Self, D>,
@ -273,6 +302,29 @@ pub trait Poseidon: PrimeField64 {
result
}
/// Same as `mds_layer` for `PackedField`.
fn mds_layer_packed_field<
F: RichField + Extendable<D>,
const D: usize,
FE,
P,
const D2: usize,
>(
state: &[P; SPONGE_WIDTH],
) -> [P; SPONGE_WIDTH]
where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>,
{
let mut result = [P::ZEROS; SPONGE_WIDTH];
for r in 0..SPONGE_WIDTH {
result[r] = Self::mds_row_shf_packed_field(r, state);
}
result
}
/// Recursive version of `mds_layer`.
fn mds_layer_circuit<const D: usize>(
builder: &mut CircuitBuilder<Self, D>,
@ -320,6 +372,29 @@ pub trait Poseidon: PrimeField64 {
}
}
/// Same as `partial_first_constant_layer` for `PackedField`.
#[inline(always)]
#[unroll_for_loops]
fn partial_first_constant_layer_packed_field<
F: RichField + Extendable<D>,
const D: usize,
FE,
P,
const D2: usize,
>(
state: &mut [P; SPONGE_WIDTH],
) where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>,
{
for i in 0..12 {
if i < SPONGE_WIDTH {
state[i] +=
P::Scalar::from_canonical_u64(Self::FAST_PARTIAL_FIRST_ROUND_CONSTANT[i]);
}
}
}
/// Recursive version of `partial_first_constant_layer`.
fn partial_first_constant_layer_circuit<const D: usize>(
builder: &mut CircuitBuilder<Self, D>,
@ -365,6 +440,46 @@ pub trait Poseidon: PrimeField64 {
result
}
/// Same as `mds_partial_layer_init` for `PackedField`.
#[inline(always)]
#[unroll_for_loops]
fn mds_partial_layer_init_packed_field<
F: RichField + Extendable<D>,
const D: usize,
FE,
P,
const D2: usize,
>(
state: &[P; SPONGE_WIDTH],
) -> [P; SPONGE_WIDTH]
where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>,
{
let mut result = [P::ZEROS; SPONGE_WIDTH];
// Initial matrix has first row/column = [1, 0, ..., 0];
// c = 0
result[0] = state[0];
for r in 1..12 {
if r < SPONGE_WIDTH {
for c in 1..12 {
if c < SPONGE_WIDTH {
// NB: FAST_PARTIAL_ROUND_INITIAL_MATRIX is stored in
// row-major order so that this dot product is cache
// friendly.
let t = P::Scalar::from_canonical_u64(
Self::FAST_PARTIAL_ROUND_INITIAL_MATRIX[r - 1][c - 1],
);
result[c] += state[r] * t;
}
}
}
}
result
}
/// Recursive version of `mds_partial_layer_init`.
fn mds_partial_layer_init_circuit<const D: usize>(
builder: &mut CircuitBuilder<Self, D>,
@ -449,6 +564,39 @@ pub trait Poseidon: PrimeField64 {
result
}
/// Same as `mds_partial_layer_fast` for `PackedField.
fn mds_partial_layer_fast_packed_field<
F: RichField + Extendable<D>,
const D: usize,
FE,
P,
const D2: usize,
>(
state: &[P; SPONGE_WIDTH],
r: usize,
) -> [P; SPONGE_WIDTH]
where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>,
{
let s0 = state[0];
let mds0to0 = Self::MDS_MATRIX_CIRC[0] + Self::MDS_MATRIX_DIAG[0];
let mut d = s0 * P::Scalar::from_canonical_u64(mds0to0);
for i in 1..SPONGE_WIDTH {
let t = P::Scalar::from_canonical_u64(Self::FAST_PARTIAL_ROUND_W_HATS[r][i - 1]);
d += state[i] * t;
}
// result = [d] concat [state[0] * v + state[shift up by 1]]
let mut result = [P::ZEROS; SPONGE_WIDTH];
result[0] = d;
for i in 1..SPONGE_WIDTH {
let t = P::Scalar::from_canonical_u64(Self::FAST_PARTIAL_ROUND_VS[r][i - 1]);
result[i] = state[0] * t + state[i];
}
result
}
/// Recursive version of `mds_partial_layer_fast`.
fn mds_partial_layer_fast_circuit<const D: usize>(
builder: &mut CircuitBuilder<Self, D>,
@ -502,6 +650,26 @@ pub trait Poseidon: PrimeField64 {
}
}
/// Same as `constant_layer` for PackedFields.
fn constant_layer_packed_field<
F: RichField + Extendable<D>,
const D: usize,
FE,
P,
const D2: usize,
>(
state: &mut [P; SPONGE_WIDTH],
round_ctr: usize,
) where
FE: FieldExtension<D2, BaseField = F>,
P: PackedField<Scalar = FE>,
{
for i in 0..SPONGE_WIDTH {
state[i] +=
P::Scalar::from_canonical_u64(ALL_ROUND_CONSTANTS[i + SPONGE_WIDTH * round_ctr]);
}
}
/// Recursive version of `constant_layer`.
fn constant_layer_circuit<const D: usize>(
builder: &mut CircuitBuilder<Self, D>,