#!/usr/bin/env python3
# SPDX-License-Identifier: AGPL-3.0-or-later
# Copyright (C) 2026 David (Corey) Herrera
"""ptx_modules.py — per-module md5 of the embedded PTX in cuda/fc_cuda_ptx.h,
or of the copy embedded in a built binary.

THE DEPLOY GATE. Kernel changes are inert until fc_cuda_ptx.h is regenerated, and
two compile-time gates (FC_MOE_SHX_STAGE, FC_MOE_DEDUP_COMPACT) can change or omit
a kernel with nothing logged. The only sufficient check is: extract the MODULE,
md5 it, compare to the intended build's module md5.

Header-wide counts are NOT sufficient and have twice produced a confident wrong
answer. All 49 ld.global.nc.v4.f32 in the pre-2026-08-11 header belong to
MATVEC_Q3K; scoped to MOE_EXPERT_GROUP the count is 0. Scope every count to the
module.

    ./ptx_modules.py ../cuda/fc_cuda_ptx.h                    # table, all modules
    ./ptx_modules.py ../cuda/fc_cuda_ptx.h MOE_EXPERT_GROUP   # one module, dumped
    ./ptx_modules.py /opt/pantheon/pantheon                  # from a BINARY
    ./ptx_modules.py /proc/<pid>/exe                          # from a RUNNING serve
    ./ptx_modules.py --verify /proc/<pid>/exe                 # ★ THE GATE, below

★ --verify TARGET is the gate. It recompiles every kernel from source with THE SAME
FLAG SET THE HEADER GENERATION USES (CANONICAL_* below, mirroring
cuda_kernels/Makefile) and compares module-by-module against TARGET — a header, a
built binary, or a running process. Needs nvcc, so a K80 host only.

Deliberately NOT a pinned md5 constant. A hardcoded hash rots the moment another
compile-time flag lands, and a rotted constant reads exactly like a real failure —
which is the trap that cost a full session. The comparison is against a BUILD, so
it stays true as the source moves. The hashes below are historical annotations for
reading old logs; they are NEVER the pass criterion.

    cb2e87d100eed7f49fe5eaeb5586fa9a  2026-08-11 shipped: compact + VEC4, 14 .entry
    a0c5e2da07a4af74501ce7c64c72c31d  VEC4 with FC_MOE_DEDUP_COMPACT=0, 12 .entry
    a32ae4d144b4cb9111f88e3096e36468  plain make before foundry 5eee7ce (FLAT), 12

The one criterion needing no build at all, for the specific defect that bit twice:
**ld.global.nc.v4.f32 must be > 0 in the MoE module.** 0 means the FLAT x-staging
body shipped. Header-wide the count is meaningless — all 49 v4 loads in the
pre-2026-08-11 header belong to MATVEC_Q3K.

Extraction notes, both of which have bitten:
  * In a HEADER the PTX is C-escaped across many quoted lines. PTX instructions
    end in ';', so looking for a terminating semicolon after the '=' lands inside
    the string literal and returns ~603 B instead of ~180 KB. Collect consecutive
    quoted literals from the marker line until a line with no literal, join, then
    unicode_escape. Sanity-check the size.
  * In a BINARY the same text is unescaped and NUL-terminated, so a raw
    `strings | grep` behaves differently from the header and has produced one
    false negative and one false positive on the same fact. Extract NUL-to-NUL.
"""
import hashlib
import json
import os
import re
import subprocess
import sys
import tempfile

