diff --git a/src/FFTA.jl b/src/FFTA.jl index 3309f17..d49fd6b 100644 --- a/src/FFTA.jl +++ b/src/FFTA.jl @@ -12,6 +12,7 @@ using Reexport: @reexport include("callgraph.jl") include("singleton_twiddle.jl") include("codelets.jl") +include("odd_codelets.jl") include("algos.jl") include("plan.jl") @@ -22,7 +23,7 @@ using PrecompileTools: @setup_workload, @compile_workload @setup_workload begin @compile_workload begin for T in (Float64, Float32) - for n in (8, 16, 32, 64, 128, 256) + for n in (8, 16, 32, 64, 128, 256, 5, 7, 11, 13, 25, 49) x = ones(Complex{T}, n) y = AbstractFFTs.fft(x) AbstractFFTs.bfft(y) diff --git a/src/algos.jl b/src/algos.jl index 9dae269..ecb8dc4 100644 --- a/src/algos.jl +++ b/src/algos.jl @@ -18,13 +18,19 @@ function fft_kernel!( N = root.sz tw = g.twiddles[idx] if t === DFT - fft_dft!(out, in, N, start_out, s_out, start_in, s_in, tw) + # small odd primes have straight-line codelets (see odd_codelets.jl) + _odd_codelet!(out, in, N, start_out, s_out, start_in, s_in, d) || + fft_dft!(out, in, N, start_out, s_out, start_in, s_in, tw) elseif t === POW2RADIX4_FFT fft_pow2_radix4!(out, in, N, start_out, s_out, start_in, s_in, d, tw, 0) elseif t === POW3_FFT _m_120 = cispi(T(2) / 3) m_120 = d === FFT_FORWARD ? _m_120 : conj(_m_120) fft_pow3!(out, in, N, start_out, s_out, start_in, s_in, m_120, d, tw, 0) + elseif t === POW5_FFT + fft_powr!(out, in, N, start_out, s_out, start_in, s_in, d, tw, 0, Val(5)) + elseif t === POW7_FFT + fft_powr!(out, in, N, start_out, s_out, start_in, s_in, d, tw, 0, Val(7)) elseif t === BLUESTEIN fft_bluestein!(out, in, d, N, start_out, s_out, start_in, s_in, g.bluestein[g.blue_index[idx]]) else @@ -327,6 +333,54 @@ fft_pow3!(out::AbstractVector{T}, in::AbstractVector, N::Int, start_out::Int, st fft_pow3!(out, in, N, start_out, stride_out, start_in, stride_in, minus120, d, pow3_twiddles(T, N, d), 0) +""" +$(TYPEDSIGNATURES) +Radix-`R` FFT for powers of 5 and 7 (`Float32`/`Float64` elements), in place: +the `R` decimated sub-transforms are computed recursively, then each group of +`R` outputs is multiplied by its twiddles and combined with the `R`-point +codelet applied in place (see `odd_codelets.jl`). Same structure as +`fft_pow3!`; no composite step and no workspace. + +# Arguments +- `out`: Output vector +- `in`: Input vector (real or complex) +- `N`: Size of the transform (a power of `R`) +- `start_out`, `stride_out`, `start_in`, `stride_in`: as in `fft_pow2_radix4!` +- `d`: Direction of the transform +- `tw`: Twiddle table, see `powr_twiddles` +- `toff`: Offset of the current recursion level in `tw` +""" +function fft_powr!( + out::AbstractVector{T}, in::AbstractVector{U}, + N::Int, + start_out::Int, stride_out::Int, + start_in::Int, stride_in::Int, + d::Direction, + tw::AbstractVector{T}, toff::Int, + ::Val{R} +) where {T<:CodeletEltype, U, R} + if N == R + _odd_codelet!(out, in, R, start_out, stride_out, start_in, stride_in, d) + return + end + m = N ÷ R + toff_next = toff + (R - 1) * m + for r in 0:R-1 + fft_powr!(out, in, m, start_out + r * m * stride_out, stride_out, start_in + r * stride_in, stride_in * R, d, tw, toff_next, Val(R)) + end + # k = 0: all twiddles are 1 + _odd_codelet!(out, out, R, start_out, m * stride_out, start_out, m * stride_out, d) + @inbounds for k in 1:m-1 + base = start_out + k * stride_out + tb = toff + (R - 1) * k + for r in 1:R-1 + out[base + r * m * stride_out] *= tw[tb + r] + end + _odd_codelet!(out, out, R, base, m * stride_out, base, m * stride_out, d) + end + return nothing +end + """ $(TYPEDSIGNATURES) Bluestein's algorithm, still O(N * log(N)) for large primes, diff --git a/src/callgraph.jl b/src/callgraph.jl index ab119bb..0449800 100644 --- a/src/callgraph.jl +++ b/src/callgraph.jl @@ -1,6 +1,6 @@ @enum Direction FFT_FORWARD=-1 FFT_BACKWARD=1 @enum Pow24 POW2 POW4 -@enum FFTEnum COMPOSITE_FFT DFT POW3_FFT POW2RADIX4_FFT BLUESTEIN +@enum FFTEnum COMPOSITE_FFT DFT POW3_FFT POW2RADIX4_FFT BLUESTEIN POW5_FFT POW7_FFT @inline function direction_sign(d::Direction) Int(d) @@ -125,6 +125,15 @@ function CallGraphNode!( push!(workspace, T[]) push!(nodes, CallGraphNode(0, 0, POW3_FFT, N, s_in, s_out)) return 1 + elseif T <: CodeletEltype && N % 5 == 0 && nextpow(5, N) == N + # radix-5/7 kernels use the odd codelets as butterflies (see fft_powr!) + push!(workspace, T[]) + push!(nodes, CallGraphNode(0, 0, POW5_FFT, N, s_in, s_out)) + return 1 + elseif T <: CodeletEltype && N % 7 == 0 && nextpow(7, N) == N + push!(workspace, T[]) + push!(nodes, CallGraphNode(0, 0, POW7_FFT, N, s_in, s_out)) + return 1 elseif N == 1 || Primes.isprime(N) push!(workspace, T[]) # use Bluestein's algorithm for big primes @@ -292,6 +301,32 @@ function pow3_twiddles(::Type{T}, N::Int, dir::Direction) where {T} return tw end +""" +$(TYPEDSIGNATURES) +Twiddle table for the radix-`R` kernel `fft_powr!` (`R` = 5 or 7), laid out +like `pow3_twiddles`: for every level of size `M = N, N/R, …` (excluding the +`R`-point base case) the tuples `(w^k, w^2k, …, w^(R-1)k)`, `w = exp(dir · 2πi/M)`, +for `k = 0..M/R-1`. +""" +function powr_twiddles(::Type{T}, N::Int, R::Int, dir::Direction) where {T} + N > R || return T[] + W = unit_roots(T, N, dir) + tw = Vector{T}(undef, N) # (R-1)N/R + (R-1)N/R² + ... < N + i = 1 + M = N + while M > R + m = M ÷ R + s = N ÷ M + for k in 0:m-1, r in 1:R-1 + tw[i] = W[(r * s * k) % N + 1] + i += 1 + end + M = m + end + resize!(tw, i - 1) + return tw +end + """ $(TYPEDSIGNATURES) Twiddle table of the node at index `idx` of `nodes`, see `dft_twiddles`, @@ -311,6 +346,10 @@ function node_twiddles(::Type{T}, nodes::Vector{CallGraphNode}, idx::Int, dir::D return pow2_twiddles(T, N, dir) elseif node.type === POW3_FFT return pow3_twiddles(T, N, dir) + elseif node.type === POW5_FFT + return powr_twiddles(T, N, 5, dir) + elseif node.type === POW7_FFT + return powr_twiddles(T, N, 7, dir) else return T[] end diff --git a/src/odd_codelets.jl b/src/odd_codelets.jl new file mode 100644 index 0000000..cb45755 --- /dev/null +++ b/src/odd_codelets.jl @@ -0,0 +1,122 @@ +# Straight-line codelets for the small odd prime leaves (5, 7, 11, 13). +# +# The generic `fft_dft!` leaf is an O(N²) loop over a twiddle table: for a +# 5-point transform that is 16 complex multiplications from memory-loaded +# twiddles. These codelets use the symmetry of the odd-length DFT instead: +# with `a_j = x_j + x_{N-j}` and `b_j = x_j - x_{N-j}` for `j = 1..(N-1)/2`, +# +# X_k = x_0 + Σ_j a_j cos(2πjk/N) + D·i Σ_j b_j sin(2πjk/N) +# X_{N-k} = x_0 + Σ_j a_j cos(2πjk/N) - D·i Σ_j b_j sin(2πjk/N) +# +# so each output pair costs `(N-1)/2` real-by-complex products for the cosine +# sum and as many for the sine sum: `(N-1)²` real multiplications per +# transform instead of `4(N-1)²`, with every constant folded and no table +# loads. Like the power-of-two codelets they are compiled once per +# `(N, element type, direction)` (see the `PrecompileTools` workload) and are +# only used for `Float32`/`Float64` elements; other types keep `fft_dft!`. +# The input may be real (the odd-length real transform runs the complex +# kernel on real input), in which case the sums are real and the same +# expressions are emitted. + +const ODD_CODELET_SIZES = (5, 7, 11, 13) + +# real coefficient × (real or complex) value + (real or complex) accumulator +@inline _fma(a::T, b::T, c::T) where {T<:Real} = fma(a, b, c) +@inline _fma(a::T, b::Complex{T}, c::Complex{T}) where {T<:Real} = Complex(fma(a, real(b), real(c)), fma(a, imag(b), imag(c))) +@inline _fma(a::T, b::T, c::Complex{T}) where {T<:Real} = Complex(fma(a, b, real(c)), imag(c)) +@inline _fma(a::T, b::Complex{T}, c::T) where {T<:Real} = Complex(fma(a, real(b), c), a * imag(b)) + +# Emit the statements of the length-`N` DFT of the symbols `xs` in direction +# `dir` (`-1` forward), returning the output symbols. +function _gen_odd!(stmts::Vector{Any}, xs::Vector{Symbol}, ::Type{T}, dir::Int, counter::Ref{Int}) where {T} + N = length(xs) + h = (N - 1) ÷ 2 + newsym() = Symbol(:t, counter[] += 1) + as = Vector{Symbol}(undef, h); bs = Vector{Symbol}(undef, h) + for j in 1:h + as[j] = newsym(); bs[j] = newsym() + push!(stmts, :($(as[j]) = $(xs[j + 1]) + $(xs[N - j + 1])), :($(bs[j]) = $(xs[j + 1]) - $(xs[N - j + 1]))) + end + outs = Vector{Symbol}(undef, N) + # X_0 + acc = xs[1] + for j in 1:h + t = newsym(); push!(stmts, :($t = $acc + $(as[j]))); acc = t + end + outs[1] = acc + for k in 1:h + # cosine sum (real coefficients), starting from x_0; explicit fma (not + # muladd) so that rounding does not depend on whether LLVM contracts + # for a particular array type + c = xs[1] + for j in 1:h + coef = T(cospi(2 * (j * k % N) / N)) + t = newsym(); push!(stmts, :($t = _fma($coef, $(as[j]), $c))); c = t + end + # sine sum, with the direction folded into the coefficients + s = nothing + for j in 1:h + coef = T(dir * sinpi(2 * (j * k % N) / N)) + t = newsym() + push!(stmts, s === nothing ? :($t = $coef * $(bs[j])) : :($t = _fma($coef, $(bs[j]), $s))) + s = t + end + # ± i·s + is = newsym(); push!(stmts, :($is = Complex(-imag($s), real($s)))) + o1 = newsym(); o2 = newsym() + push!(stmts, :($o1 = $c + $is), :($o2 = $c - $is)) + outs[k + 1] = o1; outs[N - k + 1] = o2 + end + return outs +end + +""" +$(TYPEDSIGNATURES) +Straight-line length-`N` DFT (odd `N`) of `in[start_in + k*stride_in]` into +`out[start_out + k*stride_out]`, `k = 0..N-1`, in direction `D` (`-1` +forward, `+1` backward). The input may be real or complex. +""" +@generated function fft_odd_codelet!( + out::AbstractVector{Complex{T}}, in::AbstractVector{<:Union{T,Complex{T}}}, + ::Val{N}, + start_out::Int, stride_out::Int, + start_in::Int, stride_in::Int, + ::Val{D} +) where {T,N,D} + counter = Ref(0) + stmts = Any[] + xs = [Symbol(:x, i) for i in 1:N] + for i in 1:N + push!(stmts, :($(xs[i]) = in[start_in + $(i - 1) * stride_in])) + end + outs = _gen_odd!(stmts, xs, T, D, counter) + for i in 1:N + push!(stmts, :(out[start_out + $(i - 1) * stride_out] = $(outs[i]))) + end + body = Expr(:block, stmts...) + return quote + @inbounds $body + return nothing + end +end + +# Dispatch an odd leaf to its codelet; `false` when there is none for this +# `N` or element type (the caller then uses `fft_dft!`). +@inline function _odd_codelet!( + out::AbstractVector{Complex{S}}, in::AbstractVector{<:Union{S,Complex{S}}}, N::Int, + start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction +) where {S<:Union{Float32,Float64}} + if d === FFT_FORWARD + N == 5 && (fft_odd_codelet!(out, in, Val(5), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + N == 7 && (fft_odd_codelet!(out, in, Val(7), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + N == 11 && (fft_odd_codelet!(out, in, Val(11), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + N == 13 && (fft_odd_codelet!(out, in, Val(13), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + else + N == 5 && (fft_odd_codelet!(out, in, Val(5), start_out, stride_out, start_in, stride_in, Val(1)); return true) + N == 7 && (fft_odd_codelet!(out, in, Val(7), start_out, stride_out, start_in, stride_in, Val(1)); return true) + N == 11 && (fft_odd_codelet!(out, in, Val(11), start_out, stride_out, start_in, stride_in, Val(1)); return true) + N == 13 && (fft_odd_codelet!(out, in, Val(13), start_out, stride_out, start_in, stride_in, Val(1)); return true) + end + return false +end +_odd_codelet!(out, in, N, start_out, stride_out, start_in, stride_in, d) = false