Run-time pass plans for juliac --trim (type depends on T and N only)

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`,

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

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!
 ```

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,48 +43,94 @@ 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]
     while m + 1 < d && st[m+1] == L
       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}

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])

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, "×"))")

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),