Merge branch 'main' into gd/enzymecore_test

.buildkite/pipeline.yml

@@ -4,6 +4,7 @@ steps:
       setup:
         version:
           - "1.10"
+          - "1.11"
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.version}}"
@@ -25,55 +26,55 @@ steps:
     env:
       JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
 
-  - label: "AMDGPU Julia v{{matrix.version}}"
-    matrix:
-      setup:
-        version:
-          - "1.10"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "{{matrix.version}}"
-    agents:
-      queue: "juliagpu"
-      rocm: "*"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
-    commands: |
-      echo "--- Setup Julia packages"
-      julia --color=yes -e '
-          using Pkg
-          pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")]
-          push!(pkgs, PackageSpec(; name="AMDGPU"))
-          Pkg.develop(pkgs)' || exit 3
-
-      echo "+++ Run tests"
-      julia --color=yes test/amdgpu.jl
-    env:
-      JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
-
-  - label: "Metal Julia v{{matrix.version}}"
-    matrix:
-      setup:
-        version:
-          - "1.10"
-    plugins:
-      - JuliaCI/julia#v1:
-          version: "{{matrix.version}}"
-    agents:
-      queue: "juliaecosystem"
-      os: "macos"
-      arch: "aarch64"
-    if: build.message !~ /\[skip tests\]/
-    timeout_in_minutes: 60
-    commands: |
-      echo "--- Setup Julia packages"
-      julia --color=yes -e '
-          using Pkg
-          pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")]
-          push!(pkgs, PackageSpec(; name="Metal"))
-          Pkg.develop(pkgs)' || exit 3
-
-      echo "+++ Run tests"
-      julia --color=yes test/metal.jl
-    env:
-      JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
+#   - label: "AMDGPU Julia v{{matrix.version}}"
+#     matrix:
+#       setup:
+#         version:
+#           - "1.10"
+#     plugins:
+#       - JuliaCI/julia#v1:
+#           version: "{{matrix.version}}"
+#     agents:
+#       queue: "juliagpu"
+#       rocm: "*"
+#     if: build.message !~ /\[skip tests\]/
+#     timeout_in_minutes: 60
+#     commands: |
+#       echo "--- Setup Julia packages"
+#       julia --color=yes -e '
+#           using Pkg
+#           pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")]
+#           push!(pkgs, PackageSpec(; name="AMDGPU"))
+#           Pkg.develop(pkgs)' || exit 3
+# 
+#       echo "+++ Run tests"
+#       julia --color=yes test/amdgpu.jl
+#     env:
+#       JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager
+# 
+#   - label: "Metal Julia v{{matrix.version}}"
+#     matrix:
+#       setup:
+#         version:
+#           - "1.10"
+#     plugins:
+#       - JuliaCI/julia#v1:
+#           version: "{{matrix.version}}"
+#     agents:
+#       queue: "juliaecosystem"
+#       os: "macos"
+#       arch: "aarch64"
+#     if: build.message !~ /\[skip tests\]/
+#     timeout_in_minutes: 60
+#     commands: |
+#       echo "--- Setup Julia packages"
+#       julia --color=yes -e '
+#           using Pkg
+#           pkgs = [PackageSpec(; path) for path in (".", "lib/EnzymeCore", "lib/EnzymeTestUtils")]
+#           push!(pkgs, PackageSpec(; name="Metal"))
+#           Pkg.develop(pkgs)' || exit 3
+# 
+#       echo "+++ Run tests"
+#       julia --color=yes test/metal.jl
+#     env:
+#       JULIA_PKG_SERVER_REGISTRY_PREFERENCE: eager

.github/workflows/CI.yml

@@ -22,7 +22,7 @@ jobs:
       matrix:
         version:
           - '1.10'
-          - ~1.11.0-0
+          - '1.11'
           - 'nightly'
         os:
           - ubuntu-20.04
@@ -137,7 +137,7 @@ jobs:
       matrix:
         version:
           - '1.10'
