"""
gpu_model.py —— "纸面 GPU"：公开规格、占用率、波次、延迟隐藏、服务器里的数据流（第 2 讲；第 3、4 讲也会用）
=====================================================================

【这个文件要回答的问题】
    1. 一块 GPU 的峰值算力是怎么来的？（SM 数 × 每个 SM 的核心数 × 2 × 频率）
    2. 一个 SM 上能同时驻留多少个线程束？被什么卡住？（占用率 occupancy）
    3. 网格里的块是怎样一波一波分到 SM 上的？块数不是 SM 数的整数倍时浪费多少？（波次与尾部效应）
    4. 为什么 GPU 要同时驻留那么多线程？（用别的线程束的计算掩盖访存延迟——延迟隐藏，以及 Little 定律）
    5. 一台 8 卡服务器里，数据在 NVMe、内存、PCIe、NVLink、显存之间怎样流动，各要多久？

【本文件提供】
    1. 规格      : GPUSpec、A6000、A100、H100、B200、GPUS；CPUSpec、EPYC_9754
    2. 占用率    : occupancy(gpu, threads_per_block, regs_per_thread, smem_per_block)
    3. 波次      : waves(num_blocks, slots)
    4. 延迟隐藏  : simulate_scheduler（一个线程束调度器的逐周期模拟）、scheduler_bound（对应的公式）
    5. 服务器    : LAB_SERVER_LINKS（原课 slides 里那台 8 × A6000 服务器的连接带宽）、transfer_seconds、ring_allreduce_seconds

【怎样读这里的数字】
    规格都是厂商公开的标称值（dense，不含 2:4 稀疏加速；频率取 boost 频率）。
    原课 slides 的 GPU 对比表里，张量核心一栏取自 NVIDIA 产品页，是"含稀疏加速"的数字，正好是 dense 的 2 倍，
    02 号脚本会把两者并排列出。占用率的算法是简化版（忽略了寄存器和共享内存的分配粒度等细节），
    和 NVIDIA 官方的 Occupancy Calculator 在常见情况下一致，极端情况可能差一点。
"""

import math
from dataclasses import dataclass


# =====================================================================
# 1. 规格
# =====================================================================

@dataclass(frozen=True)
class GPUSpec:
    """
    一块 GPU 的规格（标称值）。

    【属性】
        name, arch           : 名字、架构（Ampere / Hopper / Blackwell）
        cc                   : 计算能力（compute capability），决定了下面这些"每个 SM 的上限"
        sms                  : SM（流式多处理器）的个数
        fp32_lanes_per_sm    : 每个 SM 的 FP32 CUDA 核心数（一个周期能做的 FP32 乘加个数）
        boost_ghz            : boost 频率（GHz）
        fp32_tflops          : 规格表上的 FP32（非张量核心）峰值
        tensor_dense_tflops  : 张量核心的 dense 峰值，字典 {"TF32": .., "BF16": .., "FP8": .., "FP4": ..}；没有的格式不列
        mem_gb, mem_bw       : 显存容量（GB）、显存带宽（字节/秒）
        l2_mb                : L2 缓存（MB）
        max_threads_per_sm, max_blocks_per_sm : 每个 SM 最多驻留的线程数 / 线程块数
        regs_per_sm          : 每个 SM 的 32 位寄存器个数（64K = 65536，即 256 KB）
        smem_per_sm          : 每个 SM 最多可配置的共享内存（字节）
        smem_per_block       : 每个块最多能申请的共享内存（字节，需要显式申请超过 48 KB 的部分）
        l1_smem_kb           : 每个 SM 的 L1 + 共享内存合计（KB；slides 里写的 128K / 192K / 256K 指的是它）
        tdp_w                : 热设计功耗（瓦）
    【边界】B200 只录入了原课 slides 给出的数字，SM 级参数未知（为 None）。
    """
    name: str
    arch: str
    cc: str
    sms: int
    fp32_lanes_per_sm: int
    boost_ghz: float
    fp32_tflops: float
    tensor_dense_tflops: dict
    mem_gb: float
    mem_bw: float
    l2_mb: float
    max_threads_per_sm: int
    max_blocks_per_sm: int
    regs_per_sm: int
    smem_per_sm: int
    smem_per_block: int
    l1_smem_kb: int
    tdp_w: float

    @property
    def max_warps_per_sm(self):
        return self.max_threads_per_sm // 32 if self.max_threads_per_sm else None

    @property
    def fp32_cores(self):
        return self.sms * self.fp32_lanes_per_sm if self.sms else None

    def fp32_peak_from_cores(self):
        """峰值 FP32 算力 = SM 数 × 每 SM 核心数 × 2（一次乘加 FMA 算 2 次运算）× 频率。单位 FLOP/s。"""
        return self.sms * self.fp32_lanes_per_sm * 2 * self.boost_ghz * 1e9

    def ridge(self, fmt="FP32"):
        """平衡点（ridge point）= 峰值算力 ÷ 显存带宽，单位 FLOP/字节。fmt 为 "FP32" 或张量核心的格式名。"""
        peak = self.fp32_tflops if fmt == "FP32" else self.tensor_dense_tflops[fmt]
        return peak * 1e12 / self.mem_bw


