diff --git a/src/array.jl b/src/array.jl index 57eb781c..3fb46b89 100644 --- a/src/array.jl +++ b/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) diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index bcc6f80b..9c0a3cf5 100644 --- a/src/compiler/execution.jl +++ b/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...) diff --git a/src/memory.jl b/src/memory.jl index 4f724b57..ab3d6994 100644 --- a/src/memory.jl +++ b/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 diff --git a/test/execution.jl b/test/execution.jl index af6b627b..451edbf1 100644 --- a/test/execution.jl +++ b/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()