"""
=============================================================================
CS349E · 第 12 讲共享模块：graph_lab.py —— 一个迷你的 ML 编译器前端
=============================================================================

【这个模块要解决什么问题】
    PyTorch / JAX 的程序是 Python。编译器要优化它，第一步是把它变成一张**计算图**（中间表示，IR），
    然后在图上做优化（融合、常量折叠、内存规划……），最后生成代码。这个模块把这条路完整地走一遍：

      Python 函数 ──追踪（运算符重载）──▶ 图 IR（SSA）──pass──▶ 更好的图 ──▶ 融合分区 + 代价 ──▶ numpy 代码 / 时间估计

    - 追踪（tracing）：用一个"假张量" Tensor 代替真数据调用你的函数；每次运算不真算，而是往图里记一个节点。
      Tensor 可以只带形状（只看形状的追踪，像 jax.eval_shape，用来分析 80 亿参数的模型），也可以同时带一个真值
      （像 torch.jit.trace 那样边算边记）。形状可以是符号（"S"），用来演示动态形状。
    - 图 IR：每个节点 = (运算, 输入节点, 属性, 形状)，按拓扑顺序排列，每个节点只定义一次（SSA，前置讲 P4 §3）。
    - pass：分解（decompose）、公共子表达式（CSE）、常量折叠、代数化简、死代码删除（DCE）、
      把 RMSNorm 的缩放折进下一个矩阵乘法的权重、把注意力的子图换成一个 FlashAttention 节点（模式匹配）。
    - 融合（fuse）：把图分成"内核"，每个内核读它的外部输入、写被外部用到的结果；按 roofline 估计每个内核的时间。
    - 内存规划：激活的生存期 + 贪心分配偏移，比较"不复用""贪心""下界"。
    - 动态形状与 graph break：缓存（guard）、按桶补齐、符号形状；数据依赖的分支在三种前端下的三种行为。

【运行】python3 graph_lab.py    # 自检：追踪 → 优化 → 执行的结果和第 2 讲的 numpy Transformer 一致
【依赖】numpy；自检导入 ../lecture02 的 transformer_lab（只读）。
"""

import math
import os
import sys
from dataclasses import dataclass, field

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))

# =============================================================================
# 1. 运算的分类
# =============================================================================
POINTWISE = {"add", "sub", "mul", "div", "maximum", "neg", "exp", "log", "tanh", "sqrt", "rsqrt", "recip",
             "rot_pairs", "gt", "lt"}
REDUCE = {"sum", "max"}
VIEW = {"reshape", "transpose"}
LEAF = {"input", "param", "const"}
GENERATOR = {"causal_mask"}                  # 不读数据、按下标现算的张量（像 iota 比较），在内核里"免费"生成
COMPOSITE = {"softmax", "rms_norm", "silu"}  # 高层运算：追踪时原样记下，decompose pass 把它们展开成基本运算
SPECIAL = {"matmul", "kv_write", "flash_attn"}


class TraceError(RuntimeError):
    """追踪时遇到无法记录的东西（例如依赖数据值的分支）。"""


# =============================================================================
# 2. 图 IR
# =============================================================================
@dataclass
class Node:
    """
    图里的一个节点（SSA 值）。

    【字段】id：编号（也是在 nodes 列表里的位置，删除节点后会重新编号）；op：运算名；inputs：输入节点的 id 元组；
        attrs：属性（如 axis、perm、shape、eps）；shape：输出形状（维度可以是 int 或符号 str）；name：可读的名字
    """
    id: int
    op: str
    inputs: tuple
    attrs: dict
    shape: tuple
    name: str = ""


@dataclass
class Graph:
    """
    计算图。

    【字段】nodes：按拓扑顺序的节点列表；values：常量 / 参数节点的值（id → numpy 数组或标量）；
        outputs：输出节点的 id 列表；input_names：输入名 → 节点 id；guards：追踪时记下的假设（见 §8）
    """
    nodes: list = field(default_factory=list)
    values: dict = field(default_factory=dict)
    outputs: list = field(default_factory=list)
    input_names: dict = field(default_factory=dict)
    guards: list = field(default_factory=list)

    def add(self, op, inputs, attrs, shape, name=""):
        n = Node(len(self.nodes), op, tuple(inputs), dict(attrs), tuple(shape), name)
        self.nodes.append(n)
        return n.id

    def users(self):
        """每个节点被哪些节点用到（id → 列表）。"""
        u = {n.id: [] for n in self.nodes}
        for n in self.nodes:
            for i in n.inputs:
                u[i].append(n.id)
        return u

    def count(self, kinds=None):
        """非叶子节点的个数（kinds 给定时只数这些运算）。"""
        return sum(1 for n in self.nodes if n.op not in LEAF and (kinds is None or n.op in kinds))

    def __str__(self):
        lines = []
        for n in self.nodes:
            ins = ", ".join(f"%{i}" for i in n.inputs)
            at = ", ".join(f"{k}={v}" for k, v in n.attrs.items() if k != "value")
            shp = "×".join(str(d) for d in n.shape) or "标量"
            nm = f"  # {n.name}" if n.name else ""
            lines.append(f"%{n.id:<3} = {n.op}({ins}{', ' if ins and at else ''}{at})  : [{shp}]{nm}")
        lines.append("返回 " + ", ".join(f"%{o}" for o in self.outputs))
        return "\n".join(lines)


def numel(shape, bind=None):
    """元素个数；符号维度按 bind（符号 → 整数）取值。"""
    out = 1
    for d in shape:
        out *= bind[d] if isinstance(d, str) else d
    return out


# =============================================================================
# 3. 形状推断（允许符号维度）
# =============================================================================
def broadcast(a, b):
    """numpy 的广播规则，维度可以是符号：相等、或其中一个是 1。"""
    a, b = tuple(a), tuple(b)
    n = max(len(a), len(b))
    a = (1,) * (n - len(a)) + a
    b = (1,) * (n - len(b)) + b
    out = []
    for x, y in zip(a, b):
        if x == y or y == 1:
            out.append(x)
        elif x == 1:
            out.append(y)
        else:
            raise ValueError(f"形状不能广播：{a} 和 {b}")
    return tuple(out)


def matmul_shape(a, b):
    """矩阵乘法：最后两维 [..., m, k] @ [..., k, n]，前面的维度广播。"""
    if a[-1] != b[-2]:
        raise ValueError(f"矩阵乘法的内维不一致：{a} @ {b}")
    batch = broadcast(a[:-2], b[:-2]) if (len(a) > 2 or len(b) > 2) else ()
    return batch + (a[-2], b[-1])


