"""
=============================================================================
CS349E · 第 10 讲 · 03_register_tiling.py —— 寄存器分块、向量化与 bank 冲突：从 3 TFLOPS 到 18 TFLOPS
=============================================================================

【这个脚本要回答的问题】
    共享内存分块之后，瓶颈变成"每次乘加从共享内存读两个数"。让一个线程算多个输出（寄存器分块），
    从共享内存读进寄存器的数就能被用好几次。这一步把性能提高了 3–6 倍。
      演示 1：k3 → k4 → k5 → k6 的事件账本（每 1000 次乘加），并在两个矩阵大小上验证"计数和乘加次数成正比"。
      演示 2：bank 冲突：k5、k6 里哪条指令有冲突、几路；padding 和"warp 分块式"的线程映射怎样消除它。
      演示 3：整条阶梯：模型的上限 vs Boehm 的实测（k1–k6），以及模型低于实测时说明了什么。
      演示 4：寄存器分块多大才够：每次乘加的共享内存字节 vs 硬件能供应的字节；为什么 A6000 和 A100 的最优分块不同。

【运行】python3 03_register_tiling.py      （约 10 秒）
【依赖】numpy；本目录的 kernel_lab；../lecture01 的 console。
"""

import os
import sys

import numpy as np

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

from kernel_lab import (A100, A6000, BOEHM_LATER, KERNELS, PERF_A100, PERF_A6000, PERF_H100,  # noqa: E402
                        per_fma, predict, run_kernel, scale_counts, smem_wavefronts)
from simt_lab import occupancy  # noqa: E402
from console import fmt_time, print_header, print_observations, print_table  # noqa: E402

N_BIG = 4096


def _mats(n, seed=0):
    rng = np.random.default_rng(seed)
    A = rng.standard_normal((n, n)).astype(np.float32)
    B = rng.standard_normal((n, n)).astype(np.float32)
    return A, B, A.astype(np.float64) @ B.astype(np.float64)


VARIANTS = [(KERNELS[2], {}), (KERNELS[3], {}), (KERNELS[4], {}), (KERNELS[4], {"vec_b": True}), (KERNELS[5], {})]


def _name(spec, kw):
    return spec.name + ("（Bs 用 LDS.128）" if kw.get("vec_b") else "")


def demo1_ledger():
    print_header("演示 1：k3 → k6 的事件账本（每 1000 次乘加；128³ 和 256³ 两种大小）")
    A, B, ref = _mats(128)
    A2, B2, ref2 = _mats(256, seed=1)
    rows, stats, worst = [], {}, 0.0
    for spec, kw in VARIANTS:
        C, st = run_kernel(spec, A, B, **kw)
        C2, st2 = run_kernel(spec, A2, B2, **kw)
        stats[_name(spec, kw)] = st
        p1, p2 = per_fma(st), per_fma(st2)
        main = ("ffma_instr", "lds_instr", "lds_wavefronts", "gld_instr", "gld_lines")
        dev = max(abs(p1[f] - p2[f]) / p1[f] for f in main)
        worst = max(worst, dev)
        err = max(np.abs(C - ref).max(), np.abs(C2 - ref2).max())
        rows.append([_name(spec, kw), f"{err:.1e}", f"{p1['warp_instr'] * 1000:.1f}", f"{p1['ffma_instr'] * 1000:.2f}",
                     f"{p1['lds_instr'] * 1000:.2f}", f"{p1['lds_wavefronts'] * 1000:.2f}",
                     f"{p1['ideal_smem_wavefronts'] * 1000:.2f}", f"{p1['gld_instr'] * 1000:.2f}",
                     f"{p1['gld_lines'] * 1000:.2f}", f"{dev:.1%}"])
    print_table(["内核", "最大误差", "warp 指令", "FFMA", "LDS 指令", "LDS 轮数", "无冲突时的共享内存轮数（含 STS）",
                 "全局 load 指令", "L1 line", "128³ vs 256³ 的最大偏差"], rows)
    print(f"\n最后一列比较的是主循环里的 FFMA、LDS、轮数、全局 load、L1 line（每 1000 次乘加），两种大小最多相差 {worst:.1%}；"
          "\"写 C\"这类每个输出只做一次的指令按 1/n 变化，没有放进比较。")
    print_observations([
        "k3 → k4：一个线程算一列 8 个输出，每读一个 Bs 的数配 8 个 As 的数，LDS 指令从每 1000 次乘加 62.5 条降到 35.2 条。",
        "k4 → k5：一个线程算 8×8 的输出块（外积），每读 16 个数做 64 次乘加，LDS 指令降到 7.8 条；"
        "但轮数是 23.4——比指令数多了 3 倍：k5 有严重的 bank 冲突（演示 2）。",
        "k6 把 As 转置、用 LDS.128 一次读 4 个数：LDS 指令只剩约 2 条、warp 指令总数接近 FFMA 本身；"
        "轮数降到 11.7，但仍高于无冲突的 8.3。",
        "FFMA 一直是 31.25 条/1000 次乘加（一条指令 32 个线程各做一次）；优化的是 FFMA 之外的那些指令和轮数。",
    ])
    return stats


