"""
01_memory_bound_matmul.py —— 为什么朴素矩阵乘法只能用到 A100 峰值的 2.5%？（第 4 讲 §1）
=====================================================================

【这个脚本要回答的问题】
    1. slides 的计算：朴素矩阵乘法每做 2 次运算读 8 字节，计算–全局内存访问比 0.25 FLOP/字节，
       在 A100 上被带宽限制在 483.75 GFLOPS——FP32 峰值的 2.48%。这个数怎么来的？
    2. 为什么"读数据比算还慢"？（一次全局内存读约 500 个周期，一次加法 1 个周期）
    3. CUDA 里不同的变量声明放在哪种内存、谁能看见、活多久？
    4. 在模拟器上数一数朴素矩阵乘法：线程要求读的字节、实际搬运的 32 字节事务、数据本身有多少——
       三者之间的差距，就是"重复读取"，也就是分块（02 号脚本）能省下来的东西。

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

import numpy as np

from accel import A100_PCIE, H100_SXM, matmul_naive, run_matmul
from console import print_table


def part1_slides_arithmetic():
    print("=" * 78)
    print("第 1 部分：slides 的算术——计算–全局内存访问比与峰值")
    print("=" * 78)
    ratio = 2 / (2 * 4)
    rows = []
    for g in (A100_PCIE, H100_SXM):
        bound = g.bw * ratio
        rows.append([g.name, f"{g.bw / 1e9:,.0f}", f"{g.fp32 / 1e12:g}", f"{g.tf32 / 1e12:g}", f"{bound / 1e9:,.2f}",
                     f"{bound / g.fp32:.2%}", f"{bound / g.tf32:.2%}", f"{g.fp32 / g.bw:.1f}", f"{g.tf32 / g.bw:.1f}"])
    print("  循环体 Pvalue += a[row*N + k] * b[k*N + col]：1 次乘 + 1 次加 = 2 FLOPs；读 a、b 各一个 float = 8 字节")
    print(f"  计算–全局内存访问比 = 2 / 8 = {ratio} FLOP/字节")
    print_table(["GPU", "带宽 GB/s", "FP32 TFLOPS", "TF32 TFLOPS", "带宽能喂饱的 GFLOPS", "占 FP32 峰值", "占 TF32 峰值",
                 "FP32 需要的比值", "TF32 需要的比值"], rows)
    print("\n观察 Observations：")
    a = A100_PCIE
    print(f"  - A100：{a.bw / 1e9:,.0f} GB/s × 0.25 FLOP/字节 = {a.bw * ratio / 1e9:.2f} GFLOPS，是 FP32 峰值的 {a.bw * ratio / a.fp32:.2%}、"
          f"TF32 张量核心峰值的 {a.bw * ratio / a.tf32:.2%}——和 slides 的数字一致。这个内核是访存受限的：算得再快也没用，数据喂不进来。")
    print(f"  - 要让 FP32 核心跑满，比值要提高到 {a.fp32 / a.bw:.1f} FLOP/字节（平衡点）；要让 TF32 张量核心跑满，要 {a.tf32 / a.bw:.0f}。"
          "也就是说，每从显存读一个字节，要设法让它参与几十次运算。\n")


def part2_load_vs_compute():
    print("=" * 78)
    print("第 2 部分：读数据比算还慢——一条加法背后的指令（slides 第 8–9 页）")
    print("=" * 78)
    rows = [["ld.global.f32 %f1, [%rd1];", "读 A[i]（全局内存）", "约 500"],
            ["ld.global.f32 %f2, [%rd2];", "读 B[i]（全局内存）", "约 500"],
            ["add.f32 %f3, %f1, %f2;", "FP32 加法", "约 1"],
            ["st.global.f32 [%rd3], %f3;", "写 C[i]（全局内存）", "（不需要等待）"]]
    print("  C[i] = A[i] + B[i] 编译成的 PTX（GPU 的中间汇编）：")
    print_table(["指令", "做什么", "延迟（周期）"], rows, align=["left", "left", "right"])
    rows = [["寄存器", "≈ 1", "每个线程私有"], ["共享内存 / L1 缓存 / 常量缓存", "≈ 5", "每个 SM 一份，块内共享"],
            ["全局内存（显存，经过 L2）", "≈ 500", "所有 SM 共享"]]
    print("\n  slides 第 8 页的 CUDA 存储层次（访问延迟，量级）：")
    print_table(["存储", "延迟（周期）", "范围"], rows, align=["left", "right", "left"])
    print("\n观察 Observations：")
    print("  - 一次加法 1 个周期，等两个操作数却要约 500 个周期。第 2 讲 §4.7 讲过 GPU 怎样用大量线程束掩盖这段等待；"
          "但掩盖得再好，带宽也是上限（第 1 部分）。")
    print("  - 共享内存比全局内存快约 100 倍，而且就在 SM 上——如果能把要反复用的数据先搬进共享内存，就能少去显存很多次。"
          "这就是分块的出发点。\n")


def part3_memory_types():
    print("=" * 78)
    print("第 3 部分：变量声明 → 放在哪种内存（slides 第 11 页）")
    print("=" * 78)
    rows = [["int var;", "寄存器", "线程", "内核这一次执行（随线程结束）"],
            ["int varArr[N];", "本地内存（local，实际在显存里）", "线程", "同上"],
            ["__device__ __shared__ int SharedVar;", "共享内存", "线程块", "同上（随块结束）"],
            ["__device__ int GlobalVar;", "全局内存", "网格（所有线程）", "整个应用程序"],
            ["__device__ __constant__ int constVar;", "常量内存", "网格（所有线程）", "整个应用程序"]]
    print_table(["变量声明", "内存", "作用域", "生命周期"], rows, align=["left", "left", "left", "left"])
    print("\n观察 Observations：")
    print("  - 内核里普通的局部变量放在寄存器里（最快）；但局部数组（int varArr[N]）通常放不进寄存器，会被放到\"本地内存\"——"
          "名字叫本地，其实在显存里，和全局内存一样慢。写内核时要小心大的局部数组。")
    print("  - 共享内存是程序员手动管理的缓存：要自己决定搬什么进去、什么时候搬、什么时候可以覆盖（__syncthreads）。"
          "常量内存适合所有线程读同一个值的情况（有专门的广播缓存）。\n")


def part4_reuse():
    print("=" * 78)
    print("第 4 部分：在模拟器上数朴素矩阵乘法——要求读的、实际搬的、数据本身")
    print("=" * 78)
    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)
    C, s = run_matmul(matmul_naive, A, B, (16, 16), (2, 2))
    flops = 2 * N ** 3
    req = s.bytes_requested
    sec = s.sector_bytes
    uniq = s.unique_sectors * 32
    rows = [["线程要求读写的字节（每个线程每次访问 4 字节）", f"{req:,}", f"{flops / req:.2f}"],
            ["按 32 字节事务实际搬运的字节（线程束合并之后）", f"{sec:,}", f"{flops / sec:.2f}"],
            ["数据本身的字节（A、B、C 各 N² 个 float；缓存无限大时的最少流量）", f"{uniq:,}", f"{flops / uniq:.2f}"]]
    print(f"  N = {N}，块 16 × 16，网格 2 × 2；一共 {flops:,} FLOPs；结果与 NumPy 的最大差 {np.abs(C - A @ B).max():.1e}")
    print_table(["口径", "字节数", "FLOP / 字节"], rows, align=["left", "right", "right"])
    print(f"\n  线程读次数 = {s.thread_loads:,} = 2 × N³；A 的每个元素被读了 N = {N} 次（同一行的 {N} 个输出都要它），B 同理。")
    print("\n观察 Observations：")
    print(f"  - 按\"每次读都去显存\"算（slides 的口径），比值就是 0.25。线程束合并之后，A 的读是广播（一个线程束里只有 2 行）、"
          f"B 的读是相邻的，实际搬运的字节少了 {req / sec:.0f} 倍，比值 {flops / sec:.2f}——但这些事务仍然要从 L1 / L2 或显存里取。")
    print(f"  - 数据本身只有 {uniq:,} 字节：如果每个数只从显存读一次，比值就是 {flops / uniq:.2f}（= N / 6 量级，N 越大越高）。"
          f"\"要求读的\"和\"数据本身\"差了 {req / uniq:.0f} 倍——这就是重复读取，也是优化的空间。")
    print("  - 真实 GPU 的 L1 / L2 缓存会自动接住一部分重复读取，但容量有限、命中与否不受程序员控制。"
          "分块（02 号脚本）用共享内存把\"重复使用\"变成程序里明确的、有保证的事情。\n")


def main():
    part1_slides_arithmetic()
    part2_load_vs_compute()
    part3_memory_types()
    part4_reuse()


if __name__ == "__main__":
    main()
"""
02_tiled_matmul.py —— 分块矩阵乘法：用共享内存把全局内存读取减少到 1/TILE（第 4 讲 §2）
=====================================================================

【这个脚本要回答的问题】
    1. slides 的 MatMulTiledKernel：块大小 TILE 从 1 到 32，全局内存读取怎样变化？计算–访存比和 A100 上能达到的峰值比例呢？
    2. 两个 __syncthreads() 各防的是什么？漏掉一个会怎样？
    3. N 不是 TILE 的倍数时，slides 的写法会出什么事？怎样修？
    4. TILE 能无限大吗？——共享内存、每块线程数、寄存器的限制，以及对占用率的影响。
    5. 线程粗化：一个块算横向相邻的几个输出块，A 的读取还能再省多少？

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

