Deepcopy of struct (#2510)

* Deepcopy of struct

* fix

* Update runtests.jl

Author: wsmoses

src/internal_rules.jl

@@ -160,36 +160,40 @@ end
 @inline function deepcopy_rtact(
     copied::RT,
     primal::RT,
-    seen::IdDict,
-    shadow::RT,
-) where {RT<:Union{Integer,Char}}
-    return Base.deepcopy_internal(shadow, seen)
-end
-@inline function deepcopy_rtact(
-    copied::RT,
-    primal::RT,
-    seen::IdDict,
-    shadow::RT,
-) where {RT<:AbstractFloat}
-    return Base.deepcopy_internal(shadow, seen)
-end
-@inline function deepcopy_rtact(
-    copied::RT,
-    primal::RT,
-    seen::IdDict,
+    seen::Union{IdDict,Nothing},
     shadow::RT,
-) where {RT<:Array}
-    if !haskey(seen, shadow)
+)::RT where RT
+    rt = Enzyme.Compiler.active_reg_inner(RT, (), nothing)
+    if rt == Enzyme.Compiler.ActiveState || rt == Enzyme.Compiler.AnyState
+        if seen === nothing
+            return Base.deepcopy(shadow)
+        else
+            return Base.deepcopy_internal(shadow, seen)
+        end
+    else
+        if seen !== nothing && haskey(seen, shadow)
+            return seen[shadow]
+        end
         if primal === shadow
-            return seen[shadow] = copied
+            if seen !== nothing
+                seen[shadow] = copied
+            end
+            return copied
         end
-        newa = RT(undef, size(shadow))
-        seen[shadow] = newa
-        for i in eachindex(shadow)
-            @inbounds newa[i] = deepcopy_rtact(copied[i], primal[i], seen, shadow[i])
+
+        if RT <: Array
+            newa = similar(primal, size(shadow))
+            if seen === nothing
+                seen = IdDict()
+            end
+            seen[shadow] = newa
+            for i in eachindex(shadow)
+                @inbounds newa[i] = deepcopy_rtact(copied[i], primal[i], seen, shadow[i])
+            end
+            return newa
         end
+        throw(AssertionError("Unimplemented deepcopy with runtime activity for type $RT"))
     end
-    return seen[shadow]
 end
 
 function EnzymeRules.forward(
@@ -199,7 +203,7 @@ function EnzymeRules.forward(
     x::Duplicated,
 )
     primal = func.val(x.val)
-    return Duplicated(primal, deepcopy_rtact(primal, x.val, IdDict(), x.dval))
+    return Duplicated(primal, deepcopy_rtact(primal, x.val, nothing, x.dval))
 end
 
 function EnzymeRules.forward(
@@ -210,7 +214,7 @@ function EnzymeRules.forward(
 ) where {T,N}
     primal = func.val(x.val)
     return BatchDuplicated(primal, ntuple(Val(N)) do i
-        deepcopy_rtact(primal, x.val, IdDict(), x.dval[i])
+        deepcopy_rtact(primal, x.val, nothing, x.dval[i])
     end)
 end
 
@@ -226,15 +230,14 @@ function EnzymeRules.augmented_primal(
         nothing
     end
 
-    @assert !(typeof(x) <: Active)
-
     source = if EnzymeRules.needs_primal(config)
         primal
     else
         x.val
     end
 
     shadow = if EnzymeRules.needs_shadow(config)
+        @assert !(x isa Active)
         if EnzymeRules.width(config) == 1
             Enzyme.make_zero(
                 source,
@@ -304,6 +307,8 @@ function EnzymeRules.reverse(
     shadow,
     x::Annotation{Ty},
 ) where {RT,Ty}
+    @assert !(x isa Active)
+
     if EnzymeRules.needs_shadow(config)
         if EnzymeRules.width(config) == 1
             accumulate_into(x.dval, IdDict(), shadow)
@@ -317,6 +322,17 @@ function EnzymeRules.reverse(
     return (nothing,)
 end
 
+function EnzymeRules.reverse(
+    config::EnzymeRules.RevConfig,
+    func::Const{typeof(Base.deepcopy)},
+    dret::Active,
+    shadow,
+    x::Annotation,
+)
+    return (dret.val,)
+end
+
+
 @inline function pmap_fwd(
     idx,
     tapes::Vector,

test/runtests.jl

@@ -443,6 +443,19 @@ make3() = (1.0, 2.0, 3.0)
 
 end
 
+function named_deepcopy(x, nt)
+    nt2 = deepcopy(nt)
+    return nt2.a + x[1]
+end
+
+@testset "Deepcopy" begin
+    nt = (a = 0.0,)
+    x = [0.5]
+
+    @test Enzyme.gradient(Forward, named_deepcopy, x, Const(nt))[1] ≈ [1.0]
+    @test Enzyme.gradient(Reverse, named_deepcopy, x, Const(nt))[1] ≈ [1.0]
+end
+
 @testset "Deferred and deferred thunk" begin
     function dot(A)
         return A[1] * A[1] + A[2] * A[2]