github.com/consensys/gnark-crypto@v0.14.0/internal/generator/fft/template/fft.go.tmpl

import (
	"math/bits"
	"github.com/consensys/gnark-crypto/ecc"
	"github.com/consensys/gnark-crypto/internal/parallel"
	"math/big"
	{{ template "import_fr" . }}
)

{{- /* these params set the size of the kernel we generate & unroll */}}
{{ $sizeKernelLog2 := 8}}
{{ $sizeKernel := shl 1 $sizeKernelLog2}}

// Decimation is used in the FFT call to select decimation in time or in frequency
type Decimation uint8

const (
	DIT Decimation = iota
	DIF
)

// parallelize threshold for a single butterfly op, if the fft stage is not parallelized already
const butterflyThreshold = 16

// FFT computes (recursively) the discrete Fourier transform of a and stores the result in a
// if decimation == DIT (decimation in time), the input must be in bit-reversed order
// if decimation == DIF (decimation in frequency), the output will be in bit-reversed order
func (domain *Domain) FFT(a []fr.Element, decimation Decimation, opts ...Option) {

	opt := fftOptions(opts...)

	// find the stage where we should stop spawning go routines in our recursive calls
	// (ie when we have as many go routines running as we have available CPUs)
	maxSplits := bits.TrailingZeros64(ecc.NextPowerOfTwo(uint64(opt.nbTasks)))
	if opt.nbTasks == 1 {
		maxSplits = -1
	}

	// if coset != 0, scale by coset table
	if opt.coset {
		if decimation == DIT {
			// scale by coset table (in bit reversed order)
			cosetTable := domain.cosetTable
			if !domain.withPrecompute {
				// we need to build the full table or do a bit reverse dance.
				cosetTable = make([]fr.Element, len(a))
				BuildExpTable(domain.FrMultiplicativeGen, cosetTable)
			}
			parallel.Execute(len(a), func(start, end int) {
				n := uint64(len(a))
				nn := uint64(64 - bits.TrailingZeros64(n))
				for i := start; i < end; i++ {
					irev := int(bits.Reverse64(uint64(i)) >> nn)
					a[i].Mul(&a[i], &cosetTable[irev])
				}
			}, opt.nbTasks)
		} else {
			if domain.withPrecompute {
				parallel.Execute(len(a), func(start, end int) {
					for i := start; i < end; i++ {
						a[i].Mul(&a[i], &domain.cosetTable[i])
					}
				}, opt.nbTasks)
			} else {
				c := domain.FrMultiplicativeGen
				parallel.Execute(len(a), func(start, end int) {
					var at fr.Element
					at.Exp(c, big.NewInt(int64(start)))
					for i := start; i < end; i++ {
						a[i].Mul(&a[i], &at)
						at.Mul(&at, &c)
					}
				}, opt.nbTasks)
			}
	
		}
	}

	twiddles := domain.twiddles
	twiddlesStartStage := 0
	if !domain.withPrecompute {
		twiddlesStartStage = 3
		nbStages := int(bits.TrailingZeros64(domain.Cardinality))
		if nbStages - twiddlesStartStage > 0 {
			twiddles = make([][]fr.Element, nbStages - twiddlesStartStage)
			w := domain.Generator
			w.Exp(w, big.NewInt(int64(1 << twiddlesStartStage)))
			buildTwiddles(twiddles, w, uint64(nbStages - twiddlesStartStage))
		} // else, we don't need twiddles
	}

	switch decimation {
	case DIF:
		difFFT(a, domain.Generator, twiddles, twiddlesStartStage, 0, maxSplits, nil, opt.nbTasks)
	case DIT:
		ditFFT(a, domain.Generator, twiddles, twiddlesStartStage, 0, maxSplits, nil, opt.nbTasks)
	default:
		panic("not implemented")
	}
}



