Merge pull request #519 from JuliaGPU/ecore

Adapt to pending Enzymecore changes

Author: wsmoses

.buildkite/pipeline.yml

@@ -30,8 +30,6 @@ steps:
     matrix:
       setup:
         version:
-          - "1.8"
-          - "1.9"
           - "1.10"
           - "1.11"
     plugins:
@@ -42,8 +40,12 @@ steps:
     command: |
       julia -e 'println("--- :julia: Instantiating project")
                 using Pkg
-                Pkg.develop(; path=pwd())
-                Pkg.add(["CUDA", "Enzyme"])' || exit 3
+                try
+                   Pkg.develop([PackageSpec(; path=pwd()), PackageSpec("Enzyme"), PackageSpec("CUDA")])
+                catch err
+                   Pkg.develop(; path=pwd())
+                  Pkg.add(["CUDA", "Enzyme"])
+                end' || exit 3
 
       julia -e 'println("+++ :julia: Running tests")
                 using CUDA

Project.toml

@@ -1,7 +1,7 @@
 name = "KernelAbstractions"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 authors = ["Valentin Churavy <[email protected]> and contributors"]
-version = "0.9.26"
+version = "0.9.27"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -21,7 +21,7 @@ UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
 [compat]
 Adapt = "0.4, 1.0, 2.0, 3.0, 4"
 Atomix = "0.1"
-EnzymeCore = "0.7.5"
+EnzymeCore = "0.8.1"
 InteractiveUtils = "1.6"
 LinearAlgebra = "1.6"
 MacroTools = "0.5"

ext/EnzymeExt.jl

@@ -44,17 +44,18 @@ EnzymeRules.inactive(::Type{StaticSize}, x...) = nothing
 # https://github.com/EnzymeAD/Enzyme.jl/issues/1516
 # On the CPU `autodiff_deferred` can deadlock.
 # Hence a specialized CPU version
-function cpu_fwd(ctx, f, args...)
-    EnzymeCore.autodiff(Forward, Const(f), Const{Nothing}, Const(ctx), args...)
+function cpu_fwd(ctx, config, f, args...)
+    EnzymeCore.autodiff(EnzymeCore.set_runtime_activity(Forward, config), Const(f), Const{Nothing}, Const(ctx), args...)
     return nothing
 end
 
-function gpu_fwd(ctx, f, args...)
-    EnzymeCore.autodiff_deferred(Forward, Const(f), Const{Nothing}, Const(ctx), args...)
+function gpu_fwd(ctx, config, f, args...)
+    EnzymeCore.autodiff_deferred(EnzymeCore.set_runtime_activity(Forward, config), Const(f), Const{Nothing}, Const(ctx), args...)
     return nothing
 end
 
 function EnzymeRules.forward(
+	config,
         func::Const{<:Kernel{CPU}},
         ::Type{Const{Nothing}},
         args...;
@@ -65,10 +66,11 @@ function EnzymeRules.forward(
     f = kernel.f
     fwd_kernel = similar(kernel, cpu_fwd)
 
-    fwd_kernel(f, args...; ndrange, workgroupsize)
+    fwd_kernel(config, f, args...; ndrange, workgroupsize)
 end
 
 function EnzymeRules.forward(
+	config,
         func::Const{<:Kernel{<:GPU}},
         ::Type{Const{Nothing}},
         args...;
@@ -79,7 +81,7 @@ function EnzymeRules.forward(
     f = kernel.f
     fwd_kernel = similar(kernel, gpu_fwd)
 
-    fwd_kernel(f, args...; ndrange, workgroupsize)
+    fwd_kernel(config, f, args...; ndrange, workgroupsize)
 end
 
 _enzyme_mkcontext(kernel::Kernel{CPU}, ndrange, iterspace, dynamic) =
@@ -253,7 +255,7 @@ function gpu_rev(
 end
 
 function EnzymeRules.augmented_primal(
-        config::Config,
+        config::RevConfig,
         func::Const{<:Kernel},
         ::Type{Const{Nothing}},
         args::Vararg{Any, N};
@@ -311,7 +313,7 @@ function EnzymeRules.augmented_primal(
 end
 
 function EnzymeRules.reverse(
-        config::Config,
+        config::RevConfig,
         func::Const{<:Kernel},
         ::Type{<:EnzymeCore.Annotation},
         tape,
@@ -364,7 +366,7 @@ end
 #       synchronize rule and then synchronize where the launch was. However, with the current
 #       kernel semantics this ensures correctness for now.
 function EnzymeRules.augmented_primal(
-        config::Config,
+        config::RevConfig,
         func::Const{typeof(synchronize)},
         ::Type{Const{Nothing}},
         backend::T,
@@ -374,7 +376,7 @@ function EnzymeRules.augmented_primal(
 end
 
 function EnzymeRules.reverse(
-        config::Config,
+        config::RevConfig,
         func::Const{typeof(synchronize)},
         ::Type{Const{Nothing}},
         tape,

test/Project.toml

@@ -10,4 +10,4 @@ StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40"
 
 [compat]
-Enzyme = "0.12"
+Enzyme = "0.13"