Work around a Julia 1.12 codegen bug that crashes ParamsWithStats

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.

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"

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

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