ZK12: Accelerating Client-Side Cryptography with WebGPU - Tal Derei

Transcript

[Music] hi uh my name is tall with penumbra and today we'll be talking about accelerating clientside cryptography with web GPU um so web assembly has performance limitations as a platform uh it has a 4 gig upper memory limit because of the 32-bit addressing uh in modern browsers um this 4 gig memory limit can easily be consumed with larger circuits and proving Keys it also has a lack of easily accessible threading it doesn't natively support

multithreading um the way you get around this is you spawn multiple web workers each uh instantiating the wasum module and sharing memory via a shared array buffer uh it also lacks support for a 128bit multiplier um it natively only supports a 64 bit multiplier um in the case of penumbra this specifically causes a Slowdown uh since in P at at penumbra we use arcworks which is rust based and we use 64 bit limbs in the back end which again expects a 128bit multiplier and that's being emulated when lowered to wasum so we use arcworks to generate proofs uh very efficiently in the CLI

but when you take the rust based code and compile it to WM you have this huge emulation blowup cost by a factor of about 25 to 30X um in wasm compared to the native rust code uh CPS are also suboptimal Computing platforms um cryptographic operations uh involve integer computation uh and that's performed uh mainly in the arithmetic logic unit the ALU uh which only accounts for about 5% of the cor's area uh in the Silicon Dy used in the computation uh the ALU is not optimized for cryptographic operations uh the finite field mod Prime

modulus uh has a bitwidth that's much larger than the native 32 or 64bit Target of the device um modular arithmetic um is also performed over large Prime fields and that's about 10 times slower on existing processors uh CPUs also have a physical limitation of having a limited limited number of computing cores available on the chip um so all of these factors collectively result in large overheads in the browser um and CPUs have limited limited throughput and parallelization uh but with parallel architectures uh with thousands of

computing cores like gpus you can dramatically accelerate these clientside ZK workloads uh there's been a lot of work in the GPU acceleration space in the in the last few years uh ingonyama built iCal spark buil uh supernational build spark these are open- Source GPU libraries for uh Exel accelerating common ZK snark workloads like msms and NTTS uh Z prise hosted by Alo has also been uh um a game changer in the GPU acceleration space uh NY Emirate from formly ured now sarfi and matter matter

Labs uh focused on accelerating msms for large constraint counts uh two to the 26 constraints using gpus um also the core there's been a lot of improvements in core to CH chip density in modern gpus um and we expect to see those same quti chip density improvements on mobile gpus um in the last few years uh the M share has been collectively um focusing on like Z kvms like sp1 uh succinct uh succinct sp1 Nexus um polygon zkm and so forth and um

the natural question is with all of this like server side acceleration going on and zkm construction happening are there fast CLI side Z kvms and the answer is not really um if you look at the server side Hardware acceleration space u most zkv Ms are server side optimized so these are gigantic 2 to the 26 constraint VM circuits uh these are um they they have large memory requirements uh sp1 uh sp1 with plony 3 has a 128 gig memory requirement um so the these uh these

aren't feasible to run on a client device um Nexus notably has been working on a Nova folding um zkv uh with uh IVC um which keep the memory requirements under 4 gigs but the trade-off is larger client side proofs uh so the natural question is can web GPU be used to push Beyond uh these boundaries and accelerate browser based cryptography um and we think the answer is yes uh so introducing web GPU a client side GPU API uh client side ZK acceleration is uh in the browser is relatively

underexplored uh and web GPU essentially uh Taps into uh the client's GPU to tackle this problem and accelerate these workloads um Chrome uh shipped a native Chrome shipped native web G GPU support in v113 in 2023 uh before v113 uh the developer had to enable experimental web GPU Flags uh which made the web GPU uh platform very unstable and hard to work with um web GPU is also a successor to the web webgl standard uh it supports GP GPU for general purpose compute rather

than just Graphics uh like webg uh you can think of webgpu as a general abstraction layer um it it's an API that lets you natively access the devices GPU um from the browser uh it's a small abstraction layer between the physical uh discret and integrated GPU and the web browser running the web GPU API in terms of compiling web GPU code uh it's not possible to have fine grained Hardware control instead you write some highle uh Shader code in

