Passing a function as an argument via closure

General Usage Performance
, , ,
1

I am creating some numerical methods, and I am running into some issues with closures.

Here’s an example of what I’m doing:

Function to drive myPkg:

function main()
    #allocate temporaries

    function linear()
        # Use temps and do stuff
    end

    function nonlinear()
        # Use temps and do stuff
    end

    myPkg.solver(retVal, linear, nonlinear, otherParams)

end

Function inside my module:

function solver(retval, linear, nonlinear, otherParams)
    for i = 1:N
        linear(retVal, otherParams)
        nonlinear(retVal, otherParams)
    end

    return retVal
end

My problem is that an excessive amount of allocations take place and the function is over ten times slower than if I encapsulate the functions inside of the module. I’m not certain if my concept of a closure is completely wrong, or if the way I have it causes additional temporaries to be created that I don’t see.

If you want a more concrete example checkout my GitHub for comparison of the module that encapsulates the functions (https://github.com/cpross90/BenjaminOnoSolver.jl) and the module that accepts external functions (https://github.com/cpross90/SplitStep.jl). I’m looking to have the module work similarly to how the Optim.jl package receives objective functions for a cost function/gradient/hessian.

If there are any suggestions to how I can tackle this better I am all ears. Just an undergraduate applied math researcher trying to get a hold of Julia.

Thanks!

2

You could be running into performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub … try defining any captured variables in a let block as described in that issue.

3 Likes
3

Here’s a concrete example of what @stevengj is describing:

julia> function solver(f)
         minimum(f, 1:10)
       end
solver (generic function with 1 method)

julia> function slow()
         # allocate temporaries
         x = 1
         
         # define inner function
         function inner(y)
           x = y
           y
         end
         
         solver(inner)
       end
slow (generic function with 1 method)

The telltale sign of issue 15276 is the Core.Box in @code_warntype:

julia> @code_warntype slow()
Body::Int64
1 ─ %1 = %new(Core.Box)::Core.Box
│        (Core.setfield!)(%1, :contents, 1)
│   %3 = %new(Main.:(#inner#12), %1)::getfield(Main, Symbol("#inner#12"))
│   %4 = Base.min::typeof(min)
│   %5 = invoke Base._mapreduce(%3::getfield(Main, Symbol("#inner#12")), %4::typeof(min), $(QuoteNode(IndexLinear()))::IndexLinear, $(QuoteNode(1:10))::UnitRange{Int64})::Int64
└──      return %5

We can fix this by using a Ref, like so:

julia> function fast1()
         x = Ref(1)
         
         function inner(y)
           x[] = y
           y
         end
         
         solver(inner)
       end
fast1 (generic function with 1 method)

julia> @code_warntype fast1()
Body::Int64
1 ─ %1 = %new(Base.RefValue{Int64}, 1)::Base.RefValue{Int64}
│   %2 = %new(getfield(Main, Symbol("#inner#14")){Base.RefValue{Int64}}, %1)::getfield(Main, Symbol("#inner#14")){Base.RefValue{Int64}}
│   %3 = Base.min::typeof(min)
│   %4 = invoke Base._mapreduce(%2::getfield(Main, Symbol("#inner#14")){Base.RefValue{Int64}}, %3::typeof(min), $(QuoteNode(IndexLinear()))::IndexLinear, $(QuoteNode(1:10))::UnitRange{Int64})::Int64
└──      return %4

or using a let block as described in performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub

2 Likes
4

As a general coding pattern to teach to people, do you always avoid this issue with 100% certainty if you have a closure over a struct, Tuple or a NamedTuple?

5

I’m not sure exactly what you mean, but I think the answer is no: the issue is just as likely to occur with a struct as with an Int.

6

Well, with the “Global scope debacle” things are different for immutables from vectors, etc. The issue is that a name binding changes above in the

 function inner(y)
           x = y
           y
         end

function. This means that the compiler has a lot of trouble determining types because it knows the names are rebound. I guess the essence of my quetsion is: if you avoid changes in the named bindings, then does performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub bind in practice?

7

You mean that assignment (a = ...) is different from mutation (foo!(a))?

Anyway, avoiding rebinding variables that are closed over can help with that issue, yes.

8

A let block is necessary when you see a Core.Box around a variable which is read but whose binding does not change. If you need to rebind a variable, then you can instead use a Ref and mutate that Ref. My general experience has been that the first issue (Boxing of variables which are not re-bound) is less common now, but the second issue (Boxing of variables whose bindings need to change) is still likely.

2 Likes
9

Yes. My issue is trying to come up with coding patterns for beginners to avoid this. They have enough trouble understanding basic scoping, and won’t be able to interpret warntype.

OK, that is helpful. When you say “less common”, any sense of the frequency? Assuming that you have a set of users who will never be able to reliably run @code_warntype, do you think they can safely avoid this in most cases these days if they don’t rebind? I think I can teach people not to rebind, as it is easy to visually parse the x = ...

10

Wow, this is awesome, thanks to everyone who’s left something here.

It appears that I might be having other type stability issues elsewhere. I’m not seeing anything in @code_warntype that would indicate type instability in the closure, so I’m going to comb through the module more carefully. This definitely helps give me a better idea of how to implement closures. Thanks again!