ROCm/Windows gfx1151: per-process GPU VA space cap reproducer (hipMemAddressReserve / hipMemCreate share one ~256 GiB budget)

hip_va_cap_repro.py

"""Reproducer: per-process GPU virtual address space cap on ROCm/Windows.

Three phases:

  1. VMM control      - reserve 96 GiB, fill it entirely with 32 MiB
                        hipMemCreate+Map+SetAccess. Expected: succeeds.
  2. VA-only          - reserve 16 GiB windows in a loop until failure.
                        Observed: stops at 240 GiB.
  3. Combined         - reserve 229.2 GiB, then map 32 MiB pages into the last
                        window. Observed: hipErrorOutOfMemory once
                        reserved+mapped reaches ~256 GiB, with ~82 GB free.

Run each phase in a fresh process:
    python hip_va_cap_repro.py 1
    python hip_va_cap_repro.py 2
    python hip_va_cap_repro.py 3

Phase 2 may terminate the process with an access violation rather than
returning an error on runtimes that lack ROCm/rocm-systems#6051.

Only dependency is a HIP runtime. torch is imported solely to pull
amdhip64_7.dll into the process; replace with an explicit path if preferred.
"""

import ctypes
import os
import site
import sys

MB = 1 << 20
GB = 1 << 30
PAGE = 32 * MB


def load_hip():
    try:
        import torch  # noqa: F401
    except ImportError:
        pass
    names = ["amdhip64_7.dll", "amdhip64.dll", "amdhip64_6.dll"]
    for pkgs in site.getsitepackages():
        names.insert(0, os.path.join(pkgs, "_rocm_sdk_core", "bin",
                                     "amdhip64_7.dll"))
    for name in names:
        try:
            lib = ctypes.CDLL(name)
            print(f"HIP runtime: {name}\n")
            return lib
        except OSError:
            continue
    sys.exit("could not load a HIP runtime")


hip = load_hip()


class Loc(ctypes.Structure):
    _fields_ = [("type", ctypes.c_int), ("id", ctypes.c_int)]


class Flags(ctypes.Structure):
    _fields_ = [("compressionType", ctypes.c_ubyte),
                ("gpuDirectRDMACapable", ctypes.c_ubyte),
                ("usage", ctypes.c_ushort),
                ("reserved", ctypes.c_ubyte * 4)]


class Prop(ctypes.Structure):
    _fields_ = [("type", ctypes.c_int), ("requestedHandleTypes", ctypes.c_int),
                ("location", Loc), ("win32HandleMetaData", ctypes.c_void_p),
                ("allocFlags", Flags)]


class AccessDesc(ctypes.Structure):
    _fields_ = [("location", Loc), ("flags", ctypes.c_int)]


hip.hipMemAddressReserve.argtypes = [
    ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t, ctypes.c_size_t,
    ctypes.c_uint64, ctypes.c_ulonglong]
hip.hipMemCreate.argtypes = [ctypes.POINTER(ctypes.c_uint64), ctypes.c_size_t,
                             ctypes.POINTER(Prop), ctypes.c_ulonglong]
hip.hipMemMap.argtypes = [ctypes.c_uint64, ctypes.c_size_t, ctypes.c_size_t,
                          ctypes.c_uint64, ctypes.c_ulonglong]
hip.hipMemSetAccess.argtypes = [ctypes.c_uint64, ctypes.c_size_t,
                                ctypes.POINTER(AccessDesc), ctypes.c_size_t]
hip.hipMemGetInfo.argtypes = [ctypes.POINTER(ctypes.c_size_t),
                              ctypes.POINTER(ctypes.c_size_t)]


def free_mb():
    f, t = ctypes.c_size_t(), ctypes.c_size_t()
    hip.hipMemGetInfo(ctypes.byref(f), ctypes.byref(t))
    return f.value // MB, t.value // MB


def reserve(size_bytes):
    p = ctypes.c_uint64()
    rc = hip.hipMemAddressReserve(ctypes.byref(p), ctypes.c_size_t(size_bytes),
                                  ctypes.c_size_t(0), ctypes.c_uint64(0),
                                  ctypes.c_ulonglong(0))
    return (p.value if rc == 0 else None), rc


def make_prop_desc():
    prop = Prop()
    prop.type = 1                 # hipMemAllocationTypePinned
    prop.location.type = 1        # hipMemLocationTypeDevice
    prop.location.id = 0
    desc = AccessDesc()
    desc.location.type = 1
    desc.location.id = 0
    desc.flags = 3                # ReadWrite
    return prop, desc


