update

.buildkite/pipeline.yml

@@ -6,6 +6,7 @@ steps:
           - "1.6"
           - "1.7"
           - "1.8"
+          - "1.9"
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.version}}"
@@ -19,7 +20,35 @@ steps:
 
       julia -e 'println("+++ :julia: Running tests")
                 using Pkg
-                Pkg.test("CUDA"; coverage=true, test_args=["kernelabstractions"])'
+                Pkg.test("CUDA"; coverage=true, test_args=["base/kernelabstractions"])'
+    agents:
+      queue: "juliagpu"
+      cuda: "*"
+    timeout_in_minutes: 120
+    soft_fail:
+      - exit_status: 3
+
+  - label: "CUDA Enzyme Julia {{matrix.version}}"
+    matrix:
+      setup:
+        version:
+          - "1.8"
+          - "1.9"
+    plugins:
+      - JuliaCI/julia#v1:
+          version: "{{matrix.version}}"
+      - JuliaCI/julia-coverage#v1:
+          codecov: true
+    command: |
+      julia -e 'println("--- :julia: Instantiating project")
+                using Pkg
+                Pkg.develop(; path=pwd())
+                Pkg.add(["CUDA", "Enzyme"])' || exit 3
+
+      julia -e 'println("+++ :julia: Running tests")
+                using CUDA
+                include("test/extensions/enzyme.jl")
+                enzyme_testsuite(CUDABackend, CuArray, false)'
     agents:
       queue: "juliagpu"
       cuda: "*"
@@ -32,6 +61,7 @@ steps:
       setup:
         version:
           - "1.8"
+          - "1.9"
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.version}}"
@@ -59,6 +89,7 @@ steps:
       setup:
         version:
           - "1.8"
+          - "1.9"
     plugins:
       - JuliaCI/julia#v1:
           version: "{{matrix.version}}"

.github/workflows/ci.yml

@@ -25,6 +25,7 @@ jobs:
           - '1.6'
           - '1.7'
           - '1.8'
+          - '1.9'
           - 'nightly'
         os:
           - ubuntu-latest

Project.toml

@@ -1,7 +1,7 @@
 name = "KernelAbstractions"
 uuid = "63c18a36-062a-441e-b654-da1e3ab1ce7c"
 authors = ["Valentin Churavy <v.churavy@gmail.com> and contributors"]
-version = "0.9.4"
+version = "0.9.7"
 
 [deps]
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
@@ -10,6 +10,7 @@ InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
 LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 PrecompileTools = "aea7be01-6a6a-4083-8856-8a6e6704d82a"
+Requires = "ae029012-a4dd-5104-9daa-d747884805df"
 SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
 UUIDs = "cf7118a7-6976-5b1a-9a39-7adc72f591a4"
@@ -19,9 +20,20 @@ UnsafeAtomicsLLVM = "d80eeb9a-aca5-4d75-85e5-170c8b632249"
 [compat]
 Adapt = "0.4, 1.0, 2.0, 3.0"
 Atomix = "0.1"
+EnzymeCore = "0.5"
 MacroTools = "0.5"
+PrecompileTools = "1"
+Requires = "1.3"
 StaticArrays = "0.12, 1.0"
 UnsafeAtomics = "0.2.1"
 UnsafeAtomicsLLVM = "0.1"
-PrecompileTools = "1"
 julia = "1.6"
+
+[extensions]
+EnzymeExt = "EnzymeCore"
+
+[extras]
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"
+
+[weakdeps]
+EnzymeCore = "f151be2c-9106-41f4-ab19-57ee4f262869"

ext/EnzymeExt.jl

