Make zero in place (#1518)

* Make zero in place

* add make_zero!

* more fixes and tests

Author: wsmoses

Project.toml

@@ -19,7 +19,7 @@ Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c"
 [compat]
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
-EnzymeCore = "0.7.3"
+EnzymeCore = "0.7.4"
 Enzyme_jll = "0.0.119"
 GPUCompiler = "0.21, 0.22, 0.23, 0.24, 0.25, 0.26"
 LLVM = "6.1, 7"

examples/custom_rule.jl

@@ -134,7 +134,7 @@ function forward(func::Const{typeof(f)}, RT::Type{<:Union{Const, DuplicatedNoNee
     if !(x isa Const) && !(y isa Const)
         y.dval .= 2 .* x.val .* x.dval
     elseif !(y isa Const) 
-        y.dval .= 0
+        make_zero!(y.dval)
     end
     dret = !(y isa Const) ? sum(y.dval) : zero(eltype(y.val))
     if RT <: Const
@@ -211,7 +211,7 @@ function reverse(config::ConfigWidth{1}, func::Const{typeof(f)}, dret::Active, t
     x.dval .+= 2 .* xval .* dret.val 
     ## also accumulate any derivative in y's shadow into x's shadow. 
     x.dval .+= 2 .* xval .* y.dval
-    y.dval .= 0
+    make_zero!(y.dval)
     return (nothing, nothing)
 end
 
@@ -251,8 +251,8 @@ end
 
 x  = [3.0, 1.0]
 y  = [0.0, 0.0]
-dx .= 0
-dy .= 0
+make_zero!(dx)
+make_zero!(dy)
 
 autodiff(Reverse, h, Duplicated(y, dy), Duplicated(x, dx))
 @show dx # derivative of h w.r.t. x

lib/EnzymeCore/Project.toml

@@ -1,7 +1,7 @@
 name = "EnzymeCore"
 uuid = "f151be2c-9106-41f4-ab19-57ee4f262869"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.7.3"
+version = "0.7.4"
 
 [compat]
 Adapt = "3, 4"

lib/EnzymeCore/src/EnzymeCore.jl

@@ -228,6 +228,13 @@ function autodiff_deferred_thunk end
 """
 function make_zero end
 
+"""
+    make_zero!(val::T, seen::IdSet{Any}=IdSet())::Nothing
+
+    Recursively set a variables differentiable fields to zero. Only applicable for mutable types `T`.
+"""
+function make_zero! end
+
 """
     make_zero(prev::T)
 

src/Enzyme.jl

@@ -14,8 +14,8 @@ export BatchDuplicatedFunc
 import EnzymeCore: batch_size, get_func 
 export batch_size, get_func
 
-import EnzymeCore: autodiff, autodiff_deferred, autodiff_thunk, autodiff_deferred_thunk, tape_type, make_zero
-export autodiff, autodiff_deferred, autodiff_thunk, autodiff_deferred_thunk, tape_type, make_zero
+import EnzymeCore: autodiff, autodiff_deferred, autodiff_thunk, autodiff_deferred_thunk, tape_type, make_zero, make_zero!
+export autodiff, autodiff_deferred, autodiff_thunk, autodiff_deferred_thunk, tape_type, make_zero, make_zero!
 
 export jacobian, gradient, gradient!
 export markType, batch_size, onehot, chunkedonehot
@@ -1007,7 +1007,7 @@ gradient!(Reverse, dx, f, [2.0, 3.0])
 ```
 """
 @inline function gradient!(::ReverseMode, dx::X, f::F, x::X) where {X<:Array, F}
-    dx .= 0
+    make_zero!(dx)
     autodiff(Reverse, f, Active, Duplicated(x, dx))
     dx
 end

src/api.jl

@@ -104,7 +104,7 @@ struct CFnTypeInfo
 end
 
 
-@static if isdefined(LLVM, :InstructionMetadataDict)
+@static if !isdefined(LLVM, :ValueMetadataDict)
 Base.haskey(md::LLVM.InstructionMetadataDict, kind::String) =
 	ccall((:EnzymeGetStringMD, libEnzyme), Cvoid, (LLVM.API.LLVMValueRef, Cstring), md.inst, kind) != C_NULL
 

src/compiler.jl

@@ -1298,7 +1298,7 @@ end
                 xi = getfield(prev, i)
                 T = Core.Typeof(xi)
                 xi = EnzymeCore.make_zero(T, seen, xi, Val(copy_if_inactive))
-                ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), y, i-1, xi)
+                setfield!(y, i, xi)
             end
         end
         return y
@@ -1324,6 +1324,204 @@ end
     return y
 end
 
+function make_zero_immutable!(prev::T, seen::S)::T where {T <: AbstractFloat, S}
+    zero(T)
+end
+
+function make_zero_immutable!(prev::Complex{T}, seen::S)::Complex{T} where {T <: AbstractFloat, S}
+    zero(T)
+end
+
+function make_zero_immutable!(prev::T, seen::S)::T where {T <: Tuple, S}
+    ntuple(Val(length(T.parameters))) do i
+        Base.@_inline_meta
+        make_zero_immutable!(prev[i], seen)
+    end
+end
+
+function make_zero_immutable!(prev::NamedTuple{a, b}, seen::S)::NamedTuple{a, b} where {a,b, S}
+    NamedTuple{a, b}(
+    ntuple(Val(length(T.parameters))) do i
+        Base.@_inline_meta
+        make_zero_immutable!(prev[a[i]], seen)
+    end
+    )
+end
+
+
+function make_zero_immutable!(prev::T, seen::S)::T where {T, S}
+    if guaranteed_const_nongen(T, nothing)
+        return prev
+    end
+    @assert !ismutable(T)
+
+    @assert !Base.isabstracttype(RT)
+    @assert Base.isconcretetype(RT)
+    nf = fieldcount(RT)
+
+    flds = Vector{Any}(undef, nf)
+    for i in 1:nf
+        if isdefined(prev, i)
+            xi = getfield(prev, i)
+            ST = Core.Typeof(xi)
+            flds[i] = if active_reg_inner(ST, (), nothing, #=justActive=#Val(true)) == ActiveState
+                make_zero_immutable!(xi, seen)
+            else
+                EnzymeCore.make_zero!(xi, seen)
+                xi
+            end
+        else
+            nf = i - 1 # rest of tail must be undefined values
+            break
+        end
+    end
+    ccall(:jl_new_structv, Any, (Any, Ptr{Any}, UInt32), RT, flds, nf)::T
+end
+
+@inline function EnzymeCore.make_zero!(prev::Base.RefValue{T}, seen::ST)::Nothing where {T <: AbstractFloat, ST}
+    T[] = zero(T)
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Base.RefValue{Complex{T}}, seen::ST)::Nothing where {T <: AbstractFloat, ST}
+    T[] = zero(Complex{T})
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Array{T, N}, seen::ST)::Nothing where {T <: AbstractFloat, N, ST}
+    fill!(prev, zero(T))
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Array{Complex{T}, N}, seen::ST)::Nothing where {T <: AbstractFloat, N, ST}
+    fill!(prev, zero(Complex{T}))
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Base.RefValue{T})::Nothing where {T <: AbstractFloat}
+    EnzymeCore.make_zero!(prev, nothing)
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Base.RefValue{Complex{T}})::Nothing where {T <: AbstractFloat}
+    EnzymeCore.make_zero!(prev, nothing)
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Array{T, N})::Nothing where {T <: AbstractFloat, N}
+    EnzymeCore.make_zero!(prev, nothing)
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Array{Complex{T}, N})::Nothing where {T <: AbstractFloat, N}
+    EnzymeCore.make_zero!(prev, nothing)
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Array{T, N}, seen::ST)::Nothing where {T, N, ST}
+    if guaranteed_const_nongen(T, nothing)
+        return
+    end
+    if in(seen, prev)
+        return
+    end
+    push!(seen, prev)
+
+    for I in eachindex(prev)
+        if isassigned(prev, I)
+            pv = prev[I]
+            SBT = Core.Typeof(pv)
+            if active_reg_inner(SBT, (), nothing, #=justActive=#Val(true)) == ActiveState
+                @inbounds prev[I] = make_zero_immutable!(pv, seen)
+                nothing
+            else
+                EnzymeCore.make_zero!(pv, seen)
+                nothing
+            end
+        end
+    end
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Base.RefValue{T}, seen::ST)::Nothing where {T, ST}
+    if guaranteed_const_nongen(T, nothing)
+        return
+    end
+    if in(seen, prev)
+        return
+    end
+    push!(seen, prev)
+
+    pv = prev[]
+    SBT = Core.Typeof(pv)
+    if active_reg_inner(SBT, (), nothing, #=justActive=#Val(true)) == ActiveState
+        prev[] = make_zero_immutable!(pv, seen)
+        nothing
+    else
+        EnzymeCore.make_zero!(pv, seen)
+        nothing
+    end
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::Core.Box, seen::ST)::Nothing where {ST}
+    pv = prev.contents
+    T = Core.Typeof(pv)
+    if guaranteed_const_nongen(T, nothing)
+        return
+    end
+    if in(seen, prev)
+        return
+    end
+    push!(seen, prev)
+    SBT = Core.Typeof(pv)
+    if active_reg_inner(SBT, (), nothing, #=justActive=#Val(true)) == ActiveState
+        prev.contents = EnzymeCore.make_zero_immutable!(pv, seen)
+        nothing
+    else
+        EnzymeCore.make_zero!(pv, seen)
+        nothing
+    end
+    nothing
+end
+
+@inline function EnzymeCore.make_zero!(prev::T, seen::S=Base.IdSet{Any}())::Nothing where {T, S}
+    if guaranteed_const_nongen(T, nothing)
+        return
+    end
+    if in(seen, prev)
+        return
+    end
+    @assert !Base.isabstracttype(T)
+    @assert Base.isconcretetype(T)
+    nf = fieldcount(T)
+    
+
+    if nf == 0
+        return
+    end
+
+    push!(seen, prev)
+
+    for i in 1:nf
+        if isdefined(prev, i)
+            xi = getfield(prev, i)
+            SBT = Core.Typeof(xi)
+            if guaranteed_const_nongen(SBT, nothing)
+                continue
+            end
+            if active_reg_inner(SBT, (), nothing, #=justActive=#Val(true)) == ActiveState
+                setfield!(prev, i, make_zero_immutable!(xi, seen))
+                nothing
+            else
+                EnzymeCore.make_zero!(xi, seen)
+                nothing
+            end
+        end
+    end
+    return
+end
+
 struct EnzymeRuntimeException <: Base.Exception
     msg::Cstring
 end
@@ -5536,7 +5734,7 @@ end
                 @assert ismutable(x)
                 yi = getfield(y, i)
                 nexti = recursive_add(xi, yi, f, mutable_register)
-                ccall(:jl_set_nth_field, Cvoid, (Any, Csize_t, Any), x, i-1, nexti)
+                setfield!(x, i, nexti)
             end
         end
     end

test/runtests.jl

@@ -181,6 +181,28 @@ end
     # @test thunk_split.primal !== C_NULL
     # @test thunk_split.primal !== thunk_split.adjoint
     # @test thunk_a.adjoint !== thunk_split.adjoint
+    #
+    z = ([3.14, 21.5, 16.7], [0,1], [5.6, 8.9])
+    Enzyme.make_zero!(z)
+    @test z[1] ≈ [0.0, 0.0, 0.0]
+    @test z[2][1] == 0
+    @test z[2][2] == 1
+    @test z[3] ≈ [0.0, 0.0]
+    
+    z2 = ([3.14, 21.5, 16.7], [0,1], [5.6, 8.9])
+    Enzyme.make_zero!(z2)
+    @test z2[1] ≈ [0.0, 0.0, 0.0]
+    @test z2[2][1] == 0
+    @test z2[2][2] == 1
+    @test z2[3] ≈ [0.0, 0.0]
+    
+    z3 = [3.4, "foo"]
+    Enzyme.make_zero!(z3)
+    @test z3[1] ≈ 0.0
+    @test z3[2] == "foo"
+
+    z4 = sin
+    Enzyme.make_zero!(z4)
 end
 
 @testset "Reflection" begin