Is Box.Core too aggressive?

[cscherrer]

Hi, I've been able to get a little more detail on an issue I've been seeing in Soss and Tilde. As usual, it's entirely possible the bug may be in one of my packages, but I hope this example might help narrow down the problem.

Say we have this in Tilde:

m = @model n begin
    p ~ Uniform()
    x ~ For(n) do j
        Bernoulli(p)
    end
end

x = rand(Bool, 3)
logdensityof(m(3) | (;x), (p=0.2,))

From that, Tilde produces this code:

:(@inline function (_mc, _cfg, _ctx, _pars, _retfun)
    local _retn
    _args = (Tilde.argvals)(_mc)
    _obs = (Tilde.observations)(_mc)
    _cfg = merge(_cfg, (args = _args, obs = _obs, pars = _pars))
    n::Int64 = _args.n
    x::Vector{Bool} = _obs.x
    p::Float64 = _pars.p
    (p, _ctx, _retn) = (Tilde.tilde)(DensityInterface.logdensityof, (Accessors.opticcompose)(), static(:p), (Tilde.Unobserved)(p), Uniform(), _cfg, _ctx)
    _retn isa Tilde.ReturnNow && return _retn.value
    (x, _ctx, _retn) = (Tilde.tilde)(DensityInterface.logdensityof, (Accessors.opticcompose)(), static(:x), (Tilde.Observed)(x), For(function (j,)
                    Bernoulli(p)
                end, n), _cfg, _ctx)
    _retn isa Tilde.ReturnNow && return _retn.value
    return (Tilde.var"#61#63"())(_ctx, _ctx)
end)

In this code, p is declared local, then assigned, and then used in a closure. But it should be easy to guarantee that p does not change its type. We usually have this, which is why I have lines like p::Float64 = _pars.p. I got this from the Julia docs, which say

If captured variables are used in a performance-critical section of the code, then the following tips help ensure that their use is performant. First, if it is known that a captured variable does not change its type, then this can be declared explicitly with a type annotation (on the variable, not the right-hand side)

Anyway, calling this expression q, we then call

q = from_type(_get_gg_func_body(mk_function(M, q))) |> MacroTools.flatten

which yields

quote
    p = Core.Box()
    $(Expr(:meta, :((Main).inline)))
    local _retn
    _args = (Tilde.argvals)(_mc)
    _obs = (Tilde.observations)(_mc)
    _cfg = (Main).merge(_cfg, (args = _args, obs = _obs, pars = _pars))
    n::Int64 = _args.n
    x::Vector{Bool} = _obs.x
    p.contents::Float64 = _pars.p
    (p.contents, _ctx, _retn) = (Tilde.tilde)(DensityInterface.logdensityof, (Accessors.opticcompose)(), static(:p), (Tilde.Unobserved)(p.contents), (Main).Uniform(), _cfg, _ctx)
    (Main).isa(_retn, (Main).Tilde.ReturnNow) && return _retn.value
    (x, _ctx, _retn) = (Tilde.tilde)(DensityInterface.logdensityof, (Accessors.opticcompose)(), static(:x), (Tilde.Observed)(x), (Main).For(begin
                    let freevars = (p,)
                        (GeneralizedGenerated.Closure){function = (p, j;) -> begin
    begin
        (Main).Bernoulli(p.contents)
    end
end, Base.typeof(freevars)}(freevars)
                    end
                end, n), _cfg, _ctx)
    (Main).isa(_retn, (Main).Tilde.ReturnNow) && return _retn.value
    return (Tilde.var"#61#63"())(_ctx, _ctx)
end

This ends up boxing p. But that's not needed - working with a modified GG I can just remove the lines that create the box, and it works great.

The boxing leads to problems, because the assignment p.contents::Float64 = _pars.p is not valid.

Do you see any changes I could make to my codegen to correct for this, or it does this require a change in GG?