# =============================================================================
# 4. 追踪用的张量
# =============================================================================
class Tensor:
    """
    追踪用的"假张量"：记录 (图, 节点 id, 形状)，可选地带一个真值（边算边记）。

    【运算符重载】a + b、a * 2.0、a @ W、-a、a.sum(axis=-1, keepdims=True)、a.reshape(...)、a.transpose(...)
        都不真算，而是往图里加一个节点、返回新的 Tensor。如果带真值，同时用 numpy 算出新的真值。
    【分支】bool(t)（if t: ...）会按追踪器的 mode 处理：见 trace() 和 §8。
    【Python 小知识】__add__ / __radd__ 等"魔法方法"让 a + b 调用我们的函数；__radd__ 处理 2.0 + a 这种左边是数字的情况。
    """
    __array_priority__ = 1000       # 让 numpy 数组 @ Tensor 也调用我们的 __rmatmul__

    def __init__(self, tracer, nid, shape, value=None):
        self.tr, self.id, self.shape, self.value = tracer, nid, tuple(shape), value

    # ---- 二元运算 ----
    def _bin(self, other, op, rev=False):
        o = self.tr.lift(other)
        a, b = (o, self) if rev else (self, o)
        return self.tr.emit(op, [a, b], {}, broadcast(a.shape, b.shape))

    def __add__(self, o): return self._bin(o, "add")
    def __radd__(self, o): return self._bin(o, "add", True)
    def __sub__(self, o): return self._bin(o, "sub")
    def __rsub__(self, o): return self._bin(o, "sub", True)
    def __mul__(self, o): return self._bin(o, "mul")
    def __rmul__(self, o): return self._bin(o, "mul", True)
    def __truediv__(self, o): return self._bin(o, "div")
    def __rtruediv__(self, o): return self._bin(o, "div", True)
    def __neg__(self): return self.tr.emit("neg", [self], {}, self.shape)
    def __gt__(self, o): return self._bin(o, "gt")          # 比较的结果是 0 / 1 的张量
    def __lt__(self, o): return self._bin(o, "lt")

    def __matmul__(self, o):
        o = self.tr.lift(o)
        return self.tr.emit("matmul", [self, o], {}, matmul_shape(self.shape, o.shape))

    def __rmatmul__(self, o):
        o = self.tr.lift(o)
        return self.tr.emit("matmul", [o, self], {}, matmul_shape(o.shape, self.shape))

    # ---- 归约与视图 ----
    def sum(self, axis=-1, keepdims=False):
        return self.tr.emit("sum", [self], {"axis": axis, "keepdims": keepdims}, _reduced(self.shape, axis, keepdims))

    def max(self, axis=-1, keepdims=False):
        return self.tr.emit("max", [self], {"axis": axis, "keepdims": keepdims}, _reduced(self.shape, axis, keepdims))

    def reshape(self, *shape):
        return self.tr.emit("reshape", [self], {"shape": tuple(shape)}, tuple(shape))

    def transpose(self, *perm):
        return self.tr.emit("transpose", [self], {"perm": tuple(perm)}, tuple(self.shape[p] for p in perm))

    def __bool__(self):
        return self.tr.on_bool(self)

    def __repr__(self):
        return f"Tensor(%{self.id}, shape={self.shape})"


def _reduced(shape, axis, keepdims):
    ax = axis % len(shape)
    return tuple(1 if i == ax else d for i, d in enumerate(shape)) if keepdims else \
        tuple(d for i, d in enumerate(shape) if i != ax)


# =============================================================================
# 5. 运算的 numpy 语义（解释器、带真值的追踪、常量折叠共用）
# =============================================================================
def rot_pairs_np(x):
    """相邻两个数一对旋转 90°：(a, b) → (−b, a)。RoPE = x·cos + rot_pairs(x)·sin（第 2 讲的相邻配对写法）。"""
    out = np.empty_like(x)
    out[..., 0::2] = -x[..., 1::2]
    out[..., 1::2] = x[..., 0::2]
    return out


def causal_mask_np(shape, attrs):
    """[..., 行, 列] 的因果掩码：第 r 行的 query 位置 = pos0 + r % T，允许看列号 ≤ 位置的 key；不允许的地方是 −∞。"""
    rows, cols = shape[-2], shape[-1]
    T, pos0 = attrs["T"], attrs.get("pos0", 0)
    qpos = pos0 + (np.arange(rows) % T)
    m = np.where(np.arange(cols)[None, :] <= qpos[:, None], 0.0, -np.inf)
    return np.broadcast_to(m, shape).copy()


def softmax_np(x, axis):
    m = np.max(x, axis=axis, keepdims=True)
    e = np.exp(x - m)
    return e / np.sum(e, axis=axis, keepdims=True)


def eval_node(n, args, bind=None):
    """
    用 numpy 计算一个节点。args：输入节点的值列表。

    【边界】rsqrt = 1/√x；kv_write 是函数式的"把一行写进缓存"（返回新数组；真实系统里是原地写，见 §6 的 IO 计数）。
    """
    op, a = n.op, n.attrs
    if op == "add":
        return args[0] + args[1]
    if op == "sub":
        return args[0] - args[1]
    if op == "mul":
        return args[0] * args[1]
    if op == "div":
        return args[0] / args[1]
    if op == "maximum":
        return np.maximum(args[0], args[1])
    if op == "gt":
        return (args[0] > args[1]).astype(np.float64)
    if op == "lt":
        return (args[0] < args[1]).astype(np.float64)
    if op == "neg":
        return -args[0]
    if op == "exp":
        return np.exp(args[0])
    if op == "log":
        return np.log(args[0])
    if op == "tanh":
        return np.tanh(args[0])
    if op == "sqrt":
        return np.sqrt(args[0])
    if op == "rsqrt":
        return 1.0 / np.sqrt(args[0])
    if op == "recip":
        return 1.0 / args[0]
    if op == "rot_pairs":
        return rot_pairs_np(args[0])
    if op == "sum":
        return np.sum(args[0], axis=a["axis"], keepdims=a["keepdims"])
    if op == "max":
        return np.max(args[0], axis=a["axis"], keepdims=a["keepdims"])
    if op == "matmul":
        return args[0] @ args[1]
    if op == "reshape":
        shape = tuple(bind[d] if isinstance(d, str) else d for d in a["shape"])
        return np.reshape(args[0], shape)
    if op == "transpose":
        return np.transpose(args[0], a["perm"])
    if op == "causal_mask":
        shape = tuple(bind[d] if isinstance(d, str) else d for d in n.shape)
        return causal_mask_np(shape, a)
    if op == "softmax":
        return softmax_np(args[0], a["axis"])
    if op == "rms_norm":
        x, w = args
        return x / np.sqrt(np.mean(x * x, axis=-1, keepdims=True) + a["eps"]) * w
    if op == "silu":
        return args[0] / (1.0 + np.exp(-args[0]))
    if op == "kv_write":
        cache, new = args
        out = cache.copy()
        pos = a["pos"]
        pos = bind[pos] if isinstance(pos, str) else pos           # 符号的写入位置在运行时给出
        out[..., pos:pos + new.shape[-2], :] = new
        return out
    if op == "flash_attn":
        q, k, v = args[0], args[1], args[2]
        s = (q @ np.swapaxes(k, -1, -2)) * a["scale"]
        if len(args) > 3:
            s = s + args[3]
        return softmax_np(s, -1) @ v
    raise NotImplementedError(op)


# =============================================================================
# 6. 追踪
# =============================================================================
class Tracer:
    """
    追踪器：把一次函数调用记成一张图。

    【mode】数据依赖的分支（bool(张量)）怎样处理：
        "error"     ：抛出 TraceError（像 jax.jit：追踪时不知道数据的值）
        "specialize"：用示例输入的真值决定走哪条路，**不留任何检查**（像 torch.jit.trace：换一个输入可能悄悄算错）
        "guard"     ：用真值决定走哪条路，并把"这个条件是真 / 假"记成一个 guard，运行时检查（演示用；
                      TorchDynamo 对张量的值不这样做，而是在这里切断图——graph break，见 LazyRunner）
    """

    def __init__(self, mode="error"):
        self.g = Graph()
        self.mode = mode
        self.concrete = False          # 是否边算边记（有示例输入时为 True）；只看形状时不生成掩码这类大常量的值
        self._consts = {}

    def lift(self, x):
        """把 Python 数字 / numpy 数组变成常量节点（同一个标量只建一个节点）。"""
        if isinstance(x, Tensor):
            return x
        arr = np.asarray(x, dtype=np.float64)
        if arr.ndim == 0:
            key = float(arr)
            if key in self._consts:
                nid = self._consts[key]
                return Tensor(self, nid, (), self.g.values[nid])
            nid = self.g.add("const", [], {"value": key}, ())
            self.g.values[nid] = arr
            self._consts[key] = nid
            return Tensor(self, nid, (), arr)
        nid = self.g.add("const", [], {}, arr.shape)
        self.g.values[nid] = arr
        return Tensor(self, nid, arr.shape, arr)

    def emit(self, op, ins, attrs, shape, name=""):
        nid = self.g.add(op, [t.id for t in ins], attrs, shape, name)
        val = None
        if all(t.value is not None for t in ins) and op not in GENERATOR:
            val = eval_node(self.g.nodes[nid], [t.value for t in ins])
        elif op in GENERATOR and self.concrete and all(isinstance(d, int) for d in shape):
            val = causal_mask_np(shape, attrs)
        return Tensor(self, nid, shape, val)

    def on_bool(self, t):
        if self.mode == "error" or t.value is None:
            raise TraceError("追踪时遇到依赖数据值的分支（if 张量:）：图里没法表示\"看数据决定走哪条路\"")
        v = bool(t.value)
        if self.mode == "guard":
            self.g.guards.append(("value", t.id, v))
        return v


