Strange allocation in loop when calling function

Hello!

I’ve been working on writing some very simple ODE solvers as an exercise to improve my understanding of Julia and how it can be used to write fast code. In doing so, I’ve arrived at some code that exhibits some strange behavior regarding unwanted allocations (at least with my current understanding of the language).

Below I’ve included an MWE. Within it, you’ll see that I define an ODE function (two_body), a step function (euler_step), and two test functions. The first, test_with_stepfun, performs several integration steps using the step function. The second, test_no_stepfun, does the same, except I’ve replaced the call to euler_step with the function’s contents.

It appears that test_with_stepfun allocates each time euler_step is called whereas test_no_stepfun does not. Could anyone describe why I’m seeing this behavior?

Thank you!

Here’s the MWE:

using StaticArrays
using BenchmarkTools

# Two body equations of motion
function two_body(u, t)
    r = sqrt(u[1]*u[1] + u[2]*u[2] + u[3]*u[3])
    rinv3 = 1.0 / (r*r*r)
    a = SA[-u[1]*rinv3, -u[2]*rinv3, -u[3]*rinv3]
    return SA[u[4], u[5], u[6], a[1], a[2], a[3]]
end

# Define step function
function euler_step(ode, u, t1, t2)
    h = t2 - t1
    return u + h*ode(u, t1)
end

# Test function with step function
function test_with_stepfun(ode, u0)
    ts = 0.0:0.1:1.0
    u = u0
    for i in 1:length(ts) - 1
        u = euler_step(ode, u, ts[i], ts[i+1])
    end
    return nothing
end

# Test function without contents of step function
function test_no_stepfun(ode, u0)
    ts = 0.0:0.1:1.0
    u = u0
    for i in 1:length(ts) - 1
        h = ts[i+1] - ts[i]
        u = u + h*ode(u, ts[i]) 
    end
    return nothing
end

# Initial state and time info
u0 = SA[1.0, 0.0, 0.0, 0.1, 0.9, 0.1]

@btime test_with_stepfun($two_body, $u0)
@btime test_no_stepfun($two_body, $u0)

Also, here’s what I observe when I use Profile.jl

        - using StaticArrays
        - using BenchmarkTools
        - using Profile
        - 
        - # Two body equations of motion
        - function two_body(u, t)
        -     r = sqrt(u[1]*u[1] + u[2]*u[2] + u[3]*u[3])
        -     rinv3 = 1.0 / (r*r*r)
        -     a = SA[-u[1]*rinv3, -u[2]*rinv3, -u[3]*rinv3]
        -     return SA[u[4], u[5], u[6], a[1], a[2], a[3]]
        - end
        - 
        - # Define step function
        - function euler_step(ode, u, t1, t2)
        0     h = t2 - t1
        0     return u + h*ode(u, t1)
        - end
        - 
        - # Test function with step function
        - function test_with_stepfun(ode, u0)
        0     ts = 0.0:0.1:1.0
        -     u = u0
       64     for i in 1:length(ts) - 1
      960         u = euler_step(ode, u, ts[i], ts[i+1])
        0     end
        0     return nothing
        - end
        - 
        - # Test function without contents of step function
        - function test_no_stepfun(ode, u0)
        0     ts = 0.0:0.1:1.0
        -     u = u0
        0     for i in 1:length(ts) - 1
        0         h = ts[i+1] - ts[i]
        0         u = u + h*ode(u, ts[i]) 
        0     end
        0     return nothing
        - end
        - 
        - # Initial state and time info
        - u0 = SA[1.0, 0.0, 0.0, 0.1, 0.9, 0.1]
        - 
        - 
        - #@btime test_with_stepfun($two_body, $u0)
        - #@btime test_no_stepfun($two_body, $u0)
        - test_with_stepfun(two_body, u0)
        - test_no_stepfun(two_body, u0)
        - Profile.clear_malloc_data()
        - test_with_stepfun(two_body, u0)
        - test_no_stepfun(two_body, u0)

This is: Performance Tips · The Julia Language

To fix, add this type signature to the function:

julia> function test_with_stepfun(ode::F, u0) where {F<:Function}
           ts = 0.0:0.1:1.0
           u = u0
           for i in 1:length(ts) - 1
               u = euler_step(ode, u, ts[i], ts[i+1])
           end
           return nothing
       end
       
       # Test function without contents of step function
test_with_stepfun (generic function with 2 methods)

julia> @btime test_with_stepfun($two_body, $u0)
  313.339 ns (0 allocations: 0 bytes)

3 Likes

Thank you! :grinning:

1 Like