hao-ai-lab-FastVideo/v2/platform/platform.py

"""Platform — the detected ``(device, arch)`` that resolves both registries (design_v3 §17).

One ``Platform`` per pool. It owns:
  * the **device fallback chain** (e.g. ``accel → cpu``) — so a component/kernel a backend hasn't
    implemented falls back to the numpy reference (the terminal rung + parity oracle);
  * the **arch fallback chain** per device (e.g. cuda ``sm100→sm90→sm80→ptx→generic``) — so a kernel
    built for an older arch still resolves on a newer card;
  * ``build_component(spec, instance)`` — the single materialization seam (replaces a bare
    ``spec.factory`` call); and
  * ``kernels`` — a per-platform ``KernelTable`` that resolves an op once and caches the callable.

``Platform.detect()`` is honest: it returns a CPU/numpy platform unless torch+CUDA are actually
present (they are not in this environment). The (device, arch) machinery is fully exercised on the
CPU box by the pure-python ``accel`` stand-in backend — proving the dispatch is generic without
faking a GPU.
"""
from __future__ import annotations

from typing import Any, Iterator

from .registry import COMPONENTS, KERNELS

_BACKENDS_LOADED = False


def ensure_backends_loaded() -> None:
    """Import the backend modules once so they self-register. Lazy → import-cycle-free."""
    global _BACKENDS_LOADED
    if _BACKENDS_LOADED:
        return
    _BACKENDS_LOADED = True
    # imported for side effects (each module calls register_component/register_kernel at import)
    from .backends import accel, cpu, torch_cuda  # noqa: F401


# Arch fallback chains (newest → oldest → portable terminal). Resolution walks these in order.
_CUDA_ARCH_FALLBACK = ("sm100", "sm90", "sm80", "ptx", "generic")
_ACCEL_ARCH_FALLBACK = ("sm90", "sm80", "generic")
_PORTABLE_ARCHS = ("ptx", "generic")   # vendor-neutral terminals — valid on any device


def _arch_rank(arch: str) -> int:
    """``smXY`` → integer ``XY`` for comparison; portable terminals (ptx/generic/unknown) → -1."""
    if arch.startswith("sm") and arch[2:].isdigit():
        return int(arch[2:])
    return -1


def _suffix_from(chain: tuple[str, ...], arch: str) -> tuple[str, ...]:
    """The arch fallback tail for ``arch``: itself, then ONLY older real archs from ``chain``, then
    portable terminals. Never falls back to a *newer* arch — a kernel built for sm90 is binary-
    incompatible on an sm70 device, so resolution must only ever degrade to older/more-portable code.
    (A newer device, e.g. sm120, may still run older-arch kernels — those are kept.)
    """
    rank = _arch_rank(arch)
    tail: list[str] = [arch]
    for a in chain:
        if a in tail:
            continue
        ar = _arch_rank(a)
        if ar < 0:                              # portable terminal (ptx/generic): always valid
            tail.append(a)
        elif rank >= 0 and ar < rank:           # strictly-older real arch: a safe fallback
            tail.append(a)
        # else: newer-or-equal real arch → skip (would be binary-incompatible on this device)
    return tuple(tail)


class KernelTable:
    """A platform's resolved op→callable map. Resolves each op once (over the fallback chain),
    then caches — so the hot loop pays the registry walk at most once per (op, variant)."""

    def __init__(self, platform: "Platform"):
        self.platform = platform
        self._cache: dict[tuple[str, str], Any] = {}

    def get(self, op: str, variant: str = "default") -> Any:
        ck = (op, variant)
        if ck not in self._cache:
            reg = self.platform.resolve_kernel(op, variant)
            if reg is None:
                raise KeyError(
                    f"no kernel for op={op!r} variant={variant!r} on device chain "
                    f"{self.platform.device_chain} (and no numpy terminal registered)")
            self._cache[ck] = reg.fn
        return self._cache[ck]

    def resolved_source(self, op: str, variant: str = "default") -> str:
        """Which registered cell an op resolves to — used by parity/diagnostics tests."""
        reg = self.platform.resolve_kernel(op, variant)
        return "" if reg is None else f"{reg.key}@{reg.source}"


