enable fully explicit AD method

Author: omlins

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