"""
=============================================================================
CS349E · 第 8 讲共享模块：hw_engines.py —— 矩阵引擎的模型：脉动阵列、TPU v1、Tensor Core、异步流水线
=============================================================================

【这个模块要解决什么问题】
    前七讲一直把"峰值算力"和"显存带宽"当作芯片规格表上的两个数。这一讲打开芯片看：
    这两个数从哪里来？为什么矩阵乘法能跑到峰值、别的运算跑不到？为什么推理常常受限于内存？
    本模块提供四件工具：

    1. SystolicArray：一个**逐周期**模拟的权重驻留（weight-stationary）脉动阵列——TPU v1 的矩阵单元就是 256×256 的这种阵列。
       真的做乘加、真的让数据一格一格地流动，输出和 X @ W 逐位比较；统计周期数和利用率。
       tiled_matmul_time：大矩阵切成阵列大小的块，加上从 DRAM 读权重的时间（Jouppi 论文 §7 的 600×600 例子）。
    2. TPU v1 论文（Jouppi et al., ISCA 2017）的数据：三块芯片的规格（表 2）和六个生产应用（表 1、表 3），
       用来复现论文的 roofline 分析。
    3. GPU 的 SM 与 Tensor Core：V100 / A100 / H100 每个 SM 每个周期能做多少乘加；一条指令做多少乘加、要读多少操作数。
    4. pipeline_sim：一个 SM 上"读数据块 → 算数据块"循环的事件模拟——同步读、双缓冲、多级流水线（TMA + warp 专门化）。

【运行】python3 hw_engines.py   # 自检
【依赖】numpy。本模块把 ../lecture01 加进 sys.path，方便脚本导入 console 和 hw_specs（Horowitz 能量表、exp 吞吐）。
"""

import math
import os
import sys
from dataclasses import dataclass

import numpy as np

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