def trace(fn, inputs, params=None, mode="error", example=None):
    """
    追踪 fn(**inputs_as_tensors, p=params_as_tensors)，返回 Graph。

    【参数】
        inputs ：{名字: 形状}（形状里可以有符号维度，如 ("S", 64)）
        params ：{名字: numpy 数组 或 形状元组}——模型的权重，记成 "param" 节点（值在编译时已知，可以做常量折叠）；
                 给形状元组时只记形状（分析 80 亿参数的模型时不必真的分配权重）
        mode   ：数据依赖分支的处理方式（见 Tracer）
        example：{名字: numpy 数组}——给了就边算边记（像 torch.jit.trace），没给就只看形状
    【返回】Graph；graph.input_names 记录输入名 → 节点 id
    【边界】fn 的返回值可以是一个 Tensor 或 Tensor 的元组。
    """
    tr = Tracer(mode)
    tr.concrete = example is not None
    args = {}
    for name, shape in inputs.items():
        nid = tr.g.add("input", [], {}, shape, name)
        tr.g.input_names[name] = nid
        args[name] = Tensor(tr, nid, shape, None if example is None else np.asarray(example[name], dtype=np.float64))
    p = {}
    for name, val in (params or {}).items():
        if isinstance(val, tuple):
            nid = tr.g.add("param", [], {}, val, name)
            p[name] = Tensor(tr, nid, val, None)
            continue
        val = np.asarray(val, dtype=np.float64)
        nid = tr.g.add("param", [], {}, val.shape, name)
        tr.g.values[nid] = val
        p[name] = Tensor(tr, nid, val.shape, val)
    out = fn(**args, p=p) if params is not None else fn(**args)
    outs = out if isinstance(out, tuple) else (out,)
    tr.g.outputs = [o.id for o in outs]
    return tr.g


# 高层运算：对 Tensor 记一个复合节点，对 numpy 数组直接算（同一份模型代码既能"立即执行"也能被追踪）
class F:
    """函数式的 API（像 torch.nn.functional）：参数是 Tensor 时记节点，是 numpy 数组时立即计算。"""

    @staticmethod
    def _un(op, x, attrs=None):
        if isinstance(x, Tensor):
            return x.tr.emit(op, [x], attrs or {}, x.shape)
        return eval_node(Node(0, op, (), attrs or {}, np.shape(x)), [np.asarray(x)])

    exp = staticmethod(lambda x: F._un("exp", x))
    tanh = staticmethod(lambda x: F._un("tanh", x))
    rsqrt = staticmethod(lambda x: F._un("rsqrt", x))
    silu = staticmethod(lambda x: F._un("silu", x))
    rot_pairs = staticmethod(lambda x: F._un("rot_pairs", x))

    @staticmethod
    def softmax(x, axis=-1):
        return F._un("softmax", x, {"axis": axis})

    @staticmethod
    def rms_norm(x, w, eps):
        if isinstance(x, Tensor):
            w = x.tr.lift(w)
            return x.tr.emit("rms_norm", [x, w], {"eps": eps}, x.shape)
        return eval_node(Node(0, "rms_norm", (), {"eps": eps}, np.shape(x)), [np.asarray(x), np.asarray(w)])

    @staticmethod
    def causal_mask(like, T, pos0=0):
        """和 like 同形状的因果掩码（0 / −∞），在内核里按下标现算，不占显存。"""
        if isinstance(like, Tensor):
            return like.tr.emit("causal_mask", [], {"T": T, "pos0": pos0}, like.shape)
        return causal_mask_np(np.shape(like), {"T": T, "pos0": pos0})

    @staticmethod
    def kv_write(cache, new, pos):
        """把新 token 的 K（或 V）写进缓存的第 pos 行（pos 可以是符号，如 "P"，运行时由 bind 给出行号）。"""
        if isinstance(cache, Tensor):
            new = cache.tr.lift(new)
            return cache.tr.emit("kv_write", [cache, new], {"pos": pos}, cache.shape)
        return eval_node(Node(0, "kv_write", (), {"pos": pos}, np.shape(cache)), [np.asarray(cache), np.asarray(new)])


# =============================================================================
# 7. Llama 风格的 Transformer 层（用上面的 API 写；和第 2 讲 transformer_lab 的数学相同）
# =============================================================================
@dataclass(frozen=True)
class LayerCfg:
    """一层的结构参数。d：隐藏维度；H、G：query 头 / KV 头；dh：头维度；f：FFN 维度；eps：RMSNorm 的 ε。"""
    d: int
    H: int
    G: int
    dh: int
    f: int
    eps: float = 1e-5
    theta: float = 10000.0

    @property
    def R(self):
        return self.H // self.G


LLAMA8B = LayerCfg(4096, 32, 8, 128, 14336, 1e-5, 500000.0)


def rope_tables(T, dh, pos0=0, theta=10000.0):
    """RoPE 的 cos、sin 表，形状 [1, T, dh]：第 2i、2i+1 列都是第 i 对的角度（相邻配对，和第 2 讲一样）。"""
    inv = theta ** (-np.arange(0, dh, 2) / dh)
    ang = np.arange(pos0, pos0 + T)[:, None] * inv[None, :]
    return np.repeat(np.cos(ang), 2, axis=1)[None], np.repeat(np.sin(ang), 2, axis=1)[None]


def layer_params(cfg, seed=0):
    """一层的随机参数（字典），名字和第 2 讲 transformer_lab 的一层相同。"""
    rng = np.random.default_rng(seed)
    d, H, G, dh, f = cfg.d, cfg.H, cfg.G, cfg.dh, cfg.f
    mat = lambda i, o: rng.standard_normal((i, o)) / math.sqrt(i)  # noqa: E731
    return {"attn_norm": 1 + 0.1 * rng.standard_normal(d), "wq": mat(d, H * dh), "wk": mat(d, G * dh),
            "wv": mat(d, G * dh), "wo": mat(H * dh, d), "mlp_norm": 1 + 0.1 * rng.standard_normal(d),
            "w_gate": mat(d, f), "w_up": mat(d, f), "w_down": mat(f, d)}


