Radix-5 and radix-7 kernels for powers of 5 and 7

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)

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,

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

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