[feature] integrate zeknox GPU-acceleration library into gnark

.gitignore

@@ -55,4 +55,5 @@ gnarkd/circuits/**
 go.work
 go.work.sum
 
-examples/gbotrel/**
+examples/gbotrel/**
+build/

README.md

@@ -161,6 +161,40 @@ func main() {
 
 ### GPU Support
 
+#### Zeknox Library
+Unlock free GPU acceleration with [OKX zeknox library](https://github.com/okx/zeknox).
+
+##### Download prebuilt binaries
+```sh
+curl -L -o libzeknox.a https://github.com/okx/zeknox/releases/download/v1.0.0/bn254-msm-86-89-90-libzeknox.a
+curl -L -o libblst.a https://github.com/okx/zeknox/releases/download/v1.0.0/libblst.a
+sudo cp libblst.a libzeknox.a /usr/local/lib/
+```
+
+If you want to build from source, see guide in https://github.com/okx/zeknox.
+
+##### Enjoy GPU
+
+`groth16.Prove(r1cs, pk, witnessData, backend.WithZeknoxAcceleration())`
+
+```sh
+go run -tags=zeknox examples/zeknox/main.go
+# (place -tags before the filename)
+```
+
+##### Test
+Add following code to [mimc_test.go](examples/mimc/mimc_test.go)
+```go
+assert.ProverSucceeded(&mimcCircuit, &Circuit{
+		PreImage: "16130099170765464552823636852555369511329944820189892919423002775646948828469",
+		Hash:     "12886436712380113721405259596386800092738845035233065858332878701083870690753",
+	}, test.WithCurves(ecc.BN254), test.WithProverOpts(backend.WithZeknoxAcceleration()))
+```
+
+```sh
+go test github.com/consensys/gnark/examples/mimc -tags=prover_checks,zeknox
+```
+
 #### Icicle Library
 
 The following schemes and curves support experimental use of Ingonyama's Icicle GPU library for low level zk-SNARK primitives such as MSM, NTT, and polynomial operations:
@@ -178,7 +212,7 @@ You can then toggle on or off icicle acceleration by providing the `WithIcicleAc
 ```go
     // toggle on
     proofIci, err := groth16.Prove(ccs, pk, secretWitness, backend.WithIcicleAcceleration())
-    
+
     // toggle off
     proof, err := groth16.Prove(ccs, pk, secretWitness)
 ```

backend/backend.go

@@ -133,6 +133,19 @@ func WithProverKZGFoldingHashFunction(hFunc hash.Hash) ProverOption {
 	}
 }
 
+// WithZeknoxAcceleration requests to use [ZEKNOX] GPU proving backend for the
+// prover. This option requires that the program is compiled with `zeknox` build
+// tag and the ZEKNOX dependencies are properly installed. See [ZEKNOX] for
+// installation description.
+//
+// [ZEKNOX]: https://github.com/okx/zeknox
+func WithZeknoxAcceleration() ProverOption {
+	return func(pc *ProverConfig) error {
+		pc.Accelerator = "zeknox"
+		return nil
+	}
+}
+
 // WithIcicleAcceleration requests to use [ICICLE] GPU proving backend for the
 // prover. This option requires that the program is compiled with `icicle` build
 // tag and the ICICLE dependencies are properly installed. See [ICICLE] for

backend/groth16/bn254/zeknox/doc.go

@@ -0,0 +1,2 @@
+// Package zeknox_bn254 implements zeknox acceleration for BN254 Groth16 backend.
+package zeknox_bn254

backend/groth16/bn254/zeknox/marshal_test.go

@@ -0,0 +1,67 @@
+package zeknox_bn254_test
+
+import (
+	"bytes"
+	"testing"
+
+	"github.com/consensys/gnark-crypto/ecc"
+	"github.com/consensys/gnark/backend/groth16"
+	groth16_bn254 "github.com/consensys/gnark/backend/groth16/bn254"
+	zeknox_bn254 "github.com/consensys/gnark/backend/groth16/bn254/zeknox"
+	cs_bn254 "github.com/consensys/gnark/constraint/bn254"
+	"github.com/consensys/gnark/frontend"
+	"github.com/consensys/gnark/frontend/cs/r1cs"
+	"github.com/consensys/gnark/test"
+)
+
+type circuit struct {
+	A, B frontend.Variable `gnark:",public"`
+	Res  frontend.Variable
+}
+
+func (c *circuit) Define(api frontend.API) error {
+	api.AssertIsEqual(api.Mul(c.A, c.B), c.Res)
+	return nil
+}
+
+func TestMarshalNative(t *testing.T) {
+	assert := test.NewAssert(t)
+	ccs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &circuit{})
+	assert.NoError(err)
+	tCcs := ccs.(*cs_bn254.R1CS)
+	nativePK := groth16_bn254.ProvingKey{}
+	nativeVK := groth16_bn254.VerifyingKey{}
+	err = groth16_bn254.Setup(tCcs, &nativePK, &nativeVK)
+	assert.NoError(err)
+
+	pk := groth16.NewProvingKey(ecc.BN254)
+	buf := new(bytes.Buffer)
+	_, err = nativePK.WriteTo(buf)
+	assert.NoError(err)
+	_, err = pk.ReadFrom(buf)
+	assert.NoError(err)
+	if pk.IsDifferent(&nativePK) {
+		t.Error("marshal output difference")
+	}
+}
+
+func TestMarshalZeknox(t *testing.T) {
+	assert := test.NewAssert(t)
+	ccs, err := frontend.Compile(ecc.BN254.ScalarField(), r1cs.NewBuilder, &circuit{})
+	assert.NoError(err)
+	tCcs := ccs.(*cs_bn254.R1CS)
+	zePK := zeknox_bn254.ProvingKey{}
+	VK := groth16_bn254.VerifyingKey{}
+	err = zeknox_bn254.Setup(tCcs, &zePK, &VK)
+	assert.NoError(err)
+
+	nativePK := groth16_bn254.ProvingKey{}
+	buf := new(bytes.Buffer)
+	_, err = zePK.WriteTo(buf)
+	assert.NoError(err)
+	_, err = nativePK.ReadFrom(buf)
+	assert.NoError(err)
+	if zePK.IsDifferent(&nativePK) {
+		t.Error("marshal output difference")
+	}
+}

backend/groth16/bn254/zeknox/nozeknox.go

@@ -0,0 +1,34 @@
+//go:build !zeknox
+
+package zeknox_bn254
+
+import (
+	"fmt"
+
+	"github.com/consensys/gnark/backend"
+	groth16_bn254 "github.com/consensys/gnark/backend/groth16/bn254"
+	"github.com/consensys/gnark/backend/witness"
+	cs "github.com/consensys/gnark/constraint/bn254"
+)
+
+type ProvingKey struct {
+	groth16_bn254.ProvingKey
+}
+
+const HasZeknox = false
+
+func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...backend.ProverOption) (*groth16_bn254.Proof, error) {
+	return nil, fmt.Errorf("zeknox backend requested but program compiled without 'zeknox' build tag")
+}
+
+func Setup(r1cs *cs.R1CS, pk *ProvingKey, vk *groth16_bn254.VerifyingKey) error {
+	return groth16_bn254.Setup(r1cs, &pk.ProvingKey, vk)
+}
+
+func DummySetup(r1cs *cs.R1CS, pk *ProvingKey) error {
+	return groth16_bn254.DummySetup(r1cs, &pk.ProvingKey)
+}
+
+func (pk *ProvingKey) Free() {
+	// nothing to do here
+}

backend/groth16/bn254/zeknox/provingkey.go

@@ -0,0 +1,55 @@
+//go:build zeknox
+
+package zeknox_bn254
+
+import (
+	"github.com/consensys/gnark-crypto/ecc/bn254"
+	groth16_bn254 "github.com/consensys/gnark/backend/groth16/bn254"
+	cs "github.com/consensys/gnark/constraint/bn254"
+	"github.com/okx/zeknox/wrappers/go/device"
+)
+
+type deviceInfo struct {
+	G1Device struct {
+		A, B, K, Z DevicePoints[bn254.G1Affine]
+	}
+	G2Device struct {
+		B DevicePoints[bn254.G2Affine]
+	}
+	InfinityPointIndicesK []int
+}
+
+type ProvingKey struct {
+	groth16_bn254.ProvingKey
+	*deviceInfo
+}
+
+type DevicePoints[T bn254.G1Affine | bn254.G2Affine] struct {
+	*device.HostOrDeviceSlice[T]
+	// Gnark points are in Montgomery form
+	// After 1 GPU MSM, points in GPU are converted to affine form
+	// Pass it to MSM config
+	Mont bool
+}
+
+func Setup(r1cs *cs.R1CS, pk *ProvingKey, vk *groth16_bn254.VerifyingKey) error {
+	return groth16_bn254.Setup(r1cs, &pk.ProvingKey, vk)
+}
+
+func DummySetup(r1cs *cs.R1CS, pk *ProvingKey) error {
+	return groth16_bn254.DummySetup(r1cs, &pk.ProvingKey)
+}
+
+// You should call this method to free the GPU memory
+//
+// pk := groth16.NewProvingKey(ecc.BN254)
+// defer pk.(*zeknox_bn254.ProvingKey).Free()
+func (pk *ProvingKey) Free() {
+	if pk.deviceInfo != nil {
+		pk.deviceInfo.G1Device.A.Free()
+		pk.deviceInfo.G1Device.B.Free()
+		pk.deviceInfo.G1Device.K.Free()
+		pk.deviceInfo.G1Device.Z.Free()
+		pk.deviceInfo.G2Device.B.Free()
+	}
+}