hao-ai-lab-FastVideo/fastvideo-kernel/wan22_vae_decode/tools/audit.py

"""The stack rules of a solution, checked on its sources; and who launched each kernel of a call, found at run time.

The stack is CUDA C++ (with inline PTX), compiled ahead of time into one extension module (see loader.py), with
communication through the ProcessGroup it is given (NCCL) or CUDA peer access and IPC between the ranks.

Static rules (check): solution.json declares the languages in LANGUAGES, and RULES maps a rule's name to what it
looks for in the C/C++/CUDA sources and why. Kinds of match:

    include   the target of an #include
    text      the source with comments removed
    flag      solution.json: a compile option
    vendored  an #include target that must resolve to a file of the solution itself

A task's spec.py may extend the table: RULES = {name: rule, or None to drop a default rule}.

Run time (trace): one call runs under CUPTI. Every kernel is attributed to the library that launched it, from the
native call stack of its launch: kernels launched from inside cuBLAS/cuBLASLt, cuDNN or ATen (libtorch) are library
kernels, and so are Triton's (a stack the rules exclude); NCCL's are communication, which the rules allow, counted
apart; kernels launched from the solution's own module or by driver calls are the solution's, whatever their names,
so kernels built from CUTLASS/CuTe templates copied into the solution count as its own. Kernels whose launch cannot be
traced (graph replays, device-side launches) are classified by name. Library kernels are allowed, as placeholders:
their count and their share of the call's GPU time show where the solution still depends on them. Only kernel
launches call back into this module (enabling the whole driver API would call back on NCCL's watchdog thread, which
queries events while it holds a lock the launching thread needs: a deadlock under the GIL).

    python tools/audit.py [solution]    the static rules
"""

import argparse
import collections
import ctypes
import json
import os
import re
import sys

CPP_FILES = (".cu", ".cpp", ".cc", ".cxx", ".c", ".h", ".hpp", ".cuh", ".hh", ".hxx", ".inl", ".inc")
LANGUAGES = ["cuda_cpp"]

RULES = {
    "vendored-cutlass": {
        "vendored": r"(cutlass|cute)/.*",
        "why": "CUTLASS/CuTe headers are copied into the solution and maintained there (cutlass/..., cute/... next to "
               "the entry file), not taken from the installed CUTLASS"},
    "runtime-build": {
        "include": r"(.*/)?(nvrtc|nvJitLink|nvPTXCompiler)\.h",
        "text": r"\b(nvrtc[A-Z]\w*|nvJitLink[A-Z]\w*|nvPTXCompiler\w*|cuModuleLoad\w*|"
                r"cuLink(Create|AddData|AddFile|Complete)\w*|cuLibraryLoad\w*|cudaLibraryLoad\w*)\s*\(",
        "flag": r"-l(nvrtc|nvJitLink|nvptxcompiler)\w*",
        "why": "nothing is compiled or loaded at run time (NVRTC, nvJitLink, PTX or cubins loaded through the driver): "
               "device code, inline PTX included, is compiled by nvcc with the solution"},
    "python-from-cpp": {
        "text": r"\b(py(bind11)?::module_?::import|py(bind11)?::(exec|eval)(_file)?|PyImport_\w+|PyRun_\w+)\b",
        "why": "run() stays in C++ and runs no Python: Triton, torch.compile, torch.utils.cpp_extension and other "
               "kernel toolkits would come in that way"},
}


def rule_table(extra=None):
    """The rule table in force: the defaults, changed by a task's RULES (None drops a rule)."""
    table = dict(RULES)
    for name, rule in (extra or {}).items():
        if rule is None:
            table.pop(name, None)
        else:
            table[name] = rule
    return table


def _match(pattern, s):
    return re.fullmatch(pattern, s) is not None


def _strip_cpp_comments(text):
    """Removes // and /* */ comments (string literals are not parsed; good enough for these rules)."""
    text = re.sub(r"/\*.*?\*/", " ", text, flags=re.S)
    return re.sub(r"//[^\n]*", "", text)


