I have the following MWE using a memoized function:
using Memoization
using FFTW
struct Operators{F}
K::F
# More operators
end
function run(N, α)
M = 512
ω = rand(M)
# Cache the kinetic factor
function K(α)
fun = if α == 0
@memoize function K_cubic(dx::Real)
ifftshift(cis.(dx*ω.^2/2))
end
K_cubic
elseif α > 0
@memoize function K_hirota(dx::Real)
ifftshift(cis.(dx*(ω.^2/2 - α*ω.^3)))
end
K_hirota
end
return fun
end
ops = Operators(K(α))
ψ = Array{Complex{Float64}}(undef, M, N)
dx = 1/N
myloop(ψ, dx, ops, N)
end
function myloop(ψ, dx, ops, N)
for i = 1:N-1
@time @views T2A!(ψ[:, i+1], ψ[:, i], dx, ops)
# or
# @time @views T4A!(ψ[:, i+1], ψ[:, i], dx, ops)
end
end
function T2A!(ψₒ, ψᵢ, dx, ops) # 3 allocs total
# Do some non allocating operations using stuff in ops
ψₒ .= ops.K(dx) .* ψᵢ # 3 allocs from memoized K(dx)
# Do some more non allocating operations using stuff in ops
end
function T4A!(ψₒ, ψᵢ, dx, ops) # 3 allocs total
T2A!(ψₒ, ψᵢ, 0.123*dx, ops) # 3 allocs from memoized K(dx)
T2A!(ψₒ, ψₒ, -0.456*dx, ops) # 3 allocs from memoized K(dx)
T2A!(ψₒ, ψₒ, 0.123*dx, ops) # 3 allocs from memoized K(dx)
end
run(1000, 0.1)
After the first run of T2A!, every iteration of the loop has 3 allocations coming from the line ψₒ .= ops.K(dx) .* ψᵢ. Is it possible to get rid of them? The allocations are quite small (a few tens or hundreds of bytes).
Note that I can’t just cache K at the given dx and pass it as an array to T2A! due to how T4A! works (my production code has a function T6A! that calls T4A! with different values of dx, and functions that call T6A! etc). Without the “higher order” functions that change the dx, this approach worked fantastically, but now I’m stuck.