EnzymeCoreExt: overlay autodiff/deferred_codegen for kernel codegen

[gbaraldi]

Companion to EnzymeAD/Enzyme.jl#3112 — addresses EnzymeAD/Enzyme.jl#3091.

Summary

Two new overlays in `CUDA.method_table` inside `ext/EnzymeCoreExt.jl`:

1. `EnzymeCore.deferred_codegen(id::UInt) -> Ptr{Cvoid}` — the body holds the `ccall("extern deferred_codegen", llvmcall, …)` marker that GPUCompiler's scanner picks up to link the inner Enzyme adjoint into the kernel module. EnzymeCore declares `deferred_codegen` as a method-less generic function so that `jl_compile_all_defs` can't despecialize the llvmcall into a host sysimage. CUDA's overlay supplies the body during kernel compilation; Enzyme.jl provides a matching overlay in `GPUCompiler.GLOBAL_METHOD_TABLE` for its own interpreter.

2. `EnzymeCore.autodiff(::ReverseMode|::ForwardMode, …)` redirect — inside CUDA's `GPUInterpreter` this forwards to `autodiff_deferred`, so user kernels can write plain `Enzyme.autodiff(...)` and get the deferred-codegen path automatically. Removes the manual `autodiff` vs `autodiff_deferred` choice that KernelAbstractions's `gpu_fwd` makes today.

The overlays only use EnzymeCore types, so no new weakdep is needed.

Test plan

• CI green
• Verified locally with a CUDA kernel calling plain `Enzyme.autodiff(Reverse, kernel, ...)` — gradient comes out correct
• Existing `test/extensions/enzyme.jl` keeps passing

[github-actions]

CUDA.jl Benchmarks

Details

