feat: add parallel prefix product for vector e4 (#750)

Author: gbotrel

field/babybear/extensions/e4_test.go

@@ -6,6 +6,7 @@ import (
 	"math/big"
 	"reflect"
 	"sort"
+	"runtime"
 	"bytes"
 
 	fr "{{ .FieldPackagePath }}"
@@ -600,6 +601,49 @@ func TestVectorExp(t *testing.T) {
 
 }
 
+// prefixProductGeneric computes the prefix product of the vector in place (single-threaded).
+func prefixProductGeneric(vector Vector) {
+	if len(vector) == 0 {
+		return
+	}
+	for i := 1; i < len(vector); i++ {
+		vector[i].Mul(&vector[i-1], &vector[i])
+	}
+}
+
+func randomVector(size int) Vector {
+	v := make(Vector, size)
+	for i := range v {
+		v[i].MustSetRandom()
+	}
+	return v
+}
+
+func TestPrefixProduct_EmptyVector(t *testing.T) {
+	assert := require.New(t)
+	v := make(Vector, 0)
+	expected := make(Vector, 0)
+	prefixProductGeneric(expected)
+	v.PrefixProduct()
+	assert.Equal(expected, v)
+}
+
+func TestPrefixProduct_VariousNbTasks(t *testing.T) {
+	assert := require.New(t)
+	sizes := []int{1, 2, 256, 1024}
+	nbTasksList := []int{1, 16, 32, runtime.NumCPU()}
+	for _, size := range sizes {
+		for _, nbTasks := range nbTasksList {
+			v := randomVector(size)
+			expected := make(Vector, size)
+			copy(expected, v)
+			prefixProductGeneric(expected)
+			v.PrefixProduct(nbTasks)
+			assert.Equal(expected, v, "size=%d nbTasks=%d", size, nbTasks)
+		}
+	}
+}
+
 
 func TestVectorEmptyOps(t *testing.T) {
 	assert := require.New(t)
@@ -869,6 +913,28 @@ func BenchmarkVectorOps(b *testing.B) {
 	}
 }
 
+func BenchmarkPrefixProduct(b *testing.B) {
+	const N = 1 << 19
+	a1 := make(Vector, N)
+	for i := 0; i < N; i++ {
+		a1[i].MustSetRandom()
+	}
+
+	b.Run("generic", func(b *testing.B) {
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			prefixProductGeneric(a1)
+		}
+	})
+
+	b.Run("PrefixProduct", func(b *testing.B) {
+		b.ResetTimer()
+		for i := 0; i < b.N; i++ {
+			a1.PrefixProduct()
+		}
+	})
+
+}
 
 
 func BenchmarkVectorSerialization(b *testing.B) {

field/babybear/extensions/vector.go

@@ -7,6 +7,9 @@ import (
 	"bytes"
 	"slices"
 	"encoding/binary"
+	"runtime"
+	"github.com/consensys/gnark-crypto/internal/parallel"
+	"sync"
 
 	fr "{{ .FieldPackagePath }}"
 	{{- if .IsKoalaBear}}
@@ -537,6 +540,77 @@ func (vector *Vector) ReadFrom(r io.Reader) (int64, error) {
 	return n, <-chErr
 }
 
+// PrefixProduct computes the prefix product of the vector in place.
+// i.e. vector[i] = vector[0] * vector[1] * ... * vector[i]
+// If nbTasks > 1, it uses nbTasks goroutines to compute the prefix product in parallel.
+// If nbTasks is not provided, it uses the number of CPU cores.
+func (vector Vector) PrefixProduct(nbTasks ...int) {
+	N := len(vector)
+	if N < 2 {
+		return
+	}
+
+	if N < 512 {
+		vector.prefixProductGeneric()
+		return
+	}
+
+	// Use one worker per available CPU core.
+	numWorkers := runtime.GOMAXPROCS(0)
+	if len(nbTasks) == 1 && nbTasks[0] > 0 && nbTasks[0] < numWorkers {
+		numWorkers = nbTasks[0]
+	}
+
+	for N / numWorkers < 64 && numWorkers > 1 {
+		numWorkers >>= 1
+	}
+	numWorkers = max(1, numWorkers)
+
+		// --- PASS 1: Calculate prefix product for each chunk independently ---
+	parallel.Execute(N, func(start, stop int) {
+		// This is the original sequential algorithm applied to the smaller chunk.
+		for j := start + 1; j < stop; j++ {
+			vector[j].Mul(&vector[j], &vector[j-1])
+		}
+	}, numWorkers)
+
+	// get the chunk indices
+	chunks := parallel.Chunks(N, numWorkers)
+
+	// Compute multipliers for each chunk (product of all previous chunks)
+	multipliers := make([]E4, len(chunks))
+	multipliers[0].SetOne()
+	for i := 1; i < len(chunks); i++ {
+		multipliers[i].SetOne()
+		for j := 0; j < i; j++ {
+			prevChunkEnd := chunks[j][1] - 1
+			multipliers[i].Mul(&multipliers[i], &vector[prevChunkEnd])
+		}
+	}
+
+	// propagate the multipliers to each chunk in parallel
+	// note: the first chunk is not modified (multiplier is 1)
+	var wg sync.WaitGroup
+	wg.Add(len(chunks) - 1)
+	for i := 1; i < len(chunks); i++ {
+		go func(i int) {
+			defer wg.Done()
+			start, stop := chunks[i][0], chunks[i][1]
+			subVector := vector[start:stop]
+			subVector.ScalarMul(subVector, &multipliers[i])
+		}(i)
+	}
+	wg.Wait()
+
+}
+
+func (vector Vector) prefixProductGeneric() {
+	for i := 1; i < len(vector); i++ {
+		vector[i].Mul(&vector[i], &vector[i-1])
+	}
+}
+
+
 
 func vectorAddGeneric(res, a, b Vector) {
 	for i := 0; i < len(res); i++ {