diff --git a/lib/level-zero/module.jl b/lib/level-zero/module.jl index 6e0635df..0dbdf9a1 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/src/compiler/execution.jl b/src/compiler/execution.jl index 6dd47093..129a13c6 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 @@ -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 diff --git a/src/context.jl b/src/context.jl index 521adb3e..16a5af27 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 19d71ae6..9032adb1 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.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{})