Fix type-inference failure for nested submodels

HISTORY.md

@@ -1,3 +1,9 @@
+# 0.42.1
+
+Fixed a type-inference failure that made nested submodels (a `~ to_submodel(...)` statement inside a model that is itself evaluated as a submodel) very slow.
+
+Previously, evaluating a submodel recursed through the shared `_evaluate!!(::Model, ::AbstractVarInfo)` method. Each level of nesting adds another context layer to the `Model` type, which tripped Julia's type-inference recursion limit: from the first level of nesting onwards the return type was inferred as `Any` and evaluation fell back to runtime dispatch, slowing down both primal and gradient evaluation (the primal slowdown was roughly 15x, growing with nesting depth). Submodel evaluation now calls the model function directly, keeping nested submodels type-stable. See [Turing.jl#2844](https://github.com/TuringLang/Turing.jl/issues/2844).
+
 # 0.42.0
 
 `LogDensityFunction` now performs AD preparation through AbstractPPL's `prepare` / `value_and_gradient!!` interface instead of calling DifferentiationInterface directly. Internally this removes the `_use_closure` heuristic and the explicit `DI.Constant` plumbing; the choice between closure and constants now lives in AbstractPPL.

Project.toml

@@ -1,6 +1,6 @@
 name = "DynamicPPL"
 uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8"
-version = "0.42"
+version = "0.42.1"
 
 [deps]
 ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b"

src/submodel.jl

