github.com/consensys/gnark-crypto@v0.14.0/utils/parallel.go (about)

     1  package utils
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  )
     7  
     8  // Would normally put this in internal/parallel; but it may be desirable for it to be accessible from gnark
     9  
    10  type Task func(start, end int)
    11  
    12  type job struct {
    13  	start, end int
    14  	task       Task
    15  	done       *sync.WaitGroup
    16  }
    17  
    18  type WorkerPool struct {
    19  	chJobs    chan job
    20  	nbWorkers int
    21  }
    22  
    23  func NewWorkerPool() *WorkerPool {
    24  	p := &WorkerPool{}
    25  	p.nbWorkers = runtime.NumCPU() + 2
    26  	p.chJobs = make(chan job, 40*p.nbWorkers)
    27  	for i := 0; i < p.nbWorkers; i++ {
    28  		go func() {
    29  			for j := range p.chJobs {
    30  				j.task(j.start, j.end)
    31  				j.done.Done()
    32  			}
    33  		}()
    34  	}
    35  	return p
    36  }
    37  
    38  // Stop (but does not wait) the pool. It frees the worker.
    39  func (wp *WorkerPool) Stop() {
    40  	close(wp.chJobs)
    41  }
    42  
    43  func (wp *WorkerPool) Submit(n int, work func(int, int), minBlock int) *sync.WaitGroup {
    44  	var wg sync.WaitGroup
    45  
    46  	// we have an interval [0,n)
    47  	// that we split in minBlock sizes.
    48  	for start := 0; start < n; start += minBlock {
    49  		start := start
    50  		end := start + minBlock
    51  		if end > n {
    52  			end = n
    53  		}
    54  		wg.Add(1)
    55  		wp.chJobs <- job{
    56  			task:  work,
    57  			start: start,
    58  			end:   end,
    59  			done:  &wg,
    60  		}
    61  	}
    62  
    63  	return &wg
    64  }