"""
=============================================================================
CS349E · 第 8 讲 · 04_async_pipelines.py —— 异步拷贝、多级流水线与 warp 专门化
=============================================================================

【这个脚本要回答的问题】
    03 说 Tensor Core 要靠大块数据喂。可数据从显存搬到共享内存要几百纳秒——等数据的时候 Tensor Core 闲着吗？
    Ampere 的异步拷贝（cp.async）、Hopper 的 TMA（Tensor Memory Accelerator）和 warp 专门化（warp specialization）
    都是为了让"搬"和"算"同时进行。
      演示 1：Little 定律——要喂饱 3.35 TB/s 的显存，任何时刻要有多少字节在路上；用普通的线程读 vs 用 TMA 读。
      演示 2：一个 SM 的"读 → 算"循环：同步、双缓冲、多级流水线，Tensor Core 忙的比例（含延迟抖动）。
      演示 3：warp 专门化的另一个用处：让矩阵乘法和 softmax 的指数运算重叠（FlashAttention-3 的 ping-pong 调度）。

【模型的假设】显存延迟 700 ns（满载时，假设值）；每个 SM 从 L2 读的带宽按每周期 64 字节（假设值）、从显存读按平分 3.35 TB/s；
    Tensor Core 按 H100 每个 SM 每周期 4096 次运算、1.83 GHz。这些是量级，不是测量值。

【运行】python3 04_async_pipelines.py      （< 1 秒）
【依赖】本目录的 hw_engines；../lecture01 的 console、hw_specs。
"""

import os
import sys

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

from hw_engines import H100, pipeline_sim  # noqa: E402
from console import fmt_bytes, fmt_time, print_header, print_observations, print_table  # noqa: E402
from hw_specs import H100 as H100_SPEC  # noqa: E402

LAT = 700e-9
CLK = H100.clock_ghz * 1e9


# -----------------------------------------------------------------------------
def demo1_littles_law():
    print_header("演示 1：Little 定律——要喂饱显存，任何时刻要有多少字节在路上")
    bw = H100.hbm_tbs * 1e12
    inflight = bw * LAT
    per_sm = inflight / H100.sms
    print(f"在路上的字节 = 带宽 × 延迟 = {H100.hbm_tbs} TB/s × {LAT * 1e9:.0f} ns = {fmt_bytes(inflight)}；"
          f"平均每个 SM {fmt_bytes(per_sm)}")
    rows = []
    for label, per_req, issuers in (("每个线程读 4 字节（ld.global.f32）", 4, "线程"),
                                    ("每个线程读 16 字节（向量化 ld.global.v4）", 16, "线程"),
                                    ("cp.async 每个线程 16 字节（Ampere，不经过寄存器）", 16, "线程"),
                                    ("TMA：一条指令搬一整块 32 KiB（Hopper）", 32 * 1024, "一个线程发一条")):
        n = per_sm / per_req
        rows.append([label, f"{n:,.0f}", f"{n / 32:,.1f}" if issuers == "线程" else "-", issuers])
    print_table(["读的方式", "每个 SM 要同时挂着几个请求", "相当于几个 warp 各发一条", "谁来发"], rows)
    print_observations([
        f"要让显存一直忙，每个 SM 平均要有约 {fmt_bytes(per_sm)} 的读请求同时在路上（前置讲 P1 §4、P2 §7 的 Little 定律）。",
        "用普通的线程读，每个线程一次只读几个字节，要成百上千个请求同时挂着：每个线程读 4 字节时相当于 139 个 warp 各挂一条——"
        "超过了一个 SM 最多驻留 64 个 warp 的上限，只能让每个线程发好几条独立的读（ILP）。读回来的数据还要先进寄存器、再写进共享内存。",
        "cp.async（Ampere）让数据直接从显存进共享内存、不经过寄存器；TMA（Hopper）更进一步：一个线程发一条指令，"
        "硬件自己按张量的形状（多维、带步长）搬一整块，还能多播给线程块簇里的几个 SM。读数据不再需要很多线程。",
    ])