@@ -181,9 +181,17 @@ function _evaluate!!(
     # (4) Finally, we need to store that context inside the submodel.
     model = contextualize(submodel.model, eval_context)
 
-    # Once that's all set up nicely, we can just _evaluate!! the wrapped model. This
-    # returns a tuple of submodel.model's return value and the new varinfo.
-    return _evaluate!!(model, vi)
+    # Evaluate the wrapped model. These two lines are a verbatim copy of the body of
+    # `_evaluate!!(model::Model, ::AbstractVarInfo)` (in `model.jl`), and the duplication is
+    # deliberate: DO NOT replace them with `return _evaluate!!(model, vi)`. Each level of
+    # submodel nesting grows the contextualised `Model`'s context type, and routing the
+    # recursion through the shared `_evaluate!!(::Model, ...)` method trips Julia's recursion
+    # limiter, which widens the `Model` argument to abstract and collapses the return type to
+    # `Any`. Calling `model.f` directly avoids that. See
+    # https://github.com/TuringLang/DynamicPPL.jl/pull/1427 and
+    # https://github.com/TuringLang/Turing.jl/issues/2844 for the full explanation.
+    args, kwargs = make_evaluate_args_and_kwargs(model, vi)
+    return model.f(args...; kwargs...)
 end
 
 function tilde_observe!!(

test/logdensityfunction.jl

@@ -17,6 +17,30 @@ using DifferentiationInterface: DifferentiationInterface
 using ForwardDiff: ForwardDiff
 using Mooncake: Mooncake
 
+@model function issue_2844_nested_inner()
+    p ~ Normal()
+    return (; p)
+end
+@model function issue_2844_nested_middle()
+    inner ~ to_submodel(issue_2844_nested_inner())
+    return (; p=inner.p)
+end
+@model function issue_2844_nested_outer()
+    middle ~ to_submodel(issue_2844_nested_middle())
+    return 0.0 ~ Normal(middle.p)
+end
+# One submodel boundary deeper than `issue_2844_nested_outer`, so the inference test
+# guards against a regression that merely pushes the recursion limit out by one level
+# rather than removing it.
+@model function issue_2844_nested_middle2()
+    middle ~ to_submodel(issue_2844_nested_middle())
+    return (; p=middle.p)
+end
+@model function issue_2844_nested_outer_deep()
+    middle2 ~ to_submodel(issue_2844_nested_middle2())
+    return 0.0 ~ Normal(middle2.p)
+end
+
 @testset "LogDensityFunction: constructors" begin
     dist = Beta(2, 2)
     @model f() = x ~ dist
@@ -420,6 +444,21 @@ end
             end skip = skip
         end
     end
+
+    # See https://github.com/TuringLang/DynamicPPL.jl/pull/1427.
+    @testset "nested to_submodel return values" begin
+        @testset "$(nameof(model.f))" for model in (
+            issue_2844_nested_outer(), issue_2844_nested_outer_deep()
+        )
+            @testset "$tfm_strategy" for tfm_strategy in (UnlinkAll(), LinkAll())
+                ldf = DynamicPPL.LogDensityFunction(
+                    model, DynamicPPL.getlogjoint_internal, tfm_strategy
+                )
+                x = rand(ldf)
+                @test @inferred(LogDensityProblems.logdensity(ldf, x)) isa Float64
+            end
+        end
+    end
 end
 
 @testset "LogDensityFunction: performance" begin
@@ -442,8 +481,18 @@ end
         y ~ Normal(params.m, params.s)
         return 1.0 ~ Normal(y)
     end
-    @testset for model in
-                 (f(), submodel_inner() | (; s=0.0), submodel_outer(submodel_inner()))
+    # A submodel nested inside another submodel; see
+    # https://github.com/TuringLang/DynamicPPL.jl/pull/1427.
+    @model function submodel_middle(inner)
+        mid ~ to_submodel(inner)
+        return (m=mid.m, s=mid.s)
+    end
+    @testset for model in (
+        f(),
+        submodel_inner() | (; s=0.0),
+        submodel_outer(submodel_inner()),
+        submodel_outer(submodel_middle(submodel_inner())),
+    )
         @testset for tfm_strategy in (UnlinkAll(), LinkAll())
             ldf = DynamicPPL.LogDensityFunction(model, getlogjoint_internal, tfm_strategy)
             x = rand(ldf)

test/submodels.jl

@@ -17,6 +17,15 @@ function get_logp_and_rawval_accs(model::Model)
     return accs
 end
 
+# Models for the nested-submodel type-stability tests; see
+# https://github.com/TuringLang/DynamicPPL.jl/pull/1427. Each level wraps the previous one in
+# a `to_submodel`. They must be defined at module scope: a model defined in local (testset)
+# scope is not type-inferrable, which would mask the property under test.
+@model t2844_inner() = (x ~ Normal(); return (; x))
+@model t2844_middle() = (a ~ to_submodel(t2844_inner()); return (; x=a.x))
+@model t2844_outer() = (b ~ to_submodel(t2844_middle()); return (; x=b.x))
+@model t2844_deeper() = (c ~ to_submodel(t2844_outer()); return (; x=c.x))
+
 @testset "submodels.jl" begin
     @testset "$op with AbstractPPL API" for op in [condition, fix]
         x_val = 1.0
@@ -338,6 +347,22 @@ end
         @test vnt.data.a.data.b.data.x.data isa Matrix{Float64}
         @test size(vnt.data.a.data.b.data.x.data) == (2, 2)
     end
+
+    @testset "type stability of nested submodels (issue #2844)" begin
+        # See https://github.com/TuringLang/DynamicPPL.jl/pull/1427.
+        @testset "$(nameof(model.f))" for model in (
+            t2844_inner(), t2844_middle(), t2844_outer(), t2844_deeper()
+        )
+            # The fast evaluation path: `init!!` into an `OnlyAccsVarInfo`, under both
+            # transform strategies.
+            @testset "$tfm" for tfm in (UnlinkAll(), LinkAll())
+                accs = setacc!!(OnlyAccsVarInfo(), LogPriorAccumulator())
+                @test @inferred(init!!(model, accs, InitFromPrior(), tfm)) isa Tuple
+            end
+            # Evaluating a pre-populated `VarInfo` must also stay type stable.
+            @test @inferred(DynamicPPL.evaluate_nowarn!!(model, VarInfo(model))) isa Tuple
+        end
+    end
 end
 
 end