def prefill_layer(cfg, T, pos0=0):
    """
    返回一个函数 f(x, p)：一层 Transformer 的 prefill（批 1，T 个 token，因果注意力）。x：[T, d] → [T, d]。

    【写法】RMSNorm → Q/K/V 投影 → 分头 → RoPE（x·cos + rot_pairs(x)·sin）→ GQA：把同一组的 R 个 query 头
        排成 [G, R·T, dh] 和该组的 K、V 做注意力 → 合并头 → O 投影 + 残差 → RMSNorm → SwiGLU MLP + 残差。
    """
    H, G, dh = cfg.H, cfg.G, cfg.dh
    cos, sin = rope_tables(T, dh, pos0, cfg.theta)

    def f(x, p):
        h = F.rms_norm(x, p["attn_norm"], cfg.eps)
        q = (h @ p["wq"]).reshape(T, H, dh).transpose(1, 0, 2)          # [H, T, dh]
        k = (h @ p["wk"]).reshape(T, G, dh).transpose(1, 0, 2)          # [G, T, dh]
        v = (h @ p["wv"]).reshape(T, G, dh).transpose(1, 0, 2)
        q = q * cos + F.rot_pairs(q) * sin
        k = k * cos + F.rot_pairs(k) * sin
        qg = q.reshape(G, cfg.R * T, dh)                                   # 同一组的 R 个头排在一起
        s = (qg @ k.transpose(0, 2, 1)) * (1.0 / math.sqrt(dh))
        s = s + F.causal_mask(s, T, pos0)
        a = F.softmax(s, axis=-1) @ v                                      # [G, R·T, dh]
        a = a.reshape(H, T, dh).transpose(1, 0, 2).reshape(T, H * dh)
        x = x + a @ p["wo"]
        h = F.rms_norm(x, p["mlp_norm"], cfg.eps)
        return x + (F.silu(h @ p["w_gate"]) * (h @ p["w_up"])) @ p["w_down"]
    return f


def decode_layer(cfg, pos, masked=False):
    """
    返回 f(x, k_cache, v_cache, cos, sin[, kmask], p)：一层的 decode（1 个新 token 写在缓存的第 pos 行，再对整个缓存做注意力）。

    【参数】pos：整数，或符号（如 "P"，运行时用 bind 给出行号）；RoPE 的 cos、sin 作为输入传进来（随位置变）。
        masked=True 时多一个输入 kmask [1, 1, 缓存长度]（0 / −∞）：缓存按"桶"补齐时，补出来的位置要屏蔽掉。
    【返回】(x_out [1, d], 新的 k_cache, 新的 v_cache)
    """
    H, G, dh = cfg.H, cfg.G, cfg.dh

    def f(x, k_cache, v_cache, cos, sin, p, kmask=None):
        h = F.rms_norm(x, p["attn_norm"], cfg.eps)
        q = (h @ p["wq"]).reshape(H, 1, dh)
        k = (h @ p["wk"]).reshape(G, 1, dh)
        v = (h @ p["wv"]).reshape(G, 1, dh)
        q = q * cos + F.rot_pairs(q) * sin
        k = k * cos + F.rot_pairs(k) * sin
        K = F.kv_write(k_cache, k, pos)
        V = F.kv_write(v_cache, v, pos)
        qg = q.reshape(G, cfg.R, dh)
        s = (qg @ K.transpose(0, 2, 1)) * (1.0 / math.sqrt(dh))
        if masked:
            s = s + kmask
        a = (F.softmax(s, axis=-1) @ V).reshape(1, H * dh)
        x = x + a @ p["wo"]
        h = F.rms_norm(x, p["mlp_norm"], cfg.eps)
        return x + (F.silu(h @ p["w_gate"]) * (h @ p["w_up"])) @ p["w_down"], K, V
    return f


# =============================================================================
# 8. 解释器与代码生成
# =============================================================================
def run(g, inputs, bind=None, profile=None):
    """
    逐个节点解释执行一张图。

    【参数】inputs：{输入名: numpy 数组}；bind：符号维度的取值；profile：给一个字典时记录每种运算调用了几次
    【返回】输出的列表（只有一个输出时也是列表）
    """
    env = {}
    for n in g.nodes:
        if n.op == "input":
            env[n.id] = np.asarray(inputs[n.name], dtype=np.float64)
        elif n.op in ("param", "const"):
            env[n.id] = g.values[n.id]
        else:
            env[n.id] = eval_node(n, [env[i] for i in n.inputs], bind)
            if profile is not None:
                profile[n.op] = profile.get(n.op, 0) + 1
    return [env[o] for o in g.outputs]


def check_guards(g, inputs, bind=None):
    """检查追踪时记下的 guard（"这个条件当时是真 / 假"）在新输入上是否还成立。返回 True / False。"""
    if not g.guards:
        return True
    env = {}
    for n in g.nodes:
        if n.op == "input":
            env[n.id] = np.asarray(inputs[n.name], dtype=np.float64)
        elif n.op in ("param", "const"):
            env[n.id] = g.values[n.id]
        else:
            env[n.id] = eval_node(n, [env[i] for i in n.inputs], bind)
    return all(bool(env[nid]) == v for kind, nid, v in g.guards if kind == "value")


_NP_EXPR = {"add": "{0} + {1}", "sub": "{0} - {1}", "mul": "{0} * {1}", "div": "{0} / {1}", "neg": "-{0}",
            "exp": "np.exp({0})", "log": "np.log({0})", "tanh": "np.tanh({0})", "sqrt": "np.sqrt({0})",
            "rsqrt": "1.0 / np.sqrt({0})", "recip": "1.0 / {0}", "maximum": "np.maximum({0}, {1})",
            "matmul": "{0} @ {1}"}


def codegen(g, fn_name="compiled"):
    """
    把图生成为一个"直线型"的 Python 函数（每个节点一行 numpy 调用），返回 (源代码, 函数)。

    【和解释器的区别】解释器每执行一个节点要查字典、分派到 eval_node；生成的代码直接调用 numpy——
        省掉的是"框架开销"（05 脚本演示 4 在 CPU 上实测）。
    【Python 小知识】exec(src, ns) 在字典 ns 里执行一段源代码，定义出的函数可以从 ns 里取出来。
    """
    lines = [f"def {fn_name}(inputs, V):"]
    for n in g.nodes:
        v = f"v{n.id}"
        a = [f"v{i}" for i in n.inputs]
        if n.op == "input":
            lines.append(f"    {v} = inputs[{n.name!r}]")
        elif n.op in ("param", "const"):
            lines.append(f"    {v} = V[{n.id}]")
        elif n.op in _NP_EXPR:
            lines.append(f"    {v} = " + _NP_EXPR[n.op].format(*a))
        elif n.op in ("sum", "max"):
            lines.append(f"    {v} = np.{n.op}({a[0]}, axis={n.attrs['axis']}, keepdims={n.attrs['keepdims']})")
        elif n.op == "reshape":
            lines.append(f"    {v} = {a[0]}.reshape({n.attrs['shape']})")
        elif n.op == "transpose":
            lines.append(f"    {v} = {a[0]}.transpose({n.attrs['perm']})")
        else:
            lines.append(f"    {v} = _eval(N[{n.id}], [{', '.join(a)}])")
    lines.append("    return [" + ", ".join(f"v{o}" for o in g.outputs) + "]")
    src = "\n".join(lines)
    ns = {"np": np, "_eval": eval_node, "N": g.nodes}
    exec(src, ns)
    fn = ns[fn_name]
    vals = dict(g.values)
    return src, (lambda inputs: fn({k: np.asarray(x, dtype=np.float64) for k, x in inputs.items()}, vals))


# =============================================================================
# 9. 图上的 pass
# =============================================================================
def rebuild(g, keep, remap=None):
    """
    按 keep（要保留的节点 id 列表，拓扑顺序）重新编号，生成新图。remap：旧 id → 替代它的旧 id（用于"用 B 代替 A"）。
    """
    remap = remap or {}

    def res(i):
        while i in remap:
            i = remap[i]
        return i
    new = Graph(input_names={}, guards=list(g.guards))
    idmap = {}
    for old in keep:
        n = g.nodes[old]
        nid = new.add(n.op, [idmap[res(i)] for i in n.inputs], n.attrs, n.shape, n.name)
        idmap[old] = nid
        if old in g.values:
            new.values[nid] = g.values[old]
        if n.op == "input":
            new.input_names[n.name] = nid
    new.outputs = [idmap[res(o)] for o in g.outputs]
    new.guards = [(k, idmap.get(res(nid), nid), v) for k, nid, v in g.guards]
    return new


