"""
=============================================================================
CS349E · 第 12 讲 · 03_fusion.py —— 算子融合：自动地少读写显存、少启动内核
=============================================================================

【这个脚本要回答的问题】
    图编译器最大的收益来自融合（PyTorch 2 论文的消融实验）。融合把相邻的运算放进同一个内核，
    中间结果留在片上；同时内核变少，启动开销也少了。
      演示 1：一个小层上的分区：每个内核里有哪些运算、读写多少字节。
      演示 2：Llama-3.1-8B 的一层，prefill 2048 个 token：eager（复合运算各自一个内核）/ 分解但不融合 / 融合 / 融合 + 注意力模式匹配。
      演示 3：同一层的 decode（KV 长度 4096）：融合省的主要是内核个数，而不是字节。
      演示 4：PyTorch 2 论文的消融实验：去掉融合、内联、CUDA Graphs 各掉多少。

【关于数字】时间是 roofline 模型（H100：BF16 矩阵乘法 989 TFLOPS、其他运算 67 TFLOPS、3.35 TB/s、每个内核固定开销 4 µs），
    不是实测；用来比较几种分区之间的差别。融合规则是 TorchInductor 风格的简化版，见 graph_lab.fuse 的说明。
【运行】python3 03_fusion.py      （约 1 秒）
【依赖】numpy；本目录的 graph_lab；../lecture01 的 console。
"""

import os
import sys

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 (H100, LLAMA8B, LayerCfg, constant_fold, cse, decode_layer, decompose, fold_norm_weight,  # noqa: E402
                       fuse, graph_cost, kernel_time, layer_params, match_attention, prefill_layer, simplify, trace)
from console import fmt_bytes, fmt_time, print_header, print_observations, print_table  # noqa: E402


def optimize(g, attention=True):
    """标准的 pass 流水线（02 脚本）：[模式匹配] → 分解 → CSE → 常量折叠 → 化简 → 权重折叠。"""
    if attention:
        g, _ = match_attention(g)
    g = decompose(g)
    g = cse(g)
    g, _ = constant_fold(g)
    g, _ = simplify(g)
    g, _ = fold_norm_weight(g)
    return g


def param_shapes(cfg):
    return {"attn_norm": (cfg.d,), "wq": (cfg.d, cfg.H * cfg.dh), "wk": (cfg.d, cfg.G * cfg.dh),
            "wv": (cfg.d, cfg.G * cfg.dh), "wo": (cfg.H * cfg.dh, cfg.d), "mlp_norm": (cfg.d,),
            "w_gate": (cfg.d, cfg.f), "w_up": (cfg.d, cfg.f), "w_down": (cfg.f, cfg.d)}


def variants(g):
    """四种分区：(名字, 内核列表的生成函数)。"""
    return [("eager（softmax、RMSNorm、SiLU 各自一个内核）", lambda b: fuse(g, enable=False, bind=b)),
            ("分解成基本运算，不融合", lambda b: fuse(decompose(g), enable=False, bind=b)),
            ("分解 + 融合", lambda b: fuse(optimize(g, attention=False), bind=b)),
            ("分解 + 融合 + 注意力换成 FlashAttention", lambda b: fuse(optimize(g, attention=True), bind=b))]


def demo1_small():
    print_header("演示 1：一个小层（d = 128，T = 16）的分区")
    cfg = LayerCfg(128, 8, 2, 16, 344)
    p = layer_params(cfg)
    g = trace(prefill_layer(cfg, 16), {"x": (16, cfg.d)}, p)
    ks = fuse(optimize(g))
    rows = [[i, k.kind, k.label, len(k.nodes), fmt_bytes(k.reads), fmt_bytes(k.writes)] for i, k in enumerate(ks)]
    print_table(["#", "类型", "包含的运算", "节点数", "读", "写"], rows)
    k0 = fuse(decompose(g), enable=False)
    print(f"\n不融合时 {len(k0)} 个内核、读写 {fmt_bytes(sum(k.reads + k.writes for k in k0))}；"
          f"融合后 {len(ks)} 个、{fmt_bytes(sum(k.reads + k.writes for k in ks))}。")
    print_observations([
        "第 0 个内核是 rowwise：RMSNorm 的平方、求和、加 ε、rsqrt、乘回去都在一个内核里，整行留在片上，只读 x、写 x̂。",
        "Q、K 投影的内核把 RoPE（乘 cos、rot_pairs、乘 sin、相加）作为 epilogue 并了进去：矩阵乘法的结果在寄存器里就旋转好再写回。"
        "O 投影 + 残差、down 投影 + 残差同理；gate 投影 + SiLU 一个内核，up 投影 + 乘上 SiLU 的结果一个内核。",
        "kv_write 不在 prefill 里；注意力是一个 FlashAttention 内核。视图（分头、合并头）不占内核：消费者按新的下标读。",
    ])


