"""
=============================================================================
CS349E · 前置讲 P1 · 05_parallelism_and_power.py —— 三种并行与功耗：为什么 GPU 是"又宽又慢"的
=============================================================================

【这个脚本要回答的问题】
    Dennard scaling 结束以后（第 1 讲 §2.3），单核的频率不再上涨，性能只能靠**并行**：
    一条指令处理多个数据（SIMD）、多个核同时干活（多核）、以及一个核里同时有多条指令在执行（ILP，01 脚本）。
    而功耗随电压的平方增长，又让"更多、更慢的单元"比"更少、更快的单元"省电。本脚本：

      演示 1：峰值算力 = 单元数 × 频率 × 每周期每单元的运算数——拆开一个假想 CPU 和 H100 看看。
      演示 2：SIMD 与数据类型（实测）——数据放得进 cache 时，类型越窄，一条 SIMD 指令处理的元素越多；
              数据很大时，时间只看字节数。顺便发现：numpy 的 float16 反而很慢。
      演示 3：多核（实测）——计算密集（exp）和带宽密集（复制）的任务，加线程后的加速比完全不同；用 Amdahl 定律拟合。
      演示 4：功耗模型——P = C·V²·f，电压随频率升高；同样的吞吐，"多核低频"比"单核高频"省多少电。

【怎么读输出】演示 1、4 是公式（每次相同）；演示 2、3 是墙钟测量（只看趋势；编写时用的是 Apple M4 Pro，8 个性能核 + 4 个能效核）。

【运行】python3 05_parallelism_and_power.py          （约 2 秒）
        python3 05_parallelism_and_power.py --quick  （更快，数组更小）
【依赖】numpy；从 ../lecture01 导入 console、hw_specs。
"""

import os
import sys
from concurrent.futures import ThreadPoolExecutor

import numpy as np

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_num  # noqa: E402
from hw_specs import H100  # noqa: E402
from arch_lab import best_time  # noqa: E402

QUICK = "--quick" in sys.argv


# ---------------------------------------------------------------------------
# 演示 1：峰值算力的分解
# ---------------------------------------------------------------------------
def demo1_peak_decomposition():
    """
    峰值算力 = 核（SM）数 × 每核的运算单元数 × 每单元每周期的运算数 × 频率。

    【假想 CPU】8 个核、4 GHz、每核 4 条 128 位 SIMD 乘加流水线（每条 4 个 FP32 lane，每次乘加算 2 FLOP）。
        这些参数是"有代表性的假设"，不对应某一款具体芯片。
    【H100 的 CUDA core】132 个 SM，每个 SM 128 个 FP32 单元（CUDA 编程指南），乘加算 2 FLOP，加速频率 1.98 GHz
        （How to Scale Your Model 第 12 章）→ 约 67 TFLOP/s，和数据手册一致。
    【H100 的 Tensor Core】每个 SM 4 个；How to Scale Your Model 估计每个每周期约 1024 次 BF16 运算。
        由 989 TFLOP/s 反推对应的频率。
    """
    print_header("演示 1：峰值算力 = 单元数 × 每单元每周期的运算 × 频率")
    rows = []
    cpu = 8 * 4 * 4 * 2 * 4.0e9
    rows.append(["假想 CPU（8 核，SIMD FP32）", "8 核 × 4 流水线 × 4 lane", "2（乘加）", "4.0 GHz", fmt_num(cpu) + "FLOP/s"])
    cuda = 132 * 128 * 2 * 1.98e9
    rows.append(["H100 CUDA core（FP32）", "132 SM × 128 单元", "2（乘加）", "1.98 GHz", fmt_num(cuda) + "FLOP/s"])
    tc_per_cycle = 132 * 4 * 1024
    f_tc = H100.peak_flops("bf16") / tc_per_cycle
    rows.append(["H100 Tensor Core（BF16）", "132 SM × 4 Tensor Core", "≈1024（一个小矩阵乘法）",
                 f"反推 {f_tc / 1e9:.2f} GHz", fmt_num(H100.peak_flops("bf16")) + "FLOP/s"])
    print_table(["处理器", "并行单元", "每单元每周期 FLOP", "频率", "峰值"], rows)
    print_observations([
        f"H100 的频率比 CPU 低一半，CUDA core 的 FP32 峰值却是假想 CPU 的 {cuda / cpu:.0f} 倍：靠的是并行单元多得多（1.7 万个 lane）。",
        f"Tensor Core 又比 CUDA core 高 {H100.peak_flops('bf16') / cuda:.0f} 倍：每个单元每周期做的不是 1 次乘加，而是一个小矩阵乘法"
        "（第 1 讲 §2.5 的'一条指令干更多活'）。",
        "反推出的频率约 1.83 GHz，和 H100 的加速频率（1.98 GHz）同一量级——说明'每 Tensor Core 每周期约 1024 FLOP'是合理的估计"
        "（NVIDIA 没有公开这个数，见 How to Scale Your Model 第 12 章）。",
    ])


