From a854a361b837e1e4ba26d054cae61a9e3783f938 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Mon, 16 Sep 2024 00:15:31 +0000 Subject: [PATCH 1/7] Adapt to pending Enzymecore changes --- Project.toml | 2 +- ext/EnzymeExt.jl | 12 +++++++----- 2 files changed, 8 insertions(+), 6 deletions(-) diff --git a/Project.toml b/Project.toml index c65e5fc30..98f0094d1 100644 --- a/Project.toml +++ b/Project.toml @@ -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" InteractiveUtils = "1.6" LinearAlgebra = "1.6" MacroTools = "0.5" diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index ab9d0d3fb..866632346 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(config, ctx, 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...) + 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...; @@ -63,12 +64,13 @@ function EnzymeRules.forward( ) kernel = func.val f = kernel.f - fwd_kernel = similar(kernel, cpu_fwd) + fwd_kernel = similar(config, kernel, cpu_fwd) fwd_kernel(f, args...; ndrange, workgroupsize) end function EnzymeRules.forward( + config, func::Const{<:Kernel{<:GPU}}, ::Type{Const{Nothing}}, args...; @@ -77,7 +79,7 @@ function EnzymeRules.forward( ) kernel = func.val f = kernel.f - fwd_kernel = similar(kernel, gpu_fwd) + fwd_kernel = similar(config, kernel, gpu_fwd) fwd_kernel(f, args...; ndrange, workgroupsize) end From 7516ccbd3e84c8edb5d32c2f2cd61b3d052775a1 Mon Sep 17 00:00:00 2001 From: Billy Moses Date: Mon, 16 Sep 2024 00:24:21 +0000 Subject: [PATCH 2/7] fixup --- ext/EnzymeExt.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 866632346..d293e54b0 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -255,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}; @@ -313,7 +313,7 @@ function EnzymeRules.augmented_primal( end function EnzymeRules.reverse( - config::Config, + config::RevConfig, func::Const{<:Kernel}, ::Type{<:EnzymeCore.Annotation}, tape, @@ -366,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, @@ -376,7 +376,7 @@ function EnzymeRules.augmented_primal( end function EnzymeRules.reverse( - config::Config, + config::RevConfig, func::Const{typeof(synchronize)}, ::Type{Const{Nothing}}, tape, From 925e8f36b153c1675920a6bd1e030b80ba15fcf7 Mon Sep 17 00:00:00 2001 From: William Moses Date: Mon, 16 Sep 2024 20:30:14 -0500 Subject: [PATCH 3/7] Update EnzymeExt.jl --- ext/EnzymeExt.jl | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index d293e54b0..7cf7ef731 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -44,12 +44,12 @@ 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(config, ctx, f, 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...) +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 @@ -66,7 +66,7 @@ function EnzymeRules.forward( f = kernel.f fwd_kernel = similar(config, kernel, cpu_fwd) - fwd_kernel(f, args...; ndrange, workgroupsize) + fwd_kernel(config, f, args...; ndrange, workgroupsize) end function EnzymeRules.forward( @@ -81,7 +81,7 @@ function EnzymeRules.forward( f = kernel.f fwd_kernel = similar(config, 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) = From 4b1e22a38b341a98eef0b8f1ba52f7e12c36a8e6 Mon Sep 17 00:00:00 2001 From: William Moses Date: Thu, 19 Sep 2024 11:29:43 -0500 Subject: [PATCH 4/7] Update Project.toml --- test/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/test/Project.toml b/test/Project.toml index b407094bb..c8e674aba 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" From fe1bec574bb21ad84652fed186ceccc81defed93 Mon Sep 17 00:00:00 2001 From: William Moses Date: Thu, 19 Sep 2024 15:08:50 -0400 Subject: [PATCH 5/7] fix --- ext/EnzymeExt.jl | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ext/EnzymeExt.jl b/ext/EnzymeExt.jl index 7cf7ef731..da316f12d 100644 --- a/ext/EnzymeExt.jl +++ b/ext/EnzymeExt.jl @@ -64,7 +64,7 @@ function EnzymeRules.forward( ) kernel = func.val f = kernel.f - fwd_kernel = similar(config, kernel, cpu_fwd) + fwd_kernel = similar(kernel, cpu_fwd) fwd_kernel(config, f, args...; ndrange, workgroupsize) end @@ -79,7 +79,7 @@ function EnzymeRules.forward( ) kernel = func.val f = kernel.f - fwd_kernel = similar(config, kernel, gpu_fwd) + fwd_kernel = similar(kernel, gpu_fwd) fwd_kernel(config, f, args...; ndrange, workgroupsize) end From 8e1a8945bd129a038d2a0b14cd70911f6b7263c7 Mon Sep 17 00:00:00 2001 From: William Moses Date: Thu, 19 Sep 2024 15:12:23 -0400 Subject: [PATCH 6/7] Enzyme buildkite --- .buildkite/pipeline.yml | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index bff31d981..e4dc9ea02 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 From e4d702f77a86080b5098bcd61cfc9e1c5efd073b Mon Sep 17 00:00:00 2001 From: William Moses Date: Thu, 19 Sep 2024 15:51:38 -0400 Subject: [PATCH 7/7] fix --- Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/Project.toml b/Project.toml index 98f0094d1..dddeefb96 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.8" +EnzymeCore = "0.8.1" InteractiveUtils = "1.6" LinearAlgebra = "1.6" MacroTools = "0.5"