"""
=============================================================================
CS349E · 第 12 讲 · 05_dynamic_shapes.py —— 动态形状、graph break 与框架开销
=============================================================================

【这个脚本要回答的问题】
    LLM 推理里形状天天变：decode 时 KV 的长度每一步加 1，批大小随请求进出变化。编译器为一个形状编译的图，
    换一个形状就不能直接用。怎么办？另外，图被切断、或者干脆不编译时，Python 本身的开销有多大？
      演示 1：缓存与 guard：decode 512 步，按"精确形状 / 按桶补齐 / 符号形状"编译各要几次，补齐浪费多少。
      演示 2：符号形状：用符号 S 追踪一次，在不同的 S 上和精确形状的图逐位比较；guard 与 0/1 特化。
      演示 3：数据依赖的循环（迭代到收敛为止）：惰性执行会切成几张图。
      演示 4：框架开销（在你的 CPU 上实测）：立即执行 / 解释执行图 / 生成的直线代码；CUDA Graphs 的道理。

【关于数字】演示 1 的"编译时间"和演示 4 都是墙钟时间，随机器变化（观察按当次运行的数字生成）；其余每次运行都相同。
【运行】python3 05_dynamic_shapes.py      （约 5 秒）
【依赖】numpy；本目录的 graph_lab；../lecture01 的 console。
"""

import math
import os
import sys
import time

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 (F, LayerCfg, LazyRunner, codegen, constant_fold, cse, decode_layer, decompose, fold_norm_weight,  # noqa: E402
                       layer_params, rope_tables, run, simplify, trace)
from console import fmt_time, print_header, print_observations, print_table  # noqa: E402

CFG = LayerCfg(128, 8, 2, 16, 344)


def optimize(g):
    g = decompose(g)
    g = cse(g)
    g, _ = constant_fold(g)
    g, _ = simplify(g)
    g, _ = fold_norm_weight(g)
    return g


