Sampling variables used on rhs inside `For` seems to break GeneralizedGenerated about soss.jl

[thautwarm]

Some reason of this slow has been detected, it is, the variables whose mutations merely occurred in the scope it's defined in, are still wrapped by Core.Box, say,

_ℓ = 0.0
_ℓ += logpdf(Uniform(0, 50), σ)
_ℓ += logpdf(Normal(0, 10), β)
_ℓ += logpdf(Normal(0, 100), α)
...
return _ℓ
...

becomes

_ℓ = Core.Box(0.0)
_ℓ.contents += logpdf(Uniform(0, 50), σ)
_ℓ.contents += logpdf(Normal(0, 10), β)
_ℓ.contents += logpdf(Normal(0, 100), α)
return _ℓ.contents

It could hurt the performance a bit, but not sure if it's the main cause.

EDIT: This issue has been fixed, try out GG.jl's master branch with updating its dependencies.

[cscherrer]

This is great, thank you!

[cscherrer]

the ERROR: type Float64 has no field contents is gone, but no I'm getting this:

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(σ = 15.9,));
ERROR: 'where' not implemented
Stacktrace:
 [1] error(::String) at ./error.jl:33
 [2] ##GeneralizedGenerated 75#479 at /home/chad/.julia/packages/GeneralizedGenerated/3Ob2C/src/closure_conv.jl:62 [inlined]
 [3] (::getfield(GeneralizedGenerated, Symbol("#inject_freesyms_as_arg!#46")){getfield(GeneralizedGenerated, Symbol("#build_arguments#45")){getfield(GeneralizedGenerated, Symbol("#build_argument#34"))}})(::Array{Symbol,1}, ::Expr, ::Base.RefValue{Union{Nothing, JuliaVariables.ScopedVar}}) at /home/chad/.julia/packages/GeneralizedGenerated/3Ob2C/src/closure_conv.jl:61
 [4] (::getfield(GeneralizedGenerated, Symbol("###GeneralizedGenerated 124#528#59")){Array{NameResolution.LocalVar,1},Array{Symbol,1},getfield(GeneralizedGenerated, Symbol("#inject_freesyms_as_arg!#46")){getfield(GeneralizedGenerated, Symbol("#build_arguments#45")){getfield(GeneralizedGenerated, Symbol("#build_argument#34"))}},getfield(GeneralizedGenerated, Symbol("#process_mutable_cells!#56"))})(::Array{Any,1}) at /home/chad/.julia/packages/GeneralizedGenerated/3Ob2C/src/closure_conv.jl:94
 [5] (::getfield(GeneralizedGenerated, Symbol("#closure_conv#57")){getfield(GeneralizedGenerated, Symbol("#inject_freesyms_as_arg!#46")){getfield(GeneralizedGenerated, Symbol("#build_arguments#45")){getfield(GeneralizedGenerated, Symbol("#build_argument#34"))}},getfield(GeneralizedGenerated, Symbol("#process_mutable_cells!#56")),Module})(::JuliaVariables.ScopedFunc) at ./none:0
 [6] top_level_closure_conv(::Module, ::JuliaVariables.ScopedFunc) at /home/chad/.julia/packages/GeneralizedGenerated/3Ob2C/src/closure_conv.jl:134
 [7] #s42#105(::Any, ::Any, ::Any, ::Any, ::Any, ::Any, ::Any) at /home/chad/.julia/packages/GeneralizedGenerated/3Ob2C/src/closure_conv.jl:155
 [8] (::Core.GeneratedFunctionStub)(::Any, ::Vararg{Any,N} where N) at ./boot.jl:524
 [9] xform at /home/chad/git/jl/Soss/src/xform.jl:13 [inlined]

Here's the generated code:

julia> Soss.sourceXform()(cars_formula)
quote
    _result = NamedTuple()
    σ = rand(Uniform(0, 50))
    _t = xform(Uniform(0, 50))
    _result = merge(_result, (σ = _t,))
    β = rand(Normal(0, 10))
    _t = xform(Normal(0, 10))
    _result = merge(_result, (β = _t,))
    α = rand(Normal(0, 100))
    _t = xform(Normal(0, 100))
    _result = merge(_result, (α = _t,))
    μ = α .+ β .* speed
    N = length(speed)
    dist = rand(For(((n,)->begin
                        #= REPL[4]:8 =#
                        Normal(μ[n], σ)
                    end), 1:N))
    _t = xform(For(((n,)->begin
                        #= REPL[4]:8 =#
                        Normal(μ[n], σ)
                    end), 1:N))
    _result = merge(_result, (dist = _t,))
    as(_result)