# ---------------------------------------------------------------------------
# 演示 2：SIMD 与数据类型
# ---------------------------------------------------------------------------
def demo2_dtypes():
    """
    对 5 种数据类型做逐元素加法：一次在 cache 里（小数组反复算），一次在内存里（大数组）。

    【预期】
        - 在 cache 里：瓶颈是计算。128 位的 SIMD 寄存器一次能装 2 个 float64、4 个 float32、16 个 int8，
          类型越窄每条指令处理的元素越多，每个元素越快。
        - 在内存里：瓶颈是带宽。每个元素要读 2 个、写 1 个，时间 ∝ 字节数。
    【float16】numpy 在大多数 CPU 上没有原生的 float16 算术路径（要先转换成 float32 再算再转回来），所以会很慢。
        这不是 float16 本身的问题，而是"硬件和库支不支持"的问题。
    """
    print_header("演示 2：SIMD 与数据类型——同样的逐元素加法，换不同的类型")
    small, rep = 8192, (300 if QUICK else 2000)
    big = 1 << (22 if QUICK else 25)
    rows = []
    res = {}
    for dt in (np.float64, np.float32, np.float16, np.int32, np.int8):
        a = np.ones(small, dtype=dt)
        o = np.empty_like(a)
        t_small = best_time(lambda: [np.add(a, a, out=o) for _ in range(rep)], repeats=3) / (small * rep)
        A = np.ones(big, dtype=dt)
        O = np.empty_like(A)
        t_big = best_time(lambda: np.add(A, A, out=O), repeats=3) / big
        name = np.dtype(dt).name
        res[name] = (t_small, t_big)
        rows.append([name, np.dtype(dt).itemsize, f"{t_small * 1e9:.3f} ns", f"{t_big * 1e9:.3f} ns",
                     f"{3 * np.dtype(dt).itemsize / t_big / 1e9:.0f} GB/s"])
    print_table(["类型", "字节/元素", "在 cache 里：每元素", "在内存里：每元素", "在内存里：等效带宽（读 2 写 1）"], rows)
    f32s, f32b = res["float32"]
    i8s, i8b = res["int8"]
    f16s, _ = res["float16"]
    print_observations([
        f"数据很大（在内存里）时，int8 每个元素比 float32 快约 {f32b / i8b:.1f} 倍，接近字节数之比 4：时间主要取决于搬多少字节。"
        "这就是低精度推理在 memory-bound 阶段（decode）最直接的好处。",
        f"数据在 cache 里时，int8 比 float32 快约 {f32s / i8s:.1f} 倍：一条 SIMD 指令能同时处理更多个窄元素。",
        f"float16 反而比 float32 慢约 {f16s / f32s:.0f} 倍：numpy 没有用原生的半精度算术，每个元素都要转换。"
        "**数值格式再小，硬件和软件不支持就没有用**——第 3 讲的 NVFP4、FP8 能加速，前提是 Tensor Core 原生支持它们。",
    ])


# ---------------------------------------------------------------------------
# 演示 3：多核
# ---------------------------------------------------------------------------
def amdahl_speedup(serial_frac, n):
    """Amdahl 定律：串行部分占 s，其余部分完美并行到 n 份，加速比 = 1 / (s + (1 − s)/n)。"""
    return 1.0 / (serial_frac + (1.0 - serial_frac) / n)


def fit_serial_fraction(speedup, n):
    """
    由"n 个线程的实测加速比"反解 Amdahl 定律里的串行比例 s：s = (n / S − 1) / (n − 1)。

    【边界】n = 1 时无法求解，返回 0；结果会被截到 [0, 1]。
    """
    if n <= 1:
        return 0.0
    s = (n / speedup - 1.0) / (n - 1.0)
    return min(1.0, max(0.0, s))