def dce(g):
    """死代码删除：只保留从输出能够回溯到的节点（输入节点总是保留）。"""
    live = set(g.outputs) | {nid for _, nid, _ in g.guards}
    for n in reversed(g.nodes):
        if n.id in live:
            live.update(n.inputs)
    keep = [n.id for n in g.nodes if n.id in live or n.op == "input"]
    return rebuild(g, keep)


def cse(g):
    """
    公共子表达式消除（值编号）：运算、输入、属性都相同的两个节点只保留一个。

    【边界】常量按值比较（同一个标量）；参数和输入从不合并；有副作用的运算（我们的 IR 里没有）不能合并。
    """
    seen, remap, keep = {}, {}, []

    def res(i):
        while i in remap:
            i = remap[i]
        return i
    for n in g.nodes:
        if n.op in ("input", "param"):
            keep.append(n.id)
            continue
        if n.op == "const":
            key = ("const", n.attrs.get("value", None) if n.shape == () else id(g.values[n.id]))
        else:
            key = (n.op, tuple(res(i) for i in n.inputs), tuple(sorted((k, str(v)) for k, v in n.attrs.items())))
        if key in seen:
            remap[n.id] = seen[key]
        else:
            seen[key] = n.id
            keep.append(n.id)
    return rebuild(g, keep, remap)


def constant_fold(g):
    """
    常量折叠：输入全部是常量 / 参数的节点，在编译时算出来，换成一个常量。

    【例子】参数的转置、两个参数相乘、标量的运算……推理时权重是固定的，这类计算每次前向都做一遍是浪费。
    【返回】(新图, 折叠掉的节点数)
    """
    vals = dict(g.values)
    const_ids = {n.id for n in g.nodes if n.op in ("param", "const") and n.id in vals}
    folded = 0
    new_nodes = []
    for n in g.nodes:
        if n.op not in LEAF and n.op not in GENERATOR and n.inputs and all(i in const_ids for i in n.inputs):
            vals[n.id] = eval_node(n, [vals[i] for i in n.inputs])
            const_ids.add(n.id)
            new_nodes.append(Node(n.id, "const", (), {}, n.shape, n.name or f"folded {n.op}"))
            folded += 1
        else:
            new_nodes.append(n)
    g2 = Graph(nodes=new_nodes, values=vals, outputs=list(g.outputs), input_names=dict(g.input_names),
               guards=list(g.guards))
    return dce(g2), folded


def simplify(g):
    """
    代数化简（只做"结果逐位不变"的那些）：x·1 → x，x + 0 → x，x / 1 → x，−(−x) → x，
    transpose(transpose(x)) 的两个排列复合为恒等时 → x，reshape(reshape(x)) → reshape(x)。

    【没有做的】x / c → x·(1/c)、(a + b) + c → a + (b + c)：数学上相等，浮点结果不同（前置讲 P4 §4）。
    【返回】(新图, 化简的次数)
    """
    remap, count = {}, 0
    nodes = list(g.nodes)

    def res(i):
        while i in remap:
            i = remap[i]
        return i

    def is_const(i, v):
        n = nodes[res(i)]
        return n.op == "const" and n.shape == () and float(g.values[n.id]) == v
    for idx, n in enumerate(nodes):
        ins = [res(i) for i in n.inputs]
        if n.op == "mul" and (is_const(ins[1], 1.0) or is_const(ins[0], 1.0)):
            other = ins[0] if is_const(ins[1], 1.0) else ins[1]
            if tuple(nodes[other].shape) == n.shape:
                remap[n.id] = other
                count += 1
        elif n.op in ("add", "sub") and is_const(ins[1], 0.0) and tuple(nodes[ins[0]].shape) == n.shape:
            remap[n.id] = ins[0]
            count += 1
        elif n.op == "div" and is_const(ins[1], 1.0):
            remap[n.id] = ins[0]
            count += 1
        elif n.op == "neg" and nodes[ins[0]].op == "neg":
            remap[n.id] = res(nodes[ins[0]].inputs[0])
            count += 1
        elif n.op == "transpose" and nodes[ins[0]].op == "transpose":
            p1, p2 = nodes[ins[0]].attrs["perm"], n.attrs["perm"]
            if tuple(p1[p] for p in p2) == tuple(range(len(p2))):
                remap[n.id] = res(nodes[ins[0]].inputs[0])
                count += 1
        elif n.op == "reshape" and nodes[ins[0]].op == "reshape":
            nodes[idx] = Node(n.id, "reshape", (res(nodes[ins[0]].inputs[0]),), n.attrs, n.shape, n.name)
            count += 1
    g2 = Graph(nodes=nodes, values=dict(g.values), outputs=list(g.outputs), input_names=dict(g.input_names),
               guards=list(g.guards))
    keep = [n.id for n in nodes if n.id not in remap]
    return dce(rebuild(g2, keep, remap)), count


def decompose(g):
    """
    把复合运算展开成基本运算（像 PyTorch 的 decompositions / prims）：
        softmax(x)      → m = max(x)；e = exp(x − m)；e / sum(e)
        rms_norm(x, w)  → x · rsqrt(sum(x·x)/d + ε) · w
        silu(x)         → x / (1 + exp(−x))
    【为什么】后端只需要实现少数几十个基本运算；展开后，融合 pass 可以把它们和周围的运算重新组合（§5）。
    """
    tr = Tracer()
    tr.g.values = {}
    idmap = {}
    for n in g.nodes:
        ins = [Tensor(tr, idmap[i], g.nodes[i].shape) for i in n.inputs]
        if n.op in LEAF:
            nid = tr.g.add(n.op, [], n.attrs, n.shape, n.name)
            if n.id in g.values:
                tr.g.values[nid] = g.values[n.id]
                if n.op == "const" and n.shape == ():
                    tr._consts.setdefault(float(g.values[n.id]), nid)
            if n.op == "input":
                tr.g.input_names[n.name] = nid
            idmap[n.id] = nid
            continue
        if n.op == "softmax":
            x = ins[0]
            ax = n.attrs["axis"]
            e = F.exp(x - x.max(axis=ax, keepdims=True))
            out = e / e.sum(axis=ax, keepdims=True)
        elif n.op == "rms_norm":
            x, w = ins
            dd = x.shape[-1]
            out = x * F.rsqrt((x * x).sum(axis=-1, keepdims=True) * (1.0 / dd) + n.attrs["eps"]) * w
        elif n.op == "silu":
            x = ins[0]
            out = x / (1.0 + F.exp(-x))
        else:
            out = tr.emit(n.op, ins, n.attrs, n.shape, n.name)
        idmap[n.id] = out.id
    tr.g.outputs = [idmap[o] for o in g.outputs]
    tr.g.guards = [(k, idmap[nid], v) for k, nid, v in g.guards]
    return tr.g


