FEAT: Add experimental support for Icicle GPU acceleration behind build tag

[jeremyfelder]

Description

This PR adds experimental support for the Icicle GPU acceleration library using the gpu build tag for Groth16 backend and curve BN254.

Type of change

• New feature (non-breaking change which adds functionality)
• This change requires a documentation update

How has this been tested?

• Using Ingonyama's fork of Celer Network's bls-sig circuit.

1. Follow the instructions for setting up Icicle in the linked fork above
2. From brevis-cricuits/fabric/bls-sig run:

go run main.go -tags=gpu

Checklist:

• I have performed a self-review of my code
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• I did not modify files generated from templates
• golangci-lint does not output errors locally
• New and existing unit tests pass locally with my changes
• Any dependent changes have been merged and published in downstream modules

[CLAassistant]

All committers have signed the CLA.

[ivokub]

See a few comments.

On a high level I think it is a very good addition to gnark. However, I would restructure the GPU prover a bit. Instead of having the implementation inside backend/groth16/bn254, I would instead create a new backend backend/groth16/icicle_bn254 and a backend.ProverOption which would direct to use icicle backend instead of the default one. And in icicle backend I would try to reuse existing structures from CPU backend as much as possible (including marshaling and setup).

Additionally, I would rename the build tag to be icicle instead of being gpu - I think it is a bit more descriptive and maybe allows in the future to have similarly other GPU backends unambiguously.

If possible, I would also change the ./setup_icicle.sh script in the other repo - right now it assumes that $GOPATH is always ~/go/, but this may not be the case, you should refer to $GOPATH explicitly. The setup.sh script also checks if the user is root or not, but does not fail if the user is not. Additionally, it doesn't really install libbn254.so, but instead changes $LD_LIBRARY_PATH temporarily. It is not important in the context of this PR, but I think it would make GPU bindings definitely easier to use. Right now I had to debug the setup scripts to figure out why in other terminal the example was not running.

As a side-note - I ran pairing and scalarmul circuits on Ryzen7 and RTX 3090 machine - the prover was approx 4x faster, but it doesn't take into account the proving key copy from RAM to GPU.

If you would be interested then I could propose some code suggestions about refactoring GPU prover into separate backend.

[ivokub]

Suggested edit:

diff --git a/go.mod b/go.mod
index 627710bf..8df92783 100644
--- a/go.mod
+++ b/go.mod
@@ -19,11 +19,10 @@ require (
 	golang.org/x/sync v0.3.0
 )
 
