ZK7: A Lightning-fast Poseidon Implementation - Jakub Nabaglo - Polygon Zero

Transcript

[Music] um i got my friend to draw this title slide for me and i swear that talk just goes downhill from here so um my name is jacob i work at polygon zero um i'm an engineer and my job is basically to make things fast um one of the things that we're really really proud of at polygon zero is plonky two which is like a stupid fast zero knowledge proof library and we can do recursive proofs in less than 200 milliseconds which is really really fast well how did we get there well there's a lot of mathematical tricks like involved in in doing this but there's also a lot of

low-level hardware work that we had to do in order to bring it to this particular speed um in particular with hashing even now we still spend roughly 50 percent of our proof time just on hashes um and this is after like all the stuff like that i'm about to talk about it used to be way way higher so um i'd like to introduce to some of the tricks that we used in order to make hashing particularly fast in order to make plonky 2 particularly fast the hash function that we use is poseidon and the way it works if you haven't seen it before is you have some sort of an input so so s

is for state um and this is a vector of length 12 because like in our particular case it can be a different length but we have technical reasons to make it 12. and there's a few um transformations that happen to this before hello um yeah there's a few things that we do it in order to scramble it up um so it's difficult to just find things that will collide so the first thing um is sewickle is the constant layer right so so we take the um the state vector and we just do element wise addition by

constant uh by some random constants and this is kind of this is fairly easy then we get to s boxes um so that stands for substitution box and this is basically exponentiation this is still element wise um so but it's an expensive operation that adds a little bit of security and um the mds layer the mbs matrix is a type of matrix this is basically just matrix multiplication and the point is to mix all the elements up nicely so every index affects every other index so the state vector makes its way through all the layers and then back up again

a bunch of times each set of all three rounds of three all three layers is called a round and eventually once we're tired and we don't feel like doing any more rounds we got some result out which is also length 12. so one thing that we really need if we want to do these kinds of algebraic hashes quickly is a field to work with that's that's very fast and we need to be able to do you know your exponentiation and multiplication matrix multiplication probably not addition like that that's always easy very fast in order to do hashing quickly

we use something called the called the goldilocks field um let me explain what a field is like in a finite field you basically take all your integers right and you kind of roll them up into a little circle so when you go you kind of wrap around so you you five plus one is six six plus one becomes zero z and um it turns out that uh multiplication works like a similar way um and interestingly it's from this rolling up that the term zika roll-up comes from um

now the number of elements in this circle has to be prime so um it was an interesting challenge that hamish here mainly did is to find a really really good prime um to use here like it's seven here but like we obviously can't use seven we use the gold deluxe prime my favorite prime is 57 but the goldilocks prime is my number two favorite prime and this is the prime that we use for the goldilocks field 57 um it's that which looks a little bit ugly fair enough

um but then you see this and it's kind of nice or if you express it in hexadecimal it's that and it turns out that since we have machines that are 60 for bit um but also often support 32-bit operations um this particular structure ends up oh that's supposed to be another zero there ends up um just permitting some tricks that happen to work on this particular prime field and it turns out that a lot of algebraic operations are fast each goldilocks prime

goldilocks field element fits in a 64-bit word so for example this is an example of a goldlocks field member here's one but that's also one so we have redundant representations addition well we take two things and we add do the normal addition but if it overflows a 64-bit word then we need to compensate for that so in this case like we have a carry that kind of came out of here um and if we get a carry that turns into this this is just 2 to the 32 minus 1. so you can think of it as a as the 32-bit integer that's like just a ones and you just add that on

um in order to compensate for the overflow occasionally this might overflow again in which case you just kind of repeat and that can't overflow in assembly it looks like this 10 points if you can read this um it's nice and short and you don't actually have to keep this constant in like a register anywhere it falls out of here i'm not going to explain it multiplication if we want to multiply two elements of the goldlocks field we first start with like a full 64-bit by 64-bit 228 with multiplication and then we need to

