DeepInfer/deepinfer/neuron_ops/fp8_guard.py

"""Debug guard: report every float8 concatenation / advanced-indexing call site.

torch-neuronx 2.14 returns wrong data on the ``neuron`` device for ``torch.cat`` of float8 tensors and for
advanced indexing on non-contiguous float8 views, while the same ops on ``uint8`` views of the same bytes are
correct (tests/neuron/device/test_device_primitives.py). Kernel and loader code must therefore move float8 bytes
through ``uint8`` views. This guard finds the sites that do not, **without a NeuronCore**: it also works with
``DEEPINFER_NEURON_DEVICE=cpu``, where the ops are correct and only the call sites are recorded.

    DEEPINFER_NEURON_FP8_GUARD=1      record, print a deduplicated summary at exit
    DEEPINFER_NEURON_FP8_GUARD=raise  raise RuntimeError at the first offending call
    DEEPINFER_NEURON_FP8_GUARD_LOG=<path>  also append the summary to <path> (one block per process)

Installed by ``PlatformNeuron.__init__`` only when the variable is set. It monkey-patches the Python entry points
(``torch.cat``/``stack``/``concat``/..., ``Tensor.__getitem__``/``__setitem__`` with an advanced index,
``gather``/``index_select``/``take*``/``embedding``, ``index_put(_)``/``scatter*``/``index_copy_``/..., ``where``)
and records, for calls that involve a 1-byte float tensor, the first stack frame outside torch and outside this
module together with dtype / shape / contiguity. Basic indexing (ints, slices, None, Ellipsis) and dtype casts are
views or copies that behave correctly and are not reported. Debug only: the overhead is not tuned.
"""

from __future__ import annotations

import atexit
import numbers
import os
import sys
import threading
from collections import OrderedDict
from typing import Any, Callable

import torch

_MODE: str | None = None  # None (not installed), "record" or "raise"
_LOCK = threading.Lock()
_TLS = threading.local()
_SITES: "OrderedDict[tuple[str, str], dict[str, Any]]" = OrderedDict()
_MAX_SIGS = 4  # distinct dtype/shape/contiguity signatures kept per site
_SKIP_DIRS = (os.path.dirname(torch.__file__) + os.sep,)
_THIS_FILE = os.path.abspath(__file__)
_SKIP_FILES = (_THIS_FILE, os.path.join(os.path.dirname(_THIS_FILE), "transfer_to_neuron.py"))  # wrappers, not sites


def _is_fp8(x: Any) -> bool:
    return isinstance(x, torch.Tensor) and x.dtype.is_floating_point and x.dtype.itemsize == 1


def _call_site() -> str:
    f = sys._getframe(2)
    while f is not None:
        fn = os.path.abspath(f.f_code.co_filename)
        if fn not in _SKIP_FILES and not fn.startswith(_SKIP_DIRS):
            return f"{fn}:{f.f_lineno} ({f.f_code.co_name})"
        f = f.f_back
    return "<unknown>"


def _desc(t: torch.Tensor) -> str:
    dt = str(t.dtype).replace("torch.", "")
    return f"{dt}{list(t.shape)}{'' if t.is_contiguous() else ' non-contig'}"


def _report(op: str, tensors: list[torch.Tensor]) -> None:
    site = _call_site()
    sig = ", ".join(_desc(t) for t in tensors[:4]) + (", ..." if len(tensors) > 4 else "")
    with _LOCK:
        entry = _SITES.setdefault((op, site), {"count": 0, "sigs": []})
        entry["count"] += 1
        if sig not in entry["sigs"] and len(entry["sigs"]) < _MAX_SIGS:
            entry["sigs"].append(sig)
    if _MODE == "raise":
        raise RuntimeError(f"DEEPINFER_NEURON_FP8_GUARD: float8 {op} at {site}: {sig} "
                           "(route the bytes through .view(torch.uint8))")


def _tensors_in(obj: Any) -> list[torch.Tensor]:
    if isinstance(obj, torch.Tensor):
        return [obj]
    if isinstance(obj, (list, tuple)):
        return [t for t in obj if isinstance(t, torch.Tensor)]
    return []


def _is_advanced_index(idx: Any) -> bool:
    """Tensor / list / array index anywhere in the index tuple (a top-level tuple is basic indexing)."""
    items = idx if isinstance(idx, tuple) else (idx,)
    for it in items:
        if isinstance(it, (torch.Tensor, list)):
            return True
        if hasattr(it, "__array__") and not isinstance(it, numbers.Integral) and getattr(it, "ndim", 1) > 0:
            return True  # numpy array (numpy integer scalars are basic indexing)
    return False


def _wrap(fn: Callable, op: str, pick: Callable[..., list[torch.Tensor]]) -> Callable:
    if getattr(fn, "_fp8_guard_wrapped", False):
        return fn

    def wrapper(*args, **kwargs):
        if not getattr(_TLS, "busy", False):
            _TLS.busy = True
            try:
                involved = [t for t in pick(*args, **kwargs) if _is_fp8(t)]
            except Exception:  # never let the guard change behaviour of a bad call; the real call raises below
                involved = []
            finally:
                _TLS.busy = False
            if involved:
                _report(op, involved)
        return fn(*args, **kwargs)

    wrapper._fp8_guard_wrapped = True  # type: ignore[attr-defined]
    wrapper.__wrapped__ = fn  # type: ignore[attr-defined]
    wrapper.__name__ = getattr(fn, "__name__", op)
    wrapper.__doc__ = getattr(fn, "__doc__", None)
    return wrapper


