ZK13: Circle STARK GPU Acceleration: An Analysis of Performance and Implementation

Transcript

[Music] Hello everyone. How are you doing? Hope you're enjoying the event till now. Um my name is Julian. I'm an applied cryptographer and computer scientist. And here at the front row with me is Carlo. He's a mathematician, PhD and computer scient not not computer scientist but applied cryptographer too. Um we work at ERICS where we do zero knowledge uh recently blockchain computer vision LLMs and purpose-built software development and today we wanted to share some uh like knowledge we gathered by doing something cool. Um that's the whole goal of the talk like

sharing something we did with you. Um so yeah thanks to Sik Summit and podcast for giving us the space and let's go. We'll be going over some GPU optimization basics. So bear with me. Uh this talk won't be too technical in terms of GPU, but I'll get into some math which shouldn't be like hard in the context of zero knowledge at least. So GPUs can be used for many things, right? Their original purpose was um uh image and signal processing. That's the gene GPU graphic. But uh people have since harnessed their power for algebraic operations and data intensive computations in general. So um they bring some problems. Uh when one

wants to use the GPU uh for its power uh one finds out that GPUs and CPUs have different memory spaces. Uh so one cannot access the other's memory and copying data from one memory space to the other is expensive. It's costly. um more in the developer quality of life side uh the presence of different memory types uh it makes stuff difficult it it kind of makes it complex so managing memory access is not easy and there's also the fact that they're not CK oriented components so as a developer uh especially trying to build CK over GPUs one finds out that there are no primitives there are no like um uh API facilitations for CK K in

general. So that's like the problems that GPU brings. But there's also benefits. Of course, we wouldn't be having this talk if there weren't. Um when we launch threads to processor data, uh GPUs parallelize this very well. So even if they don't run at the literal same time, uh they do run in parallel. So or in concurrently, I mean. Um, so the GPU is very good at making this uh this code you're running um run as fast as possible by parallelizing stuff and handling its own resources to like uh you know uh help them uh use the the resources of the GPU when they're available.

And apart from that they're increasingly available components. So, every laptop has a GPU nowadays, and even if it's not like specifically the best GPU there is or like the latest Nvidia 1590, whatever, uh you can get um your optimizations running on commodity hardware pretty well. Um we have run some of these uh benchmarks and like laptops and the results were very good, similar even. So, yeah. Uh there's also the fact that uh MacBooks uh for example uh that are closed hardware have the added benefit of us as developers being able to benchmark our programs um in a more consistent way. So everyone that

has uh a specific MacBook cheap chip will have the same GPU maybe with varying like uh VRAM or or or something like that but the parameters are controlled so you know that maybe I don't know whatever generation of MacBook has these three options for VRAM for instance that's another benefit of closed server solutions and in this talk we may focus a little bit on CUDA specifically CUDA is like um uh the Nvidia API, you might have heard of it. Um and there are a lot of architecture specifics about CUDA that help us uh optimize our code even further apart from like the implementation in GPU. For example, uh blocks and warps, memory banks, low

memory, etc. This list is just for you to look at and see that there are possibilities. Uh don't don't bother remembering this. I mean, it's not important. It's just for you to know that there exists uh there exists like an alternative for even more optimization. So um since I will be going u into some examples of optimization, I'll first go into some preliminaries for you to understand these examples or to have context on them. Uh first is stew. Stew uh some of you may know is a prover a CK prover based on circle starks uh and made by starware right and our challenge uh last year was to build a GPU implementation of sue. So sue boasted um at the time to

be the fastest prover there was in terms of hashes per second. And now when we built 2GPU uh for nethermind um we faced the challenge of first understanding the protocol that was a big part of it. Um and since there was no GPU implementation yet uh whatever we were making was like the first iteration of of like STGPU. So that in itself brought value and since it was the first iteration you know um the phrase make it uh make it it work make it right make it fast um of course we didn't make make it fast at first and we just focused on making it work and by

making it work it was very fast I mean we didn't we didn't try optimizing it and it still was very fast you'll see right now um so yeah no optimizations were planned which means that some of you may even think of better optimizations for some stuff we're going to show today. Um, in any case, did it work? I mean, was it faster? Well, I already spoiled that. Like, it was faster. Um, and this is how much faster it was. This graph, in order to help you read it, um, goes like this. So at the very bottom the x-axis uh has some of the steps that stood uh like takes in order to prove a trace right the execution of a trace. So for each step

