FastLDF / InitContext unified (#1132)

* Fast Log Density Function

* Make it work with AD

* Optimise performance for identity VarNames

* Mark `get_range_and_linked` as having zero derivative

* Update comment

* make AD testing / benchmarking use FastLDF

* Fix tests

* Optimise away `make_evaluate_args_and_kwargs`

* const func annotation

* Disable benchmarks on non-typed-Metadata-VarInfo

* Fix `_evaluate!!` correctly to handle submodels

* Actually fix submodel evaluate

* Document thoroughly and organise code

* Support more VarInfos, make it thread-safe (?)

* fix bug in parsing ranges from metadata/VNV

* Fix get_param_eltype for TSVI

* Disable Enzyme benchmark

* Don't override _evaluate!!, that breaks ForwardDiff (sometimes)

* Move FastLDF to experimental for now

* Fix imports, add tests, etc

* More test fixes

* Fix imports / tests

* Remove AbstractFastEvalContext

* Changelog and patch bump

* Add correctness tests, fix imports

* Concretise parameter vector in tests

* Add zero-allocation tests

* Add Chairmarks as test dep

* Disable allocations tests on multi-threaded

* Fast InitContext (#1125)

* Make InitContext work with OnlyAccsVarInfo

* Do not convert NamedTuple to Dict

* remove logging

* Enable InitFromPrior and InitFromUniform too

* Fix `infer_nested_eltype` invocation

* Refactor FastLDF to use InitContext

* note init breaking change

* fix logjac sign

* workaround Mooncake segfault

* fix changelog too

* Fix get_param_eltype for context stacks

* Add a test for threaded observe

* Export init

* Remove dead code

* fix transforms for pathological distributions

* Tidy up loads of things

* fix typed_identity spelling

* fix definition order

* Improve docstrings

* Remove stray comment

* export get_param_eltype (unfortunatley)

* Add more comment

* Update comment

* Remove inlines, fix OAVI docstring

* Improve docstrings

* Simplify InitFromParams constructor

* Replace map(identity, x[:]) with [i for i in x[:]]

* Simplify implementation for InitContext/OAVI

* Add another model to allocation tests

Co-authored-by: Markus Hauru <[email protected]>

* Revert removal of dist argument (oops)

* Format

* Update some outdated bits of FastLDF docstring

* remove underscores

---------

Co-authored-by: Markus Hauru <[email protected]>

Author: penelopeysm

HISTORY.md

@@ -21,6 +21,21 @@ You should not need to use these directly, please use `AbstractPPL.condition` an
 
 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.

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
@@ -517,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,
@@ -113,6 +113,8 @@ export AbstractVarInfo,
     InitFromPrior,
     InitFromUniform,
     InitFromParams,
+    init,
+    get_param_eltype,
     # Pseudo distributions
     NamedDist,
     NoDist,
@@ -193,6 +195,7 @@ 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/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