diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index bff31d98..e4dc9ea0 100644 --- a/.buildkite/pipeline.yml +++ b/.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 diff --git a/Project.toml b/Project.toml index c65e5fc3..dddeefb9 100644 --- a/Project.toml +++ b/Project.toml @@ -1,7 +1,7 @@ name = "KernelAbstractions" uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c" authors = ["Valentin Churavy 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" diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index ab9d0d3f..da316f12 100644 --- a/ext/EnzymeExt.jl +++ b/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, diff --git a/test/Project.toml b/test/Project.toml index b407094b..c8e674ab 100644 --- a/test/Project.toml +++ b/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"