Specialize kernel launches on their argument types

lib/level-zero/module.jl

@@ -209,9 +209,19 @@ end
 
 arguments(kernel::ZeKernel) = ZeKernelArgumentList(kernel)
 
-function Base.setindex!(args::ZeKernelArgumentList, value::Any, index::Integer)
-    @assert isbits(value)
-    zeKernelSetArgumentValue(args.kernel, index-1, Core.sizeof(value), Base.RefValue(value))
+function Base.setindex!(args::ZeKernelArgumentList, value::T, index::Integer) where {T}
+    @assert isbitstype(T)
+    ref = Base.RefValue(value)
+    GC.@preserve ref begin
+        ptr = Base.unsafe_convert(Ptr{Cvoid}, ref)
+        set_argument_value!(args.kernel, Int(index), Core.sizeof(T), ptr)
+    end
+    return args
+end
+
+# kept out of `setindex!`, which is compiled for every argument type
+@noinline function set_argument_value!(kernel::ZeKernel, index::Int, size::Int, ptr::Ptr{Cvoid})
+    zeKernelSetArgumentValue(kernel, index-1, size, ptr)
 end
 
 

src/compiler/execution.jl

@@ -196,9 +196,9 @@ kernel_convert(arg) = adapt(KernelAdaptor(), arg)
 
 abstract type AbstractKernel{F,TT} end
 
-@inline @generated function call(kernel::AbstractKernel{F,TT}, args...; call_kwargs...) where {F,TT}
+@inline @generated function call(kernel::AbstractKernel{F,TT}, args::Tuple; call_kwargs...) where {F,TT}
     sig = Tuple{F, TT.parameters...}    # Base.signature_type with a function type
-    args = (:(kernel.f), (:( args[$i] ) for i in 1:length(args))...)
+    args = (:(kernel.f), (:( args[$i] ) for i in 1:fieldcount(args))...)
 
     # filter out ghost arguments that shouldn't be passed
     predicate = dt -> isghosttype(dt) || Core.Compiler.isconstType(dt)
@@ -222,7 +222,7 @@ abstract type AbstractKernel{F,TT} end
     call_tt = Base.to_tuple_type(call_t)
 
     quote
-        onecall(kernel.fun, $call_tt, $(call_args...); call_kwargs...)
+        onecall(kernel.fun, $call_tt, ($(call_args...),); call_kwargs...)
     end
 end
 
@@ -288,11 +288,17 @@ end
 const zefunction_lock = ReentrantLock()
 
 function zefunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
+    fun = zefunction_link(methodinstance(F, tt); kwargs...)
+    return HostKernel{F,tt}(f, fun)
+end
+
+# the parts of `zefunction` that don't depend on the type of the function, kept separate so
+# that they are compiled once rather than for every kernel
+function zefunction_link(source::Core.MethodInstance; kwargs...)
     Base.@lock zefunction_lock begin
         ctx = context()
         dev = device()
         config = compiler_config(dev; kwargs...)::oneAPICompilerConfig
-        source = methodinstance(F, tt)
         job = CompilerJob(source, config)
 
         res = compile_or_lookup(job)::oneAPIResults
@@ -316,12 +322,7 @@ function zefunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
             end
         end
 
-        # create a callable object that captures the function instance. we don't need to think
-        # about world age here, as GPUCompiler already does and will return a different object
-        h = hash(fun, hash(f, hash(tt)))
-        get!(_kernel_instances, h) do
-            HostKernel{F,tt}(f, fun)
-        end::HostKernel{F,tt}
+        return fun::ZeKernel
     end
 end
 
@@ -346,21 +347,26 @@ end
     return res
 end
 
