vram_probe.py - Empirically measure how much memory a GPU will actually allocate, using OpenCL.

vram_probe.py

#!/usr/bin/env python3
"""
vram_probe.py - Empirically measure how much memory a GPU will actually
allocate, using OpenCL.

Why this exists: on an integrated GPU with unified memory (like your Intel
chip), there is no fixed VRAM number. The GPU borrows from system RAM, so the
real question is "how big a device buffer can I get before allocation fails?"

Two things this script does that a naive version wouldn't:
  1. It WRITES to every buffer it allocates. Drivers allocate lazily, so a
     buffer you never touch may not be backed by physical memory and won't
     count. We fill each chunk to force real backing.
  2. It CAPS itself. Because GPU memory here is system RAM, probing all of it
     can invoke the Linux OOM killer and take down other processes. By default
     we stop at 85% of total RAM. Use --force to override (at your own risk).

Requirements:
  sudo apt install intel-opencl-icd ocl-icd-libopencl1 clinfo mesa-opencl-icd
  sudo apt install python3-pyopencl python3-numpy

  Plus an OpenCL runtime for your GPU. On Intel/Linux, one of:
    - intel-opencl-icd            (Intel Compute Runtime, best for Intel)
    - mesa-opencl-icd / rusticl   (Mesa's OpenCL)
  and the loader: ocl-icd-libopencl1

Usage:
  python3 vram_probe.py                 # probe with defaults
  python3 vram_probe.py --chunk-mb 128  # smaller step = finer resolution
  python3 vram_probe.py --max-gb 8      # hard ceiling regardless of RAM
  python3 vram_probe.py --force         # allow probing past 85% of RAM
"""

import argparse
import os
import sys

# For intel cards
os.environ["RUSTICL_ENABLE"] = "iris"

import numpy as np

try:
    import pyopencl as cl
except ImportError:
    sys.exit("PyOpenCL not found. Install with: pip install pyopencl numpy")


def total_ram_bytes():
    """Read MemTotal from /proc/meminfo (Linux)."""
    try:
        with open("/proc/meminfo") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    return int(line.split()[1]) * 1024  # kB -> bytes
    except OSError:
        pass
    return None


def pick_device():
    """Let the user choose a device if there's more than one; else auto-pick."""
    devices = []
    for platform in cl.get_platforms():
        for dev in platform.get_devices():
            devices.append((platform, dev))
    if not devices:
        sys.exit("No OpenCL devices found. Is an OpenCL runtime installed?")
    if len(devices) == 1:
        return devices[0][1]
    print("Multiple OpenCL devices found:")
    for i, (plat, dev) in enumerate(devices):
        print(f"  [{i}] {dev.name.strip()}  ({plat.name.strip()})")
    choice = input(f"Select device [0-{len(devices) - 1}]: ").strip() or "0"
    return devices[int(choice)][1]


def human(nbytes):
    return f"{nbytes / (1024 ** 3):.2f} GB ({nbytes / (1024 ** 2):.0f} MB)"


def main():
    ap = argparse.ArgumentParser(description="Probe max GPU memory allocation.")
    ap.add_argument("--chunk-mb", type=int, default=256,
                    help="Allocation step size in MB (default: 256).")
    ap.add_argument("--max-gb", type=float, default=None,
                    help="Hard ceiling in GB. Default: 85%% of system RAM.")
    ap.add_argument("--force", action="store_true",
                    help="Allow probing past 85%% of system RAM (risky: OOM).")
    args = ap.parse_args()

    dev = pick_device()
    print(f"\nDevice: {dev.name.strip()}")
    print(f"Vendor: {dev.vendor.strip()}")
    print(f"Reported global memory:  {human(dev.global_mem_size)}")
    print(f"Reported max single alloc: {human(dev.max_mem_alloc_size)}")

    ram = total_ram_bytes()
    if ram:
        print(f"System RAM:               {human(ram)}")

    # Work out the ceiling.
    chunk = args.chunk_mb * 1024 * 1024
    chunk -= chunk % 4  # keep it a multiple of 4 for uint32 fills
    if args.max_gb is not None:
        ceiling = int(args.max_gb * 1024 ** 3)
    elif ram is not None:
        ceiling = int(ram * (1.0 if args.force else 0.85))
    else:
        ceiling = dev.global_mem_size
    if not args.force and ram is not None:
        ceiling = min(ceiling, int(ram * 0.85))

    print(f"\nProbing in {args.chunk_mb} MB steps, ceiling {human(ceiling)}.")
    print("(Ctrl+C to stop early; all buffers are freed on exit.)\n")

    ctx = cl.Context(devices=[dev])
    queue = cl.CommandQueue(ctx)
    pattern = np.uint32(0)

    buffers = []
    allocated = 0
    try:
        while allocated + chunk <= ceiling:
            try:
                buf = cl.Buffer(ctx, cl.mem_flags.READ_WRITE, size=chunk)
                # Force the driver to actually back this memory, then wait.
                cl.enqueue_fill_buffer(queue, buf, pattern, 0, chunk)
                queue.finish()
            except (cl.MemoryError, cl.LogicError, cl.RuntimeError, MemoryError) as e:
                print(f"\nAllocation failed at {human(allocated)}: "
                      f"{type(e).__name__}")
                break
            buffers.append(buf)
            allocated += chunk
            print(f"\r  Committed: {human(allocated)}", end="", flush=True)
        else:
            print(f"\n\nReached ceiling without failing.")
    except KeyboardInterrupt:
        print("\n\nInterrupted by user.")

    print(f"\nMax successfully committed: {human(allocated)}")
    if ram:
        print(f"That's {100 * allocated / ram:.0f}% of system RAM.")
    if not args.force:
        print("Note: capped at 85% of RAM by default. Re-run with --force to "
              "push further (may trigger the OOM killer).")

    # Explicit cleanup so the memory is released promptly.
    for buf in buffers:
        buf.release()


if __name__ == "__main__":
    main()