Files
icicle/wrappers/golang/curves/bw6761/tests/msm_test.go
nonam3e f8d15e2613 update imports in golang bindings (#498)
## Describe the changes

This PR updates imports in golang bindings to the v2 version
2024-04-25 03:46:14 +07:00

283 lines
9.1 KiB
Go

package tests
import (
"fmt"
"sync"
"testing"
"github.com/stretchr/testify/assert"
"github.com/consensys/gnark-crypto/ecc"
bw6761 "github.com/consensys/gnark-crypto/ecc/bw6-761"
"github.com/consensys/gnark-crypto/ecc/bw6-761/fp"
"github.com/consensys/gnark-crypto/ecc/bw6-761/fr"
"github.com/ingonyama-zk/icicle/v2/wrappers/golang/core"
cr "github.com/ingonyama-zk/icicle/v2/wrappers/golang/cuda_runtime"
icicleBw6_761 "github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bw6761"
"github.com/ingonyama-zk/icicle/v2/wrappers/golang/curves/bw6761/msm"
)
func projectiveToGnarkAffine(p icicleBw6_761.Projective) bw6761.G1Affine {
px, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)((&p.X).ToBytesLittleEndian()))
py, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)((&p.Y).ToBytesLittleEndian()))
pz, _ := fp.LittleEndian.Element((*[fp.Bytes]byte)((&p.Z).ToBytesLittleEndian()))
zInv := new(fp.Element)
x := new(fp.Element)
y := new(fp.Element)
zInv.Inverse(&pz)
x.Mul(&px, zInv)
y.Mul(&py, zInv)
return bw6761.G1Affine{X: *x, Y: *y}
}
func testAgainstGnarkCryptoMsm(scalars core.HostSlice[icicleBw6_761.ScalarField], points core.HostSlice[icicleBw6_761.Affine], out icicleBw6_761.Projective) bool {
scalarsFr := make([]fr.Element, len(scalars))
for i, v := range scalars {
slice64, _ := fr.LittleEndian.Element((*[fr.Bytes]byte)(v.ToBytesLittleEndian()))
scalarsFr[i] = slice64
}
pointsFp := make([]bw6761.G1Affine, len(points))
for i, v := range points {
pointsFp[i] = projectiveToGnarkAffine(v.ToProjective())
}
return testAgainstGnarkCryptoMsmGnarkCryptoTypes(scalarsFr, pointsFp, out)
}
func testAgainstGnarkCryptoMsmGnarkCryptoTypes(scalarsFr core.HostSlice[fr.Element], pointsFp core.HostSlice[bw6761.G1Affine], out icicleBw6_761.Projective) bool {
var msmRes bw6761.G1Jac
msmRes.MultiExp(pointsFp, scalarsFr, ecc.MultiExpConfig{})
var icicleResAsJac bw6761.G1Jac
proj := projectiveToGnarkAffine(out)
icicleResAsJac.FromAffine(&proj)
return msmRes.Equal(&icicleResAsJac)
}
func convertIcicleAffineToG1Affine(iciclePoints []icicleBw6_761.Affine) []bw6761.G1Affine {
points := make([]bw6761.G1Affine, len(iciclePoints))
for index, iciclePoint := range iciclePoints {
xBytes := ([fp.Bytes]byte)(iciclePoint.X.ToBytesLittleEndian())
fpXElem, _ := fp.LittleEndian.Element(&xBytes)
yBytes := ([fp.Bytes]byte)(iciclePoint.Y.ToBytesLittleEndian())
fpYElem, _ := fp.LittleEndian.Element(&yBytes)
points[index] = bw6761.G1Affine{
X: fpXElem,
Y: fpYElem,
}
}
return points
}
func TestMSM(t *testing.T) {
cfg := msm.GetDefaultMSMConfig()
cfg.IsAsync = true
for _, power := range []int{2, 3, 4, 5, 6, 7, 8, 10, 18} {
size := 1 << power
scalars := icicleBw6_761.GenerateScalars(size)
points := icicleBw6_761.GenerateAffinePoints(size)
stream, _ := cr.CreateStream()
var p icicleBw6_761.Projective
var out core.DeviceSlice
_, e := out.MallocAsync(p.Size(), p.Size(), stream)
assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")
cfg.Ctx.Stream = &stream
e = msm.Msm(scalars, points, &cfg, out)
assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
outHost := make(core.HostSlice[icicleBw6_761.Projective], 1)
outHost.CopyFromDeviceAsync(&out, stream)
out.FreeAsync(stream)
cr.SynchronizeStream(&stream)
// Check with gnark-crypto
assert.True(t, testAgainstGnarkCryptoMsm(scalars, points, outHost[0]))
}
}
func TestMSMGnarkCryptoTypes(t *testing.T) {
cfg := msm.GetDefaultMSMConfig()
for _, power := range []int{3} {
size := 1 << power
scalars := make([]fr.Element, size)
var x fr.Element
for i := 0; i < size; i++ {
x.SetRandom()
scalars[i] = x
}
scalarsHost := (core.HostSlice[fr.Element])(scalars)
points := icicleBw6_761.GenerateAffinePoints(size)
pointsGnark := convertIcicleAffineToG1Affine(points)
pointsHost := (core.HostSlice[bw6761.G1Affine])(pointsGnark)
var p icicleBw6_761.Projective
var out core.DeviceSlice
_, e := out.Malloc(p.Size(), p.Size())
assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")
cfg.ArePointsMontgomeryForm = true
cfg.AreScalarsMontgomeryForm = true
e = msm.Msm(scalarsHost, pointsHost, &cfg, out)
assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
outHost := make(core.HostSlice[icicleBw6_761.Projective], 1)
outHost.CopyFromDevice(&out)
out.Free()
// Check with gnark-crypto
assert.True(t, testAgainstGnarkCryptoMsmGnarkCryptoTypes(scalarsHost, pointsHost, outHost[0]))
}
}
func TestMSMBatch(t *testing.T) {
cfg := msm.GetDefaultMSMConfig()
for _, power := range []int{10, 16} {
for _, batchSize := range []int{1, 3, 16} {
size := 1 << power
totalSize := size * batchSize
scalars := icicleBw6_761.GenerateScalars(totalSize)
points := icicleBw6_761.GenerateAffinePoints(totalSize)
var p icicleBw6_761.Projective
var out core.DeviceSlice
_, e := out.Malloc(batchSize*p.Size(), p.Size())
assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")
e = msm.Msm(scalars, points, &cfg, out)
assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
outHost := make(core.HostSlice[icicleBw6_761.Projective], batchSize)
outHost.CopyFromDevice(&out)
out.Free()
// Check with gnark-crypto
for i := 0; i < batchSize; i++ {
scalarsSlice := scalars[i*size : (i+1)*size]
pointsSlice := points[i*size : (i+1)*size]
out := outHost[i]
assert.True(t, testAgainstGnarkCryptoMsm(scalarsSlice, pointsSlice, out))
}
}
}
}
func TestPrecomputeBase(t *testing.T) {
cfg := msm.GetDefaultMSMConfig()
const precomputeFactor = 8
for _, power := range []int{10, 16} {
for _, batchSize := range []int{1, 3, 16} {
size := 1 << power
totalSize := size * batchSize
scalars := icicleBw6_761.GenerateScalars(totalSize)
points := icicleBw6_761.GenerateAffinePoints(totalSize)
var precomputeOut core.DeviceSlice
_, e := precomputeOut.Malloc(points[0].Size()*points.Len()*int(precomputeFactor), points[0].Size())
assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for PrecomputeBases results failed")
e = msm.PrecomputeBases(points, precomputeFactor, 0, &cfg.Ctx, precomputeOut)
assert.Equal(t, e, cr.CudaSuccess, "PrecomputeBases failed")
var p icicleBw6_761.Projective
var out core.DeviceSlice
_, e = out.Malloc(batchSize*p.Size(), p.Size())
assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")
cfg.PrecomputeFactor = precomputeFactor
e = msm.Msm(scalars, precomputeOut, &cfg, out)
assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
outHost := make(core.HostSlice[icicleBw6_761.Projective], batchSize)
outHost.CopyFromDevice(&out)
out.Free()
precomputeOut.Free()
// Check with gnark-crypto
for i := 0; i < batchSize; i++ {
scalarsSlice := scalars[i*size : (i+1)*size]
pointsSlice := points[i*size : (i+1)*size]
out := outHost[i]
assert.True(t, testAgainstGnarkCryptoMsm(scalarsSlice, pointsSlice, out))
}
}
}
}
func TestMSMSkewedDistribution(t *testing.T) {
cfg := msm.GetDefaultMSMConfig()
for _, power := range []int{2, 3, 4, 5, 6, 7, 8, 10, 18} {
size := 1 << power
scalars := icicleBw6_761.GenerateScalars(size)
for i := size / 4; i < size; i++ {
scalars[i].One()
}
points := icicleBw6_761.GenerateAffinePoints(size)
for i := 0; i < size/4; i++ {
points[i].Zero()
}
var p icicleBw6_761.Projective
var out core.DeviceSlice
_, e := out.Malloc(p.Size(), p.Size())
assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")
e = msm.Msm(scalars, points, &cfg, out)
assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
outHost := make(core.HostSlice[icicleBw6_761.Projective], 1)
outHost.CopyFromDevice(&out)
out.Free()
// Check with gnark-crypto
assert.True(t, testAgainstGnarkCryptoMsm(scalars, points, outHost[0]))
}
}
func TestMSMMultiDevice(t *testing.T) {
numDevices, _ := cr.GetDeviceCount()
numDevices = 1 // TODO remove when test env is fixed
fmt.Println("There are ", numDevices, " devices available")
orig_device, _ := cr.GetDevice()
wg := sync.WaitGroup{}
for i := 0; i < numDevices; i++ {
wg.Add(1)
cr.RunOnDevice(i, func(args ...any) {
defer wg.Done()
cfg := msm.GetDefaultMSMConfig()
cfg.IsAsync = true
for _, power := range []int{2, 3, 4, 5, 6, 7, 8, 10, 18} {
size := 1 << power
scalars := icicleBw6_761.GenerateScalars(size)
points := icicleBw6_761.GenerateAffinePoints(size)
stream, _ := cr.CreateStream()
var p icicleBw6_761.Projective
var out core.DeviceSlice
_, e := out.MallocAsync(p.Size(), p.Size(), stream)
assert.Equal(t, e, cr.CudaSuccess, "Allocating bytes on device for Projective results failed")
cfg.Ctx.Stream = &stream
e = msm.Msm(scalars, points, &cfg, out)
assert.Equal(t, e, cr.CudaSuccess, "Msm failed")
outHost := make(core.HostSlice[icicleBw6_761.Projective], 1)
outHost.CopyFromDeviceAsync(&out, stream)
out.FreeAsync(stream)
cr.SynchronizeStream(&stream)
// Check with gnark-crypto
assert.True(t, testAgainstGnarkCryptoMsm(scalars, points, outHost[0]))
}
})
}
wg.Wait()
cr.SetDevice(orig_device)
}