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/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/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 a159ed8c1..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 @@ -103,6 +108,191 @@ 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. +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) + +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