"""
=============================================================================
CS349E · 第 13 讲 · 03_gpu_schedule_space.py —— 一个 GPU 矩阵乘法内核的调度空间
=============================================================================

【这个脚本要回答的问题】
    第 10 讲手写的矩阵乘法里，块多大、几个 warp、几级流水线、块按什么顺序启动，都是"调度"。
    Triton / TVM 把它们暴露成参数，由自动调优来选。这个空间长什么样？最优点对形状有多敏感？
      演示 1：调度 → 参数：一个 Triton 风格的配置有哪些旋钮；6000 个组合里多少能启动，为什么不能。
      演示 2：4096³ 上的"地形"：最快的配置、离最快 5% 以内的有几个、随便选一个会慢多少。
      演示 3：形状一变，最优配置就变：四种形状各自的最优配置，以及"把 A 形状的最优用在 B 上"的代价。
      演示 4：在最优点附近一次只改一个参数：哪些参数敏感。

【关于数字】时间来自 sched_lab.gemm_time——一个**解析模型**（H100 式的参数 + 第 8、10 讲的资源账 + 若干假设），
    不是实测。它的作用是给调度空间一个"有合理结构"的代价，用来研究搜索；具体的 TFLOPS 不要当真。
【运行】python3 03_gpu_schedule_space.py      （约 3 秒）
【依赖】numpy；本目录的 sched_lab（它导入 ../lecture10 的 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 sched_lab import SPACE, GemmConfig, H100_TC, all_configs, gemm_time, static_valid  # noqa: E402
from console import fmt_time, print_header, print_observations, print_table  # noqa: E402

SHAPES = [((4096, 4096, 4096), "方阵 4096³（大 GEMM）"), ((512, 14336, 4096), "prefill 512 个 token × Llama-8B 的 up 投影"),
          ((2048, 2048, 512), "小的方阵（K 短）"), ((16, 4096, 4096), "decode 批 16 × 4096² 的投影")]


def reasons(c, hw=H100_TC):
    """一个配置不能启动的原因（可能有多个）。"""
    bm, bn, bk, w, st, _ = c.as_tuple()
    wy, wx = {1: (1, 1), 2: (1, 2), 4: (2, 2), 8: (4, 2)}[w]
    r = []
    if st * (bm + bn) * bk * 2 > hw.smem_limit:
        r.append("共享内存超过 227 KB")
    if bm // wy < 16 or bn // wx < 8:
        r.append("warp 的块小于一条 mma（16×8）")
    elif (bm // wy) * (bn // wx) // 32 + 40 > 255:
        r.append("每线程寄存器超过 255")
    return r


def demo1_space():
    print_header("演示 1：调度的旋钮和配置空间")
    rows = [["BLOCK_M × BLOCK_N", "每个程序实例（线程块）算 C 的多大一块", " / ".join(map(str, SPACE["bm"])) + "（两维各取）", "第 10 讲 §6、§8"],
            ["BLOCK_K", "每一步沿 K 读多少", " / ".join(map(str, SPACE["bk"])), "共享内存、流水线每级的大小"],
            ["num_warps", "每个块几个 warp（warp 排成网格，各算一个子块）", " / ".join(map(str, SPACE["warps"])), "warp 分块、Tensor Core 的利用"],
            ["num_stages", "共享内存里几级流水线", " / ".join(map(str, SPACE["stages"])), "第 10 讲 §8.4 的 Little 定律"],
            ["GROUP_M", "块的启动顺序（L2 分组）", " / ".join(map(str, SPACE["group"])), "第 10 讲 §10.2"]]
    print_table(["参数", "含义", "取值", "对应"], rows)
    cfgs = all_configs()
    valid = [c for c in cfgs if static_valid(c)]
    cnt = {}
    for c in cfgs:
        for r in reasons(c):
            cnt[r] = cnt.get(r, 0) + 1
    print(f"\n一共 {len(cfgs)} 个组合，能启动的 {len(valid)} 个（{len(valid) / len(cfgs):.0%}）。不能启动的原因（一个配置可能有几个）：")
    print_table(["原因", "配置数"], [[k, v] for k, v in sorted(cnt.items(), key=lambda kv: -kv[1])])
    print_observations([
        "在 Halide / TVM 里，这些是调度原语的参数：split 的大小（块）、bind 到线程（warp 数）、cache_read 到共享内存再加流水线（stages）、"
        "块的顺序（GROUP_M）。在 Triton 里，算法（块级的程序，第 10 讲 §10）只写一次，这些作为 tl.constexpr 的元参数交给 @triton.autotune。",
        "很大一部分组合根本不能启动：共享内存放不下（stages × (BM + BN) × BK × 2 字节）、warp 分到的块比一条 mma 还小、累加器把寄存器撑爆。"
        "这些不用测量就能判断——调优器先把它们剪掉（Triton 也会剪掉共享内存超限的配置）。",
    ])


def landscape(shape):
    pool = [c for c in all_configs() if static_valid(c)]
    ts = np.array([gemm_time(c, *shape) for c in pool])
    return pool, ts


def demo2_landscape():
    shape = (4096, 4096, 4096)
    print_header("演示 2：4096³ FP16 矩阵乘法的\"地形\"（模型时间）")
    pool, ts = landscape(shape)
    opt = ts.min()
    order = np.argsort(ts)
    rows = [[GemmConfig(*pool[i].as_tuple()).as_tuple(), fmt_time(ts[i]), f"{2 * 4096 ** 3 / ts[i] / 1e12:.0f}", f"{ts[i] / opt:.3f}"]
            for i in order[:6]]
    print_table(["(BM, BN, BK, warps, stages, GROUP_M)", "时间", "TFLOPS", "相对最快"], rows)
    q = np.percentile(ts / opt, [10, 25, 50, 75, 90])
    print(f"\n能启动的 {len(pool)} 个配置里：离最快 5% 以内的 {int((ts <= 1.05 * opt).sum())} 个、2 倍以内的 {int((ts <= 2 * opt).sum())} 个；"
          f"相对最快的分位数：10% {q[0]:.2f}、25% {q[1]:.2f}、中位数 {q[2]:.2f}、75% {q[3]:.2f}、90% {q[4]:.2f}。")
    print_observations([
        "最快的一组配置都是 256×128 的块、8 个 warp；BK 和 stages 只要合起来能藏住延迟（BK = 32 配 4–5 级、64 配 3 级、128 配 2 级）就一样快；"
        "GROUP_M 在这里不起作用（这个形状受限于计算）。和第 10 讲 §8 的分析一致：块要大到每次乘加分到的共享内存 / L2 字节够少，"
        "流水线要深到藏住延迟，又不能大到共享内存放不下。",
        f"\"好\"的配置很稀少：3744 个里只有 {int((ts <= 1.05 * opt).sum())} 个在最快的 5% 以内，随便选一个（中位数）要慢 {q[2]:.1f} 倍。"
        "这就是为什么要搜索，而且搜索要聪明（04 脚本）。",
        "这个模型给出的最高约 900 TFLOPS，是 H100 稠密峰值的九成——比真实的 cuBLAS（大约七成多）乐观：模型没有算功耗导致的降频、指令开销、"
        "TMA / WGMMA 的细节。用它研究搜索策略没问题，不要拿它的绝对数字去预测真机。",
    ])


def demo3_shapes():
    print_header("演示 3：形状一变，最优配置就变")
    best = {}
    tables = {}
    for shape, name in SHAPES:
        pool, ts = landscape(shape)
        i = int(np.argmin(ts))
        best[shape] = pool[i]
        tables[shape] = dict(zip(pool, ts))
    rows = []
    for shape, name in SHAPES:
        c = best[shape]
        t = tables[shape][c]
        rows.append([name, "×".join(map(str, shape)), c.as_tuple(), fmt_time(t), f"{2 * np.prod(shape) / t / 1e12:.0f}"])
    print_table(["形状", "M×N×K", "最优配置", "时间", "TFLOPS"], rows)
    print("\n把每一行形状的最优配置用在每一列的形状上（相对那一列自己的最优）：")
    head = ["最优配置来自 \\ 用在"] + [n.split("（")[0][:10] for _, n in SHAPES]
    rows = []
    for sa, na in SHAPES:
        r = [na.split("（")[0][:14]]
        for sb, _ in SHAPES:
            t = tables[sb].get(best[sa], np.inf)
            r.append(f"{t / tables[sb][best[sb]]:.2f}")
        rows.append(r)
    print_table(head, rows)
    print_observations([
        "大 GEMM 喜欢大块、深流水线；decode 那种 M = 16 的\"细长\"矩阵乘法受限于读权重（第 1 讲），最优是很小的块（BM = 16）和更多的块"
        "——要让所有 SM 都在读权重，而不是让少数几个块各自算一大块补齐的 0。",
        "前三种形状的最优恰好是同一个配置；把它用在 decode 形状上慢约 3 倍（M = 16 要补齐到 256 行，而且只有 32 个块，大部分 SM 闲着）；"
        "反过来，decode 的最优（16×16 的小块）用在另外三种形状上慢 6.6–14 倍。"
        "所以 cuBLAS 为不同的形状准备了很多个内核和一套启发式，Triton 的 @triton.autotune 按 key（例如 M、N、K）分别调优并缓存结果。",
        "推理里的形状很多：decode 的批大小一直在变（M = 1 … 几百），prefill 的长度也在变，每层还有 QKV、O、gate/up、down 四种 N×K。"
        "实际的做法是按桶调优（第 12 讲 §8 的桶）、加缓存，或者用学到的代价模型直接预测（04 脚本演示 1）。",
    ])


def demo4_sensitivity():
    shape = (4096, 4096, 4096)
    print_header("演示 4：在最优点附近一次只改一个参数（4096³）")
    pool, ts = landscape(shape)
    i = int(np.argmin(ts))
    b = pool[i]
    t0 = ts[i]
    rows = []
    for k in SPACE:
        cells = []
        for v in SPACE[k]:
            d = dict(zip(SPACE, b.as_tuple()))
            d[k] = v
            c = GemmConfig(*d.values())
            t = gemm_time(c, *shape)
            cells.append(f"{v}:{'—' if not np.isfinite(t) else f'{t / t0:.2f}'}" + ("*" if v == getattr(b, k) else ""))
        rows.append([k, "  ".join(cells)])
    print(f"最优配置 {b.as_tuple()}；每格是\"取值:相对最优的时间\"，— 表示不能启动，* 是最优点本身的取值。")
    print_table(["参数", "取值 : 时间比"], rows)
    print_observations([
        "块大小最敏感：BM 或 BN 小一号就慢 1.5–3 倍（复用不够，第 10 讲 §8.3 的账）；BM 再小（16、32），8 个 warp 分到的块比一条 mma 还小，"
        "不能启动。warp 数反过来：少于 8 个时每个 warp 分到的块太大，累加器放不进 255 个寄存器，也不能启动。",
        "stages 越多越好（1 → 4 级：3.67 → 1.00），直到够藏住延迟；5 级不再变好；BK = 128 时 4 级就放不下共享内存了。",
        "参数之间有很强的相互作用（块越大，能用的 stages 越少；warp 数决定每个 warp 的块），所以\"一次调一个参数\"的坐标下降容易卡在局部最优——"
        "这是自动调优要面对的地形。",
    ])


if __name__ == "__main__":
    demo1_space()
    demo2_landscape()
    demo3_shapes()
    demo4_sensitivity()