"""
=============================================================================
CS349E · 第 13 讲 · 02_polyhedral.py —— 多面体模型（二维版）：合法的变换、倾斜、分块与波前并行
=============================================================================

【这个脚本要回答的问题】
    前置讲 P4 用依赖距离向量判断"交换两层循环是否合法"。多面体模型把这件事推广：迭代空间是整数点的集合，
    循环变换是这些点上的仿射（线性）映射，合法 ⟺ 每个依赖变换后仍然"先产生、后使用"。
    在这个框架里，编译器可以**搜索**一个变换，让循环可以分块（局部性）或者内层可以并行。
      演示 1：一维 Jacobi 的时间迭代：依赖 (1, −1)、(1, 0)、(1, 1)；枚举所有幺模变换，哪些合法、哪些能分块、哪些内层能并行。
      演示 2：按变换后的顺序真的执行：交换（非法）、直接分块（非法）、倾斜后分块（合法）；倾斜分块的 cache 缺失。
      演示 3：二维 Gauss–Seidel（原地更新）：两层循环都带依赖，没有能直接并行的循环；波前变换让内层可以并行。
      演示 4：矩阵乘法：依赖只有 (0, 0, 1)——为什么它的所有循环顺序都合法（复习前置讲 P4）。

【运行】python3 02_polyhedral.py      （约 3 秒）
【依赖】numpy；本目录的 sched_lab；../prereq4 的 loop_ir.lru_misses（全相联 LRU cache 的缺失数）；../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 classify, jacobi_run, order_points, search_transforms, seidel_run, wavefront_parallelism  # noqa: E402
from loop_ir import lru_misses  # noqa: E402  （sched_lab 已把 ../prereq4 加进 sys.path）
from console import print_header, print_observations, print_table  # noqa: E402

JACOBI = [(1, -1), (1, 0), (1, 1)]


def fmt_T(T):
    return f"[[{T[0][0]}, {T[0][1]}], [{T[1][0]}, {T[1][1]}]]"


def demo1_search():
    print_header("演示 1：一维 Jacobi：A[t+1][i] = (A[t][i−1] + A[t][i] + A[t][i+1]) / 3，迭代点 (t, i)")
    print("依赖距离向量：(1, −1)、(1, 0)、(1, 1)（第 t+1 步的 i 用到第 t 步的 i−1、i、i+1）\n")
    named = [("原样 (t, i)", ((1, 0), (0, 1))), ("交换 (i, t)", ((0, 1), (1, 0))), ("倾斜 (t, i + t)", ((1, 0), (1, 1))),
             ("倾斜 + 交换 (i + t, t)", ((1, 1), (1, 0))), ("波前 (2t + i, t)", ((2, 1), (1, 0)))]
    rows = []
    for name, T in named:
        c = classify(T, JACOBI)
        rows.append([name, fmt_T(T), " ".join(str(d) for d in c["dist"]), "是" if c["legal"] else "否",
                     "是" if c["tileable"] else "否", "是" if c["parallel_inner"] else "否"])
    print_table(["变换", "矩阵 T", "变换后的依赖", "合法", "可以矩形分块", "内层可以并行"], rows)
    res = search_transforms(JACOBI, max_coef=2)
    n_til = sum(1 for _, c in res if c["tileable"])
    best = [r for r in res if r[1]["tileable"]][:3]
    print(f"\n枚举元素在 [−2, 2] 的全部 2×2 幺模矩阵：合法 {len(res)} 个，其中可分块 {n_til} 个。依赖距离最短的可分块变换：")
    print_table(["T", "变换后的依赖", "依赖分量之和"],
                [[fmt_T(T), " ".join(str(d) for d in c["dist"]), sum(sum(abs(x) for x in d) for d in c["dist"])] for T, c in best])
    print_observations([
        "合法 ⟺ 每个依赖向量变换后字典序为正（第一个非零分量 > 0）：被依赖的迭代仍然先执行。交换 t 和 i 把 (1, −1) 变成 (−1, 1)，非法——"
        "第 t+1 步的 i 会在第 t 步的 i+1 之前算。",
        "矩形分块要求更强：所有依赖的**每个分量**都 ≥ 0（这两层循环\"完全可交换\"）。原样的 (1, −1) 有负分量，直接分块非法；"
        "倾斜 i' = i + t 把依赖变成 (1, 0)、(1, 1)、(1, 2)，全部非负，就能分块了。",
        "内层并行要求外层循环\"带走\"所有依赖：变换后每个依赖的第一个分量 > 0。原样的 t 循环就满足（同一个时间步内的 i 互不依赖）；"
        "倾斜后再把 2t + i 放在外层，得到一个\"波前\"：同一条斜线上的点互不依赖。",
        "多面体编译器（Pluto、ISL / PPCG、MLIR 的 affine 方言）就是这样做的，只是用整数线性规划在所有仿射变换里找"
        "\"合法、可分块、依赖距离最短\"的那个（Pluto 的目标函数），而不是像这里暴力枚举。",
    ])


def demo2_execute():
    print_header("演示 2：按变换后的顺序真的执行（N = 2000，200 个时间步，块 50×50）")
    N, S, B = 2000, 200, 50
    init = np.random.default_rng(0).random(N + 2)
    pts = [(t, i) for t in range(S) for i in range(1, N + 1)]
    ref, _, _ = jacobi_run(pts, N, S, init)
    rows, miss_of = [], {}
    for name, T, tile in (("原样（t 外、i 内）", ((1, 0), (0, 1)), None), ("交换（i 外、t 内）", ((0, 1), (1, 0)), None),
                          ("直接分块 50×50（不倾斜）", ((1, 0), (0, 1)), B), ("倾斜后分块 50×50", ((1, 0), (1, 1)), B)):
        order = order_points(pts, T, tile)
        out, bad, addrs = jacobi_run(order, N, S, init)
        miss = lru_misses(addrs, 8 * 1024, 64)
        miss_of[name] = miss
        rows.append([name, bad, f"{np.abs(out - ref).max():.1e}", f"{miss:,}", f"{miss / len(pts):.3f}"])
    print_table(["执行顺序", "读到还没算出来的值", "和原顺序结果的最大差", "cache 缺失（8 KB 全相联 LRU，64 B 一行）", "每个迭代点"], rows)
    print_observations([
        "非法的顺序真的会算错：交换以后，第 t+1 步的 i 要用的第 t 步的 i+1 还没算出来；直接分块也一样——块的边上有依赖指向\"右边还没执行的块\"。"
        "（数组的每个时间步存一行，没有复用存储，所以只有流依赖；用 NaN 标记\"还没算出来\"，读到就计数。）",
        f"倾斜后分块的结果和原顺序逐位相同，cache 缺失从每点 {miss_of['原样（t 外、i 内）'] / len(pts):.3f} 次降到 "
        f"{miss_of['倾斜后分块 50×50'] / len(pts):.3f} 次。其中写新的一行总是缺失（每个时间步存一行、不复用存储，每 8 个点约 1 次 = 0.125）；"
        f"读的缺失从每点约 {miss_of['原样（t 外、i 内）'] / len(pts) - 0.125:.3f} 降到约 {miss_of['倾斜后分块 50×50'] / len(pts) - 0.125:.3f}："
        "原顺序每个时间步扫过整行 2000 个数（16 KB，比 cache 大），下一个时间步再读时早已被挤出去；倾斜的块在一小块时空区域里连续推进几十个时间步，"
        "数据一直在 cache 里。这就是\"时间分块\"（time tiling）。",
        "真实的模板计算（以及把多层 RNN 的时间步和层一起分块）都用这种变换；它也是多面体编译器最典型的成功案例。",
    ])


def demo3_seidel():
    print_header("演示 3：二维 Gauss–Seidel（原地更新）：A[i][j] = (A[i−1][j] + A[i][j−1] + A[i+1][j] + A[i][j+1]) / 4")
    N = 64
    init = np.random.default_rng(1).random((N + 2, N + 2))
    pts = [(i, j) for i in range(1, N + 1) for j in range(1, N + 1)]
    deps = [(1, 0), (0, 1)]
    ref = seidel_run(pts, N, init)
    rows = []
    for name, T in (("原样 (i, j)", ((1, 0), (0, 1))), ("交换 (j, i)", ((0, 1), (1, 0))), ("波前 (i + j, i)", ((1, 1), (1, 0)))):
        c = classify(T, deps)
        steps, avg, mx = wavefront_parallelism(pts, T)
        out = seidel_run(order_points(pts, T), N, init)
        # "并行"执行：同一个外层值的点按相反的顺序执行（任意顺序都应该得到同样的结果，如果内层真的可以并行）
        Q = np.asarray(pts) @ np.asarray(T).T
        idx = np.lexsort((-Q[:, 1], Q[:, 0]))
        out_rev = seidel_run([pts[k] for k in idx], N, init)
        rows.append([name, " ".join(str(d) for d in c["dist"]), "是" if c["legal"] else "否",
                     "是" if c["parallel_inner"] else "否", f"{steps} 步 × 平均 {avg:.1f}（最多 {mx}）",
                     f"{np.abs(out - ref).max():.1e}", f"{np.abs(out_rev - ref).max():.1e}"])
    print_table(["变换", "变换后的依赖", "合法", "内层可以并行", "外层步数 × 内层宽度", "和原顺序的差",
                 "内层倒着执行的差"], rows)
    print_observations([
        "原地更新带来流依赖（用到刚更新的 A[i−1][j]、A[i][j−1]）和反依赖（A[i+1][j]、A[i][j+1] 要在被覆盖之前读到）——距离都是 (1, 0)、(0, 1)。"
        "两层循环都带着依赖：原样执行时内层不能并行（\"内层倒着执行\"结果就变了）。",
        f"波前变换 (i + j, i)：同一条反对角线 i + j = 常数上的点互不依赖。外层走 {2 * N - 1} 步，每步内层平均有约 {N * N / (2 * N - 1):.0f} 个点可以并行，"
        "倒着执行结果也完全一样。",
        "代价是并行度不均匀（从 1 涨到 64 再降回 1）和更复杂的循环边界——这些边界正是多面体代码生成器（如 CLooG、ISL 的 AST 生成）要算的东西。",
    ])


def demo4_matmul():
    print_header("演示 4：矩阵乘法 C[i][j] += A[i][k]·B[k][j] 的依赖")
    deps3 = [(0, 0, 1)]
    rows = []
    import itertools
    for perm in itertools.permutations(range(3)):
        names = "".join("ijk"[p] for p in perm)
        d = tuple(deps3[0][p] for p in perm)
        rows.append([names, d, "是" if any(x != 0 for x in d) and next(x for x in d if x != 0) > 0 else "否"])
    print_table(["循环顺序（外 → 内）", "变换后的依赖", "合法"], rows)
    print_observations([
        "矩阵乘法只有一个依赖：同一个 C[i][j] 沿 k 累加，距离 (0, 0, 1)。它在任何排列下都是\"非负、字典序为正\"的，所以 6 种循环顺序都合法、"
        "任意分块都合法（前置讲 P4 §7 演示过）；i、j 两层没有依赖，可以任意并行——这就是为什么矩阵乘法的调度空间（第 10 讲、03 脚本）里"
        "几乎全是性能问题，没有正确性问题。",
        "k 方向的依赖来自\"加法的顺序\"：如果允许重新结合（浮点结果会略有不同），k 也可以切开并行——split-K"
        "（第 11 讲的 flash-decoding 就是在注意力的 KV 方向上做 split-K，再用 log-sum-exp 合并）。",
    ])


if __name__ == "__main__":
    demo1_search()
    demo2_execute()
    demo3_seidel()
    demo4_matmul()