wgsl compile it into an intermediate representation that the underlying system expects and that IR gets compiled to to the machine code um so you trust the blackbox compiler to find like the best optimization path uh but all you're writing as a developer is this highle wgsl code and the whole compilation is black booxed uh web GPU has a dedicated Shader language called wgsl it uses is compute shaders under the hood uh compute shaders are the actual GPU code that run on the GPU

uh again web GPU compiles to the native compilation Target of the GPU so the browser takes the Shader wgsl code and compiles it to uh hlsl for direct X12 uh MSL for metal uh spear for Vulcan so these are our IR targets and then compiles the IR to the GPU machine code uh in terms of the threading model it's very similar to Cuda it's a three-dimensional hierarchy of threads uh that you can index to uh span the vector and matrices that encode the computation uh we uh my teammate uh coj and myself competed in this year Z prize

competition uh and we built a web GPU multiscalar multiplication MSM solution for the browser uh this was uh this implementation supports 2 to the 16 to 2 to the 20 constraints uh for both the Twisted Edwards and the BLS 377 elliptic curve um the implementation was inspired by a paper called kozie K uh which implements a sparse Ma matx construction uh so we took the sparse Matrix construction from Koozie K and modified it with highly optimized memory management and threading uh

optimizations catered to the web GPU environment uh so instead of using like traditional pipp Anders algorithm our MSM implementation uses sparse matrices as the foundation of the of the implementation uh the MSM essentially takes the main operations from pippinger and transforms them into a collection of sparse Matrix operations like sparse Matrix transpositions and sparse Matrix Vector multiplications uh all of the elliptic curve points and scalers uh are stored in sparse matrices uh in a memory

efficient storage format called uh compressed sparse row CSR so essentially what we do is we take the highle pippinger operations compile uh rewrite them in sparse Matrix form and compute the full MSM with sparse Matrix uh algorithms uh the vast majority of the operations execute on the GPU uh we wanted to avoid the costly data transfers between the CPU and the GPU so we only transfer data to the GPU when it's uh when it's necessary so we at the beginning of the computation we we take the aine points and scalers convert them to bytes and then uh store those bytes

in storage buffers and then offload that to the GPU um every one of these white boxes in the GPU section is a dedicated dedicated uh Shader um so it's machine code that runs on the GPU and then we do a dat transfer back to the CPU um at the end of the comput of the computation we we transfer uh a small number of bucket sums and finalize the computation on the CPU so most of the computation happens on the GPU and we do a small data transfer at the beginning and at the end um uh in terms of performance we had a 45x average improvement over the

Baseline um so 1.4 seconds versus 62 seconds uh this was for 2 to 20 constraints over the Twisted Edwards elliptic curve um we benchmarked on Arm based machines uh M1 M2 M3 and Linux machines uh for 2 to the 16 to 2 to the 20 constraints um and then 45x was the average improvement over all of those machines uh we I I won't go into detail about the implementation details of the MSM itself but we wrote a blog post uh on the penumbra website that links to um the

implementation the GitHub repository and the documentation that describes the MSM implementation in detail um some reflection and learnings about working in webgpu um we we came to the conclusion that uh simmed Vector instructions um in wasum are a more production ready platform compared to web GPU at the moment uh web wasm has has access to lowlevel wide Sim of assembly instructions and you don't have the same fine grain control with the web GPU abstraction layer so you can do deeper optimizations in the wasm context that

you can't do in webgpu um but the benefit of web GPU is you can actually um web GPU outperforms simmed for larger circuits uh that's our operating thesis that for larger constraint counts webgpu actually beats uh wasm and lowl simed um simmed May intro May introduce a performance regression in the browser um way did some benchmarking and found a 2X performance regression when implementing Montgomery multiplication in Sim compared to non- simmed in Rust um

so the the idea is if you if you do two 32-bit multiplications um you can you can do two 32bit by 32-bit multiplications per simmed operation and um in theory that should be faster than doing non- simed but what what we found was uh wasm has simmed instructions but neither x86 or arm have matching native instructions so that that means Chrome emulates uh simed with multiple instructions which actually ends up being slower than non-sim instructions uh if you look at the wasum

do documentation it says the same thing um that there are semantic differences between the host hardware and the web assembly semantics so you implement SYM code expecting it to be faster but in fact Chrome emulates all of these instructions and you have a 2X performance regression uh when when implementing a a simple Montgomery multiplication um the way you get around this is by using floating Point math so fp64 and fp3 2 instructions do actually run in simmed mode um for instance a sample FMA instruction a fused multiply and AD that lets you do

