Harden the test suite for memory-constrained GPUs

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

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.

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"

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