[maleadt]

Fixes #675.

Calling a HostKernel took args... and keyword arguments, so it was compiled for Vararg{Any}: Julia doesn't specialize a method on varargs it only passes on, and the lowered keyword wrapper splats them. Every launch, through @oneapi or KernelAbstractions, then dispatched at run time, and each new signature compiled its launch path piecemeal during the first launch. This follows what CUDA.jl (JuliaGPU/CUDA.jl#3309, JuliaGPU/CUDA.jl#3359) and Metal.jl (JuliaGPU/Metal.jl#1015) did:

• the HostKernel call takes Vararg{Any,N}, with an explicit Core.kwcall method, and passes the arguments on as a tuple; the type-independent part of the launch is compiled once;
• zefunction only constructs the HostKernel, and the compilation and linking moved into zefunction_link(::MethodInstance). The cache of HostKernel objects is gone, and driver(), device() and context() now have return-type assertions, since they used to infer as Any;
• kernel arguments are set through a setter that isn't specialized on the argument type.

The launch_configuration instability the issue mentions was already fixed on main. The KernelAbstractions launch method it mentions now lives in KernelAbstractions, which already specializes it. Kernels with 32 or more arguments still dispatch, because map over such tuples isn't inferred; that predates this PR and doesn't come up often.

Measured on an Iris Xe (Julia 1.12.7) with a 64-element y[i] = x[i] * a kernel. Host compilation is the median over five new element types, with device compilation excluded:

	main	PR
host compilation per new signature, @oneapi	285–550 ms	100–115 ms
host compilation per new signature, KA	470–1280 ms	235–280 ms
allocations per warm launch, @oneapi	1520 B	208 B
allocations per warm launch, KA	3632 B	304 B

The machine was busy at the time (load around 10 on 8 cores), so the times are ranges over three interleaved runs, and the warm launch times were too noisy to compare.

A test checks that the HostKernel call isn't specialized for Vararg{Any}, and that launching kernels with new functions and argument types doesn't compile the shared launch code again.

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/lib/level-zero/module.jl b/lib/level-zero/module.jl
index 0dbdf9a..d34509d 100644
--- a/lib/level-zero/module.jl
+++ b/lib/level-zero/module.jl
@@ -221,7 +221,7 @@ 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)
+    return zeKernelSetArgumentValue(kernel, index - 1, size, ptr)
 end
 
 
diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl
index 129a13c..7d8907a 100644
--- a/src/compiler/execution.jl
+++ b/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::Tuple; 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:fieldcount(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)
@@ -289,7 +289,7 @@ 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)
+    return HostKernel{F, tt}(f, fun)
 end
 
 # the parts of `zefunction` that don't depend on the type of the function, kept separate so
@@ -347,7 +347,8 @@ end
     return res
 end
 
-@inline function onecall(kernel::ZeKernel, tt, args::Tuple; 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
         # unrolled, as iterating a heterogeneous tuple would dispatch at run time
@@ -361,10 +362,12 @@ 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})
+@noinline function launch_kernel(
+        kernel::ZeKernel, groups::ZeDim, items::ZeDim,
+        queue::Union{oneStream, ZeCommandQueue}
+    )
     groupsize!(kernel, items)
-    launch!(queue, kernel, groups)
+    return launch!(queue, kernel, groups)
 end
 
 @inline function launch!(s::oneStream, kernel::ZeKernel, groups::ZeDim)
@@ -411,9 +414,9 @@ 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} =
+@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} =
+@inline Core.kwcall(kwargs::NamedTuple, kernel::HostKernel, args::Vararg{Any, N}) where {N} =
     call(kernel, map(kernel_convert, args); kwargs...)
 
 
diff --git a/test/execution.jl b/test/execution.jl
index 9032adb..559bf30 100644
--- a/test/execution.jl
+++ b/test/execution.jl
@@ -68,44 +68,44 @@ 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))
+    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, 1.0f0)
+        k = @oneapi launch = false store_kernel(a, 1.0f0)
+        k(a, 2.0f0)
+        @test Array(a)[1] == 2.0f0
+        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)
+        # 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
-    @test nspecializations() == before
-end
 
 
 @testset "reflection" begin

[codecov]

Codecov Report

❌ Patch coverage is 96.15385% with 1 line in your changes missing coverage. Please review.
✅ Project coverage is 80.66%. Comparing base (fa0f23a) to head (882c3a9).

Files with missing lines	Patch %	Lines
src/compiler/execution.jl	94.11%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #687      +/-   ##
==========================================
- Coverage   81.17%   80.66%   -0.52%     
==========================================
  Files          52       51       -1     
  Lines        4085     3972     -113     
==========================================
- Hits         3316     3204     -112     
+ Misses        769      768       -1     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.