end

[thautwarm]

I've found out the cause.
The conversion of GG.jl didn't support where or annotating return type, but for the macro gg it is really a limitation due to implementation.
I'll solve this today.

[cscherrer]

@thautwarm I thought I remember that this was fixed, but it's not working. Can you remind me? Maybe it was a slack or GG discussion?

[thautwarm]

Oh, I haven't fixed it.
Sorry that I'm at a ddl, could you please wait until Sep 22?

[cscherrer]

Of course, no problem. Just wanted to be sure it wasn't supposed to be already done :)

[cscherrer]

Is this related?
https://discourse.julialang.org/t/surprising-capture-boxing-behavior-in-closure/20254

[thautwarm]

@cscherrer In some degree a bit related. It's a good thread to discuss closure conversions of Julia. However it's not the reason for producing your issues.

Also, could you try with the latest GG?I think 'where' not implemented should be resolved.

[cscherrer]

It works! I ran once to compile, and timed the second run. Here are the results:

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(dist=cars.Dist,)) ;

  0.194770 seconds (914.91 k allocations: 104.413 MiB, 16.13% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(α = -17.2,));
 36.730053 seconds (62.18 M allocations: 86.757 GiB, 26.35% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(β = 3.91,));
  4.034226 seconds (6.77 M allocations: 9.451 GiB, 26.45% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(σ = 15.9,));
  7.807692 seconds (13.38 M allocations: 18.699 GiB, 26.88% gc time)

I'm surprised α and β are so different. I'll open a new issue for this and then close this one. Thank you!

[thautwarm]

You're welcome! I'm also curious about the performance difference bewteen them.

[cscherrer]

I think this performance issues may be ok, or at least not related to this. Here's the first run of each (so it includes compile time):

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(dist=cars.Dist,)) ;
 18.641762 seconds (44.17 M allocations: 2.274 GiB, 6.61% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(α = -17.2,));
 31.530739 seconds (56.22 M allocations: 71.003 GiB, 22.81% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(β = 3.91,));
 14.497372 seconds (25.80 M allocations: 30.295 GiB, 21.95% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(σ = 15.9,));
  9.550008 seconds (18.23 M allocations: 19.614 GiB, 21.74% gc time)

And the second runs:

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(dist=cars.Dist,)) ;
  0.101172 seconds (905.21 k allocations: 103.306 MiB, 19.32% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(α = -17.2,));
 28.589496 seconds (50.79 M allocations: 71.019 GiB, 26.30% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(β = 3.91,));
 32.164850 seconds (56.28 M allocations: 78.441 GiB, 26.45% gc time)

julia> @time dynamicHMC(cars_formula(speed=cars.Speed),(σ = 15.9,));
  6.939472 seconds (12.39 M allocations: 17.322 GiB, 27.65% gc time)

I think it's likely that some of the long times here may have more to do with the model than with GG. So I'll just close this issue.

[thautwarm]

Could we extract the source generation processes and evaluate them in the top level REPL, and benchmark them and compare the results against those functions generated by GG?
I think time cost of 10+ seconds is not trivial...

[cscherrer]

Oh that's a great idea

[cscherrer]

Just realized something...

Here's the code for dynamicHMC:

function dynamicHMC(m :: JointDistribution, _data, N=1000::Int) 
    ℓ(pars) = logpdf(m, merge(pars, _data))

    t = xform(m,_data)
    P = TransformedLogDensity(t, ℓ)
    ∇P = ADgradient(Val(:ForwardDiff), P)
    results = mcmc_with_warmup(MersenneTwister(), ∇P, N; reporter=DynamicHMC.NoProgressReport());
    samples = TransformVariables.transform.(parent(∇P).transformation, results.chain)
end

This doesn't call GG. logpdf does, but it calls it on merged NamedTuples. So it's really the same call for all of them.

And

julia> @btime logpdf(cars_formula(speed=cars.Speed), truth)
  789.865 ns (5 allocations: 576 bytes)
-270.95393182681335

The xform calls each take about 10 μs, but only run once.