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

     1  package parallel
     2  
     3  import (
     4  	"runtime"
     5  	"sync"
     6  )
     7  
     8  // Execute process in parallel the work function
     9  func Execute(nbIterations int, work func(int, int), maxCpus ...int) {
    10  
    11  	nbTasks := runtime.NumCPU()
    12  	if len(maxCpus) == 1 {
    13  		nbTasks = maxCpus[0]
    14  		if nbTasks < 1 {
    15  			nbTasks = 1
    16  		} else if nbTasks > 512 {
    17  			nbTasks = 512
    18  		}
    19  	}
    20  
    21  	if nbTasks == 1 {
    22  		// no go routines
    23  		work(0, nbIterations)
    24  		return
    25  	}
    26  
    27  	nbIterationsPerCpus := nbIterations / nbTasks
    28  
    29  	// more CPUs than tasks: a CPU will work on exactly one iteration
    30  	if nbIterationsPerCpus < 1 {
    31  		nbIterationsPerCpus = 1
    32  		nbTasks = nbIterations
    33  	}
    34  
    35  	var wg sync.WaitGroup
    36  
    37  	extraTasks := nbIterations - (nbTasks * nbIterationsPerCpus)
    38  	extraTasksOffset := 0
    39  
    40  	for i := 0; i < nbTasks; i++ {
    41  		wg.Add(1)
    42  		_start := i*nbIterationsPerCpus + extraTasksOffset
    43  		_end := _start + nbIterationsPerCpus
    44  		if extraTasks > 0 {
    45  			_end++
    46  			extraTasks--
    47  			extraTasksOffset++
    48  		}
    49  		go func() {
    50  			work(_start, _end)
    51  			wg.Done()
    52  		}()
    53  	}
    54  
    55  	wg.Wait()
    56  }