From 03b35992f36da7bba3ddbbc16d8f0958cd65f7d3 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 12:12:04 +0200 Subject: [PATCH 1/3] Specialize kernel launches on their argument types 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 keyword wrapper splats them. Every launch, through `@oneapi` or KernelAbstractions, then dispatched at run time, and a new signature compiled the rest of its launch path piecemeal while launching. Take `Vararg{Any,N}` with an explicit `Core.kwcall` method, pass the arguments on as a tuple, set them without iterating a heterogeneous tuple, and move the part of the launch that doesn't depend on the arguments into a function that's compiled once. --- src/compiler/execution.jl | 35 ++++++++++++++++++++++------------ test/execution.jl | 40 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 12 deletions(-) diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 6dd470934..74b946bf0 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...; 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 @@ -349,18 +349,26 @@ 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 +411,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 diff --git a/test/execution.jl b/test/execution.jl index 19d71ae65..e3158c514 100644 --- a/test/execution.jl +++ b/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.launch_kernel,) + 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{}) From 2640b14c99c748b3fa14a1af14a3bae6ccc7fea0 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 12:15:38 +0200 Subject: [PATCH 2/3] Compile the type-independent part of zefunction once `zefunction` is specialized on the kernel function, but nearly all of its body (compiling or looking up the image, linking it for the active context and device) only needs a MethodInstance. Split that off into `zefunction_link`, so that it's compiled once instead of for every kernel. Also drop the cache of HostKernel objects: a HostKernel is just the function and the ZeKernel, and the `get!` closure was itself compiled per kernel. And assert the return types of `driver()`, `device()` and `context()`, which inferred as `Any` and made the compiler configuration and the kernel lookup dispatch at run time. --- src/compiler/execution.jl | 18 ++++++++---------- src/context.jl | 6 +++--- test/execution.jl | 2 +- 3 files changed, 12 insertions(+), 14 deletions(-) diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 74b946bf0..129a13c6f 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -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,9 +347,6 @@ end return res end -# cache of kernel instances -const _kernel_instances = Dict{UInt, Any}() - @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 diff --git a/src/context.jl b/src/context.jl index 521adb3ec..16a5af27c 100644 --- a/src/context.jl +++ b/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 """ diff --git a/test/execution.jl b/test/execution.jl index e3158c514..52fd60fe1 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -93,7 +93,7 @@ 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.launch_kernel,) + shared = (oneAPI.zefunction_link, oneAPI.launch_kernel) nspecializations() = sum(f -> sum(m -> length(collect(Base.specializations(m))), methods(f)), shared) before = nspecializations() From 882c3a98bb26837a9bb7849d16de8b246fcc7b25 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 12:26:45 +0200 Subject: [PATCH 3/3] Set kernel arguments through an unspecialized setter `setindex!(::ZeKernelArgumentList, ...)` is compiled for every argument type, and had the checked `zeKernelSetArgumentValue` wrapper inlined into it. Only compute the size and address of the argument there, and leave the API call to a function that's compiled once. --- lib/level-zero/module.jl | 16 +++++++++++++--- test/execution.jl | 2 +- 2 files changed, 14 insertions(+), 4 deletions(-) diff --git a/lib/level-zero/module.jl b/lib/level-zero/module.jl index 6e0635df6..0dbdf9a1d 100644 --- a/lib/level-zero/module.jl +++ b/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 diff --git a/test/execution.jl b/test/execution.jl index 52fd60fe1..9032adb1d 100644 --- a/test/execution.jl +++ b/test/execution.jl @@ -93,7 +93,7 @@ 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) + 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()