Unexpected memory allocation

Hello,

I have the following code:

using ForwardDiff

function logistic(x, p)
    return p*x*(1-x)
end

function nth_composite(func, x0, p, n)
    x = x0
    for _ in 1:n
        x = func(x, p)
    end
    return x
end

function nth_map(func, n)
    return (x, p) -> nth_composite(func, x, p, n)
end

function jacobian(map, x, p)
    return ForwardDiff.derivative(x0->map(x0, p), x)
end

function Q_rule(map, params, c)
    return x -> (c-jacobian(map, x, params))/(jacobian(map, x, params)-1)
end

function x_rule(map, x0, params, Q)
    return map(x0, params)+Q(x0)*(map(x0, params)-x0)
end

function bwj(map, n, p, seeds, c=0.7; max_iter=500, disttol=1e-8, bintol=1e-7)
    nmap = nth_map(map, n)
    Q = Q_rule(nmap, p, c)
    fps = Float64[]
    for seed in seeds
        x0 = seed
        for _ = 1:max_iter
            x1 = x_rule(nmap, x0, p, Q)
            if abs(nmap(x1, p)-x1) < disttol
                push!(fps, x1)
                break
            end
            x0 = x1
        end
    end
    return fps
end

function example()
    param = 3.828
    n = 5
    nfunc = nth_map(logistic, n)
    c = 0.7
    seeds = LinRange(0.0, 1.0, 30*n)
    maxiter = 500
    tol = 1e-8

    fps = bwj(logistic, n, param, seeds)
    return fps
end
example()

Running this code creates a lot of allocations:

julia> @time example()
  0.001914 seconds (80.90 k allocations: 1.236 MiB)
julia> using BenchmarkTools
julia> @btime example()
  820.244 μs (80903 allocations: 1.24 MiB)

I am aware that I am not preallocating Vector fps but that shouldn’t be the main source of allocations.
Based on profiling tools I identified the following:

  Total:           0     485418 (flat, cum) 200.00%
    109            .          .            
    110            .          .           function bwj(map, n, p, seeds, c=0.7; max_iter=500, disttol=1e-8, bintol=1e-7) 
    111            .          .               nmap = nth_map(map, n) 
    112            .     485418               Q = Q_rule(nmap, p, c) 

  Total:           0     173034 (flat, cum) 71.29%
    120            .     173025                           break 
    121            .          9                       end 
    122            .          .                       x0 = x1 
    123            .          .                   end 
    124            .          .               end 
    125            .          .               return fps 
    126            .          .           end 

This makes no sense to me. Why are there allocations identified on a line where there is end or break? I used this commands for profiling:

julia> using Profile
julia> using PProf
julia> Profile.Allocs.@profile sample_rate=1 example()
julia> PProf.Allocs.pprof(from_c=false)

I have read that closures can cause type-instability but I haven’t been able to improve my code with the suggestions I found.

Any help would be appreciated.

Jonas

ps. I am using julia v1.10.0

Julia does not specialize, by default, on functions that are just passing around. You need to parameterize their type (see Performance Tips · The Julia Language).

julia> function bwj(map::F, n, p, seeds, c=0.7; max_iter=500, disttol=1e-8, bintol=1e-7) where {F<:Function}
           nmap = nth_map(map, n)
           Q = Q_rule(nmap, p, c)
           fps = Float64[]
           for seed in seeds
               x0 = seed
               for _ = 1:max_iter
                   x1 = x_rule(nmap, x0, p, Q)
                   if abs(nmap(x1, p)-x1) < disttol
                       push!(fps, x1)
                       break
                   end
                   x0 = x1
               end
           end
           return fps
       end
bwj (generic function with 4 methods)

julia> @btime example()
  227.235 μs (4 allocations: 1.92 KiB)
3 Likes

Thank you so much! I spent way too much time trying to figure this out.

2 Likes