diff --git a/docs/src/explanation/design-decisions.md b/docs/src/explanation/design-decisions.md index 50a4d52c..ca1600e2 100644 --- a/docs/src/explanation/design-decisions.md +++ b/docs/src/explanation/design-decisions.md @@ -106,6 +106,27 @@ shrinks from 50 MB to 32 MB. Making 11 and 13 lazy as well would save 4 MB more; lengths with those factors are common. The list is the preference `lazy_radices`; the trimmed C library sets it to empty, because `juliac` compiles nothing at run time. +## Run-time pass plans for trimmed programs + +**Decision (2026-11-10).** A pass plan (`plan_pass`, `plan_r2c_pass`, `plan_c2r_pass`) built with +`_runtime` (`plan_pass_runtime`, `plan_r2c_pass_runtime`, `plan_c2r_pass_runtime`) uses run-time kernels +only: RChain, Bluestein and composite transforms over RChain, the four-step algorithm over RChain. Its +type depends on the element type and the number of dimensions, not on the length. The specialised plans +are unchanged. + +**Why.** A library that holds passes of many lengths in one container (as NFFT.jl does for its boxes) +has a container of an abstract or length-dependent type, and `juliac --trim` cannot resolve the calls +through it. The first attempt was a keyword `runtime = true`; the verifier then still checks the +specialised branch, whose kernel types depend on the length, so the separate functions were needed. +A probe library compiled with `juliac --trim=safe` holds run-time pass plans in a vector, applies them, +and matches the direct DFT to 3·10⁻¹⁴ when called from C. Planning on 64 × 64 × 64 views is 1.8 µs for +`d = 1` and 0.9 µs for `d = 2` on the M1 (the specialised plan: 4.7 µs and 2.8 µs). Applying the +specialised and the run-time plans is within noise of each other on the same views. + +**Costs.** The run-time kernels dispatch per stage, which costs a little on short lines (the reason the +specialised plans stay the default). For real lengths with large prime factors, the run-time kernel is a +Bluestein transform, without the composite split of the specialised plan. + ## The interface: FFTW.jl's names, FFTL's own functions, AbstractFFTs as a backend **Decision (2026-11-06).** FFTL exports FFTW.jl's names (`fft`, `bfft`, `ifft`, `rfft`, `brfft`, diff --git a/docs/src/howto/trim.md b/docs/src/howto/trim.md index 530592f6..6cd9e1d5 100644 --- a/docs/src/howto/trim.md +++ b/docs/src/howto/trim.md @@ -12,6 +12,8 @@ For trimmed programs, FFTL has plans whose radix sequences and sizes are data: | `FFTL.RPlan{T}(sz, region, σ)` | complex DFT, sign `σ` | `mul!`, `*`, `\`, or `FFTL.execute!` on pointers | | `FFTL.RRFFTPlan{T}(sz, region, σ)` | r2c (`σ = -1`) or unnormalised c2r (`σ = +1`) | `FFTL.execute_r2c!`, `FFTL.execute_c2r!` on pointers | | `FFTL.RR2RPlan{T}(sz, kinds, region)` | real-to-real, one kind per dimension | `FFTL.execute_r2r!` on pointers | +| `FFTL.plan_pass_runtime(v, d)` | batched 1D DFT along `d` of a strided view | `mul!`, `apply!` | +| `FFTL.plan_r2c_pass_runtime(v, d)`, `FFTL.plan_c2r_pass_runtime(v, d)` | batched r2c and c2r along `d` | `mul!` | `T` is `Float32` or `Float64`; sizes, regions and kinds are vectors. All three take the `threads` keyword. They compute the same transforms with the same strategies as the normal plans @@ -37,6 +39,19 @@ GC.@preserve r out FFTL.execute_r2r!(q, pointer(out), pointer(r)) out ≈ FFTL.r2r(r, FFTL.REDFT10) ``` +The pass plans of the last two rows are the run-time versions of `plan_pass`, `plan_r2c_pass` and +`plan_c2r_pass`. Their type depends on the element type and the number of dimensions only (for example +`PassPlan{Float64,2,RPassKern{Float64}}`), not on the length, so plans of any size can be stored in one +container of a trimmed program. The specialised plans of the normal functions are faster for short lines. + +```@example trim +V = randn(ComplexF64, 64, 5) +pass = FFTL.plan_pass_runtime(V, 1) +W = similar(V) +mul!(W, pass, V) +W ≈ fft(V, 1) +``` + The pointer functions expect column-major arrays of the plan's size, which the caller keeps alive (`GC.@preserve`). The source of libfftl (`libfftl/src/libfftl.jl` and the FFTW3 interface in `libfftl/src/fftw3api.jl`) shows the rules the calling code follows to stay trimmable: no closures that capture types, concrete element types diff --git a/docs/src/reference/api.md b/docs/src/reference/api.md index ff15cc31..66412733 100644 --- a/docs/src/reference/api.md +++ b/docs/src/reference/api.md @@ -58,6 +58,11 @@ Batched one-dimensional transforms of strided arrays, the operations NFFT.jl bui ```@docs plan_pass +plan_pass_runtime +plan_r2c_pass +plan_r2c_pass_runtime +plan_c2r_pass +plan_c2r_pass_runtime plan_r2r_pass apply! ``` diff --git a/src/pass.jl b/src/pass.jl index 4ec290b8..f7843c54 100644 --- a/src/pass.jl +++ b/src/pass.jl @@ -1,6 +1,24 @@ # Batched 1D transforms along one dimension of a strided array: the operation an NFFT needs on # the oversampled grid (boxes of a padded array, chosen and threaded by the caller). +# Kernels of run-time pass plans (`plan_pass_runtime`). Their type depends on the element type only: +# a smooth length runs on RChain, a leftover of large prime factors is a composite of an RChain and a +# Bluestein transform over RChain, and a single long line on several threads runs the four-step +# algorithm (RChain kernels). A library can therefore hold plans of any size in one container and +# compile them ahead of time (juliac --trim). +const RComposite{T} = Composite{T,RChain{T},RBluestein{T}} +const RPassKern{T} = Union{RChain{T},RBluestein{T},RComposite{T},RFourStep{T}} + +function runtime_kernel(::Type{T}, n::Int, σ::Int, S::Int) where {T} + smooth(n) && return RChain{T}(n, σ, radices(n, S; esize = sizeof(Complex{T}))) + s, q = composite_split(T, n) + s == 1 && return Bluestein{T}(n, σ, RChain) + chain = RChain{T}(s, σ, radices(s, S * q; esize = sizeof(Complex{T}))) + blue = Bluestein{T}(q, σ, RChain) + tw = cached(() -> Complex{T}.(unit_roots(σ, n)), CacheKey(:composite, T, n, σ, Int[]))::Vector{Complex{T}} + RComposite{T}(n, s, q, chain, blue, tw, Vector{Complex{T}}(undef, S * n)) +end + """ PassPlan{T} @@ -25,31 +43,27 @@ struct PassPlan{T<:FFTLFloat,N,K<:Kernel{T}} end # Memory offsets (in elements) of the lines of an array with size sz and strides st, one per -# index of the dimensions `outer`. -line_offsets(sz, st, outer) = - vec([sum(((I[j] - 1) * st[outer[j]] for j in eachindex(outer)); init = 0) - for I in CartesianIndices(ntuple(j -> sz[outer[j]], length(outer)))]) - -""" - plan_pass(v::StridedArray{<:Complex}, d; sign = -1, threads = 1) +# index of the dimensions `outer` (first index fastest). Plain loops over vectors: the code does not +# depend on the number of dimensions at compile time (juliac --trim). +function line_offsets(sz::NTuple{N,Int}, st::NTuple{N,Int}, outer::Vector{Int}) where {N} + offs = Int[0] + for j in eachindex(outer) + stride = st[outer[j]]; c = sz[outer[j]]; L = length(offs) + next = Vector{Int}(undef, L * c) + for i in 0:c-1, k in 1:L + next[i * L + k] = offs[k] + i * stride + end + offs = next + end + offs +end -Plan for the batched 1D DFT along dimension `d` of arrays with the size and strides of `v` -(`sign = -1` forward, `+1` unnormalised backward). Apply with `mul!(w, p, v)` (w with the same -size and strides, `w === v` in place) or `apply!(p, v)`. With `threads > 1` the lines are split -among threads, and a single long line (d = 1) uses the four-step algorithm; callers that thread -over many arrays themselves keep the default 1. -""" -function plan_pass(v::StridedArray{Complex{T},N}, d::Integer; sign::Integer = -1, threads::Integer = 1) where {T<:FFTLFloat,N} - 1 <= d <= N || throw(ArgumentError("dimension $d is not within 1:$N")) - sign in (-1, 1) || throw(ArgumentError("sign must be -1 or 1")) - all(>=(1), size(v)) || throw(ArgumentError("array size must be positive")) - sz = size(v); st = strides(v) - (sz[1] == 1 || st[1] == 1) || throw(ArgumentError("dimension 1 must be contiguous (stride 1)")) - n = sz[d] - nt = max(1, Int(threads)) +# Layout of the lines of a pass along dimension d (dimension 1 contiguous): lanes per group, the +# element stride along d, and the offsets of the groups. +function pass_layout(sz::NTuple{N,Int}, st::NTuple{N,Int}, d::Int) where {N} if d == 1 lanes = 1; Sd = 1 - outer = ntuple(i -> i + 1, N - 1) + outer = collect(2:N) else # merge dimensions 1, 2, … into the lane index while they are contiguous and before d m = 1; L = sz[1] @@ -57,16 +71,66 @@ function plan_pass(v::StridedArray{Complex{T},N}, d::Integer; sign::Integer = -1 m += 1; L *= sz[m] end lanes = L; Sd = st[d] - outer = Tuple(i for i in m+1:N if i != d) + outer = Int[i for i in m+1:N if i != d] end - offs = line_offsets(sz, st, outer) + lanes, Sd, line_offsets(sz, st, outer) +end + +function pass_checks(v::AbstractArray, d::Integer, sign::Integer) + 1 <= d <= ndims(v) || throw(ArgumentError("dimension $d is not within 1:$(ndims(v))")) + sign in (-1, 1) || throw(ArgumentError("sign must be -1 or 1")) + all(>=(1), size(v)) || throw(ArgumentError("array size must be positive")) + (size(v, 1) == 1 || strides(v)[1] == 1) || throw(ArgumentError("dimension 1 must be contiguous (stride 1)")) + nothing +end + +# The plan with the kernel k of type K. +function pass_build(::Type{K}, v::StridedArray{Complex{T},N}, d::Int, σ::Int, nt::Int, k, + lanes::Int, Sd::Int, offs::Vector{Int}) where {K,T,N} + sz = size(v); st = strides(v); n = sz[d] + B = d == 1 ? 1 : lane_block(lanes, n, k isa AnyChain ? nt : 1) + nw = d == 1 ? n : B * n + PassPlan{T,N,K}(sz, st, d, σ, k, lanes, Sd, offs, B, nt, thread_buffers(T, nw, 2nt)) +end + +""" + plan_pass(v::StridedArray{<:Complex}, d; sign = -1, threads = 1) + +Plan for the batched 1D DFT along dimension `d` of arrays with the size and strides of `v` +(`sign = -1` forward, `+1` unnormalised backward). Apply with `mul!(w, p, v)` (w with the same +size and strides, `w === v` in place) or `apply!(p, v)`. With `threads > 1` the lines are split +among threads, and a single long line (d = 1) uses the four-step algorithm; callers that thread +over many arrays themselves keep the default 1. + +The kernels are specialised to the length. For plans that a library holds in one container and +compiles ahead of time, see [`plan_pass_runtime`](@ref). +""" +function plan_pass(v::StridedArray{Complex{T},N}, d::Integer; sign::Integer = -1, threads::Integer = 1) where {T<:FFTLFloat,N} + pass_checks(v, d, sign) + sz = size(v); st = strides(v); n = sz[d]; nt = max(1, Int(threads)); σ = Int(sign) + lanes, Sd, offs = pass_layout(sz, st, Int(d)) # a single long line on several threads: four-step - σ = Int(sign) k = d == 1 && length(offs) == 1 && use_fourstep(n, nt) ? FourStep{T}(n, σ, nt) : d == 1 ? kernel(T, n, σ) : kernel(T, n, σ, Sd) - B = d == 1 ? 1 : lane_block(lanes, n, k isa AnyChain ? nt : 1) - nw = d == 1 ? n : B * n - PassPlan{T,N,typeof(k)}(sz, st, d, sign, k, lanes, Sd, offs, B, nt, thread_buffers(T, nw, 2nt)) + pass_build(typeof(k), v, Int(d), σ, nt, k, lanes, Sd, offs) +end + +""" + plan_pass_runtime(v::StridedArray{<:Complex}, d; sign = -1, threads = 1) + +The same transform as [`plan_pass`](@ref), with run-time kernels only (RChain, Bluestein and composite +over RChain, four-step over RChain). The type of the plan is `PassPlan{T,N,RPassKern{T}}`, the same for +every length, so that a library can hold plans of any size in one container and compile them ahead of +time (`juliac --trim`). The kernels are the same as those of `plan_pass` for lengths with run-time +kernels; `plan_pass` is faster for short lines. +""" +function plan_pass_runtime(v::StridedArray{Complex{T},N}, d::Integer; sign::Integer = -1, threads::Integer = 1) where {T<:FFTLFloat,N} + pass_checks(v, d, sign) + sz = size(v); st = strides(v); n = sz[d]; nt = max(1, Int(threads)); σ = Int(sign) + lanes, Sd, offs = pass_layout(sz, st, Int(d)) + k::RPassKern{T} = d == 1 && length(offs) == 1 && use_fourstep(n, nt) ? FourStep{T}(n, σ, nt, RChain) : + runtime_kernel(T, n, σ, d == 1 ? 1 : Sd) + pass_build(RPassKern{T}, v, Int(d), σ, nt, k, lanes, Sd, offs) end function Base.show(io::IO, p::PassPlan{T}) where {T} diff --git a/src/r2r.jl b/src/r2r.jl index cceb7125..3df0ec15 100644 --- a/src/r2r.jl +++ b/src/r2r.jl @@ -819,7 +819,7 @@ function plan_r2r_pass(v::StridedArray{T,N}, d::Integer, kind::R2RKind; threads: 1 <= d <= N || throw(ArgumentError("dimension $d is not within 1:$N")) sz = size(v); st = strides(v) outer = Tuple(i for i in 1:N if i != d) - offs = line_offsets(sz, st, outer) + offs = line_offsets(sz, st, collect(Int, outer)) line = R2RLine{T}(kind, sz[d]) nt = max(1, Int(threads)) R2RPassPlan{T,N,typeof(line)}(sz, st, Int(d), line, offs, nt, [R2RBuffers{T}(sz[d]) for _ in 1:nt]) diff --git a/src/realpass.jl b/src/realpass.jl index adf00897..4f209a6d 100644 --- a/src/realpass.jl +++ b/src/realpass.jl @@ -26,12 +26,20 @@ end # Plan for dimension d of an input of size sz_in; σ = -1 (r2c) reads n reals, σ = +1 (c2r) reads # nhalf(n) complex values along d. function realpass_plan(::Type{T}, σ::Int, sz_in::NTuple{N,Int}, d::Int, n::Int, nt::Int) where {T<:FFTLFloat,N} + realpass_build(RealLine{T}(n, σ), σ, sz_in, d, n, nt) +end + +# Run-time version (RChain kernels only): the plan type depends on T and N only. +function realpass_plan_runtime(::Type{T}, σ::Int, sz_in::NTuple{N,Int}, d::Int, n::Int, nt::Int) where {T<:FFTLFloat,N} + realpass_build(RealLine{T}(n, σ, RChain), σ, sz_in, d, n, nt) +end + +function realpass_build(line::L, σ::Int, sz_in::NTuple{N,Int}, d::Int, n::Int, nt::Int) where {T<:FFTLFloat,N,L<:RealLine{T}} nin = σ == -1 ? n : nhalf(n) sz_in[d] == nin || throw(DimensionMismatch("input has length $(sz_in[d]) along dimension $d, expected $nin")) sz_out = Base.setindex(sz_in, σ == -1 ? nhalf(n) : n, d) - line = RealLine{T}(n, σ) nlines = prod(sz_in) ÷ nin - RealPassPlan{T,N,typeof(line)}(σ, d, n, sz_in, sz_out, + RealPassPlan{T,N,L}(σ, d, n, sz_in, sz_out, Vector{Int}(undef, nlines), Vector{Int}(undef, nlines), line, thread_buffers(T, n, 5nt), nt) end @@ -69,6 +77,33 @@ function plan_c2r_pass(v::StridedArray{Complex{T},N}, d::Integer; n::Integer = 2 realpass_plan(T, 1, size(v), Int(d), Int(n), max(1, Int(threads))) end +""" + plan_r2c_pass_runtime(v::StridedArray{<:Real}, d; threads = 1) + +The same transform as [`plan_r2c_pass`](@ref), with run-time kernels (RChain, and Bluestein over RChain +for lengths with large prime factors). The plan type is `RealPassPlan{T,N,RealLine{T,RKern{T}}}`, the same +for every length, so a library can compile it ahead of time (`juliac --trim`). +""" +function plan_r2c_pass_runtime(v::StridedArray{T,N}, d::Integer; threads::Integer = 1) where {T<:FFTLFloat,N} + 1 <= d <= N || throw(ArgumentError("dimension $d is not within 1:$N")) + all(>=(1), size(v)) || throw(ArgumentError("array size must be positive")) + realpass_plan_runtime(T, -1, size(v), Int(d), size(v, d), max(1, Int(threads))) +end + +""" + plan_c2r_pass_runtime(v::StridedArray{<:Complex}, d; n = 2(size(v, d) - 1), threads = 1) + +The same transform as [`plan_c2r_pass`](@ref) with run-time kernels, as [`plan_r2c_pass_runtime`](@ref). +""" +function plan_c2r_pass_runtime(v::StridedArray{Complex{T},N}, d::Integer; n::Integer = 2(size(v, d) - 1), + threads::Integer = 1) where {T<:FFTLFloat,N} + 1 <= d <= N || throw(ArgumentError("dimension $d is not within 1:$N")) + all(>=(1), size(v)) || throw(ArgumentError("array size must be positive")) + n >= 1 || throw(ArgumentError("real length must be positive")) + nhalf(n) == size(v, d) || throw(DimensionMismatch("half spectrum of length $(size(v, d)) does not belong to length $n")) + realpass_plan_runtime(T, 1, size(v), Int(d), Int(n), max(1, Int(threads))) +end + function Base.show(io::IO, p::RealPassPlan{T}) where {T} print(io, "FFTL ", p.σ == -1 ? "r2c" : "c2r", " pass plan, dim $(p.d) of $(join(p.sz_in, "×")) ", "$T -> $(join(p.sz_out, "×"))") diff --git a/test/runtests.jl b/test/runtests.jl index 81f6e37e..7cfe227b 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -348,6 +348,48 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end end + @testset "run-time pass plans, $T" for T in (Float64, Float32) + # the plan type depends on the element type and the dimension count only + for (sz, d) in (((100, 3), 1), ((1009, 3), 1), ((64, 5), 1), ((37, 6), 2), ((6, 130), 2)) + p = FFTL.plan_pass_runtime(randn(Complex{T}, sz), d) + @test p isa FFTL.PassPlan{T,2,FFTL.RPassKern{T}} + end + # against the length-specialised plan, every kernel kind: smooth, Bluestein, composite, four-step + for (n, sz_other, nt) in ((64, 5, 1), (100, 7, 1), (1000, 3, 2), (1009, 4, 1), (2 * 37 * 41, 2, 1), + (2^17, 1, 3)) + x = randn(Complex{T}, n, sz_other) + for sign in (-1, 1) + ref = FFTL.plan_pass(x, 1; sign, threads = nt); rt = FFTL.plan_pass_runtime(x, 1; sign, threads = nt) + a = similar(x); b = similar(x) + mul!(a, ref, x); mul!(b, rt, x) + @test relerr(b, a) < 10tol(T) + end + end + # strided views of a padded array, dimensions 2 and 3 + G0 = randn(Complex{T}, 40, 36, 30) + for d in (2, 3), (b1, b2, b3) in ((3:34, 2:33, 2:25), (1:40, 1:36, 1:30)) + v = view(G0, b1, b2, b3) + ref = FFTL.plan_pass(v, d); rt = FFTL.plan_pass_runtime(v, d) + W1 = similar(G0); W2 = similar(G0) + w1 = view(W1, b1, b2, b3); w2 = view(W2, b1, b2, b3) + mul!(w1, ref, v); mul!(w2, rt, v) + @test relerr(w2, w1) < 10tol(T) + end + # a small case against the direct DFT + x = randn(Complex{T}, 60, 2) + w = similar(x); mul!(w, FFTL.plan_pass_runtime(x, 1; sign = -1), x) + @test relerr(w, mapslices(u -> dft(ComplexF64.(u), -1), ComplexF64.(x); dims = 1)) < 10tol(T) + # real passes: run-time plans have one type for all lengths; the results agree with the specialised ones + for (sz, d) in (((100, 3), 1), ((1009, 3), 1), ((30, 2), 1)) + x = randn(T, sz) + @test FFTL.plan_r2c_pass_runtime(x, d) isa FFTL.RealPassPlan{T,2,FFTL.RealLine{T,Union{FFTL.RChain{T},FFTL.Bluestein{T,FFTL.RChain{T},FFTL.RChain{T}}}}} + ref = FFTL.plan_r2c_pass(x, d); rt = FFTL.plan_r2c_pass_runtime(x, d) + ya = zeros(Complex{T}, ref.sz_out); yb = similar(ya) + mul!(ya, ref, x); mul!(yb, rt, x) + @test relerr(yb, ya) < 10tol(T) + end + end + @testset "real transforms, $T" for T in (Float64, Float32) for (sz, region) in pick([((8,), 1), ((9,), 1), ((10,), 1), ((15,), 1), ((16,), 1), ((37,), 1), ((100,), 1), ((101,), 1), ((1024,), 1), ((1023,), 1), ((3 * 2^11,), 1),