# --- operand pickers: return the tensors whose float8-ness makes the call suspicious -------------------------------
def _p_seq(tensors=None, *a, **k):  # cat / stack / concat(tensors, dim)
    return _tensors_in(tensors if tensors is not None else k.get("tensors"))


def _p_self(self, *a, **k):  # gather / index_select / take ...: the source tensor
    return [self] if isinstance(self, torch.Tensor) else []


def _p_getitem(self, idx):
    return [self] if _is_advanced_index(idx) else []


def _p_setitem(self, idx, value):
    return [self, *_tensors_in(value)] if _is_advanced_index(idx) else []


def _p_all(*a, **k):  # every tensor operand (index_put / scatter / where ...)
    out: list[torch.Tensor] = []
    for x in list(a) + list(k.values()):
        out.extend(_tensors_in(x))
    return out


def _p_embedding(input, weight, *a, **k):
    return [weight]


_TORCH_FNS = {
    "cat": _p_seq, "concat": _p_seq, "concatenate": _p_seq, "stack": _p_seq,
    "hstack": _p_seq, "vstack": _p_seq, "dstack": _p_seq, "column_stack": _p_seq,
    "gather": _p_self, "index_select": _p_self, "take": _p_self, "take_along_dim": _p_self,
    "index_put": _p_all, "index_put_": _p_all, "scatter": _p_all, "scatter_add": _p_all,
    "scatter_reduce": _p_all, "index_copy": _p_all, "index_add": _p_all, "index_fill": _p_all,
    "masked_scatter": _p_all, "masked_select": _p_all, "where": _p_all,
}
_TENSOR_METHODS = {
    "gather": _p_self, "index_select": _p_self, "take": _p_self, "take_along_dim": _p_self,
    "index_put": _p_all, "index_put_": _p_all, "scatter": _p_all, "scatter_": _p_all,
    "scatter_add": _p_all, "scatter_add_": _p_all, "scatter_reduce": _p_all, "scatter_reduce_": _p_all,
    "index_copy": _p_all, "index_copy_": _p_all, "index_add": _p_all, "index_add_": _p_all,
    "index_fill": _p_all, "index_fill_": _p_all, "masked_scatter": _p_all, "masked_scatter_": _p_all,
    "masked_select": _p_all, "where": _p_all,
}


def _summary() -> str:
    lines = [f"===== DEEPINFER_NEURON_FP8_GUARD (pid {os.getpid()}): {len(_SITES)} float8 cat/advanced-index site(s)"]
    for (op, site), e in _SITES.items():
        lines.append(f"  {op:<20} x{e['count']:<7} {site}")
        for sig in e["sigs"]:
            lines.append(f"      {sig}")
    return "\n".join(lines)


def _at_exit() -> None:
    text = _summary()
    print(text, file=sys.stderr, flush=True)
    path = os.getenv("DEEPINFER_NEURON_FP8_GUARD_LOG")
    if path:
        with open(path, "a") as f:
            f.write(text + "\n")


def sites() -> dict[tuple[str, str], dict[str, Any]]:
    """Recorded (op, call site) -> {count, sigs}; for tests."""
    with _LOCK:
        return {k: {"count": v["count"], "sigs": list(v["sigs"])} for k, v in _SITES.items()}


def reset() -> None:
    with _LOCK:
        _SITES.clear()


def install(mode: str | None = None) -> None:
    """Install the guard. ``mode``: ``"raise"`` or anything truthy for record mode (default: from the env var)."""
    global _MODE
    if mode is None:
        mode = os.getenv("DEEPINFER_NEURON_FP8_GUARD", "")
    if not mode or mode == "0" or _MODE is not None:
        return
    _MODE = "raise" if mode == "raise" else "record"
    for name, pick in _TORCH_FNS.items():
        if hasattr(torch, name):
            setattr(torch, name, _wrap(getattr(torch, name), f"torch.{name}", pick))
    for name, pick in _TENSOR_METHODS.items():
        if hasattr(torch.Tensor, name):
            setattr(torch.Tensor, name, _wrap(getattr(torch.Tensor, name), f"Tensor.{name}", pick))
    torch.Tensor.__getitem__ = _wrap(torch.Tensor.__getitem__, "getitem[adv]", _p_getitem)
    torch.Tensor.__setitem__ = _wrap(torch.Tensor.__setitem__, "setitem[adv]", _p_setitem)
    import torch.nn.functional as F

    F.embedding = _wrap(F.embedding, "F.embedding", _p_embedding)
    atexit.register(_at_exit)
    print(f"[deepinfer] DEEPINFER_NEURON_FP8_GUARD={_MODE}: recording float8 cat/advanced-index sites",
          file=sys.stderr, flush=True)