perf(poseidon): eliminate heap allocations in poseidon (#712)

This commit is contained in:
Ömer Faruk Irmak 2024-04-24 11:12:27 +03:00 committed by GitHub
parent 79c356002c
commit 629f14ea88
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 45 additions and 44 deletions

View file

@ -16,6 +16,7 @@ import (
const NROUNDSF = 8 //nolint:golint
var NROUNDSP = []int{56, 57, 56, 60, 60, 63, 64, 63, 60, 66, 60, 65, 70, 60, 64, 68} //nolint:golint
const MAX_WIDTH = 18 // len(NROUNDSP)+2
func zero() *ff.Element {
return ff.NewElement()
@ -28,37 +29,38 @@ func exp5(a *ff.Element) {
}
// exp5state perform exp5 for whole state
func exp5state(state []*ff.Element) {
for i := 0; i < len(state); i++ {
exp5(state[i])
func exp5state(state []ff.Element, t int) {
for i := 0; i < t; i++ {
exp5(&state[i])
}
}
// ark computes Add-Round Key, from the paper https://eprint.iacr.org/2019/458.pdf
func ark(state []*ff.Element, c []*ff.Element, it int) {
for i := 0; i < len(state); i++ {
state[i].Add(state[i], c[it+i])
func ark(state []ff.Element, c []*ff.Element, it int, t int) {
for i := 0; i < t; i++ {
state[i].Add(&state[i], c[it+i])
}
}
// mix returns [[matrix]] * [vector]
func mix(state []*ff.Element, t int, m [][]*ff.Element) []*ff.Element {
func mix(state []ff.Element, t int, m [][]*ff.Element) []ff.Element {
mul := zero()
newState := make([]*ff.Element, t)
newState := make([]ff.Element, MAX_WIDTH)
for i := 0; i < t; i++ {
newState[i] = zero()
}
for i := 0; i < len(state); i++ {
newState[i].SetUint64(0)
for j := 0; j < len(state); j++ {
mul.Mul(m[j][i], state[j])
newState[i].Add(newState[i], mul)
for j := 0; j < t; j++ {
mul.Mul(m[j][i], &state[j])
newState[i].Add(&newState[i], mul)
}
}
return newState
for i := 0; i < t; i++ {
state[i].Set(&newState[i])
}
return state
}
func permute(state []*ff.Element, t int) []*ff.Element {
func permute(state []ff.Element, t int) []ff.Element {
nRoundsF := NROUNDSF
nRoundsP := NROUNDSP[t-2]
@ -67,41 +69,40 @@ func permute(state []*ff.Element, t int) []*ff.Element {
M := c.m[t-2]
P := c.p[t-2]
ark(state, C, 0)
ark(state, C, 0, t)
for i := 0; i < nRoundsF/2-1; i++ {
exp5state(state)
ark(state, C, (i+1)*t)
exp5state(state, t)
ark(state, C, (i+1)*t, t)
state = mix(state, t, M)
}
exp5state(state)
ark(state, C, (nRoundsF/2)*t)
exp5state(state, t)
ark(state, C, (nRoundsF/2)*t, t)
state = mix(state, t, P)
for i := 0; i < nRoundsP; i++ {
exp5(state[0])
state[0].Add(state[0], C[(nRoundsF/2+1)*t+i])
exp5(&state[0])
state[0].Add(&state[0], C[(nRoundsF/2+1)*t+i])
mul := zero()
newState0 := zero()
for j := 0; j < len(state); j++ {
mul.Mul(S[(t*2-1)*i+j], state[j])
for j := 0; j < t; j++ {
mul.Mul(S[(t*2-1)*i+j], &state[j])
newState0.Add(newState0, mul)
}
for k := 1; k < t; k++ {
mul = zero()
state[k] = state[k].Add(state[k], mul.Mul(state[0], S[(t*2-1)*i+t+k-1]))
state[k].Add(&state[k], mul.Mul(&state[0], S[(t*2-1)*i+t+k-1]))
}
state[0] = newState0
state[0].Set(newState0)
}
for i := 0; i < nRoundsF/2-1; i++ {
exp5state(state)
ark(state, C, (nRoundsF/2+1)*t+nRoundsP+i*t)
exp5state(state, t)
ark(state, C, (nRoundsF/2+1)*t+nRoundsP+i*t, t)
state = mix(state, t, M)
}
exp5state(state)
exp5state(state, t)
return mix(state, t, M)
}
@ -116,8 +117,8 @@ func HashWithCap(inpBI []*big.Int, width int, nBytes int64) (*big.Int, error) {
if width < 2 {
return nil, fmt.Errorf("width must be ranged from 2 to 16")
}
if width-2 > len(NROUNDSP) {
return nil, fmt.Errorf("invalid inputs width %d, max %d", width, len(NROUNDSP)+1) //nolint:gomnd,lll
if width > MAX_WIDTH {
return nil, fmt.Errorf("invalid inputs width %d, max %d", width, MAX_WIDTH) //nolint:gomnd,lll
}
// capflag = nBytes * 2^64
@ -127,11 +128,8 @@ func HashWithCap(inpBI []*big.Int, width int, nBytes int64) (*big.Int, error) {
capflag.Mul(capflag, ff.NewElement().SetBigInt(pow64))
// initialize the state
state := make([]*ff.Element, width)
state[0] = capflag
for i := 1; i < width; i++ {
state[i] = zero()
}
state := make([]ff.Element, MAX_WIDTH)
state[0] = *capflag
rate := width - 1
i := 0
@ -139,7 +137,7 @@ func HashWithCap(inpBI []*big.Int, width int, nBytes int64) (*big.Int, error) {
for {
// each round absorb at most `rate` elements from `inpBI`
for j := 0; j < rate && i < len(inpBI); i, j = i+1, j+1 {
state[j+1].Add(state[j+1], ff.NewElement().SetBigInt(inpBI[i]))
state[j+1].Add(&state[j+1], ff.NewElement().SetBigInt(inpBI[i]))
}
state = permute(state, width)
if i == len(inpBI) {
@ -164,10 +162,13 @@ func HashFixedWithDomain(inpBI []*big.Int, domain *big.Int) (*big.Int, error) {
if !utils.CheckBigIntArrayInField(inpBI[:]) {
return nil, errors.New("inputs values not inside Finite Field")
}
inp := utils.BigIntArrayToElementArray(inpBI[:])
inp := make([]ff.Element, MAX_WIDTH)
for idx, bi := range inpBI {
inp[idx].SetBigInt(bi)
}
state := make([]*ff.Element, t)
state[0] = ff.NewElement().SetBigInt(domain)
state := make([]ff.Element, MAX_WIDTH)
state[0] = *ff.NewElement().SetBigInt(domain)
copy(state[1:], inp[:])
state = permute(state, t)

View file

@ -173,6 +173,6 @@ func BenchmarkPoseidonHash(b *testing.B) {
bigArray4 := []*big.Int{b1, b2, b0, b0, b0, b0}
for i := 0; i < b.N; i++ {
HashFixed(bigArray4) //nolint:errcheck,gosec
HashFixedWithDomain(bigArray4, b0) //nolint:errcheck,gosec
}
}

View file

@ -24,7 +24,7 @@ import (
const (
VersionMajor = 5 // Major version component of the current release
VersionMinor = 3 // Minor version component of the current release
VersionPatch = 2 // Patch version component of the current release
VersionPatch = 3 // Patch version component of the current release
VersionMeta = "mainnet" // Version metadata to append to the version string
)