# ---------------------------------------------------------------------------
# The canonical header-generation recipe. MUST mirror cuda_kernels/Makefile.
# CANONICAL_NVCC_FLAGS is EMPTY on purpose: every shipped compile-time gate is a
# SOURCE default (FC_MOE_SHX_STAGE=VEC4 in moe_dedup_shx_staging.cuh,
# FC_MOE_DEDUP_COMPACT=1 in moe_expert_group.cu, FC_NX_STAGE_NESTED=EOUTER in
# matvec_nx_staging.cuh), so a plain `make` IS the shipped build. If you ever find
# yourself adding a -D here to make the gate pass, the fix belongs in the source
# default instead — that is the whole lesson of 2026-08-11.
# ---------------------------------------------------------------------------
CANONICAL_NVCC = os.environ.get("NVCC", "/usr/local/cuda-11.4/bin/nvcc")
CANONICAL_ARCH = "sm_37"
CANONICAL_HOST_CC = os.environ.get("HOST_CC", "gcc-11")
CANONICAL_NVCC_FLAGS = []
KERNELS = ["elementwise", "normalization", "matvec_q8", "matvec_q3k",
           "matvec_kquants", "rope", "moe_routing", "moe_expert_group",
           "ssm_gated_delta", "fused_layer", "topk_logits", "attention"]

LIT = re.compile(r'"((?:[^"\\]|\\.)*)"')
MARK = re.compile(r"FC_PTX_([A-Z0-9_]+)\[\]\s*=")
TOKENS = ("ld.global.nc.v4.f32", "st.shared.v4", ".shared .align 16", ".entry")

# A module always opens with the nvcc banner; used to find module starts in a binary.
BANNER = b"//\n// Generated by NVIDIA NVVM Compiler\n"


def from_header(path):
    """{module_name: ptx_text} out of a generated fc_cuda_ptx.h."""
    lines = open(path, encoding="utf-8", errors="replace").read().split("\n")
    out = {}
    for i, line in enumerate(lines):
        m = MARK.search(line)
        if not m:
            continue
        parts, j, seen = [], i, False
        while j < len(lines):
            q = LIT.findall(lines[j])
            if q:
                parts += q
                seen = True
            elif seen:
                break
            j += 1
        out[m.group(1)] = "".join(parts).encode().decode("unicode_escape")
    return out


def from_binary(path):
    """{synthetic_name: ptx_text} for every embedded PTX module in a binary.

    Named by the first .entry it contains, since the FC_PTX_* symbol names are
    compile-time only and do not survive into the artifact.
    """
    blob = open(path, "rb").read()
    out, pos = {}, 0
    while True:
        s = blob.find(BANNER, pos)
        if s < 0:
            break
        e = blob.find(b"\x00", s)
        if e < 0:
            break
        mod = blob[s:e]
        pos = e + 1
        txt = mod.decode("utf-8", "replace")
        ents = re.findall(r"\.entry\s+(\S+?)\(", txt)
        if not ents:
            continue
        out["<%s...>" % ents[0]] = txt
    return out


def canonical_build(kdir, outdir):
    """Compile every kernel with the canonical recipe. -> {kernel: ptx_text}.

    The .ptx a kernel compiles to IS the module text the header embeds, so this is
    directly comparable to an extracted module with no post-processing.
    """
    out = {}
    for k in KERNELS:
        src = os.path.join(kdir, k + ".cu")
        if not os.path.exists(src):
            raise SystemExit("missing kernel source: %s" % src)
        dst = os.path.join(outdir, k + ".ptx")
        cmd = [CANONICAL_NVCC, "-ptx", "-arch=" + CANONICAL_ARCH,
               "-ccbin", CANONICAL_HOST_CC, "-Wno-deprecated-gpu-targets",
               "-isystem", os.path.join(kdir, "nvcc_shim")] + \
              CANONICAL_NVCC_FLAGS + ["-o", dst, src]
        p = subprocess.run(cmd, capture_output=True, text=True)
        if p.returncode != 0:
            sys.stderr.write(p.stderr[-3000:])
            raise SystemExit("canonical build FAILED on %s" % k)
        out[k] = open(dst).read()
    return out


