diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 1e44154..8aac4a9 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -189,23 +189,34 @@ statically — `_make_hvp_fn(_hvp_strategy(backend), ...)` resolves to one concrete method (and one concrete return type) at compile time, without relying on constant propagation through `===`. -Mooncake/Zygote: We route them through -`ReverseOnGrad` as the default. both CPU and GPU run this path fine (see the -Mooncake GPU test). `AutoEnzyme` stays on `ForwardOnGrad`: its reverse mode -hits the gc-transition abort on GPU (see `ext/EnzymeExt.jl`), whereas -`Enzyme.Forward` is robust on CuArrays once the matmul is wrapped. - -Note: forward mode is also supported for these backends via their forward counterparts. +The choice derives from DI's `pushforward_performance` trait +mirroring how DI's `hvp_mode` picks its composition: +backends with a fast pushforward (forward-capable or mode-agnostic, e.g. +AutoForwardDiff, AutoMooncakeForward, plain AutoEnzyme) take `ForwardOnGrad`; +reverse-only backends (AutoMooncake, AutoZygote, AutoReverseDiff, AutoTracker, +AutoEnzyme pinned to Reverse) take `ReverseOnGrad`. Both paths are exercised +on CPU and GPU (see the Mooncake/Zygote GPU tests). Plain +`AutoEnzyme()` reports `ForwardOrReverseMode` and thus lands on +`ForwardOnGrad`. This is fine, since Enzyme's reverse mode hits the +gc-transition abort on GPU (see `ext/EnzymeExt.jl`) whereas `Enzyme.Forward` is +robust on CuArrays once the matmul is wrapped. + +For a `DI.SecondOrder` backend the outer backend differentiates `gradlogp`, +so the strategy comes from the outer backend's trait. =# abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end struct ReverseOnGrad <: HVPStrategy end -_hvp_strategy(::AbstractADType) = ForwardOnGrad() -_hvp_strategy(::ADTypes.AutoMooncake) = ReverseOnGrad() -_hvp_strategy(::ADTypes.AutoZygote) = ReverseOnGrad() -_hvp_strategy(::ADTypes.AutoReverseDiff) = ReverseOnGrad() -_hvp_strategy(::ADTypes.AutoTracker) = ReverseOnGrad() +_strategy_from(::DI.PushforwardFast) = ForwardOnGrad() +_strategy_from(::DI.PushforwardSlow) = ReverseOnGrad() + +function _hvp_strategy(backend::AbstractADType) + return _strategy_from(DI.pushforward_performance(backend)) +end +function _hvp_strategy(backend::DI.SecondOrder) + return _strategy_from(DI.pushforward_performance(DI.outer(backend))) +end #= Hooks for backend-specific normalization of the user's `backend`. diff --git a/src/interface.jl b/src/interface.jl index 41d1de7..d7d7451 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -372,14 +372,16 @@ function _build_mala_deer_rec( #= Use a model-provided HVP when available. Otherwise compute Hv via AD, - picking the path from the user's `backend` (see `DEER._hvp_strategy`): - - ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Used for forward- - capable backends (AutoEnzyme, AutoForwardDiff, AutoMooncakeForward). - ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Default - routing for AutoMooncake / AutoZygote / AutoReverseDiff - / AutoTracker (which also have forward counterparts, - see `DEER._hvp_strategy` for the nuance). + picking the path from the user's `backend` via DI's + `pushforward_performance` trait (see `DEER._hvp_strategy`): + + ForwardOnGrad() — `pushforward(gradlogp, x, v)`. Used for backends with + a fast pushforward (AutoEnzyme, AutoForwardDiff, + AutoMooncakeForward). + ReverseOnGrad() — `gradient(x -> pmcmc_dot(gradlogp(x), v))`. Used for + reverse-only backends (AutoMooncake / AutoZygote / + AutoReverseDiff), which also have forward counterparts — see + `DEER._hvp_strategy` for the nuance. Either path can be bypassed by providing an analytical `hvp` / `hvp_batch` on the `DensityModel`. diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl new file mode 100644 index 0000000..8a7cf94 --- /dev/null +++ b/test/test-HVP-Strategy.jl @@ -0,0 +1,46 @@ +using Test +using ADTypes +using Enzyme: Enzyme +using ParallelMCMC + +const DEER_STRAT = ParallelMCMC.DEER +const DI_STRAT = ParallelMCMC.DEER.DI + +#= +The AD-HVP fallback strategy is derived from DI's `pushforward_performance` +trait (see #38): backends with a fast pushforward take ForwardOnGrad, and +reverse-only backends take ReverseOnGrad — no per-backend enumeration. +=# +@testset "HVP strategy from DI mode traits" begin + @testset "forward-capable backends → ForwardOnGrad" begin + @test DEER_STRAT._hvp_strategy(AutoForwardDiff()) isa DEER_STRAT.ForwardOnGrad + @test DEER_STRAT._hvp_strategy(AutoEnzyme()) isa DEER_STRAT.ForwardOnGrad + @test DEER_STRAT._hvp_strategy(AutoEnzyme(; mode=Enzyme.Forward)) isa + DEER_STRAT.ForwardOnGrad + @test DEER_STRAT._hvp_strategy(AutoMooncakeForward()) isa DEER_STRAT.ForwardOnGrad + end + + @testset "reverse-only backends → ReverseOnGrad" begin + @test DEER_STRAT._hvp_strategy(AutoMooncake()) isa DEER_STRAT.ReverseOnGrad + @test DEER_STRAT._hvp_strategy(AutoZygote()) isa DEER_STRAT.ReverseOnGrad + @test DEER_STRAT._hvp_strategy(AutoReverseDiff()) isa DEER_STRAT.ReverseOnGrad + @test DEER_STRAT._hvp_strategy(AutoTracker()) isa DEER_STRAT.ReverseOnGrad + @test DEER_STRAT._hvp_strategy(AutoEnzyme(; mode=Enzyme.Reverse)) isa + DEER_STRAT.ReverseOnGrad + end + + @testset "SecondOrder picks the strategy from the outer backend" begin + so_fwd_outer = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()) + @test DEER_STRAT._hvp_strategy(so_fwd_outer) isa DEER_STRAT.ForwardOnGrad + + so_rev_outer = DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff()) + @test DEER_STRAT._hvp_strategy(so_rev_outer) isa DEER_STRAT.ReverseOnGrad + end + + @testset "strategy resolution is type-stable" begin + @test @inferred(DEER_STRAT._hvp_strategy(AutoForwardDiff())) isa + DEER_STRAT.ForwardOnGrad + @test @inferred(DEER_STRAT._hvp_strategy(AutoMooncake())) isa + DEER_STRAT.ReverseOnGrad + end +end