partr thread support (#105)

* partr thread support

* switch to Yggdrasil build

* Fix problem in build_fftw.jl

* move threads check to end (since threads are now initialized in __init__) and explicitly call set_num_threads(1) for no-threads check since multithreaded Julia now uses multiple FFTW threads by default

.travis.yml

@@ -11,6 +11,7 @@ env:
   matrix:
     - JULIA_FFTW_PROVIDER=FFTW
     - JULIA_FFTW_PROVIDER=MKL
+    - JULIA_NUM_THREADS=2
 notifications:
   email: false
 # uncomment the following lines to override the default test script

deps/build_fftw.jl

@@ -4,44 +4,46 @@ using BinaryProvider # requires BinaryProvider 0.3.0 or later
 const verbose = "--verbose" in ARGS
 const prefix = Prefix(get([a for a in ARGS if a != "--verbose"], 1, joinpath(@__DIR__, "usr")))
 products = [
-    LibraryProduct(prefix, String["libfftw3"], :libfftw3),
-    LibraryProduct(prefix, String["libfftw3f"], :libfftw3f),
+    LibraryProduct(prefix, ["libfftw3"], :libfftw3),
+    LibraryProduct(prefix, ["libfftw3f"], :libfftw3f),
 ]
 
 # Download binaries from hosted location
-bin_prefix = "https://github.com/JuliaMath/FFTWBuilder/releases/download/v3.3.8+1"
+bin_prefix = "https://github.com/JuliaBinaryWrappers/FFTW_jll.jl/releases/download/FFTW-v3.3.9+0"
 
 # Listing of files generated by BinaryBuilder:
 download_info = Dict(
-    Linux(:aarch64, :glibc) => ("$bin_prefix/FFTW.aarch64-linux-gnu.tar.gz", "4296ad9af20d4441fd809c6aaa3ee5fa36818b7a2eb3372da7d2ead454b4e570"),
-    Linux(:aarch64, :musl) => ("$bin_prefix/FFTW.aarch64-linux-musl.tar.gz", "ef6d4e56bd9e405ef2895a857ffbc07cb7abcf450040a2335b83a95f4a431392"),
-    Linux(:armv7l, :glibc, :eabihf) => ("$bin_prefix/FFTW.arm-linux-gnueabihf.tar.gz", "8f8c69a6eca468465734e1fd58801519cea0f7a8f9e08bba93e39315758f7a7c"),
-    Linux(:armv7l, :musl, :eabihf) => ("$bin_prefix/FFTW.arm-linux-musleabihf.tar.gz", "48d137fddab6888bdc59893d22728f081578ff0884954f0f6f5df51afff53ece"),
-    Linux(:i686, :glibc) => ("$bin_prefix/FFTW.i686-linux-gnu.tar.gz", "76d85d81a81752a0e08bc2eec51a568a6000928a550c37a181b51340452b1b5f"),
-    Linux(:i686, :musl) => ("$bin_prefix/FFTW.i686-linux-musl.tar.gz", "5ee42df3aa002e9511c3cc808f728429e4930bb24df8b461dc137bf49aa71b8f"),
-    Windows(:i686) => ("$bin_prefix/FFTW.i686-w64-mingw32.tar.gz", "28b96cb5d78c87d16d305a63a838c5027d319c0292f00c925e14f21c744535a8"),
-    Linux(:powerpc64le, :glibc) => ("$bin_prefix/FFTW.powerpc64le-linux-gnu.tar.gz", "53d305eebb3a152df093d637fe8a4d6288a2b7175bf91ab9242dad62e5e0853a"),
-    MacOS(:x86_64) => ("$bin_prefix/FFTW.x86_64-apple-darwin14.tar.gz", "7562aed6279ea965435c8a388be1494b9a18f7e00058d0fa260a711bffde1bd5"),
-    Linux(:x86_64, :glibc) => ("$bin_prefix/FFTW.x86_64-linux-gnu.tar.gz", "70dcc7ad2697121564d5d91da9f2544b0e68b026779a45063039a39cd5585711"),
-    Linux(:x86_64, :musl) => ("$bin_prefix/FFTW.x86_64-linux-musl.tar.gz", "4bf1c1e7489241c38788bc061f2a091fe72605f67e58411b2933313fa0923877"),
-    FreeBSD(:x86_64) => ("$bin_prefix/FFTW.x86_64-unknown-freebsd11.1.tar.gz", "ad70aca12821f6df1c67da74fc2f1b4fa009ac14d8570ff1f912876e731185af"),
-    Windows(:x86_64) => ("$bin_prefix/FFTW.x86_64-w64-mingw32.tar.gz", "6726bff25faeca8e29dfce8be5b0fb7da0a380faa9fdb5a5a6c98ea76d009b2f"),
+    Linux(:aarch64, libc=:glibc) => ("$bin_prefix/FFTW.v3.3.9.aarch64-linux-gnu.tar.gz", "70a68ce1e89536a8ee8df427de28f527c69d31f28cd223ebf8cf7402f3b45e50"),
+    Linux(:aarch64, libc=:musl) => ("$bin_prefix/FFTW.v3.3.9.aarch64-linux-musl.tar.gz", "a51a44344e0e99ebde48b20495a2c44aeefa51b39a2160932c9cf267d14aa6cf"),
+    Linux(:armv7l, libc=:glibc, call_abi=:eabihf) => ("$bin_prefix/FFTW.v3.3.9.arm-linux-gnueabihf.tar.gz", "2ee40e4d2561f656366eb147213ed483bfd4d3b63ef07359a361ccda48baca94"),
+    Linux(:armv7l, libc=:musl, call_abi=:eabihf) => ("$bin_prefix/FFTW.v3.3.9.arm-linux-musleabihf.tar.gz", "c869cdf7bf12c4f6d4d3359a1c6f14ec59ac84dc96672495f730a30800489b9e"),
+    Linux(:i686, libc=:glibc) => ("$bin_prefix/FFTW.v3.3.9.i686-linux-gnu.tar.gz", "651f7d53dea2b95ae26799f352bea8bd4c17e4371c8faeb3cbb37db9d17b2a84"),
+    Linux(:i686, libc=:musl) => ("$bin_prefix/FFTW.v3.3.9.i686-linux-musl.tar.gz", "9f731365440b19edb8e462a59174a20c2f27df83ab55a2c1e6e4edcbb40a132d"),
+    Windows(:i686) => ("$bin_prefix/FFTW.v3.3.9.i686-w64-mingw32.tar.gz", "644422a04bfa8c74a8cc7b3750083e645a632c80ce03f69a7cc912ee6cf91552"),
+    Linux(:powerpc64le, libc=:glibc) => ("$bin_prefix/FFTW.v3.3.9.powerpc64le-linux-gnu.tar.gz", "ab5f703b110f8af01796e7f31cd11082c7bb27d7184b07f5e8da939dd1b0f2fd"),
+    MacOS(:x86_64) => ("$bin_prefix/FFTW.v3.3.9.x86_64-apple-darwin14.tar.gz", "1f4f99aadb78adc4a0d7df10097a1291c478355255bdd61979b5cdd116f56dac"),
+    Linux(:x86_64, libc=:glibc) => ("$bin_prefix/FFTW.v3.3.9.x86_64-linux-gnu.tar.gz", "ed7fbfe6abd31ba5e47569ff1ebc6b754d7adfdbc9f27478634de495dbc54d32"),
+    Linux(:x86_64, libc=:musl) => ("$bin_prefix/FFTW.v3.3.9.x86_64-linux-musl.tar.gz", "35babebbb69d0bdcb8865372dcbeee3f5cd690cd470be148afd4ea6cec5015c0"),
+    FreeBSD(:x86_64) => ("$bin_prefix/FFTW.v3.3.9.x86_64-unknown-freebsd11.1.tar.gz", "6ff38bce55886dc4d7cf875f17bbec2b2ffa5b5eb54389ef9c5d8461f4e40c82"),
+    Windows(:x86_64) => ("$bin_prefix/FFTW.v3.3.9.x86_64-w64-mingw32.tar.gz", "6793204f0d51a99948fa7983d798c2cbfa696df6e2e5452253ad874fe47b3143"),
 )
 
 # Install unsatisfied or updated dependencies:
 unsatisfied = any(!satisfied(p; verbose=verbose) for p in products)
-if haskey(download_info, platform_key())
-    url, tarball_hash = download_info[platform_key()]
-    if unsatisfied || !isinstalled(url, tarball_hash; prefix=prefix)
-        # Download and install binaries
-        install(url, tarball_hash; prefix=prefix, force=true, verbose=verbose)
-    end
-elseif unsatisfied
-    # If we don't have a BinaryProvider-compatible .tar.gz to download, complain.
+dl_info = choose_download(download_info, platform_key_abi())
+if dl_info === nothing && unsatisfied
+    # If we don't have a compatible .tar.gz to download, complain.
     # Alternatively, you could attempt to install from a separate provider,
     # build from source or something even more ambitious here.
-    error("Your platform $(triplet(platform_key())) is not supported by this package!")
+    error("Your platform (\"$(Sys.MACHINE)\", parsed as \"$(triplet(platform_key_abi()))\") is not supported by this package!")
+end
+
+# If we have a download, and we are unsatisfied (or the version we're
+# trying to install is not itself installed) then load it up!
+if unsatisfied || !isinstalled(dl_info...; prefix=prefix)
+    # Download and install binaries
+    install(dl_info...; prefix=prefix, force=true, verbose=verbose)
 end
 
 # Write out a deps.jl file that will contain mappings for our products
-write_deps_file(joinpath(@__DIR__, "deps.jl"), products)
+write_deps_file(joinpath(@__DIR__, "deps.jl"), products, verbose=verbose)

src/FFTW.jl

@@ -25,6 +25,17 @@ end
 # MKL provides its own FFTW
 const fftw_vendor = occursin("libmkl_rt", libfftw3) ? :mkl : :fftw
 
+# Use Julia partr threading backend if present
+@static if fftw_vendor == :fftw && isdefined(Threads, Symbol("@spawn"))
+    # callback function that FFTW uses to launch `num` parallel
+    # tasks (FFTW/fftw3#175):
+    function spawnloop(f::Ptr{Cvoid}, fdata::Ptr{Cvoid}, elsize::Csize_t, num::Cint, callback_data::Ptr{Cvoid})
+        @sync for i = 0:num-1
+            Threads.@spawn ccall(f, Ptr{Cvoid}, (Ptr{Cvoid},), fdata + elsize*i)
+        end
+    end
+end
+
 # If FFTW was built with threads, then they must be initialized before any FFTW planning routine.
 #   -- This initializes FFTW's threads support (defaulting to 1 thread).
 #      If this isn't called before the FFTW planner is created, then
@@ -33,11 +44,25 @@ const fftw_vendor = occursin("libmkl_rt", libfftw3) ? :mkl : :fftw
 #       plans, causing Base Julia issue #19892.)
 function __init__()
     check_deps()
-    stat = ccall((:fftw_init_threads, libfftw3), Int32, ())
+    stat =  ccall((:fftw_init_threads,   libfftw3), Int32, ())
     statf = ccall((:fftwf_init_threads, libfftw3f), Int32, ())
     if stat == 0 || statf == 0
         error("could not initialize FFTW threads")
     end
+    @static if fftw_vendor == :fftw
+        if Threads.nthreads() > 1 # number of Julia threads is set when Julia is launched
+            ccall((:fftw_make_planner_thread_safe,  libfftw3),  Cvoid, ())
+            ccall((:fftwf_make_planner_thread_safe, libfftw3f), Cvoid, ())
+        end
+        @static if isdefined(Threads, Symbol("@spawn"))
+            if Threads.nthreads() > 1 # partr will give us our threads
+                cspawnloop = @cfunction(spawnloop, Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}, Csize_t, Cint, Ptr{Cvoid}))
+                ccall((:fftw_threads_set_callback,  libfftw3),  Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}), cspawnloop, C_NULL)
+                ccall((:fftwf_threads_set_callback, libfftw3f), Cvoid, (Ptr{Cvoid}, Ptr{Cvoid}), cspawnloop, C_NULL)
+                set_num_threads(Threads.nthreads() * 4) # spawn more tasks than threads to help load-balancing
+            end
+        end
+    end
 end
 
 include("fft.jl")