Benchmark suite	Current: e874f57	Previous: d96bb17	Ratio
array/accumulate/Float32/1d	101249 ns	99804 ns	1.01
array/accumulate/Float32/dims=1	77191 ns	74695 ns	1.03
array/accumulate/Float32/dims=1L	1586424.5 ns	1575172 ns	1.01
array/accumulate/Float32/dims=2	144038 ns	140050 ns	1.03
array/accumulate/Float32/dims=2L	658149 ns	652746 ns	1.01
array/accumulate/Int64/1d	118341 ns	116667 ns	1.01
array/accumulate/Int64/dims=1	80145 ns	78138 ns	1.03
array/accumulate/Int64/dims=1L	1706952 ns	1683352 ns	1.01
array/accumulate/Int64/dims=2	156969 ns	151632 ns	1.04
array/accumulate/Int64/dims=2L	962119 ns	958723 ns	1.00
array/broadcast	20427 ns	19844 ns	1.03
array/construct	1239 ns	1172.2 ns	1.06
array/copy	17978 ns	16637 ns	1.08
array/copyto!/cpu_to_gpu	215916 ns	213824 ns	1.01
array/copyto!/gpu_to_cpu	284942 ns	278790 ns	1.02
array/copyto!/gpu_to_gpu	10826 ns	10365 ns	1.04
array/iteration/findall/bool	134850 ns	130581 ns	1.03
array/iteration/findall/int	148813.5 ns	144724 ns	1.03
array/iteration/findfirst/bool	82415.5 ns	79128 ns	1.04
array/iteration/findfirst/int	84072.5 ns	80847 ns	1.04
array/iteration/findmin/1d	82650 ns	66285 ns	1.25
array/iteration/findmin/2d	114182 ns	101434 ns	1.13
array/iteration/logical	201744 ns	188254 ns	1.07
array/iteration/scalar	66299 ns	64586 ns	1.03
array/permutedims/2d	52586 ns	49748 ns	1.06
array/permutedims/3d	52731 ns	50091 ns	1.05
array/permutedims/4d	53030.5 ns	49706 ns	1.07
array/random/rand/Float32	12572 ns	12116 ns	1.04
array/random/rand/Int64	36650 ns	23371 ns	1.57
array/random/rand!/Float32	8443.333333333334 ns	7960.333333333333 ns	1.06
array/random/rand!/Int64	34121 ns	20391 ns	1.67
array/random/randn/Float32	37407 ns	34710 ns	1.08
array/random/randn!/Float32	30936 ns	23894 ns	1.29
array/reductions/mapreduce/Float32/1d	33335.5 ns	32647 ns	1.02
array/reductions/mapreduce/Float32/dims=1	39372 ns	37819 ns	1.04
array/reductions/mapreduce/Float32/dims=1L	51253 ns	50203 ns	1.02
array/reductions/mapreduce/Float32/dims=2	56414 ns	55479 ns	1.02
array/reductions/mapreduce/Float32/dims=2L	69109 ns	67023 ns	1.03
array/reductions/mapreduce/Int64/1d	42869 ns	39410 ns	1.09
array/reductions/mapreduce/Int64/dims=1	51434 ns	40857 ns	1.26
array/reductions/mapreduce/Int64/dims=1L	87004 ns	86396 ns	1.01
array/reductions/mapreduce/Int64/dims=2	59396.5 ns	57852 ns	1.03
array/reductions/mapreduce/Int64/dims=2L	84511 ns	82413 ns	1.03
array/reductions/reduce/Float32/1d	34084 ns	32758 ns	1.04
array/reductions/reduce/Float32/dims=1	48998 ns	38011 ns	1.29
array/reductions/reduce/Float32/dims=1L	51206 ns	50349 ns	1.02
array/reductions/reduce/Float32/dims=2	56383 ns	55642 ns	1.01
array/reductions/reduce/Float32/dims=2L	69066 ns	67383 ns	1.02
array/reductions/reduce/Int64/1d	42820 ns	39537 ns	1.08
array/reductions/reduce/Int64/dims=1	41986 ns	40656 ns	1.03
array/reductions/reduce/Int64/dims=1L	86968 ns	86280 ns	1.01
array/reductions/reduce/Int64/dims=2	59257 ns	57865 ns	1.02
array/reductions/reduce/Int64/dims=2L	83898 ns	82279 ns	1.02
array/reverse/1d	17704 ns	16600 ns	1.07
array/reverse/1dL	68278 ns	67647 ns	1.01
array/reverse/1dL_inplace	65712 ns	65318 ns	1.01
array/reverse/1d_inplace	10193.5 ns	8340.333333333334 ns	1.22
array/reverse/2d	21246 ns	19869 ns	1.07
array/reverse/2dL	73393.5 ns	71763 ns	1.02
array/reverse/2dL_inplace	65729 ns	65099 ns	1.01
array/reverse/2d_inplace	10213 ns	9659 ns	1.06
array/sorting/1d	2735901 ns	2713733 ns	1.01
array/sorting/2d	1068724.5 ns	1062891 ns	1.01
array/sorting/by	3303971 ns	3280831 ns	1.01
cuda/synchronization/context/auto	1154.9 ns	1138.4 ns	1.01
cuda/synchronization/context/blocking	938.974358974359 ns	962.5652173913044 ns	0.98
cuda/synchronization/context/nonblocking	7578.6 ns	5944.2 ns	1.27
cuda/synchronization/stream/auto	1013.3 ns	997.8333333333334 ns	1.02
cuda/synchronization/stream/blocking	824.1359223300971 ns	821.8235294117648 ns	1.00
cuda/synchronization/stream/nonblocking	7164.4 ns	5912.4 ns	1.21
integration/byval/reference	143670 ns	143240 ns	1.00
integration/byval/slices=1	145773 ns	145261 ns	1.00
integration/byval/slices=2	284454 ns	283622 ns	1.00
integration/byval/slices=3	423151 ns	421946 ns	1.00
integration/cudadevrt	102401 ns	101870 ns	1.01
integration/volumerhs	9960003 ns	9906591 ns	1.01
kernel/indexing	13275 ns	12574 ns	1.06
kernel/indexing_checked	14035 ns	13327 ns	1.05
kernel/launch	2156.1111111111113 ns	2037.111111111111 ns	1.06
kernel/occupancy	677.404458598726 ns	696.7364864864865 ns	0.97
kernel/rand	18050.5 ns	14332 ns	1.26
latency/import	3548878692 ns	3854422204 ns	0.92
latency/precompile	45248138979 ns	4627097887 ns	9.78
latency/ttfp	13051637302 ns	4514039137 ns	2.89

This comment was automatically generated by workflow using github-action-benchmark.