KB = 1024

# NVIDIA RTX A6000（GA102，原课作者实验室服务器里的卡）：84 SM × 128 核，1.80 GHz，38.7 TFLOPS FP32；
# 张量核心 FP16/BF16 dense 154.8 TFLOPS（规格表的 309.7 是含稀疏的）；48 GB GDDR6，768 GB/s；L2 6 MB；300 W
A6000 = GPUSpec("RTX A6000", "Ampere (GA102)", "8.6", 84, 128, 1.80, 38.7,
                {"TF32": 77.4, "BF16": 154.8}, 48, 768e9, 6,
                1536, 16, 65536, 100 * KB, 99 * KB, 128, 300)

# NVIDIA A100 SXM4 80GB（GA100）：108 SM × 64 核，1.41 GHz，19.5 TFLOPS FP32；TF32 156、BF16 312（dense）；
# 80 GB HBM2e，2.039 TB/s（PCIe 版 1.935 TB/s，原课第 4 讲用的是它）；L2 40 MB；400 W
A100 = GPUSpec("A100 SXM 80GB", "Ampere (GA100)", "8.0", 108, 64, 1.41, 19.5,
               {"TF32": 156, "BF16": 312}, 80, 2.039e12, 40,
               2048, 32, 65536, 164 * KB, 163 * KB, 192, 400)

# NVIDIA H100 SXM5 80GB（GH100）：132 SM × 128 核，1.98 GHz，67 TFLOPS FP32；TF32 494.7、BF16 989.4、FP8 1978.9（dense）；
# 80 GB HBM3，3.35 TB/s；L2 50 MB；700 W
H100 = GPUSpec("H100 SXM 80GB", "Hopper (GH100)", "9.0", 132, 128, 1.98, 67,
               {"TF32": 494.7, "BF16": 989.4, "FP8": 1978.9}, 80, 3.35e12, 50,
               2048, 32, 65536, 228 * KB, 227 * KB, 256, 700)

# NVIDIA B200（Blackwell）：只录入原课 slides 的数字（slides 张量核心一栏是含稀疏的，这里除以 2 得到 dense）
B200 = GPUSpec("B200", "Blackwell", "10.0", None, None, None, 75,
               {"TF32": 1100, "BF16": 2250, "FP8": 4500, "FP4": 9000}, 192, 7.7e12, None,
               None, None, None, None, None, None, 1000)

GPUS = [A6000, A100, H100, B200]

# 原课第 2 讲 slides 的 GPU 对比表（照抄，单位 TFLOPS；张量核心一栏含 2:4 稀疏加速）
SLIDES_LINEUP = {
    "B200": {"FP64": 37, "FP64 Tensor": 37, "FP32": 75, "TF32 Tensor": 2200, "BF16 Tensor": 4500,
             "INT8 Tensor": 9000, "FP8 Tensor": 9000, "FP4 Tensor": 18000, "mem": "192GB HBM3e", "bw": "7.7TB/s"},
    "H100": {"FP64": 34, "FP64 Tensor": 67, "FP32": 67, "TF32 Tensor": 989, "BF16 Tensor": 1979,
             "INT8 Tensor": 3958, "FP8 Tensor": 3958, "FP4 Tensor": None, "mem": "80GB HBM3", "bw": "3.2TB/s"},
    "A100": {"FP64": 9.7, "FP64 Tensor": 19.5, "FP32": 19.5, "TF32 Tensor": 312, "BF16 Tensor": 624,
             "INT8 Tensor": 1248, "FP8 Tensor": None, "FP4 Tensor": None, "mem": "80GB HBM2e", "bw": "2TB/s"},
}


