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

import (
	"math/bits"
	"runtime"
	{{ template "import_fr" . }}
)

// BitReverse applies the bit-reversal permutation to v.
// len(v) must be a power of 2
func BitReverse(v []fr.Element) {
	n := uint64(len(v))
	if bits.OnesCount64(n) != 1 {
		panic("len(a) must be a power of 2")
	}

	if runtime.GOARCH == "arm64" {
		bitReverseNaive(v)
	} else {
		bitReverseCobra(v)
	}
}

// bitReverseNaive applies the bit-reversal permutation to v.
// len(v) must be a power of 2
func bitReverseNaive(v []fr.Element) {
	n := uint64(len(v))
	nn := uint64(64 - bits.TrailingZeros64(n))

	for i := uint64(0); i < n; i++ {
		iRev := bits.Reverse64(i) >> nn
		if iRev > i {
			v[i], v[iRev] = v[iRev], v[i]
		}
	}
}


// bitReverseCobraInPlace applies the bit-reversal permutation to v.
// len(v) must be a power of 2
// This is derived from:
//
// - Towards an Optimal Bit-Reversal Permutation Program
//   Larry Carter and Kang Su Gatlin, 1998
//   https://csaws.cs.technion.ac.il/~itai/Courses/Cache/bit.pdf
//
// - Practically efficient methods for performing bit-reversed
//   permutation in C++11 on the x86-64 architecture
//   Knauth, Adas, Whitfield, Wang, Ickler, Conrad, Serang, 2017
//   https://arxiv.org/pdf/1708.01873.pdf
// 
// - and more specifically, constantine implementation:
//	 https://github.com/mratsim/constantine/blob/d51699248db04e29c7b1ad97e0bafa1499db00b5/constantine/math/polynomials/fft.nim#L205
// 	 by Mamy Ratsimbazafy (@mratsim).
//
func bitReverseCobraInPlace(v []fr.Element) {
	logN := uint64(bits.Len64(uint64(len(v))) - 1)
	logTileSize := deriveLogTileSize(logN)
	logBLen := logN - 2*logTileSize
	bLen := uint64(1) << logBLen
	bShift := logBLen + logTileSize
	tileSize := uint64(1) << logTileSize

	// rough idea;
	// bit reversal permutation naive implementation may have some cache associativity issues,
	// since we are accessing elements by strides of powers of 2.
	// on large inputs, this is noticeable and can be improved by using a t buffer.
	// idea is for t buffer to be small enough to fit in cache.
	// in the first inner loop, we copy the elements of v into t in a bit-reversed order.
	// in the subsequent inner loops, accesses have much better cache locality than the naive implementation.
	// hence even if we apparently do more work (swaps / copies), we are faster.
	// 
	// on arm64 (and particularly on M1 macs), this is not noticeable, and the naive implementation is faster,
	// in most cases.
	// on x86 (and particularly on aws hpc6a) this is noticeable, and the t buffer implementation is faster (up to 3x).
	//
	// optimal choice for the tile size is cache dependent; in theory, we want the t buffer to fit in the L1 cache;
	// in practice, a common size for L1 is 64kb, a field element is 32bytes or more.
	// hence we can fit 2k elements in the L1 cache, which corresponds to a tile size of 2**5 with some margin for cache conflicts.
	//
	// for most sizes of interest, this tile size choice doesn't yield good results;
	// we find that a tile size of 2**9 gives best results for input sizes from 2**21 up to 2**27+.
	t := make([]fr.Element, tileSize*tileSize)


	// see https://csaws.cs.technion.ac.il/~itai/Courses/Cache/bit.pdf
	// for a detailed explanation of the algorithm.
	for b := uint64(0); b < bLen; b++ {

		for a := uint64(0); a < tileSize; a++ {
			aRev :=( bits.Reverse64(a) >> (64 - logTileSize)) << logTileSize
			for c := uint64(0); c < tileSize; c++ {
				idx := (a << bShift) | (b << logTileSize) | c
				t[aRev | c] = v[idx]
			}
		}

		bRev := (bits.Reverse64(b) >> (64 - logBLen)) << logTileSize

		for c := uint64(0); c < tileSize; c++ {
			cRev := ((bits.Reverse64(c) >> (64 - logTileSize))  << bShift) | bRev
			for aRev := uint64(0); aRev < tileSize; aRev++ {
				a := bits.Reverse64(aRev) >> (64 - logTileSize)
				idx := (a << bShift) | (b << logTileSize) | c
				idxRev := cRev | aRev
				if idx < idxRev {
					tIdx := (aRev << logTileSize) | c
					v[idxRev], t[tIdx] = t[tIdx], v[idxRev]
				}
			}
		}

		for a := uint64(0); a < tileSize; a++ {
			aRev := bits.Reverse64(a) >> (64 - logTileSize)
			for c := uint64(0); c < tileSize; c++ {
				cRev := (bits.Reverse64(c) >> (64 - logTileSize)) << bShift
				idx := (a << bShift) | (b << logTileSize) | c
				idxRev := cRev | bRev | aRev
				if idx < idxRev {
					tIdx := (aRev << logTileSize) | c
					v[idx], t[tIdx] = t[tIdx], v[idx]
				}
			}
		}
	}
}