// FFTInverse computes (recursively) the inverse discrete Fourier transform of a and stores the result in a
// if decimation == DIT (decimation in time), the input must be in bit-reversed order
// if decimation == DIF (decimation in frequency), the output will be in bit-reversed order
// coset sets the shift of the fft (0 = no shift, standard fft)
// len(a) must be a power of 2, and w must be a len(a)th root of unity in field F.
func (domain *Domain) FFTInverse(a []fr.Element, decimation Decimation, opts ...Option) {
	opt := fftOptions(opts...)

	// find the stage where we should stop spawning go routines in our recursive calls
	// (ie when we have as many go routines running as we have available CPUs)
	maxSplits := bits.TrailingZeros64(ecc.NextPowerOfTwo(uint64(opt.nbTasks)))
	if opt.nbTasks == 1 {
		maxSplits = -1
	}

	twiddlesInv := domain.twiddlesInv
	twiddlesStartStage := 0
	if !domain.withPrecompute {
		twiddlesStartStage = 3
		nbStages := int(bits.TrailingZeros64(domain.Cardinality))
		if nbStages - twiddlesStartStage > 0 {
			twiddlesInv = make([][]fr.Element, nbStages - twiddlesStartStage)
			w := domain.GeneratorInv
			w.Exp(w, big.NewInt(int64(1 << twiddlesStartStage)))
			buildTwiddles(twiddlesInv, w, uint64(nbStages - twiddlesStartStage))
		} // else, we don't need twiddles
	}

	switch decimation {
	case DIF:
		difFFT(a, domain.GeneratorInv, twiddlesInv, twiddlesStartStage, 0, maxSplits, nil, opt.nbTasks)
	case DIT:
		ditFFT(a, domain.GeneratorInv, twiddlesInv, twiddlesStartStage, 0, maxSplits, nil, opt.nbTasks)
	default:
		panic("not implemented")
	}

	// scale by CardinalityInv
	if !opt.coset {
		parallel.Execute(len(a), func(start, end int) {
			for i := start; i < end; i++ {
				a[i].Mul(&a[i], &domain.CardinalityInv)
			}
		}, opt.nbTasks)
		return
	}


	if decimation == DIT {
		if domain.withPrecompute {
			parallel.Execute(len(a), func(start, end int) {
				for i := start; i < end; i++ {
					a[i].Mul(&a[i], &domain.cosetTableInv[i]).
						Mul(&a[i], &domain.CardinalityInv)
				}
			}, opt.nbTasks)
		} else {
			c := domain.FrMultiplicativeGenInv
			parallel.Execute(len(a), func(start, end int) {
				var at fr.Element
				at.Exp(c, big.NewInt(int64(start)))
				at.Mul(&at, &domain.CardinalityInv)
				for i := start; i < end; i++ {
					a[i].Mul(&a[i], &at)
					at.Mul(&at, &c)
				}
			}, opt.nbTasks)
		}
		return
	}

	// decimation == DIF, need to access coset table in bit reversed order.
	cosetTableInv := domain.cosetTableInv
	if !domain.withPrecompute {
		// we need to build the full table or do a bit reverse dance.
		cosetTableInv = make([]fr.Element, len(a))
		BuildExpTable(domain.FrMultiplicativeGenInv, cosetTableInv)
	}
	parallel.Execute(len(a), func(start, end int) {
		n := uint64(len(a))
		nn := uint64(64 - bits.TrailingZeros64(n))
		for i := start; i < end; i++ {
			irev := int(bits.Reverse64(uint64(i)) >> nn)
			a[i].Mul(&a[i], &cosetTableInv[irev]).
				Mul(&a[i], &domain.CardinalityInv)
		}
	}, opt.nbTasks)

}