@dataclass(frozen=True)
class CPUSpec:
    """
    一颗 CPU 的规格（标称值）。

    【属性】cores、threads（硬件线程，开超线程后是核数的 2 倍）、base_ghz、l3_mb、tdp_w、
        fp32_per_cycle_simd（每核每周期能做的 FP32 运算数：两个 256 位 FMA 单元 → 2 × 8 × 2 = 32）、
        mem_channels、mem_mts（内存每通道的速率，MT/s）
    """
    name: str
    cores: int
    threads: int
    base_ghz: float
    l3_mb: int
    tdp_w: float
    fp32_per_cycle_simd: int
    mem_channels: int
    mem_mts: int

    def scalar_peak(self):
        """每核每周期 1 次乘加（2 次运算）——原课 slides 的 576 GFLOPS 就是这样算的。"""
        return self.cores * self.base_ghz * 1e9 * 2

    def simd_peak(self):
        """用满 SIMD（AVX-512 / AVX2 的 FMA）时的峰值，按基础频率估算。"""
        return self.cores * self.base_ghz * 1e9 * self.fp32_per_cycle_simd

    def mem_bw(self):
        """内存带宽 = 通道数 × 速率 × 8 字节（每通道 64 位）。"""
        return self.mem_channels * self.mem_mts * 1e6 * 8


# AMD EPYC 9754（Zen 4c，原课 slides 用来和 A6000 对比）：128 核 / 256 线程，2.25 GHz，256 MB L3，360 W，
# 12 通道 DDR5-4800
EPYC_9754 = CPUSpec("AMD EPYC 9754", 128, 256, 2.25, 256, 360, 32, 12, 4800)


# =====================================================================
# 2. 占用率
# =====================================================================