def compile_decode(p, S, masked=False):
    """追踪 + 优化 + 生成代码：一层 decode 的"编译"。S 是缓存长度（整数或符号 "S"）。返回 (可调用的函数, 图)。"""
    inputs = {"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)}
    if masked:
        inputs["kmask"] = (1, 1, S)
    g = optimize(trace(decode_layer(CFG, "P", masked=masked), inputs, p))
    return g


class CompileCache:
    """
    一个最简单的"编译缓存 + guard"：按策略把输入的缓存长度 S 映射成一个 key；key 没见过就编译一次。

    【策略】
        "exact"  ：key = S（每个新长度都重新编译——像只有静态形状的编译器）
        "pow2"   ：key = 不小于 S 的 2 的幂（按桶补齐，补出来的位置用 kmask 屏蔽）
        "mult128"：key = 128 的倍数
        "symbolic"：S ≥ 2 时 key = "S"（一张符号形状的图）；S 是 0 或 1 时单独特化（TorchDynamo 的 0/1 特化）
    【统计】compiles：编译次数；compile_time：编译花的墙钟时间；padded / real：补齐后的总长度和真实总长度
    """

    def __init__(self, p, policy):
        self.p, self.policy = p, policy
        self.cache = {}
        self.compiles = 0
        self.compile_time = 0.0
        self.padded = 0
        self.real = 0

    def key(self, S):
        if self.policy == "exact":
            return S
        if self.policy == "pow2":
            return 1 << (S - 1).bit_length()
        if self.policy == "mult128":
            return -(-S // 128) * 128
        return "S" if S >= 2 else S

    def __call__(self, x, k_cache, v_cache, cos, sin):
        S = k_cache.shape[1]
        k = self.key(S)
        if k not in self.cache:
            t0 = time.perf_counter()
            masked = self.policy in ("pow2", "mult128")
            self.cache[k] = compile_decode(self.p, k, masked)
            self.compile_time += time.perf_counter() - t0
            self.compiles += 1
        g = self.cache[k]
        L = S if isinstance(k, str) else k
        self.real += S
        self.padded += L
        ins = {"x": x, "k_cache": k_cache, "v_cache": v_cache, "cos": cos, "sin": sin}
        if L != S:
            pad = ((0, 0), (0, L - S), (0, 0))
            ins["k_cache"], ins["v_cache"] = np.pad(k_cache, pad), np.pad(v_cache, pad)
        if self.policy in ("pow2", "mult128"):
            ins["kmask"] = np.where(np.arange(L) < S, 0.0, -np.inf)[None, None, :]
        out = run(g, ins, bind={"S": L, "P": S - 1})
        return out[0], out[1][:, :S], out[2][:, :S]


def decode_loop(policy, p, prompt=100, steps=512, check_every=None):
    """从长度 prompt 的缓存开始 decode steps 步（缓存里是随机的 K、V），返回 (CompileCache, 最后的输出)。"""
    rng = np.random.default_rng(0)
    cc = CompileCache(p, policy)
    k = rng.standard_normal((CFG.G, prompt, CFG.dh))
    v = rng.standard_normal((CFG.G, prompt, CFG.dh))
    xs = rng.standard_normal((steps, 1, CFG.d))            # 每一步的输入事先生成（不把输出喂回去，见 README 的说明）
    outs = []
    for t in range(steps):
        S = prompt + t + 1
        k = np.concatenate([k, np.zeros((CFG.G, 1, CFG.dh))], axis=1)
        v = np.concatenate([v, np.zeros((CFG.G, 1, CFG.dh))], axis=1)
        cos, sin = rope_tables(1, CFG.dh, S - 1, CFG.theta)
        y, k, v = cc(xs[t], k, v, cos, sin)
        outs.append(y)
    return cc, outs


def demo1_policies(p):
    print_header("演示 1：decode 512 步（KV 长度 101 → 612），四种编译策略")
    rows, res = [], {}
    for policy, name in (("exact", "精确形状（每个长度编译一次）"), ("pow2", "按 2 的幂补齐"), ("mult128", "按 128 的倍数补齐"),
                         ("symbolic", "符号形状 S（+ 0/1 特化）")):
        cc, outs = decode_loop(policy, p)
        res[policy] = (cc, outs)
        rows.append([name, cc.compiles, fmt_time(cc.compile_time), f"{cc.padded / cc.real - 1:.1%}"])
    ref = res["exact"][1]
    for r, policy in zip(rows, ("exact", "pow2", "mult128", "symbolic")):
        r.append(f"{max(np.abs(a - b).max() for a, b in zip(res[policy][1], ref)):.1e}")
    print_table(["策略", "编译次数", "编译花的时间（本机，我们的小编译器）", "补齐多读的 KV", "和精确形状的最大差"], rows)
    per = res["exact"][0].compile_time / res["exact"][0].compiles
    print_observations([
        f"精确形状：每一步 KV 长度都不同，512 步编译 512 次。我们的小编译器每次约 {fmt_time(per)}；真实的编译器（TorchInductor 生成并编译 Triton 内核、"
        "还可能做自动调优）每次要几秒到几分钟——每一步都重新编译是完全不可行的。",
        f"按桶补齐：长度向上取整到 2 的幂只要编译 {res['pow2'][0].compiles} 次，代价是多读、多算补出来的位置"
        f"（这里多 {res['pow2'][0].padded / res['pow2'][0].real - 1:.0%}），还要多一个掩码把补出来的位置屏蔽掉；"
        f"128 的倍数编译 {res['mult128'][0].compiles} 次、只多 {res['mult128'][0].padded / res['mult128'][0].real - 1:.0%}。"
        "桶越细，编译越多、浪费越少——典型的折中。vLLM、TensorRT-LLM 为一组批大小 / 长度预先编译（或捕获 CUDA Graph），实际的请求向上补齐到最近的那一档。",
        "符号形状：用符号 S 追踪一次，所有 S ≥ 2 共用一张图。四种策略的结果相同（最大差在 1e-15 量级以内）：补齐 + 掩码在数学上不改变注意力，"
        "只是求和的长度不同、舍入顺序略有不同。",
    ])


def demo2_symbolic(p):
    print_header("演示 2：符号形状——一张图服务所有长度")
    gs = compile_decode(p, "S")
    shp = [n.shape for n in gs.nodes if n.op == "kv_write"][0]
    print(f"符号图里 kv_write 的形状：{shp}；图里没有任何地方写死了 S。")
    rng = np.random.default_rng(1)
    rows = []
    for S in (2, 7, 64, 500):
        ins = {"x": rng.standard_normal((1, CFG.d)), "k_cache": rng.standard_normal((CFG.G, S, CFG.dh)),
               "v_cache": rng.standard_normal((CFG.G, S, CFG.dh))}
        ins["cos"], ins["sin"] = rope_tables(1, CFG.dh, S - 1, CFG.theta)
        ge = compile_decode(p, S)
        a = run(gs, ins, bind={"S": S, "P": S - 1})[0]
        b = run(ge, ins, bind={"P": S - 1})[0]
        rows.append([S, f"{np.abs(a - b).max():.1e}"])
    print_table(["S", "符号形状的图 vs 为这个 S 专门编译的图"], rows)
    print_observations([
        "形状推断对符号维度照样成立：矩阵乘法 [2, 4, 16] @ [2, 16, S] → [2, 4, S]，softmax 沿 S 归约，再 @ [2, S, 16] → [2, 4, 16]——S 消掉了，"
        "后面的层和 S 无关。生成的代码在运行时从输入的形状读出 S。",
        "代价：编译器不知道 S 的具体值，有些优化做不了（按 S 选分块、把 S 整除的循环完全展开……），所以符号形状的内核可能比专门编译的慢一点。"
        "TorchDynamo 默认先按具体形状编译，第二次看到同一个维度变了，才把它标成符号（\"automatic dynamic shapes\"）。",
        "0/1 特化（PyTorch 2 论文 §5.1）：大小为 0 或 1 的维度总是按常数处理（1 决定了广播的语义），并记一个 guard；"
        "符号变量于是可以假定 ≠ 0、≠ 1，省掉很多 guard。我们的缓存对 S < 2 单独编译。",
    ])


def demo3_graph_breaks():
    print_header("演示 3：数据依赖的循环——迭代 y ← tanh(y @ W)，直到变化小于 1e-6")
    rng = np.random.default_rng(2)
    W = rng.standard_normal((32, 32)) * (0.5 / math.sqrt(32))
    y0 = rng.standard_normal((4, 32))
    lr = LazyRunner()
    y = lr.tensor(y0)
    Wt = lr.tensor(W)
    it = 0
    while True:
        y_new = F.tanh(y @ Wt)
        diff = (y_new - y) * (y_new - y)
        it += 1
        if diff.max(axis=-1).max(axis=-1) < 1e-12:        # 这里需要具体的值 → graph break
            break
        y = y_new
        if it > 200:
            break
    out = lr.materialize(y_new)
    ref = y0.copy()
    for _ in range(it):
        ref = np.tanh(ref @ W)
    nonempty = [n for n in lr.graphs if n > 0]
    print(f"迭代了 {it} 次；惰性执行切成了 {len(lr.graphs)} 张图：{len(nonempty)} 张各有 {sorted(set(nonempty))} 个运算，"
          f"最后取结果时的那张是空的（y_new 已经算好了）；结果和直接计算的差 {np.abs(out - ref).max():.1e}")
    print_observations([
        f"每次循环都要看 diff 的值决定是否继续，于是每一轮都切断一次图：{len(nonempty)} 张小图，每张只有几个运算，图之间回到 Python。"
        "融合只能在一张图内部做；每次 graph break 还要把结果写回、再作为下一张图的输入读进来。",
        "如果迭代次数事先知道，追踪会把循环展开成一张大图；如果用 lax.while_loop / torch.while_loop 这类结构化的控制流，"
        "整个循环（包括判断条件）就在一张图里，编译器可以把条件判断也放到 GPU 上，不用每轮回到 CPU。",
        "推理里典型的数据依赖：投机解码的\"接受到第几个 token\"、采样后检查是否生成了结束符、MoE 里每个专家分到几个 token——"
        "这些都是引擎和编译器之间的边界：要么在图外用 Python 处理，要么用专门的算子（第 6 讲的向量化接受、MoE 的分组 GEMM）。",
    ])


def demo4_overhead(p):
    print_header("演示 4：框架开销（本机实测）——一层 decode（d = 128，KV 长度 64）")
    rng = np.random.default_rng(3)
    S = 64
    ins = {"x": rng.standard_normal((1, CFG.d)), "k_cache": rng.standard_normal((CFG.G, S, CFG.dh)),
           "v_cache": rng.standard_normal((CFG.G, S, CFG.dh))}
    ins["cos"], ins["sin"] = rope_tables(1, CFG.dh, S - 1, CFG.theta)
    f_eager = decode_layer(CFG, S - 1)
    g = trace(decode_layer(CFG, S - 1), {k: v.shape for k, v in ins.items()}, p)
    g_opt = optimize(g)
    _, fn = codegen(g_opt)
    n_ops = g_opt.count()

    def timeit(f, reps=300):
        best = float("inf")
        for _ in range(5):
            t0 = time.perf_counter()
            for _ in range(reps):
                f()
            best = min(best, (time.perf_counter() - t0) / reps)
        return best
    t_eager = timeit(lambda: f_eager(**ins, p=p))
    t_interp = timeit(lambda: run(g_opt, ins))
    t_code = timeit(lambda: fn(ins))
    flops = 2 * (CFG.d * (CFG.H + 2 * CFG.G) * CFG.dh + CFG.H * CFG.dh * CFG.d + 3 * CFG.d * CFG.f)
    rows = [["立即执行（Python 逐个调用 numpy）", fmt_time(t_eager), ""], ["解释执行优化后的图", fmt_time(t_interp), f"{n_ops}"],
            ["生成的直线代码", fmt_time(t_code), f"{n_ops}"]]
    print_table(["方式", "每次调用", "图里的运算数"], rows)
    print(f"\n这一层的矩阵乘法约 {flops / 1e6:.2f} MFLOP；按 numpy 每秒约 50 GFLOP 估计只要约 {fmt_time(flops / 50e9)}。"
          f"生成的代码平均每个运算 {fmt_time(t_code / n_ops)}。")
    print_observations([
        f"小形状下，时间几乎全是\"开销\"：每个运算一次 Python 调用 + numpy 的分派 + 分配输出，约 {fmt_time(t_code / n_ops)}；"
        "真正的计算只占一小部分。GPU 上的对应物是\"每个内核几微秒的 CPU 发射开销\"（PyTorch eager 的算子分派、CUDA 的内核启动，第 11 讲 §8）。",
        f"解释执行比生成的代码慢（{t_interp / t_code:.1f} 倍）：每个节点多一次字典查找和分派——所以编译器最后要\"生成代码\"，而不是在运行时解释图。"
        f"立即执行是 {t_eager / t_code:.1f} 倍：它直接调用 numpy、没有图的分派，但跑的是没优化过的模型（没有权重折叠），还要执行 Python 的模型代码；"
        "生成的代码两头都占了便宜。",
        "CUDA Graphs 是同一个想法在 GPU 上的版本：第一次运行时把整串内核的启动\"录下来\"，以后一次 API 调用重放整张图，CPU 不再逐个发射内核。"
        "前提是形状和显存地址都不变——所以 vLLM 为一组固定的批大小各捕获一张图（实际的批向上补齐，演示 1 的桶），"
        "PyTorch 2 论文的消融实验里它在推理上贡献约 0.1×（03 脚本演示 4）。",
    ])


if __name__ == "__main__":
    params = layer_params(CFG)
    demo1_policies(params)
    demo2_symbolic(params)
    demo3_graph_breaks()
    demo4_overhead(params)