[maleadt]

Arrays backed by shared or host memory can be accessed from the CPU, so host accesses have to wait for GPU work on the memory. Until now, that only happened in scalar indexing (since #649) and in copies to and from Arrays, and only for the calling task's stream. Work submitted by another task or to an explicitly created queue was never waited for, and neither was work preceding unsafe_wrap(Array, a):

a = oneArray{Float32,1,oneL0.SharedBuffer}(zeros(Float32, 1024))
wait(@async @oneapi items=256 groups=4 kernel!(a))
a[1]   # 0.0: the kernel is still running

At the same time, each scalar access cost a task-local lookup and a driver call, ~920 ns per element on an Iris Xe.

This PR tracks outstanding work per allocation, as CUDA.jl and OpenCL.jl do. oneArray memory is now wrapped in a Managed object that records the stream (or queue) that last used it, and whether that use may still be in flight. Submitting work takes ownership, first waiting for the previous owner if it differs, which orders work across tasks. Taking a host pointer (scalar indexing, unsafe_wrap, host copies, CPU ccalls) waits for the owner only if the memory is dirty; the clean path is a single atomic load, bringing scalar indexing down to under 10 ns per element.

Ownership is taken transactionally, as in JuliaGPU/OpenCL.jl#472 and JuliaGPU/CUDA.jl#3229: kernel launches and array copies lock the allocations involved (in a stable order) and hold the locks until the operation has been submitted, so another task never sees ownership without the work it describes. Argument conversion therefore only collects allocations, and the launch records the actual target, including @oneapi queue=.... Waiting for another task's stream leaves its oneMKL bookkeeping (mkl_dirty) alone, and the companion queue field is now atomic, since other tasks read it.

Limitations: oneMKL calls and other raw ZePtr conversions take ownership on the task's stream without a transaction. That is sufficient because the oneMKL wrappers wait for completion before returning. Finalizers are unchanged and never wait on ownership; the LTS workaround in release still drains all streams.

New tests cover broadcast followed by scalar indexing, a kernel from another task read by the parent, unsafe_wrap and host copies after a kernel, and an explicit queue. On main, the cross-task, unsafe_wrap and copy cases fail.

Fixes #652

[github-actions]

Your PR requires formatting changes to meet the project's style guidelines.
Please consider running Runic (git runic main) to apply these changes.

Click here to view the suggested changes.

diff --git a/src/array.jl b/src/array.jl
index c3aaeb7..2fb6303 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{Managed{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,16 @@ 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(Managed(buf)) do managed
-          release(managed.mem)
+            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{Managed{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
@@ -227,11 +228,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[].mem)
+    return oneL0.context(A.data[].mem)
 end
 
 function device(A::oneArray)
-  return oneL0.device(A.data[].mem)
+    return oneL0.device(A.data[].mem)
 end
 
 buftype(x::oneArray) = buftype(typeof(x))
@@ -356,22 +357,22 @@ function Base.unsafe_convert(::Type{Ptr{T}}, x::oneArray{T}) where {T}
   if is_device(x)
     throw(ArgumentError("cannot take the CPU address of a $(typeof(x))"))
   end
-  managed = x.data[]
-  maybe_synchronize(managed)
-  convert(Ptr{T}, managed.mem) + x.offset
+    managed = x.data[]
+    maybe_synchronize(managed)
+    return 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}
-  managed = x.data[]
-  @lock managed.lock take_ownership!(managed, global_stream(context(x), device()))
-  device_pointer(x)
+    managed = x.data[]
+    @lock managed.lock take_ownership!(managed, global_stream(context(x), device()))
+    return 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)
+device_pointer(x::oneArray{T}, i::Integer = 1) where {T} =
+    convert(ZePtr{T}, x.data[].mem) + x.offset + (i - 1) * Base.elsize(x)
 
 
 ## indexing
@@ -404,7 +405,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[].mem) + 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,
@@ -452,14 +453,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)
+    stream = global_stream(ctx, dev)
   GC.@preserve src dest begin
-    with_ownership(stream, Managed[dest.data[]]) do
-      unsafe_copyto!(ctx, dev, device_pointer(dest, doffs), pointer(src, soffs), n)
-    end
+        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(stream)
+        synchronize(stream)
   end
   if Base.isbitsunion(T)
     # copy selector bytes
@@ -471,30 +472,30 @@ end
 
 function Base.unsafe_copyto!(ctx::ZeContext, dev::ZeDevice,
                              dest::Array{T}, doffs, src::oneDenseArray{T}, soffs, n) where T
-  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)
+    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
-  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(stream)
+    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 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)
+    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
-  end
   if Base.isbitsunion(T)
     # copy selector bytes
     error("oneArray does not yet support isbits-union arrays")
@@ -505,30 +506,32 @@ end
 # between Array and host-accessible oneArray
 
 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
+        dest::oneDenseArray{T, <:Any, <:Union{oneL0.SharedBuffer, oneL0.HostBuffer}}, doffs, src::Array{T}, soffs, n
+    ) where {T}
   if Base.isbitsunion(T)
     # copy selector bytes
     error("oneArray does not yet support isbits-union arrays")
   end
