diff --git a/CHANGELOG.md b/CHANGELOG.md index 8b29749..d2d204a 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -50,12 +50,33 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 its strategy was routed on `DI.outer(backend)`, so a `DifferentiationInterface.SecondOrder` ran the wrong half of the pair. A `SecondOrder` now goes to the true second-order path instead of either - strategy, and the half-selecting helpers it still uses agree: normalization - applies to the outer, the pass whose mode the backend extensions care about. - Unwrapping to that half happens before the normalization hooks are dispatched - on, so a `SecondOrder(AutoEnzyme(), ...)` still reaches `EnzymeExt` and gets - its mode and function annotation pinned rather than running as a bare - `AutoEnzyme()` (which aborts on GPU). + strategy, and the half-selecting helper it still uses agrees: normalization + applies to the outer pass. Unwrapping to that half happens before the + normalization hook is dispatched on, so a `SecondOrder(AutoEnzyme(), ...)` still + reaches `EnzymeExt` and gets its function annotation filled in rather than + running as a bare `AutoEnzyme()`. +- Backend normalization no longer picks a differentiation mode on the user's + behalf (#62). `EnzymeExt` pinned `mode=Enzyme.Forward` (with + `set_runtime_activity`) onto an `AutoEnzyme()` left mode-agnostic, on the + grounds that reverse mode hit a gc-transition abort on GPU and that composed + `pmcmc_matmul` calls needed runtime activity. The `pmcmc_*` Enzyme rules keep + Enzyme off both paths on their own now, so the pin bought nothing — and it cost + correctness, because it silently rewrote the direction of a `SecondOrder`'s + outer half. `SecondOrder(AutoEnzyme(), AutoForwardDiff())` is + reverse-over-forward to `hvp_mode`, its inner half being forward-only, and came + out forward-over-forward. Normalization now fills in only + `function_annotation=Enzyme.Const`, which is about this package's own read-only + HVP wrappers rather than about Enzyme's mode, and leaves `mode` exactly as given + — unset included, for DI to resolve from the operator it runs. `hvp_mode` is + therefore identical before and after normalization for every backend pair. + + A mode set explicitly was never overridden, so only mode-agnostic backends were + affected, and the HVP was a correct HVP either way; what changes is that the + composition asked for is the one that runs. The two normalization hooks + (`_hvp_forward_backend`, `_hvp_closure_backend`) collapse into a single + `_normalized_backend`, since without a mode to choose they no longer differ. + Users relying on a plain `AutoEnzyme()` being run forward should now pass + `AutoEnzyme(; mode=Enzyme.Forward)` explicitly. ### Changed diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 5c0859f..82e63f8 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -20,46 +20,26 @@ using Enzyme.EnzymeCore.EnzymeRules: # TODO: Implement matmul overloads upstream in Enzyme. See: https://github.com/EnzymeAD/Enzyme.jl/issues/3122 #= -Tell DEER's forward-on-grad HVP path how to normalize a plain `AutoEnzyme()`: -pin `mode=Enzyme.Forward` and `function_annotation=Enzyme.Const`. Pinning -Forward is load-bearing on GPU — without it DI defaults to reverse mode, which -hits the gc-transition abort documented on the Enzyme rules below. +Normalization of the user's `AutoEnzyme` for DEER's AD-HVP paths: fill in +`function_annotation=Enzyme.Const` when they left it open, so Enzyme doesn't +throw `EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / +`_BatchHvpReverseClosure` wrappers, which capture `gradlogp`. Those wrapper types +belong to this package, so declaring them constant is this package's business. -`mode` and `function_annotation` are normalized independently — a user who -sets one keeps that choice, but still gets the default for the other. +`mode` is passed through exactly as given, unset included. Choosing a direction +on the user's behalf is not our call: a mode they set is a decision, and an unset +one is DI's to resolve from the operator it runs. -`set_runtime_activity` is load-bearing for composed `pmcmc_matmul` calls -(e.g. `pmcmc_matmul(transpose(X), pmcmc_matmul(X, β))`). Static activity -analysis can't prove the outer call's `transpose(X)` shadow is safe to -reuse, and Enzyme aborts with `EnzymeRuntimeActivityError`. With runtime -activity, the shadow is tracked dynamically. +An earlier version pinned `mode=Enzyme.Forward` here (with +`set_runtime_activity`) against the gc-transition abort on GPU and +`EnzymeRuntimeActivityError` on composed `pmcmc_matmul` calls. The rules below +keep Enzyme off both paths on their own, so the pin bought nothing and cost +correctness: it silently rewrote the direction of a `SecondOrder`'s outer half +(see `DEER._normalized_backend`). =# -function DEER._hvp_forward_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} - mode = if backend.mode === nothing - Enzyme.set_runtime_activity(Enzyme.Forward) - else - backend.mode - end - annotation = A === Nothing ? Enzyme.Const : A - return ADTypes.AutoEnzyme(; mode=mode, function_annotation=annotation) -end - -#= -Tell DEER's reverse-on-grad HVP path how to normalize a plain `AutoEnzyme()`: -fill in `function_annotation=Enzyme.Const` so Enzyme doesn't throw -`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / -`_BatchHvpReverseClosure` wrappers, and default `mode` to reverse with -runtime activity. As in `_hvp_forward_backend`, the two fields are -normalized independently. -=# -function DEER._hvp_closure_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} - mode = if backend.mode === nothing - Enzyme.set_runtime_activity(Enzyme.Reverse) - else - backend.mode - end - annotation = A === Nothing ? Enzyme.Const : A - return ADTypes.AutoEnzyme(; mode=mode, function_annotation=annotation) +function DEER._normalized_backend(backend::ADTypes.AutoEnzyme{M,A}) where {M,A} + A === Nothing || return backend + return ADTypes.AutoEnzyme(; mode=backend.mode, function_annotation=Enzyme.Const) end #= diff --git a/src/DEER/DEER.jl b/src/DEER/DEER.jl index 03896ef..f0d677b 100644 --- a/src/DEER/DEER.jl +++ b/src/DEER/DEER.jl @@ -93,7 +93,7 @@ function _prepare_hvp(f, backend::AbstractADType, x_template::AbstractVector) v_template = similar(x_template) fill!(v_template, zero(eltype(x_template))) return DI.prepare_pushforward( - f, _hvp_forward_backend(backend), x_template, (v_template,) + f, _normalized_backend(backend), x_template, (v_template,) ) end @@ -116,14 +116,14 @@ function _hvp_prepared( ) x_exec = _materialize_ad_vector(x) v_exec = _tangent_like(x_exec, v) - res = DI.pushforward(f, prep, _hvp_forward_backend(backend), x_exec, (v_exec,)) + res = DI.pushforward(f, prep, _normalized_backend(backend), x_exec, (v_exec,)) return res isa Tuple ? first(res) : res end function _hvp_nopre(f, backend::AbstractADType, x::AbstractVector, v::AbstractVector) x_exec = _materialize_ad_vector(x) v_exec = _tangent_like(x_exec, v) - res = DI.pushforward(f, _hvp_forward_backend(backend), x_exec, (v_exec,)) + res = DI.pushforward(f, _normalized_backend(backend), x_exec, (v_exec,)) return res isa Tuple ? first(res) : res end @@ -133,7 +133,7 @@ function _prepare_batch_hvp_from_grad( V_template = similar(X_template) fill!(V_template, zero(eltype(X_template))) return DI.prepare_pushforward( - grad_batch, _hvp_forward_backend(backend), X_template, (V_template,) + grad_batch, _normalized_backend(backend), X_template, (V_template,) ) end @@ -142,7 +142,7 @@ function _batch_hvp_from_grad_prepared( ) X_exec = _materialize_ad_matrix(X) V_exec = _tangent_like(X_exec, V) - res = DI.pushforward(grad_batch, prep, _hvp_forward_backend(backend), X_exec, (V_exec,)) + res = DI.pushforward(grad_batch, prep, _normalized_backend(backend), X_exec, (V_exec,)) return res isa Tuple ? first(res) : res end @@ -196,9 +196,8 @@ relying on constant propagation through `===`. The routing follows DI's `hvp_mode`: a forward outer pass (`DI.ForwardOverAnything`) takes `ForwardOnGrad`, anything else `ReverseOnGrad`. Only the outer direction matters since we differentiate -the already-built `gradlogp`. Plain `AutoEnzyme()` lands on `ForwardOnGrad`, -which we need: Enzyme reverse hits the gc-transition abort on GPU (see -`ext/EnzymeExt.jl`). +the already-built `gradlogp`. The mode is whatever the user's backend carries; +nothing here substitutes one (see `_normalized_backend`). =# abstract type HVPStrategy end struct ForwardOnGrad <: HVPStrategy end @@ -207,38 +206,31 @@ struct ReverseOnGrad <: HVPStrategy end _strategy_from(::DI.ForwardOverAnything) = ForwardOnGrad() _strategy_from(::DI.HVPMode) = ReverseOnGrad() -function _hvp_strategy(backend::Union{AbstractADType,DI.SecondOrder}) - return _strategy_from(DI.hvp_mode(backend)) -end +_hvp_strategy(backend::AbstractADType) = _strategy_from(DI.hvp_mode(backend)) #= -Hooks for backend-specific normalization of the user's `backend`. - -`_hvp_forward_backend` is for the forward-on-grad pushforward path -(differentiates the user's `gradlogp` directly). EnzymeExt specializes it -to pin `mode=Enzyme.Forward` and `function_annotation=Enzyme.Const` when -the user passed plain `AutoEnzyme()`, without pinning Forward, DI lowers -through reverse mode and hits the gc-transition abort on GPU (see -`ext/EnzymeExt.jl`). - -`_hvp_closure_backend` is for the reverse-on-grad gradient path on the -read-only `_HvpReverseClosure` / `_BatchHvpReverseClosure` wrappers. -EnzymeExt specializes it to set `function_annotation=Enzyme.Const` so -Enzyme doesn't throw `EnzymeMutabilityException` on a closure that captures -`gradlogp`. - -A `SecondOrder` is unwrapped to its outer half first, since the pass we are -about to run is the outer one i.e., the inner derivative is whatever -`gradlogp` already is. That keeps the half `_hvp_strategy` routed on, so the -strategy and the backend that carries it out can't end up disagreeing. The -unwrapping recurses rather than calling `DI.outer` in the generic method, so -that a wrapped backend still reaches its own normalization: dispatch happens -on what comes out of `DI.outer`, not on the `SecondOrder` around it. +Hook for backend-specific normalization of the user's `backend`, applied on every +AD-HVP path before the backend reaches DI. + +It supplies what the wrappers DEER differentiates need, and nothing else. Those +wrapper types are ours, so annotating them is ours to do: EnzymeExt specializes +this to fill `function_annotation=Enzyme.Const`, without which Enzyme throws +`EnzymeMutabilityException` on the read-only `_HvpReverseClosure` / +`_BatchHvpReverseClosure`, which capture `gradlogp`. + +It deliberately does not choose a differentiation mode. Which direction a pass +runs is the user's call when they state one and DI's to resolve from the operator +when they don't; this package is not an AD package and has no business overriding +either. Picking one here also used to corrupt a `SecondOrder`, whose halves carry +directions of their own (see `_normalized_second_order`). + +Callers hand this a single pass, never a `SecondOrder`: the strategy paths below +run one AD pass over a hand-written `gradlogp`, and `_resolve_hvp` sends every +`SecondOrder` to `_make_hvp_fn_second_order` before they are reached. +`_normalized_second_order` is the one caller that starts from a pair, and it +selects the outer half itself. =# -_hvp_forward_backend(backend::DI.SecondOrder) = _hvp_forward_backend(DI.outer(backend)) -_hvp_closure_backend(backend::DI.SecondOrder) = _hvp_closure_backend(DI.outer(backend)) -_hvp_forward_backend(backend::AbstractADType) = backend -_hvp_closure_backend(backend::AbstractADType) = backend +_normalized_backend(backend::AbstractADType) = backend function _prepare_hvp_via_grad_reverse( gradlogp, backend::AbstractADType, x_template::AbstractVector @@ -246,7 +238,7 @@ function _prepare_hvp_via_grad_reverse( v_template = similar(x_template) fill!(v_template, zero(eltype(x_template))) f = _HvpReverseClosure(gradlogp) - eff_backend = _hvp_closure_backend(backend) + eff_backend = _normalized_backend(backend) prep = DI.prepare_gradient(f, eff_backend, x_template, DI.Constant(v_template)) return (f, prep, eff_backend) end @@ -262,7 +254,7 @@ function _prepare_batch_hvp_via_grad_reverse( V_template = similar(X_template) fill!(V_template, zero(eltype(X_template))) f = _BatchHvpReverseClosure(grad_batch) - eff_backend = _hvp_closure_backend(backend) + eff_backend = _normalized_backend(backend) prep = DI.prepare_gradient(f, eff_backend, X_template, DI.Constant(V_template)) return (f, prep, eff_backend) end @@ -315,8 +307,12 @@ takes both passes over the log-density, so these never touch the gradient slot. Preferred over pushing tangents through a prepared DI gradient, which drops out of its preparation once the outer pass hands it an unexpected tangent type. -Only the outer half is normalized: `_hvp_forward_backend` selects it out of the -pair. The inner is a plain first-order gradient and needs no pinning. +Both halves are passed to `DI.hvp` as the user composed them, so the direction +each one runs in is theirs and DI's, not ours. Normalization touches only the +outer half, and only to fill in annotations for the wrappers being +differentiated; because it never substitutes a mode, `DI.hvp_mode` of the pair is +the same before and after. The inner half is a plain first-order gradient over +the user's own `logdensity` and is passed straight through. The batched form differentiates `sum(logdensity_batch(X))`, whose Hessian is block-diagonal by column independence, so its HVP along `V` is the columnwise @@ -324,7 +320,7 @@ HVP. Same argument the batched gradient rests on. --------------------------------------------------------------------------- =# function _normalized_second_order(backend::DI.SecondOrder) - return DI.SecondOrder(_hvp_forward_backend(backend), DI.inner(backend)) + return DI.SecondOrder(_normalized_backend(DI.outer(backend)), DI.inner(backend)) end function _make_hvp_fn_second_order( diff --git a/src/interface.jl b/src/interface.jl index 9568b0a..69ede25 100644 --- a/src/interface.jl +++ b/src/interface.jl @@ -27,8 +27,8 @@ DI's second-order operator. Passing a `SecondOrder` yourself always means the latter, and bypasses the gradient slot even when you wrote it by hand. - `logdensity(x::AbstractVector) -> Real` -- `grad_logdensity` — callable `x -> AbstractVector`, or a backend applied to - `logdensity`. +- `grad_logdensity` — callable `x -> AbstractVector`, or a backend to + differentiate `logdensity` with. - `hvp` — optional callable `(x, v) -> AbstractVector`, or a backend. If `nothing`, DEER builds the HVP from the sampler's `backend`. - `logdensity_batch(X::AbstractMatrix) -> AbstractVector` — optional batched @@ -37,7 +37,7 @@ latter, and bypasses the gradient slot even when you wrote it by hand. A batched gradient derived from this is one gradient of its sum, so coupling between columns would go unnoticed and give wrong derivatives. - `grad_logdensity_batch` — optional callable `X -> AbstractMatrix`, or a - backend applied to `logdensity_batch`. Left out alongside a + backend to differentiate `logdensity_batch` with. Left out alongside a `logdensity_batch`, it is derived when `grad_logdensity` is a backend. - `hvp_batch` — optional callable `(X, V) -> AbstractMatrix`, or a backend, resolved against `grad_logdensity_batch` the same way `hvp` is against diff --git a/test/test-HVP-Strategy.jl b/test/test-HVP-Strategy.jl index 8bfea83..e56997e 100644 --- a/test/test-HVP-Strategy.jl +++ b/test/test-HVP-Strategy.jl @@ -36,26 +36,56 @@ const DI_STRAT = ParallelMCMC.DEER.DI @test DEER_STRAT._hvp_strategy(so_agnostic_outer) isa DEER_STRAT.ReverseOnGrad end - @testset "the backend that runs is the one routed on" begin - # both paths differentiate the already-built gradlogp, so both take the outer - so_fwd = DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()) - @test DEER_STRAT._hvp_forward_backend(so_fwd) === AutoForwardDiff() + @testset "normalization supplies Const but never a mode" begin + #= The wrappers DEER differentiates are its own types, so annotating them + `Const` is its business. The mode is not: one the user set is a decision, + and an unset one is DI's to resolve from the operator it runs. =# + bare = DEER_STRAT._normalized_backend(AutoEnzyme()) + @test bare isa AutoEnzyme{<:Any,Enzyme.Const} + @test bare.mode === nothing - so_rev = DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff()) - @test DEER_STRAT._hvp_closure_backend(so_rev) === AutoZygote() + for mode in (Enzyme.Forward, Enzyme.Reverse) + normalized = DEER_STRAT._normalized_backend(AutoEnzyme(; mode=mode)) + @test normalized.mode === mode + @test normalized isa AutoEnzyme{<:Any,Enzyme.Const} + end + + # An annotation the user chose is left alone. + annotated = AutoEnzyme(; function_annotation=Enzyme.Duplicated) + @test DEER_STRAT._normalized_backend(annotated) === annotated + + # Backends with no specialization pass straight through. + @test DEER_STRAT._normalized_backend(AutoForwardDiff()) === AutoForwardDiff() + @test DEER_STRAT._normalized_backend(AutoZygote()) === AutoZygote() end - @testset "unwrapping a SecondOrder still reaches backend normalization" begin - #= The outer half has to be taken before the backend-specific hook is - dispatched on, or a wrapped `AutoEnzyme()` comes out bare: unnormalized, - it lowers through reverse mode and aborts on GPU (see ext/EnzymeExt.jl). =# - so = DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()) - @test DEER_STRAT._hvp_forward_backend(so) === - DEER_STRAT._hvp_forward_backend(AutoEnzyme()) - @test DEER_STRAT._hvp_closure_backend(so) === - DEER_STRAT._hvp_closure_backend(AutoEnzyme()) - @test DEER_STRAT._hvp_forward_backend(so).mode isa Enzyme.ForwardMode - @test DEER_STRAT._hvp_closure_backend(so) isa AutoEnzyme{<:Any,Enzyme.Const} + @testset "normalizing a SecondOrder keeps the composition DI resolved" begin + #= Regression for #62. Normalization used to route the outer half through + a forward-only hook, which pinned `Enzyme.Forward` onto it. For a pair + `hvp_mode` resolves to reverse -- `SecondOrder(AutoEnzyme(), + AutoForwardDiff())` is reverse-over-forward, its inner half being + forward-only -- that silently made it forward-over-forward. =# + for so in ( + DI_STRAT.SecondOrder(AutoEnzyme(), AutoForwardDiff()), + DI_STRAT.SecondOrder(AutoEnzyme(), AutoZygote()), + DI_STRAT.SecondOrder(AutoEnzyme(; mode=Enzyme.Reverse), AutoForwardDiff()), + DI_STRAT.SecondOrder(AutoEnzyme(; mode=Enzyme.Forward), AutoZygote()), + DI_STRAT.SecondOrder(AutoForwardDiff(), AutoZygote()), + DI_STRAT.SecondOrder(AutoZygote(), AutoForwardDiff()), + ) + normalized = DEER_STRAT._normalized_second_order(so) + @test DI_STRAT.hvp_mode(normalized) == DI_STRAT.hvp_mode(so) + # The inner half is the user's own first-order gradient, untouched. + @test DI_STRAT.inner(normalized) === DI_STRAT.inner(so) + # And no mode is invented for the outer half either. Only `AutoEnzyme` + # carries a `mode`, so this is the case that could regress. + if DI_STRAT.outer(so) isa AutoEnzyme + @test DI_STRAT.outer(normalized).mode === DI_STRAT.outer(so).mode + #= The outer half still reaches EnzymeExt's specialization rather + than passing through as a bare backend. =# + @test DI_STRAT.outer(normalized) isa AutoEnzyme{<:Any,Enzyme.Const} + end + end end @testset "strategy resolution is type-stable" begin