Implement allocation tests for FastLDF

[penelopeysm]

I am actually quite stumped as to why @allocated behaves differently compared to Chairmarks.

module Foo

using DynamicPPL, Distributions, Chairmarks, LogDensityProblems
@model function f()
    x ~ Normal()
    1.0 ~ Normal(x)
end
@model function submodel_inner()
    m ~ Normal(0, 1)
    s ~ Exponential()
    return (m=m, s=s)
end
@model function submodel_outer()
    params ~ to_submodel(submodel_inner())
    y ~ Normal(params.m, params.s)
    1.0 ~ Normal(y)
end
for model in (f(), submodel_inner(), submodel_outer())
    vi = VarInfo(model)
    fldf = DynamicPPL.Experimental.FastLDF(model, DynamicPPL.getlogjoint_internal, vi)
    x = vi[:]
    LogDensityProblems.logdensity(fldf, x) # force compilation
    println(median(@be LogDensityProblems.logdensity(fldf, x)).allocs)
    println(@allocations LogDensityProblems.logdensity(fldf, x))
end

end

prints 0.0, 2, 0.0, 2, 0.0, 2.

The 2 allocations come from not interpolating fldf and x (verifiable by running @allocations fldf and @allocations x, which each return 1). So far so good.

But then if you replace the module with @testset, it prints: 0.0, 2, 0.0, 2, 28.0, 29!

Not sure if this relates to #1127.

[penelopeysm]

This works, if both the arguments are declared as const. The problem with this is that you can't declare const inside a @testset, because const declarations have to go at top-level:

using DynamicPPL, LogDensityProblems, Distributions
function foo end
const y = [1.0]
@model foo() = x ~ Normal()
model = foo()
const vi = VarInfo(model)
const fldf = DynamicPPL.Experimental.FastLDF(model, DynamicPPL.getlogjoint_internal, vi)
let 
    LogDensityProblems.logdensity(fldf, y) # force compilation
    println(@allocations LogDensityProblems.logdensity(fldf, y))
    # 0
end

[penelopeysm]

Passing submodel_inner() as an argument to submodel_outer actually makes it work correctly. This was not fun

@testset "allocations" begin
@model function f()
    x ~ Normal()
    1.0 ~ Normal(x)
end
@model function submodel_inner()
    m ~ Normal(0, 1)
    s ~ Exponential()
    return (m=m, s=s)
end
@model function submodel_outer(inner)
    params ~ to_submodel(inner)
    y ~ Normal(params.m, params.s)
    1.0 ~ Normal(y)
end
@testset for model in (f(), submodel_inner(), submodel_outer(submodel_inner()))
    vi = VarInfo(model)
    fldf = DynamicPPL.Experimental.FastLDF(model, DynamicPPL.getlogjoint_internal, vi)
    x = vi[:]
    LogDensityProblems.logdensity(fldf, x) # force compilation
    expected_allocs = (@allocations fldf) + (@allocations x)
    actual_allocs = @allocations LogDensityProblems.logdensity(fldf, x)
    @test actual_allocs == expected_allocs
end
end

[penelopeysm]

Done in #1113