# =============================================================================
# 1. 权重驻留的脉动阵列
# =============================================================================
class SystolicArray:
    """
    R 行 × C 列的乘加单元（PE，processing element）阵列，权重驻留（weight-stationary）：

        - 计算 Y = X · W，X 是 [B, R]（B 个输入向量，每个长 R），W 是 [R, C]（预先装进阵列，PE(i, j) 存 W[i, j]）。
        - 输入从左边进来：第 i 行在第 b + i 个周期收到 X[b, i]（"斜着"送进来，skew），之后每个周期向右传一格。
        - 部分和从上往下流：PE(i, j) 收到上面传来的部分和，加上 X[b, i] × W[i, j]，下个周期传给下面。
        - 第 j 列的底部在第 b + (R − 1) + j 个周期输出 Y[b, j]。

    【为什么这样设计】（Jouppi §4：读一个大 SRAM 的能量比做一次运算高得多）
        每个输入从缓冲区读一次，在一行里被 C 个 PE 用；每个权重读进来一次，被 B 个输入用；
        部分和在 PE 之间直接传，不写回缓冲区。阵列内部只有相邻 PE 之间的短导线，没有大 SRAM 的读写。
    【方法】run(X, W)：逐周期模拟，返回 (Y, 统计字典)
    【统计】cycles（从第一个输入进入到最后一个输出离开）、useful_macs、utilization（有用的乘加 / (周期 × R × C)）、
        buffer_reads（从输入缓冲区读了多少个数）、weight_loads（装了多少个权重）
    【边界】X 的列数必须等于 R、W 的形状必须是 [R, C]（小于阵列时先补零——补出来的 PE 白占着，见 tiled_matmul_time）。
    """

    def __init__(self, R, C):
        self.R, self.C = R, C

    def run(self, X, W, trace=False):
        """
        逐周期模拟。trace=True 时额外返回每个周期"哪些 PE 在做有用的乘加"的布尔矩阵列表（演示数据流用）。

        【Python 小知识】x_reg[:, 1:] = x_reg[:, :-1] 一次把整块数组向右平移一格（numpy 的切片赋值会先复制右边）。
        """
        R, C = self.R, self.C
        B = X.shape[0]
        if X.shape[1] != R or W.shape != (R, C):
            raise ValueError(f"形状不对：X {X.shape}、W {W.shape}，阵列 {R}×{C}")
        x_reg = np.zeros((R, C))            # 每个 PE 里正在用的输入
        x_valid = np.zeros((R, C), dtype=bool)
        p_reg = np.zeros((R, C))            # 每个 PE 算完的部分和（下个周期交给下面的 PE）
        Y = np.zeros((B, C))
        total = B + R + C - 2               # 最后一个输出在第 (B−1) + (R−1) + (C−1) 个周期产生
        macs = 0
        frames = []
        rows = np.arange(R)
        for t in range(total + 1):
            # 1. 输入向右流一格；第 0 列从左边收到新的输入（第 i 行收到 X[t − i, i]）
            x_reg[:, 1:] = x_reg[:, :-1]
            x_valid[:, 1:] = x_valid[:, :-1]
            b = t - rows
            ok = (b >= 0) & (b < B)
            x_reg[:, 0] = np.where(ok, X[np.clip(b, 0, B - 1), rows], 0.0)
            x_valid[:, 0] = ok
            # 2. 部分和向下流一格，再加上这个 PE 的乘积
            above = np.vstack([np.zeros((1, C)), p_reg[:-1]])
            p_reg = above + x_reg * W
            macs += int(x_valid.sum())
            if trace:
                frames.append(x_valid.copy())
            # 3. 底部输出：第 j 列在第 b + (R − 1) + j 个周期产生 Y[b, j]
            bj = t - (R - 1) - np.arange(C)
            okj = (bj >= 0) & (bj < B)
            Y[bj[okj], np.arange(C)[okj]] = p_reg[R - 1, okj]
        cycles = total + 1
        stats = {"cycles": cycles, "useful_macs": macs, "utilization": macs / (cycles * R * C),
                 "buffer_reads": B * R, "weight_loads": R * C, "outputs": B * C}
        return (Y, stats, frames) if trace else (Y, stats)


def systolic_cycles(B, R, C):
    """不模拟、直接按公式算：B 个输入流过 R×C 的阵列要 B + R + C − 2 + 1 个周期（填充 + 流水 + 排空）。"""
    return B + R + C - 1


def tiled_matmul_time(B, K, N, R=256, C=256, clock=700e6, weight_bw=34e9, w_bytes=1, double_buffer=True):
    """
    [B, K] × [K, N] 的矩阵乘法在 R×C 的权重驻留阵列上要多久（TPU v1 风格的粗模型）。

    【思路】W 被切成 ⌈K/R⌉ × ⌈N/C⌉ 个 R×C 的块（不够的补零——补零的部分照样占 PE、照样要从 DRAM 读进来）。
        每个块：从 DRAM 读权重（R·C·w_bytes / weight_bw 秒）+ 让 B 个输入流过（B 个周期）。
        double_buffer=True 时，下一块的权重在算这一块时读进来（TPU 的矩阵单元有两份权重存储），
        每块的时间 = max(读权重, 算)；否则相加。填充和排空（R + C 个周期）在相邻块之间流水掉，只算一次。
    【参数】clock：阵列的时钟（TPU v1 是 700 MHz）；weight_bw：权重存储的带宽（TPU v1 的 DDR3 是 34 GB/s）
    【返回】字典：time（秒）、tiles、padded_weight_bytes、load_time、compute_time、utilization（有用乘加 / 峰值）
    """
    tiles = math.ceil(K / R) * math.ceil(N / C)
    t_load = R * C * w_bytes / weight_bw
    t_comp = B / clock
    per_tile = max(t_load, t_comp) if double_buffer else t_load + t_comp
    time = tiles * per_tile + (R + C) / clock + (t_load if double_buffer else 0.0)
    useful = B * K * N
    return {"time": time, "tiles": tiles, "padded_weight_bytes": tiles * R * C * w_bytes,
            "load_time": tiles * t_load, "compute_time": tiles * t_comp,
            "utilization": useful / (time * clock * R * C)}


