"""
=============================================================================
CS349E · 第 14 讲 · 04_alphaevolve_tiling.py —— 迷你 AlphaEvolve：进化一段"按形状选分块"的启发式代码
=============================================================================

【这个脚本要回答的问题】
    AlphaEvolve 在 Gemini 的矩阵乘法内核上做的事：不是给每个形状单独调优，而是进化一段**对所有形状都适用**的
    选分块的代码。训练形状上进化，留出的形状上检验。这样的代码比专家写的好多少？离"每个形状都调到最优"还差多少？
    要花多少评估？
      演示 1：问题——形状、每个形状的最优配置、专家启发式，以及几条基线（一个配置走天下、按 M 分档的查表、逐形状最优）。
      演示 2：一次进化：分数怎样上升、哪些修改起了作用、最后的程序长什么样、产生它的那个 SEARCH/REPLACE。
      演示 3：泛化：随机对半分 vs 按模型分（8B 训练、70B 分片测试），5 个随机种子。
      演示 4：成本账：进化用了多少次"内核运行"，和逐形状调优比；AlphaEvolve 论文报告的数字。

【关于数字】"硬件"是第 13 讲的解析模型（H100 式参数，不是实测）；"LLM"是 evolve.MockLLM（模板化的修改 + 一点内核常识），
    没有任何真正的语言模型。这里看的是循环的结构和各个部件的作用，不是 AlphaEvolve 的真实效果。
【运行】python3 04_alphaevolve_tiling.py      （约 25 秒，其中约 12 秒在算"所有形状 × 所有配置"的时间表）
【依赖】numpy；本目录的 evolve（它导入 ../lecture13 的 sched_lab 和 ../lecture10 的 kernel_lab）；../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 evolve import (EXPERT, LAYERS, TOKENS, VALID, Evaluator, Program, Rule, Shape, TimeTable, all_shapes,  # noqa: E402
                    best_single_config, choose, evolve, format_diff, geomean_speedup, oracle_speedup, render,
                    simplify, split_shapes)
from console import print_header, print_observations, print_table  # noqa: E402

SEEDS = 5
BUDGET = 1000


def per_m_table(train, table):
    """
    "按 M 分档的查表"：对每个 M（token 数）在训练形状里挑一个最好的配置，写成一串 if M <= ... 的规则。
    这是工程师拿一个自动调优器、按常见的 M 分档调一遍就能得到的东西（Triton 的 autotune 常用 M 当 key）。
    【成本】要在每个训练形状上把 3744 个配置都跑一遍。
    """
    rules = []
    ms = sorted({s.M for s in train})
    for M in ms[:-1]:
        sh = [s for s in train if s.M == M]
        best = max(VALID, key=lambda c: geomean_speedup(Program((), c), sh, table))
        rules.append(Rule((("M", M),), best))
    sh = [s for s in train if s.M == ms[-1]]
    default = max(VALID, key=lambda c: geomean_speedup(Program((), c), sh, table))
    return simplify(Program(tuple(rules), default))


def demo1_problem(table, train, test):
    print_header("演示 1：问题——88 个形状、每个形状的最优配置、专家启发式与几条基线")
    print("形状：" + "、".join(f"{k} (N={n}, K={k2})" for k, (n, k2) in LAYERS.items()))
    print(f"      × M ∈ {TOKENS}（decode 的批大小到 prefill 的 token 数）= {len(LAYERS) * len(TOKENS)} 个；随机对半分成训练 {len(train)} 个、测试 {len(test)} 个")
    rows = []
    for name, M in (("8B gate_up", 1), ("8B gate_up", 16), ("8B gate_up", 128), ("8B gate_up", 1024),
                    ("8B gate_up", 8192), ("70B/8 qkv", 128), ("70B/8 qkv", 8192)):
        N, K = LAYERS[name]
        s = Shape(name, M, N, K)
        cfg, t = table.oracle(s)
        te = table.time(choose(EXPERT, s), s)
        rows.append([s.label(), str(cfg.as_tuple()), f"{t * 1e6:.1f} µs", f"{2 * M * N * K / t / 1e12:.0f}", f"{te / t:.2f}×"])
    print_table(["形状", "最优配置 (BM,BN,BK,warps,stages,group)", "时间", "TFLOPS", "专家启发式慢几倍"], rows)
    print("\n专家启发式（进化的起点）：\n" + render(EXPERT))
    one = best_single_config(train, table)
    pmt = per_m_table(train, table)
    o_tr, o_te = oracle_speedup(train, table), oracle_speedup(test, table)
    rows = []
    for name, prog, cost in (("专家启发式", EXPERT, "0（手写）"), (f"一个配置走天下 {one.default.as_tuple()}", one, f"{len(VALID) * len(train):,}"),
                             (f"按 M 分档查表（{pmt.n_rules()} 条规则）", pmt, f"{len(VALID) * len(train):,}")):
        tr, te = geomean_speedup(prog, train, table), geomean_speedup(prog, test, table)
        rows.append([name, f"{tr:.3f}×", f"{te:.3f}×", f"{te / o_te:.0%}", cost])
    rows.append(["逐形状最优（每个形状都调）", f"{o_tr:.3f}×", f"{o_te:.3f}×", "100%", f"{len(VALID)} × 每个新形状"])
    print_table(["方法", "训练形状（相对专家）", "测试形状（相对专家）", "测试上达到最优的", "要跑多少次内核"], rows)
    print_observations([
        "最优配置随形状变化很大：decode（M = 1、16）的时间几乎全花在读权重上，很多小块配置一样快，专家启发式已经是最优；"
        "prefill（M 几千）要 256×128 的大块、8 个 warp、4 级流水线，专家用的 128×128×64 / 4 warp / 3 级慢 30% 左右；"
        "差得最多的是中间的 M（例如 70B 分片的 qkv 在 M = 128 时慢 2.5 倍）。",
        f"逐形状最优相对专家快 {o_te:.2f}×（测试形状的几何平均）——这是上限，但每来一个新形状都要再调一遍。"
        f"按 M 分档的查表在训练形状上调好，到测试形状上达到最优的 {geomean_speedup(pmt, test, table) / o_te:.0%}：它只看 M，同一个 M 下不同的 N、K 用同一个配置。",
        "一个启发式（一段代码）的价值在于**对没见过的形状也能直接用**，不用每次重新调优——AlphaEvolve 论文正是这样说的："
        "基于搜索的调优\"打断研究的工作流，每次输入形状一变就要重调\"。",
    ])
    return pmt


def demo2_run(table, train, test):
    print_header(f"演示 2：一次进化（随机种子 0，最多提出 {BUDGET} 个程序）")
    ev = Evaluator(train, table)
    out = evolve(ev, BUDGET, seed=0, log_every=200)
    best, res = out["best"]
    print("\n全局最好分数被刷新的时刻（前 8 次和最后 4 次）：")
    imp = out["improvements"]
    show = imp[:8] + ([("…", None, None)] if len(imp) > 12 else []) + imp[-4:] if len(imp) > 12 else imp
    print_table(["第几个程序", "新的最好分数", "哪种修改"],
                [[str(i), f"{s:.3f}" if s else "", op or ""] for i, s, op in show])
    rows = [[op, out["tried"].get(op, 0), out["accepted"].get(op, 0),
             f"{out['accepted'].get(op, 0) / max(1, out['tried'].get(op, 0)):.0%}"] for op in sorted(out["tried"])]
    print_table(["修改的种类", "提出次数", "比父程序好的次数", "成功率"], rows)
    simp = simplify(best)
    print(f"\n最好的程序：训练形状 {res['score']:.3f}×、测试形状 {geomean_speedup(best, test, table):.3f}×；"
          f"{best.n_rules()} 条规则，删掉死规则后 {simp.n_rules()} 条（行为完全不变）：\n" + render(simp))
    print("产生它的那次修改（父程序 → 它，SEARCH/REPLACE 格式）：\n" + format_diff(res["diff"]))
    print_observations([
        f"分数在前几百个程序里上升最快，之后变慢：{BUDGET} 个程序后训练形状上 {res['score']:.3f}×（上限 {oracle_speedup(train, table):.3f}×）。"
        "这和 AlphaEvolve 论文图 8 的消融曲线形状一样：早期的改进容易，越往后越难。",
        "\"按内核常识的定向修改\"（targeted）的成功率远高于单字段的小改动：单字段改动常常跨不过\"中间状态不能启动\"的山谷——"
        "从 128×128 / 4 warp 走到 256×128 / 8 warp，中间的 256×128 / 4 warp 每线程要 296 个寄存器，根本启动不了。"
        "一次改好几个字段、改得\"连贯\"，正是大模型比小模型、比随机变异强的地方（05 脚本的 \"small LLM\" 消融）。",
        "进化出来的代码常有\"死规则\"（被前面的规则完全挡住、永远走不到），遗传编程里叫 bloat。simplify 做的是第 12 讲 DCE 的同一件事；"
        "AlphaEvolve 可以把\"简洁\"当作额外的分数，或者让 LLM 给可读性打分。",
    ])
    return out, ev


def demo3_generalization(table, shapes):
    print_header(f"演示 3：泛化——两种训练 / 测试划分，各 {SEEDS} 个随机种子")
    rows, summ = [], {}
    for how, label in (("random", "随机对半分"), ("by_model", "8B 训练、70B 分片测试")):
        train, test = split_shapes(shapes, how)
        pmt = per_m_table(train, table)
        trs, tes = [], []
        for seed in range(SEEDS):
            out = evolve(Evaluator(train, table), BUDGET, seed=seed)
            bp, br = out["best"]
            trs.append(br["score"])
            tes.append(geomean_speedup(bp, test, table))
        o_te = oracle_speedup(test, table)
        p_te = geomean_speedup(pmt, test, table)
        summ[how] = (float(np.median(trs)), float(np.median(tes)), float(np.median(tes)) / o_te, p_te / o_te)
        rows.append([label, "进化（中位数，最小–最大）", f"{np.median(trs):.3f}× ({min(trs):.3f}–{max(trs):.3f})",
                     f"{np.median(tes):.3f}× ({min(tes):.3f}–{max(tes):.3f})", f"{np.median(tes) / o_te:.0%}"])
        rows.append(["", "按 M 分档查表", f"{geomean_speedup(pmt, train, table):.3f}×", f"{p_te:.3f}×", f"{p_te / o_te:.0%}"])
        rows.append(["", "逐形状最优（上限）", f"{oracle_speedup(train, table):.3f}×", f"{o_te:.3f}×", "100%"])
    print_table(["划分", "方法", "训练形状", "测试形状", "测试上达到最优的"], rows)
    r, m = summ["random"], summ["by_model"]
    print_observations([
        f"随机对半分时，测试形状和训练形状来自同样的 8 种投影、同样的 11 种 M，进化出来的启发式在测试形状上（{r[1]:.3f}×）几乎和训练形状上"
        f"（{r[0]:.3f}×）一样好——它学到的主要是\"M 决定块的大小\"这类规律。",
        f"按模型分时，测试形状的 N、K 都没见过，训练和测试之间的落差更大（{m[0]:.3f}× → {m[1]:.3f}×）；不过相对各自的上限，"
        f"两种划分都达到了最优的 {min(r[2], m[2]):.0%}–{max(r[2], m[2]):.0%}。AlphaEvolve 的做法是从内核的真实用户那里收集形状，一半训练、一半检验——"
        "测试集要像将来真的会遇到的形状。",
        f"进化出的启发式和\"按 M 分档查表\"（{min(r[3], m[3]):.0%}–{max(r[3], m[3]):.0%}）差不多。查表要在每个训练形状上把 3744 个配置都跑一遍；"
        "进化只跑它选中的那些（演示 4），还可以用 N、K 写条件。在这个问题上，规律主要由 M 决定，所以简单的查表已经很强——"
        "这提醒我们：给进化搜索找一个**强**的基线，才能知道它到底带来了什么。",
    ])


def demo4_cost(out, ev, train, shapes):
    print_header("演示 4：成本账（演示 2 的那次进化）")
    rows = [["进化：提出的程序", f"{BUDGET:,}"],
            ["  被静态检查拒掉的（有分支返回不能启动的配置）", f"{ev.n_static_reject:,}"],
            ["  只做了级联第一级（6 个探针形状）就淘汰的", f"{ev.n_eval_probe:,}"],
            ["  完整评估的", f"{ev.n_eval_full:,}"],
            ["进化：一共\"跑了几次内核\"（每个被评估的程序 × 它被评估的形状）", f"{ev.bench_runs:,}"],
            ["逐形状穷举调优训练形状", f"{len(VALID) * len(train):,}"],
            ["逐形状穷举调优全部 88 个形状", f"{len(VALID) * len(shapes):,}"],
            ["来了一个新形状：启发式", "0（直接用）"],
            ["来了一个新形状：逐形状调优", f"{len(VALID):,}（穷举）或几十到几百（第 13 讲的代价模型引导）"]]
    print_table(["项目", "次数"], rows)
    print_observations([
        f"这次进化一共\"跑了\" {ev.bench_runs:,} 次内核，比在训练形状上穷举（{len(VALID) * len(train):,} 次）少，但不是少几个数量级。"
        "它买到的是一段**可以直接用在新形状上**的代码。真实系统里每次\"跑内核\"都要上 TPU / GPU；AlphaEvolve 的评估器是并行的"
        "（论文说每个新方案可以花上百个计算小时去评估），整个流水线为吞吐优化。",
        "AlphaEvolve 论文报告的真实结果：进化出的分块启发式相对专家设计的启发式，所有内核平均快 23%，Gemini 的整体训练时间少 1%；"
        "内核优化的时间从\"几个月的专门工程\"缩短到\"几天的自动实验\"；这个启发式已经部署到生产。正确性\"由构造保证\"：只改分块，不改算的东西。",
        "我们的迷你版相对\"专家\"的倍数（约 1.3–1.4×）不能和 23% 比：专家启发式是我写的，好坏由我决定。"
        "更有意义的比较是\"离逐形状最优还差多少\"（演示 3 的最后一列）。",
    ])


if __name__ == "__main__":
    shapes = all_shapes()
    table = TimeTable()
    for s in shapes:
        table.oracle(s)          # 先把"所有形状 × 所有配置"的时间都算出来（约 12 秒）
    train, test = split_shapes(shapes)
    demo1_problem(table, train, test)
    o, e = demo2_run(table, train, test)
    demo3_generalization(table, shapes)
    demo4_cost(o, e, train, shapes)