Minimizing allocations from memoized functions

ash
Thu, 29 Oct 2020 05:23:28 +0000
Performance

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.

11 posts - 4 participants