def _includes(text):
    return [(m.group(1), text.count("\n", 0, m.start()) + 1)
            for m in re.finditer(r'^\s*#\s*include\s*[<"]([^>"]+)[>"]', text, flags=re.M)]


def check(solution_dir, rules, languages=None):
    """Violations of `rules` by a solution's sources: a list of "file:line: what (rule: why)" strings.

    languages: the languages the solution may declare in solution.json (default LANGUAGES)."""
    languages = LANGUAGES if languages is None else languages
    out = []
    try:
        with open(os.path.join(solution_dir, "solution.json")) as f:
            spec = json.load(f)
    except (OSError, ValueError) as e:
        return [f"solution.json: unreadable ({e})"]
    if not isinstance(spec, dict):
        return ["solution.json: must hold an object"]
    declared = spec.get("languages")
    if not isinstance(declared, list) or not declared or any(lang not in languages for lang in declared):
        out.append(f"solution.json: languages {declared} are not allowed here (allowed: {languages})")
    options = spec.get("compile_options") if isinstance(spec.get("compile_options"), dict) else {}
    flags = [f for v in options.values() if isinstance(v, list) for f in v]
    for name, rule in rules.items():
        for flag in flags:
            if "flag" in rule and _match(rule["flag"], str(flag)):
                out.append(f"solution.json: compile option {flag!r} ({name}: {rule['why']})")
    for root, dirs, files in os.walk(solution_dir):
        dirs[:] = sorted(d for d in dirs if d != "__pycache__")
        for file in sorted(files):
            if not file.endswith(CPP_FILES):
                continue
            path = os.path.join(root, file)
            try:
                with open(path, encoding="utf-8") as f:
                    text = f.read()
            except (OSError, UnicodeDecodeError):
                continue
            out += _check_cpp(solution_dir, os.path.relpath(path, solution_dir), text, rules)
    return out


def _check_cpp(solution_dir, rel, text, rules):
    out = []
    code = _strip_cpp_comments(text)
    includes = _includes(code)
    for name, rule in rules.items():
        for target, line in includes:
            if "include" in rule and _match(rule["include"], target):
                out.append(f"{rel}:{line}: #include <{target}> ({name}: {rule['why']})")
            if "vendored" in rule and _match(rule["vendored"], target):
                here = os.path.dirname(os.path.join(solution_dir, rel))
                if not any(os.path.isfile(os.path.join(d, target)) for d in (here, solution_dir)):
                    out.append(f"{rel}:{line}: #include <{target}> is not in the solution ({name}: {rule['why']})")
        if "text" in rule:
            for m in re.finditer(rule["text"], code):
                line = code.count("\n", 0, m.start()) + 1
                out.append(f"{rel}:{line}: {m.group(0)!r} ({name}: {rule['why']})")
                break
    return out


# Run-time audit.

# The library a kernel was launched from, by the path of the innermost library on the launch's native stack.
LAUNCHERS = [
    (r".*/libcublas(Lt)?[^/]*", "cuBLAS"),
    (r".*/libcudnn[^/]*", "cuDNN"),
    (r".*/lib(torch|c10)[^/]*", "ATen"),
    (r".*triton.*", "Triton"),
    (r".*/libnccl[^/]*", "NCCL"),
    (r".*/lib(cufft|curand|cusparse|cusolver|cutensor|nvshmem|npp)[^/]*", "CUDA library"),
]
# Frames skipped on the way out: CUPTI and this module's callback, the Python interpreter, the CUDA runtime/driver.
SKIP = r"(.*/)?(libffi[^/]*|_ctypes[^/]*|python[^/]*|libpython[^/]*|callback_wrappers[^/]*|cupti[^/]*|libcupti[^/]*|" \
       r"libcuda\.so[^/]*|libcudart[^/]*|libc\.so[^/]*|ld-linux[^/]*|\?)"