def fill(base, cap_bytes, label):
    prop, desc = make_prop_desc()
    n = 0
    while n * PAGE < cap_bytes:
        h = ctypes.c_uint64()
        rc = hip.hipMemCreate(ctypes.byref(h), ctypes.c_size_t(PAGE),
                              ctypes.byref(prop), ctypes.c_ulonglong(0))
        if rc != 0:
            return n, "hipMemCreate", rc
        va = base + n * PAGE
        rc = hip.hipMemMap(ctypes.c_uint64(va), ctypes.c_size_t(PAGE),
                           ctypes.c_size_t(0), h, ctypes.c_ulonglong(0))
        if rc != 0:
            return n, "hipMemMap", rc
        rc = hip.hipMemSetAccess(ctypes.c_uint64(va), ctypes.c_size_t(PAGE),
                                 ctypes.byref(desc), ctypes.c_size_t(1))
        if rc != 0:
            return n, "hipMemSetAccess", rc
        n += 1
        if n % 200 == 0:
            print(f"    {label}: {n:5d} pages, {n * PAGE // MB:6d} MB mapped, "
                  f"free={free_mb()[0]} MB")
    return n, None, 0


def init():
    hip.hipInit(0)
    hip.hipSetDevice(0)
    s = ctypes.c_uint64()
    hip.hipMalloc(ctypes.byref(s), ctypes.c_size_t(MB))
    hip.hipFree(s)
    f, t = free_mb()
    print(f"device: free={f} MB total={t} MB\n")


def phase1():
    print("PHASE 1: VMM control, reserve 96 GiB and fill it completely")
    base, rc = reserve(96 * GB)
    if base is None:
        sys.exit(f"  reserve failed rc={rc}")
    n, stage, rc = fill(base, 96 * GB, "fill")
    f, _ = free_mb()
    if stage:
        print(f"\n  UNEXPECTED: stopped at {stage} rc={rc} after "
              f"{n * PAGE // MB} MB, free={f} MB")
    else:
        print(f"\n  OK: {n} handles, {n * PAGE // MB} MB mapped, free={f} MB")
        print("  -> neither handle count nor physical memory is the limit")


def phase2():
    print("PHASE 2: VA-only, reserve 16 GiB windows until failure")
    total = 0
    while total < 2048 * GB:
        p, rc = reserve(16 * GB)
        if p is None:
            print(f"\n  VA CEILING: {total // GB} GiB reserved, "
                  f"next 16 GiB failed rc={rc}")
            return
        total += 16 * GB
        print(f"  {total // GB:5d} GiB reserved (window at 0x{p:x})")


def phase3():
    print("PHASE 3: combined, 229.2 GiB reserved then map into the last window")
    windows = [("A", 13856), ("B", 110432), ("C", 110432)]  # MB
    reserved_mb = 0
    last = None
    for name, mb in windows:
        p, rc = reserve(mb * MB)
        if p is None:
            sys.exit(f"  window {name} ({mb} MB) failed rc={rc}")
        reserved_mb += mb
        last = (p, mb)
        print(f"  window {name}: {mb:7d} MB   total reserved "
              f"{reserved_mb:7d} MB ({reserved_mb / 1024:.1f} GiB)")

    base, cap_mb = last
    print(f"\n  mapping into window C (free={free_mb()[0]} MB)...")
    n, stage, rc = fill(base, cap_mb * MB, "map")
    f, t = free_mb()
    mapped = n * PAGE // MB
    print("\n  " + "-" * 58)
    print(f"  stopped at    : {stage} rc={rc} "
          f"({'hipErrorOutOfMemory' if rc == 2 else 'other'})")
    print(f"  mapped        : {mapped} MB ({n} pages)")
    print(f"  VA reserved   : {reserved_mb} MB ({reserved_mb / 1024:.1f} GiB)")
    print(f"  reserved+map  : {reserved_mb + mapped} MB "
          f"({(reserved_mb + mapped) / 1024:.2f} GiB)")
    print(f"  hipMemGetInfo : free={f} MB of {t} MB")
    print("  " + "-" * 58)


if __name__ == "__main__":
    phase = sys.argv[1] if len(sys.argv) > 1 else "3"
    init()
    {"1": phase1, "2": phase2, "3": phase3}[phase]()