diff --git a/Project.toml b/Project.toml index a846f4a..7febbda 100644 --- a/Project.toml +++ b/Project.toml @@ -8,6 +8,7 @@ AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" MuladdMacro = "46d2c3a1-f734-5fdb-9937-b9b9aeba4221" +Polyester = "f517fe37-dbe3-4b94-8317-1923a5111588" PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a" Primes = "27ebfcd6-29c5-5fa9-bf4b-fb8fc14df3ae" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" @@ -21,6 +22,7 @@ ExplicitImports = "1.12" FFTW = "1.8" LinearAlgebra = "<0.0.1, 1" MuladdMacro = "0.2" +Polyester = "0.7" PrecompileTools = "1" Primes = "0.5" Random = "<0.0.1, 1" diff --git a/src/FFTA.jl b/src/FFTA.jl index 3309f17..87ce278 100644 --- a/src/FFTA.jl +++ b/src/FFTA.jl @@ -4,6 +4,7 @@ using AbstractFFTs: AbstractFFTs using DocStringExtensions: TYPEDEF, TYPEDSIGNATURES using LinearAlgebra: LinearAlgebra using MuladdMacro: @muladd +using Polyester: @batch using Primes: Primes using Reexport: @reexport diff --git a/src/plan.jl b/src/plan.jl index e519246..85dc76b 100644 --- a/src/plan.jl +++ b/src/plan.jl @@ -282,13 +282,13 @@ function _foreach_pencil(f::F, A::AbstractArray{<:Any,N}, ::Val{dim}, workers::V total = npre * length(Rpost) nt = min(length(workers), total ÷ 2) if nt > 1 && total * size(A, dim) >= THREAD_THRESHOLD - Base.@sync for c in 1:nt + # one chunk per worker, run on Polyester's static thread pool (no task + # allocation, the calling thread takes a chunk itself) + @batch for c in 1:nt lo = (c - 1) * total ÷ nt + 1 hi = c * total ÷ nt - # (the task's worker gets its own name: a variable shared with the - # serial branch below would be captured by every task) wc = workers[c] - Threads.@spawn for k in lo:hi + for k in lo:hi ipost, ipre = divrem(k - 1, npre) f(wc, Rpre[ipre + 1], Rpost[ipost + 1]) end