"""
=============================================================================
CS349E · 第 1 讲共享模块：roofline.py —— 屋顶线（roofline）模型的全部公式
=============================================================================

【这个模块要解决什么问题】
    一段计算在一块芯片上至少要多久？How to Scale Your Model 第 1 章的答案是：只看三件事。

        T_math  = 运算次数 FLOPs        / 芯片峰值算力 FLOP/s       （"算"要多久）
        T_comms = 要搬运的字节数 Bytes  / 带宽 byte/s               （"搬"要多久）
        容量    ：数据放不放得下（放不下就只能切分到多块芯片，或者换出到更慢的存储）

    如果"算"和"搬"能完美重叠（一边搬下一块数据，一边算这一块），总时间的**下界**是两者中的较大者；
    如果完全不重叠，总时间是两者之和——这是**上界**：

        T_lower = max(T_math, T_comms)          T_upper = T_math + T_comms  ≤ 2 · T_lower

    哪个更大，就说这段计算是 compute-bound（受限于算力）还是 memory/communication-bound
    （受限于带宽）。判断它不用算时间，只要比较两个"强度"：

        算术强度 I = FLOPs / Bytes                 （这段计算每搬 1 字节做几次运算）
        临界强度 I* = 峰值算力 / 带宽               （芯片的"屋脊点"）
        I > I*  ⇔  T_math > T_comms  ⇔  compute-bound

    "带宽"可以是芯片内部（HBM → 计算单元），也可以是芯片之间（NVLink、ICI）：公式完全一样，
    只是换一个带宽。本模块把这些公式写成函数，并提供一个字符画版的 roofline 图。

【运行】python3 roofline.py   # 自检：书中几个例题的数字
【依赖】只用 Python 标准库（math）。画 PNG 的函数会尝试导入 matplotlib，没装就跳过。
"""

import math


# ---------------------------------------------------------------------------
# 1. 最基本的三件事：时间上下界、可达算力、强度
# ---------------------------------------------------------------------------
def time_bounds(flops, nbytes, peak, bw):
    """
    给定一段计算的运算量和搬运量，算出 roofline 时间上下界。

    【参数】
        flops ：运算次数（FLOPs）
        nbytes：要通过"瓶颈通道"搬运的字节数
        peak  ：峰值算力 FLOP/s
        bw    ：这条通道的带宽 byte/s
    【返回】字典：
        t_math、t_comms：分别是"只算"和"只搬"要多久（秒）
        lower = max(两者)（完美重叠），upper = 两者之和（完全不重叠）
        bound = "compute" 或 "memory"：哪一项更大
        intensity = flops / nbytes
    【边界】nbytes = 0 时强度记为无穷大（纯计算）；flops = 0 时强度为 0（纯搬运）。
    """
    t_math = flops / peak
    t_comms = nbytes / bw
    intensity = math.inf if nbytes == 0 else flops / nbytes
    return {
        "t_math": t_math,
        "t_comms": t_comms,
        "lower": max(t_math, t_comms),
        "upper": t_math + t_comms,
        "bound": "compute" if t_math >= t_comms else "memory",
        "intensity": intensity,
    }


def attainable_flops(intensity, peak, bw):
    """
    roofline 图上的"屋顶"：强度为 intensity 的计算最多能跑到多少 FLOP/s。

        attainable = min(peak, intensity × bw)

    【直觉】带宽每秒送来 bw 个字节，每个字节最多支撑 intensity 次运算，所以运算速度不超过
        intensity × bw；同时也不可能超过芯片峰值。两条线的交点就是临界强度 peak / bw。
    【边界】intensity = inf 时返回 peak。
    """
    if math.isinf(intensity):
        return peak
    return min(peak, intensity * bw)


