From 2f58ee09e5d8250fd63cd02b30fc1c8b7592f740 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 18:30:07 +0000 Subject: [PATCH 1/7] Fix method ambiguity with FFTW.jl for plan_rfft/plan_brfft With FFTW.jl loaded alongside FFTA, plan_rfft(::Vector{Float64}, ::Int) was ambiguous between FFTW's StridedArray method and FFTA's method annotated with region::RegionTypes, turning rfft(x) into a MethodError. Leave region unannotated on the AbstractFFTs entry points (as plan_fft already does) and normalise it in an internal function, so FFTW's methods are strictly more specific and take over as AbstractFFTs intends. A coexistence test runs in a subprocess (loading FFTW in the test process would make every other test exercise FFTW). --- Project.toml | 4 +++- src/plan.jl | 19 +++++++++++++++++-- test/qa/fftw_coexistence.jl | 10 ++++++++++ test/qa/fftw_coexistence_body.jl | 26 ++++++++++++++++++++++++++ test/runtests.jl | 1 + 5 files changed, 57 insertions(+), 3 deletions(-) create mode 100644 test/qa/fftw_coexistence.jl create mode 100644 test/qa/fftw_coexistence_body.jl 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/src/plan.jl b/src/plan.jl index a71cbc8..92553f3 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -99,7 +99,22 @@ function _plan_fft( 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... @@ -124,7 +139,7 @@ function AbstractFFTs.plan_rfft( end end -function AbstractFFTs.plan_brfft( +function _plan_brfft( x::AbstractArray{T,N}, len::Int, region::RegionTypes; 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/runtests.jl b/test/runtests.jl index 9aeb1e6..5d3b156 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") From c9d0a991eacced06193f9a8cc46e5f78657cc933 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:08:59 +0000 Subject: [PATCH 2/7] Real plans: mul!, allocation-free execution, dims path, real 2D path Real-input/real-output plans only implemented *, so every rfft/irfft allocated, mul!(y, p, x) with a preallocated output was a MethodError, rfft along one dimension of an N-d array went through mapslices, and 2D real plans ran a full complex transform and discarded half of it. FFTAPlan_re now carries a scratch buffer and two pencil kernels (_rfft_pencil!/_brfft_pencil!) implement the even-length half-size trick and the odd-length transform on AbstractVectors, so views work as input and output. mul! is defined for 1D plans on 1D and N-d arrays (looping over pencils along the region) and for 2D plans on N-d arrays (real transform along the first region dimension, complex along the second), and * allocates the output and calls mul!. The 2D plan's first call graph is built for the half length like the 1D plan's. --- src/plan.jl | 364 +++++++++++++++++++++++++++-------------------- test/real_mul.jl | 93 ++++++++++++ test/runtests.jl | 3 + 3 files changed, 309 insertions(+), 151 deletions(-) create mode 100644 test/real_mul.jl diff --git a/src/plan.jl b/src/plan.jl index 92553f3..539627f 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -24,13 +24,25 @@ struct FFTAPlan_re{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} region::R dir::Direction flen::Int + buf::Vector{T} # scratch for the real<->complex packing, see `_re_buflen` pinv::FFTAInvPlan{T,N} end function FFTAPlan_re{T,N}( cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, flen::Int ) where {T,N,R<:RegionTypes{N}} - FFTAPlan_re{T,N,R}(cg, r, dir, flen, FFTAInvPlan{T,N}()) + buf = Vector{T}(undef, _re_buflen(flen, dir)) + FFTAPlan_re{T,N,R}(cg, r, dir, flen, buf, FFTAInvPlan{T,N}()) +end + +# 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 +58,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) @@ -131,9 +141,11 @@ function _plan_rfft( return FFTAPlan_re{Complex{T},1}((g,), R1, FFT_FORWARD, n) 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) else throw(ArgumentError("only supports 1D and 2D FFTs")) end @@ -156,7 +168,8 @@ function _plan_brfft( return FFTAPlan_re{T,1}((g,), R1, FFT_BACKWARD, len) 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) else @@ -399,200 +412,249 @@ 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) + 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}, p::FFTAPlan_re{T}) where {T<:Complex} n = p.flen - p_c = complex(p) + R = real(T) + cg = p.callgraph[1] + buf = p.buf 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)) - end - + # 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!(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!(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) - +# 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}, p::FFTAPlan_re{T}) where {T<:Complex} n = p.flen - p_c = complex(p) - # See explanation of this approach in the method for the FORWARD transform + R = real(T) + cg = p.callgraph[1] + buf = p.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!(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!(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) +# Apply `kernel!(y_pencil, x_pencil, p)` 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} + Rpre = CartesianIndices(ntuple(Base.Fix1(size, x), Val(R - 1))) + Rpost = CartesianIndices(ntuple(i -> size(x, R + i), Val(N - R))) + for Ipost in Rpost, Ipre in Rpre + @views kernel!(y[Ipre, :, Ipost], x[Ipre, :, Ipost], p) + end + return y +end + +# 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}, cg::CallGraph{T}, dir::Direction, d::Int) where {N,T} + n = size(A, d) + ibuf = Vector{T}(undef, n) + obuf = Vector{T}(undef, n) + if @generated + quote + Base.Cartesian.@nif $N dim -> (d == dim) dim -> (fft_along_dim!(A, ibuf, obuf, cg, 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, ibuf, obuf, cg, 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) + 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) + 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.callgraph[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)) - - flen = p.flen - tlen = flen ÷ 2 + 1 - t_sz = (tlen, size(p, 2)) - - 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")) + 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.callgraph[2], FFT_BACKWARD, d2) + _re_along_dim!(_brfft_pencil!, y, tmp, p, d1) + return y +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) - - 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) - - y = similar(x_full) - LinearAlgebra.mul!(y, complex(p), x_full) +## * +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 - return real(y) +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 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 5d3b156..fbf20d1 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -90,6 +90,9 @@ Random.seed!(1) end end end + @testset verbose = true "Real plans: mul!" begin + include("real_mul.jl") + end @testset verbose = true "N-D" begin @testset verbose = true "Minimal tests" begin include("ndim/minimal_complex.jl") From 2d9697a7d6d185fb7d1f294007fcb7c475efda50 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:20:07 +0000 Subject: [PATCH 3/7] Inverse plans (inv, \, ldiv!) and in-place plans (plan_fft!, plan_bfft!) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit inv(p) threw a TypeError for every FFTA plan: no plan_inv method existed and the dummy pinv::FFTAInvPlan field made AbstractFFTs' pinv_type resolve to Union{}. The plan structs are now mutable with an initially undefined pinv field, as in FFTW.jl, so AbstractFFTs.inv caches the result of the new plan_inv methods for complex and real plans; p \ x, ldiv!(y, p, x), plan_ifft and plan_irfft work through them. plan_fft!/plan_bfft! return an FFTAPlan_inplace wrapping an ordinary plan plus a buffer: when input and output alias, the input is copied to the buffer and transformed out of place (FFTA's kernels are out of place, and the 1D pencil path is not alias-safe); otherwise the wrapped plan is used directly. fft!, bfft! and ifft! from AbstractFFTs now work — the internal kernel that shadowed AbstractFFTs.fft! is renamed fft_kernel!. --- docs/src/dev.md | 2 +- src/algos.jl | 6 +- src/plan.jl | 118 ++++++++++++++++++++++++++++++++-------- test/inverse_inplace.jl | 111 +++++++++++++++++++++++++++++++++++++ test/runtests.jl | 3 + 5 files changed, 213 insertions(+), 27 deletions(-) create mode 100644 test/inverse_inplace.jl 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 539627f..c7d9ae5 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -2,37 +2,42 @@ 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} - callgraph::NTuple{N,CallGraph{T}} - region::R - dir::Direction - pinv::FFTAInvPlan{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. + +mutable struct FFTAPlan_cx{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} + const callgraph::NTuple{N,CallGraph{T}} + const region::R + const dir::Direction + pinv::AbstractFFTs.ScaledPlan + FFTAPlan_cx{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction) where {T,N,R<:RegionTypes{N}} = + new{T,N,R}(cg, r, dir) end function FFTAPlan_cx{T,N}( cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction ) 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) end -struct FFTAPlan_re{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} - callgraph::NTuple{N,CallGraph{T}} - region::R - dir::Direction - flen::Int - buf::Vector{T} # scratch for the real<->complex packing, see `_re_buflen` - pinv::FFTAInvPlan{T,N} +mutable struct FFTAPlan_re{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} + const callgraph::NTuple{N,CallGraph{T}} + const region::R + const dir::Direction + const flen::Int + const buf::Vector{T} # scratch for the real<->complex packing, see `_re_buflen` + pinv::AbstractFFTs.ScaledPlan + FFTAPlan_re{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, flen::Int, buf::Vector{T}) where {T,N,R<:RegionTypes{N}} = + new{T,N,R}(cg, r, dir, flen, buf) end function FFTAPlan_re{T,N}( cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, flen::Int ) where {T,N,R<:RegionTypes{N}} buf = Vector{T}(undef, _re_buflen(flen, dir)) - FFTAPlan_re{T,N,R}(cg, r, dir, flen, buf, FFTAInvPlan{T,N}()) + FFTAPlan_re{T,N,R}(cg, r, dir, flen, buf) end # Scratch length needed by the real-transform pencil kernels for a plan of @@ -190,7 +195,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 @@ -230,7 +235,7 @@ function _mul_loop!( 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) + @views fft_kernel!(y[Ipre,:,Ipost], x[Ipre,:,Ipost], 1, 1, p.dir, t, cg, 1) end end @@ -368,7 +373,7 @@ function fft_along_dim!( for j in cols ibuf[j] = A[Ipre, j, Ipost] end - fft!(obuf, ibuf, 1, 1, d, t, cg, 1) + fft_kernel!(obuf, ibuf, 1, 1, d, t, cg, 1) for j in cols A[Ipre, j, Ipost] = obuf[j] end @@ -455,7 +460,7 @@ function _rfft_pencil!(y::AbstractVector{T}, x::AbstractVector{<:Real}, p::FFTAP @inbounds for j in 1:m buf[j] = T(x[2j - 1], x[2j]) end - fft!(view(y, 1:m), buf, 1, 1, FFT_FORWARD, cg[1].type, cg, 1) + 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, @@ -480,7 +485,7 @@ function _rfft_pencil!(y::AbstractVector{T}, x::AbstractVector{<:Real}, p::FFTAP # 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!(buf, x, 1, 1, FFT_FORWARD, cg[1].type, cg, 1) + 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 @@ -512,7 +517,7 @@ function _brfft_pencil!(x::AbstractVector{<:Real}, y::AbstractVector{T}, p::FFTA wj = singleton_step(wj, z1) end end - fft!(out, tmp, 1, 1, FFT_BACKWARD, cg[1].type, cg, 1) + 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]) @@ -528,7 +533,7 @@ function _brfft_pencil!(x::AbstractVector{<:Real}, y::AbstractVector{T}, p::FFTA @inbounds for j in h + 1:n tmp[j] = conj(y[n - j + 2]) end - fft!(out, tmp, 1, 1, FFT_BACKWARD, cg[1].type, cg, 1) + 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 @@ -658,3 +663,70 @@ function Base.:*(p::FFTAPlan_re{T,M}, x::AbstractArray{T,N}) where {T<:Complex,M y = similar(x, real(T), ntuple(i -> i == d1 ? p.flen : size(x, i), Val(N))) return LinearAlgebra.mul!(y, p, x) end + + +# 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)`. + +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) + AbstractFFTs.ScaledPlan(q, _normalization(p)) +end + +function AbstractFFTs.plan_inv(p::FFTAPlan_re{T,N,R}) where {T,N,R} + 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, Vector{T}(undef, _re_buflen(n, dir))) + 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) + +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..2b04fba --- /dev/null +++ b/test/inverse_inplace.jl @@ -0,0 +1,111 @@ +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) + @test _inplace_allocs(copy(X), p!) == 0 +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/runtests.jl b/test/runtests.jl index fbf20d1..a107216 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -93,6 +93,9 @@ Random.seed!(1) @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 "N-D" begin @testset verbose = true "Minimal tests" begin include("ndim/minimal_complex.jl") From a153e0c8027dc73c2b337ac435545998c2a751e3 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:23:05 +0000 Subject: [PATCH 4/7] test: only require allocation-free in-place execution for 1D plans; allow AbstractFFTs backend hooks in the ExplicitImports check --- test/inverse_inplace.jl | 4 +++- test/qa/explicit_imports.jl | 5 +++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/test/inverse_inplace.jl b/test/inverse_inplace.jl index 2b04fba..eeeb648 100644 --- a/test/inverse_inplace.jl +++ b/test/inverse_inplace.jl @@ -83,7 +83,9 @@ end @test W ≈ X bfft!(W, r) @test W ≈ bfft(X, r) - @test _inplace_allocs(copy(X), p!) == 0 + 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]) 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(...)`) From a0d37e6350319b9a1b5e495eda8194ac7f6b8423 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:32:58 +0000 Subject: [PATCH 5/7] test: treat plan_fft!/plan_bfft! as own in the Aqua piracy check, like the other AbstractFFTs entry points --- test/qa/aqua.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) 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 From 17ed942580cb106370c9decb6656d754c5a51dd6 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 19:57:45 +0000 Subject: [PATCH 6/7] Plan-owned per-thread workers, direct strided pencil reads, threading across pencils Multidimensional mul! allocated two pencil buffers per call and copied every pencil in and out of them; a 1D plan applied along one dimension of an array transformed its pencils one after the other; and the single call-graph workspace of a plan made it impossible to use threads. A plan now owns one Worker per thread it may use (num_threads keyword of the plan_* functions, default Threads.nthreads(), kept by inv): call graphs sharing the plan's nodes but with their own workspace, plus the pencil buffers. One driver, _foreach_pencil, runs every pencil loop (1D plans on N-d arrays, N-d plans, real pencil kernels): serially and allocation-free below THREAD_THRESHOLD or with one worker, otherwise as one Threads.@spawn task per worker over a contiguous chunk of pencils. Each chunk uses its own worker, so results are bit-identical whatever the thread count. Kernels read the strided pencil views directly; only the copy-out from the worker's output buffer remains. --- src/plan.jl | 252 +++++++++++++++++--------------- test/onedim/complex_backward.jl | 2 +- test/onedim/complex_forward.jl | 3 +- test/onedim/real_backward.jl | 2 +- test/onedim/real_forward.jl | 2 +- test/qa/aqua.jl | 2 +- test/runtests.jl | 3 + test/threads.jl | 67 +++++++++ test/twodim/complex_forward.jl | 2 +- test/twodim/real_backward.jl | 4 +- test/twodim/real_forward.jl | 4 +- 11 files changed, 219 insertions(+), 124 deletions(-) create mode 100644 test/threads.jl diff --git a/src/plan.jl b/src/plan.jl index c7d9ae5..a1d3356 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -7,19 +7,51 @@ const RegionTypes{N} = Union{Int,AbstractVector{Int},NTuple{N,Int}} # 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} + callgraph::NTuple{N,CallGraph{T}} + 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) +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) + +function _workers(cg::Tuple{CallGraph{T},Vararg{CallGraph{T}}}, ::Val{N}, num_threads::Int, buflen::Int) where {T,N} + num_threads >= 1 || throw(ArgumentError("num_threads must be at least 1")) + obuflen = maximum(g -> first(g.nodes).sz, cg) + workers = [Worker{T,N}(cg, Vector{T}(undef, obuflen), Vector{T}(undef, buflen))] + for _ in 2:num_threads + push!(workers, Worker{T,N}(map(_clone_workspace, cg), Vector{T}(undef, obuflen), Vector{T}(undef, buflen))) + 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}} <: FFTAPlan{T,N} const callgraph::NTuple{N,CallGraph{T}} const region::R const dir::Direction + const workers::Vector{Worker{T,N}} pinv::AbstractFFTs.ScaledPlan - FFTAPlan_cx{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction) where {T,N,R<:RegionTypes{N}} = - new{T,N,R}(cg, r, dir) + FFTAPlan_cx{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, workers::Vector{Worker{T,N}}) where {T,N,R<:RegionTypes{N}} = + new{T,N,R}(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) + FFTAPlan_cx{T,N,R}(cg, r, dir, _workers(cg, Val(N), num_threads, 0)) end mutable struct FFTAPlan_re{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} @@ -27,19 +59,20 @@ mutable struct FFTAPlan_re{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} const region::R const dir::Direction const flen::Int - const buf::Vector{T} # scratch for the real<->complex packing, see `_re_buflen` + const workers::Vector{Worker{T,N}} pinv::AbstractFFTs.ScaledPlan - FFTAPlan_re{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, flen::Int, buf::Vector{T}) where {T,N,R<:RegionTypes{N}} = - new{T,N,R}(cg, r, dir, flen, buf) + FFTAPlan_re{T,N,R}(cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, flen::Int, workers::Vector{Worker{T,N}}) where {T,N,R<:RegionTypes{N}} = + new{T,N,R}(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}} - buf = Vector{T}(undef, _re_buflen(flen, dir)) - FFTAPlan_re{T,N,R}(cg, r, dir, flen, buf) + FFTAPlan_re{T,N,R}(cg, r, dir, flen, _workers(cg, Val(N), num_threads, _re_buflen(flen, dir))) 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) @@ -93,23 +126,23 @@ 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 @@ -132,7 +165,7 @@ _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 @@ -143,14 +176,14 @@ function _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) 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, n) + return FFTAPlan_re{Complex{T},2}((g1, g2), R2, FFT_FORWARD, n; num_threads) else throw(ArgumentError("only supports 1D and 2D FFTs")) end @@ -160,7 +193,7 @@ 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 @@ -170,13 +203,13 @@ function _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) 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 @@ -224,18 +257,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_kernel!(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 @@ -254,35 +321,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 @@ -304,78 +354,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_kernel!(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 @@ -444,11 +474,10 @@ 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}, p::FFTAPlan_re{T}) where {T<:Complex} - n = p.flen +function _rfft_pencil!(y::AbstractVector{T}, x::AbstractVector{<:Real}, w::Worker{T}, n::Int) where {T<:Complex} R = real(T) - cg = p.callgraph[1] - buf = p.buf + 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 @@ -495,11 +524,10 @@ end # 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}, p::FFTAPlan_re{T}) where {T<:Complex} - n = p.flen +function _brfft_pencil!(x::AbstractVector{<:Real}, y::AbstractVector{T}, w::Worker{T}, n::Int) where {T<:Complex} R = real(T) - cg = p.callgraph[1] - buf = p.buf + cg = w.callgraph[1] + buf = w.buf if iseven(n) # Inverse of the even-length trick in `_rfft_pencil!`. m = n >> 1 @@ -541,12 +569,11 @@ function _brfft_pencil!(x::AbstractVector{<:Real}, y::AbstractVector{T}, p::FFTA return x end -# Apply `kernel!(y_pencil, x_pencil, p)` along dimension `R` of `x` and `y`. +# 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} - Rpre = CartesianIndices(ntuple(Base.Fix1(size, x), Val(R - 1))) - Rpost = CartesianIndices(ntuple(i -> size(x, R + i), Val(N - R))) - for Ipost in Rpost, Ipre in Rpre - @views kernel!(y[Ipre, :, Ipost], x[Ipre, :, Ipost], p) + n = p.flen + _foreach_pencil(x, Val(R), p.workers) do w, Ipre, Ipost + @views kernel!(y[Ipre, :, Ipost], x[Ipre, :, Ipost], w, n) end return y end @@ -562,16 +589,13 @@ function _re_along_dim!(kernel!::F, y::AbstractArray{<:Any,N}, x::AbstractArray{ end end -function _cx_along_dim!(A::AbstractArray{<:Any,N}, cg::CallGraph{T}, dir::Direction, d::Int) where {N,T} - n = size(A, d) - ibuf = Vector{T}(undef, n) - obuf = Vector{T}(undef, n) +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, ibuf, obuf, cg, dir, Val(dim))) + Base.Cartesian.@nif $N dim -> (d == dim) dim -> (fft_along_dim!(A, workers, k, dir, Val(dim))) end else - fft_along_dim!(A, ibuf, obuf, cg, dir, Val(d)) + fft_along_dim!(A, workers, k, dir, Val(d)) end end @@ -586,7 +610,7 @@ function LinearAlgebra.mul!(y::AbstractArray{T,N}, p::FFTAPlan_re{T,1}, x::Abstr d1 = only(p.region) _check_re_dims(y, p, x, d1, true) if N == 1 - _rfft_pencil!(y, x, p) + _rfft_pencil!(y, x, p.workers[1], p.flen) else _re_along_dim!(_rfft_pencil!, y, x, p, d1) end @@ -602,7 +626,7 @@ function LinearAlgebra.mul!(y::AbstractArray{<:Real,N}, p::FFTAPlan_re{T,1}, x:: d1 = only(p.region) _check_re_dims(y, p, x, d1, false) if N == 1 - _brfft_pencil!(y, x, p) + _brfft_pencil!(y, x, p.workers[1], p.flen) else _re_along_dim!(_brfft_pencil!, y, x, p, d1) end @@ -624,7 +648,7 @@ function LinearAlgebra.mul!(y::AbstractArray{T,N}, p::FFTAPlan_re{T,2}, x::Abstr 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.callgraph[2], FFT_FORWARD, d2) + _cx_along_dim!(y, p.workers, 2, FFT_FORWARD, d2) return y end @@ -640,7 +664,7 @@ function LinearAlgebra.mul!(y::AbstractArray{<:Real,N}, p::FFTAPlan_re{T,2}, x:: 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.callgraph[2], FFT_BACKWARD, d2) + _cx_along_dim!(tmp, p.workers, 2, FFT_BACKWARD, d2) _re_along_dim!(_brfft_pencil!, y, tmp, p, d1) return y end @@ -673,7 +697,7 @@ 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) + q = FFTAPlan_cx{T,N,R}(cg, p.region, dir, _workers(cg, Val(N), _num_threads(p), 0)) AbstractFFTs.ScaledPlan(q, _normalization(p)) end @@ -683,7 +707,7 @@ function AbstractFFTs.plan_inv(p::FFTAPlan_re{T,N,R}) where {T,N,R} 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, Vector{T}(undef, _re_buflen(n, dir))) + q = FFTAPlan_re{T,N,R}(cg, p.region, dir, n, _workers(cg, Val(N), _num_threads(p), _re_buflen(n, dir))) AbstractFFTs.ScaledPlan(q, _normalization(p)) 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/runtests.jl b/test/runtests.jl index a107216..fc3a9d2 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -96,6 +96,9 @@ Random.seed!(1) @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 From 600d633a349bb815b2f12726ebea3279260385e8 Mon Sep 17 00:00:00 2001 From: Panagiotis Georgakopoulos Date: Sat, 29 Aug 2026 22:52:54 +0000 Subject: [PATCH 7/7] Real plans: copy strided pencils to contiguous buffers before the kernels Replacing mapslices with strided views made real transforms along dims=2 of wide matrices (64 x N) 1.2-1.3x slower on x86-64: the mapslices copy had been an unlabelled copy-in that turned a stride-of-a-cache-line gather into one contiguous pass before the kernel. Pencils whose parent arrays are unit-stride along the transform dimension still go to the kernels directly (the dims=1 gain stays); any other pencil is copied to two plan-owned contiguous buffers first and copied back after, so execution stays allocation-free and the dims=2 result is now identical to the mapslices one. --- src/plan.jl | 35 ++++++++++++++++++++++++++++++----- 1 file changed, 30 insertions(+), 5 deletions(-) diff --git a/src/plan.jl b/src/plan.jl index 539627f..1237e77 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -19,20 +19,23 @@ function FFTAPlan_cx{T,N}( FFTAPlan_cx{T,N,R}(cg, r, dir, FFTAInvPlan{T,N}()) end -struct FFTAPlan_re{T,N,R<:RegionTypes{N}} <: FFTAPlan{T,N} +struct FFTAPlan_re{T,N,R<:RegionTypes{N},S<:Real} <: FFTAPlan{T,N} callgraph::NTuple{N,CallGraph{T}} region::R dir::Direction flen::Int - buf::Vector{T} # scratch for the real<->complex packing, see `_re_buflen` + buf::Vector{T} # scratch for the real<->complex packing, see `_re_buflen` + rbuf::Vector{S} # contiguous copy of a strided real pencil, see `_re_pencil_loop!` + cbuf::Vector{T} # contiguous copy of a strided complex pencil pinv::FFTAInvPlan{T,N} end function FFTAPlan_re{T,N}( cg::NTuple{N,CallGraph{T}}, r::R, dir::Direction, flen::Int ) where {T,N,R<:RegionTypes{N}} + S = real(T) buf = Vector{T}(undef, _re_buflen(flen, dir)) - FFTAPlan_re{T,N,R}(cg, r, dir, flen, buf, FFTAInvPlan{T,N}()) + FFTAPlan_re{T,N,R,S}(cg, r, dir, flen, buf, Vector{S}(undef, flen), Vector{T}(undef, flen ÷ 2 + 1), FFTAInvPlan{T,N}()) end # Scratch length needed by the real-transform pencil kernels for a plan of @@ -536,15 +539,37 @@ function _brfft_pencil!(x::AbstractVector{<:Real}, y::AbstractVector{T}, p::FFTA return x end +# Unit-stride pencils are handed to the kernels directly; strided ones are +# copied to the plan'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. +_unit_stride(v::StridedArray) = stride(v, 1) == 1 +_unit_stride(v) = false + # Apply `kernel!(y_pencil, x_pencil, p)` 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} Rpre = CartesianIndices(ntuple(Base.Fix1(size, x), Val(R - 1))) Rpost = CartesianIndices(ntuple(i -> size(x, R + i), Val(N - R))) - for Ipost in Rpost, Ipre in Rpre - @views kernel!(y[Ipre, :, Ipost], x[Ipre, :, Ipost], p) + contiguous = R == 1 && _unit_stride(x) && _unit_stride(y) + if contiguous + for Ipost in Rpost, Ipre in Rpre + @views kernel!(y[Ipre, :, Ipost], x[Ipre, :, Ipost], p) + end + else + xin, yout = _pencil_buffers(kernel!, p) + for Ipost in Rpost, Ipre in Rpre + xv = @view x[Ipre, :, Ipost] + yv = @view y[Ipre, :, Ipost] + copyto!(xin, xv) + kernel!(yout, xin, p) + copyto!(yv, yout) + end end return y end +# 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}