a a multiplication and and an addition with a single rounding error will run in simed mode in terms of the rounding um rounding behavior is non-deterministic and is platform dependent you can't actually control the rounding um it uses something called rounding to nearest so some sometimes it rounds up sometimes it rounds down uh but the rounding mode can't be uh controlled uh n Nile Emirate uh had a 2018 paper where he implemented uh 52 bit X 52- bit um Montgomery multiplication so it's a 104 104 bit uh

multiplier um but the paper assumed a rounding toward zero because of of cuda's FM RZ instruction so the paper does this floating Point multiplication and then rounds the result down to zero in in his Z prise implementation he referenced the same paper but used 51 bit limbs uh and with FMA simmed op codes uh in wasum uh this f64 relaxed mad is a relaxed simmed op code it's a

wasm enhancement and essentially by moving from 52 bit limbs to 51 bit Limbs and and some clever rounding tricks he solved the non-deterministic rounding and turned non-deterministic rounding errors to be deterministic so the whole idea is instead of doing 52 bit by 52- bit multiplication you you can do 51 bit by 51 bit multiplication and essentially solve this non-deterministic grounding with some some clever rounding tricks um the thing we're looking at now is floating points in web GPU uh web GPU

has a dedicated dedicated built-in FMA instruction uh which we can use to implement Montgomery multiplication using floating points um the way you implement floating points in web GP um is one of two ways you can either avoid the rounding errors or you can handle the rounding errors the way you avoid the rounding errors is you essentially don't overflow the mantisa so in our current zise implementation we implemented the MSM with 13 bit limbs um 13bit limbs was the optimal limb size to avoid uh to minimize the

carry operations and Montgomery multiplications and we store that 13 bit x 13 bit multiplication in a in a u32 so the accumulator is a u3 is a 32bit value um but fp32 will work without any round any rounding necessary if you don't overflow the 2 to the 23 uh if you don't overflow 2 to the 23 because there are 23 explicit bits in the Mena um so if we take our 13 bit Limbs and reduce it to 10 bit limbs we automatically gain uh access to efficient floating Point 32 operations

without any rounding So in theory you can take all of the u32 uh u32 data types and replace them with F32 without handling any of the rounding uh because again the multiplications don't overflow the mantisa um if we keep 13 bit limbs uh we need to handle the rounding explicitly and we need we need to reference Nile's implementation to implement the correct rounding modes uh in terms of pushing the limits of web web GPU um there are differences in in uh the GPU Hardware capability of the device so performance will vary

depending on the capabilities of of the underlying GPU device uh there are also Shader limitations uh if you increase the depths uh the depth and complexity of the wgsl Shader uh that can decrease GPU performance uh and lead to crashes uh we we we implemented a very uh dense in our MSM scheme that crashed an M3 Pro um so the the way you get around that is you take the complex uh Shader and you break it up into multiple Shader Shader components uh uh web GPU also has uh compilation

time so executing GPU code uh requires compiling the wgsl shaders once and then the code is cached so you pay a a large compilation cost upfront uh performance also varies between brow browser versions um as the web GPU standard matures the uh web GPU standard will get faster but older versions older versions of of webgpu will be slower um in terms of squeezing performance uh there are a few tips uh use fewer wgsl shaders uh this isn't the

case with Cuda because a kernel and Cuda zero cost uh but that's not the case in web G GPU uh don't overload the shaders uh if the complexity or depth of the Shader gets too large again it can crash the machine uh efficiently manage storage buffer allocations uh and be aware of the GPU device limits if you go beyond those limits you can also crash the machine uh we expect the web GPU API to evolve over the the coming years uh but there are some areas that have not been explored uh one of them is Multi-Device proof generation where you take multiple

web GPU instances uh peer-to-peer networked across multiple devices to share and compute a workload um also new low-level features like storing uh big in integer Limbs and floating points um we also see uh we also hope to see U more low-level features uh implemented in web GPU uh for example a dedicated uh Assembly Language for wgsl to do low-level optimizations and uh we we we'd ideally like to see a lot of the Cuda

optimizations find their way to web GPU for example dynamically setting work group sizes uh that's all I have [Applause]