"""
=============================================================================
CS349E · 第 12 讲 · 02_graph_passes.py —— 图级优化：分解、CSE、常量折叠、化简、权重折叠、模式匹配
=============================================================================

【这个脚本要回答的问题】
    抓到图以后，编译器在图上做一系列"pass"，每个 pass 把图改写成一张等价但更好的图。
      演示 1：追踪一层 Llama 风格的 Transformer（和第 2 讲的模型相同尺寸），结果和第 2 讲的 numpy 实现逐位比较。
      演示 2：一条 pass 流水线：模式匹配 → 分解 → CSE → 常量折叠 → 化简 → 权重折叠，每一步之后的节点数和误差；pass 的顺序为什么重要。
      演示 3：常量折叠与代数化简在一个小例子上；为什么 x / c 不改写成 x · (1/c)。
      演示 4：把 RMSNorm 的缩放折进下一个矩阵乘法的权重。

【运行】python3 02_graph_passes.py      （< 1 秒）
【依赖】numpy；本目录的 graph_lab；../lecture02 的 transformer_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 graph_lab import (LayerCfg, _reference_prefill, constant_fold, cse, decompose, dce, fold_norm_weight,  # noqa: E402
                       layer_params, match_attention, prefill_layer, run, simplify, trace, POINTWISE, REDUCE)
from console import print_header, print_observations, print_table  # noqa: E402

CFG = LayerCfg(128, 8, 2, 16, 344)       # 第 2 讲 TINY 的一层：d = 128、8 个 query 头、2 组 KV、头维度 16、FFN 344
T = 16


def stats(g):
    ops = [n.op for n in g.nodes]
    return {"节点": g.count(), "matmul": ops.count("matmul"), "逐元素": sum(o in POINTWISE for o in ops),
            "归约": sum(o in REDUCE for o in ops), "视图": ops.count("reshape") + ops.count("transpose"),
            "复合": sum(o in ("softmax", "rms_norm", "silu") for o in ops), "flash_attn": ops.count("flash_attn"),
            "常量": ops.count("const"), "参数": ops.count("param")}


def demo1_trace_layer():
    print_header(f"演示 1：追踪一层 Llama 风格的 Transformer（d = {CFG.d}，{CFG.H} 个 query 头 / {CFG.G} 组 KV，T = {T}）")
    p = layer_params(CFG)
    x = np.random.default_rng(0).standard_normal((T, CFG.d))
    ref = _reference_prefill(CFG, p, x)
    g = trace(prefill_layer(CFG, T), {"x": (T, CFG.d)}, p)
    eager = prefill_layer(CFG, T)(x=x, p=p)
    st = stats(g)
    print_table(list(st), [list(st.values())])
    print(f"\n图执行 vs 第 2 讲 transformer_lab 的实现：最大差 {np.abs(run(g, {'x': x})[0] - ref).max():.1e}；"
          f"同一份代码立即执行：{np.abs(eager - ref).max():.1e}")
    print_observations([
        "第 2 讲的模型用 einsum 和 [批, T, 头, d] 的四维数组写成；这里用\"假张量\"能追踪的运算重写了一遍"
        "（RoPE 写成 x·cos + rot_pairs(x)·sin，GQA 写成把同一组的头排在一起），结果相同（1e-15 量级）。",
        "一层的图有几十个节点：9 个矩阵乘法（7 个投影 + 注意力里的 2 个）、RoPE 和残差的逐元素运算、11 个 reshape / transpose（分头、合并头），"
        "以及三个复合运算（两个 RMSNorm、一个 softmax）和一个 SiLU。",
        "视图（reshape / transpose）在 PyTorch 里通常不复制数据，只改下标的算法；编译器把它们当成\"读的时候换个下标\"，不生成内核（03 脚本）。",
    ])
    return g, p, x, ref


def demo2_pipeline(g, x, ref):
    print_header("演示 2：pass 流水线（每一步之后的图和误差）")
    rows = []
    stages = [("追踪得到的图", g)]
    g1, n1 = match_attention(g)
    stages.append((f"模式匹配：注意力 → flash_attn（{n1} 处）", g1))
    g2 = decompose(g1)
    stages.append(("分解：softmax / rms_norm / silu → 基本运算", g2))
    g3 = cse(g2)
    stages.append(("公共子表达式消除", g3))
    g4, n4 = constant_fold(g3)
    stages.append((f"常量折叠（{n4} 个节点）", g4))
    g5, n5 = simplify(g4)
    stages.append((f"代数化简（{n5} 处）", g5))
    g6, n6 = fold_norm_weight(g5)
    stages.append((f"RMSNorm 的缩放折进权重（{n6} 处）", g6))
    for name, gg in stages:
        st = stats(gg)
        rows.append([name] + list(st.values()) + [f"{np.abs(run(gg, {'x': x})[0] - ref).max():.1e}"])
    print_table(["阶段"] + list(stats(g)) + ["和第 2 讲的最大差"], rows)
    _, nb = match_attention(decompose(g))
    print(f"\n换个顺序：先分解、再做模式匹配——找到 {nb} 处注意力（softmax 已经被拆成 max、sub、exp、sum、div，模式认不出来了）。")
    print_observations([
        "模式匹配把\"matmul → 乘常数 → 加因果掩码 → softmax → matmul\"这一串换成一个 flash_attn 节点（像 TorchInductor 把注意力的子图换成调用 FlashAttention）："
        "两个 matmul、一个 softmax、掩码都不见了。",
        "分解把复合运算展开成基本运算，节点数变多：后端只要实现几十种基本运算；展开以后融合 pass 可以按自己的方式把它们重新组合（03 脚本）。",
        "CSE 合并了两个完全相同的常量（q 和 k 的 RoPE 用同一张 cos 表、同一张 sin 表）。常量折叠和代数化简在这张图上什么也没做："
        "追踪时 Python 已经把 1/√d 这类算术算掉了，模型代码里也没有\"参数的转置\"这类可以预先算好的东西（演示 3 换一个例子）。",
        "pass 的顺序很重要：先分解再匹配，就认不出注意力了。真实的编译器（XLA、Inductor）里 pass 的排序和反复迭代是大量工程经验的结晶。",
        "每一步之后结果都和第 2 讲相差 1e-15 量级：每个 pass 都必须保持语义——这也是编译器测试的核心（差分测试，第 14 讲 KernelBench 的正确性检查也是这个）。",
    ])


def demo3_fold_and_simplify():
    print_header("演示 3：常量折叠与代数化简（一个小例子）")
    rng = np.random.default_rng(1)
    p = {"W": rng.standard_normal((16, 8)), "b": rng.standard_normal(16)}

    def f(x, p):
        W = p["W"].transpose(1, 0) * 2.0           # 参数的变换：可以在编译时算好
        bias = p["b"] * 0.5
        y = (x @ W) * 1.0 + 0.0 + bias             # ·1、+0：可以删掉
        z = y.transpose(1, 0).transpose(1, 0)      # 两次转置：抵消
        return -(-z)                               # 两次取负：抵消

    g = trace(f, {"x": (4, 8)}, p)
    x = rng.standard_normal((4, 8))
    ref = run(g, {"x": x})[0]
    g1, nf = constant_fold(g)
    g2, ns = simplify(g1)
    g3 = dce(g2)
    rows = [[name, gg.count(), " ".join(n.op for n in gg.nodes if n.op not in ("input", "param", "const")),
             f"{np.abs(run(gg, {'x': x})[0] - ref).max():.1e}"]
            for name, gg in (("追踪", g), (f"常量折叠（{nf} 个）", g1), (f"化简（{ns} 处）+ DCE", g3))]
    print_table(["阶段", "运算节点", "运算", "误差"], rows)
    v = rng.standard_normal(1_000_000)
    c = 3.0
    diff = np.count_nonzero(v / c != v * (1.0 / c))
    print(f"\n另：100 万个随机数，x / 3 和 x · (1/3) 有 {diff:,} 个（{diff / len(v):.1%}）结果不同（最后一位）。")
    print_observations([
        "常量折叠：W.transpose · 2 和 b · 0.5 只依赖参数，推理时参数不变，编译时算一次就够了——每次前向省掉一次 8×16 的转置和乘法。"
        "真实模型里常见的例子：权重的转置 / 重排（为某种内核布局准备）、量化参数的预计算、把 BatchNorm 折进卷积。",
        "化简：·1、+0、两次转置、两次取负都被删掉；DCE（死代码删除）再把没人用的常量清掉。这些模式在\"人写的代码\"里不多，"
        "但在**自动生成**的图里（分解、自动微分、别的 pass 之后）非常常见——化简 pass 是给其他 pass 收拾残局的。",
        "我们**没有**把 x / c 改成 x · (1/c)：两者在数学上相等，但有约三成的数结果不同（1/3 本身就要舍入）。编译器默认只做逐位不变的改写；"
        "愿意接受误差的改写（重新结合、乘倒数、近似的 exp）要用户显式允许（-ffast-math，前置讲 P4 §8.3）。",
    ])


def demo4_fold_norm(g, p, x, ref):
    print_header("演示 4：把 RMSNorm 的缩放 w 折进下一个矩阵乘法：(x̂ · w) @ W = x̂ @ (diag(w) · W)")
    g0 = decompose(g)
    g1, n = fold_norm_weight(g0)
    out0, out1 = run(g0, {"x": x})[0], run(g1, {"x": x})[0]
    params = [nd.name for nd in g1.nodes if nd.op == "param"]
    print(f"折叠了 {n} 处；新图的参数：{', '.join(params)}")
    print(f"逐元素乘法：{sum(nd.op == 'mul' for nd in g0.nodes)} → {sum(nd.op == 'mul' for nd in g1.nodes)}；"
          f"折叠前后最大差 {np.abs(out1 - out0).max():.1e}（和第 2 讲 {np.abs(out1 - ref).max():.1e}）")
    print_observations([
        "条件：RMSNorm 的输出 x̂ · w 只被\"@ 参数矩阵\"使用。注意力前的 RMSNorm 被 Q、K、V 三个投影使用——三个权重矩阵各折进一份 w；"
        "MLP 前的被 gate、up 两个使用。折叠以后每个 token 少做两次 d 维的逐元素乘法，RMSNorm 的参数也不用再读。",
        "数学上完全相等，浮点结果只差最后几位（1e-15）：乘法的顺序变了。",
        "同一个技巧在量化里很常见：SmoothQuant 把激活的逐通道缩放转移到权重上、QuaRot 把旋转矩阵折进相邻的权重（第 3 讲）——"
        "都是\"图级的代数改写 + 常量折叠\"。",
    ])


if __name__ == "__main__":
    g, p, x, ref = demo1_trace_layer()
    demo2_pipeline(g, x, ref)
    demo3_fold_and_simplify()
    demo4_fold_norm(g, p, x, ref)