Source code for simvx.graphics.gpu.pipeline_cache

"""Persistent ``VkPipelineCache``: warm ``vkCreate*Pipelines`` across runs.

A pipeline cache never changes rendering output; it only lets the driver reuse
previously-compiled pipeline binaries instead of recompiling SPIR-V to GPU ISA on
every launch. On the render-modernization examples the first-frame stall is almost
entirely that recompilation (measured ~0.4 s recurring per run plus ~1.2 s of
first-run pipeline construction); seeding the driver from a saved cache removes the
recurring cost and warms the first run.

One cache per logical device, seeded from and flushed to a per-GPU file under the
user cache dir (``$XDG_CACHE_HOME/simvx/pipeline_cache`` or ``~/.cache/...``). The
file is keyed + header-validated by ``pipelineCacheUUID`` + vendor/device id, so a
foreign-GPU or driver-mismatched file is discarded rather than fed to the driver.
Every failure path degrades to no-cache (``NULL``): a missing, corrupt, unreadable,
or unwritable file is never fatal, and the render result is identical either way.

The :class:`~simvx.graphics.engine.Engine` owns the lifecycle: it calls
:func:`init_pipeline_cache` right after the device is created (before any pipeline
is built) and :func:`destroy_pipeline_cache` (which flushes to disk) at teardown.
Every ``vkCreate*Pipelines`` call site passes :func:`pipeline_cache_for` as its
``pipelineCache`` argument; when no cache is registered for that device (e.g. an
isolated pass test with no Engine) it returns ``NULL`` and behaviour is unchanged.
"""

from __future__ import annotations

import logging
import os
import struct
import tempfile
from pathlib import Path
from typing import Any

import vulkan as vk

log = logging.getLogger(__name__)

# device-pointer -> {"cache": VkPipelineCache, "path": Path}
_caches: dict[int, dict[str, Any]] = {}

# VkPipelineCacheHeaderVersionOne prefix: headerSize, headerVersion, vendorID,
# deviceID, pipelineCacheUUID[16]. Little-endian per the Vulkan spec.
_HEADER = struct.Struct("<III I 16s")


def _device_key(device: Any) -> int:
    try:
        return int(vk.ffi.cast("uintptr_t", device))
    except Exception:  # noqa: BLE001 -- exotic binding handle; fall back to identity
        return id(device)


def _cache_dir() -> Path:
    base = os.environ.get("XDG_CACHE_HOME") or (Path.home() / ".cache")
    return Path(base) / "simvx" / "pipeline_cache"


def _gpu_signature(physical_device: Any) -> tuple[int, int, bytes]:
    props = vk.vkGetPhysicalDeviceProperties(physical_device)
    uuid = bytes(bytearray(props.pipelineCacheUUID)[:16])
    return int(props.vendorID), int(props.deviceID), uuid


def _cache_path(physical_device: Any) -> Path:
    vendor, dev, uuid = _gpu_signature(physical_device)
    return _cache_dir() / f"{vendor:08x}-{dev:08x}-{uuid.hex()}.bin"


def _read_valid_blob(path: Path, physical_device: Any) -> bytes:
    """Return the on-disk cache bytes iff they match this GPU/driver, else empty."""
    try:
        data = path.read_bytes()
    except OSError:
        return b""
    if len(data) < _HEADER.size:
        return b""
    _length, version, vendor, dev, uuid = _HEADER.unpack_from(data, 0)
    want_vendor, want_dev, want_uuid = _gpu_signature(physical_device)
    if (
        version != vk.VK_PIPELINE_CACHE_HEADER_VERSION_ONE
        or vendor != want_vendor
        or dev != want_dev
        or uuid != want_uuid
    ):
        return b""  # foreign / stale cache; the driver would ignore it anyway
    return data


[docs] def init_pipeline_cache(device: Any, physical_device: Any) -> None: """Create the device's pipeline cache, seeded from disk when a valid file exists. Idempotent and never fatal: any failure leaves no cache registered, so :func:`pipeline_cache_for` returns ``NULL`` and pipeline creation is unchanged. """ key = _device_key(device) if key in _caches: return try: path = _cache_path(physical_device) blob = _read_valid_blob(path, physical_device) create_info = vk.VkPipelineCacheCreateInfo( initialDataSize=len(blob), pInitialData=blob if blob else None, ) cache = vk.vkCreatePipelineCache(device, create_info, None) _caches[key] = {"cache": cache, "path": path} log.debug("pipeline cache: seeded %d bytes from %s", len(blob), path) except Exception as exc: # noqa: BLE001 -- cache is an optimisation, never fatal log.debug("pipeline cache init skipped: %s", exc)
[docs] def pipeline_cache_for(device: Any) -> Any: """The device's ``VkPipelineCache`` handle, or ``NULL`` when none is registered.""" entry = _caches.get(_device_key(device)) return entry["cache"] if entry else vk.ffi.NULL
def _get_cache_data(device: Any, cache: Any) -> bytes: """Two-call ``vkGetPipelineCacheData`` (query size, then fetch) via the low API. The high-level ``vulkan`` binding does not wrap this call, so go through ``_callApi`` against the raw symbol. """ ffi = vk.ffi size_p = ffi.new("size_t*") if vk._vulkan._callApi(vk._vulkan.lib.vkGetPipelineCacheData, device, cache, size_p, ffi.NULL) != vk.VK_SUCCESS: return b"" n = int(size_p[0]) if n == 0: return b"" buf = ffi.new("char[]", n) if vk._vulkan._callApi(vk._vulkan.lib.vkGetPipelineCacheData, device, cache, size_p, buf) != vk.VK_SUCCESS: return b"" return bytes(ffi.buffer(buf, int(size_p[0])))
[docs] def save_pipeline_cache(device: Any) -> None: """Flush the (grown) pipeline cache to its per-GPU file. Never fatal.""" entry = _caches.get(_device_key(device)) if not entry: return try: blob = _get_cache_data(device, entry["cache"]) if not blob: return path: Path = entry["path"] path.parent.mkdir(parents=True, exist_ok=True) fd, tmp = tempfile.mkstemp(dir=str(path.parent), suffix=".tmp") try: with os.fdopen(fd, "wb") as handle: handle.write(blob) os.replace(tmp, path) # atomic swap so a crash never leaves a torn file finally: if os.path.exists(tmp): os.unlink(tmp) log.debug("pipeline cache: wrote %d bytes to %s", len(blob), path) except Exception as exc: # noqa: BLE001 -- cache is an optimisation, never fatal log.debug("pipeline cache save skipped: %s", exc)
[docs] def destroy_pipeline_cache(device: Any) -> None: """Flush the cache to disk and destroy the handle. Idempotent, never fatal.""" key = _device_key(device) entry = _caches.get(key) if not entry: return save_pipeline_cache(device) try: vk.vkDestroyPipelineCache(device, entry["cache"], None) except Exception as exc: # noqa: BLE001 log.debug("pipeline cache destroy skipped: %s", exc) _caches.pop(key, None)