def demo2_bank_conflicts():
    print_header("演示 2：bank 冲突——一个 warp（32 个线程）的访问模式")
    lane = np.arange(32)
    w0 = np.zeros(32, dtype=np.int64)
    trow, tcol = lane // 16, lane % 16          # k5 / k6：256 线程，每行 16 个线程（BN/TN = 16）
    pats = [
        ("k5 读 As：(trow·8 + i)·8 + d（i = d = 0）", (trow * 8) * 8, 1,
         "两个 trow 相差 64 个字 → 同一个 bank 不同的字"),
        ("k5 读 As，每行补 1 个字（行宽 9）", (trow * 8) * 9, 1, "相差 72 个字 → bank 相差 8，不冲突"),
        ("k5 读 Bs：tcol·8 + j（j = 0）", tcol * 8, 1, "16 个 tcol 步长 8 → 只用 4 个 bank"),
        ("k6 读 Bs（LDS.128）：tcol·8", tcol * 8, 4, "每个相位 8 个线程 × 4 个字，步长 8 → 2 路"),
        ("warp 分块式：tcol·4（两组：0 和 64 起）", tcol * 4, 4, "每个相位 8 个线程读连续 32 个字"),
        ("k6 读 As（LDS.128）：trow·8", trow * 8, 4, "同一相位的 8 个线程地址相同（广播）"),
        ("k6 转置写 As：(iAc·4 + q)·128 + iAr", (lane % 2) * 4 * 128 + lane // 2, 1, "iAc = 0/1 相差 512 个字"),
    ]
    rows = []
    for name, word, w, why in pats:
        wf, ideal = smem_wavefronts(word.astype(np.int64), w0, lane, w)
        rows.append([name, {1: "32 位", 4: "128 位"}[w], wf, ideal, f"{wf / ideal:.0f} 路", why])
    print_table(["访问", "宽度", "实际轮数", "无冲突轮数", "冲突", "原因"], rows)
    print_observations([
        "bank = 字号 mod 32。冲突的根源总是\"同一相位里有地址相差 32 的倍数个字、但不是同一个字\"："
        "k5 的 As 行宽 8 个字，两个 trow 相差 64；Bs 每个线程占 8 列，16 个线程只覆盖 4 个 bank 的位置。",
        "两种常用的解法：(1) padding——每行多留一个字，让相邻行错开一个 bank（浪费一点共享内存）；"
        "(2) 改线程映射——让同一相位的线程读连续的地址。Boehm 的 warp 分块（kernel 10）把每个线程的 8 列拆成两组 4 列"
        "（tcol·4 和 64 + tcol·4），LDS.128 就没有冲突了。",
        "LDS.128 即使完全没有冲突也要 4 轮：向量化省的是指令，不是共享内存的带宽。",
        "Tensor Core 的内核里同样的问题用 swizzle（地址异或）解决，见 04 脚本演示 2。",
    ])


def demo3_ladder(stats):
    print_header("演示 3：整条阶梯——模型的上限 vs Boehm 的实测（A6000，4096³ FP32）")
    A, B, _ = _mats(128)
    for spec in KERNELS[:2]:
        stats[spec.name] = run_kernel(spec, A, B)[1]
    order = [(s, {}) for s in KERNELS[:2]] + VARIANTS
    rows = []
    for spec, kw in order:
        p = predict(spec, scale_counts(stats[_name(spec, kw)], N_BIG), N_BIG, PERF_A6000, A6000)
        rows.append([_name(spec, kw), fmt_time(p["t_issue"]), fmt_time(p["t_l1"]), fmt_time(p["t_smem"]),
                     fmt_time(p["t_fma"]), p["bound"], f"{p['gflops']:.0f}", f"{spec.boehm:.0f}",
                     f"{spec.boehm / p['gflops']:.2f}", f"{p['occ']['occupancy']:.0%}"])
    for name, g in BOEHM_LATER.items():
        rows.append([name, "", "", "", "", "", "", f"{g:.0f}", "", ""])
    print_table(["内核", "指令发射", "L1", "共享内存", "FP32 运算", "模型瓶颈", "模型 GFLOP/s", "Boehm 实测", "实测 / 模型",
                 "占用率（假设的寄存器数）"], rows)
    print_observations([
        "模型对每一步给出的瓶颈和 Boehm 用 profiler 看到的一致：k1 卡在不合并的访问，k3 卡在共享内存（MIO throttle），"
        "k4 之后共享内存的压力逐步下降。k1、k4 的实测和模型几乎一样；k2、k3、k6 的实测是模型的 40%–75%（延迟、地址计算没进模型）。",
        "k5 按源码逐句翻译（Bs 用 8 条标量 LDS）时，模型给出约 12.6 TFLOPS，**低于**实测的 16.0——一个吞吐模型的上限不应该低于实测，"
        "说明某个假设错了。最可能的是编译器把 Bs 的 8 个连续读合并成了两条 LDS.128（对齐可以证明）；按这个假设模型是约 19 TFLOPS，"
        "实测是它的 85%。我们看不到 SASS，这只是推断——在真机上用 cuobjdump -sass 一看便知。",
        "k6 之后剩下的差距（18.2 → cuBLAS 23.2）来自：Bs 的 2 路冲突（演示 2）、没有双缓冲（04 脚本）、"
        "以及 warp 分块带来的更好的寄存器复用。Boehm 的 kernel 10（warp 分块）到了 cuBLAS 的 94%。",
    ])


def demo4_how_big():
    print_header("演示 4：寄存器分块多大才够？（不计冲突的理想情况，256 线程的块）")
    gpus = [("RTX A6000", PERF_A6000, A6000), ("A100", PERF_A100, A100), ("H100（FP32）", PERF_H100, None)]
    supply = {name: 128 / perf.fma_per_clk for name, perf, _ in gpus}       # 共享内存 128 字节/周期 ÷ 每周期乘加数
    print("硬件每次乘加最多能从共享内存供应的字节 = 128 字节/周期 ÷ 每 SM 每周期的 FP32 乘加："
          + "，".join(f"{k} {v:.0f} 字节" for k, v in supply.items()))
    rows = []
    for tm, tn in ((1, 1), (2, 2), (4, 4), (8, 4), (8, 8), (16, 8)):
        need = 4 * (tm + tn) / (tm * tn)                                  # 每次乘加从共享内存读的字节
        regs = tm * tn + tm + tn + 24                                     # 累加器 + 操作数 + 地址等（粗估）
        occ6 = occupancy(256, regs, 8192, A6000)
        occ1 = occupancy(256, regs, 8192, A100)
        rows.append([f"{tm}×{tn}", f"{need:.2f}", *[("够" if need <= supply[n] + 1e-9 else f"差 {need / supply[n]:.1f} 倍")
                                                   for n, _, _ in gpus],
                     regs, f"{occ6['occupancy']:.0%}", f"{occ1['occupancy']:.0%}"])
    print_table(["每线程输出 TM×TN", "每次乘加读共享内存（字节）", "A6000（128 乘加/周期）", "A100（64 乘加/周期）",
                 "H100 FP32（128）", "每线程寄存器（粗估）", "占用率 A6000", "占用率 A100"], rows)
    print_observations([
        "每次乘加要从共享内存读 4·(TM + TN)/(TM·TN) 字节；A6000 每个 SM 每周期 128 次 FP32 乘加，共享内存只能给 128 字节，"
        "所以要 TM×TN ≥ 8×8 才喂得饱。A100 每 SM 只有 64 个 FP32 单元，4×4 就够了。",
        "这正好解释了 Boehm 自动调优的结果：A6000 上最好的是 TM = TN = 8（BM = BN = 128、BK = 16），"
        "A100 上最好的是 TM = TN = 4（BM = BN = 64、BK = 16）。更大的分块在 A100 上没有好处，只会多占寄存器、降低占用率。",
        "寄存器分块的代价是寄存器：8×8 要 100 多个寄存器，占用率降到三成左右；16×8 只剩一成多，延迟就难藏了。"
        "\"最优\"是带宽需求和占用率之间的折中，而且随硬件变——这就是要自动调优的原因（第 13 讲）。",
        "Tensor Core 把每 SM 每周期的乘加提高了一个数量级，这张表的\"需求\"就要重新算了（04 脚本演示 3）。",
    ])


if __name__ == "__main__":
    st = demo1_ledger()
    demo2_bank_conflicts()
    demo3_ladder(st)
    demo4_how_big()