def occupancy(gpu, threads_per_block, regs_per_thread=32, smem_per_block=0):
    """
    一个 SM 上能同时驻留几个线程块、几个线程束（简化版的占用率计算器）。

    【思路】驻留的块数取下面四个上限里最小的一个：
        1. 线程束上限 : 每个 SM 最多 max_warps 个线程束 → ⌊max_warps / 每块线程束数⌋
        2. 块数上限   : 每个 SM 最多 max_blocks_per_sm 个块
        3. 寄存器     : 每个线程束占 ⌈regs_per_thread × 32 / 256⌉ × 256 个寄存器（按 256 个一份分配）
                        → ⌊regs_per_sm / (每线程束寄存器 × 每块线程束数)⌋
        4. 共享内存   : 每块占 smem_per_block + 1 KB（计算能力 8.0 起，系统给每个块保留 1 KB）
                        → ⌊smem_per_sm / 每块共享内存⌋（smem_per_block 为 0 时不受限）
    【参数】gpu：GPUSpec；threads_per_block：每块线程数；regs_per_thread：每个线程用的寄存器数（编译器决定）；
        smem_per_block：每块申请的共享内存（字节）
    【返回】字典：blocks（每 SM 驻留的块数）、warps、occupancy（warps / 最大线程束数）、limiter（哪个上限起作用）、
        以及四个上限各自的值
    【边界】每线程寄存器超过 255、每块线程超过 1024、共享内存超过每块上限时，块根本发射不了：返回 blocks = 0。
    """
    wpb = math.ceil(threads_per_block / 32)
    if threads_per_block > 1024 or regs_per_thread > 255 or smem_per_block > gpu.smem_per_block:
        return {"blocks": 0, "warps": 0, "occupancy": 0.0, "limiter": "无法发射"}
    limits = {"线程束数": gpu.max_warps_per_sm // wpb, "块数": gpu.max_blocks_per_sm}
    regs_per_warp = math.ceil(regs_per_thread * 32 / 256) * 256
    limits["寄存器"] = gpu.regs_per_sm // (regs_per_warp * wpb)
    if smem_per_block > 0:
        limits["共享内存"] = gpu.smem_per_sm // (smem_per_block + 1024)
    limiter = min(limits, key=lambda k: limits[k])
    blocks = limits[limiter]
    warps = blocks * wpb
    return {"blocks": blocks, "warps": warps, "occupancy": warps / gpu.max_warps_per_sm,
            "limiter": limiter, "limits": limits}


# =====================================================================
# 3. 波次
# =====================================================================

def waves(num_blocks, slots):
    """
    假设每个块耗时相同：num_blocks 个块分到 slots 个"位置"（= SM 数 × 每 SM 驻留块数）上，要几波？

    【返回】(波数, 效率)。效率 = 实际的块数 ÷ (波数 × 位置数)：最后一波没坐满的位置都在空等（尾部效应）。
    【例子】132 个 SM、每个 SM 1 个块：132 个块 → 1 波，效率 100%；133 个块 → 2 波，效率 50.4%。
    """
    w = math.ceil(num_blocks / slots)
    return w, num_blocks / (w * slots)


# =====================================================================
# 4. 延迟隐藏：一个线程束调度器的逐周期模拟
# =====================================================================

def simulate_scheduler(n_warps, compute_per_load, mem_latency=500, load_interval=0, cycles=40_000):
    """
    一个线程束调度器（SM 的一个分区）上 n_warps 个线程束的逐周期模拟。

    【模型】
        每个线程束反复执行同一段程序：[一条访存指令 LD，然后 compute_per_load 条计算指令]。
        - LD 之后的第一条计算指令要用到读回来的数据，必须等 mem_latency 个周期；
        - 计算指令之间不需要等待（每周期可以发射下一条）；
        - 调度器每个周期最多发射一条指令：从"已经准备好"的线程束里按轮转顺序选一个；
        - load_interval > 0 时模拟显存带宽的限制：两次 LD 之间至少隔 load_interval 个周期
          （一个线程束一次读 128 字节；带宽 = 128 / load_interval 字节/周期）。
    【参数】n_warps：驻留的线程束数；compute_per_load：每次访存之后的计算指令数（越大越"计算密集"）；
        mem_latency：访存延迟（周期，slides：全局内存约 500）；cycles：模拟多少个周期
    【返回】字典：utilization（发射了指令的周期占比）、loads_in_flight（平均有几个 LD 在路上）、load_rate（每周期发射的 LD 数）
    【边界】compute_per_load = 0 表示只读不算（纯搬运）：每次 LD 之后要等数据回来才能发下一条 LD。
        真实的线程束可以连发几条互不依赖的 LD（指令级并行），这里不模拟。
    """
    W, c = n_warps, compute_per_load
    ready = [0] * W                # 每个线程束下一条指令最早能发射的周期
    pc = [0] * W                   # 0 表示下一条是 LD；1..c 表示下一条是第几条计算
    issued = loads = 0
    inflight_sum = 0
    pending = []                   # 正在路上的 LD 的完成时间
    next_load_ok = 0
    rr = 0
    for cyc in range(cycles):
        chosen = -1
        for j in range(W):
            w = (rr + j) % W
            if ready[w] <= cyc and not (pc[w] == 0 and cyc < next_load_ok):
                chosen = w
                break
        if chosen >= 0:
            w = chosen
            issued += 1
            if pc[w] == 0:                                   # 发射 LD
                loads += 1
                pending.append(cyc + mem_latency)
                next_load_ok = cyc + load_interval
                ready[w] = cyc + mem_latency                 # 下一条指令要用这次读回的数据：等 L 个周期
                pc[w] = 1 if c > 0 else 0
            else:                                            # 发射一条计算
                ready[w] = cyc + 1
                pc[w] = pc[w] + 1 if pc[w] < c else 0
            rr = (w + 1) % W
        pending = [t for t in pending if t > cyc]
        inflight_sum += len(pending)
    return {"utilization": issued / cycles, "loads_in_flight": inflight_sum / cycles, "load_rate": loads / cycles}


def scheduler_bound(n_warps, compute_per_load, mem_latency=500, load_interval=0):
    """
    simulate_scheduler 的近似公式：利用率 ≈ min(1, W·(1+c)/(L+c), (1+c)/load_interval)。

    【推导】一个线程束一轮要 L + c 个周期（等数据 L，再连发 c 条计算），其间发射 1 + c 条指令；
        W 个线程束交错执行，发射率是单个的 W 倍，但不超过每周期 1 条；
        带宽限制下每 load_interval 个周期最多一次 LD，每次 LD 带 c 条计算。
    """
    a = n_warps * (1 + compute_per_load) / (mem_latency + compute_per_load)
    b = (1 + compute_per_load) / load_interval if load_interval else float("inf")
    return min(1.0, a, b)


# =====================================================================
# 5. 服务器里的数据流
# =====================================================================

# 原课第 2 讲 slides 里那台服务器（8 × A6000、2 × AMD EPYC 9354）上各段连接的带宽（字节/秒，单向，标称值的量级）
LAB_SERVER_LINKS = {
    "NVMe SSD（单块 PCIe 4.0 盘，顺序读）": 7e9,
    "CPU 内存（DDR5，一颗 CPU）": 300e9,                 # slides 画的是 64–512 GB/s
    "PCIe 4.0 x16（CPU ↔ GPU，或经 PCIe 的 GPU ↔ GPU）": 32e9,   # slides：16 条通道 × 2 GB/s
    "NVLink 桥（A6000 两两相连，单向）": 56.25e9,        # slides 写的 112.5 GB/s 是双向之和
    "A6000 显存（GDDR6）": 768e9,
}


def transfer_seconds(nbytes, bandwidth, latency=10e-6):
    """α-β 模型（P4 §5）：时间 = 固定延迟 + 字节数 / 带宽。"""
    return latency + nbytes / bandwidth


def ring_allreduce_seconds(nbytes, n_gpus, slowest_bw):
    """环形 all-reduce 的带宽项：每张卡发送和接收 2(N−1)/N × 数据量，速度由环上最慢的一段决定（第 14 讲推导）。"""
    return 2 * (n_gpus - 1) / n_gpus * nbytes / slowest_bw