# =============================================================================
# 2. Jouppi et al. 2017 的数据
# =============================================================================
@dataclass(frozen=True)
class Chip2015:
    """
    论文表 2 的一块芯片（按"一个 die"计）。

    【字段】peak_tops：峰值（TOPS，乘和加各算一次；TPU 是 8 位整数，CPU / GPU 是浮点）；
        mem_gbs：存权重的存储的带宽（GB/s）；onchip_mib：片上存储（MiB）；tdp_w：功耗；nm、mhz：工艺、时钟
    """
    name: str
    peak_tops: float
    mem_gbs: float
    onchip_mib: float
    tdp_w: float
    nm: int
    mhz: int

    def ridge(self):
        """
        roofline 的拐点，单位用论文的"每读一个权重字节做几次乘加（MAC）"。

        【注意单位】论文的 TOPS 把乘和加各算一次，roofline 横轴的 "operations" 却是乘加次数（表 1 第 10 列：
            MLP0 批大小 200 → 每个权重字节 200 次乘加）。所以拐点 = 峰值 / 2 / 带宽：
            TPU 92e12 / 2 / 34e9 ≈ 1353（论文："1350 operations per byte"）。
        """
        return self.peak_tops * 1e12 / 2 / (self.mem_gbs * 1e9)


TPU_V1 = Chip2015("TPU v1", 92.0, 34.0, 28.0, 75.0, 28, 700)
K80_DIE = Chip2015("K80（一个 die）", 2.8, 160.0, 8.0, 150.0, 28, 560)
HASWELL = Chip2015("Haswell E5-2699 v3", 1.3, 51.0, 51.0, 145.0, 22, 2300)

# 论文表 1、表 3、表 6：应用名 → (类型, 每个权重字节的乘加次数, TPU 批大小, 权重个数, 部署占比, TPU 实测 TOPS, TPU 相对 CPU, GPU 相对 CPU)
TPU_APPS = {
    "MLP0": ("MLP", 200, 200, 20e6, 0.61 / 2, 12.3, 41.0, 2.5),
    "MLP1": ("MLP", 168, 168, 5e6, 0.61 / 2, 9.7, 18.5, 0.3),
    "LSTM0": ("LSTM", 64, 64, 52e6, 0.29 / 2, 3.7, 3.5, 0.4),
    "LSTM1": ("LSTM", 96, 96, 34e6, 0.29 / 2, 2.8, 1.2, 1.2),
    "CNN0": ("CNN", 2888, 8, 8e6, 0.05 / 2, 86.0, 40.3, 1.6),
    "CNN1": ("CNN", 1750, 32, 100e6, 0.05 / 2, 14.1, 71.0, 2.7),
}


def roofline_tops(chip, intensity):
    """roofline 上界（TOPS，乘加算 2 次）：min(峰值, 2 × 每字节乘加 × 带宽)。"""
    return min(chip.peak_tops, 2 * intensity * chip.mem_gbs * 1e9 / 1e12)