func bitReverseCobra(v []fr.Element) {
	switch len(v) {
	case 1 << 21:
		bitReverseCobraInPlace_9_21(v)
	case 1 << 22:
		bitReverseCobraInPlace_9_22(v)
	case 1 << 23:
		bitReverseCobraInPlace_9_23(v)
	case 1 << 24:
		bitReverseCobraInPlace_9_24(v)
	case 1 << 25:
		bitReverseCobraInPlace_9_25(v)
	case 1 << 26:
		bitReverseCobraInPlace_9_26(v)
	case 1 << 27:
		bitReverseCobraInPlace_9_27(v)
	default:
		if len(v) > 1<<27 {
			bitReverseCobraInPlace(v)
		} else {
			bitReverseNaive(v)
		}
	}
}


func deriveLogTileSize(logN uint64) uint64 {
	q := uint64(9) // see bitReverseCobraInPlace for more details

	for int(logN)-int(2*q) <= 0 {
		q--
	}

	return q
}


{{bitReverseCobraInPlace 9 21}}
{{bitReverseCobraInPlace 9 22}}
{{bitReverseCobraInPlace 9 23}}
{{bitReverseCobraInPlace 9 24}}
{{bitReverseCobraInPlace 9 25}}
{{bitReverseCobraInPlace 9 26}}
{{bitReverseCobraInPlace 9 27}}


{{define "bitReverseCobraInPlace logTileSize logN"}}

// bitReverseCobraInPlace_{{.logTileSize}}_{{.logN}} applies the bit-reversal permutation to v.
// len(v) must be 1 << {{.logN}}.
// see bitReverseCobraInPlace for more details; this function is specialized for {{.logTileSize}},
// as it declares the t buffer and various constants statically for performance.
func bitReverseCobraInPlace_{{.logTileSize}}_{{.logN}}(v []fr.Element) {
	const (
		logTileSize = uint64({{.logTileSize}})
		tileSize = uint64(1) << logTileSize
		logN = {{.logN}}
		logBLen = logN - 2*logTileSize
		bShift = logBLen + logTileSize
		bLen = uint64(1) << logBLen
	)

	var t [tileSize * tileSize]fr.Element
	{{$k := sub 64  .logTileSize}}
	{{$l := .logTileSize}}
	{{$tileSize := shl 1 .logTileSize}}
	
	for b := uint64(0); b < bLen; b++ {
		
		for a := uint64(0); a < tileSize; a++ {
			aRev := (bits.Reverse64(a) >> {{$k}}) << logTileSize
			for c := uint64(0); c < tileSize; c++ {
				idx := (a << bShift) | (b << logTileSize) | c
				t[aRev | c] = v[idx]
			}
		}

		bRev := (bits.Reverse64(b) >> (64 - logBLen)) << logTileSize

		for c := uint64(0); c < tileSize; c++ {
			cRev := ((bits.Reverse64(c) >> {{$k}}) << bShift) | bRev
			for aRev := uint64(0); aRev < tileSize; aRev++ {
				a := bits.Reverse64(aRev) >> {{$k}}
				idx := (a << bShift) | (b << logTileSize) | c
				idxRev := cRev | aRev
				if idx < idxRev {
					tIdx := (aRev << logTileSize) | c
					v[idxRev], t[tIdx] = t[tIdx], v[idxRev]
				}
			}
		}

		for a := uint64(0); a < tileSize; a++ {
			aRev := bits.Reverse64(a) >> {{$k}}
			for c := uint64(0); c < tileSize; c++ {
				cRev := (bits.Reverse64(c) >> {{$k}}) << bShift
				idx := (a << bShift) | (b << logTileSize) | c
				idxRev := cRev | bRev  | aRev
				if idx < idxRev {
					tIdx := (aRev << logTileSize) | c
					v[idx], t[tIdx] = t[tIdx], v[idx]
				}
			}
		}
	}

	
}

{{- end}}