func difFFT(a []fr.Element, w fr.Element, twiddles [][]fr.Element, twiddlesStartStage, stage, maxSplits int, chDone chan struct{}, nbTasks int) {
	if chDone != nil {
		defer close(chDone)
	}

	n := len(a)
	if n == 1 {
		return
	} else if n == {{$sizeKernel}} && stage >= twiddlesStartStage  {
		kerDIFNP_{{$sizeKernel}}(a, twiddles, stage-twiddlesStartStage)
		return
	}
	m := n >> 1

	parallelButterfly := (m > butterflyThreshold) && (stage < maxSplits)

	if stage < twiddlesStartStage {
		if parallelButterfly {
			w := w
			parallel.Execute(m, func(start, end int) {
				if start == 0 {
					fr.Butterfly(&a[0], &a[m])
					start++
				}
				var at fr.Element
				at.Exp(w, big.NewInt(int64(start)))
				innerDIFWithoutTwiddles(a, at,w, start, end, m)
			}, nbTasks / (1 << (stage))) // 1 << stage == estimated used CPUs
		} else {
			innerDIFWithoutTwiddles(a, w, w, 0, m, m)
		}
		// compute next twiddle
		w.Square(&w) 
	} else {
		if parallelButterfly {
			parallel.Execute(m, func(start, end int) {
				innerDIFWithTwiddles(a, twiddles[stage-twiddlesStartStage], start, end, m)
			}, nbTasks / (1 << (stage)))
		} else {
			innerDIFWithTwiddles(a, twiddles[stage-twiddlesStartStage], 0, m, m)
		}
	}

	if m == 1 {
		return
	}
	
	nextStage := stage + 1
	if stage < maxSplits {
		chDone := make(chan struct{}, 1)
		go difFFT(a[m:n], w, twiddles, twiddlesStartStage, nextStage, maxSplits, chDone, nbTasks)
		difFFT(a[0:m], w, twiddles, twiddlesStartStage, nextStage, maxSplits, nil, nbTasks)
		<-chDone
	} else {
		difFFT(a[0:m], w, twiddles, twiddlesStartStage, nextStage, maxSplits, nil, nbTasks)
		difFFT(a[m:n], w, twiddles, twiddlesStartStage, nextStage, maxSplits, nil, nbTasks)
	}

}


func innerDIFWithTwiddles(a []fr.Element, twiddles []fr.Element, start, end, m int) {
	if start == 0 {
		fr.Butterfly(&a[0], &a[m])
		start++
	}
	for i := start; i < end; i++ {
		fr.Butterfly(&a[i], &a[i+m])
		a[i+m].Mul(&a[i+m], &twiddles[i])
	}
}

func innerDIFWithoutTwiddles(a []fr.Element, at, w fr.Element, start, end, m int) {
	if start == 0 {
		fr.Butterfly(&a[0], &a[m])
		start++
	}
	for i := start; i < end; i++ {
		fr.Butterfly(&a[i], &a[i+m])
		a[i+m].Mul(&a[i+m], &at)
		at.Mul(&at, &w)
	}
}


func ditFFT(a []fr.Element, w fr.Element, twiddles [][]fr.Element, twiddlesStartStage, stage, maxSplits int, chDone chan struct{}, nbTasks int) {
	if chDone != nil {
		defer close(chDone)
	}
	n := len(a)
	if n == 1 {
		return
	} else if n == {{$sizeKernel}} && stage >= twiddlesStartStage  {
		kerDITNP_{{$sizeKernel}}(a, twiddles, stage-twiddlesStartStage)
		return
	}
	m := n >> 1

	nextStage := stage + 1
	nextW := w 
	nextW.Square(&nextW)

	if stage < maxSplits {
		// that's the only time we fire go routines
		chDone := make(chan struct{}, 1)
		go ditFFT(a[m:],nextW, twiddles, twiddlesStartStage, nextStage, maxSplits, chDone, nbTasks)
		ditFFT(a[0:m],nextW, twiddles, twiddlesStartStage, nextStage, maxSplits, nil, nbTasks)
		<-chDone
	} else {
		ditFFT(a[0:m],nextW, twiddles, twiddlesStartStage, nextStage, maxSplits, nil, nbTasks)
		ditFFT(a[m:n],nextW, twiddles, twiddlesStartStage, nextStage, maxSplits, nil, nbTasks)
	}

	parallelButterfly := (m > butterflyThreshold) && (stage < maxSplits)

	if stage < twiddlesStartStage {
		// we need to compute the twiddles for this stage on the fly.
		if  parallelButterfly {
			w := w
			parallel.Execute(m, func(start, end int) {
				if start == 0 {
					fr.Butterfly(&a[0], &a[m])
					start++
				}
				var at fr.Element
				at.Exp(w, big.NewInt(int64(start)))
				innerDITWithoutTwiddles(a, at,w, start, end, m)
			}, nbTasks / (1 << (stage))) // 1 << stage == estimated used CPUs

		} else {
			innerDITWithoutTwiddles(a, w,w, 0, m, m)
		}
		return
	}
	if parallelButterfly {
		parallel.Execute(m, func(start, end int) {
			innerDITWithTwiddles(a, twiddles[stage-twiddlesStartStage], start, end, m)
		}, nbTasks / (1 << (stage)))
	} else {
		innerDITWithTwiddles(a, twiddles[stage-twiddlesStartStage], 0, m, m)
	}
}