# =============================================================================
# 3. GPU 的 SM 与 Tensor Core
# =============================================================================
@dataclass(frozen=True)
class GPUArch:
    """
    一代 NVIDIA 数据中心 GPU 的 SM 级参数（公开的架构白皮书 / 博客；FP16 稠密、不含稀疏加速）。

    【字段】
        sms：SM 个数；tc_per_sm：每个 SM 的 Tensor Core 个数；
        tc_fma_per_clk：每个 Tensor Core 每个周期做几次 FP16 乘加（FMA）；
        fp32_fma_per_sm：每个 SM 每个周期的 FP32 FMA（CUDA core）；
        clock_ghz：算峰值用的时钟；smem_kb：每个 SM 最多可用的共享内存；l2_mb；hbm_tbs：显存带宽
    """
    name: str
    sms: int
    tc_per_sm: int
    tc_fma_per_clk: int
    fp32_fma_per_sm: int
    clock_ghz: float
    smem_kb: int
    l2_mb: int
    hbm_tbs: float

    def tc_flops_per_sm_clk(self):
        return 2 * self.tc_per_sm * self.tc_fma_per_clk

    def tc_peak(self):
        """Tensor Core 的 FP16 稠密峰值 = SM 数 × 每 SM 每周期 FLOPs × 时钟。"""
        return self.sms * self.tc_flops_per_sm_clk() * self.clock_ghz * 1e9

    def fp32_peak(self):
        return self.sms * 2 * self.fp32_fma_per_sm * self.clock_ghz * 1e9


# V100：80 SM × 8 个第一代 Tensor Core，每个每周期 64 次 FP16 FMA（4×4×4），1.53 GHz → 125 TFLOPS
# A100：108 SM × 4 个第三代，每个每周期 256 次，1.41 GHz → 312 TFLOPS
# H100 SXM5：132 SM × 4 个第四代，每个每周期 512 次（Hopper 博客："每个 SM 的 MMA 速率是 A100 的 2 倍"），
#           1.83 GHz → 989 TFLOPS；每个 SM 128 个 FP32 单元
V100 = GPUArch("V100", 80, 8, 64, 64, 1.53, 96, 6, 0.9)
A100 = GPUArch("A100", 108, 4, 256, 64, 1.41, 164, 40, 1.555)
H100 = GPUArch("H100 SXM5", 132, 4, 512, 128, 1.83, 228, 50, 3.35)
GPU_GENS = [V100, A100, H100]

# 一条指令做多少次 FMA、要从寄存器 / 共享内存读多少个操作数（元素个数）——04 / 03 脚本的"指令与操作数账"
#   FFMA      ：一个 warp（32 个线程）一条指令，每个线程 1 次 FMA，读 3 个寄存器（a、b、c）
#   HMMA      ：mma.sync.m16n8k16（Ampere 起），一个 warp：A 16×16、B 16×8 从寄存器读，C/D 16×8 的累加器
#   WGMMA     ：wgmma.m64n256k16（Hopper），一个 warpgroup（4 个 warp）：A 64×16、B 16×256（B 直接从共享内存读），
#               累加器 64×256 留在寄存器里
INSTR_SHAPES = {
    "FFMA（CUDA core）": {"fma": 32, "operands": 3 * 32, "acc": 0},
    "HMMA m16n8k16（Ampere，一个 warp）": {"fma": 16 * 8 * 16, "operands": 16 * 16 + 16 * 8, "acc": 16 * 8},
    "WGMMA m64n256k16（Hopper，一个 warpgroup）": {"fma": 64 * 256 * 16, "operands": 64 * 16 + 16 * 256,
                                                "acc": 64 * 256},
}


def tile_intensity(bm, bn, elem_bytes=2):
    """
    一个 SM 计算 bm×bn 的输出块：每读进来 k 列 A（bm×k）和 k 行 B（k×bn），做 2·bm·bn·k 次运算。
    强度 = 2·bm·bn·k / ((bm + bn)·k·elem_bytes) = bm·bn / (bm + bn) FLOP/字节（elem_bytes = 2 时）。
    """
    return 2 * bm * bn / ((bm + bn) * elem_bytes)


