"""
01_layers_to_operators.py —— 一个神经网络层会被拆成哪些底层运算？（第 2 讲 §1）
=====================================================================

【这个脚本要回答的问题】
    1. 原课 slides 里的影评分类网络（Embedding → Linear → ReLU → Linear → Avg → Softmax），
       落到底层，只剩下五种运算：查表搬运（gather）、矩阵乘法（matmul）、逐元素一元运算（map）、
       逐元素二元运算（zip）、归约（reduce）。后四种正是原课 HW1 要写的四个 CUDA 内核。
    2. 每种运算的 FLOPs 和读写字节数各是多少？谁是"算得多"的，谁是"搬得多"的？
    3. 运算的顺序可以换吗？Linear 和 Avg 都是线性的，可以交换，计算量差 T 倍；ReLU 夹在中间就不行。

【怎么运行】
    python3 01_layers_to_operators.py      （只需要 numpy 和 torch；约 2 秒）

【说明】
    第 1 部分用 P3 的影评小语料（ReviewWorld）真的跑一遍，并和 PyTorch 的 nn 模块逐项对比；
    第 2 部分换成接近 GPT-2 small 的尺寸（d = 768、隐藏层 3072、32 条 × 128 token）重新数一遍，
    并用 H100 的公开规格（FP32 67 TFLOPS、3.35 TB/s）估算每个运算的理想耗时（roofline）。
"""

import os
import sys

import numpy as np
import torch

from console import print_table
from gpu_model import H100

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "prereq3"))
from nn_lab import ReviewWorld  # noqa: E402

torch.set_num_threads(1)


# =====================================================================
# 五种底层运算（NumPy 实现 + 记账）
# =====================================================================

class OpLog:
    """
    记录每一次底层运算的类型、FLOPs 和读写字节数。

    【约定】
        - 一个元素级的运算（加、乘、比较、exp、max……）算 1 次 FLOP（exp 在硬件上其实要好几条指令，这里只数"运算次数"）；
        - 矩阵乘法 (M×K)·(K×N) 算 2MKN 次（每个乘加 = 1 次乘 + 1 次加）；
        - 字节数假设每个运算单独执行：输入从显存读一遍、输出写回显存一遍（不融合，float32 = 4 字节，下标 int64 = 8 字节）。
    """

    def __init__(self):
        self.rows = []

    def add(self, layer, op, kind, flops, bytes_read, bytes_written, shape):
        self.rows.append(dict(layer=layer, op=op, kind=kind, flops=int(flops),
                              bytes=int(bytes_read + bytes_written), shape=shape))


def gather(log, layer, table, ids):
    """嵌入查表：out[b, t, :] = table[ids[b, t], :]。没有任何算术，纯数据搬运。"""
    out = table[ids]
    log.add(layer, "gather", "数据搬运", 0, out.nbytes + ids.nbytes, out.nbytes, out.shape)
    return out


def matmul(log, layer, x, W):
    """x 的最后一维和 W (K×N) 相乘。x 可以是 (…, K)，先展平成 (M, K) 再乘。"""
    M = int(np.prod(x.shape[:-1]))
    K, N = W.shape
    out = (x.reshape(M, K) @ W).reshape(*x.shape[:-1], N)
    log.add(layer, "matmul", "矩阵乘法", 2 * M * K * N, x.nbytes + W.nbytes, out.nbytes, out.shape)
    return out


def zip_(log, layer, name, fn, a, b):
    """逐元素二元运算（带广播）：out = fn(a, b)。"""
    out = fn(a, b)
    log.add(layer, f"zip({name})", "逐元素", out.size, np.asarray(a).nbytes + np.asarray(b).nbytes, out.nbytes, out.shape)
    return out


def map_(log, layer, name, fn, a):
    """逐元素一元运算：out = fn(a)。"""
    out = fn(a)
    log.add(layer, f"map({name})", "逐元素", out.size, a.nbytes, out.nbytes, out.shape)
    return out


def reduce_(log, layer, name, fn, a, axis):
    """沿 axis 归约（保留这一维，长度变成 1，方便之后广播）。"""
    out = fn(a, axis=axis, keepdims=True)
    log.add(layer, f"reduce({name})", "归约", a.size, a.nbytes, out.nbytes, out.shape)
    return out


def forward_ops(log, params, ids, mask, avg_before_linear2=False, avg_before_relu=False):
    """
    slides 里的网络，用五种底层运算一步步写出来。

    【参数】params：E（V×d）、W1（d×h）、b1（h）、W2（h×C）、b2（C）；ids：(B, T) 词编号；mask：(B, T)，真实 token 为 1
        avg_before_linear2 ：先平均再做第二个 Linear（代数上等价，见第 3 部分）
        avg_before_relu    ：先平均再做 ReLU（代数上不等价，用来做对照）
    【返回】(B, C) 的概率
    """
    m = mask[..., None].astype(np.float32)                               # (B, T, 1)

    def masked_avg(layer, x):
        s = reduce_(log, layer, "sum", np.sum, zip_(log, layer, "mul", np.multiply, x, m), axis=1)  # (B, 1, F)
        cnt = reduce_(log, layer, "sum", np.sum, m, axis=1)                                         # (B, 1, 1)
        return zip_(log, layer, "div", np.divide, s, cnt)

    x = gather(log, "Embedding", params["E"], ids)                       # (B, T, d)
    x = matmul(log, "Linear1", x, params["W1"])                           # (B, T, h)
    x = zip_(log, "Linear1", "add", np.add, x, params["b1"])
    if avg_before_relu:
        x = masked_avg("Avg", x)
    x = map_(log, "ReLU", "relu", lambda a: np.maximum(a, np.float32(0)), x)
    if avg_before_linear2 and not avg_before_relu:
        x = masked_avg("Avg", x)
    x = matmul(log, "Linear2", x, params["W2"])                           # (B, ·, C)
    x = zip_(log, "Linear2", "add", np.add, x, params["b2"])
    if not avg_before_linear2 and not avg_before_relu:
        x = masked_avg("Avg", x)                                          # (B, 1, C)
    x = x[:, 0, :]                                                        # (B, C)：去掉长度为 1 的那一维（只改形状，不搬数据）
    mx = reduce_(log, "Softmax", "max", np.max, x, axis=1)
    e = map_(log, "Softmax", "exp", np.exp, zip_(log, "Softmax", "sub", np.subtract, x, mx))
    s = reduce_(log, "Softmax", "sum", np.sum, e, axis=1)
    return zip_(log, "Softmax", "div", np.divide, e, s)


def init_params(V, d, h, C, seed=0):
    """随机初始化参数（float32）。"""
    rng = np.random.default_rng(seed)
    f = lambda *s: rng.standard_normal(s).astype(np.float32)
    return {"E": f(V, d) * np.float32(0.5), "W1": f(d, h) / np.float32(np.sqrt(d)), "b1": f(h) * np.float32(0.1),
            "W2": f(h, C) / np.float32(np.sqrt(h)), "b2": f(C) * np.float32(0.1)}


def torch_forward(params, ids, mask):
    """同一个网络用 PyTorch 的 nn 模块写一遍，作为参考答案。"""
    V, d = params["E"].shape
    h, C = params["W2"].shape
    emb, l1, l2 = torch.nn.Embedding(V, d), torch.nn.Linear(d, h), torch.nn.Linear(h, C)
    with torch.no_grad():
        emb.weight.copy_(torch.from_numpy(params["E"]))
        l1.weight.copy_(torch.from_numpy(params["W1"].T.copy()))          # nn.Linear 存的是 W 的转置（out × in）
        l1.bias.copy_(torch.from_numpy(params["b1"]))
        l2.weight.copy_(torch.from_numpy(params["W2"].T.copy()))
        l2.bias.copy_(torch.from_numpy(params["b2"]))
        x = l2(torch.relu(l1(emb(torch.from_numpy(ids)))))
        m = torch.from_numpy(mask.astype(np.float32))[..., None]
        x = (x * m).sum(1) / m.sum(1)
        return torch.softmax(x, dim=1).numpy()


def summarize(log, title, gpu=H100):
    """把 OpLog 打印成表，并按 H100 的 FP32 规格估算理想耗时。【返回】按运算类型汇总的字典。"""
    peak, bw = gpu.fp32_tflops * 1e12, gpu.mem_bw
    rows, total_f, total_b, total_t = [], 0, 0, 0.0
    by_kind = {}
    for r in log.rows:
        t = max(r["flops"] / peak, r["bytes"] / bw)
        ai = r["flops"] / r["bytes"] if r["bytes"] else 0
        rows.append([r["layer"], r["op"], r["kind"], str(tuple(r["shape"])), f"{r['flops']:,}", f"{r['bytes']:,}",
                     f"{ai:.2f}", f"{t * 1e6:.3f}"])
        total_f += r["flops"]
        total_b += r["bytes"]
        total_t += t
        k = by_kind.setdefault(r["kind"], [0, 0, 0.0, 0])
        k[0] += r["flops"]
        k[1] += r["bytes"]
        k[2] += t
        k[3] += 1
    print(f"  {title}")
    print_table(["层", "底层运算", "类型", "输出形状", "FLOPs", "读写字节", "FLOP/字节", "H100 理想 μs"], rows,
                align=["left", "left", "left", "left", "right", "right", "right", "right"])
    print()
    krows = [[k, v[3], f"{v[0] / total_f:.2%}", f"{v[1] / total_b:.2%}", f"{v[2] / total_t:.2%}"] for k, v in by_kind.items()]
    print_table(["运算类型", "次数", "占 FLOPs", "占字节", "占理想时间"], krows)
    return {k: {"n": v[3], "flops": v[0] / total_f, "bytes": v[1] / total_b, "time": v[2] / total_t} for k, v in by_kind.items()}, total_f


# =====================================================================
# 第 1 部分：在影评小语料上跑一遍
# =====================================================================

def part1_toy():
    print("=" * 78)
    print("第 1 部分：slides 里的影评分类网络 = 五种底层运算（在影评小语料上真的跑一遍）")
    print("=" * 78)
    world = ReviewWorld()
    sents, labels, kinds = world.dataset(4, seed=3)
    ids, mask = world.encode(sents)
    for s in sents:
        print("   ", " ".join(s))
    V, d, h, C = len(world.itos), 16, 32, 2
    params = init_params(V, d, h, C)
    log = OpLog()
    p = forward_ops(log, params, ids, mask)
    ref = torch_forward(params, ids, mask)
    print(f"\n  批大小 B = {ids.shape[0]}，补齐后的长度 T = {ids.shape[1]}，词表 V = {V}，d = {d}，h = {h}，类别 C = {C}")
    summary, _ = summarize(log, "逐个底层运算（一次前向）：")
    diff = float(np.abs(p - ref).max())
    print(f"\n  我们的五种运算拼出来的概率 vs PyTorch nn 模块：最大差 {diff:.2e}")
    print("\n观察 Observations：")
    print(f"  - 一次前向一共调用了 {len(log.rows)} 次底层运算，种类只有五种。其中 map / zip / reduce / matmul 正是 HW1 要写的"
          "四个 CUDA 内核（mapKernel、zipKernel、reduceKernel、MatrixMultiplyKernel）；gather 是纯数据搬运。")
    print(f"  - 即使在这个玩具网络里，两次矩阵乘法也占了 {summary['矩阵乘法']['flops']:.1%} 的 FLOPs；"
          f"另外 {summary['逐元素']['n'] + summary['归约']['n']} 次逐元素和归约运算几乎不做运算，读写的字节却占 "
          f"{summary['逐元素']['bytes'] + summary['归约']['bytes']:.1%}。第 2 部分换成真实尺寸再看。")
    print(f"  - 和 PyTorch 的结果相差 {diff:.1e}，是 float32 的舍入误差：两边做的是同一组运算，只是求和顺序等细节不同。\n")


# =====================================================================
# 第 2 部分：换成接近 GPT-2 small 的尺寸
# =====================================================================

