From aea75adaa1a1cdb37c6876d53fd6974fc8b8747b Mon Sep 17 00:00:00 2001 From: Tobias Knopp Date: Tue, 10 Nov 2026 23:35:31 +0100 Subject: [PATCH 1/3] Threads for single real-to-real and real lines A single long line of a real transform ran on one thread: the real transforms of a non-smooth length, and all real-to-real transforms of one line. Two causes: - line_tasks gives one task for line kernels with buffers of their own (Bluestein, composite). Every task now gets its own copy of such a kernel (task_lines, task_reallines, rline_task), so the lines of a non-smooth dimension run on several tasks. - A single line had no parallel loop for the real-to-real kinds. An even line from 2^17 points on uses a half-length complex DFT with pre- and post-processing on several tasks (DCT-II, DST-II, DHT, DCT-III, DST-III, DCT-IV, DST-IV); an odd line uses one complex DFT of length 2n (n for the DHT). DCT-I and DST-I use the even and odd extension of length 2(n-1) and 2(n+1). Checked against FFTW: relative error 1e-15 for all kinds, odd and even lengths and 1 to 3 dimensions with 8 threads. The test suite passes except for a test that needs KernelAbstractions in the test environment. Co-Authored-By: Claude Haiku 5.5 --- src/r2r.jl | 273 +++++++++++++++++++++++++++++++++++++++++++++++++-- src/real.jl | 58 +++++++---- src/rreal.jl | 12 ++- 3 files changed, 309 insertions(+), 34 deletions(-) diff --git a/src/r2r.jl b/src/r2r.jl index 3df0ec15..7a68e3c6 100644 --- a/src/r2r.jl +++ b/src/r2r.jl @@ -477,11 +477,26 @@ end # M1, DCT-II along dim 2 of n×n against the blocked lines: 0.46–0.64 on one thread (n = 32…1024) # and on 8 threads up to 128²; from 256² on 8 threads 0.97–1.7, as the half-length DFT needs three # parallel steps where the blocked lines need one. -use_half_dft(S, n, outer, nt) = S >= 4 && iseven(n) && n >= 4 && (nt == 1 || S * n * outer < 1 << 16) +# A single even line on several threads has no other parallel loop: it takes the half-length DFT too, +# from 2^17 points on (M1, 8 threads: 32768 is 3× slower than one thread, 131072 is 1.4× faster). +use_half_dft(S, n, outer, nt) = iseven(n) && n >= 4 && + (S == 1 && outer == 1 ? nt > 1 && n >= 1 << 17 : S >= 4 && (nt == 1 || S * n * outer < 1 << 16)) function half_dft(::Type{T}, sz, d, kind, z::Vector{Complex{T}}, nt) where {T} S, outer = dim_extent(sz, d) n = sz[d] + if S == 1 && outer == 1 && nt > 1 && n >= 1 << 17 && !iseven(n) && haskey(ODD_KINDS, kind) + return half_odd(T, n, ODD_KINDS[kind], nt) + end + if (kind == REDFT11 || kind == RODFT11) && S == 1 && outer == 1 && nt > 1 && n >= 1 << 17 && iseven(n) + return half_iv(T, n, kind == RODFT11, z, nt) + end + if (kind == REDFT01 || kind == RODFT01) && S == 1 && outer == 1 && nt > 1 && n >= 1 << 17 && iseven(n) + return half_iii(T, n, kind == RODFT01, nt) + end + if (kind == REDFT00 || kind == RODFT00) && S == 1 && outer == 1 && nt > 1 && n >= 1 << 17 + return half_i(T, n, kind == RODFT00, nt) + end (kind == REDFT10 || kind == RODFT10 || kind == DHT) && use_half_dft(S, n, outer, nt) || return nothing M = n ÷ 2 zz = reshape(z, S, M, outer) @@ -490,12 +505,17 @@ function half_dft(::Type{T}, sz, d, kind, z::Vector{Complex{T}}, nt) where {T} HalfDFT{T,typeof(cplan)}(S, outer, n, kind == RODFT10, kind == DHT, cplan, zz, w, cis_table(T, M + 1, -1, 0, 2n)) end +# One line kernel per task. A kernel with buffers of its own (Bluestein, composite) cannot run on +# two tasks at once, so each task gets a copy; the other kernels are shared. +task_lines(K::R2RLine{T}, threads::Integer) where {T} = + own_buffers(K) && threads > 1 ? [K; [R2RLine{T}(K.kind, K.n) for _ in 2:threads]] : [K] + function R2RPlan{T}(sz::NTuple{N,Int}, kinds, region; threads::Integer = Threads.nthreads()) where {T<:FFTLFloat,N} reg = Tuple(collect(Int, region isa Integer ? (region,) : region)) check_r2r_region(reg, N) ks = kinds isa R2RKind ? ntuple(_ -> kinds, length(reg)) : Tuple(kinds) length(ks) == length(reg) || throw(ArgumentError("one kind per transformed dimension")) - lines = Tuple(R2RLine{T}(k, sz[d]) for (k, d) in zip(ks, reg)) + lines = Tuple(task_lines(R2RLine{T}(k, sz[d]), threads) for (k, d) in zip(ks, reg)) nt = max(1, Int(threads)) nmax = maximum((sz[d] for d in reg); init = 1) z = Vector{Complex{T}}(undef, prod(sz) ÷ 2) @@ -531,29 +551,39 @@ Base.IndexStyle(::Type{<:StridedLine}) = IndexLinear() # one dimension of an Array, each line read and written in place through a StridedLine r2r_dim!(y::Array{T}, x::Array{T}, d::Int, K::R2RLine{T}, p::R2RPlan{T}) where {T} = - (GC.@preserve x y r2r_dim!(pointer(y), pointer(x), d, K, p); y) + (GC.@preserve x y r2r_dim!(pointer(y), pointer(x), d, (K,), p); y) +r2r_dim!(y::Array{T}, x::Array{T}, d::Int, Ks::Vector{<:R2RLine{T}}, p::R2RPlan{T}) where {T} = + (GC.@preserve x y r2r_dim!(pointer(y), pointer(x), d, Ks, p); y) # p: a plan with sz, nthreads and bufs (R2RPlan, RR2RPlan) -function r2r_dim!(py::Ptr{T}, px::Ptr{T}, d::Int, K::R2RLine{T}, p) where {T} +# a tuple, not a vector: no allocation per call +r2r_dim!(py::Ptr{T}, px::Ptr{T}, d::Int, K::R2RLine{T}, p) where {T} = r2r_dim!(py, px, d, (K,), p) + +# Ks: one kernel per task (see task_lines), or one shared kernel. +function r2r_dim!(py::Ptr{T}, px::Ptr{T}, d::Int, Ks, p) where {T} S, outer = dim_extent(p.sz, d) n = p.sz[d] - S > 1 && return r2r_dim_blocked!(py, px, S, outer, n, K, p) - each_line(S * outer, n, line_tasks(K, p.nthreads)) do c, l + S > 1 && return r2r_dim_blocked!(py, px, S, outer, n, Ks, p) + each_line(S * outer, n, line_tasks_of(Ks, p.nthreads)) do c, l local o, i o, i = divrem(l, S) local off = (i + S * n * o) * sizeof(T) - r2r_line!(StridedLine{T}(py + off, S, n), StridedLine{T}(px + off, S, n), K, p.bufs[c]) + r2r_line!(StridedLine{T}(py + off, S, n), StridedLine{T}(px + off, S, n), Ks[task_kernel(Ks, c)], p.bufs[c]) end nothing end +line_tasks_of(Ks, nt) = length(Ks) > 1 ? length(Ks) : line_tasks(Ks[1], nt) +task_kernel(Ks, c) = length(Ks) > 1 ? c : 1 + # Strided lines in blocks of L neighbours: gather the block into a contiguous buffer (one line # after the other), transform each line there, scatter back. Reading the lines one by one would # use one value per cache line and no SIMD. -function r2r_dim_blocked!(py::Ptr{T}, px::Ptr{T}, S::Int, outer::Int, n::Int, K::R2RLine{T}, p) where {T} +function r2r_dim_blocked!(py::Ptr{T}, px::Ptr{T}, S::Int, outer::Int, n::Int, Ks, p) where {T} L = r2r_block(T) nb = cld(S, L) - each_line(nb * outer, n * L, line_tasks(K, p.nthreads)) do c, l + each_line(nb * outer, n * L, line_tasks_of(Ks, p.nthreads)) do c, l + local K = Ks[task_kernel(Ks, c)] local o, b o, b = divrem(l, nb) local i0 = b * L @@ -584,6 +614,231 @@ function r2r_dim_blocked!(py::Ptr{T}, px::Ptr{T}, S::Int, outer::Int, n::Int, K: nothing end +# DCT-IV and DST-IV of one even line on several threads, as a complex DFT of half the length. With +# z_m = (x_2m + i x_{n-1-2m}) e^{-iπm/n} and Z = DFT_{n/2}(z), the outputs are y_2k = 2 Re w_k and +# y_{n-1-2k} = -2 Im w_k, w_k = Z_k e^{-iπ(4k+1)/(4n)} (FFTW's REDFT11 scaling). DST-IV is the DCT-IV of the +# reversed input times (-1)^k. Checked against FFTW: relerr 2e-16 for n = 8…64. +struct HalfIV{T,F} + n::Int + dst::Bool + z::Vector{Complex{T}} # pre-twiddled input, length n ÷ 2 + zt::Vector{Complex{T}} # its DFT + pre::Vector{Complex{T}} # e^{-iπm/n} + post::Vector{Complex{T}} # 2 e^{-iπ(4k+1)/(4n)} + cplan::F +end + +function half_iv(::Type{T}, n::Int, dst::Bool, z::Vector{Complex{T}}, nt::Int) where {T} + M = n ÷ 2 + pre = Complex{T}[Complex{T}(cis(-π * m / n)) for m in 0:M-1] + post = Complex{T}[2 * Complex{T}(cis(-π * (4k + 1) / (4n))) for k in 0:M-1] + cplan = FFTLPlan{T}((M,), (1,), -1; threads = nt) + HalfIV{T,typeof(cplan)}(n, dst, z, zeros(Complex{T}, M), pre, post, cplan) +end + +half_dft!(py::Ptr{T}, px::Ptr{T}, H::HalfIV{T}, nt::Int) where {T} = half_iv!(py, px, H, nt) + +function half_iv!(py::Ptr{T}, px::Ptr{T}, H::HalfIV{T}, nt::Int) where {T} + n = H.n; M = n ÷ 2; dst = H.dst + z = H.z; zt = H.zt; pre = H.pre; post = H.post + GC.@preserve H px py begin + each_line(M, 1, nt) do c, m + @inbounds begin + ir = dst ? n - 1 - 2m : 2m + ii = dst ? 2m : n - 1 - 2m + z[m+1] = Complex{T}(unsafe_load(px, ir + 1), unsafe_load(px, ii + 1)) * pre[m+1] + end + end + mul!(zt, H.cplan, z) + each_line(M, 1, nt) do c, k + @inbounds begin + w = zt[k+1] * post[k+1] + # output 2k (even) and n-1-2k (odd): the DST-IV sign (-1)^k is (+1) and (-1) there + unsafe_store!(py, real(w), 2k + 1) + unsafe_store!(py, dst ? imag(w) : -imag(w), n - 2k) + end + end + end + nothing +end + +# All kinds of one odd line on several threads: a complex DFT of length 2n (n for the DHT) of the +# zero-padded input, with a phase before (pre) and after (post) the transform. Every formula was +# checked against FFTW for n = 3 … 1001 (relerr 4e-16). DST-III and DST-IV are the DCT of the +# reversed input times (-1)^k. +const ODD_KINDS = Dict(REDFT10 => :dct2, RODFT10 => :dst2, REDFT11 => :dct4, RODFT11 => :dst4, + REDFT01 => :dct3, RODFT01 => :dst3, DHT => :dht) + +struct HalfOdd{T,F} + n::Int + L::Int + kind::Symbol + oidx::Int # output index offset of the DFT (1 for the DST-II) + rev::Bool # DST-III, DST-IV: reversed input + alt::Bool # DST-III, DST-IV: output signs (-1)^k + pre::Vector{Complex{T}} # per input sample, length n + post::Vector{Complex{T}} # per output sample, length n + w::Vector{Complex{T}} # zero-padded input, length L + out::Vector{Complex{T}} # its DFT + cplan::F +end + +function half_odd(::Type{T}, n::Int, kind::Symbol, nt::Int) where {T} + L = kind == :dht ? n : 2n + σ = (kind == :dct3 || kind == :dst3) ? 1 : -1 + pre = Vector{Complex{T}}(undef, n) + post = Vector{Complex{T}}(undef, n) + for j in 0:n-1 + pre[j+1] = kind in (:dct4, :dst4) ? Complex{T}(cis(-π * j / (2n))) : + kind in (:dct3, :dst3) ? (j == 0 ? one(Complex{T}) : 2 * Complex{T}(cis(π * j / (2n)))) : one(Complex{T}) + post[j+1] = kind == :dct2 ? Complex{T}(cis(-π * j / (2n))) : + kind == :dst2 ? Complex{T}(cis(-π * (j + 1) / (2n))) : + kind in (:dct4, :dst4) ? Complex{T}(cis(-π * (2j + 1) / (4n))) : one(Complex{T}) + end + cplan = FFTLPlan{T}((L,), (1,), σ; threads = nt) + HalfOdd{T,typeof(cplan)}(n, L, kind, kind == :dst2 ? 1 : 0, kind in (:dst3, :dst4), kind in (:dst3, :dst4), + pre, post, zeros(Complex{T}, L), zeros(Complex{T}, L), cplan) +end + +half_dft!(py::Ptr{T}, px::Ptr{T}, H::HalfOdd{T}, nt::Int) where {T} = half_odd!(py, px, H, nt) + +function half_odd!(py::Ptr{T}, px::Ptr{T}, H::HalfOdd{T}, nt::Int) where {T} + n = H.n; L = H.L; kind = H.kind; w = H.w; out = H.out; pre = H.pre; post = H.post + GC.@preserve H px py begin + each_line(L, 1, nt) do c, j + @inbounds begin + if j < n + src = H.rev ? n - 1 - j : j + w[j+1] = unsafe_load(px, src + 1) * pre[j+1] + else + w[j+1] = zero(Complex{T}) + end + end + end + mul!(out, H.cplan, w) + each_line(n, 1, nt) do c, k + @inbounds begin + v = post[k+1] * out[k + H.oidx + 1] + y = kind == :dst2 ? -2 * imag(v) : kind == :dht ? real(v) - imag(v) : + kind in (:dct3, :dst3) ? real(v) : 2 * real(v) + H.alt && isodd(k) && (y = -y) + unsafe_store!(py, y, k + 1) + end + end + end + nothing +end + +# DCT-III and DST-III of one even line on several threads, from two half-length complex DFTs: +# E_p = Re DFT(x_2q e^{-iπq/n})_p and O_p = Re(e^{-2πip/n} DFT(x_2q+1 e^{-iπ(2q+1)/(2n)})_p) give +# y_2p = 2(E_p + O_p) - x_0 and y_{n-1-2p} = 2(E_p - O_p) - x_0. DST-III is the DCT-III of the reversed +# input times (-1)^k. Checked against FFTW: relerr 3e-16 for n = 8…64. +struct HalfIII{T,F} + n::Int + dst::Bool + a::Vector{Complex{T}} # pre-twiddled even samples, then their DFT + b::Vector{Complex{T}} # pre-twiddled odd samples, then their DFT + pre_a::Vector{Complex{T}} # e^{-iπq/n} + pre_b::Vector{Complex{T}} # e^{-iπ(2q+1)/(2n)} + post::Vector{Complex{T}} # e^{-2πip/n} + x0::Base.RefValue{T} # the first input value (after reversal), kept for the output pass + cplan::F +end + +function half_iii(::Type{T}, n::Int, dst::Bool, nt::Int) where {T} + M = n ÷ 2 + pre_a = Complex{T}[Complex{T}(cis(-π * q / n)) for q in 0:M-1] + pre_b = Complex{T}[Complex{T}(cis(-π * (2q + 1) / (2n))) for q in 0:M-1] + post = Complex{T}[Complex{T}(cis(-2π * p / n)) for p in 0:M-1] + cplan = FFTLPlan{T}((M,), (1,), -1; threads = nt) + HalfIII{T,typeof(cplan)}(n, dst, zeros(Complex{T}, M), zeros(Complex{T}, M), pre_a, pre_b, post, + Ref(zero(T)), cplan) +end + +half_dft!(py::Ptr{T}, px::Ptr{T}, H::HalfIII{T}, nt::Int) where {T} = half_iii!(py, px, H, nt) + +function half_iii!(py::Ptr{T}, px::Ptr{T}, H::HalfIII{T}, nt::Int) where {T} + n = H.n; M = n ÷ 2; dst = H.dst + a = H.a; b = H.b; pa = H.pre_a; pb = H.pre_b; post = H.post + GC.@preserve H px py begin + # input index of the (possibly reversed) sample j + rd(j) = dst ? n - 1 - j : j + each_line(M, 1, nt) do c, q + @inbounds begin + a[q+1] = unsafe_load(px, rd(2q) + 1) * pa[q+1] + b[q+1] = unsafe_load(px, rd(2q + 1) + 1) * pb[q+1] + end + end + H.x0[] = unsafe_load(px, rd(0) + 1) + mul!(a, H.cplan, a) + mul!(b, H.cplan, b) + x0 = H.x0[] + each_line(M, 1, nt) do c, p + @inbounds begin + E = real(a[p+1]) + O = real(post[p+1] * b[p+1]) + unsafe_store!(py, 2 * (E + O) - x0, 2p + 1) + v = 2 * (E - O) - x0 + unsafe_store!(py, dst ? -v : v, n - 2p) # output index n-1-2p is odd for even n + end + end + end + nothing +end + +# DCT-I and DST-I of one line on several threads: the real even (DCT-I, length 2(n-1)) or odd (DST-I, +# length 2(n+1)) extension of the input, transformed by one complex DFT of that length. The real or the +# negated imaginary part of the DFT is the result. Checked against FFTW for n = 16 … 300000. +struct HalfI{T,F} + n::Int + dst::Bool + w::Vector{Complex{T}} # the extended input, length m + out::Vector{Complex{T}} # its DFT + cplan::F +end + +function half_i(::Type{T}, n::Int, dst::Bool, nt::Int) where {T} + m = 2 * (dst ? n + 1 : n - 1) + cplan = FFTLPlan{T}((m,), (1,), -1; threads = nt) + HalfI{T,typeof(cplan)}(n, dst, zeros(Complex{T}, m), zeros(Complex{T}, m), cplan) +end + +half_dft!(py::Ptr{T}, px::Ptr{T}, H::HalfI{T}, nt::Int) where {T} = half_i!(py, px, H, nt) + +function half_i!(py::Ptr{T}, px::Ptr{T}, H::HalfI{T}, nt::Int) where {T} + n = H.n; m = length(H.w); dst = H.dst; w = H.w + GC.@preserve H px py begin + if dst + # odd extension about 0 and n+1: w_j = x_{j-1} for 1 ≤ j ≤ n, w_{m-j} = -w_j, zeros at 0 and n+1 + each_line(m, 1, nt) do c, j + @inbounds begin + if j == 0 || j == n + 1 + w[j+1] = 0 + elseif j <= n + w[j+1] = unsafe_load(px, j) + 0im + elseif j < m + w[j+1] = -(unsafe_load(px, m - j - 1 + 1) + 0im) + end + end + end + else + # even extension about 0 and n-1: w_j = x_j for j ≤ n-1, w_{m-j} = w_j + each_line(m, 1, nt) do c, j + @inbounds begin + src = j < n ? j : m - j + w[j+1] = unsafe_load(px, src + 1) + 0im + end + end + end + mul!(H.out, H.cplan, w) + out = H.out + each_line(n, 1, nt) do c, k + @inbounds unsafe_store!(py, dst ? -imag(out[k+2]) : real(out[k+1]), k + 1) + end + end + nothing +end + function half_dft!(py::Ptr{T}, px::Ptr{T}, H::HalfDFT{T}, nt::Int) where {T} S, n, outer = H.S, H.n, H.outer M = n ÷ 2 diff --git a/src/real.jl b/src/real.jl index 5708d559..6caf7bc4 100644 --- a/src/real.jl +++ b/src/real.jl @@ -186,6 +186,7 @@ mutable struct RFFTLPlan{T<:FFTLFloat,N,K,C,D,F} <: AbstractFFTs.Plan{T} sz::NTuple{N,Int} region::NTuple{D,Int} line::K + lines::Vector{K} # one line kernel per task (copies of `line` if it has buffers of its own) four::F # one long even line on several threads: four-step, or nothing cplan::C # complex passes on the half spectrum, or nothing S::Int @@ -196,8 +197,8 @@ mutable struct RFFTLPlan{T<:FFTLFloat,N,K,C,D,F} <: AbstractFFTs.Plan{T} nthreads::Int bufs::Vector{Vector{Complex{T}}} # three per thread, plus line buffers for strided lines pinv::AbstractFFTs.Plan - RFFTLPlan{T,N,K,C,D,F}(sz, region, line, four, cplan, S, offs_r, offs_c, nt, bufs) where {T,N,K,C,D,F} = - new{T,N,K,C,D,F}(sz, region, line, four, cplan, S, offs_r, offs_c, similar(offs_r), similar(offs_c), nt, bufs) + RFFTLPlan{T,N,K,C,D,F}(sz, region, line, lines, four, cplan, S, offs_r, offs_c, nt, bufs) where {T,N,K,C,D,F} = + new{T,N,K,C,D,F}(sz, region, line, lines, four, cplan, S, offs_r, offs_c, similar(offs_r), similar(offs_c), nt, bufs) end """ @@ -211,6 +212,7 @@ mutable struct BRFFTLPlan{T<:FFTLFloat,N,K,C,D,F} <: AbstractFFTs.Plan{Complex{T d::Int region::NTuple{D,Int} line::K + lines::Vector{K} # one line kernel per task, as in RFFTLPlan four::F cplan::C S::Int @@ -222,8 +224,8 @@ mutable struct BRFFTLPlan{T<:FFTLFloat,N,K,C,D,F} <: AbstractFFTs.Plan{Complex{T bufs::Vector{Vector{Complex{T}}} copy::Array{Complex{T},N} # the input, transformed along the complex dimensions pinv::AbstractFFTs.Plan - BRFFTLPlan{T,N,K,C,D,F}(sz, d, region, line, four, cplan, S, offs_r, offs_c, nt, bufs, copy) where {T,N,K,C,D,F} = - new{T,N,K,C,D,F}(sz, d, region, line, four, cplan, S, offs_r, offs_c, similar(offs_r), similar(offs_c), nt, bufs, copy) + BRFFTLPlan{T,N,K,C,D,F}(sz, d, region, line, lines, four, cplan, S, offs_r, offs_c, nt, bufs, copy) where {T,N,K,C,D,F} = + new{T,N,K,C,D,F}(sz, d, region, line, lines, four, cplan, S, offs_r, offs_c, similar(offs_r), similar(offs_c), nt, bufs, copy) end sortregion(region, N) = Tuple(normregion(region, N)) @@ -246,7 +248,8 @@ function RFFTLPlan{T}(sz::NTuple{N,Int}, region; threads::Integer = Threads.nthr nt = max(1, Int(threads)) bufs = thread_buffers(T, n, 5nt) four = length(offs_r) == 1 && use_real_fourstep(T, n, nt) ? FourStep{T}(n ÷ 2, -1, nt) : nothing - RFFTLPlan{T,N,typeof(line),typeof(cplan),length(reg),typeof(four)}(sz, reg, line, four, cplan, S, offs_r, offs_c, nt, bufs) + lines = task_reallines(line, n, -1, nt) + RFFTLPlan{T,N,typeof(line),typeof(cplan),length(reg),typeof(four)}(sz, reg, line, lines, four, cplan, S, offs_r, offs_c, nt, bufs) end function BRFFTLPlan{T}(sz::NTuple{N,Int}, d::Int, region; threads::Integer = Threads.nthreads()) where {T<:FFTLFloat,N} @@ -262,7 +265,8 @@ function BRFFTLPlan{T}(sz::NTuple{N,Int}, d::Int, region; threads::Integer = Thr nt = max(1, Int(threads)) bufs = thread_buffers(T, d, 5nt) four = length(offs_r) == 1 && use_real_fourstep(T, d, nt) ? FourStep{T}(d ÷ 2, 1, nt) : nothing - BRFFTLPlan{T,N,typeof(line),typeof(cplan),length(reg),typeof(four)}(sz, d, reg, line, four, cplan, S, offs_r, offs_c, nt, bufs, + lines = task_reallines(line, d, 1, nt) + BRFFTLPlan{T,N,typeof(line),typeof(cplan),length(reg),typeof(four)}(sz, d, reg, line, lines, four, cplan, S, offs_r, offs_c, nt, bufs, Array{Complex{T},N}(undef, sz)) end @@ -291,6 +295,13 @@ own_buffers(k) = k isa Union{Bluestein,Composite} own_buffers(K::RealLine) = K.inner isa Union{Bluestein,Composite} line_tasks(K, nt) = own_buffers(K) ? 1 : nt +# One line kernel per task. A kernel with buffers of its own (Bluestein, composite) cannot run on two +# tasks at once, so each task gets a copy; the other kernels are shared. +task_reallines(line::RealLine{T}, n::Int, σ::Int, nt::Int) where {T} = + own_buffers(line) && nt > 1 ? [line; [RealLine{T}(n, σ) for _ in 2:nt]] : [line] +rline_tasks(L, nt) = length(L) > 1 ? length(L) : line_tasks(L[1], nt) +rline_task(L, c) = L[length(L) > 1 ? c : 1] + # Runs f(c, i) for the lines i = 0…nlines-1 on the plan's tasks. function each_line(f::F, nlines, n, nt) where {F} ntasks = nt > 1 && nlines * n >= TUNING[].thread_min_work ? nt : 1 @@ -346,23 +357,30 @@ function r2c_contiguous!(py::Ptr{Complex{T}}, px::Ptr{T}, p) where {T} F === nothing || return r2c_line_fourstep!(py, px, F, p.line, p.nthreads) return r2c_line!(py, px, cptr(b[1]), cptr(b[2]), cptr(b[3]), p.line) end - each_line(length(p.offs_r), p.line.n, line_tasks(p.line, p.nthreads)) do c, i + each_line(length(p.offs_r), p.line.n, rline_tasks(p.lines, p.nthreads)) do c, i local b = p.bufs r2c_line!(py + p.offs_c[i+1] * sizeof(Complex{T}), px + p.offs_r[i+1] * sizeof(T), - cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), p.line) + cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), rline_task(p.lines, c)) end end +# one kernel for all tasks +# a tuple, not a vector: no allocation per call +r2c_lines!(py, px, offs_y, offs_x, Sy, Sx, line::RealLine, bufs, nt::Int) = + r2c_lines!(py, px, offs_y, offs_x, Sy, Sx, (line,), bufs, nt) +c2r_lines!(py, px, offs_y, offs_x, Sy, Sx, line::RealLine, bufs, nt::Int) = + c2r_lines!(py, px, offs_y, offs_x, Sy, Sx, (line,), bufs, nt) + # strided line: gather into a contiguous buffer, transform, scatter function r2c_lines!(py::Ptr{Complex{T}}, px::Ptr{T}, offs_y::Vector{Int}, offs_x::Vector{Int}, Sy::Int, Sx::Int, - line::L, bufs::Vector{Vector{Complex{T}}}, nt::Int) where {T,L<:RealLine{T}} - n = line.n; h = nhalf(n) - each_line(length(offs_x), n, line_tasks(line, nt)) do c, i + lines::Union{Tuple,Vector}, bufs::Vector{Vector{Complex{T}}}, nt::Int) where {T} + n = lines[1].n; h = nhalf(n) + each_line(length(offs_x), n, rline_tasks(lines, nt)) do c, i local lr = Ptr{T}(cptr(bufs[5c-1])); local lc = Ptr{Complex{T}}(cptr(bufs[5c])) for j in 0:n-1 unsafe_store!(lr, unsafe_load(px, offs_x[i+1] + Sx * j + 1), j + 1) end - r2c_line!(lc, lr, cptr(bufs[5c-4]), cptr(bufs[5c-3]), cptr(bufs[5c-2]), line) + r2c_line!(lc, lr, cptr(bufs[5c-4]), cptr(bufs[5c-3]), cptr(bufs[5c-2]), rline_task(lines, c)) for k in 0:h-1 unsafe_store!(py, unsafe_load(lc, k + 1), offs_y[i+1] + Sy * k + 1) end @@ -429,22 +447,22 @@ function c2r_contiguous!(py::Ptr{T}, px::Ptr{Complex{T}}, p) where {T} F === nothing || return c2r_line_fourstep!(py, px, Ptr{Complex{T}}(cptr(b[1])), F, p.line, p.nthreads) return c2r_line!(py, px, cptr(b[1]), cptr(b[2]), cptr(b[3]), p.line) end - each_line(length(p.offs_r), p.line.n, line_tasks(p.line, p.nthreads)) do c, i + each_line(length(p.offs_r), p.line.n, rline_tasks(p.lines, p.nthreads)) do c, i local b = p.bufs c2r_line!(py + p.offs_r[i+1] * sizeof(T), px + p.offs_c[i+1] * sizeof(Complex{T}), - cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), p.line) + cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), rline_task(p.lines, c)) end end function c2r_lines!(py::Ptr{T}, px::Ptr{Complex{T}}, offs_y::Vector{Int}, offs_x::Vector{Int}, Sy::Int, Sx::Int, - line::L, bufs::Vector{Vector{Complex{T}}}, nt::Int) where {T,L<:RealLine{T}} - n = line.n; h = nhalf(n) - each_line(length(offs_x), n, line_tasks(line, nt)) do c, i + lines::Union{Tuple,Vector}, bufs::Vector{Vector{Complex{T}}}, nt::Int) where {T} + n = lines[1].n; h = nhalf(n) + each_line(length(offs_x), n, rline_tasks(lines, nt)) do c, i local lr = Ptr{T}(cptr(bufs[5c-1])); local lc = Ptr{Complex{T}}(cptr(bufs[5c])) for k in 0:h-1 unsafe_store!(lc, unsafe_load(px, offs_x[i+1] + Sx * k + 1), k + 1) end - c2r_line!(lr, lc, cptr(bufs[5c-4]), cptr(bufs[5c-3]), cptr(bufs[5c-2]), line) + c2r_line!(lr, lc, cptr(bufs[5c-4]), cptr(bufs[5c-3]), cptr(bufs[5c-2]), rline_task(lines, c)) for j in 0:n-1 unsafe_store!(py, unsafe_load(lr, j + 1), offs_y[i+1] + Sy * j + 1) end @@ -454,9 +472,9 @@ end # the strided batches of RFFTLPlan and RRFFTPlan: input and output with the same stride S r2c_strided!(py::Ptr{Complex{T}}, px::Ptr{T}, p) where {T} = - r2c_lines!(py, px, p.offs_c, p.offs_r, p.S, p.S, p.line, p.bufs, p.nthreads) + r2c_lines!(py, px, p.offs_c, p.offs_r, p.S, p.S, p.lines, p.bufs, p.nthreads) c2r_strided!(py::Ptr{T}, px::Ptr{Complex{T}}, p) where {T} = - c2r_lines!(py, px, p.offs_r, p.offs_c, p.S, p.S, p.line, p.bufs, p.nthreads) + c2r_lines!(py, px, p.offs_r, p.offs_c, p.S, p.S, p.lines, p.bufs, p.nthreads) Base.:*(p::RFFTLPlan{T,N}, x::StridedArray{T,N}) where {T,N} = mul!(Array{Complex{T},N}(undef, AbstractFFTs.output_size(p)), p, x) diff --git a/src/rreal.jl b/src/rreal.jl index 7965b6fb..686db644 100644 --- a/src/rreal.jl +++ b/src/rreal.jl @@ -19,6 +19,7 @@ mutable struct RRFFTPlan{T<:FFTLFloat} region::Vector{Int} σ::Int line::RealLine{T,RKern{T}} + lines::Vector{RealLine{T,RKern{T}}} # one line kernel per task (copies if the kernel has buffers of its own) cplan::Union{Nothing,RPlan{T}} # complex passes on the half spectrum S::Int offs_r::Vector{Int} @@ -44,7 +45,8 @@ function RRFFTPlan{T}(sz::AbstractVector{<:Integer}, region::AbstractVector{<:In cplan = length(reg) > 1 ? RPlan{T}(csz, reg[2:end], Int(σ); threads = nt) : nothing S, offs_r, offs_c = line_layout(dims, d1) copybuf = Vector{Complex{T}}(undef, σ == 1 && cplan !== nothing ? prod(csz) : 0) - RRFFTPlan{T}(dims, reg, Int(σ), line, cplan, S, offs_r, offs_c, nt, thread_buffers(T, n, 5nt), copybuf, T[]) + RRFFTPlan{T}(dims, reg, Int(σ), line, task_reallines(line, n, Int(σ), nt), cplan, S, offs_r, offs_c, nt, + thread_buffers(T, n, 5nt), copybuf, T[]) end Base.show(io::IO, p::RRFFTPlan{T}) where {T} = @@ -93,10 +95,10 @@ end function r2c_inplace!(p::RRFFTPlan{T}, py::Ptr{Complex{T}}) where {T} if p.S == 1 - each_line(length(p.offs_c), p.line.n, line_tasks(p.line, p.nthreads)) do c, i + each_line(length(p.offs_c), p.line.n, rline_tasks(p.lines, p.nthreads)) do c, i local b = p.bufs local q = py + p.offs_c[i+1] * sizeof(Complex{T}) - r2c_line!(q, Ptr{T}(q), cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), p.line) + r2c_line!(q, Ptr{T}(q), cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), rline_task(p.lines, c)) end else length(p.rbuf) == prod(p.sz) || resize!(p.rbuf, prod(p.sz)) @@ -112,10 +114,10 @@ function c2r_inplace!(p::RRFFTPlan{T}, px::Ptr{Complex{T}}) where {T} cp = p.cplan cp === nothing || execute!(cp, px, px) if p.S == 1 - each_line(length(p.offs_c), p.line.n, line_tasks(p.line, p.nthreads)) do c, i + each_line(length(p.offs_c), p.line.n, rline_tasks(p.lines, p.nthreads)) do c, i local b = p.bufs local q = px + p.offs_c[i+1] * sizeof(Complex{T}) - c2r_line!(Ptr{T}(q), q, cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), p.line) + c2r_line!(Ptr{T}(q), q, cptr(b[5c-4]), cptr(b[5c-3]), cptr(b[5c-2]), rline_task(p.lines, c)) end else length(p.rbuf) == prod(p.sz) || resize!(p.rbuf, prod(p.sz)) From 32d622c7de65f83418dbaa920e30ec95b4bc600f Mon Sep 17 00:00:00 2001 From: Tobias Knopp Date: Wed, 11 Nov 2026 08:02:24 +0100 Subject: [PATCH 2/3] Threads for non-smooth real lines of any parity A single real line of a non-smooth length ran on one thread: the half-length DFT (the full length for odd n) has buffers of its own, and only four-step (smooth lengths) was threaded. The line now uses the composite or Bluestein kernel on several threads, like the complex transform of the same length. Non-smooth lines start at real_bluestein_bytes = 2^17 (Float32 gains from about 130 KB, Float64 from about 240 KB on the measured machine); smooth lines keep real_fourstep_bytes = 2^18. Co-Authored-By: Claude Haiku 5.5 --- src/real.jl | 81 +++++++++++++++++++++++++++++++++++++++--------- src/tuning.jl | 4 ++- test/runtests.jl | 9 ++++++ 3 files changed, 79 insertions(+), 15 deletions(-) diff --git a/src/real.jl b/src/real.jl index 6caf7bc4..19078538 100644 --- a/src/real.jl +++ b/src/real.jl @@ -137,9 +137,10 @@ end # One long even line on several threads: the complex DFT of length n/2 by the four-step # algorithm, the pairs in chunks on the tasks. -function r2c_line_fourstep!(py::Ptr{Complex{T}}, px::Ptr{T}, F::FourStep{T}, K::RealLine{T}, nt::Int) where {T} +function r2c_line_fourstep!(py::Ptr{Complex{T}}, px::Ptr{T}, F::Kernel{T}, K::RealLine{T}, nt::Int, b1::Ptr{T}, b2::Ptr{T}) where {T} + iseven(K.n) || return r2c_line_odd_threads!(py, px, F, K, nt, b1, b2) M = K.n ÷ 2 - run_fourstep!(Ptr{T}(py), px, F, nt) + run_half!(Ptr{T}(py), px, F, nt) r2c_ends!(py, M) np = M ÷ 2; nc = min(np, BLOCKS_PER_THREAD * nt) parallel_items(nc, nt) do c, g @@ -149,7 +150,8 @@ function r2c_line_fourstep!(py::Ptr{Complex{T}}, px::Ptr{T}, F::FourStep{T}, K:: nothing end -function c2r_line_fourstep!(py::Ptr{T}, px::Ptr{Complex{T}}, cb1::Ptr{Complex{T}}, F::FourStep{T}, K::RealLine{T}, nt::Int) where {T} +function c2r_line_fourstep!(py::Ptr{T}, px::Ptr{Complex{T}}, cb1::Ptr{Complex{T}}, F::Kernel{T}, K::RealLine{T}, nt::Int, b2::Ptr{T}) where {T} + iseven(K.n) || return c2r_line_odd_threads!(py, px, cb1, F, K, nt, b2) M = K.n ÷ 2 c2r_ends!(cb1, px, M) np = M ÷ 2; nc = min(np, BLOCKS_PER_THREAD * nt) @@ -157,13 +159,18 @@ function c2r_line_fourstep!(py::Ptr{T}, px::Ptr{Complex{T}}, cb1::Ptr{Complex{T} local r = chunk_range(g + 1, nc, np) isempty(r) || c2r_pairs!(cb1, px, K, M, first(r) + 1, last(r) + 1) end - run_fourstep!(py, Ptr{T}(cb1), F, nt) + run_half!(py, Ptr{T}(cb1), F, nt) nothing end -# the four-step plan of a single long line, or nothing (also for the run-time plans) +# The threaded kernel of the half-length DFT of a single long line, or nothing (also for the run-time plans) line_fourstep(p) = nothing +# Half-length DFT of one line on several threads, by the kernel type. +run_half!(py::Ptr{T}, px::Ptr{T}, K::FourStep{T}, nt::Int) where {T} = run_fourstep!(py, px, K, nt) +run_half!(py::Ptr{T}, px::Ptr{T}, K::CompositeMT{T}, nt::Int) where {T} = run_composite_mt!(py, px, K, nt) +run_half!(py::Ptr{T}, px::Ptr{T}, K::BluesteinMT{T}, nt::Int) where {T} = run_bluestein_mt!(py, px, K, nt) + # Lines along dimension d of an array of size `sz` (real side), as (lane stride S, line offsets # in elements of the real side and of the complex side). function line_layout(sz, d) @@ -230,11 +237,57 @@ end sortregion(region, N) = Tuple(normregion(region, N)) -# One real line of length n on several tasks: four-step for the half-length DFT. The extra -# parallel pass for the pairs pays off from 256 KB per line on (M1, 8 threads: Float64 32768 and -# Float32 65536 are 0.6-0.9 of one task, Float64 16384 and Float32 16384-32768 up to 3x slower). -use_real_fourstep(::Type{T}, n, nt) where {T} = - iseven(n) && n * sizeof(T) >= TUNING[].real_fourstep_bytes && use_fourstep(n ÷ 2, nt) +# One real line of length n on several tasks: the half-length DFT on several threads (four-step for +# a smooth half length, composite or Bluestein for the others). The extra parallel pass for the +# pairs needs about 256 KB per line for smooth lengths and 128 KB for non-smooth ones (M1, 8 threads, +# gain of the threaded over the serial line: Float32 smooth 0.7 at 128 KB, 1.5 at 256 KB; non-smooth +# 1.05 at 120-130 KB, 1.8 at 147 KB; Float64 non-smooth 1.5 at 240 KB). +use_real_threads(::Type{T}, n, nt) where {T} = + (use_fourstep(inner_length(n), nt) && n * sizeof(T) >= TUNING[].real_fourstep_bytes) || + ((use_composite_mt(T, inner_length(n), nt) || use_bluestein_mt(inner_length(n), nt)) && + n * sizeof(T) >= TUNING[].real_bluestein_bytes) + +# Odd n: the full-length DFT on several threads, in the line's buffers. Forward: the real line is +# the real part of a complex line; backward: the half spectrum is completed by conjugate symmetry. +function r2c_line_odd_threads!(py::Ptr{Complex{T}}, px::Ptr{T}, F::Kernel{T}, K::RealLine{T}, nt::Int, b1::Ptr{T}, b2::Ptr{T}) where {T} + n = K.n + cb1 = Ptr{Complex{T}}(b1) + for j in 1:n + unsafe_store!(cb1, Complex(unsafe_load(px, j), zero(T)), j) + end + run_half!(b2, b1, F, nt) + unsafe_copyto!(py, Ptr{Complex{T}}(b2), nhalf(n)) + nothing +end + +function c2r_line_odd_threads!(py::Ptr{T}, px::Ptr{Complex{T}}, cb1::Ptr{Complex{T}}, F::Kernel{T}, K::RealLine{T}, nt::Int, b2::Ptr{T}) where {T} + n = K.n + h = nhalf(n) + for k in 1:h + unsafe_store!(cb1, unsafe_load(px, k), k) + end + unsafe_store!(cb1, Complex(real(unsafe_load(px, 1)), zero(T)), 1) + for k in 1:n-h + unsafe_store!(cb1, conj(unsafe_load(px, k + 1)), n - k + 1) + end + run_half!(Ptr{T}(b2), Ptr{T}(cb1), F, nt) + cb2 = Ptr{Complex{T}}(b2) + for j in 1:n + unsafe_store!(py, real(unsafe_load(cb2, j)), j) + end + nothing +end + +# Length of the DFT that a line kernel runs: n ÷ 2 for even n (pairs of reals), n for odd n. +inner_length(n) = iseven(n) ? n ÷ 2 : n + +# The threaded kernel of the DFT of length h (σ = -1 forward, +1 backward), or nothing. +function half_threads(::Type{T}, h, σ, nt) where {T} + use_fourstep(h, nt) && return FourStep{T}(h, σ, nt) + use_composite_mt(T, h, nt) && return CompositeMT{T}(h, σ, nt) + use_bluestein_mt(h, nt) && return BluesteinMT{T}(h, σ, nt) + nothing +end function RFFTLPlan{T}(sz::NTuple{N,Int}, region; threads::Integer = Threads.nthreads()) where {T<:FFTLFloat,N} reg = sortregion(region, N) @@ -247,7 +300,7 @@ function RFFTLPlan{T}(sz::NTuple{N,Int}, region; threads::Integer = Threads.nthr S, offs_r, offs_c = line_layout(sz, d1) nt = max(1, Int(threads)) bufs = thread_buffers(T, n, 5nt) - four = length(offs_r) == 1 && use_real_fourstep(T, n, nt) ? FourStep{T}(n ÷ 2, -1, nt) : nothing + four = length(offs_r) == 1 && use_real_threads(T, n, nt) ? half_threads(T, inner_length(n), -1, nt) : nothing lines = task_reallines(line, n, -1, nt) RFFTLPlan{T,N,typeof(line),typeof(cplan),length(reg),typeof(four)}(sz, reg, line, lines, four, cplan, S, offs_r, offs_c, nt, bufs) end @@ -264,7 +317,7 @@ function BRFFTLPlan{T}(sz::NTuple{N,Int}, d::Int, region; threads::Integer = Thr S, offs_r, offs_c = line_layout(rsz, d1) nt = max(1, Int(threads)) bufs = thread_buffers(T, d, 5nt) - four = length(offs_r) == 1 && use_real_fourstep(T, d, nt) ? FourStep{T}(d ÷ 2, 1, nt) : nothing + four = length(offs_r) == 1 && use_real_threads(T, d, nt) ? half_threads(T, inner_length(d), 1, nt) : nothing lines = task_reallines(line, d, 1, nt) BRFFTLPlan{T,N,typeof(line),typeof(cplan),length(reg),typeof(four)}(sz, d, reg, line, lines, four, cplan, S, offs_r, offs_c, nt, bufs, Array{Complex{T},N}(undef, sz)) @@ -354,7 +407,7 @@ function r2c_contiguous!(py::Ptr{Complex{T}}, px::Ptr{T}, p) where {T} b = p.bufs if length(p.offs_r) == 1 F = line_fourstep(p) - F === nothing || return r2c_line_fourstep!(py, px, F, p.line, p.nthreads) + F === nothing || return r2c_line_fourstep!(py, px, F, p.line, p.nthreads, cptr(b[1]), cptr(b[2])) return r2c_line!(py, px, cptr(b[1]), cptr(b[2]), cptr(b[3]), p.line) end each_line(length(p.offs_r), p.line.n, rline_tasks(p.lines, p.nthreads)) do c, i @@ -444,7 +497,7 @@ function c2r_contiguous!(py::Ptr{T}, px::Ptr{Complex{T}}, p) where {T} b = p.bufs if length(p.offs_r) == 1 F = line_fourstep(p) - F === nothing || return c2r_line_fourstep!(py, px, Ptr{Complex{T}}(cptr(b[1])), F, p.line, p.nthreads) + F === nothing || return c2r_line_fourstep!(py, px, Ptr{Complex{T}}(cptr(b[1])), F, p.line, p.nthreads, cptr(b[2])) return c2r_line!(py, px, cptr(b[1]), cptr(b[2]), cptr(b[3]), p.line) end each_line(length(p.offs_r), p.line.n, rline_tasks(p.lines, p.nthreads)) do c, i diff --git a/src/tuning.jl b/src/tuning.jl index 742e80da..1eb1d992 100644 --- a/src/tuning.jl +++ b/src/tuning.jl @@ -14,6 +14,7 @@ Size thresholds of FFTL (see `TUNING`): - `factor_twiddles`: bytes above which the four-step uses factor twiddle tables - `bluestein_mt_min`: length from which Bluestein's algorithm runs on several threads - `real_fourstep_bytes`: bytes of a real line from which it runs as a four-step on several threads +- `real_bluestein_bytes`: bytes of a real line from which a non-smooth one runs on several threads - `plane_budget`: complex values of a plane (dimensions 1 and 2) that may be fused """ struct Tuning @@ -23,10 +24,11 @@ struct Tuning factor_twiddles::Int bluestein_mt_min::Int real_fourstep_bytes::Int + real_bluestein_bytes::Int plane_budget::Int end -const DEFAULT_TUNING = Tuning(1 << 15, 1 << 16, 1 << 13, 1 << 23, 1 << 14, 1 << 18, 1 << 17) +const DEFAULT_TUNING = Tuning(1 << 15, 1 << 16, 1 << 13, 1 << 23, 1 << 14, 1 << 18, 1 << 17, 1 << 17) # Read at plan construction, except thread_min_work, block_budget and plane_budget, which a pass # reads when it starts. Assigning `FFTL.TUNING[] = …` changes the settings for plans made afterwards diff --git a/test/runtests.jl b/test/runtests.jl index 2967db23..2d5116ec 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -847,6 +847,15 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] @test relerr(X, FFTL.plan_rfft(x; threads = 1) * x) < 10tol(T) @test relerr(FFTL.plan_brfft(X, n; threads = 3) * X, n .* x) < 10tol(T) end + # non-smooth lengths, even and odd: the half-length (even) or the full-length (odd) DFT on threads + for (T, n) in ((Float64, 37888), (Float64, 92935), (Float64, 2 * 1009 * 37), (Float32, 33490), (Float32, 92935), (Float32, 100003)) + x = randn(T, n) + p = FFTL.plan_rfft(x; threads = 3) + @test p.four isa Union{FFTL.CompositeMT,FFTL.BluesteinMT} + X = p * x + @test relerr(X, FFTL.plan_rfft(x; threads = 1) * x) < 10tol(T) + @test relerr(FFTL.plan_brfft(X, n; threads = 3) * X, n .* x) < 10tol(T) + end # below 256 KB per line one task does the line @test FFTL.plan_rfft(randn(Float64, 2^14); threads = 3).four === nothing @test FFTL.plan_rfft(randn(Float32, 2^15); threads = 3).four === nothing From a971729bc225c0c11f05c345125f662641cfaf2e Mon Sep 17 00:00:00 2001 From: Tobias Knopp Date: Wed, 11 Nov 2026 09:43:30 +0100 Subject: [PATCH 3/3] Run-time real plans: copies of the line kernels have the run-time kernel type task_reallines built the copies of a line with buffers of their own (a non-smooth length on several tasks) with the default kernel type, but RRFFTPlan stores lines of RChain kernels. Non-smooth lengths on more than one thread then failed: the r2c and c2r plans of length 101 with 4 threads gave wrong results in libfftl (FFTW3 interface test). The copies now take the kernel type of the plan. A test covers run-time plans of non-smooth lengths on 4 threads. Co-Authored-By: Claude Haiku 5.5 --- src/real.jl | 5 +++-- src/rreal.jl | 2 +- test/runtests.jl | 3 ++- 3 files changed, 6 insertions(+), 4 deletions(-) diff --git a/src/real.jl b/src/real.jl index 19078538..399f8402 100644 --- a/src/real.jl +++ b/src/real.jl @@ -350,8 +350,9 @@ line_tasks(K, nt) = own_buffers(K) ? 1 : nt # One line kernel per task. A kernel with buffers of its own (Bluestein, composite) cannot run on two # tasks at once, so each task gets a copy; the other kernels are shared. -task_reallines(line::RealLine{T}, n::Int, σ::Int, nt::Int) where {T} = - own_buffers(line) && nt > 1 ? [line; [RealLine{T}(n, σ) for _ in 2:nt]] : [line] +# The copies have the kernel type of the line (C = RChain for the run-time plans). +task_reallines(line::RealLine{T}, n::Int, σ::Int, nt::Int, ::Type{C} = Chain) where {T,C} = + own_buffers(line) && nt > 1 ? [line; [RealLine{T}(n, σ, C) for _ in 2:nt]] : [line] rline_tasks(L, nt) = length(L) > 1 ? length(L) : line_tasks(L[1], nt) rline_task(L, c) = L[length(L) > 1 ? c : 1] diff --git a/src/rreal.jl b/src/rreal.jl index 686db644..87d5bddb 100644 --- a/src/rreal.jl +++ b/src/rreal.jl @@ -45,7 +45,7 @@ function RRFFTPlan{T}(sz::AbstractVector{<:Integer}, region::AbstractVector{<:In cplan = length(reg) > 1 ? RPlan{T}(csz, reg[2:end], Int(σ); threads = nt) : nothing S, offs_r, offs_c = line_layout(dims, d1) copybuf = Vector{Complex{T}}(undef, σ == 1 && cplan !== nothing ? prod(csz) : 0) - RRFFTPlan{T}(dims, reg, Int(σ), line, task_reallines(line, n, Int(σ), nt), cplan, S, offs_r, offs_c, nt, + RRFFTPlan{T}(dims, reg, Int(σ), line, task_reallines(line, n, Int(σ), nt, RChain), cplan, S, offs_r, offs_c, nt, thread_buffers(T, n, 5nt), copybuf, T[]) end diff --git a/test/runtests.jl b/test/runtests.jl index 2d5116ec..a8d1a458 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -774,7 +774,8 @@ thin(T, v) = FULL || T == Float64 ? v : v[1:2:end] end @testset "run-time real plans, $T" for T in pick((Float64, Float32), (Float64,)) - for (sz, reg, nt) in [((16,), [1], 1), ((101,), [1], 1), ((12, 10), [1, 2], 3), ((12, 10), [2], 1), ((9, 8, 7), [1, 3], 3)] + for (sz, reg, nt) in [((16,), [1], 1), ((101,), [1], 1), ((12, 10), [1, 2], 3), ((12, 10), [2], 1), ((9, 8, 7), [1, 3], 3), + ((101,), [1], 4), ((7, 101), [1, 2], 4)] # non-smooth lines on several tasks x = randn(T, sz) p = FFTL.RRFFTPlan{T}(collect(sz), reg, -1; threads = nt) ref = FFTL.plan_rfft(x, Tuple(reg)) * x