Add Enzyme reverse rules

Project.toml

@@ -6,6 +6,12 @@ version = "2.9.4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 
+[weakdeps]
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+
+[extensions]
+QuadGKEnzymeExt = "Enzyme"
+
 [compat]
 DataStructures = "0.11, 0.12, 0.13, 0.14, 0.15, 0.16, 0.17, 0.18, 0.19"
 julia = "1.2"

ext/QuadGKEnzymeExt.jl

@@ -0,0 +1,100 @@
+
+module QuadGKEnzymeExt
+
+using QuadGK, Enzyme, LinearAlgebra
+
+function Enzyme.EnzymeRules.augmented_primal(config, ofunc::Const{typeof(quadgk)}, ::Type{RT}, f, segs::Annotation{T}...; kws...) where {RT, T<:Real}
+    prims = map(x->x.val, segs)
+
+    retres, segbuf = if f isa Const
+        if EnzymeRules.needs_primal(config)
+            quadgk(f.val, prims...; kws...), nothing
+        else
+            nothing
+        end
+    else
+        I, E, segbuf = quadgk_segbuf(f.val, prims...; kws...)
+        if EnzymeRules.needs_primal(config)
+            (I, E), segbuf
+        else
+            nothing, segbuf
+        end
+    end
+
+    dres = if !Enzyme.EnzymeRules.needs_shadow(config)
+        nothing
+    elseif EnzymeRules.width(config) == 1
+        zero.(res...)
+    else
+        ntuple(Val(EnzymeRules.width(config))) do i
+            Base.@_inline_meta
+            zero.(res...)
+        end
+    end
+
+    cache = if RT <: Duplicated || RT <: DuplicatedNoNeed || RT <: BatchDuplicated || RT <: BatchDuplicatedNoNeed
+        dres
+    else
+        nothing
+    end
+    cache2 = segbuf, cache
+
+    return Enzyme.EnzymeRules.AugmentedReturn{
+        Enzyme.EnzymeRules.needs_primal(config) ? eltype(RT) : Nothing,
+        Enzyme.EnzymeRules.needs_shadow(config) ? (Enzyme.EnzymeRules.width(config) == 1 ? eltype(RT) : NTuple{Enzyme.EnzymeRules.width(config), eltype(RT)}) : Nothing,
+        typeof(cache2)
+    }(retres, dres, cache2)
+end
+
+function call(f, x)
+    f(x)
+end
+
+struct ClosureVector{F}
+    f::F
+end
+
+@inline function guaranteed_nonactive(::Type{T}) where T
+    rt = Enzyme.Compiler.active_reg_inner(T, (), nothing)
+    return rt == Enzyme.Compiler.AnyState || rt == Enzyme.Compiler.DupState
+end
+
+function Base.:+(a::ClosureVector, b::ClosureVector)
+    Enzyme.Compiler.recursive_add(a, b, identity, guaranteed_nonactive)
+end
+
+function Base.:-(a::ClosureVector, b::ClosureVector)
+    Enzyme.Compiler.recursive_add(a, b, x->-x, guaranteed_nonactive)
+end
+
+function Base.:*(a::Number, b::ClosureVector)
+    # b + (a-1) * b = a * b
+    Enzyme.Compiler.recursive_add(b, b, x->(a-1)*x, guaranteed_nonactive)
+end
+
+function Base.:*(a::ClosureVector, b::Number)
+    return b*a
+end
+
+function Enzyme.EnzymeRules.reverse(config, ofunc::Const{typeof(quadgk)}, dres::Active, cache, f, segs::Annotation{T}...; kws...) where {T<:Real}
+    df = if f isa Const
+        nothing
+    else
+        segbuf = cache[1]
+        fwd, rev = Enzyme.autodiff_thunk(ReverseSplitNoPrimal, Const{typeof(call)}, Active, typeof(f), Const{T})
+        _df, _ = quadgk(map(x->x.val, segs)...; kws..., eval_segbuf=segbuf, maxevals=0, norm=f->0) do x
+            tape, prim, shad = fwd(Const(call), f, Const(x))
+            drev = rev(Const(call), f, Const(x), dres.val[1], tape)
+            return ClosureVector(drev[1][1])
+        end
+        _df.f
+    end
+    dsegs1 = segs[1] isa Const ? nothing : -LinearAlgebra.dot(f.val(segs[1].val), dres.val[1])
+    dsegsn = segs[end] isa Const ? nothing : LinearAlgebra.dot(f.val(segs[end].val), dres.val[1])
+    return (df, # f
+            dsegs1,
+            ntuple(i -> nothing, Val(length(segs)-2))...,
+            dsegsn)
+end
+
+end # module

test/Project.toml

@@ -0,0 +1,4 @@
+[deps]
+Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
+Enzyme = "7da242da-08ed-463a-9acd-ee780be4f1d9"
+LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"

test/runtests.jl

@@ -426,3 +426,17 @@ quadgk_segbuf_printnull(args...; kws...) = quadgk_segbuf_print(devnull, args...;
     @inferred QuadGK.to_segbuf([0,1])
     @inferred QuadGK.to_segbuf([(0,1+3im)])
 end
+
+# Extension package only supported in 1.9+
+@static if VERSION >= v"1.9"
+    using Enzyme
+    f1(x) = quadgk(cos, 0., x)[1]
+    f2(x) = quadgk(cos, x, 1)[1]
+    f3(x) = quadgk(y->cos(x * y), 0., 1.)[1]
+
+    @testset "Enzyme" begin
+        @test cos(0.3) ≈ Enzyme.autodiff(Reverse, f1, Active(0.3))[1][1]
+        @test -cos(0.3) ≈ Enzyme.autodiff(Reverse, f2, Active(0.3))[1][1]
+        @test (0.3 * cos(0.3) - sin(0.3))/(0.3*0.3) ≈ Enzyme.autodiff(Reverse, f3, Active(0.3))[1][1]
+    end
+end