test/runtests.jl

@@ -5,20 +5,6 @@ using AbstractFFTs: Plan, plan_inv
 using Test
 using LinearAlgebra
 
-# Base Julia issue #19892
-# (test this first to make sure it happens before set_num_threads)
-let a = randn(10^5,1), p1 = plan_rfft(a, flags=FFTW.ESTIMATE)
-    FFTW.set_num_threads(2)
-    p2 = plan_rfft(a, flags=FFTW.ESTIMATE)
-    @test p1*a ≈ p2*a
-    # make sure threads are actually being used for p2
-    # (tests #21163).
-    if FFTW.has_sprint_plan
-        @test !occursin("dft-thr", string(p1))
-        @test occursin("dft-thr", string(p2))
-    end
-end
-
 # fft
 a = rand(8) + im*rand(8)
 @test norm(ifft(fft(a)) - a) < 1e-8
@@ -516,3 +502,18 @@ let A = rand(Float32, 35), Ac = rand(Complex{Float32}, 35)
         @test_throws ArgumentError plan_fft(Array{Complex{Float32}}(undef, 32)) * view(Ac, 2:33)
     end
 end
+
+# Base Julia issue #19892
+let a = randn(10^5,1)
+    FFTW.set_num_threads(1)
+    p1 = plan_rfft(a, flags=FFTW.ESTIMATE)
+    FFTW.set_num_threads(2)
+    p2 = plan_rfft(a, flags=FFTW.ESTIMATE)
+    @test p1*a ≈ p2*a
+    # make sure threads are actually being used for p2
+    # (tests #21163).
+    if FFTW.has_sprint_plan
+        @test !occursin("dft-thr", string(p1))
+        @test occursin("dft-thr", string(p2))
+    end
+end