Additional functions, tests, and fixes

Author: wsmoses

src/enzyme.jl

@@ -1,9 +1,12 @@
 import EnzymeCore
 
-function EnzymeCore.EnzymeRules.augmented_primal(config, func::EnzymeCore.Const{typeof(NNlib.conv!)}, ::Type{RT}, y::OutType, x, w, cdims; kwargs...) where {OutType, RT}
+for name in (typeof(NNlib.conv!), typeof(NNlib.depthwiseconv!))
+    @eval begin
+
+function EnzymeCore.EnzymeRules.augmented_primal(config, func::EnzymeCore.Const{$name}, ::Type{RT}, y::OutType, x, w, cdims; kwargs...) where {OutType, RT}
 
     @assert !(OutType <: EnzymeCore.Const)
-    if OutType <: EnzymeCore.Duplicated || OutType <: EnzymeCore.DuplicatedNoNeed
+    if OutType <: EnzymeCore.Duplicated || OutType <: EnzymeCore.BatchDuplicated
         func.val(y.val, x.val, w.val, cdims.val; kwargs...)
     end
 
@@ -29,7 +32,7 @@ function EnzymeCore.EnzymeRules.augmented_primal(config, func::EnzymeCore.Const{
     return EnzymeCore.EnzymeRules.AugmentedReturn(primal, shadow, cache)
 end
 
-function EnzymeCore.EnzymeRules.reverse(config, func::EnzymeCore.Const{typeof(NNlib.conv!)}, ::Type{RT}, cache, y, x, w, cdims; kwargs...) where {RT}
+function EnzymeCore.EnzymeRules.reverse(config, func::EnzymeCore.Const{$name}, ::Type{RT}, cache, y, x, w, cdims; kwargs...) where {RT}
     cache_x, cache_w = cache
 
     # Don't cache x if not overwritten and w is active (and thus required)
@@ -71,11 +74,13 @@ function EnzymeCore.EnzymeRules.reverse(config, func::EnzymeCore.Const{typeof(NN
     return (nothing, nothing, nothing, nothing)
 end
 
+end
+end
 
 function EnzymeCore.EnzymeRules.augmented_primal(config, func::EnzymeCore.Const{typeof(NNlib.gather!)}, ::Type{RT}, dst::OutType, src, idx::EnzymeCore.Const) where {OutType, RT}
 
     @assert !(OutType <: EnzymeCore.Const)
-    if OutType <: EnzymeCore.Duplicated || OutType <: EnzymeCore.DuplicatedNoNeed
+    if OutType <: EnzymeCore.Duplicated || OutType <: EnzymeCore.BatchDuplicated
         func.val(dst.val, src.val, idx.val)
     end
 
@@ -114,14 +119,76 @@ function EnzymeCore.EnzymeRules.reverse(config, func::EnzymeCore.Const{typeof(NN
     end
 
     for (ddst, dsrc) in zip(ddsts, dsrcs)
-        if !(typeof(src) <: EnzymeCore.Const) && ddst !== dst.val
-            src_size = size(src.val)
-            NNlib.∇gather_src(ddst, src_size, cache_idx)
+        if !(typeof(src) <: EnzymeCore.Const) && dsrc !== src.val &&
+           !(typeof(dst) <: EnzymeCore.Const) && ddst !== dst.val
+            NNlib.scatter!(+, dsrc, ddst, cache_idx)
         end
-        if !(typeof(w) <: EnzymeCore.Const) && dw !== w
+        if !(typeof(dst) <: EnzymeCore.Const) && ddst !== dst.val
             ddst .= 0
         end
     end
 
+    return (nothing, nothing, nothing)
+end
+
+
+
+function EnzymeCore.EnzymeRules.augmented_primal(config, func::EnzymeCore.Const{typeof(NNlib.scatter!)}, ::Type{RT}, op::EnzymeCore.Const, dst::OutType, src, idx::EnzymeCore.Const) where {OutType, RT}
+
+    @assert !(OutType <: EnzymeCore.Const)
+    if OutType <: EnzymeCore.Duplicated || OutType <: EnzymeCore.BatchDuplicated
+        func.val(op.val, dst.val, src.val, idx.val)
+    end
+
+    primal = if EnzymeCore.EnzymeRules.needs_primal(config)
+        dst.val
+    else
+        nothing
+    end
+    shadow = if EnzymeCore.EnzymeRules.needs_shadow(config)
+        dst.dval
+    else
+        nothing
+    end
+
+    # Cache idx if its overwritten
+    cache_idx = ( EnzymeCore.EnzymeRules.overwritten(config)[4] && !(typeof(src) <: EnzymeCore.Const) ) ? copy(idx.val) : nothing
+
+    return EnzymeCore.EnzymeRules.AugmentedReturn(primal, shadow, cache_idx)
+end
+
+function EnzymeCore.EnzymeRules.reverse(config, func::EnzymeCore.Const{typeof(NNlib.scatter!)}, ::Type{RT}, cache_idx, op::Union{EnzymeCore.Const{typeof(+)},EnzymeCore.Const{typeof(-)}}, dst::OutType, src, idx::EnzymeCore.Const) where {OutType, RT}
+
+    # Don't cache idx if not overwritten
+    if !(typeof(src) <: EnzymeCore.Const)
+        if !EnzymeCore.EnzymeRules.overwritten(config)[4]
+            cache_idx = idx.val
+        end
+    end
+
+    ddsts = dst.dval
+    dsrcs = src.dval
+
+    if EnzymeCore.EnzymeRules.width(config) == 1
+        ddsts = (ddsts,)
+        dsrcs = (dsrcs,)
+    end
+
+    for (ddst, dsrc) in zip(ddsts, dsrcs)
+        if !(typeof(src) <: EnzymeCore.Const) && dsrc !== src.val &&
+            !(typeof(dst) <: EnzymeCore.Const) && ddst !== dst.val
+
+            if eltype(typeof(op)) == typeof(+)
+                dsrc .+= NNlib.gather(ddst, cache_idx)
+            else
+                @assert eltype(typeof(op)) == typeof(-)
+                dsrc .-= NNlib.gather(ddst, cache_idx)
+            end
+        end
+    end
+
     return (nothing, nothing, nothing, nothing)
-end
+end
+
+
+

test/conv.jl

@@ -861,7 +861,7 @@ end
   w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
   cdims = DenseConvDims(x, w)
   gradtest((x, w) -> conv(x, w, cdims), x, w)
-  gradtest((x, w) -> sum(conv(x, w, cdims)), x, w; check_enzyme_rrule=true)  # https://github.com/FluxML/Flux.jl/issues/1055
+  gradtest((x, w) -> sum(conv(x, w, cdims)), x, w)  # https://github.com/FluxML/Flux.jl/issues/1055
 
   y = conv(x, w, cdims)
   gradtest((y, w) -> ∇conv_data(y, w, cdims), y, w)
@@ -877,3 +877,25 @@ end
   gradtest((y, w) -> ∇depthwiseconv_data(y, w, dcdims), y, w)
   gradtest((y, w) -> sum(∇depthwiseconv_data(y, w, dcdims)), y, w)
 end
+
+@testset "EnzymeRules: conv! spatial_rank=$spatial_rank" for spatial_rank in (1, 2, 3)
+  x = rand(rng, repeat([5], spatial_rank)..., 3, 2)
+  w = rand(rng, repeat([3], spatial_rank)..., 3, 3)
+  cdims = DenseConvDims(x, w)
+
+  for name in (:conv, :depthwiseconv)
+    curconv = @eval $(Symbol("$(name)"))
+    curconv! = @eval $(Symbol("$(name)!"))
+    dst = curconv(x, w, cdims)
+
+    for Tret in (EnzymeCore.Const, EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
+        Tdst in (EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
+        Tx in (EnzymeCore.Const, EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
+        Tw in (EnzymeCore.Const, EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated)
+
+        EnzymeTestUtils.are_activities_compatible(Tret, Tdst, Tsrc) || continue
+
+        EnzymeTestUtils.test_reverse(curconv!, Tret, (dst, Tdst), (x, Tx), (x, Tw), (idx, EnzymeCore.Const))
+    end
+  end
+end

test/gather.jl

@@ -152,16 +152,23 @@ function gather_testsuite(Backend)
         Backend == CPU ?
             gradtest_fn(xs -> gather(xs, idx), src) :
             gradtest_fn((s, i) -> gather(s, i), src, idx)
+    end
 
-        if Backend == CPU
-            for Tret in (EnzymeCore.Const, EnzymeCore.Duplicated),
-                Tdst in (EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
-                Tsrc in (EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated)
 
-                EnzymeTestUtils.are_activities_compatible(Tret, Tdst, Tsrc) || continue
+    @testset "EnzymeRules: gather! gradient for scalar index" begin
+        src = device(Float64[3, 4, 5, 6, 7])
+        idx = device([
+            1 2 3 4;
+            4 2 1 3;
+            3 5 5 3])
+        dst = gather(src, idx)
+        for Tret in (EnzymeCore.Const, EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
+            Tdst in (EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
+            Tsrc in (EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated)
+
+            EnzymeTestUtils.are_activities_compatible(Tret, Tdst, Tsrc) || continue
 
-                EnzymeTestUtils.test_reverse(fun, Tret, (dst, Tdst), (src, Tsrc), (idx, EnzymeCore.Const))
-            end
+            EnzymeTestUtils.test_reverse(gather!, Tret, (dst, Tdst), (src, Tsrc), (idx, EnzymeCore.Const))
         end
     end
 

test/scatter.jl

@@ -207,5 +207,24 @@ function scatter_testsuite(Backend)
                     gradtest_fn((xs, i) -> scatter(op, xs, i), src, idx)
             end
         end
+
+        @testset "EnzymeRules" begin
+            idx = device([2, 2, 3, 4, 4])
+            src = device(ones(T, 3, 5))
+
+            for op in (+, -)
+
+                dst = scatter(op, src, idx)
+
+                for Tret in (EnzymeCore.Const, EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
+                    Tdst in (EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
+                    Tsrc in (EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated)
+
+                    EnzymeTestUtils.are_activities_compatible(Tret, Tdst, Tsrc) || continue
+
+                    EnzymeTestUtils.test_reverse(scatter!, Tret, (op, EnzymeCore.Const), (dst, Tdst), (src, Tsrc), (idx, EnzymeCore.Const))
+                end
+            end
+        end
     end
 end

test/test_utils.jl

@@ -12,30 +12,14 @@ Applies also `ChainRulesTestUtils.test_rrule` if the rrule for `f` is explicitly
 """
 function gradtest(
     f, xs...; atol = 1e-6, rtol = 1e-6, fkwargs = NamedTuple(),
-    check_rrule = false, check_enzyme_rrule = false, fdm = :central, check_broadcast = false,
+    check_rrule = false, fdm = :central, check_broadcast = false,
     skip = false, broken = false,
 )
     # TODO: revamp when https://github.com/JuliaDiff/ChainRulesTestUtils.jl/pull/166
     # is merged
     if check_rrule
         test_rrule(f, xs...; fkwargs = fkwargs)
     end
-    if check_enzyme_rrule
-        if len(xs) == 2
-            for Tret in (EnzymeCore.Const, EnzymeCore.Active),
-                Tx in (EnzymeCore.Const, EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated),
-                Ty in (EnzymeCore.Const, EnzymeCore.Duplicated, EnzymeCore.BatchDuplicated)
-
-                EnzymeTestUtils.are_activities_compatible(Tret, Tx, Ty) || continue
-
-                EnzymeTestUtils.test_reverse(fun, Tret, (xs[1], Tx), (ys[1], Ty); atol, rtol)
-            end
-        else
-            throw(AssertionError("Unsupported arg count for testing"))
-        end
-
-        EnzymeTestUtils.test_rrule(f, xs...; fkwargs = fkwargs)
-    end
 
     if check_broadcast
         length(fkwargs) > 0 && @warn("CHECK_BROADCAST: dropping keywords args")