"""
=============================================================================
CS349E · 第 8 讲 · 03_tensor_cores.py —— GPU 的 Tensor Core：峰值从哪来，为什么省能量，怎样喂饱它
=============================================================================

【这个脚本要回答的问题】
    H100 的 BF16 矩阵乘法峰值 989 TFLOPS，普通的 FP32 运算只有约 60 TFLOPS——差 16 倍。
      演示 1：峰值从哪来：SM 数 × 每个 SM 的 Tensor Core 数 × 每个 Tensor Core 每周期的乘加数 × 时钟（V100 → A100 → H100）。
      演示 2：为什么矩阵指令省能量：一条指令做多少乘加、每次乘加要读几个寄存器（FFMA vs HMMA vs WGMMA）。
      演示 3：怎样喂饱它：一个 SM 每周期能做 4096 次运算，共享内存每周期只能给 128 字节——数据块要多大；
              共享内存装得下几级流水线；为什么只靠显存喂不饱。

【运行】python3 03_tensor_cores.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 GPU_GENS, H100, INSTR_SHAPES, tile_intensity  # noqa: E402
from console import print_header, print_observations, print_table  # noqa: E402
from hw_specs import energy_pj  # noqa: E402

OFFICIAL = {"V100": (125, 15.7), "A100": (312, 19.5), "H100 SXM5": (989, 67)}   # 官方：FP16 Tensor Core、FP32（TFLOPS）


# -----------------------------------------------------------------------------
def demo1_peak():
    print_header("演示 1：峰值从哪来——SM 数 × Tensor Core 数 × 每周期乘加数 × 时钟（FP16 稠密）")
    rows = []
    for g in GPU_GENS:
        off_tc, off_fp32 = OFFICIAL[g.name]
        rows.append([g.name, g.sms, g.tc_per_sm, g.tc_fma_per_clk, g.tc_flops_per_sm_clk(), f"{g.clock_ghz:.2f} GHz",
                     f"{g.tc_peak() / 1e12:.0f}", off_tc, f"{g.fp32_peak() / 1e12:.1f}", off_fp32,
                     f"{g.tc_peak() / g.fp32_peak():.0f}×"])
    print_table(["GPU", "SM 数", "TC/SM", "每个 TC 每周期 FMA", "每个 SM 每周期 FLOPs", "时钟", "算出的 TC 峰值",
                 "官方", "算出的 FP32 峰值", "官方", "TC / FP32"], rows)
    print("（TC = Tensor Core。H100 的官方 FP32 峰值 67 TFLOPS 按更高的加速时钟算；我们统一用 1.83 GHz。）")
    print_observations([
        "三代 GPU 的峰值都能从\"SM 数 × 每个 SM 每周期 FLOPs × 时钟\"精确算出来。每个 SM 每周期的矩阵运算能力"
        "从 V100 的 1024 翻到 A100 的 2048、H100 的 4096（Hopper 博客：\"每个 SM 的 MMA 速率是 A100 的 2 倍\"）。",
        "同一个 SM 里，Tensor Core 的运算能力是 CUDA core（FP32）的 8 倍（V100）到 16 倍（A100、H100）：芯片面积越来越多地给了矩阵运算。"
        "所以\"能写成矩阵乘法\"的计算和\"不能\"的计算（softmax、归一化、逐元素运算）速度差一个数量级以上（第 1 讲 §5.5）。",
    ])


# -----------------------------------------------------------------------------
def demo2_instructions():
    print_header("演示 2：一条指令做多少乘加、每次乘加读几个操作数")
    e_math = energy_pj("fp16_mult") + energy_pj("fp32_add")
    e_reg32 = energy_pj("instr_regfile") / 3              # Horowitz：一条指令读寄存器堆约 6 pJ（约 3 个 32 位操作数）
    e_instr = 50.0                                          # 每条 warp 指令的取指 / 译码 / 调度（假设值）
    rows = []
    for name, d in INSTR_SHAPES.items():
        fma = d["fma"]
        op_bytes = 4 if "FFMA" in name else 2               # FFMA 读 FP32，矩阵指令读 FP16
        operand_e = d["operands"] / fma * e_reg32 * op_bytes / 4
        acc_e = 2 * d["acc"] / fma * e_reg32                # 累加器（FP32）读一次、写一次
        instr_e = e_instr / fma
        total = e_math + operand_e + acc_e + instr_e
        rows.append([name, f"{fma:,}", f"{d['operands'] / fma:.3f}", f"{2 * d['acc'] / fma:.3f}",
                     f"{instr_e:.4f}", f"{total:.2f}", f"{e_math / total:.0%}"])
    print(f"能量按 45 nm 的量级估计：一次 FP16 乘 + FP32 加 ≈ {e_math:.1f} pJ；读一个 32 位寄存器 ≈ {e_reg32:.1f} pJ"
          f"（Horowitz：一条指令访问寄存器堆约 6 pJ）；每条 warp 指令的取指 / 译码 / 调度 ≈ {e_instr:.0f} pJ（假设值）")
    print_table(["指令", "一条指令的 FMA 数", "每次 FMA 读输入操作数", "每次 FMA 读写累加器", "每次 FMA 的指令开销 pJ",
                 "每次 FMA 的总能量 pJ", "其中运算本身"], rows)
    print_observations([
        "FFMA：一条 warp 指令只做 32 次乘加，每次乘加要从寄存器读 3 个 32 位数——取操作数的能量是运算本身的好几倍。",
        "HMMA（一个 warp 做 16×8×16 的矩阵乘）：一条指令 2048 次乘加，A 的每个元素被用 8 次、B 的每个元素被用 16 次，"
        "每次乘加只读 0.19 个输入操作数；指令开销摊到每次乘加上几乎为零。",
        "WGMMA（Hopper，一个 warpgroup 做 64×256×16）：一条指令 26 万次乘加，B 直接从共享内存读、被 64 行共用，"
        "输入操作数降到每次乘加 0.02 个。剩下的主要开销是累加器的读写——Blackwell 把累加器搬进了 Tensor Core 旁边专门的"
        "Tensor Memory（TMEM），连这一项也减掉了。",
        "这就是第 1 讲 §2.5–2.6 说的\"专用化\"：同样的工艺，运算本身占能量的比例从约两成提高到八九成。能量数字只是量级估计。",
    ])


# -----------------------------------------------------------------------------
def demo3_feeding():
    print_header("演示 3：喂饱 Tensor Core——数据块要多大（H100 的一个 SM）")
    flops_clk = H100.tc_flops_per_sm_clk()
    smem_b_clk = 128                                         # 32 个 bank × 4 字节
    need_smem = flops_clk / smem_b_clk
    hbm_ridge = H100.tc_peak() / (H100.hbm_tbs * 1e12)
    print(f"每个 SM 每周期 {flops_clk} 次运算（BF16），共享内存每周期最多给 {smem_b_clk} 字节 → 数据块的强度至少要 "
          f"{need_smem:.0f} FLOP/字节；如果所有数据都直接来自显存，要 {hbm_ridge:.0f} FLOP/字节（989 TFLOPS / 3.35 TB/s）")
    rows = []
    for bm, bn in ((64, 64), (128, 128), (128, 256), (256, 256)):
        inten = tile_intensity(bm, bn)
        stage = (bm + bn) * 64 * 2
        stages = (H100.smem_kb * 1024) // stage
        rows.append([f"{bm}×{bn}", f"{inten:.0f}", "够" if inten >= need_smem else "不够",
                     f"{min(1.0, inten / hbm_ridge):.0%}", f"{stage // 1024} KiB", stages])
    print_table(["输出块 bm×bn", "强度 FLOP/字节", "共享内存喂得饱吗", "只靠显存能跑到峰值的", "每级（k = 64）的 A、B 块",
                 f"{H100.smem_kb} KB 共享内存装得下几级"], rows)
    print_observations([
        "共享内存这一层：64×64 的输出块刚好够（32 FLOP/字节），更大的块更从容。A100 每个 SM 只有一半的运算能力，"
        "同样的共享内存带宽下 16 FLOP/字节就够——H100 把 Tensor Core 翻倍，逼着内核用更大的块（WGMMA 的 64×N 形状、"
        "warpgroup 级的指令都是为此设计的）。",
        "显存这一层：单个 SM 的块再大也到不了 295 FLOP/字节——256×256 的块也只有 128。剩下的复用要靠 50 MB 的 L2（不同 SM 读同一块"
        "A / B 时命中 L2）和线程块簇（cluster）里的 TMA 多播（一次读、发给几个 SM）。这是第 10 讲写矩阵乘法内核时要处理的问题。",
        "块越大，每一级流水线占的共享内存越多，能放的级数越少（256×256 只放得下 3 级）——块大小和流水线深度在抢同一块 228 KB（04 脚本）。",
    ])


if __name__ == "__main__":
    demo1_peak()
    demo2_instructions()
    demo3_feeding()