DeepInfer/deepinfer/model_executor/layers/tilelang_ops/ops.py

"""Kernel dispatcher: interface modules do
``from deepinfer.model_executor.layers.tilelang_ops import ops`` and call
``ops.xxx_tl``; backend selection happens only in this file.

Re-exports all ``xxx_tl`` factories of the selected backend.

Backend resolution:

* ``DEEPINFER_OPS_BACKEND`` picks the backend package for every kernel. Unset means the platform
  default, which is what the engine has always used, so nothing changes unless asked for.
* ``DEEPINFER_OPS_FALLBACK`` is a comma-separated list of kernel names that must come from the
  platform default even when another backend is selected. It exists so a partially ported backend
  can still run end to end: the ported kernels exercise the new code, the rest keep working, and
  the comparison stays meaningful throughout the port.

When the selected backend does not export a kernel at all, the default backend's implementation
is used and a line naming it is printed once, so a partial port is never mistaken for a complete
one.
"""

import importlib
import os
import pkgutil
import sys
from pathlib import Path
from types import ModuleType
from typing import Any

from deepinfer import PLATFORM

_TRACE = bool(os.getenv("DEEPINFER_TRACE_KERNELS"))


def _record(name: str, args: tuple = (), kwargs: dict | None = None) -> None:
    """Log one use of a backend kernel factory, so the real surface can be measured.

    Called once per kernel *construction* rather than per launch: TileLang caches compiled
    kernels, so construction happens at warmup and the count stays meaningful.

    With `DEEPINFER_TRACE_KERNEL_ARGS=1` the construction arguments are logged too. Several
    factories take a dozen flags (norm, fused_transform, fused_sample) and only a handful of
    combinations are ever used, so the arguments say which code paths a port actually needs.
    """
    path = Path(os.environ["DEEPINFER_TRACE_KERNELS"])
    with path.open("a") as f:
        if os.getenv("DEEPINFER_TRACE_KERNEL_ARGS"):
            f.write(f"{name}|args={args!r}|kwargs={sorted((kwargs or {}).items())!r}\n")
        else:
            f.write(f"{name}\n")


class _KernelTrace:
    """Proxy a `*_tl` factory, recording its first use.

    A plain function wrapper is not enough: callers read attributes off the factory (the TileLang
    `@tilelang.jit` decorator attaches `pass_configs`, for one), so attribute access has to keep
    working. Everything unknown is forwarded to the real object.
    """

    def __init__(self, name: str, factory: Any) -> None:
        self._name = name
        self._factory = factory
        self._seen = False

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if not self._seen:
            self._seen = True
            _record(self._name, args, kwargs)
        return self._factory(*args, **kwargs)

    def __getattr__(self, item: str) -> Any:
        return getattr(self._factory, item)

    def __repr__(self) -> str:
        return f"<traced {self._name}>"


class _FallbackNotice:
    """Report the first *use* of a kernel the selected backend does not implement.

    The import-time summary says how many names are missing, which is a static count: it says
    nothing about whether the running code actually reaches any of them. This records that, so an
    end-to-end run proves which kernels still came from the default backend instead of leaving it
    to be inferred from a name list.
    """

    def __init__(self, name: str, factory: Any, selected: str, default: str) -> None:
        self._name = name
        self._factory = factory
        self._selected = selected
        self._default = default
        self._seen = False

    def __call__(self, *args: Any, **kwargs: Any) -> Any:
        if not self._seen:
            self._seen = True
            print(
                f"[ops] kernel {self._name!r} used from {self._default!r}: {self._selected!r} does not implement it",
                file=sys.stderr,
                flush=True,
            )
        return self._factory(*args, **kwargs)

    def __getattr__(self, item: str) -> Any:
        return getattr(self._factory, item)


def _collect(backend: str) -> dict[str, Any]:
    """Collect every ``*_tl`` factory exported by a backend package, by name."""
    pkg = importlib.import_module(f"{__package__}.{backend}")
    found: dict[str, Any] = {}
    for info in pkgutil.iter_modules(pkg.__path__):
        mod = importlib.import_module(f"{pkg.__name__}.{info.name}")
        for name, obj in vars(mod).items():
            if name.startswith("_") or isinstance(obj, ModuleType) or not name.endswith("_tl"):
                continue
            if name in found:
                raise RuntimeError(f"backend {backend!r} exports {name!r} twice")
            found[name] = obj
    return found


_DEFAULT_BACKEND = PLATFORM.tilelang_ops_backend
_BACKEND = os.getenv("DEEPINFER_OPS_BACKEND") or _DEFAULT_BACKEND
_FALLBACK = (
    {name.strip() for name in os.getenv("DEEPINFER_OPS_FALLBACK", "").split(",") if name.strip()} if _BACKEND != _DEFAULT_BACKEND else set()
)

_SELECTED = _collect(_BACKEND)
_DEFAULT: dict[str, Any] = _SELECTED if _BACKEND == _DEFAULT_BACKEND else _collect(_DEFAULT_BACKEND)

_unknown = sorted(_FALLBACK - set(_SELECTED))
if _unknown:
    raise RuntimeError(f"DEEPINFER_OPS_FALLBACK names kernel(s) backend {_BACKEND!r} does not export: {_unknown}")

for _name, _obj in _SELECTED.items():
    globals()[_name] = _KernelTrace(_name, _obj) if _TRACE else _obj

for _name in _FALLBACK:
    globals()[_name] = _DEFAULT[_name]

_INHERITED = sorted(set(_DEFAULT) - set(_SELECTED) - _FALLBACK)
for _name in _INHERITED:
    globals()[_name] = _FallbackNotice(_name, _DEFAULT[_name], _BACKEND, _DEFAULT_BACKEND)

# Kept short: with data parallelism this prints once per rank, and the full list is dozens of
# names that are not on the caller's path. `DEEPINFER_OPS_VERBOSE=1` prints them all.
if _INHERITED:
    print(
        f"[ops] backend {_BACKEND!r} lacks {len(_INHERITED)} kernel(s), falling back to "
        f"{_DEFAULT_BACKEND!r} ({_INHERITED if os.getenv('DEEPINFER_OPS_VERBOSE') else 'first: ' + _INHERITED[0]})",
        file=sys.stderr,
        flush=True,
    )


class KernelUnavailableError(AttributeError):
    """Selected backend lacks this kernel factory; conftest.py turns it into
    a test skip on non-cuda backends. Subclasses AttributeError to keep the
    module ``__getattr__`` / ``hasattr`` protocol intact."""


def __getattr__(name: str) -> Any:  # for static analyzers
    raise KernelUnavailableError(f"tilelang_ops backend {PLATFORM.tilelang_ops_backend!r} has no kernel {name!r}")