@ka_code_typed implementation for KernelAbstractions

Author: mjulian31

Project.toml

@@ -7,6 +7,8 @@ version = "0.4.1"
 Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e"
 CUDA = "052768ef-5323-5732-b1bb-66c8b64840ba"
 Cassette = "7057c7e9-c182-5462-911a-8362d720325c"
+InteractiveUtils = "b77e0a4c-d291-57a0-90e8-8db25a27a240"
+LLVM = "929cbde3-209d-540e-8aea-75f648917ca0"
 MacroTools = "1914dd2f-81c6-5fcd-8719-6d5c9610ff09"
 SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
 StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"

src/KernelAbstractions.jl

@@ -146,6 +146,9 @@ instead.
 See also [`@uniform`](@ref).
 """
 macro private(T, dims)
+    if dims isa Integer
+        dims = (dims,)
+    end
     quote
         $Scratchpad($(esc(T)), Val($(esc(dims))))
     end
@@ -384,7 +387,7 @@ function partition(kernel, ndrange, workgroupsize)
     if static_ndrange <: StaticSize
         static_blocks = StaticSize{blocks}
         blocks = nothing
-    else 
+    else
         static_blocks = DynamicSize
         blocks = CartesianIndices(blocks)
     end
@@ -466,4 +469,6 @@ include("backends/cuda.jl")
 ###
 
 include("extras/extras.jl")
+
+include("reflection.jl")
 end #module

src/backends/cpu.jl

@@ -38,7 +38,7 @@ function wait(cpu::CPU, ev::MultiEvent, progress=nothing)
         alldone = ntuple(N) do i
             if alldone[i]
                 true
-            else 
+            else
                 isdone(events[i])
             end
         end
@@ -48,7 +48,7 @@ function wait(cpu::CPU, ev::MultiEvent, progress=nothing)
             progress()
         end
     end
-   
+
     if any(failed, events)
         ex = CompositeException()
         for event in events
@@ -78,31 +78,38 @@ function async_copy!(::CPU, A, B; dependencies=nothing, progress=nothing)
 end
 
 function (obj::Kernel{CPU})(args...; ndrange=nothing, workgroupsize=nothing, dependencies=nothing, progress=nothing)
+    ndrange, workgroupsize, iterspace, dynamic = launch_config(obj, ndrange, workgroupsize)
+
+    if dependencies isa Event
+        dependencies = (dependencies,)
+    end
+
+    if length(blocks(iterspace)) == 0
+        return MultiEvent(dependencies)
+    end
+
+    Event(__run, obj, ndrange, iterspace, args, dynamic,
+          dependencies=dependencies, progress=progress)
+end
+
+function launch_config(kernel::Kernel{CPU}, ndrange, workgroupsize)
     if ndrange isa Integer
         ndrange = (ndrange,)
     end
     if workgroupsize isa Integer
         workgroupsize = (workgroupsize, )
     end
-    if dependencies isa Event
-        dependencies = (dependencies,)
-    end
 
-    if KernelAbstractions.workgroupsize(obj) <: DynamicSize && workgroupsize === nothing
+    if KernelAbstractions.workgroupsize(kernel) <: DynamicSize && workgroupsize === nothing
         workgroupsize = (1024,) # Vectorization, 4x unrolling, minimal grain size
     end
-    iterspace, dynamic = partition(obj, ndrange, workgroupsize)
+    iterspace, dynamic = partition(kernel, ndrange, workgroupsize)
     # partition checked that the ndrange's agreed
-    if KernelAbstractions.ndrange(obj) <: StaticSize
+    if KernelAbstractions.ndrange(kernel) <: StaticSize
         ndrange = nothing
     end
 
-    if length(blocks(iterspace)) == 0
-        return MultiEvent(dependencies)
-    end
-
-    Event(__run, obj, ndrange, iterspace, args, dynamic,
-          dependencies=dependencies, progress=progress)
+    return ndrange, workgroupsize, iterspace, dynamic
 end
 
 # Inference barriers

src/reflection.jl

@@ -0,0 +1,42 @@
+import InteractiveUtils
+export @ka_code_typed
+
+function ka_code_typed(kernel, argtypes; ndrange=nothing, workgroupsize=nothing, dependencies=nothing, kwargs...)
+    # get the iterspace and dynamic of a kernel
+    ndrange, workgroupsize, iterspace, dynamic = KernelAbstractions.launch_config(kernel, ndrange, workgroupsize)
+    # get the first block
+    block = @inbounds KernelAbstractions.blocks(iterspace)[1]
+    # get a context of the kernel based on the first block
+    ctx = KernelAbstractions.mkcontext(kernel, block, ndrange, iterspace, dynamic)
+    # reformat
+    if argtypes isa Type
+        argtypes = argtypes.parameters
+    end
+    # use code_typed
+    return InteractiveUtils.code_typed(KernelAbstractions.Cassette.overdub, (typeof(ctx), typeof(kernel.f), argtypes...); kwargs...)
+end
+
+
+macro ka_code_typed(ex0...)
+    ex = ()
+    for i = 1:length(ex0)
+        if ex0[i].head == :call
+            while length(ex0[i].args) > 2
+                kw = ex0[i].args[end]
+                @assert kw.head == :kw
+                kw.args[2] = esc(kw.args[2])
+                kw.head = Symbol("=")
+                resize!(ex0[i].args, length(ex0[i].args) - 1)
+                ex = (kw,)..., ex...
+            end
+        end
+        ex = ex..., ex0[i]
+    end
+
+    thecall = InteractiveUtils.gen_call_with_extracted_types_and_kwargs(__module__, :ka_code_typed, ex)
+
+    quote
+        local results = $thecall
+        length(results) == 1 ? results[1] : results
+    end
+end