-  # 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)
+    # 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)
   end
 
   return dest
 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
+        dest::Array{T}, doffs, src::oneDenseArray{T, <:Any, <:Union{oneL0.SharedBuffer, oneL0.HostBuffer}}, soffs, n
+    ) where {T}
   if Base.isbitsunion(T)
     # copy selector bytes
     error("oneArray does not yet support isbits-union arrays")
   end
-  # 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)
+    # 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)
   end
 
   return dest
@@ -626,7 +629,7 @@ 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 = pointer(arr; type=oneL0.HostBuffer)
+    ptr = pointer(arr; type = oneL0.HostBuffer)
   unsafe_wrap(Array, ptr, size(arr))
 end
 
diff --git a/src/compiler/execution.jl b/src/compiler/execution.jl
index 614538c..7e582c5 100644
--- a/src/compiler/execution.jl
+++ b/src/compiler/execution.jl
@@ -154,7 +154,7 @@ end
 # 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}}
+    managed::Union{Nothing, Vector{Managed}}
 end
 KernelAdaptor() = KernelAdaptor(nothing)
 
@@ -162,10 +162,10 @@ KernelAdaptor() = KernelAdaptor(nothing)
 Adapt.adapt_storage(to::KernelAdaptor, p::ZePtr{T}) where {T} = reinterpret(Ptr{T}, p)
 
 # convert oneAPI host arrays to device arrays
-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)
+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))
+    return oneDeviceArray{T, N, AS.CrossWorkgroup}(size(xs), ptr, xs.maxsize - xs.offset)
 end
 
 # Base.RefValue isn't GPU compatible, so provide a compatible alternative.
@@ -206,7 +206,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, managed=nothing) = adapt(KernelAdaptor(managed), arg)
+kernel_convert(arg, managed = nothing) = adapt(KernelAdaptor(managed), arg)
 
 
 ## abstract kernel functionality
@@ -255,10 +255,10 @@ struct HostKernel{F,TT} <: AbstractKernel{F,TT}
     # 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)
+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)
+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.
 #
@@ -447,13 +447,15 @@ end
 # 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...)
+@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
+    return GC.@preserve source args with_ownership(queue, managed) do
         call(kernel, converted; queue, kwargs...)
     end
 end
diff --git a/src/context.jl b/src/context.jl
index 1ad4ae3..565bdb0 100644
--- a/src/context.jl
+++ b/src/context.jl
@@ -267,8 +267,10 @@ end
 function stream_queue(s::oneStream)
     q = Base.@atomic s.queue
     q === nothing || return q
-    q = ZeCommandQueue(s.ctx, s.dev; flags = oneL0.ZE_COMMAND_QUEUE_FLAG_IN_ORDER,
-                       priority = s.priority)
+    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
diff --git a/src/memory.jl b/src/memory.jl
index 24ecb33..23365d9 100644
--- a/src/memory.jl
+++ b/src/memory.jl
@@ -65,7 +65,7 @@ end
 # 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}
+mutable struct Managed{B <: oneL0.AbstractBuffer}
     const mem::B
     const lock::ReentrantLock
 
@@ -116,15 +116,15 @@ 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)
+    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)
+    locked = sort!(unique!(managed); by = objectid)
     nlocked = 0
-    try
+    return try
         for m in locked
             lock(m.lock)
             nlocked += 1
@@ -139,7 +139,7 @@ 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
+    return with_managed_locks(managed) do
         for m in managed
             take_ownership!(m, owner)
         end
diff --git a/test/array.jl b/test/array.jl
index 8534f1e..a8cbd70 100644
--- a/test/array.jl
+++ b/test/array.jl
@@ -160,63 +160,63 @@ end
 # 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
+    for _ in 1:n
+        x = x * 1.0000001f0 + 1.0f0
+    end
+    return x
 end
 function slow_increment!(a, n)
-  i = get_global_id()
-  @inbounds a[i] = slow_increment(a[i], n)
-  return
+    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
-
-  # 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))
+    n = 1_000_000
+    launch(a) = @oneapi items = 256 groups = cld(length(a), 256) slow_increment!(a, n)
+    once = slow_increment(0.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
-  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))
-  @oneapi items=256 groups=cld(length(a), 256) queue=queue slow_increment!(a, n)
-  @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
+
+    # 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))
+    @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,

[codecov]

Codecov Report

❌ Patch coverage is 96.11650% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 80.84%. Comparing base (218dc51) to head (ff98651).
⚠️ Report is 1 commits behind head on main.

Files with missing lines	Patch %	Lines
src/array.jl	91.66%	3 Missing ⚠️
src/compiler/execution.jl	94.44%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #686      +/-   ##
==========================================
- Coverage   81.35%   80.84%   -0.51%     
==========================================
  Files          52       51       -1     
  Lines        4092     4026      -66     
==========================================
- Hits         3329     3255      -74     
- Misses        763      771       +8     

☔ View full report in Codecov by Harness.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.