import os
import sys

import numpy as np

from accel import A100_PCIE, matmul_naive, matmul_tiled, matmul_tiled_coarse, run_matmul
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 A100, H100, occupancy  # noqa: E402
from simt import IllegalAddress  # noqa: E402


def mats(N, seed=0):
    rng = np.random.default_rng(seed)
    return rng.standard_normal((N, N)).astype(np.float32), rng.standard_normal((N, N)).astype(np.float32)


def part1_sweep():
    print("=" * 78)
    print("第 1 部分：TILE 从 1 到 32——全局内存读取、共享内存读取、计算–访存比")
    print("=" * 78)
    N = 32
    A, B = mats(N)
    C0, s0 = run_matmul(matmul_naive, A, B, (16, 16), (2, 2))
    rows = [["朴素（不用共享内存）", s0.thread_loads, s0.ld_sectors, 0, 0, "—", "0.25",
             f"{min(1, 0.25 * A100_PCIE.bw / A100_PCIE.fp32):.1%}"]]
    res = {}
    for T in (2, 4, 8, 16, 32):
        C, s = run_matmul(matmul_tiled, A, B, (T, T), (N // T, N // T), T)
        ratio = 2 * N ** 3 / (s.thread_loads * 4)
        res[T] = s
        rows.append([f"TILE = {T}（块 {T} × {T}）", s.thread_loads, s.ld_sectors, s.thread_shared, s.barriers,
                     np.array_equal(C, C0), f"{ratio:g}", f"{min(1, ratio * A100_PCIE.bw / A100_PCIE.fp32):.1%}"])
    print(f"  N = {N}，C = A · B，一共 {2 * N ** 3:,} FLOPs：")
    print_table(["写法", "线程读全局内存（次）", "全局读事务", "共享内存访问（次）", "栅栏次数", "与朴素版逐位相同",
                 "计算–全局访存比 FLOP/字节", "A100 FP32 能达到的峰值比例"], rows)
    print("\n观察 Observations：")
    print(f"  - 线程读全局内存的次数正好是 2N³ / TILE：TILE = 16 时从 {s0.thread_loads:,} 次降到 {res[16].thread_loads:,} 次。"
          "原因：块里同一行的 TILE 个线程都要 A 的同一个元素，同一列的 TILE 个线程都要 B 的同一个元素——"
          "现在由一个线程把它搬进共享内存，大家从共享内存里读。")
    sh_reads = res[16].thread_shared - res[16].thread_loads
    print(f"  - 运算的总量没变，只是读的地方变了：TILE = 16 时共享内存的读取 {sh_reads:,} 次（和朴素版的全局读取一样多），"
          f"外加搬进来时写共享内存的 {res[16].thread_loads:,} 次。共享内存比全局内存快约 100 倍，这笔交换很划算。")
    print(f"  - 一个反常的数：TILE = 2 时线程读全局内存的次数减半，全局读事务却是朴素版的 {res[2].ld_sectors / s0.ld_sectors:.0f} 倍"
          f"（{res[2].ld_sectors:,} vs {s0.ld_sectors:,}）。2 × 2 的块只有 4 个线程，凑不满一个线程束，每条读指令只取 2 个数，"
          "合并访存完全失效。块太小，省下的读取次数被浪费的事务吃掉了——块的一行至少要有 8 个线程（第 3 讲 §2.5）。")
    print("  - 分块之后每个线程仍然按 k = 0, 1, …, N−1 的顺序累加，所以结果和朴素版逐位相同——优化没有改变数值。")
    print(f"  - 计算–访存比 = 0.25 × TILE：TILE = 16 时 4 FLOP/字节，A100 上能用到 FP32 峰值的 {min(1, 4 * A100_PCIE.bw / A100_PCIE.fp32):.0%}；"
          f"TILE = 32 时 8 FLOP/字节，{min(1, 8 * A100_PCIE.bw / A100_PCIE.fp32):.0%}。要超过平衡点（约 10）还需要更多手段（第 5 部分、05 号脚本）。"
          f"代价是栅栏：TILE = 2 时一共 {res[2].barriers:,} 次，TILE = 32 时只有 {res[32].barriers} 次。\n")


def part2_sync_bugs():
    print("=" * 78)
    print("第 2 部分：两个 __syncthreads() 各防什么？")
    print("=" * 78)
    N, T = 32, 8
    A, B = mats(N, 1)
    C0, _ = run_matmul(matmul_naive, A, B, (16, 16), (2, 2))
    rows = []
    for name, s1, s2 in (("两个栅栏都在（slides 的写法）", True, True), ("漏掉第二个（算完之后）", True, False),
                         ("漏掉第一个（搬完之后）", False, True)):
        C, s = run_matmul(matmul_tiled, A, B, (T, T), (N // T, N // T), T, True, s1, s2)
        wrong = int(np.sum(C != C0))                       # NaN 和任何数都不相等，也算错
        rows.append([name, f"{wrong} / {N * N}", int(np.isnan(C).sum()), s.race_count, s.races[0] if s.races else "—"])
    print_table(["写法", "错误的元素", "其中 NaN", "检测到的竞争", "第一条竞争记录"], rows,
                align=["left", "right", "right", "right", "left"])
    print("\n观察 Observations：")
    print("  - 第一个栅栏（搬完之后）保证\"读之前大家都写完了\"：漏掉它，线程会读到同伴还没搬进来的位置——"
          "模拟器的共享内存初始是 NaN，所以结果里出现了 NaN（真 GPU 上是上一相留下的旧值或垃圾值，错得更隐蔽）。"
          "1024 个元素里只有 16 个碰巧是对的：模拟器按编号依次执行线程，每个块（共 16 个）里编号最大的那个线程运行时，它要的数都已经被前面的线程搬好了。")
    print("  - 第二个栅栏（算完之后）保证\"写之前大家都读完了\"：漏掉它，跑得快的线程会进入下一相、覆盖 As / Bs，"
          "而别的线程还在用旧的那一块。结果是错的数，而不是 NaN。")
    print("  - 这两类竞争分别叫\"写后读\"（read-after-write）和\"读后写\"（write-after-read）。模拟器按固定顺序执行线程，"
          "真 GPU 上的执行顺序不确定，所以同样的 bug 在真 GPU 上可能时有时无——竞争检测器在两种情况下都报告了冲突。\n")


def part3_bounds():
    print("=" * 78)
    print("第 3 部分：N 不是 TILE 的倍数")
    print("=" * 78)
    N, T = 10, 4
    A, B = mats(N, 2)
    ref = A @ B
    rows = []
    try:
        run_matmul(matmul_tiled, A, B, (T, T), (3, 3), T, False)
        r = "没有报错"
    except IllegalAddress as e:
        r = f"越界：{e}"
    rows.append([f"slides 的写法，网格 ⌈{N}/{T}⌉ = 3", r])
    C, _ = run_matmul(matmul_tiled, A, B, (T, T), (2, 2), T, False)
    filled = C[:8, :8]
    rows.append([f"slides 的写法，网格 {N}//{T} = 2", f"没有报错；只算了左上角 8 × 8，而且每个元素只加了前 {(N // T) * T} 项："
                 f"左上角与正确结果的最大差 {np.abs(filled - ref[:8, :8]).max():.2f}，其余 {N * N - 64} 个元素是 cudaMalloc 之后的垃圾值（NaN）"])
    C, _ = run_matmul(matmul_tiled, A, B, (T, T), (3, 3), T, True)
    rows.append([f"修正：相数 ⌈{N}/{T}⌉ = 3，越界位置补 0，写之前检查", f"与 NumPy 的最大差 {np.abs(C - ref).max():.1e}"])
    print_table(["写法", "结果"], rows, align=["left", "left"])
    print("\n观察 Observations：")
    print(f"  - slides 的 for (ph = 0; ph < N/TILE_WIDTH; ++ph) 用整数除法，N = {N}、TILE = {T} 时只做 {N // T} 相，"
          f"漏掉了最后 {N - (N // T) * T} 项；而且没有越界检查，网格向上取整后边上的线程会读到矩阵外面。")
    print("  - 修正的写法：相数向上取整；搬运时如果行或列超出 N，就往共享内存里写 0（0 乘任何数都是 0，不影响结果）；"
          "最后只让在矩阵内的线程写 C。注意越界的线程不能提前 return——它们还要参与搬运和栅栏（第 2 讲的 BarrierError）。\n")


def part4_limits():
    print("=" * 78)
    print("第 4 部分：TILE 能有多大？——资源限制与占用率")
    print("=" * 78)
    rows = []
    for T in (8, 16, 32, 64):
        threads = T * T
        smem = 2 * T * T * 4
        row = [T, threads, f"{smem / 1024:g} KB"]
        if threads > 1024:
            row += ["超过每块 1024 个线程，无法发射", "—", "—"]
        else:
            for g in (A100, H100):
                o = occupancy(g, threads, 32, smem)
                row.append(f"{o['blocks']} 块 / {o['warps']} 束 = {o['occupancy']:.0%}（{o['limiter']}）")
            row.append(f"{0.25 * T:g}")
        rows.append(row)
    print("  每个线程算一个输出、块大小 TILE × TILE，共享内存 As + Bs = 2 × TILE² × 4 字节，假设每线程 32 个寄存器：")
    print_table(["TILE", "每块线程", "每块共享内存", "A100 每 SM 驻留", "H100 每 SM 驻留", "计算–访存比"], rows)
    print("\n观察 Observations：")
    print("  - TILE = 32 时每块 1024 个线程、8 KB 共享内存（slides：As、Bs 各 32 × 32 × 4 = 4 KB），离每块的共享内存上限还很远；"
          "真正先撞上的是\"每块最多 1024 个线程\"——TILE = 64 根本发射不了。")
    print("  - 所以要让每个块覆盖更大的输出块，就不能再\"一个线程一个输出\"：要让每个线程算好几个输出（第 5 部分的粗化、"
          "以及 cuBLAS 的寄存器分块），用寄存器再做一层复用。")
    print("  - slides 第 21 页说\"1024 个线程时每个线程只能用 16 个寄存器\"：那是用一个分区的 16384 个寄存器去除。"
          "整个 SM 有 65536 个，1024 个线程的块分布在 4 个分区上，每个线程最多 64 个；SM 驻留满 2048 个线程时每个 32 个（第 2 讲 §3.5）。\n")


def part5_coarsening():
    print("=" * 78)
    print("第 5 部分：线程粗化——一个块算横向相邻的几个输出块")
    print("=" * 78)
    N, T = 32, 8
    A, B = mats(N, 3)
    C0, _ = run_matmul(matmul_naive, A, B, (16, 16), (2, 2))
    rows = []
    res = {}
    for CF in (1, 2, 4):
        C, s = run_matmul(matmul_tiled_coarse, A, B, (T, T), (N // (T * CF), N // T), T, CF)
        a, b = s.by_array["A"]["thread_ld"], s.by_array["B"]["thread_ld"]
        res[CF] = (a, b, s)
        rows.append([CF, s.blocks, CF, a, b, a + b, f"{2 * N ** 3 / ((a + b) * 4):.2f}", np.array_equal(C, C0)])
    print(f"  N = {N}，TILE = {T}；COARSE = 每个线程算几个输出（每个块覆盖 1 × COARSE 个输出块）：")
    print_table(["COARSE", "块数", "每线程的累加器", "读 A（次）", "读 B（次）", "全局读合计", "计算–全局访存比", "与朴素版逐位相同"], rows)
    print("\n观察 Observations：")
    print(f"  - 不粗化时，横向相邻的块要的是同一条 A 的横带，却各自从全局内存读一遍。COARSE = 4 时 A 的读取从 {res[1][0]} 次降到 "
          f"{res[4][0]} 次（1/4），B 不变，合计少了 {1 - (res[4][0] + res[4][1]) / (res[1][0] + res[1][1]):.0%}。")
    print(f"  - 代价：每个线程要 COARSE 个累加器（更多寄存器 → 占用率可能下降），块数从 {res[1][2].blocks} 降到 {res[4][2].blocks}"
          "（并行度降低；N 小时 GPU 可能喂不饱）。粗化是\"用更少的线程做更多的事\"，在并行度过剩时才划算。")
    print("  - 这已经接近 cuBLAS 的思路：块在共享内存里分块，线程在寄存器里再分块（每个线程算一个 8 × 8 之类的小块），"
          "让每个从显存读来的字节参与几十次运算（05 号脚本）。\n")


def main():
    part1_sweep()
    part2_sync_bugs()
    part3_bounds()
    part4_limits()
    part5_coarsening()


if __name__ == "__main__":
    main()
"""
03_coalescing_transpose.py —— 内存并行：合并访存、矩阵转置的三个版本、共享内存的 bank 冲突（第 4 讲 §3）
=====================================================================

【这个脚本要回答的问题】
    1. 显存每次至少搬 32 字节。一个线程束的 32 个线程按不同的步长、不同的起点读 float，各要几个事务？
    2. 矩阵转置为什么难？slides 的三个版本：
         naive        ：读合并、写不合并
         共享内存      ：读写都合并，但写共享内存时有 32 路 bank 冲突
         共享内存 + 1 ：数组多开一列，bank 冲突消失
    3. bank 冲突到底是怎么回事？"+1"为什么能消除它？
    4. 放到 A100 上的真实尺寸（8192 × 8192），三个版本各要搬多少字节、多少时间？

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

import numpy as np

from accel import A100_PCIE, transpose_naive, transpose_smem, transpose_smem_padded
from console import print_table
from simt import Device, launch


def strided_read(t, A, OUT, stride, offset):
    """线程 lane 读 A[offset + lane * stride]（只发射一个线程束）。"""
    lane = t.threadIdx.x
    OUT[lane] = A[offset + lane * stride]


def part1_coalescing():
    print("=" * 78)
    print("第 1 部分：一个线程束读 32 个 float——步长和起点决定事务数")
    print("=" * 78)
    dev = Device()
    A = dev.to_device(np.arange(32 * 64 + 64, dtype=np.float32), "A")
    OUT = dev.malloc(32, np.float32, "OUT")
    rows = []
    for stride, offset in ((1, 0), (1, 1), (1, 8), (2, 0), (4, 0), (8, 0), (16, 0), (32, 0), (0, 0)):
        s = launch(strided_read, 1, 32, A, OUT, stride, offset)
        rec = s.by_array["A"]
        useful = 4 if stride == 0 else 128
        rows.append([stride, offset, rec["ld_sectors"], rec["ld_sectors"] * 32, f"{useful / (rec['ld_sectors'] * 32):.1%}"])
    print_table(["步长（元素）", "起点偏移（元素）", "事务（32 字节）", "搬运的字节", "有用字节占比"], rows)
    print("\n观察 Observations：")
    print("  - 步长 1、起点对齐：32 个 float = 128 字节 = 4 个 sector，一个字节都不浪费——这就是合并访存（coalesced）。"
          "slides：连续的访问被合并成一次，\"让 DRAM 突发传输成为可能\"。")
    print("  - 起点偏移 1 个元素（4 字节）：128 字节跨了 5 个 sector；偏移 8 个元素（32 字节）又对齐了。"
          "cudaMalloc 返回的地址至少 256 字节对齐，所以\"从数组开头按线程编号读\"总是对齐的。")
    print("  - 步长 2、4、8：事务数 8、16、32，有用的字节降到 50%、25%、12.5%；步长 ≥ 8 之后每个线程独占一个 sector，"
          "再大也不会更糟（但会碰到更多不同的缓存行）。步长 0（所有线程读同一个数）是广播，1 个事务。\n")


def part2_transpose():
    print("=" * 78)
    print("第 2 部分：slides 的三个转置内核")
    print("=" * 78)
    m = 64
    rng = np.random.default_rng(0)
    a = rng.standard_normal((m, m)).astype(np.float32)
    rows = []
    res = {}
    for name, k in (("naive_cuda_transpose", transpose_naive), ("smem_cuda_transpose", transpose_smem),
                    ("smem_cuda_transpose（[TX][TY + 1]）", transpose_smem_padded)):
        dev = Device()
        da, dc = dev.to_device(a, "a"), dev.malloc(m * m, np.float32, "c")
        s = launch(k, (m // 32, m // 32), (32, 32), m, da, dc)
        ok = np.array_equal(dev.to_host(dc, (m, m)), a.T)
        ra, rc = s.by_array["a"], s.by_array["c"]
        res[name] = s
        rows.append([name, ok, f"{ra['ld_sectors'] / ra['ld_requests']:.0f}", f"{rc['st_sectors'] / rc['st_requests']:.0f}",
                     s.global_sectors, s.sh_requests, s.sh_wavefronts, f"{s.bank_conflict_factor:.1f}" if s.sh_requests else "—"])
    print(f"  {m} × {m} 的 float32 矩阵，块 32 × 32（每个块转置一个 32 × 32 的小块）：")
    print_table(["内核", "结果正确", "读 a：事务/请求", "写 c：事务/请求", "全局事务合计", "共享内存请求", "共享内存轮数",
                 "每请求轮数"], rows)
    n, sm, pd = (res[k] for k in res)
    print("\n观察 Observations：")
    print("  - naive：读 a 时一个线程束读一行的 32 个相邻元素（4 个事务）；写 c 时 c[myCol][myRow] 的 myCol 随 threadIdx.x 变化，"
          f"相邻线程写同一列、相隔 {m} × 4 字节，每个线程独占一个 sector（32 个事务）。转置的本质矛盾：读和写，总有一边是按列走的。")
    print(f"  - 共享内存版：先把一个 32 × 32 的小块按行（合并地）读进共享内存，__syncthreads()，再按行（合并地）写出它的转置——"
          f"\"按列走\"挪到了共享内存里。全局事务从 {n.global_sectors} 降到 {sm.global_sectors}。")
    print(f"  - 但写共享内存 smemArray[threadIdx.x][threadIdx.y] 时，一个线程束（threadIdx.x = 0..31）写的是同一列的 32 个位置，"
          f"它们落在同一个 bank 上（第 3 部分），只能排成 32 轮：写的 128 个请求每个 32 轮，加上读的 128 个请求每个 1 轮，"
          f"一共 {sm.sh_requests} 个请求、{sm.sh_wavefronts} 轮（平均 {sm.bank_conflict_factor:.1f}）。")
    print(f"  - 声明成 [32][33]（多一列，不用）之后，同一列的 32 个位置错开到 32 个不同的 bank：{pd.sh_wavefronts} 轮，"
          "每个请求 1 轮。多花 32 × 4 = 128 字节的共享内存，换掉了 32 路冲突。\n")


def part3_banks():
    print("=" * 78)
    print("第 3 部分：bank 冲突的算术——第 k 个 4 字节字在第 k mod 32 号 bank")
    print("=" * 78)
    T = 32
    rows = []
    for name, width in (("smemArray[32][32]，写 [tx][ty]（ty = 0）", T), ("smemArray[32][33]，写 [tx][ty]（ty = 0）", T + 1),
                        ("smemArray[32][32]，读 [ty][tx]（ty = 0）", None)):
        if width is None:
            words = [0 * T + tx for tx in range(32)]
        else:
            words = [tx * width + 0 for tx in range(32)]
        banks = [w % 32 for w in words]
        rounds = max(banks.count(b) for b in set(banks))
        rows.append([name, ", ".join(str(w) for w in words[:6]) + ", …", ", ".join(str(b) for b in banks[:6]) + ", …",
                     len(set(banks)), rounds])
    print_table(["访问", "线程 0–5 访问的字号", "对应的 bank", "用到的不同 bank", "需要几轮"], rows, align=["left", "left", "left", "right", "right"])
    print("\n观察 Observations：")
    print("  - 共享内存分成 32 个 bank，每个 bank 每轮能提供一个 4 字节字；连续的字依次落在 bank 0、1、……、31、0、1……（slides 第 31 页）。"
          "一个线程束的 32 个线程如果访问 32 个不同 bank 的字，一轮就完成；落在同一个 bank 的不同字要排队。")
    print("  - [32][32] 的数组按列写：线程 tx 访问第 tx × 32 号字，全都是 bank 0——32 路冲突。"
          "多开一列变成 [32][33] 后，线程 tx 访问第 tx × 33 号字，bank = tx × 33 mod 32 = tx，32 个线程正好错开。")
    print("  - 按行读 [ty][tx]：连续的 32 个字，天然没有冲突。规律和全局内存的合并访存正好对应：线程束里相邻的线程访问相邻的字，"
          "两种内存都最高兴。\n")


def part4_paper():
    print("=" * 78)
    print("第 4 部分：放到 A100 上——8192 × 8192 的转置要多久（纸面估算）")
    print("=" * 78)
    m_small, m_big = 64, 8192
    rng = np.random.default_rng(1)
    a = rng.standard_normal((m_small, m_small)).astype(np.float32)
    scale = (m_big / m_small) ** 2
    data_bytes = 2 * m_big * m_big * 4
    rows = []
    for name, k in (("naive", transpose_naive), ("共享内存", transpose_smem), ("共享内存 + 1", transpose_smem_padded)):
        dev = Device()
        da, dc = dev.to_device(a, "a"), dev.malloc(m_small * m_small, np.float32, "c")
        s = launch(k, (2, 2), (32, 32), m_small, da, dc)
        moved = s.sector_bytes * scale
        t = moved / A100_PCIE.bw
        rows.append([name, f"{moved / 1e6:,.0f}", f"{t * 1e3:.3f}", f"{data_bytes / t / 1e9:,.0f}"])
        if name == "naive":
            naive_moved, naive_bw = moved, data_bytes / t
    print(f"  每个线程束的事务数按 64 × 64 的模拟结果等比放大；时间 = 搬运的字节 ÷ {A100_PCIE.bw / 1e9:,.0f} GB/s（不计缓存、不计 bank 冲突）")
    print_table(["版本", "搬运的字节 MB", "时间 ms", "有效带宽 GB/s（读 + 写的数据量 ÷ 时间）"], rows)
    print("\n观察 Observations：")
    print(f"  - 转置本身没有任何运算，数据量是读写各一遍：{data_bytes / 1e6:,.0f} MB。naive 版本因为写不合并，"
          f"要搬 {naive_moved / data_bytes:.1f} 倍的字节，有效带宽只有峰值的 {naive_bw / A100_PCIE.bw:.0%}。")
    print("  - 两个共享内存版本搬的字节一样多（都完全合并）；它们的差别在共享内存的 bank 冲突上，这个纸面模型没有计入——"
          "真实 GPU 上，带冲突的版本会因为共享内存的 32 路排队而明显变慢，这就是 slides 最后一步加 +1 的原因。")
    print("  - 这里假设没有缓存。真实的 L2 会接住 naive 版本的一部分写（同一个 sector 后来又被别的线程束写到），"
          "所以实测差距通常没有这么大，但趋势一样：转置这类\"纯搬运\"的内核，性能完全取决于访存模式。\n")


def main():
    part1_coalescing()
    part2_transpose()
    part3_banks()
    part4_paper()


if __name__ == "__main__":
    main()
"""
04_sparse_csr.py —— 稀疏矩阵：CSR 格式、SpMV 内核、什么时候稀疏才划算、2:4 结构化稀疏（第 4 讲 §4）
=====================================================================

【这个脚本要回答的问题】
    1. CSR（压缩稀疏行）用哪三个数组存一个稀疏矩阵？slides 的 SpMV 循环怎样用它们？
    2. 非零元占多少比例以下，CSR 才比稠密存储省空间、省时间？
    3. slides 的 SpMVCSRKernel（一个线程一行）在 GPU 上有什么问题？换成"一个线程束一行"呢？
       各行长度很不均匀时（真实的稀疏矩阵常常如此）会怎样？
    4. 第 2 讲规格表里"含稀疏加速"的数字，指的是哪种稀疏？

【怎么运行】
    python3 04_sparse_csr.py       （只需要 numpy；约 4 秒，第 3 部分开了指令计数）
"""

import numpy as np

from accel import A100_PCIE, dense_to_csr, spmv_csr_thread, spmv_csr_warp
from console import print_table
from simt import Device, ceil_div, launch


def spmv_sequential(data, col_index, row_ptr, x):
    """slides 第 42 页的串行 SpMV：逐行把非零元乘上 x 的对应位置再加起来（float32，按存储顺序累加）。"""
    n = len(row_ptr) - 1
    y = np.zeros(n, np.float32)
    for row in range(n):
        dot = np.float32(0)
        for el in range(row_ptr[row], row_ptr[row + 1]):
            dot += x[col_index[el]] * data[el]
        y[row] += dot
    return y


def run_spmv(kernel, M, x, warp_per_row=False, **kw):
    """把稠密矩阵 M 转成 CSR、拷到模拟的显存里、发射 SpMV 内核。【返回】(y, LaunchStats)"""
    data, col, ptr = dense_to_csr(M)
    dev = Device()
    d = [dev.to_device(a, n) for a, n in ((data, "data"), (col, "col_index"), (ptr, "row_ptr"), (x, "x"))]
    dy = dev.to_device(np.zeros(M.shape[0], np.float32), "y")
    rows = M.shape[0]
    if warp_per_row:
        s = launch(kernel, ceil_div(rows, 4), 128, *d, dy, rows, **kw)        # 每块 4 个线程束 = 4 行
    else:
        s = launch(kernel, ceil_div(rows, 64), 64, *d, dy, rows, **kw)
    return dev.to_host(dy), s


def part1_format():
    print("=" * 78)
    print("第 1 部分：CSR 的三个数组")
    print("=" * 78)
    M = np.array([[3, 0, 1, 0], [0, 0, 0, 0], [0, 2, 4, 1], [1, 0, 0, 1]], np.float32)
    data, col, ptr = dense_to_csr(M)
    x = np.array([1, 2, 3, 4], np.float32)
    print("  稠密矩阵 M：")
    for r in M:
        print("    ", r.astype(int).tolist())
    print(f"  data      = {data.astype(int).tolist()}     （非零元的值，按行优先顺序）")
    print(f"  col_index = {col.tolist()}     （每个非零元在第几列）")
    print(f"  row_ptr   = {ptr.tolist()}           （第 r 行的非零元是 data[row_ptr[r] : row_ptr[r+1]]）")
    y_seq = spmv_sequential(data, col, ptr, x)
    y_gpu, _ = run_spmv(spmv_csr_thread, M, x)
    print(f"\n  y = M · x，x = {x.astype(int).tolist()}：串行循环 {y_seq.astype(int).tolist()}，模拟器上的内核 {y_gpu.astype(int).tolist()}，"
          f"稠密乘法 {(M @ x).astype(int).tolist()}")
    print("\n观察 Observations：")
    print(f"  - 16 个元素只存了 {len(data)} 个非零元。row_ptr 比行数多一个元素：第 1 行是空行，row_ptr[1] = row_ptr[2] = {ptr[1]}，"
          "内层循环一次也不执行——空行不需要任何特殊处理。")
    print("  - SpMV 里 x 是按 col_index 间接取的（gather）：x[col_index[el]]。这是稀疏计算和稠密计算最大的不同——"
          "访问哪里由数据决定，编译器和硬件都没法提前知道。\n")


def part2_breakeven():
    print("=" * 78)
    print("第 2 部分：多稀疏才划算？——存储与 SpMV 时间（n = 4096，纸面估算）")
    print("=" * 78)
    n = 4096
    rows = []
    for d in (0.001, 0.01, 0.1, 0.25, 0.5, 0.75):
        nnz = int(d * n * n)
        dense_b = 4 * n * n
        csr_b = 8 * nnz + 4 * (n + 1)
        t_dense = (dense_b + 8 * n) / A100_PCIE.bw
        t_csr_best = (csr_b + 8 * n) / A100_PCIE.bw                       # x 只读一次（全部命中缓存）
        t_csr_worst = (csr_b + 32 * nnz + 4 * n) / A100_PCIE.bw          # 每次 gather 都单独搬一个 32 字节 sector
        rows.append([f"{d:.1%}", f"{nnz:,}", f"{dense_b / 2 ** 20:.0f}", f"{csr_b / 2 ** 20:.1f}", f"{csr_b / dense_b:.0%}",
                     f"{t_dense * 1e6:.1f}", f"{t_csr_best * 1e6:.1f}", f"{t_csr_worst * 1e6:.1f}"])
    print_table(["非零元比例", "非零元个数", "稠密 MiB", "CSR MiB", "CSR / 稠密", "稠密 matvec μs", "CSR SpMV μs（x 全命中缓存）",
                 "CSR SpMV μs（每次 gather 一个 sector）"], rows)
    print("\n观察 Observations：")
    print("  - CSR 每个非零元要存值和列号两样（各 4 字节），所以非零元比例超过 50% 时反而比稠密存储更大。")
    print("  - SpMV 和稠密的矩阵-向量乘法都是访存受限的（每个矩阵元素只用一次），时间≈搬运的字节 ÷ 带宽。"
          "x 能全部留在缓存里时，CSR 的分界点也在 50% 左右；x 太大、gather 次次落空时，每个非零元要多搬一个 32 字节的 sector，"
          "分界点降到 10% 左右。")
    print("  - 所以\"稀疏\"要足够稀疏才值得：图、推荐系统的交互矩阵常常是 0.1% 以下；神经网络剪枝后 50% 稀疏的权重，"
          "用 CSR 反而更慢——这就是为什么硬件厂商改用固定模式的结构化稀疏（第 4 部分）。\n")


def make_matrix(rows, cols, lengths, seed):
    """生成一个每行非零元个数为 lengths[r] 的随机稀疏矩阵（列位置随机，值为标准正态）。"""
    rng = np.random.default_rng(seed)
    M = np.zeros((rows, cols), np.float32)
    for r, L in enumerate(lengths):
        c = rng.choice(cols, size=L, replace=False)
        M[r, c] = rng.standard_normal(L).astype(np.float32)
    return M


def part3_kernels():
    print("=" * 78)
    print("第 3 部分：一个线程一行 vs 一个线程束一行（模拟器）")
    print("=" * 78)
    R, Ccols = 256, 1024
    rng = np.random.default_rng(0)
    skew = np.full(R, 2)
    skew[rng.choice(R, 8, replace=False)] = 200                       # 248 行各 2 个非零元，8 行各 200 个
    cases = [("均匀：每行 8 个", np.full(R, 8)), ("均匀：每行 64 个", np.full(R, 64)), ("不均匀：多数 2 个、少数 200 个", skew)]
    rows = []
    res = {}
    for name, lengths in cases:
        M = make_matrix(R, Ccols, lengths, seed=1)
        x = rng.standard_normal(Ccols).astype(np.float32)
        ref64 = M.astype(np.float64) @ x.astype(np.float64)
        data, col, ptr = dense_to_csr(M)
        for kname, k, wpr in (("一个线程一行（slides）", spmv_csr_thread, False), ("一个线程束一行", spmv_csr_warp, True)):
            y, s = run_spmv(k, M, x, warp_per_row=wpr, count_instructions=True)
            exact = np.array_equal(y, spmv_sequential(data, col, ptr, x)) if not wpr else "—"
            dc = s.by_array["data"]["ld_sectors"] + s.by_array["col_index"]["ld_sectors"]
            res[(name, kname)] = (s, dc)
            rows.append([name, kname, int(lengths.sum()), exact, f"{np.abs(y - ref64).max():.1e}",
                         f"{dc:,}", f"{s.ld_sectors:,}", f"{s.warp_issues:,}", f"{s.simt_efficiency:.0%}"])
    print(f"  {R} 行 × {Ccols} 列，x 是随机向量：")
    print_table(["矩阵", "内核", "非零元", "与串行循环逐位相同", "与 float64 的最大差", "读 data + col_index 的事务",
                 "全部读事务", "线程束发射次数", "SIMT 效率"], rows)
    (t8, t8d), (w8, w8d) = res[("均匀：每行 8 个", "一个线程一行（slides）")], res[("均匀：每行 8 个", "一个线程束一行")]
    (t64, t64d), (w64, w64d) = res[("均匀：每行 64 个", "一个线程一行（slides）")], res[("均匀：每行 64 个", "一个线程束一行")]
    (ts, _), (ws, _) = res[("不均匀：多数 2 个、少数 200 个", "一个线程一行（slides）")], res[("不均匀：多数 2 个、少数 200 个", "一个线程束一行")]
    print("\n观察 Observations：")
    print(f"  - 访存：一个线程一行时，相邻线程同时读各自那一行的第 k 个非零元，位置相隔一整行——不合并；一个线程束一行时，"
          f"32 个线程读同一行连续的 32 个非零元。每行 64 个时，读 data 和 col_index 的事务 {t64d:,} vs {w64d:,}（{t64d / w64d:.0f} 倍）；"
          f"算上 x 的 gather（两种写法都是随机访问，谁也合并不了），全部读事务 {t64.ld_sectors:,} vs {w64.ld_sectors:,}（{t64.ld_sectors / w64.ld_sectors:.1f} 倍）。")
    print(f"  - 指令：一个线程束一行的发射次数多得多——每行 64 个时 {w64.warp_issues:,} vs {t64.warp_issues:,}，每行 8 个时 "
          f"{w8.warp_issues:,} vs {t8.warp_issues:,}。每一行都要一个线程束、都要做一次 5 轮的树形归约；行短时大部分车道闲着"
          f"（每行 8 个只用 8 个车道读数）。")
    print("  - SpMV 是访存受限的，真 GPU 上通常是访存那一项说了算，所以长行一般用一个线程束一行；行很短时，事务只省了 "
          f"{t8.ld_sectors / w8.ld_sectors:.1f} 倍，却要多付几十倍的指令，这时常用\"几个线程一行\"（例如 4 或 8 个）的折中。")
    print(f"  - 行长不均匀时，一个线程一行的线程束要等它里面最长的那一行：SIMT 效率只有 {ts.simt_efficiency:.0%}（第 2 讲的工作量不均）；"
          f"一个线程束一行把长行分给 32 个线程，效率 {ws.simt_efficiency:.0%}。实际的库（cuSPARSE）会按行长分组、"
          "对不同的组用不同的内核，或者换别的格式（ELL、COO、混合格式）。没有一种分工在所有矩阵上都最好。")
    print("  - 数值：一个线程一行和串行循环的累加顺序相同，逐位一致；一个线程束一行先分 32 份再树形合并，顺序不同，"
          "和 float64 的差都在 1e-6 到 1e-5 量级。\n")


def part4_structured():
    print("=" * 78)
    print("第 4 部分：2:4 结构化稀疏——规格表上\"含稀疏\"的数字指的是它")
    print("=" * 78)
    rng = np.random.default_rng(3)
    W = rng.standard_normal((4, 8)).astype(np.float32)
    groups = W.reshape(4, 2, 4)
    keep = np.argsort(-np.abs(groups), axis=2)[:, :, :2]
    keep.sort(axis=2)
    pruned = np.zeros_like(groups)
    np.put_along_axis(pruned, keep, np.take_along_axis(groups, keep, axis=2), axis=2)
    pruned = pruned.reshape(4, 8)
    values = np.take_along_axis(groups, keep, axis=2).reshape(4, 4)
    print("  每 4 个连续的权重里只保留绝对值最大的 2 个（其余置 0）：")
    print("  剪枝后的 W（4 × 8）：")
    for r in pruned:
        print("    ", " ".join(f"{v:6.2f}" for v in r))
    print(f"  压缩存储：值（4 × 4）+ 每个值在组内的位置（2 位）：位置 = {keep.reshape(4, 4).tolist()}")
    x = rng.standard_normal(8).astype(np.float32)
    y_dense = pruned @ x
    y_sparse = np.array([sum(values[r, j] * x[(j // 2) * 4 + keep.reshape(4, 4)[r, j]] for j in range(4)) for r in range(4)],
                        np.float32)
    print(f"  用压缩形式算 W·x 和用剪枝后的稠密矩阵算，最大差 {np.abs(y_dense - y_sparse).max():.1e}")
    bits16 = 16
    print(f"\n  FP16 权重的存储：稠密 16 位 / 个；2:4 压缩后每 4 个存 2 个值（32 位）+ 2 个 2 位的位置（4 位）= 每个原始权重 "
          f"{(2 * bits16 + 2 * 2) / 4:g} 位，是稠密的 {(2 * bits16 + 2 * 2) / 4 / bits16:.1%}")
    print("\n观察 Observations：")
    print("  - 2:4 稀疏的非零元比例固定是 50%——按第 2 部分的算账，CSR 在这个比例下毫无优势。但它的模式是固定的："
          "每 4 个里 2 个，位置只要 2 位就能记下，硬件（Ampere 起的张量核心）可以直接跳过一半的乘法。")
    print("  - 这就是第 2 讲规格表里\"张量核心峰值含稀疏加速\"的来源：同样的时间做 2 倍的有效运算。代价是模型要按这个模式剪枝"
          "（通常还要再微调一下以恢复精度），而且只对权重矩阵适用。")
    print("  - 稀疏的两条路：非结构化稀疏（CSR 等）适合极稀疏的数据（图、推荐、检索），靠省存储和省访存取胜；"
          "结构化稀疏（2:4、块稀疏）适合神经网络，靠和硬件配合取胜。第 17 讲的 MoE 可以看作\"按专家的块稀疏\"。\n")


def main():
    part1_format()
    part2_breakeven()
    part3_kernels()
    part4_structured()


if __name__ == "__main__":
    main()
"""
05_cublas_and_roofline.py —— cuBLAS：列优先、leading dimension、行优先的"交换技巧"，以及从朴素到库的 roofline 阶梯（第 4 讲 §5）
=====================================================================

【这个脚本要回答的问题】
    1. cuBLAS（和所有 BLAS）按列优先存矩阵，参数里还有 lda、ldb、ldc。这些是什么意思？
    2. 我们的张量都是行优先的，怎样用一个列优先的库算 C = A · B，而不做任何转置拷贝？
    3. 从朴素内核到分块、再到寄存器分块，计算–全局访存比和 A100 上能达到的峰值比例是怎样一级级上去的？
       为什么实际工作中几乎总是调用 cuBLAS，而不自己写 GEMM？

【怎么运行】
    python3 05_cublas_and_roofline.py      （只需要 numpy；不到 1 秒）

【说明】这里用 NumPy 按 BLAS 的语义写了 sgemm / sgemv / sdot（参数和 cuBLAS 一一对应，只是没有 handle），
    用来验证"列优先 + leading dimension"的含义；roofline 部分是纸面估算。
"""

import numpy as np

from accel import A100_PCIE
from console import print_table


# =====================================================================
# 按 BLAS 语义实现的三个函数（数据都放在一维缓冲区里，和 C 指针一样）
# =====================================================================

def col_major_view(buf, offset, rows, cols, ld):
    """
    把一维缓冲区 buf 从 offset 开始看成一个 rows × cols 的列优先矩阵：元素 (i, j) 在 buf[offset + i + j·ld]。

    【参数】ld（leading dimension）：相邻两列的起点相隔几个元素。整块矩阵时 ld = 行数；
        取一个更大矩阵的子块时，ld = 那个大矩阵的行数（子块的列之间要跳过大矩阵的整列）。
    【边界】ld < rows 时列与列会重叠，BLAS 会报参数错误；这里抛 ValueError。
    【Python 小知识】np.lib.stride_tricks.as_strided 按给定的形状和字节步长构造一个视图，不复制数据。
    """
    if ld < max(1, rows):
        raise ValueError(f"leading dimension {ld} 小于行数 {rows}（BLAS 会报 'parameter had an illegal value'）")
    item = buf.itemsize
    base = buf[offset:]
    return np.lib.stride_tricks.as_strided(base, (rows, cols), (item, ld * item), writeable=True)


def sgemm(transa, transb, m, n, k, alpha, A, offA, lda, B, offB, ldb, beta, C, offC, ldc):
    """
    cublasSgemm 的语义：C = alpha · op(A) · op(B) + beta · C，op(X) = X（'N'）或 Xᵀ（'T'）。

    【形状】op(A) 是 m × k，op(B) 是 k × n，C 是 m × n——都按列优先存放。
        transa = 'N' 时 A 存成 m × k（lda ≥ m）；transa = 'T' 时 A 存成 k × m（lda ≥ k）。B 同理。
    【效果】原地修改缓冲区 C（和 C 指针一样）。
    """
    opA = col_major_view(A, offA, m, k, lda) if transa == "N" else col_major_view(A, offA, k, m, lda).T
    opB = col_major_view(B, offB, k, n, ldb) if transb == "N" else col_major_view(B, offB, n, k, ldb).T
    Cv = col_major_view(C, offC, m, n, ldc)
    Cv[...] = (alpha * (opA.astype(np.float64) @ opB.astype(np.float64)) + beta * Cv.astype(np.float64)).astype(np.float32)


def sgemv(trans, m, n, alpha, A, lda, x, incx, beta, y, incy):
    """cublasSgemv：y = alpha · op(A) · x + beta · y；A 是列优先的 m × n，x、y 按 incx、incy 的步长取元素。"""
    Av = col_major_view(A, 0, m, n, lda)
    opA = Av if trans == "N" else Av.T
    xs = x[::incx][:opA.shape[1]]
    ys = y[::incy][:opA.shape[0]]
    y[::incy][:opA.shape[0]] = (alpha * (opA.astype(np.float64) @ xs) + beta * ys).astype(np.float32)


def sdot(n, x, incx, y, incy):
    """cublasSdot：Σ x[i·incx] · y[i·incy]，i = 0..n−1。"""
    return float(np.dot(x[::incx][:n].astype(np.float64), y[::incy][:n].astype(np.float64)))


# =====================================================================
# 第 1 部分：列优先与 leading dimension
# =====================================================================

def part1_layout():
    print("=" * 78)
    print("第 1 部分：同一段内存，行优先读 vs 列优先读；leading dimension")
    print("=" * 78)
    buf = np.arange(1, 7, dtype=np.float32)
    print(f"  一维缓冲区：{buf.astype(int).tolist()}")
    print(f"  按行优先读成 2 × 3：{buf.reshape(2, 3).astype(int).tolist()}")
    print(f"  按列优先读成 3 × 2（ld = 3）：{col_major_view(buf, 0, 3, 2, 3).astype(int).tolist()}   ← 正好是上面那个的转置")
    big = np.arange(20, dtype=np.float32)
    sub = col_major_view(big, 1 + 1 * 5, 3, 2, 5)
    print(f"\n  一个 5 × 4 的列优先矩阵（ld = 5），元素 (i, j) = buf[i + 5j]：")
    for i in range(5):
        print("    ", col_major_view(big, 0, 5, 4, 5)[i].astype(int).tolist())
    print(f"  它从 (1, 1) 开始的 3 × 2 子块：指针偏移 1 + 1 × 5 = 6，ld 仍然是 5 → {sub.astype(int).tolist()}")
    print("\n观察 Observations：")
    print("  - BLAS 诞生于 Fortran 的年代，Fortran 的数组是列优先的，所以 cuBLAS 也按列优先解释矩阵。"
          "同一段内存，行优先看是 X，列优先看就是 Xᵀ——这个事实是第 3 部分\"交换技巧\"的全部基础。")
    print("  - leading dimension（lda）是\"相邻两列的起点相隔几个元素\"。它让你不用拷贝就能对一个大矩阵的子块做运算："
          "指针指向子块的起点，ld 给大矩阵的行数。这和 P1 §5 的步长是同一个思想（列优先时，行步长 1、列步长 ld）。\n")


# =====================================================================
# 第 2 部分：验证 sgemm / sgemv / sdot 的语义
# =====================================================================

def part2_verify():
    print("=" * 78)
    print("第 2 部分：按 BLAS 语义调用 sgemm、sgemv、sdot，与 NumPy 对照")
    print("=" * 78)
    rng = np.random.default_rng(0)
    m, n, k = 4, 5, 3
    rows = []
    for ta in ("N", "T"):
        for tb in ("N", "T"):
            Am = rng.standard_normal((m, k)).astype(np.float32)          # 数学上的 op(A)
            Bm = rng.standard_normal((k, n)).astype(np.float32)
            Cm = rng.standard_normal((m, n)).astype(np.float32)
            stA = Am if ta == "N" else Am.T                                # 实际存储的矩阵
            stB = Bm if tb == "N" else Bm.T
            A = np.asfortranarray(stA).ravel(order="F")                    # 列优先地放进一维缓冲区
            B = np.asfortranarray(stB).ravel(order="F")
            C = np.asfortranarray(Cm).ravel(order="F").copy()
            alpha, beta = np.float32(1.5), np.float32(-0.5)
            sgemm(ta, tb, m, n, k, alpha, A, 0, stA.shape[0], B, 0, stB.shape[0], beta, C, 0, m)
            ref = alpha * Am.astype(np.float64) @ Bm + beta * Cm
            got = col_major_view(C, 0, m, n, m)
            rows.append([f"sgemm('{ta}', '{tb}')", f"lda = {stA.shape[0]}，ldb = {stB.shape[0]}", f"{np.abs(got - ref).max():.1e}"])
    big = rng.standard_normal((6, 7)).astype(np.float32)
    Bm = rng.standard_normal((3, 2)).astype(np.float32)
    bufA = np.asfortranarray(big).ravel(order="F")
    C = np.zeros(4 * 2, np.float32)
    sgemm("N", "N", 4, 2, 3, np.float32(1), bufA, 1 + 2 * 6, 6, np.asfortranarray(Bm).ravel(order="F"), 0, 3, np.float32(0), C, 0, 4)
    rows.append(["sgemm 用 6 × 7 矩阵的子块 A[1:5, 2:5]", "偏移 13，lda = 6",
                 f"{np.abs(col_major_view(C, 0, 4, 2, 4) - big[1:5, 2:5].astype(np.float64) @ Bm).max():.1e}"])
    Ag = rng.standard_normal((4, 3)).astype(np.float32)
    x = rng.standard_normal(6).astype(np.float32)
    y = np.zeros(8, np.float32)
    sgemv("N", 4, 3, np.float32(1), np.asfortranarray(Ag).ravel(order="F"), 4, x, 2, np.float32(0), y, 2)
    rows.append(["sgemv('N')，incx = 2、incy = 2", "x 取 x[0], x[2], x[4]", f"{np.abs(y[::2] - Ag.astype(np.float64) @ x[::2]).max():.1e}"])
    d = sdot(3, x, 2, x, 2)
    rows.append(["sdot(n = 3, incx = 2)", "x[0]² + x[2]² + x[4]²", f"{abs(d - float(np.sum(x[::2].astype(np.float64) ** 2))):.1e}"])
    try:
        sgemm("N", "N", 4, 5, 3, np.float32(1), np.zeros(12, np.float32), 0, 3, np.zeros(15, np.float32), 0, 3,
              np.float32(0), np.zeros(20, np.float32), 0, 4)
        msg = "没有报错"
    except ValueError as e:
        msg = str(e)
    rows.append(["sgemm 的 lda 写成 3（A 有 4 行）", "参数错误", msg])
    print_table(["调用", "参数", "与 NumPy 的最大差 / 结果"], rows, align=["left", "left", "left"])
    print("\n观察 Observations：")
    print("  - 四种 transa / transb 组合都对：'T' 的意思是\"内存里存的是它的转置\"，库在内部按转置的方式取数，不需要你先转置一份。")
    print("  - incx、incy 是向量的步长（隔几个元素取一个），lda 是矩阵的列步长：BLAS 的接口从一开始就是\"指针 + 步长\"，"
          "和第 3 讲的张量表示同出一源。传错一个整数，库只会报\"第几个参数非法\"，或者更糟——静默地算错。\n")


# =====================================================================
# 第 3 部分：行优先的交换技巧
# =====================================================================

def part3_rowmajor_trick():
    print("=" * 78)
    print("第 3 部分：我们的矩阵是行优先的，怎样用列优先的 sgemm 算 C = A · B？")
    print("=" * 78)
    rng = np.random.default_rng(1)
    M, K, N = 3, 4, 5
    A = rng.standard_normal((M, K)).astype(np.float32)
    B = rng.standard_normal((K, N)).astype(np.float32)
    bufA, bufB = A.ravel(), B.ravel()                    # 行优先的一维缓冲区（就是 PyTorch 张量的存储）
    C = np.zeros(M * N, np.float32)
    sgemm("N", "N", N, M, K, np.float32(1), bufB, 0, N, bufA, 0, K, np.float32(0), C, 0, N)
    C_rowmajor = C.reshape(M, N)
    print("  调用：sgemm('N', 'N', m = N, n = M, k = K, B, ldb = N, A, lda = K, C, ldc = N)  ——注意 A、B 交换了位置")
    print(f"  把结果缓冲区按行优先读成 {M} × {N}，与 NumPy 的 A @ B 最大差 {np.abs(C_rowmajor - A.astype(np.float64) @ B).max():.1e}")
    wrong = np.zeros(M * N, np.float32)
    sgemm("N", "N", M, N, K, np.float32(1), bufA, 0, K, bufB, 0, N, np.float32(0), wrong, 0, M)
    print(f"  对照：直接按\"m = M, n = N\"调用（把行优先缓冲区当成列优先）：结果按行优先读出来与正确答案的最大差 "
          f"{np.abs(wrong.reshape(M, N) - A @ B).max():.2f}（错了）")
    print("\n观察 Observations：")
    print("  - 行优先存储的 A 在列优先的库眼里就是 Aᵀ（K × M），B 就是 Bᵀ，我们想要的 C 存成行优先，在库眼里就是 Cᵀ。"
          "而 Cᵀ = (A·B)ᵀ = Bᵀ·Aᵀ——所以只要让库算\"Bᵀ 乘 Aᵀ\"，也就是把 A、B 的位置交换、m 和 n 交换，结果自动就是行优先的 C。")
    print("  - 一次转置拷贝都不需要。PyTorch 在 GPU 上调用 cuBLAS 做矩阵乘法时，就是这样处理行优先张量的；"
          "张量是转置视图时，就改用 'T' 参数，同样不拷贝。\n")


# =====================================================================
# 第 4 部分：roofline 阶梯
# =====================================================================

def part4_ladder():
    print("=" * 78)
    print("第 4 部分：从朴素到库——计算–全局访存比的阶梯（A100 80GB PCIe，FP32 数据，纸面估算）")
    print("=" * 78)
    g = A100_PCIE
    rows = []
    entries = [("朴素：每线程一个输出，直接读全局内存", 0.25),
               ("共享内存分块 TILE = 16", 0.25 * 16),
               ("共享内存分块 TILE = 32", 0.25 * 32)]
    for bm, bn in ((64, 64), (128, 128), (256, 128)):
        entries.append((f"块分块 {bm} × {bn} + 寄存器分块（每线程算一个小块）", 2 * bm * bn / ((bm + bn) * 4)))
    for name, ai in entries:
        rows.append([name, f"{ai:.2f}", f"{min(1, ai * g.bw / g.fp32):.1%}", f"{min(1, ai * g.bw / g.tf32):.1%}"])
    print(f"  FP32 平衡点 {g.fp32 / g.bw:.1f} FLOP/字节，TF32 张量核心平衡点 {g.tf32 / g.bw:.1f} FLOP/字节")
    print_table(["写法", "计算–全局访存比 FLOP/字节", "FP32 峰值的上限", "TF32 峰值的上限"], rows)
    print("\n  块分块的比值：一个块算 BM × BN 的输出，每一步读 A 的 BM × BK 和 B 的 BK × BN，做 2·BM·BN·BK 次运算 →")
    print("  比值 = 2·BM·BN / ((BM + BN) × 4 字节)，和 BK 无关；块越大、越接近正方形，比值越高。")
    print("\n观察 Observations：")
    print(f"  - 朴素版只能用到 FP32 峰值的 {0.25 * g.bw / g.fp32:.1%}；共享内存分块把它提高到 40%–80%；"
          f"让每个块覆盖 128 × 128 的输出（每个线程用寄存器算 8 × 8 之类的小块），比值 32，FP32 已经计算受限。")
    print(f"  - 换成 TF32 张量核心，平衡点升到 {g.tf32 / g.bw:.0f}：同样的分块只能用到约 40%–53%，还要更大的块、"
          "加上 L2 缓存在块之间的复用、异步拷贝和多级流水线，才能把张量核心喂饱。")
    print("  - 这就是为什么实际工作中几乎总是调用 cuBLAS（或 CUTLASS、Triton 生成的内核）：它为每一代 GPU、每一种数据类型、"
          "每一类矩阵形状挑选分块大小（第 2 讲的波次量化也在考虑之内），用张量核心指令、双缓冲 / TMA 异步拷贝把访存和计算重叠，"
          "还能把偏置、激活函数融合进结果的写回（cuBLASLt 的 epilogue）。自己写 GEMM 是为了理解它，不是为了替代它。\n")


def main():
    part1_layout()
    part2_verify()
    part3_rowmajor_trick()
    part4_ladder()


if __name__ == "__main__":
    main()
"""
accel.py —— 第 4 讲的内核：朴素 / 分块 / 线程粗化的矩阵乘法、三种矩阵转置、CSR 稀疏矩阵-向量乘法（在 SIMT 模拟器上运行）
=====================================================================

【这个文件要回答的问题】
    原课第 4 讲讲了四种 GPU 加速技术：分块（tiling）、内存并行（合并访存、bank 冲突）、稀疏矩阵乘法、cuBLAS。
    前三种都是"同一个计算，换一种让线程分工和取数的方式"。这里把 slides 里的每个内核都写成模拟器能执行的版本，
    让 01–04 号脚本去数它们各自省下了多少次访存。

【本文件提供】
    矩阵乘法 : matmul_naive（每线程一个输出，直接读全局内存）、matmul_tiled（slides 的 MatMulTiledKernel，可选是否做越界处理、
               可以故意漏掉 __syncthreads 做反面教材）、matmul_tiled_coarse（线程粗化：一个块算横向相邻的几个输出块）
    转置     : transpose_naive、transpose_smem（共享内存，有 bank 冲突）、transpose_smem_padded（+1 填充，无冲突）
    稀疏     : dense_to_csr、spmv_csr_thread（slides 的 SpMVCSRKernel：一个线程一行）、spmv_csr_warp（一个线程束一行）
    工具     : run_matmul（分配 → 拷贝 → 发射 → 拷回，返回结果和统计）、A100_PCIE（slides 用的 A100 规格）

【约定】矩阵都是行优先、float32、方阵（转置和矩阵乘法）；线程的 x 方向对应列、y 方向对应行（第 3 讲 §2.3）。
"""

import os
import sys
from dataclasses import dataclass

import numpy as np

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


@dataclass(frozen=True)
class Spec:
    """roofline 估算用的规格：名字、FP32 峰值、TF32 张量核心峰值（dense）、显存带宽（字节/秒）。"""
    name: str
    fp32: float
    tf32: float
    bw: float


# 原课第 4 讲 slides 用的 A100 80GB PCIe：FP32 19.5 TFLOPS、TF32 156 TFLOPS、显存带宽 1,935 GB/s
A100_PCIE = Spec("A100 80GB PCIe", 19.5e12, 156e12, 1935e9)
H100_SXM = Spec("H100 SXM", 67e12, 494.7e12, 3.35e12)


# =====================================================================
# 1. 矩阵乘法
# =====================================================================

def matmul_naive(t, A, B, C, N):
    """
    slides 第 6 页的 MatMulKernel：每个线程算 C 的一个元素，循环里每一步从全局内存读 A、B 各一个数。
        for (k = 0; k < N; k++) Pvalue += a[row * N + k] * b[k * N + col];
    """
    row = t.blockIdx.y * t.blockDim.y + t.threadIdx.y
    col = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if row >= N or col >= N:
        return
    acc = np.float32(0)
    for k in range(N):
        acc += A[row * N + k] * B[k * N + col]
    C[row * N + col] = acc


def matmul_tiled(t, A, B, C, N, TILE, bounds=True, sync1=True, sync2=True):
    """
    slides 第 36 页的 MatMulTiledKernel（共享内存分块）。块大小 = TILE × TILE，每个线程算 C 的一个元素。

    【思路】把 C 的这一块所需要的 A 的"一条横带"和 B 的"一条竖带"切成 TILE × TILE 的小块，一相（phase）处理一对：
        1. 块内每个线程从全局内存搬一个 A 的元素、一个 B 的元素到共享内存 As、Bs
        2. __syncthreads()：等大家都搬完（否则有人会读到还没写的位置）
        3. 每个线程用 As 的一行和 Bs 的一列做 TILE 次乘加（只读共享内存）
        4. __syncthreads()：等大家都算完（否则有人会在别人还在读时覆盖 As、Bs）
    每个从全局内存读进来的数在共享内存里被用了 TILE 次 → 全局内存读取减少到 1/TILE。
    【参数】bounds：True 时用 ⌈N/TILE⌉ 相、越界位置补 0（正确处理 N 不是 TILE 倍数的情况）；
        False 时照搬 slides：相数 = N / TILE（整数除法）、不检查越界。
        sync1 / sync2：是否保留两次 __syncthreads（故意去掉一个用来演示数据竞争）
    【数值】每个线程仍然按 k = 0, 1, …, N−1 的顺序累加，所以结果和 matmul_naive 逐位相同。
    """
    As = t.shared("As", (TILE, TILE))
    Bs = t.shared("Bs", (TILE, TILE))
    tx, ty = t.threadIdx.x, t.threadIdx.y
    row = t.blockIdx.y * TILE + ty
    col = t.blockIdx.x * TILE + tx
    acc = np.float32(0)
    phases = ceil_div(N, TILE) if bounds else N // TILE
    for ph in range(phases):
        ka, kb = ph * TILE + tx, ph * TILE + ty
        if bounds:
            As[ty, tx] = A[row * N + ka] if (row < N and ka < N) else np.float32(0)
            Bs[ty, tx] = B[kb * N + col] if (kb < N and col < N) else np.float32(0)
        else:
            As[ty, tx] = A[row * N + ka]
            Bs[ty, tx] = B[kb * N + col]
        if sync1:
            yield t.syncthreads()
        for k in range(TILE):
            acc += As[ty, k] * Bs[k, tx]
        if sync2:
            yield t.syncthreads()
    if not bounds or (row < N and col < N):
        C[row * N + col] = acc


def matmul_tiled_coarse(t, A, B, C, N, TILE, COARSE):
    """
    线程粗化（thread coarsening，PMPP 6.3 节）：一个块负责横向相邻的 COARSE 个输出块，每个线程算 COARSE 个输出。

    【为什么】在 matmul_tiled 里，横向相邻的块要的是同一条 A 的横带，却各自从全局内存读一遍。
        粗化后 A 的小块每相只读一次，被 COARSE 个输出块共用 → A 的全局读取再减少到 1/COARSE。
    【代价】每个线程要 COARSE 个累加器（更多寄存器），块数减少到 1/COARSE（并行度降低）。
    【边界】要求 N 是 TILE × COARSE 的倍数（简化）。
    """
    As = t.shared("As", (TILE, TILE))
    Bs = t.shared("Bs", (TILE, TILE))
    tx, ty = t.threadIdx.x, t.threadIdx.y
    row = t.blockIdx.y * TILE + ty
    col0 = t.blockIdx.x * TILE * COARSE + tx
    acc = [np.float32(0)] * COARSE
    for ph in range(N // TILE):
        As[ty, tx] = A[row * N + ph * TILE + tx]                       # A 的小块：每相只读一次
        for c in range(COARSE):
            col = col0 + c * TILE
            Bs[ty, tx] = B[(ph * TILE + ty) * N + col]
            yield t.syncthreads()
            for k in range(TILE):
                acc[c] += As[ty, k] * Bs[k, tx]
            yield t.syncthreads()
    for c in range(COARSE):
        C[row * N + col0 + c * TILE] = acc[c]


def run_matmul(kernel, A, B, block, grid, *extra, **launch_kw):
    """
    主机端：分配 → 拷入 → 发射 → 拷回 → 释放。【返回】(C 的 NumPy 数组, LaunchStats)
    """
    N = A.shape[0]
    dev = Device()
    dA, dB = dev.to_device(A, "A"), dev.to_device(B, "B")
    dC = dev.malloc(N * N, np.float32, "C")
    s = launch(kernel, grid, block, dA, dB, dC, N, *extra, **launch_kw)
    C = dev.to_host(dC, (N, N))
    return C, s


# =====================================================================
# 2. 转置（slides 第 25–33 页）
# =====================================================================

def transpose_naive(t, m, a, c):
    """
    slides 第 26 页的 naive_cuda_transpose：
        c[INDX(myCol, myRow, m)] = a[INDX(myRow, myCol, m)];
    读 a 时一个线程束读同一行的相邻元素（合并）；写 c 时相邻线程写同一列、相隔 m 个元素（不合并）。
    """
    my_col = t.blockDim.x * t.blockIdx.x + t.threadIdx.x
    my_row = t.blockDim.y * t.blockIdx.y + t.threadIdx.y
    if my_row < m and my_col < m:
        c[my_col * m + my_row] = a[my_row * m + my_col]


def transpose_smem(t, m, a, c, pad=0):
    """
    slides 第 30 页的 smem_cuda_transpose：先按行（合并地）把一块读进共享内存，再按行（合并地）写出转置后的块。
        smemArray[threadIdx.x][threadIdx.y] = a[INDX(tileRow + threadIdx.y, tileCol + threadIdx.x, m)];
        __syncthreads();
        c[INDX(tileCol + threadIdx.y, tileRow + threadIdx.x, m)] = smemArray[threadIdx.y][threadIdx.x];
    pad = 1 时就是 slides 第 33 页的版本：共享数组声明成 [TX][TY + 1]，消除 bank 冲突。
    【边界】要求 m 是块大小的倍数（和 slides 一样，不做越界检查）。
    """
    TX, TY = t.blockDim.x, t.blockDim.y
    smem = t.shared("smemArray", (TX, TY + pad))
    tile_col = TX * t.blockIdx.x
    tile_row = TY * t.blockIdx.y
    tx, ty = t.threadIdx.x, t.threadIdx.y
    smem[tx, ty] = a[(tile_row + ty) * m + tile_col + tx]
    yield t.syncthreads()
    c[(tile_col + ty) * m + tile_row + tx] = smem[ty, tx]


def transpose_smem_padded(t, m, a, c):
    """slides 第 33 页：smemArray[THREADS_PER_BLOCK_X][THREADS_PER_BLOCK_Y + 1]。"""
    yield from transpose_smem(t, m, a, c, pad=1)


# =====================================================================
# 3. 稀疏矩阵（slides 第 41–43 页）
# =====================================================================

def dense_to_csr(M):
    """
    稠密矩阵 → CSR（Compressed Sparse Row，压缩稀疏行）三个数组：
        data[k]      : 第 k 个非零元的值（按行优先顺序排列）
        col_index[k] : 第 k 个非零元所在的列
        row_ptr[r]   : 第 r 行的非零元在 data 里从哪里开始；row_ptr[r+1] − row_ptr[r] = 第 r 行的非零元个数
                       （长度是行数 + 1，最后一个元素 = 非零元总数）
    【返回】(data float32, col_index int32, row_ptr int32)
    """
    rows, cols = np.nonzero(M)
    data = M[rows, cols].astype(np.float32)
    row_ptr = np.zeros(M.shape[0] + 1, np.int32)
    np.add.at(row_ptr, rows + 1, 1)
    return data, cols.astype(np.int32), np.cumsum(row_ptr).astype(np.int32)


def spmv_csr_thread(t, data, col_index, row_ptr, x, y, num_rows):
    """
    slides 第 43 页的 SpMVCSRKernel：一个线程负责一行，y[row] += Σ data[elem] · x[col_index[elem]]。
    【问题】相邻线程处理相邻的行，同一时刻读的是各自那一行的第 k 个非零元——位置相隔一整行的非零元个数（不合并）；
        各行长度不同时，线程束要等最长的那一行（分化）。
    """
    row = t.blockIdx.x * t.blockDim.x + t.threadIdx.x
    if row < num_rows:
        dot = np.float32(0)
        for elem in range(row_ptr[row], row_ptr[row + 1]):
            dot += x[col_index[elem]] * data[elem]
        y[row] = y[row] + dot


def spmv_csr_warp(t, data, col_index, row_ptr, x, y, num_rows):
    """
    改进（PMPP 第 14 章的"向量 CSR"）：一个线程束负责一行。32 个线程读这一行连续的 32 个非零元（合并），
    各自累加部分和，最后在共享内存里树形归约成一个数，由 0 号车道写出。
    【适用】每行非零元很多时；每行只有几个非零元时，线程束里大部分车道闲着。
    """
    lane = t.threadIdx.x % 32
    warp_in_block = t.threadIdx.x // 32
    row = t.blockIdx.x * (t.blockDim.x // 32) + warp_in_block
    part = t.shared("part", t.blockDim.x)
    dot = np.float32(0)
    if row < num_rows:
        for elem in range(row_ptr[row] + lane, row_ptr[row + 1], 32):
            dot += x[col_index[elem]] * data[elem]
    part[t.threadIdx.x] = dot
    yield t.syncthreads()
    s = 16
    while s > 0:
        if lane < s:
            part[t.threadIdx.x] = part[t.threadIdx.x] + part[t.threadIdx.x + s]
        yield t.syncthreads()
        s //= 2
    if lane == 0 and row < num_rows:
        y[row] = y[row] + part[t.threadIdx.x]
"""
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)))