@@ -0,0 +1,81 @@
+module EnzymeExt
+    if isdefined(Base, :get_extension)
+        using EnzymeCore
+        using EnzymeCore.EnzymeRules
+    else
+        using ..EnzymeCore
+        using ..EnzymeCore.EnzymeRules
+    end
+    import KernelAbstractions: Kernel, StaticSize, launch_config, __groupsize, __groupindex, blocks, mkcontext, CompilerMetadata, CPU
+
+    EnzymeRules.inactive(::typeof(StaticSize), x...) = nothing
+
+    function fwd(ctx, f, args...)
+        EnzymeCore.autodiff_deferred(Forward, Const(f), Const, Const(ctx), args...)
+        return nothing
+    end
+
+    function aug_fwd(ctx, f::FT, ::Val{ModifiedBetween}, subtape, args...) where {ModifiedBetween, FT}
+        forward, reverse = EnzymeCore.autodiff_deferred_thunk(ReverseSplitModified(ReverseSplitWithPrimal, Val(ModifiedBetween)), Const{Core.Typeof(f)}, Const, Const{Core.Typeof(ctx)}, map(Core.Typeof, args)...)
+        subtape[__groupindex(ctx)] = forward(Const(f), Const(ctx), args...)[1]
+        return nothing
+    end
+
+    function rev(ctx, f::FT, ::Val{ModifiedBetween}, subtape, args...) where {ModifiedBetween, FT}
+        forward, reverse = EnzymeCore.autodiff_deferred_thunk(ReverseSplitModified(ReverseSplitWithPrimal, Val(ModifiedBetween)), Const{Core.Typeof(f)}, Const, Const{Core.Typeof(ctx)}, map(Core.Typeof, args)...)
+        tp = subtape[__groupindex(ctx)]
+        reverse(Const(f), Const(ctx), args..., tp)
+        return nothing
+    end
+
+    function EnzymeRules.forward(func::Const{<:Kernel}, ::Type{Const{Nothing}}, args...; ndrange=nothing, workgroupsize=nothing)
+        kernel = func.val
+        f = kernel.f
+        fwd_kernel = similar(kernel, fwd)
+
+        fwd_kernel(f, args...; ndrange, workgroupsize)
+    end
+
+    function EnzymeRules.augmented_primal(config::Config, func::Const{<:Kernel{CPU}}, ::Type{Const{Nothing}}, args...; ndrange=nothing, workgroupsize=nothing)
+        kernel = func.val
+        f = kernel.f
+
+        ndrange, workgroupsize, iterspace, dynamic = launch_config(kernel, ndrange, workgroupsize)
+        block = first(blocks(iterspace))
+
+        ctx = mkcontext(kernel, block, ndrange, iterspace, dynamic)
+        ctxTy = Core.Typeof(ctx) # CompilerMetadata{ndrange(kernel), Core.Typeof(dynamic)}
+
+        # TODO autodiff_deferred on the func.val
+        ModifiedBetween = Val((overwritten(config)[1], false, overwritten(config)[2:end]...))
+
+        FT = Const{Core.Typeof(f)}
+
+        # TODO in KA backends like CUDAKernels, etc have a version with a parent job type
+        TapeType = EnzymeCore.tape_type(ReverseSplitModified(ReverseSplitWithPrimal, ModifiedBetween), FT, Const,  Const{ctxTy}, map(Core.Typeof, args)...)
+
+        subtape = Array{TapeType}(undef, __groupsize(ctx))
+
+        aug_kernel = similar(kernel, aug_fwd)
+
+        aug_kernel(f, ModifiedBetween, subtape, args...; ndrange, workgroupsize)
+
+        # TODO the fact that ctxTy is type unstable means this is all type unstable.
+        # Since custom rules require a fixed return type, explicitly cast to Any, rather
+        # than returning a AugmentedReturn{Nothing, Nothing, T} where T.
+
+        res =  AugmentedReturn{Nothing, Nothing, Vector}(nothing, nothing, subtape)
+        return res
+    end
+
+    function EnzymeRules.reverse(config::Config, func::Const{<:Kernel}, ::Type{<:EnzymeCore.Annotation}, subtape, args...; ndrange=nothing, workgroupsize=nothing)
+        kernel = func.val
+        f = kernel.f
+
+        ModifiedBetween = Val((overwritten(config)[1], false, overwritten(config)[2:end]...))
+
+        rev_kernel = similar(func.val, rev)
+        rev_kernel(f, ModifiedBetween, subtape, args...; ndrange, workgroupsize)
+        return ((nothing for a in args)...,)
+    end
+end