From b10667a3a8281a148b03fb4eb10e0356824a412c Mon Sep 17 00:00:00 2001
From: "Steven G. Johnson" <stevenj@mit.edu>
Date: Tue, 28 Mar 2017 10:05:58 -0400
Subject: [PATCH] updated fix for #19892 (FFTW threads initialization) (#21169)

* updated fix for #19892; initialize FFTW threads the first time the planner is called (#21127 incorrectly prevented threads from being used at all)

* add test for #21163

(cherry picked from commit 378ed8a75a690f5e5b6cb04201a563533d64d452)
---
 base/fft/FFTW.jl | 46 ++++++++++++++++++++++++++++++++--------------
 test/fft.jl      |  8 ++++++--
 2 files changed, 38 insertions(+), 16 deletions(-)

diff --git a/base/fft/FFTW.jl b/base/fft/FFTW.jl
index b40a766e40e3e5..ca1e8f06f3873f 100644
--- a/base/fft/FFTW.jl
+++ b/base/fft/FFTW.jl
@@ -89,6 +89,27 @@ alignment_of(A::FakeArray) = Int32(0)
 
 ## Julia wrappers around FFTW functions
 
+# _init_() must be called before any FFTW planning routine.
+#   -- Once FFTW is split into its own module, this can be called
+#      in the module __init__(), but for now we must call it lazily
+#      in every routine that might initialize the FFTW planner.
+#   -- This initializes FFTW's threads support (defaulting to 1 thread).
+#      If this isn't called before the FFTW planner is created, then
+#      FFTW's threads algorithms won't be registered or used at all.
+#      (Previously, we called fftw_cleanup, but this invalidated existing
+#       plans, causing issue #19892.)
+const threads_initialized = Ref(false)
+function _init_()
+    if !threads_initialized[]
+        stat = ccall((:fftw_init_threads,libfftw), Int32, ())
+        statf = ccall((:fftwf_init_threads,libfftwf), Int32, ())
+        if stat == 0 || statf == 0
+            error("could not initialize FFTW threads")
+        end
+        threads_initialized[] = true
+    end
+end
+
 # Wisdom
 
 # Import and export wisdom to/from a single file for all precisions,
@@ -101,6 +122,7 @@ alignment_of(A::FakeArray) = Int32(0)
 # FFTW's api/import-wisdom-from-file.c file].
 
 function export_wisdom(fname::AbstractString)
+    _init_()
     f = ccall(:fopen, Ptr{Void}, (Cstring,Cstring), fname, :w)
     systemerror("could not open wisdom file $fname for writing", f == C_NULL)
     ccall((:fftw_export_wisdom_to_file,libfftw), Void, (Ptr{Void},), f)
@@ -110,6 +132,7 @@ function export_wisdom(fname::AbstractString)
 end
 
 function import_wisdom(fname::AbstractString)
+    _init_()
     f = ccall(:fopen, Ptr{Void}, (Cstring,Cstring), fname, :r)
     systemerror("could not open wisdom file $fname for reading", f == C_NULL)
     if ccall((:fftw_import_wisdom_from_file,libfftw),Int32,(Ptr{Void},),f)==0||
@@ -120,6 +143,7 @@ function import_wisdom(fname::AbstractString)
 end
 
 function import_system_wisdom()
+    _init_()
     if ccall((:fftw_import_system_wisdom,libfftw), Int32, ()) == 0 ||
        ccall((:fftwf_import_system_wisdom,libfftwf), Int32, ()) == 0
         error("failed to import system wisdom")
@@ -127,25 +151,15 @@ function import_system_wisdom()
 end
 
 function forget_wisdom()
+    _init_()
     ccall((:fftw_forget_wisdom,libfftw), Void, ())
     ccall((:fftwf_forget_wisdom,libfftwf), Void, ())
 end
 
 # Threads
 
-const threads_initialized = Ref(false)
 function set_num_threads(nthreads::Integer)
-    if !threads_initialized[]
-        # must forget wisdom if any FFTW routines have been called
-        # (don't call fftw_cleanup, since that would invalidate existing plans)
-        forget_wisdom()
-        stat = ccall((:fftw_init_threads,libfftw), Int32, ())
-        statf = ccall((:fftwf_init_threads,libfftwf), Int32, ())
-        if stat == 0 || statf == 0
-            error("could not initialize FFTW threads")
-        end
-        threads_initialized[] = true
-    end
+    _init_()
     ccall((:fftw_plan_with_nthreads,libfftw), Void, (Int32,), nthreads)
     ccall((:fftwf_plan_with_nthreads,libfftwf), Void, (Int32,), nthreads)
 end
@@ -159,11 +173,15 @@ typealias PlanPtr Ptr{fftw_plan_struct}
 
 const NO_TIMELIMIT = -1.0 # from fftw3.h
 
-set_timelimit(precision::fftwTypeDouble,seconds) =
+function set_timelimit(precision::fftwTypeDouble,seconds)
+    _init_()
     ccall((:fftw_set_timelimit,libfftw), Void, (Float64,), seconds)
+end
 
-set_timelimit(precision::fftwTypeSingle,seconds) =
+function set_timelimit(precision::fftwTypeSingle,seconds)
+    _init_()
     ccall((:fftwf_set_timelimit,libfftwf), Void, (Float64,), seconds)
+end
 
 # Array alignment mod 16:
 #   FFTW plans may depend on the alignment of the array mod 16 bytes,
diff --git a/test/fft.jl b/test/fft.jl
index 7de9fa967cfa96..eaeea06ab20639 100644
--- a/test/fft.jl
+++ b/test/fft.jl
@@ -2,10 +2,14 @@
 
 # issue #19892
 # (test this first to make sure it happens before set_num_threads)
-let a = randn(10^5,1), p1 = plan_rfft(a)
+let a = randn(10^5,1), p1 = plan_rfft(a, flags=FFTW.ESTIMATE)
     FFTW.set_num_threads(2)
-    p2 = plan_rfft(a)
+    p2 = plan_rfft(a, flags=FFTW.ESTIMATE)
     @test p1*a ≈ p2*a
+    # make sure threads are actually being used for p2
+    # (tests #21163).
+    @test !contains(string(p1), "dft-thr")
+    @test contains(string(p2), "dft-thr")
 end
 
 # fft