Add gc and mem free (#23)

backend/groth16/bls12-377/icicle/icicle.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math/big"
 	"math/bits"
+	"runtime"
 	"sync"
 	"time"
 	"unsafe"
@@ -50,6 +51,7 @@ func (pk *ProvingKey) setupDevicePointers() error {
 	if pk.deviceInfo != nil {
 		return nil
 	}
+	runtime.GC()
 	pk.deviceInfo = &deviceInfo{}
 	n := int(pk.Domain.Cardinality)
 	sizeBytes := n * fr.Bytes
@@ -109,17 +111,24 @@ func (pk *ProvingKey) setupDevicePointers() error {
 	pk.DomainDevice.CosetTable = <-copyCosetDone
 	pk.DenDevice = <-copyDenDone
 
+	// free
+	pk.Domain = fft.Domain{Cardinality: pk.Domain.Cardinality}
+	runtime.GC()
 	/*************************  Start G1 Device Setup  ***************************/
 	/*************************     A      ***************************/
 	pointsBytesA := len(pk.G1.A) * fp.Bytes * 2
 	copyADone := make(chan unsafe.Pointer, 1)
 	go iciclegnark.CopyPointsToDevice(pk.G1.A, pointsBytesA, copyADone) // Make a function for points
-
+	pk.G1Device.A = <-copyADone
+	pk.G1.A = nil
+	runtime.GC()
 	/*************************     B      ***************************/
 	pointsBytesB := len(pk.G1.B) * fp.Bytes * 2
 	copyBDone := make(chan unsafe.Pointer, 1)
 	go iciclegnark.CopyPointsToDevice(pk.G1.B, pointsBytesB, copyBDone) // Make a function for points
-
+	pk.G1Device.B = <-copyBDone
+	pk.G1.B = nil
+	runtime.GC()
 	/*************************     K      ***************************/
 	var pointsNoInfinity []curve.G1Affine
 	for i, gnarkPoint := range pk.G1.K {
@@ -133,33 +142,23 @@ func (pk *ProvingKey) setupDevicePointers() error {
 	pointsBytesK := len(pointsNoInfinity) * fp.Bytes * 2
 	copyKDone := make(chan unsafe.Pointer, 1)
 	go iciclegnark.CopyPointsToDevice(pointsNoInfinity, pointsBytesK, copyKDone) // Make a function for points
-
+	pk.G1Device.K = <-copyKDone
+	pk.G1.K = nil
+	runtime.GC()
 	/*************************     Z      ***************************/
 	pointsBytesZ := len(pk.G1.Z) * fp.Bytes * 2
 	copyZDone := make(chan unsafe.Pointer, 1)
 	go iciclegnark.CopyPointsToDevice(pk.G1.Z, pointsBytesZ, copyZDone) // Make a function for points
-
-	/*************************  End G1 Device Setup  ***************************/
-	pk.G1Device.A = <-copyADone
-	pk.G1Device.B = <-copyBDone
-	pk.G1Device.K = <-copyKDone
 	pk.G1Device.Z = <-copyZDone
-
+	pk.G1.Z = make([]curve.G1Affine, 1)
+	runtime.GC()
 	/*************************  Start G2 Device Setup  ***************************/
 	pointsBytesB2 := len(pk.G2.B) * fp.Bytes * 4
 	copyG2BDone := make(chan unsafe.Pointer, 1)
 	go iciclegnark.CopyG2PointsToDevice(pk.G2.B, pointsBytesB2, copyG2BDone) // Make a function for points
 	pk.G2Device.B = <-copyG2BDone
-
-	/*************************  End G2 Device Setup  ***************************/
-
-	/*free pk data*/
-	pk.G1.A = nil
-	pk.G1.B = nil
-	pk.G1.Z = make([]curve.G1Affine, 1)
-	pk.G1.K = nil
-	pk.Domain = fft.Domain{Cardinality: pk.Domain.Cardinality}
-
+	pk.G2.B = nil
+	runtime.GC()
 	return nil
 }
 
@@ -451,11 +450,15 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...b
 	log.Debug().Dur("took", time.Since(start)).Msg("prover done")
 
 	// free device/GPU memory that is not needed for future proofs (scalars/hpoly)
-	go func() {
+	/*go func() {
 		iciclegnark.FreeDevicePointer(wireValuesADevice.P)
 		iciclegnark.FreeDevicePointer(wireValuesBDevice.P)
 		iciclegnark.FreeDevicePointer(h)
-	}()
+	}()*/
+	// for mem safe
+	iciclegnark.FreeDevicePointer(wireValuesADevice.P)
+	iciclegnark.FreeDevicePointer(wireValuesBDevice.P)
+	iciclegnark.FreeDevicePointer(h)
 
 	return proof, nil
 }

backend/groth16/bn254/icicle/icicle.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math/big"
 	"math/bits"
+	"runtime"
 	"sync"
 	"time"
 	"unsafe"
@@ -48,6 +49,7 @@ func (pk *ProvingKey) setupDevicePointers() error {
 	if pk.deviceInfo != nil {
 		return nil
 	}
+	runtime.GC()
 	pk.deviceInfo = &deviceInfo{}
 	n := int(pk.Domain.Cardinality)
 	sizeBytes := n * fr.Bytes
@@ -157,7 +159,7 @@ func (pk *ProvingKey) setupDevicePointers() error {
 	pk.G1.Z = make([]curve.G1Affine, 1)
 	pk.G1.K = nil
 	pk.Domain = fft.Domain{Cardinality: pk.Domain.Cardinality}
-
+	runtime.GC()
 	return nil
 }
 
@@ -445,11 +447,9 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...b
 	log.Debug().Dur("took", time.Since(start)).Msg("prover done")
 
 	// free device/GPU memory that is not needed for future proofs (scalars/hpoly)
-	go func() {
-		iciclegnark.FreeDevicePointer(wireValuesADevice.P)
-		iciclegnark.FreeDevicePointer(wireValuesBDevice.P)
-		iciclegnark.FreeDevicePointer(h)
-	}()
+	iciclegnark.FreeDevicePointer(wireValuesADevice.P)
+	iciclegnark.FreeDevicePointer(wireValuesBDevice.P)
+	iciclegnark.FreeDevicePointer(h)
 
 	return proof, nil
 }

backend/groth16/bw6-761/icicle/icicle.go

@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"math/big"
 	"math/bits"
+	"runtime"
 	"sync"
 	"time"
 	"unsafe"
@@ -48,6 +49,7 @@ func (pk *ProvingKey) setupDevicePointers() error {
 	if pk.deviceInfo != nil {
 		return nil
 	}
+	runtime.GC()
 	pk.deviceInfo = &deviceInfo{}
 	n := int(pk.Domain.Cardinality)
 	sizeBytes := n * fr.Bytes
@@ -158,6 +160,7 @@ func (pk *ProvingKey) setupDevicePointers() error {
 	pk.G1.K = nil
 	pk.Domain = fft.Domain{Cardinality: pk.Domain.Cardinality}
 
+	runtime.GC()
 	return nil
 }
 
@@ -446,11 +449,9 @@ func Prove(r1cs *cs.R1CS, pk *ProvingKey, fullWitness witness.Witness, opts ...b
 	log.Debug().Dur("took", time.Since(start)).Msg("prover done")
 
 	// free device/GPU memory that is not needed for future proofs (scalars/hpoly)
-	go func() {
-		iciclegnark.FreeDevicePointer(wireValuesADevice.P)
-		iciclegnark.FreeDevicePointer(wireValuesBDevice.P)
-		iciclegnark.FreeDevicePointer(h)
-	}()
+	iciclegnark.FreeDevicePointer(wireValuesADevice.P)
+	iciclegnark.FreeDevicePointer(wireValuesBDevice.P)
+	iciclegnark.FreeDevicePointer(h)
 
 	return proof, nil
 }