From ef0827e891e367e7d573091b31ea455385e6f58f Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 21 Jul 2026 09:57:26 -0500 Subject: [PATCH 1/2] Add ONEAPI_MEMORY_LIMIT to budget device memory across processes MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit The proactive-GC heuristic in the memory pool computes its pressure thresholds against the device's total memory, but _allocated_bytes only tracks the current process, and GC in one process cannot free buffers held by another. When several processes share one device — the parallel test suite being the prime example — each process happily grows to 80% of the card before feeling any pressure, so their combined footprint exhausts physical memory and allocations start failing with OutOfGPUMemoryError (hundreds of such errors on an 8GB Arc A750 with --jobs=2). Introduce ONEAPI_MEMORY_LIMIT (a byte count, or a percentage of device memory like "50%") as a per-process soft budget: _maybe_gc now scales its 40%/80% thresholds by the limit instead of the full card. This is deliberately soft — oversized single allocations still go through, with the existing retry_reclaim ladder as the reactive backstop. --- docs/src/memory.md | 18 ++++++++++++++++++ src/pool.jl | 29 ++++++++++++++++++++++++++++- 2 files changed, 46 insertions(+), 1 deletion(-) diff --git a/docs/src/memory.md b/docs/src/memory.md index f97146e8..0e4ea1e9 100644 --- a/docs/src/memory.md +++ b/docs/src/memory.md @@ -49,3 +49,21 @@ oneAPI.unsafe_free!(a) **Warning**: Only use `unsafe_free!` if you are sure the array is no longer used, including by any pending GPU operations. +## Memory Limit + +When multiple processes share one GPU, each process only tracks its own allocations, +and garbage collection in one process cannot free memory held by another. To keep the +processes from collectively exhausting the device, give each a budget with the +`ONEAPI_MEMORY_LIMIT` environment variable, either as a byte count or as a percentage +of device memory: + +``` +ONEAPI_MEMORY_LIMIT=50% # this process budgets half the device +ONEAPI_MEMORY_LIMIT=4000000000 # ... or an absolute byte count +``` + +This is a soft limit: it scales the thresholds at which oneAPI.jl proactively runs the +garbage collector to free stale GPU buffers, before falling back to a full collection +when an allocation actually fails. Allocations larger than the limit are still allowed +to succeed. The test suite sets this automatically for its parallel workers. + diff --git a/src/pool.jl b/src/pool.jl index ed3b75ba..fcc6a3a4 100644 --- a/src/pool.jl +++ b/src/pool.jl @@ -13,10 +13,37 @@ function _get_total_mem(dev) return _total_mem_cache[] end +# Per-process memory budget for the proactive-GC thresholds below. `_allocated_bytes` +# only tracks this process, and GC in this process cannot free buffers held by another, +# so when several processes share one device (e.g. parallel test workers) each must +# budget against its share of the card, not the whole card — otherwise their combined +# usage exceeds physical memory before any of them feels pressure. ONEAPI_MEMORY_LIMIT +# accepts a byte count ("4000000000") or a percentage of device memory ("50%"); unset +# means the full device. +const _memory_limit_cache = Threads.Atomic{Int64}(-1) + +function _memory_limit(dev) + cached = _memory_limit_cache[] + cached >= 0 && return cached + total = _get_total_mem(dev) + str = get(ENV, "ONEAPI_MEMORY_LIMIT", "") + limit = if isempty(str) + total + elseif endswith(str, "%") + pct = parse(Int, chop(str)) + 0 < pct <= 100 || error("ONEAPI_MEMORY_LIMIT percentage must be in (0, 100]") + total * pct ÷ 100 + else + parse(Int64, str) + end + Threads.atomic_cas!(_memory_limit_cache, Int64(-1), Int64(limit)) + return _memory_limit_cache[] +end + function _maybe_gc(dev, bytes) allocated = _allocated_bytes[] allocated <= 0 && return - total_mem = _get_total_mem(dev) + total_mem = _memory_limit(dev) return 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. From cb9ee324689fcde3c9b6d74313c1bda2d30a79b5 Mon Sep 17 00:00:00 2001 From: Michel Schanen Date: Tue, 21 Jul 2026 09:57:26 -0500 Subject: [PATCH 2/2] Harden the test suite for memory-constrained GPUs MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Three changes, motivated by a full-suite run on an 8GB Arc A750 with --jobs=2 failing 520 tests (366 of them OutOfGPUMemoryError, cascading across ~20 test files): - give each test worker an equal ONEAPI_MEMORY_LIMIT share of device memory, so the pool's proactive GC starts collecting before the workers collectively exhaust the card; - pass recycle_on_failure/retries to ParallelTestRunner (once it supports them, >= 2.7): a failing test file can leave its worker — or the driver — in a state where every subsequent allocation fails, so workers are recycled on failure and failed files get one retry, sequentially, on a fresh worker with an otherwise-idle device; - load BFloat16s.jl on the workers (Core.BFloat16 alone lacks the host-side conversions and rand sampling the testsuite's CPU reference path needs; this alone accounted for ~150 errors), and only include BFloat16 in the generic test element types when the host-side support is actually complete — probed at runtime, so the eltype enables itself once BFloat16s.jl implements the missing div/rem family. Device-side BFloat16 (the bfloat16.jl example) is exercised regardless. --- test/Project.toml | 1 + test/runtests.jl | 40 ++++++++++++++++++++++++++++++++++++++-- 2 files changed, 39 insertions(+), 2 deletions(-) diff --git a/test/Project.toml b/test/Project.toml index 4dd3fee0..788581e1 100644 --- a/test/Project.toml +++ b/test/Project.toml @@ -1,6 +1,7 @@ [deps] AbstractFFTs = "621f4979-c628-5d54-868e-fcf4e3e8185c" Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" +BFloat16s = "ab4f0b2a-ad5b-11e8-123f-65d77653426b" Dates = "ade2ca70-3891-5945-98fb-dc099432e06a" Distributed = "8ba89e20-285c-5b6f-9357-94700520ee1b" FFTW = "7a1cc6ca-52ef-59f5-83cd-3a7055c09341" diff --git a/test/runtests.jl b/test/runtests.jl index b0d9f555..37bf165d 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -35,6 +35,13 @@ args = parse_args(ARGS) # contention/oversubscription bugs. const spread_gpus = lowercase(get(ENV, "ONEAPI_TEST_SPREAD_GPUS", "")) in ("1", "true", "yes") worker_env = Vector{Pair{String, String}}() + +# Workers share one device but cannot see each other's allocations, so give each an equal +# share of device memory: the pool's proactive GC then kicks in before the workers +# collectively exhaust the card (important on 8GB consumer GPUs). +njobs = something(args.jobs, Some(ParallelTestRunner.default_njobs()))[] +push!(worker_env, "ONEAPI_MEMORY_LIMIT" => "$(max(1, 100 ÷ njobs))%") + device_claim_code = :() if spread_gpus ndev = length(oneAPI.devices()) @@ -64,6 +71,10 @@ end init_worker_code = quote $device_claim_code using oneAPI, Adapt + # BFloat16s.jl supplies the host-side numeric support (conversions, rand sampling) + # that Core.BFloat16 lacks on its own; the testsuite's CPU reference path needs it + # when BFloat16 is part of the tested element types. + using BFloat16s import GPUArrays include($gpuarrays_testsuite) @@ -83,7 +94,22 @@ init_worker_code = quote append!(eltypes, [Float64, ComplexF64]) end @static if isdefined(Core, :BFloat16) - const bfloat16_supported = oneAPI._device_supports_bfloat16() + # Exercising BFloat16 as a generic element type needs the testsuite's CPU + # reference path to work, which requires more host-side BFloat16 support than + # BFloat16s.jl currently implements (e.g. the div/rem family throws "rem not + # defined for BFloat16"). Probe the operations the testsuite relies on, so the + # element type enables itself automatically once BFloat16s.jl catches up. + # Device-side BFloat16 (oneMKL gemm/gemv, the bfloat16.jl example) is tested + # regardless. + bfloat16_host_support() = try + x, y = BFloat16(3), BFloat16(2) + mod(x, y); fld(x, y); rand(BFloat16) + true + catch + false + end + const bfloat16_supported = oneAPI._device_supports_bfloat16() && + bfloat16_host_support() if bfloat16_supported push!(eltypes, Core.BFloat16) end @@ -151,4 +177,14 @@ init_code = quote ..@grab_output, ..@on_device, ..sink end -runtests(oneAPI, args; testsuite, init_code, init_worker_code, env = worker_env) +# On memory-pressured GPUs (e.g. 8GB cards) a failing test can leave the worker — or even +# the driver — in a state where every subsequent allocation fails, so recycle workers on +# failure and give failed files one exclusive retry on an otherwise-idle device. +failure_handling = if pkgversion(ParallelTestRunner) >= v"2.7.0" + (; recycle_on_failure = true, retries = 1) +else + (;) +end + +runtests(oneAPI, args; testsuite, init_code, init_worker_code, env = worker_env, + failure_handling...)