-          - ~1.11.0-0
+          - '1.11'
           - 'nightly'
         os:
           - ubuntu-latest

.github/workflows/benchmark_pr.yml

@@ -42,7 +42,7 @@ jobs:
                 mkdir -p plots
                 benchpkgplot ${{ steps.extract-package-name.outputs.package_name }} --rev="${{github.event.repository.default_branch}},${{github.event.pull_request.head.sha}}" --npart=10 --format=png --input-dir=results/ --output-dir=plots/
             - name: Upload plot as artifact
-              uses: actions/upload-artifact@v2
+              uses: actions/upload-artifact@v4
               with:
                 name: plots
                 path: plots

Project.toml

@@ -1,7 +1,7 @@
 name = "Enzyme"
 uuid = "7da242da-08ed-463a-9acd-ee780be4f1d9"
 authors = ["William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.13.11"
+version = "0.13.12"
 
 [deps]
 CEnum = "fa961155-64e5-5f13-b03f-caf6b980ea82"
@@ -36,7 +36,7 @@ BFloat16s = "0.2, 0.3, 0.4, 0.5"
 CEnum = "0.4, 0.5"
 ChainRulesCore = "1"
 EnzymeCore = "0.8.4"
-Enzyme_jll = "0.0.154"
+Enzyme_jll = "0.0.157"
 GPUCompiler = "0.21, 0.22, 0.23, 0.24, 0.25, 0.26, 0.27, 1"
 LLVM = "6.1, 7, 8, 9"
 LogExpFunctions = "0.3"

lib/EnzymeTestUtils/Project.toml

@@ -1,7 +1,7 @@
 name = "EnzymeTestUtils"
 uuid = "12d8515a-0907-448a-8884-5fe00fdf1c5a"
 authors = ["Seth Axen <[email protected]>", "William Moses <[email protected]>", "Valentin Churavy <[email protected]>"]
-version = "0.2.0"
+version = "0.2.1"
 
 [deps]
 ConstructionBase = "187b0558-2788-49d3-abe0-74a17ed4e7c9"

lib/EnzymeTestUtils/src/finite_difference_calls.jl

@@ -29,13 +29,13 @@ function _fd_forward(fdm, f, rettype, y, activities)
     # vectorize inputs and outputs of function
     f_vec = first ∘ to_vec ∘ Base.splat(f_sig_args) ∘ from_vec_in
     if rettype <: Union{Duplicated,DuplicatedNoNeed}
-        all(ignores) && return zero_tangent(y)
+        all(ignores) && return Enzyme.make_zero(y)
         sig_arg_dval_vec, _ = to_vec(ẋs[.!ignores])
         ret_deval_vec = FiniteDifferences.jvp(fdm, f_vec,
                                               (sig_arg_val_vec, sig_arg_dval_vec))
         return from_vec_out(ret_deval_vec)
     elseif rettype <: Union{BatchDuplicated,BatchDuplicatedNoNeed}
-        all(ignores) && return (var"1"=zero_tangent(y),)
+        all(ignores) && return (var"1"=Enzyme.make_zero(y),)
         ret_dvals = map(ẋs[.!ignores]...) do sig_args_dvals...
             sig_args_dvals_vec, _ = to_vec(sig_args_dvals)
             ret_dval_vec = FiniteDifferences.jvp(fdm, f_vec,
@@ -67,13 +67,13 @@ function _fd_reverse(fdm, f, ȳ, activities, active_return)
     xs = map(x -> x.val, activities)
     ignores = map(a -> a isa Const, activities)
     f_sig_args = _wrap_reverse_function(active_return, f, xs, ignores)
-    all(ignores) && return map(zero_tangent, xs)
+    all(ignores) && return map(Enzyme.make_zero, xs)
     ignores = collect(ignores)
     is_batch = _any_batch_duplicated(map(typeof, activities)...)
     batch_size = is_batch ? _batch_size(map(typeof, activities)...) : 1
     x̄s = map(collect(activities)) do a
         if a isa Union{Const,Active}
-            dval = ntuple(_ -> zero_tangent(a.val), batch_size)
+            dval = ntuple(_ -> Enzyme.make_zero(a.val), batch_size)
             return is_batch ? dval : dval[1]
         else
             return a.dval
@@ -178,7 +178,7 @@ function _wrap_reverse_function(active_return, f, xs, ignores)
         # zero, if the input and output alias.
         if active_return
             for k in keys(zeros)
-                zeros[k] = zero_tangent(k)
+                zeros[k] = Enzyme.make_zero(k)
             end
         end
 

lib/EnzymeTestUtils/src/generate_tangent.jl

@@ -25,13 +25,6 @@ function rand_tangent(rng, x)
     return from_vec(v_new)
 end
 
-# differs from Enzyme.make_zero primarily in that reshaped Arrays in the argument will share
-# the same memory in the output.
-function zero_tangent(x)
-    v, from_vec = to_vec(x)
-    return from_vec(zero(v))
-end
-
 auto_activity(arg) = auto_activity(Random.default_rng(), arg)
 function auto_activity(rng, arg::Tuple)
     if length(arg) == 2 && arg[2] isa Type && arg[2] <: Annotation

lib/EnzymeTestUtils/src/test_reverse.jl

@@ -92,12 +92,12 @@ function test_reverse(
         y = fcopy(args_copy...; deepcopy(fkwargs)...)
         # generate tangent for output
         if !_any_batch_duplicated(ret_activity, map(typeof, activities)...)
-            ȳ = ret_activity <: Const ? zero_tangent(y) : rand_tangent(rng, y)
+            ȳ = ret_activity <: Const ? Enzyme.make_zero(y) : rand_tangent(rng, y)
         else
             batch_size = _batch_size(ret_activity, map(typeof, activities)...)
             ks = ntuple(Symbol ∘ string, batch_size)
             ȳ = ntuple(batch_size) do _
-                return ret_activity <: Const ? zero_tangent(y) : rand_tangent(y)
+                return ret_activity <: Const ? Enzyme.make_zero(y) : rand_tangent(y)
             end
         end
         # call finitedifferences, avoid mutating original arguments

lib/EnzymeTestUtils/src/to_vec.jl

@@ -89,6 +89,52 @@ function to_vec(x::RT, seen_vecs::AliasDict) where {RT<:Array}
     end
     return x_vec, Array_from_vec
 end
+
+@static if VERSION < v"1.11-"
+else
+# basic containers: loop over defined elements, recursively converting them to vectors
+function to_vec(x::RT, seen_vecs::AliasDict) where {RT<:GenericMemory}
+    has_seen = haskey(seen_vecs, x)
+    is_const = Enzyme.Compiler.guaranteed_const_nongen(RT, nothing)
+    if has_seen || is_const
+        x_vec = Float32[]
+    else
+        x_vecs = Vector{<:AbstractFloat}[]
+        from_vecs = []
+        subvec_inds = UnitRange{Int}[]
+        l = 0
+        for i in eachindex(x)
+            isassigned(x, i) || continue
+            xi_vec, xi_from_vec = to_vec(x[i], seen_vecs)
+            push!(x_vecs, xi_vec)
+            push!(from_vecs, xi_from_vec)
+            push!(subvec_inds, (l + 1):(l + length(xi_vec)))
+            l += length(xi_vec)
+        end
+        x_vec = reduce(vcat, x_vecs; init=Float32[])
+        seen_vecs[x] = x_vec
+    end
+    function Memory_from_vec(x_vec_new::Vector{<:AbstractFloat}, seen_xs::AliasDict)
+        if xor(has_seen, haskey(seen_xs, x))
+            throw(ErrorException("Arrays must be reconstructed in the same order as they are vectorized."))
+        end
+        has_seen && return reshape(seen_xs[x], size(x))
+        is_const && return x
+        x_new = typeof(x)(undef, size(x))
+        k = 1
+        for i in eachindex(x)
+            isassigned(x, i) || continue
+            xi = from_vecs[k](x_vec_new[subvec_inds[k]], seen_xs)
+            x_new[i] = xi
+            k += 1
+        end
+        seen_xs[x] = x_new
+        return x_new
+    end
+    return x_vec, Memory_from_vec
+end
+end
+
 function to_vec(x::Tuple, seen_vecs::AliasDict)
     x_vec, from_vec = to_vec(collect(x), seen_vecs)
     function Tuple_from_vec(x_vec_new::Vector{<:AbstractFloat}, seen_xs::AliasDict)

lib/EnzymeTestUtils/test/generate_tangent.jl

@@ -1,6 +1,6 @@
 using Test
 using EnzymeTestUtils
-using EnzymeTestUtils: rand_tangent, zero_tangent
+using EnzymeTestUtils: rand_tangent
 using Enzyme
 using Quaternions
 
@@ -42,30 +42,6 @@ using Quaternions
         @test y.a != x.a
     end
 
-    @testset "zero_tangent" begin
-        @test zero_tangent(1) == 1
-        @test zero_tangent(true) == true
-        @test zero_tangent(false) == false
-        @test zero_tangent(:foo) === :foo
-        @test zero_tangent("bar") === "bar"
-        @testset for T in (
-            Float32, Float64, ComplexF32, ComplexF64, QuaternionF32, QuaternionF64
-        )
-            x = randn(T)
-            @test zero_tangent(x) === zero(T)
-            y = randn(T, 5)
-            @test zero_tangent(y) == zero(y)
-            @test zero_tangent(y) isa typeof(y)
-        end
-        x = TestStruct(TestStruct(:foo, TestStruct(1, 3.0f0 + 1im)), [4.0, 5.0])
-        y = zero_tangent(x)
-        @test y.x.x == :foo
-        @test y.x.a.x == 1
-        @test y.x.a.a === zero(ComplexF32)
-        @test y.a isa Vector{Float64}
-        @test y.a == zero(x.a)
-    end
-
     @testset "auto_activity" begin
         @test EnzymeTestUtils.auto_activity((1.0, Const)) === Const(1.0)
         @test EnzymeTestUtils.auto_activity((1.0, Active)) === Active(1.0)

lib/EnzymeTestUtils/test/test_forward.jl

@@ -90,7 +90,8 @@ end
                         x = TestStruct(randn(T, 5), randn(T))
                     end
                     atol = rtol = sqrt(eps(real(T)))
-                    test_forward(fun, Tret, (x, Tx); atol, rtol)
+                    runtime_activity = TT <: TestStruct && (Tret <: Const)
+                    test_forward(fun, Tret, (x, Tx); atol, rtol, runtime_activity)
                 end
             end
         end

lib/EnzymeTestUtils/test/test_reverse.jl

@@ -131,13 +131,7 @@ end
                          Tx in (Const, Duplicated, BatchDuplicated)
 
                 are_activities_compatible(Tret, Tx) || continue
-                if Tx <: Const
-                    test_reverse(f, Tret, (x, Tx))
-                else
-                    @test_broken !fails() do
-                        return test_reverse(f, Tret, (x, Tx))
-                    end
-                end
+                test_reverse(f, Tret, (x, Tx))
             end
         end
 

src/Enzyme.jl

@@ -1559,7 +1559,7 @@ end
         Base.@_inline_meta
         ntuple(Val(N)) do idx
             Base.@_inline_meta
-            return (i == idx) ? 1.0 : 0.0
+            return (i == idx) ? T(1) : T(0)
         end
     end
 end
@@ -1571,7 +1571,7 @@ end
         Base.@_inline_meta
         ntuple(Val(N)) do idx
             Base.@_inline_meta
-            return (i + start - 1 == idx) ? 1.0 : 0.0
+            return (i + start - 1 == idx) ? T(1) : T(0)
         end
     end
 end