281 lines
8.3 KiB
Nim
281 lines
8.3 KiB
Nim
|
# Constantine
|
|||
|
|
# Copyright (c) 2018-2019 Status Research & Development GmbH
|
|||
|
|
# Copyright (c) 2020-Present Mamy André-Ratsimbazafy
|
|||
|
|
# Licensed and distributed under either of
|
|||
|
|
# * MIT license (license terms in the root directory or at http://opensource.org/licenses/MIT).
|
|||
|
|
# * Apache v2 license (license terms in the root directory or at http://www.apache.org/licenses/LICENSE-2.0).
|
|||
|
|
# at your option. This file may not be copied, modified, or distributed except according to those terms.
|
|||
|
|
|
|||
|
|
import
|
|||
|
|
../../constantine/config/curves,
|
|||
|
|
../../constantine/[arithmetic, primitives],
|
|||
|
|
../../constantine/io/io_fields,
|
|||
|
|
# Research
|
|||
|
|
./strided_views,
|
|||
|
|
./fft_lut
|
|||
|
|
|
|||
|
|
# See: https://github.com/ethereum/research/blob/master/kzg_data_availability/fft.py
|
|||
|
|
# Quirks of the Python impl:
|
|||
|
|
# - no tests of FFT alone?
|
|||
|
|
# - a lot of "if type(x) == tuple else"
|
|||
|
|
#
|
|||
|
|
# See: https://github.com/protolambda/go-kate/blob/7bb4684/fft_fr.go#L19-L21
|
|||
|
|
# The go port uses stride+offset to deal with skip iterator.
|
|||
|
|
#
|
|||
|
|
# Other readable FFTs includes:
|
|||
|
|
# - https://github.com/kwantam/fffft
|
|||
|
|
# - https://github.com/ConsenSys/gnark/blob/master/internal/backend/bls381/fft/fft.go
|
|||
|
|
# - https://github.com/poanetwork/threshold_crypto/blob/8820c11/src/poly_vals.rs#L332-L370
|
|||
|
|
# - https://github.com/zkcrypto/bellman/blob/10c5010/src/domain.rs#L272-L315
|
|||
|
|
# - Modern Computer Arithmetic, Brent and Zimmermann, p53 algorithm 2.2
|
|||
|
|
# https://members.loria.fr/PZimmermann/mca/mca-cup-0.5.9.pdf
|
|||
|
|
|
|||
|
|
# ############################################################
|
|||
|
|
#
|
|||
|
|
# Finite-Field Fast Fourier Transform
|
|||
|
|
#
|
|||
|
|
# ############################################################
|
|||
|
|
#
|
|||
|
|
# This is a research, unoptimized implementation of
|
|||
|
|
# Finite Field Fast Fourier Transform
|
|||
|
|
|
|||
|
|
# In research phase we tolerate using
|
|||
|
|
# - garbage collected types
|
|||
|
|
# - and exceptions for fast prototyping
|
|||
|
|
#
|
|||
|
|
# In particular, in production all signed integers
|
|||
|
|
# must be verified not to overflow
|
|||
|
|
# and should not throw (or use unsigned)
|
|||
|
|
|
|||
|
|
# FFT Context
|
|||
|
|
# ----------------------------------------------------------------
|
|||
|
|
|
|||
|
|
type
|
|||
|
|
FFTStatus = enum
|
|||
|
|
FFTS_Success
|
|||
|
|
FFTS_TooManyValues = "Input length greater than the field 2-adicity (number of roots of unity)"
|
|||
|
|
FFTS_SizeNotPowerOfTwo = "Input must be of a power of 2 length"
|
|||
|
|
|
|||
|
|
FFTDescriptor[F] = object
|
|||
|
|
## Metadata for FFT on field F
|
|||
|
|
maxWidth: int
|
|||
|
|
rootOfUnity: F
|
|||
|
|
## The root of unity that generates all roots
|
|||
|
|
expandedRootsOfUnity: seq[F]
|
|||
|
|
## domain, starting and ending with 1
|
|||
|
|
|
|||
|
|
func isPowerOf2(n: SomeUnsignedInt): bool =
|
|||
|
|
(n and (n - 1)) == 0
|
|||
|
|
|
|||
|
|
func nextPowerOf2(n: uint64): uint64 =
|
|||
|
|
## Returns x if x is a power of 2
|
|||
|
|
## or the next biggest power of 2
|
|||
|
|
1'u64 shl (log2(n-1) + 1)
|
|||
|
|
|
|||
|
|
func expandRootOfUnity[F](rootOfUnity: F): seq[F] =
|
|||
|
|
## From a generator root of unity
|
|||
|
|
## expand to width + 1 values.
|
|||
|
|
## (Last value is 1 for the reverse array)
|
|||
|
|
# For a field of order q, there are gcd(n, q−1)
|
|||
|
|
# nth roots of unity, a.k.a. solutions to xⁿ ≡ 1 (mod q)
|
|||
|
|
# but it's likely too long to compute bigint GCD
|
|||
|
|
# so embrace heap (re-)allocations.
|
|||
|
|
# Figuring out how to do to right size the buffers
|
|||
|
|
# in production will be fun.
|
|||
|
|
result.setLen(2)
|
|||
|
|
result[0].setOne()
|
|||
|
|
result[1] = rootOfUnity
|
|||
|
|
while not result[^1].isOne().bool:
|
|||
|
|
result.setLen(result.len + 1)
|
|||
|
|
result[^1].prod(result[^2], rootOfUnity)
|
|||
|
|
|
|||
|
|
# FFT Algorithm
|
|||
|
|
# ----------------------------------------------------------------
|
|||
|
|
|
|||
|
|
# TODO: research Decimation in Time and Decimation in Frequency
|
|||
|
|
# and FFT butterflies
|
|||
|
|
|
|||
|
|
func simpleFT[F](
|
|||
|
|
output: var View[F],
|
|||
|
|
vals: View[F],
|
|||
|
|
rootsOfUnity: View[F]
|
|||
|
|
) =
|
|||
|
|
# FFT is a recursive algorithm
|
|||
|
|
# This is the base-case using a O(n²) algorithm
|
|||
|
|
|
|||
|
|
let L = output.len
|
|||
|
|
var last {.noInit.}, v {.noInit.}: F
|
|||
|
|
|
|||
|
|
for i in 0 ..< L:
|
|||
|
|
last.prod(vals[0], rootsOfUnity[0])
|
|||
|
|
for j in 1 ..< L:
|
|||
|
|
v.prod(vals[j], rootsOfUnity[(i*j) mod L])
|
|||
|
|
last += v
|
|||
|
|
output[i] = last
|
|||
|
|
|
|||
|
|
func fft_internal[F](
|
|||
|
|
output: var View[F],
|
|||
|
|
vals: View[F],
|
|||
|
|
rootsOfUnity: View[F]
|
|||
|
|
) =
|
|||
|
|
if output.len <= 4:
|
|||
|
|
simpleFT(output, vals, rootsOfUnity)
|
|||
|
|
return
|
|||
|
|
|
|||
|
|
# Recursive Divide-and-Conquer
|
|||
|
|
let (evenVals, oddVals) = vals.splitAlternate()
|
|||
|
|
var (outLeft, outRight) = output.splitMiddle()
|
|||
|
|
let halfROI = rootsOfUnity.skipHalf()
|
|||
|
|
|
|||
|
|
fft_internal(outLeft, evenVals, halfROI)
|
|||
|
|
fft_internal(outRight, oddVals, halfROI)
|
|||
|
|
|
|||
|
|
let half = outLeft.len
|
|||
|
|
var y_times_root{.noinit.}: F
|
|||
|
|
|
|||
|
|
for i in 0 ..< half:
|
|||
|
|
# FFT Butterfly
|
|||
|
|
y_times_root .prod(output[i+half], rootsOfUnity[i])
|
|||
|
|
output[i+half] .diff(output[i], y_times_root)
|
|||
|
|
output[i] += y_times_root
|
|||
|
|
|
|||
|
|
func fft*[F](
|
|||
|
|
desc: FFTDescriptor[F],
|
|||
|
|
output: var openarray[F],
|
|||
|
|
vals: openarray[F]): FFT_Status =
|
|||
|
|
if vals.len > desc.maxWidth:
|
|||
|
|
return FFTS_TooManyValues
|
|||
|
|
if not vals.len.uint64.isPowerOf2():
|
|||
|
|
return FFTS_SizeNotPowerOfTwo
|
|||
|
|
|
|||
|
|
let rootz = desc.expandedRootsOfUnity
|
|||
|
|
.toView()
|
|||
|
|
.slice(0, desc.maxWidth-1, desc.maxWidth div vals.len)
|
|||
|
|
|
|||
|
|
var voutput = output.toView()
|
|||
|
|
fft_internal(voutput, vals.toView(), rootz)
|
|||
|
|
return FFTS_Success
|
|||
|
|
|
|||
|
|
func ifft*[F](
|
|||
|
|
desc: FFTDescriptor[F],
|
|||
|
|
output: var openarray[F],
|
|||
|
|
vals: openarray[F]): FFT_Status =
|
|||
|
|
## Inverse FFT
|
|||
|
|
if vals.len > desc.maxWidth:
|
|||
|
|
return FFTS_TooManyValues
|
|||
|
|
if not vals.len.uint64.isPowerOf2():
|
|||
|
|
return FFTS_SizeNotPowerOfTwo
|
|||
|
|
|
|||
|
|
let rootz = desc.expandedRootsOfUnity
|
|||
|
|
.toView()
|
|||
|
|
.reversed()
|
|||
|
|
.slice(0, desc.maxWidth-1, desc.maxWidth div vals.len)
|
|||
|
|
|
|||
|
|
var voutput = output.toView()
|
|||
|
|
fft_internal(voutput, vals.toView(), rootz)
|
|||
|
|
|
|||
|
|
var invLen {.noInit.}: F
|
|||
|
|
invLen.fromUint(vals.len.uint64)
|
|||
|
|
invLen.inv()
|
|||
|
|
|
|||
|
|
for i in 0..< output.len:
|
|||
|
|
output[i] *= invLen
|
|||
|
|
|
|||
|
|
return FFTS_Success
|
|||
|
|
|
|||
|
|
# FFT Descriptor
|
|||
|
|
# ----------------------------------------------------------------
|
|||
|
|
|
|||
|
|
proc init*(T: type FFTDescriptor, maxScale: uint8): T =
|
|||
|
|
result.maxWidth = 1 shl maxScale
|
|||
|
|
result.rootOfUnity = scaleToRootOfUnity(T.F.C)[maxScale]
|
|||
|
|
result.expandedRootsOfUnity =
|
|||
|
|
result.rootOfUnity.expandRootOfUnity()
|
|||
|
|
# Aren't you tired of reading about unity?
|
|||
|
|
|
|||
|
|
# ############################################################
|
|||
|
|
#
|
|||
|
|
# Sanity checks
|
|||
|
|
#
|
|||
|
|
# ############################################################
|
|||
|
|
|
|||
|
|
{.experimental: "views".}
|
|||
|
|
|
|||
|
|
when isMainModule:
|
|||
|
|
import
|
|||
|
|
std/[times, monotimes, strformat],
|
|||
|
|
../../helpers/prng_unsafe
|
|||
|
|
|
|||
|
|
proc roundtrip() =
|
|||
|
|
let fftDesc = FFTDescriptor[Fr[BLS12_381]].init(maxScale = 4)
|
|||
|
|
var data = newSeq[Fr[BLS12_381]](fftDesc.maxWidth)
|
|||
|
|
for i in 0 ..< fftDesc.maxWidth:
|
|||
|
|
data[i].fromUint i.uint64
|
|||
|
|
|
|||
|
|
var coefs = newSeq[Fr[BLS12_381]](data.len)
|
|||
|
|
let fftOk = fft(fftDesc, coefs, data)
|
|||
|
|
doAssert fftOk == FFTS_Success
|
|||
|
|
# display("coefs", 0, coefs)
|
|||
|
|
|
|||
|
|
var res = newSeq[Fr[BLS12_381]](data.len)
|
|||
|
|
let ifftOk = ifft(fftDesc, res, coefs)
|
|||
|
|
doAssert ifftOk == FFTS_Success
|
|||
|
|
# display("res", 0, coefs)
|
|||
|
|
|
|||
|
|
for i in 0 ..< res.len:
|
|||
|
|
if bool(res[i] != data[i]):
|
|||
|
|
echo "Error: expected ", data[i].toHex(), " but got ", res[i].toHex()
|
|||
|
|
quit 1
|
|||
|
|
|
|||
|
|
echo "FFT round-trip check SUCCESS"
|
|||
|
|
|
|||
|
|
proc warmup() =
|
|||
|
|
# Warmup - make sure cpu is on max perf
|
|||
|
|
let start = cpuTime()
|
|||
|
|
var foo = 123
|
|||
|
|
for i in 0 ..< 300_000_000:
|
|||
|
|
foo += i*i mod 456
|
|||
|
|
foo = foo mod 789
|
|||
|
|
|
|||
|
|
# Compiler shouldn't optimize away the results as cpuTime rely on sideeffects
|
|||
|
|
let stop = cpuTime()
|
|||
|
|
echo &"Warmup: {stop - start:>4.4f} s, result {foo} (displayed to avoid compiler optimizing warmup away)\n"
|
|||
|
|
|
|||
|
|
|
|||
|
|
proc bench() =
|
|||
|
|
echo "Starting benchmark ..."
|
|||
|
|
const NumIters = 100
|
|||
|
|
|
|||
|
|
var rng: RngState
|
|||
|
|
rng.seed 0x1234
|
|||
|
|
# TODO: view types complain about mutable borrow
|
|||
|
|
# in `random_unsafe` due to pseudo view type LimbsViewMut
|
|||
|
|
# (which was views before Nim properly supported them)
|
|||
|
|
|
|||
|
|
warmup()
|
|||
|
|
|
|||
|
|
for scale in 4 ..< 16:
|
|||
|
|
# Setup
|
|||
|
|
|
|||
|
|
let desc = FFTDescriptor[Fr[BLS12_381]].init(uint8 scale)
|
|||
|
|
var data = newSeq[Fr[BLS12_381]](desc.maxWidth)
|
|||
|
|
for i in 0 ..< desc.maxWidth:
|
|||
|
|
# data[i] = rng.random_unsafe(data[i].typeof())
|
|||
|
|
data[i].fromUint i.uint64
|
|||
|
|
|
|||
|
|
var coefsOut = newSeq[Fr[BLS12_381]](data.len)
|
|||
|
|
|
|||
|
|
# Bench
|
|||
|
|
let start = getMonotime()
|
|||
|
|
for i in 0 ..< NumIters:
|
|||
|
|
let status = desc.fft(coefsOut, data)
|
|||
|
|
doAssert status == FFTS_Success
|
|||
|
|
let stop = getMonotime()
|
|||
|
|
|
|||
|
|
let ns = inNanoseconds((stop-start) div NumIters)
|
|||
|
|
echo &"FFT scale {scale:>2} {ns:>8} ns/op"
|
|||
|
|
|
|||
|
|
roundtrip()
|
|||
|
|
warmup()
|
|||
|
|
bench()
|