# ---------------------------------------------------------------------------
# 2. 矩阵乘法：X[B, D] @ W[D, F] -> Z[B, F]
# ---------------------------------------------------------------------------
def matmul_counts(B, D, F, x_bytes=2.0, w_bytes=2.0, out_bytes=2.0):
    """
    一次矩阵乘法 X[B, D] · W[D, F] → Z[B, F] 的运算量和（最少的）搬运量。

    【参数】
        B, D, F  ：三个维度。在 Transformer 里 B 通常是"这一步处理的 token 数"，
                   D、F 是权重矩阵的两个维度（几千到几万）
        x_bytes  ：X（激活）每个元素几个字节（bf16 = 2）
        w_bytes  ：W（权重）每个元素几个字节（bf16 = 2，int8/fp8 = 1，fp4 = 0.5）
        out_bytes：Z（输出）每个元素几个字节
    【返回】(flops, nbytes)
        flops  = 2·B·D·F：B·F 个输出，每个是 D 次乘法 + (D-1) 次加法 ≈ 2D 次运算
        nbytes = B·D·x + D·F·w + B·F·out：X、W 各读一遍，Z 写一遍（假设片上存储足够大、
                 每个数只搬一次——这是搬运量的**下界**，真实的分块实现会多搬一些，见 tiled_matmul_bytes）
    【边界】任何一维为 0 时 flops = 0。
    """
    flops = 2.0 * B * D * F
    nbytes = B * D * x_bytes + D * F * w_bytes + B * F * out_bytes
    return flops, nbytes


def matmul_intensity(B, D, F, x_bytes=2.0, w_bytes=2.0, out_bytes=2.0):
    """
    矩阵乘法的算术强度（FLOP/byte）。

    【直觉】bf16 且 B ≪ D、F 时：2BDF / (2BD + 2DF + 2BF) ≈ 2BDF / 2DF = B。
        也就是"强度 ≈ 一次处理的 token 数"——权重 W 搬进来一次，被 B 个 token 复用。
    """
    flops, nbytes = matmul_counts(B, D, F, x_bytes, w_bytes, out_bytes)
    return flops / nbytes


def critical_batch(D, F, critical_intensity, x_bytes=2.0, w_bytes=2.0, out_bytes=2.0):
    """
    精确求出：B 至少多大，矩阵乘法 X[B,D]·W[D,F] 才变成 compute-bound。

    【推导】要求 2BDF / (B·D·x + D·F·w + B·F·o) ≥ R（R = 临界强度），整理得
            B · (2DF − R·(D·x + F·o)) ≥ R · D·F·w
        若括号里 > 0：B ≥ R·D·F·w / (2DF − R·(D·x + F·o))，向上取整；
        若括号里 ≤ 0：B 再大强度也到不了 R（强度上限是 2DF/(D·x + F·o)），返回 None。
    【参数】critical_intensity：芯片的临界强度 R（FLOP/byte）
    【返回】最小的整数 B，或 None（永远到不了 compute-bound）
    【边界】结果至少是 1。
    """
    R = critical_intensity
    denom = 2.0 * D * F - R * (D * x_bytes + F * out_bytes)
    if denom <= 0:
        return None
    return max(1, math.ceil(R * D * F * w_bytes / denom))


def batched_matvec_counts(B, D, F, elem_bytes=2.0):
    """
    "每个样本有自己的矩阵"的批量矩阵-向量乘法：x[b] · W[b] 对 b = 0..B-1 各做一次。

    【为什么单独列出来】How to Scale Your Model 第 1 章练习 4：总运算量和普通矩阵乘法一样是 2BDF，
        但搬运量变成 B·D + **B·D·F** + B·F（每个样本的矩阵都要搬一遍），强度 ≈ 2/elem_bytes，
        **和 B 无关**。批处理救不了它。
    【系统联系】decode 阶段的注意力就是这种形状：每个序列都有**自己的** KV cache，
        批里的序列再多，也不能共享同一份 K、V。这是第 1 讲 §3.8 的关键点。
    """
    flops = 2.0 * B * D * F
    nbytes = elem_bytes * (B * D + B * D * F + B * F)
    return flops, nbytes


