"""
=============================================================================
CS349E · 前置讲 P2 · 04_occupancy_latency_waves.py —— 占用率、延迟隐藏与波次
=============================================================================

【这个脚本要回答的问题】
    一个内核启动后，块被分配到 SM 上执行。一个 SM 同时能放几个块？放得越多越好吗？块的个数和 SM 的个数
    对不上时会怎样？本脚本用 simt_lab 里的三个小公式回答：

      演示 1：占用率计算器——寄存器、共享内存、线程数、块数，哪一个先卡住？
      演示 2：延迟隐藏——Little 定律放到一个 SM 上：要多少个 warp（或者每个 warp 多发几个请求）才能喂饱带宽？
      演示 3：波次（wave）量化——块数比 SM 数多一点点时，最后一波几乎是空的；分块矩阵乘法里"块越大越好"的反例。

【怎么读输出】全部是公式（H100 的公开参数 + 标明的假设），每次运行完全相同。

【运行】python3 04_occupancy_latency_waves.py      （不到 1 秒）
【依赖】本目录的 simt_lab；../lecture01 的 console。
"""

import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture01"))

from console import print_header, print_observations, print_table, fmt_bytes  # noqa: E402
from simt_lab import H100, occupancy, waves, latency_hiding, cdiv  # noqa: E402

LIMITER_NAMES = {"warps": "warp 数", "blocks": "块数", "registers": "寄存器", "shared_mem": "共享内存", "无法启动": "无法启动"}


