diff --git a/README.md b/README.md index 5c2ee8bc..e451eac0 100644 --- a/README.md +++ b/README.md @@ -103,10 +103,9 @@ only dependency is the Julia runtime shipped next to it. On macOS arm64 the libr the runtime 10 MB. The run-time plans the library uses (`FFTL.RPlan`, `FFTL.RRFFTPlan`, `FFTL.RR2RPlan`) are public, so Julia programs that use them can be trimmed the same way. -**Entry points for special variants.** `FFTL.plan_pass(v, d; nonzero, needed)` transforms all -lines of a strided view along one dimension and can skip zero inputs and unneeded outputs -(pruning); `FFTL.plan_r2r_pass` does the same for real-to-real kinds. Libraries that build their -own transforms from one-dimensional passes use these; NFFT.jl is one of them. +**Entry points for special variants.** `FFTL.plan_pass(v, d)` transforms all lines of a strided view +along one dimension, and `FFTL.plan_r2r_pass` does the same for real-to-real kinds. Libraries that build +their own transforms from one-dimensional passes use these; NFFT.jl is one of them. ## Performance @@ -157,7 +156,7 @@ describes the setup and what a change must come with. ## Tests -`] test FFTL` runs a quick selection that covers every feature (about 1.5 min). Set +`] test FFTL` runs a quick selection that covers every feature (about 5 min with two threads). Set `FFTL_TEST_FULL=true` for the full lists and `FFTL_TEST_VERBOSE=true` for the time of every test set. CI runs the quick tests on Linux and macOS for pushes to `main` and pull requests, the full lists and the oldest supported Julia (1.10) once a week, and builds and tests the C library. diff --git a/docs/AGENTS.md b/docs/AGENTS.md index 2dc4b87e..8f2fc521 100644 --- a/docs/AGENTS.md +++ b/docs/AGENTS.md @@ -29,12 +29,12 @@ rules there apply to everything written here as well. Everything in this folder with the DCT, Poisson with the DST (second-order convergence shown), fast convolution. Each has at least one figure (CairoMakie) next to its printed results. - **How-to Guides** (`src/howto/`): coming from FFTW.jl, plans and reuse, transforms along dimensions - (including `plan_pass` with pruning, which NFFT.jl uses), real data, cosine/sine/Hartley transforms, + (including `plan_pass`, which NFFT.jl uses), real data, cosine/sine/Hartley transforms, threads, number types, precompilation, GPU, C library (libfftl), compiling ahead of time (the run-time plans `RPlan`, `RRFFTPlan`, `RR2RPlan`), contribute. - **Explanation** (`src/explanation/`): the DFT and its conventions, FFTW and FFTL (history, papers, what FFTL takes from FFTW; written from papers and documentation only), algorithms (Stockham chains, - codelets, pruning, Bluestein, multidimensional without transposition, four-step, real and real-to-real + codelets, Bluestein, multidimensional without transposition, four-step, real and real-to-real reductions), implementation (type specialisation, SIMD, kernel cache, cache sizes, the two chain forms, why libfftl needs `RChain`), design decisions (what was decided and why, with the evidence), package size and load time, performance (measurement protocol). diff --git a/docs/src/explanation/algorithms.md b/docs/src/explanation/algorithms.md index de61e59a..ac848b16 100644 --- a/docs/src/explanation/algorithms.md +++ b/docs/src/explanation/algorithms.md @@ -38,13 +38,6 @@ The generated code runs on complex SIMD vectors (`src/complexvec.jl`), with the the element type. On aarch64 the complex add-and-rotate and the complex multiply are single instructions (`FCADD`, `FCMLA`). -## Pruned codelets - -When the input is known to be zero outside some range, or only some outputs are used, the -generator drops the terms that are zero and the statements that feed no needed output. A pass -that transforms a zero-padded array therefore does not read the padding, and an NFFT that needs -only part of the spectrum does not compute the rest. See [transforms along dimensions](@ref "Transforms along dimensions"). - ## Bluestein's algorithm Lengths with a prime factor above 31 go through Bluestein's identity. With ``jk = (j^2 + k^2 - (k-j)^2)/2``, @@ -156,6 +149,6 @@ listed on [the cosine, sine and Hartley page](@ref "Cosine, sine and Hartley tra ## Passes `plan_pass` applies one 1D transform to all lines of a strided array with the same kernel. It -lets a library assemble its own multidimensional transform one dimension at a time and declare, -per dimension, which inputs are zero and which outputs it needs; both declarations use the -pruned codelets above. NFFT.jl transforms its oversampled grid this way. +lets a library assemble its own multidimensional transform one dimension at a time. NFFT.jl +transforms its oversampled grid this way; it skips the boxes of the grid that it does not need +itself, so the passes themselves carry no pruning. diff --git a/docs/src/explanation/image-size.md b/docs/src/explanation/image-size.md index c35a27ce..3b73cd81 100644 --- a/docs/src/explanation/image-size.md +++ b/docs/src/explanation/image-size.md @@ -45,8 +45,8 @@ later Julia versions, so re-run the count below after an upgrade. 1. **Plain methods for the fixed set.** The stage kernels that `RChain` dispatches to (radix × sign × vector type, empty zero-input mask or full output mask) and the compiled chains of the smooth lengths up to `SMALL_N` are defined in `src/kernel_methods.jl` with `@eval`, from the same body-building - functions the generators call. The `@generated` methods stay as the fallback for everything else - (pruned stages, other chains), so behaviour is unchanged. A plain method keeps its (compressed) source + functions the generators call. The `@generated` methods stay as the fallback for every other combination + (other radices, vector types and chains), so behaviour is unchanged. A plain method keeps its (compressed) source and its machine code, not its optimised IR. 2. **Lazy radices.** The codelets of the radices in `LAZY_RADICES` (17, 19, 23, 29, 31 by default; the preference `lazy_radices`, read in `src/rchain.jl`) are 54 % of the statements of the stage kernels and are rarely needed. They are not @@ -90,8 +90,8 @@ Further steps that were considered: | no `Float32` in the workload | 16 MB (at 73 MB) | 6 s at the first `Float32` call | Replacing the preference by environment variables did not shorten the load, and Julia does not notice a -changed variable, so a stale image would stay in use. Removing the pruned kernels from the workload was not -needed: no pruned kernel is in the image. +changed variable, so a stale image would stay in use. The pass plans do not take zero masks any more (their +kernels were `@generated` with a mask per segment), so no pruned kernel exists to store. ## Checking it diff --git a/docs/src/howto/contribute.md b/docs/src/howto/contribute.md index 545cae7a..bea83aec 100644 --- a/docs/src/howto/contribute.md +++ b/docs/src/howto/contribute.md @@ -33,9 +33,9 @@ describe what the code does now, not how it got there. using Pkg; Pkg.test("FFTL") ``` -This runs a quick selection that covers every feature. On a 256-thread AMD EPYC 7702 with two Julia -threads it takes about 40 minutes, so run it in the background or on the CI runner. Two environment -variables change it: +This runs a quick selection that covers every feature. With two Julia threads it takes about 5 minutes +on a 256-thread AMD EPYC 7702, and it is the check to run before every push. Two environment variables +change it: | variable | effect | |:--|:--| diff --git a/docs/src/howto/dimensions.md b/docs/src/howto/dimensions.md index ddcf39a9..55c3e4e9 100644 --- a/docs/src/howto/dimensions.md +++ b/docs/src/howto/dimensions.md @@ -38,27 +38,8 @@ a 3D array costs little more than one for a 1D line. `FFTL.plan_pass` transforms along one dimension of a strided view, applied to all lines at once. It is the building block of the transforms above, and the entry point for libraries that build -special variants of the transform themselves, such as NFFT.jl on its oversampled grid. Two -optional declarations let it skip work: - -- `nonzero` says that each line is zero outside the given ranges of indices along `d`. The first - stage then does not read the zeros. -- `needed` says that only the given output indices are used. The last stage computes only those. - -The result on the declared outputs is identical to the full transform. Here, the input is zero -outside its first 16 of 64 entries, and only the first 8 outputs are needed: - -```@example dims -n = 64 -V = zeros(ComplexF64, n, 10) -V[1:16, :] .= randn(ComplexF64, 16, 10) - -full = fft(V, 1) -pass = FFTL.plan_pass(V, 1; nonzero = (1:16,), needed = (1:8,)) -W = similar(V) -mul!(W, pass, V) -maximum(abs, W[1:8, :] - full[1:8, :]) -``` +special variants of the transform themselves, such as NFFT.jl on its oversampled grid. The plan is +made for the size and strides of the view, and applied with `mul!` or `apply!`. Passes are single-threaded by default. A caller that already threads over many arrays keeps `threads = 1`; see [Threads](@ref "Threads"). diff --git a/docs/src/howto/fftw.md b/docs/src/howto/fftw.md index 40a79501..8a26379b 100644 --- a/docs/src/howto/fftw.md +++ b/docs/src/howto/fftw.md @@ -83,7 +83,7 @@ roundtrip(p, x) ≈ x - Number types beyond `Float32` and `Float64`, computed in their own precision ([Number types](@ref "Number types")). - GPU arrays with the same functions ([GPU](gpu.md)). -- Passes over one dimension of a strided array, with pruning +- Passes over one dimension of a strided array ([Transforms along dimensions](@ref "Passes: one dimension of a strided array")). - A C library compiled from the Julia code, with FFTW3's C and Fortran interfaces ([libfftl](@ref "C library (libfftl)")). diff --git a/docs/src/index.md b/docs/src/index.md index da195126..0b1c4567 100644 --- a/docs/src/index.md +++ b/docs/src/index.md @@ -45,8 +45,8 @@ library needs no Julia installation. The run-time plans it is built from are pub programs that use them can be trimmed as well. **Entry points for special variants.** [Passes](@ref "Passes: one dimension of a strided array") -transform one dimension of a strided array and can skip zero inputs and unneeded outputs. -Libraries that assemble their own transforms from such passes, such as NFFT.jl, use them. +transform one dimension of a strided array. Libraries that assemble their own transforms from such +passes, such as NFFT.jl, use them. **The interface of the ecosystem.** FFTL exports FFTW.jl's names with FFTW.jl's conventions and scalings, so `using FFTL` can take the place of `using FFTW`, and it is a backend of diff --git a/docs/src/reference/platforms.md b/docs/src/reference/platforms.md index 6ed3e9d5..25df9c92 100644 --- a/docs/src/reference/platforms.md +++ b/docs/src/reference/platforms.md @@ -15,7 +15,7 @@ available. | DST-II | ✓ | ✓ | ✓ | ✓ | | DST-I, DST-III, DST-IV | ✓ | ✓ | ◐ partly tuned | ✓ | | DHT, `R2HC`, `HC2R` | ✓ | ✓ | ◐ not tuned | ✓ | -| batched 1D passes (`plan_pass`, `plan_r2r_pass`) with pruning | ✓ | – | – | – | +| batched 1D passes (`plan_pass`, `plan_r2r_pass`) | ✓ | – | – | – | | any rank and any subset of dimensions | ✓ | ✓ | ✓ | ✓ | | any length (Bluestein above prime factor 31) | ✓ | ✓ | ✓ | ✓ | diff --git a/src/codegen.jl b/src/codegen.jl index 7a8f3254..c00d951e 100644 --- a/src/codegen.jl +++ b/src/codegen.jl @@ -11,8 +11,6 @@ # Odd primes use the symmetric form # X_k, X_{r-k} = A ± σ·i·B, A = a_0 + Σ_j (a_j + a_{r-j}) cos(2πjk/r), B = Σ_j (a_j - a_{r-j}) sin(2πjk/r). # -# Pruned codelets: inputs known to be zero are the term `nothing` (sums and products simplify), -# and only the needed outputs are computed (statements they do not depend on are dropped). mutable struct CodeGen stmts::Vector{Any} @@ -29,13 +27,9 @@ function emit!(g::CodeGen, ex) end rotate(t, k) = (t[1], mod(t[2] + k, 4)) -rotate(::Nothing, k) = nothing # rotation units of one factor σ·i: -i = (-i)^1, +i = (-i)^3 unit(g::CodeGen) = g.σ == -1 ? 1 : 3 -plus!(g, ::Nothing, b) = b -plus!(g, a, ::Nothing) = a -plus!(g, ::Nothing, ::Nothing) = nothing function plus!(g, a, b) d = mod(b[2] - a[2], 4) ex = d == 0 ? :($(a[1]) + $(b[1])) : @@ -46,7 +40,6 @@ function plus!(g, a, b) end minus!(g, a, b) = plus!(g, a, rotate(b, 2)) -force!(g, ::Nothing) = emit!(g, :(zero(VT))) function force!(g, t) t[2] == 0 && return t[1] t[2] == 2 && return emit!(g, :(-$(t[1]))) @@ -55,7 +48,6 @@ function force!(g, t) end # t · W_N^e with constant e, N -twiddle!(g, ::Nothing, e, N) = nothing function twiddle!(g, t, e, N) e = mod(e, N) e == 0 && return t @@ -70,37 +62,34 @@ function twiddle!(g, t, e, N) (emit!(g, :(ccmul($(t[1]), cconst(VT, $c, $s)))), t[2]) end -# sums of plain symbols where `nothing` is zero -radd!(g, x, y) = x === nothing ? y : y === nothing ? x : emit!(g, :($x + $y)) -rsub!(g, x, y) = y === nothing ? x : x === nothing ? emit!(g, :(-$y)) : emit!(g, :($x - $y)) -rmuladd!(g, x, c, acc) = x === nothing ? acc : acc === nothing ? emit!(g, :($c * $x)) : emit!(g, :(muladd($x, $c, $acc))) +radd!(g, x, y) = emit!(g, :($x + $y)) +rsub!(g, x, y) = emit!(g, :($x - $y)) +rmuladd!(g, x, c, acc) = emit!(g, :(muladd($x, $c, $acc))) +scale!(g, c, x) = emit!(g, :($c * $x)) function prime!(g, a) r = length(a); h = (r - 1) ÷ 2 - a = [x === nothing ? nothing : force!(g, x) for x in a] + a = [force!(g, x) for x in a] sp = [radd!(g, a[j+1], a[r-j+1]) for j in 1:h] dm = [rsub!(g, a[j+1], a[r-j+1]) for j in 1:h] X = Vector{Any}(undef, r) - acc = nothing - for j in 1:h + acc = sp[1] + for j in 2:h acc = radd!(g, acc, sp[j]) end - x0 = radd!(g, a[1], acc) - X[1] = x0 === nothing ? nothing : (x0, 0) + X[1] = (radd!(g, a[1], acc), 0) for k in 1:h A = a[1] for j in 1:h A = rmuladd!(g, sp[j], cospi(2j * k / r), A) end - B = nothing - for j in 1:h + B = scale!(g, sinpi(2 * k / r), dm[1]) + for j in 2:h B = rmuladd!(g, dm[j], sinpi(2j * k / r), B) end # X_k = A + σ·i·B and X_(r-k) = A - σ·i·B; -i·B is the term (B, 1), +i·B the term (B, 3) - tA = A === nothing ? nothing : (A, 0) - tB = B === nothing ? nothing : (B, 0) - X[k+1] = plus!(g, tA, rotate(tB, g.σ == -1 ? 1 : 3)) - X[r-k+1] = plus!(g, tA, rotate(tB, g.σ == -1 ? 3 : 1)) + X[k+1] = plus!(g, (A, 0), rotate((B, 0), g.σ == -1 ? 1 : 3)) + X[r-k+1] = plus!(g, (A, 0), rotate((B, 0), g.σ == -1 ? 3 : 1)) end X end @@ -133,33 +122,15 @@ function dft!(g::CodeGen, a::Vector) end """ - codelet(r, σ, zin = 0, need = all) -> (stmts, inputs, outputs) + codelet(r, σ) -> (stmts, inputs, outputs) Statements computing a size-`r` DFT with sign `σ` from the input symbols `a1…ar` into the -returned output symbols. Inputs whose bit is set in `zin` (bit j-1 for input j) are zero and not -read; outputs whose bit is clear in `need` are not computed (their entry is `nothing`). +returned output symbols. """ -function codelet(r::Int, σ::Int, zin::Integer = 0, need::Integer = (1 << r) - 1) +function codelet(r::Int, σ::Int) g = CodeGen(σ) a = [Symbol("a", j) for j in 1:r] - X = dft!(g, [((zin >> (j - 1)) & 1 == 1 ? nothing : (s, 0)) for (j, s) in enumerate(a)]) - out = [(need >> (k - 1)) & 1 == 1 ? force!(g, X[k]) : nothing for k in 1:r] - prune_statements!(g.stmts, out), a, out + X = dft!(g, [(s, 0) for s in a]) + out = [force!(g, X[k]) for k in 1:r] + g.stmts, a, out end - -# Drop statements that no needed output depends on. -function prune_statements!(stmts, out) - live = Set{Symbol}(o for o in out if o isa Symbol) - keep = falses(length(stmts)) - for i in length(stmts):-1:1 - lhs, rhs = stmts[i].args - lhs in live || continue - keep[i] = true - collect_symbols!(live, rhs) - end - stmts[keep] -end - -collect_symbols!(set, s::Symbol) = push!(set, s) -collect_symbols!(set, ex::Expr) = foreach(x -> collect_symbols!(set, x), ex.args) -collect_symbols!(set, x) = set diff --git a/src/kernel_methods.jl b/src/kernel_methods.jl index 58437c9b..e90b555d 100644 --- a/src/kernel_methods.jl +++ b/src/kernel_methods.jl @@ -1,56 +1,53 @@ # Compiled-ahead methods of the stage kernels for the fixed codelet set (radix × sign × vector -# type, empty zero-input mask or full output mask) that `RChain` dispatches to. The `@generated` -# kernels in kernels.jl handle every other combination (pruned stages); a package image keeps -# the complete optimised IR of every precompiled `@generated` specialization but none of a -# plain method, so the fixed set is defined as plain methods with the same bodies. The codelets -# of LAZY_RADICES are not in the set. +# type) that `RChain` dispatches to, and of the chains of the small lengths (rchain.jl). The +# `@generated` kernels in kernels.jl define every other combination. A package image keeps the +# complete optimised IR of each precompiled `@generated` specialization but none of a plain +# method, so the fixed set is defined as plain methods with the same bodies. The codelets of +# LAZY_RADICES are not in the set. kernel_subst(x, d) = x isa Symbol ? get(d, x, x) : x isa Expr ? Expr(x.head, map(a -> kernel_subst(a, d), x.args)...) : x const KERNEL_HEADS = ( (:stage!, :(stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{mask})), :zin, :all), + ::Val{R}, ::Val{σ}, ::Type{VT})), :all), (:stage_range!, :(stage_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, qlo::Int, qhi::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{mask})), :zin, :all), + plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :all), (:stage_last!, :(stage_last!(py::Ptr{T}, px::Ptr{T}, s::Int, qlo::Int, qhi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{mask})), :need, :all), + ::Val{R}, ::Val{σ}, ::Type{VT})), :all), (:stage_lanes!, :(stage_lanes!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s0::Int, Sx::Int, Sy::Int, - B::Int, plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{mask})), :zin, :all), + B::Int, plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :all), (:stage_last_lanes!, :(stage_last_lanes!(py::Ptr{T}, px::Ptr{T}, s0::Int, Sx::Int, Sy::Int, B::Int, - qlo::Int, qhi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{mask})), :need, :all), + qlo::Int, qhi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :all), # the grouped first stage exists for 2 or 4 complex values per vector only; its scalar tail # runs on one complex value (:stage_first_group_range!, :(stage_first_group_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{mask})), :zin, :group), + plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :group), (:stage_first_group_1!, :(stage_first_group_1!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{mask})), :zin, :tail)) + plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT})), :tail)) -for T in (Float32, Float64), (name, head, kind, vts) in KERNEL_HEADS +for T in (Float32, Float64), (name, head, vts) in KERNEL_HEADS body = getfield(@__MODULE__, Symbol(string(name)[1:end-1], "_body")) vts === :all ? (types = vec_ladder(T)) : vts === :group ? (types = filter(V -> ncomplex(V) in (2, 4), vec_ladder(T))) : (types = any(V -> ncomplex(V) in (2, 4), vec_ladder(T)) ? (Vec{2,T},) : ()) for σ in (-1, 1), VT in types, R in RADICES R in LAZY_RADICES && continue - mask = kind === :need ? (1 << R) - 1 : 0 - h = kernel_subst(head, Dict{Symbol,Any}(:T => T, :R => R, :σ => σ, :VT => VT, :mask => mask)) + h = kernel_subst(head, Dict{Symbol,Any}(:T => T, :R => R, :σ => σ, :VT => VT)) # the codelet statements name the type parameters T and VT, which a plain method does not bind - b = kernel_subst(body(T, R, σ, VT, mask), Dict{Symbol,Any}(:T => T, :VT => VT)) + b = kernel_subst(body(T, R, σ, VT), Dict{Symbol,Any}(:T => T, :VT => VT)) @eval $(Expr(:function, h, b)) end end # The compiled chains of the small lengths (rchain.jl) run through `run!`/`run_lanes!` methods of -# plain definition for the same reason; any other chain type, and every pruned call, uses the -# @generated methods. +# plain definition for the same reason; any other chain type uses the generic methods. for T in (Float32, Float64), (n, σ) in SMALL_ENTRIES Rs = Tuple(radices(n, 1; esize = sizeof(Complex{T}))) C = Chain{T,Rs,σ,length(Rs)} - Pr = typeof(NOPRUNE) - @eval $(Expr(:function, :(run!(py::Ptr{$T}, px::Ptr{$T}, pw1::Ptr{$T}, pw2::Ptr{$T}, P::$C, S::Int, Pr::$Pr)), - run_body(T, Rs, σ, length(Rs), (), ()))) + @eval $(Expr(:function, :(run!(py::Ptr{$T}, px::Ptr{$T}, pw1::Ptr{$T}, pw2::Ptr{$T}, P::$C, S::Int)), + run_body(T, Rs, σ, length(Rs)))) @eval $(Expr(:function, :(run_lanes!(py::Ptr{$T}, px::Ptr{$T}, pw1::Ptr{$T}, pw2::Ptr{$T}, P::$C, - Sx::Int, Sy::Int, B::Int, Pr::$Pr)), - run_lanes_body(T, Rs, σ, length(Rs), (), ()))) + Sx::Int, Sy::Int, B::Int)), + run_lanes_body(T, Rs, σ, length(Rs)))) end diff --git a/src/kernels.jl b/src/kernels.jl index b76b8656..8120237b 100644 --- a/src/kernels.jl +++ b/src/kernels.jl @@ -8,24 +8,17 @@ # # A transform along a dimension with stride S (batch of S interleaved transforms) is the same # chain started with s = S: the S lanes become part of the contiguous index q. -# -# Pruned stages (zero inputs, unneeded outputs of a line): first stages take a range plo:phi-1 of -# p and a bit mask `zin` of inputs j that are zero there (not loaded, not computed with); last -# stages take a range of q (or q0) and a bit mask `need` of the outputs k to compute and store. -# The unpruned methods cover the full range with empty masks. - -zero_input(zin, j) = (zin >> (j - 1)) & 1 == 1 @inline cptr(x::Array{Complex{T}}) where {T} = Ptr{T}(pointer(x)) # General stage, vectorised over q: s must be a multiple of the complex lanes of VT. stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage!(py, px, pt, m, s, 0, m, Val(R), Val(σ), VT, Val(0)) + stage!(py, px, pt, m, s, 0, m, Val(R), Val(σ), VT) -function stage_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) where {T,VT} - stmts, a, X = codelet(R, σ, zin) +function stage_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} + stmts, a, X = codelet(R, σ) C = ncomplex(VT); B = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (q + s * (p + $(j - 1) * m)) * $B)) for j in 1:R if !zero_input(zin, j)] + loads = [:($(a[j]) = vload($VT, px + (q + s * (p + $(j - 1) * m)) * $B)) for j in 1:R] twl = [:($(Symbol("w", k)) = cbroadcast($VT, pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T)))) for k in 1:R-1] stores = Any[:(vstore($(X[1]), py + o * $B))] for k in 1:R-1 @@ -46,8 +39,8 @@ function stage_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) where {T,V end @generated function stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{zin}) where {T,R,σ,VT,zin} - stage_body(T, R, σ, VT, zin) + ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} + stage_body(T, R, σ, VT) end # Stage for strides that are not a multiple of the complex lanes of VT (odd batch sizes): the @@ -60,17 +53,17 @@ end if rem == 0 stage!(py, px, pt, m, s, Val(R), Val(σ), VT) else - stage_range!(py, px, pt, m, s, 0, s - rem, 0, m, Val(R), Val(σ), VT, Val(0)) - stage_range!(py, px, pt, m, s, s - rem, s, 0, m, Val(R), Val(σ), Vec{2,T}, Val(0)) + stage_range!(py, px, pt, m, s, 0, s - rem, 0, m, Val(R), Val(σ), VT) + stage_range!(py, px, pt, m, s, s - rem, s, 0, m, Val(R), Val(σ), Vec{2,T}) end end # Stage over a q range: the kernel behind the split of a single line, whose q extent is known # at generation time so the vectorised range and the scalar tail are emitted directly. -function stage_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) where {T,VT} - stmts, a, X = codelet(R, σ, zin) +function stage_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} + stmts, a, X = codelet(R, σ) C = ncomplex(VT); B = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (q + s * (p + $(j - 1) * m)) * $B)) for j in 1:R if !zero_input(zin, j)] + loads = [:($(a[j]) = vload($VT, px + (q + s * (p + $(j - 1) * m)) * $B)) for j in 1:R] twl = [:($(Symbol("w", k)) = cbroadcast($VT, pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T)))) for k in 1:R-1] stores = Any[:(vstore($(X[1]), py + o * $B))] for k in 1:R-1 @@ -91,19 +84,19 @@ function stage_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) wher end @generated function stage_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, qlo::Int, qhi::Int, - plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{zin}) where {T,R,σ,VT,zin} - stage_range_body(T, R, σ, VT, zin) + plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} + stage_range_body(T, R, σ, VT) end # Last stage (m = 1): all twiddles are 1. stage_last!(py::Ptr{T}, px::Ptr{T}, s::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_last!(py, px, s, 0, s, Val(R), Val(σ), VT, Val((1 << R) - 1)) + stage_last!(py, px, s, 0, s, Val(R), Val(σ), VT) -function stage_last_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, need::Int) where {T,VT} - stmts, a, X = codelet(R, σ, 0, need) +function stage_last_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} + stmts, a, X = codelet(R, σ) C = ncomplex(VT); B = 2sizeof(T) loads = [:($(a[j]) = vload($VT, px + (q + s * $(j - 1)) * $B)) for j in 1:R] - stores = [:(vstore($(X[k+1]), py + (q + $k * s) * $B)) for k in 0:R-1 if X[k+1] !== nothing] + stores = [:(vstore($(X[k+1]), py + (q + $k * s) * $B)) for k in 0:R-1] quote for q in qlo:$C:qhi-1 $(loads...) @@ -115,8 +108,8 @@ function stage_last_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, need::Int) wher end @generated function stage_last!(py::Ptr{T}, px::Ptr{T}, s::Int, qlo::Int, qhi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{need}) where {T,R,σ,VT,need} - stage_last_body(T, R, σ, VT, need) + ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} + stage_last_body(T, R, σ, VT) end # Last stage for strides that are not a multiple of the complex lanes of VT (odd batch sizes): @@ -128,8 +121,8 @@ end if rem == 0 stage_last!(py, px, s, Val(R), Val(σ), VT) else - stage_last!(py, px, s, 0, s - rem, Val(R), Val(σ), VT, Val((1 << R) - 1)) - stage_last!(py, px, s, s - rem, s, Val(R), Val(σ), Vec{2,T}, Val((1 << R) - 1)) + stage_last!(py, px, s, 0, s - rem, Val(R), Val(σ), VT) + stage_last!(py, px, s, s - rem, s, Val(R), Val(σ), Vec{2,T}) end end @@ -139,26 +132,26 @@ end # outputs of one butterfly; remaining outputs (R not a multiple of C) are stored one complex # value at a time. Butterflies beyond a multiple of C (m % C) run one complex value at a time. stage_first_group!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_first_group!(py, px, pt, m, 0, m, Val(R), Val(σ), VT, Val(0)) + stage_first_group!(py, px, pt, m, 0, m, Val(R), Val(σ), VT) @inline function stage_first_group!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{zin}) where {T,R,σ,VT,zin} + ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} C = ncomplex(VT) rem = (phi - plo) % C if rem == 0 - stage_first_group_range!(py, px, pt, m, plo, phi, Val(R), Val(σ), VT, Val(zin)) + stage_first_group_range!(py, px, pt, m, plo, phi, Val(R), Val(σ), VT) else - stage_first_group_range!(py, px, pt, m, plo, phi - rem, Val(R), Val(σ), VT, Val(zin)) - stage_first_group_1!(py, px, pt, m, phi - rem, phi, Val(R), Val(σ), Vec{2,T}, Val(zin)) + stage_first_group_range!(py, px, pt, m, plo, phi - rem, Val(R), Val(σ), VT) + stage_first_group_1!(py, px, pt, m, phi - rem, phi, Val(R), Val(σ), Vec{2,T}) end end -function stage_first_group_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) where {T,VT} +function stage_first_group_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} C = ncomplex(VT) C in (2, 4) || error("the grouped first stage needs 2 or 4 complex values per vector") - stmts, a, X = codelet(R, σ, zin) + stmts, a, X = codelet(R, σ) N = 2C; CB = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (p + $(j - 1) * m) * $CB)) for j in 1:R if !zero_input(zin, j)] + loads = [:($(a[j]) = vload($VT, px + (p + $(j - 1) * m) * $CB)) for j in 1:R] outs = Any[X[1]] twm = Any[] for k in 1:R-1 @@ -197,17 +190,17 @@ function stage_first_group_range_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zi end @generated function stage_first_group_range!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{zin}) where {T,R,σ,VT,zin} - stage_first_group_range_body(T, R, σ, VT, zin) + ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} + stage_first_group_range_body(T, R, σ, VT) end # Scalar tail of the grouped first stage: one complex value per vector, the twiddle table in # the grouped layout (k-1)m + p, one store per output. -function stage_first_group_1_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) where {T,VT} +function stage_first_group_1_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} ncomplex(VT) == 1 || error("the scalar tail of the grouped first stage needs one complex value per vector") - stmts, a, X = codelet(R, σ, zin) + stmts, a, X = codelet(R, σ) CB = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (p + $(j - 1) * m) * $CB)) for j in 1:R if !zero_input(zin, j)] + loads = [:($(a[j]) = vload($VT, px + (p + $(j - 1) * m) * $CB)) for j in 1:R] twm = Any[] stores = Any[:(vstore($(X[1]), py + $R * p * $CB))] for k in 1:R-1 @@ -227,8 +220,8 @@ function stage_first_group_1_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::I end @generated function stage_first_group_1!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, plo::Int, phi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{zin}) where {T,R,σ,VT,zin} - stage_first_group_1_body(T, R, σ, VT, zin) + ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} + stage_first_group_1_body(T, R, σ, VT) end # Stages on a block of B lanes (cache blocking of a dimension pass, and its unit of work among @@ -238,12 +231,12 @@ end # i + S·q0 with i = 0…B-1, q0 = 0…s0-1. B must be a multiple of the complex lanes of VT. stage_lanes!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s0::Int, Sx::Int, Sy::Int, B::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_lanes!(py, px, pt, m, s0, Sx, Sy, B, 0, m, Val(R), Val(σ), VT, Val(0)) + stage_lanes!(py, px, pt, m, s0, Sx, Sy, B, 0, m, Val(R), Val(σ), VT) -function stage_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) where {T,VT} - stmts, a, X = codelet(R, σ, zin) +function stage_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} + stmts, a, X = codelet(R, σ) C = ncomplex(VT); Bt = 2sizeof(T) - loads = [:($(a[j]) = vload($VT, px + (i + Sx * (q0 + s0 * (p + $(j - 1) * m))) * $Bt)) for j in 1:R if !zero_input(zin, j)] + loads = [:($(a[j]) = vload($VT, px + (i + Sx * (q0 + s0 * (p + $(j - 1) * m))) * $Bt)) for j in 1:R] twl = [:($(Symbol("w", k)) = cbroadcast($VT, pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T)))) for k in 1:R-1] stores = Any[:(vstore($(X[1]), py + (i + Sy * o) * $Bt))] for k in 1:R-1 @@ -266,19 +259,19 @@ function stage_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, zin::Int) wher end @generated function stage_lanes!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s0::Int, Sx::Int, Sy::Int, - B::Int, plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{zin}) where {T,R,σ,VT,zin} - stage_lanes_body(T, R, σ, VT, zin) + B::Int, plo::Int, phi::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} + stage_lanes_body(T, R, σ, VT) end stage_last_lanes!(py::Ptr{T}, px::Ptr{T}, s0::Int, Sx::Int, Sy::Int, B::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} = - stage_last_lanes!(py, px, s0, Sx, Sy, B, 0, s0, Val(R), Val(σ), VT, Val((1 << R) - 1)) + stage_last_lanes!(py, px, s0, Sx, Sy, B, 0, s0, Val(R), Val(σ), VT) -function stage_last_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, need::Int) where {T,VT} - stmts, a, X = codelet(R, σ, 0, need) +function stage_last_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}) where {T,VT} + stmts, a, X = codelet(R, σ) C = ncomplex(VT); Bt = 2sizeof(T) loads = [:($(a[j]) = vload($VT, px + (i + Sx * (q0 + s0 * $(j - 1))) * $Bt)) for j in 1:R] - stores = [:(vstore($(X[k+1]), py + (i + Sy * (q0 + s0 * $k)) * $Bt)) for k in 0:R-1 if X[k+1] !== nothing] + stores = [:(vstore($(X[k+1]), py + (i + Sy * (q0 + s0 * $k)) * $Bt)) for k in 0:R-1] quote for q0 in qlo:qhi-1 for i in 0:$C:B-1 @@ -292,6 +285,6 @@ function stage_last_lanes_body(::Type{T}, R::Int, σ::Int, ::Type{VT}, need::Int end @generated function stage_last_lanes!(py::Ptr{T}, px::Ptr{T}, s0::Int, Sx::Int, Sy::Int, B::Int, qlo::Int, qhi::Int, - ::Val{R}, ::Val{σ}, ::Type{VT}, ::Val{need}) where {T,R,σ,VT,need} - stage_last_lanes_body(T, R, σ, VT, need) + ::Val{R}, ::Val{σ}, ::Type{VT}) where {T,R,σ,VT} + stage_last_lanes_body(T, R, σ, VT) end diff --git a/src/pass.jl b/src/pass.jl index a237e9b3..4ec290b8 100644 --- a/src/pass.jl +++ b/src/pass.jl @@ -1,5 +1,5 @@ # Batched 1D transforms along one dimension of a strided array: the operation an NFFT needs on -# the oversampled grid (boxes of a padded array, pruned and threaded by the caller). +# the oversampled grid (boxes of a padded array, chosen and threaded by the caller). """ PassPlan{T} @@ -10,7 +10,7 @@ the lanes: dimension 1 (and the following dimensions as long as they are contigu the vectorised lane index, the element stride is the stride of dimension d. `offs` lists the offsets of the line groups (complex values from the first element). """ -struct PassPlan{T<:FFTLFloat,N,K<:Kernel{T},Pr<:Pruning} +struct PassPlan{T<:FFTLFloat,N,K<:Kernel{T}} sz::NTuple{N,Int} st::NTuple{N,Int} d::Int @@ -22,7 +22,6 @@ struct PassPlan{T<:FFTLFloat,N,K<:Kernel{T},Pr<:Pruning} B::Int # lane block nthreads::Int bufs::Vector{Vector{Complex{T}}} # two per thread - prune::Pr end # Memory offsets (in elements) of the lines of an array with size sz and strides st, one per @@ -32,21 +31,15 @@ line_offsets(sz, st, outer) = for I in CartesianIndices(ntuple(j -> sz[outer[j]], length(outer)))]) """ - plan_pass(v::StridedArray{<:Complex}, d; sign = -1, threads = 1, nonzero = nothing, needed = nothing) + 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. - -`nonzero` (ranges of indices along d, e.g. `(1:a, b:n)`) declares that each line is zero -outside these ranges; `needed` that only these outputs of each line are used (the others may -hold anything afterwards). The first resp. last stage then skips the zero inputs resp. the -unneeded outputs. """ -function plan_pass(v::StridedArray{Complex{T},N}, d::Integer; sign::Integer = -1, threads::Integer = 1, - nonzero = nothing, needed = nothing) where {T<:FFTLFloat,N} +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")) @@ -73,71 +66,7 @@ function plan_pass(v::StridedArray{Complex{T},N}, d::Integer; sign::Integer = -1 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 - pr = k isa Chain ? pruning(k, nonzero, needed; lanes = d > 1) : NOPRUNE - PassPlan{T,N,typeof(k),typeof(pr)}(sz, st, d, sign, k, lanes, Sd, offs, B, nt, - thread_buffers(T, nw, 2nt), pr) -end - -# Segments of constant mask over units of `cu` consecutive indices 0…len-1; mask(range) gives the -# mask of one unit. -function segments(mask, len, cu) - ranges = UnitRange{Int}[]; masks = Int[] - for u0 in 0:cu:len-1 - r = u0:min(u0 + cu, len)-1 - mk = mask(r) - if !isempty(masks) && masks[end] == mk - ranges[end] = first(ranges[end]):last(r) - else - push!(ranges, r); push!(masks, mk) - end - end - ranges, masks -end - -inranges(ranges, n) = (b = falses(n); foreach(r -> (b[r] .= true), ranges); b) - -# Pruning only touches the first or last of K stages, so it saves at most about (1 - f)/K of a -# pass for a nonzero (needed) fraction f; below PRUNE_MIN_ZERO zeros the segment overhead -# outweighs that (M1: σ = 1.25 lines gained 0–3 %, short Float32 lines lost up to 13 %). -const PRUNE_MIN_ZERO = 1 / 3 - -""" - pruning(chain, nonzero, needed; lanes) - -Segments and masks of the first and last stage of `chain` for lines that are zero outside -`nonzero` and of which only `needed` is used (index ranges along the line; `nothing`: no -pruning). `lanes`: the chain runs in lane blocks (units of one butterfly); otherwise units follow -the vectorisation over butterflies of the single-line kernels. -""" -function pruning(c::Chain{T,Rs}, nonzero, needed; lanes::Bool) where {T,Rs} - n = c.n; K = length(Rs) - K >= 2 || return NOPRUNE - worth(ranges) = ranges !== nothing && 1 - count(inranges(ranges, n)) / n >= PRUNE_MIN_ZERO - nonzero = worth(nonzero) ? nonzero : nothing - needed = worth(needed) ? needed : nothing - ladder = vec_ladder(T) - fr = UnitRange{Int}[]; fm = Int[]; lr = UnitRange{Int}[]; lm = Int[] - if nonzero !== nothing - nz = inranges(nonzero, n) - r1 = Rs[1]; m = n ÷ r1 - g = findfirst(V -> ncomplex(V) in (2, 4) && m % ncomplex(V) == 0, ladder) - cu = lanes || g === nothing ? 1 : ncomplex(ladder[g]) - # input j of butterfly p is x[p + (j-1)m]; zero for the unit if zero for all its p - fr, fm = segments(m, cu) do r - sum((all(p -> !nz[p + (j - 1) * m + 1], r) ? 1 << (j - 1) : 0) for j in 1:r1) - end - end - if needed !== nothing - nd = inranges(needed, n) - rK = Rs[K]; sl = n ÷ rK - V = ladder[findfirst(V -> sl % ncomplex(V) == 0, ladder)] - cu = lanes ? 1 : ncomplex(V) - # output k of position q is y[q + (k-1)·sl]; needed for the unit if needed for any q - lr, lm = segments(sl, cu) do r - sum((any(q -> nd[q + (k - 1) * sl + 1], r) ? 1 << (k - 1) : 0) for k in 1:rK) - end - end - Pruning{Tuple(fm),Tuple(lm)}(fr, lr) + PassPlan{T,N,typeof(k)}(sz, st, d, sign, k, lanes, Sd, offs, B, nt, thread_buffers(T, nw, 2nt)) end function Base.show(io::IO, p::PassPlan{T}) where {T} @@ -174,7 +103,7 @@ function run_pass!(py::Ptr{T}, px::Ptr{T}, p::PassPlan{T}) where {T} # lines in groups on the tasks (kernels with buffers of their own on one task) each_line(length(p.offs), p.sz[1], k isa AnyChain ? p.nthreads : 1) do c, i local off = p.offs[i + 1] * CT - run!(py + off, px + off, cptr(p.bufs[2c-1]), cptr(p.bufs[2c]), k, 1, p.prune) + run!(py + off, px + off, cptr(p.bufs[2c-1]), cptr(p.bufs[2c]), k, 1) end else B = p.B; nb = cld(p.lanes, B) @@ -184,14 +113,8 @@ function run_pass!(py::Ptr{T}, px::Ptr{T}, p::PassPlan{T}) where {T} local g, b; g, b = divrem(it, nb) local i0 = b * B local off = (p.offs[g + 1] + i0) * CT - run_lanes!(py + off, px + off, cptr(p.bufs[2c-1]), cptr(p.bufs[2c]), k, p.Sd, p.Sd, min(B, p.lanes - i0), p.prune) + run_lanes!(py + off, px + off, cptr(p.bufs[2c-1]), cptr(p.bufs[2c]), k, p.Sd, p.Sd, min(B, p.lanes - i0)) end end nothing end - -# kernels without pruning support (Bluestein, four-step) ignore it -run!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::Kernel{T}, S::Int, ::Pruning) where {T} = - run!(py, px, pw1, pw2, P, S) -run_lanes!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::Kernel{T}, Sx::Int, Sy::Int, B::Int, ::Pruning) where {T} = - run_lanes!(py, px, pw1, pw2, P, Sx, Sy, B) diff --git a/src/plan.jl b/src/plan.jl index faec158b..7c65e909 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -161,20 +161,6 @@ function first_twiddles(::Type{T}, W::Vector{ComplexF64}, r::Int) where {T} t end -""" - Pruning{FM,LM} - -Pruned first and last stage of a chain on lines with zero inputs or unneeded outputs: the first -stage runs segment i over p ∈ `first[i]` with zero-input mask FM[i], the last stage segment i over -q ∈ `last[i]` (q0 in lane blocks) with needed-output mask LM[i]. The masks are type parameters so -that each segment calls its own generated codelet; an empty tuple means no pruning there. -""" -struct Pruning{FM,LM} - first::Vector{UnitRange{Int}} - last::Vector{UnitRange{Int}} -end -const NOPRUNE = Pruning{(),()}(UnitRange{Int}[], UnitRange{Int}[]) - # Stage calls of a chain from `src` through `dsts` (one destination per stage). `base` is a # number of complex values known to divide S (1 if nothing is known); stage i has stride # S·(R_1⋯R_(i-1)) and takes the widest vector type whose complex count divides base·R_1⋯R_(i-1). @@ -182,37 +168,36 @@ const NOPRUNE = Pruning{(),()}(UnitRange{Int}[], UnitRange{Int}[]) # Stage calls with the q (or butterfly) remainder known at generation time: divisible stages # take the plain kernels (the same calls as before the tail batching), single lines with a # scalar remainder get the vectorised range plus the scalar tail emitted directly. -# rem === nothing (a runtime stride, e.g. an odd batch size, or pruned ranges) takes the +# rem === nothing (a runtime stride, e.g. an odd batch size) takes the # splitting wrappers, which divide at run time. function mid_call(dst, src, tw, r, σ, V, V1, rem) rem === nothing && return :(stage_split!($dst, $src, $tw, m, s, Val($r), Val($σ), $V)) rem == 0 && return :(stage!($dst, $src, $tw, m, s, Val($r), Val($σ), $V)) quote - stage_range!($dst, $src, $tw, m, s, 0, s - $rem, 0, m, Val($r), Val($σ), $V, Val(0)) - stage_range!($dst, $src, $tw, m, s, s - $rem, s, 0, m, Val($r), Val($σ), $V1, Val(0)) + stage_range!($dst, $src, $tw, m, s, 0, s - $rem, 0, m, Val($r), Val($σ), $V) + stage_range!($dst, $src, $tw, m, s, s - $rem, s, 0, m, Val($r), Val($σ), $V1) end end function last_call(dst, src, r, σ, V, V1, rem) - need = (1 << r) - 1 rem === nothing && return :(stage_last_split!($dst, $src, s, Val($r), Val($σ), $V)) rem == 0 && return :(stage_last!($dst, $src, s, Val($r), Val($σ), $V)) quote - stage_last!($dst, $src, s, 0, s - $rem, Val($r), Val($σ), $V, Val($need)) - stage_last!($dst, $src, s, s - $rem, s, Val($r), Val($σ), $V1, Val($need)) + stage_last!($dst, $src, s, 0, s - $rem, Val($r), Val($σ), $V) + stage_last!($dst, $src, s, s - $rem, s, Val($r), Val($σ), $V1) end end function group_call(dst, src, tw1, r, σ, GV, V1, rem) rem === nothing && return :(stage_first_group!($dst, $src, $tw1, m, Val($r), Val($σ), $GV)) - rem == 0 && return :(stage_first_group_range!($dst, $src, $tw1, m, 0, m, Val($r), Val($σ), $GV, Val(0))) + rem == 0 && return :(stage_first_group_range!($dst, $src, $tw1, m, 0, m, Val($r), Val($σ), $GV)) quote - stage_first_group_range!($dst, $src, $tw1, m, 0, m - $rem, Val($r), Val($σ), $GV, Val(0)) - stage_first_group_1!($dst, $src, $tw1, m, m - $rem, m, Val($r), Val($σ), $V1, Val(0)) + stage_first_group_range!($dst, $src, $tw1, m, 0, m - $rem, Val($r), Val($σ), $GV) + stage_first_group_1!($dst, $src, $tw1, m, m - $rem, m, Val($r), Val($σ), $V1) end end -function chain_body(T, Rs, σ, src, dsts, base, single, FM = (), LM = ()) +function chain_body(T, Rs, σ, src, dsts, base, single) K = length(Rs) n = prod(Rs; init = 1) ladder = vec_ladder(T) @@ -233,40 +218,20 @@ function chain_body(T, Rs, σ, src, dsts, base, single, FM = (), LM = ()) ladder[findfirst(V -> (base * sst) % ncomplex(V) == 0, ladder)] # the q remainder of this stage: a compile-time constant for single lines, provably zero # for batched lines (their stride is a multiple of the vector type), runtime for odd batch - # sizes and pruned ranges + # sizes rem = single ? sst % ncomplex(V) : (base > 1 ? 0 : nothing) if i == K - if isempty(LM) || K == 1 - push!(body, last_call(dst, src, r, σ, V, single_vec(T), rem)) - else - for (j, need) in enumerate(LM) - push!(body, :(stage_last!($dst, $src, s, first(Pr.last[$j]), last(Pr.last[$j]) + 1, Val($r), Val($σ), $V, Val($need)))) - end - end + push!(body, last_call(dst, src, r, σ, V, single_vec(T), rem)) elseif single && i == 1 && group !== nothing && ncomplex(V) == 1 # the first stage of a single line: vectorised over the butterflies (scalar remainder # handled by the kernel) push!(body, :(m = L ÷ $r)) GV = ladder[group] - if isempty(FM) - push!(body, group_call(dst, src, :(cptr(P.tw1)), r, σ, GV, single_vec(T), (n ÷ r) % ncomplex(GV))) - else - for (j, zin) in enumerate(FM) - push!(body, :(stage_first_group!($dst, $src, cptr(P.tw1), m, first(Pr.first[$j]), last(Pr.first[$j]) + 1, - Val($r), Val($σ), $(ladder[group]), Val($zin)))) - end - end + push!(body, group_call(dst, src, :(cptr(P.tw1)), r, σ, GV, single_vec(T), (n ÷ r) % ncomplex(GV))) push!(body, :(L = m), :(s *= $r)) else push!(body, :(m = L ÷ $r)) - if i > 1 || isempty(FM) - push!(body, mid_call(dst, src, :(cptr(P.tw[$i])), r, σ, V, single_vec(T), rem)) - else - for (j, zin) in enumerate(FM) - push!(body, :(stage!($dst, $src, cptr(P.tw[$i]), m, s, first(Pr.first[$j]), last(Pr.first[$j]) + 1, - Val($r), Val($σ), $V, Val($zin)))) - end - end + push!(body, mid_call(dst, src, :(cptr(P.tw[$i])), r, σ, V, single_vec(T), rem)) push!(body, :(L = m), :(s *= $r)) end src = dst @@ -292,10 +257,7 @@ end Transform S interleaved sequences of length `P.n` from `px` into `py` (which may equal `px`), using the work buffers `pw1`, `pw2` (S·n complex values each). """ -run!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::Chain{T}, S::Int) where {T} = - run!(py, px, pw1, pw2, P, S, NOPRUNE) - -function run_body(::Type{T}, Rs, σ::Int, K::Int, FM, LM) where {T} +function run_body(::Type{T}, Rs, σ::Int, K::Int) where {T} K == 0 && return :(py == px || unsafe_copyto!(py, px, 2S); nothing) # Out of place: the stages alternate between py and pw1 and end in py; px is only read. oop = [iseven(K - i) ? :py : :pw1 for i in 1:K] @@ -307,7 +269,7 @@ function run_body(::Type{T}, Rs, σ::Int, K::Int, FM, LM) where {T} ip[i] = isodd(K - i) ? :pw1 : :pw2 end dispatch(dsts) = stride_dispatch(T, (base, single) -> quote - $(chain_body(T, Rs, σ, :px, dsts, base, single, FM, LM)...) + $(chain_body(T, Rs, σ, :px, dsts, base, single)...) end) inplace = dispatch(ip) quote @@ -322,9 +284,8 @@ function run_body(::Type{T}, Rs, σ::Int, K::Int, FM, LM) where {T} end end -@generated function run!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::Chain{T,Rs,σ,K}, S::Int, - Pr::Pruning{FM,LM}) where {T,Rs,σ,K,FM,LM} - run_body(T, Rs, σ, K, FM, LM) +@generated function run!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::Chain{T,Rs,σ,K}, S::Int) where {T,Rs,σ,K} + run_body(T, Rs, σ, K) end """ @@ -608,10 +569,7 @@ end Transform a block of B lanes; element j of lane i at `px + i + Sx·j` and `py + i + Sy·j` (py may equal px). The work buffers hold B·n complex values each. """ -run_lanes!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::Chain{T}, Sx::Int, Sy::Int, B::Int) where {T} = - run_lanes!(py, px, pw1, pw2, P, Sx, Sy, B, NOPRUNE) - -function run_lanes_body(::Type{T}, Rs, σ::Int, K::Int, FM, LM) where {T} +function run_lanes_body(::Type{T}, Rs, σ::Int, K::Int) where {T} function body(V) ex = Any[:(L = P.n), :(s0 = 1)] src, srcS = :px, :Sx @@ -620,24 +578,10 @@ function run_lanes_body(::Type{T}, Rs, σ::Int, K::Int, FM, LM) where {T} last = i == K dst, dstS = last ? (:py, :Sy) : (isodd(i) ? :pw1 : :pw2, :B) if last - if isempty(LM) - push!(ex, :(stage_last_lanes!($dst, $src, s0, $srcS, $dstS, B, Val($r), Val($σ), $V))) - else - for (j, need) in enumerate(LM) - push!(ex, :(stage_last_lanes!($dst, $src, s0, $srcS, $dstS, B, first(Pr.last[$j]), last(Pr.last[$j]) + 1, - Val($r), Val($σ), $V, Val($need)))) - end - end + push!(ex, :(stage_last_lanes!($dst, $src, s0, $srcS, $dstS, B, Val($r), Val($σ), $V))) else push!(ex, :(m = L ÷ $r)) - if i > 1 || isempty(FM) - push!(ex, :(stage_lanes!($dst, $src, cptr(P.tw[$i]), m, s0, $srcS, $dstS, B, Val($r), Val($σ), $V))) - else - for (j, zin) in enumerate(FM) - push!(ex, :(stage_lanes!($dst, $src, cptr(P.tw[$i]), m, s0, $srcS, $dstS, B, first(Pr.first[$j]), last(Pr.first[$j]) + 1, - Val($r), Val($σ), $V, Val($zin)))) - end - end + push!(ex, :(stage_lanes!($dst, $src, cptr(P.tw[$i]), m, s0, $srcS, $dstS, B, Val($r), Val($σ), $V))) push!(ex, :(L = m), :(s0 *= $r)) end src, srcS = dst, dstS @@ -670,8 +614,8 @@ function run_lanes_body(::Type{T}, Rs, σ::Int, K::Int, FM, LM) where {T} end @generated function run_lanes!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::Chain{T,Rs,σ,K}, - Sx::Int, Sy::Int, B::Int, Pr::Pruning{FM,LM}) where {T,Rs,σ,K,FM,LM} - run_lanes_body(T, Rs, σ, K, FM, LM) + Sx::Int, Sy::Int, B::Int) where {T,Rs,σ,K} + run_lanes_body(T, Rs, σ, K) end # y[i + Sy·j] = x[i + Sx·j] for B lanes and n elements diff --git a/src/rchain.jl b/src/rchain.jl index 612fe70d..5e8b5b8f 100644 --- a/src/rchain.jl +++ b/src/rchain.jl @@ -54,15 +54,14 @@ for T in (Float32, Float64) ladder = vec_ladder(T) LTYPES = Tuple(ladder) NR = length(RADICES); NL = length(ladder) - need(r) = (1 << r) - 1 kinds = ( - (:cstage!, (:py, :px, :pt, :m, :s), (r, σ, V, nd) -> :(stage!(py, px, pt, m, s, Val($r), Val($σ), $V))), - (:cstage_range!, (:py, :px, :pt, :m, :s, :qlo, :qhi), (r, σ, V, nd) -> :(stage_range!(py, px, pt, m, s, qlo, qhi, 0, m, Val($r), Val($σ), $V, Val(0)))), - (:cstage_last!, (:py, :px, :s), (r, σ, V, nd) -> :(stage_last!(py, px, s, Val($r), Val($σ), $V))), - (:cstage_last_range!, (:py, :px, :s, :qlo, :qhi), (r, σ, V, nd) -> :(stage_last!(py, px, s, qlo, qhi, Val($r), Val($σ), $V, Val($nd)))), - (:cstage_group!, (:py, :px, :pt, :m), (r, σ, V, nd) -> :(stage_first_group!(py, px, pt, m, Val($r), Val($σ), $V))), - (:cstage_lanes!, (:py, :px, :pt, :m, :s0, :sx, :sy, :B), (r, σ, V, nd) -> :(stage_lanes!(py, px, pt, m, s0, sx, sy, B, Val($r), Val($σ), $V))), - (:cstage_last_lanes!, (:py, :px, :s0, :sx, :sy, :B), (r, σ, V, nd) -> :(stage_last_lanes!(py, px, s0, sx, sy, B, Val($r), Val($σ), $V)))) + (:cstage!, (:py, :px, :pt, :m, :s), (r, σ, V) -> :(stage!(py, px, pt, m, s, Val($r), Val($σ), $V))), + (:cstage_range!, (:py, :px, :pt, :m, :s, :qlo, :qhi), (r, σ, V) -> :(stage_range!(py, px, pt, m, s, qlo, qhi, 0, m, Val($r), Val($σ), $V))), + (:cstage_last!, (:py, :px, :s), (r, σ, V) -> :(stage_last!(py, px, s, Val($r), Val($σ), $V))), + (:cstage_last_range!, (:py, :px, :s, :qlo, :qhi), (r, σ, V) -> :(stage_last!(py, px, s, qlo, qhi, Val($r), Val($σ), $V))), + (:cstage_group!, (:py, :px, :pt, :m), (r, σ, V) -> :(stage_first_group!(py, px, pt, m, Val($r), Val($σ), $V))), + (:cstage_lanes!, (:py, :px, :pt, :m, :s0, :sx, :sy, :B), (r, σ, V) -> :(stage_lanes!(py, px, pt, m, s0, sx, sy, B, Val($r), Val($σ), $V))), + (:cstage_last_lanes!, (:py, :px, :s0, :sx, :sy, :B), (r, σ, V) -> :(stage_last_lanes!(py, px, s0, sx, sy, B, Val($r), Val($σ), $V)))) for (name, args, call) in kinds typed = [a in (:py, :px, :pt) ? :($a::Ptr{$T}) : :($a::Int) for a in args] cases = Any[] @@ -73,7 +72,7 @@ for T in (Float32, Float64) name === :cstage_group! && !(ncomplex(V) in (2, 4)) && continue r in LAZY_RADICES && continue c = ((si - 1) * NL + vi - 1) * NR + ri - push!(cases, :(c == $c && return $(call(r, σ, V, need(r))))) + push!(cases, :(c == $c && return $(call(r, σ, V)))) end # Without lazy radices there is no run-time call at all (the trimmed library needs none). fallback = isempty(LAZY_RADICES) ? :(throw(ArgumentError("no codelet for this index"))) : @@ -91,7 +90,7 @@ for T in (Float32, Float64) table = Symbol(name, :_, T, :_codelets) @eval struct $callable{R,σ,VT} end @eval function (::$callable{R,σ,VT})($(typed...)) where {R,σ,VT} - return $(call(:R, :σ, :VT, :((1 << R) - 1))) + return $(call(:R, :σ, :VT)) end @eval const $table = Any[nothing for _ in 1:2 * $NL * $NR] # Index c = ((si - 1) * NL + vi - 1) * NR + ri, with ri = 1 … NR into RADICES. diff --git a/test/inference.jl b/test/inference.jl index 48b35175..f9a0872a 100644 --- a/test/inference.jl +++ b/test/inference.jl @@ -44,7 +44,7 @@ @test (@inferred mul!(y, p, x)) === y @test (@allocated mul!(y, p, x)) == 0 end - # pass plans on views of a padded array: unpruned in place and out of place, pruned in place + # pass plans on views of a padded array, in place and out of place let G = randn(Complex{T}, 40, 36, 30), W = similar(G) v = view(G, 3:34, 2:33, 2:25) w = view(W, 3:34, 2:33, 2:25) @@ -57,19 +57,6 @@ mul!(w, p, v) @test (@inferred mul!(w, p, v)) === w @test (@allocated mul!(w, p, v)) == 0 - n = 64 - N = 10 - rng = (1:(N - N ÷ 2), (n - N ÷ 2 + 1):n) - P = randn(Complex{T}, n + 4, 37) - u = view(P, 2 .+ (1:n), 3:35) - pnz = FFTL.plan_pass(u, 1; sign = -1, threads = 1, nonzero = rng) - FFTL.apply!(pnz, u) - @test (@inferred FFTL.apply!(pnz, u)) === u - @test (@allocated FFTL.apply!(pnz, u)) == 0 broken = VERSION < v"1.11" - pne = FFTL.plan_pass(u, 1; sign = 1, threads = 1, needed = rng) - FFTL.apply!(pne, u) - @test (@inferred FFTL.apply!(pne, u)) === u - @test (@allocated FFTL.apply!(pne, u)) == 0 broken = VERSION < v"1.11" @test_broken @inferred FFTL.plan_pass(v, 2) end # real transforms diff --git a/test/runtests.jl b/test/runtests.jl index b3a8351e..81f6e37e 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -311,7 +311,7 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end @testset "passes on strided views, $T" for T in (Float64, Float32) - # views of a padded array as an NFFT uses them: interior boxes and pruned boxes + # views of a padded array as an NFFT uses them: interior boxes and boxes at the edge G0 = randn(Complex{T}, 40, 36, 30) boxes = pick([(3:34, 2:33, 2:25), (3:6, 2:33, 2:25), (3:34, 5:9, 2:25), (1:40, 1:36, 1:30), (2:38, 2:36, 3:3), (5:40, 1:36, 4:28)], @@ -348,34 +348,6 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end end - @testset "pruned passes, $T" for T in (Float64, Float32) - # NFFT pattern: N modes in a line of length n, the first N - N÷2 and the last N÷2 entries - for (n, N) in pick([(128, 64), (100, 40), (96, 70), (240, 180), (1024, 600), (162, 120), (64, 10)], - thin(T, [(128, 64), (96, 70), (64, 10)])), d in 1:2, nt in pick((1, 3), (3,)) - a = N - N ÷ 2; b = N ÷ 2 - rng = (1:a, n-b+1:n) - dims = d == 1 ? (n + 4, 37) : (37, n + 4) - idx(r) = d == 1 ? (2 .+ r, 3:35) : (3:35, 2 .+ r) - G0 = randn(Complex{T}, dims) - line_zero = trues(n); for r in rng; line_zero[r] .= false; end - for sign in (-1, 1) - # zero inputs outside the ranges - G = copy(G0); v = view(G, idx(1:n)...) - sel = d == 1 ? (findall(line_zero), :) : (:, findall(line_zero)) - v[sel...] .= 0 - ref = mapslices(u -> dft(ComplexF64.(u), sign), Array(v); dims = d) - FFTL.apply!(FFTL.plan_pass(v, d; sign, threads = nt, nonzero = rng), v) - @test relerr(v, ref) < 10tol(T) - # only the outputs in the ranges are needed - G = copy(G0); v = view(G, idx(1:n)...) - ref = mapslices(u -> dft(ComplexF64.(u), sign), Array(v); dims = d) - FFTL.apply!(FFTL.plan_pass(v, d; sign, threads = nt, needed = rng), v) - keep = d == 1 ? (findall(.!line_zero), :) : (:, findall(.!line_zero)) - @test relerr(v[keep...], ref[keep...]) < 10tol(T) - end - 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),