def tiled_matmul_bytes(m, k, n, bm, bk, bn, elem_bytes=2.0):
    """
    分块（tiling）矩阵乘法 (m, k)·(k, n) 从主存搬运的字节数（How to Scale Your Model 第 1 章的公式）。

    【思路】把输出切成 (bm × bn) 的小块，每个输出块沿 k 方向分 tk = k/bk 步，每步把 A 的一块
        (bm × bk) 和 B 的一块 (bk × bn) 搬进片上存储（SMEM/VMEM），算完再写回输出块。
        总块数 tm·tn，所以：
            bytes = elem · tm·tn · ( tk·(bm·bk + bk·bn) + bm·bn )
        忽略最后的写回项，强度 = 2·m·k·n / bytes ≈ (2/elem) · bm·bn/(bm+bn)。
        bf16 时正好是书里的 bm·bn/(bm+bn)。
    【参数】m, k, n：矩阵维度；bm, bk, bn：块大小（这里要求能整除，便于讲清楚）
    【返回】(flops, nbytes)
    【边界】不能整除时抛出 ValueError——真实内核会处理"边角块"，这里不展开。
    """
    if m % bm or k % bk or n % bn:
        raise ValueError("为了公式清楚，这里要求块大小能整除矩阵维度")
    tm, tk, tn = m // bm, k // bk, n // bn
    flops = 2.0 * m * k * n
    nbytes = elem_bytes * tm * tn * (tk * (bm * bk + bk * bn) + bm * bn)
    return flops, nbytes


# ---------------------------------------------------------------------------
# 3. 网络 roofline：把矩阵乘法切到两块芯片上
# ---------------------------------------------------------------------------
def sharded_matmul_two_chips(B, D, F, peak, link_bw, elem_bytes=2.0):
    """
    How to Scale Your Model 第 1 章的例子：X[B, D]·W[D, F] 沿 D 维切成两半放在两块芯片上。

    【做法】芯片 0 算 X[:, :D/2]·W[:D/2, :]，芯片 1 算另一半，各得到一份完整形状的"部分和" [B, F]，
        再把部分和发给对方相加。
    【代价】
        T_math  = 2BDF / 2 / peak = BDF / peak            （每块芯片只做一半）
        T_comms = elem·B·F / link_bw                       （把自己的部分和发出去）
        强度（相对网络）= BDF / (elem·B·F) = D / elem
    【关键】临界条件是 D/elem > peak/link_bw，只和 D 有关、和 B 无关——
        因为算的量和发的量都正比于 B，B 在比值里约掉了。
    【返回】time_bounds 的结果字典（t_math、t_comms、lower、upper、bound、intensity）
    """
    flops_per_chip = 2.0 * B * D * F / 2.0
    comm_bytes = elem_bytes * B * F
    return time_bounds(flops_per_chip, comm_bytes, peak, link_bw)