def demo3_multicore():
    """
    k 个线程各处理数组的 1/k：一个计算密集（np.exp，每个元素要算几十条指令），一个带宽密集（复制）。

    【为什么要两个】多核能加速的前提是瓶颈在"核"上。计算密集的任务，核越多越快；带宽密集的任务，
        所有核共用同一个内存系统，加核很快就没用了。
    """
    print_header("演示 3：多核——计算密集 vs 带宽密集（每个线程处理一段）")
    n = 1 << (22 if QUICK else 25)
    x = np.random.default_rng(0).random(n, dtype=np.float32)
    y = np.empty_like(x)
    results = {}
    for kind in ("exp", "copy"):
        base = None
        results[kind] = {}
        for k in (1, 2, 4, 8, 12):
            bounds = np.linspace(0, n, k + 1).astype(int)

            def work(i, kind=kind, bounds=bounds):
                lo, hi = bounds[i], bounds[i + 1]
                if kind == "exp":
                    np.exp(x[lo:hi], out=y[lo:hi])
                else:
                    np.copyto(y[lo:hi], x[lo:hi])

            with ThreadPoolExecutor(k) as ex:
                t = best_time(lambda: list(ex.map(work, range(k))), repeats=3)
            base = base or t
            results[kind][k] = base / t
    rows = []
    for k in (1, 2, 4, 8, 12):
        rows.append([k, f"{results['exp'][k]:.2f}×", f"{results['copy'][k]:.2f}×", f"{amdahl_speedup(0.05, k):.2f}×"])
    print_table(["线程数", "exp（计算密集）加速比", "复制（带宽密集）加速比", "参考：Amdahl s = 5%"], rows)
    s_exp = fit_serial_fraction(results["exp"][8], 8)
    print_observations([
        f"计算密集的 exp：4 个线程约 {results['exp'][4]:.1f} 倍，接近线性；8 个线程约 {results['exp'][8]:.1f} 倍，"
        f"按 Amdahl 定律反解出的'等效串行比例'约 {s_exp * 100:.0f}%。",
        "4 个线程之后变差的原因这里没有拆开测，可能的因素有：多个核同时满载时频率会降低、"
        "这台机器的 8 个性能核分成两个簇、线程可能被调度到较慢的能效核上、以及线程池本身的开销。"
        "真实系统里'加核不线性'几乎总是同时有好几个原因。",
        f"带宽密集的复制：加到 4 个线程以后基本不再变快（{results['copy'][12]:.1f} 倍封顶）——瓶颈是共享的内存系统，不是核。"
        "这和 decode 在 GPU 上的处境一样：加算力没用，要加带宽。",
        "Amdahl 定律：哪怕只有 5% 的部分不能并行，12 个核最多也只有约 7.7 倍。P2 §1、§8 会在 GPU 的尺度（上万个 lane）上重新算这笔账。",
    ])


# ---------------------------------------------------------------------------
# 演示 4：功耗与"又宽又慢"
# ---------------------------------------------------------------------------
def voltage_at(f_ghz, f_lo=1.0, v_lo=0.6, f_hi=4.0, v_hi=1.1):
    """
    频率 f 时需要的电压：在 (f_lo, v_lo) 和 (f_hi, v_hi) 之间线性插值（一个常用的一阶近似）。

    【直觉】晶体管要在更短的时间里翻转完成，就需要更高的电压来提供更大的电流。
    【边界】这组端点是**示意值**，不对应某一款芯片；超出范围时线性外推。
    """
    return v_lo + (v_hi - v_lo) * (f_ghz - f_lo) / (f_hi - f_lo)


def demo4_power():
    """
    动态功耗 P = C·V²·f（取 C 为常数 1，只看相对值）。每次运算的能量 = P / f = C·V²。

    【问题】要达到"1 个核 × 4 GHz"的吞吐，用 n 个核、每个跑 4/n GHz（假设任务能完美并行），总功耗是多少？
    【简化】只算动态功耗，不算漏电（静态功耗）；漏电大的时候，"跑快一点然后关掉"（race to idle）有时反而更省。
    """
    print_header("演示 4：功耗——P = C·V²·f，电压随频率上升（示意参数）")
    rows = []
    for f in (1.0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0):
        v = voltage_at(f)
        p = v * v * f
        rows.append([f"{f:.1f} GHz", f"{v:.2f} V", f"{p:.2f}", f"{v * v:.3f}"])
    print_table(["频率", "所需电压", "相对功耗 V²·f", "每次运算的相对能量 V²"], rows)

    print("\n同样的总吞吐（= 1 个核跑 4 GHz），用更多、更慢的核：")
    rows = []
    p_ref = voltage_at(4.0) ** 2 * 4.0
    for n in (1, 2, 4):
        f = 4.0 / n
        p = n * voltage_at(f) ** 2 * f
        rows.append([n, f"{f:.1f} GHz", f"{p:.2f}", f"{p / p_ref * 100:.0f}%"])
    print_table(["核数", "每核频率", "总相对功耗", "相对 1 核 4 GHz"], rows)
    print_observations([
        "频率从 1 GHz 提到 4 GHz，吞吐 ×4，功耗却 ×13 左右：每次运算的能量随电压的平方上涨。",
        "反过来，同样的吞吐用 4 个 1 GHz 的核来做，总功耗只有 1 个 4 GHz 核的三成左右——前提是任务能并行。",
        "这就是 GPU 的设计哲学：上万个运算单元、跑在 2 GHz 左右，而不是几个跑在 5 GHz 的复杂核。"
        "LLM 的矩阵乘法有海量的并行度，正好吃得下这种设计。",
        "它也解释了第 1 讲的'性能 = 能效 × 功率'：功耗预算固定时，降低电压和频率、换来更多并行单元，是提高总性能的主要手段。",
    ])


if __name__ == "__main__":
    demo1_peak_decomposition()
    demo2_dtypes()
    demo3_multicore()
    demo4_power()