Atomic accumulators

HISTORY.md

@@ -1,5 +1,41 @@
 # DynamicPPL Changelog
 
+## 0.39.0
+
+### Breaking changes
+
+#### Parent and leaf contexts
+
+The `DynamicPPL.NodeTrait` function has been removed.
+Instead of implementing this, parent contexts should subtype `DynamicPPL.AbstractParentContext`.
+This is an abstract type which requires you to overload two functions, `DynamicPPL.childcontext` and `DynamicPPL.setchildcontext`.
+
+There should generally be few reasons to define your own parent contexts (the only one we are aware of, outside of DynamicPPL itself, is `Turing.Inference.GibbsContext`), so this change should not really affect users.
+
+Leaf contexts require no changes, apart from a removal of the `NodeTrait` function.
+
+`ConditionContext` and `PrefixContext` are no longer exported.
+You should not need to use these directly, please use `AbstractPPL.condition` and `DynamicPPL.prefix` instead.
+
+#### Miscellaneous
+
+Removed the method `returned(::Model, values, keys)`; please use `returned(::Model, ::AbstractDict{<:VarName})` instead.
+
+The method `DynamicPPL.init` (for implementing `AbstractInitStrategy`) now has a different signature: it must return a tuple of the generated value, plus a transform function that maps it back to unlinked space.
+This is a generalisation of the previous behaviour, where `init` would always return an unlinked value (in effect forcing the transform to be the identity function).
+
+### Other changes
+
+#### FastLDF
+
+Added `DynamicPPL.Experimental.FastLDF`, a version of `LogDensityFunction` that provides performance improvements on the order of 2–10× for both model evaluation as well as automatic differentiation.
+Exact speedups depend on the model size: larger models have less significant speedups because the bulk of the work is done in calls to `logpdf`.
+
+Please note that `FastLDF` is currently considered internal and its API may change without warning.
+We intend to replace `LogDensityFunction` with `FastLDF` in a release in the near future, but until then we recommend not using it.
+
+For more information about `FastLDF`, please see https://github.com/TuringLang/DynamicPPL.jl/pull/1113 as well as the `src/fasteval.jl` file, which contains extensive comments.
+
 ## 0.38.9
 
 Remove warning when using Enzyme as the AD backend.

Project.toml

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

benchmarks/Project.toml

@@ -23,7 +23,7 @@ DynamicPPL = {path = "../"}
 ADTypes = "1.14.0"
 BenchmarkTools = "1.6.0"
 Distributions = "0.25.117"
-DynamicPPL = "0.38"
+DynamicPPL = "0.39"
 Enzyme = "0.13"
 ForwardDiff = "0.10.38, 1"
 LogDensityProblems = "2.1.2"

docs/Project.toml

@@ -21,7 +21,7 @@ Accessors = "0.1"
 Distributions = "0.25"
 Documenter = "1"
 DocumenterMermaid = "0.1, 0.2"
-DynamicPPL = "0.38"
+DynamicPPL = "0.39"
 FillArrays = "0.13, 1"
 ForwardDiff = "0.10, 1"
 JET = "0.9, 0.10, 0.11"

docs/src/api.md

@@ -170,6 +170,12 @@ DynamicPPL.prefix
 
 ## Utilities
 
+`typed_identity` is the same as `identity`, but with an overload for `with_logabsdet_jacobian` that ensures that it never errors.
+
+```@docs
+typed_identity
+```
+
 It is possible to manually increase (or decrease) the accumulated log likelihood or prior from within a model function.
 
 ```@docs
@@ -352,13 +358,6 @@ Base.empty!
 SimpleVarInfo
 ```
 
-### Tilde-pipeline
-
-```@docs
-tilde_assume!!
-tilde_observe!!
-```
-
 ### Accumulators
 
 The subtypes of [`AbstractVarInfo`](@ref) store the cumulative log prior and log likelihood, and sometimes other variables that change during executing, in what are called accumulators.
