Use init!! instead of fast_evaluate!!

Author: penelopeysm

docs/src/api.md

@@ -66,11 +66,10 @@ The [LogDensityProblems.jl](https://github.com/tpapp/LogDensityProblems.jl) inte
 LogDensityFunction
 ```
 
-Internally, this is accomplished using:
+Internally, this is accomplished using [`init!!`](@ref) on:
 
 ```@docs
 OnlyAccsVarInfo
-fast_evaluate!!
 ```
 
 ## Condition and decondition
@@ -517,7 +516,7 @@ The function `init!!` is used to initialise, or overwrite, values in a VarInfo.
 It is really a thin wrapper around using `evaluate!!` with an `InitContext`.
 
 ```@docs
-DynamicPPL.init!!
+init!!
 ```
 
 To accomplish this, an initialisation _strategy_ is required, which defines how new values are to be obtained.

src/DynamicPPL.jl

@@ -92,9 +92,11 @@ export AbstractVarInfo,
     getargnames,
     extract_priors,
     values_as_in_model,
+    # evaluation
+    evaluate!!,
+    init!!,
     # LogDensityFunction and fasteval
     LogDensityFunction,
-    fast_evaluate!!,
     OnlyAccsVarInfo,
     # Leaf contexts
     AbstractContext,

src/chains.jl

@@ -137,7 +137,7 @@ end
 """
     ParamsWithStats(
         param_vector::AbstractVector,
-        ldf::DynamicPPL.Experimental.FastLDF,
+        ldf::DynamicPPL.LogDensityFunction,
         stats::NamedTuple=NamedTuple();
         include_colon_eq::Bool=true,
         include_log_probs::Bool=true,
@@ -156,7 +156,7 @@ via `unflatten` plus re-evaluation. It is faster for two reasons:
 """
 function ParamsWithStats(
     param_vector::AbstractVector,
-    ldf::DynamicPPL.Experimental.FastLDF,
+    ldf::DynamicPPL.LogDensityFunction,
     stats::NamedTuple=NamedTuple();
     include_colon_eq::Bool=true,
     include_log_probs::Bool=true,
@@ -174,9 +174,7 @@ function ParamsWithStats(
     else
         (DynamicPPL.ValuesAsInModelAccumulator(include_colon_eq),)
     end
-    _, vi = DynamicPPL.Experimental.fast_evaluate!!(
-        ldf.model, strategy, AccumulatorTuple(accs)
-    )
+    _, vi = DynamicPPL.init!!(ldf.model, OnlyAccsVarInfo(AccumulatorTuple(accs)), strategy)
     params = DynamicPPL.getacc(vi, Val(:ValuesAsInModel)).values
     if include_log_probs
         stats = merge(

src/fasteval.jl

@@ -29,60 +29,6 @@ using LogDensityProblems: LogDensityProblems
 import DifferentiationInterface as DI
 using Random: Random
 
-"""
-    DynamicPPL.fast_evaluate!!(
-        [rng::Random.AbstractRNG,]
-        model::Model,
-        strategy::AbstractInitStrategy,
-        accs::AccumulatorTuple,
-    )
-
-Evaluate a model using parameters obtained via `strategy`, and only computing the results in
-the provided accumulators.
-
-It is assumed that the accumulators passed in have been initialised to appropriate values,
-as this function will not reset them. The default constructors for each accumulator will do
-this for you correctly.
-
-Returns a tuple of the model's return value, plus an `OnlyAccsVarInfo`. Note that the `accs`
-argument may be mutated (depending on how the accumulators are implemented); hence the `!!`
-in the function name.
-"""
-@inline function fast_evaluate!!(
-    # Note that this `@inline` is mandatory for performance. If it's not inlined, it leads
-    # to extra allocations (even for trivial models) and much slower runtime.
-    rng::Random.AbstractRNG,
-    model::Model,
-    strategy::AbstractInitStrategy,
-    accs::AccumulatorTuple,
-)
-    ctx = InitContext(rng, strategy)
-    model = DynamicPPL.setleafcontext(model, ctx)
-    # Calling `evaluate!!` would be fine, but would lead to an extra call to resetaccs!!,
-    # which is unnecessary. So we shortcircuit this by simply calling `_evaluate!!`
-    # directly. To preserve thread-safety we need to reproduce the ThreadSafeVarInfo logic
-    # here.
-    # TODO(penelopeysm): This should _not_ check Threads.nthreads(). I still don't know what
-    # it _should_ do, but this is wrong regardless.
-    # https://github.com/TuringLang/DynamicPPL.jl/issues/1086
-    vi = if Threads.nthreads() > 1
-        param_eltype = DynamicPPL.get_param_eltype(strategy)
-        accs = map(accs) do acc
-            DynamicPPL.convert_eltype(float_type_with_fallback(param_eltype), acc)
-        end
-        ThreadSafeVarInfo(OnlyAccsVarInfo(accs))
-    else
-        OnlyAccsVarInfo(accs)
-    end
-    return DynamicPPL._evaluate!!(model, vi)
-end
-@inline function fast_evaluate!!(
-    model::Model, strategy::AbstractInitStrategy, accs::AccumulatorTuple
-)
-    # This `@inline` is also mandatory for performance
-    return fast_evaluate!!(Random.default_rng(), model, strategy, accs)
-end
-
 """
     DynamicPPL.LogDensityFunction(
         model::Model,
@@ -274,7 +220,7 @@ function (f::LogDensityAt)(params::AbstractVector{<:Real})
         VectorWithRanges(f.iden_varname_ranges, f.varname_ranges, params), nothing
     )
     accs = fast_ldf_accs(f.getlogdensity)
-    _, vi = DynamicPPL.fast_evaluate!!(f.model, strategy, accs)
+    _, vi = DynamicPPL.init!!(f.model, OnlyAccsVarInfo(accs), strategy)
     return f.getlogdensity(vi)
 end
 

src/model.jl

@@ -881,30 +881,56 @@ end
         [init_strategy::AbstractInitStrategy=InitFromPrior()]
     )
 
-Evaluate the `model` and replace the values of the model's random variables
-in the given `varinfo` with new values, using a specified initialisation strategy.
-If the values in `varinfo` are not set, they will be added
-using a specified initialisation strategy.
+Evaluate the `model` and replace the values of the model's random variables in the given
+`varinfo` with new values, using a specified initialisation strategy. If the values in
+`varinfo` are not set, they will be added using a specified initialisation strategy.
 
 If `init_strategy` is not provided, defaults to `InitFromPrior()`.
 
 Returns a tuple of the model's return value, plus the updated `varinfo` object.
 """
-function init!!(
+@inline function init!!(
+    # Note that this `@inline` is mandatory for performance, especially for
+    # LogDensityFunction. If it's not inlined, it leads to extra allocations (even for
+    # trivial models) and much slower runtime.
     rng::Random.AbstractRNG,
     model::Model,
-    varinfo::AbstractVarInfo,
-    init_strategy::AbstractInitStrategy=InitFromPrior(),
+    vi::AbstractVarInfo,
+    strategy::AbstractInitStrategy=InitFromPrior(),
 )
-    new_model = setleafcontext(model, InitContext(rng, init_strategy))
-    return evaluate!!(new_model, varinfo)
+    ctx = InitContext(rng, strategy)
+    model = DynamicPPL.setleafcontext(model, ctx)
+    # TODO(penelopeysm): This should _not_ check Threads.nthreads(). I still don't know what
+    # it _should_ do, but this is wrong regardless.
+    # https://github.com/TuringLang/DynamicPPL.jl/issues/1086
+    return if Threads.nthreads() > 1
+        # TODO(penelopeysm): The logic for setting eltype of accs is very similar to that
+        # used in `unflatten`. The reason why we need it here is because the VarInfo `vi`
+        # won't have been filled with parameters prior to `init!!` being called.
+        #
+        # Note that this eltype promotion is only needed for threadsafe evaluation. In an
+        # ideal world, this code should be handled inside `evaluate_threadsafe!!` or a
+        # similar method. In other words, it should not be here, and it should not be inside
+        # `unflatten` either. The problem is performance. Shifting this code around can have
+        # massive, inexplicable, impacts on performance. This should be investigated
+        # properly.
+        param_eltype = DynamicPPL.get_param_eltype(strategy)
+        accs = map(vi.accs) do acc
+            DynamicPPL.convert_eltype(float_type_with_fallback(param_eltype), acc)
+        end
+        vi = DynamicPPL.setaccs!!(vi, accs)
+        tsvi = ThreadSafeVarInfo(resetaccs!!(vi))
+        retval, tsvi_new = DynamicPPL._evaluate!!(model, tsvi)
+        return retval, setaccs!!(vi, DynamicPPL.getaccs(tsvi_new))
+    else
+        return DynamicPPL._evaluate!!(model, resetaccs!!(vi))
+    end
 end
-function init!!(
-    model::Model,
-    varinfo::AbstractVarInfo,
-    init_strategy::AbstractInitStrategy=InitFromPrior(),
+@inline function init!!(
+    model::Model, vi::AbstractVarInfo, strategy::AbstractInitStrategy=InitFromPrior()
 )
-    return init!!(Random.default_rng(), model, varinfo, init_strategy)
+    # This `@inline` is also mandatory for performance
+    return init!!(Random.default_rng(), model, vi, strategy)
 end
 
 """