diff --git a/Project.toml b/Project.toml index a846f4a..b22277c 100644 --- a/Project.toml +++ b/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" diff --git a/src/FFTA.jl b/src/FFTA.jl index 3309f17..384b7e5 100644 --- a/src/FFTA.jl +++ b/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) diff --git a/src/algos.jl b/src/algos.jl index 9dae269..6dbadf3 100644 --- a/src/algos.jl +++ b/src/algos.jl @@ -18,13 +18,19 @@ function fft_kernel!( N = root.sz tw = g.twiddles[idx] if t === DFT - fft_dft!(out, in, N, start_out, s_out, start_in, s_in, tw) + # small odd primes have straight-line codelets (see odd_codelets.jl) + _odd_codelet!(out, in, N, start_out, s_out, start_in, s_in, d) || + fft_dft!(out, in, N, start_out, s_out, start_in, s_in, tw) elseif t === POW2RADIX4_FFT - fft_pow2_radix4!(out, in, N, start_out, s_out, start_in, s_in, d, tw, 0) + 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, diff --git a/src/callgraph.jl b/src/callgraph.jl index ab119bb..c70cf41 100644 --- a/src/callgraph.jl +++ b/src/callgraph.jl @@ -1,6 +1,6 @@ @enum Direction FFT_FORWARD=-1 FFT_BACKWARD=1 @enum Pow24 POW2 POW4 -@enum FFTEnum COMPOSITE_FFT DFT POW3_FFT POW2RADIX4_FFT BLUESTEIN +@enum FFTEnum COMPOSITE_FFT DFT POW3_FFT POW2RADIX4_FFT BLUESTEIN POW5_FFT POW7_FFT @inline function direction_sign(d::Direction) Int(d) @@ -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 diff --git a/src/codelets.jl b/src/codelets.jl index 5b22de2..9bbf9eb 100644 --- a/src/codelets.jl +++ b/src/codelets.jl @@ -14,6 +14,8 @@ const CODELET_MAX = 64 const CodeletEltype = Union{ComplexF32,ComplexF64} +_simd_width(::Type{ComplexF64}) = 2 # complex values per SIMD vector +_simd_width(::Type{ComplexF32}) = 4 # Emit the statements computing the length-`length(xs)` DFT of the symbols # `xs` (radix-2 DIT), returning the symbols holding the outputs. @@ -98,3 +100,115 @@ end return false end _pow2_codelet!(out, in, N, start_out, stride_out, start_in, stride_in, d) = false + +# --------------------------------------------------------------------------- +# Lockstep codelets: `W` leaves at once on SIMD vectors +# +# At the level of the recursion whose four children are codelets (blocks of +# 128 or 256 points), the four sibling leaves are independent transforms of +# the same size with the same strides. They are computed together, each leaf +# in its own complex lane of a `Vec{2W}` (`W = 2` for `Float64`, `4` for +# `Float32`): the same radix-2 DIT statements as `fft_pow2_codelet!`, on +# vectors instead of scalars, with the twiddle products written so that each +# lane rounds exactly as the scalar codelet does (`muladd(wr, o, swap(o)·wi)` +# is `fma(wr, re, -wi·im)` / `fma(wr, im, wi·re)` on FMA hardware). Inputs +# are gathered lane by lane (the leaves' inputs are interleaved in memory) +# and outputs scattered; the arithmetic in between is `W`× narrower. +# Measured on Neoverse-N1 for 64 leaves of 64 points: 1.35× (`Float64`) +# and 2.4× (`Float32`) over the scalar codelets. +# --------------------------------------------------------------------------- + +@inline _vswap(v::Vec{L}) where {L} = shufflevector(v, Val(ntuple(i -> isodd(i) ? i : i - 2, L))) + +function _gen_dit_vec!(stmts::Vector{Any}, xs::Vector{Symbol}, ::Type{T}, L::Int, dir::Int, counter::Ref{Int}) where {T} + N = length(xs) + N == 1 && return xs + newsym() = Symbol(:t, counter[] += 1) + E = _gen_dit_vec!(stmts, xs[1:2:end], T, L, dir, counter) + O = _gen_dit_vec!(stmts, xs[2:2:end], T, L, dir, counter) + outs = Vector{Symbol}(undef, N) + for k in 0:N÷2-1 + o = O[k + 1] + t = newsym() + if k == 0 + push!(stmts, :($t = $o)) + elseif 4k == N # multiply by ∓i + push!(stmts, :($t = _vswap($o) * $(dir < 0 ? :sign_pm : :sign_mp))) + else + w = cispi(dir * 2 * k / N) + wr, wi = T(real(w)), T(imag(w)) + wiv = ntuple(i -> isodd(i) ? -wi : wi, L) + push!(stmts, :($t = muladd(Vec{$L,$T}($wr), $o, _vswap($o) * Vec{$L,$T}($wiv)))) + end + a, b = newsym(), newsym() + push!(stmts, :($a = $(E[k + 1]) + $t), :($b = $(E[k + 1]) - $t)) + outs[k + 1] = a + outs[k + 1 + N÷2] = b + end + return outs +end + +""" +$(TYPEDSIGNATURES) +`W` length-`N` transforms in lockstep: leaf `q = 0..W-1` reads +`in[start_in + q*stride_in + j*4stride_in]` and writes +`out[start_out + q*N*stride_out + j*stride_out]`, `j = 0..N-1` — the strides +of four sibling leaves of the radix-4 recursion. Direction `D` as in +`fft_pow2_codelet!`. +""" +@generated function fft_pow2_codelet_lockstep!( + out::AbstractVector{Complex{T}}, in::AbstractVector{Complex{T}}, + ::Val{N}, ::Val{W}, + start_out::Int, stride_out::Int, + start_in::Int, stride_in::Int, + ::Val{D} +) where {T,N,W,D} + L = 2W + counter = Ref(0) + stmts = Any[] + xs = [Symbol(:x, i) for i in 1:N] + for i in 1:N + parts = Any[] + for q in 0:W-1 + c = Symbol(:c, i, :_, q) + push!(stmts, :($c = in[start_in + $q * stride_in + $(4(i - 1)) * stride_in])) + push!(parts, :(real($c)), :(imag($c))) + end + push!(stmts, :($(xs[i]) = Vec{$L,$T}($(Expr(:tuple, parts...))))) + end + outs = _gen_dit_vec!(stmts, xs, T, L, D, counter) + for i in 1:N, q in 0:W-1 + push!(stmts, :(out[start_out + $(q * N) * stride_out + $(i - 1) * stride_out] = Complex{$T}($(outs[i])[$(2q + 1)], $(outs[i])[$(2q + 2)]))) + end + body = Expr(:block, stmts...) + return quote + sign_pm = Vec{$L,$T}($(ntuple(i -> isodd(i) ? one(T) : -one(T), L))) + sign_mp = -sign_pm + @inbounds $body + return nothing + end +end + +# The four codelet-sized children of a block of `N` points (`N` = 4 × 32 or +# 4 × 64) in lockstep; `false` when not applicable (the caller recurses). +@inline function _pow2_lockstep!( + out::AbstractVector{T}, in::AbstractVector{T}, N::Int, + start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction +) where {T<:CodeletEltype} + m = N >> 2 + (m == 32 || m == 64) || return false + vd = d === FFT_FORWARD ? Val(-1) : Val(1) + if T === ComplexF32 + m == 64 ? fft_pow2_codelet_lockstep!(out, in, Val(64), Val(4), start_out, stride_out, start_in, stride_in, vd) : + fft_pow2_codelet_lockstep!(out, in, Val(32), Val(4), start_out, stride_out, start_in, stride_in, vd) + else + for r0 in (0, 2) + so = start_out + r0 * m * stride_out + si = start_in + r0 * stride_in + m == 64 ? fft_pow2_codelet_lockstep!(out, in, Val(64), Val(2), so, stride_out, si, stride_in, vd) : + fft_pow2_codelet_lockstep!(out, in, Val(32), Val(2), so, stride_out, si, stride_in, vd) + end + end + return true +end +_pow2_lockstep!(out, in, N, start_out, stride_out, start_in, stride_in, d) = false diff --git a/src/leaffirst.jl b/src/leaffirst.jl new file mode 100644 index 0000000..e68644d --- /dev/null +++ b/src/leaffirst.jl @@ -0,0 +1,210 @@ +# 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 + + +# --------------------------------------------------------------------------- +# Threaded leaves-first order for a single large transform +# +# The two stages above are embarrassingly parallel: the sub-transforms are +# independent (each chunk of groups uses its own gather buffer, one per +# worker), and a butterfly pass is independent across blocks and, within a +# block, across butterflies. Passes over few large blocks are split by +# butterfly range in multiples of 64 (SIMD-friendly), so the operations and +# their order per element are the same as in the serial code: results do not +# depend on the number of threads. +# --------------------------------------------------------------------------- + +""" +$(TYPEDSIGNATURES) +`_pow2_leaffirst!` on `nt = length(gathers)` threads (Polyester), `gathers` +being one gather buffer per worker (see `Worker.gathers`). +""" +function _pow2_leaffirst_threaded!( + out::AbstractVector{T}, in::AbstractVector{U}, + N::Int, start_out::Int, start_in::Int, stride_in::Int, + d::Direction, tw::AbstractVector{T}, toff::Int, gathers::Vector{Vector{T}} +) where {T<:Complex, U} + nt = length(gathers) + G = _leaffirst_group(T) + B = N + toffB = toff + while B > LEAFFIRST_BLOCK + toffB += 3 * (B ÷ 4) + B ÷= 4 + end + P = N ÷ B + _lf_subtransforms!(out, in, B, P, G, start_out, start_in, stride_in, d, tw, toffB, gathers) + M = 4B + while M <= N + toffM = toff + L = N + while L > M + toffM += 3 * (L ÷ 4) + L ÷= 4 + end + _lf_level!(out, N, M, start_out, d, tw, toffM, nt) + M *= 4 + end + return nothing +end + +# (chunks are large here — a whole transform of ≥ 2^18 elements split nt +# ways — so plain tasks are used; Polyester's `@batch` cannot pass a vector +# of buffers to its threads on every platform) + +# stage 1: the P sub-transforms of size B in groups of G, one chunk of groups per thread +function _lf_subtransforms!(out::AbstractVector{T}, in::AbstractVector, B::Int, P::Int, G::Int, + start_out::Int, start_in::Int, stride_in::Int, d::Direction, + tw::AbstractVector{T}, toffB::Int, gathers::Vector{Vector{T}}) where {T} + nt = length(gathers) + digits = trailing_zeros(P) ÷ 2 + ngroups = P ÷ G + Base.@sync for c in 1:nt + buf = gathers[c] + g0 = (c - 1) * ngroups ÷ nt + g1 = c * ngroups ÷ nt - 1 + Threads.@spawn for g in g0:g1 + q0 = g * G + @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 + end + return nothing +end + +# stage 2: all passes of the level whose blocks have size M +function _lf_level!(out::AbstractVector{T}, N::Int, M::Int, start_out::Int, d::Direction, + tw::AbstractVector{T}, toffM::Int, nt::Int) where {T} + m = M ÷ 4 + nblocks = N ÷ M + if nblocks >= nt + # whole blocks per thread + Base.@sync for c in 1:nt + b0 = (c - 1) * nblocks ÷ nt + b1 = c * nblocks ÷ nt - 1 + Threads.@spawn for b in b0:b1 + _pow2_pass!(out, m, start_out + b * M, 1, d, tw, toffM) + end + end + else + # few blocks: split the butterflies of every block, in multiples of 64 + Base.@sync for c in 1:nt + k0 = ((c - 1) * m ÷ nt) ÷ 64 * 64 + k1 = c == nt ? m : (c * m ÷ nt) ÷ 64 * 64 + k1 > k0 || continue + Threads.@spawn for b in 0:nblocks-1 + _pow2_pass!(out, m, start_out + b * M, 1, d, tw, toffM, k0, k1) + end + end + end + return nothing +end + +# Whether the root transform of `g` (size `N`) can take the threaded +# leaves-first path with the given gather buffers. +_threaded_1d_ok(g::CallGraph, gathers::Vector) = + length(gathers) > 1 && g[1].type === POW2RADIX4_FFT && g[1].sz >= LEAFFIRST_MIN && + !isempty(gathers[1]) && g[1].sz ÷ LEAFFIRST_BLOCK >= 4 * length(gathers) ÷ 4 + 1 diff --git a/src/odd_codelets.jl b/src/odd_codelets.jl new file mode 100644 index 0000000..cb45755 --- /dev/null +++ b/src/odd_codelets.jl @@ -0,0 +1,122 @@ +# Straight-line codelets for the small odd prime leaves (5, 7, 11, 13). +# +# The generic `fft_dft!` leaf is an O(N²) loop over a twiddle table: for a +# 5-point transform that is 16 complex multiplications from memory-loaded +# twiddles. These codelets use the symmetry of the odd-length DFT instead: +# with `a_j = x_j + x_{N-j}` and `b_j = x_j - x_{N-j}` for `j = 1..(N-1)/2`, +# +# X_k = x_0 + Σ_j a_j cos(2πjk/N) + D·i Σ_j b_j sin(2πjk/N) +# X_{N-k} = x_0 + Σ_j a_j cos(2πjk/N) - D·i Σ_j b_j sin(2πjk/N) +# +# so each output pair costs `(N-1)/2` real-by-complex products for the cosine +# sum and as many for the sine sum: `(N-1)²` real multiplications per +# transform instead of `4(N-1)²`, with every constant folded and no table +# loads. Like the power-of-two codelets they are compiled once per +# `(N, element type, direction)` (see the `PrecompileTools` workload) and are +# only used for `Float32`/`Float64` elements; other types keep `fft_dft!`. +# The input may be real (the odd-length real transform runs the complex +# kernel on real input), in which case the sums are real and the same +# expressions are emitted. + +const ODD_CODELET_SIZES = (5, 7, 11, 13) + +# real coefficient × (real or complex) value + (real or complex) accumulator +@inline _fma(a::T, b::T, c::T) where {T<:Real} = fma(a, b, c) +@inline _fma(a::T, b::Complex{T}, c::Complex{T}) where {T<:Real} = Complex(fma(a, real(b), real(c)), fma(a, imag(b), imag(c))) +@inline _fma(a::T, b::T, c::Complex{T}) where {T<:Real} = Complex(fma(a, b, real(c)), imag(c)) +@inline _fma(a::T, b::Complex{T}, c::T) where {T<:Real} = Complex(fma(a, real(b), c), a * imag(b)) + +# Emit the statements of the length-`N` DFT of the symbols `xs` in direction +# `dir` (`-1` forward), returning the output symbols. +function _gen_odd!(stmts::Vector{Any}, xs::Vector{Symbol}, ::Type{T}, dir::Int, counter::Ref{Int}) where {T} + N = length(xs) + h = (N - 1) ÷ 2 + newsym() = Symbol(:t, counter[] += 1) + as = Vector{Symbol}(undef, h); bs = Vector{Symbol}(undef, h) + for j in 1:h + as[j] = newsym(); bs[j] = newsym() + push!(stmts, :($(as[j]) = $(xs[j + 1]) + $(xs[N - j + 1])), :($(bs[j]) = $(xs[j + 1]) - $(xs[N - j + 1]))) + end + outs = Vector{Symbol}(undef, N) + # X_0 + acc = xs[1] + for j in 1:h + t = newsym(); push!(stmts, :($t = $acc + $(as[j]))); acc = t + end + outs[1] = acc + for k in 1:h + # cosine sum (real coefficients), starting from x_0; explicit fma (not + # muladd) so that rounding does not depend on whether LLVM contracts + # for a particular array type + c = xs[1] + for j in 1:h + coef = T(cospi(2 * (j * k % N) / N)) + t = newsym(); push!(stmts, :($t = _fma($coef, $(as[j]), $c))); c = t + end + # sine sum, with the direction folded into the coefficients + s = nothing + for j in 1:h + coef = T(dir * sinpi(2 * (j * k % N) / N)) + t = newsym() + push!(stmts, s === nothing ? :($t = $coef * $(bs[j])) : :($t = _fma($coef, $(bs[j]), $s))) + s = t + end + # ± i·s + is = newsym(); push!(stmts, :($is = Complex(-imag($s), real($s)))) + o1 = newsym(); o2 = newsym() + push!(stmts, :($o1 = $c + $is), :($o2 = $c - $is)) + outs[k + 1] = o1; outs[N - k + 1] = o2 + end + return outs +end + +""" +$(TYPEDSIGNATURES) +Straight-line length-`N` DFT (odd `N`) of `in[start_in + k*stride_in]` into +`out[start_out + k*stride_out]`, `k = 0..N-1`, in direction `D` (`-1` +forward, `+1` backward). The input may be real or complex. +""" +@generated function fft_odd_codelet!( + out::AbstractVector{Complex{T}}, in::AbstractVector{<:Union{T,Complex{T}}}, + ::Val{N}, + start_out::Int, stride_out::Int, + start_in::Int, stride_in::Int, + ::Val{D} +) where {T,N,D} + counter = Ref(0) + stmts = Any[] + xs = [Symbol(:x, i) for i in 1:N] + for i in 1:N + push!(stmts, :($(xs[i]) = in[start_in + $(i - 1) * stride_in])) + end + outs = _gen_odd!(stmts, xs, T, D, counter) + for i in 1:N + push!(stmts, :(out[start_out + $(i - 1) * stride_out] = $(outs[i]))) + end + body = Expr(:block, stmts...) + return quote + @inbounds $body + return nothing + end +end + +# Dispatch an odd leaf to its codelet; `false` when there is none for this +# `N` or element type (the caller then uses `fft_dft!`). +@inline function _odd_codelet!( + out::AbstractVector{Complex{S}}, in::AbstractVector{<:Union{S,Complex{S}}}, N::Int, + start_out::Int, stride_out::Int, start_in::Int, stride_in::Int, d::Direction +) where {S<:Union{Float32,Float64}} + if d === FFT_FORWARD + N == 5 && (fft_odd_codelet!(out, in, Val(5), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + N == 7 && (fft_odd_codelet!(out, in, Val(7), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + N == 11 && (fft_odd_codelet!(out, in, Val(11), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + N == 13 && (fft_odd_codelet!(out, in, Val(13), start_out, stride_out, start_in, stride_in, Val(-1)); return true) + else + N == 5 && (fft_odd_codelet!(out, in, Val(5), start_out, stride_out, start_in, stride_in, Val(1)); return true) + N == 7 && (fft_odd_codelet!(out, in, Val(7), start_out, stride_out, start_in, stride_in, Val(1)); return true) + N == 11 && (fft_odd_codelet!(out, in, Val(11), start_out, stride_out, start_in, stride_in, Val(1)); return true) + N == 13 && (fft_odd_codelet!(out, in, Val(13), start_out, stride_out, start_in, stride_in, Val(1)); return true) + end + return false +end +_odd_codelet!(out, in, N, start_out, stride_out, start_in, stride_in, d) = false diff --git a/src/plan.jl b/src/plan.jl index e519246..ae29f49 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -21,6 +21,20 @@ struct Worker{T,N,S<:Real} buf::Vector{T} # real<->complex packing scratch, see `_re_buflen` (real plans only) rbuf::Vector{S} # contiguous copy of a strided real pencil, see `_re_pencil_loop!` (real plans only) cbuf::Vector{T} # contiguous copy of a strided complex pencil (real plans only) + rtw::Vector{T} # twiddles of the even-length real pre/post-processing, see `_re_twiddles` (shared) + gathers::Vector{Vector{T}} # every worker's root-node workspace: gather buffers for the threaded leaves-first path +end + +""" +$(TYPEDSIGNATURES) +`w_n^j = exp(-2πi j/n)` for `j = 1..n÷4`, the twiddles of the even-length +real-to-complex post-processing in `_rfft_pencil!` (conjugated in +`_brfft_pencil!`); empty for odd `n`. Computed once per plan with `cispi`. +""" +function _re_twiddles(::Type{T}, n::Int) where {T<:Complex} + (n == 0 || isodd(n)) && return T[] + R = real(T) + return T[T(cispi(-2 * R(j) / R(n))) for j in 1:n÷4] end # A call graph that shares the nodes and twiddle tables but has its own @@ -38,12 +52,14 @@ function _workers(cg::Tuple{CallGraph{T},Vararg{CallGraph{T}}}, ::Val{N}, num_th S = real(T) obuflen = maximum(g -> first(g.nodes).sz, cg) clen = rlen == 0 ? 0 : rlen ÷ 2 + 1 - mk(graphs) = Worker{T,N,S}(graphs, Vector{T}(undef, obuflen), Vector{T}(undef, buflen), Vector{S}(undef, rlen), Vector{T}(undef, clen)) - workers = [mk(cg)] + rtw = _re_twiddles(T, rlen) + graphs = [cg] for _ in 2:num_threads - push!(workers, mk(map(_clone_workspace, cg))) + push!(graphs, map(_clone_workspace, cg)) end - return workers + gathers = [first(g)[1].type === POW2RADIX4_FFT ? first(g).workspace[1] : T[] for g in graphs] + mk(graphs) = Worker{T,N,S}(graphs, Vector{T}(undef, obuflen), Vector{T}(undef, buflen), Vector{S}(undef, rlen), Vector{T}(undef, clen), rtw, gathers) + return [mk(g) for g in graphs] end # Transforms with fewer elements than this are not split over threads. @@ -237,7 +253,18 @@ function LinearAlgebra.mul!(y::AbstractVector{U}, p::FFTAPlan_cx{T,1}, x::Abstra if size(p) != size(x) throw(DimensionMismatch("plan has axes $(size(p)), but input array has axes $(size(x))")) end - fft_kernel!(y, x, 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1) + _kernel_1d!(y, x, p.dir, p.callgraph[1], p.workers[1].gathers) + return y +end + +# One whole transform: the threaded leaves-first path for large powers of two +# when the plan has several workers, the kernel otherwise. +function _kernel_1d!(y::AbstractVector, x::AbstractVector, d::Direction, g::CallGraph{T}, gathers::Vector{Vector{T}}) where {T} + if _threaded_1d_ok(g, gathers) + _pow2_leaffirst_threaded!(y, x, g[1].sz, 1, 1, 1, d, g.twiddles[1], 0, gathers) + else + fft_kernel!(y, x, 1, 1, d, g[1].type, g, 1) + end return y end @@ -282,13 +309,13 @@ function _foreach_pencil(f::F, A::AbstractArray{<:Any,N}, ::Val{dim}, workers::V total = npre * length(Rpost) nt = min(length(workers), total ÷ 2) if nt > 1 && total * size(A, dim) >= THREAD_THRESHOLD - Base.@sync for c in 1:nt + # one chunk per worker, run on Polyester's static thread pool (no task + # allocation, the calling thread takes a chunk itself) + @batch for c in 1:nt lo = (c - 1) * total ÷ nt + 1 hi = c * total ÷ nt - # (the task's worker gets its own name: a variable shared with the - # serial branch below would be captured by every task) wc = workers[c] - Threads.@spawn for k in lo:hi + for k in lo:hi ipost, ipre = divrem(k - 1, npre) f(wc, Rpre[ipre + 1], Rpost[ipost + 1]) end @@ -495,29 +522,27 @@ function _rfft_pencil!(y::AbstractVector{T}, x::AbstractVector{<:Real}, w::Worke # algorithms." IEEE Transactions on acoustics, speech, and signal # processing 35, no. 6 (2003): 849-863. m = n >> 1 - @inbounds for j in 1:m - buf[j] = T(x[2j - 1], x[2j]) - end - fft_kernel!(view(y, 1:m), buf, 1, 1, FFT_FORWARD, cg[1].type, cg, 1) + _pack_pairs!(buf, x, m) + _kernel_1d!(view(y, 1:m), buf, FFT_FORWARD, cg, w.gathers) # Construct the result by first constructing the elements of the # real and imaginary part, followed by the usual radix-2 assembly, - # see eq (9). The twiddle is for `n`, not `m`, so it is recomputed. - z1 = singleton_params(-one(R) / n) - wj = cispi(-R(2) / n) + # see eq (9). The twiddles are for `n`, not `m`: the plan keeps them + # in `rtw` (a recurrence here would be a serial dependency chain). + rtw = w.rtw @inbounds begin y1 = y[1] y[1] = real(y1) + imag(y1) y[m + 1] = real(y1) - imag(y1) - for j in 2:((m >> 1) + 1) - yj = y[j] - ymj = y[m - j + 2] - XX = R(0.5) * ( yj + conj(ymj)) - XY = R(0.5) * (-yj + conj(ymj)) * im - y[j] = XX + wj * XY - y[m - j + 2] = conj(XX - wj * XY) - wj = singleton_step(wj, z1) - end + end + _rfft_post_simd!(y, m, rtw) || @inbounds for j in 2:((m >> 1) + 1) + yj = y[j] + ymj = y[m - j + 2] + wj = rtw[j - 1] + XX = R(0.5) * ( yj + conj(ymj)) + XY = R(0.5) * (-yj + conj(ymj)) * im + y[j] = XX + wj * XY + y[m - j + 2] = conj(XX - wj * XY) end else # Odd length: run the full transform on the real input (the kernels @@ -542,23 +567,17 @@ function _brfft_pencil!(x::AbstractVector{<:Real}, y::AbstractVector{T}, w::Work m = n >> 1 tmp = view(buf, 1:m) out = view(buf, m + 1:2m) - z1 = singleton_params(one(R) / n) - wj = cispi(R(2) / n) - @inbounds begin - tmp[1] = T(real(y[1]) + real(y[m + 1]), real(y[1]) - real(y[m + 1])) - for j in 2:((m >> 1) + 1) - XX = y[j] + conj(y[m - j + 2]) - XY = wj * (y[j] - conj(y[m - j + 2])) - tmp[j] = XX + im * XY - tmp[m - j + 2] = conj(XX - im * XY) - wj = singleton_step(wj, z1) - end - end - fft_kernel!(out, tmp, 1, 1, FFT_BACKWARD, cg[1].type, cg, 1) - @inbounds for j in 1:m - x[2j - 1] = real(out[j]) - x[2j] = imag(out[j]) + rtw = w.rtw + @inbounds tmp[1] = T(real(y[1]) + real(y[m + 1]), real(y[1]) - real(y[m + 1])) + _brfft_pre_simd!(tmp, y, m, rtw) || @inbounds for j in 2:((m >> 1) + 1) + wj = conj(rtw[j - 1]) + XX = y[j] + conj(y[m - j + 2]) + XY = wj * (y[j] - conj(y[m - j + 2])) + tmp[j] = XX + im * XY + tmp[m - j + 2] = conj(XX - im * XY) end + _kernel_1d!(out, tmp, FFT_BACKWARD, cg, w.gathers) + _unpack_pairs!(x, out, m) else # Odd length: rebuild the conjugate-symmetric spectrum and transform. h = n ÷ 2 + 1 @@ -584,8 +603,6 @@ end # stride of a cache line or more the kernel is otherwise 1.2–1.3× slower. The # test is sufficient rather than exact (a 1×N array's dim-2 pencils are # contiguous but still take the copy): a wasted copy, never a wrong answer. -_unit_stride(v::StridedArray) = stride(v, 1) == 1 -_unit_stride(v) = false # Apply `kernel!(y_pencil, x_pencil, worker, flen)` along dimension `R` of `x` and `y`. function _re_pencil_loop!(kernel!::F, y::AbstractArray{<:Any,N}, x::AbstractArray{<:Any,N}, p::FFTAPlan_re, ::Val{R}) where {F,N,R} diff --git a/src/real_simd.jl b/src/real_simd.jl new file mode 100644 index 0000000..a6fdb3c --- /dev/null +++ b/src/real_simd.jl @@ -0,0 +1,141 @@ +# Vectorised pre/post-processing of the even-length real transforms. +# +# `_rfft_pencil!` turns the half-length complex transform `Y` into the real +# transform with, for `j = 2..m/2+1`, +# +# XX = (Y[j] + conj(Y[m-j+2]))/2, XY = i (conj(Y[m-j+2]) - Y[j])/2 +# X[j] = XX + w^j XY, X[m-j+2] = conj(XX - w^j XY) +# +# and `_brfft_pencil!` the inverse. On scalar `Complex` values these loops are +# compute bound; here `W` values of `j` are handled per iteration on vectors +# of `2W` reals (`W = 2` for `Float64`, `4` for `Float32`): the front block is +# loaded in order, the mirrored back block is loaded and reversed lane-wise, +# conjugation is a sign flip of the imaginary lanes and `i·z` a lane swap with +# a sign flip. A scalar tail handles the remainder and the middle element. +# Used for unit-stride vectors of `ComplexF32`/`ComplexF64`; other cases keep +# the scalar loops. + +# `x[2j-1], x[2j]` pairs of a unit-stride real vector are the memory layout of +# the complex vector `buf`: pack and unpack with a copy where possible. +function _pack_pairs!(buf::AbstractVector{Complex{R}}, x::AbstractVector{R}, m::Int) where {R<:Real} + if _unit_stride(x) && _unit_stride(buf) + GC.@preserve buf x unsafe_copyto!(Ptr{R}(pointer(buf)), pointer(x), 2m) + else + @inbounds for j in 1:m + buf[j] = Complex{R}(x[2j - 1], x[2j]) + end + end + return buf +end +function _unpack_pairs!(x::AbstractVector{R}, out::AbstractVector{Complex{R}}, m::Int) where {R<:Real} + if _unit_stride(x) && _unit_stride(out) + GC.@preserve x out unsafe_copyto!(pointer(x), Ptr{R}(pointer(out)), 2m) + else + @inbounds for j in 1:m + x[2j - 1] = real(out[j]) + x[2j] = imag(out[j]) + end + end + return x +end +_pack_pairs!(buf, x, m) = (@inbounds for j in 1:m; buf[j] = eltype(buf)(x[2j - 1], x[2j]); end; buf) +_unpack_pairs!(x, out, m) = (@inbounds for j in 1:m; x[2j - 1] = real(out[j]); x[2j] = imag(out[j]); end; x) + +# unit-stride dense storage that `pointer` can address (also used by the +# pencil loops of plan.jl) +_unit_stride(v::StridedArray) = stride(v, 1) == 1 +_unit_stride(v) = false + +@inline _swap_ri(v::Vec{L}) where {L} = shufflevector(v, Val(ntuple(i -> isodd(i) ? i : i - 2, L))) +@inline _rev_complex(v::Vec{L}) where {L} = shufflevector(v, Val(ntuple(i -> isodd(i) ? L - i - 1 : L - i + 1, L))) +@inline _dup_re(v::Vec{L}) where {L} = shufflevector(v, Val(ntuple(i -> 2 * ((i - 1) ÷ 2), L))) +@inline _dup_im(v::Vec{L}) where {L} = shufflevector(v, Val(ntuple(i -> 2 * ((i - 1) ÷ 2) + 1, L))) + +""" +$(TYPEDSIGNATURES) +The post-processing loop of `_rfft_pencil!` for `j = 2..m÷2+1` in place on +`y`, with the twiddles `rtw[j-1] = w^(j-1)`. Returns `false` (nothing done) +when the arrays are not dense vectors of a SIMD element type. +""" +function _rfft_post_simd!(y::AbstractVector{T}, m::Int, rtw::AbstractVector{T}) where {T<:CodeletEltype} + (_unit_stride(y) && rtw isa Vector{T}) || return false + R = real(T) + W = _simd_width(T) + L = 2W + V = Vec{L,R} + sz = sizeof(R) + conjs = Vec{L,R}(ntuple(i -> isodd(i) ? one(R) : -one(R), L)) + half = Vec{L,R}(R(0.5)) + py = Ptr{R}(pointer(y)); pt = Ptr{R}(pointer(rtw)) + jlast = (m >> 1) + 1 + j = 2 + GC.@preserve y rtw begin + @inbounds while j + W - 1 < jlast + pf = py + (j - 1) * 2sz + pb = py + (m - j + 2 - W) * 2sz + yj = vload(V, pf) + ymj = _rev_complex(vload(V, pb)) * conjs # conj(y[m-j+2]) in j order + w = vload(V, pt + (j - 2) * 2sz) + wr = _dup_re(w); wi = _dup_im(w) * -conjs # (-wi, wi) + XX = half * (yj + ymj) + XY = _swap_ri(half * (ymj - yj)) * -conjs # i (conj(ymj) - yj)/2 + t = muladd(_swap_ri(XY), wi, XY * wr) # w * XY + vstore(XX + t, pf) + vstore(_rev_complex((XX - t) * conjs), pb) + j += W + end + @inbounds for jj in j:jlast + yj = y[jj]; ymj = y[m - jj + 2]; wj = rtw[jj - 1] + XX = R(0.5) * ( yj + conj(ymj)) + XY = R(0.5) * (-yj + conj(ymj)) * im + y[jj] = XX + wj * XY + y[m - jj + 2] = conj(XX - wj * XY) + end + end + return true +end +_rfft_post_simd!(y, m, rtw) = false + +""" +$(TYPEDSIGNATURES) +The pre-processing loop of `_brfft_pencil!` for `j = 2..m÷2+1`, from the +spectrum `y` into `tmp` (a dense vector), with the conjugated twiddles. +""" +function _brfft_pre_simd!(tmp::AbstractVector{T}, y::AbstractVector{T}, m::Int, rtw::AbstractVector{T}) where {T<:CodeletEltype} + (_unit_stride(tmp) && _unit_stride(y) && rtw isa Vector{T}) || return false + R = real(T) + W = _simd_width(T) + L = 2W + V = Vec{L,R} + sz = sizeof(R) + conjs = Vec{L,R}(ntuple(i -> isodd(i) ? one(R) : -one(R), L)) + py = Ptr{R}(pointer(y)); po = Ptr{R}(pointer(tmp)); pt = Ptr{R}(pointer(rtw)) + jlast = (m >> 1) + 1 + j = 2 + GC.@preserve tmp y rtw begin + @inbounds while j + W - 1 < jlast + pf = py + (j - 1) * 2sz + pb = py + (m - j + 2 - W) * 2sz + yj = vload(V, pf) + ymj = _rev_complex(vload(V, pb)) * conjs + w = vload(V, pt + (j - 2) * 2sz) + wr = _dup_re(w); wi = _dup_im(w) * conjs # conj(w): (wi, -wi) + XX = yj + ymj + d = yj - ymj + XY = muladd(_swap_ri(d), wi, d * wr) # conj(w) * (yj - conj(ymj)) + iXY = _swap_ri(XY) * -conjs # i·XY + vstore(XX + iXY, po + (j - 1) * 2sz) + vstore(_rev_complex((XX - iXY) * conjs), po + (m - j + 2 - W) * 2sz) + j += W + end + @inbounds for jj in j:jlast + wj = conj(rtw[jj - 1]) + XX = y[jj] + conj(y[m - jj + 2]) + XY = wj * (y[jj] - conj(y[m - jj + 2])) + tmp[jj] = XX + im * XY + tmp[m - jj + 2] = conj(XX - im * XY) + end + end + return true +end +_brfft_pre_simd!(tmp, y, m, rtw) = false diff --git a/src/simd_pass.jl b/src/simd_pass.jl new file mode 100644 index 0000000..1a8c5ea --- /dev/null +++ b/src/simd_pass.jl @@ -0,0 +1,92 @@ +# 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. + +# (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, k0::Int = 0, k1::Int = m +) where {T<:CodeletEltype} + W = _simd_width(T) + (stride_out == 1 && m >= W && k0 % W == 0 && (k1 - k0) % W == 0 && _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 k0:W:k1-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, k0 = 0, k1 = m) = false diff --git a/test/leaffirst.jl b/test/leaffirst.jl new file mode 100644 index 0000000..c411d20 --- /dev/null +++ b/test/leaffirst.jl @@ -0,0 +1,45 @@ +# 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; num_threads = 1) + 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; num_threads = 1) + 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 + # threaded leaves-first path (several workers): same operations, identical output + if Threads.nthreads() > 1 + pt = plan_fft(x; num_threads = Threads.nthreads()) + @test pt * x == y + @test (@allocated mul!(y, pt, x)) <= 64 * 1024 # task objects only + prt = plan_rfft(xr; num_threads = Threads.nthreads()) + @test prt * xr == yr + @test brfft(prt * xr, n) ≈ n .* xr rtol = rtol + end +end diff --git a/test/runtests.jl b/test/runtests.jl index 23f1319..7edc5f7 100644 --- a/test/runtests.jl +++ b/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