Weird type instability

I am having trouble understanding why fun1 has a type instability, and not fun2, is this a bug?

function fun1(N::Integer, a::Number, test::Bool=false)
    if test
        b = exp(-abs2(a)/2)
        x = [(3b^n+2)/(n+1) for n = 0:N-1]
        return x
    else
        return rand(N)
    end
end

function fun2(N::Integer, a::Number, test::Bool=false)
    b = exp(-abs2(a)/2)
    x = [(3b^n+2)/(n+1) for n = 0:N-1]
    return x
end
julia> @code_warntype fun1(6,2,true)
Variables:
  #self#::#fun1
  N::Int64
  a::Int64
  test::Bool
  b::Core.Box
  x::Any
  #13::##13#14
...

I don’t have a great answer, but if you move b before the if statement, the code is type-stable.

This is the @code_warntype output for fun1(6,2,true) in 0.6:

Variables:
  #self#::#fun1
  N::Int64
  a::Int64
  test::Bool
  b::Core.Box
  x::Array{_,1} where _
  #3::##3#4

Body:
  begin
      b::Core.Box = $(Expr(:new, :(Core.Box)))
      NewvarNode(:(x::Array{_,1} where _))
...

The Core.Box is highlighted red.

1 Like

This is performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub See various workarounds there…

A type assertion did not help, but moving b did, thanks @anon94023334.

Moving the type assertion over to b helps (using the example in the issue (performance of captured variables in closures · Issue #15276 · JuliaLang/julia · GitHub), if we do:

function fun(N,a,t)
           if t
               b::Float64 = (a/2)
               x = [b for n = 1:N]
               return x
           else
               return Vector{Float64}(N)
           end
end

The type-instability goes away. But b is still boxed.

Helping Julia think about b looks to me like the best option. Consider,

function fun1(N::Integer, a::Number, test::Bool=false)
           if test let b = exp(-abs2(a)/2)
               x = [(3b^n+2)/(n+1) for n = 0:N-1]
               return x
           end else
               return rand(N)
           end
end

This way, not only is fun1 type-stable, but b is not boxed and not allocated when test==false. Also, having the let on the same line as the if and the end with the else makes the code almost un-ugly.

2 Likes

That’s a nice syntax. Would be cool if the let block did not need an end here, but I guess if there wasn’t a bug it wouldn’t be needed anyway…

Forgive me if I am asking a silly question. How is this problem related to 15276? I don’t see an inner function here.

List comprehension has similar issues.