-require github.com/rogpeppe/go-internal v1.11.0 // indirect
-
 require (
 	github.com/davecgh/go-spew v1.1.1 // indirect
 	github.com/ingonyama-zk/icicle v0.0.0-20230928131117-97f0079e5c71 // indirect
+	github.com/rogpeppe/go-internal v1.11.0 // indirect
 	github.com/mattn/go-colorable v0.1.13 // indirect
 	github.com/mattn/go-isatty v0.0.19 // indirect
 	github.com/mmcloughlin/addchain v0.4.0 // indirect

[jeremyfelder]

@ivokub Thanks for the review 💪🏻 🚀

On a high level I think it is a very good addition to gnark. However, I would restructure the GPU prover a bit. Instead of having the implementation inside backend/groth16/bn254, I would instead create a new backend backend/groth16/icicle_bn254 and a backend.ProverOption which would direct to use icicle backend instead of the default one. And in icicle backend I would try to reuse existing structures from CPU backend as much as possible (including marshaling and setup).

Wouldn't this cause there to be a lot of duplicated code between backend/groth16/icicle_bn254 and backend/groth16/bn254?
Would this be instead of the build tag or in addition to the build tag?

Additionally, I would rename the build tag to be icicle instead of being gpu - I think it is a bit more descriptive and maybe allows in the future to have similarly other GPU backends unambiguously.

Updated this 👍🏻

If possible, I would also change the ./setup_icicle.sh script in the other repo - right now it assumes that $GOPATH is always ~/go/, but this may not be the case, you should refer to $GOPATH explicitly. The setup.sh script also checks if the user is root or not, but does not fail if the user is not. Additionally, it doesn't really install libbn254.so, but instead changes $LD_LIBRARY_PATH temporarily. It is not important in the context of this PR, but I think it would make GPU bindings definitely easier to use. Right now I had to debug the setup scripts to figure out why in other terminal the example was not running.

Thanks for the suggestions! This was mainly a quick way to set it up but yes we will make the process more reliable

[ivokub]

@ivokub Thanks for the review 💪🏻 🚀

On a high level I think it is a very good addition to gnark. However, I would restructure the GPU prover a bit. Instead of having the implementation inside backend/groth16/bn254, I would instead create a new backend backend/groth16/icicle_bn254 and a backend.ProverOption which would direct to use icicle backend instead of the default one. And in icicle backend I would try to reuse existing structures from CPU backend as much as possible (including marshaling and setup).

Wouldn't this cause there to be a lot of duplicated code between backend/groth16/icicle_bn254 and backend/groth16/bn254? Would this be instead of the build tag or in addition to the build tag?

Additionally, I would rename the build tag to be icicle instead of being gpu - I think it is a bit more descriptive and maybe allows in the future to have similarly other GPU backends unambiguously.

Updated this 👍🏻

If possible, I would also change the ./setup_icicle.sh script in the other repo - right now it assumes that $GOPATH is always ~/go/, but this may not be the case, you should refer to $GOPATH explicitly. The setup.sh script also checks if the user is root or not, but does not fail if the user is not. Additionally, it doesn't really install libbn254.so, but instead changes $LD_LIBRARY_PATH temporarily. It is not important in the context of this PR, but I think it would make GPU bindings definitely easier to use. Right now I had to debug the setup scripts to figure out why in other terminal the example was not running.

Thanks for the suggestions! This was mainly a quick way to set it up but yes we will make the process more reliable

Can you have a look at my attempt at branch feat/icicle-integration. It is mostly copy-paste, but on a high level what I did:

• first, I moved all the GPU implementation to separate package backend/groth16/bn254/icicle. Additionally, I separated the added stuff in the ProvingKey as deviceInfo and created new ProvingKey which embeds both "native" PK and deviceInfo. By embedding, we can retain the existing marshal/unmarshal methods.
• then, I added a prover option which allows to optionally use Icicle backend. This allows to use both CPU and GPU prover without needing to change build tags. If -tags=icicle is not set then the option is a no-op and we use native prover.
• made the CPU copies lazy within Prove method. It is not perfect, but we could also overload the methods icicle_bn254.ProvingKey.ReadFrom and icicle_bn254.Setup so that they would first call corresponding method on the embedded fields and then copy.

Now, for example if I have:

       proof, err := groth16.Prove(ccs, pk, secretWitness)
        if err != nil {
                panic(err)
        }
        // lazy precompute and copy
        proofIci, err := groth16.Prove(ccs, pk, secretWitness, backend.WithIcicleAcceleration())
        if err != nil {
                panic(err)
        }
        // already precomputed and copied
        proofIci2, err := groth16.Prove(ccs, pk, secretWitness, backend.WithIcicleAcceleration())
        if err != nil {
                panic(err)
        }
        err = groth16.Verify(proof, vk, publicWitness)
        if err != nil {
                panic(err)
        } 
        err = groth16.Verify(proofIci, vk, publicWitness)
        if err != nil {
                panic(err)
        }
        err = groth16.Verify(proofIci2, vk, publicWitness)
        if err != nil {
                panic(err)
        }

Then in case we run it without icicle build tag we run it three times with CPU backend. But with icicle backend we first run with CPU, second time with GPU but with precompute and third time GPU without precompute (using existing memory).

[ivokub]

And if this approach looks good then you can merge my feat/icicle-integration into your branch. I think I cannot work on your branch :) Lets try using this PR to keep all the contributions intact 👍

[jeremyfelder]

Can you have a look at my attempt at branch feat/icicle-integration.

Looks great 💪🏻 . Love the option to toggle acceleration without rebuilding 🔥

I find it a bit strange to use the HasIcicle boolean in addition to the build tag and prover option but I see that its the simplest way to reduce refactoring groth16.go to accept ProverOptions

I do think it might make adding additional acceleration options difficult but I guess that can be a future problem 😅

[ivokub]

Yeah, it isn't ideal to have HasIcicle constant, but I couldn't figure out a nice way to otherwise keep the GPU pointers. One option is to have inside icicle package a local map from *ProvingKey to *deviceInfo, but this affects GC which cannot collect ProvingKey structs anymore nicely. Alternatively, we could amend vanilla (i.e. the structure we embed into icicle.ProvingKey) bn254.ProvingKey so that it can have additional backend-specific context. But as you said - I think it is difficult to figure out right now how it should look like and we can try to generalise in the future when we have more backends and intuition how we should store backend-related data.

Otherwise, looks good from my side, but will just ask @gbotrel opinion.

[gbotrel]

lgtm 👍 amazing work :)