Memory allocation with function passed as arguments

General Usage Performance
1

When I call test() twice after compiling code

fn = (x -> 2*exp(x), x -> 2*sin(x))

function test_calc(x, fn1, fn2)
    @. x = x + fn1(x) + fn2(x)
end

function test()
    A = Float64.(collect(1:10))
    @time test_calc(A,x -> 2*exp(x), x -> 2*sin(x))
    fn1, fn2 = fn
    A = Float64.(collect(1:10))
    @time test_calc(A, fn1, fn2)
    A = Float64.(collect(1:10))
    @time test_calc(A, fn...)
    return nothing
end

I get

julia> test()
  0.055331 seconds (3.96 k allocations: 209.977 KiB, 99.92% compilation time)
  0.121727 seconds (251.28 k allocations: 14.448 MiB, 99.14% compilation time)
  0.000023 seconds (7 allocations: 144 bytes)

julia> test()
  0.000018 seconds (6 allocations: 128 bytes)
  0.000085 seconds (6 allocations: 128 bytes)
  0.000019 seconds (7 allocations: 144 bytes)
  • Is there any way to pass functions as arguments without allocating memory (this is quite important to the code I’m developing)?
  • Why do the 1st two calculation methods allocate more memory after being compiled (after the 1st test() call)?
  • Why is the 3rd calculation method allocating more memory than the other two after the 2nd test() call?
2

The solution is rather simple: add @inline to test_calc. Then allocation is gone.

fn = (x -> 2*exp(x), x -> 2*sin(x))

@inline function test_calc(x, fn1, fn2)
    @. x = x + fn1(x) + fn2(x)
end

function test()
    A = Float64.(collect(1:10))
    @time test_calc(A,x -> 2*exp(x), x -> 2*sin(x))
    fn1, fn2 = fn
    A = Float64.(collect(1:10))
    @time test_calc(A, fn1, fn2)
    A = Float64.(collect(1:10))
    @time test_calc(A, fn...)
    return nothing
end

Run test() twice time, you get:

julia> test()
  0.000006 seconds
  0.000003 seconds
  0.000003 seconds

I haven’t observe that case in Julia 1.6 and 1.8. They both have 6 allocations (without @inline):

julia> test()
  0.000036 seconds (6 allocations: 128 bytes)
  0.000026 seconds (6 allocations: 128 bytes)
  0.000040 seconds (6 allocations: 128 bytes)

Also, if your anonymous function is complicated, you can add @inline to them to force inlining the function. like @inline(x -> 2*exp(x)).

1 Like
3

Note that in general I would use @btime, and then just interpolate global variables with $:

using BenchmarkTools
begin
    fn = (x -> 2*exp(x), x -> 2*sin(x))
    @inline test_calc(x, fn1, fn2) = @. x = x + fn1(x) + fn2(x)
    @btime test_calc(A, $(x -> 2*exp(x)), $(x -> 2*sin(x))) setup=(A = Float64.(collect(1:10))) evals=1
    fn1, fn2 = fn
    @btime test_calc(A, $fn1, $fn2) setup=(A = Float64.(collect(1:10))) evals=1
    @btime test_calc(A, $fn...) setup=(A = Float64.(collect(1:10))) evals=1
end;

which gives

  238.000 ns (0 allocations: 0 bytes)
  239.000 ns (0 allocations: 0 bytes)
  238.000 ns (0 allocations: 0 bytes)
1 Like
4

Note also that you can also eliminate the allocations by using a type parameter for the function types, rather than @inline:

test_calc(x, fn1::F1, fn2::F2) where {F1,F2} = @. x = x + fn1(x) + fn2(x)

This is covered in the manual’s performance tips: Julia doesn’t always specialize functions on function arguments, but type parameters are a workaround: Be aware of when Julia avoids specializing

2 Likes
5

Wow! Thank you very much. I am definitely going to read the manual so that I understand what’s going on.

1 Like