Threaded leaves-first path for single large power-of-two transforms

Project.toml

@@ -8,10 +8,12 @@ AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c"
 DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221"
+Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
 Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae"
 Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 Reexport = "189a3867-3050-52da-a836-e630ba90ab69"
+SIMD = "fdea26ae-647d-5447-a871-4b548cad5224"
 
 [compat]
 AbstractFFTs = "1"
@@ -21,10 +23,12 @@ ExplicitImports = "1.12"
 FFTW = "1.8"
 LinearAlgebra = "<0.0.1, 1"
 MuladdMacro = "0.2"
+Polyester = "0.7"
 PrecompileTools = "1"
 Primes = "0.5"
 Random = "<0.0.1, 1"
 Reexport = "1"
+SIMD = "3"
 Test = "<0.0.1, 1"
 julia = "1.6.7"
 

src/FFTA.jl

@@ -4,15 +4,21 @@ using AbstractFFTs: AbstractFFTs
 using DocStringExtensions: TYPEDEF, TYPEDSIGNATURES
 using LinearAlgebra: LinearAlgebra
 using MuladdMacro: @muladd
+using Polyester: @batch
 using Primes: Primes
 using Reexport: @reexport
+using SIMD: Vec, vload, vstore, shufflevector
 
 @reexport using AbstractFFTs
 
 include("callgraph.jl")
 include("singleton_twiddle.jl")
 include("codelets.jl")
+include("odd_codelets.jl")
+include("simd_pass.jl")
+include("leaffirst.jl")
 include("algos.jl")
+include("real_simd.jl")
 include("plan.jl")
 
 # Compile the codelets (and the common plan/execute paths) at precompile time
@@ -22,7 +28,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)
+            fft_pow2_radix4!(out, in, N, start_out, s_out, start_in, s_in, d, tw, 0, g.workspace[idx])
         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
@@ -185,6 +191,8 @@ Radix-4 FFT for powers of 2, in place
 - `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`
+- `buf`: Gather buffer for the leaves-first order of large transforms (see
+  `leaffirst_buflen`; `nothing` or empty to use the plain recursion)
 
 """
 function fft_pow2_radix4!(
@@ -193,8 +201,15 @@ function fft_pow2_radix4!(
     start_out::Int, stride_out::Int,
     start_in::Int, stride_in::Int,
     d::Direction,
-    tw::AbstractVector{T}, toff::Int
+    tw::AbstractVector{T}, toff::Int,
+    buf::Union{Nothing,AbstractVector{T}} = nothing
 ) where {T<:Complex, U}
+    # Large transforms: leaves first, gathered through `buf` (see leaffirst.jl)
+    if buf !== nothing && !isempty(buf) && N >= LEAFFIRST_MIN && stride_out == 1
+        _pow2_leaffirst!(out, in, N, start_out, start_in, stride_in, d, tw, toff, buf)
+        return
+    end
+
     # If N is 2, compute the size two DFT
     @inbounds if N == 2
         out[start_out]              = in[start_in] + in[start_in + stride_in]
@@ -232,12 +247,39 @@ function fft_pow2_radix4!(
     m = N ÷ 4
     toff_next = toff + 3m   # the next level's table follows this level's
 
+    # four codelet-sized children: computed in lockstep on SIMD vectors
+    if _pow2_lockstep!(out, in, N, start_out, stride_out, start_in, stride_in, d)
+        _pow2_pass!(out, m, start_out, stride_out, d, tw, toff)
+        return
+    end
+
     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
+    _pow2_pass!(out, m, start_out, stride_out, d, tw, toff)
+end
+
+"""
+$(TYPEDSIGNATURES)
+One radix-4 butterfly pass combining the four quarter transforms of size `m`
+stored at `out[start_out + k*stride_out]`, `k = 0..4m-1`, with the twiddles of
+this level at `tw[toff+1:toff+3m]`.
+"""
+_pow2_pass!(out::AbstractVector{T}, m::Int, start_out::Int, stride_out::Int, d::Direction,
+            tw::AbstractVector{T}, toff::Int) where {T} =
+    _pow2_pass!(out, m, start_out, stride_out, d, tw, toff, 0, m)
+
+# the butterflies `k = k0..k1-1` of the pass (a chunk, for threading)
+function _pow2_pass!(out::AbstractVector{T}, m::Int, start_out::Int, stride_out::Int, d::Direction,
+                     tw::AbstractVector{T}, toff::Int, k0::Int, k1::Int) where {T}
+    dir = direction_sign(d)
+    minusi = -dir * im
+    # vectorised butterfly pass for the floating-point types (see simd_pass.jl)
+    _pow2_pass_simd!(out, m, start_out, stride_out, d, tw, toff, k0, k1) && return
+
+    @inbounds for k in k0:k1-1
         wkoe = tw[toff + 3k + 1]
         wkeo = tw[toff + 3k + 2]
         wkoo = tw[toff + 3k + 3]
@@ -327,6 +369,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)
@@ -118,13 +118,22 @@ function CallGraphNode!(
         throw(DimensionMismatch("Array length must be strictly positive"))
     end
     if iseven(N) && ispow2(N)
-        push!(workspace, T[])
+        push!(workspace, Vector{T}(undef, leaffirst_buflen(T, N)))
         push!(nodes, CallGraphNode(0, 0, POW2RADIX4_FFT, N, s_in, s_out))
         return 1
     elseif N % 3 == 0 && nextpow(3, N) == N
         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