[maleadt]

Allocating a oneArray called into the driver every time, and on an Iris Xe that costs ~120 µs for allocating and freeing a small array's memory, growing to ~700 µs with a thousand arrays alive: zeMemAllocDevice is slow, and the zeContextMakeMemoryResident call per allocation gets slower with every live allocation. Since every broadcast or reduction allocates its result or temporaries, that dominates the cost of operations on small arrays. CUDA.jl avoids this with CUDA's stream-ordered memory pools, which Level Zero doesn't have. NEO has an internal pool for small allocations, but it's only enabled for Battlemage, Lunar Lake and newer; it's disabled for Gen12LP (Iris Xe), Arc A-series, Meteor Lake and PVC.

Separately allocated arrays also slow down kernel launches on the i915 KMD: NEO backs every allocation with a buffer object of at least 64 KiB, and because kernels declare indirect access (#662), every live one goes into every submission, adding ~0.45 µs per live array to each submission.

This sub-allocates device arrays of up to 1 MiB from 2 MiB allocations, in power-of-two sizes from 256 bytes, as SYCL's Unified Runtime does for device and host USM by default. Slabs are made resident once. A freed chunk keeps the ownership state of the memory that used it (#686) and passes it on to the next array allocated there, so that array's first use from another stream waits for earlier work on the memory, and reuse on the same stream is ordered by the stream. That's also why we can't rely on NEO's pool where it exists: it makes freed memory reusable right away. Allocation prefers chunks that don't need that wait. Empty slabs are freed when memory runs low. ONEAPI_MEMORY_POOL=0 disables sub-allocation, e.g. to have out-of-bounds accesses fault rather than corrupt another array. Shared and host memory are not sub-allocated, as neither NEO nor SYCL do.

Measured on an Iris Xe (median; the machine was shared, so ranges are over repeated runs):

	#686	this PR
allocate and free a small array's memory	122 µs	0.9 µs
... with 1000 arrays alive	691 µs	0.8 µs
launch a kernel, 100 arrays alive	41–59 µs	12–16 µs
launch a kernel, 1000 arrays alive	497–661 µs	17–18 µs

Constructing a oneArray still costs ~25 µs on top of that, because of uncached device queries (#690). With launch batching (#693), a submission pays for the live allocations once per batch instead of once per launch, which takes the launch cost with 1000 live arrays to 23–44 µs without this PR and ~7 µs with it.

Arrays now go through one allocation entry point, which resize! uses as well.

The test suite passes on an Iris Xe with and without sub-allocation.

[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 eba51c8..11554c8 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
@@ -125,13 +125,14 @@ mutable struct oneArray{T,N,B} <: AbstractGPUArray{T,N}
     dev = device()
     alignment = Base.datatype_alignment(T)
     data = GPUArrays.cached_alloc((oneArray, B, ctx, dev, bufsize, alignment)) do
-        DataRef(release, alloc_managed(B, ctx, dev, bufsize, alignment))
+            DataRef(release, alloc_managed(B, ctx, dev, bufsize, alignment))
     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
@@ -224,11 +225,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))
@@ -353,22 +354,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
@@ -401,7 +402,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,
@@ -449,14 +450,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
@@ -468,30 +469,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")
@@ -502,30 +503,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
@@ -623,7 +626,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
 
@@ -651,8 +654,12 @@ function Base.resize!(a::oneVector{T}, n::Integer) where {T}
     # as a result, we can safely support resizing unowned buffers.
     ctx = context(a)
     dev = device(a)
-    new_data = DataRef(release, alloc_managed(buftype(a), ctx, dev, bufsize,
-                                              Base.datatype_alignment(T)))
+    new_data = DataRef(
+        release, alloc_managed(
+            buftype(a), ctx, dev, bufsize,
+            Base.datatype_alignment(T)
+        )
+    )
     m = min(length(a), n)
     if m > 0
         with_ownership(global_stream(ctx, dev), Managed[new_data[], a.data[]]) do
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 8208026..45afe24 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
 
@@ -76,7 +76,7 @@ mutable struct Managed{B<:oneL0.AbstractBuffer}
     # accesses can check it without taking the lock.
     Base.@atomic dirty::Bool
 
-    Managed(mem::B, owner=nothing, dirty::Bool=false) where {B} =
+    Managed(mem::B, owner = nothing, dirty::Bool = false) where {B} =
         new{B}(mem, ReentrantLock(), owner, dirty)
 end
 
@@ -117,15 +117,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
@@ -140,7 +140,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/src/pool.jl b/src/pool.jl
index 1e12599..1c65bac 100644
--- a/src/pool.jl
+++ b/src/pool.jl
@@ -123,7 +123,7 @@ function release(managed::Managed)
     if mem isa oneL0.DeviceBuffer && 0 < sizeof(mem) <= MAX_CHUNK
         pool_free(managed) && return
     end
-    release(mem)
+    return release(mem)
 end
 
 
@@ -170,8 +170,10 @@ struct DevicePool
     slabs::Dict{UInt, Slab}
 end
 const NCLASSES = trailing_zeros(MAX_CHUNK) - trailing_zeros(MIN_CHUNK) + 1
-DevicePool() = DevicePool([Slab[] for _ in 1:NCLASSES],
-                          [FreeChunk[] for _ in 1:NCLASSES], Dict{UInt, Slab}())
+DevicePool() = DevicePool(
+    [Slab[] for _ in 1:NCLASSES],
+    [FreeChunk[] for _ in 1:NCLASSES], Dict{UInt, Slab}()
+)
 
 const pools = Dict{Tuple{ZeContext, ZeDevice}, DevicePool}()
 
@@ -182,7 +184,7 @@ const pool_lock = Threads.SpinLock()
 function with_pool_lock(f)
     GC.enable_finalizers(false)
     lock(pool_lock)
-    try
+    return try
         f()
     finally
         unlock(pool_lock)
@@ -260,7 +262,7 @@ function pool_free(managed::Managed{oneL0.DeviceBuffer})
     ptr = pointer(mem)
     owner = managed.owner
     dirty = Base.@atomic managed.dirty
-    with_pool_lock() do
+    return with_pool_lock() do
         pool = get(pools, (oneL0.context(mem), oneL0.device(mem)), nothing)
         pool === nothing && return false
         slab = get(pool.slabs, UInt(ptr) & ~UInt(SLAB_SIZE - 1), nothing)
diff --git a/test/array.jl b/test/array.jl
index 22bb29d..6cbefe2 100644
--- a/test/array.jl
+++ b/test/array.jl
@@ -160,142 +160,144 @@ 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
 
 @testset "sub-allocation" begin
-  n = 1_000_000
-  once = slow_increment(0f0, n)
-  base(x) = pointer(oneL0.lookup_alloc(context(), pointer(x)))
-
-  if oneAPI.POOL[]
-    # small arrays share driver allocations, which every kernel launch would otherwise make
-    # resident one by one
-    a = oneArray{Float32}(undef, 16)
-    b = oneArray{Float32}(undef, 16)
-    @test base(a) == base(b)
-
-    c = oneArray{Float32}(undef, 1 << 20)
-    @test base(c) == pointer(c)
-  end
+    n = 1_000_000
+    once = slow_increment(0.0f0, n)
+    base(x) = pointer(oneL0.lookup_alloc(context(), pointer(x)))
+
+    if oneAPI.POOL[]
+        # small arrays share driver allocations, which every kernel launch would otherwise make
+        # resident one by one
+        a = oneArray{Float32}(undef, 16)
+        b = oneArray{Float32}(undef, 16)
+        @test base(a) == base(b)
+
+        c = oneArray{Float32}(undef, 1 << 20)
+        @test base(c) == pointer(c)
+    end
 
-  # memory reused on the same stream is only written after earlier work on it
-  a = oneArray(zeros(Float32, 1024))
-  @oneapi items=256 groups=4 slow_increment!(a, n)
-  ptr = pointer(a)
-  oneAPI.unsafe_free!(a)
-  b = oneArray{Float32}(undef, 1024)
-  oneAPI.POOL[] && @test pointer(b) == ptr
-  copyto!(b, ones(Float32, 1024))
-  @test Array(b) == ones(Float32, 1024)
-
-  # and from another task, only after waiting for it. that task only gets the memory if
-  # there are no other chunks of that size, so use them up.
-  n_chunk = oneAPI.MAX_CHUNK ÷ sizeof(Float32)
-  a = oneArray(zeros(Float32, n_chunk))
-  @oneapi items=256 groups=4 slow_increment!(a, n)
-  others = []
-  if oneAPI.POOL[]
-    class = oneAPI.chunk_class(oneAPI.MAX_CHUNK)
-    available() = oneAPI.with_pool_lock() do
-      pool = oneAPI.pools[(context(), device())]
-      !isempty(pool.fresh[class]) || !isempty(pool.free[class])
+    # memory reused on the same stream is only written after earlier work on it
+    a = oneArray(zeros(Float32, 1024))
+    @oneapi items = 256 groups = 4 slow_increment!(a, n)
+    ptr = pointer(a)
+    oneAPI.unsafe_free!(a)
+    b = oneArray{Float32}(undef, 1024)
+    oneAPI.POOL[] && @test pointer(b) == ptr
+    copyto!(b, ones(Float32, 1024))
+    @test Array(b) == ones(Float32, 1024)
+
+    # and from another task, only after waiting for it. that task only gets the memory if
+    # there are no other chunks of that size, so use them up.
+    n_chunk = oneAPI.MAX_CHUNK ÷ sizeof(Float32)
+    a = oneArray(zeros(Float32, n_chunk))
+    @oneapi items = 256 groups = 4 slow_increment!(a, n)
+    others = []
+    if oneAPI.POOL[]
+        class = oneAPI.chunk_class(oneAPI.MAX_CHUNK)
+        available() = oneAPI.with_pool_lock() do
+            pool = oneAPI.pools[(context(), device())]
+            !isempty(pool.fresh[class]) || !isempty(pool.free[class])
+        end
+        while available()
+            push!(others, oneArray{Float32}(undef, n_chunk))
+        end
     end
-    while available()
-      push!(others, oneArray{Float32}(undef, n_chunk))
+    ptr = oneAPI.device_pointer(a)
+    oneAPI.unsafe_free!(a)
+    stream = oneAPI.global_stream(context(), device())
+    inherited, b = fetch(
+        @async begin
+            b = oneArray{Float32}(undef, n_chunk)
+            # the new array is owned by the stream that used the memory, so that using it here
+            # waits for the kernel first (which this GPU might do anyway, executing work in order)
+            inherited = b.data[].owner === stream
+            copyto!(b, ones(Float32, n_chunk))
+            inherited, b
+        end
+    )
+    if oneAPI.POOL[]
+        @test oneAPI.device_pointer(b) == ptr
+        @test inherited
     end
-  end
-  ptr = oneAPI.device_pointer(a)
-  oneAPI.unsafe_free!(a)
-  stream = oneAPI.global_stream(context(), device())
-  inherited, b = fetch(@async begin
-    b = oneArray{Float32}(undef, n_chunk)
-    # the new array is owned by the stream that used the memory, so that using it here
-    # waits for the kernel first (which this GPU might do anyway, executing work in order)
-    inherited = b.data[].owner === stream
-    copyto!(b, ones(Float32, n_chunk))
-    inherited, b
-  end)
-  if oneAPI.POOL[]
-    @test oneAPI.device_pointer(b) == ptr
-    @test inherited
-  end
-  synchronize()
-  @test Array(b) == ones(Float32, n_chunk)
+    synchronize()
+    @test Array(b) == ones(Float32, n_chunk)
 
-  # many small arrays keep their contents, also when unused memory is given back
-  xs = [oneArray(fill(Float32(i), i)) for i in 1:300]
-  for i in 1:2:300
-    oneAPI.unsafe_free!(xs[i])
-  end
-  if oneAPI.POOL[]
-    before = oneAPI._allocated_bytes[]
-    ys = [oneArray{Float32}(undef, 1 << 16) for _ in 1:8]   # a slab of their own
-    foreach(oneAPI.unsafe_free!, ys)
-    oneAPI.trim_pools!()   # (which also frees other empty slabs)
-    @test oneAPI._allocated_bytes[] <= before
-  end
-  @test all(i -> Array(xs[i]) == fill(Float32(i), i), 2:2:300)
-  for i in 2:2:300
-    xs[i] = oneArray(fill(Float32(i), i))
-  end
-  @test all(i -> Array(xs[i]) == fill(Float32(i), i), 2:2:300)
+    # many small arrays keep their contents, also when unused memory is given back
+    xs = [oneArray(fill(Float32(i), i)) for i in 1:300]
+    for i in 1:2:300
+        oneAPI.unsafe_free!(xs[i])
+    end
+    if oneAPI.POOL[]
+        before = oneAPI._allocated_bytes[]
+        ys = [oneArray{Float32}(undef, 1 << 16) for _ in 1:8]   # a slab of their own
+        foreach(oneAPI.unsafe_free!, ys)
+        oneAPI.trim_pools!()   # (which also frees other empty slabs)
+        @test oneAPI._allocated_bytes[] <= before
+    end
+    @test all(i -> Array(xs[i]) == fill(Float32(i), i), 2:2:300)
+    for i in 2:2:300
+        xs[i] = oneArray(fill(Float32(i), i))
+    end
+    @test all(i -> Array(xs[i]) == fill(Float32(i), i), 2:2:300)
 end
 
 # https://github.com/JuliaGPU/oneAPI.jl/issues/661: kernels have to declare indirect access,
@@ -357,14 +359,14 @@ end
   resize!(c, 0)
   oneAPI.unsafe_free!(c)
 
-  # (with buffers big enough not to be sub-allocated, which `_allocated_bytes` doesn't track)
+    # (with buffers big enough not to be sub-allocated, which `_allocated_bytes` doesn't track)
   GC.enable(false)
   try
-    d = oneArray{Float32}(undef, 1 << 19)
+        d = oneArray{Float32}(undef, 1 << 19)
     before = oneAPI._allocated_bytes[]
     resize!(d, 1 << 20)
     oneAPI.unsafe_free!(d)
-    @test oneAPI._allocated_bytes[] == before - sizeof(Float32) * (1 << 19)
+        @test oneAPI._allocated_bytes[] == before - sizeof(Float32) * (1 << 19)
   finally
     GC.enable(true)
   end

[codecov]

Codecov Report

❌ Patch coverage is 96.07843% with 4 lines in your changes missing coverage. Please review.
✅ Project coverage is 81.18%. Comparing base (ff98651) to head (e56b908).

Files with missing lines	Patch %	Lines
src/pool.jl	95.83%	4 Missing ⚠️

Additional details and impacted files

@@                  Coverage Diff                  @@
##           tb/managed-memory     #692      +/-   ##
=====================================================
+ Coverage              80.84%   81.18%   +0.33%     
=====================================================
  Files                     51       51              
  Lines                   4026     4119      +93     
=====================================================
+ Hits                    3255     3344      +89     
- Misses                   771      775       +4     

☔ 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.