Correctness of Libtask with submodels depends on whether submodels' functions are inlined or not

[penelopeysm]

Consider the following models, which are entirely equivalent:

using Turing, Random

@model function inner(y, x)
    y ~ Normal(x)
end
@model function nested(y)
    x ~ Normal()
    a ~ to_submodel(inner(y, x))
end
m1 = nested(1.0)
mean(sample(Xoshiro(468), m1, PG(10), 1000))

@model function single(y)
    x ~ Normal()
    y ~ Normal(x)
end
m2 = single(1.0)
mean(sample(Xoshiro(468), m2, PG(10), 1000))

It turns out that (at least on my system) both of these sample correctly with PG (giving a mean x of 0.46).

If you think about it carefully though, there's no reason why the submodel one should work. The call to Libtask.produce is inside the y ~ Normal(x) statement, which is in turn inside the inner submodel. So, in order to get all the way to the produce call, we need to mark everything along the call stack with Libtask.might_produce. This is all basic Libtask functionality.

The problem is that the method inner(model, varinfo, ...) itself has never been marked as produce-able! So how does this sampling actually work correctly, then?

The answer is that if the submodel is small enough, then the method inner(model, varinfo, ...) will just be inlined.

We can demonstrate what happens when this fails to inline, by changing the output of the DynamicPPL compiler. If we were to apply this patch on the current release of DPPL (v0.39.13)

diff --git a/src/compiler.jl b/src/compiler.jl
index 84a9a485..604e0b32 100644
--- a/src/compiler.jl
+++ b/src/compiler.jl
@@ -710,6 +710,7 @@ function build_output(modeldef, linenumbernode)
     # element in the returned value is always the most up-to-date `__varinfo__`.
     # See the docstrings of `replace_returns` for more info.
     evaluatordef[:body] = MacroTools.@q begin
+        @noinline
         $(linenumbernode)
         $(replace_returns(add_return_to_last_statment(modeldef[:body])))
     end

and rerun the PG sampling for the submodel, we get a mean x of 0.01, which is consistent with the call to produce never actually being triggered.

So essentially, we are pretty much at the mercy of the Julia compiler here. If it decides to inline the submodel, we are good. If it doesn't, we will silently give wrong results.

[penelopeysm]

I don't immediately see a very clean answer to this. The most direct solution is that the user will have to mark all their models as might_produce themselves (just like how we force them to do it if the model has keyword arguments). The other option would be to make Libtask a DynamicPPL dependency, and modify the @model macro to additionally emit extra code to mark all the newly defined methods with might_produce.

[penelopeysm]

Another possibility would be to write a DynamicPPLLibtask extension, that extends check_model to also check if any inner submodels are marked as might_produce (we can do this quite easily I think by checking if right isa Submodel and calling Libtask.might_produce on the relevant method, which we can get from methods(right.model.f)). In that case, the remedy is still the user's responsibility (they will get an error that will tell them to mark the submodel as might_produce) but at least we don't silently give wrong results (which is my nightmare scenario).

I might try to experiment with that and see if we can make it work, as that seems like a less intrusive option than hardcoding it into the output of @model.

[0AnshuAditya0]

PR #198 in Libtask adds @might_produce - is that directly intended to fix this? because even if Libtask exposes the macro, DynamicPPL's compiler still needs to emit it on the evaluator function for submodels. so the fix probably needs changes in both places - Libtask side to expose the mechanism, and compiler.jl to actually annotate the generated evaluator with @might_produce whenever it's building a submodel evaluator.

If we go the @inline route, we need to be careful. @inline is technically just a hint; the compiler will ignore it if a submodel exceeds certain complexity limits or contains specific control flow. If the compiler decides not to inline a large submodel, It will fail

Also the test suite concern - if all existing submodel + PG tests pass, it's probably because every test submodel happens to be small enough to inline.So tests give a false sense of security. a more complex submodel in production would silently break without anyone noticing.

[penelopeysm]

is that directly intended to fix this?

@might_produce is just a thin wrapper around might_produce, which has existed since the beginning of (this version of) Libtask.

DynamicPPL's compiler still needs to emit it on the evaluator function for submodels

DynamicPPL doesn't currently depend on Libtask, so this has to be handled at the Turing level, which #2778 now does.

[penelopeysm]

Closed by #2780