Merge pull request #202 from omlins/ad_deferred

AD: throw error if gpu compiler incompatible argument

Author: omlins

README.md

@@ -16,21 +16,31 @@ A particularity of ParallelStencil is that it enables writing a single high-leve
 Beyond traditional high-performance computing, ParallelStencil supports automatic differentiation of architecture-agnostic parallel kernels relying on [Enzyme.jl], enabling both high-level and generic syntax for maximal flexibility.
 
 ## Contents
-* [Parallelization and optimization with one macro call](#parallelization-with-one-macro-call)
-* [Stencil computations with math-close notation](#stencil-computations-with-math-close-notation)
-* [50-lines example deployable on GPU and CPU](#50-lines-example-deployable-on-GPU-and-CPU)
-* [50-lines multi-xPU example](#50-lines-multi-xpu-example)
-* [Seamless interoperability with communication packages and hiding communication](#seamless-interoperability-with-communication-packages-and-hiding-communication)
-* [Support for architecture-agnostic low level kernel programming](#support-for-architecture-agnostic-low-level-kernel-programming)
-* [Support for logical arrays of small arrays / structs](#support-for-logical-arrays-of-small-arrays--structs)
-* [Support for automatic differentiation of architecture-agnostic parallel kernels](#support-for-automatic-differentiation-of-architecture-agnostic-parallel-kernels)
-* [Module documentation callable from the Julia REPL / IJulia](#module-documentation-callable-from-the-julia-repl--ijulia)
-* [Concise single/multi-xPU miniapps](#concise-singlemulti-xpu-miniapps)
-* [Dependencies](#dependencies)
-* [Installation](#installation)
-* [Questions, comments and discussions](#questions-comments-and-discussions)
-* [Your contributions](#your-contributions)
-* [References](#references)
+- [Contents](#contents)
+- [Parallelization and optimization with one macro call](#parallelization-and-optimization-with-one-macro-call)
+- [Stencil computations with math-close notation](#stencil-computations-with-math-close-notation)
+- [50-lines example deployable on GPU and CPU](#50-lines-example-deployable-on-gpu-and-cpu)
+- [50-lines multi-xPU example](#50-lines-multi-xpu-example)
+- [Seamless interoperability with communication packages and hiding communication](#seamless-interoperability-with-communication-packages-and-hiding-communication)
+- [Support for architecture-agnostic low level kernel programming](#support-for-architecture-agnostic-low-level-kernel-programming)
+- [Support for logical arrays of small arrays / structs](#support-for-logical-arrays-of-small-arrays--structs)
+- [Support for automatic differentiation of architecture-agnostic parallel kernels](#support-for-automatic-differentiation-of-architecture-agnostic-parallel-kernels)
+- [Module documentation callable from the Julia REPL / IJulia](#module-documentation-callable-from-the-julia-repl--ijulia)
+- [Concise single/multi-xPU miniapps](#concise-singlemulti-xpu-miniapps)
+    - [Performance metric](#performance-metric)
+    - [Miniapp content](#miniapp-content)
+    - [Thermo-mechanical convection 2-D app](#thermo-mechanical-convection-2-d-app)
+    - [Viscous Stokes 2-D app](#viscous-stokes-2-d-app)
+    - [Viscous Stokes 3-D app](#viscous-stokes-3-d-app)
+    - [Acoustic wave 2-D app](#acoustic-wave-2-d-app)
+    - [Acoustic wave 3-D app](#acoustic-wave-3-d-app)
+    - [Scalar porosity waves 2-D app](#scalar-porosity-waves-2-d-app)
+    - [Hydro-mechanical porosity waves 2-D app](#hydro-mechanical-porosity-waves-2-d-app)
+- [Dependencies](#dependencies)
+- [Installation](#installation)
+- [Questions, comments and discussions](#questions-comments-and-discussions)
+- [Your contributions](#your-contributions)
+- [References](#references)
 
 ## Parallelization and optimization with one macro call
 A simple call to `@parallel` is enough to parallelize and optimize a function and to launch it. The package used underneath for parallelization is defined in a call to `@init_parallel_stencil` beforehand. Supported are [CUDA.jl], [AMDGPU.jl] and [Metal.jl] for running on GPU and [Base.Threads] for CPU. The following example outlines how to run parallel computations on a GPU using the native kernel programming capabilities of [CUDA.jl] underneath (omitted lines are represented with `#(...)`, omitted arguments with `...`):
@@ -318,7 +328,7 @@ import ParallelStencil.AD
 using Enzyme
 #(...)
 @parallel f!(A, B, a)                                                                                                                 # normal call of f!
-@parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, DuplicatedNoNeed(A, Ā), DuplicatedNoNeed(B, B̄), Const(a))  # call to the gradient of f!, differentiated with respect to A and B
+@parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, DuplicatedNoNeed(A, Ā), DuplicatedNoNeed(B, B̄), a)  # call to the gradient of f!, differentiated with respect to A and B
 ```
 The submodule `ParallelStencil.AD` contains GPU-compatible wrappers of Enzyme functions (returning always `nothing` as required by the backend packages [CUDA.jl] and [AMDGPU.jl]); the wrapper `AD.autodiff_deferred!` maps, e.g., `Enzyme.autodiff_deferred`. The keyword argument `configcall` makes it trivial to call these generic functions for automatic differentiation with the right launch parameters.
 

src/AD.jl

@@ -7,8 +7,9 @@ Provides GPU-compatible wrappers for automatic differentiation functions of the
     import ParallelStencil.AD
 
 # Functions
-- `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).
+- `autodiff_deferred!`: wraps function `Enzyme.autodiff_deferred`, and supports the same arguments, with the exception of the return type activity (3rd argument) which must be omitted and will be automatically inserted (inserting it explicitly will raise a GPU compiler error when targeting a GPU). Additionally, it promotes all arguments that are not `Enzyme.Annotation` to `Enzyme.Const`. As a result, the function can be conveniently called as, e.g., `autodiff_deferred!(Enzyme.Reverse, f!,...)`, instead of the fully explicit form, e.g., `autodiff_deferred!(Enzyme.Reverse, Enzyme.Const(f!), Enzyme.Const,...)` (which is not supported).
+- `autodiff_deferred_thunk!`: wraps function `Enzyme.autodiff_deferred_thunk`, in the same way as `autodiff_deferred!` wraps `Enzyme.autodiff_deferred` (see above).
+
 
 # Examples
     const USE_GPU = true
@@ -36,7 +37,7 @@ Provides GPU-compatible wrappers for automatic differentiation functions of the
 
         @info "running on CPU/GPU"
         @parallel f!(A, B, a) # normal call of f!
-        @parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, Duplicated(A, Ā), DuplicatedNoNeed(B, B̄), Const(a)) # automatic differentiation of f!
+        @parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, Duplicated(A, Ā), DuplicatedNoNeed(B, B̄), a) # automatic differentiation of f!
         
         return
     end

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!`: 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).
+- `autodiff_deferred!`: wraps function `Enzyme.autodiff_deferred`, and supports the same arguments, with the exception of the return type activity (3rd argument) which must be omitted and will be automatically inserted (inserting it explicitly will raise a GPU compiler error when targeting a GPU). Additionally, it promotes all arguments that are not `Enzyme.Annotation` to `Enzyme.Const`. As a result, the function can be conveniently called as, e.g., `autodiff_deferred!(Enzyme.Reverse, f!,...)`, instead of the fully explicit form, e.g., `autodiff_deferred!(Enzyme.Reverse, Enzyme.Const(f!), Enzyme.Const,...)` (which is not supported).
+- `autodiff_deferred_thunk!`: wraps function `Enzyme.autodiff_deferred_thunk`, in the same way as `autodiff_deferred!` wraps `Enzyme.autodiff_deferred` (see above).
 
 To see a description of a function type `?<functionname>`.
 """

src/ParallelKernel/EnzymeExt/autodiff_gpu.jl

@@ -1,5 +1,6 @@
 import ParallelStencil
 import ParallelStencil: PKG_THREADS, PKG_POLYESTER
+using ParallelStencil.ParallelKernel.Exceptions
 import Enzyme
 
 # NOTE: package specific initialization of Enzyme could be done as follows (not needed in the currently supported versions of Enzyme)
@@ -27,9 +28,7 @@ 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...)
+    @ArgumentError("AD.autodiff_deferred!: explicit insertion of the return type activity (third argument) is not supported as not GPU compiler compatible; call without it instead.") 
     return
 end
 
@@ -42,9 +41,7 @@ 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...)
+    @ArgumentError("AD.autodiff_deferred_thunk!: explicit insertion of the return type activity (third argument) is not supported as not GPU compiler compatible; call without it instead.")
     return
 end
 

test/ParallelKernel/test_parallel.jl

@@ -140,15 +140,14 @@ eval(:(
                         @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.
+                        @parallel configcall=f!(A, B, a) AD.autodiff_deferred!(Enzyme.Reverse, f!, DuplicatedNoNeed(A, Ā), DuplicatedNoNeed(B, B̄), a) # NOTE: f! and a are automatically promoted to Const(f!) and Const(a) and the return type Const is inserted.
                         @test Array(Ā) ≈ Ā_ref
                         @test Array(B̄) ≈ B̄_ref
                     end;
+                    @testset "AD.autodiff_deferred! (GPU compiler error)" begin
+                        @test_throws Exception @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_throws Exception @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.
+                    end;
                     @testset "@parallel ∇ (numerical)" begin
                         Ā = @ones(N)
                         B̄ = @ones(N)