def fold_norm_weight(g):
    """
    把 RMSNorm 的缩放向量 w 折进后面的矩阵乘法：(x̂ · w) @ W = x̂ @ (diag(w) · W)。

    【条件】mul(x̂, w) 的 w 是一维参数（沿最后一维广播），而且这个 mul 的**所有**使用者都是"@ 参数矩阵"。
        于是每个使用者换成 x̂ @ W'，W' = w[:, None] · W 在编译时算好（常量折叠）。
    【效果】每个 token 少做 d 次乘法、少一个逐元素运算；不改变数学结果（浮点结果有最后一位的差别）。
        量化流水线（SmoothQuant、QuaRot 等，第 3 讲）也用同样的"把逐通道的缩放折进相邻的权重"。
    【返回】(新图, 折叠了几处)
    """
    users = g.users()
    nodes = list(g.nodes)
    vals = dict(g.values)
    count = 0
    for n in g.nodes:
        if n.op != "mul":
            continue
        a, b = n.inputs
        wn = nodes[b]
        if wn.op != "param" or len(wn.shape) != 1:
            continue
        us = users[n.id]
        if not us or not all(nodes[u].op == "matmul" and nodes[u].inputs[0] == n.id and nodes[nodes[u].inputs[1]].op == "param"
                             for u in us):
            continue
        for u in us:
            W = nodes[nodes[u].inputs[1]]
            nid = len(nodes)
            nodes.append(Node(nid, "param", (), {}, W.shape, f"{W.name}·diag({wn.name})"))
            if wn.id in vals and W.id in vals:                      # 只看形状的图里没有值，只改结构
                vals[nid] = vals[wn.id][:, None] * vals[W.id]
            nodes[u] = Node(u, "matmul", (a, nid), {}, nodes[u].shape, nodes[u].name)
        count += 1
    # 新参数节点追加在末尾，需要按拓扑顺序重排：把参数节点移到最前面
    order = [n.id for n in nodes if n.op in LEAF] + [n.id for n in nodes if n.op not in LEAF]
    g2 = Graph(nodes=nodes, values=vals, outputs=list(g.outputs), input_names=dict(g.input_names), guards=list(g.guards))
    return dce(rebuild(g2, order)), count


def match_attention(g):
    """
    模式匹配：matmul(Q, transpose(K)) → mul(·, 常数) → [add(·, 因果掩码)] → softmax → matmul(·, V)
    这一串换成一个 flash_attn(Q, K, V[, mask]) 节点——就像 TorchInductor 把注意力的子图换成调用 FlashAttention。

    【为什么要在 decompose 之前做】softmax 被展开成 max、sub、exp、sum、div 以后，模式就难认了：pass 的顺序很重要。
    【返回】(新图, 替换了几处)
    """
    users = g.users()
    nodes = list(g.nodes)
    remap, count = {}, 0
    for n in g.nodes:
        if n.op != "softmax":
            continue
        x = nodes[n.inputs[0]]
        mask = None
        if x.op == "add" and nodes[x.inputs[1]].op == "causal_mask":
            mask, x = x.inputs[1], nodes[x.inputs[0]]
        if x.op != "mul" or nodes[x.inputs[1]].op != "const":
            continue
        scale = float(g.values[x.inputs[1]])
        mm = nodes[x.inputs[0]]
        if mm.op != "matmul" or nodes[mm.inputs[1]].op != "transpose":
            continue
        us = users[n.id]
        if len(us) != 1 or nodes[us[0]].op != "matmul" or nodes[us[0]].inputs[0] != n.id:
            continue
        pv = nodes[us[0]]
        q = mm.inputs[0]
        k = nodes[mm.inputs[1]].inputs[0]
        v = pv.inputs[1]
        ins = (q, k, v) + ((mask,) if mask is not None else ())
        nodes[pv.id] = Node(pv.id, "flash_attn", ins, {"scale": scale}, pv.shape, "flash_attn")
        count += 1
    g2 = Graph(nodes=nodes, values=dict(g.values), outputs=list(g.outputs), input_names=dict(g.input_names),
               guards=list(g.guards))
    return dce(rebuild(g2, [n.id for n in nodes], remap)), count


# =============================================================================
# 10. 融合与代价
# =============================================================================
def node_flops(n, shapes, bind):
    """一个节点的 FLOPs：逐元素 / 归约每个元素 1 次；矩阵乘法 2·批·m·n·k；flash_attn 4·批·行·列·d。"""
    if n.op in POINTWISE or n.op in REDUCE or n.op in COMPOSITE:
        src = shapes[n.inputs[0]] if n.op in REDUCE else n.shape
        return numel(src, bind)
    if n.op == "matmul":
        a = shapes[n.inputs[0]]
        return 2 * numel(n.shape, bind) * (bind[a[-1]] if isinstance(a[-1], str) else a[-1])
    if n.op == "flash_attn":
        q, k = shapes[n.inputs[0]], shapes[n.inputs[1]]
        rows = numel(q, bind) // (bind[q[-1]] if isinstance(q[-1], str) else q[-1])
        kl = k[-2]
        kl = bind[kl] if isinstance(kl, str) else kl
        return 4 * rows * kl * q[-1]
    return 0


@dataclass
class Kernel:
    """融合之后的一个内核：nodes（节点 id）、kind（matmul / pointwise / rowwise / attention / write）、reads、writes（字节）、flops。"""
    nodes: list
    kind: str
    reads: float = 0.0
    writes: float = 0.0
    flops: float = 0.0
    label: str = ""


def _base(g, i):
    """视图（reshape / transpose）看穿到底层的节点。"""
    while g.nodes[i].op in VIEW:
        i = g.nodes[i].inputs[0]
    return i


