Sub-allocate small device arrays

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.

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)

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)

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

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
 

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