# -----------------------------------------------------------------------------
def demo2_pipeline():
    print_header("演示 2：一个 SM 的\"读 → 算\"循环（每块 A 128×64 + B 64×128，BF16 共 32 KiB；算 128×128×64）")
    tile_bytes = (128 + 128) * 64 * 2
    t_comp = 2 * 128 * 128 * 64 / (H100.tc_flops_per_sm_clk() * CLK)
    t_l2 = tile_bytes / (64 * CLK)
    t_hbm = tile_bytes / (H100.hbm_tbs * 1e12 / H100.sms)
    print(f"算一块 {fmt_time(t_comp)}；从 L2 传一块 {fmt_time(t_l2)}（每周期 64 字节），只从显存传（平分带宽）{fmt_time(t_hbm)}；"
          f"延迟 {LAT * 1e9:.0f} ns")
    rows = []
    for src, t_x in (("数据在 L2 里", t_l2), ("每块都要从显存读", t_hbm)):
        for jitter in (0.0, 1.0):
            cells = [src, "固定" if jitter == 0 else "指数抖动（平均 +100%）"]
            for S in (1, 2, 3, 4, 6, 8):
                r = pipeline_sim(256, t_comp, t_x, LAT, S, jitter=jitter, seed=1, issue_ahead=S > 1)
                cells.append(f"{r['busy']:.0%}")
            rows.append(cells)
    print_table(["数据从哪来", "延迟", "同步（1 级）", "双缓冲（2 级）", "3 级", "4 级", "6 级", "8 级"], rows)
    need = (LAT + t_l2 + t_comp) / t_comp
    print(f"\n（表里是 Tensor Core 忙的比例。按公式，固定延迟时要 (延迟 + 传输 + 计算) / 计算 ≈ {need:.1f} 级才能完全藏住延迟。）")
    print_observations([
        "同步地读一块、算一块，Tensor Core 大部分时间在等数据（数据在 L2 里时只忙 22%）。双缓冲（cp.async 的典型用法）只藏住一块的计算时间，"
        "延迟比计算长得多时还不够；要 4–5 级才能完全藏住 700 ns 的延迟。",
        "延迟有抖动时（满载的显存、L2 冲突），需要更多级来\"吸收\"偶尔的慢请求——这就是 Hopper 内核常用 4–8 级流水线的原因，"
        "而级数受共享内存大小的限制（03 演示 3）。",
        "如果每块都要从显存读（没有 L2 复用），传输本身比计算慢约 4.6 倍，流水线再深也只能到约 21%：流水线藏得住延迟，藏不住带宽。"
        "要靠更大的块、L2 和多播提高复用（03 演示 3），这已经是内核设计的问题（第 10 讲）。",
        "warp 专门化：一个 producer warp 只负责发 TMA 读请求、管理缓冲区，几个 consumer warpgroup 只负责算——"
        "两类工作不在同一条指令流里互相等待，流水线才能稳定地保持很多级。",
    ])


# -----------------------------------------------------------------------------
def demo3_pingpong():
    print_header("演示 3：让矩阵乘法和指数运算重叠（注意力，d = 128，整块 H100）")
    exp_rate = H100_SPEC.exp_per_s
    rows = []
    for prec, peak in (("BF16", H100.tc_peak()), ("FP8", 2 * H100.tc_peak())):
        t_mma = 4 * 128 / peak              # 每个注意力分数：QKᵀ 2d + PV 2d 次运算
        t_exp = 1 / exp_rate                # 每个分数一次 exp
        serial = t_mma + t_exp
        overlap = max(t_mma, t_exp)
        rows.append([prec, f"{t_mma * 1e12:.3f} ps", f"{t_exp * 1e12:.3f} ps", f"{t_exp / t_mma:.0%}",
                     f"{serial * 1e12:.3f} ps", f"{overlap * 1e12:.3f} ps", f"{serial / overlap:.2f}×"])
    print(f"（每秒 {exp_rate / 1e12:.1f} T 次 exp：第 1 讲 hw_specs 里 H100 的数字，取自 Blackwell Ultra 博客的表。）")
    print_table(["精度", "每个分数的矩阵运算时间", "每个分数的 exp 时间", "exp / 矩阵", "串行（一个 warpgroup）",
                 "重叠（两个 warpgroup 交替）", "重叠的收益"], rows)
    print_observations([
        "注意力里每个分数要做 512 次矩阵运算（Tensor Core）和 1 次 exp（特殊函数单元 SFU）。两者在不同的硬件单元上，"
        "可以同时进行；但如果同一个 warpgroup 先算矩阵、再算 softmax，它们就串行了。",
        "FlashAttention-3 的 ping-pong 调度：两个 consumer warpgroup 交替——一个在做矩阵乘法时，另一个在做 softmax。"
        "BF16 时 exp 占矩阵运算的四成多，重叠能快约 1.4 倍；FP8 时 Tensor Core 快一倍，exp 几乎和矩阵运算一样长，重叠更关键（约 1.9 倍）。",
        "FlashAttention-3 报告 BF16 前向约 740 TFLOPS（75% 利用率）、FP8 接近 1.2 PFLOPS。Blackwell Ultra 把 SFU 的 exp 吞吐翻倍"
        "（第 1 讲 §5.5），也是同一个问题的硬件解法——第 11 讲会回到 FlashAttention。",
    ])


if __name__ == "__main__":
    demo1_littles_law()
    demo2_pipeline()
    demo3_pingpong()