Leaves-first order for large power-of-two transforms

Project.toml

@@ -12,6 +12,7 @@ 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"
@@ -25,6 +26,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

@@ -6,12 +6,15 @@ using LinearAlgebra: LinearAlgebra
 using MuladdMacro: @muladd
 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("simd_pass.jl")
+include("leaffirst.jl")
 include("algos.jl")
 include("plan.jl")
 

src/algos.jl

@@ -20,7 +20,7 @@ function fft_kernel!(
         if t === DFT
             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)
@@ -185,6 +185,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 +195,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]
@@ -237,6 +246,22 @@ function fft_pow2_radix4!(
     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)
 
+    _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]`.
+"""
+function _pow2_pass!(out::AbstractVector{T}, m::Int, start_out::Int, stride_out::Int, d::Direction,
+                     tw::AbstractVector{T}, toff::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) && return
+
     @inbounds for k in 0:m-1
         wkoe = tw[toff + 3k + 1]
         wkeo = tw[toff + 3k + 2]

src/callgraph.jl

@@ -118,7 +118,7 @@ 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

src/leaffirst.jl

@@ -0,0 +1,98 @@
+# Leaves-first order for large power-of-two transforms.
+#
+# The depth-first radix-4 recursion of `fft_pow2_radix4!` reads each leaf's
+# `CODELET_MAX` inputs at stride `N ÷ CODELET_MAX`: consecutive leaves in
+# recursion order are a quarter of the array apart in the input, so every
+# cache line fetched for a leaf is used for one element and evicted before
+# the leaves that need its neighbours run. Once the array is out of the last
+# cache level that costs the leaves 2–2.5× their in-cache time (measured at
+# 2^20–2^22 elements, ComplexF64) and they become the largest stage of the
+# transform.
+#
+# Above `LEAFFIRST_MIN` the transform is therefore computed as `P = N ÷ B`
+# sub-transforms of size `B` (`LEAFFIRST_BLOCK` or half of it, a level of the
+# recursion; decimated inputs at stride `P`), taken in input
+# order in groups of `G` — one cache line of consecutive inputs — which are
+# gathered into a contiguous buffer and transformed in cache, followed by the
+# `log4(P)` remaining butterfly passes over the whole array. The output is
+# identical to the recursion's (same operations, same order per element).
+# Measured on a Neoverse-N1 (ComplexF64, with the SIMD butterfly pass):
+# 2^20 43 → 26 ms, 2^22 188 → 124 ms.
+
+const LEAFFIRST_MIN = 1 << 18    # transforms with fewer elements keep the recursion
+const LEAFFIRST_BLOCK = 1 << 12  # size of the contiguous sub-transforms
+
+# pencils gathered together: one 64-byte cache line of consecutive inputs
+_leaffirst_group(::Type{T}) where {T} = max(4, 64 ÷ sizeof(T))
+
+"""
+$(TYPEDSIGNATURES)
+Length of the gather buffer a `POW2RADIX4_FFT` node of size `N` keeps in its
+workspace: `0` below `LEAFFIRST_MIN`.
+"""
+leaffirst_buflen(::Type{T}, N::Int) where {T} =
+    N >= LEAFFIRST_MIN ? _leaffirst_group(T) * LEAFFIRST_BLOCK : 0
+
+# base-4 digit reversal of `q` over `digits` digits
+@inline function _rev4(q::Int, digits::Int)
+    r = 0
+    for _ in 1:digits
+        r = 4r + (q & 3)
+        q >>= 2
+    end
+    return r
+end
+
+function _pow2_leaffirst!(
+    out::AbstractVector{T}, in::AbstractVector{U},
+    N::Int, start_out::Int, start_in::Int, stride_in::Int,
+    d::Direction, tw::AbstractVector{T}, toff::Int, buf::AbstractVector{T}
+) where {T<:Complex, U}
+    # the sub-transform size is the recursion's own block size at the level
+    # nearest LEAFFIRST_BLOCK (LEAFFIRST_BLOCK or half of it, depending on the
+    # parity of log2 N), so that P = N ÷ B is a power of 4
+    G = _leaffirst_group(T)
+    B = N
+    toffB = toff
+    while B > LEAFFIRST_BLOCK
+        toffB += 3 * (B ÷ 4)
+        B ÷= 4
+    end
+    P = N ÷ B
+    digits = trailing_zeros(P) ÷ 2
+    # (P ≥ G is guaranteed by LEAFFIRST_MIN ≥ 4·G·LEAFFIRST_BLOCK)
+    # 1. sub-transforms of size B, in input order, G at a time through `buf`
+    for q0 in 0:G:P-1
+        @inbounds for j in 0:B-1
+            src = start_in + (q0 + j * P) * stride_in
+            for r in 0:G-1
+                buf[r * B + j + 1] = in[src + r * stride_in]
+            end
+        end
+        for r in 0:G-1
+            fft_pow2_radix4!(out, buf, B, start_out + B * _rev4(q0 + r, digits), 1, 1 + r * B, 1, d, tw, toffB)
+        end
+    end
+    # 2. the remaining butterfly passes, one level at a time (`_pow2_level!`
+    #    descends from the top level's table offset)
+    M = 4B
+    while M <= N
+        _pow2_level!(out, N, M, start_out, d, tw, toff)
+        M *= 4
+    end
+    return nothing
+end
+
+# all radix-4 passes of the level whose blocks have size `L`, inside the
+# block of size `N` at `start_out` (unit stride)
+function _pow2_level!(out::AbstractVector{T}, N::Int, L::Int, start_out::Int, d::Direction,
+                      tw::AbstractVector{T}, toff::Int) where {T}
+    if N == L
+        _pow2_pass!(out, N ÷ 4, start_out, 1, d, tw, toff)
+        return
+    end
+    m = N ÷ 4
+    for q in 0:3
+        _pow2_level!(out, m, L, start_out + q * m, d, tw, toff + 3m)
+    end
+end

src/simd_pass.jl

@@ -0,0 +1,95 @@
+# SIMD radix-4 butterfly pass for the power-of-two kernel.
+#
+# `fft_pow2_radix4!` combines the four quarter transforms of a block with one
+# pass of radix-4 butterflies. Written on scalar `Complex` values, that pass
+# is compute bound on the complex multiplications (LLVM does not vectorise
+# across butterflies). Here `W` butterflies are processed per iteration on
+# vectors of `2W` reals (`W = 2` for `Float64`, `4` for `Float32`, one or two
+# NEON/SSE registers): a complex product `a·w` becomes
+# `a * (wr, wr) + swap(a) * (-wi, wi)`. The twiddle table keeps its compact
+# `(w^k, w^2k, w^3k)` layout (see `pow2_twiddles`); the `W` triplets an
+# iteration needs are loaded as three vectors and rearranged in registers,
+# which measured as fast as an expanded table in cache and faster out of it.
+#
+# Used when the output block is contiguous in memory (`stride_out == 1` on a
+# dense vector or contiguous view) and `N ÷ 4 >= W`; otherwise the scalar loop
+# in `fft_pow2_radix4!` runs. Results agree with the scalar loop to rounding.
+
+_simd_width(::Type{ComplexF64}) = 2
+_simd_width(::Type{ComplexF32}) = 4
+
+# (ai, ar) from (ar, ai) for every complex lane
+@inline _swap(v::Vec{L}) where {L} = shufflevector(v, Val(ntuple(i -> isodd(i) ? i : i - 2, L)))
+# a * w with wr = (re w, re w, ...) and wi = (-im w, im w, ...)
+@inline _cmul(a, wr, wi) = muladd(_swap(a), wi, a * wr)
+
+# The three twiddle vectors of one group of `W` butterflies: the compact table
+# holds `w1 w2 w3` for each `k`, i.e. `6W` reals `(r1 i1 r2 i2 r3 i3)_k` per
+# group, loaded as `u, v, w`. Returns `(wr1, wi1, wr2, wi2, wr3, wi3)` with
+# `wi` already carrying the `(-, +)` sign pattern.
+@inline function _twiddle_vectors(u::Vec{4,Float64}, v::Vec{4,Float64}, w::Vec{4,Float64}, sign)
+    # u = (r1 i1 r2 i2)  v = (r3 i3 r1' i1')  w = (r2' i2' r3' i3')
+    wr1 = shufflevector(u, v, Val((0, 0, 6, 6))); wi1 = shufflevector(u, v, Val((1, 1, 7, 7))) * sign
+    wr2 = shufflevector(u, w, Val((2, 2, 4, 4))); wi2 = shufflevector(u, w, Val((3, 3, 5, 5))) * sign
+    wr3 = shufflevector(v, w, Val((0, 0, 6, 6))); wi3 = shufflevector(v, w, Val((1, 1, 7, 7))) * sign
+    return wr1, wi1, wr2, wi2, wr3, wi3
+end
+@inline function _twiddle_vectors(u::Vec{8,Float32}, v::Vec{8,Float32}, w::Vec{8,Float32}, sign)
+    # u = (r1 i1 r2 i2 r3 i3 r1' i1')  v = (r2' i2' r3' i3' r1'' i1'' r2'' i2'')  w = (r3'' i3'' r1''' i1''' r2''' i2''' r3''' i3''')
+    w1 = shufflevector(shufflevector(u, v, Val((0, 1, 6, 7, 12, 13, 12, 13))), w, Val((0, 1, 2, 3, 4, 5, 10, 11)))
+    w2 = shufflevector(shufflevector(u, v, Val((2, 3, 8, 9, 14, 15, 14, 15))), w, Val((0, 1, 2, 3, 4, 5, 12, 13)))
+    w3 = shufflevector(shufflevector(u, v, Val((4, 5, 10, 11, 4, 5, 10, 11))), w, Val((0, 1, 2, 3, 8, 9, 14, 15)))
+    dupr(x) = shufflevector(x, Val((0, 0, 2, 2, 4, 4, 6, 6)))
+    dupi(x) = shufflevector(x, Val((1, 1, 3, 3, 5, 5, 7, 7)))
+    return dupr(w1), dupi(w1) * sign, dupr(w2), dupi(w2) * sign, dupr(w3), dupi(w3) * sign
+end
+
+# unit-stride dense storage that `pointer` can address
+_simd_contiguous(out::StridedVector) = stride(out, 1) == 1
+_simd_contiguous(out) = false
+
+"""
+$(TYPEDSIGNATURES)
+The radix-4 butterfly pass of `fft_pow2_radix4!` over the `4m` outputs starting
+at `out[start_out]` (unit stride), `W` butterflies per iteration. Returns
+`false` without touching `out` when the pass cannot be vectorised (strided or
+non-contiguous output, or `m < W`), in which case the caller runs the scalar
+loop.
+"""
+@inline function _pow2_pass_simd!(
+    out::AbstractVector{T}, m::Int, start_out::Int, stride_out::Int, d::Direction,
+    tw::AbstractVector{T}, toff::Int
+) where {T<:CodeletEltype}
+    W = _simd_width(T)
+    (stride_out == 1 && m >= W && _simd_contiguous(out) && tw isa Vector{T}) || return false
+    R = real(T)
+    L = 2W
+    V = Vec{L,R}
+    sz = sizeof(R)
+    # (-, +) pattern for the imaginary parts of the twiddles; the ∓i rotation
+    # of the last butterfly leg uses the opposite pattern in the forward
+    # direction and the same one backward
+    wsign = Vec{L,R}(ntuple(i -> isodd(i) ? -one(R) : one(R), L))
+    esign = d === FFT_FORWARD ? -wsign : wsign
+    po = Ptr{R}(pointer(out)) + (start_out - 1) * 2sz
+    pt = Ptr{R}(pointer(tw)) + toff * 2sz
+    GC.@preserve out tw begin
+        @inbounds for k in 0:W:m-1
+            p0 = po + k * 2sz
+            p1 = p0 + m * 2sz
+            p2 = p0 + 2m * 2sz
+            p3 = p0 + 3m * 2sz
+            y0 = vload(V, p0); y1 = vload(V, p1); y2 = vload(V, p2); y3 = vload(V, p3)
+            tb = pt + 3k * 2sz
+            wr1, wi1, wr2, wi2, wr3, wi3 = _twiddle_vectors(vload(V, tb), vload(V, tb + L * sz), vload(V, tb + 2L * sz), wsign)
+            t1 = _cmul(y1, wr1, wi1)
+            t2 = _cmul(y2, wr2, wi2)
+            t3 = _cmul(y3, wr3, wi3)
+            a = y0 + t2; b = y0 - t2
+            c = t1 + t3; e = _swap(t1 - t3) * esign
+            vstore(a + c, p0); vstore(b + e, p1); vstore(a - c, p2); vstore(b - e, p3)
+        end
+    end
+    return true
+end
+_pow2_pass_simd!(out, m, start_out, stride_out, d, tw, toff) = false

test/leaffirst.jl

@@ -0,0 +1,36 @@
+# Large power-of-two transforms take the leaves-first path (src/leaffirst.jl).
+# Checked without FFTW (loading it here would capture `plan_*` dispatch for the
+# rest of the suite): a length-n transform is rebuilt from the two half-length
+# transforms of its even and odd elements (which, for the smallest n here, take
+# the plain recursion), real transforms against the complex one, and the
+# backward transform against the identity.
+using FFTA, LinearAlgebra, Test
+
+function rebuilt_fft(x::AbstractVector{Complex{T}}) where {T}
+    n = length(x)
+    E = fft(x[1:2:end]); O = fft(x[2:2:end])
+    m = n ÷ 2
+    X = similar(x)
+    for k in 0:n-1
+        w = Complex{T}(cispi(-2 * T(k) / T(n)))
+        X[k + 1] = E[k % m + 1] + w * O[k % m + 1]
+    end
+    return X
+end
+
+@testset "leaves-first order, n = 2^$k, $T" for k in 18:21, T in (Float64, Float32)
+    n = 1 << k
+    rtol = T === Float64 ? 1e-9 : 1e-3
+    x = randn(Complex{T}, n)
+    p = plan_fft(x)
+    y = p * x
+    @test y ≈ rebuilt_fft(x) rtol = rtol
+    @test bfft(y) ≈ n .* x rtol = rtol
+    @test (@allocated mul!(y, p, x)) == 0
+    xr = randn(T, n)
+    pr = plan_rfft(xr)
+    yr = pr * xr
+    @test yr ≈ fft(complex(xr))[1:n÷2+1] rtol = rtol
+    @test brfft(yr, n) ≈ n .* xr rtol = rtol
+    @test (@allocated mul!(yr, pr, xr)) == 0
+end

test/runtests.jl

@@ -51,6 +51,9 @@ Random.seed!(1)
     @testset verbose = true "Twiddle tables" begin
         include("twiddles.jl")
     end
+    @testset "Leaves-first order (large powers of two)" begin
+        include("leaffirst.jl")
+    end
     @testset verbose = true "Argument checking" begin
         include("argument_checking.jl")
     end