@@ -463,15 +462,48 @@ By default, it does not perform any actual sampling: it only evaluates the model
 If you wish to sample new values, see the section on [VarInfo initialisation](#VarInfo-initialisation) just below this.
 
 The behaviour of a model execution can be changed with evaluation contexts, which are a field of the model.
-Contexts are subtypes of `AbstractPPL.AbstractContext`.
+
+All contexts are subtypes of `AbstractPPL.AbstractContext`.
+
+Contexts are split into two kinds:
+
+**Leaf contexts**: These are the most important contexts as they ultimately decide how model evaluation proceeds.
+For example, `DefaultContext` evaluates the model using values stored inside a VarInfo's metadata, whereas `InitContext` obtains new values either by sampling or from a known set of parameters.
+DynamicPPL has more leaf contexts which are used for internal purposes, but these are the two that are exported.
 
 ```@docs
 DefaultContext
-PrefixContext
-ConditionContext
 InitContext
 ```
 
+To implement a leaf context, you need to subtype `AbstractPPL.AbstractContext` and implement the `tilde_assume!!` and `tilde_observe!!` methods for your context.
+
+```@docs
+tilde_assume!!
+tilde_observe!!
+```
+
+**Parent contexts**: These essentially act as 'modifiers' for leaf contexts.
+For example, `PrefixContext` adds a prefix to all variable names during evaluation, while `ConditionContext` marks certain variables as observed.
+
+To implement a parent context, you have to subtype `DynamicPPL.AbstractParentContext`, and implement the `childcontext` and `setchildcontext` methods.
+If needed, you can also implement `tilde_assume!!` and `tilde_observe!!` for your context.
+This is optional; the default implementation is to simply delegate to the child context.
+
+```@docs
+AbstractParentContext
+childcontext
+setchildcontext
+```
+
+Since contexts form a tree structure, these functions are automatically defined for manipulating context stacks.
+They are mainly useful for modifying the fundamental behaviour (i.e. the leaf context), without affecting any of the modifiers (i.e. parent contexts).
+
+```@docs
+leafcontext
+setleafcontext
+```
+
 ### VarInfo initialisation
 
 The function `init!!` is used to initialise, or overwrite, values in a VarInfo.
@@ -491,10 +523,12 @@ InitFromParams
 ```
 
 If you wish to write your own, you have to subtype [`DynamicPPL.AbstractInitStrategy`](@ref) and implement the `init` method.
+In very rare situations, you may also need to implement `get_param_eltype`, which defines the element type of the parameters generated by the strategy.
 
 ```@docs
-DynamicPPL.AbstractInitStrategy
-DynamicPPL.init
+AbstractInitStrategy
+init
+get_param_eltype
 ```
 
 ### Choosing a suitable VarInfo

ext/DynamicPPLEnzymeCoreExt.jl

@@ -1,16 +1,15 @@
 module DynamicPPLEnzymeCoreExt
 
-if isdefined(Base, :get_extension)
-    using DynamicPPL: DynamicPPL
-    using EnzymeCore
-else
-    using ..DynamicPPL: DynamicPPL
-    using ..EnzymeCore
-end
+using DynamicPPL: DynamicPPL
+using EnzymeCore
 
 # Mark is_transformed as having 0 derivative. The `nothing` return value is not significant, Enzyme
 # only checks whether such a method exists, and never runs it.
 @inline EnzymeCore.EnzymeRules.inactive(::typeof(DynamicPPL.is_transformed), args...) =
     nothing
+# Likewise for get_range_and_linked.
+@inline EnzymeCore.EnzymeRules.inactive(
+    ::typeof(DynamicPPL._get_range_and_linked), args...
+) = nothing
 
 end

ext/DynamicPPLMooncakeExt.jl

@@ -5,5 +5,8 @@ using Mooncake: Mooncake
 
 # This is purely an optimisation.
 Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{typeof(is_transformed),Vararg}
+Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{
+    typeof(DynamicPPL._get_range_and_linked),Vararg
+}
 
 end # module

src/DynamicPPL.jl

@@ -84,8 +84,8 @@ export AbstractVarInfo,
     # Compiler
     @model,
     # Utilities
-    init,
     OrderedDict,
+    typed_identity,
     # Model
     Model,
     getmissings,
@@ -94,20 +94,27 @@ export AbstractVarInfo,
     values_as_in_model,
     # LogDensityFunction
     LogDensityFunction,
-    # Contexts
+    # Leaf contexts
+    AbstractContext,
     contextualize,
     DefaultContext,
-    PrefixContext,
-    ConditionContext,
+    InitContext,
+    # Parent contexts
+    AbstractParentContext,
+    childcontext,
+    setchildcontext,
+    leafcontext,
+    setleafcontext,
     # Tilde pipeline
     tilde_assume!!,
     tilde_observe!!,
     # Initialisation
-    InitContext,
     AbstractInitStrategy,
     InitFromPrior,
     InitFromUniform,
     InitFromParams,
+    init,
+    get_param_eltype,
     # Pseudo distributions
     NamedDist,
     NoDist,
@@ -183,11 +190,13 @@ include("distribution_wrappers.jl")
 include("submodel.jl")
 include("varnamedvector.jl")
 include("accumulators.jl")
-include("default_accumulators.jl")
+include("accumulators/logp.jl")
+include("accumulators/atomic.jl")
 include("abstract_varinfo.jl")
 include("threadsafe.jl")
 include("varinfo.jl")
 include("simple_varinfo.jl")
+include("onlyaccs.jl")
 include("compiler.jl")
 include("pointwise_logdensities.jl")
 include("logdensityfunction.jl")

src/abstract_varinfo.jl

@@ -204,7 +204,7 @@ Jacobian here is taken with respect to the forward (link) transform.
 
 See also: [`setlogjac!!`](@ref).
 """
-getlogjac(vi::AbstractVarInfo) = getacc(vi, Val(:LogJacobian)).logjac
+getlogjac(vi::AbstractVarInfo) = getacc(vi, Val(:LogJacobian)).logp
 
 """
     getloglikelihood(vi::AbstractVarInfo)

src/accumulators/atomic.jl

@@ -0,0 +1,42 @@
+abstract type AtomicLogProbAccumulator{T} <: LogProbAccumulator{T} end
+
+function acclogp(acc::AtomicLogProbAccumulator, val)
+    @atomic :monotonic acc.logp + val
+    return acc
+end
+logp(acc::AtomicLogProbAccumulator) = acc.logp
+
+mutable struct AtomicLogPriorAccumulator{T<:Real} <: AtomicLogProbAccumulator{T}
+    @atomic logp::T
+end
+accumulator_name(::Type{<:AtomicLogPriorAccumulator}) = :LogPrior
+function accumulate_assume!!(acc::AtomicLogPriorAccumulator, val, logjac, vn, right)
+    return acclogp(acc, logpdf(right, val))
+end
+accumulate_observe!!(acc::AtomicLogPriorAccumulator, right, left, vn) = acc
+
+mutable struct AtomicLogJacobianAccumulator{T<:Real} <: AtomicLogProbAccumulator{T}
+    @atomic logp::T
+end
+accumulator_name(::Type{<:AtomicLogJacobianAccumulator}) = :LogJacobian
+function accumulate_assume!!(acc::AtomicLogJacobianAccumulator, val, logjac, vn, right)
+    return acclogp(acc, logjac)
+end
+accumulate_observe!!(acc::AtomicLogJacobianAccumulator, right, left, vn) = acc
+
+mutable struct AtomicLogLikelihoodAccumulator{T<:Real} <: AtomicLogProbAccumulator{T}
+    @atomic logp::T
+end
+accumulator_name(::Type{<:AtomicLogLikelihoodAccumulator}) = :LogLikelihood
+accumulate_assume!!(acc::AtomicLogLikelihoodAccumulator, val, logjac, vn, right) = acc
+function accumulate_observe!!(acc::AtomicLogLikelihoodAccumulator, right, left, vn)
+    return acclogp(acc, Distributions.loglikelihood(right, left))
+end
+
+function default_atomic_accumulators(::Type{FloatT}=LogProbType) where {FloatT}
+    return AccumulatorTuple(
+        AtomicLogPriorAccumulator{FloatT}(),
+        AtomicLogJacobianAccumulator{FloatT}(),
+        AtomicLogLikelihoodAccumulator{FloatT}(),
+    )
+end

src/default_accumulators.jl → src/accumulators/logp.jl

@@ -128,10 +128,10 @@ $(TYPEDFIELDS)
 """
 struct LogJacobianAccumulator{T<:Real} <: LogProbAccumulator{T}
     "the logabsdet of the link transform Jacobian"
-    logjac::T
+    logp::T
 end
 
-logp(acc::LogJacobianAccumulator) = acc.logjac
+logp(acc::LogJacobianAccumulator) = acc.logp
 
 accumulator_name(::Type{<:LogJacobianAccumulator}) = :LogJacobian
 

src/compiler.jl

@@ -718,14 +718,15 @@ end
 # TODO(mhauru) matchingvalue has methods that can accept both types and values. Why?
 # TODO(mhauru) This function needs a more comprehensive docstring.
 """
-    matchingvalue(vi, value)
+    matchingvalue(param_eltype, value)
 
-Convert the `value` to the correct type for the `vi` object.
+Convert the `value` to the correct type, given the element type of the parameters
+being used to evaluate the model.
 """
-function matchingvalue(vi, value)
+function matchingvalue(param_eltype, value)
     T = typeof(value)
     if hasmissing(T)
-        _value = convert(get_matching_type(vi, T), value)
+        _value = convert(get_matching_type(param_eltype, T), value)
         # TODO(mhauru) Why do we make a deepcopy, even though in the !hasmissing branch we
         # are happy to return `value` as-is?
         if _value === value
@@ -738,29 +739,30 @@ function matchingvalue(vi, value)
     end
 end
 
-function matchingvalue(vi, value::FloatOrArrayType)
-    return get_matching_type(vi, value)
+function matchingvalue(param_eltype, value::FloatOrArrayType)
+    return get_matching_type(param_eltype, value)
 end
-function matchingvalue(vi, ::TypeWrap{T}) where {T}
-    return TypeWrap{get_matching_type(vi, T)}()
+function matchingvalue(param_eltype, ::TypeWrap{T}) where {T}
+    return TypeWrap{get_matching_type(param_eltype, T)}()
 end
 
 # TODO(mhauru) This function needs a more comprehensive docstring. What is it for?
 """
-    get_matching_type(vi, ::TypeWrap{T}) where {T}
+    get_matching_type(param_eltype, ::TypeWrap{T}) where {T}
 
-Get the specialized version of type `T` for `vi`.
+Get the specialized version of type `T`, given an element type of the parameters
+being used to evaluate the model.
 """
 get_matching_type(_, ::Type{T}) where {T} = T
-function get_matching_type(vi, ::Type{<:Union{Missing,AbstractFloat}})
-    return Union{Missing,float_type_with_fallback(eltype(vi))}
+function get_matching_type(param_eltype, ::Type{<:Union{Missing,AbstractFloat}})
+    return Union{Missing,float_type_with_fallback(param_eltype)}
 end
-function get_matching_type(vi, ::Type{<:AbstractFloat})
-    return float_type_with_fallback(eltype(vi))
+function get_matching_type(param_eltype, ::Type{<:AbstractFloat})
+    return float_type_with_fallback(param_eltype)
 end
-function get_matching_type(vi, ::Type{<:Array{T,N}}) where {T,N}
-    return Array{get_matching_type(vi, T),N}
+function get_matching_type(param_eltype, ::Type{<:Array{T,N}}) where {T,N}
+    return Array{get_matching_type(param_eltype, T),N}
 end
-function get_matching_type(vi, ::Type{<:Array{T}}) where {T}
-    return Array{get_matching_type(vi, T)}
+function get_matching_type(param_eltype, ::Type{<:Array{T}}) where {T}
+    return Array{get_matching_type(param_eltype, T)}
 end