in the y-axis we have the accumulated time taken to reach that step. So in total a proof um using CPU and AVX which is the orange line takes about it takes about two and a half seconds right our like GPU implementation takes way less than that maybe in the order of like 100 to 200 milliseconds for the same proof which is about 20 times faster and that's just implementing stewing GPU no optimizations no like nitpicking um that's just how fast it Just for completeness, um this is not over uh a specific trace in particular.

We made this over um a simulated trace that that was that's named the white Fibonacci which is made to be big and wide for us to like um test the prover with as much data as possible. Right? This talk will be focusing on examples in these parts which are uh first aerical commitment and then uh some stuff having to do with fry. Um we won't go into detail uh again on on on suss protocol and stuff but just so you know this is where uh today's talk will focus on.

Now the first case study batch inverse. Okay. So these are the parts that uh this first case study uh focuses on, right? Um we'll first review the oh did it skip a slide? No. Okay, perfect. We'll first review the sequential algorithm uh that you know batch inverses as it is implemented in STO. Then we'll parallelize it as it is implemented in 2GPU and then we'll compare them both to see how they fare against each other. So batch inverse uh is an algorithm with the objective of inverting a bunch of field elements.

That's all. So um you have a very big list of field elements and you want their multiplicative inverse. Easy as that. Now um for the sequential algorithm the one that's implemented in regulars too using CPU um what they did was uh since field inversion is a very complicated uh operation and computationally heavy uh what we did was um apply something called Montgomery's trick. Montgomery's trick allows us to circumvent somehow this um expensive operation that is the field inversion uh usually done with the extended ukleon

algorithm. For an example of Montgomery's trick, we'll take these four field elements A1 through to A4 and we'll start inverting them by computing these um little products that I'll show here. First one being beta 1, then beta 2 being the accumulated product from A1 through to A2, then beta 3 being the accumulated product from A1 through to A3 and so on so forth. Um, I'll refer to these as either accumulated or partial products.

Uh, so you know what they are. And just for completion, there's a more efficient way of doing this, which is using the previous accumulated product um to save some multiplications. But um at the end of the day, but calculating beta 4 takes the same time here and here. Uh but you save some computations. So after we have our accumulated products, we want to invert um we want to use the extended algorithm as little as possible. So we use it only once here to invert uh the final accumulated product beta beta sub4. And yeah that's the only uh point here in which we will use this inversion. Then um with our already computed inversion which is this

one we use another partial product which is beta sub3 to cancel out the terms that don't matter to us in order to obtain the inversion of phase sub 4. So intuitively what we're doing here is this all right the first fraction like the the fraction the one over a1 through to a4 is beta 4 to the minus one this is inversion of betas of four easy as that and then beta sub 3 is used to cancel out the terms so you cancel out a1 through to a3 and they're left with a4 in the denominator so it's inverted now we compute uh the same way we compute beta sub3 inverted by doing something very similar. Um we use betas

of four inverted once again but this time we going to cancel out the a4 term this way. So beta sub 4 inverted is this whole thing as I said before and then we cancel out the a4 term in order to get this this is beta sub3 inverted like intuitively right so now we know how to calculate the inversion of one of the elements the final one and one of the partial products each of these calculations only took one multiplication one field multiplication With the same idea, we can compute the inversion of a sub3 using the previous

results. Right? Beta sub 3 inverted is used here. We calculate it here and we use it and we use it below. Then there's also the fact that um we calculate the next beta like the next accumulated product for calculating the next uh inversion. Once again, beta sub 2 uh inverted is used uh to calculate a sub 2 inverted and the idea is the same up until we reach our final element which is a sub one inverted which as you remember is equal to the to the to the

first accumulated product. So you you just invert that and you have your result. I hope everything's clear uh till now, but if anyone has any questions, please save them up for later. I'll make sure to go like in depth with whatever math you missed. But in any case, I think this is easier than most CK math. So yeah, um what we did is for an array of size n, we replaced n field inversions with three n field multiplications and only one field inversion. I don't have the data of how many multiplications are needed for the field inversion, but I can tell you this is a fair trade. like we like this uh this is perfect. So we need a multiplications to to build the accumulated products or the partial products. Then we need to invert

