Type instability in Julia Generator

Viewed 89

While trying to reduce the number of allocations generated by a function computing a likelihood by e.g. using Generator expressions, I came across the following behavior which I do not quite understand. Take the following two functions:

function testMax!(x,X,β)
    xmax = 0.0
    @inbounds for i ∈ eachindex(x)
        x[i] = X[i,2] * β[1] + X[i,2] * β[2]
        if x[i] > xmax 
            xmax = x[i]
        end 
    end 
    y = 0.0 
    for i ∈ eachindex(x)
        y += exp(x[i]-xmax)
    end
    return xmax, y
end

function testMaxWeird!(x,X,β)
    xmax = 0.0
    @inbounds for i ∈ eachindex(x)
        x[i] = X[i,2] * β[1] + X[i,2] * β[2]
        if x[i] > xmax 
            xmax = x[i]
        end 
    end 
    y = sum(exp(x[j]-xmax) for j ∈ eachindex(x))
    return xmax, y
end

Both generate the same output

using Random
Random.seed!(1234)
H = 10000;
X = rand(H,2);
β = rand(2);

x = zeros(H);
testMax!(x,X,β)
x = zeros(H);
testMaxWeird!(x,X,β)

returns (1.0772897308017204, 6101.682959406999). However, the first one is type stable, while the second one is not (and therefore much slower).

@code_warntype testMax!(x,X,β)
@code_warntype testMaxWeird!(x,X,β)

In particular, the problem lies with the type of y and xmax, the difference in the outputs being in the @code_warntype lines

y::Float64
xmax::Float64

versus

 y::Any
 xmax@_9::Core.Box

I am just confused as of why exactly this occurs, and whether it is due to bad practice on how I am defining xmax multiple times within the function, or to the way I am using the Generator expression?

Edit/Follow-up

The references and solutions provided are very helpful. I am still somewhat confused as to when exactly this can be expected to happen -- is it due to the way that xmax is updated within the for-loop, is its scope different from any other local variable defined within the function? Why does e.g. the (less efficient) way of computing the max not lead to the same closure issue?

function testMax2!(x,X,β)
    @inbounds for i ∈ eachindex(x)
        x[i] = X[i,2] * β[1] + X[i,2] * β[2]
    end 
    xmax = maximum(x)
    y = sum(exp(x[j]-xmax) for j ∈ eachindex(x))            
    return xmax, y
end

Edit 2: Nevermind, I think the Performance Tips explain this: "The parser, when translating it into lower-level instructions, substantially reorganizes the above code by extracting the inner function to a separate code block." I assume that this means that it comes from the variable being assigned multiple times such that the "reordering" of the code might lead to confusion.

1 Answers
Related