From 9b93341f807ba861aa05ed61e8bb4a44f1fa7659 Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Sun, 21 Jun 2026 17:00:15 +0100 Subject: [PATCH 1/4] Fix type-inference failure for nested submodels Evaluating a submodel recursed through the shared `_evaluate!!(::Model, ::AbstractVarInfo)` method. Each level of submodel nesting adds another context layer to the `Model` type, so this tripped Julia's type-inference recursion limit (`limit_type_size`): from the first level of nesting onwards the model's return type was inferred as `Any`, evaluation fell back to runtime dispatch, and both primal and gradient evaluation slowed down (~15x primal, growing with nesting depth). Submodel evaluation now calls the generated model function directly instead of going through `_evaluate!!(::Model, ...)`. This keeps the growing context type out of the recursion's dispatch path (`model.f` is a distinct function type per `@model`), so inference stays type stable for nested submodels. Adds type-stability (`@inferred`) and allocation (`iszero(allocs)`) regression tests at both the evaluation and `LogDensityFunction` levels, covering several nesting depths under `UnlinkAll`/`LinkAll`. The tests were verified to fail on the unfixed code. Also bumps the version to 0.42.1 and adds a HISTORY.md entry. See https://github.com/TuringLang/Turing.jl/issues/2844 --- HISTORY.md | 6 ++++ Project.toml | 2 +- src/model.jl | 4 +++ src/submodel.jl | 28 ++++++++++++++++-- test/logdensityfunction.jl | 60 ++++++++++++++++++++++++++++++++++++-- test/submodels.jl | 34 +++++++++++++++++++++ 6 files changed, 128 insertions(+), 6 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index f1245ed48..1dac30587 100644 --- a/HISTORY.md +++ b/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. Because each level of submodel nesting adds another context layer to the `Model` type, this tripped Julia's type-inference recursion limit: from the first level of nesting onwards the model's return type was inferred as `Any`, the evaluation fell back to runtime dispatch, and both primal and gradient evaluation slowed down (the primal slowdown was roughly 15x and grew with nesting depth). Submodel evaluation now calls the generated model function directly, which keeps 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. diff --git a/Project.toml b/Project.toml index 07fdb81ea..cbabc7812 100644 --- a/Project.toml +++ b/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" diff --git a/src/model.jl b/src/model.jl index 6d6a23c28..141910a67 100644 --- a/src/model.jl +++ b/src/model.jl @@ -906,6 +906,10 @@ function _evaluate!!(model::Model, varinfo::AbstractVarInfo) args, kwargs = make_evaluate_args_and_kwargs(model, varinfo) return model.f(args...; kwargs...) end +# NOTE: The two lines above are intentionally duplicated in the submodel evaluation +# path (`_evaluate!!(::Submodel, ...)` in `submodel.jl`). That path must NOT call this +# method, because doing so reintroduces a type-inference recursion limit for nested +# submodels. See the long comment there before refactoring either copy. is_splat_symbol(s::Symbol) = startswith(string(s), "#splat#") diff --git a/src/submodel.jl b/src/submodel.jl index 97618db66..77ddc8550 100644 --- a/src/submodel.jl +++ b/src/submodel.jl @@ -181,9 +181,31 @@ 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) + # Finally, evaluate the wrapped model. The two lines below are a verbatim copy of + # the body of `_evaluate!!(model::Model, ::AbstractVarInfo)`, and morally that is + # what we are doing here. The duplication is deliberate and load-bearing. + # + # DO NOT replace this with `return _evaluate!!(model, vi)`. Submodel evaluation is + # recursive: `model.f` runs the wrapped model, which may itself contain + # `~ to_submodel(...)` statements that re-enter this method. Each level of nesting + # wraps the evaluation context in another `PrefixContext`/leaf-context layer, so the + # `Ctx` type parameter of `model::Model{...,Ctx}` grows by one layer per level. If + # that recursion passed through the shared method `_evaluate!!(::Model, ...)`, type + # inference would see the *same* method recursing with an ever-growing dispatch type, + # trip its recursion-limiting heuristic (`limit_type_size`), and widen `model` back + # to `::Model`. That makes `model.f` unresolvable, collapses the return type to `Any`, + # and turns the whole nested evaluation into runtime dispatch -- the ~15x primal + # slowdown reported in https://github.com/TuringLang/Turing.jl/issues/2844. + # + # Calling `model.f` directly keeps the growing `Ctx` out of the recursion's dispatch + # path. `model.f` is a distinct function type per `@model`, and the only methods left + # in the cycle (`tilde_assume!!`/`_evaluate!!` on `::Submodel`) dispatch on the + # un-grown submodel type, with the growing context confined to a non-dispatched + # `::AbstractContext` argument that inference can safely widen. Note that `@inline`-ing + # `_evaluate!!(::Model, ...)` does NOT help: the recursion limit is applied during + # abstract interpretation, before inlining. + args, kwargs = make_evaluate_args_and_kwargs(model, vi) + return model.f(args...; kwargs...) end function tilde_observe!!( diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index 8a5d37ed3..14f6583e8 100644 --- a/test/logdensityfunction.jl +++ b/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,25 @@ end end skip = skip end end + + # Nested submodels used to break type inference of the log density, because the + # recursive submodel evaluation passed through the shared `_evaluate!!(::Model, ...)` + # method whose `Model` dispatch type grows by one context layer per level. See + # https://github.com/TuringLang/Turing.jl/issues/2844 and the comment in + # `src/submodel.jl`. + @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 +485,21 @@ 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. Before + # https://github.com/TuringLang/Turing.jl/issues/2844 was fixed, evaluating this + # model was not type-stable and allocated (and got worse with each level of + # nesting), because the recursive submodel evaluation passed through the shared + # `_evaluate!!(::Model, ...)` method. See the comment in `src/submodel.jl`. + @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) diff --git a/test/submodels.jl b/test/submodels.jl index b0fe8184c..82f0471c5 100644 --- a/test/submodels.jl +++ b/test/submodels.jl @@ -17,6 +17,16 @@ function get_logp_and_rawval_accs(model::Model) return accs end +# Models for the nested-submodel type-stability tests (issue #2844). Each level wraps the +# previous one in a `to_submodel` and re-exports its return value, so the depth of nesting +# can be varied. They must be defined at module scope: a model defined in local (testset) +# scope captures that scope and is not type-inferrable, which would mask the property under +# test. See the "type stability of nested submodels" testset below. +@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 +348,30 @@ 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 + # Evaluating a submodel recurses into `model.f`, which may contain further + # `to_submodel` statements. Each level of nesting grows the `Model`'s context + # type; if that recursion goes through the shared `_evaluate!!(::Model, ...)` + # method, Julia's recursion-limit heuristic widens the model type and the result + # type collapses to `Any` from the first level of nesting onwards. This made + # nested submodels much slower to evaluate (and to differentiate). See + # https://github.com/TuringLang/Turing.jl/issues/2844 and the comment in + # `src/submodel.jl`. These tests check that evaluation stays type stable at every + # level of nesting. + @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 From 8d9ec4e63a4a64d091d10f6734102488b7e65789 Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Mon, 22 Jun 2026 17:36:28 +0100 Subject: [PATCH 2/4] Clarify nested submodel inference note --- nested-submodel-inference-investigation.md | 132 +++++++++++++++++++++ src/submodel.jl | 46 ++++--- 2 files changed, 159 insertions(+), 19 deletions(-) create mode 100644 nested-submodel-inference-investigation.md diff --git a/nested-submodel-inference-investigation.md b/nested-submodel-inference-investigation.md new file mode 100644 index 000000000..fa8c27346 --- /dev/null +++ b/nested-submodel-inference-investigation.md @@ -0,0 +1,132 @@ +# Nested Submodel Inference Note + +This note explains why `_evaluate!!(::Submodel, ...)` calls `model.f` directly instead of +calling `_evaluate!!(model, vi)`. + +## The Issue + +The problem is not that `_evaluate!!(::Submodel, ...)` is recursive. It is recursive, but +Julia's recursion limiter does not widen that call. + +The problem is the extra recursive call through `_evaluate!!(::Model, vi)`: + +```julia +_evaluate!!(::Submodel, vi, parent_context, left_vn) + -> _evaluate!!(::Model, vi) +``` + +At that call site, Julia infers an argument tuple type for `_evaluate!!(model, vi)`. +The second argument type is a concrete `Model{..., Ctx}`. With nested submodels, the +`Ctx` parameter grows at each level: + +```julia +Model{..., PrefixContext{c, InitContext{...}}} +Model{..., PrefixContext{c, PrefixContext{b, InitContext{...}}}} +``` + +Julia sees the same `_evaluate!!(::Model, vi)` method recurring with a more complex call +signature, so it widens the signature to force convergence: + +```julia +Tuple{typeof(_evaluate!!), Model, OnlyAccsVarInfo{...}} +``` + +Once `model` is only known as abstract `Model`, inference can no longer resolve +`model.f` precisely, and the return type collapses. + +## Main vs Branch Call Chain + +The old path was: + +```text +_evaluate!!(::Model, vi) # top-level model + -> model.f(...) + -> tilde_assume!!(..., ::Submodel, ...) + -> _evaluate!!(::Submodel, vi, context, vn) + -> contextualize(submodel.model, eval_context) + -> _evaluate!!(::Model, vi) # problematic recursive call + -> model.f(...) + -> tilde_assume!!(..., ::Submodel, ...) + -> _evaluate!!(::Submodel, ...) +``` + +The fixed path is: + +```text +_evaluate!!(::Model, vi) # top-level model + -> model.f(...) + -> tilde_assume!!(..., ::Submodel, ...) + -> _evaluate!!(::Submodel, vi, context, vn) + -> contextualize(submodel.model, eval_context) + -> make_evaluate_args_and_kwargs(model, vi) + -> model.f(...) # direct call, no _evaluate!!(::Model) + -> tilde_assume!!(..., ::Submodel, ...) + -> _evaluate!!(::Submodel, ...) +``` + +The fix removes the nested `_evaluate!!(::Model, vi)` call edge. It does not remove +submodel recursion. + +## Why The Submodel Recursion Is Fine + +For a submodel tilde statement, the caller already has the concrete context and submodel +in its own signature: + +```julia +tilde_assume!!(context, submodel, vn, NoTemplate, vi) + -> _evaluate!!(submodel, vi, context, vn) +``` + +The callee signature mostly reorders values that were already present in the caller +signature. In Julia's `Compiler/src/typelimits.jl`, `type_more_complex` has an early exit +for this case: + +```julia +elseif is_derived_type_from_any(unwrap_unionall(t), sources, depth) + return false # t isn't something new +end +``` + +That is why a recursive `_evaluate!!(::Submodel, ...)` call can be eligible for limiting +but still not be widened: the growing context is not newly packaged into another dispatch +type. It is a call argument that was already visible in the caller. + +## Why The Model Recursion Is Not Fine + +The old implementation did this inside `_evaluate!!(::Submodel, ...)`: + +```julia +model = contextualize(submodel.model, eval_context) +return _evaluate!!(model, vi) +``` + +That constructs a fresh `Model{..., Ctx}` type and then passes it to the shared +`_evaluate!!(::Model, vi)` method. The full contextualized model type was not already +present in the caller signature; it was synthesized inside the submodel evaluator. + +So Julia compares recursive `_evaluate!!(::Model, vi)` signatures like: + +```julia +Tuple{typeof(_evaluate!!), Model{..., PrefixContext{c, InitContext{...}}}, VI} +Tuple{typeof(_evaluate!!), Model{..., PrefixContext{c, PrefixContext{b, InitContext{...}}}}, VI} +``` + +The second call is more complex because the `Ctx` parameter inside `Model` is more deeply +nested. `limit_type_size` therefore widens the `Model` argument. + +## Why Manual Inlining Works + +The body of `_evaluate!!(::Model, vi)` is: + +```julia +args, kwargs = make_evaluate_args_and_kwargs(model, vi) +return model.f(args...; kwargs...) +``` + +Doing those two lines directly inside `_evaluate!!(::Submodel, ...)` avoids creating the +recursive `_evaluate!!(::Model, vi)` call edge. The contextualized `model` value still +exists, but it is no longer passed through a recursive call to the shared +`_evaluate!!(::Model, vi)` method. + +An `@inline` annotation on `_evaluate!!(::Model, vi)` is not enough. Julia applies this +recursion-limiting heuristic during abstract interpretation, before optimizer inlining. diff --git a/src/submodel.jl b/src/submodel.jl index 77ddc8550..5351f63a9 100644 --- a/src/submodel.jl +++ b/src/submodel.jl @@ -181,29 +181,37 @@ function _evaluate!!( # (4) Finally, we need to store that context inside the submodel. model = contextualize(submodel.model, eval_context) - # Finally, evaluate the wrapped model. The two lines below are a verbatim copy of - # the body of `_evaluate!!(model::Model, ::AbstractVarInfo)`, and morally that is - # what we are doing here. The duplication is deliberate and load-bearing. + # Finally, evaluate the wrapped model. These two lines are a verbatim copy of the body + # of `_evaluate!!(model::Model, ::AbstractVarInfo)`, so calling `model.f` here is + # behaviour-identical to `_evaluate!!(model, vi)`. The duplication is deliberate. # # DO NOT replace this with `return _evaluate!!(model, vi)`. Submodel evaluation is # recursive: `model.f` runs the wrapped model, which may itself contain - # `~ to_submodel(...)` statements that re-enter this method. Each level of nesting - # wraps the evaluation context in another `PrefixContext`/leaf-context layer, so the - # `Ctx` type parameter of `model::Model{...,Ctx}` grows by one layer per level. If - # that recursion passed through the shared method `_evaluate!!(::Model, ...)`, type - # inference would see the *same* method recursing with an ever-growing dispatch type, - # trip its recursion-limiting heuristic (`limit_type_size`), and widen `model` back - # to `::Model`. That makes `model.f` unresolvable, collapses the return type to `Any`, - # and turns the whole nested evaluation into runtime dispatch -- the ~15x primal - # slowdown reported in https://github.com/TuringLang/Turing.jl/issues/2844. + # `~ to_submodel(...)` statements that re-enter this method. The recursive + # `_evaluate!!(::Submodel, ...)` edge itself is fine: Julia sees the edge, but the + # limiter keeps its signature unchanged. The problematic edge is the old intermediate + # call to the shared `_evaluate!!(::Model, ::AbstractVarInfo)` method. # - # Calling `model.f` directly keeps the growing `Ctx` out of the recursion's dispatch - # path. `model.f` is a distinct function type per `@model`, and the only methods left - # in the cycle (`tilde_assume!!`/`_evaluate!!` on `::Submodel`) dispatch on the - # un-grown submodel type, with the growing context confined to a non-dispatched - # `::AbstractContext` argument that inference can safely widen. Note that `@inline`-ing - # `_evaluate!!(::Model, ...)` does NOT help: the recursion limit is applied during - # abstract interpretation, before inlining. + # Each submodel level constructs a contextualized `Model{...,Ctx}` whose `Ctx` + # parameter nests one `PrefixContext` deeper. If we call `_evaluate!!(model, vi)`, + # inference sees a recursive `_evaluate!!(::Model, ...)` call whose argument tuple + # contains that growing `Model{...,Ctx}` type. Julia's recursion limiter then widens + # the call signature to `Tuple{typeof(_evaluate!!), Model, ...}`. Once `model` is only + # known as abstract `Model`, `model.f` is no longer resolvable and the return type + # collapses to `Any`, causing the slowdown in + # https://github.com/TuringLang/Turing.jl/issues/2844. + # + # Calling `model.f` directly removes only that bad `_evaluate!!(::Model, ...)` edge. + # The remaining recursive submodel edge still exists, but its growing context is a + # normal argument already present in the `tilde_assume!!` caller signature, so the + # limiter does not widen it. + # + # Marking `_evaluate!!(::Model, ...)` `@inline` does NOT help: the recursion limit is + # applied during abstract interpretation, before inlining. + # + # This keeps *distinct* nested submodels type stable to arbitrary depth. A model that + # recurses into the *same* submodel `@model` can still hit the limiter; that is a + # separate, pre-existing limitation. args, kwargs = make_evaluate_args_and_kwargs(model, vi) return model.f(args...; kwargs...) end From 9152951e16f25e3d5b7f18396966436ee3fdcf9d Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Mon, 22 Jun 2026 17:50:27 +0100 Subject: [PATCH 3/4] Remove nested-submodel inference investigation note The explanation lives in the in-code comments in src/submodel.jl and src/model.jl, so the standalone markdown note is redundant. Claude-Session: https://claude.ai/code/session_014Noe4UB3FFurwpt2BNUE5E --- nested-submodel-inference-investigation.md | 132 --------------------- 1 file changed, 132 deletions(-) delete mode 100644 nested-submodel-inference-investigation.md diff --git a/nested-submodel-inference-investigation.md b/nested-submodel-inference-investigation.md deleted file mode 100644 index fa8c27346..000000000 --- a/nested-submodel-inference-investigation.md +++ /dev/null @@ -1,132 +0,0 @@ -# Nested Submodel Inference Note - -This note explains why `_evaluate!!(::Submodel, ...)` calls `model.f` directly instead of -calling `_evaluate!!(model, vi)`. - -## The Issue - -The problem is not that `_evaluate!!(::Submodel, ...)` is recursive. It is recursive, but -Julia's recursion limiter does not widen that call. - -The problem is the extra recursive call through `_evaluate!!(::Model, vi)`: - -```julia -_evaluate!!(::Submodel, vi, parent_context, left_vn) - -> _evaluate!!(::Model, vi) -``` - -At that call site, Julia infers an argument tuple type for `_evaluate!!(model, vi)`. -The second argument type is a concrete `Model{..., Ctx}`. With nested submodels, the -`Ctx` parameter grows at each level: - -```julia -Model{..., PrefixContext{c, InitContext{...}}} -Model{..., PrefixContext{c, PrefixContext{b, InitContext{...}}}} -``` - -Julia sees the same `_evaluate!!(::Model, vi)` method recurring with a more complex call -signature, so it widens the signature to force convergence: - -```julia -Tuple{typeof(_evaluate!!), Model, OnlyAccsVarInfo{...}} -``` - -Once `model` is only known as abstract `Model`, inference can no longer resolve -`model.f` precisely, and the return type collapses. - -## Main vs Branch Call Chain - -The old path was: - -```text -_evaluate!!(::Model, vi) # top-level model - -> model.f(...) - -> tilde_assume!!(..., ::Submodel, ...) - -> _evaluate!!(::Submodel, vi, context, vn) - -> contextualize(submodel.model, eval_context) - -> _evaluate!!(::Model, vi) # problematic recursive call - -> model.f(...) - -> tilde_assume!!(..., ::Submodel, ...) - -> _evaluate!!(::Submodel, ...) -``` - -The fixed path is: - -```text -_evaluate!!(::Model, vi) # top-level model - -> model.f(...) - -> tilde_assume!!(..., ::Submodel, ...) - -> _evaluate!!(::Submodel, vi, context, vn) - -> contextualize(submodel.model, eval_context) - -> make_evaluate_args_and_kwargs(model, vi) - -> model.f(...) # direct call, no _evaluate!!(::Model) - -> tilde_assume!!(..., ::Submodel, ...) - -> _evaluate!!(::Submodel, ...) -``` - -The fix removes the nested `_evaluate!!(::Model, vi)` call edge. It does not remove -submodel recursion. - -## Why The Submodel Recursion Is Fine - -For a submodel tilde statement, the caller already has the concrete context and submodel -in its own signature: - -```julia -tilde_assume!!(context, submodel, vn, NoTemplate, vi) - -> _evaluate!!(submodel, vi, context, vn) -``` - -The callee signature mostly reorders values that were already present in the caller -signature. In Julia's `Compiler/src/typelimits.jl`, `type_more_complex` has an early exit -for this case: - -```julia -elseif is_derived_type_from_any(unwrap_unionall(t), sources, depth) - return false # t isn't something new -end -``` - -That is why a recursive `_evaluate!!(::Submodel, ...)` call can be eligible for limiting -but still not be widened: the growing context is not newly packaged into another dispatch -type. It is a call argument that was already visible in the caller. - -## Why The Model Recursion Is Not Fine - -The old implementation did this inside `_evaluate!!(::Submodel, ...)`: - -```julia -model = contextualize(submodel.model, eval_context) -return _evaluate!!(model, vi) -``` - -That constructs a fresh `Model{..., Ctx}` type and then passes it to the shared -`_evaluate!!(::Model, vi)` method. The full contextualized model type was not already -present in the caller signature; it was synthesized inside the submodel evaluator. - -So Julia compares recursive `_evaluate!!(::Model, vi)` signatures like: - -```julia -Tuple{typeof(_evaluate!!), Model{..., PrefixContext{c, InitContext{...}}}, VI} -Tuple{typeof(_evaluate!!), Model{..., PrefixContext{c, PrefixContext{b, InitContext{...}}}}, VI} -``` - -The second call is more complex because the `Ctx` parameter inside `Model` is more deeply -nested. `limit_type_size` therefore widens the `Model` argument. - -## Why Manual Inlining Works - -The body of `_evaluate!!(::Model, vi)` is: - -```julia -args, kwargs = make_evaluate_args_and_kwargs(model, vi) -return model.f(args...; kwargs...) -``` - -Doing those two lines directly inside `_evaluate!!(::Submodel, ...)` avoids creating the -recursive `_evaluate!!(::Model, vi)` call edge. The contextualized `model` value still -exists, but it is no longer passed through a recursive call to the shared -`_evaluate!!(::Model, vi)` method. - -An `@inline` annotation on `_evaluate!!(::Model, vi)` is not enough. Julia applies this -recursion-limiting heuristic during abstract interpretation, before optimizer inlining. From 9a0dd67b0a20054dfb4f2c33c989429150dcd8f0 Mon Sep 17 00:00:00 2001 From: Xianda Sun Date: Mon, 22 Jun 2026 18:00:08 +0100 Subject: [PATCH 4/4] Condense nested-submodel comments and changelog Keep the full mechanism explanation in src/submodel.jl; reduce the duplicated explanations in the tests and HISTORY to concise pointers to the PR. Remove the redundant NOTE block in src/model.jl. Claude-Session: https://claude.ai/code/session_014Noe4UB3FFurwpt2BNUE5E --- HISTORY.md | 2 +- src/model.jl | 4 ---- src/submodel.jl | 40 +++++++++----------------------------- test/logdensityfunction.jl | 13 +++---------- test/submodels.jl | 19 +++++------------- 5 files changed, 18 insertions(+), 60 deletions(-) diff --git a/HISTORY.md b/HISTORY.md index 1dac30587..7f90d3af7 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -2,7 +2,7 @@ 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. Because each level of submodel nesting adds another context layer to the `Model` type, this tripped Julia's type-inference recursion limit: from the first level of nesting onwards the model's return type was inferred as `Any`, the evaluation fell back to runtime dispatch, and both primal and gradient evaluation slowed down (the primal slowdown was roughly 15x and grew with nesting depth). Submodel evaluation now calls the generated model function directly, which keeps nested submodels type-stable. See [Turing.jl#2844](https://github.com/TuringLang/Turing.jl/issues/2844). +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 diff --git a/src/model.jl b/src/model.jl index 141910a67..6d6a23c28 100644 --- a/src/model.jl +++ b/src/model.jl @@ -906,10 +906,6 @@ function _evaluate!!(model::Model, varinfo::AbstractVarInfo) args, kwargs = make_evaluate_args_and_kwargs(model, varinfo) return model.f(args...; kwargs...) end -# NOTE: The two lines above are intentionally duplicated in the submodel evaluation -# path (`_evaluate!!(::Submodel, ...)` in `submodel.jl`). That path must NOT call this -# method, because doing so reintroduces a type-inference recursion limit for nested -# submodels. See the long comment there before refactoring either copy. is_splat_symbol(s::Symbol) = startswith(string(s), "#splat#") diff --git a/src/submodel.jl b/src/submodel.jl index 5351f63a9..f8acf83c0 100644 --- a/src/submodel.jl +++ b/src/submodel.jl @@ -181,37 +181,15 @@ function _evaluate!!( # (4) Finally, we need to store that context inside the submodel. model = contextualize(submodel.model, eval_context) - # Finally, evaluate the wrapped model. These two lines are a verbatim copy of the body - # of `_evaluate!!(model::Model, ::AbstractVarInfo)`, so calling `model.f` here is - # behaviour-identical to `_evaluate!!(model, vi)`. The duplication is deliberate. - # - # DO NOT replace this with `return _evaluate!!(model, vi)`. Submodel evaluation is - # recursive: `model.f` runs the wrapped model, which may itself contain - # `~ to_submodel(...)` statements that re-enter this method. The recursive - # `_evaluate!!(::Submodel, ...)` edge itself is fine: Julia sees the edge, but the - # limiter keeps its signature unchanged. The problematic edge is the old intermediate - # call to the shared `_evaluate!!(::Model, ::AbstractVarInfo)` method. - # - # Each submodel level constructs a contextualized `Model{...,Ctx}` whose `Ctx` - # parameter nests one `PrefixContext` deeper. If we call `_evaluate!!(model, vi)`, - # inference sees a recursive `_evaluate!!(::Model, ...)` call whose argument tuple - # contains that growing `Model{...,Ctx}` type. Julia's recursion limiter then widens - # the call signature to `Tuple{typeof(_evaluate!!), Model, ...}`. Once `model` is only - # known as abstract `Model`, `model.f` is no longer resolvable and the return type - # collapses to `Any`, causing the slowdown in - # https://github.com/TuringLang/Turing.jl/issues/2844. - # - # Calling `model.f` directly removes only that bad `_evaluate!!(::Model, ...)` edge. - # The remaining recursive submodel edge still exists, but its growing context is a - # normal argument already present in the `tilde_assume!!` caller signature, so the - # limiter does not widen it. - # - # Marking `_evaluate!!(::Model, ...)` `@inline` does NOT help: the recursion limit is - # applied during abstract interpretation, before inlining. - # - # This keeps *distinct* nested submodels type stable to arbitrary depth. A model that - # recurses into the *same* submodel `@model` can still hit the limiter; that is a - # separate, pre-existing limitation. + # 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 diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index 14f6583e8..66fe01ed1 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -445,11 +445,7 @@ end end end - # Nested submodels used to break type inference of the log density, because the - # recursive submodel evaluation passed through the shared `_evaluate!!(::Model, ...)` - # method whose `Model` dispatch type grows by one context layer per level. See - # https://github.com/TuringLang/Turing.jl/issues/2844 and the comment in - # `src/submodel.jl`. + # 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() @@ -485,11 +481,8 @@ end y ~ Normal(params.m, params.s) return 1.0 ~ Normal(y) end - # A submodel nested inside another submodel. Before - # https://github.com/TuringLang/Turing.jl/issues/2844 was fixed, evaluating this - # model was not type-stable and allocated (and got worse with each level of - # nesting), because the recursive submodel evaluation passed through the shared - # `_evaluate!!(::Model, ...)` method. See the comment in `src/submodel.jl`. + # 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) diff --git a/test/submodels.jl b/test/submodels.jl index 82f0471c5..d75f4f99d 100644 --- a/test/submodels.jl +++ b/test/submodels.jl @@ -17,11 +17,10 @@ function get_logp_and_rawval_accs(model::Model) return accs end -# Models for the nested-submodel type-stability tests (issue #2844). Each level wraps the -# previous one in a `to_submodel` and re-exports its return value, so the depth of nesting -# can be varied. They must be defined at module scope: a model defined in local (testset) -# scope captures that scope and is not type-inferrable, which would mask the property under -# test. See the "type stability of nested submodels" testset below. +# 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)) @@ -350,15 +349,7 @@ end end @testset "type stability of nested submodels (issue #2844)" begin - # Evaluating a submodel recurses into `model.f`, which may contain further - # `to_submodel` statements. Each level of nesting grows the `Model`'s context - # type; if that recursion goes through the shared `_evaluate!!(::Model, ...)` - # method, Julia's recursion-limit heuristic widens the model type and the result - # type collapses to `Any` from the first level of nesting onwards. This made - # nested submodels much slower to evaluate (and to differentiate). See - # https://github.com/TuringLang/Turing.jl/issues/2844 and the comment in - # `src/submodel.jl`. These tests check that evaluation stays type stable at every - # level of nesting. + # See https://github.com/TuringLang/DynamicPPL.jl/pull/1427. @testset "$(nameof(model.f))" for model in ( t2844_inner(), t2844_middle(), t2844_outer(), t2844_deeper() )