class Platform:
    def __init__(self, device: str, arch: str, device_chain: tuple[str, ...],
                 arch_chains: dict[str, tuple[str, ...]]):
        self.device = device
        self.arch = arch
        self.device_chain = tuple(device_chain)
        self.arch_chains = dict(arch_chains)
        self._kernels: KernelTable | None = None

    # --- the kernel table (lazy, cached) ------------------------------------- #
    @property
    def kernels(self) -> KernelTable:
        if self._kernels is None:
            ensure_backends_loaded()
            self._kernels = KernelTable(self)
        return self._kernels

    def arch_chain(self, device: str) -> tuple[str, ...]:
        return self.arch_chains.get(device, ("generic",))

    # --- candidate-key generators (the fallback chains, in priority order) --- #
    def _component_keys(self, kind: str, variant: str) -> Iterator[tuple]:
        for dev in self.device_chain:
            yield (kind, dev, variant)
            if variant != "default":
                yield (kind, dev, "default")

    def _kernel_keys(self, op: str, variant: str) -> Iterator[tuple]:
        for dev in self.device_chain:
            for arch in self.arch_chain(dev):
                yield (op, dev, arch, variant)
                if variant != "default":
                    yield (op, dev, arch, "default")

    # --- the two resolution seams -------------------------------------------- #
    def build_component(self, spec: Any, instance: Any, variant: str = "default") -> Any:
        """Materialize a component. Registry first (device chain), then ``spec.factory`` as the
        terminal numpy rung — so existing cards (factory-only) are unchanged, while a GPU/accel
        backend overrides by registering ``(kind, device, …)``."""
        ensure_backends_loaded()
        reg = COMPONENTS.resolve_first(self._component_keys(spec.kind, variant))
        if reg is not None:
            return reg.fn(spec, instance, self)
        if getattr(spec, "factory", None) is not None:
            return spec.factory(instance)          # the (kind, "cpu", "numpy") terminal rung
        raise RuntimeError(
            f"component kind={spec.kind!r} ({getattr(spec, 'component_id', '?')!r}) has no impl on "
            f"device chain {self.device_chain} and no factory terminal")

    def resolve_kernel(self, op: str, variant: str = "default"):
        ensure_backends_loaded()
        return KERNELS.resolve_first(self._kernel_keys(op, variant))

    # --- constructors -------------------------------------------------------- #
    @classmethod
    def cpu(cls) -> "Platform":
        return cls("cpu", "numpy", ("cpu",), {"cpu": ("numpy",)})

    @classmethod
    def accel(cls, arch: str = "sm90") -> "Platform":
        """A pure-python stand-in accelerator (CPU-resident) used to prove cross-device dispatch,
        arch fallback, and the parity oracle without a real GPU."""
        return cls("accel", arch, ("accel", "cpu"),
                   {"accel": _suffix_from(_ACCEL_ARCH_FALLBACK, arch), "cpu": ("numpy",)})

    @classmethod
    def cuda(cls, arch: str = "sm90") -> "Platform":
        """A real CUDA platform (resolves torch adapters + fastvideo-kernel cells). Only usable on a
        box where those cells are ``available``; declared here so detection has a target."""
        return cls("cuda", arch, ("cuda", "cpu"),
                   {"cuda": _suffix_from(_CUDA_ARCH_FALLBACK, arch), "cpu": ("numpy",)})

    @classmethod
    def detect(cls) -> "Platform":
        """Honest detection: CUDA iff torch+CUDA are actually importable, else CPU/numpy."""
        ensure_backends_loaded()
        import importlib.util
        if importlib.util.find_spec("torch") is not None:
            try:
                import torch  # type: ignore
                if torch.cuda.is_available():
                    major, minor = torch.cuda.get_device_capability()
                    return cls.cuda(f"sm{major}{minor}")
            except Exception:
                pass
        return cls.cpu()

    def __repr__(self) -> str:
        return f"Platform(device={self.device!r}, arch={self.arch!r}, chain={self.device_chain})"


# --------------------------------------------------------------------------- #
# Matrix dumps — enumerate the full (op/kind × device × arch × variant) grid,  #
# including unavailable backends, without importing torch or touching a GPU.   #
# --------------------------------------------------------------------------- #
def kernel_matrix() -> list[dict[str, Any]]:
    ensure_backends_loaded()
    return KERNELS.manifest()


def component_matrix() -> list[dict[str, Any]]:
    ensure_backends_loaded()
    return COMPONENTS.manifest()