"""
=============================================================================
CS349E · 第 14 讲 · 01_kernelbench_harness.py —— 一个迷你 KernelBench：任务、评估器与 fast_p
=============================================================================

【这个脚本要回答的问题】
    KernelBench 把"LLM 会不会写内核"变成了一个可以自动判分的考试。考卷长什么样？怎么判分？fast_p 怎么算？
      演示 1：解剖一道题（diag(a) @ B）：参考实现、一个候选、评估器的每一步、加速比。
      演示 2：候选库（模拟 LLM 写出的 28 个内核，覆盖 9 道题）分别交给"朴素"和"加固"的评估器判分；
              失败的原因分成"执行失败"（编译 / 运行出错）和"结果不对"两类（KernelBench 图 2 的分法）。
      演示 3：fast_p 曲线：一个"从候选库里随机抽一个交卷"的模型，在不同的 p 上能拿多少分；重复采样 k 次的 fast_1@k。
      演示 4：正确性测试要测几组、测什么：同一个 bug 在不同输入分布上的表现；KernelBench 观察到的"要么全对、要么全错"。

【关于数字】时间是本机 numpy 的墙钟时间（每次运行会有几个百分点的波动）；观察是按这次运行的数字生成的。
    候选库是我写的，它的"成分"（多少编译错误、多少慢的正确实现）是仿照 KernelBench 论文报告的现象设计的，
    所以 fast_p 曲线的形状只能和论文**定性**地比较。
【运行】python3 01_kernelbench_harness.py      （约 40 秒）
【依赖】numpy；本目录的 kbench；../lecture01 的 console。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, HERE)
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from kbench import (DISTS, DIST_NOTES, LIBRARY, TASK, TASKS, check_output, evaluate, fast_p_at_k,  # noqa: E402
                    library_for, load_candidate)
from console import print_header, print_observations, print_table  # noqa: E402


def fmt_speed(r):
    """加速比的显示：不正确的写状态；正确的写 "x.xx×"。"""
    if not r.correct:
        return {"compile_error": "编译错误", "runtime_error": "运行出错", "wrong": "结果不对",
                "rejected": "拒绝"}[r.status]
    return f"{r.speedup:.2f}×" if r.speedup < 100 else f"{r.speedup:.0f}×"


def demo1_anatomy():
    print_header("演示 1：解剖一道题——L1_diag_matmul：diag(a) @ B，N = 2048")
    t = TASK["L1_diag_matmul"]
    cand = next(c for c in LIBRARY if c.task == t.name and c.name == "row_scale")
    print("参考实现（\"LLM\"拿到的题目）：" + t.ref_src)
    print("候选（\"LLM\"交的卷）：" + cand.src)
    r = evaluate(t, cand.src, "hardened", cand_name=cand.name)
    print(f"评估器：静态检查 → 编译 → 5 组输入（{'、'.join(DISTS)}，参数 1 组默认 + 2 组随机）逐个和参考比较（atol = rtol = 1e-2）"
          f"→ 计时（每次换新输入、和参考交替测、取中位数）")
    print(f"结果：{r.status}；参考 {r.t_ref * 1e3:.2f} ms，候选 {r.t_cand * 1e3:.3f} ms，加速比 {r.speedup:.1f}×")
    flops_ref = 2 * 2048 ** 3
    print(f"参考实现做了 2N³ = {flops_ref / 1e9:.1f} GFLOP 的稠密矩阵乘法；候选只做 N² = {2048 ** 2 / 1e6:.1f} M 次乘法——"
          f"运算量差 {flops_ref / 2048 ** 2:.0f} 倍，但时间只差 {r.speedup:.0f} 倍：候选受限于读写 B 和 C 的 {2 * 2048 ** 2 * 4 / 1e6:.0f} MB。")
    print_observations([
        f"这道题的\"优化\"不是写更好的内核，而是换算法：对角矩阵不用真的造出来。KernelBench 报告 Claude 3.5 Sonnet 在这道题上拿到 13×；"
        f"我们在 numpy 上量到 {r.speedup:.1f}×（硬件不同、参考实现的开销不同，倍数不能直接比）。",
        "判分只看两件事：输出对不对（和参考比，容差 1e-2）、比参考快多少。它不管候选\"怎么做到的\"——这是 KernelBench 的优点"
        "（任何技巧都可以用：融合、换算法、专用指令），也是它的弱点（演示 2、02 脚本：有些\"技巧\"是钻评估器的空子）。",
        "正确性只能靠随机测试来\"估计\"：判断两个程序在所有输入上等价是不可判定的（KernelBench 附录 B.2 引用停机问题）。",
    ])


def demo2_library():
    print_header("演示 2：候选库交给两种评估器——朴素（KernelBench 最初的做法）vs 加固")
    rows, results = [], {"naive": {}, "hardened": {}}
    for c in LIBRARY:
        rn = evaluate(TASK[c.task], c.src, "naive", cand_name=c.name)
        rh = evaluate(TASK[c.task], c.src, "hardened", cand_name=c.name)
        results["naive"][(c.task, c.name)] = rn
        results["hardened"][(c.task, c.name)] = rh
        rows.append([c.task, c.name, c.intent[:30] + ("…" if len(c.intent) > 30 else ""), fmt_speed(rn), fmt_speed(rh),
                     rh.msg[:34] if not rh.correct else ""])
    print_table(["任务", "候选", "想做什么", "朴素", "加固", "加固评估器的说明"], rows,
                aligns=["<", "<", "<", ">", ">", "<"])
    print()
    cats = {"执行失败（编译 / 运行出错）": ("compile_error", "runtime_error"), "结果不对": ("wrong",),
            "被拒绝": ("rejected",), "正确": ("correct",)}
    srows = []
    for name, sts in cats.items():
        srows.append([name] + [sum(1 for r in results[m].values() if r.status in sts) for m in ("naive", "hardened")])
    print_table(["状态", "朴素", "加固"], srows)
    flips = [k for k in results["naive"] if results["naive"][k].correct and not results["hardened"][k].correct]
    sp = sorted(r.speedup for r in results["hardened"].values() if r.correct)
    med = float(np.median(sp))
    slow = sum(1 for s in sp if s < 1)
    gu = results["hardened"][("L3_mlp_block", "fused_gate_up")].speedup
    gu_text = (f"L3 的 gate/up 合并反而更慢（{gu:.2f}×）——切片出来的不连续视图拖慢了后面的逐元素运算" if gu < 1
               else f"L3 的 gate/up 合并也几乎没变（{gu:.2f}×）")
    print_observations([
        f"{len(LIBRARY)} 个候选里，加固的评估器判对 {len(sp)} 个；正确的候选里有 {slow} 个比参考还慢，加速比的中位数 {med:.2f}×。"
        f"KernelBench 也报告：多数模型写出的**正确**内核，加速比的中位数在 1 以下或刚过 1（附录 B.3）。",
        "执行失败（语法错误、维度写错、形状不匹配）一跑就暴露，错误信息也很具体；\"结果不对\"只有一个\"不等于参考\"的信号。"
        "KernelBench 的结论相同：推理模型主要是执行失败少，\"结果不对\"的比例各家差不多（图 2）。",
        "朴素判对、加固判错的有 " + "、".join(f"{t}/{n}" for t, n in flips) + "：不减最大值的 softmax 在对数正态输入（值上百）上溢出；"
        "\"一遍算方差\"在均值 1000 的输入上抵消误差很大；把缩放因子折进权重只在缩放因子为正时成立，参数一随机化就错。"
        "它们都不是\"故意作弊\"，而是只在默认的标准正态输入、默认初始化的参数上才对。",
        "L1_layernorm/no_eps（忘了 eps）连加固的评估器也判对了：只有某一行几乎是常数（方差 ≈ 0）时它才会出错，而我们的 5 种分布都没有这种输入。"
        "测试永远不完备，加一种分布只能多抓一类 bug。",
        "numpy 上的\"内核\"想变快，要么换算法（diag、矩阵乘法后求和 → 矩阵向量乘法），要么去掉昂贵的运算（GELU 里的 x**3 是慢的 pow）。"
        "\"少分配几个临时数组\"在这里几乎没用（softmax、epilogue 的就地版本都在 1.0× 附近），" + gu_text +
        "。LLM 认为的\"优化\"在某个平台上不一定是优化（KernelBench §4.4：跨 GPU 的加速比也不稳定）。",
    ])
    return results


def demo3_fast_p(results):
    print_header("演示 3：fast_p——一个\"从候选库里随机抽一个交卷\"的模型能拿多少分")
    hard = results["hardened"]
    ps = (0.0, 0.5, 1.0, 1.2, 2.0, 5.0, 10.0)
    q = {p: [] for p in ps}
    best = {p: [] for p in ps}
    for t in TASKS:
        rs = [hard[(t.name, c.name)] for c in library_for(t.name)]
        for p in ps:
            ok = [r.correct and r.speedup > p for r in rs]
            q[p].append(sum(ok) / len(ok))
            best[p].append(any(ok))
    rows = [[f"p = {p:g}", f"{np.mean(q[p]):.2f}", f"{np.mean(best[p]):.2f}"] for p in ps]
    print_table(["阈值", "fast_p（随机抽 1 个，期望值）", "候选库里最好的那个（\"先知\"）"], rows)
    print()
    ks = (1, 2, 5, 10, 100)
    print_table(["采样次数 k"] + [str(k) for k in ks],
                [["fast_1@k"] + [f"{fast_p_at_k(q[1.0], k):.2f}" for k in ks],
                 ["fast_0@k（正确率）"] + [f"{fast_p_at_k(q[0.0], k):.2f}" for k in ks]])
    print_observations([
        f"fast_0（正确率）{np.mean(q[0.0]):.2f}，fast_1 只有 {np.mean(q[1.0]):.2f}，p 越大分越低——和 KernelBench 图 3 的形状一样："
        "大部分正确的内核并不比参考快。KernelBench 论文里，one-shot 的 fast_1（相对 PyTorch eager）在 21 个\"模型 × 级别\"组合里最高 36%"
        "（DeepSeek-R1 在 Level 2），21 个里有 17 个不超过 10%（表 1）。",
        f"重复采样：每题抽 k 个、有一个成功就算成功（需要一个能自动判分的评估器来挑）。k = 10 时 fast_1@k = {fast_p_at_k(q[1.0], 10):.2f}，"
        f"但再多也到不了\"先知\"的 {np.mean(best[1.0]):.2f} 以上——没有好候选的题采多少次都没用。"
        "KernelBench：DeepSeek-V3 在 Level 2 上采 100 次 fast_1 从 4% 升到 37%，但 34 个卷积变体采 100 次也没有一个正确的。",
        "fast_p 是\"按题计数\"：一道题快 100× 和快 1.01× 在 fast_1 里一样算 1。它不受个别巨大加速比的影响（03 脚本演示 5 对比几种汇总方式），"
        "代价是看不出\"快了多少\"。",
    ])


def demo4_tests():
    print_header("演示 4：正确性测试要测什么——同一个候选在不同输入分布上的表现")
    picks = [("L1_softmax", "no_max"), ("L1_layernorm", "one_pass_var"), ("L1_layernorm", "no_eps"),
             ("L1_gelu", "sigmoid_approx"), ("L1_softmax", "inplace")]
    rows = []
    for tn, cn in picks:
        t = TASK[tn]
        c = next(c for c in LIBRARY if c.task == tn and c.name == cn)
        fwd, init, _ = load_candidate(c.src, {"np": np, "math": __import__("math")})
        rng = np.random.default_rng(0)
        params = t.make_params(rng, False)
        pc = init(dict(params)) if init else params
        row = [f"{tn}/{cn}"]
        for d in DISTS:
            passed = 0
            for _ in range(5):
                x = t.make_inputs(rng, d)
                with np.errstate(all="ignore"):
                    ok, _ = check_output(fwd(x, pc, lambda s, dt=np.float32: np.empty(s, dt)),
                                         np.asarray(t.ref_fn()(x, params, None), np.float32), True)
                passed += ok
            row.append(f"{passed}/5")
        rows.append(row)
    print_table(["候选"] + list(DISTS), rows)
    print("分布：" + "；".join(f"{d} = {DIST_NOTES[d]}" for d in DISTS))
    print()
    # "全对或全错"：100 组标准正态输入
    allrows = []
    for c in LIBRARY:
        t = TASK[c.task]
        fwd, init, err = load_candidate(c.src, {"np": np, "math": __import__("math")})
        if err:
            continue
        rng = np.random.default_rng(1)
        params = t.make_params(rng, False)
        try:
            pc = init(dict(params)) if init else params
        except Exception:  # noqa: BLE001
            continue
        passed = 0
        n = 20
        for _ in range(n):
            x = t.make_inputs(rng, "normal")
            try:
                with np.errstate(all="ignore"):
                    ok, _ = check_output(fwd(x, pc, lambda s, dt=np.float32: np.empty(s, dt)),
                                         np.asarray(t.ref_fn()(x, params, None), np.float32), True)
            except Exception:  # noqa: BLE001
                ok = False
            passed += ok
        allrows.append(passed)
    hist = {"0/20": sum(1 for p in allrows if p == 0), "20/20": sum(1 for p in allrows if p == 20),
            "其他": sum(1 for p in allrows if 0 < p < 20)}
    print(f"能运行的 {len(allrows)} 个候选，在 20 组标准正态输入上：全错 {hist['0/20']} 个，全对 {hist['20/20']} 个，部分对 {hist['其他']} 个。")
    print_observations([
        "\"错\"的方式决定了要测什么输入：不稳定的 softmax 在对数正态和 offset（值上百、上千，exp 溢出）上出错；一遍算方差只在 offset（均值 1000）上出错；"
        "GELU 的粗糙近似只在有 −2 附近取值的分布（正态、拉普拉斯）上超出容差，在 (−1, 1) 的均匀分布和全正的分布上都\"通过\"。"
        "只用一种分布的测试，总有一些 bug 发现不了。",
        f"在同一种分布上，候选几乎都是\"要么全对、要么全错\"（这里 {hist['其他']} 个部分对）——KernelBench 用 100 个生成的内核做实验，"
        "5 组和 100 组输入的判决完全一致，所以只测 5 组。但这个结论只对\"同一种分布\"成立：换分布才是抓 bug 的关键"
        "（KernelBench v0.1 加了 rand_mix：正态、均匀、拉普拉斯、指数、对数正态）。",
        "no_eps 在所有分布上都是 5/5：方差为 0 的行才会让它除以 0。要抓住它，得专门构造\"边界输入\"（常数行、全零、极大值……）"
        "——这是人写测试时会想到、随机测试很难碰到的东西。",
    ])


if __name__ == "__main__":
    demo1_anatomy()
    res = demo2_library()
    demo3_fast_p(res)
    demo4_tests()