Add EnzymeCore extension for parent_job (#2281)

Co-authored-by: Tim Besard <tim.besard@gmail.com>

Author: vchuravy

Project.toml

@@ -12,6 +12,7 @@ CUDA_Runtime_Discovery = "1af6417a-86b4-443c-805f-a4643ffb695f"
 CUDA_Runtime_jll = "76a88914-d11a-5bdc-97e0-2f5a05c973a2"
 Crayons = "a8cc5b0e-0ffa-5ad4-8c14-923d3ee1735f"
 DataFrames = "a93c6f00-e57d-5684-b7b6-d8193f3e46c0"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 ExprTools = "e2ba6199-217a-4e67-a87a-7c52f15ade04"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
 GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
@@ -37,10 +38,12 @@ Statistics = "10745b16-79ce-11e8-11f9-7d13ad32a3b2"
 
 [weakdeps]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 
 [extensions]
 ChainRulesCoreExt = "ChainRulesCore"
+EnzymeCoreExt = "EnzymeCore"
 SpecialFunctionsExt = "SpecialFunctions"
 
 [compat]
@@ -55,6 +58,7 @@ ChainRulesCore = "1"
 Crayons = "4"
 DataFrames = "1"
 ExprTools = "0.1"
+EnzymeCore = "0.7.1"
 GPUArrays = "10.0.1"
 GPUCompiler = "0.24, 0.25, 0.26"
 KernelAbstractions = "0.9.2"
@@ -81,4 +85,5 @@ julia = "1.8"
 
 [extras]
 ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"

ext/EnzymeCoreExt.jl

@@ -0,0 +1,16 @@
+# compatibility with EnzymeCore
+
+module EnzymeCoreExt
+
+using CUDA
+import CUDA: GPUCompiler, CUDABackend
+
+isdefined(Base, :get_extension) ? (import EnzymeCore) : (import ..EnzymeCore)
+
+function EnzymeCore.compiler_job_from_backend(::CUDABackend, @nospecialize(F::Type), @nospecialize(TT::Type))
+    mi = GPUCompiler.methodinstance(F, TT)
+    return GPUCompiler.CompilerJob(mi, CUDA.compiler_config(CUDA.device()))
+end
+
+end # module
+

src/initialization.jl

@@ -155,6 +155,9 @@ function __init__()
         @require SpecialFunctions="276daf66-3868-5448-9aa4-cd146d93841b" begin
             include("../ext/SpecialFunctionsExt.jl")
         end
+        @require EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869" begin
+            include("../ext/EnzymeCoreExt.jl")
+        end
     end
 
     # ensure that operations executed by the REPL back-end finish before returning,

test/Project.toml

@@ -7,8 +7,10 @@ ChainRulesCore = "d360d2e6-b24c-11e9-a2a3-2a2ae2dbcce4"
 DataStructures = "864edb3b-99cc-5e75-8d2d-829cb0a9cfe8"
 Dates = "ade2ca70-3891-5945-98fb-dc099432e06a"
 Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b"
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
 FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341"
 GPUArrays = "0c68f7d7-f131-5f86-a1c3-88cf8149b2d7"
+GPUCompiler = "61eb1bfa-7361-4325-ad38-22787b887f55"
 InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 Interpolations = "a98d9a8b-a2ab-59e6-89dd-64a1c18fca59"
 KernelAbstractions = "63c18a36-062a-441e-b654-da1e3ab1ce7c"

test/libraries/enzyme.jl

@@ -0,0 +1,6 @@
+using EnzymeCore
+using GPUCompiler
+
+@testset "compiler_job_from_backend" begin
+    @test EnzymeCore.compiler_job_from_backend(CUDABackend(), typeof(()->nothing), Tuple{}) isa GPUCompiler.CompilerJob
+end