diff --git a/Project.toml b/Project.toml index a0b9cf0b0..d1d254386 100644 --- a/Project.toml +++ b/Project.toml @@ -12,7 +12,6 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" Compat = "34da2185-b29b-5c13-b0c7-acf172513d20" ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" DocStringExtensions = "ffbed154-4ef7-542d-bbb7-c09d3a79fcae" FillArrays = "1a297f60-69ca-5386-bcde-b61e274b549b" @@ -29,6 +28,8 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" [weakdeps] +DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c" @@ -38,11 +39,13 @@ Mooncake = "da2b9cff-9c12-43a0-ae48-6db2b0edb7d6" ReverseDiff = "37e2e3b7-166d-5795-8a7a-e32c996b4267" [extensions] +DynamicPPLDifferentiationInterfaceExt = ["DifferentiationInterface"] DynamicPPLEnzymeCoreExt = ["EnzymeCore"] +DynamicPPLEnzymeExt = ["Enzyme"] DynamicPPLForwardDiffExt = ["ForwardDiff"] DynamicPPLMCMCChainsExt = ["MCMCChains"] DynamicPPLMarginalLogDensitiesExt = ["MarginalLogDensities"] -DynamicPPLMooncakeExt = ["Mooncake", "DifferentiationInterface"] +DynamicPPLMooncakeExt = ["Mooncake"] DynamicPPLReverseDiffExt = ["ReverseDiff"] [compat] @@ -58,6 +61,7 @@ ConstructionBase = "1.5.4" DifferentiationInterface = "0.6.41, 0.7" Distributions = "0.25" DocStringExtensions = "0.9" +Enzyme = "0.13" EnzymeCore = "0.6 - 0.8" FillArrays = "1.16.0" ForwardDiff = "0.10.12, 1" diff --git a/ext/DynamicPPLDifferentiationInterfaceExt.jl b/ext/DynamicPPLDifferentiationInterfaceExt.jl new file mode 100644 index 000000000..c76fa6f0b --- /dev/null +++ b/ext/DynamicPPLDifferentiationInterfaceExt.jl @@ -0,0 +1,65 @@ +module DynamicPPLDifferentiationInterfaceExt + +import DifferentiationInterface as DI +using DynamicPPL: + DynamicPPL, + AccumulatorTuple, + LogDensityAt, + Model, + VarNamedTuple, + AbstractTransformStrategy, + _use_closure, + logdensity_at +using ADTypes: ADTypes + +# Fallback only: backend-specific extensions with more specific AD types should take precedence. +function DynamicPPL._prepare_gradient( + adtype::ADTypes.AbstractADType, + x::AbstractVector{<:Real}, + model::Model, + getlogdensity::Any, + varname_ranges::VarNamedTuple, + transform_strategy::AbstractTransformStrategy, + accs::AccumulatorTuple, +) + args = (model, getlogdensity, varname_ranges, transform_strategy, accs) + return if _use_closure(adtype) + DI.prepare_gradient(LogDensityAt(args...), adtype, x) + else + DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...) + end +end + +function DynamicPPL._value_and_gradient( + adtype::ADTypes.AbstractADType, + prep, + params::AbstractVector{<:Real}, + model::Model, + getlogdensity::Any, + varname_ranges::VarNamedTuple, + transform_strategy::AbstractTransformStrategy, + accs::AccumulatorTuple, +) + return if _use_closure(adtype) + DI.value_and_gradient( + LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs), + prep, + adtype, + params, + ) + else + DI.value_and_gradient( + logdensity_at, + prep, + adtype, + params, + DI.Constant(model), + DI.Constant(getlogdensity), + DI.Constant(varname_ranges), + DI.Constant(transform_strategy), + DI.Constant(accs), + ) + end +end + +end # module diff --git a/ext/DynamicPPLEnzymeExt.jl b/ext/DynamicPPLEnzymeExt.jl new file mode 100644 index 000000000..591c54e38 --- /dev/null +++ b/ext/DynamicPPLEnzymeExt.jl @@ -0,0 +1,91 @@ +module DynamicPPLEnzymeExt + +using DynamicPPL: ADTypes, DynamicPPL, logdensity_at +using Enzyme: Enzyme + +_enzyme_gradient_mode(::ADTypes.AutoEnzyme{Nothing}) = Enzyme.ReverseWithPrimal +function _enzyme_gradient_mode(adtype::ADTypes.AutoEnzyme) + return Enzyme.EnzymeCore.WithPrimal(adtype.mode) +end + +_cache_enzyme_gradient(::ADTypes.AutoEnzyme{Nothing}) = true +_cache_enzyme_gradient(::ADTypes.AutoEnzyme{<:Enzyme.EnzymeCore.ReverseMode}) = true +_cache_enzyme_gradient(::ADTypes.AutoEnzyme) = false + +function _extract_value_and_gradient(result::NamedTuple{(:derivs, :val)}) + return result.val, first(result.derivs) +end + +function DynamicPPL._prepare_gradient( + adtype::ADTypes.AutoEnzyme, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + return _cache_enzyme_gradient(adtype) ? (; dx=similar(x)) : nothing +end + +function DynamicPPL._value_and_gradient( + adtype::ADTypes.AutoEnzyme, + prep::NamedTuple{(:dx,)}, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + isempty(params) && return logdensity_at( + params, model, getlogdensity, varname_ranges, transform_strategy, accs + ), + copy(params) + dx = prep.dx + fill!(dx, zero(eltype(dx))) + _, val = Enzyme.autodiff( + _enzyme_gradient_mode(adtype), + logdensity_at, + Enzyme.Active, + Enzyme.Duplicated(params, dx), + Enzyme.Const(model), + Enzyme.Const(getlogdensity), + Enzyme.Const(varname_ranges), + Enzyme.Const(transform_strategy), + Enzyme.Const(accs), + ) + return val, copy(dx) +end + +function DynamicPPL._value_and_gradient( + adtype::ADTypes.AutoEnzyme, + ::Nothing, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + isempty(params) && return logdensity_at( + params, model, getlogdensity, varname_ranges, transform_strategy, accs + ), + copy(params) + # Pass the plain function plus Const arguments; Enzyme is brittle with closure-like callables. + val, dx = _extract_value_and_gradient( + Enzyme.gradient( + _enzyme_gradient_mode(adtype), + logdensity_at, + params, + Enzyme.Const(model), + Enzyme.Const(getlogdensity), + Enzyme.Const(varname_ranges), + Enzyme.Const(transform_strategy), + Enzyme.Const(accs), + ), + ) + return val, copy(dx) +end + +end # module diff --git a/ext/DynamicPPLForwardDiffExt.jl b/ext/DynamicPPLForwardDiffExt.jl index 4b6d3fb41..b825b00a6 100644 --- a/ext/DynamicPPLForwardDiffExt.jl +++ b/ext/DynamicPPLForwardDiffExt.jl @@ -2,6 +2,9 @@ module DynamicPPLForwardDiffExt using DynamicPPL: ADTypes, DynamicPPL, LogDensityProblems using ForwardDiff +# DiffResults is a direct dependency of ForwardDiff; access it through ForwardDiff's namespace +# rather than listing it as a separate (weak)dep of DynamicPPL. +const DiffResults = ForwardDiff.DiffResults # check if the AD type already has a tag use_dynamicppl_tag(::ADTypes.AutoForwardDiff{<:Any,Nothing}) = true @@ -12,10 +15,6 @@ function DynamicPPL.tweak_adtype( ) where {chunk_size} # Use DynamicPPL tag to improve stack traces # https://www.stochasticlifestyle.com/improved-forwarddiff-jl-stacktraces-with-package-tags/ - # NOTE: DifferentiationInterface disables tag checking if the - # tag inside the AutoForwardDiff type is not nothing. See - # https://github.com/JuliaDiff/DifferentiationInterface.jl/blob/1df562180bdcc3e91c885aa5f4162a0be2ced850/DifferentiationInterface/ext/DifferentiationInterfaceForwardDiffExt/onearg.jl#L338-L350. - # So we don't currently need to override ForwardDiff.checktag as well. tag = if use_dynamicppl_tag(ad) ForwardDiff.Tag(DynamicPPL.DynamicPPLTag(), eltype(params)) else @@ -32,4 +31,42 @@ function DynamicPPL.tweak_adtype( return ADTypes.AutoForwardDiff(; chunksize=ForwardDiff.chunksize(chunk), tag=tag) end +function DynamicPPL._prepare_gradient( + adtype::ADTypes.AutoForwardDiff{chunk_size}, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) where {chunk_size} + f = DynamicPPL.LogDensityAt( + model, getlogdensity, varname_ranges, transform_strategy, accs + ) + # chunk_size is already set to a concrete positive integer by tweak_adtype + chunk = ForwardDiff.Chunk(length(x), chunk_size) + cfg = ForwardDiff.GradientConfig(f, x, chunk, adtype.tag) + result = DiffResults.GradientResult(similar(x)) + return (; cfg, result) +end + +function DynamicPPL._value_and_gradient( + ::ADTypes.AutoForwardDiff, + prep, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = DynamicPPL.LogDensityAt( + model, getlogdensity, varname_ranges, transform_strategy, accs + ) + # Val{false}() skips tag checking, since our DynamicPPLTag is reused across calls + # with different LogDensityAt instances. + ForwardDiff.gradient!(prep.result, f, params, prep.cfg, Val{false}()) + return DiffResults.value(prep.result), copy(DiffResults.gradient(prep.result)) +end + end # module diff --git a/ext/DynamicPPLMooncakeExt.jl b/ext/DynamicPPLMooncakeExt.jl index 9760d9f4b..03dfd2df7 100644 --- a/ext/DynamicPPLMooncakeExt.jl +++ b/ext/DynamicPPLMooncakeExt.jl @@ -1,7 +1,16 @@ module DynamicPPLMooncakeExt using DynamicPPL: DynamicPPL, is_transformed -using Mooncake: Mooncake +using Mooncake: + Mooncake, + Dual, + NoTangent, + prepare_derivative_cache, + prepare_gradient_cache, + primal, + tangent, + value_and_derivative!!, + value_and_gradient!! # These are purely optimisations (although quite significant ones sometimes, especially for # _get_range_and_linked). @@ -13,11 +22,92 @@ Mooncake.@zero_derivative Mooncake.DefaultCtx Tuple{ typeof(Base.haskey),DynamicPPL.VarInfo,DynamicPPL.VarName } -using DynamicPPL: @model, LinkAll, getlogjoint_internal, LogDensityFunction -using ADTypes: AutoMooncake -import DifferentiationInterface +using DynamicPPL: @model, LinkAll, LogDensityAt, getlogjoint_internal, LogDensityFunction +using ADTypes: AutoMooncake, AutoMooncakeForward +using Accessors: Accessors using Distributions: Normal, InverseGamma, Beta using PrecompileTools: @setup_workload, @compile_workload + +function _cache_config(::Union{AutoMooncake{Nothing},AutoMooncakeForward{Nothing}}) + return Mooncake.Config(; friendly_tangents=false) +end +function _cache_config(adtype::Union{AutoMooncake,AutoMooncakeForward}) + # Use Accessors to set friendly_tangents=false while preserving all other config fields. + return Accessors.@set adtype.config.friendly_tangents = false +end + +# LogDensityAt is a constant w.r.t. differentiation; NoTangent avoids tangent allocation. +Mooncake.tangent_type(::Type{<:DynamicPPL.LogDensityAt}) = NoTangent + +function DynamicPPL._prepare_gradient( + adtype::AutoMooncake, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + return prepare_gradient_cache(f, x; config=_cache_config(adtype)) +end + +function DynamicPPL._prepare_gradient( + adtype::AutoMooncakeForward, + x::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + cache = prepare_derivative_cache(f, x; config=_cache_config(adtype)) + return (; cache, dx=similar(x), grad=similar(x)) +end + +# Inline this hook so the `(value, grad)` result stays as an sret tuple instead of boxing. +@inline function DynamicPPL._value_and_gradient( + ::AutoMooncake, + prep, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + value, (_, grad) = value_and_gradient!!(prep, f, params; args_to_zero=(false, true)) + return value, copy(grad) +end + +@inline function DynamicPPL._value_and_gradient( + ::AutoMooncakeForward, + prep, + params::AbstractVector{<:Real}, + model::DynamicPPL.Model, + getlogdensity::Any, + varname_ranges::DynamicPPL.VarNamedTuple, + transform_strategy::DynamicPPL.AbstractTransformStrategy, + accs::DynamicPPL.AccumulatorTuple, +) + f = LogDensityAt(model, getlogdensity, varname_ranges, transform_strategy, accs) + (; cache, dx, grad) = prep + # Handle empty parameter vector: value_and_derivative!! loop won't execute. + isempty(params) && return f(params), copy(grad) + value = zero(eltype(grad)) + fill!(dx, zero(eltype(dx))) + @inbounds for i in eachindex(grad, dx) + dx[i] = one(eltype(dx)) + result = value_and_derivative!!(cache, Dual(f, NoTangent()), Dual(params, dx)) + value = primal(result) + grad[i] = tangent(result) + dx[i] = zero(eltype(dx)) + end + return value, copy(grad) +end + @setup_workload begin @compile_workload begin for dist in (Normal(), InverseGamma(2, 3), Beta(2, 2)) diff --git a/src/logdensityfunction.jl b/src/logdensityfunction.jl index b1ae82f2b..081ef54bb 100644 --- a/src/logdensityfunction.jl +++ b/src/logdensityfunction.jl @@ -23,7 +23,6 @@ using ADTypes: ADTypes using BangBang: BangBang using AbstractPPL: AbstractPPL, VarName using LogDensityProblems: LogDensityProblems -import DifferentiationInterface as DI using Random: Random """ @@ -178,7 +177,7 @@ struct LogDensityFunction{ L<:AbstractTransformStrategy, F, VNT<:VarNamedTuple, - ADP<:Union{Nothing,DI.GradientPrep}, + ADP, # unconstrained: backends may store any prep object via _prepare_gradient # type of the vector passed to logdensity functions X<:AbstractVector, AC<:AccumulatorTuple, @@ -246,12 +245,9 @@ struct LogDensityFunction{ else # Make backend-specific tweaks to the adtype adtype = DynamicPPL.tweak_adtype(adtype, model, x) - args = (model, getlogdensity, all_ranges, transform_strategy, accs) - if _use_closure(adtype) - DI.prepare_gradient(LogDensityAt(args...), adtype, x) - else - DI.prepare_gradient(logdensity_at, adtype, x, map(DI.Constant, args)...) - end + DynamicPPL._prepare_gradient( + adtype, x, model, getlogdensity, all_ranges, transform_strategy, accs + ) end return new{ typeof(model), @@ -407,6 +403,22 @@ function (f::LogDensityAt)(params::AbstractVector{<:Real}) ) end +# Extensible hooks: backends provide methods via package extensions. +function _prepare_gradient end +function _value_and_gradient end + +function _value_and_gradient(adtype::ADTypes.AbstractADType, args...) + throw( + ArgumentError( + "No gradient implementation found for AD backend $adtype. " * + "If you intended to use the default (ForwardDiff), ensure that ForwardDiff is " * + "loaded (e.g. `using ForwardDiff`). For other backends, load the corresponding " * + "package (e.g. `using Mooncake`, `using Enzyme`) or load " * + "DifferentiationInterface as a fallback.", + ), + ) +end + function LogDensityProblems.logdensity( ldf::LogDensityFunction, params::AbstractVector{<:Real} ) @@ -426,32 +438,16 @@ function LogDensityProblems.logdensity_and_gradient( # `params` has to be converted to the same vector type that was used for AD preparation, # otherwise the preparation will not be valid. params = convert(get_input_vector_type(ldf), params) - return if _use_closure(ldf.adtype) - DI.value_and_gradient( - LogDensityAt( - ldf.model, - ldf._getlogdensity, - ldf._varname_ranges, - ldf.transform_strategy, - ldf._accs, - ), - ldf._adprep, - ldf.adtype, - params, - ) - else - DI.value_and_gradient( - logdensity_at, - ldf._adprep, - ldf.adtype, - params, - DI.Constant(ldf.model), - DI.Constant(ldf._getlogdensity), - DI.Constant(ldf._varname_ranges), - DI.Constant(ldf.transform_strategy), - DI.Constant(ldf._accs), - ) - end + return DynamicPPL._value_and_gradient( + ldf.adtype, + ldf._adprep, + params, + ldf.model, + ldf._getlogdensity, + ldf._varname_ranges, + ldf.transform_strategy, + ldf._accs, + ) end function LogDensityProblems.capabilities(::Type{<:LogDensityFunction{M,Nothing}}) where {M} diff --git a/src/test_utils/ad.jl b/src/test_utils/ad.jl index 8c9f96491..820a5093c 100644 --- a/src/test_utils/ad.jl +++ b/src/test_utils/ad.jl @@ -2,7 +2,6 @@ module AD using ADTypes: AbstractADType, AutoForwardDiff using Chairmarks: @be -import DifferentiationInterface as DI using DocStringExtensions using DynamicPPL: DynamicPPL, diff --git a/test/Project.toml b/test/Project.toml index 94c624616..6d0028c6d 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -9,11 +9,11 @@ Bijectors = "76274a88-744f-5084-9051-94815aaf08c4" Chairmarks = "0ca39b1e-fe0b-4e98-acfc-b1656634c4de" Combinatorics = "861a8166-3701-5b0c-9a16-15d98fcdc6aa" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" -DifferentiationInterface = "a0c0ee7d-e4b9-4e03-894e-1c5f64a51d63" DimensionalData = "0703355e-b756-11e9-17c0-8b28908087d0" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" Distributions = "31c24e10-a181-5473-b8eb-7969acd0382f" Documenter = "e30172f5-a6a5-5a46-863b-614d45cd2de4" +Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9" ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210" InvertedIndices = "41ab1584-1d38-5bbf-9106-f11c6c58b48f" LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e" @@ -43,19 +43,18 @@ BangBang = "0.4" Bijectors = "0.15.17" Chairmarks = "1" Combinatorics = "1" -DifferentiationInterface = "0.6.41, 0.7" DimensionalData = "0.30" Distributions = "0.25" Documenter = "1" +Enzyme = "0.13" ForwardDiff = "0.10.12, 1" InvertedIndices = "1" LogDensityProblems = "2" MCMCChains = "7.2.1" MacroTools = "0.5.6" -MarginalLogDensities = "0.4" -Mooncake = "0.4, 0.5" -OrderedCollections = "1" +MarginalLogDensities = "0.4.3" OffsetArrays = "1" +OrderedCollections = "1" ReverseDiff = "1" SpecialFunctions = "2.6.1" StableRNGs = "1" diff --git a/test/logdensityfunction.jl b/test/logdensityfunction.jl index b7f80a07c..d2b50aea9 100644 --- a/test/logdensityfunction.jl +++ b/test/logdensityfunction.jl @@ -13,9 +13,10 @@ using LogDensityProblems: LogDensityProblems using Random: Xoshiro using StableRNGs: StableRNG +using Enzyme: Enzyme using ForwardDiff: ForwardDiff -using ReverseDiff: ReverseDiff using Mooncake: Mooncake +using ReverseDiff: ReverseDiff @testset "LogDensityFunction: constructors" begin dist = Beta(2, 2) @@ -445,6 +446,7 @@ end AutoReverseDiff(; compile=false), AutoReverseDiff(; compile=true), AutoMooncake(; config=nothing), + AutoMooncakeForward(; config=nothing), ] @testset "Correctness" begin