diff --git a/Project.toml b/Project.toml index 868f974..19a72e2 100644 --- a/Project.toml +++ b/Project.toml @@ -17,6 +17,7 @@ AbstractFFTs = "1" Aqua = "0.8" DocStringExtensions = "0.9" ExplicitImports = "1.12" +FFTW = "1.8" LinearAlgebra = "<0.0.1, 1" MuladdMacro = "0.2" Primes = "0.5" @@ -28,7 +29,8 @@ julia = "1.6.7" [extras] Aqua = "4c88cf16-eb10-579e-8560-4a9242c79595" ExplicitImports = "7d51a73a-1435-4ff3-83d9-f097790105c7" +FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [targets] -test = ["Aqua", "ExplicitImports", "Test"] +test = ["Aqua", "ExplicitImports", "FFTW", "Test"] diff --git a/docs/src/dev.md b/docs/src/dev.md index 2e35f08..a7796fb 100644 --- a/docs/src/dev.md +++ b/docs/src/dev.md @@ -10,7 +10,7 @@ Here is the documentation for key parts of the development side of the package. CallGraphNode CallGraph CallGraphNode! -fft! +fft_kernel! fft_dft! fft_pow3! ``` diff --git a/src/algos.jl b/src/algos.jl index f67cdca..7374da1 100644 --- a/src/algos.jl +++ b/src/algos.jl @@ -2,7 +2,7 @@ Int(d) end -function fft!( +function fft_kernel!( out::AbstractVector{T}, in::AbstractVector{<:Number}, start_out::Int, start_in::Int, d::Direction, @@ -85,7 +85,7 @@ function fft_composite!( 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) + fft_kernel!(tmp, in, R_start_out, R_start_in, d, Rt, g, right_idx) end if j1 > 0 @@ -112,7 +112,7 @@ function fft_composite!( 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) + fft_kernel!(out, tmp, L_start_out, L_start_in, d, Lt, g, left_idx) end end end diff --git a/src/plan.jl b/src/plan.jl index a71cbc8..e4de7a2 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -2,35 +2,92 @@ abstract type FFTAPlan{T,N} <: AbstractFFTs.Plan{T} end -struct FFTAInvPlan{_T,_N} <: FFTAPlan{_T,_N} end - const RegionTypes{N} = Union{Int,AbstractVector{Int},NTuple{N,Int}} -struct FFTAPlan_cx{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} +# The plan types are mutable so that `AbstractFFTs.inv` can cache the inverse +# plan in the initially undefined `pinv` field (see `plan_inv`), as FFTW.jl does. + +""" +$(TYPEDEF) +Per-thread state of a plan: call graphs (sharing the nodes of the plan's call +graphs but with their own workspace) and pencil buffers. A plan owns one +`Worker` per thread it may use (see the `num_threads` keyword of the `plan_*` +functions); pencils of a multidimensional or batched transform are +distributed over them. +""" +struct Worker{T,N,S<:Real} callgraph::NTuple{N,CallGraph{T}} - region::R - dir::Direction - pinv::FFTAInvPlan{T,N} + obuf::Vector{T} # transform output of one pencil before it is copied back + 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) +end + +# A call graph that shares the (immutable) nodes but has its own workspace. +_clone_workspace(g::CallGraph{T}) where {T} = + CallGraph{T}(g.nodes, [Vector{T}(undef, length(w)) for w in g.workspace], g.BLUESTEIN_CUTOFF) + +# `buflen` is the packing scratch length and `rlen` the real length of a real +# plan (0 for complex plans, which need neither). +function _workers(cg::Tuple{CallGraph{T},Vararg{CallGraph{T}}}, ::Val{N}, num_threads::Int, buflen::Int, rlen::Int = 0) where {T,N} + num_threads >= 1 || throw(ArgumentError("num_threads must be at least 1")) + 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)] + for _ in 2:num_threads + push!(workers, mk(map(_clone_workspace, cg))) + end + return workers +end + +# Transforms with fewer elements than this are not split over threads. +const THREAD_THRESHOLD = 1 << 15 + +mutable struct FFTAPlan_cx{T,N,R<:RegionTypes{N},S<:Real} <: FFTAPlan{T,N} + const callgraph::NTuple{N,CallGraph{T}} + const region::R + const dir::Direction + const workers::Vector{Worker{T,N,S}} + pinv::AbstractFFTs.ScaledPlan + FFTAPlan_cx{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, workers::Vector{Worker{T,N,S}}) where {T,N,R<:RegionTypes{N},S} = + new{T,N,R,S}(cg, r, dir, workers) end function FFTAPlan_cx{T,N}( cg::NTuple{N,CallGraph{T}}, r::R, - dir::Direction + dir::Direction; num_threads::Int=Threads.nthreads() ) where {T,N,R<:RegionTypes{N}} - FFTAPlan_cx{T,N,R}(cg, r, dir, FFTAInvPlan{T,N}()) + FFTAPlan_cx{T,N,R}(cg, r, dir, _workers(cg, Val(N), num_threads, 0)) end -struct FFTAPlan_re{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} - callgraph::NTuple{N,CallGraph{T}} - region::R - dir::Direction - flen::Int - pinv::FFTAInvPlan{T,N} +mutable struct FFTAPlan_re{T,N,R<:RegionTypes{N},S<:Real} <: FFTAPlan{T,N} + const callgraph::NTuple{N,CallGraph{T}} + const region::R + const dir::Direction + const flen::Int + const workers::Vector{Worker{T,N,S}} + pinv::AbstractFFTs.ScaledPlan + FFTAPlan_re{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, flen::Int, workers::Vector{Worker{T,N,S}}) where {T,N,R<:RegionTypes{N},S} = + new{T,N,R,S}(cg, r, dir, flen, workers) end function FFTAPlan_re{T,N}( cg::NTuple{N,CallGraph{T}}, r::R, - dir::Direction, flen::Int + dir::Direction, flen::Int; num_threads::Int=Threads.nthreads() ) where {T,N,R<:RegionTypes{N}} - FFTAPlan_re{T,N,R}(cg, r, dir, flen, FFTAInvPlan{T,N}()) + FFTAPlan_re{T,N,R}(cg, r, dir, flen, _workers(cg, Val(N), num_threads, _re_buflen(flen, dir), flen)) +end + +_num_threads(p::Union{FFTAPlan_cx,FFTAPlan_re}) = length(p.workers) + +# Scratch length needed by the real-transform pencil kernels for a plan of +# real length `n` (see `_rfft_pencil!` / `_brfft_pencil!`). +function _re_buflen(n::Int, dir::Direction) + if iseven(n) + dir === FFT_FORWARD ? n >> 1 : n + else + dir === FFT_FORWARD ? n : 2n + end end function Base.size(p::FFTAPlan{<:Any,N}, i::Int) where N @@ -46,8 +103,6 @@ function Base.size(p::FFTAPlan{<:Any,N}, i::Int) where N end Base.size(p::FFTAPlan{<:Any,N}) where N = ntuple(Base.Fix1(size, p), Val{N}()) -Base.complex(p::FFTAPlan_re{T,N,R}) where {T,N,R} = FFTAPlan_cx{T,N,R}(p.callgraph, p.region, p.dir, p.pinv) - function _sort(region::T)::T where {N,T<:NTuple{N,Int}} @static if VERSION >= v"1.12" sort(region) @@ -78,31 +133,46 @@ function _plan_fft( x::AbstractArray{T,N}, region::RegionTypes, dir::Direction; - BLUESTEIN_CUTOFF=DEFAULT_BLUESTEIN_CUTOFF, _kwargs... + BLUESTEIN_CUTOFF=DEFAULT_BLUESTEIN_CUTOFF, num_threads::Int=Threads.nthreads(), _kwargs... ) where {T<:Complex,N} M = length(region) if M == 1 R1 = Int(region[1]) g = CallGraph{T}(size(x, R1), BLUESTEIN_CUTOFF) - return FFTAPlan_cx{T,1}((g,), R1, dir) + return FFTAPlan_cx{T,1}((g,), R1, dir; num_threads) elseif M == 2 R2 = _sort(region) g1 = CallGraph{T}(size(x, R2[1]), BLUESTEIN_CUTOFF) g2 = CallGraph{T}(size(x, R2[2]), BLUESTEIN_CUTOFF) - return FFTAPlan_cx{T,2}((g1, g2), R2, dir) + return FFTAPlan_cx{T,2}((g1, g2), R2, dir; num_threads) else RM = _sort(region) return FFTAPlan_cx{T,M}( ntuple(i -> CallGraph{T}(size(x, RM[i]), BLUESTEIN_CUTOFF), Val(M)), - RM, dir + RM, dir; num_threads ) end end -function AbstractFFTs.plan_rfft( +# The `AbstractFFTs` entry points deliberately leave `region` unannotated: an +# annotation such as `region::RegionTypes` makes these methods ambiguous with +# FFTW.jl's `plan_rfft(::StridedArray{Float64}, region)` when both packages +# are loaded (neither is more specific in every argument), which turns +# `rfft(x)` into a `MethodError`. With the annotation on the array only, +# FFTW's methods are strictly more specific and win, as AbstractFFTs intends. +AbstractFFTs.plan_rfft(x::AbstractArray{T,N}, region; kwargs...) where {T<:Real,N} = + _plan_rfft(x, _region(region); kwargs...) + +AbstractFFTs.plan_brfft(x::AbstractArray{T,N}, len::Integer, region; kwargs...) where {T,N} = + _plan_brfft(x, Int(len), _region(region); kwargs...) + +_region(r::RegionTypes) = r +_region(r) = collect(Int, r) + +function _plan_rfft( x::AbstractArray{T,N}, region::RegionTypes; - BLUESTEIN_CUTOFF=DEFAULT_BLUESTEIN_CUTOFF, _kwargs... + BLUESTEIN_CUTOFF=DEFAULT_BLUESTEIN_CUTOFF, num_threads::Int=Threads.nthreads(), _kwargs... ) where {T<:Real,N} M = length(region) if M == 1 @@ -113,22 +183,24 @@ function AbstractFFTs.plan_rfft( # problems, we just solve the problem as a single complex nn = iseven(n) ? n >> 1 : n g = CallGraph{Complex{T}}(nn, BLUESTEIN_CUTOFF) - return FFTAPlan_re{Complex{T},1}((g,), R1, FFT_FORWARD, n) + return FFTAPlan_re{Complex{T},1}((g,), R1, FFT_FORWARD, n; num_threads) elseif M == 2 R2 = _sort(region) - g1 = CallGraph{Complex{T}}(size(x, R2[1]), BLUESTEIN_CUTOFF) + n = size(x, R2[1]) + nn = iseven(n) ? n >> 1 : n + g1 = CallGraph{Complex{T}}(nn, BLUESTEIN_CUTOFF) g2 = CallGraph{Complex{T}}(size(x, R2[2]), BLUESTEIN_CUTOFF) - return FFTAPlan_re{Complex{T},2}((g1, g2), R2, FFT_FORWARD, size(x, R2[1])) + return FFTAPlan_re{Complex{T},2}((g1, g2), R2, FFT_FORWARD, n; num_threads) else throw(ArgumentError("only supports 1D and 2D FFTs")) end end -function AbstractFFTs.plan_brfft( +function _plan_brfft( x::AbstractArray{T,N}, len::Int, region::RegionTypes; - BLUESTEIN_CUTOFF=DEFAULT_BLUESTEIN_CUTOFF, _kwargs... + BLUESTEIN_CUTOFF=DEFAULT_BLUESTEIN_CUTOFF, num_threads::Int=Threads.nthreads(), _kwargs... ) where {T,N} M = length(region) if M == 1 @@ -138,12 +210,13 @@ function AbstractFFTs.plan_brfft( R1 = Int(region[1]) nn = iseven(len) ? len >> 1 : len g = CallGraph{T}(nn, BLUESTEIN_CUTOFF) - return FFTAPlan_re{T,1}((g,), R1, FFT_BACKWARD, len) + return FFTAPlan_re{T,1}((g,), R1, FFT_BACKWARD, len; num_threads) elseif M == 2 R2 = _sort(region) - g1 = CallGraph{T}(len, BLUESTEIN_CUTOFF) + nn = iseven(len) ? len >> 1 : len + g1 = CallGraph{T}(nn, BLUESTEIN_CUTOFF) g2 = CallGraph{T}(size(x, R2[2]), BLUESTEIN_CUTOFF) - return FFTAPlan_re{T,2}((g1, g2), R2, FFT_BACKWARD, len) + return FFTAPlan_re{T,2}((g1, g2), R2, FFT_BACKWARD, len; num_threads) else throw(ArgumentError("only supports 1D and 2D FFTs")) end @@ -162,7 +235,7 @@ 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!(y, x, 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1) + fft_kernel!(y, x, 1, 1, p.dir, p.callgraph[1][1].type, p.callgraph[1], 1) return y end @@ -191,18 +264,52 @@ function LinearAlgebra.mul!(y::AbstractArray{U,N}, p::FFTAPlan_cx{T,1}, x::Abstr return y end +""" +$(TYPEDSIGNATURES) +Call `f(worker, Ipre, Ipost)` for every pencil `A[Ipre, :, Ipost]` along +dimension `dim` of `A`. When the plan has several workers and +the transform is large enough (`THREAD_THRESHOLD` elements), the pencils are +split into one contiguous chunk per worker and the chunks run as parallel +tasks, each using its own worker (so its own call-graph workspace and +buffers). Results do not depend on the number of threads. +""" +function _foreach_pencil(f::F, A::AbstractArray{<:Any,N}, ::Val{dim}, workers::Vector{W}) where {F,N,dim,W} + Rpre = CartesianIndices(ntuple(Base.Fix1(size, A), Val(dim - 1))) + Rpost = CartesianIndices(ntuple(i -> size(A, dim + i), Val(N - dim))) + npre = length(Rpre) + 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 + 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 + ipost, ipre = divrem(k - 1, npre) + f(wc, Rpre[ipre + 1], Rpost[ipost + 1]) + end + end + else + w1 = workers[1] + for Ipost in Rpost, Ipre in Rpre + f(w1, Ipre, Ipost) + end + end + return nothing +end + function _mul_loop!( y::AbstractArray{U,N}, x::AbstractArray{T,N}, p::FFTAPlan_cx{T,1}, ::Val{R} ) where {T,U,N,R} - Rpre = CartesianIndices(ntuple(Base.Fix1(size, x), Val(R - 1))) - Rpost = CartesianIndices(ntuple(i -> size(x, R + i), Val(N - R))) - cg = p.callgraph[1] - t = cg[1].type - for Ipost in Rpost, Ipre in Rpre - @views fft!(y[Ipre,:,Ipost], x[Ipre,:,Ipost], 1, 1, p.dir, t, cg, 1) + dir = p.dir + _foreach_pencil(x, Val(R), p.workers) do w, Ipre, Ipost + cg = w.callgraph[1] + @views fft_kernel!(y[Ipre,:,Ipost], x[Ipre,:,Ipost], 1, 1, dir, cg[1].type, cg, 1) end end @@ -221,35 +328,18 @@ function LinearAlgebra.mul!( throw(DimensionMismatch("Plan region is outside array dimensions.")) end - sz = size(X) - max_sz = maximum(sz) - obuf = Vector{T}(undef, max_sz) - ibuf = Vector{T}(undef, max_sz) - sizehint!(obuf, max_sz) # not guaranteed but hopefully prevents allocations - sizehint!(ibuf, max_sz) dir = p.dir + workers = p.workers copyto!(out, X) # operate in-place on output array if @generated quote - Base.Cartesian.@nexprs $N dim -> begin - n = size(out, dim) - resize!(obuf, n) - resize!(ibuf, n) - cg = p.callgraph[dim] - - fft_along_dim!(out, ibuf, obuf, cg, dir, Val(dim)) - end + Base.Cartesian.@nexprs $N dim -> fft_along_dim!(out, workers, dim, dir, Val(dim)) end else for dim in 1:N - n = size(out, dim) - resize!(obuf, n) - resize!(ibuf, n) - cg = p.callgraph[dim] - - fft_along_dim!(out, ibuf, obuf, cg, dir, Val(dim)) + fft_along_dim!(out, workers, dim, dir, Val(dim)) end end @@ -271,78 +361,58 @@ function LinearAlgebra.mul!( throw(DimensionMismatch("Plan region is outside array dimensions.")) end - max_sz = maximum(Base.Fix1(size, out), p.region) - obuf = Vector{T}(undef, max_sz) - ibuf = Vector{T}(undef, max_sz) - sizehint!(obuf, max_sz) # not guaranteed but hopefully prevents allocations - sizehint!(ibuf, max_sz) - copyto!(out, X) # operate in-place on output array - _execute_mdfft!(out, ibuf, obuf, p.dir, p.region, p.callgraph) + _execute_mdfft!(out, p.workers, p.dir, p.region) return out end @noinline function _execute_mdfft!( out::AbstractArray{U,N}, - ibuf::Vector{T}, obuf::Vector{T}, + workers::Vector{<:Worker{T,M}}, dir::Direction, @nospecialize(region::RegionTypes), - @nospecialize(callgraphs::NTuple) -) where {T,U,N} - - M = length(region) +) where {T,U,N,M} if @generated quote k = 1 # region is assumed to be pre-sorted during planning Base.Cartesian.@nexprs $N dim -> begin if region[k] == dim - n = size(out, dim) - resize!(obuf, n) - resize!(ibuf, n) - cg = callgraphs[k] - - fft_along_dim!(out, ibuf, obuf, cg, dir, Val(dim)) - + fft_along_dim!(out, workers, k, dir, Val(dim)) k = min(k + 1, M) end end return nothing end else - for dim in 1:M - pdim = region[dim] - n = size(out, pdim) - resize!(obuf, n) - resize!(ibuf, n) - cg = callgraphs[dim] - - fft_along_dim!(out, ibuf, obuf, cg, dir, Val(pdim)) + for k in 1:M + fft_along_dim!(out, workers, k, dir, Val(Int(region[k]))) end end end +""" +$(TYPEDSIGNATURES) +Transform every pencil of `A` along dimension `dim` in place, using the `k`-th +call graph of the workers. The kernels read the (strided) pencil directly and +write the worker's output buffer, which is then copied back. +""" function fft_along_dim!( A::AbstractArray{U,N}, - ibuf::Vector{T}, obuf::Vector{T}, - cg::CallGraph{T}, d::Direction, + workers::Vector{<:Worker{T,M}}, + k::Int, d::Direction, ::Val{dim} -) where {T <: Complex{<:AbstractFloat}, U, N, dim} - - Rpre = CartesianIndices(ntuple(Base.Fix1(size, A), Val(dim - 1))) - Rpost = CartesianIndices(ntuple(i -> size(A, dim + i), Val(N - dim))) - t = cg[1].type - cols = eachindex(axes(A, dim), ibuf, obuf) - - for Ipost in Rpost, Ipre in Rpre - for j in cols - ibuf[j] = A[Ipre, j, Ipost] - end - fft!(obuf, ibuf, 1, 1, d, t, cg, 1) - for j in cols - A[Ipre, j, Ipost] = obuf[j] +) where {T <: Complex{<:AbstractFloat}, U, N, M, dim} + n = size(A, dim) + _foreach_pencil(A, Val(dim), workers) do w, Ipre, Ipost + cg = w.callgraph[k] + obuf = w.obuf + pencil = @view A[Ipre, :, Ipost] + fft_kernel!(obuf, pencil, 1, 1, d, cg[1].type, cg, 1) + @inbounds for j in 1:n + pencil[j] = obuf[j] end end end @@ -384,200 +454,334 @@ function Base.:*(p::FFTAPlan_cx{T,N1}, x::AbstractArray{T,N2}) where {T<:Complex end ### Real -# By converting the problem to complex and back to real -#### 1D plan 1D array -##### Forward -function Base.:*(p::FFTAPlan_re{Complex{T},1}, x::AbstractVector{T}) where {T<:Real} - if p.dir !== FFT_FORWARD - throw(ArgumentError("only FFT_FORWARD supported for real vectors")) +# Real transforms are computed from complex transforms of half (even `n`) or +# full (odd `n`) length; see `_rfft_pencil!`/`_brfft_pencil!`. All entry points +# funnel into `mul!`, which is allocation-free for 1D plans (the scratch space +# lives in the plan) and applies the same 1D kernel to every pencil along the +# transform dimension of an N-d array. + +function _check_re_dims(y, p::FFTAPlan_re, x, d1::Int, fwd::Bool) + # `x` is the input, `y` the output; the real-length dimension is `d1`. + rlen, clen = p.flen, p.flen ÷ 2 + 1 + xr, yr = fwd ? (rlen, clen) : (clen, rlen) + if ndims(x) != ndims(y) + throw(DimensionMismatch("input has $(ndims(x)) dimensions, output has $(ndims(y))")) + elseif size(x, d1) != xr + throw(DimensionMismatch("real 1D plan has size $rlen. Dimension of input array along region $d1 should have size $xr, but has size $(size(x, d1))")) + elseif size(y, d1) != yr + throw(DimensionMismatch("output array should have size $yr along region $d1, but has size $(size(y, d1))")) end - Base.require_one_based_indexing(x) - - n = p.flen - p_c = complex(p) - if iseven(n) - # For problems of even size, we solve the rfft problem by splitting the - # problem into the even and odd part and solving them simultaneously as - # a single (complex) fft of half the size, see equations (6)-(8) of - # Sorensen, H. V., D. Jones, Michael Heideman, and C. Burrus. - # "Real-valued fast Fourier transform algorithms." - # IEEE Transactions on acoustics, speech, and signal processing 35, no. 6 (2003): 849-863. - if x isa Vector && isbitstype(T) - # For a vector of bits, we can just reinterpret the bits to get the - # appropriate representation of even (zero based) elements as the real - # part and the odd as the complex part - x_c = reinterpret(Complex{T}, x) - else - # for non-bits, we'd have to copy to a new array - x_c = complex.(view(x, 1:2:n), view(x, 2:2:n)) + for i in 1:ndims(x) + if i != d1 && size(x, i) != size(y, i) + throw(DimensionMismatch("input array has size $(size(x)), but output array has size $(size(y))")) end + end + return nothing +end +# Forward real-to-complex transform of one pencil: `x` real of length `n`, +# `y` complex of length `n ÷ 2 + 1`. +function _rfft_pencil!(y::AbstractVector{T}, x::AbstractVector{<:Real}, w::Worker{T}, n::Int) where {T<:Complex} + R = real(T) + cg = w.callgraph[1] + buf = w.buf + if iseven(n) + # Solve the rfft problem by splitting the input into even and odd parts + # and solving them simultaneously as a single (complex) fft of half + # the size, see equations (6)-(8) of Sorensen, H. V., D. Jones, Michael + # Heideman, and C. Burrus. "Real-valued fast Fourier transform + # algorithms." IEEE Transactions on acoustics, speech, and signal + # processing 35, no. 6 (2003): 849-863. m = n >> 1 - # Allocate complex result vector of half the input size plus one - y = similar(x_c, m + 1) - # Solve the complex fft of half the size - LinearAlgebra.mul!(view(y, 1:m), p_c, x_c) - - # The w stored in the plan is for m, not n, so probably cheapest to - # just recompute it instead of taking a square root - z1 = singleton_params(-one(T) / n) - wj = cispi(-T(2) / n) + @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) # Construct the result by first constructing the elements of the # real and imaginary part, followed by the usual radix-2 assembly, - # see eq (9) - y1 = y[1] - y[1] = real(y1) + imag(y1) - y[end] = real(y1) - imag(y1) - - @inbounds for j in 2:((m >> 1) + 1) - yj = y[j] - ymj = y[m-j+2] - XX = T(0.5) * ( yj + conj(ymj)) - XY = T(0.5) * (-yj + conj(ymj)) * im - y[j] = XX + wj * XY - y[m-j+2] = conj(XX - wj * XY) - wj = singleton_step(wj, z1) + # 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) + @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 - return y else - # when the problem cannot be split in two equal size chunks we - # convert the problem to a complex fft and truncate the redundant - # part of the result vector - if size(p_c) != size(x) - throw(DimensionMismatch("plan and input array axes do not match")) + # Odd length: run the full transform on the real input (the kernels + # accept real input; the DFT leaf exploits its symmetry) and keep the + # first half. + fft_kernel!(buf, x, 1, 1, FFT_FORWARD, cg[1].type, cg, 1) + @inbounds for j in 1:(n ÷ 2 + 1) + y[j] = buf[j] end - y = similar(x, Complex{T}) - fft!(y, x, 1, 1, p_c.dir, p_c.callgraph[1][1].type, p_c.callgraph[1], 1) - return y[1:end÷2+1] end + return y end -##### Backward -function Base.:*(p::FFTAPlan_re{T,1}, x::AbstractVector{T}) where {T<:Complex} - if p.dir !== FFT_BACKWARD - throw(ArgumentError("only FFT_BACKWARD supported for complex vectors")) - end - Base.require_one_based_indexing(x) - - n = p.flen - p_c = complex(p) - # See explanation of this approach in the method for the FORWARD transform +# Backward complex-to-real transform of one pencil: `y` complex of length +# `n ÷ 2 + 1`, `x` real of length `n`. +function _brfft_pencil!(x::AbstractVector{<:Real}, y::AbstractVector{T}, w::Worker{T}, n::Int) where {T<:Complex} + R = real(T) + cg = w.callgraph[1] + buf = w.buf if iseven(n) + # Inverse of the even-length trick in `_rfft_pencil!`. m = n >> 1 - - R = real(T) + 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]) + end + else + # Odd length: rebuild the conjugate-symmetric spectrum and transform. + h = n ÷ 2 + 1 + tmp = view(buf, 1:n) + out = view(buf, n + 1:2n) + @inbounds for j in 1:h + tmp[j] = y[j] + end + @inbounds for j in h + 1:n + tmp[j] = conj(y[n - j + 2]) + end + fft_kernel!(out, tmp, 1, 1, FFT_BACKWARD, cg[1].type, cg, 1) + @inbounds for j in 1:n + x[j] = real(out[j]) + end + end + return x +end - x_tmp = similar(x, length(x) - 1) - x_tmp[1] = complex( - (real(x[1]) + real(x[end])), - (real(x[1]) - real(x[end])) - ) - for j in 2:((m >> 1) + 1) - XX = x[j] + conj(x[m-j+2]) - XY = wj * (x[j] - conj(x[m-j+2])) - x_tmp[j] = XX + im * XY - x_tmp[m-j+2] = conj(XX - im * XY) - wj = singleton_step(wj, z1) +# Unit-stride pencils are handed to the kernels directly; strided ones are +# copied to the worker's contiguous buffers first. The copy costs one pass over +# the pencil but keeps the kernel's scattered reads within cache lines — for a +# 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} + n = p.flen + if R == 1 && _unit_stride(x) && _unit_stride(y) + _foreach_pencil(x, Val(R), p.workers) do w, Ipre, Ipost + @views kernel!(y[Ipre, :, Ipost], x[Ipre, :, Ipost], w, n) end + else + _foreach_pencil(x, Val(R), p.workers) do w, Ipre, Ipost + xin, yout = _pencil_buffers(kernel!, w) # per worker: no sharing between tasks + copyto!(xin, @view x[Ipre, :, Ipost]) + kernel!(yout, xin, w, n) + copyto!(@view(y[Ipre, :, Ipost]), yout) + end + end + return y +end +# forward: real in, complex out; backward: complex in, real out +_pencil_buffers(::typeof(_rfft_pencil!), w::Worker) = (w.rbuf, w.cbuf) +_pencil_buffers(::typeof(_brfft_pencil!), w::Worker) = (w.cbuf, w.rbuf) +# forward: real in, complex out; backward: complex in, real out +_pencil_buffers(::typeof(_rfft_pencil!), p::FFTAPlan_re) = (p.rbuf, p.cbuf) +_pencil_buffers(::typeof(_brfft_pencil!), p::FFTAPlan_re) = (p.cbuf, p.rbuf) + +# Dispatch on the transform dimension so that the loop above is type-stable. +function _re_along_dim!(kernel!::F, y::AbstractArray{<:Any,N}, x::AbstractArray{<:Any,N}, p::FFTAPlan_re, d::Int) where {F,N} + if @generated + quote + Base.Cartesian.@nif $N dim -> (d == dim) dim -> (_re_pencil_loop!(kernel!, y, x, p, Val(dim))) + end + else + _re_pencil_loop!(kernel!, y, x, p, Val(d)) + end +end - y_c = p_c * x_tmp - if isbitstype(T) - return copy(reinterpret(R, y_c)) - else - y_re = similar(y_c, R, 2 * length(y_c)) - for i in eachindex(y_c) - y_re[2i-1], y_re[2i] = reim(y_c[i]) - end - return y_re +function _cx_along_dim!(A::AbstractArray{<:Any,N}, workers::Vector{<:Worker{T,M}}, k::Int, dir::Direction, d::Int) where {N,T,M} + if @generated + quote + Base.Cartesian.@nif $N dim -> (d == dim) dim -> (fft_along_dim!(A, workers, k, dir, Val(dim))) end else - x_tmp = similar(x, n) - x_tmp[1:end÷2+1] .= x - x_tmp[end÷2+2:end] .= @views conj.(x[end-iseven(n):-1:2]) - y = similar(x_tmp) - LinearAlgebra.mul!(y, p_c, x_tmp) - return real(y) + fft_along_dim!(A, workers, k, dir, Val(d)) end end -#### 1D plan ND array +## mul! +#### 1D plan ##### Forward -function Base.:*(p::FFTAPlan_re{Complex{T},1}, x::AbstractArray{T,N}) where {T<:Real,N} +function LinearAlgebra.mul!(y::AbstractArray{T,N}, p::FFTAPlan_re{T,1}, x::AbstractArray{<:Real,N}) where {T<:Complex,N} if p.dir !== FFT_FORWARD - throw(ArgumentError("only FFT_FORWARD supported for real arrays")) + throw(ArgumentError("only FFT_FORWARD supported for real $(N == 1 ? "vectors" : "arrays")")) + end + Base.require_one_based_indexing(x, y) + d1 = only(p.region) + _check_re_dims(y, p, x, d1, true) + if N == 1 + _rfft_pencil!(y, x, p.workers[1], p.flen) + else + _re_along_dim!(_rfft_pencil!, y, x, p, d1) end - Base.require_one_based_indexing(x) - return mapslices(Base.Fix1(*, p), x; dims=only(p.region)) + return y end ##### Backward -function Base.:*(p::FFTAPlan_re{T,1}, x::AbstractArray{T,N}) where {T<:Complex,N} +function LinearAlgebra.mul!(y::AbstractArray{<:Real,N}, p::FFTAPlan_re{T,1}, x::AbstractArray{T,N}) where {T<:Complex,N} if p.dir !== FFT_BACKWARD - throw(ArgumentError("only FFT_BACKWARD supported for complex arrays")) + throw(ArgumentError("only FFT_BACKWARD supported for complex $(N == 1 ? "vectors" : "arrays")")) end - Base.require_one_based_indexing(x) - dim1 = only(p.region) - rlen = p.flen ÷ 2 + 1 - if rlen != size(x, dim1) - throw(DimensionMismatch("real 1D plan has size $(p.flen). Dimension of input array along region $dim1 should have size $rlen, but has size $(size(x, dim1))")) + Base.require_one_based_indexing(x, y) + d1 = only(p.region) + _check_re_dims(y, p, x, d1, false) + if N == 1 + _brfft_pencil!(y, x, p.workers[1], p.flen) + else + _re_along_dim!(_brfft_pencil!, y, x, p, d1) end - return mapslices(Base.Fix1(*, p), x; dims=dim1) + return y end -#### 2D plan ND array +#### 2D plan +# The real transform is taken along the first region dimension, then a complex +# transform along the second (forward), or the reverse (backward). ##### Forward -function Base.:*(p::FFTAPlan_re{Complex{T},2}, x::AbstractArray{T,N}) where {T<:Real,N} +function LinearAlgebra.mul!(y::AbstractArray{T,N}, p::FFTAPlan_re{T,2}, x::AbstractArray{<:Real,N}) where {T<:Complex,N} if p.dir !== FFT_FORWARD throw(ArgumentError("only FFT_FORWARD supported for real arrays")) end - Base.require_one_based_indexing(x) - half_1 = 1:(p.flen÷2+1) - x_c = complex(x) - y = similar(x_c) - LinearAlgebra.mul!(y, complex(p), x_c) - return copy(selectdim(y, first(p.region), half_1)) + Base.require_one_based_indexing(x, y) + d1, d2 = p.region + _check_re_dims(y, p, x, d1, true) + if size(x, d2) != size(p, 2) + throw(DimensionMismatch("real 2D plan has size $(size(p)). Transform dimensions of input array are $((size(x, d1), size(x, d2))) but should be $(size(p))")) + end + _re_along_dim!(_rfft_pencil!, y, x, p, d1) + _cx_along_dim!(y, p.workers, 2, FFT_FORWARD, d2) + return y end ##### Backward -function Base.:*(p::FFTAPlan_re{T,2}, x::AbstractArray{T,N}) where {T<:Complex,N} +function LinearAlgebra.mul!(y::AbstractArray{<:Real,N}, p::FFTAPlan_re{T,2}, x::AbstractArray{T,N}) where {T<:Complex,N} if p.dir !== FFT_BACKWARD throw(ArgumentError("only FFT_BACKWARD supported for complex arrays")) end - Base.require_one_based_indexing(x) - - dim1 = first(p.region) - dim2 = last(p.region) - x_sz = (xrows, xcols) = (size(x, dim1), size(x, dim2)) + Base.require_one_based_indexing(x, y) + d1, d2 = p.region + _check_re_dims(y, p, x, d1, false) + if size(x, d2) != size(p, 2) + throw(DimensionMismatch("real 2D plan has size $(size(p)). Transform dimensions of input array are $((size(x, d1), size(x, d2))) but should be $((size(p, 1) ÷ 2 + 1, size(p, 2)))")) + end + tmp = copy(x) # the complex pass must not modify the input + _cx_along_dim!(tmp, p.workers, 2, FFT_BACKWARD, d2) + _re_along_dim!(_brfft_pencil!, y, tmp, p, d1) + return y +end - flen = p.flen - tlen = flen ÷ 2 + 1 - t_sz = (tlen, size(p, 2)) +## * +function Base.:*(p::FFTAPlan_re{T,M}, x::AbstractArray{<:Real,N}) where {T<:Complex,M,N} + if p.dir !== FFT_FORWARD + throw(ArgumentError("only FFT_FORWARD supported for real $(N == 1 ? "vectors" : "arrays")")) + end + d1 = first(p.region) + y = similar(x, T, ntuple(i -> i == d1 ? p.flen ÷ 2 + 1 : size(x, i), Val(N))) + return LinearAlgebra.mul!(y, p, x) +end - if t_sz != x_sz - throw(DimensionMismatch("real 2D plan has size $(size(p)). Transform dimensions of input array are $x_sz but should be $t_sz")) +function Base.:*(p::FFTAPlan_re{T,M}, x::AbstractArray{T,N}) where {T<:Complex,M,N} + if p.dir !== FFT_BACKWARD + throw(ArgumentError("only FFT_BACKWARD supported for complex $(N == 1 ? "vectors" : "arrays")")) end + d1 = first(p.region) + y = similar(x, real(T), ntuple(i -> i == d1 ? p.flen : size(x, i), Val(N))) + return LinearAlgebra.mul!(y, p, x) +end - res_size = ntuple(i -> ifelse(i == dim1, flen, size(x, i)), Val(N)) - # for the inverse transformation we have to reconstruct the full array - half_1 = 1:tlen - half_2 = tlen+1:flen - x_full = similar(x, res_size) - # use first half as is - copy!(selectdim(x_full, dim1, half_1), x) - # the second half in the first transform dimension is reversed and conjugated - x_half_2 = selectdim(x_full, dim1, half_2) # view to the second half of x - start_reverse = xrows - iseven(flen) +# Inverse plans +# `AbstractFFTs.inv(p)` calls `plan_inv(p)` once and caches it in `p.pinv`; +# `p \\ x` and `ldiv!(y, p, x)` go through `inv(p)`. - map!(conj, x_half_2, selectdim(x, dim1, start_reverse:-1:2)) - # for the 2D transform we have to reverse index 2:end of the same block in the second transform dimension as well - reverse!(selectdim(x_half_2, dim2, 2:xcols), dims=dim2) +function AbstractFFTs.plan_inv(p::FFTAPlan_cx{T,N,R}) where {T,N,R} + dir = p.dir === FFT_FORWARD ? FFT_BACKWARD : FFT_FORWARD + cutoff = p.callgraph[1].BLUESTEIN_CUTOFF + cg = ntuple(i -> CallGraph{T}(size(p, i), cutoff), Val(N)) + q = FFTAPlan_cx{T,N,R}(cg, p.region, dir, _workers(cg, Val(N), _num_threads(p), 0)) + AbstractFFTs.ScaledPlan(q, _normalization(p)) +end - y = similar(x_full) - LinearAlgebra.mul!(y, complex(p), x_full) +function AbstractFFTs.plan_inv(p::FFTAPlan_re{T,N,R,S}) where {T,N,R,S} + dir = p.dir === FFT_FORWARD ? FFT_BACKWARD : FFT_FORWARD + cutoff = p.callgraph[1].BLUESTEIN_CUTOFF + n = p.flen + nn = iseven(n) ? n >> 1 : n + cg = (CallGraph{T}(nn, cutoff), ntuple(i -> CallGraph{T}(size(p, i + 1), cutoff), Val(N - 1))...) + q = FFTAPlan_re{T,N,R}(cg, p.region, dir, n, _workers(cg, Val(N), _num_threads(p), _re_buflen(n, dir), n)) + AbstractFFTs.ScaledPlan(q, _normalization(p)) +end + +# 1 / (product of the transform lengths); `size(p)` lists exactly those. +_normalization(p::FFTAPlan{T,N}) where {T,N} = + AbstractFFTs.normalization(real(T), size(p), ntuple(identity, Val(N))) + +# In-place plans +# (`AbstractFFTs.fft!`/`bfft!`/`ifft!` build these; the internal kernel is +# `fft_kernel!` so that it does not shadow `AbstractFFTs.fft!`.) +# FFTA's kernels are out of place, so an in-place plan wraps an ordinary plan +# and, when the input and output alias, transforms a copy of the input held +# in the plan's buffer. + +mutable struct FFTAPlan_inplace{T,N,M,P<:FFTAPlan_cx{T,M}} <: FFTAPlan{T,M} + const p::P + const buf::Array{T,N} # copy of the input when it aliases the output + pinv::AbstractFFTs.ScaledPlan + FFTAPlan_inplace(p::P, buf::Array{T,N}) where {T,N,M,P<:FFTAPlan_cx{T,M}} = new{T,N,M,P}(p, buf) +end +Base.size(p::FFTAPlan_inplace) = size(p.p) +Base.size(p::FFTAPlan_inplace, i::Int) = size(p.p, i) + +AbstractFFTs.plan_fft!(x::AbstractArray{T,N}, region; kwargs...) where {T<:Complex,N} = + FFTAPlan_inplace(_plan_fft(x, region, FFT_FORWARD; kwargs...), Array{T,N}(undef, size(x))) +AbstractFFTs.plan_bfft!(x::AbstractArray{T,N}, region; kwargs...) where {T<:Complex,N} = + FFTAPlan_inplace(_plan_fft(x, region, FFT_BACKWARD; kwargs...), Array{T,N}(undef, size(x))) + +function LinearAlgebra.mul!(y::AbstractArray{T,N}, ip::FFTAPlan_inplace{T,N}, x::AbstractArray{T,N}) where {T,N} + if y === x + if size(ip.buf) != size(x) + throw(DimensionMismatch("in-place plan was created for size $(size(ip.buf)), input has size $(size(x))")) + end + copyto!(ip.buf, x) + LinearAlgebra.mul!(y, ip.p, ip.buf) + else + LinearAlgebra.mul!(y, ip.p, x) + end + return y +end +Base.:*(ip::FFTAPlan_inplace{T}, x::AbstractArray{T}) where {T} = LinearAlgebra.mul!(x, ip, x) - return real(y) +function AbstractFFTs.plan_inv(ip::FFTAPlan_inplace) + s = AbstractFFTs.plan_inv(ip.p) + AbstractFFTs.ScaledPlan(FFTAPlan_inplace(s.p, ip.buf), s.scale) end diff --git a/test/inverse_inplace.jl b/test/inverse_inplace.jl new file mode 100644 index 0000000..eeeb648 --- /dev/null +++ b/test/inverse_inplace.jl @@ -0,0 +1,113 @@ +using FFTA, Test, LinearAlgebra + +# Inverse plans (`inv`, `\`, `ldiv!`, `plan_ifft`, `plan_irfft`) and in-place +# plans (`plan_fft!`, `plan_bfft!`, `plan_ifft!`, `fft!`, `bfft!`, `ifft!`). + +_inplace_allocs(w, p) = (mul!(w, p, w); @test_allocations mul!(w, p, w)) + +@testset "1D, n=$n, $T" for n in (1, 2, 8, 9, 15, 64, 100, 101, 1000), T in (ComplexF64, ComplexF32) + x = randn(T, n) + p = plan_fft(x) + y = p * x + @test inv(p) * y ≈ x + @test p \ y ≈ x + z = similar(x) + @test ldiv!(z, p, y) === z + @test z ≈ x + @test inv(p) === inv(p) # cached in the plan + @test inv(inv(p)) * x ≈ y + @test plan_ifft(x) * y ≈ x + pb = plan_bfft(x) + @test inv(pb) * (pb * x) ≈ x + + xr = randn(real(T), n) + pr = plan_rfft(xr) + yr = pr * xr + @test inv(pr) * yr ≈ xr + @test pr \ yr ≈ xr + zr = similar(xr) + @test ldiv!(zr, pr, yr) === zr + @test zr ≈ xr + pbr = plan_brfft(yr, n) + @test inv(pbr) * (pbr * yr) ≈ yr + @test plan_irfft(yr, n) * yr ≈ xr + + p! = plan_fft!(x) + w = copy(x) + @test (p! * w) === w + @test w ≈ y + w = copy(x) + @test mul!(w, p!, w) === w + @test w ≈ y + v = similar(x) + @test mul!(v, p!, x) ≈ y # out of place use of an in-place plan + pb! = plan_bfft!(x) + w = copy(y) + pb! * w + @test w ≈ n * x + w = copy(y) + @test (inv(p!) * w) === w + @test w ≈ x + @test inv(p!) === inv(p!) + w = copy(x) + fft!(w) + @test w ≈ y + ifft!(w) + @test w ≈ x + bfft!(w) + @test w ≈ bfft(x) + @test plan_ifft!(x) * copy(y) ≈ x + if n < FFTA.DEFAULT_BLUESTEIN_CUTOFF + @test _inplace_allocs(copy(x), p!) == 0 + end + @test_throws DimensionMismatch mul!(zeros(T, n + 1), p!, zeros(T, n + 1)) +end + +@testset "N-d complex, region $r" for r in (1, 2, 3, (1, 2), (2, 3), (1, 2, 3), 1:3, [1, 3]) + X = randn(ComplexF64, 6, 7, 8) + p = plan_fft(X, r) + Y = p * X + @test inv(p) * Y ≈ X + @test p \ Y ≈ X + Z = similar(X) + ldiv!(Z, p, Y) + @test Z ≈ X + p! = plan_fft!(X, r) + W = copy(X) + @test (p! * W) === W + @test W ≈ Y + W = copy(X) + fft!(W, r) + @test W ≈ Y + ifft!(W, r) + @test W ≈ X + bfft!(W, r) + @test W ≈ bfft(X, r) + if r isa Int # plans over several dimensions still allocate their pencil buffers per call + @test _inplace_allocs(copy(X), p!) == 0 + end +end + +@testset "N-d real, region $r" for r in (1, 2, (1, 2), (2, 3), [1, 2]) + Xr = randn(8, 6, 5) + p = plan_rfft(Xr, r) + Y = p * Xr + @test inv(p) * Y ≈ Xr + @test p \ Y ≈ Xr + Z = similar(Xr) + ldiv!(Z, p, Y) + @test Z ≈ Xr + pb = plan_brfft(Y, size(Xr, first(r)), r) + @test inv(pb) * (pb * Y) ≈ Y +end + +@testset "in-place plan, plan and inverse reused (DSP.jl pattern)" begin + x = randn(ComplexF64, 32) + p! = plan_fft!(x) + ip! = inv(p!) + @test ip!.p isa FFTA.FFTAPlan_inplace + buf = copy(x) + p! * buf + ip! * buf + @test buf ≈ x +end diff --git a/test/onedim/complex_backward.jl b/test/onedim/complex_backward.jl index f4c5785..1e2a492 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)) <= 47 * Threads.nthreads() end end diff --git a/test/onedim/complex_forward.jl b/test/onedim/complex_forward.jl index 0a7d35e..537e8f4 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 builds one worker per thread + @test (@test_allocations fft(x)) <= 47 * Threads.nthreads() end end diff --git a/test/onedim/real_backward.jl b/test/onedim/real_backward.jl index 05599aa..dc007b5 100644 --- a/test/onedim/real_backward.jl +++ b/test/onedim/real_backward.jl @@ -22,7 +22,7 @@ end end @testset "allocation regression" begin - @test (@test_allocations brfft(y, n)) <= 55 + @test (@test_allocations brfft(y, n)) <= 55 * Threads.nthreads() end end diff --git a/test/onedim/real_forward.jl b/test/onedim/real_forward.jl index 00b251b..3c65fcb 100644 --- a/test/onedim/real_forward.jl +++ b/test/onedim/real_forward.jl @@ -25,7 +25,7 @@ end end @testset "allocation regression" begin - @test (@test_allocations rfft(x)) <= 51 + @test (@test_allocations rfft(x)) <= 51 * Threads.nthreads() end end diff --git a/test/qa/aqua.jl b/test/qa/aqua.jl index 3adebe2..b612acb 100644 --- a/test/qa/aqua.jl +++ b/test/qa/aqua.jl @@ -6,6 +6,6 @@ import Aqua FFTA; # This type piracy is caused by the problematic design of AbstractFFTs.jl # Ref https://github.com/JuliaMath/AbstractFFTs.jl/issues/32 - piracies = (; treat_as_own = [plan_bfft, plan_brfft, plan_fft, plan_rfft]), + piracies = (; treat_as_own = [plan_bfft, plan_brfft, plan_fft, plan_rfft, plan_fft!, plan_bfft!]), ) end diff --git a/test/qa/explicit_imports.jl b/test/qa/explicit_imports.jl index 428672f..4d11751 100644 --- a/test/qa/explicit_imports.jl +++ b/test/qa/explicit_imports.jl @@ -18,11 +18,12 @@ import ExplicitImports @test ExplicitImports.check_all_qualified_accesses_via_owners(FFTA) === nothing # No non-public accesses in FFTA (ie. no `... MyPkg._non_public_internal_func(...)`) - # AbstractFFTs requires subtyping of `Plan` but it is not public + # AbstractFFTs requires subtyping of `Plan` and implementing `plan_inv` + # (with `ScaledPlan` and `normalization`) but none of them is public # This is an upstream bug in AbstractFFTs.jl @test ExplicitImports.check_all_qualified_accesses_are_public( FFTA; - ignore=(:Plan, :require_one_based_indexing, :Fix1, :Cartesian, :peel) + ignore=(:Plan, :plan_inv, :ScaledPlan, :normalization, :require_one_based_indexing, :Fix1, :Cartesian, :peel) ) === nothing # No self-qualified accesses in FFTA (ie. no `... FFTA.func(...)`) diff --git a/test/qa/fftw_coexistence.jl b/test/qa/fftw_coexistence.jl new file mode 100644 index 0000000..745a86f --- /dev/null +++ b/test/qa/fftw_coexistence.jl @@ -0,0 +1,10 @@ +using Test + +# Loading FFTW.jl makes its methods take over every `fft`/`rfft` call in the +# process, so the coexistence checks run in a separate Julia process to leave +# the rest of the test suite exercising FFTA. +@testset "coexistence with FFTW.jl (subprocess)" begin + body = joinpath(@__DIR__, "fftw_coexistence_body.jl") + cmd = `$(Base.julia_cmd()) --startup-file=no --project=$(Base.active_project()) $body` + @test success(pipeline(cmd; stdout = stdout, stderr = stderr)) +end diff --git a/test/qa/fftw_coexistence_body.jl b/test/qa/fftw_coexistence_body.jl new file mode 100644 index 0000000..5256ee3 --- /dev/null +++ b/test/qa/fftw_coexistence_body.jl @@ -0,0 +1,26 @@ +using FFTA, FFTW, Test + +# When FFTW.jl is loaded alongside FFTA, AbstractFFTs' design is that FFTW's +# `StridedArray` methods take over. Make sure FFTA's methods do not make the +# generic entry points ambiguous instead (this used to turn `rfft(x)` into a +# `MethodError` when `region` was annotated with `RegionTypes`). +@testset "coexistence with FFTW.jl" begin + x = randn(16) + y = randn(ComplexF64, 16) + @test rfft(x) ≈ FFTW.rfft(x) + @test irfft(rfft(x), 16) ≈ x + @test fft(y) ≈ FFTW.fft(y) + @test plan_rfft(x) isa FFTW.FFTWPlan + @test plan_brfft(rfft(x), 16) isa FFTW.FFTWPlan + @test plan_fft(y) isa FFTW.FFTWPlan + @test plan_bfft(y) isa FFTW.FFTWPlan + X = randn(8, 6) + @test rfft(X, 2) ≈ FFTW.rfft(X, 2) + @test irfft(rfft(X, 1:2), 8, 1:2) ≈ X + # FFTA's own methods are still reachable for non-strided arrays + @test rfft(view(x, [1:16;])) ≈ FFTW.rfft(x) + @test plan_rfft(view(x, [1:16;])) isa FFTA.FFTAPlan + @test isempty(filter(Test.detect_ambiguities(FFTA, FFTW)) do (m1, m2) + m1.name in (:plan_fft, :plan_bfft, :plan_rfft, :plan_brfft) + end) +end diff --git a/test/real_mul.jl b/test/real_mul.jl new file mode 100644 index 0000000..a35c629 --- /dev/null +++ b/test/real_mul.jl @@ -0,0 +1,93 @@ +using FFTA, Test, LinearAlgebra + +# `mul!` on real plans (forward: real -> complex, backward: complex -> real), +# including N-d arrays with a `dims` argument, 2D plans, views, and the +# absence of allocations for 1D plans. + +@testset "1D real plans, mul!, n=$n" for n in (1, 2, 3, 4, 8, 9, 15, 16, 31, 64, 100, 101, 1000) + x = randn(n) + p = plan_rfft(x) + y = p * x + y2 = similar(y) + @test mul!(y2, p, x) === y2 + @test y2 == y + @test y ≈ naive_1d_fourier_transform(x, FFTA.FFT_FORWARD)[1:(n ÷ 2 + 1)] + pb = plan_brfft(y, n) + xb = similar(x) + @test mul!(xb, pb, y) === xb + @test xb ≈ n * x + @test xb == pb * y + if n < FFTA.DEFAULT_BLUESTEIN_CUTOFF # Bluestein still allocates scratch per call + @test (@test_allocations mul!(y2, p, x)) == 0 + @test (@test_allocations mul!(xb, pb, y)) == 0 + end + # views as input and output + X = randn(n, 3) + Y = zeros(ComplexF64, n ÷ 2 + 1, 3) + mul!(view(Y, :, 2), p, view(X, :, 2)) + @test Y[:, 2] ≈ rfft(X[:, 2]) + Xb = zeros(n, 3) + mul!(view(Xb, :, 3), pb, view(Y, :, 2)) + @test Xb[:, 3] ≈ n * X[:, 2] + # wrong sizes + @test_throws DimensionMismatch mul!(similar(y, n ÷ 2 + 2), p, x) + @test_throws DimensionMismatch mul!(similar(x, n + 1), pb, y) + @test_throws ArgumentError mul!(similar(x), p, y) + @test_throws ArgumentError mul!(similar(y), pb, x) +end + +@testset "1D real plans on N-d arrays, mul!, size $sz" for sz in ((7, 4), (8, 5), (6, 7, 8), (5, 8, 9)) + x = randn(sz) + for d in 1:length(sz) + n = size(x, d) + p = plan_rfft(x, d) + y = p * x + y2 = similar(y) + @test mul!(y2, p, x) == y + @test y ≈ mapslices(rfft, x; dims = d) + @test (@test_allocations mul!(y2, p, x)) == 0 + pb = plan_brfft(y, n, d) + xb = similar(x) + @test mul!(xb, pb, y) ≈ n * x + @test (@test_allocations mul!(xb, pb, y)) == 0 + @test_throws DimensionMismatch mul!(similar(x, ntuple(i -> i == d ? n + 1 : sz[i], length(sz))), pb, y) + end +end + +@testset "2D real plans, mul!, size $sz" for sz in ((8, 6), (9, 6), (8, 7), (9, 7), (1, 1), (2, 3), (64, 64), (100, 101)) + x = randn(sz) + p = plan_rfft(x) + y = p * x + @test y ≈ naive_2d_fourier_transform(x, FFTA.FFT_FORWARD)[1:(sz[1] ÷ 2 + 1), :] + y2 = similar(y) + @test mul!(y2, p, x) == y + pb = plan_brfft(y, sz[1]) + xb = similar(x) + @test mul!(xb, pb, y) ≈ prod(sz) * x + @test xb == pb * y + @test_throws DimensionMismatch mul!(similar(y, sz[1] ÷ 2 + 1, sz[2] + 1), p, x) + @test_throws DimensionMismatch mul!(xb, pb, similar(y, sz[1] ÷ 2 + 1, sz[2] + 1)) +end + +@testset "2D real plans on 3D arrays, mul!, region $r" for r in ((1, 2), (1, 3), (2, 3), [1, 2], 2:3) + x = randn(6, 7, 8) + p = plan_rfft(x, r) + y = p * x + @test y ≈ mapslices(rfft, x; dims = r) + y2 = similar(y) + @test mul!(y2, p, x) == y + n1 = size(x, first(r)) + pb = plan_brfft(y, n1, r) + xb = similar(x) + @test mul!(xb, pb, y) ≈ prod(size(x, i) for i in r) * x +end + +@testset "Float32 real plans" begin + x = randn(Float32, 48) + y = rfft(x) + @test eltype(y) == ComplexF32 + @test y ≈ rfft(Float64.(x)) + @test irfft(y, 48) ≈ x + X = randn(Float32, 12, 10) + @test irfft(rfft(X), 12) ≈ X +end diff --git a/test/runtests.jl b/test/runtests.jl index 9aeb1e6..fc3a9d2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -46,6 +46,7 @@ Random.seed!(1) @testset verbose = true "QA" begin include("qa/aqua.jl") include("qa/explicit_imports.jl") + include("qa/fftw_coexistence.jl") end @testset verbose = true "Argument checking" begin include("argument_checking.jl") @@ -89,6 +90,15 @@ Random.seed!(1) end end end + @testset verbose = true "Real plans: mul!" begin + include("real_mul.jl") + end + @testset verbose = true "Inverse and in-place plans" begin + include("inverse_inplace.jl") + end + @testset verbose = true "Workers and threads" begin + include("threads.jl") + end @testset verbose = true "N-D" begin @testset verbose = true "Minimal tests" begin include("ndim/minimal_complex.jl") diff --git a/test/threads.jl b/test/threads.jl new file mode 100644 index 0000000..865ba90 --- /dev/null +++ b/test/threads.jl @@ -0,0 +1,67 @@ +using FFTA, Test, LinearAlgebra + +# Plans own one worker (call-graph workspace + pencil buffers) per thread they +# may use; pencils of multidimensional and batched transforms are distributed +# over them. Results must not depend on the worker count. The suite normally +# runs with one Julia thread, in which case the parallel tasks simply run one +# after the other; the chunking logic is exercised either way. + +@testset "num_threads keyword" begin + x = randn(ComplexF64, 16) + @test length(plan_fft(x; num_threads = 1).workers) == 1 + @test length(plan_fft(x; num_threads = 3).workers) == 3 + @test length(plan_fft(x).workers) == Threads.nthreads() + @test length(plan_rfft(randn(16); num_threads = 2).workers) == 2 + @test_throws ArgumentError plan_fft(x; num_threads = 0) + p = plan_fft(x; num_threads = 2) + @test length(inv(p).p.workers) == 2 # inverse keeps the worker count + @test p.workers[1].callgraph[1] === p.callgraph[1] + @test p.workers[2].callgraph[1] !== p.callgraph[1] + @test p.workers[2].callgraph[1].nodes === p.callgraph[1].nodes # nodes are shared +end + +# sizes above and below THREAD_THRESHOLD, contiguous and strided pencils, +# composite lengths (which use the call-graph workspace) +@testset "results independent of worker count, size $sz, region $r" for sz in ((64, 64), (513, 40), (30, 700, 4), (300, 300), (8, 8, 8)), + r in (1, 2, (1, 2)) + X = randn(ComplexF64, sz) + p1 = plan_fft(X, r; num_threads = 1) + Y1 = p1 * X + @test Y1 ≈ mapslices(fft, X; dims = r) + for nt in (2, 3, 5) + p = plan_fft(X, r; num_threads = nt) + @test p * X == Y1 + Y = similar(Y1) + @test mul!(Y, p, X) == Y1 + @test inv(p) * Y1 ≈ X + pb = plan_bfft(X, r; num_threads = nt) + @test pb * X == plan_bfft(X, r; num_threads = 1) * X + end + Xr = randn(Float64, sz) + Yr1 = plan_rfft(Xr, r; num_threads = 1) * Xr + @test Yr1 ≈ mapslices(rfft, Xr; dims = r) + for nt in (2, 3) + pr = plan_rfft(Xr, r; num_threads = nt) + @test pr * Xr == Yr1 + @test inv(pr) * Yr1 ≈ Xr + pbr = plan_brfft(Yr1, size(Xr, first(r)), r; num_threads = nt) + @test pbr * Yr1 == plan_brfft(Yr1, size(Xr, first(r)), r; num_threads = 1) * Yr1 + end +end + +@testset "single-worker execution does not allocate" begin + X = randn(ComplexF64, 64, 64) + for r in (1, 2, (1, 2)) + p = plan_fft(X, r; num_threads = 1) + Y = p * X + mul!(Y, p, X) + @test (@test_allocations mul!(Y, p, X)) == 0 + end + Xr = randn(64, 64) + for r in (1, 2) + p = plan_rfft(Xr, r; num_threads = 1) + Y = p * Xr + mul!(Y, p, Xr) + @test (@test_allocations mul!(Y, p, Xr)) == 0 + end +end diff --git a/test/twodim/complex_forward.jl b/test/twodim/complex_forward.jl index 8e850d8..3b9d7af 100644 --- a/test/twodim/complex_forward.jl +++ b/test/twodim/complex_forward.jl @@ -21,7 +21,7 @@ end end @testset "allocations" begin - @test (@test_allocations fft(X)) <= 116 + @test (@test_allocations fft(X)) <= 116 * Threads.nthreads() end end end diff --git a/test/twodim/real_backward.jl b/test/twodim/real_backward.jl index 7f87faa..de5056a 100644 --- a/test/twodim/real_backward.jl +++ b/test/twodim/real_backward.jl @@ -19,7 +19,7 @@ end end @testset "allocations" begin - @test (@test_allocations brfft(Y, m)) <= 12050 + @test (@test_allocations brfft(Y, m)) <= 12050 * Threads.nthreads() end end end @@ -41,5 +41,5 @@ end X = randn(256, 256) Y = rfft(X) brfft(Y, 256) # compile - @test (@test_allocations brfft(Y, 256)) <= 82 + @test (@test_allocations brfft(Y, 256)) <= 82 * Threads.nthreads() end diff --git a/test/twodim/real_forward.jl b/test/twodim/real_forward.jl index 35d71c9..c01acff 100644 --- a/test/twodim/real_forward.jl +++ b/test/twodim/real_forward.jl @@ -22,7 +22,7 @@ end end @testset "allocations" begin - @test (@test_allocations rfft(X)) <= 132 + @test (@test_allocations rfft(X)) <= 132 * Threads.nthreads() end end end @@ -38,5 +38,5 @@ end @testset "allocations" begin X = randn(256, 256) rfft(X) # compile - @test (@test_allocations rfft(X)) <= 63 + @test (@test_allocations rfft(X)) <= 63 * Threads.nthreads() end