JuliaGPU · vchuravy · Mar 12, 2025 · Mar 12, 2025 · Mar 12, 2025 · Jul 10, 2026
diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml
@@ -48,12 +48,9 @@ steps:
   #   command: |
   #     julia -e 'println("--- :julia: Instantiating project")
   #               using Pkg
-  #               try
-  #                  Pkg.develop([PackageSpec(; path=pwd()), PackageSpec("Enzyme"), PackageSpec("EnzymeCore"), PackageSpec("CUDA")])
-  #               catch err
-  #                  Pkg.develop(; path=pwd())
-  #                 Pkg.add(["CUDA", "Enzyme"])
-  #               end' || exit 3
+  #               Pkg.develop(; path=pwd())
+  #               Pkg.add(["CUDA", "Enzyme"])
+  #               ' || exit 3
 
   #     julia -e 'println("+++ :julia: Running tests")
   #               using CUDA

diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
@@ -124,24 +124,6 @@ jobs:
               end'
             echo '[pocl_standalone_jll]' > test/LocalPreferences.toml
             echo 'libpocl_path="${{ github.workspace }}/target/lib/libpocl_standalone.so"' >> test/LocalPreferences.toml
-      - name: "Co-develop Enzyme and KA"
-        run: |
-          julia -e '
-            using Pkg
-            withenv("JULIA_PKG_PRECOMPILE_AUTO" => 0) do
-              Pkg.activate("test")
-              Pkg.add(["Enzyme", "EnzymeCore"])
-
-              # to check compatibility, also add Enzyme to the main environment
-              # (or Pkg.test, which merges both environments, could fail)
-              Pkg.activate(".")
-              # Try to co-develop Enzyme and KA
-              try
-                Pkg.develop([PackageSpec("Enzyme"), PackageSpec("EnzymeCore")])
-              catch err
-              end
-            end
-          '
       - uses: julia-actions/julia-buildpkg@v1
       - uses: julia-actions/julia-runtest@v1
         if: runner.os != 'Windows'

diff --git a/ext/EnzymeCore07Ext.jl b/ext/EnzymeCore07Ext.jl
@@ -1,81 +1,30 @@
-# 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...)
-    return nothing
-end
-
-function gpu_fwd(ctx, f, args...)
+function fwd(ctx, f, args...)
     EnzymeCore.autodiff_deferred(Forward, Const(f), Const{Nothing}, Const(ctx), args...)
     return nothing
 end
 
 function EnzymeRules.forward(
-        func::Const{<:Kernel{CPU}},
-        ::Type{Const{Nothing}},
-        args...;
-        ndrange = nothing,
-        workgroupsize = nothing,
-    )
-    kernel = func.val
-    f = kernel.f
-    fwd_kernel = similar(kernel, cpu_fwd)
-
-    return fwd_kernel(f, args...; ndrange, workgroupsize)
-end
-
-function EnzymeRules.forward(
-        func::Const{<:Kernel{<:GPU}},
+        func::Const{<:Kernel},
         ::Type{Const{Nothing}},
         args...;
         ndrange = nothing,
         workgroupsize = nothing,
     )
     kernel = func.val
     f = kernel.f
-    fwd_kernel = similar(kernel, gpu_fwd)
+    fwd_kernel = similar(kernel, fwd)
 
     return fwd_kernel(f, args...; ndrange, workgroupsize)
 end
 
-_enzyme_mkcontext(kernel::Kernel{CPU}, ndrange, iterspace, dynamic) =
-    mkcontext(kernel, first(blocks(iterspace)), ndrange, iterspace, dynamic)
-_enzyme_mkcontext(kernel::Kernel{<:GPU}, ndrange, iterspace, dynamic) =
+_enzyme_mkcontext(kernel::Kernel, ndrange, iterspace, dynamic) =
     mkcontext(kernel, ndrange, iterspace)
 
