Type instability of nested function

New to Julia
1

I have the following function, and it has a type instability (Core.Box).

In the real code, I use f as an argument of another function which calls f multiple times (e.g. bisection to solve f(x)=0).

How can I make this function type stable?

function test(a)
    b = a
    if a < 0
        b = -a
    end
    f(x) = x + b
    f(10)
end

julia> @code_warntype test(1)
Variables
  #self#::Core.Compiler.Const(test, false)
  a::Int64
  b::Core.Box
  f::var"#f#61"

Body::Any
1 ─       (b = Core.Box())
│         Core.NewvarNode(:(f))
│   %3  = a::Int64
│         Core.setfield!(b, :contents, %3)
│   %5  = (a < 0)::Bool
└──       goto #3 if not %5
2 ─ %7  = -a::Int64
└──       Core.setfield!(b, :contents, %7)
3 ┄       (f = %new(Main.:(var"#f#61"), b))
│   %10 = (f)(10)::Any
└──       return %10
2

I guess you hit performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub

Current solution is just to put your closure in let block

function test(a)
    b = a
    if a < 0
        b = -a
    end
    let b = b
        f(x) = x + b
        f(10)
    end
end

julia> @code_warntype test(1)
MethodInstance for test(::Int64)
  from test(a) in Main at REPL[3]:1
Arguments
  #self#::Core.Const(test)
  a::Int64
Locals
  b@_3::Int64
  f::var"#f#4"{Int64}
  b@_5::Int64
Body::Int64
4 Likes
3

The function body is probably a dummy, but it still might be helpful: conditionals blocks in Julia are expressions, not statements. So, you may write

function test(a)
    b = if a >= 0
        a
    else
        -a
    end
    f(x) = x + b
    f(10)
end

If you have to have a mutable variable captured, use a reference to avoid re-binding:

function test_ref(a)
    b = Ref(a)
    if a < 0
        b[] = -a
    end
    f(x) = x + b[]
    f(10)
end
2 Likes
4

By the way, what does (b = Core.Box()) mean in the original example? Is it due to capturing? :wink:

5

I am not an expert, but I think you are right, capturing is the most reasonable explanation. IIRC it is a scope problem: variable b is changed outside of closure, so compiler puts it in a Box to trace these changes (presumably it is the same trick as Ref). When new scope is introduced with the help of the let block, compiler is happy and not boxing anything. But I may be wrong, so take it with a grain of salt.

1 Like
6

The b in the let scope is never changed. Thus there is no boxing, right? It is the first time I knew that closure capturing can lead to type instability.