Repeated allocations for fft in pseudospectral code

General Usage Performance
,
2

Use a plan:

julia> using FFTW, LinearAlgebra, BenchmarkTools

julia> u = randn(32);

julia> function apply_n_fft(u, n)
           uhat = complex.(similar(u));
           for _ in 1:n
               uhat .= fft(u);
           end
       end
apply_n_fft (generic function with 1 method)

julia> @btime apply_n_fft($(u), 5);
  7.727 μs (49 allocations: 7.83 KiB)

julia> @btime apply_n_fft($(u), 50);
  81.314 μs (454 allocations: 70.41 KiB)

julia> function apply_n_fft_plan(u, n)
           uc = complex(u)
           uhat = similar(uc);
           plan = plan_fft(uc)
           for _ in 1:n
               mul!(uhat, plan, uc)
           end
       end
apply_n_fft_plan (generic function with 1 method)

julia> @btime apply_n_fft_plan($(u), 5);
  1.713 μs (8 allocations: 1.34 KiB)

julia> @btime apply_n_fft_plan($(u), 50);
  2.932 μs (8 allocations: 1.34 KiB)

julia> @btime apply_n_fft_plan($(u), 500);
  15.597 μs (8 allocations: 1.34 KiB)

Without a precomputed plan, fft internally needs to (a) re-compute the same plan all over again which is also a very time-consuming task which you don’t really want to repeat, (b) convert the input array to complex, and (c) allocate the output array, every single time. With the pre-computed plan, and having the input array already as complex numbers, lets you save all these repeated allocations, and the mul! would be entirely in-place, with the number of total allocations not depending on n.

6 Likes