Merge pull request #199 from omlins/ad_deferred

Extend AD wrappers

Author: omlins

src/AD.jl

@@ -7,8 +7,8 @@ Provides GPU-compatible wrappers for automatic differentiation functions of the
     import ParallelStencil.AD
 
 # Functions
-- `autodiff_deferred!(mode, f, args...)`: wraps function `Enzyme.autodiff_deferred`, calling it as `Enzyme.autodiff_deferred(mode, f, Const, args...)`, promoting all arguments of `args` and `f` that are not `Enzyme.Annotation` to `Enzyme.Const`. Important: the return type activity is automatically inserted as 3rd argument as it is always `Enzyme.Const`, given that `@parallel` and `@parallel_indices` functions must return nothing.
-- `autodiff_deferred_thunk!(mode, f, args...)`: wraps function `Enzyme.autodiff_deferred_thunk`, calling it as `Enzyme.autodiff_deferred_thunk(mode, f, Const, args...)`, promoting all arguments of `args` and `f` that are not `Enzyme.Annotation` to `Enzyme.Const`. Important: the return type activity is automatically inserted as 3rd argument as it is always `Enzyme.Const`, given that `@parallel` and `@parallel_indices` functions must return nothing.
+- `autodiff_deferred!`: wraps function `Enzyme.autodiff_deferred`, and supports the same arguments. Additionally, it promotes all arguments that are not `Enzyme.Annotation` to `Enzyme.Const` and automatically inserts the return type activity as 3rd argument if omitted; as a result, the function can be called also as, e.g., `autodiff_deferred!(Enzyme.Reverse, f!,...)`, besides the fully explicit form, e.g., `autodiff_deferred!(Enzyme.Reverse, Enzyme.Const(f!), Enzyme.Const,...)`.
+- `autodiff_deferred_thunk!`: wraps function `Enzyme.autodiff_deferred_thunk`, and supports the same arguments. Additionally, it applies the same argument promotion and return type activity insertion as `autodiff_deferred!` (see above).
 
 # Examples
     const USE_GPU = true

src/ParallelKernel/EnzymeExt/AD.jl

@@ -7,8 +7,8 @@ Provides GPU-compatible wrappers for automatic differentiation functions of the
     import ParallelKernel.AD
 
 # Functions
