"""Lab 0: which engine runs each shape? The tables are read straight out of submission.py.

Run:   python3 lab0_dispatch.py                 (uses ../../submission.py)
       python3 lab0_dispatch.py my_copy.py      (any edited copy)
Needs: Python 3 only.

It re-implements the decision logic of kernel_fn() (submission.py lines 774-807) in plain Python
and prints, for the 30 scored shapes and the 9 test shapes:
  - which engine runs (Triton split-K, compressed GEMV, plain GEMV, cuBLAS)
  - the kernel's template parameters and its launch geometry (CTAs and threads per CTA)
  - for the GEMV kernels, how K is divided (per CTA, per warp, per loop iteration)

Things to try
  - Before running, predict the engine for m=5 unpooler, m=9 router and m=33 o. Then check.
  - Copy submission.py, delete the ('pooler', 1) line from ZTABLE, run this lab on the copy.
"""
import ast
import os
import re
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
PATH = sys.argv[1] if len(sys.argv) > 1 else os.path.join(HERE, "..", "..", "submission.py")
SRC = open(PATH).read()
SMS = 132  # streaming multiprocessors on an H200


def table(name):
    """Parse 'NAME: dict = {...}' from submission.py (each table ends at the first '}' + newline)."""
    m = re.search(r"^" + name + r": dict = (\{.*?\})\n", SRC, re.S | re.M)
    return ast.literal_eval(m.group(1))


CTABLE, ZTABLE, TABLE = table("CTABLE"), table("ZTABLE"), table("TABLE")
LAYERS = {"qkv_gate": (2304, 8192), "o": (8192, 1024), "pooler": (4096, 8192),
          "unpooler": (8192, 4096), "router": (256, 8192)}


def nb_of(m):
    return 1 if m <= 8 else 2 if m <= 16 else 4


def k_split(k, s, nw, step, unroll):
    """How the GEMV kernels divide K: per CTA (K/S), per warp (K/S/NW), loop trips per warp."""
    blocks = k // step                      # 32-wide K blocks (gemv) or 64-wide K groups (gemvz)
    ok = blocks % (s * nw) == 0 and (blocks // (s * nw)) % unroll == 0
    per_warp = blocks // (s * nw)
    note = "" if ok else "   <-- DOES NOT DIVIDE: part of K is skipped or read twice (wrong answer)"
    return f"K per CTA {k // s}, per warp {per_warp * step}, {per_warp // unroll if ok else '?'} loop trips{note}"


def dispatch(layer, m):
    """Mirror of kernel_fn: returns (engine, detail, ctas, threads_per_cta)."""
    n = 1024 if layer == "unpooler" and m <= 8 else LAYERS[layer][0]
    k = LAYERS[layer][1]
    router = layer == "router"
    if (layer, m) in TABLE:                                  # 1) Triton split-K GEMM
        cfg = TABLE[(layer, m)]
        bn, bm, bk, split, nw, ns, tma = cfg[:7]
        atom = len(cfg) > 7 and cfg[7] and split > 1
        grid = (n // bn, split, -(-m // bm))
        detail = (f"_proj_tma_kernel BN={bn} BM={bm} BK={bk} SPLIT={split} warps={nw} stages={ns} "
                  f"{'atomic' if atom else 'serial'} epilogue{', DUAL hi+lo' if router else ''}; grid {grid}")
        return "Triton", detail, grid[0] * grid[1] * grid[2], nw * 32
    zc = ZTABLE.get((layer, m))
    if zc is not None:                                       # 2) compressed-weight GEMV
        nw, s, ug, p = zc[:4]
        detail = f"gemvz_kernel<NW={nw}, UG={ug}, NB={nb_of(m)}, PIPE={p}> S={s}; {k_split(k, s, nw, 64, ug)}"
        return "gemvz", detail, (n // 16) * s, (nw + 1) * 32
    if m <= 32 and CTABLE.get((layer, m), 0) is not None:    # 3) plain BF16 GEMV
        nw, s, u, p = CTABLE.get((layer, m)) or ((8, 8, 2, 0) if router else (4, 1, 2, 0))
        dflt = "" if (layer, m) in CTABLE else " (no table entry: default config)"
        detail = (f"gemv_kernel<NW={nw}, U={u}, NB={nb_of(m)}, DUAL={'true' if router else 'false'}, PIPE={p}> "
                  f"S={s}{dflt}; {k_split(k, s, nw, 32, u)}")
        return "gemv", detail, (n // 16) * s, nw * 32
    if router:                                               # 4) cuBLAS fallbacks
        return "cuBLAS", "two BF16 GEMMs (hi, lo) with FP32 output, then an add", None, None
    return "cuBLAS", "x @ W[:n].T", None, None


def show(title, cases):
    print(title)
    print(f"{'case':<26}{'engine':<8}{'CTAs':>6}{'thr/CTA':>8}  details")
    for label, layer, m in cases:
        eng, detail, ctas, thr = dispatch(layer, m)
        waves = f"{ctas:>6}" if ctas else f"{'-':>6}"
        print(f"{label:<26}{eng:<8}{waves}{thr or '-':>8}  {detail}")
    print()


if __name__ == "__main__":
    loads = [("decode 1", 1), ("decode 16", 16), ("decode 32", 32), ("verify 1", 8), ("verify 16", 128), ("verify 32", 256)]
    scored = [(f"{ld} {layer} m={m}", layer, m) for ld, m in loads for layer in LAYERS]
    show("== the 30 scored shapes ==", scored)
    tests = [(1, "qkv_gate"), (7, "pooler"), (33, "o"), (130, "unpooler"), (5, "unpooler"), (8, "router"),
             (9, "router"), (512, "qkv_gate"), (1000, "unpooler")]
    show("== the 9 correctness tests ==", [(f"test {layer} m={m}", layer, m) for m, layer in tests])
    print(f"(an H200 has {SMS} SMs: a grid with fewer CTAs than that leaves SMs idle; more than that needs\n"
          " several CTAs per SM at once, or a second 'wave')")