From 5436df45bbb4deb9b7231042f98dfb419750de69 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 12:19:54 +0200 Subject: [PATCH 1/4] Make the stream's companion queue field atomic The oneMKL companion queue of a task's stream is created lazily by the owning task. Upcoming per-allocation ownership tracking has other tasks wait for work on a stream they don't own, which means reading this field from another thread; make it atomic so the queue is safely published. --- src/context.jl | 18 ++++++++++-------- 1 file changed, 10 insertions(+), 8 deletions(-) diff --git a/src/context.jl b/src/context.jl index 16a5af27c..1ad4ae340 100644 --- a/src/context.jl +++ b/src/context.jl @@ -211,7 +211,8 @@ mutable struct oneStream const ctx::ZeContext const dev::ZeDevice const list::oneL0.ZeImmediateCommandList - queue::Union{Nothing, ZeCommandQueue} + # atomic because other tasks read it when waiting for work on memory this stream owns + Base.@atomic queue::Union{Nothing, ZeCommandQueue} mkl_dirty::Bool # high-water mark of per-thread spill (bytes) among kernels submitted to this # stream, maintained by the scratch hedge (see src/compiler/execution.jl) @@ -264,11 +265,12 @@ end # the companion command queue of a stream, created on first use. Only the SYCL/oneMKL # interop path needs one; pure-Julia workloads never create a queue. function stream_queue(s::oneStream) - q = s.queue + q = Base.@atomic s.queue q === nothing || return q - s.queue = ZeCommandQueue(s.ctx, s.dev; flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, - priority = s.priority) - return s.queue::ZeCommandQueue + q = ZeCommandQueue(s.ctx, s.dev; flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER, + priority = s.priority) + Base.@atomic s.queue = q + return q end """ @@ -356,7 +358,7 @@ function synchronize_all_streams(ctx::ZeContext, dev::Union{ZeDevice, Nothing}) # lists/queues finalized before their stream was registered stale. for stream in streams stream.list.handle == C_NULL || oneL0.synchronize(stream.list) - q = stream.queue + q = Base.@atomic stream.queue (q === nothing || q.handle == C_NULL) || oneL0.synchronize(q) end # retire drained streams of dead tasks, allowing their list and queue to be @@ -404,7 +406,7 @@ See also: [`global_stream`](@ref), [`context`](@ref), [`device`](@ref) function oneL0.synchronize(s::oneStream; blocking::Bool=false) sync = blocking ? oneL0.synchronize : oneL0.nonblocking_synchronize sync(s.list) - q = s.queue + q = Base.@atomic s.queue if q !== nothing sync(q) s.mkl_dirty = false @@ -435,7 +437,7 @@ end # creation the flag does not force). @inline function mkl_wait!(s::oneStream) if s.mkl_dirty - q = s.queue + q = Base.@atomic s.queue q === nothing || oneL0.synchronize(q) s.mkl_dirty = false end From 5aaced038d8091d1810b7f0ab0ac275ddf4d6786 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 12:38:45 +0200 Subject: [PATCH 2/4] Track outstanding GPU work per allocation Host accesses to shared and host arrays only waited for the calling task's stream, and only in scalar indexing and in copies to and from Arrays. Work submitted by another task was not waited for, nor was work preceding `unsafe_wrap(Array, a)`, and every element access cost a driver call. Wrap allocations in a `Managed` object that records the stream (or queue) that last used the memory and whether that use may still be outstanding, as CUDA.jl and OpenCL.jl do. Converting an array to a device pointer takes ownership on the task's stream, waiting for the previous owner if it differs; converting to a host pointer waits for the owner only if the memory is dirty, so repeated host accesses cost a single atomic load. Array copies take ownership of their arrays while holding their locks until the copy has been appended, so that another task never observes ownership without the work it describes. Waiting for another task's stream leaves its `mkl_dirty` flag alone, which only the owning task may modify. Fixes #652. --- src/array.jl | 112 +++++++++++++++++++++++++++----------------------- src/memory.jl | 88 +++++++++++++++++++++++++++++++++++++++ test/array.jl | 49 ++++++++++++++++++++++ 3 files changed, 198 insertions(+), 51 deletions(-) diff --git a/src/array.jl b/src/array.jl index a0f2d520d..c3aaeb7b8 100644 --- a/src/array.jl +++ b/src/array.jl @@ -105,7 +105,7 @@ end See also: [`oneVector`](@ref), [`oneMatrix`](@ref), [`is_device`](@ref), [`is_shared`](@ref) """ mutable struct oneArray{T,N,B} <: AbstractGPUArray{T,N} - data::DataRef{B} + data::DataRef{Managed{B}} maxsize::Int # maximum data size; excluding any selector bytes offset::Int # offset of the data in the buffer, in bytes @@ -126,15 +126,15 @@ mutable struct oneArray{T,N,B} <: AbstractGPUArray{T,N} alignment = Base.datatype_alignment(T) data = GPUArrays.cached_alloc((oneArray, B, ctx, dev, bufsize, alignment)) do buf = allocate(B, ctx, dev, bufsize, alignment) - data = DataRef(buf) do buf - release(buf) + data = DataRef(Managed(buf)) do managed + release(managed.mem) end end obj = new{T,N,B}(data, maxsize, 0, dims) finalizer(unsafe_free!, obj) end - function oneArray{T,N}(data::DataRef{B}, dims::Dims{N}; + function oneArray{T,N}(data::DataRef{Managed{B}}, dims::Dims{N}; maxsize::Int=prod(dims) * sizeof(T), offset::Int=0) where {T,N,B} check_eltype(T) if sizeof(T) == 0 @@ -152,7 +152,7 @@ GPUArrays.storage(a::oneArray) = a.data ## alias detection # GPUArrays implements `Base.dataids` and `Base.mightalias` from where an array lives -GPUArrays.memory_location(A::oneArray) = (UInt(convert(ZePtr{Cvoid}, A.data[])), A.offset) +GPUArrays.memory_location(A::oneArray) = (UInt(convert(ZePtr{Cvoid}, A.data[].mem)), A.offset) Base.unaliascopy(A::oneArray) = copy(A) @@ -227,11 +227,11 @@ Base.size(x::oneArray) = x.dims Base.sizeof(x::oneArray) = Base.elsize(x) * length(x) function context(A::oneArray) - return oneL0.context(A.data[]) + return oneL0.context(A.data[].mem) end function device(A::oneArray) - return oneL0.device(A.data[]) + return oneL0.device(A.data[].mem) end buftype(x::oneArray) = buftype(typeof(x)) @@ -246,7 +246,7 @@ Device memory provides the fastest GPU access but cannot be directly accessed fr See also: [`is_shared`](@ref), [`is_host`](@ref) """ -is_device(a::oneArray) = isa(a.data[], oneL0.DeviceBuffer) +is_device(a::oneArray) = isa(a.data[].mem, oneL0.DeviceBuffer) """ is_shared(a::oneArray) -> Bool @@ -257,7 +257,7 @@ Shared memory is accessible from both CPU and GPU with unified virtual addressin See also: [`is_device`](@ref), [`is_host`](@ref) """ -is_shared(a::oneArray) = isa(a.data[], oneL0.SharedBuffer) +is_shared(a::oneArray) = isa(a.data[].mem, oneL0.SharedBuffer) """ is_host(a::oneArray) -> Bool @@ -268,7 +268,7 @@ Host memory resides on the CPU but is visible to the GPU, useful for staging dat See also: [`is_device`](@ref), [`is_shared`](@ref) """ -is_host(a::oneArray) = isa(a.data[], oneL0.HostBuffer) +is_host(a::oneArray) = isa(a.data[].mem, oneL0.HostBuffer) ## derived types @@ -350,34 +350,41 @@ Base.convert(::Type{T}, x::T) where T <: oneArray = x ## interop with libraries +# host pointers wait for outstanding work on the memory, so that the host sees its results +# and doesn't race with it function Base.unsafe_convert(::Type{Ptr{T}}, x::oneArray{T}) where {T} - buf = x.data[] if is_device(x) throw(ArgumentError("cannot take the CPU address of a $(typeof(x))")) end - convert(Ptr{T}, x.data[]) + x.offset + managed = x.data[] + maybe_synchronize(managed) + convert(Ptr{T}, managed.mem) + x.offset end +# device pointers record the task's stream as the memory's user, assuming the pointer is +# used by work submitted there. operations that know their submission target take +# ownership themselves, and use `device_pointer` instead. function Base.unsafe_convert(::Type{ZePtr{T}}, x::oneArray{T}) where {T} - convert(ZePtr{T}, x.data[]) + x.offset + managed = x.data[] + @lock managed.lock take_ownership!(managed, global_stream(context(x), device())) + device_pointer(x) end +device_pointer(x::oneArray{T}, i::Integer=1) where {T} = + convert(ZePtr{T}, x.data[].mem) + x.offset + (i-1) * Base.elsize(x) + ## indexing # Host-accessible arrays can be indexed from CPU, bypassing GPUArrays restrictions. -# Wait for queued work first, e.g., the kernel computing the result of a reduction. This -# only synchronizes the current task's stream; work submitted by other tasks, or to an -# explicitly created queue, needs to be synchronized explicitly. +# Taking the host pointer waits for outstanding work on the memory. function Base.getindex(x::oneArray{<:Any, <:Any, <:Union{oneL0.HostBuffer, oneL0.SharedBuffer}}, I::Int) @boundscheck checkbounds(x, I) - synchronize(global_stream(context(x), device())) return unsafe_load(pointer(x, I; type = oneL0.HostBuffer)) end function Base.setindex!(x::oneArray{<:Any, <:Any, <:Union{oneL0.HostBuffer, oneL0.SharedBuffer}}, v, I::Int) @boundscheck checkbounds(x, I) - synchronize(global_stream(context(x), device())) return unsafe_store!(pointer(x, I; type = oneL0.HostBuffer), v) end @@ -397,7 +404,7 @@ function typetagdata(a::oneArray, i=1) # for zero-size element types (e.g. singleton unions), the byte offset # is always zero, so the corresponding element offset is also zero elem_offset = iszero(Base.elsize(a)) ? 0 : a.offset ÷ Base.elsize(a) - return convert(ZePtr{UInt8}, a.data[]) + a.maxsize + elem_offset + i - 1 + return convert(ZePtr{UInt8}, a.data[].mem) + a.maxsize + elem_offset + i - 1 end function Base.copyto!(dest::oneArray{T}, doffs::Integer, src::Array{T}, soffs::Integer, @@ -445,11 +452,14 @@ Base.copyto!(dest::oneDenseArray{T}, src::oneDenseArray{T}) where {T} = function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dest::oneDenseArray{T}, doffs, src::Array{T}, soffs, n) where T + stream = global_stream(ctx, dev) GC.@preserve src dest begin - unsafe_copyto!(ctx, dev, pointer(dest, doffs), pointer(src, soffs), n) + with_ownership(stream, Managed[dest.data[]]) do + unsafe_copyto!(ctx, dev, device_pointer(dest, doffs), pointer(src, soffs), n) + end # Keep pageable host memory alive until the queued copy completes. - synchronize(global_stream(ctx, dev)) + synchronize(stream) end if Base.isbitsunion(T) # copy selector bytes @@ -461,21 +471,30 @@ end function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dest::Array{T}, doffs, src::oneDenseArray{T}, soffs, n) where T - GC.@preserve src dest unsafe_copyto!(ctx, dev, pointer(dest, doffs), pointer(src, soffs), n) + stream = global_stream(ctx, dev) + GC.@preserve src dest begin + with_ownership(stream, Managed[src.data[]]) do + unsafe_copyto!(ctx, dev, pointer(dest, doffs), device_pointer(src, soffs), n) + end + end if Base.isbitsunion(T) # copy selector bytes error("oneArray does not yet support isbits-union arrays") end # copies to the host are synchronizing - synchronize(global_stream(context(src), device())) + synchronize(stream) return dest end function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dest::oneDenseArray{T}, doffs, src::oneDenseArray{T}, soffs, n) where T - GC.@preserve src dest unsafe_copyto!(ctx, dev, pointer(dest, doffs), pointer(src, soffs), n) + GC.@preserve src dest begin + with_ownership(global_stream(ctx, dev), Managed[dest.data[], src.data[]]) do + unsafe_copyto!(ctx, dev, device_pointer(dest, doffs), device_pointer(src, soffs), n) + end + end if Base.isbitsunion(T) # copy selector bytes error("oneArray does not yet support isbits-union arrays") @@ -487,20 +506,14 @@ end function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dest::oneDenseArray{T,<:Any,<:Union{oneL0.SharedBuffer,oneL0.HostBuffer}}, doffs, src::Array{T}, soffs, n) where T - # maintain queue-ordered semantics - synchronize(global_stream(ctx, dev)) - if Base.isbitsunion(T) # copy selector bytes error("oneArray does not yet support isbits-union arrays") end - GC.@preserve src dest begin - ptr = pointer(dest, doffs) + # the host pointer waits for outstanding work; hold the lock so that no new work is + # submitted while copying + GC.@preserve src dest with_managed_locks(Managed[dest.data[]]) do unsafe_copyto!(pointer(dest, doffs; type=oneL0.HostBuffer), pointer(src, soffs), n) - if Base.isbitsunion(T) - # copy selector bytes - error("oneArray does not yet support isbits-union arrays") - end end return dest @@ -508,20 +521,14 @@ end function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice, dest::Array{T}, doffs, src::oneDenseArray{T,<:Any,<:Union{oneL0.SharedBuffer,oneL0.HostBuffer}}, soffs, n) where T - # maintain queue-ordered semantics - synchronize(global_stream(ctx, dev)) - if Base.isbitsunion(T) # copy selector bytes error("oneArray does not yet support isbits-union arrays") end - GC.@preserve src dest begin - ptr = pointer(dest, doffs) + # the host pointer waits for outstanding work; hold the lock so that no new work is + # submitted while copying + GC.@preserve src dest with_managed_locks(Managed[src.data[]]) do unsafe_copyto!(pointer(dest, doffs), pointer(src, soffs; type=oneL0.HostBuffer), n) - if Base.isbitsunion(T) - # copy selector bytes - error("oneArray does not yet support isbits-union arrays") - end end return dest @@ -609,15 +616,17 @@ a shared buffer (as created by `oneArray{T}(undef, ...)`) or a host buffer. The returned `Array` does **not** keep `arr` alive. The caller has to keep a reference to `arr` for as long as the `Array`, or anything derived from it, is used; otherwise - the `Array` may end up referring to freed memory. Device operations execute - asynchronously, so call `synchronize()` before accessing the returned array after - using `arr` on the device. + the `Array` may end up referring to freed memory. Wrapping waits for device + operations that use `arr`, but operations submitted afterwards execute + asynchronously: before accessing the returned array again, wait for them by + synchronizing the stream or queue they were submitted to, or by taking the host + pointer of `arr`, e.g., with `pointer(arr; type=oneL0.HostBuffer)`. """ function Base.unsafe_wrap(::Type{Array}, arr::oneArray{T,N,<:Union{oneL0.SharedBuffer,oneL0.HostBuffer}}) where {T,N} # TODO: can we make this more convenient by increasing the buffer's refcount and using # a finalizer on the Array? does that work when taking views etc of the Array? - ptr = reinterpret(Ptr{T}, pointer(arr)) + ptr = pointer(arr; type=oneL0.HostBuffer) unsafe_wrap(Array, ptr, size(arr)) end @@ -646,13 +655,14 @@ function Base.resize!(a::oneVector{T}, n::Integer) where {T} ctx = context(a) dev = device(a) buf = allocate(buftype(a), ctx, dev, bufsize, Base.datatype_alignment(T)) - ptr = convert(ZePtr{T}, buf) + new_data = DataRef(Managed(buf)) do managed + release(managed.mem) + end m = min(length(a), n) if m > 0 - unsafe_copyto!(ctx, dev, ptr, pointer(a), m) - end - new_data = DataRef(buf) do buf - release(buf) + with_ownership(global_stream(ctx, dev), Managed[new_data[], a.data[]]) do + unsafe_copyto!(ctx, dev, convert(ZePtr{T}, buf), device_pointer(a), m) + end end unsafe_free!(a) diff --git a/src/memory.jl b/src/memory.jl index 4cb619706..24ecb3335 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -58,3 +58,91 @@ function unsafe_fill!(ctx::ZeContext, dev::ZeDevice, ptr::Union{Ptr{T},ZePtr{T}} append_fill!(list, ptr, pattern, sizeof(T), bytes) end end + + +## managed memory + +# An allocation together with the submission target that last used it, so that host +# accesses only wait when there may be outstanding work, and work from another task first +# waits for the previous user. Shared by all arrays derived from the same allocation. +mutable struct Managed{B<:oneL0.AbstractBuffer} + const mem::B + const lock::ReentrantLock + + # the stream or explicitly created queue that last used the memory + owner::Union{Nothing, oneStream, ZeCommandQueue} + + # whether the owner may still have work using the memory. atomic so that host + # accesses can check it without taking the lock. + Base.@atomic dirty::Bool + + Managed(mem::B) where {B} = new{B}(mem, ReentrantLock(), nothing, false) +end + +Base.sizeof(managed::Managed) = sizeof(managed.mem) + +# wait for all work on a stream, possibly from another task. unlike `synchronize`, this +# leaves `mkl_dirty` alone, which only the stream's own task may modify. +function drain(s::oneStream) + oneL0.nonblocking_synchronize(s.list) + q = Base.@atomic s.queue + q === nothing || oneL0.nonblocking_synchronize(q) + return +end +drain(q::ZeCommandQueue) = oneL0.nonblocking_synchronize(q) + +# the caller must hold `managed.lock` +function synchronize_owner!(managed::Managed) + if Base.@atomic managed.dirty + drain(managed.owner) + Base.@atomic managed.dirty = false + end + return +end + +# record `owner` as the user of the memory, waiting for the previous one if it differs. +# the caller must hold `managed.lock` until the operation has been submitted to `owner`, +# so that other tasks never observe the ownership without the work it describes. +function take_ownership!(managed::Managed, owner::Union{oneStream, ZeCommandQueue}) + sizeof(managed) == 0 && return + if managed.owner !== owner + synchronize_owner!(managed) + managed.owner = owner + end + Base.@atomic managed.dirty = true + return +end + +# wait for outstanding work before accessing the memory from the host +function maybe_synchronize(managed::Managed) + Base.@atomic(managed.dirty) || return + @lock managed.lock synchronize_owner!(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}) + locked = sort!(unique!(managed); by=objectid) + nlocked = 0 + try + for m in locked + lock(m.lock) + nlocked += 1 + end + f() + finally + for i in nlocked:-1:1 + unlock(locked[i].lock) + end + end +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() + end +end diff --git a/test/array.jl b/test/array.jl index c09cfbd69..14bdfb254 100644 --- a/test/array.jl +++ b/test/array.jl @@ -156,6 +156,55 @@ end end end +# https://github.com/JuliaGPU/oneAPI.jl/issues/652: host accesses wait for outstanding work +# on the memory, also when it was submitted by another task. The kernel takes long enough +# for a missing wait to be observed. +function slow_increment(x, n) + for _ in 1:n + x = x * 1.0000001f0 + 1f0 + end + return x +end +function slow_increment!(a, n) + i = get_global_id() + @inbounds a[i] = slow_increment(a[i], n) + return +end +@testset "host access to $B arrays waits for GPU work" for B in (oneL0.SharedBuffer, oneL0.HostBuffer) + n = 1_000_000 + launch(a) = @oneapi items=256 groups=cld(length(a), 256) slow_increment!(a, n) + once = slow_increment(0f0, n) + twice = slow_increment(once, n) + + # broadcast, then scalar indexing + a = oneArray{Float32, 1, B}(undef, 1024) + a .= 42 + @test a[1] == a[end] == 42 + + # a kernel submitted by another task + a = oneArray{Float32, 1, B}(zeros(Float32, 1024)) + wait(@async launch(a)) + @test a[1] ≈ a[end] ≈ once + + # wrapping the memory in an Array + a = oneArray{Float32, 1, B}(zeros(Float32, 1024)) + launch(a) + A = unsafe_wrap(Array, a) + @test A[1] ≈ A[end] ≈ once + + # copying to an Array, which happens on the CPU + a = oneArray{Float32, 1, B}(zeros(Float32, 1024)) + wait(@async launch(a)) + A = Array(a) + @test A[1] ≈ A[end] ≈ once + + # work from another task has to wait for this task's work on the same memory + a = oneArray{Float32, 1, B}(zeros(Float32, 1024)) + launch(a) + wait(@async launch(a)) + @test a[1] ≈ a[end] ≈ twice +end + # https://github.com/JuliaGPU/oneAPI.jl/issues/661: kernels have to declare indirect access, # or an allocation that reuses the address of a freed one isn't visible to them. Whether a # reused shared allocation reads as zeros depends on allocator state, so churn through From 5eaad195653e631f9a533d49bd14b9cc74748b0f Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 12:59:45 +0200 Subject: [PATCH 3/4] Take ownership of kernel arguments at launch Kernel arguments took ownership of their memory while being converted, which happens before the kernel is submitted and without knowing where it will be submitted to. Another task could take ownership in between, losing the dependency on this kernel, and a kernel submitted to an explicitly created queue recorded the task's stream as the owner, so host accesses didn't wait for it. Make argument conversion collect the allocations instead, and have the launch take ownership on the actual queue while holding their locks until the kernel has been submitted, as OpenCL.jl does. The callable a kernel was compiled from is converted again at launch to collect the memory it captures; `@oneapi` and KernelInterface launches pass it along. --- src/compiler/execution.jl | 39 ++++++++++++++++++++++++++++++++------- src/oneAPIKernels.jl | 2 +- test/array.jl | 6 ++++++ 3 files changed, 39 insertions(+), 8 deletions(-) diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 961184779..2c02b0dc7 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -139,7 +139,7 @@ end kernel_tt = Tuple{map(Core.Typeof, kernel_args)...} kernel = zefunction(kernel_f, kernel_tt; compiler_kwargs...) if launch - kernel(args...; call_kwargs...) + launch_managed(kernel, f, args; call_kwargs...) end kernel end @@ -148,14 +148,23 @@ end ## argument conversion -struct KernelAdaptor end +# conversion doesn't take ownership of the arrays' memory, since the launch target isn't +# known yet. instead, the allocations are collected, if requested, so that the launch can +# take ownership of them right before submitting the kernel. +struct KernelAdaptor + managed::Union{Nothing, Vector{Managed}} +end +KernelAdaptor() = KernelAdaptor(nothing) # convert oneAPI host pointers to device pointers Adapt.adapt_storage(to::KernelAdaptor, p::ZePtr{T}) where {T} = reinterpret(Ptr{T}, p) # convert oneAPI host arrays to device arrays -Adapt.adapt_storage(::KernelAdaptor, xs::oneArray{T,N}) where {T,N} = - Base.unsafe_convert(oneDeviceArray{T,N,AS.CrossWorkgroup}, xs) +function Adapt.adapt_storage(to::KernelAdaptor, xs::oneArray{T,N}) where {T,N} + to.managed === nothing || push!(to.managed, xs.data[]) + ptr = reinterpret(LLVMPtr{T,AS.CrossWorkgroup}, device_pointer(xs)) + oneDeviceArray{T,N,AS.CrossWorkgroup}(size(xs), ptr, xs.maxsize - xs.offset) +end # Base.RefValue isn't GPU compatible, so provide a compatible alternative. # TODO: port improvements from CUDA.jl @@ -195,7 +204,7 @@ input object `x` as-is. Do not add methods to this function, but instead extend the underlying Adapt.jl package and register methods for the the `oneAPI.KernelAdaptor` type. """ -kernel_convert(arg) = adapt(KernelAdaptor(), arg) +kernel_convert(arg, managed=nothing) = adapt(KernelAdaptor(managed), arg) ## abstract kernel functionality @@ -418,9 +427,25 @@ 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)) + launch_managed(kernel, nothing, args) @inline Core.kwcall(kwargs::NamedTuple, kernel::HostKernel, args::Vararg{Any,N}) where {N} = - call(kernel, map(kernel_convert, args); kwargs...) + launch_managed(kernel, nothing, args; kwargs...) + +# launch a kernel, taking ownership of the memory its arguments use, and of the memory +# captured by `source`, the callable the kernel was compiled from (if known): its converted +# form, which the kernel holds, only contains pointers. ownership is taken with the +# allocations locked until the kernel has been submitted, so that other tasks never observe +# it without the work it describes. +@inline function launch_managed(kernel::HostKernel, source, args::Tuple; + queue::Union{oneStream, ZeCommandQueue}=global_stream(context(), device()), + kwargs...) + managed = Managed[] + source === nothing || kernel_convert(source, managed) + converted = map(arg -> kernel_convert(arg, managed), args) + GC.@preserve source args with_ownership(queue, managed) do + call(kernel, converted; queue, kwargs...) + end +end ## TODO: device-side kernels diff --git a/src/oneAPIKernels.jl b/src/oneAPIKernels.jl index 664aed0ac..6923e14b7 100644 --- a/src/oneAPIKernels.jl +++ b/src/oneAPIKernels.jl @@ -139,7 +139,7 @@ function KI.launch(obj::KI.Kernel{oneAPIBackend}, groups::Dims{3}, items::Dims{3 # stream for the active ones (; f, host) = obj.kern check_target(host.fun.mod, get(kwargs, :queue, nothing)) - GC.@preserve f host(args...; items, groups, kwargs...) + oneAPI.launch_managed(host, f, args; items, groups, kwargs...) return end diff --git a/test/array.jl b/test/array.jl index 14bdfb254..d589af2f6 100644 --- a/test/array.jl +++ b/test/array.jl @@ -203,6 +203,12 @@ end launch(a) wait(@async launch(a)) @test a[1] ≈ a[end] ≈ twice + + # a kernel submitted to an explicitly created queue + queue = oneL0.ZeCommandQueue(context(), device(); mode=oneL0.ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS) + a = oneArray{Float32, 1, B}(zeros(Float32, 1024)) + @oneapi items=256 groups=cld(length(a), 256) queue=queue slow_increment!(a, n) + @test a[1] ≈ a[end] ≈ once end # https://github.com/JuliaGPU/oneAPI.jl/issues/661: kernels have to declare indirect access, From ff98651ddb3e35130f7edca324e417071dc9fd81 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 15:06:09 +0200 Subject: [PATCH 4/4] Take ownership of memory captured by kernels compiled without launching A kernel returned by `@oneapi launch=false` only held the converted callable, so launching it didn't take ownership of the arrays the callable captures, and host accesses to them didn't wait for it. On main, scalar indexing synchronized the task's stream, so this was a regression for such kernels. Those kernels now keep the callable they were compiled from. The cached kernel objects can't, since they're shared by every callable with the same converted form. --- src/compiler/execution.jl | 17 ++++++++++++++--- test/array.jl | 8 ++++++++ 2 files changed, 22 insertions(+), 3 deletions(-) diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl index 2c02b0dc7..614538c89 100644 --- a/src/compiler/execution.jl +++ b/src/compiler/execution.jl @@ -140,8 +140,10 @@ end kernel = zefunction(kernel_f, kernel_tt; compiler_kwargs...) if launch launch_managed(kernel, f, args; call_kwargs...) + kernel + else + with_source(kernel, f) end - kernel end end @@ -247,7 +249,16 @@ end struct HostKernel{F,TT} <: AbstractKernel{F,TT} f::F fun::ZeKernel + + # the callable the kernel was compiled from, whose memory launches take ownership of. + # only set on the kernels `@oneapi launch=false` returns, as the cached ones are shared + # by every callable that converts to `f`. + source::Any end +HostKernel{F,TT}(f, fun::ZeKernel) where {F,TT} = HostKernel{F,TT}(f, fun, nothing) + +with_source(kernel::HostKernel{F,TT}, source) where {F,TT} = + HostKernel{F,TT}(kernel.f, kernel.fun, source) # Upper bound on the spill (scratch) memory a single work-group may require, in bytes. # @@ -427,9 +438,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} = - launch_managed(kernel, nothing, args) + launch_managed(kernel, kernel.source, args) @inline Core.kwcall(kwargs::NamedTuple, kernel::HostKernel, args::Vararg{Any,N}) where {N} = - launch_managed(kernel, nothing, args; kwargs...) + launch_managed(kernel, kernel.source, args; kwargs...) # launch a kernel, taking ownership of the memory its arguments use, and of the memory # captured by `source`, the callable the kernel was compiled from (if known): its converted diff --git a/test/array.jl b/test/array.jl index d589af2f6..8534f1e6a 100644 --- a/test/array.jl +++ b/test/array.jl @@ -204,6 +204,14 @@ end wait(@async launch(a)) @test a[1] ≈ a[end] ≈ twice + # a kernel compiled without launching, accessing memory it captures + let a = oneArray{Float32, 1, B}(zeros(Float32, 1024)) + f = () -> (slow_increment!(a, n); return) + kernel = @oneapi launch=false f() + kernel(; items=256, groups=cld(length(a), 256)) + @test a[1] ≈ a[end] ≈ once + end + # a kernel submitted to an explicitly created queue queue = oneL0.ZeCommandQueue(context(), device(); mode=oneL0.ZE_COMMAND_QUEUE_MODE_ASYNCHRONOUS) a = oneArray{Float32, 1, B}(zeros(Float32, 1024))