def demo2_prefill():
    cfg = LLAMA8B
    T = 2048
    print_header(f"演示 2：Llama-3.1-8B 的一层，prefill {T} 个 token（BF16，H100 roofline + 每个内核 4 µs）")
    g = trace(prefill_layer(cfg, T, 0), {"x": (T, cfg.d)}, param_shapes(cfg))
    rows, res = [], {}
    for name, mk in variants(g):
        c = graph_cost(mk({}))
        res[name] = c
        rows.append([name, c["kernels"], fmt_bytes(c["bytes"]), f"{c['flops'] / 1e9:.0f}", fmt_time(c["t"]),
                     fmt_time(c["t_total"])])
    base = res[variants(g)[0][0]]["t_total"]
    for r, (name, _) in zip(rows, variants(g)):
        r.append(f"{base / res[name]['t_total']:.2f}×")
    print_table(["分区", "内核数", "显存读写", "GFLOPs", "roofline 时间", "+ 固定开销", "相对 eager"], rows)
    ks = fuse(optimize(g, attention=False))
    top = sorted(ks, key=lambda k: -kernel_time(k))[:4]
    print("\n融合但没有 FlashAttention 时最慢的 4 个内核：")
    print_table(["类型", "运算", "读", "写", "时间"],
                [[k.kind, k.label, fmt_bytes(k.reads), fmt_bytes(k.writes), fmt_time(kernel_time(k))] for k in top])
    print_observations([
        "只分解不融合，比 eager 还慢：eager 里 softmax、RMSNorm 本来就是手写的融合内核，分解以后变成五六个内核、每个都读写一遍整张量。"
        "这正是 PyTorch 2 论文观察到的：Inductor 先把大算子分解成小的基本运算，再靠融合把它们组合回来——关掉融合和内联，编译后反而变慢（演示 4）。",
        "融合把读写从约 5.8 GB 降到约 2.1 GB，主要省在逐元素和 RMSNorm 的中间结果上；但注意力的分数矩阵（32 个头 × 2048 × 2048，268 MB）"
        "仍要写出、读回、再写出、再读回——softmax 是按行归约，我们的规则没法把它和前后两个矩阵乘法融成一个内核。",
        "模式匹配换成 FlashAttention 后，读写再降到约 1.0 GB，时间里剩下的基本就是矩阵乘法本身。编译器自己\"发明\"不出 FlashAttention："
        "它需要在线 softmax 这种算法上的改写（第 11 讲），所以 Inductor、XLA 都是把注意力匹配出来交给手写的内核。",
        "prefill 的矩阵乘法受限于计算（T = 2048 时强度很高），固定开销只占几个百分点；融合省下的主要是逐元素和归约那部分受限于带宽的时间。",
    ])


def demo3_decode():
    cfg = LLAMA8B
    S = 4096
    print_header(f"演示 3：同一层的 decode（1 个新 token，KV 长度 S = {S}，形状用符号 S 追踪）")
    g = trace(decode_layer(cfg, "P"), {"x": (1, cfg.d), "k_cache": (cfg.G, "S", cfg.dh), "v_cache": (cfg.G, "S", cfg.dh),
                                        "cos": (1, 1, cfg.dh), "sin": (1, 1, cfg.dh)}, param_shapes(cfg))
    rows = []
    for name, mk in variants(g):
        c = graph_cost(mk({"S": S}))
        over = c["kernels"] * H100.launch
        rows.append([name, c["kernels"], fmt_bytes(c["bytes"]), fmt_time(c["t"]), fmt_time(over),
                     f"{over / c['t_total']:.0%}", f"{1 / (32 * c['t_total']):.0f}"])
    print_table(["分区", "内核数", "显存读写", "roofline 时间", "固定开销", "固定开销占比", "32 层：每秒 token 数"], rows)
    print_observations([
        "decode 的读写几乎全是权重（一层约 436 MB）和 KV cache（S = 4096 时约 17 MB）：融合几乎不减少字节，roofline 时间都在 136 µs 左右。",
        "省下的是内核个数：从约 46 个降到 12 个，固定开销从比 roofline 时间还多降到约四分之一。批 1 的 decode 里\"每个内核只有几微秒的活\"，"
        "内核个数直接决定了延迟——第 11 讲的 megakernel 把这一点推到极致（一层 1 个内核），CUDA Graphs 把每个内核的 CPU 开销降下来。",
        "注意力换成 FlashAttention 在 decode 时省得不多：分数只有 32 × 4096 个数，本来就小；decode 注意力的问题是并行度（flash-decoding，第 11 讲 §6）。",
    ])


def demo4_pt2_ablation():
    print_header("演示 4：PyTorch 2 论文的消融实验（表 4：HuggingFace 模型、float16、A100，几何平均的加速比）")
    rows = [["全部优化", "1.91×", "1.45×"], ["去掉循环 / 布局重排", "1.91×", "1.28×"], ["去掉矩阵乘法模板", "1.85×", "1.41×"],
            ["去掉参数冻结", "1.85×", "1.45×"], ["去掉模式匹配", "1.83×", "1.45×"], ["去掉 CUDA Graphs", "1.81×", "1.37×"],
            ["去掉融合", "1.68×", "1.27×"], ["去掉内联", "1.58×", "1.31×"], ["去掉融合和内联", "0.80×", "0.59×"]]
    print_table(["TorchInductor 的设置", "推理", "训练"], rows)
    print_observations([
        "去掉融合和内联以后，编译后的代码比 eager 还慢（推理 0.80×、训练 0.59×）：Inductor 把大算子分解成基本运算，全靠融合把它们合回去。"
        "我们的演示 2 是同一个现象（只分解不融合比 eager 慢）。",
        "融合（调度时把内核合并，含\"横向\"合并两个互不依赖的内核）和内联（降低时把逐元素运算的计算直接复制进每个使用者）功能重叠，"
        "单独去掉一个只掉一部分。",
        "CUDA Graphs 推理时贡献约 0.1×：它省的是 CPU 端发射内核的开销，对每个内核很短的小模型 / 小批量最明显（05 脚本、第 11 讲 §8）。",
        "参数冻结（把权重当常量，允许常量折叠，02 脚本）和模式匹配（02 脚本的注意力匹配就是一例）贡献较小，但在特定模型上可以很关键。",
    ])


if __name__ == "__main__":
    demo1_small()
    demo2_prefill()
    demo3_decode()
    demo4_pt2_ablation()