# Fallback by kernel name, for launches that could not be traced.
NAMES = [
    (r"(void )?at::.*|.*at::native::.*|_ZN2at.*|.*at_cuda_detail.*|_ZN3c10.*", "ATen"),
    (r".*(nvjet|cublas|gemv2?[NT]?_|gemmSN|splitKreduce|xmma|sgemm|hgemm|dgemm).*|cutlass3x_.*|cutlass_\d.*", "cuBLAS"),
    (r".*cudnn.*", "cuDNN"),
    (r"triton_.*", "Triton"),
    (r"(void )?nccl.*|_Z\d+nccl.*", "NCCL"),
]
# Libraries whose kernels are communication: allowed outright, and counted apart from the placeholders.
COMMUNICATION = ("NCCL",)
# The CUPTI activity buffers: few and large, since a decode launches tens of thousands of kernels.
BUFFER_BYTES = 32 << 20

_libc = None
_demangler = None


class _DlInfo(ctypes.Structure):
    _fields_ = [("fname", ctypes.c_char_p), ("fbase", ctypes.c_void_p), ("sname", ctypes.c_char_p),
                ("saddr", ctypes.c_void_p)]


def _launcher():
    """Path of the innermost library on the current native stack that is not CUPTI, Python or the CUDA runtime."""
    global _libc
    if _libc is None:
        _libc = ctypes.CDLL("libc.so.6")
        _libc.backtrace.argtypes = [ctypes.POINTER(ctypes.c_void_p), ctypes.c_int]
        _libc.dladdr.argtypes = [ctypes.c_void_p, ctypes.POINTER(_DlInfo)]
    buf = (ctypes.c_void_p * 96)()
    n = _libc.backtrace(buf, 96)
    for i in range(n):
        info = _DlInfo()
        path = info.fname.decode() if buf[i] and _libc.dladdr(buf[i], ctypes.byref(info)) and info.fname else "?"
        if not re.fullmatch(SKIP, path):
            return path
    return ""


def demangle(name):
    global _demangler
    if _demangler is None:
        try:
            lib = ctypes.CDLL("libstdc++.so.6")
            lib.__cxa_demangle.argtypes = [ctypes.c_char_p, ctypes.c_char_p, ctypes.POINTER(ctypes.c_size_t),
                                           ctypes.POINTER(ctypes.c_int)]
            lib.__cxa_demangle.restype = ctypes.c_char_p
            _demangler = lib
        except OSError:
            _demangler = False
    if not _demangler or not name.startswith("_Z"):
        return name
    status = ctypes.c_int()
    out = _demangler.__cxa_demangle(name.encode(), None, None, ctypes.byref(status))
    return out.decode() if status.value == 0 and out else name


def short_name(name, limit=110):
    """A kernel's demangled name without its return type, parameter list and the at::native/anonymous namespaces."""
    text = demangle(name).replace("(anonymous namespace)::", "")
    depth, start, end = 0, 0, len(text)
    for i, c in enumerate(text):
        if c in "<(":
            if c == "(" and depth == 0 and i > 0:
                end = i
                break
            depth += 1
        elif c in ">)":
            depth -= 1
        elif c == " " and depth == 0:
            start = i + 1
    text = text[start:end] or text
    for noise in ("at::native::", "at::", "std::"):
        text = text.replace(noise, "")
    return text if len(text) <= limit else text[:limit - 3] + "..."


def classify(name, launcher, traced):
    """The library a kernel belongs to (None: the solution's own)."""
    if traced:
        for pattern, label in LAUNCHERS:
            if re.fullmatch(pattern, launcher):
                return label
        return None
    for pattern, label in NAMES:
        if re.fullmatch(pattern, name) or re.fullmatch(pattern, demangle(name)):
            return label
    return None


Op = collections.namedtuple("Op", "kind name bytes launcher library ns")


def launch_callbacks(cupti):
    """The driver API callbacks of kernel launches (cuLaunchKernel, cuLaunchKernelEx, cuGraphLaunch, ...)."""
    return [int(getattr(cbid, "value", cbid)) for name, cbid in cupti.driver_api_trace_cbid.__members__.items()
            if "Launch" in name and "HostFunc" not in name]


