"""
=============================================================================
CS349E · 第 13 讲 · 01_halide_schedules.py —— 算法与调度分离：生产者在哪里算？
=============================================================================

【这个脚本要回答的问题】
    Halide 把"算什么"（算法）和"怎么算"（调度）分开。最重要的一个调度决策是：一个中间结果（生产者）
    相对于用它的地方（消费者）在哪一层循环里算——这决定了并行性、局部性和重复计算之间的三角取舍。
      演示 1：算法：两级的模糊（先水平、再竖直）；边界推断——每一级要算多大的区域。
      演示 2：四种调度（root / inline / 分块 / 滑动窗口）的工作量、中间存储、局部性和并行性（数出来的，和机器无关）。
      演示 3：在你的 CPU 上实测：同一个算法，不同的调度生成不同的 C 代码，快慢差几倍——而且最快的是哪个取决于机器。
      演示 4：更深的流水线（4 级 3×3 模糊）：分块融合的重复计算随级数和块大小怎样变；全部内联会怎样。

【关于数字】演示 3、4 的时间是墙钟时间（单线程 C，-O3），随机器变化；观察按当次运行的数字生成。
【运行】python3 01_halide_schedules.py      （约 10 秒）
【依赖】numpy；本目录的 sched_lab；../prereq4 的 cbuild（需要 C 编译器，没有时跳过实测）；../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 blur_pipeline, bounds, box_chain, evaluations, gen_c, numpy_reference, run_pipeline_c  # noqa: E402
from cbuild import best_time, find_cc  # noqa: E402
from console import fmt_bytes, fmt_time, print_header, print_observations, print_table  # noqa: E402


def demo1_algorithm():
    print_header("演示 1：算法与边界推断")
    pipe, _ = blur_pipeline()
    inp, bx, by = pipe
    print("算法（Halide 的写法，只说每个点怎么算）：")
    print("    blurx(x, y) = (in(x−1, y) + in(x, y) + in(x+1, y)) / 3")
    print("    blury(x, y) = (blurx(x, y−1) + blurx(x, y) + blurx(x, y+1)) / 3\n")
    rows = []
    for region in ((0, 64, 0, 64), (0, 4096, 0, 4096)):
        b = bounds(pipe, region)
        for f in (by, bx, inp):
            r = b[f]
            rows.append([f"输出 {region[1]}×{region[3]}", f.name, f"[{r[0]}, {r[1]}) × [{r[2]}, {r[3]})", (r[1] - r[0]) * (r[3] - r[2])])
    print_table(["要算的输出", "Func", "需要的区域 x × y", "点数"], rows)
    print_observations([
        "算法里没有循环、没有缓冲区、没有\"先算哪个\"：每个 Func 是定义在整个平面上的函数。调度另外说明怎样把它变成循环。",
        "边界推断从输出往回推：blury 的每个点要 blurx 上、中、下三个点，所以要算 64×64 的输出，blurx 要多算上下各一行（64×66），"
        "输入还要再多左右各一列（66×66）。这些区域是编译器自动算出来的——换一种调度，区域跟着变。",
    ])


def demo2_tradeoff():
    print_header("演示 2：四种调度的取舍（4096×4096 的输出；数出来的，和机器无关）")
    pipe, _ = blur_pipeline()
    W = H = 4096
    rows = []
    for name, sched, tile, loc, par in (
            ("root：先算完整个 blurx，再算 blury", "root", (64, 64), "差：两次用到同一个 blurx 点之间隔着整张图", "好：每级内所有点独立，但两级之间要等"),
            ("inline：blury 每用一次 blurx 就现算", "inline", (64, 64), "最好：算完马上用", "好：所有点独立"),
            ("分块 64×64：每块里算这一块要的 blurx", "tile", (64, 64), "好：在一个块内", "好：块之间独立"),
            ("分块 4096×32（整行的条带）", "tile", (4096, 32), "好：在一个条带内", "好：条带之间独立"),
            ("滑动窗口：按行推进，只留最近 3 行 blurx", "sliding", (64, 64), "好：3 行以内", "差：行与行之间串行")):
        ev, st, ops = evaluations(pipe, {"blurx": sched}, W, H, tile)
        rows.append([name, f"{ev[pipe[1]] / (W * H):.3f}", f"{ops / (W * H):.2f}", fmt_bytes(st[pipe[1]] * 4), loc, par])
    print_table(["调度", "blurx 每个输出点算几次", "每个输出点的运算", "blurx 的存储", "局部性", "并行性"], rows)
    print_observations([
        "root 不重复计算，但要存一整张 blurx（64 MB），而且算 blurx 和用 blurx 隔着整张图——数据早已被挤出 cache。",
        "inline 局部性最好、不用存储，但 blurx 的每个点被算了 3 次（每个输出点要上中下三个）：运算从每点 6 次涨到 12 次。",
        "分块：每块多算上下各一行的 blurx，64×64 的块多算 3%；块越小、多算得越多。条带（整行宽 × 32 行）多算 6%。",
        "滑动窗口：每个 blurx 只算一次、只存 3 行，局部性和工作量都最好——代价是行与行之间有依赖，不能按行并行（Halide 的做法是先按条带切开，"
        "每个条带内部再滑动：用一点重复计算换回并行性）。",
        "这就是 Halide 论文说的三角取舍：并行性、局部性、重复计算——没有一种调度三者都最好。",
    ])


def demo3_measure():
    print_header("演示 3：本机实测（单线程 C，-O3）")
    if find_cc() is None:
        print("没有 C 编译器，跳过。")
        return
    pipe, R = blur_pipeline()
    rng = np.random.default_rng(0)
    rows, res = [], {}
    for W in (2048, 8192):
        H = W
        img = rng.random((H + 2 * R, W + 2 * R)).astype(np.float32)
        ref = numpy_reference(pipe, img, R, W, H)
        for name, s, tile in (("root", "root", (64, 64)), ("inline", "inline", (64, 64)), ("分块 64×64", "tile", (64, 64)),
                              ("分块 256×32", "tile", (256, 32)), (f"条带 {W}×32", "tile", (W, 32)), ("滑动窗口", "sliding", (64, 64))):
            out, call = run_pipeline_c(pipe, {"blurx": s}, img, R, W, H, tile=tile)
            t = best_time(call, 3)
            res[(W, name if not name.startswith("条带") else "条带")] = t
            rows.append([f"{W}×{W}", name, fmt_time(t), f"{np.abs(out - ref).max():.1e}"])
    print_table(["图像", "调度", "时间", "和 numpy 参照的最大差"], rows)
    print("\n生成的 C 代码（分块 64×64）：")
    print(gen_c(pipe, {"blurx": "tile"}, R, (64, 64)))
    best_big = min((k for k in res if k[0] == 8192), key=res.get)[1]
    worst_big = max((k for k in res if k[0] == 8192), key=res.get)[1]
    print_observations([
        "同一个算法、同样的数学结果（差在 1e-7 量级，只是 inline 的加法顺序不同），只换调度，速度差几倍。"
        f"8192×8192 时最快的是\"{best_big}\"、最慢的是\"{worst_big}\"（{res[(8192, worst_big)] / res[(8192, best_big)]:.1f} 倍）。",
        f"inline 每个点多算一倍的运算：在 8192 上是 root 的 {res[(8192, 'inline')] / res[(8192, 'root')]:.1f} 倍。"
        f"滑动窗口和条带都避免了把整张 blurx（8192 时 268 MB）写出去再读回来：8192 时比 root 快 "
        f"{res[(8192, 'root')] / res[(8192, '滑动窗口')]:.2f} 倍 / {res[(8192, 'root')] / res[(8192, '条带')]:.2f} 倍；"
        f"2048 时 blurx 只有 16 MB，这台机器的 cache 和内存带宽又都很大，root 的代价小得多（我们没有测 cache 命中率）：条带是 root 的 {res[(2048, '条带')] / res[(2048, 'root')]:.2f} 倍、"
        f"滑动窗口是 {res[(2048, '滑动窗口')] / res[(2048, 'root')]:.2f} 倍。",
        f"方块分块在这台机器上反而很慢（64×64 是 root 的 {res[(8192, '分块 64×64')] / res[(8192, 'root')]:.1f} 倍）："
        "每块只有几十个连续的数，一行一行地跳着读，硬件预取和 TLB 都不友好；生成的代码虽然被向量化了（编译器的报告），每块的开销仍然很大。"
        "整行的条带保留了连续读，又把 blurx 留在 cache 里。我们没有进一步分析原因——这正说明调度的好坏很难靠直觉判断。",
        "Halide 论文用自动调优在调度空间里搜索，找到的往往是\"反直觉\"的调度；不同的机器上最优调度也不同（第 10 讲的寄存器分块是同一个教训）。"
        "这个单线程的实验也没有体现并行性——多核时滑动窗口的串行依赖就成了劣势。",
    ])


def demo4_chain():
    print_header("演示 4：4 级 3×3 模糊串成的流水线（2048×2048）")
    ch, R = box_chain(4)
    W = H = 2048
    rows = []
    for name, s, tile in (("全部 root", "root", (64, 64)), ("全部 inline", "inline", (64, 64)), ("全部分块 16×16", "tile", (16, 16)),
                          ("全部分块 32×32", "tile", (32, 32)), ("全部分块 64×64", "tile", (64, 64)), ("全部分块 128×128", "tile", (128, 128))):
        sched = {f.name: s for f in ch[1:-1]}
        _, st, ops = evaluations(ch, sched, W, H, tile)
        rows.append([name, s, tile, f"{ops / (W * H):.1f}", fmt_bytes(sum(st.values()) * 4)])
    t_meas = {}
    if find_cc() is not None:
        rng = np.random.default_rng(1)
        img = rng.random((H + 2 * R, W + 2 * R)).astype(np.float32)
        ref = numpy_reference(ch, img, R, W, H)
        for r in rows:
            sched = {f.name: r[1] for f in ch[1:-1]}
            out, call = run_pipeline_c(ch, sched, img, R, W, H, tile=r[2])
            assert np.abs(out - ref).max() < 1e-4
            t = best_time(call, 2 if r[1] == "inline" else 3)
            t_meas[r[0]] = t
            r.append(fmt_time(t))
    print_table(["调度", "每个输出点的运算", "中间存储"] + (["本机时间"] if t_meas else []), [[r[0]] + r[3:] for r in rows])
    print_observations([
        "每级 9 个点的平均（8 次加法 + 1 次乘法），4 级不重复计算时每个输出点 36 次运算（root）。",
        "全部 inline：第 4 级的每个点要第 3 级的 9 个点，每个又要第 2 级的 9 个……运算按 9 的幂增长到每点 7380 次——深的模板流水线不能全部内联。"
        + (f"实测是 root 的 {t_meas['全部 inline'] / t_meas['全部 root']:.0f} 倍。" if t_meas else ""),
        "分块融合：每一级的区域都比下一级大一圈，块越小，\"多算的那一圈\"占比越大：16×16 的块每点 51.5 次运算，128×128 只多 4%。"
        "块越大，中间存储越多（缓冲放不进 L1 / 共享内存时，融合的好处就没了）——在重复计算和存储之间又是一个折中。"
        + (f"在这台机器上，全部 root（{fmt_time(t_meas['全部 root'])}）比任何一种分块融合都快（最快的分块 "
           f"{fmt_time(min(v for k, v in t_meas.items() if k.startswith('全部分块')))}）：中间结果 50 MB，和演示 3 的 2048×2048 一样，在这台机器上写出去再读回来的代价不高。"
           if t_meas else ""),
        "这和 GPU 上把几个卷积层融合（层越多、感受野越大、重叠的\"光晕\"越宽）、FlashAttention 的重算（第 11 讲）是同一类决定。",
    ])


if __name__ == "__main__":
    demo1_algorithm()
    demo2_tradeoff()
    demo3_measure()
    demo4_chain()