Better launch interface

Author: christiangnrd

src/KernelAbstractions.jl

@@ -194,6 +194,11 @@ function unsafe_free! end
 
 unsafe_free!(::AbstractArray) = return
 
+"""
+Abstract type for all KernelAbstractions backends.
+"""
+abstract type Backend end
+
 include("intrinsics.jl")
 import .KernelIntrinsics
 export KernelIntrinsics
@@ -500,11 +505,6 @@ constify(arg) = adapt(ConstAdaptor(), arg)
 # Backend hierarchy
 ###
 
-"""
-
-Abstract type for all KernelAbstractions backends.
-"""
-abstract type Backend end
 
 """
 Abstract type for all GPU based KernelAbstractions backends.

src/intrinsics.jl

@@ -1,5 +1,8 @@
 module KernelIntrinsics
 
+import ..KernelAbstractions: Backend
+import GPUCompiler: split_kwargs, assign_args!
+
 """
     get_global_size()::@NamedTuple{x::Int, y::Int, z::Int}
 
@@ -100,7 +103,6 @@ kernel on the host.
 !!! note
     Backend implementations **must** implement:
     ```
-    KI.KIKernel(::NewBackend, f, args...; kwargs...)
     (kernel::KIKernel{<:NewBackend})(args...; numworkgroups=nothing, workgroupsize=nothing, kwargs...)
     ```
     As well as the on-device functionality.
@@ -157,4 +159,88 @@ Used for certain algorithm optimizations.
     As well as the on-device functionality.
 """
 multiprocessor_count(_) = 0
+
+# TODO: docstring
+# kiconvert(::NewBackend, arg)
+function kiconvert end
+
+# TODO: docstring
+# KI.kifunction(::NewBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
+function kifunction end
+
+const MACRO_KWARGS = [:launch, :backend]
+const COMPILER_KWARGS = [:kernel, :name, :always_inline]
+const LAUNCH_KWARGS = [:numworkgroups, :workgroupsize]
+
+macro kikernel(backend, ex...)
+    call = ex[end]
+    kwargs = map(ex[1:end-1]) do kwarg
+        if kwarg isa Symbol
+            :($kwarg = $kwarg)
+        elseif Meta.isexpr(kwarg, :(=))
+            kwarg
+        else
+            throw(ArgumentError("Invalid keyword argument '$kwarg'"))
+        end
+    end
+
+    # destructure the kernel call
+    Meta.isexpr(call, :call) || throw(ArgumentError("final argument to @kikern should be a function call"))
+    f = call.args[1]
+    args = call.args[2:end]
+
+    code = quote end
+    vars, var_exprs = assign_args!(code, args)
+
+    # group keyword argument
+    macro_kwargs, compiler_kwargs, call_kwargs, other_kwargs =
+        split_kwargs(kwargs, MACRO_KWARGS, COMPILER_KWARGS, LAUNCH_KWARGS)
+    if !isempty(other_kwargs)
+        key,val = first(other_kwargs).args
+        throw(ArgumentError("Unsupported keyword argument '$key'"))
+    end
+
+    # handle keyword arguments that influence the macro's behavior
+    launch = true
+    for kwarg in macro_kwargs
+        key,val = kwarg.args
+        if key === :launch
+            isa(val, Bool) || throw(ArgumentError("`launch` keyword argument to @kikern should be a Bool"))
+            launch = val::Bool
+        else
+            throw(ArgumentError("Unsupported keyword argument '$key'"))
+        end
+    end
+    if !launch && !isempty(call_kwargs)
+        error("@kikern with launch=false does not support launch-time keyword arguments; use them when calling the kernel")
+    end
+
+    # FIXME: macro hygiene wrt. escaping kwarg values (this broke with 1.5)
+    #        we esc() the whole thing now, necessitating gensyms...
+    @gensym f_var kernel_f kernel_args kernel_tt kernel
+
+    # convert the arguments, call the compiler and launch the kernel
+    # while keeping the original arguments alive
+    push!(code.args,
+        quote
+            $f_var = $f
+            GC.@preserve $(vars...) $f_var begin
+                $kernel_f = $kiconvert($backend, $f_var)
+                $kernel_args = map(x -> $kiconvert($backend, x), ($(var_exprs...),))
+                $kernel_tt = Tuple{map(Core.Typeof, $kernel_args)...}
+                $kernel = $kifunction($backend, $kernel_f, $kernel_tt; $(compiler_kwargs...))
+                if $launch
+                    $kernel($(var_exprs...); $(call_kwargs...))
+                end
+                $kernel
+            end
+         end)
+
+    return esc(quote
+        let
+            $code
+        end
+    end)
+end
+
 end

src/pocl/backend.jl

@@ -2,7 +2,7 @@ module POCLKernels
 
 using ..POCL
 using ..POCL: @device_override, cl, method_table
-using ..POCL: device
+using ..POCL: device, clconvert, clfunction
 
 import KernelAbstractions as KA
 
@@ -138,9 +138,38 @@ function (obj::KA.Kernel{POCLBackend})(args...; ndrange = nothing, workgroupsize
     return nothing
 end
 
+const KI = KA.KernelIntrinsics
+
+KI.kiconvert(::POCLBackend, arg) = clconvert(arg)
+
+function KI.kifunction(::POCLBackend, f::F, tt::TT=Tuple{}; name=nothing, kwargs...) where {F,TT}
+    kern = clfunction(f, tt; name, kwargs...)
+    KI.KIKernel{POCLBackend, typeof(kern)}(POCLBackend(), kern)
+end
+
+function (obj::KI.KIKernel{POCLBackend})(args...; numworkgroups=nothing, workgroupsize=nothing, kwargs...)
+    local_size = isnothing(workgroupsize) ? 1 : workgroupsize
+    global_size = if isnothing(numworkgroups)
+        1
+    else
+        numworkgroups*local_size
+    end
+
+    obj.kern(args...; local_size, global_size)
+end
+
+
+function KI.kernel_max_work_group_size(::POCLBackend, kikern::KI.KIKernel{<:POCLBackend}; max_work_items::Int=typemax(Int))::Int
+    4096
+end
+function KI.max_work_group_size(::POCLBackend)::Int
+    4096
+end
+function KI.multiprocessor_count(::POCLBackend)::Int
+    1
+end
 
 ## Indexing Functions
-const KI = KA.KernelIntrinsics
 
 @device_override @inline function KI.get_local_id()
     return (; x = Int(get_local_id(1)), y = Int(get_local_id(2)), z = Int(get_local_id(3)))