feat: add vector operation on koalabear extension

field/generator/generator_extensions.go

@@ -54,6 +54,7 @@ func generateExtensions(F *config.Field, outputDir string) error {
 		data.QInvNeg = F.QInverse[0]
 		entries_ext4 := []bavard.Entry{
 			{File: filepath.Join(outputDir, "e4.go"), Templates: []string{"e4.go.tmpl"}},
+			{File: filepath.Join(outputDir, "e4_vector.go"), Templates: []string{"e4.vector.go.tmpl"}},
 			{File: filepath.Join(outputDir, "e4_test.go"), Templates: []string{"e4_test.go.tmpl"}},
 		}
 

field/generator/internal/templates/extensions/e4.go.tmpl

@@ -2,9 +2,6 @@ import (
 	"math/big"
 
 	fr "{{ .FieldPackagePath }}"
-	{{- if .F31}}
-	"github.com/consensys/gnark-crypto/utils/cpu"
-	{{- end}}
 )
 
 // E4 is a degree two finite field extension of fr2
@@ -347,23 +344,3 @@ func (z *E4) Div(x *E4, y *E4) *E4 {
 	r.Inverse(y).Mul(x, &r)
 	return z.Set(&r)
 }
-
-{{- if .F31}}
-
-func MulAccE4(alpha *E4, scale []fr.Element, res []E4) {
-	N := len(res)
-	if N != len(scale) {
-		panic("MulAccE4: len(res) != len(scale)")
-	}
-	if !cpu.SupportAVX512 || N%4 != 0 {
-		var tmp E4
-		for i := 0; i < N; i++ {
-			tmp.MulByElement(alpha, &scale[i])
-			res[i].Add(&res[i], &tmp)
-		}
-		return
-	}
-
-	mulAccE4_avx512(alpha, &scale[0], &res[0], uint64(N))
-}
-{{- end}}

field/generator/internal/templates/extensions/e4.vector.go.tmpl

@@ -0,0 +1,94 @@
+import (
+	fr "{{ .FieldPackagePath }}"
+	"github.com/consensys/gnark-crypto/utils/cpu"
+)
+
+// Vector represents a vector of E4 elements
+type Vector []E4
+
+func (vector *Vector) Add(a, b Vector) {
+	if len(a) != len(b) || len(a) != len(*vector) {
+		panic("vector.Add: vectors don't have the same length")
+	}
+
+	for i := range a {
+		(*vector)[i].Add(&a[i], &b[i])
+	}
+}
+
+func (vector *Vector) AddMixed(a Vector, b fr.Vector) {
+	if len(a) != len(b) || len(a) != len(*vector) {
+		panic("vector.AddMixed: vectors don't have the same length")
+	}
+	for i := range a {
+		(*vector)[i].Set(&a[i])
+		(*vector)[i].B0.A0.Add(&(*vector)[i].B0.A0, &b[i])
+	}
+}
+
+func (vector *Vector) ScalarMul(a Vector, b *E4) {
+	if len(a) != len(*vector) {
+		panic("vector.ScalarMul: vectors don't have the same length")
+	}
+
+	for i := range a {
+		(*vector)[i].Mul(&a[i], b)
+	}
+}
+
+func (vector *Vector) Sub(a, b Vector) {
+	if len(a) != len(b) || len(a) != len(*vector) {
+		panic("vector.Sub: vectors don't have the same length")
+	}
+
+	for i := range a {
+		(*vector)[i].Sub(&a[i], &b[i])
+	}
+}
+
+func (vector *Vector) SubMixed(a Vector, b fr.Vector) {
+	if len(a) != len(b) || len(a) != len(*vector) {
+		panic("vector.SubMixed: vectors don't have the same length")
+	}
+	for i := range a {
+		(*vector)[i].Set(&a[i])
+		(*vector)[i].B0.A0.Sub(&(*vector)[i].B0.A0, &b[i])
+	}
+}
+
+func (vector *Vector) Mul(a, b Vector) {
+	if len(a) != len(b) || len(a) != len(*vector) {
+		panic("vector.Mul: vectors don't have the same length")
+	}
+
+	for i := range a {
+		(*vector)[i].Mul(&a[i], &b[i])
+	}
+}
+
+func (vector *Vector) MulMixed(a Vector, b fr.Vector) {
+	if len(a) != len(b) || len(a) != len(*vector) {
+		panic("vector.MulMixed: vectors don't have the same length")
+	}
+	for i := range a {
+		(*vector)[i].Set(&a[i])
+		(*vector)[i].MulByElement(&(*vector)[i], &b[i])
+	}
+}
+
+func (vector *Vector) MulAcc(scale []fr.Element, alpha *E4) {
+	N := len(*vector)
+	if N != len(scale) {
+		panic("MulAccE4: len(res) != len(scale)")
+	}
+	if !cpu.SupportAVX512 || N%4 != 0 {
+		var tmp E4
+		for i := 0; i < N; i++ {
+			tmp.MulByElement(alpha, &scale[i])
+			(*vector)[i].Add(&(*vector)[i], &tmp)
+		}
+		return
+	}
+
+	mulAccE4_avx512(alpha, &scale[0], &(*vector)[0], uint64(N))
+}