figure out how to make how to reduce this so it fits in 64 bits well what do we do firstly bear with me we break it up into three chunks so the low half goes here then the high quarter goes here and the kind of middle quarter that's not not the one that's in part of the low half goes here like we've just basically just rewritten this um but it also turns out that this is just negative 1 and this is just 2 to the 32 minus 1.

so it's just a subtraction and an addition this can be very easily computed with just a bit shift and a subtraction so i know that the example might be a little bit look a little bit complicated but it's really quite easy you do a full multiplication you split it up subtract and add and it's really really fast so let's get to the actual hashing bit we looked at the structure of the poseidon hash and let's talk about s-boxes so exponentiation

i kind of um missed a detail before because there's two kinds of sbox layers we have the one that's used in full rounds where we basically in our case take every element of the state to the power of seven and we have the one in the partial round where we only do it to the very first one and we just don't care about the rest um so let's focus on well firstly this is basically how do um power of seven uh we can break it up into um a square a multiply another square they multiply so it's just for operations to do that let's talk about the full rounds

one thing that you might notice here is that we are doing the same operation to a lot of data so you might naturally think well can we use vector instructions for this so what do i mean by that normally when you have calculations executing on a computer you use scalar instructions so your computer normally does kind of one thing at a time so say you have two inputs like just two numbers and you want to add them together um a scalar instruction will just take those two numbers and spit out one output done but a lot of architectures have vector instructions um that will let you

do um one operation but to a lot of data at once with just the one instruction so in our case we have um four numbers we want to like two vectors or phone numbers i should say add them together we get a result that's a vector of four numbers um so can we use this well the first issue is this so avx-2 which is like the vector instruction set in in x86 uh doesn't actually have a full multiplication instruction so um straight up like we know this is at least going to be difficult however it does have a 32-bit by 32-bit

full multiplication so maybe we can do something with that well if you think about full 64-bit multiplication what you can do is just split up both factors into the high and low halves do like four 32-bits full multiplies like pair-wise and put them together and in fact i've done this and the code is that um don't don't try and read this um my point is it's messy and you might say okay but if that's on four elements it's just one instruction to do it on

one element so surely this is actually slower and yes you're right we might as well go home no um yes so full multiplication is actually slower with vector instructions however remember that the full multiplication is just the very first step of gold deluxe multiplication um reductions are actually way faster in vector and in fact interestingly enough um squares are faster in vector as well um just because you get to reuse some of the like you don't have to recompute a bunch of values um so we do actually use vector instructions for the s boxes and

it does give us does give us a really nice speed up now um for the partial rounds we can't do the same thing right like we're applying one operation to kind of one thing we don't have enough work to apply vector instructions so well firstly we do have to do it in scala in order to minimize the latency that's kind of what we're bound by um but but there is another thing we can do which is to try and hide that latency a little bit if we look at the mds matrix multiplication that directly follows the sbox layer

so basically imagine that each of these things is a number i i just couldn't be bothered actually putting stuff in there um and this is just a matrix that we're multiplying this vector by what we can do is just split both of them up and perform matrix multiplication on the remainder of the state vector um and like the part of the matrix that's actually relevant and this can run in parallel with the s box um and after we actually have s naught to the power of seven then we can tack this on which is just going to be faster because it's less alrighty now let's talk about the mds matrix um

an mds matrix is a matrix that meets particular security criteria that i don't know about but we need to multiply by it the matrix that we use is this one the numbers don't mean anything in particular to us but but you might make some observations about this matrix firstly you might notice that it's succulent like each row is just the row above it just kind of rotates it by one it's also all powers of two um and all the values in it are small how is this useful well if you have a circular matrix

there's just way fewer constants for you to keep track of or like and in on architectures where you do have to keep them in registers like you don't have to use as many of them powers of two mean that you get to use bit shifts instead of multiplication which in in on some architectures might be helpful and because the values are small enough we can actually prove that we will never have a we will never need to do a reduction of from more than 96 bits so that just makes the reduction at the end a little bit cheaper um and there's a few

