diff --git a/docs/src/explanation/design-decisions.md b/docs/src/explanation/design-decisions.md index ca1600e2..437bef8a 100644 --- a/docs/src/explanation/design-decisions.md +++ b/docs/src/explanation/design-decisions.md @@ -35,6 +35,12 @@ and the large losses of an estimate-only planner are cache effects the algorithm A per-machine calibration or an optional measuring mode can follow if the rule leaves too much on the table. +One such cache effect is part of the rule: single lines on AVX-512 take radix 16 only while the 16 +streams of a stage are less than an L1 set stride (4 KB) apart, that is up to 4096 for `Float32` and +2048 for `Float64`. Beyond that the streams share one cache set and radix 8 wins. AMD EPYC 9354, one +line, all buffers at fixed page offsets, time of the radix-16 sequence over the radix-8 sequence: +`Float32` 4096: 0.88, 8192: 1.06, 16384: 1.06; `Float64` 4096: 1.33, 8192: 1.03, 16384: 1.11. + ## Rules for lengths are derived on some lengths and checked on others **Decision (2026-11-04).** The radix choice and the Bluestein length are rules on the prime @@ -95,6 +101,34 @@ bodies the `@generated` methods build. The `@generated` methods remain as the fa **Why.** A package image keeps the complete optimised IR of every precompiled `@generated` specialisation, 130 MB of the 146 MB image. See [Package size and load time](@ref "Package size and load time"). +## The first stage of a single line runs at full vector width + +**Decision (2026-11-12).** Where the widest vector holds 8 complex values (`Float32` on AVX-512), the +first stage of a single line runs `stage_small!`: vectorised over the butterflies, with the outputs +interleaved in registers into contiguous stores. It exists for the power-of-two radices and belongs to +the fixed set of plain methods. + +**Why.** The grouped first stage handles at most 4 complex values per vector. Intel Core i7-1185G7, +one line, `Float32`, lengths 128–16384 with radices of 16s and 8s: 12–26 % faster (the same code +against itself differs by up to 3 %). + +**Not done.** The same for later stages with a stride of 2 or 4. It needs a per-lane twiddle table and +one kernel per stride, and only `Float32` lengths whose radices start with 4, 4 have such a stage. On an +AMD EPYC 9354 it gained 6–8 % at 192 and 320 and nothing measurable at other lengths. + +## Stage kernels keep twiddles in memory and address streams from few registers + +**Decision (2026-11-12).** In the stages vectorised over q, the twiddle multiply reads re(w) and im(w) +from the table (`ccmul_at`; memory operands on x86), and the R input and R output streams are addressed +as base + odd·stride·2^e from one base per 8 streams (`stream_addr`). The base pointer is hidden from +LLVM's loop optimiser (`opaque_ptr`), which otherwise goes back to one pointer per stream. + +**Why.** A radix-16 stage has 15 twiddles and 32 streams, more than fit in registers, so its inner loop +reloaded vectors and pointers from the stack. AMD EPYC 9354 (AVX-512), instructions per butterfly: +radix 16 middle stage 284 → 198, radix 8 100 → 79. One line, `plan_fft`, medians of 8 interleaved +rounds: lengths 512–16384 take 0.79–0.94 (`Float32`) and 0.82–0.92 (`Float64`) of the time, lengths +from 32768 on are unchanged (the same code against itself differs by 1–5 %). Not verified on aarch64. + ## Codelets of the radices 17–31 are compiled on first use **Decision (2026-11-10).** The codelets of the radices 17, 19, 23, 29 and 31 are not part of the package diff --git a/src/complexvec.jl b/src/complexvec.jl index c02fdf86..3e7bd9ea 100644 --- a/src/complexvec.jl +++ b/src/complexvec.jl @@ -98,6 +98,7 @@ if HAS_FCMA @inline fcmla90(c::Vec{$N,$T}, a::Vec{$N,$T}, b::Vec{$N,$T}) = Vec(ccall($("llvm.aarch64.neon.vcmla.rot90." * suf), llvmcall, $VT, ($VT, $VT, $VT), c.data, a.data, b.data)) @inline ccmul(a::Vec{$N,$T}, w::Vec{$N,$T}) = fcmla90(fcmla0(zero(Vec{$N,$T}), a, w), a, w) + @inline ccmul_at(a::Vec{$N,$T}, p::Ptr{$T}) = ccmul(a, cbroadcast(Vec{$N,$T}, p)) end end else @@ -129,11 +130,17 @@ else end end end - # a·w = (a.re·w.re - a.im·w.im, a.re·w.im + a.im·w.re) - @inline function ccmul(a::Vec{N,T}, w::Vec{N,T}) where {N,T} - t = dup_im(a) * swap_reim(w) - blend_reim(muladd(dup_re(a), w, -t), muladd(dup_re(a), w, t)) + # a·w = (a.re·w.re - a.im·w.im, a.im·w.re + a.re·w.im) from re(w) and im(w), each repeated over + # the two lanes of its complex value + @inline function ccmul_parts(a::Vec{N,T}, wre::Vec{N,T}, wim::Vec{N,T}) where {N,T} + t = swap_reim(a) * wim + blend_reim(muladd(a, wre, -t), muladd(a, wre, t)) end + @inline ccmul(a::Vec{N,T}, w::Vec{N,T}) where {N,T} = ccmul_parts(a, dup_re(w), dup_im(w)) + # a·w for the complex value w at p, repeated over all complex lanes: re(w) and im(w) are scalar + # broadcasts, which x86 folds into the multiplies as memory operands + @inline ccmul_at(a::Vec{N,T}, p::Ptr{T}) where {N,T} = + ccmul_parts(a, Vec{N,T}(unsafe_load(p)), Vec{N,T}(unsafe_load(p, 2))) end # (re, im) repeated over all complex lanes of the register diff --git a/src/factor.jl b/src/factor.jl index 8d95ffe8..869fd86f 100644 --- a/src/factor.jl +++ b/src/factor.jl @@ -7,6 +7,10 @@ # up to 4096 on the M1; on Tiger Lake (AVX-512) 7–30 % from 4096 on (S = 1) and up to 45 % in # batches, within noise below; on Zen 3 (AVX2, 16 # registers) it is 3–25 % slower from 256 on. +# Single lines on AVX-512 take radix 16 as often as possible and 8s for the rest while the 16 +# streams of a stage are less than an L1 set stride apart (Float32 up to 4096, Float64 up to +# 2048); beyond that they share one cache set and the 8s win (Zen 4: 3–33 % at 4096–16384). +# Not measured on NEON. # Otherwise radix 4 throughout, one 8 (or 2) for an odd exponent; radix 16 and above lose # (register spills, cache-set conflicts at large strides). # - aarch64: a factor 3 is merged with a 4 into 12, else with a 2 into 6; a factor 5 with a spare @@ -59,7 +63,7 @@ Radix sequence for a chain of length `n` applied to `S` interleaved sequences of `esize` bytes. """ function radices(n::Int, S::Int = 1; esize::Int = 16) - r = heuristic_radices(n, S) + r = heuristic_radices(n, S, esize) S > LARGE_STRIDE && set_conflict(r, n, S, esize) || return r alt = packed_radices(n) set_conflict(alt, n, S, esize) ? r : alt @@ -94,7 +98,7 @@ function packed_radices(n::Int) [reverse!(sort!(filter(ispow2, bins))); reverse!(sort!(filter(!ispow2, bins)))] end -function heuristic_radices(n::Int, S::Int) +function heuristic_radices(n::Int, S::Int, esize::Int = 16) n == 1 && return Int[] merge = merge_composite(S) f, rest = factorize(n) @@ -105,6 +109,11 @@ function heuristic_radices(n::Int, S::Int) end if RADIX8_POW2 && ispow2(n) && 64 < n <= 16384 && S <= LARGE_STRIDE k = trailing_zeros(n) + if HAS_AVX512 && S == 1 && (n ÷ 16) * esize < 1 << L1_SET_LOG2 + # as many 16s as possible, the rest in 8s: 4t ≤ k, t ≡ k (mod 3) + t = k % 3 + 3 * ((k ÷ 4 - k % 3) ÷ 3) + return vcat(fill(16, t), fill(8, (k - 4t) ÷ 3)) + end spare = k % 3 return vcat(spare == 2 ? [4] : spare == 1 ? [4, 4] : Int[], fill(8, (k - (spare == 2 ? 2 : spare == 1 ? 4 : 0)) ÷ 3)) end diff --git a/src/kernel_methods.jl b/src/kernel_methods.jl index e90b555d..5aa839c0 100644 --- a/src/kernel_methods.jl +++ b/src/kernel_methods.jl @@ -40,6 +40,15 @@ for T in (Float32, Float64), (name, head, vts) in KERNEL_HEADS end end +# stage_small!: the widest vector type, power-of-two radices. +for T in (Float32, Float64), σ in (-1, 1), R in RADICES + VT = first(vec_ladder(T)) + (ispow2(R) && !(R in LAZY_RADICES)) || continue + h = :(stage_small!(py::Ptr{$T}, px::Ptr{$T}, pt::Ptr{$T}, m::Int, ::Val{$R}, ::Val{$σ}, ::Type{$VT})) + b = kernel_subst(stage_small_body(T, R, σ, VT), Dict{Symbol,Any}(:T => T, :VT => VT)) + @eval $(Expr(:function, h, b)) +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 uses the generic methods. for T in (Float32, Float64), (n, σ) in SMALL_ENTRIES diff --git a/src/kernels.jl b/src/kernels.jl index 8120237b..46b26928 100644 --- a/src/kernels.jl +++ b/src/kernels.jl @@ -18,19 +18,22 @@ stage!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, s::Int, ::Val{R}, ::Val{σ}, 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] - 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))] + loads = [:($(a[j]) = vload($VT, $(stream_addr(:xb, :xr, j - 1)))) for j in 1:R] + twl = [:($(Symbol("w", k)) = pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T))) for k in 1:R-1] + stores = Any[:(vstore($(X[1]), yb0))] for k in 1:R-1 - push!(stores, :(vstore(ccmul($(X[k+1]), $(Symbol("w", k))), py + (o + $k * s) * $B))) + push!(stores, :(vstore(ccmul_at($(X[k+1]), $(Symbol("w", k))), $(stream_addr(:yb, :yr, k))))) end quote + $(stream_strides(:xr, R, :(s * m * $B))...) + $(stream_strides(:yr, R, :(s * $B))...) for p in plo:phi-1 $(twl...) for q in 0:$C:s-1 + $(stream_bases(:xb, :xr, R, :(px + (q + s * p) * $B))...) $(loads...) $(stmts...) - o = q + s * $R * p + $(stream_bases(:yb, :yr, R, :(py + (q + s * $R * p) * $B))...) $(stores...) end end @@ -60,22 +63,51 @@ 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. +# base + i·2^e as one scaled-index address computation (a GEP over an element of 2^e bytes), so +# that the scaling stays in the addressing mode instead of becoming a value of its own. +for (e, ty) in enumerate(("i8", "i16", "i32", "i64")) + ir = "define ptr @entry(ptr %0, i64 %1) #0 {\ntop:\n %p = getelementptr $ty, ptr %0, i64 %1\n ret ptr %p\n}\nattributes #0 = { alwaysinline }" + @eval @inline scaled_ptr(b::Ptr{T}, i::Int, ::Val{$(e - 1)}) where {T} = Base.llvmcall(($ir, "entry"), Ptr{T}, Tuple{Ptr{T},Int}, b, i) +end + +# The pointer p, opaque to the loop optimisations: an induction pointer passed through here is +# not rewritten into one pointer per stream. +@inline opaque_ptr(p::Ptr{T}) where {T} = Base.llvmcall( + ("declare ptr @llvm.launder.invariant.group.p0(ptr)\ndefine ptr @entry(ptr %0) #0 {\ntop:\n %p = call ptr @llvm.launder.invariant.group.p0(ptr %0)\n ret ptr %p\n}\nattributes #0 = { alwaysinline }", "entry"), + Ptr{T}, Tuple{Ptr{T}}, p) + +# Address of stream j (zero-based) of streams r bytes apart: base + odd·r·2^e with odd in +# (1, 3, 5, 7) and one base per 8 streams, so that a few registers reach every stream. +function stream_addr(base, r, j) + b = Symbol(base, j ÷ 8); k = j % 8 + k == 0 && return b + e = trailing_zeros(k) + :(scaled_ptr($b, $(Symbol(r, k >> e)), Val($e))) +end +stream_bases(base, r, R, first) = + [i == 0 ? :($(Symbol(base, 0)) = opaque_ptr($first)) : :($(Symbol(base, i)) = scaled_ptr($(Symbol(base, i - 1)), $(Symbol(r, 1)), Val(3))) for i in 0:(R-1)÷8] +stream_strides(r, R, unit) = + vcat(:($(Symbol(r, 1)) = $unit), [:($(Symbol(r, o)) = $o * $(Symbol(r, 1))) for o in (3, 5, 7) if o < R]) + 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] - 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))] + loads = [:($(a[j]) = vload($VT, $(stream_addr(:xb, :xr, j - 1)))) for j in 1:R] + twl = [:($(Symbol("w", k)) = pt + ($(R - 1) * p + $(k - 1)) * $(2sizeof(T))) for k in 1:R-1] + stores = Any[:(vstore($(X[1]), yb0))] for k in 1:R-1 - push!(stores, :(vstore(ccmul($(X[k+1]), $(Symbol("w", k))), py + (o + $k * s) * $B))) + push!(stores, :(vstore(ccmul_at($(X[k+1]), $(Symbol("w", k))), $(stream_addr(:yb, :yr, k))))) end quote + $(stream_strides(:xr, R, :(s * m * $B))...) + $(stream_strides(:yr, R, :(s * $B))...) for p in plo:phi-1 $(twl...) for q in qlo:$C:qhi-1 + $(stream_bases(:xb, :xr, R, :(px + (q + s * p) * $B))...) $(loads...) $(stmts...) - o = q + s * $R * p + $(stream_bases(:yb, :yr, R, :(py + (q + s * $R * p) * $B))...) $(stores...) end end @@ -95,12 +127,15 @@ stage_last!(py::Ptr{T}, px::Ptr{T}, s::Int, ::Val{R}, ::Val{σ}, ::Type{VT}) whe 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] + loads = [:($(a[j]) = vload($VT, $(stream_addr(:xb, :r, j - 1)))) for j in 1:R] + stores = [:(vstore($(X[k+1]), $(stream_addr(:yb, :r, k)))) for k in 0:R-1] quote + $(stream_strides(:r, R, :(s * $B))...) for q in qlo:$C:qhi-1 + $(stream_bases(:xb, :r, R, :(px + q * $B))...) $(loads...) $(stmts...) + $(stream_bases(:yb, :r, R, :(py + q * $B))...) $(stores...) end nothing @@ -194,6 +229,79 @@ end stage_first_group_range_body(T, R, σ, VT) end +# Complex values of u ++ v (2C) in chunks of S, alternating u_0 v_0 u_1 v_1 …: the lane masks of +# the first C and the last C. +function chunk_masks(C, S) + idx = [(isodd(i ÷ S) ? C : 0) + (i ÷ S ÷ 2) * S + i % S for i in 0:2C-1] + lanes(ix) = Tuple(l for c in ix for l in (2c, 2c + 1)) + lanes(idx[1:C]), lanes(idx[C+1:2C]) +end + +# Interleave R = 2^r vectors in chunks of S complex values: pairs at S, then the lower and upper +# halves recursively at 2S; chunks of a whole vector are already in place. labs tracks +# (vector, complex value) of every complex value so the generator can check the permutation. +function ilv_net!(ex, ys, labs, S, C) + (length(ys) == 1 || S >= C) && return ys, labs + ml, mh = chunk_masks(C, S) + lo = Any[]; hi = Any[]; llo = Any[]; lhi = Any[] + for i in 1:2:length(ys) + cat = vcat(labs[i], labs[i+1]) + push!(llo, cat[collect(ml[1:2:end]) .÷ 2 .+ 1]); push!(lhi, cat[collect(mh[1:2:end]) .÷ 2 .+ 1]) + a = gensym(:lo); b = gensym(:hi) + push!(ex, :($a = shufflevector($(ys[i]), $(ys[i+1]), Val($ml))), + :($b = shufflevector($(ys[i]), $(ys[i+1]), Val($mh)))) + push!(lo, a); push!(hi, b) + end + o1, l1 = ilv_net!(ex, lo, llo, 2S, C) + o2, l2 = ilv_net!(ex, hi, lhi, 2S, C) + vcat(o1, o2), vcat(l1, l2) +end + +# Outputs X_k (lane p) in the order of y: vector t, lane l holds position tC + l = Rp + k. +function chunk_interleave!(ex, ys, C) + R = length(ys) + out, labs = ilv_net!(ex, ys, [[(k, l) for l in 0:C-1] for k in 0:R-1], 1, C) + for t in 0:R-1, l in 0:C-1 + labs[t+1][l+1] == ((t * C + l) % R, (t * C + l) ÷ R) || + error("interleave of $R vectors of $C complex values is wrong") + end + out +end + +# First stage of a single line (s = 1) at any vector width, for R a power of two and m a +# multiple of C: butterflies p, …, p+C-1 share a register (twiddles in the grouped first-stage +# table) and the R outputs are interleaved in registers into contiguous stores. +function stage_small_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 + (p + $(j - 1) * m) * $B)) for j in 1:R] + outs = Any[X[1]] + twm = Any[] + for k in 1:R-1 + b = Symbol("b", k) + push!(twm, :($b = ccmul($(X[k+1]), vload($VT, pt + (p + $(k - 1) * m) * $B)))) + push!(outs, b) + end + shuf = Any[] + ys = chunk_interleave!(shuf, outs, C) + stores = [:(vstore($(ys[t+1]), py + ($R * p + $(t * C)) * $B)) for t in 0:R-1] + quote + for p in 0:$C:m-1 + $(loads...) + $(stmts...) + $(twm...) + $(shuf...) + $(stores...) + end + nothing + end +end + +@generated function stage_small!(py::Ptr{T}, px::Ptr{T}, pt::Ptr{T}, m::Int, ::Val{R}, ::Val{σ}, + ::Type{VT}) where {T,R,σ,VT} + stage_small_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}) where {T,VT} diff --git a/src/plan.jl b/src/plan.jl index 7c65e909..f5d1bcd2 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -80,6 +80,12 @@ function build_chain(::Type{T}, n::Int, σ::Int, rs::AbstractVector{Int}) where Chain{T,Tuple(rs),σ,length(rs)}(n, Tuple(tw), tw1) end +# The first of several stages of a single line runs stage_small! (the widest vector, C complex +# values, outputs interleaved in registers) where the grouped first stage (vector index group +# of the complex counts ncs, 0 for none) would be narrower. +small_first(ncs, group::Int, r::Int, n::Int) = + r < n && ispow2(r) && (n ÷ r) % ncs[1] == 0 && ncs[1] > (group == 0 ? 1 : ncs[group]) + # Stage twiddles and first-stage table of a radix sequence, shared by Chain and RChain. chain_tables(::Type{T}, n::Int, σ::Int, rs::AbstractVector{Int}) where {T} = cached(() -> build_tables(T, n, σ, rs), CacheKey(:tables, T, n, σ, collect(Int, rs)))::Tuple{Vector{Vector{Complex{T}}},Vector{Complex{T}}} @@ -220,7 +226,10 @@ function chain_body(T, Rs, σ, src, dsts, base, single) # for batched lines (their stride is a multiple of the vector type), runtime for odd batch # sizes rem = single ? sst % ncomplex(V) : (base > 1 ? 0 : nothing) - if i == K + if single && i == 1 && small_first(map(ncomplex, ladder), something(group, 0), r, n) + push!(body, :(stage_small!($dst, $src, cptr(P.tw1), $(n ÷ r), Val($r), Val($σ), $(ladder[1]))), + :(L = L ÷ $r), :(s *= $r)) + elseif i == K 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 diff --git a/src/rchain.jl b/src/rchain.jl index 5e8b5b8f..583bee35 100644 --- a/src/rchain.jl +++ b/src/rchain.jl @@ -60,6 +60,7 @@ for T in (Float32, Float64) (: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_small!, (:py, :px, :pt, :m), (r, σ, V) -> :(stage_small!(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 @@ -70,6 +71,8 @@ for T in (Float32, Float64) # never selects another vector type for it, and code compiled ahead of time must not # contain the call (its generator would fail) name === :cstage_group! && !(ncomplex(V) in (2, 4)) && continue + # stage_small! exists for the widest vector type and power-of-two radices + name === :cstage_small! && !(vi == 1 && ispow2(r)) && continue r in LAZY_RADICES && continue c = ((si - 1) * NL + vi - 1) * NR + ri push!(cases, :(c == $c && return $(call(r, σ, V)))) @@ -115,6 +118,7 @@ const OP_MID_TAIL = 2 # q in vectors up to s - rem, the scalar remainder aft const OP_LAST = 3 const OP_LAST_TAIL = 4 const OP_GROUP = 5 # first stage vectorised over the butterflies p +const OP_SMALL = 6 # first stage over the butterflies p on the widest vector (stage_small!) struct StageCall op::Int @@ -154,7 +158,7 @@ end # The stage calls of a single line, with the rules of `chain_body` for S = 1: stage i takes the # widest vector type whose complex count fits its stride s = R_1⋯R_(i-1), the remainder # s mod C runs one complex value at a time; the first stage is grouped over the butterflies -# when its stride is 1. +# when its stride is 1, with stage_small! where that is wider. function single_schedule(::Type{T}, n::Int, σ::Int, rs::Vector{Int}) where {T} ncs = ladder_ncomplex(T) nv = length(ncs) @@ -177,7 +181,9 @@ function single_schedule(::Type{T}, n::Int, σ::Int, rs::Vector{Int}) where {T} rem = s % ncs[vi] c = codelet_index(T, σ, vi, r); c1 = codelet_index(T, σ, nv, r) m = L ÷ r - if i == K + if i == 1 && small_first(ncs, group, r, n) + push!(calls, StageCall(OP_SMALL, codelet_index(T, σ, 1, r), c1, m, s, 0)) + elseif i == K push!(calls, StageCall(rem == 0 ? OP_LAST : OP_LAST_TAIL, c, c1, m, s, rem)) elseif i == 1 && group != 0 && ncs[vi] == 1 push!(calls, StageCall(OP_GROUP, codelet_index(T, σ, group, r), c1, m, s, 0)) @@ -276,6 +282,8 @@ function run_single!(py::Ptr{T}, px::Ptr{T}, pw1::Ptr{T}, pw2::Ptr{T}, P::RChain cstage!(dst, src, cptr(P.tw[i]), st.m, st.s, st.c) elseif op == OP_GROUP cstage_group!(dst, src, cptr(P.tw1), st.m, st.c) + elseif op == OP_SMALL + cstage_small!(dst, src, cptr(P.tw1), st.m, st.c) elseif op == OP_LAST cstage_last!(dst, src, st.s, st.c) elseif op == OP_MID_TAIL diff --git a/test/data/radices-x86_64-32regs.csv b/test/data/radices-x86_64-32regs.csv index 890ce426..826e8734 100644 --- a/test/data/radices-x86_64-32regs.csv +++ b/test/data/radices-x86_64-32regs.csv @@ -192,7 +192,7 @@ 125,1024,5 5 5 126,1,2 7 9 126,1024,7 6 3 -128,1,4 4 8 +128,1,16 8 128,1024,8 4 4 130,1,2 5 13 130,1024,13 10 @@ -330,7 +330,7 @@ 253,1024,23 11 255,1,3 5 17 255,1024,17 5 3 -256,1,4 8 8 +256,1,16 16 256,1024,4 4 4 4 260,1,4 5 13 260,1024,4 5 13 @@ -878,9 +878,9 @@ 1020,1024,17 12 5 1023,1,3 11 31 1023,1024,31 11 3 -1024,1,4 4 8 8 +1024,1,16 8 8 1024,1024,4 4 4 4 4 -2048,1,4 8 8 8 +2048,1,16 16 8 2048,1024,8 4 4 4 4 4096,1,8 8 8 8 4096,1024,4 4 4 4 4 4 diff --git a/test/runtests.jl b/test/runtests.jl index a8d1a458..cd5f6849 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -753,6 +753,25 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end end + @testset "small-stride stages, $T" for T in (Float64, Float32) + # stage_small!: every power-of-two radix as the first stage, in both chain forms, against + # the direct DFT + p(a) = Ptr{T}(pointer(a)) + C = FFTL.ladder_ncomplex(T)[1] + for r in (2, 4, 8, 16), σ in (-1, 1) + rs = [r, 2C] + n = prod(rs) + c = FFTL.build_chain(T, n, σ, rs); rc = FFTL.RChain{T}(n, σ, rs) + @test (rc.single[1].op == FFTL.OP_SMALL) == (C > 4) + x = randn(Complex{T}, n); w1 = similar(x); w2 = similar(x); y1 = similar(x); y2 = similar(x) + GC.@preserve x y1 y2 w1 w2 begin + FFTL.run!(p(y1), p(x), p(w1), p(w2), c, 1); FFTL.run_single!(p(y2), p(x), p(w1), p(w2), rc) + end + @test y1 == y2 + @test relerr(y2, dft(ComplexF64.(x), σ)) < tol(T) + end + end + # run-time plans (RChain kernels): same results as the reference, four-step and plane fusion @testset "RPlan, $T" for T in pick((Float64, Float32), (Float64,)) for (sz, reg, nt) in pick([((1000,), [1], 1), ((2 * 3^11,), [1], 3), ((1009,), [1], 1), ((64, 48), [1, 2], 1),