def fuse(g, enable=True, elem_bytes=2, bind=None):
    """
    把图分成内核。enable=False 时每个运算一个内核（"eager 模式"）。

    【规则（TorchInductor 风格的简化版）】
        - 视图（reshape / transpose）和掩码生成器不单独成内核：使用者按新的下标去读（零字节）。
        - 矩阵乘法、flash_attn 各自是一个内核；kv_write 是一个"原地写一行"的小内核。
        - epilogue 融合：对矩阵乘法输出（或它的视图）做的逐元素运算可以并进矩阵乘法的内核（偏置、残差、激活、RoPE），
          在累加器写回之前做；但 epilogue 里不能有归约。
        - 逐元素运算并进生产者所在的逐元素内核，只要"迭代空间"相同（同一个形状）；
        - 沿最后一维的归约（sum / max）并进生产者的逐元素内核（输入融合），之后用到"每行一个数"的逐元素运算
          也可以并进来——整行留在片上（Triton 的融合 softmax、RMSNorm 都是这样）。这样的内核叫 rowwise。
        - 合并前检查依赖：如果节点的另一个输入来自一个依赖于目标内核的内核，合并会形成环，不能合并。
        - 一个节点的结果如果被别的内核用到、或者是输出，就要写回显存。
    【返回】Kernel 的列表（按创建的顺序；每个内核的依赖都在它前面）
    【局限】不做"一个生产者复制进多个使用者"（内联）的成本判断；不考虑寄存器 / 共享内存放不放得下；
        不把 RMSNorm 并进下一个矩阵乘法的序幕（prologue）。
    """
    bind = bind or {}
    shapes = {n.id: n.shape for n in g.nodes}
    kof = {}
    kernels = []
    deps = []

    def ancestors(ki):
        seen, stack = set(), [ki]
        while stack:
            for d in deps[stack.pop()]:
                if d not in seen:
                    seen.add(d)
                    stack.append(d)
        return seen

    def compatible(k, n, iter_shape):
        if k.kind == "matmul":
            return n.op in POINTWISE and numel(n.shape, bind) == numel(k.iter_shape, bind) and not k.has_reduce
        if k.kind in ("pointwise", "rowwise"):
            if iter_shape is None:
                return False
            ks = k.iter_shape
            same = tuple(iter_shape) == tuple(ks)
            row_of = (len(iter_shape) == len(ks) and tuple(iter_shape[:-1]) == tuple(ks[:-1]) and iter_shape[-1] == 1)
            return same or (row_of and n.op in POINTWISE)
        return False

    for n in g.nodes:
        if n.op in LEAF or n.op in VIEW or n.op in GENERATOR:
            continue
        target = None
        if enable and (n.op in POINTWISE or n.op in REDUCE):
            iter_shape = shapes[n.inputs[0]] if n.op in REDUCE else n.shape
            if n.op in REDUCE and n.attrs["axis"] not in (-1, len(iter_shape) - 1):
                iter_shape = None
            in_k = list(dict.fromkeys(kof[_base(g, i)] for i in n.inputs if _base(g, i) in kof))
            for ki in reversed(in_k):                      # 先试最近创建的内核
                k = kernels[ki]
                if not compatible(k, n, iter_shape):
                    continue
                if any(ki in ancestors(j) for j in in_k if j != ki):
                    continue                               # 会形成环
                target = ki
                break
        if target is None:
            kind = {"matmul": "matmul", "flash_attn": "attention", "kv_write": "write"}.get(
                n.op, "rowwise" if n.op in REDUCE else "pointwise")
            k = Kernel([], kind)
            k.iter_shape = shapes[n.inputs[0]] if n.op in REDUCE else n.shape
            k.has_reduce = False
            kernels.append(k)
            deps.append(set())
            target = len(kernels) - 1
        k = kernels[target]
        k.nodes.append(n.id)
        kof[n.id] = target
        for i in n.inputs:
            b = _base(g, i)
            if b in kof and kof[b] != target:
                deps[target].add(kof[b])
        if n.op in REDUCE:
            k.has_reduce = True
            if k.kind == "pointwise":
                k.kind = "rowwise"
    # 内核的执行顺序：按依赖做拓扑排序（创建顺序不一定满足：一个早创建的内核可能后来吸收了依赖于晚创建内核的节点）
    order, done = [], set()

    def visit(ki):
        if ki in done:
            return
        done.add(ki)
        for d in sorted(deps[ki]):
            visit(d)
        order.append(ki)
    for ki in range(len(kernels)):
        visit(ki)
    kernels = [kernels[i] for i in order]
    pos = {old: new for new, old in enumerate(order)}
    kof = {nid: pos[ki] for nid, ki in kof.items()}
    # 读写字节与 FLOPs
    users = g.users()
    out_set = set(_base(g, o) for o in g.outputs)
    for ki, k in enumerate(kernels):
        members = set(k.nodes)
        ext = set()
        for nid in k.nodes:
            n = g.nodes[nid]
            k.flops += node_flops(n, shapes, bind)
            for i in n.inputs:
                b = _base(g, i)
                if g.nodes[b].op in GENERATOR or (g.nodes[b].op == "const" and g.nodes[b].shape == ()):
                    continue
                if b not in members:
                    ext.add(b)
        for b in ext:
            k.reads += numel(shapes[b], bind) * elem_bytes
        if k.kind == "write":
            n = g.nodes[k.nodes[0]]
            k.reads = numel(shapes[n.inputs[1]], bind) * elem_bytes          # 原地写：只读新的一行
            k.writes = numel(shapes[n.inputs[1]], bind) * elem_bytes
        else:
            for nid in k.nodes:
                stack, real = list(users[nid]), []
                while stack:                                                  # 使用者是视图时，再看视图的使用者
                    u = stack.pop()
                    if g.nodes[u].op in VIEW:
                        stack.extend(users[u])
                    elif kof.get(u) != ki:
                        real.append(u)
                if real or nid in out_set:
                    k.writes += numel(shapes[nid], bind) * elem_bytes
        k.label = "+".join(dict.fromkeys(g.nodes[i].op for i in k.nodes))
    return kernels


@dataclass(frozen=True)
class Device:
    """roofline 用的设备参数：峰值（矩阵乘法 / 其他）、显存带宽、每个内核的固定开销（启动 + 序幕，第 11 讲 §8）。"""
    name: str
    mm_flops: float
    vec_flops: float
    bw: float
    launch: float


H100 = Device("H100（BF16）", 989e12, 67e12, 3.35e12, 4e-6)


def kernel_time(k, dev=H100):
    """一个内核的 roofline 时间（不含固定开销）：max(运算 / 峰值, 读写 / 带宽)。矩阵乘法和注意力用 Tensor Core 的峰值。"""
    peak = dev.mm_flops if k.kind in ("matmul", "attention") else dev.vec_flops
    return max(k.flops / peak, (k.reads + k.writes) / dev.bw)


def graph_cost(kernels, dev=H100):
    """整张图：内核数、读写字节、FLOPs、roofline 时间之和、加上固定开销后的时间。"""
    t = sum(kernel_time(k, dev) for k in kernels)
    return {"kernels": len(kernels), "bytes": sum(k.reads + k.writes for k in kernels),
            "flops": sum(k.flops for k in kernels), "t": t, "t_total": t + len(kernels) * dev.launch}


# =============================================================================
# 11. 内存规划
# =============================================================================
def lifetimes(g, kernels, bind=None, elem_bytes=2):
    """
    每个写回显存的中间结果的 (大小, 产生它的内核序号, 最后一个使用它的内核序号)。参数和输入不算（它们一直在）。

    【返回】列表 [(节点 id, 字节, 开始, 结束)]
    """
    bind = bind or {}
    where = {}
    for ki, k in enumerate(kernels):
        for nid in k.nodes:
            where[nid] = ki
    users = g.users()
    out_set = set(_base(g, o) for o in g.outputs)
    res = []
    for ki, k in enumerate(kernels):
        if k.kind == "write":
            continue
        for nid in k.nodes:
            stack, last, ext = list(users[nid]), ki, False
            while stack:
                u = stack.pop()
                if g.nodes[u].op in VIEW:
                    stack.extend(users[u])
                elif where.get(u, ki) != ki:
                    ext = True
                    last = max(last, where[u])
            if nid in out_set:
                ext, last = True, len(kernels)
            if ext:
                res.append((nid, numel(g.nodes[nid].shape, bind) * elem_bytes, ki, last))
    return res


def plan_memory(lts, method="greedy"):
    """
    给每个中间结果分配一段显存（偏移），生存期重叠的不能重叠。

    【方法】
        "naive" ：每个结果一块独立的显存（不复用）——总量 = 所有大小之和
        "greedy"：按大小从大到小，放在"和已放好的、生存期重叠的块都不冲突"的最低偏移（TFLite 等用的 greedy-by-size）
    【返回】(峰值字节, {节点: 偏移})
    【下界】任何时刻同时活着的字节之和的最大值（lower_bound）——任何分配都不可能低于它。
    """
    if method == "naive":
        off, total = {}, 0
        for nid, size, s, e in lts:
            off[nid] = total
            total += size
        return total, off
    placed = []
    off = {}
    for nid, size, s, e in sorted(lts, key=lambda t: -t[1]):
        cands = sorted([0] + [o + sz for (o, sz, s2, e2) in placed if not (e < s2 or e2 < s)])
        for c in cands:
            if all(c + size <= o or o + sz <= c for (o, sz, s2, e2) in placed if not (e < s2 or e2 < s)):
                placed.append((c, size, s, e))
                off[nid] = c
                break
    return max((o + sz for o, sz, _, _ in placed), default=0), off


def lower_bound(lts):
    """任何时刻同时活着的中间结果字节之和的最大值。"""
    if not lts:
        return 0
    T = max(e for _, _, _, e in lts) + 1
    live = np.zeros(T + 1)
    for _, size, s, e in lts:
        live[s:e + 1] += size
    return float(live.max())


