""" ============================================================================= CS349E · 前置讲 P4 · 05_schedules_and_autotuning.py —— 算法与调度分离、代价模型与自动调优 ============================================================================= 【这个脚本要回答的问题】 04 脚本手工试了几种循环写法。真实的编译器面对的是一个巨大的"调度空间"（循环顺序 × 分块大小 × 展开 × 向量化 × 并行……），手工试不过来。两条路： (a) 代价模型（cost model）：不运行，靠分析预测哪个调度快； (b) 自动调优（autotuning）：真的编译、运行、计时，挑最快的。 Halide、TVM（AutoTVM / Ansor）、Triton 的 autotune、cuBLAS 的启发式都是这两条路的组合。本脚本： 演示 1：算法与调度分离——同一行"算法"，几种"调度"，生成的 C 代码。 演示 2：代价模型 vs 模拟器——在小规模上，解析模型（loop_ir.predict_misses）和 LRU 模拟的缺失数排序一致吗？ 演示 3：调度空间实测——36 个调度在 512³ 上编译运行；三个代价模型（只看 cache / 加上"最内层能否向量化" / 再加上"编译器会展开很短的循环"）预测的排序和实测排序的相关性（Spearman）。 演示 4：搜索策略——随机试 b 个 vs 按代价模型挑前 b 个：用多少次实测能找到接近最优的调度？ 【怎么读输出】演示 1、2 每次相同；演示 3、4 依赖墙钟测量（编写时是 Apple M4 Pro 单核），具体数字会变，看趋势。 演示 4 的搜索是在演示 3 已经测好的表上"回放"的（每种策略重复 1000 次随机顺序），不需要重新测。 【运行】python3 05_schedules_and_autotuning.py （约 20 秒） python3 05_schedules_and_autotuning.py --quick （约 5 秒：256³、空间小一些） 【依赖】numpy；从 ../lecture01 导入 console；同目录的 loop_ir、cbuild。没有 C 编译器时只做演示 1、2。 """ import itertools import os import sys import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "lecture01")) from console import print_header, print_observations, print_table # noqa: E402 from cbuild import best_time, find_cc # noqa: E402 from loop_ir import default_layout, lru_misses, matmul, predict_misses # noqa: E402 QUICK = "--quick" in sys.argv HAVE_CC = find_cc() is not None N_RUN = 256 if QUICK else 512 # 代价模型 v2 用到的粗略硬件常数（编写时的机器：Apple M4 Pro 性能核；换机器可以改，模型只用来排序） L1_BYTES, LINE = 128 * 1024, 128 # L1 数据 cache 128 KiB，cache line 128 字节 FREQ = 4.5e9 # 主频（Hz） FMA_LATENCY = 4 # 一次浮点乘加的延迟（周期）：依赖链上每步至少等这么久 SIMD_FMA_PER_CYCLE = 16 # 向量化时每周期最多做多少个 float32 乘加（4 条 128 位流水线 × 4 个 lane） LINE_REFILL_BW = 150e9 # 从 L2 补一行进 L1 的带宽（B/s），很粗的估计 # --------------------------------------------------------------------------- # 调度空间 # --------------------------------------------------------------------------- def schedule_space(N): """ 返回 [(名字, LoopNest)]：6 种不分块的循环顺序 + 分块（T × 块外顺序 × 块内顺序）。 分块只取能整除 N 的 T（不整除时的守卫会让性能断崖，见 04 脚本），名字写成 "T32 ikj|ijk"： 竖线左边是块的遍历顺序，右边是块内的循环顺序。 """ space = [("".join(o), matmul(N).reorder(*o)) for o in itertools.permutations("ijk")] Ts = (16, 32, 64) if QUICK else (8, 16, 32, 64, 128) for T in Ts: for outer in (("io", "jo", "ko"), ("io", "ko", "jo")): for inner in (("ii", "ki", "ji"), ("ii", "ji", "ki"), ("ki", "ii", "ji")): nest = matmul(N).split("i", T).split("j", T).split("k", T).reorder(*outer, *inner) name = f"T{T} {''.join(v[0] for v in outer)}|{''.join(v[0] for v in inner)}" space.append((name, nest)) return space # --------------------------------------------------------------------------- # 代价模型 v2：缓存 + 最内层循环的特征 # --------------------------------------------------------------------------- def innermost_features(nest, unroll_limit=None, level=-1): """ 最内层循环每走一步，每个访问的地址跳几个元素。返回 (是否归约, 是否可向量化)： 归约：写的目标不随最内层变量变化（所有迭代加到同一个位置，例如最内层是 k 时的 C[i,j]）； 可向量化：不是归约，而且所有访问的跳跃都是 0 或 ±1（连续访问，SIMD 能一次装 4 个）。 【unroll_limit】给定时模拟"编译器会把很短的最内层循环完全展开"：如果最内层是归约、而且范围 ≤ unroll_limit， 就改看外面一层循环（展开以后它成了真正的最内层）。这是模型 v3 的规则（见 demo3 的说明）。 """ if unroll_limit is not None and level == -1: red, _ = innermost_features(nest) if red and nest.loops[-1].extent <= unroll_limit and len(nest.loops) > 1: return innermost_features(nest, level=-2) v = nest.loops[level].var env0 = {l.var: 0 for l in nest.loops} env1 = dict(env0) env1[v] = 1 o0 = {x: eval(nest.expr[x], {}, env0) for x in nest.orig} o1 = {x: eval(nest.expr[x], {}, env1) for x in nest.orig} strides, red = [], True for name, idx, is_w in nest.accesses(): shape = nest.shapes[name] def lin(o): val = eval(idx[0], {}, o) for d in range(1, len(shape)): val = val * shape[d] + eval(idx[d], {}, o) return val st = lin(o1) - lin(o0) strides.append(st) if is_w and st != 0: red = False vec = (not red) and all(abs(s) <= 1 for s in strides) return red, vec def model_v1(nest): """代价模型 v1：只看缓存——loop_ir.predict_misses 预测的 L1 缺失数（越小越快）。""" return predict_misses(nest, L1_BYTES, LINE)[0] def model_v2(nest, N, unroll_limit=None): """ 代价模型 v2：预测时间 = 访存时间 + 计算时间。 访存：v1 的缺失数 × 一行的字节 / 补行带宽； 计算：N³ 次乘加；归约 → 依赖链，每次等 FMA_LATENCY 个周期；可向量化 → 每周期 SIMD_FMA_PER_CYCLE 个； 其他 → 每周期 1 个。 【它仍然看不见的】编译器自己的变换（例如把很短的内层循环完全展开后改为向量化外层）、组相联冲突、TLB、预取。 """ red, vec = innermost_features(nest, unroll_limit) t_mem = model_v1(nest) * LINE / LINE_REFILL_BW iters = N ** 3 t_cmp = iters * FMA_LATENCY / FREQ if red else (iters / SIMD_FMA_PER_CYCLE / FREQ if vec else iters / FREQ) return t_mem + t_cmp def model_v3(nest, N): """代价模型 v3 = v2 + "编译器会把范围 ≤ 16 的最内层归约循环完全展开、转而向量化外面一层"。""" return model_v2(nest, N, unroll_limit=16) def rank(x): """把数组变成名次（从 0 开始；相同的值取平均名次）。【Python 小知识】argsort 两次得到每个元素的名次。""" x = np.asarray(x, dtype=float) order = np.argsort(x, kind="stable") r = np.empty(len(x)) r[order] = np.arange(len(x)) for v in np.unique(x): m = x == v r[m] = r[m].mean() return r def spearman(a, b): """Spearman 等级相关系数：两组数的名次之间的皮尔逊相关。1 = 排序完全一致，0 = 无关，−1 = 完全相反。""" ra, rb = rank(a), rank(b) ra -= ra.mean() rb -= rb.mean() return float((ra * rb).sum() / np.sqrt((ra ** 2).sum() * (rb ** 2).sum())) # --------------------------------------------------------------------------- # 演示 1：算法与调度分离 # --------------------------------------------------------------------------- def demo1_algorithm_vs_schedule(): """ Halide 的口号："算法"说算什么，"调度"说怎么算，两者分开写。这里同一个算法，三种调度。 """ print_header("演示 1：算法与调度分离——同一个算法，几种调度") N = 512 algo = matmul(N) print(f"算法（只写一次）：{algo.stmts[0].text}，i、j、k ∈ [0, {N})") scheds = [ ("默认", algo.copy()), ("交换", algo.copy().reorder("i", "k", "j")), ("分块 + 展开", algo.copy().split("i", 32).split("j", 32).split("k", 32) .reorder("io", "ko", "jo", "ii", "ki", "ji").unroll("ki", 4)), ] rows = [[name, "; ".join(s.log) if s.log else "（无）", s.describe()] for name, s in scheds] print_table(["调度", "调度原语", "得到的循环（从外到内，方括号是范围）"], rows) print("\n'分块 + 展开'生成的 C 代码：") print(" " + scheds[2][1].to_c().replace("\n", "\n ")) print_observations([ "语句 C[i,j] += A[i,k] * B[k,j] 在三种调度里一字未改；变的只是循环结构和'原始变量怎么由循环变量算出来'（i = io*32 + ii）。" "这让我们可以放心地试很多调度：正确性由算法保证（只要调度合法，03 脚本），性能由调度决定。", "Halide（图像处理）、TVM（深度学习）都把这个思想做成了语言；Triton 走得更远，让程序员直接写'块级'的程序，" "块内的线程映射、向量化、共享内存交给编译器（第 13 讲）。", "unroll(ki, 4) 只是一个提示（#pragma）：要不要展开、怎么展开，最后由 C 编译器决定。调度语言能说的话越多，" "编译器自己做的决定就越少，结果越可预测。", ]) # --------------------------------------------------------------------------- # 演示 2：代价模型 vs 模拟器 # --------------------------------------------------------------------------- def demo2_model_vs_simulator(N=48, cache=4096, line=64): """ 在 N = 48 的矩阵乘法上，对一组调度比较：解析模型预测的缺失数 vs 全相联 LRU 模拟的真实缺失数。 （小规模才能模拟；cache 也取小的 4 KiB，让工作集和 cache 的关系与大规模时相似。） """ print_header(f"演示 2：代价模型 vs 模拟器——{N}³ 的矩阵乘法、{cache // 1024} KiB cache、{line} 字节一行") space = [("".join(o), matmul(N).reorder(*o)) for o in itertools.permutations("ijk")] for T in (4, 8, 16, 24): for inner in (("ii", "ki", "ji"), ("ii", "ji", "ki")): nest = matmul(N).split("i", T).split("j", T).split("k", T).reorder("io", "ko", "jo", *inner) space.append((f"T{T} ikj|{''.join(v[0] for v in inner)}", nest)) rows, pred, sim = [], [], [] lay = default_layout(matmul(N).shapes) for name, nest in space: p, L, F = predict_misses(nest, cache, line) s = lru_misses(nest.address_trace(lay), cache, line) pred.append(p) sim.append(s) rows.append([name, f"{p:,}", f"{s:,}", f"{p / s:.1f}"]) print_table(["调度", "模型预测的缺失", "LRU 模拟的缺失", "预测 / 模拟"], rows) rho = spearman(pred, sim) ratio = np.array(pred) / np.array(sim) print(f"\nSpearman 等级相关：{rho:.2f}；预测 / 模拟 的比值在 {ratio.min():.1f} 到 {ratio.max():.1f} 之间。") print_observations([ f"解析模型的**绝对数值**偏差很大（预测 / 模拟在 {ratio.min():.1f}–{ratio.max():.1f} 倍之间）：它假设'工作集放得下的那一层每执行一次，" "就把工作集整个重新搬一遍'，忽略了相邻两次执行之间的数据重叠，所以总是高估。", f"但它的**排序**和模拟相当一致（Spearman {rho:.2f}）：哪个调度局部性更好、分块能不能放进 cache，它分得清。" "对编译器来说，代价模型只需要排序对，不需要数值准。", "模型几乎不花时间（只枚举一两层循环的工作集），模拟要把几十万次访问走一遍——这就是'代价模型'存在的理由。", ]) # --------------------------------------------------------------------------- # 演示 3：调度空间实测 # --------------------------------------------------------------------------- def demo3_measure_space(): """ 在 N_RUN³ 上编译、运行调度空间里的每一个调度（-O2，每个取 3 次中最快的），同时算两个代价模型的预测。 返回 [(名字, 实测时间, v1 预测, v2 预测)]。 """ space = schedule_space(N_RUN) print_header(f"演示 3：调度空间实测——{len(space)} 个调度，{N_RUN}³，-O2") rng = np.random.default_rng(0) A = rng.standard_normal((N_RUN, N_RUN)).astype(np.float32) B = rng.standard_normal((N_RUN, N_RUN)).astype(np.float32) ref = A.astype(np.float64) @ B C = np.zeros((N_RUN, N_RUN), np.float32) def setup(): C[...] = 0.0 table = [] for name, nest in space: f = nest.compile_c() t = best_time(lambda: f({"A": A, "B": B, "C": C}), repeats=3, setup=setup) if not np.allclose(C, ref, atol=1e-2 * np.sqrt(N_RUN)): raise AssertionError(f"{name} 的结果不对") red, vec = innermost_features(nest) table.append((name, t, model_v1(nest), model_v2(nest, N_RUN), red, vec, model_v3(nest, N_RUN), nest.loops[-1].extent)) table.sort(key=lambda r: r[1]) rows = [] for i, (name, t, v1, v2, red, vec, v3, ext) in enumerate(table): if i < 10 or i >= len(table) - 4: feat = "归约" if red else ("可向量化" if vec else "跳跃访问") rows.append([i + 1, name, f"{t * 1e3:.1f} ms", f"{2 * N_RUN ** 3 / t / 1e9:.1f}", f"{v1:,}", f"{feat}（{ext} 次）", f"{v2 * 1e3:.1f} ms", f"{v3 * 1e3:.1f} ms"]) elif i == 10: rows.append(["…", "", "", "", "", "", "", ""]) print_table(["名次", "调度", "实测", "GFLOP/s", "v1：预测 L1 缺失", "最内层（范围）", "v2：预测时间", "v3：预测时间"], rows) times = [r[1] for r in table] r1 = spearman([r[2] for r in table], times) r2 = spearman([r[3] for r in table], times) r3 = spearman([r[6] for r in table], times) best = table[0] picks = {k: min(table, key=key) for k, key in (("v1", lambda r: (r[2], r[3])), ("v2", lambda r: r[3]), ("v3", lambda r: r[6]))} print(f"\nSpearman（预测 vs 实测）：v1 = {r1:.2f}，v2 = {r2:.2f}，v3 = {r3:.2f}。") print(f"实测最快：{best[0]}（{best[1] * 1e3:.1f} ms）；模型的首选：" + "；".join(f"{k} → {p[0]}（{p[1] * 1e3:.1f} ms）" for k, p in picks.items()) + "。") top = table[:10] n_vec = sum(1 for r in top if r[5]) n_short_red = sum(1 for r in top if r[4] and r[7] <= 16) slow_red = min((r for r in table if r[4] and r[0].startswith("T") and r[7] >= 32), key=lambda r: (r[2], r[1]), default=None) lines = [ f"最快和最慢的调度相差 {times[-1] / times[0]:.0f} 倍。前 10 名里，{n_vec} 个是'分块 + 最内层连续访问、可向量化'" + (f"；另外 {n_short_red} 个最内层是归约，但只有 8 或 16 次。" if n_short_red else "。"), f"只看缓存的 v1 和实测的 Spearman 只有 {r1:.2f}：它分得清'分块 vs 不分块'、'最内层连续 vs 跳跃'，" "却看不出同样分块时，块内最内层是归约（依赖链，不能向量化）还是可向量化的循环。", f"v2 加上'最内层是否归约 / 是否可向量化'后，相关提高到 {r2:.2f}；" + (f"长归约确实慢（例如 {slow_red[0]}：v1 预测的缺失是全表最少的一档，实测却要 {slow_red[1] * 1e3:.1f} ms，" f"是最快的 {slow_red[1] / best[1]:.0f} 倍）。" if slow_red else "") + "但 v2 把实测最快的几个（块很小、最内层是 8 次的 k 归约）也判成了'慢'：编译器把这么短的循环完全展开，" "转而向量化外面的 j 循环——'归约'这个特征在这里失效了。", f"v3 加了一条规则'范围 ≤ 16 的最内层归约会被完全展开'，相关提高到 {r3:.2f}。注意：这条规则是**看了实测结果之后**才加的——" "代价模型在实践中就是这样被一条条'调'出来的，也因此容易只适用于某个编译器、某台机器。" "TVM 的 Ansor 等系统干脆用实测数据**训练**代价模型。", f"三个模型的首选都是缺失最少的大块，实测却比最快的慢 {picks['v3'][1] / best[1]:.1f} 倍：T = 8 的小块快，是因为编译器把块内循环" "展开后能把数据留在寄存器里（寄存器级的复用），而我们的模型只看到了 L1。模型总是只看见它被设计去看的那一层。", ] print_observations(lines) return table # --------------------------------------------------------------------------- # 演示 4：搜索策略 # --------------------------------------------------------------------------- def demo4_search(table, trials=1000, seed=0): """ 在演示 3 测好的表上回放三种搜索策略，预算 b = 实测的次数： 随机搜索：随机挑 b 个调度实测，取最快的（重复 trials 次，取平均）； v1 / v2 引导：按代价模型的预测排序，实测前 b 个，取最快的。 指标：找到的最快时间 / 全空间真正最快的时间（1.00 = 找到了最优）。 """ print_header("演示 4：搜索策略——实测多少次才能找到接近最优的调度？") rng = np.random.default_rng(seed) times = np.array([r[1] for r in table]) best = times.min() o1 = np.lexsort(([r[3] for r in table], [r[2] for r in table])) o2 = np.argsort([r[3] for r in table], kind="stable") o3 = np.argsort([r[6] for r in table], kind="stable") budgets = [b for b in (1, 2, 4, 8, 16) if b < len(table)] + [len(table)] rows, res = [], {} for b in budgets: rand = np.mean([times[rng.choice(len(times), b, replace=False)].min() for _ in range(trials)]) / best g = [times[o[:b]].min() / best for o in (o1, o2, o3)] res[b] = (rand, *g) rows.append([b, f"{rand:.2f}"] + [f"{x:.2f}" for x in g]) print_table(["实测预算 b", "随机搜索（平均）", "v1 引导", "v2 引导", "v3 引导"], rows) b_small = 4 if 4 in res else budgets[0] b_mid = 16 if 16 in res else budgets[-2] r_s, r_m = res[b_small], res[b_mid] print_observations([ f"预算很小时（b = {b_small}），模型引导反而不如随机搜索（随机平均 {r_s[0]:.2f}，三个模型 {r_s[1]:.2f} / {r_s[2]:.2f} / {r_s[3]:.2f}）：" "模型最看好的几个调度都是'缺失最少的大块'，正好是模型系统性看错的地方——一个有偏的模型会把搜索稳稳地带到错的地方。", f"预算到 b = {b_mid} 时，最好的模型引导是 {min(r_m[1:]):.2f}，随机是 {r_m[0]:.2f}。" "模型的价值是把明显很差的调度（不分块、跳跃访问、长归约）排除在外，而不是直接选出最优。", "所以实际系统总是'模型筛选 + 实测'：模型负责缩小范围，最后的决定靠实测。真实的调度空间动辄上万甚至更多，" "纯随机搜索不现实，纯模型又不可靠。", "这就是 AutoTVM / Ansor（学习代价模型 + 进化搜索）、Triton 的 @autotune（列出几个配置、全部实测）、" "cuBLAS / CUTLASS 的启发式（为常见形状预先调好）背后的权衡。第 14 讲的'让 LLM 写内核'，" "本质上是用 LLM 代替代价模型来提出候选，再用实测来评判（AlphaEvolve 的评估器）。", ]) if __name__ == "__main__": demo1_algorithm_vs_schedule() demo2_model_vs_simulator() if HAVE_CC: results = demo3_measure_space() demo4_search(results) else: print("\n没有找到 C 编译器，跳过演示 3、4（需要编译并运行生成的代码）。")