From 14323c309ad74bfb23373e07833a44da9be6395b Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 15:07:31 +0200 Subject: [PATCH 1/2] Allocate and release array memory through a single entry point Arrays created their `Managed` allocation in two places, the constructor and `resize!`, which already diverged once (#667). Route both through `alloc_managed` and `release(::Managed)`, so that how array memory is allocated can change in one place. --- src/array.jl | 13 ++++--------- src/pool.jl | 7 +++++++ 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/src/array.jl b/src/array.jl index c3aaeb7b8..eba51c844 100644 --- a/src/array.jl +++ b/src/array.jl @@ -125,10 +125,7 @@ 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 - buf = allocate(B, ctx, dev, bufsize, alignment) - data = DataRef(Managed(buf)) do managed - release(managed.mem) - end + DataRef(release, alloc_managed(B, ctx, dev, bufsize, alignment)) end obj = new{T,N,B}(data, maxsize, 0, dims) finalizer(unsafe_free!, obj) @@ -654,14 +651,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) - buf = allocate(buftype(a), ctx, dev, bufsize, Base.datatype_alignment(T)) - new_data = DataRef(Managed(buf)) do managed - release(managed.mem) - end + 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 - unsafe_copyto!(ctx, dev, convert(ZePtr{T}, buf), device_pointer(a), m) + unsafe_copyto!(ctx, dev, convert(ZePtr{T}, new_data[].mem), device_pointer(a), m) end end unsafe_free!(a) diff --git a/src/pool.jl b/src/pool.jl index a159ed8c1..968e0a254 100644 --- a/src/pool.jl +++ b/src/pool.jl @@ -103,6 +103,13 @@ function allocate(::Type{oneL0.HostBuffer}, ctx, dev, bytes::Int, alignment::Int return buf end +# Allocate the memory backing an array, together with its ownership state, and release it. +# Arrays only allocate and free memory through these. +alloc_managed(::Type{B}, ctx, dev, bytes::Int, alignment::Int) where {B} = + Managed(allocate(B, ctx, dev, bytes, alignment)) + +release(managed::Managed) = release(managed.mem) + function release(buf::oneL0.AbstractBuffer) sizeof(buf) == 0 && return From e56b908236b64901c9d54ff0b407bfb08aa80908 Mon Sep 17 00:00:00 2001 From: Tim Besard Date: Tue, 13 Oct 2026 15:19:03 +0200 Subject: [PATCH 2/2] Sub-allocate small device arrays from bigger allocations MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Allocating a oneArray called into the driver every time. On an Iris Xe, allocating and freeing a small array's memory took ~120 µs, growing to ~700 µs with a thousand arrays alive, as zeMemAllocDevice is slow and the residency call per allocation gets slower with every live allocation. Broadcasts and reductions allocate their results and temporaries, so this dominated operations on small arrays. NEO's own pool for small allocations is disabled on Gen12LP, Arc A-series, Meteor Lake and PVC. Separate allocations also slow down submissions on i915: every live one is a buffer object of at least 64 KiB, and because kernels declare indirect access, each goes into every submission, ~0.45 µs per live allocation. Device allocations of up to 1 MiB are now chunks of 2 MiB slabs, in power-of-two sizes from 256 bytes, bringing allocation down to ~1 µs. A freed chunk keeps the ownership state of the memory that used it, and is handed out again with it: work on the same stream is ordered after the previous use, and anything else waits for it first. Allocation prefers chunks that don't need that wait. Empty slabs are freed when memory runs low. ONEAPI_MEMORY_POOL=0 disables sub-allocation. --- docs/src/memory.md | 11 +++ src/memory.jl | 3 +- src/oneAPI.jl | 4 + src/pool.jl | 191 ++++++++++++++++++++++++++++++++++++++++++++- test/array.jl | 84 +++++++++++++++++++- 5 files changed, 286 insertions(+), 7 deletions(-) diff --git a/docs/src/memory.md b/docs/src/memory.md index 18bb1d7c5..1d0fa0a3d 100644 --- a/docs/src/memory.md +++ b/docs/src/memory.md @@ -33,6 +33,17 @@ free(buf) `host_alloc` and `shared_alloc` allocate host and shared USM memory respectively, using the same signature. +## Sub-allocation + +Device arrays of up to 1 MiB are carved out of larger driver allocations of 2 MiB, which +are kept around when the arrays are freed. The driver backs every allocation with a +buffer object of at least 64 KiB, and has to make every live one resident for each kernel +launch, so with many small arrays separate allocations make allocating and launching +kernels slow. + +Set `ONEAPI_MEMORY_POOL=0` to give every array its own allocation, e.g., when debugging +out-of-bounds accesses, which may then fault instead of corrupting another array. + ## Garbage Collection Julia's garbage collector automatically manages `oneArray` objects. However, GPU memory is a limited resource. If you are running into out-of-memory errors, you might need to manually trigger garbage collection or free arrays. diff --git a/src/memory.jl b/src/memory.jl index 24ecb3335..8208026c9 100644 --- a/src/memory.jl +++ b/src/memory.jl @@ -76,7 +76,8 @@ mutable struct Managed{B<:oneL0.AbstractBuffer} # accesses can check it without taking the lock. Base.@atomic dirty::Bool - Managed(mem::B) where {B} = new{B}(mem, ReentrantLock(), nothing, false) + Managed(mem::B, owner=nothing, dirty::Bool=false) where {B} = + new{B}(mem, ReentrantLock(), owner, dirty) end Base.sizeof(managed::Managed) = sizeof(managed.mem) diff --git a/src/oneAPI.jl b/src/oneAPI.jl index f6489f6c2..15c13d8fa 100644 --- a/src/oneAPI.jl +++ b/src/oneAPI.jl @@ -136,6 +136,10 @@ function __init__() end if oneL0.functional[] + # slabs are process-local + empty!(pools) + POOL[] = oneL0.parse_env_bool("ONEAPI_MEMORY_POOL", true) + @static if Sys.iswindows() @warn """oneAPI.jl support for native Windows is experimental and incomplete. For the time being, it is recommended to use WSL or Linux instead.""" diff --git a/src/pool.jl b/src/pool.jl index 968e0a254..1e125997c 100644 --- a/src/pool.jl +++ b/src/pool.jl @@ -44,7 +44,7 @@ function _maybe_gc(dev, bytes) allocated = _allocated_bytes[] allocated <= 0 && return total_mem = _memory_limit(dev) - return if allocated + bytes > total_mem * 0.8 + if allocated + bytes > total_mem * 0.8 # Flush deferred resource releases (e.g., MKL sparse handles) from previous GC # cycles first — these are safe to release now because they were deferred earlier. # Do this BEFORE GC to avoid racing with new finalizers. @@ -53,7 +53,12 @@ function _maybe_gc(dev, bytes) GC.gc(true) elseif allocated + bytes > total_mem * 0.4 GC.gc(false) + else + return end + # the memory of collected arrays only becomes available once their slabs are empty + trim_pools!() + return end # GPUCompiler implements 8- and 16-bit atomics with 32-bit atomics on the containing aligned @@ -105,10 +110,188 @@ end # Allocate the memory backing an array, together with its ownership state, and release it. # Arrays only allocate and free memory through these. -alloc_managed(::Type{B}, ctx, dev, bytes::Int, alignment::Int) where {B} = - Managed(allocate(B, ctx, dev, bytes, alignment)) +function alloc_managed(::Type{B}, ctx, dev, bytes::Int, alignment::Int) where {B} + if B === oneL0.DeviceBuffer && POOL[] && 0 < bytes && max(bytes, alignment) <= MAX_CHUNK + managed = pool_alloc(ctx, dev, bytes, alignment) + managed === nothing || return managed + end + return Managed(allocate(B, ctx, dev, bytes, alignment)) +end + +function release(managed::Managed) + mem = managed.mem + if mem isa oneL0.DeviceBuffer && 0 < sizeof(mem) <= MAX_CHUNK + pool_free(managed) && return + end + release(mem) +end + + +## sub-allocation + +# Small device allocations are carved out of bigger ones, slabs. NEO backs every allocation +# with a buffer object of at least 64 KiB, and because kernels declare indirect access (see +# `link_kernel`), every live one is made resident for each submission, which on i915 costs +# ~0.45 µs per allocation and kernel launch. NEO's own pool doing this is disabled on Gen12LP +# and PVC. +# +# A freed chunk keeps the ownership state of the memory that used it, and is handed out +# again with that state: work on the same stream is ordered after the previous use anyway, +# and anything else waits for it first (see `take_ownership!`). Reuse thus never waits for +# the GPU, unlike freeing memory. + +# Opt-out knob, initialized from ONEAPI_MEMORY_POOL in `__init__`. +const POOL = Ref{Bool}(true) -release(managed::Managed) = release(managed.mem) +const SLAB_SIZE = 2^21 +const MIN_CHUNK = 256 +const MAX_CHUNK = 2^20 + +mutable struct Slab + const buf::oneL0.DeviceBuffer + const chunk::Int + # chunks handed out, and the first one never handed out + used::Int + next::Int +end + +struct FreeChunk + slab::Slab + ptr::ZePtr{Cvoid} + owner::Union{Nothing, oneStream, ZeCommandQueue} + dirty::Bool +end + +struct DevicePool + # per chunk size: slabs with chunks never handed out, and freed chunks + fresh::Vector{Vector{Slab}} + free::Vector{Vector{FreeChunk}} + # all slabs, by address + 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}()) + +const pools = Dict{Tuple{ZeContext, ZeDevice}, DevicePool}() + +# Protects the pools. Never held across anything that can yield, and with finalizers +# disabled, as chunks are freed from finalizers. +const pool_lock = Threads.SpinLock() + +function with_pool_lock(f) + GC.enable_finalizers(false) + lock(pool_lock) + try + f() + finally + unlock(pool_lock) + GC.enable_finalizers(true) + end +end + +chunk_class(chunk::Int) = trailing_zeros(chunk) - trailing_zeros(MIN_CHUNK) + 1 + +# how many freed chunks to consider, from the most recently freed one, looking for one that +# can be reused without waiting for another stream +const REUSE_WINDOW = 8 + +function take_chunk!(pool::DevicePool, class::Int, owner) + free = pool.free[class] + # a chunk last used by `owner`, or whose memory has no outstanding work + for i in length(free):-1:max(1, length(free) - REUSE_WINDOW + 1) + c = free[i] + if c.owner === owner || !c.dirty + deleteat!(free, i) + c.slab.used += 1 + return c + end + end + # a fresh chunk + fresh = pool.fresh[class] + if !isempty(fresh) + slab = last(fresh) + ptr = pointer(slab.buf) + slab.next * slab.chunk + slab.next += 1 + slab.next * slab.chunk == SLAB_SIZE && pop!(fresh) + slab.used += 1 + return FreeChunk(slab, ptr, nothing, false) + end + # any freed chunk + isempty(free) && return nothing + c = pop!(free) + c.slab.used += 1 + return c +end + +function pool_alloc(ctx, dev, bytes::Int, alignment::Int) + bytes, alignment = atomic_allocation(bytes, alignment) + chunk = max(MIN_CHUNK, nextpow(2, max(bytes, alignment))) + class = chunk_class(chunk) + owner = global_stream(ctx, dev) + + c = with_pool_lock() do + pool = get!(DevicePool, pools, (ctx, dev)) + take_chunk!(pool, class, owner) + end + if c === nothing + # `allocate` may collect garbage, freeing chunks, but those are only reused when + # this slab runs out + buf = allocate(oneL0.DeviceBuffer, ctx, dev, SLAB_SIZE, SLAB_SIZE) + if UInt(pointer(buf)) % SLAB_SIZE != 0 + # chunks are mapped to their slab by address + release(buf) + return nothing + end + slab = Slab(buf, chunk, 1, 1) + c = FreeChunk(slab, pointer(buf), nothing, false) + with_pool_lock() do + pool = pools[(ctx, dev)] + pool.slabs[UInt(pointer(buf))] = slab + push!(pool.fresh[class], slab) + end + end + return Managed(oneL0.DeviceBuffer(c.ptr, bytes, ctx, dev), c.owner, c.dirty) +end + +# Return a chunk to its pool, unless the memory wasn't sub-allocated. +function pool_free(managed::Managed{oneL0.DeviceBuffer}) + mem = managed.mem + ptr = pointer(mem) + owner = managed.owner + dirty = Base.@atomic managed.dirty + 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) + slab === nothing && return false + slab.used -= 1 + push!(pool.free[chunk_class(slab.chunk)], FreeChunk(slab, ptr, owner, dirty)) + return true + end +end + +# Free slabs none of whose chunks are in use. Also run when the driver runs out of memory. +function trim_pools!() + empty = with_pool_lock() do + empty = Slab[] + for pool in values(pools) + for (base, slab) in collect(pool.slabs) + slab.used == 0 || continue + push!(empty, slab) + delete!(pool.slabs, base) + class = chunk_class(slab.chunk) + filter!(s -> s !== slab, pool.fresh[class]) + filter!(c -> c.slab !== slab, pool.free[class]) + end + end + empty + end + # waits for outstanding work on the memory (see `release`) + foreach(slab -> release(slab.buf), empty) + return +end +oneL0.register_reclaim_callback!(trim_pools!) function release(buf::oneL0.AbstractBuffer) sizeof(buf) == 0 && return diff --git a/test/array.jl b/test/array.jl index 8534f1e6a..22bb29db0 100644 --- a/test/array.jl +++ b/test/array.jl @@ -219,6 +219,85 @@ end @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 + + # 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 + 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) + + # 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, # 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 @@ -278,13 +357,14 @@ end resize!(c, 0) oneAPI.unsafe_free!(c) + # (with buffers big enough not to be sub-allocated, which `_allocated_bytes` doesn't track) GC.enable(false) try - d = oneArray{Float32}(undef, 1024) + 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) * 1024 + @test oneAPI._allocated_bytes[] == before - sizeof(Float32) * (1 << 19) finally GC.enable(true) end