def part2_realistic():
    print("=" * 78)
    print("第 2 部分：换成真实尺寸（d = 768、h = 3072、B = 32、T = 128），并按 H100 的 FP32 规格估算")
    print("=" * 78)
    B, T, V, d, h, C = 32, 128, 10000, 768, 3072, 3
    rng = np.random.default_rng(1)
    ids = rng.integers(1, V, size=(B, T))
    lens = rng.integers(T // 2, T + 1, size=B)
    mask = (np.arange(T)[None, :] < lens[:, None]).astype(np.float32)
    params = init_params(V, d, h, C, seed=2)
    log = OpLog()
    forward_ops(log, params, ids, mask)
    summary, total = summarize(log, f"H100：FP32 峰值 {H100.fp32_tflops:g} TFLOPS、显存带宽 {H100.mem_bw / 1e12:g} TB/s、"
                                    f"平衡点 {H100.ridge():.0f} FLOP/字节")
    mm = summary["矩阵乘法"]
    print("\n观察 Observations：")
    print(f"  - 矩阵乘法占了 {mm['flops']:.2%} 的 FLOPs、但只占 {mm['bytes']:.1%} 的读写字节；"
          f"逐元素 + 归约运算的 FLOPs 只有 {1 - mm['flops']:.2%}，字节却占 {summary['逐元素']['bytes'] + summary['归约']['bytes']:.1%}。")
    print(f"  - 按 roofline 估算，矩阵乘法占理想时间的 {mm['time']:.1%}，其余 {1 - mm['time']:.1%} 花在几乎不做运算的搬运上。"
          "逐元素运算的算术强度不到 1 FLOP/字节，远低于平衡点 20：它们的速度只由显存带宽决定。")
    print("  - 这就是第 1 讲 §4.4 的结论在一个具体网络上的样子：要让 GPU 跑得快，矩阵乘法要算得快（第 3–4 讲），"
          "小运算要少搬数据（算子融合，第 10–11 讲）。\n")
    return total


# =====================================================================
# 第 3 部分：运算顺序的代数改写
# =====================================================================

def part3_rewrite():
    print("=" * 78)
    print("第 3 部分：Linear 和 Avg 可以交换顺序，ReLU 不行")
    print("=" * 78)
    B, T, V, d, h, C = 32, 128, 10000, 768, 3072, 3
    rng = np.random.default_rng(1)
    ids = rng.integers(1, V, size=(B, T))
    lens = rng.integers(T // 2, T + 1, size=B)
    mask = (np.arange(T)[None, :] < lens[:, None]).astype(np.float32)
    params = init_params(V, d, h, C, seed=2)
    rows = []
    outs = {}
    for name, kw in (("原顺序：… → ReLU → Linear2 → Avg", {}),
                     ("改写：… → ReLU → Avg → Linear2", {"avg_before_linear2": True}),
                     ("错误改写：… → Avg → ReLU → Linear2", {"avg_before_relu": True})):
        log = OpLog()
        outs[name] = forward_ops(log, params, ids, mask, **kw)
        lin2 = sum(r["flops"] for r in log.rows if r["layer"] == "Linear2")
        tot = sum(r["flops"] for r in log.rows)
        rows.append([name, f"{lin2:,}", f"{tot:,}"])
    base = outs["原顺序：… → ReLU → Linear2 → Avg"]
    for r, name in zip(rows, outs):
        r.append(f"{float(np.abs(outs[name] - base).max()):.2e}")
    print_table(["顺序", "Linear2 的 FLOPs", "整个网络的 FLOPs", "与原顺序的最大差（概率）"], rows)
    d_ok = float(rows[1][3])
    d_bad = float(rows[2][3])
    ratio = int(rows[0][1].replace(",", "")) / int(rows[1][1].replace(",", ""))
    print("\n观察 Observations：")
    print(f"  - Avg 是加权求和（线性），Linear2 是 x·W + b（仿射），两者交换顺序数学上完全相等：结果只差 {d_ok:.1e}（float32 舍入），"
          f"Linear2 的计算量却少了 {ratio:.0f} 倍——原来对每个 token 各做一次，现在每条影评只做一次。")
    tot0, tot1 = int(rows[0][2].replace(",", "")), int(rows[1][2].replace(",", ""))
    print(f"  - 但整个网络只省了 {1 - tot1 / tot0:.2%}：大头是 Linear1（d × h 的矩阵乘法要对每个 token 做一次）。"
          "如果 Avg 能再往前挪到 Linear1 之前，Linear1 也能省 T 倍——可惜中间隔着 ReLU。")
    print(f"  - 把 Avg 再往前挪到 ReLU 之前就不行了：max(0, ·) 不是线性的，平均之后再取 ReLU ≠ 取 ReLU 之后再平均，"
          f"概率差到 {d_bad:.2f}。")
    print("  - 原课 HW2 的网络干脆把顺序定成\"先对嵌入取平均，再过 Linear → ReLU → Linear\"：那是另一个模型（表达能力不同），"
          "但便宜 T 倍。选模型结构时，计算量本身就是一个设计目标。")
    print("  - 这类\"不改变结果、只改变计算量\"的改写，是编译器和框架做图优化的基本手段（第 5 讲的计算图、第 12 讲的 XLA）。"
          "前提是先证明它在数学上等价——这需要知道每个运算的性质（线性、可交换、可结合）。\n")


def main():
    part1_toy()
    part2_realistic()
    part3_rewrite()


if __name__ == "__main__":
    main()
"""
02_gpu_server_and_specs.py —— GPU 服务器由什么组成？GPU 的算力从哪里来？（第 2 讲 §2–§3）
=====================================================================

【这个脚本要回答的问题】
    1. 原课 slides 的 GPU 对比表（A100 / H100 / B200）怎么读？哪些数字含稀疏加速？平衡点各是多少？
    2. 峰值算力能不能从"SM 数 × 核心数 × 频率"算出来？
    3. CPU 和 GPU 差在哪里？（slides 的对比：256 个线程 vs 10752 个核心）
    4. 一个 SM 上有多少寄存器、共享内存？为什么线程束之间切换"不花时间"？
    5. 在一台 8 × A6000 的服务器上训练时，数据走哪些路、每条路要多久？为什么梯度在 GPU 之间怎么走很重要？

【怎么运行】
    python3 02_gpu_server_and_specs.py     （只用标准库；不到 1 秒）

【说明】全部是按公开规格（标称值）算的"纸面数字"，没有任何本机测量。
"""

from console import print_table
from gpu_model import (A100, A6000, B200, EPYC_9754, GPUS, H100, LAB_SERVER_LINKS, SLIDES_LINEUP,
                       ring_allreduce_seconds, transfer_seconds)


def part1_lineup():
    print("=" * 78)
    print("第 1 部分：读懂 slides 的 GPU 对比表——稀疏 vs dense、平衡点")
    print("=" * 78)
    keys = ["FP64", "FP32", "TF32 Tensor", "BF16 Tensor", "FP8 Tensor", "FP4 Tensor"]
    rows = []
    for k in keys:
        row = [k]
        for g in ("A100", "H100", "B200"):
            v = SLIDES_LINEUP[g][k]
            row.append("—" if v is None else f"{v:,g}")
        rows.append(row)
    rows.append(["显存"] + [SLIDES_LINEUP[g]["mem"] for g in ("A100", "H100", "B200")])
    rows.append(["显存带宽"] + [SLIDES_LINEUP[g]["bw"] for g in ("A100", "H100", "B200")])
    print("  原课 slides 的表（TFLOPS，照抄）：")
    print_table(["", "A100", "H100", "B200"], rows)

    print("\n  本讲义用的 dense 标称值，以及由它算出的平衡点（峰值算力 ÷ 显存带宽）：")
    rows = []
    for g in GPUS:
        bf = g.tensor_dense_tflops.get("BF16")
        rows.append([g.name, f"{g.fp32_tflops:g}", f"{bf:g}", f"{g.tensor_dense_tflops.get('FP8', 0) or '—'}",
                     f"{g.mem_bw / 1e12:.3g}", f"{g.ridge():.0f}", f"{g.ridge('BF16'):.0f}"])
    print_table(["GPU", "FP32", "BF16 张量核心", "FP8 张量核心", "带宽 TB/s", "FP32 平衡点", "BF16 平衡点"], rows)
    print("\n观察 Observations：")
    print(f"  - slides 表里 H100 的 BF16 是 {SLIDES_LINEUP['H100']['BF16 Tensor']:,}，本讲义用 {H100.tensor_dense_tflops['BF16']:g}："
          "NVIDIA 产品页的张量核心数字默认是\"含 2:4 结构化稀疏\"的（每 4 个权重里有 2 个是 0，硬件跳过它们），"
          "正好是 dense 的 2 倍。普通的（稠密）矩阵乘法只能拿到 dense 那一半。读规格表要先看脚注。")
    print(f"  - 张量核心比普通 FP32 核心快一个数量级以上（H100：BF16 {H100.tensor_dense_tflops['BF16']:g} vs FP32 {H100.fp32_tflops:g}），"
          "所以训练都用 BF16 / FP8 的矩阵乘法；代价是平衡点从约 20 FLOP/字节升到约 300：越快的算力越难喂饱。")
    print(f"  - 每一代的算力涨得比带宽快：A100 → B200，BF16 算力涨了 {B200.tensor_dense_tflops['BF16'] / A100.tensor_dense_tflops['BF16']:.1f} 倍，"
          f"显存带宽只涨了 {B200.mem_bw / A100.mem_bw:.1f} 倍。访存受限的问题一代比一代严重。\n")


def part2_peak_from_cores():
    print("=" * 78)
    print("第 2 部分：峰值算力 = SM 数 × 每 SM 的核心数 × 2 × 频率")
    print("=" * 78)
    rows = []
    for g in (A6000, A100, H100):
        rows.append([g.name, g.sms, g.fp32_lanes_per_sm, f"{g.fp32_cores:,}", f"{g.boost_ghz:g}",
                     f"{g.fp32_peak_from_cores() / 1e12:.2f}", f"{g.fp32_tflops:g}"])
    print_table(["GPU", "SM 数", "每 SM 的 FP32 核心", "FP32 核心总数", "boost GHz", "算出来的 TFLOPS", "规格表 TFLOPS"], rows)
    print("\n观察 Observations：")
    print("  - 三块卡算出来的数和规格表一致（差别在频率的四舍五入）。\"2\"来自乘加指令 FMA：a·b + c 一条指令算两次运算。")
    print(f"  - A6000 的 {A6000.fp32_cores:,} 个核心就是 slides 写的 10752 cores。A100 每个 SM 只有 {A100.fp32_lanes_per_sm} 个 FP32 核心"
          f"（每个分区 16 个），一条线程束指令要分两个周期执行；它的算力主要放在张量核心上。")
    print("  - 所谓\"CUDA 核心\"不是 CPU 那样能独立取指令的核心，而是一条 32 宽的向量流水线里的一个\"车道\"（lane）。"
          "真正能独立调度指令的单位是每个 SM 的 4 个线程束调度器——这就是下面要讲的 SIMT。\n")


def part3_cpu_vs_gpu():
    print("=" * 78)
    print("第 3 部分：CPU vs GPU（slides：AMD EPYC 9754 vs RTX A6000）")
    print("=" * 78)
    c, g = EPYC_9754, A6000
    resident = g.sms * g.max_threads_per_sm
    rows = [
        ["硬件线程 / 核心", f"{c.threads}（{c.cores} 核 × 2）", f"{g.fp32_cores:,} 个 FP32 核心；最多同时驻留 {resident:,} 个线程"],
        ["频率", f"{c.base_ghz} GHz（基础）", f"{g.boost_ghz} GHz（boost）"],
        ["FP32 峰值（slides 的算法：每核每周期 1 次乘加）", f"{c.scalar_peak() / 1e9:,.0f} GFLOPS", f"{g.fp32_tflops * 1e3:,.0f} GFLOPS"],
        ["FP32 峰值（用满 CPU 的 SIMD：每核每周期 32 次）", f"{c.simd_peak() / 1e12:.1f} TFLOPS", f"{g.fp32_tflops:g} TFLOPS"],
        ["张量核心（BF16，dense）", "—", f"{g.tensor_dense_tflops['BF16']:g} TFLOPS"],
        ["内存带宽", f"{c.mem_bw() / 1e9:.0f} GB/s（12 通道 DDR5-4800）", f"{g.mem_bw / 1e9:.0f} GB/s（GDDR6）"],
        ["最后一级缓存", f"{c.l3_mb} MB L3", f"{g.l2_mb} MB L2"],
        ["功耗", f"{c.tdp_w:g} W", f"{g.tdp_w:g} W"],
        ["FP32 GFLOPS / 瓦（SIMD 口径）", f"{c.simd_peak() / 1e9 / c.tdp_w:.1f}", f"{g.fp32_tflops * 1e3 / g.tdp_w:.1f}"],
    ]
    print_table(["", c.name, g.name], rows, align=["left", "left", "left"])
    print("\n观察 Observations：")
    print(f"  - slides 的 576 GFLOPS 是按\"每核每周期一次乘加\"算的；CPU 其实也有 SIMD（AVX-512），用满时约 {c.simd_peak() / 1e12:.1f} TFLOPS。"
          f"同口径比较，A6000 的 FP32 只快约 {g.fp32_tflops * 1e12 / c.simd_peak():.1f} 倍；加上张量核心，差距是 "
          f"{g.tensor_dense_tflops['BF16'] * 1e12 / c.simd_peak():.0f} 倍。")
    print(f"  - 更根本的差别在设计取向：CPU 有 {c.l3_mb} MB 的大缓存、复杂的乱序执行和分支预测，目标是让少数线程的每一步都尽快完成"
          f"（延迟优先）；GPU 把更多的晶体管花在算术单元和寄存器上，同时驻留 {resident:,} 个线程，靠数量换吞吐（吞吐优先）。")
    print("  - 所以 GPU 只适合\"同一段程序在海量数据上各算各的\"——而神经网络的矩阵乘法、逐元素运算正好是这样。\n")


def part4_sm_resources():
    print("=" * 78)
    print("第 4 部分：一个 SM 里有什么——寄存器和共享内存为什么是稀缺资源")
    print("=" * 78)
    rows = []
    for g in (A6000, A100, H100):
        rf_kb = g.regs_per_sm * 4 / 1024
        rows.append([g.name, g.sms, f"{rf_kb:.0f} KB", f"{g.l1_smem_kb} KB", f"{g.smem_per_sm // 1024} KB",
                     g.max_threads_per_sm, f"{g.regs_per_sm // g.max_threads_per_sm}",
                     f"{g.sms * rf_kb / 1024:.1f} MB", f"{g.l2_mb:g} MB", f"{g.sms * g.max_threads_per_sm:,}"])
    print_table(["GPU", "SM 数", "每 SM 寄存器", "每 SM L1+共享", "其中共享内存最多", "每 SM 最多线程",
                 "满驻留时每线程寄存器", "全卡寄存器合计", "L2", "全卡最多驻留线程"], rows)
    print("\n观察 Observations：")
    rf_total = H100.sms * H100.regs_per_sm * 4 / 2 ** 20
    print(f"  - 每个 SM 有 64K 个 32 位寄存器（256 KB，4 个分区各 64 KB）。H100 全卡的寄存器加起来约 {rf_total:.1f} MB，"
          f"和全卡 L1 + 共享内存的总量（{H100.sms * H100.l1_smem_kb / 1024:.1f} MB）一样大，是 L2（{H100.l2_mb:g} MB）的 "
          f"{rf_total / H100.l2_mb:.0%}。寄存器是 GPU 上容量最大的一级片上存储之一，而 CPU 每个核心只有几百个寄存器。")
    print("  - 驻留在 SM 上的每个线程束都有自己的一份寄存器，常驻不动。所以线程束之间切换不需要像 CPU 的线程切换那样"
          "保存 / 恢复寄存器——调度器每个周期都可以换一个线程束发射，\"切换是瞬间的\"（slides 原话的意思）。")
    print(f"  - 代价是寄存器要分着用：H100 一个 SM 驻留满 {H100.max_threads_per_sm} 个线程时，每个线程平均只有 "
          f"{H100.regs_per_sm // H100.max_threads_per_sm} 个寄存器。内核用的寄存器越多，能驻留的线程就越少（04 号脚本的占用率）。\n")


def part5_server():
    print("=" * 78)
    print("第 5 部分：一台 8 × A6000 服务器里的数据流（slides 里作者实验室的配置）")
    print("=" * 78)
    links = LAB_SERVER_LINKS
    print_table(["连接", "带宽（单向）"], [[k, f"{v / 1e9:g} GB/s"] for k, v in links.items()])

    pcie = links["PCIe 4.0 x16（CPU ↔ GPU，或经 PCIe 的 GPU ↔ GPU）"]
    nvl = links["NVLink 桥（A6000 两两相连，单向）"]
    nvme = links["NVMe SSD（单块 PCIe 4.0 盘，顺序读）"]
    P = 1.3e9                                 # 一个 13 亿参数的模型
    grad = 4 * P                              # FP32 梯度
    tokens = 8 * 2048                         # 每张卡一步处理 8 条 × 2048 token
    batch_bytes = tokens * 4                  # token 编号（int32）
    weights16 = 2 * P
    print(f"\n  场景：数据并行训练一个 {P / 1e9:g}B 参数的模型，每张卡每步处理 {tokens:,} 个 token。")
    rows = [
        ["从 NVMe 读一步的训练数据（token 编号）", f"{batch_bytes / 1e3:,.0f} KB", f"{transfer_seconds(batch_bytes, nvme) * 1e3:.3f}"],
        ["CPU 内存 → GPU（PCIe）送这批 token", f"{batch_bytes / 1e3:,.0f} KB", f"{transfer_seconds(batch_bytes, pcie) * 1e3:.3f}"],
        ["开训前：从 4 块 NVMe 读 BF16 权重", f"{weights16 / 1e9:.1f} GB", f"{transfer_seconds(weights16, 4 * nvme) * 1e3:,.0f}"],
        ["GPU0 → GPU1 发一份 FP32 梯度（NVLink）", f"{grad / 1e9:.1f} GB", f"{transfer_seconds(grad, nvl) * 1e3:,.1f}"],
        ["GPU0 → GPU2 发一份 FP32 梯度（PCIe）", f"{grad / 1e9:.1f} GB", f"{transfer_seconds(grad, pcie) * 1e3:,.1f}"],
        ["8 卡环形 all-reduce 梯度（环上最慢的是 PCIe）", f"{grad / 1e9:.1f} GB", f"{ring_allreduce_seconds(grad, 8, pcie) * 1e3:,.1f}"],
        ["假想：8 卡全部用 NVLink 连成环", f"{grad / 1e9:.1f} GB", f"{ring_allreduce_seconds(grad, 8, nvl) * 1e3:,.1f}"],
    ]
    print_table(["数据流", "大小", "时间 ms"], rows)
    mfu = 0.4
    compute = 6 * P * tokens / (A6000.tensor_dense_tflops["BF16"] * 1e12 * mfu)
    ar = ring_allreduce_seconds(grad, 8, pcie)
    print(f"\n  对照：每张卡这一步的计算量 6·N·tokens = {6 * P * tokens:.3g} FLOPs，按 BF16 dense 峰值的 {mfu:.0%} 算要 "
          f"{compute * 1e3:,.0f} ms。")
    print("\n观察 Observations：")
    print(f"  - 训练时 CPU 送给 GPU 的只是 token 编号（每步 {batch_bytes / 1e3:.0f} KB），几乎不花时间；权重常驻显存。"
          "真正大的流量是 GPU 之间的梯度——每一步都要把整个模型大小的数据在卡之间交换一遍。")
    print(f"  - 梯度同步要 {ar * 1e3:.0f} ms，是这一步计算时间（{compute * 1e3:.0f} ms）的 {ar / compute:.0%}。"
          "这台机器只有两两之间有 NVLink 桥：环形 all-reduce 必然要走 PCIe，被最慢的一段卡住。"
          f"如果全部卡都用 NVLink 互连（像 DGX 服务器的 NVSwitch），同样的通信只要 {ring_allreduce_seconds(grad, 8, nvl) * 1e3:.0f} ms。")
    print("  - 这就是 slides 问\"为什么服务器的结构和我们有关\"的答案：数据并行的速度取决于梯度走哪条路。"
          "第 14–15 讲会讲怎样把通信和反向计算重叠起来，把这笔时间藏住。\n")


def main():
    part1_lineup()
    part2_peak_from_cores()
    part3_cpu_vs_gpu()
    part4_sm_resources()
    part5_server()


if __name__ == "__main__":
    main()
"""
03_simt_basics.py —— 在 SIMT 模拟器上跑第一个 CUDA 内核：线程编号、线程束、越界检查、合并访存（第 2 讲 §4）
=====================================================================

【这个脚本要回答的问题】
    1. 每个线程怎样知道自己该算哪个元素？i = blockIdx.x * blockDim.x + threadIdx.x 到底给出了什么？
    2. 线程块大小怎么选？块大小不是 32 的倍数会怎样？
    3. 三个最常见的错误：忘了越界检查、块数算少了、一个块放太多线程——各会发生什么？
    4. 元素比线程多时，一个线程处理多个元素：按"跨步循环"分（grid-stride loop）还是按"每人一段"分？
       CPU 上"每人一段"对缓存最友好，GPU 上却恰恰相反——为什么？

【怎么运行】
    python3 03_simt_basics.py       （只需要 numpy；约 2 秒）

【对应的 CUDA C】
    __global__ void vecAdd(const float* A, const float* B, float* C, int n) {
        int i = blockIdx.x * blockDim.x + threadIdx.x;
        if (i < n) C[i] = A[i] + B[i];
    }
    vecAdd<<<(n + 255) / 256, 256>>>(dA, dB, dC, n);
"""

import numpy as np

from console import print_table
from simt import BarrierError, Device, IllegalAddress, LaunchError, ceil_div, launch  # noqa: F401


# =====================================================================
# 内核
# =====================================================================

def vec_add(t, A, B, C, n):
    """向量加法：每个线程算一个元素。对应上面的 CUDA C。"""
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < n:
        C[i] = A[i] + B[i]


def vec_add_no_check(t, A, B, C, n):
    """忘了越界检查的版本：线程数多于元素数时，多出来的线程会访问数组之外的内存。"""
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    C[i] = A[i] + B[i]


def who_am_i(t, BLK, THR, WARP, n):
    """把每个线程的 blockIdx.x、threadIdx.x、块内线程束号写到第 i 个位置，用来画编号表。"""
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < n:
        BLK[i] = t.blockIdx.x
        THR[i] = t.threadIdx.x
        WARP[i] = t.threadIdx.x // 32


def vec_add_grid_stride(t, A, B, C, n):
    """
    跨步循环（grid-stride loop）：线程 g 处理 g、g + 总线程数、g + 2 × 总线程数……
        for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < n; i += blockDim.x * gridDim.x)
    同一时刻，一个线程束里的 32 个线程访问的是 32 个相邻的元素。
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    stride = t.blockDim.x * t.gridDim.x
    while i < n:
        C[i] = A[i] + B[i]
        i += stride


def vec_add_chunked(t, A, B, C, n):
    """
    每人一段：线程 g 处理 [g·k, (g+1)·k) 这一段连续的元素（k = ⌈n / 总线程数⌉）。
    同一时刻，一个线程束里相邻两个线程访问的元素相隔 k 个。
    """
    g = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    total = t.blockDim.x * t.gridDim.x
    k = -(-n // total)
    for i in range(g * k, min((g + 1) * k, n)):
        C[i] = A[i] + B[i]


# =====================================================================
# 第 1 部分：线程编号
# =====================================================================

def part1_indexing():
    print("=" * 78)
    print("第 1 部分：i = blockIdx.x * blockDim.x + threadIdx.x —— 每个线程的\"身份证号\"")
    print("=" * 78)
    dev = Device()
    n, block = 10, 4
    grid = ceil_div(n, block)
    BLK, THR, WARP = (dev.malloc(n, np.int32, name) for name in ("BLK", "THR", "WARP"))
    launch(who_am_i, grid, block, BLK, THR, WARP, n)
    blk, thr = dev.to_host(BLK), dev.to_host(THR)
    print(f"  n = {n}，blockDim.x = {block}，gridDim.x = ⌈{n}/{block}⌉ = {grid}：一共发射 {grid * block} 个线程")
    rows = [[i, int(blk[i]), int(thr[i]), f"{int(blk[i])} × {block} + {int(thr[i])}"] for i in range(n)]
    print_table(["i（负责的元素）", "blockIdx.x", "threadIdx.x", "算式"], rows)
    print(f"  另外 {grid * block - n} 个线程（blockIdx.x = {grid - 1}，threadIdx.x = {n - (grid - 1) * block}..{block - 1}）"
          f"算出的 i ≥ n，被 if (i < n) 挡在外面，什么也不做。")

    a = np.arange(n, dtype=np.float32)
    b = (10 * np.arange(n)).astype(np.float32)
    dA, dB, dC = dev.to_device(a, "A"), dev.to_device(b, "B"), dev.malloc(n, name="C")
    launch(vec_add, grid, block, dA, dB, dC, n)
    c = dev.to_host(dC)
    print(f"\n  vecAdd 的结果：{c.astype(int).tolist()}，和 NumPy 的 a + b 逐位相同：{np.array_equal(c, a + b)}")
    print("\n观察 Observations：")
    print("  - 内核代码只有一份，所有线程执行的是同一个函数；唯一不同的是内置变量 blockIdx、threadIdx 的值。"
          "每个线程用它们算出自己负责的数据——这就是 SIMT（Single Instruction, Multiple Threads）。")
    print("  - 编号是\"块号 × 块大小 + 块内编号\"，和 P1 §5 的行优先下标一模一样：把线程看成一个 gridDim × blockDim 的二维数组，"
          "i 就是它展平后的位置。")
    print("  - 块数要向上取整：(n + block − 1) / block。最后一个块往往不满，所以越界检查不能省（第 3 部分）。\n")


# =====================================================================
# 第 2 部分：块大小与线程束
# =====================================================================

def part2_block_size():
    print("=" * 78)
    print("第 2 部分：块大小怎么选——线程束、空闲线程、访存事务")
    print("=" * 78)
    dev = Device()
    n = 1000
    rng = np.random.default_rng(0)
    a = rng.standard_normal(n).astype(np.float32)
    b = rng.standard_normal(n).astype(np.float32)
    dA, dB = dev.to_device(a, "A"), dev.to_device(b, "B")
    rows = []
    res = {}
    for block in (32, 100, 128, 256, 1024):
        dC = dev.malloc(n, name="C")
        grid = ceil_div(n, block)
        s = launch(vec_add, grid, block, dA, dB, dC, n)
        ok = np.array_equal(dev.to_host(dC), a + b)
        dev.free(dC)
        res[block] = s
        rows.append([block, grid, s.threads, s.threads - n, s.warps, s.idle_lanes, s.ld_requests,
                     s.ld_sectors, f"{s.ld_sectors / s.ld_requests:.2f}", ok])
    print(f"  n = {n} 个 float32，C = A + B：")
    print_table(["块大小", "块数", "发射的线程", "i ≥ n 的线程", "线程束", "线程束里的空位", "读请求", "读事务",
                 "事务 / 请求", "结果正确"], rows)
    s100, s128 = res[100], res[128]
    print("\n观察 Observations：")
    print("  - 线程束是 32 个线程一组按块内编号切出来的。块大小 100 不是 32 的倍数：每个块切成 4 个线程束，最后一个只有 4 个线程，"
          f"其余 28 个位置空着（{s100.idle_lanes} 个空位）。这些位置照样占用调度和寄存器资源，所以块大小总是取 32 的倍数。")
    print(f"  - 块大小 100 时读事务一共 {s100.ld_sectors} 个，比 128 时（{s128.ld_sectors} 个）多 "
          f"{s100.ld_sectors / s128.ld_sectors - 1:.0%}。原因有两个：每个块的最后一个线程束只有 4 个线程，也要单独发一次请求"
          f"（读请求从 {s128.ld_requests} 个变成 {s100.ld_requests} 个）；1、3、5……号块从第 400、1200……字节开始"
          "（400 = 12 × 32 + 16），没有对齐到 32 字节的边界，这些块里一个满的线程束读的 128 字节要跨 5 个 sector。对齐也是合并访存的一部分（第 4 讲）。")
    print(f"  - 块大小是 32 的倍数时，读事务总数正好是 1000 × 4 字节 × 2 个数组 ÷ 32 字节 = {s128.ld_sectors}："
          "线程束里 32 个线程读相邻的 4 字节，合起来正好是 4 个完整的 sector，一个字节都不浪费——这就是\"合并访存\"（coalesced access）。\n")


# =====================================================================
# 第 3 部分：三个常见错误
# =====================================================================

def part3_bugs():
    print("=" * 78)
    print("第 3 部分：三个常见错误")
    print("=" * 78)
    dev = Device()
    n = 1000
    a = np.ones(n, dtype=np.float32)
    dA, dB = dev.to_device(a, "A"), dev.to_device(a, "B")

    dC = dev.malloc(n, name="C")
    try:
        launch(vec_add_no_check, ceil_div(n, 256), 256, dA, dB, dC, n)
        msg1 = "没有报错"
    except IllegalAddress as e:
        msg1 = str(e)
    print(f"  错误 1：忘了 if (i < n)  →  {msg1}")

    dC2 = dev.malloc(n, name="C2")
    grid = n // 256
    launch(vec_add, grid, 256, dA, dB, dC2, n)
    c = dev.to_host(dC2)
    wrong = int(np.sum(~(c == 2)))
    print(f"  错误 2：块数写成 n / 256 = {grid}（向下取整）  →  只发射了 {grid * 256} 个线程；"
          f"结果里有 {wrong} 个元素不对（还是 cudaMalloc 之后的垃圾值 {c[-1]}）——而且不会有任何报错")

    try:
        launch(vec_add, 1, 2048, dA, dB, dC, n)
        msg3 = "没有报错"
    except LaunchError as e:
        msg3 = str(e)
    print(f"  错误 3：一个块 2048 个线程  →  {msg3}")
    print("\n观察 Observations：")
    print("  - 错误 1 在真 GPU 上不一定报错：越界的线程可能读写到别的数组的内存，结果悄悄变错；如果碰巧越过了已分配的区域，"
          "才会报 \"an illegal memory access was encountered\"。模拟器对每次越界都报错，方便定位。")
    print(f"  - 错误 2 最危险：程序正常结束、结果大部分是对的，只有最后 {wrong} 个元素错。"
          "所以要向上取整：(n + 255) / 256，并且测试时一定要用不是块大小整数倍的 n。")
    print("  - 错误 3 是\"非法的发射配置\"：每个块最多 1024 个线程（x、y 维各最多 1024，z 维最多 64），"
          "超过时内核根本不会运行。真 CUDA 里要主动调用 cudaGetLastError() 才看得到这个错误。\n")


# =====================================================================
# 第 4 部分：一个线程处理多个元素
# =====================================================================

def part4_many_elements():
    print("=" * 78)
    print("第 4 部分：元素比线程多——跨步循环 vs 每人一段")
    print("=" * 78)
    dev = Device()
    n = 4096
    rng = np.random.default_rng(1)
    a = rng.standard_normal(n).astype(np.float32)
    b = rng.standard_normal(n).astype(np.float32)
    dA, dB = dev.to_device(a, "A"), dev.to_device(b, "B")
    grid, block = 4, 64
    rows = []
    res = {}
    for name, k in (("跨步循环（grid-stride）", vec_add_grid_stride), ("每人一段（chunked）", vec_add_chunked)):
        dC = dev.malloc(n, name="C")
        s = launch(k, grid, block, dA, dB, dC, n)
        ok = np.array_equal(dev.to_host(dC), a + b)
        res[name] = s
        rows.append([name, s.thread_loads, s.ld_requests, s.ld_sectors, f"{s.ld_sectors / s.ld_requests:.1f}",
                     s.unique_sectors, f"{s.sector_bytes / 1024:.0f} KiB", ok])
    print(f"  n = {n}，只发射 {grid} 个块 × {block} 个线程 = {grid * block} 个线程，每个线程处理 {n // (grid * block)} 个元素：")
    print_table(["分法", "线程读次数", "读请求", "读事务", "事务 / 请求", "不同的 sector", "事务搬运的字节", "结果正确"], rows)
    gs, ch = res["跨步循环（grid-stride）"], res["每人一段（chunked）"]
    print("\n观察 Observations：")
    print(f"  - 两种分法读的元素完全一样（都是 {gs.thread_loads} 次），用到的不同 sector 也一样多（{gs.unique_sectors} 个），"
          f"但每人一段的读事务是跨步循环的 {ch.ld_sectors / gs.ld_sectors:.0f} 倍。")
    print("  - 原因在\"同一时刻\"：线程束里的 32 个线程是一起执行同一条指令的。跨步循环时它们读 32 个相邻元素（4 个 sector）；"
          f"每人一段时相邻线程隔 {n // (grid * block)} 个元素（64 字节），32 个线程落在 32 个不同的 sector 里。")
    print("  - CPU 上正好相反：一个 CPU 线程顺着自己那一段往下读，缓存行和预取器都很高兴（P4 §1），所以并行库通常按段切分。"
          "在 GPU 上思考访存时，要想的是\"一个线程束在同一条指令里一起访问了哪些地址\"，而不是\"一个线程依次访问了哪些地址\"。\n")


def main():
    part1_indexing()
    part2_block_size()
    part3_bugs()
    part4_many_elements()


if __name__ == "__main__":
    main()
"""
04_occupancy_latency_hiding.py —— 线程块怎样分到 SM 上？为什么 GPU 要同时驻留几万个线程？（第 2 讲 §4.4–§4.6）
=====================================================================

【这个脚本要回答的问题】
    1. 一个 SM 上能同时驻留几个线程块？被线程数、块数、寄存器、共享内存中的哪一个卡住？（占用率）
    2. slides 的问题：24 个块 × 64 个线程 = 1536 个线程，能在 A6000 上同时运行吗？
    3. 网格里的块是一波一波分到 SM 上的；块数比 SM 数多一点点时，会浪费多少？（尾部效应）
    4. 一次读显存要约 500 个周期，GPU 怎样不让计算单元干等？需要多少个线程束才够？（延迟隐藏、Little 定律）

【怎么运行】
    python3 04_occupancy_latency_hiding.py     （只用标准库；约 2 秒）

【说明】第 1–3 部分是按公开规格算的；第 4 部分是一个线程束调度器的逐周期模拟（gpu_model.simulate_scheduler），
    模型很简化（见那里的注释），用来看趋势，不代表真实 GPU 的精确数字。
"""

import math

from console import print_table
from gpu_model import A100, A6000, H100, occupancy, scheduler_bound, simulate_scheduler, waves


def part1_occupancy():
    print("=" * 78)
    print("第 1 部分：占用率——一个 SM 上能同时驻留多少个线程束")
    print("=" * 78)
    gpus = (A6000, A100, H100)
    rows = []
    for tpb in (32, 64, 128, 256, 512, 768, 1024):
        row = [tpb]
        for g in gpus:
            o = occupancy(g, tpb)
            row.append(f"{o['blocks']} 块 / {o['warps']} 束 = {o['occupancy']:.0%}（{o['limiter']}）")
        rows.append(row)
    print("  ① 只改块大小（每线程 32 个寄存器、不用共享内存）：")
    print_table(["每块线程数"] + [f"{g.name}（最多 {g.max_warps_per_sm} 束、{g.max_blocks_per_sm} 块）" for g in gpus], rows,
                align=["right", "left", "left", "left"])

    rows = []
    for regs in (32, 40, 64, 96, 128, 255):
        o = occupancy(H100, 256, regs)
        rows.append([regs, o["blocks"], o["warps"], f"{o['occupancy']:.0%}", o["limiter"]])
    print("\n  ② H100，每块 256 个线程，只改每个线程用的寄存器数：")
    print_table(["每线程寄存器", "驻留块数", "驻留线程束", "占用率", "瓶颈"], rows)

    rows = []
    for smem_kb in (0, 8, 16, 32, 48, 64, 100, 160):
        o = occupancy(H100, 256, 32, smem_kb * 1024)
        rows.append([f"{smem_kb} KB", o["blocks"], o["warps"], f"{o['occupancy']:.0%}", o["limiter"]])
    print("\n  ③ H100，每块 256 个线程、每线程 32 个寄存器，只改每块申请的共享内存：")
    print_table(["每块共享内存", "驻留块数", "驻留线程束", "占用率", "瓶颈"], rows)
    o32 = occupancy(A6000, 32)
    o1024 = occupancy(A6000, 1024)
    print("\n观察 Observations：")
    print(f"  - 块太小会被\"每个 SM 最多几个块\"卡住：A6000 上每块 32 个线程时只能驻留 {o32['blocks']} 块 = {o32['warps']} 个线程束，"
          f"占用率 {o32['occupancy']:.0%}。块太大也可能浪费：A6000 每个 SM 最多 1536 个线程，1024 个线程的块只放得下 1 个，"
          f"占用率 {o1024['occupancy']:.0%}。128–512 是常见的选择。")
    print("  - 寄存器和共享内存是\"按块整份分配\"的：每线程寄存器从 32 涨到 40，驻留块数就从 8 掉到 6；"
          "共享内存每块 48 KB 时只放得下 4 块。写内核时多用一个数组、多开一个局部变量，都可能让占用率掉一截。")
    print("  - 占用率不是越高越好（第 4 讲的分块矩阵乘法就是用更多共享内存、更低的占用率换更少的访存），"
          "但太低时没有足够的线程束来掩盖访存延迟（第 4 部分）。\n")


def part2_slides_question():
    print("=" * 78)
    print("第 2 部分：slides 的问题——<<<dim3(2,3,4), dim3(2,4,8)>>> 能在 A6000 上同时运行吗？")
    print("=" * 78)
    blocks = 2 * 3 * 4
    tpb = 2 * 4 * 8
    g = A6000
    o = occupancy(g, tpb)
    slots = g.sms * o["blocks"]
    print(f"  网格 2 × 3 × 4 = {blocks} 个块，每块 2 × 4 × 8 = {tpb} 个线程（{math.ceil(tpb / 32)} 个线程束），一共 {blocks * tpb} 个线程。")
    print(f"  A6000：{g.sms} 个 SM，每个 SM 最多驻留 {o['blocks']} 个这样的块（瓶颈：{o['limiter']}），全卡最多 {slots:,} 个块同时驻留。")
    print(f"  {blocks} ≤ {slots:,}：能。实际上 {blocks} 个块只用到 {blocks} 个 SM（每个 SM 1 块），还有 {g.sms - blocks} 个 SM 空着。")
    print("\n观察 Observations：")
    print(f"  - \"能同时驻留\"和\"能同时执行\"是两回事：每个 SM 的 4 个调度器每周期各发射一条线程束指令，"
          f"一个 SM 上只有这 {math.ceil(tpb / 32)} 个线程束，4 个调度器有一半没事干。")
    print(f"  - 这个网格对 A6000 来说太小了：全卡能驻留 {g.sms * g.max_threads_per_sm:,} 个线程，它只用了 {blocks * tpb}（"
          f"{blocks * tpb / (g.sms * g.max_threads_per_sm):.1%}）。GPU 要\"喂饱\"，网格里的线程数要比全卡能驻留的线程数多好几倍。\n")


def part3_waves():
    print("=" * 78)
    print("第 3 部分：块怎样分到 SM 上——波次与尾部效应（H100，每个 SM 同时驻留 1 个块）")
    print("=" * 78)
    rows = []
    for nb in (66, 132, 133, 200, 264, 265, 1000, 10000):
        w, eff = waves(nb, H100.sms)
        rows.append([nb, w, f"{eff:.1%}"])
    print_table(["块数", "波数", "SM 的平均利用率"], rows)
    w1, e1 = waves(133, H100.sms)
    w2, e2 = waves(10000, H100.sms)
    print("\n观察 Observations：")
    print(f"  - 假设每个块耗时相同。133 个块要 {w1} 波：第二波只有 1 个块在算，其余 131 个 SM 干等，平均利用率只有 {e1:.1%}"
          "——和 66 个块（只用一半 SM）差不多。")
    print(f"  - 块数很多时尾巴被摊薄：10000 个块是 {w2} 波，利用率 {e2:.1%}。所以一个矩阵乘法切成多少块"
          "（块数是不是 SM 数的整数倍）会影响速度；cuBLAS 会针对形状挑选分块方式（第 4 讲）。")
    print("  - 块与块之间没有执行顺序的保证，也不能互相等待（没有跨块的 __syncthreads）。这是 CUDA 能在不同大小的 GPU 上"
          "透明扩展的原因：同一个网格，SM 多就少几波，SM 少就多几波。\n")


def part4_latency_hiding():
    print("=" * 78)
    print("第 4 部分：延迟隐藏——一个线程束调度器的逐周期模拟")
    print("=" * 78)
    g = H100
    L = 500
    bytes_per_cycle_per_sched = g.mem_bw / (g.sms * 4 * g.boost_ghz * 1e9)
    interval = round(128 / bytes_per_cycle_per_sched)
    print(f"  模型：访存延迟 {L} 周期（slides：全局内存约 500 周期）；每个线程束反复做\"读 128 字节 → c 条依赖这次数据的计算\"。")
    print(f"  带宽限制：H100 每个调度器分到的显存带宽 = {g.mem_bw / 1e12:g} TB/s ÷ ({g.sms} SM × 4) ÷ {g.boost_ghz} GHz"
          f" ≈ {bytes_per_cycle_per_sched:.2f} 字节/周期，即每 {interval} 个周期才能读一次 128 字节。")
    warps = (1, 2, 4, 8, 12, 16)
    rows = []
    sims = {}
    for c in (4, 16, 64):
        row = [c]
        for W in warps:
            r = simulate_scheduler(W, c, L, interval)
            sims[(c, W)] = r
            row.append(f"{r['utilization']:.0%}（{scheduler_bound(W, c, L, interval):.0%}）")
        rows.append(row)
    print("\n  调度器利用率（发射了指令的周期占比）；括号里是公式 min(1, W(1+c)/(L+c), (1+c)/间隔) 的值：")
    print_table(["每次读之后的计算指令 c"] + [f"{W} 个线程束" for W in warps], rows)

    needed = L / interval
    inflight = sims[(4, 16)]["loads_in_flight"]
    rate = sims[(4, 16)]["load_rate"]
    print("\n  Little 定律（P4 §2）：在路上的请求数 = 发射速率 × 延迟。")
    print(f"  c = 4、16 个线程束时，模拟测得平均 {inflight:.1f} 个读请求在路上，发射速率 {rate:.4f} 个/周期，"
          f"× {L} = {rate * L:.1f}（吻合）。")
    print(f"  要用满带宽，每个调度器需要同时有 {L} ÷ {interval} ≈ {needed:.1f} 个读请求在路上；整个 SM 是 {4 * needed:.0f} 个，"
          f"全卡 {4 * needed * g.sms:,.0f} 个 × 128 字节 ≈ {4 * needed * g.sms * 128 / 2 ** 20:.1f} MiB 的数据\"在路上\"。")
    print("\n观察 Observations：")
    print(f"  - 只有 1 个线程束时，调度器绝大部分时间在等数据（c = 4 时利用率 {sims[(4, 1)]['utilization']:.0%}）。"
          "线程束越多，一个在等数据的时候另一个就能发射：利用率几乎随线程束数线性上升——这就是\"用大量线程隐藏访存延迟\"。")
    print(f"  - 计算少的内核（c = 4）在 {math.ceil(needed)} 个左右的线程束时撞上带宽上限（{sims[(4, 16)]['utilization']:.0%}），"
          "再多线程束也没用：它是访存受限的。计算多的内核（c = 64）在 16 个线程束时接近 100%：计算受限。")
    print(f"  - H100 每个 SM 最多驻留 {g.max_warps_per_sm} 个线程束（每个调度器 {g.max_warps_per_sm // 4} 个），"
          f"和用满带宽需要的约 {needed:.0f} 个同一个量级。GPU 的设计就是按\"带宽 × 延迟\"来配线程数的；"
          "占用率太低（第 1 部分），访存受限的内核就连带宽都吃不满。\n")


def main():
    part1_occupancy()
    part2_slides_question()
    part3_waves()
    part4_latency_hiding()


if __name__ == "__main__":
    main()
"""
05_warp_divergence.py —— 线程束分化：同一个线程束里的线程走了不同的分支，会发生什么？（第 2 讲 §4.3）
=====================================================================

【这个脚本要回答的问题】
    1. slides 说线程束里的线程"每个周期执行一条共同的指令"，但又"可以独立分支"。两者怎么同时成立？
       ——分支不同时，线程束把两条路依次走一遍，走某条路时不走这条路的线程空等（被"屏蔽" masked）。
    2. 同样是一半线程走 if、一半走 else，按"奇偶"分和按"线程束"分，代价差多少？
    3. 循环次数不一样的线程放在同一个线程束里，会怎样？
    4. 为什么线程束是 32 个线程？——一个玩具模型：越宽越省取指令的开销，越窄越不怕分化。

【怎么运行】
    python3 05_warp_divergence.py      （只需要 numpy；约 2 秒）

【怎样计数】
    launch(..., count_instructions=True) 时，模拟器把"执行一行 Python 代码"算作一条指令；
    同一个线程束里执行同一行（第 k 次）的线程合起来算一次"发射"。
        SIMT 效率 = 线程实际执行的指令数 ÷ (发射次数 × 32)
    = 1 表示每次发射 32 个线程都在干活；= 0.5 表示平均一半的线程在空等。
"""

import numpy as np

from console import print_table
from simt import Device, launch


# =====================================================================
# 内核
# =====================================================================

def branch_kernel(t, OUT, mode, work):
    """
    一半线程走 if、一半走 else，两条路各做 work 次运算。

    mode = "parity" : 按 threadIdx.x 的奇偶分（每个线程束里都有一半走 if、一半走 else）
    mode = "warp"   : 按线程束分（threadIdx.x // 32 为偶数的整个线程束走 if，奇数的走 else）
    mode = "none"   : 所有线程都走 if（没有分化，对照）
    """
    i = t.threadIdx.x
    if mode == "parity":
        go_left = i % 2 == 0
    elif mode == "warp":
        go_left = (i // 32) % 2 == 0
    else:
        go_left = True
    x = i
    if go_left:
        for _ in range(work):
            x = x * 3 + 1
    else:
        for _ in range(work):
            x = x * 5 - 1
    OUT[t.blockIdx.x * t.blockDim.x + i] = x


def loop_kernel(t, OUT, TRIPS):
    """每个线程循环 TRIPS[i] 次（每个线程的工作量不同）。"""
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    x = 0
    for _ in range(TRIPS[i]):
        x = x + 1
    OUT[i] = x


# =====================================================================
# 第 1 部分：分支分化
# =====================================================================

def part1_branch():
    print("=" * 78)
    print("第 1 部分：if / else 分化——按奇偶分 vs 按线程束分")
    print("=" * 78)
    dev = Device()
    block, work = 128, 20
    OUT = dev.malloc(block, np.int64, "OUT")
    rows = []
    res = {}
    for mode, desc in (("none", "所有线程都走 if（无分化）"), ("warp", "按线程束分：整束走同一边"),
                       ("parity", "按奇偶分：每束里一半一半")):
        s = launch(branch_kernel, 1, block, OUT, mode, work, count_instructions=True)
        res[mode] = s
        rows.append([desc, s.thread_instructions, s.warp_issues, f"{s.simt_efficiency:.1%}"])
    print(f"  1 个块 × {block} 个线程（{block // 32} 个线程束），两条分支各循环 {work} 次：")
    print_table(["分法", "线程执行的指令", "线程束发射次数", "SIMT 效率"], rows)
    n, w, p = res["none"], res["warp"], res["parity"]
    print("\n观察 Observations：")
    print(f"  - 按线程束分时，每个线程束内部步调一致，发射次数和无分化时一样（{w.warp_issues} vs {n.warp_issues}），效率 {w.simt_efficiency:.0%}。"
          "虽然整个块里一半线程走 if、一半走 else，但分化只发生在线程束之间，线程束之间本来就是独立调度的，不要紧。")
    print(f"  - 按奇偶分时，每个线程束都要把 if 和 else 两条路依次走一遍：发射次数涨到 {p.warp_issues}（{p.warp_issues / n.warp_issues:.2f} 倍），"
          f"走 if 时奇数号线程空等、走 else 时偶数号线程空等，效率降到 {p.simt_efficiency:.0%}。"
          f"线程执行的指令几乎一样多（{p.thread_instructions} vs {n.thread_instructions}，差别来自选择分法的那几行 if），时间却几乎翻倍。")
    print("  - 不到正好 50%，是因为分支前后的公共代码（算下标、写结果）大家一起执行，不受分化影响。"
          "写内核时尽量让分支条件\"按线程束对齐\"，或者把分支改写成不分支的算术（例如 x * mask）。\n")


# =====================================================================
# 第 2 部分：循环次数不同
# =====================================================================

def part2_loops():
    print("=" * 78)
    print("第 2 部分：工作量不均——一个线程束要等它里面最慢的那个线程")
    print("=" * 78)
    dev = Device()
    n = 128
    rng = np.random.default_rng(0)
    trips = rng.integers(0, 32, size=n)
    rows = []
    res = {}
    for name, arr in (("随机排列（线程 i 循环 trips[i] 次）", trips), ("按工作量排序后再分配", np.sort(trips)),
                      ("每个线程都循环平均次数", np.full(n, int(round(trips.mean())))) ):
        OUT = dev.malloc(n, np.int64, "OUT")
        TR = dev.to_device(arr.astype(np.int64), "TRIPS")
        s = launch(loop_kernel, n // 32, 32, OUT, TR, count_instructions=True)
        res[name] = (s, arr)
        warp_max = [int(arr[k:k + 32].max()) for k in range(0, n, 32)]
        rows.append([name, int(arr.sum()), sum(warp_max), s.thread_instructions, s.warp_issues, f"{s.simt_efficiency:.1%}"])
    print(f"  {n} 个线程（4 个线程束），每个线程循环 0–31 次中的某个次数：")
    print_table(["分配方式", "总循环次数", "Σ 每个线程束的最大次数", "线程执行的指令", "线程束发射次数", "SIMT 效率"], rows)
    print("\n观察 Observations：")
    s_rand, s_sort = res["随机排列（线程 i 循环 trips[i] 次）"][0], res["按工作量排序后再分配"][0]
    print(f"  - 线程束的循环要转到它里面次数最多的那个线程做完为止：随机排列时，每个线程束都碰上一个接近 31 次的线程，"
          f"效率只有 {s_rand.simt_efficiency:.0%}。")
    print(f"  - 把工作量相近的线程放进同一个线程束（排序），总工作量不变，发射次数从 {s_rand.warp_issues} 降到 {s_sort.warp_issues}，"
          f"效率升到 {s_sort.simt_efficiency:.0%}。")
    print("  - LLM 系统里这种\"长短不一\"到处都是：一个批里的序列长度不同、稀疏矩阵每一行的非零元个数不同（第 4 讲的 CSR）。"
          "按长度排序 / 分桶、把长任务拆碎，是 GPU 上常见的对策（第 22 讲的连续批处理也是同一个问题的系统级版本）。\n")


# =====================================================================
# 第 3 部分：为什么是 32
# =====================================================================

def simt_model(W, F, p, len_if, len_else):
    """
    玩具模型：宽度为 W 的线程束执行一段带一个分支的程序，返回总效率。

    【假设】每个线程独立地以概率 p 走 if（长 len_if 条指令），否则走 else（长 len_else 条）；
        线程束里只要有一个线程走某条路，这条路就要整束走一遍；每发射一条线程束指令，
        取指、译码、调度的固定开销相当于 F 个"车道·周期"。
    【返回】(每束的期望发射次数, 车道利用率, 开销占比, 总效率)
        总效率 = 线程真正需要的指令数 ÷ (发射次数 × W + 发射次数 × F)
    """
    p_if = 1 - (1 - p) ** W
    p_else = 1 - p ** W
    issues = len_if * p_if + len_else * p_else
    useful = W * (p * len_if + (1 - p) * len_else)
    lane = issues * W
    over = issues * F
    return issues, useful / lane, over / (lane + over), useful / (lane + over)


def part3_why_32():
    print("=" * 78)
    print("第 3 部分：为什么线程束是 32 个线程？——一个玩具模型")
    print("=" * 78)
    F = 4
    print(f"  假设：每发射一条线程束指令，取指、译码、调度的固定开销相当于 F = {F} 个\"车道·周期\"（见 simt_model 的注释）。")
    print("  场景 A：没有分支（神经网络里的大多数内核：逐元素运算、矩阵乘法），每个线程 10 条指令。")
    print("  场景 B：有一个罕见但很长的分支：每个线程以 1% 的概率走 if（100 条指令），否则走 else（10 条）。")
    widths = (1, 2, 4, 8, 16, 32, 64, 128)
    rows = []
    eff_a, eff_b = {}, {}
    for W in widths:
        _, _, ovh_a, ea = simt_model(W, F, 0.0, 0, 10)
        iss_b, lane_b, _, eb = simt_model(W, F, 0.01, 100, 10)
        eff_a[W], eff_b[W] = ea, eb
        rows.append([W, f"{ovh_a:.1%}", f"{ea:.1%}", f"{iss_b:.1f}", f"{lane_b:.1%}", f"{eb:.1%}"])
    print_table(["线程束宽度 W", "A：开销占比", "A：总效率", "B：每束期望发射次数", "B：车道利用率", "B：总效率"], rows)
    best_b = max(eff_b, key=eff_b.get)
    print("\n观察 Observations：")
    print(f"  - 场景 A 里越宽越好（W = 128 时 {eff_a[128]:.0%}）：没有分化，宽度只用来分摊每条指令的固定开销。"
          "宽度为 1（像 CPU 的标量核）时，开销比干活还多。")
    print(f"  - 场景 B 里有一个最优宽度（这组假设下是 W = {best_b}，总效率 {eff_b[best_b]:.0%}；W = 128 时只剩 {eff_b[128]:.0%}）："
          "线程束越宽，里面\"至少有一个线程走了那条长路\"的概率越高，整束都要陪它走。最优宽度取决于分支有多罕见、有多长。")
    print("  - 所以线程束宽度是一个工程折中，而 GPU 押注的是场景 A——数据并行、少分支的程序。32 还有别的理由："
          "32 个线程 × 4 字节 = 128 字节，正好是一个缓存行（合并访存以它为单位，03 号脚本）；"
          "NVIDIA 从 2006 年的 G80 起一直用 32；AMD 的 GPU 用过 64（GCN 的 wavefront），新架构同时支持 32 和 64。"
          "这个玩具模型的具体数字取决于假设的 F 和分支概率，不能当真。\n")


def main():
    part1_branch()
    part2_loops()
    part3_why_32()


if __name__ == "__main__":
    main()
"""
console.py —— 在终端里打印对齐的表格（中英文混排也能对齐）

【为什么需要它】
    终端里一个汉字占两个英文字母的宽度，Python 的 f"{s:10s}" 却按"字符个数"补空格，
    所以中英混排的表格会歪掉。这里的 pad() 按"显示宽度"补空格。

【用法】
    from console import pad, print_table
    print_table(["模型", "困惑度"], [["trigram", "3.21"], ["RNN", "2.60"]])
"""

import unicodedata


def display_width(text):
    """
    一个字符串在终端里占多宽：东亚宽字符（汉字、全角标点）算 2，其余算 1。

    【Python 小知识】unicodedata.east_asian_width(c) 返回 'W'（宽）、'F'（全角）、
        'Na'（窄）等标记；W 和 F 在终端里占两格。
    """
    return sum(2 if unicodedata.east_asian_width(c) in ("W", "F") else 1 for c in str(text))


def pad(text, width, align="left"):
    """
    把 text 补空格到显示宽度 width。

    【参数】text：任意对象（会先转成字符串）；width：目标宽度；align："left" 或 "right"
    【返回】补好空格的字符串
    【边界】text 已经比 width 宽时原样返回（不截断，宁可表格歪一点也不丢信息）。
    """
    s = str(text)
    gap = max(0, width - display_width(s))
    return s + " " * gap if align == "left" else " " * gap + s


def print_table(headers, rows, align=None):
    """
    打印一张表。

    【参数】
        headers : 表头列表
        rows    : 每行是一个列表，长度和 headers 相同；元素可以是任意对象
        align   : 每列的对齐方式列表（"left"/"right"）；None 表示第一列左对齐、其余右对齐
    【边界】rows 为空时只打印表头。
    """
    cols = len(headers)
    if align is None:
        align = ["left"] + ["right"] * (cols - 1)
    cells = [[str(h) for h in headers]] + [[str(c) for c in r] for r in rows]
    widths = [max(display_width(row[j]) for row in cells) for j in range(cols)]
    for i, row in enumerate(cells):
        print("  " + "  ".join(pad(row[j], widths[j], align[j]) for j in range(cols)))
        if i == 0:
            print("  " + "  ".join("-" * widths[j] for j in range(cols)))
"""
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
"""
simt.py —— 用 Python 写的 SIMT GPU 模拟器：按 CUDA 的方式执行内核，并把"为什么快 / 为什么慢"数出来
=====================================================================
（第 2 讲引入；第 3、4 讲大量使用；第 10、11、21 讲写融合内核时还会用）

【这个文件要回答的问题】
    没有 NVIDIA GPU，怎样学 CUDA？我们把 CUDA 的执行模型用 Python 搭一遍：
        - 内核（kernel）是一个普通的 Python 函数，第一个参数 t 提供 CUDA 的内置变量
          t.threadIdx、t.blockIdx、t.blockDim、t.gridDim（都有 .x .y .z）；
        - launch(kernel, grid, block, 参数...) 相当于 CUDA 的 kernel<<<grid, block>>>(参数...)：
          模拟器为网格里的每一个线程调用一次这个函数；
        - 显存是 Device 上分配的 DeviceArray（相当于 cudaMalloc 得到的设备指针），
          内核里对它的每一次读（A[i]）和写（A[i] = v）都会被记录下来；
        - 共享内存用 t.shared(...) 申请；__syncthreads() 写成 `yield t.syncthreads()`。
    模拟器不追求快（它比真 GPU 慢上百万倍），追求的是"能数"：
        每个线程束（warp）的每条访存指令产生了几个 32 字节的内存事务（合并访存）、
        共享内存有没有 bank 冲突、线程束里有多少线程在空转（分化）、有没有数据竞争。

【本文件提供】
    常量         : WARP_SIZE、SECTOR_BYTES、NUM_BANKS、BANK_BYTES、MAX_THREADS_PER_BLOCK 等
    Dim3 / dim3  : 三维的网格 / 线程块大小（对应 CUDA 的 dim3）
    Device       : 一块模拟的 GPU：malloc / memcpy / free（即 cudaMalloc / cudaMemcpy / cudaFree）、拷贝记账、泄漏检查
    DeviceArray  : 显存里的一维数组（相当于设备指针）：A[i] 读、A[i] = v 写、A.atomic_add(i, v)、A.atomic_max(i, v)
    launch       : 发射内核，返回 LaunchStats（线程数、访存请求、内存事务、bank 冲突、指令发射、竞争……）
    异常         : SimError 及其子类 IllegalAddress（越界）、LaunchError（配置非法）、BarrierError（栅栏用错）、MemcpyError

【模拟的规则（以及它和真 GPU 的差别）】
    1. 线程束：块内线程按 CUDA 的线性编号 x + y·Bx + z·Bx·By 排队，每 32 个一组成为一个线程束。
    2. 执行顺序：一个块一个块地执行；块内每个线程依次运行，直到结束或遇到 __syncthreads()；
       块内全部线程都到达栅栏后，再一起进入下一段（"一个栅栏周期"）。
       真 GPU 上线程是并发的、执行顺序不确定——所以"在模拟器上结果对"不能证明内核没有数据竞争，
       这就是模拟器另外做竞争检测的原因（规则 6）。
    3. 合并访存：同一个线程束里，执行"同一条访存指令"（同一行代码对同一个数组的第 k 次访问）的线程放在一起看，
       它们访问的地址落在多少个不同的 32 字节段（sector）里，这条指令就产生多少个内存事务。
       32 个线程各读相邻的 4 字节 → 128 字节 → 4 个事务（最理想）；各隔很远 → 32 个事务。
    4. bank 冲突：共享内存分成 32 个 bank，每个 bank 4 字节宽，连续的 4 字节字依次落在 bank 0、1、……、31、0……
       同一条指令里，落在同一个 bank 的不同字要排队，排得最长的那个 bank 决定这条指令要走几轮（wavefront）。
       多个线程读同一个字是广播，不算冲突。（按 4 字节字计算；8 字节类型的细节这里不模拟。）
    5. 指令计数（可选，count_instructions=True）：每执行一行 Python 代码算一条指令；
       同一个线程束里执行"同一行的第 k 次"的线程，合起来算一次"发射"（issue）。
       发射次数 × 32 和"线程实际执行的指令数"之比，就是线程束分化造成的浪费。
    6. 竞争检测：同一个栅栏周期里，一个线程写了共享内存的某个位置、另一个线程又读或写了它 → 记为一次竞争。
       全局内存的竞争检测较慢，默认关闭（check_global_races=True 打开）。原子操作不参与竞争检测。
    7. 越界一律报错（IllegalAddress）。真 GPU 上越界不一定报错（可能悄悄读到别的数据），模拟器更严格；
       共享内存还会按每一维分别检查下标。
    8. 共享内存不会被初始化：浮点数组里填 NaN、整数数组填 -999999，读到"没写过的值"一眼就能看出来。
    9. 主机代码不能直接读写 DeviceArray（真 CUDA 里这样做通常会段错误）：要用 Device.memcpy 拷来拷去。

【用法】
    dev = Device()
    dA = dev.to_device(np.arange(8, dtype=np.float32))      # cudaMalloc + cudaMemcpy(HostToDevice)
    dC = dev.malloc(8)                                       # cudaMalloc
    def scale(t, A, C, n):                                   # __global__ void scale(float* A, float* C, int n)
        i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
        if i < n:
            C[i] = A[i] * 2
    stats = launch(scale, grid=1, block=8, *(dA, dC, 8))     # 也可以写 launch(scale, 1, 8, dA, dC, 8)
    print(dev.to_host(dC), stats.summary())
"""

import inspect
import sys
from collections import defaultdict, namedtuple

import numpy as np

# =====================================================================
# 0. 硬件常量（NVIDIA 近几代 GPU 的通用值）
# =====================================================================

WARP_SIZE = 32                         # 一个线程束的线程数
SECTOR_BYTES = 32                      # 全局内存事务的粒度：一个 sector 32 字节（缓存行 128 字节 = 4 个 sector）
NUM_BANKS = 32                         # 共享内存的 bank 数
BANK_BYTES = 4                         # 每个 bank 每轮能提供 4 字节
MAX_THREADS_PER_BLOCK = 1024           # 每个线程块最多 1024 个线程
MAX_BLOCK_DIM = (1024, 1024, 64)       # 线程块每一维的上限
MAX_GRID_DIM = (2 ** 31 - 1, 65535, 65535)   # 网格每一维的上限
DEFAULT_SHARED_LIMIT = 48 * 1024       # 每个块静态共享内存的默认上限（48 KiB；更大需要显式申请，见第 4 讲）
DEVICE_ALIGNMENT = 256                 # cudaMalloc 返回的地址至少按 256 字节对齐

_THIS_FILE = __file__


# =====================================================================
# 1. dim3 与异常
# =====================================================================

class Dim3(namedtuple("Dim3", ["x", "y", "z"])):
    """
    三维大小 / 三维下标，对应 CUDA 的 dim3 / uint3。

    【Python 小知识】namedtuple 生成一个"带字段名的元组"：d.x 和 d[0] 是同一个值，元组不可修改。
        这里继承它只是为了加一个 count 属性。__slots__ = () 表示不额外占内存。
    """
    __slots__ = ()

    @property
    def count(self):
        """x·y·z：网格里的块数，或线程块里的线程数。"""
        return self.x * self.y * self.z


def dim3(x=1, y=1, z=1):
    """和 CUDA 的 dim3(x, y, z) 一样：没给的维度默认是 1。"""
    return Dim3(int(x), int(y), int(z))


def as_dim3(v):
    """
    把 int、(x,)、(x, y)、(x, y, z) 或 Dim3 统一成 Dim3。

    【例子】as_dim3(256) → Dim3(256, 1, 1)；as_dim3((16, 16)) → Dim3(16, 16, 1)
    【边界】长度超过 3 的元组报 ValueError。
    """
    if isinstance(v, Dim3):
        return v
    if isinstance(v, (int, np.integer)):
        return Dim3(int(v), 1, 1)
    v = tuple(int(a) for a in v)
    if not 1 <= len(v) <= 3:
        raise ValueError(f"dim3 只能有 1–3 维，收到 {v}")
    return Dim3(*(v + (1,) * (3 - len(v))))


class SimError(Exception):
    """模拟器检测到的错误的基类。"""


class IllegalAddress(SimError):
    """越界访问（对应 CUDA 的 cudaErrorIllegalAddress / "an illegal memory access was encountered"）。"""


class LaunchError(SimError):
    """发射配置非法（对应 cudaErrorInvalidConfiguration 等），例如一个块超过 1024 个线程、共享内存超限。"""


class BarrierError(SimError):
    """__syncthreads() 用错：同一个块里有的线程到了栅栏、有的已经结束。真 GPU 上这是未定义行为（常见后果是卡死）。"""


class MemcpyError(SimError):
    """cudaMemcpy 的参数不对：方向和指针类型不符、字节数超过数组大小等。"""


SYNC = object()          # __syncthreads() 的"信号"：内核里写 yield t.syncthreads()，模拟器收到它就让线程在栅栏处等


# =====================================================================
# 2. 显存：Device 与 DeviceArray
# =====================================================================

HostToDevice = "HostToDevice"           # 对应 cudaMemcpyHostToDevice
DeviceToHost = "DeviceToHost"           # 对应 cudaMemcpyDeviceToHost
DeviceToDevice = "DeviceToDevice"       # 对应 cudaMemcpyDeviceToDevice


class DeviceArray:
    """
    显存里的一段一维数组，相当于 CUDA 里 cudaMalloc 得到的设备指针（加上它的长度和类型）。

    【属性】name（名字，用于统计和报错）、dtype、n（元素个数）、base（起始字节地址）、nbytes
    【在内核里怎样用】
        x = A[i]           读第 i 个元素（i 必须是整数；越界抛 IllegalAddress）
        A[i] = v           写
        A.atomic_add(i, v) 原子地加上 v，返回旧值（对应 atomicAdd）
        A.atomic_max(i, v) 原子地取最大值，返回旧值（对应 atomicMax）
    【边界】在内核之外（主机代码里）读写会抛 SimError：主机只能用 Device.memcpy 访问显存。
        释放之后再访问也会报错（对应"使用已释放的指针"）。
    """

    __slots__ = ("device", "name", "dtype", "n", "base", "itemsize", "nbytes", "_data", "freed")

    def __init__(self, device, n, dtype, name, base):
        self.device = device
        self.name = name
        self.dtype = np.dtype(dtype)
        self.n = int(n)
        self.base = base
        self.itemsize = self.dtype.itemsize
        self.nbytes = self.n * self.itemsize
        # 显存刚分配时内容是"垃圾"；这里填一个容易认出来的值
        self._data = np.full(self.n, np.nan if self.dtype.kind == "f" else -999999, dtype=self.dtype)
        self.freed = False

    def __len__(self):
        return self.n

    def __repr__(self):
        return f"<DeviceArray {self.name}: {self.n} × {self.dtype}, 地址 0x{self.base:x}>"

    def _check(self, i):
        """检查下标：必须是整数、在 [0, n) 里、数组没有被释放、而且是在内核里访问。返回 int 下标。"""
        run = _RUN
        if run is None:
            raise SimError(f"主机代码不能直接读写显存数组 {self.name}：请用 cudaMemcpy（Device.memcpy / to_host）")
        if self.freed:
            raise IllegalAddress(f"访问已经 cudaFree 的数组 {self.name}")
        if not isinstance(i, (int, np.integer)):
            raise TypeError(f"{self.name}[...] 的下标必须是整数（C 指针只能按整数偏移），收到 {type(i).__name__}")
        i = int(i)
        if not 0 <= i < self.n:
            raise IllegalAddress(f"线程 {run.where()} 越界访问 {self.name}[{i}]（合法范围 0..{self.n - 1}）")
        return i

    def __getitem__(self, i):
        i = self._check(i)
        _RUN.access("global", self, "ld", i, self.base + i * self.itemsize, self.itemsize)
        return self._data[i]

    def __setitem__(self, i, v):
        i = self._check(i)
        _RUN.access("global", self, "st", i, self.base + i * self.itemsize, self.itemsize)
        self._data[i] = v

    def atomic_add(self, i, v):
        """原子加：old = A[i]; A[i] = old + v; 返回 old。模拟器里线程依次执行，天然是原子的；这里主要是计数。"""
        i = self._check(i)
        _RUN.access("global", self, "atom", i, self.base + i * self.itemsize, self.itemsize)
        old = self._data[i]
        self._data[i] = old + v
        return old

    def atomic_max(self, i, v):
        """原子最大值：old = A[i]; A[i] = max(old, v); 返回 old。"""
        i = self._check(i)
        _RUN.access("global", self, "atom", i, self.base + i * self.itemsize, self.itemsize)
        old = self._data[i]
        if v > old:
            self._data[i] = v
        return old


class Device:
    """
    一块模拟的 GPU：负责显存的分配与释放、主机 ↔ 设备的拷贝，以及拷贝时间的记账。

    【参数】
        name           : 名字
        pcie_bandwidth : 主机 ↔ 设备拷贝的带宽（字节/秒），默认 32e9（PCIe 4.0 x16 的标称值，原课 slides 同）
        pcie_latency   : 每次拷贝的固定开销（秒），默认 10 微秒（量级）
    【记账】copies（次数）、h2d_bytes、d2h_bytes、copy_seconds（按 α + 字节数 / 带宽 估算的拷贝时间）
    【边界】只模拟"同步拷贝"（cudaMemcpy 返回时拷贝已完成）；异步拷贝、页锁定内存见第 3 讲讲义。
    """

    def __init__(self, name="模拟 GPU", pcie_bandwidth=32e9, pcie_latency=10e-6):
        self.name = name
        self.pcie_bandwidth = pcie_bandwidth
        self.pcie_latency = pcie_latency
        self._next = 0x7000_0000_0000        # 第一个设备地址（只是一个"看起来像地址"的整数）
        self.live = {}                        # 基地址 → DeviceArray（还没释放的）
        self.allocated_bytes = 0
        self.peak_bytes = 0
        self.mallocs = 0
        self.copies = 0
        self.h2d_bytes = 0
        self.d2h_bytes = 0
        self.d2d_bytes = 0
        self.copy_seconds = 0.0
        self._names = defaultdict(int)

    # ---------------- 分配与释放 ----------------
    def malloc(self, n, dtype=np.float32, name=None):
        """
        cudaMalloc：在显存上申请 n 个 dtype 类型的元素，返回 DeviceArray（真 CUDA 通过 void** 参数把地址写回）。

        【边界】n 必须 ≥ 1；地址按 256 字节对齐（cudaMalloc 的保证），所以不同数组的起点总是对齐的。
        """
        if n < 1:
            raise ValueError("cudaMalloc 的大小必须至少是 1 个元素")
        if name is None:
            self._names["d"] += 1
            name = f"d{self._names['d']}"
        arr = DeviceArray(self, n, dtype, name, self._next)
        self._next += -(-arr.nbytes // DEVICE_ALIGNMENT) * DEVICE_ALIGNMENT    # 向上取整到 256 的倍数
        self.live[arr.base] = arr
        self.mallocs += 1
        self.allocated_bytes += arr.nbytes
        self.peak_bytes = max(self.peak_bytes, self.allocated_bytes)
        return arr

    def free(self, arr):
        """cudaFree。【边界】重复释放、释放不是本设备分配的数组 → SimError。"""
        if arr.freed or self.live.get(arr.base) is not arr:
            raise SimError(f"cudaFree 了一个无效的指针（{arr.name}：重复释放或不属于这块 GPU）")
        arr.freed = True
        del self.live[arr.base]
        self.allocated_bytes -= arr.nbytes

    def leaks(self):
        """还没有 cudaFree 的数组列表（程序结束时它们就是"显存泄漏"）。"""
        return list(self.live.values())

    # ---------------- 拷贝 ----------------
    def memcpy(self, dst, src, nbytes, kind):
        """
        cudaMemcpy(dst, src, nbytes, kind)：按字节数拷贝。

        【参数】
            dst, src : 一个是 NumPy 数组（主机内存），一个是 DeviceArray（显存），或两个都是 DeviceArray
            nbytes   : 拷贝的字节数（必须是元素大小的整数倍）
            kind     : HostToDevice / DeviceToHost / DeviceToDevice
        【边界】方向和参数类型对不上、字节数超出任一方的大小、两边元素类型不同 → MemcpyError。
            真 CUDA 里主机指针和设备指针在 C 的类型系统里都是 float*，编译器分不出来，这类错误要到运行时才暴露。
        """
        want = {HostToDevice: (DeviceArray, np.ndarray), DeviceToHost: (np.ndarray, DeviceArray),
                DeviceToDevice: (DeviceArray, DeviceArray)}
        if kind not in want:
            raise MemcpyError(f"未知的拷贝方向 {kind}")
        tdst, tsrc = want[kind]
        if not isinstance(dst, tdst) or not isinstance(src, tsrc):
            raise MemcpyError(f"cudaMemcpy 方向 {kind} 要求 dst 是 {tdst.__name__}、src 是 {tsrc.__name__}，"
                              f"实际是 {type(dst).__name__}、{type(src).__name__}")
        for a in (dst, src):
            if isinstance(a, DeviceArray) and a.freed:
                raise MemcpyError(f"拷贝已释放的数组 {a.name}")
        if dst.dtype != src.dtype:
            raise MemcpyError(f"两边的元素类型不同（{dst.dtype} vs {src.dtype}）；cudaMemcpy 只按字节搬，不做类型转换")
        item = dst.dtype.itemsize
        if nbytes % item:
            raise MemcpyError(f"字节数 {nbytes} 不是元素大小 {item} 的整数倍")
        count = nbytes // item
        for a, role in ((dst, "dst"), (src, "src")):
            size = a.n if isinstance(a, DeviceArray) else a.size
            if count > size:
                raise MemcpyError(f"拷贝 {count} 个元素，超出了 {role} 的大小 {size}")
        if isinstance(src, np.ndarray) and not src.flags["C_CONTIGUOUS"]:
            raise MemcpyError("主机数组必须是连续存放的（cudaMemcpy 只认一段连续的字节）")
        sdata = src._data if isinstance(src, DeviceArray) else src.reshape(-1)
        if isinstance(dst, DeviceArray):
            dst._data[:count] = sdata[:count]
        else:
            if not dst.flags["C_CONTIGUOUS"]:
                raise MemcpyError("主机数组必须是连续存放的")
            dst.reshape(-1)[:count] = sdata[:count]
        self.copies += 1
        if kind == HostToDevice:
            self.h2d_bytes += nbytes
        elif kind == DeviceToHost:
            self.d2h_bytes += nbytes
        else:
            self.d2d_bytes += nbytes
        if kind != DeviceToDevice:             # 设备内部的拷贝走显存带宽，不经过 PCIe，这里不计入 PCIe 时间
            self.copy_seconds += self.pcie_latency + nbytes / self.pcie_bandwidth

    def to_device(self, host, name=None):
        """便捷函数：malloc + memcpy(HostToDevice)。host 会先被展平成一维（按行优先）。"""
        host = np.ascontiguousarray(host)
        arr = self.malloc(host.size, host.dtype, name)
        self.memcpy(arr, host.reshape(-1), host.nbytes, HostToDevice)
        return arr

    def to_host(self, arr, shape=None):
        """便捷函数：新建一个主机数组 + memcpy(DeviceToHost)；shape 给出时把结果 reshape 成这个形状。"""
        out = np.empty(arr.n, dtype=arr.dtype)
        self.memcpy(out, arr, arr.nbytes, DeviceToHost)
        return out.reshape(shape) if shape is not None else out

    def copy_report(self):
        """拷贝记账的汇总字典。"""
        return {"copies": self.copies, "h2d_bytes": self.h2d_bytes, "d2h_bytes": self.d2h_bytes,
                "d2d_bytes": self.d2d_bytes, "copy_seconds": self.copy_seconds}


# =====================================================================
# 3. 共享内存
# =====================================================================

class SharedArray:
    """
    一个线程块的共享内存数组（对应 __shared__ float As[T][T];）。同一个块的所有线程看到的是同一个对象。

    【在内核里怎样用】As = t.shared("As", (T, T))；As[ty, tx] = v；x = As[k]（一维时）
    【边界】下标必须是整数或整数元组，个数等于维数，每一维都在范围内（比 C 更严格：C 只要不出整块内存就不报错）。
    """

    __slots__ = ("name", "shape", "dtype", "itemsize", "offset", "_data", "_strides", "n")

    def __init__(self, name, shape, dtype, offset):
        self.name = name
        self.shape = tuple(int(s) for s in shape)
        self.dtype = np.dtype(dtype)
        self.itemsize = self.dtype.itemsize
        self.offset = offset                       # 在本块共享内存里的起始字节偏移
        self.n = int(np.prod(self.shape))
        self._data = np.full(self.n, np.nan if self.dtype.kind == "f" else -999999, dtype=self.dtype)
        strides, acc = [], 1
        for s in reversed(self.shape):
            strides.append(acc)
            acc *= s
        self._strides = tuple(reversed(strides))

    @property
    def nbytes(self):
        return self.n * self.itemsize

    def _flat(self, idx):
        """多维下标 → 一维位置（行优先），同时检查每一维的范围。"""
        if not isinstance(idx, tuple):
            idx = (idx,)
        if len(idx) != len(self.shape):
            raise IllegalAddress(f"共享数组 {self.name} 是 {len(self.shape)} 维的，下标却有 {len(idx)} 个")
        flat = 0
        for i, s, st in zip(idx, self.shape, self._strides):
            if not isinstance(i, (int, np.integer)):
                raise TypeError(f"{self.name}[...] 的下标必须是整数")
            if not 0 <= i < s:
                raise IllegalAddress(f"线程 {_RUN.where()} 越界访问共享数组 {self.name}{list(idx)}（形状 {self.shape}）")
            flat += int(i) * st
        return flat

    def __getitem__(self, idx):
        f = self._flat(idx)
        _RUN.access("shared", self, "ld", f, self.offset + f * self.itemsize, self.itemsize)
        return self._data[f]

    def __setitem__(self, idx, v):
        f = self._flat(idx)
        _RUN.access("shared", self, "st", f, self.offset + f * self.itemsize, self.itemsize)
        self._data[f] = v

    def atomic_add(self, idx, v):
        """共享内存上的原子加，返回旧值。"""
        f = self._flat(idx)
        _RUN.access("shared", self, "atom", f, self.offset + f * self.itemsize, self.itemsize)
        old = self._data[f]
        self._data[f] = old + v
        return old


# =====================================================================
# 4. 线程上下文与统计
# =====================================================================

class ThreadCtx:
    """
    内核的第一个参数 t：一个线程看得到的内置变量和工具函数。

    【属性】threadIdx、blockIdx、blockDim、gridDim（Dim3，和 CUDA 同名）；
        warp（块内第几个线程束）、lane（线程束里的第几个线程，0..31）
    【方法】shared(name, shape, dtype)：申请 / 取得本块的共享数组；syncthreads()：配合 yield 使用的栅栏
    """

    __slots__ = ("threadIdx", "blockIdx", "blockDim", "gridDim", "linear", "_run")

    def __init__(self, run, thread_idx, block_idx, linear):
        self.threadIdx = thread_idx
        self.blockIdx = block_idx
        self.blockDim = run.block
        self.gridDim = run.grid
        self.linear = linear
        self._run = run

    @property
    def warp(self):
        return self.linear // WARP_SIZE

    @property
    def lane(self):
        return self.linear % WARP_SIZE

    def shared(self, name, shape, dtype=np.float32):
        """__shared__ 声明：第一个调用的线程负责"分配"，同一块里的其他线程拿到同一个数组。"""
        return self._run.get_shared(name, shape, dtype)

    @staticmethod
    def syncthreads():
        """__syncthreads()：必须写成 `yield t.syncthreads()`，否则不起作用。"""
        return SYNC


class LaunchStats:
    """
    一次内核发射的统计结果。

    【全局内存】（"请求"= 一个线程束的一条访存指令；"事务"= 一个 32 字节 sector）
        ld_requests / ld_sectors、st_requests / st_sectors、atom_requests / atom_sectors
        thread_loads / thread_stores / thread_atomics ：按线程数的访问次数
        bytes_requested ：线程们要求读写的字节数之和（每个线程每次访问 itemsize 字节）
        unique_sectors  ：整个内核里不同 sector 的个数——如果缓存无限大，显存只需要传这么多（"最少流量"）
    【共享内存】sh_requests（线程束级访问指令数）、sh_wavefronts（实际轮数；有 bank 冲突时 > 请求数）、thread_shared
    【执行】threads、blocks、warps（总线程束数）、idle_lanes（不满的线程束里空着的线程位置数）、barriers（块级栅栏次数）
        warp_issues、thread_instructions（只有 count_instructions=True 时才统计）
    【竞争】races：检测到的数据竞争（最多保存前 20 条描述），race_count：总数
    【分数组】by_array：数组名 → 各项计数
    """

    def __init__(self, grid, block):
        self.grid, self.block = grid, block
        self.blocks = grid.count
        tpb = block.count
        self.threads = self.blocks * tpb
        wpb = -(-tpb // WARP_SIZE)
        self.warps = self.blocks * wpb
        self.idle_lanes = self.blocks * (wpb * WARP_SIZE - tpb)
        self.ld_requests = self.ld_sectors = 0
        self.st_requests = self.st_sectors = 0
        self.atom_requests = self.atom_sectors = 0
        self.thread_loads = self.thread_stores = self.thread_atomics = 0
        self.bytes_requested = 0
        self.unique_sectors = 0
        self.sh_requests = self.sh_wavefronts = self.thread_shared = 0
        self.barriers = 0
        self.warp_issues = self.thread_instructions = 0
        self.counted_instructions = False
        self.races = []
        self.race_count = 0
        self.by_array = defaultdict(lambda: defaultdict(int))

    # 派生指标 ---------------------------------------------------------
    @property
    def global_requests(self):
        return self.ld_requests + self.st_requests + self.atom_requests

    @property
    def global_sectors(self):
        return self.ld_sectors + self.st_sectors + self.atom_sectors

    @property
    def sector_bytes(self):
        """按 32 字节事务计算的全局内存流量（L1 / L2 和 SM 之间实际搬运的字节数，不考虑缓存命中）。"""
        return self.global_sectors * SECTOR_BYTES

    @property
    def sectors_per_request(self):
        """平均每条线程束访存指令产生几个事务；4 字节元素完全合并时是 4，完全不合并时是 32。"""
        return self.global_sectors / self.global_requests if self.global_requests else 0.0

    @property
    def ld_sectors_per_request(self):
        """只看读：平均每条线程束读指令产生几个事务。"""
        return self.ld_sectors / self.ld_requests if self.ld_requests else 0.0

    @property
    def bank_conflict_factor(self):
        """共享内存每条指令平均要几轮：1 表示没有 bank 冲突，32 表示最坏（32 路冲突）。"""
        return self.sh_wavefronts / self.sh_requests if self.sh_requests else 0.0

    @property
    def simt_efficiency(self):
        """线程实际执行的指令数 ÷ (发射次数 × 32)。1 表示每次发射 32 个线程都在干活。"""
        if not self.warp_issues:
            return float("nan")
        return self.thread_instructions / (self.warp_issues * WARP_SIZE)

    def summary(self):
        """常用指标的字典（打印用）。"""
        d = {"threads": self.threads, "blocks": self.blocks, "warps": self.warps,
             "thread_loads": self.thread_loads, "thread_stores": self.thread_stores,
             "ld_requests": self.ld_requests, "ld_sectors": self.ld_sectors,
             "st_requests": self.st_requests, "st_sectors": self.st_sectors,
             "sectors_per_request": round(self.sectors_per_request, 3),
             "unique_sectors": self.unique_sectors, "races": self.race_count}
        if self.sh_requests:
            d.update(sh_requests=self.sh_requests, sh_wavefronts=self.sh_wavefronts,
                     bank_conflict_factor=round(self.bank_conflict_factor, 3))
        if self.counted_instructions:
            d.update(warp_issues=self.warp_issues, thread_instructions=self.thread_instructions,
                     simt_efficiency=round(self.simt_efficiency, 4))
        return d


# =====================================================================
# 5. 发射内核
# =====================================================================

_RUN = None              # 当前正在执行的发射（模块级变量：DeviceArray 通过它知道"现在是哪个线程在访问"）


class _Run:
    """一次 launch 的运行时状态（模块内部使用）。"""

    def __init__(self, grid, block, count_instructions, check_global_races, shared_limit):
        self.grid, self.block = grid, block
        self.stats = LaunchStats(grid, block)
        self.count_instructions = count_instructions
        self.check_global_races = check_global_races
        self.shared_limit = shared_limit
        self.arrays = {}                       # id(数组) → 数组对象（统计时要用名字和元素大小）
        self.unique = set()                    # 全部访问过的全局 sector
        self.global_access = {}                # 全局竞争检测：(数组 id, 下标) → [(全局线程号, 块号, 栅栏周期, 读/写)]
        self.block_linear = 0
        self.in_thread = False                 # 只有在执行某个线程的代码时才计指令

    # ---------------- 块的开始与结束 ----------------
    def start_block(self, block_idx, block_linear):
        self.block_idx = block_idx
        self.block_linear = block_linear
        self.shared = {}                       # 名字 → SharedArray
        self.shared_bytes = 0
        self.groups = defaultdict(list)        # (线程束, 访问点, 第 k 次) → [(字节地址, 字节数)]
        self.line_groups = defaultdict(int)    # (线程束, 代码行, 第 k 次) → 执行这一行的线程数
        self.occ_by_thread = {}                # 线程 → {访问点: 已经执行了几次}
        self.line_occ_by_thread = {}
        self.epoch = 0                         # 第几个栅栏周期
        self.sh_readers = defaultdict(set)     # 本栅栏周期内：(数组 id, 下标) → 读过它的线程
        self.sh_writers = defaultdict(set)

    def set_thread(self, ctx):
        """切换"当前线程"：之后的每一次访存都记在这个线程名下。"""
        self.ctx = ctx
        self.tid = ctx.linear
        self.warp = ctx.linear // WARP_SIZE
        self.gtid = self.block_linear * self.block.count + ctx.linear
        self.occ = self.occ_by_thread.setdefault(ctx.linear, {})
        self.line_occ = self.line_occ_by_thread.setdefault(ctx.linear, {})

    def where(self):
        """报错时描述当前线程的位置。"""
        c = self.ctx
        return f"blockIdx={tuple(c.blockIdx)} threadIdx={tuple(c.threadIdx)}"

    def barrier(self):
        """全块到达栅栏：进入下一个栅栏周期，清空共享内存的竞争记录。"""
        self.epoch += 1
        self.stats.barriers += 1
        self.sh_readers.clear()
        self.sh_writers.clear()

    def get_shared(self, name, shape, dtype):
        arr = self.shared.get(name)
        shape = tuple(int(s) for s in (shape if isinstance(shape, (tuple, list)) else (shape,)))
        if arr is None:
            offset = -(-self.shared_bytes // 16) * 16          # 按 16 字节对齐
            arr = SharedArray(name, shape, dtype, offset)
            if offset + arr.nbytes > self.shared_limit:
                raise LaunchError(f"共享内存超限：本块需要 {offset + arr.nbytes} 字节 > 上限 {self.shared_limit} 字节")
            self.shared_bytes = offset + arr.nbytes
            self.shared[name] = arr
        elif arr.shape != shape or arr.dtype != np.dtype(dtype):
            raise SimError(f"共享数组 {name} 被以不同的形状 / 类型声明了两次")
        return arr

    # ---------------- 访存记录（热路径） ----------------
    def access(self, space, arr, kind, index, addr, nbytes):
        """
        记录一次访问。由 DeviceArray / SharedArray 调用：
            sys._getframe(2) 是"调用 A[i] 的那一行代码"所在的栈帧（0 是本函数，1 是 __getitem__ 等）。
        【Python 小知识】sys._getframe(k) 返回调用栈上往上数第 k 层的帧对象；f_code 是那段代码，
            f_lineno 是当前行号。两者合起来唯一确定"哪条指令"。
        """
        f = sys._getframe(2)
        key = (space, id(arr), kind, f.f_code, f.f_lineno)
        occ = self.occ
        k = occ.get(key, 0)
        occ[key] = k + 1
        self.groups[(self.warp, key, k)].append((addr, nbytes))
        self.arrays[id(arr)] = arr
        if space == "global":
            s0, s1 = addr // SECTOR_BYTES, (addr + nbytes - 1) // SECTOR_BYTES
            self.unique.add(s0)
            if s1 != s0:
                self.unique.update(range(s0 + 1, s1 + 1))
            if self.check_global_races and kind != "atom":
                self._global_race(arr, index, kind, f.f_lineno)
        elif kind != "atom":
            self._shared_race(arr, index, kind, f.f_lineno)

    def _report_race(self, msg):
        self.stats.race_count += 1
        if len(self.stats.races) < 20:
            self.stats.races.append(msg)

    def _shared_race(self, arr, index, kind, lineno):
        k = (id(arr), index)
        tid = self.tid
        if kind == "st":
            others = (self.sh_writers[k] | self.sh_readers[k]) - {tid}
            self.sh_writers[k].add(tid)
        else:
            others = self.sh_writers[k] - {tid}
            self.sh_readers[k].add(tid)
        if others:
            self._report_race(f"共享 {arr.name}[{index}]：线程 {tid} 在第 {lineno} 行{'写' if kind == 'st' else '读'}，"
                              f"同一栅栏周期里线程 {sorted(others)[:3]} 也访问过它（块 {tuple(self.block_idx)}）")

    def _global_race(self, arr, index, kind, lineno):
        k = (id(arr), index)
        me = (self.gtid, self.block_linear, self.epoch, kind)
        hist = self.global_access.setdefault(k, [])
        for g, b, e, kd in hist:
            if g != self.gtid and (kind == "st" or kd == "st") and (b != self.block_linear or e == self.epoch):
                self._report_race(f"全局 {arr.name}[{index}]：线程 {self.gtid} 在第 {lineno} 行"
                                  f"{'写' if kind == 'st' else '读'}，与线程 {g} 的{'写' if kd == 'st' else '读'}冲突")
                break
        hist.append(me)

    def count_line(self, code, lineno):
        """指令计数：当前线程执行了 code 的第 lineno 行。"""
        key = (code, lineno)
        occ = self.line_occ
        k = occ.get(key, 0)
        occ[key] = k + 1
        self.line_groups[(self.warp, key, k)] += 1

    def end_block(self):
        """块执行完：把记录下的访问按"线程束 × 指令"汇总成请求、事务、bank 冲突。"""
        st = self.stats
        for (warp, key, k), lst in self.groups.items():
            space, aid, kind = key[0], key[1], key[2]
            arr = self.arrays[aid]
            rec = st.by_array[arr.name]
            if space == "global":
                secs = set()
                for addr, nb in lst:
                    s0, s1 = addr // SECTOR_BYTES, (addr + nb - 1) // SECTOR_BYTES
                    secs.add(s0)
                    if s1 != s0:
                        secs.update(range(s0 + 1, s1 + 1))
                ns = len(secs)
                st.bytes_requested += sum(nb for _, nb in lst)
                if kind == "ld":
                    st.ld_requests += 1
                    st.ld_sectors += ns
                    st.thread_loads += len(lst)
                elif kind == "st":
                    st.st_requests += 1
                    st.st_sectors += ns
                    st.thread_stores += len(lst)
                else:
                    st.atom_requests += 1
                    st.atom_sectors += ns
                    st.thread_atomics += len(lst)
                rec[kind + "_requests"] += 1
                rec[kind + "_sectors"] += ns
                rec["thread_" + kind] += len(lst)
            else:
                banks = defaultdict(set)
                for off, nb in lst:
                    for w in range(off // BANK_BYTES, (off + nb - 1) // BANK_BYTES + 1):
                        banks[w % NUM_BANKS].add(w)
                wf = max(len(v) for v in banks.values())
                st.sh_requests += 1
                st.sh_wavefronts += wf
                st.thread_shared += len(lst)
                rec["sh_requests"] += 1
                rec["sh_wavefronts"] += wf
                rec["thread_sh"] += len(lst)
        if self.count_instructions:
            st.warp_issues += len(self.line_groups)
            st.thread_instructions += sum(self.line_groups.values())


def _global_tracer(frame, event, arg):
    """sys.settrace 的全局回调：只跟踪内核及其调用的 Python 函数（跳过模拟器自身的代码）。"""
    if frame.f_code.co_filename == _THIS_FILE:
        return None
    return _line_tracer


def _line_tracer(frame, event, arg):
    """局部回调：每执行一行，记一条指令。"""
    if event == "line" and _RUN is not None and _RUN.in_thread:
        _RUN.count_line(frame.f_code, frame.f_lineno)
    return _line_tracer


def _check_config(grid, block):
    """检查发射配置是否合法（对应 CUDA 的 invalid configuration argument）。"""
    for name, d, lim in (("block", block, MAX_BLOCK_DIM), ("grid", grid, MAX_GRID_DIM)):
        for axis, v, m in zip("xyz", d, lim):
            if v < 1:
                raise LaunchError(f"{name}.{axis} = {v}：每一维至少是 1")
            if v > m:
                raise LaunchError(f"{name}.{axis} = {v} 超过上限 {m}")
    if block.count > MAX_THREADS_PER_BLOCK:
        raise LaunchError(f"一个线程块有 {block.count} 个线程，超过上限 {MAX_THREADS_PER_BLOCK}")


def launch(kernel, grid, block, *args, count_instructions=False, check_global_races=False,
           shared_limit=DEFAULT_SHARED_LIMIT):
    """
    发射内核：相当于 CUDA 的 kernel<<<grid, block>>>(*args)，返回 LaunchStats。

    【参数】
        kernel             : 内核函数 kernel(t, *args)。要用 __syncthreads() 的内核必须是生成器函数（里面有 yield）
        grid, block        : int、元组或 Dim3
        *args              : 传给内核的参数（DeviceArray、整数、元组……；对应 CUDA 的内核参数）
        count_instructions : True 时统计指令发射与线程束分化（用 sys.settrace，会慢 3–10 倍）
        check_global_races : True 时检查全局内存的数据竞争（很慢，只在小规模演示时用）
        shared_limit       : 每个块共享内存的上限（字节）
    【思路】
        对每个块：为每个线程建一个 ThreadCtx；
            普通内核：依次调用 kernel(t, ...)；
            生成器内核：反复"让每个还活着的线程跑到下一个 yield 或结束"，直到全部结束；
                如果同一轮里有的线程停在栅栏、有的已经结束 → BarrierError。
        块结束时把访存记录汇总（见 _Run.end_block）。
    【边界】内核不能有返回值（CUDA 的 __global__ 函数返回 void），返回非 None 时报 SimError；
        不支持在内核里再发射内核（动态并行）。
    【Python 小知识】生成器函数（函数体里有 yield）被调用时并不执行，而是返回一个生成器对象；
        next(gen) 让它运行到下一个 yield 并把 yield 的值交出来；函数执行完时 next 抛出 StopIteration，
        异常对象的 .value 是函数的返回值。我们正是用它来实现"线程跑到栅栏处停下"。
    """
    global _RUN
    grid, block = as_dim3(grid), as_dim3(block)
    _check_config(grid, block)
    if _RUN is not None:
        raise SimError("模拟器不支持在内核里再发射内核")
    is_gen = inspect.isgeneratorfunction(kernel)
    run = _Run(grid, block, count_instructions, check_global_races, shared_limit)
    run.stats.counted_instructions = count_instructions
    _RUN = run
    if count_instructions:
        sys.settrace(_global_tracer)
    try:
        threads_idx = [Dim3(x, y, z) for z in range(block.z) for y in range(block.y) for x in range(block.x)]
        bl = 0
        for bz in range(grid.z):
            for by in range(grid.y):
                for bx in range(grid.x):
                    bidx = Dim3(bx, by, bz)
                    run.start_block(bidx, bl)
                    ctxs = [ThreadCtx(run, ti, bidx, lin) for lin, ti in enumerate(threads_idx)]
                    if not is_gen:
                        for c in ctxs:
                            run.set_thread(c)
                            run.in_thread = True
                            ret = kernel(c, *args)
                            run.in_thread = False
                            if ret is not None:
                                raise SimError("内核不能有返回值（__global__ 函数返回 void）：结果要写进显存数组")
                    else:
                        gens = [kernel(c, *args) for c in ctxs]
                        alive = list(range(len(ctxs)))
                        while alive:
                            waiting, finished = [], []
                            for i in alive:
                                run.set_thread(ctxs[i])
                                run.in_thread = True
                                try:
                                    tok = next(gens[i])
                                except StopIteration as stop:
                                    run.in_thread = False
                                    if stop.value is not None:
                                        raise SimError("内核不能有返回值（__global__ 函数返回 void）") from None
                                    finished.append(i)
                                    continue
                                run.in_thread = False
                                if tok is not SYNC:
                                    raise SimError("内核里的 yield 只能用来表示 __syncthreads()：请写 yield t.syncthreads()")
                                waiting.append(i)
                            if waiting and finished:
                                c = ctxs[finished[0]]
                                raise BarrierError(
                                    f"块 {tuple(bidx)} 里有 {len(waiting)} 个线程停在 __syncthreads()，"
                                    f"另有 {len(finished)} 个线程（例如 threadIdx={tuple(c.threadIdx)}）已经结束——"
                                    "真 GPU 上这是未定义行为（常见后果是卡死）。不要在 if 分支里提前 return 再用栅栏。")
                            if waiting:
                                run.barrier()
                            alive = waiting
                    run.end_block()
                    bl += 1
    finally:
        if count_instructions:
            sys.settrace(None)
        _RUN = None
    run.stats.unique_sectors = len(run.unique)
    return run.stats


def ceil_div(a, b):
    """向上取整的整数除法：CUDA 里常写成 (a + b - 1) / b，用来算"需要几个块"。"""
    return -(-a // b)
"""
01_cuda_workflow.py —— 一个 CUDA 程序的五步：分配、拷入、发射、拷出、释放，以及它们的时间账（第 3 讲 §1）
=====================================================================

【这个脚本要回答的问题】
    1. slides 的 VecAddCUDA 一步步做了什么？每一步之后显存里有什么？
    2. 拷贝要多久、内核要多久？为什么"把一次向量加法搬到 GPU 上做"反而更慢？什么样的运算才值得搬过去？
    3. 连续做几个运算时，每次都拷来拷去 vs 数据一直留在显存里，差多少？
    4. 显存管理的四个经典错误：忘了释放、重复释放、拷贝方向写反、释放之后再用。
    5. __global__、__device__、__host__ 是什么意思？

【怎么运行】
    python3 01_cuda_workflow.py       （只需要 numpy；约 1 秒）

【说明】正确性在 SIMT 模拟器上验证；时间是按公开规格算的纸面数字（PCIe 带宽、显存带宽、发射开销取量级）。
"""

import os
import sys

import numpy as np

from console import print_table

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture02"))
from gpu_model import A6000, H100  # noqa: E402
from simt import (DeviceToHost, HostToDevice, IllegalAddress, MemcpyError, SimError, Device,  # noqa: E402
                  ceil_div, launch)

LAUNCH_OVERHEAD = 5e-6        # 发射一次内核的固定开销（量级：几微秒）
COPY_LATENCY = 10e-6          # 每次 cudaMemcpy 的固定开销（量级）


def vec_add_kernel(t, A, B, C, n):
    """__global__ void VecAddKernel(int* A, int* B, int* C, int n)（slides 第 20 页）"""
    i = t.blockDim.x * t.blockIdx.x + t.threadIdx.x
    if i < n:
        C[i] = A[i] + B[i]


def show(dev, step):
    """打印显存的当前状态。"""
    live = ", ".join(f"{a.name}({a.n}×{a.dtype})" for a in dev.leaks()) or "（空）"
    print(f"  {step:<44s} 显存里：{live}；已拷贝 {dev.copies} 次（{dev.h2d_bytes + dev.d2h_bytes:,} 字节）")


def vec_add_cuda(dev, A_cpu, B_cpu, verbose=False):
    """
    slides 第 29 页的 VecAddCUDA，逐行对应：
        int *dA, *dB, *dC;
        cudaMalloc(&dA, n * sizeof(int));  ……                       ← 第 1 步：分配显存
        cudaMemcpy(dA, Acpu, n * sizeof(int), cudaMemcpyHostToDevice); ← 第 2 步：拷入
        VecAddKernel<<<num_blocks, threads_per_block>>>(dA, dB, dC, n); ← 第 3 步：发射内核
        cudaMemcpy(Ccpu, dC, n * sizeof(int), cudaMemcpyDeviceToHost); ← 第 4 步：拷出
        cudaFree(dA); cudaFree(dB); cudaFree(dC);                     ← 第 5 步：释放
    【返回】C_cpu（NumPy 数组）
    """
    n = A_cpu.size
    item = A_cpu.dtype.itemsize
    dA = dev.malloc(n, A_cpu.dtype, "dA")
    dB = dev.malloc(n, A_cpu.dtype, "dB")
    dC = dev.malloc(n, A_cpu.dtype, "dC")
    if verbose:
        show(dev, "① cudaMalloc × 3 之后：")
    dev.memcpy(dA, A_cpu, n * item, HostToDevice)
    dev.memcpy(dB, B_cpu, n * item, HostToDevice)
    if verbose:
        show(dev, "② cudaMemcpy(HostToDevice) × 2 之后：")
    threads_per_block = 256
    num_blocks = (n + threads_per_block - 1) // threads_per_block
    launch(vec_add_kernel, num_blocks, threads_per_block, dA, dB, dC, n)
    if verbose:
        show(dev, f"③ VecAddKernel<<<{num_blocks}, {threads_per_block}>>> 之后：")
    C_cpu = np.empty(n, A_cpu.dtype)
    dev.memcpy(C_cpu, dC, n * item, DeviceToHost)
    if verbose:
        show(dev, "④ cudaMemcpy(DeviceToHost) 之后：")
    dev.free(dA)
    dev.free(dB)
    dev.free(dC)
    if verbose:
        show(dev, "⑤ cudaFree × 3 之后：")
    return C_cpu


def part1_five_steps():
    print("=" * 78)
    print("第 1 部分：slides 的 VecAddCUDA，一步一步")
    print("=" * 78)
    dev = Device()
    n = 1000
    A = np.arange(n, dtype=np.int32)
    B = (2 * np.arange(n)).astype(np.int32)
    C = vec_add_cuda(dev, A, B, verbose=True)
    print(f"\n  结果和 NumPy 的 A + B 逐位相同：{np.array_equal(C, A + B)}；泄漏的显存：{dev.leaks()}")
    print("\n观察 Observations：")
    print("  - 主机和设备各有一份内存：内核只能读写显存里的 dA、dB、dC，主机只能读写 Acpu、Bcpu、Ccpu；"
          "两边之间唯一的通道是 cudaMemcpy（经过 PCIe）。")
    print(f"  - 一次向量加法拷了 {dev.copies} 次、{dev.h2d_bytes + dev.d2h_bytes:,} 字节，而内核本身对每个元素只做一次加法。"
          "第 2 部分算这笔账。\n")


def part2_time_budget():
    print("=" * 78)
    print("第 2 部分：时间账——拷贝 vs 内核（纸面估算）")
    print("=" * 78)
    rows = []
    for gpu, pcie in ((A6000, 32e9), (H100, 64e9)):
        for n in (10 ** 3, 10 ** 5, 10 ** 7, 10 ** 8):
            copy_in = 2 * (COPY_LATENCY + 4 * n / pcie)
            copy_out = COPY_LATENCY + 4 * n / pcie
            kernel = LAUNCH_OVERHEAD + 12 * n / gpu.mem_bw
            total = copy_in + kernel + copy_out
            rows.append([gpu.name, f"{n:,}", f"{(copy_in + copy_out) * 1e3:.3f}", f"{kernel * 1e3:.4f}",
                         f"{(copy_in + copy_out) / total:.1%}"])
    print(f"  向量加法（float32）：拷入 2 个数组 + 内核（读 8 字节、写 4 字节 / 元素）+ 拷出 1 个数组；"
          f"每次拷贝 + {COPY_LATENCY * 1e6:g} μs、每次发射 + {LAUNCH_OVERHEAD * 1e6:g} μs；A6000 用 PCIe 4.0（32 GB/s），H100 用 PCIe 5.0（64 GB/s）")
    print_table(["GPU", "元素数", "拷贝 ms", "内核 ms", "拷贝占比"], rows)

    rows = []
    g = A6000
    for n in (256, 1024, 4096, 16384):
        flops = 2 * n ** 3
        copy = 3 * (COPY_LATENCY + 4 * n * n / 32e9)
        compute = max(flops / (g.fp32_tflops * 1e12), 12 * n * n / g.mem_bw) + LAUNCH_OVERHEAD
        rows.append([n, f"{flops / (12 * n * n):,.0f}", f"{copy * 1e3:.3f}", f"{compute * 1e3:.3f}", f"{compute / copy:.2f}"])
    print(f"\n  对照：n × n 的 float32 矩阵乘法（拷入 A、B，拷出 C），在 A6000 上按 FP32 峰值算：")
    print_table(["n", "算术强度 FLOP/字节（按拷贝的字节算）", "拷贝 ms", "计算 ms", "计算 ÷ 拷贝"], rows)
    print("\n观察 Observations：")
    print(f"  - 向量加法在 GPU 上，拷贝占了 9 成以上的时间（大数组时内核只占几个百分点）。原因很简单：内核和拷贝搬的是同样多的字节，"
          f"而显存带宽是 PCIe 的 {A6000.mem_bw / 32e9:.0f} 倍（A6000）到 {H100.mem_bw / 64e9:.0f} 倍（H100）。"
          "\"只用一次的数据\"不值得专门搬到 GPU 上。")
    r = {row[0]: float(row[4]) for row in rows}
    print(f"  - 矩阵乘法拷贝的每个字节对应约 n / 6 次运算，n 越大计算越能摊薄拷贝——但 PCIe 实在太慢：n = 4096 时计算只有拷贝的 "
          f"{r[4096]:.2f} 倍，到 n = 16384 才是 {r[16384]:.2f} 倍。如果用张量核心（BF16 快约 4 倍），拷贝的占比还会更大。")
    print("  - 结论：数据不值得为了一次运算专门搬到 GPU 上；正确的用法是数据一直留在显存里，被很多个内核反复使用（第 3 部分）。"
          "训练时权重、梯度、优化器状态全都常驻显存，正是这个道理。")
    print("  - 数字很小时（n = 1000），时间几乎全是固定开销（每次拷贝和发射都要几微秒）：小运算要攒成大运算再交给 GPU。\n")


def scale_kernel(t, X, Y, n, a, b):
    """Y = a·X + b（逐元素）。"""
    i = t.blockDim.x * t.blockIdx.x + t.threadIdx.x
    if i < n:
        Y[i] = X[i] * a + b


def part3_keep_on_device():
    print("=" * 78)
    print("第 3 部分：连续 8 个逐元素运算——每次都拷来拷去 vs 数据留在显存里")
    print("=" * 78)
    n, k = 4096, 8
    x = np.linspace(-1, 1, n, dtype=np.float32)
    coefs = [(np.float32(1.01), np.float32(0.001 * j)) for j in range(k)]

    dev1 = Device()
    h = x.copy()
    for a, b in coefs:                                     # 写法一：每个运算都"拷入 → 发射 → 拷出"
        dX = dev1.to_device(h, "X")
        dY = dev1.malloc(n, np.float32, "Y")
        launch(scale_kernel, ceil_div(n, 256), 256, dX, dY, n, a, b)
        h = dev1.to_host(dY)
        dev1.free(dX)
        dev1.free(dY)

    dev2 = Device()
    dX = dev2.to_device(x, "X")                            # 写法二：拷入一次，8 个内核都在显存里接力
    dY = dev2.malloc(n, np.float32, "Y")
    for a, b in coefs:
        launch(scale_kernel, ceil_div(n, 256), 256, dX, dY, n, a, b)
        dX, dY = dY, dX                                    # 交换指针：这次的输出是下次的输入（不拷贝）
    h2 = dev2.to_host(dX)
    dev2.free(dX)
    dev2.free(dY)

    ref = x.copy()
    for a, b in coefs:
        ref = ref * a + b
    nbig = 10 ** 8
    t1 = k * (2 * COPY_LATENCY + 8 * nbig / 32e9) + k * (LAUNCH_OVERHEAD + 8 * nbig / A6000.mem_bw)
    t2 = 2 * COPY_LATENCY + 8 * nbig / 32e9 + k * (LAUNCH_OVERHEAD + 8 * nbig / A6000.mem_bw)
    rows = [["每次都拷来拷去", dev1.copies, f"{dev1.h2d_bytes + dev1.d2h_bytes:,}", np.array_equal(h, ref), f"{t1 * 1e3:.1f}"],
            ["拷入一次，显存里接力", dev2.copies, f"{dev2.h2d_bytes + dev2.d2h_bytes:,}", np.array_equal(h2, ref), f"{t2 * 1e3:.1f}"]]
    print_table(["写法", "拷贝次数", "拷贝字节（n = 4096）", "结果与 NumPy 逐位相同", "n = 1 亿时 A6000 上的纸面时间 ms"], rows)
    print("\n观察 Observations：")
    print(f"  - 结果完全相同，拷贝次数从 {dev1.copies} 次降到 {dev2.copies} 次。n = 1 亿时纸面时间从 {t1 * 1e3:.0f} ms 降到 {t2 * 1e3:.0f} ms"
          f"（{t1 / t2:.1f} 倍）。")
    print("  - 这正是 PyTorch 的做法：张量一旦 .to('cuda')，之后所有运算的输入输出都在显存里，只有 .item()、.cpu()、print 才会拷回主机。"
          "训练循环里不小心写一个 loss.item() 就会触发一次同步拷贝。")
    print("  - 再进一步：8 个内核之间还是要把中间结果写回显存再读出来。把它们合成一个内核（算子融合）可以连这些显存读写也省掉"
          "——第 10–11 讲。\n")


def device_helper(x):
    """相当于 __device__ float helper(float x)：只能被设备上的代码（内核或别的 __device__ 函数）调用。"""
    return x * x + 1


def uses_helper(t, X, Y, n):
    i = t.blockDim.x * t.blockIdx.x + t.threadIdx.x
    if i < n:
        Y[i] = device_helper(X[i])


def returns_value(t, X, n):
    """错误示范：__global__ 函数必须返回 void。"""
    i = t.threadIdx.x
    return X[i] if i < n else 0


def part4_memory_bugs():
    print("=" * 78)
    print("第 4 部分：显存管理的经典错误（模拟器都会抓住）")
    print("=" * 78)
    rows = []
    dev = Device()
    h = np.ones(16, np.float32)

    for _ in range(3):                                     # 错误 1：循环里分配、忘了释放
        d = dev.malloc(1 << 20, np.float32, "tmp")
        dev.memcpy(d, np.zeros(1 << 20, np.float32), 4 << 20, HostToDevice)
    rows.append(["忘了 cudaFree（循环 3 次，每次 4 MiB）", f"程序结束时还有 {len(dev.leaks())} 块、{dev.allocated_bytes / 2 ** 20:.0f} MiB 没释放"])

    d = dev.malloc(16, np.float32, "d")
    dev.free(d)
    try:
        dev.free(d)
        r = "没有报错"
    except SimError as e:
        r = str(e)
    rows.append(["同一块显存 cudaFree 两次", r])

    d = dev.malloc(16, np.float32, "d2")
    try:
        dev.memcpy(h, d, 64, HostToDevice)
        r = "没有报错"
    except MemcpyError as e:
        r = str(e)
    rows.append(["拷贝方向写反（参数顺序是 dst, src）", r])

    dev.free(d)
    try:
        launch(scale_kernel, 1, 16, d, d, 16, np.float32(1), np.float32(0))
        r = "没有报错"
    except IllegalAddress as e:
        r = str(e)
    rows.append(["cudaFree 之后又在内核里用", r])

    try:
        launch(returns_value, 1, 16, dev.to_device(h, "x"), 16)
        r = "没有报错"
    except SimError as e:
        r = str(e)
    rows.append(["__global__ 函数返回一个值", r])
    print_table(["错误", "后果"], rows, align=["left", "left"])

    dX, dY = dev.to_device(np.arange(4, dtype=np.float32), "X"), dev.malloc(4, np.float32, "Y")
    launch(uses_helper, 1, 4, dX, dY, 4)
    print(f"\n  __device__ 辅助函数：Y = X² + 1 = {dev.to_host(dY).tolist()}")
    print("\n观察 Observations：")
    print("  - 显存泄漏在训练循环里最致命：每一步漏几 MiB，几千步之后就 out of memory。PyTorch 用引用计数和缓存分配器"
          "替你管理显存（第 24 讲的 PagedAttention 会再讲显存分配）。")
    print("  - 真 CUDA 里这些错误大多不会立刻报出来：重复释放、释放后使用的后果往往是之后某个不相干的地方崩溃；"
          "方向写反时两个指针在 C 里都是 float*，编译器分辨不出。养成每次调用都检查返回的 cudaError_t 的习惯。")
    print("  - __global__ 由主机调用、在设备上执行、返回 void；__device__ 由设备调用、在设备上执行（可以有返回值）；"
          "__host__ 是普通的 CPU 函数。在模拟器里，__device__ 函数就是被内核调用的普通 Python 函数。\n")


def main():
    part1_five_steps()
    part2_time_budget()
    part3_keep_on_device()
    part4_memory_bugs()


if __name__ == "__main__":
    main()
"""
02_thread_indexing.py —— 多维网格与线程编号：二维矩阵怎样映射到线程上（第 3 讲 §2）
=====================================================================

【这个脚本要回答的问题】
    1. slides 的 fullKernel 用三维网格 + 三维线程块，算出的 tid 真的把 0..1535 每个数恰好用了一次吗？
    2. slides 的 MatAddKernel 让 x 方向对应行、y 方向对应列。结果是对的——但它比"x 对应列"慢在哪里？
    3. 矩阵大小不是块大小的整数倍时，slides 的写法会出什么事？
    4. 二维线程块的形状（32 × 8、16 × 16、4 × 64……）怎样影响合并访存？

【怎么运行】
    python3 02_thread_indexing.py      （只需要 numpy；约 2 秒）
"""

import os
import sys

import numpy as np

from console import print_table

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture02"))
from simt import Device, IllegalAddress, LaunchError, ceil_div, dim3, launch  # noqa: E402


# =====================================================================
# 内核
# =====================================================================

def full_kernel(t, COUNT, BID, TID_IN_BLOCK):
    """
    slides 第 28 页的 fullKernel（把 dout[tid] = func(din[tid]) 换成了"记录谁写了 tid"）：
        int block_id = blockIdx.x + blockIdx.y * gridDim.x + blockIdx.z * gridDim.x * gridDim.y;
        int block_offset = block_id * blockDim.x * blockDim.y * blockDim.z;
        int thread_offset = threadIdx.x + threadIdx.y * blockDim.x + threadIdx.z * blockDim.x * blockDim.y;
        int tid = block_offset + thread_offset;
    """
    b, g, th = t.blockIdx, t.gridDim, t.threadIdx
    bd = t.blockDim
    block_id = b.x + b.y * g.x + b.z * g.x * g.y
    block_offset = block_id * bd.x * bd.y * bd.z
    thread_offset = th.x + th.y * bd.x + th.z * bd.x * bd.y
    tid = block_offset + thread_offset
    COUNT.atomic_add(tid, 1)
    BID[tid] = block_id
    TID_IN_BLOCK[tid] = thread_offset


def mat_add_slides(t, A, B, C, N):
    """
    slides 第 26 页的 MatAddKernel：i 来自 x 方向，当作行号；j 来自 y 方向，当作列号；没有越界检查。
        int i = blockIdx.x * blockDim.x + threadIdx.x;
        int j = blockIdx.y * blockDim.y + threadIdx.y;
        C[i * N + j] = A[i * N + j] + B[i * N + j];
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    j = t.blockIdx.y * t.blockDim.y + t.threadIdx.y
    C[i * N + j] = A[i * N + j] + B[i * N + j]


def mat_add_row_from_y(t, A, B, C, N):
    """
    改写：行号来自 y 方向、列号来自 x 方向，并加上越界检查。
    线程束由 threadIdx.x 变化最快的 32 个线程组成 → 它们算同一行的相邻列 → 读写相邻地址。
    """
    row = t.blockIdx.y * t.blockDim.y + t.threadIdx.y
    col = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if row < N and col < N:
        C[row * N + col] = A[row * N + col] + B[row * N + col]


# =====================================================================
# 第 1 部分：三维网格的编号
# =====================================================================

def part1_full_kernel():
    print("=" * 78)
    print("第 1 部分：slides 的 fullKernel——三维网格、三维线程块的全局编号")
    print("=" * 78)
    dev = Device()
    grid, block = dim3(2, 3, 4), dim3(2, 4, 8)
    total = grid.count * block.count
    COUNT = dev.to_device(np.zeros(total, np.int32), "COUNT")
    BID = dev.malloc(total, np.int32, "BID")
    TIB = dev.malloc(total, np.int32, "TIB")
    s = launch(full_kernel, grid, block, COUNT, BID, TIB, check_global_races=True)
    cnt, bid, tib = dev.to_host(COUNT), dev.to_host(BID), dev.to_host(TIB)
    print(f"  网格 {tuple(grid)} = {grid.count} 个块，线程块 {tuple(block)} = {block.count} 个线程，共 {total} 个线程")
    print(f"  每个 tid 被写了几次：最少 {cnt.min()} 次、最多 {cnt.max()} 次；数据竞争 {s.race_count} 次")
    rows = []
    for tid in (0, 1, 2, 31, 32, 63, 64, 1535):
        b, w = int(bid[tid]), int(tib[tid])
        bx, by, bz = b % grid.x, (b // grid.x) % grid.y, b // (grid.x * grid.y)
        tx, ty, tz = w % block.x, (w // block.x) % block.y, w // (block.x * block.y)
        rows.append([tid, f"({bx}, {by}, {bz})", f"({tx}, {ty}, {tz})", b, w, w // 32])
    print_table(["tid", "blockIdx", "threadIdx", "block_id", "块内编号", "块内第几个线程束"], rows)
    print("\n观察 Observations：")
    print(f"  - 1536 个 tid 每个恰好被一个线程写了一次（最少、最多都是 {cnt.max()} 次），没有竞争：这个公式是一一对应的。"
          "它就是把 (blockIdx, threadIdx) 这六个数当成一个六维下标，按行优先展平（P1 §5）——x 变化最快，z 最慢。")
    print("  - 线程束按块内编号切分：tid 0–31 是块 (0,0,0) 的第 0 个线程束，它们的 threadIdx 是 (0..1, 0..3, 0..3)。"
          "x 方向只有 2，一个线程束跨了 y 和 z 方向——维度怎么排，决定了哪些线程在同一个线程束里一起访存（第 4 部分）。\n")


# =====================================================================
# 第 2 部分：MatAddKernel 的 x / y 方向
# =====================================================================

def part2_mat_add_xy():
    print("=" * 78)
    print("第 2 部分：slides 的 MatAddKernel——x 对应行还是对应列？")
    print("=" * 78)
    dev = Device()
    N = 32
    rng = np.random.default_rng(0)
    a = rng.standard_normal((N, N)).astype(np.float32)
    b = rng.standard_normal((N, N)).astype(np.float32)
    dA, dB = dev.to_device(a, "A"), dev.to_device(b, "B")
    rows = []
    res = {}
    for name, k in (("slides：i（行）= x 方向", mat_add_slides), ("改写：行 = y 方向、列 = x 方向", mat_add_row_from_y)):
        dC = dev.malloc(N * N, np.float32, "C")
        s = launch(k, 1, (N, N), dA, dB, dC, N)
        ok = np.array_equal(dev.to_host(dC, (N, N)), a + b)
        res[name] = s
        rows.append([name, ok, s.ld_requests, s.ld_sectors, f"{s.ld_sectors_per_request:.1f}", s.st_sectors,
                     f"{s.sector_bytes:,}", f"{s.unique_sectors * 32:,}"])
        dev.free(dC)
    print(f"  N = {N}，一个 {N} × {N} 的线程块（1024 个线程 = 32 个线程束）：")
    print_table(["写法", "结果正确", "读请求", "读事务", "读事务 / 读请求", "写事务", "事务搬运的字节", "数据本身的字节"], rows)
    s1, s2 = res["slides：i（行）= x 方向"], res["改写：行 = y 方向、列 = x 方向"]
    print("\n观察 Observations：")
    print(f"  - 两种写法结果都对，但 slides 的写法读写事务是改写后的 {s1.global_sectors / s2.global_sectors:.0f} 倍"
          f"（{s1.sector_bytes:,} vs {s2.sector_bytes:,} 字节）。")
    print("  - 原因：一个线程束是 threadIdx.x = 0..31、threadIdx.y 相同的 32 个线程。slides 的写法里 x 是行号，"
          f"这 32 个线程访问 32 个不同的行（地址相隔 N × 4 = {N * 4} 字节），每个线程独占一个 sector；"
          "改写后它们访问同一行的 32 个相邻元素，4 个 sector 就够了。")
    print("  - 规则：行优先存储的矩阵，让 threadIdx.x（变化最快的那一维）对应列号（内存里连续的那一维）。"
          "缓存会挽回一部分损失（这些 sector 后来会被别的线程束用到），但不能指望它——第 4 讲会专门讲合并访存。\n")


# =====================================================================
# 第 3 部分：矩阵大小不是块大小的整数倍
# =====================================================================

def part3_bounds():
    print("=" * 78)
    print("第 3 部分：N 不是 32 的倍数——slides 的 MatAddCUDA 没有越界检查")
    print("=" * 78)
    dev = Device()
    N, THREADS = 40, 32
    BLOCKS = (N + THREADS - 1) // THREADS
    a = np.ones((N, N), np.float32)
    dA, dB, dC = dev.to_device(a, "A"), dev.to_device(a, "B"), dev.malloc(N * N, np.float32, "C")
    print(f"  N = {N}，THREADS = {THREADS}，BLOCKS = ⌈{N}/{THREADS}⌉ = {BLOCKS}：网格 {BLOCKS} × {BLOCKS} 个块，"
          f"一共 {(BLOCKS * THREADS) ** 2} 个线程，矩阵只有 {N * N} 个元素")
    try:
        launch(mat_add_slides, (BLOCKS, BLOCKS), (THREADS, THREADS), dA, dB, dC, N)
        msg = "没有报错"
    except IllegalAddress as e:
        msg = str(e)
    print(f"  slides 的内核：{msg}")
    s = launch(mat_add_row_from_y, (BLOCKS, BLOCKS), (THREADS, THREADS), dA, dB, dC, N)
    ok = np.array_equal(dev.to_host(dC, (N, N)), a + a)
    print(f"  加上 if (row < N && col < N) 之后：结果正确 {ok}；{s.threads} 个线程里有 {s.threads - N * N} 个什么也不做")
    print("\n观察 Observations：")
    print("  - slides 的例子取 N = 32、一个块，正好整除，所以不需要越界检查；换一个 N 就越界了。"
          "二维网格要对行和列分别检查，一维的教训（第 2 讲 §4.8）在二维同样成立。")
    print(f"  - 多出来的 {s.threads - N * N} 个线程（{(s.threads - N * N) / s.threads:.0%}）是\"向上取整\"的代价。"
          "N 远大于块大小时这个比例可以忽略；N 很小或形状很\"瘦\"时（例如 1 × 4096 的矩阵用 32 × 32 的块），浪费会很大。\n")


# =====================================================================
# 第 4 部分：二维线程块的形状
# =====================================================================

def part4_block_shape():
    print("=" * 78)
    print("第 4 部分：二维线程块的形状怎样影响合并访存")
    print("=" * 78)
    dev = Device()
    N = 64
    a = np.ones((N, N), np.float32)
    dA, dB = dev.to_device(a, "A"), dev.to_device(a, "B")
    rows = []
    res = {}
    for bx, by in ((32, 8), (64, 4), (16, 16), (8, 32), (4, 64), (1, 256)):
        dC = dev.malloc(N * N, np.float32, "C")
        s = launch(mat_add_row_from_y, (ceil_div(N, bx), ceil_div(N, by)), (bx, by), dA, dB, dC, N)
        ok = np.array_equal(dev.to_host(dC, (N, N)), a + a)
        dev.free(dC)
        res[(bx, by)] = s
        rows.append([f"{bx} × {by}", bx * by, ok, s.ld_requests, s.ld_sectors, f"{s.ld_sectors_per_request:.1f}"])
    try:
        launch(mat_add_row_from_y, 1, (64, 32), dA, dB, dA, N)
        msg = "没有报错"
    except LaunchError as e:
        msg = str(e)
    print(f"  {N} × {N} 的矩阵加法（行 = y、列 = x），每块 256 个线程，只改块的形状：")
    print_table(["blockDim.x × blockDim.y", "每块线程", "结果正确", "读请求", "读事务", "读事务 / 读请求"], rows)
    print(f"\n  另外，dim3(64, 32) = 2048 个线程：{msg}")
    print("\n观察 Observations：")
    print(f"  - blockDim.x = 32、64 时，一个线程束是同一行的 32 个相邻元素：每个请求 {res[(32, 8)].ld_sectors_per_request:.0f} 个事务。"
          f"blockDim.x = 16、8 时，一个线程束跨 2 行、4 行，但每一行读的仍然是整的 64、32 字节（2 个、1 个 sector），"
          f"所以每个请求还是 {res[(8, 32)].ld_sectors_per_request:.0f} 个事务——在\"按 32 字节 sector 计数\"的模型里，一行至少 8 个 float 就不浪费。")
    print(f"  - blockDim.x = 4 时每行只读 16 字节，占半个 sector：每个请求 {res[(4, 64)].ld_sectors_per_request:.0f} 个事务（翻倍）；"
          f"1 × 256 时退化成每个线程一个 sector（{res[(1, 256)].ld_sectors_per_request:.0f} 个）。")
    print("  - 实践中二维块的 x 方向通常取 16 或 32（16 × 16、32 × 8、32 × 32）；x = 32 时一个线程束只碰一行、地址连续，"
          "对缓存行（128 字节）和显存的突发传输也最友好。原课 HW1 的提示和 PMPP 的例子常用 16 × 16 或 32 × 32。\n")


def main():
    part1_full_kernel()
    part2_mat_add_xy()
    part3_bounds()
    part4_block_shape()


if __name__ == "__main__":
    main()
"""
03_hw1_kernels.py —— 原课 HW1 的四类内核：map、zip、reduce、矩阵乘法（带步长和广播），在模拟器上逐位验证（第 3 讲 §3）
=====================================================================

【这个脚本要回答的问题】
    1. 一个内核怎样处理"任意步长"的张量——连续的、转置之后的视图、广播之后的？（map、zip）
    2. 归约有两种分工：一个线程算一个输出 vs 一个线程块算一个输出。哪种更好？和沿哪一维归约有什么关系？
    3. 一个线程算一个输出的矩阵乘法，结果和 NumPy 一样吗？为什么不是逐位一样？
    4. 每种内核的访存是否合并？

【怎么运行】
    python3 03_hw1_kernels.py       （只需要 numpy；约 6 秒）

【说明】内核和主机端封装在 kernels.py 里（对应原课 HW1 的 combine.cu 和 cuda_kernel_ops.py）。
    "逐位相同"的参照有两个：NumPy 自己的结果，以及 kernels.py 里按和内核相同的运算顺序写的参考实现。
"""

import numpy as np

from console import print_table
from kernels import (HostTensor, matmul_reference, reduce_reference, tensor_map, tensor_matmul, tensor_reduce,
                     tensor_zip)
from simt import Device

ADD = lambda x, y: x + y                       # noqa: E731
MUL = lambda x, y: x * y                       # noqa: E731
MAX = lambda x, y: x if x > y else y           # noqa: E731


def part1_map():
    print("=" * 78)
    print("第 1 部分：map——连续张量、转置视图、广播")
    print("=" * 78)
    dev = Device()
    rng = np.random.default_rng(0)
    x = rng.standard_normal((32, 32)).astype(np.float32)
    X = HostTensor.from_numpy(x)
    fn = lambda v: v * v + np.float32(1)       # noqa: E731
    rows = []
    res = {}
    cases = [("连续的 32 × 32", X, None, x * x + np.float32(1)),
             ("转置视图 X.permute(1, 0)（只改步长）", X.permute(1, 0), None, x.T * x.T + np.float32(1)),
             ("把 (1, 32) 的一行广播成 (32, 32)", HostTensor.from_numpy(x[:1]), (32, 32),
              np.broadcast_to(x[:1] * x[:1] + np.float32(1), (32, 32)))]
    for name, a, shape, ref in cases:
        out, s = tensor_map(dev, fn, a, out_shape=shape)
        ok = np.array_equal(out.to_numpy(), ref)
        res[name] = s
        rec_in, rec_out = s.by_array["a"], s.by_array["out"]
        rows.append([name, str(a.strides), ok, rec_in["ld_requests"], rec_in["ld_sectors"],
                     f"{rec_in['ld_sectors'] / rec_in['ld_requests']:.1f}", rec_out["st_sectors"]])
    print_table(["输入", "输入的步长", "与 NumPy 逐位相同", "读请求", "读事务", "事务 / 请求", "写事务"], rows)
    print("\n观察 Observations：")
    print("  - 同一个 map 内核处理了三种输入：内核只看 shape 和 strides，转置和广播都不需要拷贝数据（P1 §5 的视图）。")
    s_t = res["转置视图 X.permute(1, 0)（只改步长）"]
    s_b = res["把 (1, 32) 的一行广播成 (32, 32)"]
    print(f"  - 但视图有代价：输出按行优先连续写，转置视图的输入却是按列读的——相邻线程读的地址相隔 32 × 4 = 128 字节，"
          f"每个请求 {s_t.by_array['a']['ld_sectors'] / s_t.by_array['a']['ld_requests']:.0f} 个事务（连续时 4 个）。"
          "PyTorch 里对转置后的张量调用 .contiguous() 就是先付一次拷贝的代价，换后续运算的合并访存。")
    print(f"  - 广播时每个线程束读的都是同一行的 32 个元素（步长 0 的那一维不移动地址）：每请求 4 个事务，"
          f"32 个线程束一共 {s_b.by_array['a']['ld_sectors']} 个——但它们反复读的是同一份 128 字节，有缓存时显存只需要传一次。\n")


def part2_zip():
    print("=" * 78)
    print("第 2 部分：zip——两个输入各自广播")
    print("=" * 78)
    dev = Device()
    rng = np.random.default_rng(1)
    rows = []
    for name, sa, sb in (("同形状", (4, 8), (4, 8)), ("加偏置：(16, 8) + (8,)", (16, 8), (8,)),
                         ("外积式：(5, 1) × (1, 7)", (5, 1), (1, 7)), ("三维：(2, 1, 8) + (3, 1)", (2, 1, 8), (3, 1))):
        a = rng.standard_normal(sa).astype(np.float32)
        b = rng.standard_normal(sb).astype(np.float32)
        fn, fnp = (MUL, np.multiply) if "×" in name else (ADD, np.add)
        out, s = tensor_zip(dev, fn, HostTensor.from_numpy(a), HostTensor.from_numpy(b))
        ref = fnp(a, b)
        rows.append([name, str(out.shape), np.array_equal(out.to_numpy(), ref), s.threads, s.thread_loads])
    print_table(["情况", "输出形状", "与 NumPy 逐位相同", "发射的线程", "线程读次数"], rows)
    print("\n观察 Observations：")
    print("  - 输出形状由 shape_broadcast 决定；每个线程把自己的输出下标分别\"缩\"回 a 和 b 的形状（broadcast_index），"
          "大小为 1 的维度下标取 0。广播的那个输入里同一个元素会被很多线程读到，但内存里只有一份。")
    print("  - 逐元素运算没有求和顺序的问题，所以和 NumPy 逐位相同——前提是 fn 本身做的是同样的 float32 运算。\n")


def part3_reduce():
    print("=" * 78)
    print("第 3 部分：reduce——一个线程算一个输出 vs 一个线程块算一个输出")
    print("=" * 78)
    dev = Device()
    rng = np.random.default_rng(2)
    shape = (8, 16, 64)
    x = rng.standard_normal(shape).astype(np.float32)
    X = HostTensor.from_numpy(x)
    rows = []
    res = {}
    for dim in (0, 2):
        for method in ("thread", "block"):
            out, s = tensor_reduce(dev, ADD, X, dim, 0.0, method=method)
            mine = out.to_numpy()
            ref_same = reduce_reference(x, dim, 0.0, np.add, method=method)
            exact64 = x.astype(np.float64).sum(dim, keepdims=True)
            res[(dim, method)] = s
            rows.append([dim, "一个线程 / 输出" if method == "thread" else "一个块（32 线程）/ 输出", s.blocks, s.threads,
                         np.array_equal(mine, ref_same), f"{np.abs(mine - exact64).max():.1e}",
                         s.ld_requests, f"{s.ld_sectors_per_request:.1f}", s.barriers])
    print(f"  对形状 {shape} 的 float32 张量求和：")
    print_table(["沿哪一维", "分工", "块数", "线程数", "与同顺序参考逐位相同", "与 float64 精确和的最大差", "读请求",
                 "读事务 / 读请求", "栅栏次数"], rows)

    out, _ = tensor_reduce(dev, MAX, X, 1, -np.inf, method="block")
    ok_max = np.array_equal(out.to_numpy(), x.max(1, keepdims=True))
    out, _ = tensor_reduce(dev, MUL, HostTensor.from_numpy(np.float32(1) + x[:, :, :4] * np.float32(0.01)), 2, 1.0,
                           method="thread")
    ok_prod = np.allclose(out.to_numpy(), np.prod(np.float32(1) + x[:, :, :4] * np.float32(0.01), axis=2, keepdims=True))
    print(f"\n  同一个内核换一个 fn 和初始值：最大值（初始值 -inf）与 NumPy 逐位相同 {ok_max}；乘积（初始值 1）与 NumPy 接近 {ok_prod}")
    print("\n观察 Observations：")
    print("  - 所有写法都和\"按同样顺序计算\"的参考实现逐位相同；和 float64 的精确和相比差 1e-6 量级——不同的分工意味着不同的求和顺序，"
          "float32 的舍入就不同（P2 §3）。所以测试归约内核时要用容差比较，或者和同顺序的参考比较。")
    s0t, s0b, s2t, s2b = res[(0, "thread")], res[(0, "block")], res[(2, "thread")], res[(2, "block")]
    print(f"  - 访存是否合并取决于\"沿哪一维归约\"和\"怎样分工\"的搭配：")
    print(f"      沿第 0 维（不连续的维度）：一个线程 / 输出时，相邻线程算相邻的输出、读相邻的地址，每请求 {s0t.ld_sectors_per_request:.1f} 个事务；"
          f"一个块 / 输出时，块内线程沿第 0 维读，地址相隔 16 × 64 × 4 字节，每个线程独占一个 sector"
          f"（这一维只有 8 个元素，32 个线程里只有 8 个有活干），每请求 {s0b.ld_sectors_per_request:.1f} 个事务。")
    print(f"      沿第 2 维（连续的维度）：正好反过来——一个线程 / 输出时每请求 {s2t.ld_sectors_per_request:.1f} 个事务，"
          f"一个块 / 输出时块内线程读一段连续内存，每请求 {s2b.ld_sectors_per_request:.1f} 个。")
    print("  - 一个块 / 输出的另一个好处是并行度：输出很少时（例如把整个张量求和成一个数），一个线程 / 输出只有一个线程在干活。"
          "04 号脚本专门讨论块内的树形归约。\n")


def part4_matmul():
    print("=" * 78)
    print("第 4 部分：矩阵乘法——一个线程算一个输出元素")
    print("=" * 78)
    dev = Device()
    rng = np.random.default_rng(3)
    rows = []
    stats = {}
    cases = [("二维：(24, 16) @ (16, 20)", (24, 16), (16, 20), None),
             ("批量：(3, 12, 16) @ (3, 16, 10)", (3, 12, 16), (3, 16, 10), None),
             ("批维广播：(4, 12, 16) @ (1, 16, 10)", (4, 12, 16), (1, 16, 10), None),
             ("b 是转置视图：(24, 16) @ Bᵀ", (24, 16), (20, 16), "bt")]
    for name, sa, sb, tr in cases:
        a = rng.standard_normal(sa).astype(np.float32)
        b = rng.standard_normal(sb).astype(np.float32)
        A, B = HostTensor.from_numpy(a), HostTensor.from_numpy(b)
        if tr == "bt":
            B = B.permute(1, 0)
            b = b.T
        out, s = tensor_matmul(dev, A, B)
        mine = out.to_numpy()
        same = np.array_equal(mine, matmul_reference(a, b))
        npy = a @ b
        stats[name] = s
        rows.append([name, str(out.shape), same, np.array_equal(mine, npy), f"{np.abs(mine - npy).max():.1e}",
                     s.thread_loads, f"{s.ld_sectors_per_request:.2f}"])
    print_table(["情况", "输出形状", "与同顺序参考逐位相同", "与 NumPy 逐位相同", "与 NumPy 的最大差", "线程读次数",
                 "读事务 / 读请求"], rows)
    s2 = stats["二维：(24, 16) @ (16, 20)"]
    print("\n观察 Observations：")
    print(f"  - 每个输出元素要读 a 的一行和 b 的一列，各 m = 16 个数：线程读次数 = 2 × 24 × 20 × 16 = {s2.thread_loads}，"
          "正好是 2·n·p·m。每做一次乘加就要读两个 float（8 字节）——算术强度 0.25 FLOP/字节，这是第 4 讲的出发点。")
    print("  - 和\"按 k 从 0 到 m−1 依次先乘后加\"的参考实现逐位相同；和 NumPy（调用 BLAS）只在最后几位不同："
          "BLAS 会分块、换求和顺序、用 FMA。两者都是正确的 float32 结果。")
    print("  - b 是转置视图时，同一个内核照样算对：内核只用 strides 取数，不关心数据在内存里是怎么排的——"
          "代价是读 b 的访存模式变了（第 4 讲：转置的合并访存问题）。\n")


def main():
    part1_map()
    part2_zip()
    part3_reduce()
    part4_matmul()


if __name__ == "__main__":
    main()
"""
04_reduction_tree.py —— 归约怎样并行：一个线程、一个块的树形归约、多个块 + 原子操作（第 3 讲 §4）
=====================================================================

【这个脚本要回答的问题】
    1. 把 16384 个数加成一个数：一个线程从头加到尾、一个块做树形归约、多个块各归约一段再用 atomicAdd 汇总——
       关键路径（必须依次完成的步数）各是多少？
    2. 块内树形归约的三种写法（HW1 伪代码的 interleaved、strided、sequential）：
       合并的是同样的配对，为什么线程束分化和 bank 冲突完全不同？
    3. 很多线程同时往同一个位置累加：不用原子操作会怎样？用了原子操作，结果为什么在真 GPU 上每次可能不一样？

【怎么运行】
    python3 04_reduction_tree.py      （只需要 numpy；约 6 秒，第 2 部分开了指令计数，比较慢）
"""

import math

import numpy as np

from console import print_table
from kernels import HostTensor, reduce_reference, tensor_reduce
from simt import Device, ceil_div, launch

ADD = lambda x, y: x + y        # noqa: E731


# =====================================================================
# 内核
# =====================================================================

def one_thread_sum(t, X, OUT, n):
    """只有一个线程干活：从头加到尾。"""
    if t.blockIdx.x == 0 and t.threadIdx.x == 0:
        acc = np.float32(0)
        for i in range(n):
            acc += X[i]
        OUT[0] = acc


def grid_sum_atomic(t, X, OUT, n):
    """
    两级归约：网格里每个块用跨步循环读一部分 → 块内树形归约（sequential 写法）→ 线程 0 用 atomicAdd 把部分和加到 OUT[0]。
        对应 CUDA：if (tid == 0) atomicAdd(&out[0], cache[0]);
    """
    tid = t.threadIdx.x
    nthreads = t.blockDim.x
    cache = t.shared("cache", nthreads)
    i = t.blockIdx.x * nthreads + tid
    stride = nthreads * t.gridDim.x
    acc = np.float32(0)
    while i < n:
        acc += X[i]
        i += stride
    cache[tid] = acc
    yield t.syncthreads()
    s = nthreads // 2
    while s > 0:
        if tid < s:
            cache[tid] = cache[tid] + cache[tid + s]
        yield t.syncthreads()
        s //= 2
    if tid == 0:
        OUT.atomic_add(0, cache[0])


def every_thread_atomic(t, X, OUT, n):
    """每个线程直接 atomicAdd(&OUT[0], X[i])：n 次原子操作打在同一个地址上。"""
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < n:
        OUT.atomic_add(0, X[i])


def every_thread_racy(t, X, OUT, n):
    """错误写法：OUT[0] = OUT[0] + X[i]（先读、再加、再写，三步之间别的线程可能插进来）。"""
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < n:
        OUT[0] = OUT[0] + X[i]


# =====================================================================
# 第 1 部分：关键路径
# =====================================================================

def part1_critical_path():
    print("=" * 78)
    print("第 1 部分：把 16384 个数加成一个数——三种分工的关键路径")
    print("=" * 78)
    dev = Device()
    n = 16384
    rng = np.random.default_rng(0)
    x = rng.standard_normal(n).astype(np.float32)
    dX = dev.to_device(x, "X")
    exact = float(x.astype(np.float64).sum())
    rows = []

    OUT = dev.to_device(np.zeros(1, np.float32), "OUT")
    s1 = launch(one_thread_sum, 1, 1, dX, OUT, n)
    v1 = float(dev.to_host(OUT)[0])
    rows.append(["一个线程从头加到尾", 1, 1, s1.thread_loads, 0, 0, f"{n:,}", f"{v1 - exact:+.2e}"])

    out, s2 = tensor_reduce(dev, ADD, HostTensor.from_numpy(x.reshape(1, n)), 1, 0.0, method="block", block=256)
    v2 = float(out.to_numpy().reshape(-1)[0])
    path2 = n // 256 + int(math.log2(256))
    rows.append(["一个块（256 线程）+ 树形归约", 1, 256, s2.thread_loads, s2.barriers, 0, f"{n // 256} + {int(math.log2(256))} = {path2}",
                 f"{v2 - exact:+.2e}"])

    G, B = 16, 256
    OUT = dev.to_device(np.zeros(1, np.float32), "OUT")
    s3 = launch(grid_sum_atomic, G, B, dX, OUT, n)
    v3 = float(dev.to_host(OUT)[0])
    path3 = n // (G * B) + int(math.log2(B)) + G
    rows.append([f"{G} 个块 × {B} 线程 + 树形归约 + atomicAdd", G, G * B, s3.thread_loads, s3.barriers, s3.thread_atomics,
                 f"{n // (G * B)} + {int(math.log2(B))} + {G} = {path3}", f"{v3 - exact:+.2e}"])
    print_table(["分工", "块数", "线程数", "线程读次数", "栅栏次数", "原子操作", "关键路径（依次完成的步数）", "与精确和之差"], rows)
    print("\n  关键路径的算法：每个线程先依次累加自己那份（n / 线程数 步），树形归约每轮一步（log₂256 = 8 轮），"
          "落在同一个地址上的原子操作只能一个接一个执行（16 步）。")
    print("\n观察 Observations：")
    print(f"  - 三种写法读的数一样多（{n:,} 次），总的加法次数也差不多；差别全在\"必须依次完成的步数\"：从 {n:,} 步降到 {path2} 步、再到 {path3} 步。"
          "GPU 的时间由关键路径决定（只要线程足够多），这就是并行算法和串行算法的根本区别。")
    print("  - 树形归约的思想：加法满足结合律，(a+b)+(c+d) 和 ((a+b)+c)+d 数学上相等，前者两次加法可以同时做。"
          "n 个数的树有 log₂n 层，每层的加法互不依赖。")
    print("  - 三个结果在最后几位各不相同：float32 的加法不满足结合律（P2 §3），换一种分工就换了一种求和顺序。"
          "树形求和其实比从头加到尾更准（P2 的成对求和），这里树形的误差也更小。\n")


# =====================================================================
# 第 2 部分：块内树形归约的三种写法
# =====================================================================

def part2_tree_variants():
    print("=" * 78)
    print("第 2 部分：块内树形归约的三种写法——同样的配对，不同的代价")
    print("=" * 78)
    dev = Device()
    rows_n, cols, B = 8, 1024, 256
    rng = np.random.default_rng(1)
    x = rng.standard_normal((rows_n, cols)).astype(np.float32)
    X = HostTensor.from_numpy(x)
    rows = []
    res = {}
    for mode, desc in (("interleaved", "interleaved：if (tid % (2s) == 0)（HW1 伪代码）"),
                       ("strided", "strided：index = 2s·tid"),
                       ("sequential", "sequential：if (tid < s)，s 减半")):
        out, s = tensor_reduce(dev, ADD, X, 1, 0.0, method="block", mode=mode, block=B, count_instructions=True)
        ok = np.array_equal(out.to_numpy(), reduce_reference(x, 1, 0.0, np.add, method="block", block=B, mode=mode))
        res[mode] = s
        rows.append([desc, ok, s.warp_issues, f"{s.simt_efficiency:.1%}", s.sh_requests, s.sh_wavefronts,
                     f"{s.bank_conflict_factor:.2f}"])
    print(f"  {rows_n} 行 × {cols} 列，每行一个块（{B} 个线程 = {B // 32} 个线程束），沿行求和：")
    print_table(["写法", "与同顺序参考逐位相同", "线程束发射次数", "SIMT 效率", "共享内存请求", "共享内存轮数", "每请求轮数"], rows)
    i, st, sq = res["interleaved"], res["strided"], res["sequential"]
    print("\n  三种写法每一轮做的事（以 s = 1、2 为例，一个线程束 = 线程 0–31）：")
    print("    interleaved：s=1 时线程 0,2,4,…,30 干活（每个线程束都有一半线程闲着）；s=2 时线程 0,4,8,… 干活……")
    print("    strided    ：s=1 时线程 0–127 干活，访问 cache[0,2,4,…]、cache[1,3,5,…]；s=2 时线程 0–63，访问 cache[0,4,8,…]……")
    print("    sequential ：s=128 时线程 0–127 干活，访问 cache[0..127] 和 cache[128..255]；s=64 时线程 0–63……")
    print("\n观察 Observations：")
    print(f"  - interleaved 每一轮都让每个线程束里的一部分线程干活：SIMT 效率 {i.simt_efficiency:.0%}（线程束分化）；"
          f"而且每个线程束每一轮都要发共享内存请求，一共 {i.sh_requests} 个（另两种写法只有 {sq.sh_requests} 个）。")
    print(f"  - strided 让编号连续的线程干活，分化少了（效率 {st.simt_efficiency:.0%}），共享内存请求也降到 {st.sh_requests} 个；"
          f"但线程 tid 访问 cache[2s·tid]：s = 1 时一个线程束访问 0, 2, 4, …, 62 号字，两两落在同一个 bank 上，s 越大冲突越严重。"
          f"平均每请求 {st.bank_conflict_factor:.2f} 轮，共享内存的总轮数还是 {st.sh_wavefronts}——和 interleaved 一样，省下的全被 bank 冲突吃回去了。"
          f"它的发射次数（{st.warp_issues}）也没有比 interleaved（{i.warp_issues}）少：多了一行算 index。")
    print(f"  - sequential 两个问题都没有：干活的线程编号连续（分化只在最后几轮 s < 32 时出现），访问的地址也连续（没有 bank 冲突）。"
          f"发射次数最少（{sq.warp_issues}），共享内存总轮数 {sq.sh_wavefronts}，是另两种的 {sq.sh_wavefronts / i.sh_wavefronts:.0%}。")
    print("  - 这里的计数包含了每个线程先累加 4 个数的阶段（三种写法完全相同），所以整体的发射次数看起来差别不大。")
    print("  - interleaved 和 strided 合并的配对完全相同，结果逐位相同；sequential 配对不同，结果可能差最后一位——"
          "但三者都和各自的同顺序参考逐位相同。这组写法出自 NVIDIA 的经典教程 *Optimizing Parallel Reduction in CUDA*（Mark Harris），"
          "第 4 讲会正式讲 bank 冲突。\n")


# =====================================================================
# 第 3 部分：原子操作与数据竞争
# =====================================================================

def part3_atomics():
    print("=" * 78)
    print("第 3 部分：很多线程往同一个位置累加——竞争、原子操作、不确定性")
    print("=" * 78)
    dev = Device()
    n = 2048
    rng = np.random.default_rng(2)
    x = rng.standard_normal(n).astype(np.float32)
    dX = dev.to_device(x, "X")
    rows = []
    for name, k, kw in (("OUT[0] = OUT[0] + X[i]（不用原子操作）", every_thread_racy, {"check_global_races": True}),
                        ("atomicAdd(&OUT[0], X[i])（每个线程一次）", every_thread_atomic, {"check_global_races": True}),
                        ("块内树形归约 + 每块一次 atomicAdd", grid_sum_atomic, {})):
        OUT = dev.to_device(np.zeros(1, np.float32), "OUT")
        grid = ceil_div(n, 256) if k is not grid_sum_atomic else 8
        s = launch(k, grid, 256, dX, OUT, n, **kw)
        v = float(dev.to_host(OUT)[0])
        rows.append([name, s.thread_atomics, s.race_count, f"{v:.6f}"])
    print_table(["写法", "原子操作次数", "检测到的竞争", "模拟器上的结果"], rows)

    orders = set()
    for seed in range(20):
        perm = np.random.default_rng(seed).permutation(n)
        acc = np.float32(0)
        for v in x[perm]:
            acc = np.float32(acc + v)
        orders.add(float(acc))
    print(f"\n  对照：同样 {n} 个 float32，按 20 种随机顺序依次相加，得到 {len(orders)} 个不同的结果，"
          f"范围 [{min(orders):.6f}, {max(orders):.6f}]。")
    print("\n观察 Observations：")
    print("  - 不用原子操作的写法在模拟器上算出了\"正确\"的和——因为模拟器让线程一个接一个地执行。真 GPU 上，两个线程可能同时读到旧值、"
          "各自加完再写回，一个人的加法就丢了（P4 §3 的\"读-加-写\"竞争）。竞争检测器把它抓了出来：每次访问都和别的线程冲突。")
    print("  - atomicAdd 把\"读-加-写\"变成硬件保证的一步，结果正确；代价是打在同一个地址上的原子操作只能排队执行。"
          "先在块内归约、每块只做一次原子操作，原子操作从 2048 次降到 8 次。")
    print(f"  - 原子操作只保证\"不丢\"，不保证顺序：真 GPU 上每次运行，块完成的先后不同，加法顺序就不同。"
          f"float32 不满足结合律，所以结果每次可能差最后几位（这里 20 种顺序得到 {len(orders)} 个不同的值）。"
          "这是 GPU 训练\"不可逐位复现\"的常见来源之一；需要确定性时，要用固定顺序的归约（PyTorch 的 use_deterministic_algorithms）。\n")


def main():
    part1_critical_path()
    part2_tree_variants()
    part3_atomics()


if __name__ == "__main__":
    main()
"""
console.py —— 在终端里打印对齐的表格（中英文混排也能对齐）

【为什么需要它】
    终端里一个汉字占两个英文字母的宽度，Python 的 f"{s:10s}" 却按"字符个数"补空格，
    所以中英混排的表格会歪掉。这里的 pad() 按"显示宽度"补空格。

【用法】
    from console import pad, print_table
    print_table(["模型", "困惑度"], [["trigram", "3.21"], ["RNN", "2.60"]])
"""

import unicodedata


def display_width(text):
    """
    一个字符串在终端里占多宽：东亚宽字符（汉字、全角标点）算 2，其余算 1。

    【Python 小知识】unicodedata.east_asian_width(c) 返回 'W'（宽）、'F'（全角）、
        'Na'（窄）等标记；W 和 F 在终端里占两格。
    """
    return sum(2 if unicodedata.east_asian_width(c) in ("W", "F") else 1 for c in str(text))


def pad(text, width, align="left"):
    """
    把 text 补空格到显示宽度 width。

    【参数】text：任意对象（会先转成字符串）；width：目标宽度；align："left" 或 "right"
    【返回】补好空格的字符串
    【边界】text 已经比 width 宽时原样返回（不截断，宁可表格歪一点也不丢信息）。
    """
    s = str(text)
    gap = max(0, width - display_width(s))
    return s + " " * gap if align == "left" else " " * gap + s


def print_table(headers, rows, align=None):
    """
    打印一张表。

    【参数】
        headers : 表头列表
        rows    : 每行是一个列表，长度和 headers 相同；元素可以是任意对象
        align   : 每列的对齐方式列表（"left"/"right"）；None 表示第一列左对齐、其余右对齐
    【边界】rows 为空时只打印表头。
    """
    cols = len(headers)
    if align is None:
        align = ["left"] + ["right"] * (cols - 1)
    cells = [[str(h) for h in headers]] + [[str(c) for c in r] for r in rows]
    widths = [max(display_width(row[j]) for row in cells) for j in range(cols)]
    for i, row in enumerate(cells):
        print("  " + "  ".join(pad(row[j], widths[j], align[j]) for j in range(cols)))
        if i == 0:
            print("  " + "  ".join("-" * widths[j] for j in range(cols)))
"""
kernels.py —— 原课 HW1 的四类内核（map、zip、reduce、矩阵乘法）在 SIMT 模拟器上的实现，以及主机端的封装（第 3 讲）
=====================================================================
（第 4 讲用它做对照，第 5 讲把它接成迷你框架的一个"后端"）

【这个文件要回答的问题】
    原课 HW1 要在 combine.cu 里写四个 CUDA 内核，再在 cuda_kernel_ops.py 里把它们接到 miniTorch 的张量上。
    这里用 SIMT 模拟器把同样的东西写一遍：
        内核（设备端）：每个线程算出自己负责的输出位置 → 用步长找到输入的位置 → 计算 → 写回
        封装（主机端）：分配显存 → 拷贝输入 → 算网格大小 → 发射内核 → 拷回结果 → 释放显存
    张量 = 一维存储 + 形状 + 步长（P1 §5）。内核只拿到"指针 + 几个整数元组"，
    所以同一个内核能处理连续的张量、转置之后的视图、广播之后的张量。

【本文件提供】
    HostTensor          : 主机端的张量（NumPy 一维存储 + shape + strides），from_numpy、permute、to_numpy
    内核（设备端）       : map_kernel、zip_kernel、reduce_kernel_thread、reduce_kernel_block、matmul_kernel
    封装（主机端）       : tensor_map、tensor_zip、tensor_reduce、tensor_matmul（每个都返回 (结果, LaunchStats)）
    参考实现            : reduce_reference、matmul_reference（按和内核完全相同的顺序用 float32 计算，用来逐位对比）

【约定】
    - 和 miniTorch 一样：reduce 之后被归约的那一维保留、长度变成 1；matmul 支持三维（第 0 维是批），批维可以广播。
    - 下标换算用 P1 的 tensor_layout.py（to_index、index_to_position、broadcast_index），它们在内核里相当于 __device__ 函数。
    - 形状和步长以 Python 元组的形式直接传给内核（相当于 CUDA 的内核参数 / 常量内存）。
      原课 HW1 的做法是把它们也 cudaMemcpy 到显存里再传指针，原理相同。
"""

import os
import sys

import numpy as np

HERE = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, os.path.join(HERE, "..", "lecture02"))
sys.path.insert(0, os.path.join(HERE, "..", "prereq1"))
from simt import ceil_div, launch  # noqa: E402
from tensor_layout import broadcast_index, index_to_position, shape_broadcast, strides_for, to_index  # noqa: E402

THREADS_PER_BLOCK = 32          # 一维内核每块的线程数（本讲义取 32，模拟器跑得快；真 GPU 上常用 128–256）


# =====================================================================
# 0. 主机端的张量
# =====================================================================

class HostTensor:
    """
    主机内存里的张量：一维存储 storage + 形状 shape + 步长 strides（单位是元素）。和 miniTorch 的 TensorData 一样。

    【方法】
        HostTensor.from_numpy(arr) : 从 NumPy 数组建一个连续的张量（复制一份数据）
        t.permute(*order)          : 交换维度，返回共享存储的新视图（零拷贝）
        t.to_numpy()               : 按 shape / strides 读出一个普通的 NumPy 数组（用来检查结果）
    【边界】没有 offset（起点总是存储的第 0 个元素），和 miniTorch 一致。
    """

    def __init__(self, storage, shape, strides=None):
        self.storage = storage
        self.shape = tuple(int(s) for s in shape)
        self.strides = tuple(strides) if strides is not None else strides_for(self.shape)
        self.size = int(np.prod(self.shape)) if self.shape else 1

    @classmethod
    def from_numpy(cls, arr, dtype=np.float32):
        arr = np.ascontiguousarray(arr, dtype=dtype)
        return cls(arr.reshape(-1).copy(), arr.shape)

    def permute(self, *order):
        return HostTensor(self.storage, tuple(self.shape[k] for k in order), tuple(self.strides[k] for k in order))

    def to_numpy(self):
        itemsize = self.storage.itemsize
        return np.lib.stride_tricks.as_strided(self.storage, self.shape, tuple(s * itemsize for s in self.strides)).copy()

    def __repr__(self):
        return f"HostTensor(shape={self.shape}, strides={self.strides})"


# =====================================================================
# 1. 内核（设备端）
# =====================================================================

def map_kernel(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, fn):
    """
    map：out[i] = fn(a[对应位置])，a 可以是任意步长的视图，也可以比 out 小（广播）。

    对应原课 HW1 的 mapKernel。每个线程负责一个输出元素：
        1. i = blockIdx.x * blockDim.x + threadIdx.x；i ≥ out_size 的线程什么也不做
        2. out_index = to_index(i, out_shape)              —— 第 i 个输出元素的多维下标
        3. a_index = broadcast_index(out_index, ...)      —— 广播：a 里大小为 1 的维度下标取 0
        4. 按各自的步长换算成一维位置，读 a、算 fn、写 out
    【边界】fn 是一个普通的 Python 函数（CUDA 里是模板参数 / 函数指针）。
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < out_size:
        out_index = to_index(i, out_shape)
        a_index = broadcast_index(out_index, out_shape, a_shape)
        out[index_to_position(out_index, out_strides)] = fn(a[index_to_position(a_index, a_strides)])


def zip_kernel(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, b, b_shape, b_strides, fn):
    """
    zip：out[i] = fn(a[…], b[…])，a、b 各自按广播规则对应到 out 的形状。对应原课 HW1 的 zipKernel。
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < out_size:
        out_index = to_index(i, out_shape)
        ai = broadcast_index(out_index, out_shape, a_shape)
        bi = broadcast_index(out_index, out_shape, b_shape)
        out[index_to_position(out_index, out_strides)] = fn(a[index_to_position(ai, a_strides)],
                                                             b[index_to_position(bi, b_strides)])


def reduce_kernel_thread(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, dim, start, fn):
    """
    reduce（写法一：一个线程算一个输出）：线程 i 负责第 i 个输出，自己沿 dim 从头加到尾。

    【优点】简单，不需要共享内存和同步。
    【缺点】输出很少时（例如把一个长向量求和成 1 个数）只有一个线程在干活，其余几万个线程闲着。
    """
    i = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if i < out_size:
        out_index = to_index(i, out_shape)
        a_index = list(out_index)
        acc = start
        for k in range(a_shape[dim]):
            a_index[dim] = k
            acc = fn(acc, a[index_to_position(a_index, a_strides)])
        out[index_to_position(out_index, out_strides)] = acc


def reduce_kernel_block(t, out, out_shape, out_strides, out_size, a, a_shape, a_strides, dim, start, fn, mode):
    """
    reduce（写法二：一个线程块算一个输出）：对应原课 HW1 的提示——"每个输出元素由一个块单独计算"，
    以及可选的"共享内存 + 树形归约"优化。

    【步骤】
        1. 块号 blockIdx.x = 第几个输出；块内每个线程沿 dim 以 blockDim.x 为步长各加一部分（跨步循环）
        2. 每个线程把自己的部分和写进共享数组 cache[tid]，然后 __syncthreads()
        3. 树形归约：每一轮把"活着的"部分和两两合并，直到只剩 cache[0]；每一轮之后都要 __syncthreads()
             mode = "interleaved" : 第 s 轮，tid % (2s) == 0 的线程把 cache[tid + s] 加到 cache[tid]（HW1 给的伪代码 reduce0）
             mode = "strided"     : 第 s 轮，前 blockDim/(2s) 个线程各算 index = 2s·tid，把 cache[index + s] 加到 cache[index]
                                    （合并的是和 interleaved 完全相同的两两配对，只是换了"由谁来做"）
             mode = "sequential"  : 步长从 blockDim/2 开始减半，tid < s 的线程把 cache[tid + s] 加到 cache[tid]
        4. 线程 0 把 cache[0] 写到输出
    【边界】blockDim.x 必须是 2 的幂（树形归约每轮减半）。
        注意 __syncthreads() 不能放在只有部分线程会执行的分支里——所以 if 只包住"加"那一行，栅栏在 if 外面。
    """
    tid = t.threadIdx.x
    nthreads = t.blockDim.x
    cache = t.shared("cache", nthreads, np.float32)
    out_index = to_index(t.blockIdx.x, out_shape)
    a_index = list(out_index)
    acc = start
    k = tid
    while k < a_shape[dim]:
        a_index[dim] = k
        acc = fn(acc, a[index_to_position(a_index, a_strides)])
        k += nthreads
    cache[tid] = acc
    yield t.syncthreads()
    if mode == "interleaved":
        s = 1
        while s < nthreads:
            if tid % (2 * s) == 0:
                cache[tid] = fn(cache[tid], cache[tid + s])
            yield t.syncthreads()
            s *= 2
    elif mode == "strided":
        s = 1
        while s < nthreads:
            index = 2 * s * tid
            if index < nthreads:
                cache[index] = fn(cache[index], cache[index + s])
            yield t.syncthreads()
            s *= 2
    else:
        s = nthreads // 2
        while s > 0:
            if tid < s:
                cache[tid] = fn(cache[tid], cache[tid + s])
            yield t.syncthreads()
            s //= 2
    if tid == 0:
        out[index_to_position(out_index, out_strides)] = cache[0]


def matmul_kernel(t, out, out_shape, out_strides, a, a_shape, a_strides, b, b_shape, b_strides):
    """
    批量矩阵乘法 out[z] = a[z] @ b[z]：一个线程算一个输出元素（原课 HW1 的"简单并行"提示）。

    【线程分工】三维网格：blockIdx.z = 第几个批；x 方向对应输出的列 j、y 方向对应行 i：
        j = blockIdx.x * blockDim.x + threadIdx.x，i = blockIdx.y * blockDim.y + threadIdx.y
        这样一个线程束里相邻的线程算同一行相邻的列——读 b 和写 out 都是合并的（第 2 部分会看到反过来有多糟）。
    【广播】a 或 b 的批维长度为 1 时，它的批步长当作 0：所有批共用同一个矩阵。
    【边界】越界检查：i、j 超出输出形状的线程什么也不做（网格是按块大小向上取整的）。
    【数值】acc 用 float32 累加，每步先乘、再加（两次舍入）。真 GPU 上 nvcc 默认把 a*b+c 编译成 FMA（只舍入一次），
        所以真 GPU 的结果可能和这里差最后一位。
    """
    z = t.blockIdx.z
    i = t.blockIdx.y * t.blockDim.y + t.threadIdx.y
    j = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    n, p = out_shape[1], out_shape[2]
    if i < n and j < p:
        m = a_shape[2]
        a_base = (z if a_shape[0] > 1 else 0) * a_strides[0] + i * a_strides[1]
        b_base = (z if b_shape[0] > 1 else 0) * b_strides[0] + j * b_strides[2]
        acc = np.float32(0)
        for k in range(m):
            acc += a[a_base + k * a_strides[2]] * b[b_base + k * b_strides[1]]
        out[z * out_strides[0] + i * out_strides[1] + j * out_strides[2]] = acc


# =====================================================================
# 2. 主机端封装（对应原课 HW1 的 cuda_kernel_ops.py）
# =====================================================================

def _upload(dev, x, name):
    """把一个 HostTensor 的整段存储拷到显存（视图也拷整段存储，步长照用）。"""
    return dev.to_device(x.storage, name)


def tensor_map(dev, fn, a, out_shape=None, block=THREADS_PER_BLOCK, **launch_kw):
    """
    out = map(fn, a)。out_shape 默认等于 a 的形状；给出更大的形状时 a 会被广播。

    【步骤】和原课 HW1 Part B 一样：建输出 → cudaMalloc / cudaMemcpy → 算块数 → 发射 → 拷回 → cudaFree。
    【返回】(输出 HostTensor（连续）, LaunchStats)
    """
    out_shape = tuple(out_shape or a.shape)
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_out = _upload(dev, a, "a"), dev.malloc(out.size, np.float32, "out")
    stats = launch(map_kernel, ceil_div(out.size, block), block,
                   d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, fn, **launch_kw)
    out.storage = dev.to_host(d_out)
    dev.free(d_a)
    dev.free(d_out)
    return out, stats


def tensor_zip(dev, fn, a, b, block=THREADS_PER_BLOCK, **launch_kw):
    """out = zip(fn, a, b)，输出形状是两者广播之后的形状。【返回】(输出, LaunchStats)"""
    out_shape = shape_broadcast(a.shape, b.shape)
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_b, d_out = _upload(dev, a, "a"), _upload(dev, b, "b"), dev.malloc(out.size, np.float32, "out")
    stats = launch(zip_kernel, ceil_div(out.size, block), block,
                   d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, d_b, b.shape, b.strides, fn,
                   **launch_kw)
    out.storage = dev.to_host(d_out)
    for d in (d_a, d_b, d_out):
        dev.free(d)
    return out, stats


def tensor_reduce(dev, fn, a, dim, start, method="block", mode="sequential", block=None, **launch_kw):
    """
    沿 dim 归约，输出在 dim 上的长度变成 1。

    【参数】fn：二元函数（例如加法）；start：初始值（求和 0、求积 1、求最大值 -inf）；
        method："thread"（一个线程一个输出）或 "block"（一个块一个输出 + 树形归约）；mode：树形归约的两种写法；
        block：每块线程数（"block" 方法默认 32，必须是 2 的幂）
    【返回】(输出, LaunchStats)
    """
    out_shape = tuple(1 if k == dim else s for k, s in enumerate(a.shape))
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_out = _upload(dev, a, "a"), dev.malloc(out.size, np.float32, "out")
    if method == "thread":
        block = block or THREADS_PER_BLOCK
        stats = launch(reduce_kernel_thread, ceil_div(out.size, block), block,
                       d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, dim, np.float32(start), fn,
                       **launch_kw)
    else:
        block = block or 32
        if block & (block - 1):
            raise ValueError("树形归约要求每块线程数是 2 的幂")
        stats = launch(reduce_kernel_block, out.size, block,
                       d_out, out.shape, out.strides, out.size, d_a, a.shape, a.strides, dim, np.float32(start), fn, mode,
                       **launch_kw)
    out.storage = dev.to_host(d_out)
    dev.free(d_a)
    dev.free(d_out)
    return out, stats


def tensor_matmul(dev, a, b, block=(8, 8), **launch_kw):
    """
    批量矩阵乘法：a (B, n, m) @ b (B, m, p) → (B, n, p)；二维输入会先补一个长度为 1 的批维，结果再去掉。

    【边界】a 的最后一维必须等于 b 的倒数第二维；两者的批维要么相等，要么有一个是 1（广播）。
    【返回】(输出, LaunchStats)
    """
    squeeze = len(a.shape) == 2 and len(b.shape) == 2
    if len(a.shape) == 2:
        a = HostTensor(a.storage, (1,) + a.shape, (0,) + a.strides)
    if len(b.shape) == 2:
        b = HostTensor(b.storage, (1,) + b.shape, (0,) + b.strides)
    if a.shape[2] != b.shape[1]:
        raise ValueError(f"矩阵乘法的形状不匹配：{a.shape} @ {b.shape}")
    if a.shape[0] != b.shape[0] and 1 not in (a.shape[0], b.shape[0]):
        raise ValueError(f"批维无法广播：{a.shape[0]} vs {b.shape[0]}")
    B = max(a.shape[0], b.shape[0])
    out_shape = (B, a.shape[1], b.shape[2])
    out = HostTensor(np.zeros(int(np.prod(out_shape)), np.float32), out_shape)
    d_a, d_b, d_out = _upload(dev, a, "a"), _upload(dev, b, "b"), dev.malloc(out.size, np.float32, "out")
    bx, by = block
    grid = (ceil_div(out_shape[2], bx), ceil_div(out_shape[1], by), B)
    stats = launch(matmul_kernel, grid, (bx, by), d_out, out.shape, out.strides,
                   d_a, a.shape, a.strides, d_b, b.shape, b.strides, **launch_kw)
    out.storage = dev.to_host(d_out)
    for d in (d_a, d_b, d_out):
        dev.free(d)
    if squeeze:
        out = HostTensor(out.storage, out_shape[1:])
    return out, stats


# =====================================================================
# 3. 参考实现（按和内核相同的顺序计算，用来逐位对比）
# =====================================================================

def reduce_reference(x, dim, start, fn_np, method="thread", block=32, mode="sequential"):
    """
    用 NumPy 按和内核完全相同的运算顺序做归约（float32），结果应该和内核逐位相同。

    【参数】x：NumPy 数组；fn_np：向量化的二元函数（np.add、np.maximum……）
        method="thread"：从头到尾依次合并；method="block"：先按 tid 跨步累加成 block 份，再按 mode 做树形合并
    """
    x = np.moveaxis(np.asarray(x, np.float32), dim, -1)
    L = x.shape[-1]
    if method == "thread":
        acc = np.full(x.shape[:-1], start, np.float32)
        for k in range(L):
            acc = fn_np(acc, x[..., k]).astype(np.float32)
    else:
        parts = np.full(x.shape[:-1] + (block,), start, np.float32)
        for k in range(L):
            parts[..., k % block] = fn_np(parts[..., k % block], x[..., k])
        if mode in ("interleaved", "strided"):         # 两种写法合并的是同样的配对，结果相同
            s = 1
            while s < block:
                idx = np.arange(0, block, 2 * s)
                parts[..., idx] = fn_np(parts[..., idx], parts[..., idx + s])
                s *= 2
        else:
            s = block // 2
            while s > 0:
                parts[..., :s] = fn_np(parts[..., :s], parts[..., s:2 * s])
                s //= 2
        acc = parts[..., 0]
    return np.expand_dims(acc, dim)


def matmul_reference(a, b):
    """按 k 从 0 到 m−1 依次"先乘后加"的 float32 矩阵乘法（和 matmul_kernel 的运算顺序相同）。"""
    a = np.asarray(a, np.float32)
    b = np.asarray(b, np.float32)
    acc = np.zeros(np.broadcast_shapes(a[..., :1].shape, b[..., :1, :].shape), np.float32)   # 输出形状（批维广播）
    for k in range(a.shape[-1]):
        acc = (acc + (a[..., :, k:k + 1] * b[..., k:k + 1, :]).astype(np.float32)).astype(np.float32)
    return acc