-# cache of kernel instances
-const _kernel_instances = Dict{UInt, Any}()
-
-@inline function onecall(kernel::ZeKernel, tt, args...; groups::ZeDim=1, items::ZeDim=1,
+@inline function onecall(kernel::ZeKernel, tt, args::Tuple; groups::ZeDim=1, items::ZeDim=1,
                          queue::Union{oneStream, ZeCommandQueue}=global_stream(context(), device()))
     Base.@lock kernel begin
-        for (i, arg) in enumerate(args)
-            oneL0.arguments(kernel)[i] = arg
+        # unrolled, as iterating a heterogeneous tuple would dispatch at run time
+        ntuple(Val(length(args))) do i
+            oneL0.arguments(kernel)[i] = args[i]
+            return
         end
-
-        groupsize!(kernel, items)
-        launch!(queue, kernel, groups)
+        launch_kernel(kernel, groups, items, queue)
     end
 end
 
+# kept out of `onecall`, which is compiled for every kernel signature, as none of this
+# depends on it
+@noinline function launch_kernel(kernel::ZeKernel, groups::ZeDim, items::ZeDim,
+                                 queue::Union{oneStream, ZeCommandQueue})
+    groupsize!(kernel, items)
+    launch!(queue, kernel, groups)
+end
+
 @inline function launch!(s::oneStream, kernel::ZeKernel, groups::ZeDim)
     mkl_wait!(s)
 
@@ -403,9 +409,12 @@ end
     return
 end
 
-function (kernel::HostKernel)(args...; kwargs...)
-    call(kernel, map(kernel_convert, args)...; kwargs...)
-end
+# `Vararg{Any,N}`, as Julia doesn't specialize a method on varargs it only passes on, and an
+# explicit `Core.kwcall` method, as the one lowered from keyword arguments splats them
+@inline (kernel::HostKernel)(args::Vararg{Any,N}) where {N} =
+    call(kernel, map(kernel_convert, args))
+@inline Core.kwcall(kwargs::NamedTuple, kernel::HostKernel, args::Vararg{Any,N}) where {N} =
+    call(kernel, map(kernel_convert, args); kwargs...)
 
 
 ## TODO: device-side kernels

src/context.jl

@@ -28,7 +28,7 @@ See also: `driver!`, `drivers`
 function driver()
     get!(task_local_storage(), :ZeDriver) do
         first(drivers())
-    end
+    end::ZeDriver
 end
 
 """
@@ -74,7 +74,7 @@ See also: `device!`, `devices`, `driver`
 function device()
     get!(task_local_storage(), :ZeDevice) do
         first(devices(driver()))
-    end
+    end::ZeDevice
 end
 
 """
@@ -174,7 +174,7 @@ function context()
                 ZeContext(drv)
             end
         end
-    end
+    end::ZeContext
 end
 
 """

test/execution.jl

@@ -68,6 +68,46 @@ end
 end
 
 
+function store_kernel(a, x)
+    @inbounds a[1] = x
+    return
+end
+function store_twice_kernel(a, x)
+    @inbounds a[1] = x + x
+    return
+end
+
+@testset "launch specialization" begin
+    # the launch is specialized on the argument types, so it doesn't dispatch at run time
+    a = oneArray(zeros(Float32, 1))
+    @oneapi items=1 store_kernel(a, 1f0)
+    k = @oneapi launch=false store_kernel(a, 1f0)
+    k(a, 2f0)
+    @test Array(a)[1] == 2f0
+    for m in [methods(k)..., methods(Core.kwcall, Tuple{NamedTuple, typeof(k), Vararg})...],
+        mi in Base.specializations(m)
+        params = mi.specTypes.parameters
+        any(p -> p === typeof(k), params) || continue
+        @test !Base.isvarargtype(last(params))
+    end
+
+    # creating and launching a kernel mostly doesn't depend on its signature, so launching
+    # kernels with new signatures shouldn't compile that again
+    shared = (oneAPI.zefunction_link, oneAPI.launch_kernel, oneL0.set_argument_value!)
+    nspecializations() =
+        sum(f -> sum(m -> length(collect(Base.specializations(m))), methods(f)), shared)
+    before = nspecializations()
+    for T in (Int32, Int16, UInt8)
+        b = oneArray(zeros(T, 1))
+        @oneapi items=1 store_kernel(b, one(T))
+        @test Array(b)[1] == one(T)
+        @oneapi items=1 store_twice_kernel(b, one(T))
+        @test Array(b)[1] == T(2)
+    end
+    @test nspecializations() == before
+end
+
+
 @testset "reflection" begin
     oneAPI.code_lowered(dummy, Tuple{})
     oneAPI.code_typed(dummy, Tuple{})