# =============================================================================
# 12. graph break：惰性张量（lazy tensor）
# =============================================================================
class _LazyTracer(Tracer):
    """惰性执行用的追踪器：跨越 graph break 一直存在；遇到上一张图里的张量时，把它的值作为新图的输入导入。"""

    def __init__(self, runner):
        super().__init__("error")
        self.runner = runner
        self.gen = 0

    def fix(self, t):
        return t if getattr(t, "gen", self.gen) == self.gen else self.runner._import(t)

    def lift(self, x):
        t = super().lift(x) if not isinstance(x, Tensor) else self.fix(x)
        t.gen = self.gen
        return t

    def emit(self, op, ins, attrs, shape, name=""):
        ins = [self.fix(t) for t in ins]
        out = super().emit(op, ins, attrs, shape, name)
        out.gen = self.gen
        return out

    def on_bool(self, t):
        return bool(self.runner.materialize(t))


class LazyRunner:
    """
    惰性执行：运算先记进"当前的图"，直到 Python 需要一个具体的值（if 张量:、materialize）时，
    才把当前的图编译、执行，得到值，再开一张新图继续记——每次这样就是一次 **graph break**。
    之前那张图里算出来的张量，在新图里作为输入导入（带着已经算好的值）。

    【和 TorchDynamo 的关系】Dynamo 分析字节码，在遇到数据依赖的分支时把函数切成"图 + Python + 图"，
        效果和这里一样：多张小图，中间回到 Python。PyTorch/XLA 的 Lazy Tensor 也是这个机制。
    【统计】graphs：每张执行过的图里有几个运算节点
    """

    def __init__(self):
        self.graphs = []
        self.store = {}
        self.tr = _LazyTracer(self)
        self._new()

    def _new(self):
        self.tr.g = Graph()
        self.tr._consts = {}
        self.tr.gen += 1
        self.pending = {}
        self.imported = {}

    def tensor(self, arr):
        """把一个具体的 numpy 数组变成当前图的输入。"""
        arr = np.asarray(arr, dtype=np.float64)
        name = f"in{len(self.tr.g.input_names)}"
        nid = self.tr.g.add("input", [], {}, arr.shape, name)
        self.tr.g.input_names[name] = nid
        self.pending[name] = arr
        t = Tensor(self.tr, nid, arr.shape, None)
        t.gen = self.tr.gen
        return t

    def _import(self, t):
        key = (t.gen, t.id)
        if key not in self.imported:
            self.imported[key] = self.tensor(self.store[key])
        return self.imported[key]

    def materialize(self, t):
        """执行当前的图（全部节点），记下每个节点的值；返回 t 的值；之后开一张新图。"""
        t = self.tr.fix(t)
        g = self.tr.g
        ids = [n.id for n in g.nodes if n.op not in LEAF]
        g.outputs = ids
        vals = run(g, self.pending)
        for nid, v in zip(ids, vals):
            self.store[(self.tr.gen, nid)] = v
        for name, nid in g.input_names.items():
            self.store[(self.tr.gen, nid)] = self.pending[name]
        self.graphs.append(len(ids))
        out = self.store[(self.tr.gen, t.id)]
        self._new()
        return out


# =============================================================================
# 自检
# =============================================================================
def _reference_prefill(cfg, p, x, pos0=0):
    """用第 2 讲 transformer_lab 的函数算一层 prefill（批 1），作为对照。"""
    sys.path.insert(0, os.path.join(HERE, "..", "lecture02"))
    import transformer_lab as tl
    T = x.shape[0]
    xb = x[None]
    h = tl.rms_norm(xb, p["attn_norm"], cfg.eps)
    q = (h @ p["wq"]).reshape(1, T, cfg.H, cfg.dh)
    k = (h @ p["wk"]).reshape(1, T, cfg.G, cfg.dh)
    v = (h @ p["wv"]).reshape(1, T, cfg.G, cfg.dh)
    cos, sin = tl.rope_angles(np.arange(pos0, pos0 + T), cfg.dh, cfg.theta)
    q, k = tl.apply_rope(q, cos, sin), tl.apply_rope(k, cos, sin)
    a = tl.attention_core(q, k, v, pos0)
    xb = xb + a @ p["wo"]
    h = tl.rms_norm(xb, p["mlp_norm"], cfg.eps)
    xb = xb + (tl.silu(h @ p["w_gate"]) * (h @ p["w_up"])) @ p["w_down"]
    return xb[0]


def _reference_kv(cfg, p, x):
    """prefill x 之后每层的 K、V 缓存（[G, T, dh]，RoPE 之后），用第 2 讲的函数算。"""
    sys.path.insert(0, os.path.join(HERE, "..", "lecture02"))
    import transformer_lab as tl
    T = x.shape[0]
    h = tl.rms_norm(x[None], p["attn_norm"], cfg.eps)
    k = (h @ p["wk"]).reshape(1, T, cfg.G, cfg.dh)
    v = (h @ p["wv"]).reshape(1, T, cfg.G, cfg.dh)
    cos, sin = tl.rope_angles(np.arange(T), cfg.dh, cfg.theta)
    k = tl.apply_rope(k, cos, sin)
    return k[0].transpose(1, 0, 2), v[0].transpose(1, 0, 2)


def _selftest():
    cfg = LayerCfg(64, 8, 2, 8, 128)
    p = layer_params(cfg)
    rng = np.random.default_rng(1)
    T = 12
    x = rng.standard_normal((T, cfg.d))
    ref = _reference_prefill(cfg, p, x)
    g = trace(prefill_layer(cfg, T), {"x": (T, cfg.d)}, p)
    assert np.abs(run(g, {"x": x})[0] - ref).max() < 1e-12
    g2, _ = match_attention(g)
    g3 = decompose(g2)
    g4 = cse(g3)
    g5, _ = constant_fold(g4)
    g6, _ = simplify(g5)
    g7, nf = fold_norm_weight(g6)
    assert nf == 2
    for gg in (g2, g3, g4, g5, g6, g7):
        assert np.abs(run(gg, {"x": x})[0] - ref).max() < 1e-12
    _, fn = codegen(g7)
    assert np.abs(fn({"x": x})[0] - ref).max() < 1e-12
    # 融合：不融合时每个运算一个内核
    k0 = fuse(decompose(g), enable=False)
    k1 = fuse(g7)
    assert len(k1) < len(k0)
    # decode：符号长度 S
    dg = 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)}, p)
    S = 9
    kc, vc = rng.standard_normal((cfg.G, S, cfg.dh)), rng.standard_normal((cfg.G, S, cfg.dh))
    cos, sin = rope_tables(1, cfg.dh, S - 1, cfg.theta)
    out = run(dg, {"x": x[:1], "k_cache": kc, "v_cache": vc, "cos": cos, "sin": sin}, bind={"S": S, "P": S - 1})
    assert out[1].shape == (cfg.G, S, cfg.dh)
    # decode 的最后一个 token = prefill 的最后一行（同样的 KV）
    kfull = (_reference_kv(cfg, p, x[:S]))
    ref_last = _reference_prefill(cfg, p, x[:S])[-1]
    out2 = run(dg, {"x": x[S - 1:S], "k_cache": kfull[0], "v_cache": kfull[1], "cos": cos, "sin": sin}, bind={"S": S, "P": S - 1})
    assert np.abs(out2[0][0] - ref_last).max() < 1e-12
    # 内存规划：贪心 ≥ 下界，≤ 不复用
    lts = lifetimes(g7, k1)
    lb, gr, nv = lower_bound(lts), plan_memory(lts)[0], plan_memory(lts, "naive")[0]
    assert lb <= gr <= nv
    print("graph_lab 自检通过。")


if __name__ == "__main__":
    print("== graph_lab 自检 ==")
    _selftest()