# =============================================================================
# 4. 一个 SM 上的"读 → 算"流水线
# =============================================================================
def pipeline_sim(n_tiles, t_compute, t_transfer, latency, stages, jitter=0.0, seed=0, issue_ahead=True):
    """
    一个 SM 循环处理 n_tiles 个数据块：每块先从显存读进共享内存（等 latency + t_transfer），再用 Tensor Core 算（t_compute）。
    共享内存里有 stages 个缓冲区。

    【三种写法】
        stages = 1：同步（读完一块才算，算完才读下一块）——"加载 → 计算"完全串行；
        stages = 2：双缓冲（算第 i 块时读第 i+1 块）——cp.async（Ampere）的典型用法；
        stages ≥ 3：多级流水线——Hopper 上一个 producer warp 用 TMA 不停地发读请求，consumer warpgroup 只管算
                    （warp 专门化，warp specialization），最多有 stages 个块在路上或在缓冲区里。
    【模型】读请求在缓冲区空出来时发出（issue_ahead=False 时等计算完才发，等价于同步）；
        每个读请求的延迟 = latency × (1 + 随机抖动)，传输按带宽排队（同一时间只传一块，t_transfer）；
        计算按顺序一块一块做，需要的块到了、上一块算完才开始。
    【参数】jitter：延迟的随机抖动（指数分布的均值，相对 latency 的比例；0 表示固定延迟）
    【返回】字典：total（总时间）、busy（Tensor Core 忙的比例 = n_tiles·t_compute / total）
    【边界】stages ≥ 1。
    """
    rng = np.random.default_rng(seed)
    arrive = [0.0] * n_tiles
    issue_t = [0.0] * n_tiles
    link_free = 0.0
    comp_done = [0.0] * n_tiles
    for i in range(n_tiles):
        # 缓冲区：第 i 块要等第 i − stages 块算完（它的缓冲区才空出来）才能发读请求
        free_at = comp_done[i - stages] if i >= stages else 0.0
        if not issue_ahead and i >= 1:
            free_at = max(free_at, comp_done[i - 1])
        issue_t[i] = free_at
        lat = latency * (1.0 + (rng.exponential(jitter) if jitter > 0 else 0.0))
        start_xfer = max(issue_t[i] + lat, link_free)
        link_free = start_xfer + t_transfer
        arrive[i] = link_free
        start_c = max(arrive[i], comp_done[i - 1] if i >= 1 else 0.0)
        comp_done[i] = start_c + t_compute
    total = comp_done[-1]
    return {"total": total, "busy": n_tiles * t_compute / total}


def _selftest():
    rng = np.random.default_rng(0)
    for (B, R, C) in ((5, 3, 4), (1, 8, 8), (17, 4, 6)):
        X = rng.standard_normal((B, R))
        W = rng.standard_normal((R, C))
        Y, st = SystolicArray(R, C).run(X, W)
        assert np.max(np.abs(Y - X @ W)) < 1e-12
        assert st["cycles"] == systolic_cycles(B, R, C)
        assert st["useful_macs"] == B * R * C
    assert abs(TPU_V1.ridge() - 1353) < 1, TPU_V1.ridge()
    assert abs(H100.tc_peak() / 1e12 - 989.5) < 1 and abs(A100.tc_peak() / 1e12 - 311.9) < 1
    assert abs(V100.tc_peak() / 1e12 - 125.3) < 1
    t256 = tiled_matmul_time(1, 600, 600, 256, 256)["load_time"]
    t512 = tiled_matmul_time(1, 600, 600, 512, 512)["load_time"]
    assert 16e-6 < t256 < 19e-6 and 29e-6 < t512 < 33e-6
    print("hw_engines 自检通过：脉动阵列输出 = X·W、周期数 = B + R + C − 1；TPU 拐点 1353；"
          f"V100 / A100 / H100 峰值 {V100.tc_peak() / 1e12:.0f} / {A100.tc_peak() / 1e12:.0f} / {H100.tc_peak() / 1e12:.0f} TFLOPS；"
          f"600×600 读权重 {t256 * 1e6:.1f} µs（256²）vs {t512 * 1e6:.1f} µs（512²）")


if __name__ == "__main__":
    _selftest()