# ---------------------------------------------------------------------------
# 4. 字符画 roofline 图（不依赖任何画图库）
# ---------------------------------------------------------------------------
def ascii_roofline(points, peak, bw, xmin=0.1, xmax=1e4, width=64, height=16, title=""):
    """
    在终端里画一张 log-log 的 roofline 图。

    【参数】
        points：[(标签字符, 强度, 实测或理论 FLOP/s), ...]；第三项为 None 时画在屋顶上
        peak, bw：芯片峰值算力和带宽
        xmin, xmax：横轴（强度）范围，FLOP/byte
        width, height：图的字符宽高
    【返回】字符串列表（每个元素是一行），调用方 print("\\n".join(...)) 即可
    【做法】横轴 log10(强度)，纵轴 log10(FLOP/s)；纵轴范围取"xmin 处的屋顶"到"峰值的 1.5 倍"。
        屋顶用 '/'（斜坡部分）和 '-'（平顶部分）画，屋脊点用 '+'，数据点用各自的标签字符。
    【边界】超出范围的点会被夹到边上（画在图的边缘），不会报错。
    【Python 小知识】[[" "] * width for _ in range(height)] 生成一个 height 行、width 列的二维字符列表；
        不能写成 [[" "] * width] * height——那样每一行都是同一个列表对象，改一行等于改所有行。
    """
    ymin = attainable_flops(xmin, peak, bw) / 2.0
    ymax = peak * 1.5
    lx0, lx1 = math.log10(xmin), math.log10(xmax)
    ly0, ly1 = math.log10(ymin), math.log10(ymax)

    def to_col(x):
        c = int(round((math.log10(x) - lx0) / (lx1 - lx0) * (width - 1)))
        return min(width - 1, max(0, c))

    def to_row(y):
        r = int(round((ly1 - math.log10(y)) / (ly1 - ly0) * (height - 1)))
        return min(height - 1, max(0, r))

    grid = [[" "] * width for _ in range(height)]
    ridge = peak / bw
    for col in range(width):
        x = 10 ** (lx0 + (lx1 - lx0) * col / (width - 1))
        y = attainable_flops(x, peak, bw)
        grid[to_row(y)][col] = "/" if x < ridge else "-"
    if xmin <= ridge <= xmax:
        grid[to_row(peak)][to_col(ridge)] = "+"
    for label, x, y in points:
        if y is None:
            y = attainable_flops(x, peak, bw)
        grid[to_row(y)][to_col(x)] = label

    lines = []
    if title:
        lines.append(title)
    for r in range(height):
        y = 10 ** (ly1 - (ly1 - ly0) * r / (height - 1))
        tick = f"{y:8.1e} |" if r in (0, height // 2, height - 1) else "         |"
        lines.append(tick + "".join(grid[r]))
    lines.append("         +" + "-" * width)
    left = f"{xmin:g}"
    right = f"{xmax:g}"
    mid = f"强度 FLOP/byte（对数轴），屋脊点 ≈ {ridge:.0f}"
    lines.append("          " + left + " " * max(1, width - len(left) - len(right)) + right)
    lines.append("          " + mid)
    return lines


def save_roofline_png(path, curves, points, title="roofline"):
    """
    用 matplotlib 画一张 PNG 版 roofline 图（可选；没装 matplotlib 就返回 False）。

    【参数】
        path  ：输出文件路径
        curves：[(名字, peak, bw), ...]，每块芯片一条屋顶线
        points：[(名字, 强度, FLOP/s), ...]
    【返回】True 表示画好了，False 表示没装 matplotlib
    【Python 小知识】在函数里 import 是"按需导入"：只有真的要画图时才需要 matplotlib。
        matplotlib.use("Agg") 让它不打开窗口、直接画到文件里（在服务器或脚本里很常用）。
    """
    try:
        import matplotlib
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt
    except ImportError:
        return False
    plt.rcParams["font.sans-serif"] = ["PingFang SC", "Heiti SC", "Arial Unicode MS", "DejaVu Sans"]
    plt.rcParams["axes.unicode_minus"] = False
    xs = [10 ** (i / 20) for i in range(-20, 81)]
    fig, ax = plt.subplots(figsize=(7, 4.5))
    for name, peak, bw in curves:
        ax.loglog(xs, [attainable_flops(x, peak, bw) for x in xs], label=name)
    for i, (name, x, y) in enumerate(points):
        ax.scatter([x], [y], s=18)
        # 相邻两个点的标签上下错开，避免挤在一起
        ax.annotate(name, (x, y), fontsize=8, xytext=(3, 4 if i % 2 == 0 else -11), textcoords="offset points")
    ax.set_xlabel("arithmetic intensity (FLOP/byte)")
    ax.set_ylabel("attainable FLOP/s")
    ax.set_title(title)
    ax.grid(True, which="both", alpha=0.3)
    ax.legend(fontsize=8)
    fig.tight_layout()
    fig.savefig(path, dpi=130)
    plt.close(fig)
    return True


# ---------------------------------------------------------------------------
# 自检：How to Scale Your Model 第 1 章的几个数字
# ---------------------------------------------------------------------------
if __name__ == "__main__":
    # TPU v5e：1.97e14 FLOP/s，8.2e11 B/s → 临界强度 240
    peak, bw = 1.97e14, 8.2e11
    print(f"TPU v5e 临界强度 = {peak / bw:.0f} FLOP/byte（书中：240）")
    print(f"dot product 强度（N = 1e6）= {(2 * 1e6 - 1) / (4 * 1e6 + 2):.4f}（书中：→ 1/2）")
    print(f"bf16 matmul D=F=8192 的精确临界 batch = {critical_batch(8192, 8192, peak / bw)}（书中近似：B > 240）")
    print(f"int8 权重 + bf16 激活 的临界 batch = {critical_batch(8192, 8192, peak / bw, w_bytes=1.0)}（书中近似：B > 120）")
    r = sharded_matmul_two_chips(B=256, D=8192, F=8192, peak=peak, link_bw=4.5e10)
    print(f"两芯片切分 matmul：网络强度 = D/2 = {r['intensity']:.0f}，网络临界强度 = {peak / 4.5e10:.0f}"
          f"（书中：D > 8755 才 compute-bound）")
    print("\n".join(ascii_roofline([("a", 0.5, None), ("b", 64, None), ("c", 1024, None)], peak, bw,
                                   title="TPU v5e roofline：a=点积 b=B=64 的 matmul c=B=1024 的 matmul")))