From 841186148f9aa6d26d78e0309634601bb51c95fe Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:05:02 +0000 Subject: [PATCH 1/7] Store twiddle tables and Bluestein data in the plan Twiddle factors were regenerated on every execution: every kernel seeded Singleton's recurrence with a sincospi call (per output row of the O(n^2) DFT leaf, per j1 in the composite step, per level of the radix-4/3 kernels), and fft_bluestein! allocated three pad-length buffers and recomputed the chirp and its FFT on every call. CallGraph now carries, per node, a twiddle table in the layout its kernel reads sequentially (DFT: w^k; composite: the (j1, k2) block; radix-4/3: per-level interleaved triplets/pairs addressed by a flat offset), a BluesteinScratch (chirp, its pre-scaled transform, work arrays, pow2 tables for the padded length) per Bluestein node, and the direction the tables were built for. All tables derive from one unit_roots table per node that evaluates sincospi on the first octant only when 8 | N. Planned execution is allocation-free for every size. Tables are correctly rounded, so the Float32 error no longer grows with n (~1.5 ulp at 2^22 instead of ~1000); the accuracy test grid is extended to 2^22. The old kernel signatures remain as wrappers that build tables on the fly. --- docs/src/dev.md | 19 +++ src/algos.jl | 255 +++++++++++++------------------ src/callgraph.jl | 259 +++++++++++++++++++++++++++++--- src/plan.jl | 20 +-- test/onedim/accuracy.jl | 2 + test/onedim/complex_backward.jl | 2 +- test/onedim/complex_forward.jl | 3 +- test/runtests.jl | 3 + test/twiddles.jl | 112 ++++++++++++++ 9 files changed, 495 insertions(+), 180 deletions(-) create mode 100644 test/twiddles.jl diff --git a/docs/src/dev.md b/docs/src/dev.md index 2e35f08..81988b4 100644 --- a/docs/src/dev.md +++ b/docs/src/dev.md @@ -11,6 +11,25 @@ CallGraphNode CallGraph CallGraphNode! fft! +fft_composite! fft_dft! +fft_pow2_radix4! fft_pow3! +fft_bluestein! +``` + +## Twiddle factors + +Twiddle factors are computed once, at plan time, and stored per node of the +call graph in the layout each kernel reads sequentially. + +```@docs +twiddle +unit_roots +dft_twiddles +composite_twiddles +pow2_twiddles +pow3_twiddles +node_twiddles +BluesteinScratch ``` diff --git a/src/algos.jl b/src/algos.jl index f67cdca..322159d 100644 --- a/src/algos.jl +++ b/src/algos.jl @@ -1,7 +1,3 @@ -@inline function direction_sign(d::Direction) - Int(d) -end - function fft!( out::AbstractVector{T}, in::AbstractVector{<:Number}, start_out::Int, start_in::Int, @@ -10,6 +6,9 @@ function fft!( g::CallGraph{T}, idx::Int ) where T + if d !== g.dir + throw(ArgumentError("call graph was planned for direction $(g.dir), not $d")) + end if t === COMPOSITE_FFT fft_composite!(out, in, start_out, start_in, d, g, idx) else @@ -17,16 +16,17 @@ function fft!( s_in = root.s_in s_out = root.s_out N = root.sz + tw = g.twiddles[idx] if t === DFT - fft_dft!(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) + 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) + fft_pow3!(out, in, N, start_out, s_out, start_in, s_in, m_120, d, tw, 0) elseif t === BLUESTEIN - fft_bluestein!(out, in, d, N, start_out, s_out, start_in, s_in) + fft_bluestein!(out, in, d, N, start_out, s_out, start_in, s_in, g.bluestein[g.blue_index[idx]]) else throw(ArgumentError("kernel not implemented")) end @@ -60,7 +60,6 @@ function fft_composite!( right_idx = idx + root.right left = g[left_idx] right = g[right_idx] - N = root.sz N1 = left.sz N2 = right.sz s_in = root.s_in @@ -69,51 +68,27 @@ function fft_composite!( Rt = right.type Lt = left.type - Rtype = real(T) - dir = direction_sign(d) tmp = g.workspace[idx] + tw = g.twiddles[idx] # see `composite_twiddles` - if Rt === BLUESTEIN - R_bluestein_scratchspace = prealloc_blue(N2, d, T) - end for j1 in 0:N1-1 R_start_in = start_in + j1 * s_in R_start_out = 1 + N2 * j1 - if @isdefined R_bluestein_scratchspace - R_s_in = right.s_in - R_s_out = right.s_out - fft_bluestein!(tmp, in, d, N2, R_start_out, R_s_out, R_start_in, R_s_in, R_bluestein_scratchspace) - else - fft!(tmp, in, R_start_out, R_start_in, d, Rt, g, right_idx) - end + fft!(tmp, in, R_start_out, R_start_in, d, Rt, g, right_idx) if j1 > 0 - # The composite twiddle at position (j1, k2) is `cispi(dir · 2 j1 k2 / N)`. - # Singleton's recurrence advances `wk2 = cispi(dir · 2 j1 k2 / N)` in k2 - # for fixed j1; (α, β) depend on j1 so we reset them at each outer step. - zj1 = singleton_params(dir * Rtype(j1) / Rtype(N)) - wk2 = one(T) + base = (j1 - 1) * (N2 - 1) @inbounds for k2 in 1:N2-1 - wk2 = singleton_step(wk2, zj1) - tmp[R_start_out + k2] *= wk2 + tmp[R_start_out + k2] *= tw[base + k2] end end end - if Lt === BLUESTEIN - L_bluestein_scratchspace = prealloc_blue(N1, d, T) - end for k2 in 0:N2-1 L_start_out = start_out + k2 * s_out L_start_in = 1 + k2 - if @isdefined L_bluestein_scratchspace - L_s_in = left.s_in - L_s_out = left.s_out - fft_bluestein!(out, tmp, d, N1, L_start_out, L_s_out, L_start_in, L_s_in, L_bluestein_scratchspace) - else - fft!(out, tmp, L_start_out, L_start_in, d, Lt, g, left_idx) - end + fft!(out, tmp, L_start_out, L_start_in, d, Lt, g, left_idx) end end @@ -129,7 +104,8 @@ Discrete Fourier Transform, O(N^2) algorithm, in place. - `stride_out`: Stride of the output vector - `start_in`: Index of the first element of the input vector - `stride_in`: Stride of the input vector -- `d`: Direction of the transform +- `W`: Twiddle table, see `dft_twiddles` (or `d`, the direction, in which + case the table is computed on the fly) """ function fft_dft!( @@ -137,25 +113,25 @@ function fft_dft!( N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, - d::Direction + W::AbstractVector{T} ) where {T<:Complex} - tmp = in[start_in] - @inbounds for j in 1:N-1 - tmp += in[start_in + j*stride_in] - end - out[start_out] = tmp - - Rtype = real(T) - dir = direction_sign(d) - @inbounds for j in 1:N-1 + @inbounds begin tmp = in[start_in] - zj = singleton_params(dir * Rtype(j) / Rtype(N)) - wk = one(T) - @inbounds for k in 1:N-1 - wk = singleton_step(wk, zj) - tmp += wk * in[start_in + k*stride_in] + for j in 1:N-1 + tmp += in[start_in + j*stride_in] + end + out[start_out] = tmp + + for j in 1:N-1 + tmp = in[start_in] + idx = 0 # j * k mod N + for k in 1:N-1 + idx += j + idx >= N && (idx -= N) + tmp += W[idx + 1] * in[start_in + k*stride_in] + end + out[start_out + j*stride_out] = tmp end - out[start_out + j*stride_out] = tmp end end @@ -164,30 +140,35 @@ function fft_dft!( N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, - d::Direction + W::AbstractVector{Complex{T}} ) where {T<:Real} halfN = N÷2 - tmp = Complex{T}(in[start_in]) - @inbounds for j in 1:N-1 - tmp += in[start_in + j*stride_in] - end - out[start_out] = tmp - - dir = direction_sign(d) - @inbounds for j in 1:halfN + @inbounds begin tmp = Complex{T}(in[start_in]) - zj = singleton_params(dir * T(j) / T(N)) - wk = one(Complex{T}) - @inbounds for k in 1:N-1 - wk = singleton_step(wk, zj) - tmp += wk * in[start_in + k*stride_in] + for j in 1:N-1 + tmp += in[start_in + j*stride_in] + end + out[start_out] = tmp + + for j in 1:halfN + tmp = Complex{T}(in[start_in]) + idx = 0 + for k in 1:N-1 + idx += j + idx >= N && (idx -= N) + tmp += W[idx + 1] * in[start_in + k*stride_in] + end + out[start_out + j*stride_out] = tmp + out[start_out + (N-j)*stride_out] = conj(tmp) end - out[start_out + j*stride_out] = tmp - out[start_out + (N-j)*stride_out] = conj(tmp) end end +fft_dft!(out::AbstractVector{T}, in::AbstractVector, N::Int, start_out::Int, stride_out::Int, + start_in::Int, stride_in::Int, d::Direction) where {T<:Complex} = + fft_dft!(out, in, N, start_out, stride_out, start_in, stride_in, dft_twiddles(T, N, d)) + """ $(TYPEDSIGNATURES) @@ -202,6 +183,8 @@ Radix-4 FFT for powers of 2, in place - `start_in`: Index of the first element of the input vector - `stride_in`: Stride of the input vector - `d`: Direction of the transform +- `tw`: Twiddle table, see `pow2_twiddles` (omit it to compute the table on the fly) +- `toff`: Offset of the current recursion level in `tw` """ function fft_pow2_radix4!( @@ -209,7 +192,8 @@ function fft_pow2_radix4!( N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, - d::Direction + d::Direction, + tw::AbstractVector{T}, toff::Int ) where {T<:Complex, U} # If N is 2, compute the size two DFT @inbounds if N == 2 @@ -240,23 +224,17 @@ function fft_pow2_radix4!( # ...othersize split the problem in four and recur m = N ÷ 4 + toff_next = toff + 3m # the next level's table follows this level's - fft_pow2_radix4!(out, in, m, start_out , stride_out, start_in , stride_in*4, d) - fft_pow2_radix4!(out, in, m, start_out + m*stride_out, stride_out, start_in + stride_in, stride_in*4, d) - fft_pow2_radix4!(out, in, m, start_out + 2*m*stride_out, stride_out, start_in + 2*stride_in, stride_in*4, d) - fft_pow2_radix4!(out, in, m, start_out + 3*m*stride_out, stride_out, start_in + 3*stride_in, stride_in*4, d) - - Rtype = real(T) - # Singleton recurrence for the three running twiddles `w^k`, `w^2k`, `w^3k`. - z1 = singleton_params(dir * Rtype(1) / Rtype(N)) - z2 = singleton_params(dir * Rtype(2) / Rtype(N)) - z3 = singleton_params(dir * Rtype(3) / Rtype(N)) - - wkoe = one(T) - wkeo = one(T) - wkoo = one(T) + fft_pow2_radix4!(out, in, m, start_out , stride_out, start_in , stride_in*4, d, tw, toff_next) + fft_pow2_radix4!(out, in, m, start_out + m*stride_out, stride_out, start_in + stride_in, stride_in*4, d, tw, toff_next) + fft_pow2_radix4!(out, in, m, start_out + 2*m*stride_out, stride_out, start_in + 2*stride_in, stride_in*4, d, tw, toff_next) + fft_pow2_radix4!(out, in, m, start_out + 3*m*stride_out, stride_out, start_in + 3*stride_in, stride_in*4, d, tw, toff_next) @inbounds for k in 0:m-1 + wkoe = tw[toff + 3k + 1] + wkeo = tw[toff + 3k + 2] + wkoo = tw[toff + 3k + 3] kee = start_out + k * stride_out koe = start_out + (k + m) * stride_out keo = start_out + (k + 2 * m) * stride_out @@ -273,12 +251,13 @@ function fft_pow2_radix4!( out[koe] = y_kee_m_y_keo + t_koe_m_t_koo out[keo] = y_kee_p_y_keo - t_koe_p_t_koo out[koo] = y_kee_m_y_keo - t_koe_m_t_koo - wkoe = singleton_step(wkoe, z1) - wkeo = singleton_step(wkeo, z2) - wkoo = singleton_step(wkoo, z3) end end +fft_pow2_radix4!(out::AbstractVector{T}, in::AbstractVector, N::Int, start_out::Int, stride_out::Int, + start_in::Int, stride_in::Int, d::Direction) where {T<:Complex} = + fft_pow2_radix4!(out, in, N, start_out, stride_out, start_in, stride_in, d, pow2_twiddles(T, N, d), 0) + """ $(TYPEDSIGNATURES) @@ -294,6 +273,8 @@ Power of 3 FFT, in place - `stride_in`: Stride of the input vector - `minus120`: Depending on direction, perform either ∓120° rotation - `d`: Direction of the transform +- `tw`: Twiddle table, see `pow3_twiddles` (omit it to compute the table on the fly) +- `toff`: Offset of the current recursion level in `tw` """ function fft_pow3!( @@ -302,7 +283,8 @@ function fft_pow3!( start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, minus120::T, - d::Direction + d::Direction, + tw::AbstractVector{T}, toff::Int ) where {T, U} plus120 = conj(minus120) if N == 3 @@ -314,20 +296,16 @@ function fft_pow3!( # Size of subproblem Nprime = N ÷ 3 + toff_next = toff + 2 * Nprime # Dividing into subproblems - fft_pow3!(out, in, Nprime, start_out, stride_out, start_in, stride_in*3, minus120, d) - fft_pow3!(out, in, Nprime, start_out + Nprime*stride_out, stride_out, start_in + stride_in, stride_in*3, minus120, d) - fft_pow3!(out, in, Nprime, start_out + 2*Nprime*stride_out, stride_out, start_in + 2*stride_in, stride_in*3, minus120, d) + fft_pow3!(out, in, Nprime, start_out, stride_out, start_in, stride_in*3, minus120, d, tw, toff_next) + fft_pow3!(out, in, Nprime, start_out + Nprime*stride_out, stride_out, start_in + stride_in, stride_in*3, minus120, d, tw, toff_next) + fft_pow3!(out, in, Nprime, start_out + 2*Nprime*stride_out, stride_out, start_in + 2*stride_in, stride_in*3, minus120, d, tw, toff_next) - Rtype = real(T) - dir = direction_sign(d) - - z1 = singleton_params(dir * Rtype(1) / Rtype(N)) - z2 = singleton_params(dir * Rtype(2) / Rtype(N)) - wk1 = one(T) - wk2 = one(T) - for k in 0:Nprime-1 + @inbounds for k in 0:Nprime-1 + wk1 = tw[toff + 2k + 1] + wk2 = tw[toff + 2k + 2] k0 = start_out + stride_out * k k1 = start_out + stride_out * (k + Nprime) k2 = start_out + stride_out * (k + 2 * Nprime) @@ -335,36 +313,13 @@ function fft_pow3!( @muladd out[k0] = y_k0 + y_k1 * wk1 + y_k2 * wk2 @muladd out[k1] = y_k0 + y_k1 * wk1 * plus120 + y_k2 * wk2 * minus120 @muladd out[k2] = y_k0 + y_k1 * wk1 * minus120 + y_k2 * wk2 * plus120 - wk1 = singleton_step(wk1, z1) - wk2 = singleton_step(wk2, z2) end end +fft_pow3!(out::AbstractVector{T}, in::AbstractVector, N::Int, start_out::Int, stride_out::Int, + start_in::Int, stride_in::Int, minus120::T, d::Direction) where {T} = + fft_pow3!(out, in, N, start_out, stride_out, start_in, stride_in, minus120, d, pow3_twiddles(T, N, d), 0) -function prealloc_blue(N::Int, d::Direction, ::Type{T}) where T<:Number - pad_len = nextpow(2, 2N - 1) - - b_series = Vector{T}(undef, pad_len) - a_series = Vector{T}(undef, pad_len) - tmp = Vector{T}(undef, pad_len) - - b_series[N+1:end] .= zero(T) - - sgn = -direction_sign(d) - p = 0 # n^2 - for i in 1:N - b_series[i] = cispi(sgn * p / N) - p += (2i - 1) # prevents overflow unless N is absolutely massive - p > N && (p -= 2N) - end - - # enforce periodic boundaries for b_n - for j in 0:N-1 - b_series[pad_len-j] = b_series[2+j] - end - - return (tmp, a_series, b_series, pad_len) -end """ $(TYPEDSIGNATURES) @@ -372,7 +327,8 @@ Bluestein's algorithm, still O(N * log(N)) for large primes, but with a big constant factor. Zero-pads two sequences derived from the DFT formula to a power of 2 length greater than `2N-1` and computes their convolution -with a power 2 FFT. +with a power 2 FFT. The chirp, its transform and the work arrays are +precomputed in `scratch` (see `BluesteinScratch`). # Arguments - `out`: Output vector @@ -383,7 +339,7 @@ with a power 2 FFT. - `stride_out`: Stride of the output vector - `start_in`: Index of the first element of the input vector - `stride_in`: Stride of the input vector -- `scratch` (optional): preallocated scratch space for bluestein +- `scratch`: precomputed data and scratch space (omit it to compute it on the fly) """ function fft_bluestein!( @@ -392,33 +348,30 @@ function fft_bluestein!( N::Int, start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, - scratch::Tuple{Vector{T},Vector{T},Vector{T},Int}=prealloc_blue(N, d, T) + scratch::BluesteinScratch{T}=BluesteinScratch{T}(N, d) ) where T<:Complex + (; pad_len, chirp, chirp_fft, a, tmp, tw) = scratch - (tmp, a_series, b_series, pad_len) = scratch - - a_series[N+1:end] .= zero(T) - tmp[N+1:end] .= zero(T) - - for i in 1:N - a_series[i] = in[start_in+(i-1)*stride_in] * conj(b_series[i]) + # a_n = x_n · conj(b_n), zero padded + @inbounds for i in 1:N + a[i] = in[start_in + (i-1)*stride_in] * conj(chirp[i]) + end + @inbounds for i in N+1:pad_len + a[i] = zero(T) end - # leave b_n vector alone for last step - fft_pow2_radix4!(tmp, a_series, pad_len, 1, 1, 1, 1, FFT_BACKWARD) # Fa - fft_pow2_radix4!(a_series, b_series, pad_len, 1, 1, 1, 1, FFT_BACKWARD) # Fb - - tmp .*= a_series - # convolution theorem ifft - fft_pow2_radix4!(a_series, tmp, pad_len, 1, 1, 1, 1, FFT_FORWARD) - conv_a_b = a_series - - Xk = tmp - for i in 1:N - Xk[i] = conj(b_series[i]) * conv_a_b[i] / pad_len + # Circular convolution of `a` with the periodised chirp via the forward + # transform only: conv = conj(fft(conj(fft(a) .* fft(b)))) / pad_len, with + # the 1/pad_len already folded into `chirp_fft`. + fft_pow2_radix4!(tmp, a, pad_len, 1, 1, 1, 1, FFT_FORWARD, tw, 0) + @inbounds for i in 1:pad_len + tmp[i] = conj(tmp[i] * chirp_fft[i]) end + fft_pow2_radix4!(a, tmp, pad_len, 1, 1, 1, 1, FFT_FORWARD, tw, 0) - out_inds = range(start_out; step=stride_out, length=N) - copyto!(out, CartesianIndices((out_inds,)), Xk, CartesianIndices((N,))) + # X_k = conj(b_k) · conv_k + @inbounds for i in 1:N + out[start_out + (i-1)*stride_out] = conj(chirp[i]) * conj(a[i]) + end return nothing end diff --git a/src/callgraph.jl b/src/callgraph.jl index f672157..3316737 100644 --- a/src/callgraph.jl +++ b/src/callgraph.jl @@ -2,6 +2,10 @@ @enum Pow24 POW2 POW4 @enum FFTEnum COMPOSITE_FFT DFT POW3_FFT POW2RADIX4_FFT BLUESTEIN +@inline function direction_sign(d::Direction) + Int(d) +end + """ $(TYPEDEF) Node of a call graph @@ -24,6 +28,29 @@ struct CallGraphNode s_out::Int end +""" +$(TYPEDEF) +Precomputed data and scratch space for Bluestein's algorithm on a length-`N` +transform in direction `dir` (see `fft_bluestein!`). + +# Fields +- `N`: Length of the transform +- `pad_len`: Power-of-two length ≥ 2N-1 of the padded convolution +- `chirp`: `b_n = exp(∓iπ n²/N)` for `n = 0..N-1` +- `chirp_fft`: forward FFT of the zero-padded, periodised chirp, divided by `pad_len` +- `a`, `tmp`: work arrays of length `pad_len` +- `tw`: forward power-of-two twiddle tables for `pad_len` (see `pow2_twiddles`) +""" +struct BluesteinScratch{T<:Complex} + N::Int + pad_len::Int + chirp::Vector{T} + chirp_fft::Vector{T} + a::Vector{T} + tmp::Vector{T} + tw::Vector{T} +end + """ $(TYPEDEF) Object representing a graph of FFT Calls @@ -31,6 +58,10 @@ Object representing a graph of FFT Calls # Arguments - `nodes`: Nodes keeping track of the graph - `workspace`: Preallocated Workspace +- `twiddles`: Precomputed twiddle factors of each node (see `node_twiddles`) +- `bluestein`: Precomputed data for each `BLUESTEIN` node, indexed through `blue_index` +- `blue_index`: Index into `bluestein` for each node (`0` for other node types) +- `dir`: Direction the twiddle factors were computed for - `BLUESTEIN_CUTOFF`: Minimum prime that will be FFTed with the Bluestein algorithm, below which the O(N^2) DFT is used. @@ -38,6 +69,10 @@ Object representing a graph of FFT Calls struct CallGraph{T<:Complex} nodes::Vector{CallGraphNode} workspace::Vector{Vector{T}} + twiddles::Vector{Vector{T}} + bluestein::Vector{BluesteinScratch{T}} + blue_index::Vector{Int} + dir::Direction BLUESTEIN_CUTOFF::Int end @@ -46,19 +81,6 @@ const DEFAULT_BLUESTEIN_CUTOFF = 73 # Get the node in the graph at index i Base.getindex(g::CallGraph{T}, i::Int) where {T} = g.nodes[i] -""" -$(TYPEDSIGNATURES) -Check if `N` is a power of 2 or 4 - -""" -# function _ispow24(N::Int) -# if ispow2(N) -# zero_cnt = trailing_zeros(N) -# return iseven(zero_cnt) ? POW4 : POW2 -# end -# return nothing -# end - """ $(TYPEDSIGNATURES) Recursively instantiate a set of `CallGraphNode`s @@ -84,7 +106,6 @@ function CallGraphNode!( throw(DimensionMismatch("Array length must be strictly positive")) end if iseven(N) && ispow2(N) - # _ispow24(N) push!(workspace, T[]) push!(nodes, CallGraphNode(0, 0, POW2RADIX4_FFT, N, s_in, s_out)) return 1 @@ -125,14 +146,218 @@ function CallGraphNode!( return 1 + left_len + right_len end +# --------------------------------------------------------------------------- +# Twiddle factors +# +# All twiddle factors are computed once at plan time with `cispi`, which is +# accurate to a rounding error, and stored per node in a layout chosen so that +# the kernels read them sequentially. `k` is reduced modulo `N` before the +# conversion to floating point so that large exponents lose no accuracy. +# --------------------------------------------------------------------------- + +""" +$(TYPEDSIGNATURES) +The twiddle factor `exp(dir · 2πi · k / N)` as an element of type `T`. +""" +@inline function twiddle(::Type{T}, dir::Direction, k::Integer, N::Integer) where {T<:Complex} + R = real(T) + T(cispi(2 * R(direction_sign(dir) * mod(k, N)) / R(N))) +end + +""" +$(TYPEDSIGNATURES) +The `N` unit roots `W[k + 1] = exp(dir · 2πi k / N)`, `k = 0..N-1`. When `8 | N` +only the first octant is evaluated with `sincospi` and the rest follows from +symmetry, which makes planning large power-of-two transforms cheap. +""" +function unit_roots(::Type{T}, N::Int, dir::Direction) where {T<:Complex} + W = Vector{T}(undef, N) + if N % 8 != 0 + for k in 0:N-1 + W[k + 1] = twiddle(T, dir, k, N) + end + return W + end + R = real(T) + sgn = R(direction_sign(dir)) + q = N ÷ 4 + for k in 0:N÷8 + s, c = sincospi(2 * R(k) / R(N)) + s *= sgn + # with s = sgn·sin θ, c = cos θ and the imaginary part carrying `sgn`: + W[k + 1] = T(c, s) # θ + W[q - k + 1] = T(sgn * s, sgn * c) # π/2 - θ + W[q + k + 1] = T(-sgn * s, sgn * c) # π/2 + θ + W[2q - k + 1] = T(-c, s) # π - θ + W[2q + k + 1] = T(-c, -s) # π + θ + W[3q - k + 1] = T(-sgn * s, -sgn * c) # 3π/2 - θ + W[3q + k + 1] = T(sgn * s, -sgn * c) # 3π/2 + θ + k > 0 && (W[N - k + 1] = T(c, -s)) # 2π - θ + end + return W +end + +""" +$(TYPEDSIGNATURES) +Twiddle table for the O(N²) DFT: `W[k + 1] = exp(dir · 2πi k / N)`, `k = 0..N-1`. +""" +dft_twiddles(::Type{T}, N::Int, dir::Direction) where {T} = unit_roots(T, N, dir) + +""" +$(TYPEDSIGNATURES) +Twiddle table for the Cooley-Tukey step of a composite `N = N1 · N2` transform: +the factor for output row `j1` and inner index `k2` is stored at +`(j1 - 1) * (N2 - 1) + k2` for `j1 = 1..N1-1`, `k2 = 1..N2-1`. +""" +function composite_twiddles(::Type{T}, N::Int, N1::Int, N2::Int, dir::Direction) where {T} + W = unit_roots(T, N, dir) + tw = Vector{T}(undef, (N1 - 1) * (N2 - 1)) + i = 1 + for j1 in 1:N1-1 + idx = 0 + for k2 in 1:N2-1 + idx += j1 + idx >= N && (idx -= N) + tw[i] = W[idx + 1] + i += 1 + end + end + return tw +end + +""" +$(TYPEDSIGNATURES) +Twiddle table for the radix-4 power-of-two kernel. For every recursion level of +size `M = N, N/4, N/16, …` (down to, but excluding, the 4- and 2-point base +cases) the table holds the triplets `(w^k, w^2k, w^3k)`, `w = exp(dir · 2πi/M)`, +for `k = 0..M/4-1`, one level after the other. A level of size `M` therefore +occupies `3M/4` entries and the next level starts `3M/4` entries later. +""" +function pow2_twiddles(::Type{T}, N::Int, dir::Direction) where {T} + N > 4 || return T[] + W = unit_roots(T, N, dir) + tw = Vector{T}(undef, (N - 2) ) # 3N/4 + 3N/16 + ... < N + i = 1 + M = N + while M > 4 + m = M ÷ 4 + s = N ÷ M # w_M^k = w_N^(s k) + for k in 0:m-1 + tw[i] = W[s * k + 1] + tw[i + 1] = W[2 * s * k + 1] + tw[i + 2] = W[3 * s * k + 1] + i += 3 + end + M = m + end + resize!(tw, i - 1) + return tw +end + +""" +$(TYPEDSIGNATURES) +Twiddle table for the radix-3 kernel, laid out like `pow2_twiddles`: for every +level of size `M = N, N/3, …` (excluding the 3-point base case) the pairs +`(w^k, w^2k)`, `w = exp(dir · 2πi/M)`, for `k = 0..M/3-1`. +""" +function pow3_twiddles(::Type{T}, N::Int, dir::Direction) where {T} + N > 3 || return T[] + W = unit_roots(T, N, dir) + tw = Vector{T}(undef, N) # 2N/3 + 2N/9 + ... < N + i = 1 + M = N + while M > 3 + m = M ÷ 3 + s = N ÷ M + for k in 0:m-1 + tw[i] = W[s * k + 1] + tw[i + 1] = W[2 * s * k + 1] + i += 2 + 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`, +`composite_twiddles`, `pow2_twiddles` and `pow3_twiddles`. `BLUESTEIN` nodes +keep their data in a `BluesteinScratch` instead and get an empty table. +""" +function node_twiddles(::Type{T}, nodes::Vector{CallGraphNode}, idx::Int, dir::Direction) where {T} + node = nodes[idx] + N = node.sz + if node.type === COMPOSITE_FFT + N1 = nodes[idx + node.left].sz + N2 = nodes[idx + node.right].sz + return composite_twiddles(T, N, N1, N2, dir) + elseif node.type === DFT + return dft_twiddles(T, N, dir) + elseif node.type === POW2RADIX4_FFT + return pow2_twiddles(T, N, dir) + elseif node.type === POW3_FFT + return pow3_twiddles(T, N, dir) + else + return T[] + end +end + +""" +$(TYPEDSIGNATURES) +Precompute the chirp, its transform, the work arrays and the twiddle tables +used by `fft_bluestein!` for a length-`N` transform in direction `d`. +""" +function BluesteinScratch{T}(N::Int, d::Direction) where {T<:Complex} + pad_len = nextpow(2, 2N - 1) + R = real(T) + + # chirp b_n = exp(sgn · iπ n²/N); n² is tracked modulo 2N so that the + # argument of `cispi` stays small. + chirp = Vector{T}(undef, N) + sgn = -direction_sign(d) + p = 0 # n^2 mod 2N, kept in (-N, N] + for i in 1:N + chirp[i] = T(cispi(R(sgn * p) / R(N))) + p += (2i - 1) # prevents overflow unless N is absolutely massive + p > N && (p -= 2N) + end + + tw = pow2_twiddles(T, pad_len, FFT_FORWARD) + + # periodised, zero-padded chirp and its forward transform, scaled by + # 1/pad_len so that the inverse transform in `fft_bluestein!` needs no + # further normalisation + a = zeros(T, pad_len) + copyto!(a, 1, chirp, 1, N) + for j in 0:N-2 + a[pad_len - j] = chirp[2 + j] + end + chirp_fft = Vector{T}(undef, pad_len) + fft_pow2_radix4!(chirp_fft, a, pad_len, 1, 1, 1, 1, FFT_FORWARD, tw, 0) + chirp_fft ./= R(pad_len) + + return BluesteinScratch{T}(N, pad_len, chirp, chirp_fft, a, Vector{T}(undef, pad_len), tw) +end + """ $(TYPEDSIGNATURES) -Instantiate a CallGraph from a number `N` +Instantiate a CallGraph from a number `N`, with twiddle factors for direction `dir` """ -function CallGraph{T}(N::Int, BLUESTEIN_CUTOFF::Int) where {T} +function CallGraph{T}(N::Int, BLUESTEIN_CUTOFF::Int, dir::Direction=FFT_FORWARD) where {T} nodes = CallGraphNode[] workspace = Vector{Vector{T}}() CallGraphNode!(nodes, N, workspace, BLUESTEIN_CUTOFF, 1, 1) - CallGraph(nodes, workspace, BLUESTEIN_CUTOFF) + twiddles = [node_twiddles(T, nodes, idx, dir) for idx in eachindex(nodes)] + bluestein = BluesteinScratch{T}[] + blue_index = zeros(Int, length(nodes)) + for (idx, node) in enumerate(nodes) + if node.type === BLUESTEIN + push!(bluestein, BluesteinScratch{T}(node.sz, dir)) + blue_index[idx] = length(bluestein) + end + end + CallGraph(nodes, workspace, twiddles, bluestein, blue_index, dir, BLUESTEIN_CUTOFF) end diff --git a/src/plan.jl b/src/plan.jl index a71cbc8..f380f61 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -83,17 +83,17 @@ function _plan_fft( M = length(region) if M == 1 R1 = Int(region[1]) - g = CallGraph{T}(size(x, R1), BLUESTEIN_CUTOFF) + g = CallGraph{T}(size(x, R1), BLUESTEIN_CUTOFF, dir) return FFTAPlan_cx{T,1}((g,), R1, dir) elseif M == 2 R2 = _sort(region) - g1 = CallGraph{T}(size(x, R2[1]), BLUESTEIN_CUTOFF) - g2 = CallGraph{T}(size(x, R2[2]), BLUESTEIN_CUTOFF) + g1 = CallGraph{T}(size(x, R2[1]), BLUESTEIN_CUTOFF, dir) + g2 = CallGraph{T}(size(x, R2[2]), BLUESTEIN_CUTOFF, dir) return FFTAPlan_cx{T,2}((g1, g2), R2, dir) else RM = _sort(region) return FFTAPlan_cx{T,M}( - ntuple(i -> CallGraph{T}(size(x, RM[i]), BLUESTEIN_CUTOFF), Val(M)), + ntuple(i -> CallGraph{T}(size(x, RM[i]), BLUESTEIN_CUTOFF, dir), Val(M)), RM, dir ) end @@ -112,12 +112,12 @@ function AbstractFFTs.plan_rfft( # two n/2 complex FFTs followed by a butterfly. For odd size # problems, we just solve the problem as a single complex nn = iseven(n) ? n >> 1 : n - g = CallGraph{Complex{T}}(nn, BLUESTEIN_CUTOFF) + g = CallGraph{Complex{T}}(nn, BLUESTEIN_CUTOFF, FFT_FORWARD) return FFTAPlan_re{Complex{T},1}((g,), R1, FFT_FORWARD, n) elseif M == 2 R2 = _sort(region) - g1 = CallGraph{Complex{T}}(size(x, R2[1]), BLUESTEIN_CUTOFF) - g2 = CallGraph{Complex{T}}(size(x, R2[2]), BLUESTEIN_CUTOFF) + g1 = CallGraph{Complex{T}}(size(x, R2[1]), BLUESTEIN_CUTOFF, FFT_FORWARD) + g2 = CallGraph{Complex{T}}(size(x, R2[2]), BLUESTEIN_CUTOFF, FFT_FORWARD) return FFTAPlan_re{Complex{T},2}((g1, g2), R2, FFT_FORWARD, size(x, R2[1])) else throw(ArgumentError("only supports 1D and 2D FFTs")) @@ -137,12 +137,12 @@ function AbstractFFTs.plan_brfft( # problems, we just solve the problem as a single complex R1 = Int(region[1]) nn = iseven(len) ? len >> 1 : len - g = CallGraph{T}(nn, BLUESTEIN_CUTOFF) + g = CallGraph{T}(nn, BLUESTEIN_CUTOFF, FFT_BACKWARD) return FFTAPlan_re{T,1}((g,), R1, FFT_BACKWARD, len) elseif M == 2 R2 = _sort(region) - g1 = CallGraph{T}(len, BLUESTEIN_CUTOFF) - g2 = CallGraph{T}(size(x, R2[2]), BLUESTEIN_CUTOFF) + g1 = CallGraph{T}(len, BLUESTEIN_CUTOFF, FFT_BACKWARD) + g2 = CallGraph{T}(size(x, R2[2]), BLUESTEIN_CUTOFF, FFT_BACKWARD) return FFTAPlan_re{T,2}((g1, g2), R2, FFT_BACKWARD, len) else throw(ArgumentError("only supports 1D and 2D FFTs")) diff --git a/test/onedim/accuracy.jl b/test/onedim/accuracy.jl index b4f357f..0a227dc 100644 --- a/test/onedim/accuracy.jl +++ b/test/onedim/accuracy.jl @@ -28,6 +28,8 @@ const POWERS_OF_2 = ( (1 << 16, 22.0), # 65536 = 4^8 (1 << 17, 28.0), # 131072 (1 << 18, 28.0), # 262144 = 4^9 + (1 << 20, 4.0), # 1048576 = 4^10 (with stored tables ~1.4; ~60 with the recurrence) + (1 << 22, 4.0), # 4194304 = 4^11 (with stored tables ~1.5; ~1000 with the recurrence) ) const POWERS_OF_3 = ( diff --git a/test/onedim/complex_backward.jl b/test/onedim/complex_backward.jl index f4c5785..004fd3b 100644 --- a/test/onedim/complex_backward.jl +++ b/test/onedim/complex_backward.jl @@ -20,7 +20,7 @@ end end @testset "allocation regression" begin - @test (@test_allocations bfft(y)) <= 47 + @test (@test_allocations bfft(y)) <= 80 end end diff --git a/test/onedim/complex_forward.jl b/test/onedim/complex_forward.jl index 0a7d35e..5abafc7 100644 --- a/test/onedim/complex_forward.jl +++ b/test/onedim/complex_forward.jl @@ -20,7 +20,8 @@ end end @testset "allocation regression" begin - @test (@test_allocations fft(x)) <= 47 + # plan creation allocates one twiddle table per node (and the Bluestein scratch) + @test (@test_allocations fft(x)) <= 80 end end diff --git a/test/runtests.jl b/test/runtests.jl index 9aeb1e6..00365dc 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -47,6 +47,9 @@ Random.seed!(1) include("qa/aqua.jl") include("qa/explicit_imports.jl") end + @testset verbose = true "Twiddle tables" begin + include("twiddles.jl") + end @testset verbose = true "Argument checking" begin include("argument_checking.jl") end diff --git a/test/twiddles.jl b/test/twiddles.jl new file mode 100644 index 0000000..311275f --- /dev/null +++ b/test/twiddles.jl @@ -0,0 +1,112 @@ +using FFTA, Test, LinearAlgebra + +# Twiddle tables are built at plan time (see src/callgraph.jl). These tests +# pin down their layout and accuracy against direct evaluation. + +@testset "unit_roots: N=$N, $dir, $T" for N in (1, 2, 3, 5, 8, 12, 16, 24, 64, 73, 1000, 4096), + dir in (FFTA.FFT_FORWARD, FFTA.FFT_BACKWARD), + T in (ComplexF64, ComplexF32, Complex{BigFloat}) + W = FFTA.unit_roots(T, N, dir) + @test length(W) == N + @test eltype(W) == T + ref = [FFTA.twiddle(T, dir, k, N) for k in 0:N-1] + @test maximum(abs.(W .- ref)) <= 2 * eps(real(T)) + if N > 1 + # w^k · w^(N-k) == 1 for the symmetric construction + @test maximum(abs.(W[2:end] .* reverse(W[2:end]) .- 1)) <= 4 * eps(real(T)) + end +end + +@testset "pow2_twiddles layout, N=$N" for N in (2, 4, 8, 16, 32, 64, 256, 1024) + for dir in (FFTA.FFT_FORWARD, FFTA.FFT_BACKWARD) + tw = FFTA.pow2_twiddles(ComplexF64, N, dir) + off = 0 + M = N + while M > 4 + m = M ÷ 4 + for k in 0:m-1, j in 1:3 + @test tw[off + 3k + j] ≈ FFTA.twiddle(ComplexF64, dir, j * k, M) atol=1e-15 + end + off += 3m + M = m + end + @test length(tw) == off + end +end + +@testset "pow3_twiddles layout, N=$N" for N in (3, 9, 27, 81, 729) + for dir in (FFTA.FFT_FORWARD, FFTA.FFT_BACKWARD) + tw = FFTA.pow3_twiddles(ComplexF64, N, dir) + off = 0 + M = N + while M > 3 + m = M ÷ 3 + for k in 0:m-1, j in 1:2 + @test tw[off + 2k + j] ≈ FFTA.twiddle(ComplexF64, dir, j * k, M) atol=1e-15 + end + off += 2m + M = m + end + @test length(tw) == off + end +end + +@testset "composite_twiddles layout" begin + for (N1, N2) in ((4, 5), (5, 7), (8, 125), (3, 3)) + N = N1 * N2 + tw = FFTA.composite_twiddles(ComplexF64, N, N1, N2, FFTA.FFT_FORWARD) + @test length(tw) == (N1 - 1) * (N2 - 1) + for j1 in 1:N1-1, k2 in 1:N2-1 + @test tw[(j1 - 1) * (N2 - 1) + k2] ≈ FFTA.twiddle(ComplexF64, FFTA.FFT_FORWARD, j1 * k2, N) atol=1e-15 + end + end +end + +@testset "CallGraph carries tables for its direction" begin + for dir in (FFTA.FFT_FORWARD, FFTA.FFT_BACKWARD) + g = FFTA.CallGraph{ComplexF64}(1000, FFTA.DEFAULT_BLUESTEIN_CUTOFF, dir) + @test g.dir === dir + @test length(g.twiddles) == length(g.nodes) == length(g.blue_index) + @test all(iszero, g.blue_index) + @test isempty(g.bluestein) + for (i, n) in enumerate(g.nodes) + n.type === FFTA.COMPOSITE_FFT && @test length(g.twiddles[i]) == (g.nodes[i + n.left].sz - 1) * (g.nodes[i + n.right].sz - 1) + n.type === FFTA.DFT && @test length(g.twiddles[i]) == n.sz + end + gb = FFTA.CallGraph{ComplexF64}(2 * 1009, FFTA.DEFAULT_BLUESTEIN_CUTOFF, dir) + bi = findfirst(n -> n.type === FFTA.BLUESTEIN, gb.nodes) + @test bi !== nothing + @test gb.blue_index[bi] == 1 + s = gb.bluestein[1] + @test s.N == 1009 && s.pad_len == 2048 && length(s.chirp) == 1009 && length(s.chirp_fft) == 2048 + end + # the opposite direction is rejected + g = FFTA.CallGraph{ComplexF64}(8, FFTA.DEFAULT_BLUESTEIN_CUTOFF, FFTA.FFT_FORWARD) + y = zeros(ComplexF64, 8) + @test_throws ArgumentError FFTA.fft!(y, ones(ComplexF64, 8), 1, 1, FFTA.FFT_BACKWARD, g[1].type, g, 1) +end + +@testset "kernels with tables computed on the fly" begin + # convenience methods used by the tests and for experimentation + x = randn(ComplexF64, 64) + y1 = similar(x); y2 = similar(x) + FFTA.fft_pow2_radix4!(y1, x, 64, 1, 1, 1, 1, FFTA.FFT_FORWARD) + @test y1 ≈ fft(x) + FFTA.fft_dft!(y2, x, 64, 1, 1, 1, 1, FFTA.FFT_FORWARD) + @test y2 ≈ fft(x) + x3 = randn(ComplexF64, 27); y3 = similar(x3) + m120 = cispi(2 / 3) + FFTA.fft_pow3!(y3, x3, 27, 1, 1, 1, 1, m120, FFTA.FFT_FORWARD) + @test y3 ≈ fft(x3) + xb = randn(ComplexF64, 101); yb = similar(xb) + FFTA.fft_bluestein!(yb, xb, FFTA.FFT_FORWARD, 101, 1, 1, 1, 1) + @test yb ≈ fft(xb) +end + +@testset "planned execution does not allocate, n=$n" for n in (5, 64, 73, 101, 720, 1000, 1009, 4096, 65537) + x = randn(ComplexF64, n) + p = plan_fft(x) + y = p * x + mul!(y, p, x) + @test (@test_allocations mul!(y, p, x)) == 0 +end From 33507b8df3226721df054e8522ea62404e62d4e5 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:32:56 +0000 Subject: [PATCH 2/7] Bluestein: 3-smooth padding, embedded call graph, cutoff 47 The padded convolution length is now the smallest 2^a 3^b >= 2N-1 when that is enough smaller than the next power of two to be cheaper (a length with factors of 3 costs about 1.6x per n log n compared to a power of two in FFTA), e.g. 4099 -> 8748 instead of 16384 and 65537 -> 139968 instead of 262144. The padded transforms run through a CallGraph embedded in the BluesteinScratch instead of the bare power-of-two kernel. With twiddle tables the O(N^2) DFT leaf and Bluestein cross over near N = 45 (measured for ComplexF64), so DEFAULT_BLUESTEIN_CUTOFF drops from 73 to 47. --- src/algos.jl | 11 ++++--- src/callgraph.jl | 55 +++++++++++++++++++++++++-------- test/onedim/complex_backward.jl | 2 +- test/onedim/complex_forward.jl | 2 +- test/twiddles.jl | 16 ++++++++-- 5 files changed, 64 insertions(+), 22 deletions(-) diff --git a/src/algos.jl b/src/algos.jl index 322159d..0cc4bf6 100644 --- a/src/algos.jl +++ b/src/algos.jl @@ -326,8 +326,8 @@ $(TYPEDSIGNATURES) Bluestein's algorithm, still O(N * log(N)) for large primes, but with a big constant factor. Zero-pads two sequences derived from the DFT formula to a -power of 2 length greater than `2N-1` and computes their convolution -with a power 2 FFT. The chirp, its transform and the work arrays are +3-smooth length ≥ `2N-1` (see `bluestein_pad_length`) and computes their +convolution with FFTs of that length. The chirp, its transform and the work arrays are precomputed in `scratch` (see `BluesteinScratch`). # Arguments @@ -350,7 +350,8 @@ function fft_bluestein!( start_in::Int, stride_in::Int, scratch::BluesteinScratch{T}=BluesteinScratch{T}(N, d) ) where T<:Complex - (; pad_len, chirp, chirp_fft, a, tmp, tw) = scratch + (; pad_len, chirp, chirp_fft, a, tmp, graph) = scratch + gt = graph[1].type # a_n = x_n · conj(b_n), zero padded @inbounds for i in 1:N @@ -363,11 +364,11 @@ function fft_bluestein!( # Circular convolution of `a` with the periodised chirp via the forward # transform only: conv = conj(fft(conj(fft(a) .* fft(b)))) / pad_len, with # the 1/pad_len already folded into `chirp_fft`. - fft_pow2_radix4!(tmp, a, pad_len, 1, 1, 1, 1, FFT_FORWARD, tw, 0) + fft!(tmp, a, 1, 1, FFT_FORWARD, gt, graph, 1) @inbounds for i in 1:pad_len tmp[i] = conj(tmp[i] * chirp_fft[i]) end - fft_pow2_radix4!(a, tmp, pad_len, 1, 1, 1, 1, FFT_FORWARD, tw, 0) + fft!(a, tmp, 1, 1, FFT_FORWARD, gt, graph, 1) # X_k = conj(b_k) · conv_k @inbounds for i in 1:N diff --git a/src/callgraph.jl b/src/callgraph.jl index 3316737..432b7be 100644 --- a/src/callgraph.jl +++ b/src/callgraph.jl @@ -35,20 +35,20 @@ transform in direction `dir` (see `fft_bluestein!`). # Fields - `N`: Length of the transform -- `pad_len`: Power-of-two length ≥ 2N-1 of the padded convolution +- `pad_len`: Length ≥ 2N-1 of the padded convolution (see `bluestein_pad_length`) - `chirp`: `b_n = exp(∓iπ n²/N)` for `n = 0..N-1` - `chirp_fft`: forward FFT of the zero-padded, periodised chirp, divided by `pad_len` - `a`, `tmp`: work arrays of length `pad_len` -- `tw`: forward power-of-two twiddle tables for `pad_len` (see `pow2_twiddles`) +- `graph`: forward `CallGraph` for the length-`pad_len` transforms """ -struct BluesteinScratch{T<:Complex} +struct BluesteinScratch{T<:Complex,G} N::Int pad_len::Int chirp::Vector{T} chirp_fft::Vector{T} a::Vector{T} tmp::Vector{T} - tw::Vector{T} + graph::G end """ @@ -70,13 +70,15 @@ struct CallGraph{T<:Complex} nodes::Vector{CallGraphNode} workspace::Vector{Vector{T}} twiddles::Vector{Vector{T}} - bluestein::Vector{BluesteinScratch{T}} + bluestein::Vector{BluesteinScratch{T,CallGraph{T}}} blue_index::Vector{Int} dir::Direction BLUESTEIN_CUTOFF::Int end -const DEFAULT_BLUESTEIN_CUTOFF = 73 +# Primes below this use the O(N²) DFT with a twiddle table; at and above it +# Bluestein's algorithm is cheaper (crossover measured at ~45 for ComplexF64). +const DEFAULT_BLUESTEIN_CUTOFF = 47 # Get the node in the graph at index i Base.getindex(g::CallGraph{T}, i::Int) where {T} = g.nodes[i] @@ -306,11 +308,37 @@ end """ $(TYPEDSIGNATURES) -Precompute the chirp, its transform, the work arrays and the twiddle tables -used by `fft_bluestein!` for a length-`N` transform in direction `d`. +Length of the padded convolution in Bluestein's algorithm for a length-`N` +transform: the smallest power of two ≥ 2N-1, unless a 3-smooth length +`2^a 3^b ≥ 2N-1` is enough smaller to be cheaper — a transform of a length +with factors of 3 costs about 1.6× per `n log n` in FFTA compared to a +power of two, so `m` is preferred over `p` when `1.6 m log m < p log p`. +""" +function bluestein_pad_length(N::Int) + m = 2N - 1 + p = nextpow(2, m) + best = p + best_cost = p * log2(p) + pow3 = 3 + while pow3 < p + c = pow3 * nextpow(2, cld(m, pow3)) # smallest 2^a·3^b with this power of 3 + cost = 1.6 * c * log2(c) + if c >= m && cost < best_cost + best, best_cost = c, cost + end + pow3 *= 3 + end + return best +end + +""" +$(TYPEDSIGNATURES) +Precompute the chirp, its transform, the work arrays and the call graph of the +padded transform used by `fft_bluestein!` for a length-`N` transform in +direction `d`. """ function BluesteinScratch{T}(N::Int, d::Direction) where {T<:Complex} - pad_len = nextpow(2, 2N - 1) + pad_len = bluestein_pad_length(N) R = real(T) # chirp b_n = exp(sgn · iπ n²/N); n² is tracked modulo 2N so that the @@ -324,7 +352,8 @@ function BluesteinScratch{T}(N::Int, d::Direction) where {T<:Complex} p > N && (p -= 2N) end - tw = pow2_twiddles(T, pad_len, FFT_FORWARD) + # the padded length is 3-smooth, so its graph has no Bluestein node + graph = CallGraph{T}(pad_len, 2, FFT_FORWARD) # periodised, zero-padded chirp and its forward transform, scaled by # 1/pad_len so that the inverse transform in `fft_bluestein!` needs no @@ -335,10 +364,10 @@ function BluesteinScratch{T}(N::Int, d::Direction) where {T<:Complex} a[pad_len - j] = chirp[2 + j] end chirp_fft = Vector{T}(undef, pad_len) - fft_pow2_radix4!(chirp_fft, a, pad_len, 1, 1, 1, 1, FFT_FORWARD, tw, 0) + fft!(chirp_fft, a, 1, 1, FFT_FORWARD, graph[1].type, graph, 1) chirp_fft ./= R(pad_len) - return BluesteinScratch{T}(N, pad_len, chirp, chirp_fft, a, Vector{T}(undef, pad_len), tw) + return BluesteinScratch{T,CallGraph{T}}(N, pad_len, chirp, chirp_fft, a, Vector{T}(undef, pad_len), graph) end """ @@ -351,7 +380,7 @@ function CallGraph{T}(N::Int, BLUESTEIN_CUTOFF::Int, dir::Direction=FFT_FORWARD) workspace = Vector{Vector{T}}() CallGraphNode!(nodes, N, workspace, BLUESTEIN_CUTOFF, 1, 1) twiddles = [node_twiddles(T, nodes, idx, dir) for idx in eachindex(nodes)] - bluestein = BluesteinScratch{T}[] + bluestein = BluesteinScratch{T,CallGraph{T}}[] blue_index = zeros(Int, length(nodes)) for (idx, node) in enumerate(nodes) if node.type === BLUESTEIN diff --git a/test/onedim/complex_backward.jl b/test/onedim/complex_backward.jl index 004fd3b..1741001 100644 --- a/test/onedim/complex_backward.jl +++ b/test/onedim/complex_backward.jl @@ -20,7 +20,7 @@ end end @testset "allocation regression" begin - @test (@test_allocations bfft(y)) <= 80 + @test (@test_allocations bfft(y)) <= 120 end end diff --git a/test/onedim/complex_forward.jl b/test/onedim/complex_forward.jl index 5abafc7..449da77 100644 --- a/test/onedim/complex_forward.jl +++ b/test/onedim/complex_forward.jl @@ -21,7 +21,7 @@ end @testset "allocation regression" begin # plan creation allocates one twiddle table per node (and the Bluestein scratch) - @test (@test_allocations fft(x)) <= 80 + @test (@test_allocations fft(x)) <= 120 end end diff --git a/test/twiddles.jl b/test/twiddles.jl index 311275f..facc41a 100644 --- a/test/twiddles.jl +++ b/test/twiddles.jl @@ -78,7 +78,7 @@ end @test bi !== nothing @test gb.blue_index[bi] == 1 s = gb.bluestein[1] - @test s.N == 1009 && s.pad_len == 2048 && length(s.chirp) == 1009 && length(s.chirp_fft) == 2048 + @test s.N == 1009 && s.pad_len == FFTA.bluestein_pad_length(1009) == 2048 && length(s.chirp) == 1009 && length(s.chirp_fft) == 2048 end # the opposite direction is rejected g = FFTA.CallGraph{ComplexF64}(8, FFTA.DEFAULT_BLUESTEIN_CUTOFF, FFTA.FFT_FORWARD) @@ -103,7 +103,19 @@ end @test yb ≈ fft(xb) end -@testset "planned execution does not allocate, n=$n" for n in (5, 64, 73, 101, 720, 1000, 1009, 4096, 65537) +@testset "bluestein_pad_length" begin + for N in (47, 73, 127, 1009, 4099, 65537, 120779) + m = FFTA.bluestein_pad_length(N) + @test m >= 2N - 1 + f = FFTA.Primes.factor(Dict, m) + @test all(p -> p in (2, 3), keys(f)) + end + @test FFTA.bluestein_pad_length(1009) == 2048 + @test FFTA.bluestein_pad_length(4099) in (8748, 9216) # 3-smooth, cheaper than 16384 + @test FFTA.bluestein_pad_length(65537) < 262144 +end + +@testset "planned execution does not allocate, n=$n" for n in (5, 47, 64, 73, 101, 720, 1000, 1009, 4096, 65537) x = randn(ComplexF64, n) p = plan_fft(x) y = p * x From 919628e0f4dc09186c356509f95a1aa1199b4481 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:33:41 +0000 Subject: [PATCH 3/7] docs: list bluestein_pad_length --- docs/src/dev.md | 1 + 1 file changed, 1 insertion(+) diff --git a/docs/src/dev.md b/docs/src/dev.md index 81988b4..d5e465d 100644 --- a/docs/src/dev.md +++ b/docs/src/dev.md @@ -32,4 +32,5 @@ pow2_twiddles pow3_twiddles node_twiddles BluesteinScratch +bluestein_pad_length ``` From 071712b6794ff6677730d6860b81df576fa9a893 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 20:41:55 +0000 Subject: [PATCH 4/7] Bluestein padding: calibrate the 3-smooth cost factor to 1.9 and keep powers of two below 2048 Measured against the power-of-two alternative, the 3-smooth pad was slower for small transforms (73 -> 162 took 6.7 us instead of 5.1 us with 256): the composite step's fixed overhead dominates there, and the 1.6x per-element factor was optimistic; 1.9x fits the measured pairs (162/256, 8748/16384, 17496/32768, 139968/262144). --- src/callgraph.jl | 13 ++++++++----- test/twiddles.jl | 4 +++- 2 files changed, 11 insertions(+), 6 deletions(-) diff --git a/src/callgraph.jl b/src/callgraph.jl index 432b7be..add561a 100644 --- a/src/callgraph.jl +++ b/src/callgraph.jl @@ -310,9 +310,12 @@ end $(TYPEDSIGNATURES) Length of the padded convolution in Bluestein's algorithm for a length-`N` transform: the smallest power of two ≥ 2N-1, unless a 3-smooth length -`2^a 3^b ≥ 2N-1` is enough smaller to be cheaper — a transform of a length -with factors of 3 costs about 1.6× per `n log n` in FFTA compared to a -power of two, so `m` is preferred over `p` when `1.6 m log m < p log p`. +`2^a 3^b ≥ 2N-1` is enough smaller to be cheaper. Measured on FFTA, a +transform of a length with factors of 3 costs about 1.9× per `n log n` +compared to a power of two (the composite step and the radix-3 kernel are +slower than the radix-4 one), so `m` is preferred over `p` when +`1.9 m log m < p log p`; below 2048 the fixed overhead of the composite +step dominates and the power of two is always used. """ function bluestein_pad_length(N::Int) m = 2N - 1 @@ -322,8 +325,8 @@ function bluestein_pad_length(N::Int) pow3 = 3 while pow3 < p c = pow3 * nextpow(2, cld(m, pow3)) # smallest 2^a·3^b with this power of 3 - cost = 1.6 * c * log2(c) - if c >= m && cost < best_cost + cost = 1.9 * c * log2(c) + if c >= max(m, 2048) && cost < best_cost best, best_cost = c, cost end pow3 *= 3 diff --git a/test/twiddles.jl b/test/twiddles.jl index facc41a..0318850 100644 --- a/test/twiddles.jl +++ b/test/twiddles.jl @@ -110,9 +110,11 @@ end f = FFTA.Primes.factor(Dict, m) @test all(p -> p in (2, 3), keys(f)) end + @test FFTA.bluestein_pad_length(73) == 256 # small: always a power of two @test FFTA.bluestein_pad_length(1009) == 2048 @test FFTA.bluestein_pad_length(4099) in (8748, 9216) # 3-smooth, cheaper than 16384 - @test FFTA.bluestein_pad_length(65537) < 262144 + @test FFTA.bluestein_pad_length(8443) == 17496 + @test FFTA.bluestein_pad_length(65537) == 139968 end @testset "planned execution does not allocate, n=$n" for n in (5, 47, 64, 73, 101, 720, 1000, 1009, 4096, 65537) From 42fccda30fd17de950f67df2b88d8585a8c5b190 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 20:57:49 +0000 Subject: [PATCH 5/7] Bluestein: conservative 3-smooth threshold (2.1) and cutoff 29, from measurements on two architectures On x86-64 (AVX2) a 3-smooth length costs 2.0-3.1x per n log n relative to a power of two, against 1.3-2.3x on aarch64 (NEON); with the previous 1.9 threshold the 3-smooth pad was a 1.58x win at N = 8443 on aarch64 and a 1.19x loss on x86-64. A threshold of 2.1 never loses on either machine. The DFT/Bluestein crossover is n = 23 on x86-64 and about 29 on aarch64, so DEFAULT_BLUESTEIN_CUTOFF drops from 47 to 29. --- src/callgraph.jl | 29 ++++++++++++++++++++--------- test/twiddles.jl | 9 ++++++--- 2 files changed, 26 insertions(+), 12 deletions(-) diff --git a/src/callgraph.jl b/src/callgraph.jl index add561a..8b1b396 100644 --- a/src/callgraph.jl +++ b/src/callgraph.jl @@ -77,8 +77,13 @@ struct CallGraph{T<:Complex} end # Primes below this use the O(N²) DFT with a twiddle table; at and above it -# Bluestein's algorithm is cheaper (crossover measured at ~45 for ComplexF64). -const DEFAULT_BLUESTEIN_CUTOFF = 47 +# Bluestein's algorithm is used. Measured crossovers (ComplexF64, planned +# execution): x86-64 AVX2 (Core Ultra 7 165H) n ≈ 23 — Bluestein is 1.7–2.2× +# faster than the DFT leaf for n = 41–47; aarch64 NEON (Neoverse-N1) n ≈ 29 +# for the 64-point pad, with the DFT leaf up to 1.45× faster again at n = 37–43 +# where the pad grows to 128, and Bluestein ahead from 47. 29 never loses +# much on either. +const DEFAULT_BLUESTEIN_CUTOFF = 29 # Get the node in the graph at index i Base.getindex(g::CallGraph{T}, i::Int) where {T} = g.nodes[i] @@ -310,12 +315,18 @@ end $(TYPEDSIGNATURES) Length of the padded convolution in Bluestein's algorithm for a length-`N` transform: the smallest power of two ≥ 2N-1, unless a 3-smooth length -`2^a 3^b ≥ 2N-1` is enough smaller to be cheaper. Measured on FFTA, a -transform of a length with factors of 3 costs about 1.9× per `n log n` -compared to a power of two (the composite step and the radix-3 kernel are -slower than the radix-4 one), so `m` is preferred over `p` when -`1.9 m log m < p log p`; below 2048 the fixed overhead of the composite -step dominates and the power of two is always used. +`2^a 3^b ≥ 2N-1` is enough smaller to be cheaper. + +How much cheaper it has to be is machine dependent. Measured per `n log n` +against a power of two (ComplexF64, planned execution), a length with +factors of 3 costs 1.3–2.3× on aarch64 NEON (Neoverse-N1) but 2.0–3.1× on +x86-64 AVX2 (Core Ultra 7 165H); with a threshold of 1.9 the 3-smooth pad +was a 1.58× win at N = 8443 on aarch64 and a 1.19× loss on x86-64. The +threshold is therefore set to 2.1, at which the chooser never loses on +either machine (and rarely fires: it needs a 3-smooth length well under +half the power of two). Below 2048 the fixed overhead of the composite +step dominates and the power of two is always used. The constant should +be revisited when the composite/radix-3 path gets faster. """ function bluestein_pad_length(N::Int) m = 2N - 1 @@ -325,7 +336,7 @@ function bluestein_pad_length(N::Int) pow3 = 3 while pow3 < p c = pow3 * nextpow(2, cld(m, pow3)) # smallest 2^a·3^b with this power of 3 - cost = 1.9 * c * log2(c) + cost = 2.1 * c * log2(c) if c >= max(m, 2048) && cost < best_cost best, best_cost = c, cost end diff --git a/test/twiddles.jl b/test/twiddles.jl index 0318850..4f9baa7 100644 --- a/test/twiddles.jl +++ b/test/twiddles.jl @@ -112,9 +112,12 @@ end end @test FFTA.bluestein_pad_length(73) == 256 # small: always a power of two @test FFTA.bluestein_pad_length(1009) == 2048 - @test FFTA.bluestein_pad_length(4099) in (8748, 9216) # 3-smooth, cheaper than 16384 - @test FFTA.bluestein_pad_length(8443) == 17496 - @test FFTA.bluestein_pad_length(65537) == 139968 + # with the conservative threshold these keep the power of two + @test FFTA.bluestein_pad_length(4099) == 16384 + @test FFTA.bluestein_pad_length(8443) == 32768 + @test FFTA.bluestein_pad_length(65537) == 262144 + # a 3-smooth pad is taken when it is far enough below the power of two + @test FFTA.bluestein_pad_length(2200) == 4374 # 2·3^7 = 4374 vs 8192 end @testset "planned execution does not allocate, n=$n" for n in (5, 47, 64, 73, 101, 720, 1000, 1009, 4096, 65537) From 5b96dff483917899602a662d79fc8f8ecf24c07d Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 20:58:30 +0000 Subject: [PATCH 6/7] Bluestein padding: make the cost factor a keyword and state that 2.1 currently never selects a 3-smooth length --- src/callgraph.jl | 21 ++++++++++++++------- test/twiddles.jl | 6 ++++-- 2 files changed, 18 insertions(+), 9 deletions(-) diff --git a/src/callgraph.jl b/src/callgraph.jl index 8b1b396..3e35e48 100644 --- a/src/callgraph.jl +++ b/src/callgraph.jl @@ -85,6 +85,11 @@ end # much on either. const DEFAULT_BLUESTEIN_CUTOFF = 29 +# Relative per-element cost of a 3-smooth transform length compared to a +# power of two, above which `bluestein_pad_length` keeps the power of two. +# See its docstring for the measurements behind the value. +const BLUESTEIN_SMOOTH_FACTOR = 2.1 + # Get the node in the graph at index i Base.getindex(g::CallGraph{T}, i::Int) where {T} = g.nodes[i] @@ -322,13 +327,15 @@ against a power of two (ComplexF64, planned execution), a length with factors of 3 costs 1.3–2.3× on aarch64 NEON (Neoverse-N1) but 2.0–3.1× on x86-64 AVX2 (Core Ultra 7 165H); with a threshold of 1.9 the 3-smooth pad was a 1.58× win at N = 8443 on aarch64 and a 1.19× loss on x86-64. The -threshold is therefore set to 2.1, at which the chooser never loses on -either machine (and rarely fires: it needs a 3-smooth length well under -half the power of two). Below 2048 the fixed overhead of the composite -step dominates and the power of two is always used. The constant should -be revisited when the composite/radix-3 path gets faster. +threshold (`factor`) is therefore set to 2.1, at which the chooser never +loses on either machine — which, since the smallest admissible length is +always more than half the power of two, means it currently never picks a +3-smooth length: the mechanism is kept (and the constant should be +lowered) for when the composite/radix-3 path becomes competitive with the +radix-4 kernel. Below 2048 the fixed overhead of the composite step +dominates and the power of two is always used. """ -function bluestein_pad_length(N::Int) +function bluestein_pad_length(N::Int; factor::Real = BLUESTEIN_SMOOTH_FACTOR) m = 2N - 1 p = nextpow(2, m) best = p @@ -336,7 +343,7 @@ function bluestein_pad_length(N::Int) pow3 = 3 while pow3 < p c = pow3 * nextpow(2, cld(m, pow3)) # smallest 2^a·3^b with this power of 3 - cost = 2.1 * c * log2(c) + cost = factor * c * log2(c) if c >= max(m, 2048) && cost < best_cost best, best_cost = c, cost end diff --git a/test/twiddles.jl b/test/twiddles.jl index 4f9baa7..7ae82a9 100644 --- a/test/twiddles.jl +++ b/test/twiddles.jl @@ -116,8 +116,10 @@ end @test FFTA.bluestein_pad_length(4099) == 16384 @test FFTA.bluestein_pad_length(8443) == 32768 @test FFTA.bluestein_pad_length(65537) == 262144 - # a 3-smooth pad is taken when it is far enough below the power of two - @test FFTA.bluestein_pad_length(2200) == 4374 # 2·3^7 = 4374 vs 8192 + # the mechanism, with a factor at which 3-smooth lengths are worth it + @test FFTA.bluestein_pad_length(2200; factor = 1.5) == 4374 # 2·3^7 = 4374 vs 8192 + @test FFTA.bluestein_pad_length(8443; factor = 1.5) == 17496 + @test FFTA.bluestein_pad_length(73; factor = 1.0) == 256 # below the 2048 floor end @testset "planned execution does not allocate, n=$n" for n in (5, 47, 64, 73, 101, 720, 1000, 1009, 4096, 65537) From cc6a69bb1aacb5ee1dc36a83ed306c37922fcd36 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 20:58:53 +0000 Subject: [PATCH 7/7] test: correct the expected 3-smooth pad for N = 2200 --- test/twiddles.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/twiddles.jl b/test/twiddles.jl index 7ae82a9..f40db98 100644 --- a/test/twiddles.jl +++ b/test/twiddles.jl @@ -117,7 +117,7 @@ end @test FFTA.bluestein_pad_length(8443) == 32768 @test FFTA.bluestein_pad_length(65537) == 262144 # the mechanism, with a factor at which 3-smooth lengths are worth it - @test FFTA.bluestein_pad_length(2200; factor = 1.5) == 4374 # 2·3^7 = 4374 vs 8192 + @test FFTA.bluestein_pad_length(2200; factor = 1.5) == 4608 # 2^9·3^2 (2·3^7 = 4374 < 2N-1) vs 8192 @test FFTA.bluestein_pad_length(8443; factor = 1.5) == 17496 @test FFTA.bluestein_pad_length(73; factor = 1.0) == 256 # below the 2048 floor end