and once you have an mds matrix there's a few things that you can do to actually uh make the computation faster so firstly vector instructions uh like work great on this um because there is actually a lot of parallelism in matrix multiplication just inherently you do have to be very careful about how the data is moved around because in vector instructions normally what happens is a vector instruction will do things element wise like it will it can add an element of a vector with the corresponding index of the other vector but you can't do an addition like in between different indices so you have

to use other instructions to like move the data into registers in order to be able to play your arithmetic vector instructions um in our avx avx-2 implementation there's a lot of hacks that we had to do to actually make the moving data around thing fast um and once we had that the rest was trivial you should use multiply accumulate if it's available and if it's fast so in particular um avx2 doesn't have multiplier accumulation integer and i couldn't get it to work well i tried doing this in floating point and it was a bit of a disaster however um

the apple m1 actually has a 32 by 32 to 64-bit multiply accumulate and it can do eight of those per cycle which is really really high throughput and that means that the m1 is actually really fast at computing mds matrix multiplication and the last thing that i will say if you're doing this kind of thing is uh with md with this kind of matrix multiplication you can separate out the high and low halves of the input what i mean by that is basically this equation right like just by linearity if you extract the low half of each element of x like this is just a

vector of 32-bit integers now you can apply the mds matrix you can do the same for the high half multiply that by 2 to the 32 linearity and you will get the same result out and what it means is that you can you never have to deal with values that are greater than that would take up more than 64 bits if you have a 32-bit value like if you do a bit shift or multiplication by a small enough number it will fit in 64 bits and you don't have to worry about things like okay how do i do like addition with carry in vector which is doable but it's annoying and you don't want to have to do it

so this is actually quite crucial to making mps multiplication fast um that's basically it for me um look i introduced you to a bunch of tricks that you can that we had to invent in order to make our poseidon hash fast none of them are actually like particularly profound but the thing is that each of them adds a little bit and it's all of them combined that make our hashing really fast and in turn mean that our prover can be sub 200 milliseconds timing wise it's perfect we have plenty

of time for q a so five minutes is there a mic no just uh repeat the question if you've been uh so what's this what's the speed up over standard deciding hash is it like 2x um the question is like what is the speed up over a standard implementation um uh it would depend on the architecture and but like you're looking at yeah two to three x maybe five look i haven't checked in quite a while and it also depends on how much you've optimized your like integer

implementation that doesn't have these tricks uh um is the room for future of speed ups yes i have a whole list of things to do next that i didn't want to talk about because i haven't done them yet um but there's a bunch of ways there's a bunch of assembly tricks that are like perhaps not really interesting unless you're an assembly not like i am it would actually make it faster doing the mds matrix multiplication in foyer space would actually uh yield some

improvement but you have to be very careful about it so your for your transform doesn't end up like actually costing you more yeah just to extend on that so the idea is in particular that the circular matrix has a fourier transform which is just a single vector in it so you end up having essentially a dot product of two vectors rather than a matrix modification which is part of the motivation for using the circular matrix in the first place but then if you're in fourier space you don't necessarily have that nice mds matrix form with powers of two so instead of doing the search for an mds matrix that looks like that you take a

random nds matrix and then take the fourier transform of that and then look for examples which have a very nice fourier transform and use that instead so you end up with essentially a sort of a dot product with uh in the fourier space with uh small coefficients as well have to stuck is hamish's work to be clear all right time for one more question yeah all right so all of this depends on this very particular field that you're working in is that a field that's actually very common in practice um does it depend on the particular field that we use a lot of these optimizations are enabled by the

goldilocks field and the reason why is you can do reductions with very simple operations and vector instructions tend to be well they can tend to have higher latency for example than scale instructions but they also tend to be much simpler just because you don't want to have that much like die space devoted to them and so the fact that the goldilocks field um has a simple reaction is really really useful for this and if we had just a arbitrary prime like a lot of these things just wouldn't work because you wouldn't be able to vectorize them properly like if your reduction involves a

multiplication you're screwed on the other hand having a nice nds matrix is going to be helpful yeah like a lot of the standard nds matrices are just these i think they're the koshy ones they're essentially random and triggers and then your multiplication is as slow as it's ever going to be all the time all right so thanks for coming in and you can find jack jacob afterwards