-_augmented_return(::Kernel{CPU}, subtape, arg_refs, tape_type) =
-    AugmentedReturn{Nothing, Nothing, Tuple{Array, typeof(arg_refs), typeof(tape_type)}}(
-    nothing,
-    nothing,
-    (subtape, arg_refs, tape_type),
-)
-_augmented_return(::Kernel{<:GPU}, subtape, arg_refs, tape_type) =
+_augmented_return(::Kernel, subtape, arg_refs, tape_type) =
     AugmentedReturn{Nothing, Nothing, Any}(nothing, nothing, (subtape, arg_refs, tape_type))
 
 function _create_tape_kernel(
-        kernel::Kernel{CPU},
-        ModifiedBetween,
-        FT,
-        ctxTy,
-        ndrange,
-        iterspace,
-        args2...,
-    )
-    TapeType = EnzymeCore.tape_type(
-        ReverseSplitModified(ReverseSplitWithPrimal, ModifiedBetween),
-        FT,
-        Const{Nothing},
-        Const{ctxTy},
-        map(Core.Typeof, args2)...,
-    )
-    subtape = Array{TapeType}(undef, size(blocks(iterspace)))
-    aug_kernel = similar(kernel, cpu_aug_fwd)
-    return TapeType, subtape, aug_kernel
-end
-
-function _create_tape_kernel(
-        kernel::Kernel{<:GPU},
+        kernel::Kernel,
         ModifiedBetween,
         FT,
         ctxTy,
@@ -104,60 +53,11 @@ function _create_tape_kernel(
     # Allocate per thread
     subtape = allocate(backend(kernel), TapeType, prod(ndrange))
 
-    aug_kernel = similar(kernel, gpu_aug_fwd)
+    aug_kernel = similar(kernel, aug_fwd)
     return TapeType, subtape, aug_kernel
 end
 
-_create_rev_kernel(kernel::Kernel{CPU}) = similar(kernel, cpu_rev)
-_create_rev_kernel(kernel::Kernel{<:GPU}) = similar(kernel, gpu_rev)
-
-function cpu_aug_fwd(
-        ctx,
-        f::FT,
-        ::Val{ModifiedBetween},
-        subtape,
-        ::Val{TapeType},
-        args...,
-    ) where {ModifiedBetween, FT, TapeType}
-    # A2 = Const{Nothing} -- since f->Nothing
-    forward, _ = EnzymeCore.autodiff_thunk(
-        ReverseSplitModified(ReverseSplitWithPrimal, Val(ModifiedBetween)),
-        Const{Core.Typeof(f)},
-        Const{Nothing},
-        Const{Core.Typeof(ctx)},
-        map(Core.Typeof, args)...,
-    )
-
-    # On the CPU: F is a per block function
-    # On the CPU: subtape::Vector{Vector}
-    I = __index_Group_Cartesian(ctx, CartesianIndex(1, 1)) #=fake=#
-    subtape[I] = forward(Const(f), Const(ctx), args...)[1]
-    return nothing
-end
-
-function cpu_rev(
-        ctx,
-        f::FT,
-        ::Val{ModifiedBetween},
-        subtape,
-        ::Val{TapeType},
-        args...,
-    ) where {ModifiedBetween, FT, TapeType}
-    _, reverse = EnzymeCore.autodiff_thunk(
-        ReverseSplitModified(ReverseSplitWithPrimal, Val(ModifiedBetween)),
-        Const{Core.Typeof(f)},
-        Const{Nothing},
-        Const{Core.Typeof(ctx)},
-        map(Core.Typeof, args)...,
-    )
-    I = __index_Group_Cartesian(ctx, CartesianIndex(1, 1)) #=fake=#
-    tp = subtape[I]
-    reverse(Const(f), Const(ctx), args..., tp)
-    return nothing
-end
-
-# GPU support
-function gpu_aug_fwd(
+function aug_fwd(
         ctx,
         f::FT,
         ::Val{ModifiedBetween},
@@ -184,7 +84,7 @@ function gpu_aug_fwd(
     return nothing
 end
 
-function gpu_rev(
+function rev(
         ctx,
         f::FT,
         ::Val{ModifiedBetween},
@@ -232,11 +132,7 @@ function EnzymeRules.augmented_primal(
     arg_refs = ntuple(Val(N)) do i
         Base.@_inline_meta
         if args[i] isa Active
-            if func.val isa Kernel{<:GPU}
-                error("Active kernel arguments not supported on GPU")
-            else
-                Ref(EnzymeCore.make_zero(args[i].val))
-            end
+            error("Active kernel arguments not supported")
         else
             nothing
         end
@@ -292,7 +188,7 @@ function EnzymeRules.reverse(
 
     ModifiedBetween = Val((overwritten(config)[1], false, overwritten(config)[2:end]...))
 
-    rev_kernel = _create_rev_kernel(kernel)
+    rev_kernel = similar(kernel, rev)
     rev_kernel(
         f,
         ModifiedBetween,