def trace(fn, *args):
    """Calls fn(*args) once under CUPTI; returns its GPU operations in start order, as Op(kind, name, bytes,
    launcher, library, ns): kind is kernel, memcpy or memset; library is None for the solution's own kernels; ns is
    the operation's GPU time."""
    import torch
    from cupti import cupti

    launches, records, errors = {}, [], []

    def on_api(userdata, domain, cbid, data):
        try:
            if data.callback_site == cupti.ApiCallbackSite.API_ENTER:
                graph = "Graph" in data.function_name
                launches[data.correlation_id] = "" if graph else _launcher()
        except Exception as e:  # never let an exception cross into the driver
            errors.append(repr(e))

    def requested():
        return BUFFER_BYTES, 0

    def completed(activities):
        for a in activities:
            if a.kind == cupti.ActivityKind.CONCURRENT_KERNEL:
                records.append((a.start, a.end, "kernel", a.name, 0, a.correlation_id,
                                bool(a.graph_id or a.is_device_launched)))
            elif a.kind == cupti.ActivityKind.MEMCPY:
                records.append((a.start, a.end, "memcpy", "", a.bytes, a.correlation_id, False))
            elif a.kind == cupti.ActivityKind.MEMSET:
                records.append((a.start, a.end, "memset", "", a.bytes, a.correlation_id, False))

    kinds = (cupti.ActivityKind.CONCURRENT_KERNEL, cupti.ActivityKind.MEMCPY, cupti.ActivityKind.MEMSET)
    torch.cuda.synchronize()
    sub = cupti.subscribe(on_api, 0)
    try:
        for cbid in launch_callbacks(cupti):
            cupti.enable_callback(1, sub, cupti.CallbackDomain.DRIVER_API, cbid)
        for k in kinds:
            cupti.activity_enable(k)
        cupti.activity_register_callbacks(requested, completed)
        fn(*args)
        torch.cuda.synchronize()
        cupti.activity_flush_all(1)
    finally:
        for k in kinds:
            cupti.activity_disable(k)
        cupti.unsubscribe(sub)
        cupti.finalize()
    if errors:
        raise RuntimeError(f"the launch trace failed: {errors[0]}")
    ops = []
    for start, end, kind, name, nbytes, corr, untraceable in sorted(records):
        launcher = launches.get(corr)
        traced = kind == "kernel" and not untraceable and bool(launcher)
        library = classify(name, launcher or "", traced) if kind == "kernel" else None
        ops.append(Op(kind, name, nbytes, os.path.basename(launcher or ""), library, end - start))
    return ops


def placeholder(library):
    """Whether a library's kernels are placeholders (cuBLAS, cuDNN, ATen, ...), as opposed to the solution's own or
    communication."""
    return bool(library) and library not in COMMUNICATION


def kernels(ops):
    """What a call launched, by kernel: [[library or "", name, launches, ns]] by GPU time, the largest first; copies
    and sets as the kernels "(memcpy)" and "(memset)"."""
    table = collections.defaultdict(lambda: [0, 0])
    for op in ops:
        key = (op.library or "", op.name if op.kind == "kernel" else f"({op.kind})")
        table[key][0] += 1
        table[key][1] += op.ns
    return sorted(([lib, name, n, ns] for (lib, name), (n, ns) in table.items()), key=lambda r: -r[3])


def library_share(ops):
    """The share of a call's GPU time (the summed time of its operations) spent in placeholder library kernels."""
    total = sum(op.ns for op in ops)
    return sum(op.ns for op in ops if placeholder(op.library)) / total if total else 0.0


def comm_share(ops):
    """The share of a call's GPU time spent in communication kernels (NCCL)."""
    total = sum(op.ns for op in ops)
    return sum(op.ns for op in ops if op.library in COMMUNICATION) / total if total else 0.0


def main():
    here = os.path.dirname(os.path.abspath(__file__))
    sys.path.insert(0, here)
    import spec  # noqa: E402

    p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    p.add_argument("solution", nargs="?", default="solution")
    a = p.parse_args()
    table = rule_table(getattr(spec, "RULES", None))
    found = check(a.solution, table)
    for v in found:
        print(v)
    print(f"stack rules: {'ok' if not found else f'{len(found)} violations'} ({len(table)} rules)")
    sys.exit(1 if found else 0)


if __name__ == "__main__":
    main()