diff --git a/HISTORY.md b/HISTORY.md index 5978c4337..621c0856e 100644 --- a/HISTORY.md +++ b/HISTORY.md @@ -1,3 +1,7 @@ +# 0.42.3 + +Worked around an upstream Julia code-generation bug that, under concurrent allocation load, could crash the model re-evaluation performed by `ParamsWithStats`. The affected re-evaluation is now routed through an internal function barrier that avoids the miscompilation; behaviour is unchanged. The crash was originally reported in [AbstractMCMC.jl#214](https://github.com/TuringLang/AbstractMCMC.jl/issues/214). + # 0.42.2 Widened the `Mooncake` compat bound to `0.4.147, 0.5, 0.6`. Mooncake 0.6 is a breaking release (forward-mode redesign), but the reverse-mode rule API `DynamicPPLMooncakeExt` uses (`Mooncake.@zero_derivative`) and the prepared-cache API it relies on via AbstractPPL are unchanged, so no code changes were needed. diff --git a/Project.toml b/Project.toml index 5652733e2..569a16fc9 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "DynamicPPL" uuid = "366bfd00-2699-11ea-058f-f148b4cae6d8" -version = "0.42.2" +version = "0.42.3" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" diff --git a/src/chains.jl b/src/chains.jl index 0ad70fd70..219ccb444 100644 --- a/src/chains.jl +++ b/src/chains.jl @@ -39,29 +39,21 @@ function ParamsWithStats( include_colon_eq::Bool=true, include_log_probs::Bool=true, ) - accs = if include_log_probs - ( + # Route through the same `@noinline _pws_eval` barrier as the vector-based method below, + # so the accumulator tuple -- and hence the `(retval, varinfo)` result of `init!!` -- is + # concretely typed at the call site rather than `Union`-typed (see `_pws_eval` for why + # this matters on Julia 1.12 `-O2`). + return if include_log_probs + accs = ( DynamicPPL.LogPriorAccumulator(), DynamicPPL.LogLikelihoodAccumulator(), DynamicPPL.RawValueAccumulator(include_colon_eq), ) + _pws_eval(model, accs, init_strategy, stats, true) else - (DynamicPPL.RawValueAccumulator(include_colon_eq),) + accs = (DynamicPPL.RawValueAccumulator(include_colon_eq),) + _pws_eval(model, accs, init_strategy, stats, false) end - oavi = OnlyAccsVarInfo(accs) - oavi = last(DynamicPPL.init!!(model, oavi, init_strategy, UnlinkAll())) - params = densify!!(get_raw_values(oavi)) - if include_log_probs - stats = merge( - stats, - ( - logprior=DynamicPPL.getlogprior(oavi), - loglikelihood=DynamicPPL.getloglikelihood(oavi), - logjoint=DynamicPPL.getlogjoint(oavi), - ), - ) - end - return ParamsWithStats(params, stats) end """ @@ -143,20 +135,32 @@ function pws_with_eval( include_log_probs::Bool=true, ) strategy = InitFromVector(param_vector, ldf) - accs = if include_log_probs - ( + # `_pws_eval` is called separately from each branch so that `accs` -- and hence the + # `(retval, varinfo)` tuple returned by `init!!` inside it -- has a concrete type at each + # call site. If `accs` were `Union`-typed here, Julia 1.12's `-O2` optimizer would + # heap-box that union-typed tuple and, because `retval` is unused, leave its pointer + # fields uninitialized across a GC safepoint while the box is GC-rooted; a garbage + # collection landing there scans those fields and segfaults in `gc_mark_obj8`. + return if include_log_probs + accs = ( DynamicPPL.LogPriorAccumulator(), DynamicPPL.LogLikelihoodAccumulator(), DynamicPPL.RawValueAccumulator(include_colon_eq), ) + _pws_eval(ldf.model, accs, strategy, stats, true) else - (DynamicPPL.RawValueAccumulator(include_colon_eq),) + accs = (DynamicPPL.RawValueAccumulator(include_colon_eq),) + _pws_eval(ldf.model, accs, strategy, stats, false) end +end +@noinline function _pws_eval( + model::Model, accs::Tuple, strategy, stats::NamedTuple, include_log_probs::Bool +) # UnlinkAll() actually doesn't have any impact here, because there isn't even a # LogJacobianAccumulator; consequently, it doesn't matter whether we interpret the # parameters as being in linked space or not. However, we just include it for clarity. _, vi = DynamicPPL.init!!( - ldf.model, OnlyAccsVarInfo(AccumulatorTuple(accs)), strategy, UnlinkAll() + model, OnlyAccsVarInfo(AccumulatorTuple(accs)), strategy, UnlinkAll() ) params = densify!!(get_raw_values(vi)) if include_log_probs diff --git a/test/chains.jl b/test/chains.jl index 1dc81693a..66362f0e1 100644 --- a/test/chains.jl +++ b/test/chains.jl @@ -168,6 +168,47 @@ end end end +@testset "no Union-typed VarInfo across init!! (Julia 1.12 -O2 GC miscompile guard)" begin + # Regression guard for a Julia 1.12 `-O2` codegen/GC-rooting bug. Both `ParamsWithStats` + # model-reevaluation methods choose an accumulator tuple based on `include_log_probs`. If + # that choice is made *before* `init!!`, the resulting `(retval, varinfo)` tuple is + # `Union`-typed; 1.12 heap-boxes it, leaves the dead `retval` pointer fields uninitialized + # across a GC safepoint, and a GC landing there corrupts the heap (a `gc_mark_obj8` + # segfault, only reproducible under concurrent allocation load). Routing `init!!` through + # the `@noinline _pws_eval` barrier keeps the accumulator `VarInfo` concretely typed at + # every call site. The union is a type-inference fact independent of Julia version (only + # the crash is 1.12-specific), so assert it never reappears in the inferred IR. + model = DynamicPPL.TestUtils.DEMO_MODELS[1] + ldf = LogDensityFunction(model) + param_vector = rand(ldf) + # Compile both entry points (vector + strategy) so their method bodies are in the IR cache. + ParamsWithStats(param_vector, ldf) + ParamsWithStats(DynamicPPL.InitFromPrior(), model) + + is_union_of_oavi(t) = + t isa Union && all(u -> u <: DynamicPPL.OnlyAccsVarInfo, Base.uniontypes(t)) + function offending_methods(f) + hits = String[] + for m in methods(f) + body = try + Base.bodyfunction(m) + catch + nothing + end + body === nothing && continue + for (ci, _) in code_typed(body; optimize=true) + any(is_union_of_oavi, ci.ssavaluetypes) && push!(hits, string(m)) + end + end + return hits + end + offending = [ + offending_methods(DynamicPPL.pws_with_eval) + offending_methods(ParamsWithStats) + ] + @test isempty(offending) +end + @info "Completed $(@__FILE__) in $(now() - __now__)." end # module