Understanding the output of @code_warntype

New to Julia
,
1

Hi there,

I simulated some data which I stored as an Array{Observation} where Observation is a composite type that stores the elements of each observation. I then I wrote a function to compute the log-likelihood of the data as:

struct Observation{T}
    y::T
    x1::T 
    x2::T
end

# this function will generate the logit data for Q4 
function simulate_logit_data(; n_obs = 300000)
    # define the DGP 
    X1 = 1.0:0.01:10.0
    X2 = 18:0.01:65.0
    α = -1.0
    θ_1 = 0.5
    θ_2 = 0.02
    # pre-allocate the container for the data 
    data = Array{Observation}(undef, n_obs)
    for i in eachindex(data)
        x1 = rand(X1)
        x2 = rand(X2)
        u = rand()
        y = (α + θ_1*x2 + θ_2*x2 + log(u/(1-u)) > 0) ? 1.0 : 0.0
        data[i] = Observation(y, x1, x2)
    end
    return data
end

function Xθ(θ, obs; f = fieldnames(Observation))
    result::Float64 = θ[1]
    for i in 2:length(f)
        result += θ[i]*getfield(obs, f[i])
    end
    return result
end

# this function will compute the log-likelihood given a parameter space point 
function LL(θ, data)
    ll::Float64 = 0.0
    for i in eachindex(data)
        ll += -log(1+exp(Xθ(θ, data[i]))) + getfield(data[i], :y)*Xθ(θ, data[i])
    end 
    return ll
end

function main()
    data = simulate_logit_data()
    @btime LL($(rand(3)), $data)
end 

main()

But the performance is not as expected:

  1. I am getting a lot of allocations. I don’t seem to understand why this is happening.
julia> @btime LL($rand(3), $data)
  37.647 ms (2400002 allocations: 36.62 MiB)
-21070.048454177613
  1. I tried to check for type stability using @code_warntype but I can not understand where the issues come. Here I attach a screenshot with partial output (due to space constraint) but which contains the two issues:
MethodInstance for LL(::Vector{Float64}, ::Vector{Observation})
  from LL(θ, data) @ Main REPL[12]:1
Arguments
  #self#::Core.Const(Main.LL)
  θ::Vector{Float64}
  data::Vector{Observation}
Locals
  @_4::Union{Nothing, Tuple{Int64, Int64}}
  ll::Float64
  i::Int64
  @_7::Float64
  @_8::Any
Body::Float64
1 ──       Core.NewvarNode(:(@_4))
│          Core.NewvarNode(:(ll))
│          (@_7 = 0.0)
│    %4  = @_7::Core.Const(0.0)
│    %5  = (%4 isa Main.Float64)::Core.Const(true)
└───       goto #3 if not %5
2 ──       goto #4
3 ──       Core.Const(:(@_7))
│          Core.Const(:(Base.convert(Main.Float64, %8)))
│          Core.Const(:(Main.Float64))
└───       Core.Const(:(@_7 = Core.typeassert(%9, %10)))
4 ┄─ %12 = @_7::Core.Const(0.0)
│          (ll = %12)
│    %14 = Main.eachindex(data)::Base.OneTo{Int64}
│          (@_4 = Base.iterate(%14))
│    %16 = @_4::Union{Nothing, Tuple{Int64, Int64}}
│    %17 = (%16 === nothing)::Bool
│    %18 = Base.not_int(%17)::Bool
└───       goto #10 if not %18
5 ┄─ %20 = @_4::Tuple{Int64, Int64}
│          (i = Core.getfield(%20, 1))
│    %22 = Core.getfield(%20, 2)::Int64
│    %23 = Main.:+::Core.Const(+)
│    %24 = ll::Float64
│    %25 = Main.:+::Core.Const(+)
│    %26 = Main.:-::Core.Const(-)
│    %27 = Main.log::Core.Const(log)
│    %28 = Main.:+::Core.Const(+)
│    %29 = Main.exp::Core.Const(exp)
│    %30 = Main.Xθ::Core.Const(Main.Xθ)
│    %31 = i::Int64
│    %32 = Base.getindex(data, %31)::Observation
│    %33 = (%30)(θ, %32)::Float64
│    %34 = (%29)(%33)::Float64
│    %35 = (%28)(1, %34)::Float64
│    %36 = (%27)(%35)::Float64
│    %37 = (%26)(%36)::Float64
│    %38 = Main.:*::Core.Const(*)
│    %39 = Main.getfield::Core.Const(getfield)
│    %40 = i::Int64
│    %41 = Base.getindex(data, %40)::Observation
│    %42 = (%39)(%41, :y)::Any
│    %43 = Main.Xθ::Core.Const(Main.Xθ)
│    %44 = i::Int64
│    %45 = Base.getindex(data, %44)::Observation
│    %46 = (%43)(θ, %45)::Float64
│    %47 = (%38)(%42, %46)::Any
│    %48 = (%25)(%37, %47)::Any
│    %49 = (%23)(%24, %48)::Any
│          (@_8 = %49)
│    %51 = @_8::Any
│    %52 = (%51 isa Main.Float64)::Bool
└───       goto #7 if not %52
6 ──       goto #8
7 ── %55 = @_8::Any
│    %56 = Base.convert(Main.Float64, %55)::Any
│    %57 = Main.Float64::Core.Const(Float64)
└───       (@_8 = Core.typeassert(%56, %57))
8 ┄─ %59 = @_8::Float64
│          (ll = %59)
│          (@_4 = Base.iterate(%14, %22))
│    %62 = @_4::Union{Nothing, Tuple{Int64, Int64}}
│    %63 = (%62 === nothing)::Bool
│    %64 = Base.not_int(%63)::Bool
└───       goto #10 if not %64
9 ──       goto #5
10 ┄ %67 = ll::Float64
└───       return %67