# ---------------------------------------------------------------------------
# 演示 1：占用率计算器
# ---------------------------------------------------------------------------
def demo1_occupancy():
    """
    占用率（occupancy）= 一个 SM 上同时驻留的 warp 数 / 最大 warp 数（H100：64）。

    【为什么会被卡住】一个块一旦被放到某个 SM 上，它的所有线程的寄存器、它的共享内存就一直占着，直到这个块结束。
        SM 上的寄存器（65,536 个）和共享内存（228 KB）是有限的，所以每个线程用的寄存器越多、每个块用的共享内存越多，
        能同时放的块就越少。
    """
    print_header("演示 1：占用率计算器（H100：每 SM 64 warp、32 块、65,536 寄存器、228 KB 共享内存）")
    configs = [
        (256, 32, 0, "简单逐元素内核"),
        (256, 64, 0, "每线程 64 个寄存器"),
        (256, 128, 0, "每线程 128 个寄存器"),
        (256, 255, 0, "每线程 255 个寄存器（上限）"),
        (128, 32, 48 * 1024, "每块 48 KB 共享内存"),
        (128, 64, 100 * 1024, "每块 100 KB 共享内存（大分块矩阵乘法）"),
        (1024, 32, 0, "每块 1024 线程"),
        (1024, 64, 0, "每块 1024 线程 + 64 寄存器"),
        (64, 32, 0, "每块 64 线程"),
        (32, 32, 0, "每块 32 线程"),
        (256, 32, 230 * 1024, "每块 230 KB 共享内存（超上限）"),
    ]
    rows = []
    for tpb, regs, smem, note in configs:
        o = occupancy(tpb, regs, smem)
        lim = o["limits"]
        rows.append([note, tpb, regs, fmt_bytes(smem, binary=True) if smem else "0",
                     lim.get("registers", "-"), lim.get("shared_mem", "-"), lim.get("warps", "-"),
                     o["blocks_per_sm"], o["active_warps"], f"{o['occupancy'] * 100:.1f}%", LIMITER_NAMES[o["limiter"]]])
    print_table(["配置", "线程/块", "寄存器/线程", "共享内存/块", "寄存器允许的块数", "共享内存允许的块数", "warp 上限允许的块数",
                 "每 SM 块数", "活跃 warp", "占用率", "瓶颈"], rows)
    print_observations([
        "寄存器是最常见的瓶颈：每线程 64 个寄存器就把占用率压到 50%，128 个压到 25%。复杂的内核（比如矩阵乘法、"
        "注意力）往往每线程要上百个寄存器来放累加器——第 1 讲 §3.9 说'输出块的累加器要一直待在片上'，就是待在这里。",
        "共享内存也会卡：每块 100 KB 时一个 SM 只能放 2 个块。大分块的矩阵乘法内核通常就是这种'低占用率、每个块干很多活'的风格。",
        "每块只有 32 或 64 个线程时，'每 SM 最多 32 个块'先卡住：32 线程/块最多只有 32 个 warp（50%）。块不要太小。",
        "共享内存超过每块 227 KB 的上限，内核根本启动不了。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：延迟隐藏
# ---------------------------------------------------------------------------
def demo2_latency_hiding():
    """
    一个 SM 分到的 HBM 带宽 = 3.35 TB/s ÷ 132 ≈ 25 GB/s；访存延迟假设 500 ns；
    按 Little 定律，这个 SM 需要约 12.7 KB 的读请求同时在途。

    【一个 warp 能贡献多少在途字节】= 它同时发出了几条 load（请求数）× 每条 load 搬多少字节
        （每线程 4 字节时一条 warp load 是 128 B，每线程 16 字节的向量读取是 512 B）。
    """
    print_header("演示 2：延迟隐藏——Little 定律放到一个 SM 上")
    _, need, _ = latency_hiding(1, 128)
    print(f"每个 SM 需要的在途字节 ≈ {fmt_bytes(H100.mem_bw / H100.n_sm)}/s × {H100.mem_latency_s * 1e9:.0f} ns ≈ {fmt_bytes(need)}")
    rows = []
    for warps in (8, 16, 32, 48, 64):
        cells = [warps, f"{warps / 64 * 100:.0f}%"]
        for bytes_per_req, inflight in ((128, 1), (128, 2), (512, 1), (512, 2)):
            frac, _, _ = latency_hiding(warps, bytes_per_req, inflight)
            cells.append(f"{frac * 100:.0f}%")
        rows.append(cells)
    print_table(["活跃 warp", "占用率", "4 B/线程，1 个请求", "4 B/线程，2 个请求", "16 B/线程，1 个请求", "16 B/线程，2 个请求"], rows)
    print_observations([
        "每个线程一次只读 4 字节、每个 warp 只有 1 个请求在途时，就算占满 64 个 warp（100% 占用率）也只能拿到约 65% 的带宽：在途的数据不够。",
        "两种补救：让每个 warp 同时发出多个读请求（指令级并行，循环展开就能做到），或者每个线程一次读 16 字节（向量化读取）。"
        "两者结合，只要 16 个 warp（25% 占用率）就能喂饱带宽。",
        "这就是 Volkov（GTC 2010）'Better Performance at Lower Occupancy' 的观点：占用率只是隐藏延迟的手段之一，"
        "每个线程多做几件独立的事（ILP）可以替代很多线程（TLP）。高性能的矩阵乘法和注意力内核通常就是低占用率、高 ILP 的。",
        "延迟 500 ns 是假设值；HBM 实际延迟和负载有关。但结论的方向不变：带宽越高、延迟越长，需要在途的数据就越多（P1 §4）。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：波次量化
# ---------------------------------------------------------------------------
def demo3_waves():
    """
    所有 SM 同时最多能跑 132 × (每 SM 块数) 个块，这叫一"波"（wave）。块更多时，要等前一波的块结束才能开始下一波。

    【问题】块数恰好比一波多一点时，最后一波只有几个块在跑，大部分 SM 闲着——这叫波次量化（wave quantization）损失。
    【矩阵乘法里的取舍】输出切成 tm × tn 的块（每个块由一个线程块负责）。块越大，每个块的算术强度越高
        （第 1 讲 §3.9：≈ tm·tn/(tm+tn)），但块的个数越少，越容易出现"最后一波很空"甚至"不到一波"。
    """
    print_header("演示 3：波次量化——块数和 SM 数对不上时会怎样（H100：132 个 SM）")
    rows = []
    for n_blocks in (66, 132, 133, 264, 265, 1000, 10000):
        w, last, util = waves(n_blocks, 1)
        rows.append([n_blocks, w, f"{last * 100:.0f}%", f"{util * 100:.1f}%"])
    print_table(["块数（每 SM 放 1 个）", "波数", "最后一波的 SM 利用率", "整体 SM 利用率"], rows)

    print("\n分块矩阵乘法 C[M, N]：每个线程块负责一个 tm × tn 的输出块（每 SM 放 1 个块）")
    rows = []
    for M, N in ((4096, 4096), (2048, 2048), (1024, 4096), (512, 1024)):
        for tm, tn in ((64, 64), (128, 128), (128, 256), (256, 256)):
            n_tiles = cdiv(M, tm) * cdiv(N, tn)
            w, _, util = waves(n_tiles, 1)
            intensity = tm * tn / (tm + tn)
            rows.append([f"{M}×{N}", f"{tm}×{tn}", n_tiles, w, f"{util * 100:.0f}%", f"{intensity:.0f}"])
    print_table(["输出 C", "块大小", "块数", "波数", "SM 利用率", "块的强度（bf16，FLOP/B）"], rows)
    print_observations([
        "133 个块要 2 波，第二波只有 1 个块在跑：整体利用率掉到约 50%。块数是 SM 数的整数倍或者远大于 SM 数时，这个损失才小。",
        "大矩阵（4096×4096）用大块（256×256）没问题：块的强度最高、利用率也高。",
        "小矩阵（512×1024）用 256×256 的块只有 8 个块，132 个 SM 里只用了 8 个（6%）——强度再高也没用。"
        "这时宁可用小块、多块，或者把 K 维也切开（split-K）让更多 SM 参与。",
        "LLM decode 的矩阵乘法正是'很瘦'的情形（M = 批大小，可能只有几到几十）：块数不够填满 GPU 是 decode 内核的常见问题，"
        "第 10–11 讲的 split-K、flash-decoding 就是为它设计的。",
    ])


if __name__ == "__main__":
    demo1_occupancy()
    demo2_latency_hiding()
    demo3_waves()