func innerDITWithTwiddles(a []fr.Element, twiddles []fr.Element, start, end, m int) {
	if start == 0 {
		fr.Butterfly(&a[0], &a[m])
		start++
	}
	for i := start; i < end; i++ {
		a[i+m].Mul(&a[i+m], &twiddles[i])
		fr.Butterfly(&a[i], &a[i+m])
	}
}

func innerDITWithoutTwiddles(a []fr.Element, at, w fr.Element, start, end, m int) {
	if start == 0 {
		fr.Butterfly(&a[0], &a[m])
		start++
	}
	for i := start; i < end; i++ {
		a[i+m].Mul(&a[i+m], &at)
		fr.Butterfly(&a[i], &a[i+m])
		at.Mul(&at, &w)
	}
}



func kerDIFNP_{{$sizeKernel}}(a []fr.Element, twiddles [][]fr.Element, stage int) {
	// code unrolled & generated by internal/generator/fft/template/fft.go.tmpl

	{{ $n := shl 1 $sizeKernelLog2}}
	{{ $m := div $n 2}}
	{{ $split := 1}}
	{{- range $step := iterate 0 $sizeKernelLog2}} 
		{{- $offset := 0}}

		{{- $bound := mul $split $n}}
		{{- if eq $bound $n}}
			innerDIFWithTwiddles(a[:{{$n}}], twiddles[stage + {{$step}}], 0, {{$m}}, {{$m}})
		{{- else}}
			for offset := 0; offset < {{$bound}}; offset += {{$n}} {
				{{- if eq $m 1}}
					fr.Butterfly(&a[offset], &a[offset+1])
				{{- else}}
					innerDIFWithTwiddles(a[offset:offset + {{$n}}], twiddles[stage + {{$step}}], 0, {{$m}}, {{$m}})
				{{- end}}
			}
		{{- end}}

		{{- $n = div $n 2}}
		{{- $m = div $n 2}}
		{{- $split = mul $split 2}}
	{{- end}}
}


func kerDITNP_{{$sizeKernel}}(a []fr.Element, twiddles [][]fr.Element, stage int) {
	// code unrolled & generated by internal/generator/fft/template/fft.go.tmpl

	{{ $n := 2}}
	{{ $m := div $n 2}}
	{{ $split := div (shl 1 $sizeKernelLog2) 2}}
	{{- range $step := reverse (iterate 0 $sizeKernelLog2)}} 
		{{- $offset := 0}}

		{{- $bound := mul $split $n}}
		{{- if eq $bound $n}}
			innerDITWithTwiddles(a[:{{$n}}], twiddles[stage + {{$step}}], 0, {{$m}}, {{$m}})
		{{- else}}
			for offset := 0; offset < {{$bound}}; offset += {{$n}} {
				{{- if eq $m 1}}
					fr.Butterfly(&a[offset], &a[offset+1])
				{{- else}}
					innerDITWithTwiddles(a[offset:offset + {{$n}}], twiddles[stage + {{$step}}], 0, {{$m}}, {{$m}})
				{{- end}}
			}
		{{- end}}

		{{- $n = mul $n 2}}
		{{- $m = div $n 2}}
		{{- $split = div $split 2}}
	{{- end}}
}