def verify(target, kdir):
    """Gate: TARGET's modules must be byte-identical to a canonical build."""
    tgt = modules(target)
    with tempfile.TemporaryDirectory() as td:
        can = canonical_build(kdir, td)
    cand = {hashlib.md5(v.encode()).hexdigest(): k for k, v in can.items()}
    print("canonical recipe: %s -ptx -arch=%s -ccbin %s %s"
          % (CANONICAL_NVCC, CANONICAL_ARCH, CANONICAL_HOST_CC,
             " ".join(CANONICAL_NVCC_FLAGS) or "(no extra flags - source defaults)"))
    print("target: %s\n" % target)
    matched, bad = set(), []
    for name, txt in sorted(tgt.items()):
        d = hashlib.md5(txt.encode()).hexdigest()
        k = cand.get(d)
        if k:
            matched.add(k)
            print("  OK       %-34s %s  == %s.ptx" % (name, d, k))
        else:
            bad.append(name)
            print("  MISMATCH %-34s %s  matches NO canonical module" % (name, d))
    missing = [k for k in can if k not in matched]
    for k in missing:
        print("  ABSENT   %s.ptx is in the canonical build but not in the target" % k)
    moe = [t for n, t in tgt.items() if "moe_gateup_q3k" in n or n == "MOE_EXPERT_GROUP"]
    if moe:
        n4 = moe[0].count("ld.global.nc.v4.f32")
        print("\n  MoE module ld.global.nc.v4.f32 = %d %s"
              % (n4, "(VEC4 x-staging present)" if n4 else
                 "*** ZERO - the FLAT body shipped, this is the defect ***"))
        if n4 == 0:
            bad.append("MOE ld.global.nc.v4.f32 == 0")
    ok = not bad and not missing
    print("\n%s" % ("GATE PASS - every module byte-identical to the canonical build"
                    if ok else "GATE FAIL"))
    return 0 if ok else 1


def modules(path):
    """Dispatch on content: a generated header, else a binary."""
    with open(path, "rb") as f:
        head = f.read(4096)
    if b"FC_CUDA_PTX_H" in head or b"FC_PTX_" in head:
        return from_header(path)
    return from_binary(path)


def main(argv):
    if len(argv) < 2:
        print(__doc__)
        return 2
    if argv[1] == "--verify":
        if len(argv) < 3:
            print("usage: ptx_modules.py --verify TARGET [kernel_dir]")
            return 2
        kdir = argv[3] if len(argv) > 3 else os.path.join(
            os.path.dirname(os.path.abspath(__file__)), "..", "cuda_kernels")
        return verify(argv[2], os.path.normpath(kdir))
    path = argv[1]
    only = argv[2] if len(argv) > 2 else None
    mods = modules(path)
    if not mods:
        print("no embedded PTX modules found in %s" % path)
        return 1
    rows = {}
    for name, txt in sorted(mods.items()):
        if only and only not in name:
            continue
        rows[name] = dict(
            bytes=len(txt),
            ptx_lines=txt.count("\n"),
            md5=hashlib.md5(txt.encode()).hexdigest(),
            entries=re.findall(r"\.entry\s+(\S+?)\(", txt),
            **{t: txt.count(t) for t in TOKENS}
        )
        if only:
            dump = "/tmp/mod_%s.ptx" % re.sub(r"\W+", "_", name)
            open(dump, "w").write(txt)
            print("wrote %s" % dump)
            if len(txt) < 4096:
                print("!! %d B is implausibly small for a PTX module — the "
                      "extraction is wrong, not the build" % len(txt))
    if not rows:
        print("no module matched %r; have: %s" % (only, ", ".join(sorted(mods))))
        return 1
    hdr = "%-34s %10s %8s %34s %5s %6s %6s %6s"
    print(hdr % ("module", "bytes", "ptxln", "md5", "entry", "ldv4", "stshv4", "sh16"))
    for name, r in rows.items():
        print(hdr % (name, r["bytes"], r["ptx_lines"], r["md5"], r[".entry"],
                     r["ld.global.nc.v4.f32"], r["st.shared.v4"],
                     r[".shared .align 16"]))
    if only:
        for name, r in rows.items():
            print("\n%s entries:" % name)
            for e in r["entries"]:
                print("   %s" % e)
    if "--json" in argv:
        print(json.dumps(rows, indent=1, sort_keys=True))
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv))