Implement support for finalizers (#2736)

Co-authored-by: Lucas Wilcox <lucas@swirlee.com>

Author: vchuravy

docs/src/faq.md

@@ -744,4 +744,38 @@ autodiff(Forward, f_val, Duplicated(Val(1.0), Val(1.0)))
 ERROR: Type of ghost or constant type Duplicated{Val{1.0}} is marked as differentiable.
 ```
 
+## Finalizers
+
+Julia supports attaching finalizers to objects (see the listing below for an example) 
+
+```julia
+mutable struct Obj
+    x::Float64
+    function Obj(x)
+        o = new(x)
+        finalizer(o) do o
+            # do someting with o
+        end
+        return o
+    end
+end
+```
+
+When Enzyme encounters a code like:
+
+```julia
+function f(x)
+    o = Obj(x)
+    # computations over o
+    return o.x
+end
+
+autodiff(Forward, f, Duplicated(1.0, 1.0))
+```
+
+Enzyme has to allocate a shadow object for `o` and in the process encounters the finalizer being attached to the primal object.
+Now the question is what should Enzyme do with the finalizer for the shadow objects? One option would be to simply ignore it,
+but finalizers are often used for resource management (like manually allocating memory) and thus we would leak resources that are attached
+to the shadow object. Instead, we define finalizers to be inactive (contain no instructions that are relevant with respect to AD),
+yet we must attach them to the shadow object to release resources attached to them. 
 

src/internal_rules.jl

@@ -1797,3 +1797,44 @@ function EnzymeRules.reverse(
     dxs = map(x -> _hypotreverse(x, w, dret, n), xs)
     return (dx, dy, dz, dxs...)
 end
+
+function EnzymeRules.forward(config, ::Const{typeof(Base.finalizer)}, _, f::Const, o)
+    f = f.val
+    Base.finalizer(f, o.val)
+    if EnzymeRules.width(config) == 1
+        Base.finalizer(f, o.dval)
+    else
+        foreach(o.dval) do dv
+            Base.finalizer(f, dv)
+        end
+    end
+
+    if EnzymeRules.needs_primal(config)
+        return o
+    else
+        return nothing
+    end
+end
+
+function EnzymeRules.augmented_primal(config, ::Const{typeof(Base.finalizer)}, _, f::Const, o)
+    @assert !(o isa Active)
+    f = f.val
+    Base.finalizer(f, o.val)
+    if EnzymeRules.width(config) == 1
+        Base.finalizer(f, o.dval)
+    else
+        foreach(o.dval) do dv
+            Base.finalizer(f, dv)
+        end
+    end
+
+    primal = EnzymeRules.needs_primal(config) ? o.val : nothing
+    shadow = EnzymeRules.needs_shadow(config) ? o.dval : nothing
+
+    return EnzymeRules.AugmentedReturn(primal, shadow, nothing)
+end
+
+function EnzymeRules.reverse(config, ::Const{typeof(Base.finalizer)}, dret, tape, f::Const, o)
+    # No-op
+    return (nothing, nothing)
+end

test/finalizers.jl

@@ -0,0 +1,80 @@
+using Enzyme
+using Test
+
+const FREE_LIST = Vector{Any}()
+
+mutable struct Container
+    value::Float64
+    function Container(v::Float64)
+        c = new(v)
+        finalizer(c) do c
+            # Necromance object
+            push!(FREE_LIST, c)
+        end
+        return c
+    end
+end
+
+@noinline function compute(c)
+    return c.value^2
+end
+
+function compute(x::Float64)
+    c = Container(x)
+    return compute(c)
+end
+
+@testset "primal" begin
+    x = compute(1.0)
+    @test x == 1.0
+    GC.gc()
+    @test length(FREE_LIST) == 1
+    empty!(FREE_LIST)
+end
+
+@testset "forward" begin
+    dx, x = autodiff(ForwardWithPrimal, compute, Duplicated(1.0, 2.0))
+    @test x == 1.0
+    @test dx == 4.0
+    GC.gc()
+    @test length(FREE_LIST) == 2
+    empty!(FREE_LIST)
+
+    dx, = autodiff(Forward, compute, Duplicated(1.0, 2.0))
+    @test dx == 4.0
+    GC.gc()
+    @test length(FREE_LIST) == 2
+    empty!(FREE_LIST)
+end
+
+@testset "batched forward" begin
+    dx, x = autodiff(ForwardWithPrimal, compute, BatchDuplicated(1.0, (1.0, 2.0)))
+    @test x == 1.0
+    @test dx[1] == 2.0
+    @test dx[2] == 4.0
+    GC.gc()
+    @test length(FREE_LIST) == 3
+    empty!(FREE_LIST)
+
+    dx, = autodiff(Forward, compute, BatchDuplicated(1.0, (1.0, 2.0)))
+    @test dx[1] == 2.0
+    @test dx[2] == 4.0
+    GC.gc()
+    @test length(FREE_LIST) == 3
+    empty!(FREE_LIST)
+end
+
+@testset "reverse" begin
+    ((dx,), x) = autodiff(ReverseWithPrimal, compute, Active(1.0))
+    @test x == 1.0
+    @test dx == 2.0
+    GC.gc()
+    @test length(FREE_LIST) == 2
+    empty!(FREE_LIST)
+
+    ((dx,),) = autodiff(Reverse, compute, Active(1.0))
+    @test dx == 2.0
+    GC.gc()
+    @test length(FREE_LIST) == 2
+    empty!(FREE_LIST)
+end