""" ============================================================================= CS349E · 前置讲 P4 · 04_loop_transformations.py —— 循环变换：交换、分块、展开与向量化、融合（模拟 + 实测） ============================================================================= 【这个脚本要回答的问题】 03 脚本回答了"哪些变换是合法的"；这里回答"哪些变换是**有利的**、利有多大"。主角是矩阵乘法 C[i,j] += A[i,k] * B[k,j]——同一个算法，换个写法，速度能差几十倍： 演示 1：循环交换——6 种循环顺序。先用 P1 的 cache 模拟器（组相联、LRU、3C 分类）在小矩阵上数缺失， 再生成 C 代码、编译、在 512×512 上实测。 演示 2：分块（tiling）——块大小从 8 到 256，模拟的缺失和实测时间怎样变化；块大小不整除时的"守卫"。 演示 3：编译器选项与"浮点不满足结合律"——-O0 / -O1 / -O2 / -O3；点积形式的循环为什么要 -ffast-math 才能向量化。 演示 4：循环融合——生产者 / 消费者两个循环合成一个，内存流量和时间。 【怎么读输出】模拟部分（缺失数）每次相同；实测部分是墙钟时间（编写时是 Apple M4 Pro 单核，只看趋势）。 【运行】python3 04_loop_transformations.py （约 15 秒） python3 04_loop_transformations.py --quick （约 5 秒：小矩阵、少几个配置） 【依赖】numpy；从 ../lecture01 导入 console，从 ../prereq1 导入 arch_lab（cache 模拟器）；同目录的 loop_ir、cbuild。 没有 C 编译器时，只做模拟部分。 """ 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")) sys.path.insert(0, os.path.join(HERE, "..", "prereq1")) from console import print_header, print_observations, print_table # noqa: E402 from arch_lab import Cache # noqa: E402 from cbuild import best_time, find_cc # noqa: E402 from loop_ir import LoopNest, default_layout, fuse, lru_misses, matmul # noqa: E402 QUICK = "--quick" in sys.argv HAVE_CC = find_cc() is not None N_RUN = 256 if QUICK else 512 # 实测用的矩阵大小 N_SIM = 32 if QUICK else 64 # 模拟用的矩阵大小（Python 模拟器慢，只能用小矩阵） SIM_CACHE = 4096 # 模拟的 cache：4 KiB、4 路组相联、64 字节一行（和小矩阵的比例大致对应真实 L1 和 512 矩阵） def mm_inputs(N, seed=0): rng = np.random.default_rng(seed) A = rng.standard_normal((N, N)).astype(np.float32) B = rng.standard_normal((N, N)).astype(np.float32) return A, B, A.astype(np.float64) @ B.astype(np.float64) def measure(nest, A, B, ref, flags=("-O2",), repeats=3): """编译 nest 生成的 C 代码，在 A、B 上运行（每次运行前把 C 清零），返回 (最短时间, 结果是否正确)。""" f = nest.compile_c(flags=flags) C = np.zeros(ref.shape, np.float32) def setup(): C[...] = 0.0 t = best_time(lambda: f({"A": A, "B": B, "C": C}), repeats=repeats, setup=setup) ok = np.allclose(C, ref, atol=1e-2 * np.sqrt(A.shape[1])) return t, ok def gflops(N, t): return 2 * N ** 3 / t / 1e9 # --------------------------------------------------------------------------- # 演示 1：循环交换 # --------------------------------------------------------------------------- def inner_strides(nest): """最内层循环每走一步，A、B、C 的地址各跳几个元素（0 = 不动，1 = 连续）。""" v = nest.loops[-1].var out = [] for name in ("A", "B", "C"): for s in nest.stmts: for n, idx in s.reads + [s.target]: if n == name: env0 = {x: 0 for x in nest.orig} env1 = dict(env0) env1[v] = 1 lin = [eval(idx[0], {}, e) * nest.shapes[n][1] + eval(idx[1], {}, e) for e in (env0, env1)] out.append(f"{name}:{lin[1] - lin[0]}") break else: continue break return " ".join(out) def demo1_interchange(): """ 6 种循环顺序： 模拟：N_SIM×N_SIM 的矩阵、4 KiB 的 cache（P1 的 arch_lab.Cache，4 路组相联，3C 分类）； 实测：N_RUN×N_RUN，生成的 C 代码用 -O2 编译。 """ print_header(f"演示 1：循环交换——6 种顺序（模拟 {N_SIM}³、4 KiB cache；实测 {N_RUN}³、-O2）") orders = ["ijk", "ikj", "jik", "jki", "kij", "kji"] A, B, ref = mm_inputs(N_RUN) rows, res, sim = [], {}, {} for o in orders: small = matmul(N_SIM).reorder(*o) trace = small.address_trace(default_layout(small.shapes)) c = Cache(SIM_CACHE, line=64, ways=4, name="sim") c.run(trace) miss = sum(c.misses.values()) sim[o] = (1 - c.hit_rate(), dict(c.misses)) row = [o, inner_strides(small), f"{1 - c.hit_rate():.1%}", f"{miss:,}", f"{c.misses['compulsory']:,} / {c.misses['capacity']:,} / {c.misses['conflict']:,}"] if HAVE_CC: t, ok = measure(matmul(N_RUN).reorder(*o), A, B, ref) res[o] = t row += [f"{t * 1e3:.1f} ms", f"{gflops(N_RUN, t):.2f}", "✓" if ok else "✗"] rows.append(row) heads = ["顺序", "最内层每步的跳跃（元素）", "缺失率", "缺失数", "强制 / 容量 / 冲突"] if HAVE_CC: heads += ["实测时间", "GFLOP/s", "正确"] print_table(heads, rows) jki_rate, jki_m = sim["jki"] lines = [ f"最内层循环决定了一切：ikj、kij 的最内层是 j，B 和 C 都按行连续访问（跳跃 1），A 不动（跳跃 0），" f"缺失率只有 {sim['ikj'][0]:.1%} / {sim['kij'][0]:.1%}。", f"jki、kji 的最内层是 i，A 和 C 每一步跳过一整行（跳跃 {N_SIM}）：每次访问都碰一个新的 cache line，" f"一列 {N_SIM} 个元素就要 {N_SIM} 行，A、C 两列加起来 {2 * N_SIM} 行，" f"{'超过了' if 2 * N_SIM > SIM_CACHE // 64 else '塞满了'} 4 KiB cache 的 {SIM_CACHE // 64} 行——" f"缺失率 {jki_rate:.1%}（约每 {1 / jki_rate:.0f} 次访问缺失一次），几乎全是容量缺失（{jki_m['capacity']:,} 次）。", "ijk、jik 的最内层是 k：A 连续、B 跳行、C 不动，介于两者之间。冲突缺失在这组例子里很少——" "4 路组相联已经能化解大部分'行距是 2 的幂'的冲突（对比 P1 §2.6 的直接映射例子）。", ] if HAVE_CC: best, worst = min(res, key=res.get), max(res, key=res.get) lines.append(f"实测：最快的 {best}（{res[best] * 1e3:.1f} ms）比最慢的 {worst}（{res[worst] * 1e3:.1f} ms）快 " f"{res[worst] / res[best]:.0f} 倍——同一个算法、同样的运算次数，只是换了循环顺序。") lines.append("实测的差距比模拟的缺失数差距更大：最内层连续访问时编译器还能用 SIMD 向量化（一条指令算 4 个 float），" "跳跃访问时不能；ijk 的最内层是对 C[i,j] 的累加，形成一条依赖链（P1 §1.6），也不能向量化。") print_observations(lines) return res # --------------------------------------------------------------------------- # 演示 2：分块 # --------------------------------------------------------------------------- def tiled(N, T, inner=("ii", "ki", "ji")): """把 i、j、k 都按 T 分块，外层 io, ko, jo，内层默认 ii, ki, ji（块内是 ikj 顺序）。""" return matmul(N).split("i", T).split("j", T).split("k", T).reorder("io", "ko", "jo", *inner) def demo2_tiling(untiled): """ 分块：先在块的层面遍历（io, ko, jo），再在块内遍历（ii, ki, ji）。块内的 A、B、C 子矩阵一共 3·T² 个数， 放得进 cache 时，块内的计算全部命中，每个数被复用 T 次。 模拟用全相联 LRU（loop_ir.lru_misses，比 P1 的模拟器快）；实测用 -O2。 """ print_header(f"演示 2：分块——块大小的影响（模拟 {N_SIM}³、4 KiB 全相联；实测 {N_RUN}³、-O2）") A, B, ref = mm_inputs(N_RUN) base_sim = lru_misses(matmul(N_SIM).reorder("i", "k", "j").address_trace(default_layout(matmul(N_SIM).shapes)), SIM_CACHE) rows = [["不分块（ikj）", "—", f"{base_sim:,}", ""]] if HAVE_CC: rows[0] += [f"{untiled['ikj'] * 1e3:.1f} ms", f"{gflops(N_RUN, untiled['ikj']):.2f}", "✓"] res, sims, guarded = {}, {}, set() sizes = (4, 8, 12, 16, 32) if QUICK else (8, 16, 24, 32, 48, 64, 128, 256) for T in sizes: Ts = min(T, N_SIM) small = tiled(N_SIM, Ts) m = lru_misses(small.address_trace(default_layout(small.shapes)), SIM_CACHE) sims[T] = m ws = 3 * Ts * Ts * 4 row = [f"T = {T}", f"{ws / 1024:.1f} KiB" + ("" if Ts == T else f"（模拟取 T = {Ts}）"), f"{m:,}", ""] if HAVE_CC: nest = tiled(N_RUN, T) t, ok = measure(nest, A, B, ref) res[T] = t row[3] = "有守卫" if nest.guards else "" if nest.guards: guarded.add(T) row += [f"{t * 1e3:.1f} ms", f"{gflops(N_RUN, t):.2f}", "✓" if ok else "✗"] rows.append(row) heads = ["配置", "块内工作集 3T²×4B", "模拟缺失数", "备注"] if HAVE_CC: heads += ["实测时间", "GFLOP/s", "正确"] print_table(heads, rows) fit = [T for T in sims if 3 * min(T, N_SIM) ** 2 * 4 <= SIM_CACHE] nofit = [T for T in sims if 3 * min(T, N_SIM) ** 2 * 4 > SIM_CACHE] lines = [ f"模拟里，块内工作集放得进 4 KiB 时（T = {', '.join(map(str, fit))}），缺失数最少只有不分块的 " f"{min(sims[T] for T in fit) / base_sim:.0%}；块太大、放不下以后（T ≥ {min(nofit)}），缺失数回到不分块的量级——" "块大小要和 cache 大小匹配。", ] if HAVE_CC: clean = {T: t for T, t in res.items() if T not in guarded} bt = min(clean, key=clean.get) lines.append(f"实测最快的是 T = {bt}（{res[bt] * 1e3:.1f} ms，{gflops(N_RUN, res[bt]):.1f} GFLOP/s），比不分块的 ikj 快 " f"{untiled['ikj'] / res[bt]:.1f} 倍。真实 CPU 的 L1 有 128 KiB、L2 有 16 MiB，{N_RUN}×{N_RUN} 的矩阵本来就放得进 L2，" "所以实测的收益没有模拟里那么夸张；" + (f"最快的是很小的块（T = {bt}），主要是因为块内循环短到可以被编译器完全展开、" "数据留在寄存器里（'寄存器分块'，第 10 讲）。" if bt <= 16 else "块内工作集放得进 L1 时最快。")) if guarded: g = sorted(guarded) worst = max(res[T] for T in g) lines.append(f"不整除的块大小（T = {', '.join(map(str, g))}）结果仍然正确——loop_ir.split 自动加了 if 守卫——" f"但慢得惊人（最慢 {worst * 1e3:.1f} ms，是 T = {bt} 的 {worst / res[bt]:.0f} 倍）：最内层循环里的 if " "让编译器无法向量化。实际的编译器会把'整块'和'剩下的边角'分开生成两份代码，热循环里不做判断。") print_observations(lines) return res # --------------------------------------------------------------------------- # 演示 3：编译器选项与浮点结合律 # --------------------------------------------------------------------------- def demo3_compiler_flags(tile_res): """ 1. 同一个分块的 C 代码，用 -O0 / -O1 / -O2 / -O3 编译。 2. 点积形式的矩阵乘法：先把 B 转置成 BT，C[i,j] += A[i,k] * BT[j,k]（最内层 k，A 和 BT 都连续访问）。 最内层是一个归约（所有 k 加到同一个 C[i,j]）。向量化需要把求和顺序改成"4 路分别累加、最后合并"—— 这改变了浮点舍入（P3 §4.2），所以编译器默认不敢做；-ffast-math 允许它重新结合。 """ print_header("演示 3：编译器选项——优化级别，以及浮点归约为什么需要 -ffast-math 才能向量化") if not HAVE_CC: print("没有 C 编译器，跳过。") return A, B, ref = mm_inputs(N_RUN) clean = {T: t for T, t in tile_res.items() if N_RUN % T == 0} T = min(clean, key=clean.get) if clean else 32 rows, opt_t = [], {} for flags in (("-O0",), ("-O1",), ("-O2",), ("-O3",)): t, ok = measure(tiled(N_RUN, T), A, B, ref, flags=flags, repeats=2 if flags == ("-O0",) else 3) opt_t[flags[0]] = t rows.append([f"分块 T = {T}", " ".join(flags), f"{t * 1e3:.1f} ms", f"{gflops(N_RUN, t):.2f}", "✓" if ok else "✗"]) BT = np.ascontiguousarray(B.T) dot = LoopNest([("i", N_RUN), ("j", N_RUN), ("k", N_RUN)], ["C[i,j] += A[i,k] * BT[j,k]"], {"A": (N_RUN, N_RUN), "BT": (N_RUN, N_RUN), "C": (N_RUN, N_RUN)}) dot_res = {} for flags in (("-O2",), ("-O2", "-ffast-math")): f = dot.compile_c(flags=flags) C = np.zeros((N_RUN, N_RUN), np.float32) def setup(): C[...] = 0.0 t = best_time(lambda: f({"A": A, "BT": BT, "C": C}), setup=setup) dot_res[flags] = (t, np.abs(C - ref).max()) rows.append(["点积形式（B 转置）", " ".join(flags), f"{t * 1e3:.1f} ms", f"{gflops(N_RUN, t):.2f}", f"✓（最大误差 {np.abs(C - ref).max():.1e}）"]) print_table(["代码", "编译选项", "时间", "GFLOP/s", "正确"], rows) t2, e2 = dot_res[("-O2",)] t3, e3 = dot_res[("-O2", "-ffast-math")] print_observations([ f"同一份 C 代码，-O2 比 -O0 快 {opt_t['-O0'] / opt_t['-O2']:.0f} 倍：寄存器分配、循环不变量外提、展开、向量化都是后端在做。" "我们的循环 IR 只负责'循环怎么排'，指令层面的优化全交给了 C 编译器。", f"点积形式的两个输入都连续访问，看起来很理想，但 -O2 下它用了 {t2 * 1e3:.1f} ms：最内层是对同一个 C[i,j] 的累加，" "编译器必须保持'从 k = 0 加到 k = N−1'的顺序，只能一个一个地加，还要等上一次加法完成（依赖链）。", f"加上 -ffast-math 后变成 {t3 * 1e3:.1f} ms（{t2 / t3:.1f} 倍）：编译器被允许重新结合加法，开 4 个（或更多）部分和并行累加、" "最后合并——这就是向量化的归约。" + (f"有意思的是最大误差反而从 {e2:.1e} 降到 {e3:.1e}：几个部分和分开累加，相当于部分的'两两求和'（P3 §4.4）。" if e3 < e2 else f"结果的舍入变了（最大误差 {e2:.1e} → {e3:.1e}）。") + "但结果不再和'逐个累加'逐位相同——这就是编译器默认不敢做的原因。", "这正是 P3 §4.2 说的'浮点加法不满足结合律'在编译器里的后果：能不能重排求和顺序，是语义问题，不是性能问题。" "GPU 上的归约、Tensor Core 的累加顺序，都在做同样的取舍。", ]) # --------------------------------------------------------------------------- # 演示 4：循环融合 # --------------------------------------------------------------------------- def demo4_fusion(): """ 生产者 B[i,j] = A[i,j] * 2 和消费者 C[i,j] = B[i,j] + A[i,j]（03 脚本已经证明融合合法：依赖距离 0）。 不融合：两个循环，B 写回内存后再读出来；融合：一个循环，B 刚算出来就被用掉（还在寄存器 / L1 里）。 模拟：全相联 LRU 的缺失数（矩阵比 cache 大得多）；实测：2048×2048（每个矩阵 16 MB）。 """ n_sim, n_run = (64, 1024) if QUICK else (128, 2048) print_header(f"演示 4：循环融合——生产者 / 消费者（模拟 {n_sim}²、4 KiB；实测 {n_run}²）") def nests(n): shapes = {"A": (n, n), "B": (n, n), "C": (n, n)} p = LoopNest([("i", n), ("j", n)], ["B[i,j] = A[i,j] * 2"], shapes) c = LoopNest([("i", n), ("j", n)], ["C[i,j] = B[i,j] + A[i,j]"], shapes) return p, c, fuse(p, c) p, c, fz = nests(n_sim) lay = default_layout(p.shapes) m_sep = lru_misses(np.concatenate([p.address_trace(lay), c.address_trace(lay)]), SIM_CACHE) m_fz = lru_misses(fz.address_trace(lay), SIM_CACHE) rows = [["不融合（两个循环）", f"{m_sep:,}", "读 A、写 B；读 B、读 A、写 C = 5 遍"], ["融合（一个循环）", f"{m_fz:,}", "读 A、写 B、写 C = 3 遍（B 的读、A 的第二次读都命中）"]] heads = ["方式", "模拟缺失数", "按'整矩阵遍数'数的流量"] times = {} if HAVE_CC: p, c, fz = nests(n_run) rng = np.random.default_rng(0) arrays = {"A": rng.standard_normal((n_run, n_run)).astype(np.float32), "B": np.zeros((n_run, n_run), np.float32), "C": np.zeros((n_run, n_run), np.float32)} fp, fc, ff = p.compile_c(), c.compile_c(), fz.compile_c() times["sep"] = best_time(lambda: (fp(arrays), fc(arrays))) ref_c = arrays["C"].copy() times["fused"] = best_time(lambda: ff(arrays)) same = np.array_equal(ref_c, arrays["C"]) heads += ["实测时间"] rows[0].append(f"{times['sep'] * 1e3:.1f} ms") rows[1].append(f"{times['fused'] * 1e3:.1f} ms（结果{'相同' if same else '不同'}）") print_table(heads, rows) lines = [f"模拟的缺失数从 {m_sep:,} 降到 {m_fz:,}（约 {m_sep / m_fz:.2f} 倍），和'5 遍 → 3 遍'的流量账一致：" "融合后，消费者要的 B 和 A 刚被生产者碰过，还在 cache 里。"] if times: lines.append(f"实测从 {times['sep'] * 1e3:.1f} ms 降到 {times['fused'] * 1e3:.1f} ms（{times['sep'] / times['fused']:.2f} 倍）。" "如果 B 只是一个中间结果、后面没人用，融合后甚至可以不写 B（只剩'读 A、写 C' 2 遍）——" "这就是 02 脚本里中间结果'留在寄存器'的情形。") lines.append("融合和分块是同一个思想的两面：都是让数据在离计算单元最近的存储里被尽量多地使用，再被扔掉。") print_observations(lines) if __name__ == "__main__": untiled_times = demo1_interchange() tile_times = demo2_tiling(untiled_times) demo3_compiler_flags(tile_times) demo4_fusion()