Stop compiling host code for every new kernel type

src/array.jl

@@ -523,7 +523,7 @@ function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice,
       execute!(stream) do list
         stage_upload!(list, ctx, dev, device_pointer(dest, doffs), pointer(src, soffs), bytes)
       end
-    end
+    end::Bool
     if !staged
       with_ownership(stream, Managed[dest.data[]]) do
         unsafe_copyto!(ctx, dev, device_pointer(dest, doffs), pointer(src, soffs), n)

src/compiler/execution.jl

@@ -333,12 +333,21 @@ const zefunction_lock = ReentrantLock()
 function zefunction(f::F, tt::TT=Tuple{}; kwargs...) where {F,TT}
     # the kernel also depends on the task's context and device, see `zefunction_link`
     kw = values(kwargs)
-    fun = GPUCompiler.cached_lookup(ZeKernel, F, tt, (context(), device(), kw)) do
-        zefunction_link(methodinstance(F, tt); kw...)
-    end
+    fun = GPUCompiler.cached_lookup(KernelLinker(F, tt, kw), ZeKernel, F, tt,
+                                    (context(), device(), kw))
     return HostKernel{F,tt}(f, fun)
 end
 
+# What `zefunction` calls when the kernel isn't cached: a callable struct rather than a
+# closure, whose type would depend on the kernel's, compiling it (and the inlined
+# `methodinstance`) again for every kernel.
+struct KernelLinker{K}
+    ft::Type
+    tt::Type
+    kw::K
+end
+(l::KernelLinker)() = zefunction_link(methodinstance(l.ft, l.tt); l.kw...)
+
 # 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...)

src/memory.jl

@@ -134,8 +134,10 @@ function maybe_synchronize(managed::Managed)
 end
 
 # lock the given allocations in a stable order, avoiding deadlocks between tasks that
-# submit operations involving the same allocations
-function with_managed_locks(f, managed::Vector{Managed})
+# submit operations involving the same allocations. not specialized on `f`, which is a
+# different closure for every operation, e.g. for every kernel that is launched.
+Base.@nospecializeinfer function with_managed_locks(@nospecialize(f),
+                                                  managed::Vector{Managed})
     locked = sort!(unique!(managed); by=objectid)
     nlocked = 0
     ret = try
@@ -153,11 +155,20 @@ function with_managed_locks(f, managed::Vector{Managed})
 end
 
 # submit an operation using the memory of `managed` to `owner`, as a single transaction
-function with_ownership(f, owner, managed::Vector{Managed})
-    with_managed_locks(managed) do
-        for m in managed
-            take_ownership!(m, owner)
-        end
-        f()
+Base.@nospecializeinfer with_ownership(@nospecialize(f),
+                                       owner::Union{oneStream, ZeCommandQueue},
+                                       managed::Vector{Managed}) =
+    with_managed_locks(OwnershipTransfer(f, owner, managed), managed)
+
+# (a struct rather than a closure, whose type would depend on that of `f`)
+struct OwnershipTransfer
+    f::Any
+    owner::Union{oneStream, ZeCommandQueue}
+    managed::Vector{Managed}
+end
+function (t::OwnershipTransfer)()
+    for m in t.managed
+        take_ownership!(m, t.owner)
     end
+    t.f()
 end

test/execution.jl

@@ -109,7 +109,9 @@ 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!)
+    shared = (oneAPI.zefunction_link, oneAPI.launch_kernel, oneL0.set_argument_value!,
+              oneAPI.KernelLinker(Nothing, Tuple{}, (;)), oneAPI.with_managed_locks,
+              oneAPI.with_ownership)
     nspecializations() =
         sum(f -> sum(m -> length(collect(Base.specializations(m))), methods(f)), shared)
     before = nspecializations()