-- `autodiff_deferred!(mode, f, args...)`: wraps function `Enzyme.autodiff_deferred`, calling it as `Enzyme.autodiff_deferred(mode, f, Const, args...)`, promoting all arguments of `args` and `f` that are not `Enzyme.Annotation` to `Enzyme.Const`. Important: the return type activity is automatically inserted as 3rd argument as it is always `Enzyme.Const`, given that `@parallel` and `@parallel_indices` functions must return nothing.
-- `autodiff_deferred_thunk!(mode, f, args...)`: wraps function `Enzyme.autodiff_deferred_thunk`, calling it as `Enzyme.autodiff_deferred_thunk(mode, f, Const, args...)`, promoting all arguments of `args` and `f` that are not `Enzyme.Annotation` to `Enzyme.Const`. Important: the return type activity is automatically inserted as 3rd argument as it is always `Enzyme.Const`, given that `@parallel` and `@parallel_indices` functions must return nothing.
+- `autodiff_deferred!`: wraps function `Enzyme.autodiff_deferred`, and supports the same arguments. Additionally, it promotes all arguments that are not `Enzyme.Annotation` to `Enzyme.Const` and automatically inserts the return type activity as 3rd argument if omitted; as a result, the function can be called also as, e.g., `autodiff_deferred!(Enzyme.Reverse, f!,...)`, besides the fully explicit form, e.g., `autodiff_deferred!(Enzyme.Reverse, Enzyme.Const(f!), Enzyme.Const,...)`.
+- `autodiff_deferred_thunk!`: wraps function `Enzyme.autodiff_deferred_thunk`, and supports the same arguments. Additionally, it applies the same argument promotion and return type activity insertion as `autodiff_deferred!` (see above).
 
 To see a description of a function type `?<functionname>`.
 """

src/ParallelKernel/EnzymeExt/autodiff_gpu.jl

@@ -25,19 +25,35 @@ function promote_to_const(args::Vararg{Any,N}) where N
     end
 end
 
- function ParallelStencil.ParallelKernel.AD.autodiff_deferred!(mode, f, args::Vararg{Any,N}) where N # NOTE: minimal specialization is required to avoid overwriting the default method
-     f    = promote_to_const(f)[1]
-     args = promote_to_const(args...)
-     Enzyme.autodiff_deferred(mode, f, Enzyme.Const, args...)
-     return
- end
- 
- function ParallelStencil.ParallelKernel.AD.autodiff_deferred_thunk!(mode, f, args::Vararg{Any,N}) where N # NOTE: minimal specialization is required to avoid overwriting the default method
-     f    = promote_to_const(f)[1]
-     args = promote_to_const(args...)
-     Enzyme.autodiff_deferred_thunk(mode, f, Enzyme.Const, args...)
-     return
- end
+
+function ParallelStencil.ParallelKernel.AD.autodiff_deferred!(mode, f, ::Type{T}, args::Vararg{Any,N}) where {T<:Enzyme.Annotation, N} # NOTE: minimal specialization is required to avoid overwriting the default method
+    f    = promote_to_const(f)[1]
+    args = promote_to_const(args...)
+    Enzyme.autodiff_deferred(mode, f, T, args...)
+    return
+end
+
+function ParallelStencil.ParallelKernel.AD.autodiff_deferred!(mode, f, args::Vararg{Any,N}) where N # NOTE: minimal specialization is required to avoid overwriting the default method
+    f    = promote_to_const(f)[1]
+    args = promote_to_const(args...)
+    Enzyme.autodiff_deferred(mode, f, Enzyme.Const, args...)
+    return
+end
+
+
+function ParallelStencil.ParallelKernel.AD.autodiff_deferred_thunk!(mode, f, ::Type{T}, args::Vararg{Any,N}) where {T<:Enzyme.Annotation, N} # NOTE: minimal specialization is required to avoid overwriting the default method
+    f    = promote_to_const(f)[1]
+    args = promote_to_const(args...)
+    Enzyme.autodiff_deferred_thunk(mode, f, T, args...)
+    return
+end
+
+function ParallelStencil.ParallelKernel.AD.autodiff_deferred_thunk!(mode, f, args::Vararg{Any,N}) where N # NOTE: minimal specialization is required to avoid overwriting the default method
+    f    = promote_to_const(f)[1]
+    args = promote_to_const(args...)
+    Enzyme.autodiff_deferred_thunk(mode, f, Enzyme.Const, args...)
+    return
+end
 
 
 ## FUNCTIONS TO CHECK EXTENSIONS SUPPORT

test/ParallelKernel/test_kernel_language.jl

@@ -57,7 +57,7 @@ eval(:(
                     @test @prettystring(1, @sharedMem($FloatDefault, (2,3))) == "CUDA.@cuDynamicSharedMem $(nameof($FloatDefault)) (2, 3)"
                     # @test @prettystring(1, @pk_show()) == "CUDA.@cushow"
                     # @test @prettystring(1, @pk_println()) == "CUDA.@cuprintln"
-                elseif $package == $AMDGPU
+                elseif $package == $PKG_AMDGPU
                     @test @prettystring(1, @gridDim()) == "AMDGPU.gridGroupDim()"
                     @test @prettystring(1, @blockIdx()) == "AMDGPU.workgroupIdx()"
                     @test @prettystring(1, @blockDim()) == "AMDGPU.workgroupDim()"
@@ -208,14 +208,16 @@ eval(:(
                         Bshfl_up    = similar(A)
                         Bshfl_down  = similar(A)
                         Bshfl_xor   = similar(A)
+                        Bwarpsize   = Vector{Int}(undef, N)
+                        Blaneid     = Vector{Int}(undef, N)
 
-                        @parallel_indices (ix) function kernel_semantics!(Bout_any, Bout_all, Bout_ballot, Bshfl, Bshfl_up, Bshfl_down, Bshfl_xor, A, P)
+                        @parallel_indices (ix) function kernel_semantics!(Bout_any, Bout_all, Bout_ballot, Bshfl, Bshfl_up, Bshfl_down, Bshfl_xor, Bwarpsize, Blaneid, A, P)
                             m = @active_mask()
                             w = @warpsize()
                             l = @laneid()
-                            # basic invariants under CPU model
-                            @test w == 1
-                            @test l == 1
+                            # store values for verification outside kernel
+                            Bwarpsize[ix] = w
+                            Blaneid[ix] = l
                             # shuffle identities
                             Bshfl[ix]      = @shfl_sync(m, A[ix], l)
                             Bshfl_up[ix]   = @shfl_up_sync(m, A[ix], 1)
@@ -228,8 +230,11 @@ eval(:(
                             Bout_ballot[ix] = @vote_ballot_sync(m, pa)
                             return
                         end
-                        @parallel (1:N) kernel_semantics!(Bout_any, Bout_all, Bout_ballot, Bshfl, Bshfl_up, Bshfl_down, Bshfl_xor, A, P)
+                        @parallel (1:N) kernel_semantics!(Bout_any, Bout_all, Bout_ballot, Bshfl, Bshfl_up, Bshfl_down, Bshfl_xor, Bwarpsize, Blaneid, A, P)
 
+                        # basic invariants under CPU model
+                        @test all(Bwarpsize .== 1)
+                        @test all(Blaneid .== 1)
                         @test all(Bshfl .== A)
                         @test all(Bshfl_up .== A)
                         @test all(Bshfl_down .== A)

test/ParallelKernel/test_parallel.jl

@@ -103,16 +103,16 @@ eval(:(
                     @test maxsize(BitstypeStruct(5, 6.0), 8, (x=[9 9; 9 9; 9 9], y=[9 9; 9 9; 9 9]), (x=[7 7 7; 7 7 7], y=[7 7 7; 7 7 7])) == (3, 3, 1)
                 end;
             end;
-            @testset "@parallel ∇" begin
-                @test @prettystring(1, @parallel ∇=B->B̄ f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Reverse, f!, Enzyme.Const(A), Enzyme.DuplicatedNoNeed(B, B̄), Enzyme.Const(a))"
-                @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Reverse, f!, Enzyme.DuplicatedNoNeed(A, Ā), Enzyme.DuplicatedNoNeed(B, B̄), Enzyme.Const(a))"
-                @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) ad_mode=Enzyme.Forward f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Forward, f!, Enzyme.DuplicatedNoNeed(A, Ā), Enzyme.DuplicatedNoNeed(B, B̄), Enzyme.Const(a))"
-                @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) ad_mode=Enzyme.Forward ad_annotations=(Duplicated=B) f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Forward, f!, Enzyme.DuplicatedNoNeed(A, Ā), Enzyme.Duplicated(B, B̄), Enzyme.Const(a))"
-                @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) ad_mode=Enzyme.Forward ad_annotations=(Duplicated=(B,A), Active=b) f!(A, B, a, b)) == "@parallel configcall = f!(A, B, a, b) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Forward, f!, Enzyme.Duplicated(A, Ā), Enzyme.Duplicated(B, B̄), Enzyme.Const(a), Enzyme.Active(b))"
-                @test @prettystring(1, @parallel ∇=(V.x->V̄.x, V.y->V̄.y) f!(V.x, V.y, a)) == "@parallel configcall = f!(V.x, V.y, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Reverse, f!, Enzyme.DuplicatedNoNeed(V.x, V̄.x), Enzyme.DuplicatedNoNeed(V.y, V̄.y), Enzyme.Const(a))"
-            end;
-            @testset "AD.autodiff_deferred! | @parallel ∇ (numerical)" begin
-                @static if $package == $PKG_THREADS
+            @static if $package != $PKG_POLYESTER # Enzyme does not support Polyester.
+                @testset "@parallel ∇" begin
+                    @test @prettystring(1, @parallel ∇=B->B̄ f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Reverse, f!, Enzyme.Const(A), Enzyme.DuplicatedNoNeed(B, B̄), Enzyme.Const(a))"
+                    @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Reverse, f!, Enzyme.DuplicatedNoNeed(A, Ā), Enzyme.DuplicatedNoNeed(B, B̄), Enzyme.Const(a))"
+                    @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) ad_mode=Enzyme.Forward f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Forward, f!, Enzyme.DuplicatedNoNeed(A, Ā), Enzyme.DuplicatedNoNeed(B, B̄), Enzyme.Const(a))"
+                    @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) ad_mode=Enzyme.Forward ad_annotations=(Duplicated=B) f!(A, B, a)) == "@parallel configcall = f!(A, B, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Forward, f!, Enzyme.DuplicatedNoNeed(A, Ā), Enzyme.Duplicated(B, B̄), Enzyme.Const(a))"
+                    @test @prettystring(1, @parallel ∇=(A->Ā, B->B̄) ad_mode=Enzyme.Forward ad_annotations=(Duplicated=(B,A), Active=b) f!(A, B, a, b)) == "@parallel configcall = f!(A, B, a, b) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Forward, f!, Enzyme.Duplicated(A, Ā), Enzyme.Duplicated(B, B̄), Enzyme.Const(a), Enzyme.Active(b))"
+                    @test @prettystring(1, @parallel ∇=(V.x->V̄.x, V.y->V̄.y) f!(V.x, V.y, a)) == "@parallel configcall = f!(V.x, V.y, a) ParallelStencil.ParallelKernel.AD.autodiff_deferred!(Enzyme.Reverse, f!, Enzyme.DuplicatedNoNeed(V.x, V̄.x), Enzyme.DuplicatedNoNeed(V.y, V̄.y), Enzyme.Const(a))"
+                end;
+                @testset "ad numerical" begin
                     N = 16
                     a = 6.5
                     A = @rand(N)
@@ -133,16 +133,30 @@ eval(:(
                         end
                         return
                     end
-                    @parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, DuplicatedNoNeed(A, Ā), DuplicatedNoNeed(B, B̄), Const(a)) # NOTE: f! is automatically promoted to Const.
                     Enzyme.autodiff_deferred(Enzyme.Reverse, Const(g!), Const, DuplicatedNoNeed(A_ref, Ā_ref), DuplicatedNoNeed(B_ref, B̄_ref), Const(a))
-                    @test Array(Ā) ≈ Ā_ref
-                    @test Array(B̄) ≈ B̄_ref
-                    Ā = @ones(N)
-                    B̄ = @ones(N)
-                    @parallel ∇=(A->Ā, B->B̄) f!(A, B, a)
-                    @test Array(Ā) ≈ Ā_ref
-                    @test Array(B̄) ≈ B̄_ref
-                end
+                    @testset "AD.autodiff_deferred!" begin
+                        @parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, DuplicatedNoNeed(A, Ā), DuplicatedNoNeed(B, B̄), Const(a)) # NOTE: f! is automatically promoted to Const(f!) and the return type Const is inserted.
+                        @test Array(Ā) ≈ Ā_ref
+                        @test Array(B̄) ≈ B̄_ref
+                        Ā = @ones(N)
+                        B̄ = @ones(N)
+                        @parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, Const, DuplicatedNoNeed(A, Ā), DuplicatedNoNeed(B, B̄), Const(a)) # NOTE: f! is automatically promoted to Const(f!)
+                        @test Array(Ā) ≈ Ā_ref
+                        @test Array(B̄) ≈ B̄_ref
+                        Ā = @ones(N)
+                        B̄ = @ones(N)
+                        @parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, Const(f!), Const, DuplicatedNoNeed(A, Ā), DuplicatedNoNeed(B, B̄), Const(a)) # NOTE: no automatic promotion or insertion here.
+                        @test Array(Ā) ≈ Ā_ref
+                        @test Array(B̄) ≈ B̄_ref
+                    end;
+                    @testset "@parallel ∇ (numerical)" begin
+                        Ā = @ones(N)
+                        B̄ = @ones(N)
+                        @parallel ∇=(A->Ā, B->B̄) f!(A, B, a) # NOTE: expands to the same as above
+                        @test Array(Ā) ≈ Ā_ref
+                        @test Array(B̄) ≈ B̄_ref
+                    end;
+                end;
             end;
             @testset "@parallel_indices" begin
                 @testset "inbounds" begin

test/runtests.jl

@@ -15,8 +15,10 @@ function runtests(testfiles=String[])
     istest(f) = endswith(f, ".jl") && startswith(basename(f), "test_")
     testfiles = isempty(testfiles) ? sort(filter(istest, vcat([joinpath.(root, files) for (root, dirs, files) in walkdir(testdir)]...))) : testfiles
 
-    nerror = 0
-    errorfiles = String[]
+    nabort = 0
+    nfail  = 0
+    abortfiles = String[]
+    failfiles  = String[]
     printstyled("Testing package ParallelStencil.jl\n"; bold=true, color=:white)
 
     if (PKG_CUDA in SUPPORTED_PACKAGES && !CUDA.functional())
@@ -38,25 +40,28 @@ function runtests(testfiles=String[])
             println("$f")
             continue
         end
-        cmd = `$exename -O3 --startup-file=no $(joinpath(testdir, f))`
+        cmd = `$exename --color=yes -O3 --startup-file=no $(joinpath(testdir, f))`
         stdout_path = tempname()
         stderr_path = tempname()
         stdout_content = ""
         stderr_content = ""
+        proc = nothing
         try
             open(stdout_path, "w") do stdout_io
                 open(stderr_path, "w") do stderr_io
-                    proc = run(pipeline(Cmd(cmd; ignorestatus=true), stdout=stdout_io, stderr=stderr_io); wait=false)
-                    wait(proc)
+                    proc = run(pipeline(Cmd(cmd; ignorestatus=true), stdout=stdout_io, stderr=stderr_io); wait=true)
                 end
             end
             stdout_content = read(stdout_path, String)
             stderr_content = read(stderr_path, String)
             print(stdout_content)
             print(Base.stderr, stderr_content)
         catch ex
-            println("Test Error: an exception occurred while running the test file $f :")
+            println("Test Abort: a system-level exception occurred while running the test file $f :")
             println(ex)
+            nabort += 1
+            push!(abortfiles, f)
+            continue
         finally
             if ispath(stdout_path)
                 rm(stdout_path; force=true)
@@ -66,21 +71,32 @@ function runtests(testfiles=String[])
             end
         end
         if !occursin(r"(?i)test summary", stdout_content)
-            nerror += 1
-            push!(errorfiles, f)
+            nabort += 1
+            push!(abortfiles, f)
+        elseif proc !== nothing && !success(proc)
+            nfail += 1
+            push!(failfiles, f)
         end
     end
     println("")
-    if nerror == 0
-        printstyled("Test suite: all selected test files executed (see above for results).\n"; bold=true, color=:green)
+    if nabort == 0 && nfail == 0
+        printstyled("Test suite: all selected test files executed and all tests passed.\n"; bold=true, color=:green)
     else
-        printstyled("Test suite: $nerror test file(s) aborted execution due to error (see above for details); files aborting execution:\n"; bold=true, color=:red)
-        for f in errorfiles
-            println(" - $f")
+        if nfail > 0
+            printstyled("Test suite: $nfail test files(s) have tests that failed or errored (see above for results); files with failed/errored tests:\n"; bold=true, color=:red)
+            for f in failfiles
+                println(" - $f")
+            end
+        end
+        if nabort > 0
+            printstyled("Test suite: $nabort test file(s) aborted execution due to fatal error (see above for details); files aborting execution:\n"; bold=true, color=:red)
+            for f in abortfiles
+                println(" - $f")
+            end
         end
     end
     println("")
-    return nerror
+    return nabort+nfail
 end
 
 exit(runtests(ARGS))