the the final one and then we need two multiplications to compute each element we want to invert. Now you may have observed as I told you when we were explaining the equations that um this cannot be parallelized as is because I had to show you that uh the accumulated products when inverted were used u were calculated like sequentially right so for instance calculating the accumulated products requires the previous ones and when we were deconstructing the inversions that also require the previous values so there's a strictly inherently sequential uh thing going on with the way the the the trick is built or or implemented here. So

we'll adapt it a little bit to for it to be parallelized which is what one would do when when one wants to parallelize your algorithms for GPU. We take the same elements we want to invert but now we'll make an assumption. We'll assume uh the amount of elements to invert is a power of two just for for completion like for for ease of like um building uh a tree like structure which we'll build in the next slide. But uh let's say you want to invert three elements or something that's not a power of two. Um then you just fill the list with whatever number you want to invert as a discarding number like one or whatever. So there's no loss of

generality here. And yeah, let's suppose the amount is power of two. And as I said, we build a treel like structure. First, we calculate uh this new kind of partial product which is beta sub 1 through two. It's it's like it's understandable, right? It only takes a1 and a2. In the other side of the tree, we compute beta sub3 through 4 uh which is a3 * a4. And now uh I'll show you the root which um reveals the true purpose of the notation which is um the way uh or rather the index at which the the multiplication starts or the product starts comma the index at which it ends. So beta sub 1 through 4 means

we're multiplying all the numbers between one and four. In any case even if the notation is not clear I think this is easy to understand. You only take partial products in a tree like structure, binary treel like structure. And now you may notice that in the same level of the tree each multiplication is independent. So uh we can parallelize the same level of the tree. Right? Uh beta sub 1 2 does not depend on anything calculated by beta sub34. This allows us to parallelize this uh accumulated product calculation.

Now when we have the roots of the tree, we invert it the same way as before the whole accumulated product. And now uh we want to uh deconstruct this tree the same way as we did before, but now it's more graphically like understandable. So we use the inversion of beta 1 through 4 to compute the inversion of the left side of the left accumulated product which is beta 1 through two by canceling out the terms that we don't want which are in beta sub3 4. This is how it's done intuitively the same way as before.

um beta sub one through four inverted is always uh the whole product and then we only cancel out a3 and a4 easy as before you get one over a1 and a2 which is beta sub one uh through two inverted easy as that uh we can do the same thing on the other side of the tree right and get beta sub 3a 4 and then once we have uh the partial products we can get the leaves of the of the tree um which are the result of this algorithm actually a sub one inverted uses this partial multiplication inversion and cancels out the term that it doesn't need then a sub2 inverted does the same and so on and so forth I think it's

clear once again if anyone has any questions please let me know later the point is that once again since we built this in a tree like structure we can parallelize each level which allows for the use or the harnessing of GPU power and even after this parallelization we did we could apply further uh could have specific optimizations like I said before like the user shared memory or for example consecutive memory accessing consecutive memory accessing takes advantage of the fact that your data is consecutively stored in memory and uses operations that apply um that that that are applied to these consecutive um datons at the same time for instance all of this helps

speed up our implementation even further. Now this is a performance difference. This is the amount of time taken for each algorithm and well you see it's a very big difference. We did some like heavy math here or like we had to do some work even if it's not difficult we had to do some work for the trick. Um, oh yeah, please. I was going to mention that. That's a good question. I mean, it's the the CPU AVX implementation is the one I I showed before, which was the sequential one, but it also harnesses AVX power like um you know a single interaction, multiple data. There's also

an AVX talk after this one. Um, it's great. I mean Carlo here was uh participating in that implementation but at least with like um today's hardware I mean GPU heavily outperforms CPU even when using AVX um it it's the only thing we verified is this this uh this way of building the products which is using the Merkel tree uh but but I'll give you a slightly different example which does no does not have any like math mumbo jumbo or or any like complicated tricks like Montgomery's. In any case, it's a 75 times uh improvement uh for a list of two to the power of 24