thanks a lot in advance.

2

Hard to help if you don’t share a working example of the input data and the full printout of @code_warntype (you omitted the part relevant to the possibly problematic @_8::Any). Don’t share a screenshot of the REPL, copy and paste the text. If space is an issue even with the formatted scroll bar, you can put text in a Hide Details block for us to expand and collapse manually (click gear icon when writing a post/comment).

Right off the bat though, change to @btime LL($(rand(3)), $data) or use a setup to do rand(3) so you’re not profiling that call and its allocations. LL(::Vector{Float64}, ::Vector{Observation}) also seems like a problem for type stability, Observation is an abstract type. EDIT: now you shared the full @code_warntype, type instability does indeed start at:

│    %39 = Main.getfield::Core.Const(getfield)
│    %40 = i::Int64
│    %41 = Base.getindex(data, %40)::Observation
│    %42 = (%39)(%41, :y)::Any

It can’t infer a concrete type for the y field of the abstract Observation. Do you need the vector to hold any instance of Observation? Or is there a particular value for T in Vector{Observation{T}} you can set?

2 Likes
3

Oh I see. The compiler only knows data::Vector{Observation} and since its general it can’t infer. I do not need the T flexibility do I changed it to all fields being Float64 and now the Any red part of the @code_warntype output disappeared and:

@btime LL($(rand(3)), data)
  7.757 ms (1 allocation: 16 bytes)
-58.12176896543346

Wow! Okay, I guess that gets rid of the alloctions thanks a lot.

Any clue about the @_4::Union{Nothing, Tuple{Int64, Int64}} yellow part?

4

That’s just an underlying iteration variable’s type inference. When the iterable runs out of elements, it has a value of nothing. In general you don’t have to worry about inferences of 2-3 element Unions, especially if they don’t propagate into larger ones or Any via multiple arguments in function calls. The compiler just duplicates the code into 2-3 type-stable branches if needed.

1 Like
5

Quick follow-up question if thats okay:

My gradient is type-stable but it causing a large number of allocations:

# compute the gradient 
function ∇LL(
            θ, 
            data;
            f = fieldnames(Observation)
            )
    grad = zeros(length(θ))
    X = zeros(3)
    for i in eachindex(data)
        X .= map(x -> getfield(data[i], x), f)
        grad .+= (getfield(data[i], :y) - exp(Xθ(θ, data[i]))/(1+exp(Xθ(θ, data[i])))) * X
    end
    return grad
end 

@btime ∇LL($(rand(3)), $data)
  20.653 ms (600004 allocations: 22.89 MiB)

My understanding is that I am re-using the initial containers every time so why that many allocations?

Thanks!

6

This allocates. Either use broadcasting instead of map, such that it is fused with the broadcasted assignment, or don’t use broadcasting at all and just write a loop.

Same for the grad assignment in the loop: The RHS isn’t broadcasted and contains a vector scalar multiplication IIUC. This will allocate.

2 Likes
7

Thanks a lot!

8

I’d suggest some refactor to make code more generic and avoid explicit container allocation:

using Random

function simulate_logit_data(
    ;
    X1::AbstractVector{T} = 1.0:0.01:10.0,
    X2::AbstractVector{T} = 18:0.01:65.0,
    α::T = -1.0,
    θ_1::T = 0.5,
    θ_2::T = 0.02,
    n_obs = 300000,
    rng = copy(Random.default_rng()),
) where {T<:Real}
    data = map(1:n_obs) do _
        x1 = rand(rng, X1)
        x2 = rand(rng, X2)
        u = rand(rng)
        y = (α + θ_1*x2 + θ_2*x2 + log(u/(1-u)) > 0) ? 1.0 : 0.0
        return Observation(y, x1, x2)
    end
    return data
end

function Xθ(θ, obs::Observation)
    return float(θ[1] + θ[2] * obs.x1 + θ[3] * obs.x2)
end

function LL(θ, data)
    ll = sum(data) do x
        -log(1+exp(Xθ(θ, x))) + x.y * Xθ(θ, x)
    end
    return ll
end
9

And for the gradient,

function ∇LL(
    θ, 
    data::AbstractVector{<:Observation},
)
    grad = fill!(similar(θ), 0) # ensure a mutable container
    for obs in data
        X = (obs.y, obs.x1, obs.x2)
        grad .+= (obs.y - exp(Xθ(θ, obs))/(1+exp(Xθ(θ, obs)))) .* X
    end
    return grad
end