elements. Now, this is how it scales. So the difference I showed you here or in the previous slide uh is this point and um we can see that uh it improves between 66 times and 79 times uh with the use of GPU. Maybe it's not clear why uh to some of you why the graph grows so big but it only grows from 66 to 79. And that's because the the blue line which is the GPU performance also grows a little bit only that um it's barely perceptible perceivable I mean so the next case study oh sorry you

wanted to take a picture in any case you can probably perceive it um the next case study um Merkel trees uh this is done in this step of the proof generation of um stew too which wants to um oh yeah performance yeah this one so does it mean where dot is mer commit merc tree is it where the procedure starts or where it ends sorry I don't understand the question it's cumulative right so yeah this is cumulative yeah so ending

up between 500 and 1,000 accumulated uh time taken. Is it after Merkel tree or at commitment or before? Oh, you mean if this if if this graph is accumulated or not? No. Or you mean this one? So it it happens between the the second dot and the third. Yeah. Yeah. So So really the constraint pointwise able is much more expensive. Well, I mean it's hard to see the difference in the GPU line because it's so thin like No, just on on on CPU. I'm surprised that Oh, yeah. I mean uh I would I would have expected the the huge increase is the mer commit. We weren't

the ones that implemented that, but um I think that's a good thing to look into later. We can talk about that if you're interested after the talk. Um in any case um the algorithm builds a Merkel tree to to commit a bunch of polomial evaluations. So it's not just one Merkel tree right it or I mean it is just one Merkel tree. You want the root as a commitment but it it is a lot of polinomial evaluations. Uh and we want to commit this for the purpose of the protocol right but for this implementation we won't look into the math tricks. We'll only look over some numbers because once again uh there are no math tricks. Now um as a side note we use

Blake 2s for hashing in this example uh which was also implemented in GPU but it works with any other hash function without any mathematical workarounds. Once again we achieved an incredible increment or increase in performance. Um the way we parallelize this is the same way we parallelize um the the treebased Montgomery trick. Uh if we want to parallelize uh the way we build a tree, we usually parallelize each layer of the tree. Since a layer of the tree depends on the one below it, um then we cannot like uh separate those inherently sequential steps because you

need one layer to compute the other which you need to compute the next one that's inherently sequential. Um so what we do here is we parallelize the building of each layer by themselves. The way we do that is just separating the layer into segments or the level into segments, calculating each level in parallel, each segments in parallel, and then joining the results to build the next level. Simple as that. We're just separating a list into segments, computing each segment in separate separately. That's it. No tricks. This is a performance increase. So CPU and AVX as time scales of course takes more and more. GPU is way faster here.

Compared to CPU, it almost seems constant. And once again, this is CPU using AVX. It's not like strictly sequential. It's paralyzing using AV AVX instructions. So, so that takes that takes a lot of work, too. And AVX instructions are great, by the way. Um, it goes from 11 times as fast to being 670 times as fast. Uh, when the data size or the number of valuations goes to two to a 20. So, there's a big difference, right? In any case, as a summary of what I wanted to show you today, what have we learned? Well, first of all, um, GPU paralization is very powerful to make our CK algorithms faster. A simple huristic like the one I

showed you before, separating a list into segments. It goes a long way into making this way faster. Uh, for both high performance and commodity hardware GPUs, we try this in like a 4090, in like a 1080, uh, using laptops, using a tower PC. every time it's way faster and also as a side note I mean having a mixed team of mathematicians and computer scientists was great for this since two is not an easy protocol to understand or even optimize um so yeah this added a layer of of understanding and collaboration that we didn't take into account and and it was great to have that thank you so very much once again thanks to CK summit and thank all of you for being here thanks for hearing Does anyone have any questions?

Was your um CPU baseline code multi-threaded or just AVX? Uh I don't know. That's a good question. You think it was threaded? Yeah. No. Okay. Well, we don't know. I I honestly don't know. Um I think it was not. I I don't remember seeing any of that. Um, once again, we weren't the team that builds the CPU code that was like already done and I don't remember seeing threading, but that's a good question by the way. Yeah. Okay. Anyone else? Brilliant. Well, thank you very much.

Thank you. Obviously the speakers are around all day if you want to grab them for any more questions. Um, but we will end this talk here. So, thank you Julian. [Applause]