""" ============================================================================= CS349E · 前置讲 P4 · 03_loop_nests_and_dependences.py —— 迭代空间与依赖分析：哪些循环变换是合法的？ ============================================================================= 【这个脚本要回答的问题】 循环变换（交换顺序、分块、并行、融合）能让程序快很多（04 脚本），但前提是**不改变结果**。 编译器怎样判断一个变换是否安全？答案是依赖分析：找出所有"读写同一个内存位置"的迭代对， 变换后它们的先后顺序必须不变。本脚本： 演示 1：迭代空间——一个 4×4 的二维循环，按 ij、ji、2×2 分块三种顺序执行时，每个点是第几个被执行的。 演示 2：依赖距离——五个经典例子（矩阵乘法、前缀和、二维递推、斜向依赖、Jacobi 迭代）的依赖。 演示 3：合法性——由依赖距离判断哪些循环顺序合法、哪些循环可以并行；再真的执行一遍，看不合法的变换算错了没有。 演示 4：循环融合的合法性——生产者 / 消费者两个循环，什么时候能合成一个。 【怎么读输出】全部是确定性的计算，每次运行结果相同。依赖用"小规模枚举"求出（loop_ir.dependences）， 真实编译器用整数线性规划 / 多面体模型解析地求（第 13 讲）。 【运行】python3 03_loop_nests_and_dependences.py （约 1 秒） 【依赖】numpy；从 ../lecture01 导入 console；同目录的 loop_ir。 """ 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 loop_ir import LoopNest, dependences, fuse, fusion_distances, matmul, parallel_loops, permutation_legal # noqa: E402 # --------------------------------------------------------------------------- # 演示 1：迭代空间 # --------------------------------------------------------------------------- def order_grid(nest, n): """返回 n×n 的网格：grid[i][j] = 点 (i, j) 是第几个被执行的（从 0 数）。""" P = nest.points() grid = np.zeros((n, n), dtype=int) for t, (i, j) in enumerate(zip(P["i"], P["j"])): grid[i, j] = t return grid def demo1_iteration_space(n=4): """ 二维循环 for i: for j: X[i,j] = …，迭代空间是 n×n 个点。循环顺序决定了"走过这些点的路线"。 """ print_header(f"演示 1：迭代空间——{n}×{n} 的循环，三种执行顺序（格子里的数 = 第几个执行）") base = LoopNest([("i", n), ("j", n)], ["X[i,j] = Y[i,j] + 1"], {"X": (n, n), "Y": (n, n)}) variants = [("ij（行优先）", base.copy()), ("ji（列优先）", base.copy().reorder("j", "i")), ("2×2 分块", base.copy().split("i", 2).split("j", 2).reorder("io", "jo", "ii", "ji"))] grids = [(name, order_grid(v, n), v.describe()) for name, v in variants] for name, g, desc in grids: print(f"{name}：循环 {desc}") for row in g: print(" " + " ".join(f"{x:3d}" for x in row)) print_observations([ "同一组迭代点，三种'路线'。ij 一行一行走；ji 一列一列走；分块先走完左上角的 2×2，再走下一块。", "如果 X、Y 按行存储，ij 顺序访问的地址是连续的（空间局部性好）；ji 每一步跳过一整行（P1 §2.6 的按列遍历）。" "分块则让一小块数据被集中地反复使用（时间局部性），这是 04 脚本矩阵乘法提速的关键。", "循环变换只改变'路线'，不改变'要走哪些点'——所以只要每个点的计算互不依赖，任何路线都给出同样的结果。" "麻烦在于点之间常常有依赖（演示 2）。", ]) # --------------------------------------------------------------------------- # 演示 2：依赖距离 # --------------------------------------------------------------------------- def examples(): """ 五个例子。为了让下标不越界，用"整体平移"的写法（例如前缀和写成 S[i+1] = S[i] + X[i+1]，i 从 0 开始）。 """ n = 5 return [ ("矩阵乘法", matmul(4), "C[i,j] += A[i,k] * B[k,j]"), ("前缀和", LoopNest([("i", n)], ["S[i+1] = S[i] + X[i+1]"], {"S": (n + 1,), "X": (n + 1,)}), "S[i+1] = S[i] + X[i+1]"), ("二维递推", LoopNest([("i", n), ("j", n)], ["A[i+1,j+1] = A[i,j+1] + A[i+1,j]"], {"A": (n + 1, n + 1)}), "A[i+1,j+1] = A[i,j+1] + A[i+1,j]"), ("斜向依赖", LoopNest([("i", n), ("j", n)], ["A[i+1,j] = A[i,j+1] * 0.5"], {"A": (n + 1, n + 1)}), "A[i+1,j] = A[i,j+1] * 0.5"), ("Jacobi（读旧数组写新数组）", LoopNest([("i", n), ("j", n)], ["B[i+1,j+1] = A[i,j+1] + A[i+2,j+1] + A[i+1,j] + A[i+1,j+2]"], {"A": (n + 2, n + 2), "B": (n + 2, n + 2)}), "B[i+1,j+1] = A[i,j+1] + A[i+2,j+1] + A[i+1,j] + A[i+1,j+2]"), ] def demo2_dependences(): """ 对每个例子枚举依赖，列出（数组、类型、距离向量、出现次数）。距离全为 0 的依赖发生在同一次迭代内，不影响循环顺序，单独标出。 【三种依赖】RAW（先写后读，真依赖：后者要用前者算出的值）；WAR（先读后写，反依赖：后者不能提前覆盖前者要读的值）； WAW（先写后写，输出依赖：最后留下的值必须是后写的那个）。 """ print_header("演示 2：依赖距离——谁必须在谁之前执行") rows, all_deps = [], {} for name, nest, text in examples(): deps = dependences(nest) all_deps[name] = (nest, deps) cross = {d: c for d, c in deps.items() if any(d.distance)} if not cross: rows.append([name, text, "—", "—", "（无跨迭代依赖）", ""]) for d, c in sorted(cross.items(), key=lambda kv: (kv[0].kind, kv[0].distance)): rows.append([name, text, d.array, d.kind, str(d.distance), c]) print_table(["例子", "语句", "数组", "类型", "距离（按原始循环顺序）", "次数"], rows) print_observations([ "矩阵乘法只有一种跨迭代依赖：距离 (0, 0, 1)——同一个 C[i,j] 在 k 和 k+1 两次迭代里被累加（先写后读、先写后写）。" "i、j 两个方向完全独立。", "前缀和的距离是 (1)：每一步都要用上一步的结果，这是一条依赖链，不能简单地并行（但可以用'并行前缀和'算法重写，P2 的归约是它的近亲）。", "二维递推有 (1, 0) 和 (0, 1) 两种依赖：每个点依赖上面和左边的点。", "斜向依赖的距离是 (1, −1)：第 i+1 行的第 j 个点用到第 i 行第 j+1 个点（右上方）的值。负的分量是麻烦的来源（演示 3）。", "Jacobi 读旧数组 A、写新数组 B，没有任何跨迭代依赖：所有点都可以同时算——这也是为什么 GPU 上的模板计算" "（以及神经网络里的绝大多数算子）都写成'读输入、写另一个输出'的形式。", ]) return all_deps # --------------------------------------------------------------------------- # 演示 3：合法性 # --------------------------------------------------------------------------- def run(nest, arrays): """复制一份数组、在上面执行 nest（Python 代码生成），返回结果数组字典。""" arr = {k: v.copy() for k, v in arrays.items()} nest.run_python(arr) return arr def demo3_legality(all_deps, seed=0): """ 【规则】一个循环顺序合法 ⇔ 每个依赖的距离向量按新顺序排列后，第一个非零分量仍然 > 0（字典序为正）。 一层循环可以并行 ⇔ 没有依赖由这一层"携带"（按当前顺序，距离向量的第一个非零分量不在这一层）。 【验证】对每个例子、每种循环顺序，真的执行一遍（float64，随机初值），和原始顺序的结果比较。 """ print_header("演示 3：合法性——由依赖判断，再真的执行一遍验证") rng = np.random.default_rng(seed) rows = [] for name, (nest, deps) in all_deps.items(): arrays = {k: rng.standard_normal(shape) for k, shape in nest.shapes.items()} ref = run(nest, arrays) cross = [d for d in deps if any(d.distance)] par = parallel_loops(cross, nest.orig, nest.orig) for order in itertools.permutations(nest.orig): ok, bad = permutation_legal(cross, nest.orig, list(order)) v = nest.copy().reorder(*order) out = run(v, arrays) same = all(np.array_equal(out[k], ref[k]) for k in ref) rows.append([name, "".join(order), "合法" if ok else f"不合法（{bad[0].distance}）", "相同" if same else "**不同**", ", ".join(par) if order == tuple(nest.orig) else ""]) print_table(["例子", "循环顺序", "依赖分析的判断", "实际执行的结果", "原始顺序下可并行的循环"], rows) print_observations([ "依赖分析判为合法的顺序，执行结果都和原始顺序逐位相同；判为不合法的（斜向依赖交换成 ji），结果真的不同了。", "斜向依赖 (1, −1) 交换顺序后变成 (−1, 1)：终点跑到了源头前面，读到的是还没更新的旧值。", "矩阵乘法的 6 种顺序都合法（依赖只在 k 上，而且所有顺序里 k 都是从小到大走的），i、j 两层可以并行——" "GPU 给每个输出元素（或每块输出）一个线程（块），正是在并行 i、j；k 方向的累加则在一个线程里顺序做，" "或者用 split-K 拆开以后再归约（P2 §7）。", "二维递推的两层都携带依赖，都不能直接并行；但同一条反对角线（i + j 相同）上的点互不依赖——" "'倾斜'（skewing）变换能把它们变成可以并行的一层（波前并行），这类变换是多面体编译的内容（第 13 讲）。", ]) # --------------------------------------------------------------------------- # 演示 4：循环融合 # --------------------------------------------------------------------------- def demo4_fusion(n=8, seed=0): """ 生产者循环 B[i] = A[i] * 2 和两种消费者： 消费者 1：C[i] = B[i] + 1 （用同一个位置） 消费者 2：C[i] = B[i+1] + 1 （用后一个位置） 把生产者和消费者融合成一个循环（每次迭代先生产、再消费），合法吗？ 【规则】融合前，所有生产都在所有消费之前。融合后，消费者在第 i 次迭代读 B[i+1]，但 B[i+1] 要到第 i+1 次迭代才写—— 依赖距离（消费的迭代 − 生产的迭代）为 −1 < 0，不合法。 """ print_header("演示 4：循环融合——生产者和消费者能不能合成一个循环？") rng = np.random.default_rng(seed) prod = LoopNest([("i", n)], ["B[i] = A[i] * 2"], {"A": (n + 1,), "B": (n + 1,)}) rows = [] for label, stmt in (("C[i] = B[i] + 1", "C[i] = B[i] + 1"), ("C[i] = B[i+1] + 1", "C[i] = B[i+1] + 1")): cons = LoopNest([("i", n)], [stmt], {"B": (n + 1,), "C": (n + 1,)}) dist = fusion_distances(prod, cons) legal = all(d[2][0] >= 0 for d in dist) arrays = {"A": rng.standard_normal(n + 1), "B": np.zeros(n + 1), "C": np.zeros(n + 1)} sep = run(prod, arrays) sep = run(cons, sep) fused = run(fuse(prod, cons), arrays) same = np.array_equal(sep["C"], fused["C"]) ds = sorted({d[2] for d in dist}) rows.append([label, ", ".join(f"{d[0]}" if len(d) == 1 else str(d) for d in ds), "合法" if legal else "不合法", "相同" if same else "**不同**"]) print_table(["消费者", "生产 → 消费的依赖距离", "融合的判断", "融合后实际执行的结果"], rows) print_observations([ "消费者读同一个位置时，依赖距离是 0：融合后同一次迭代里先生产、后消费，结果不变。B 的值刚算出来就被用掉，" "可以留在寄存器里，根本不用写回内存——这就是 02 脚本'融合省内存流量'的来源。", "消费者读 B[i+1] 时距离是 −1，融合后读到的是还没生产的值，结果错了。修复办法是先把消费者'平移'一步" "（第 i 次迭代计算 C[i−1]），让距离变成 0——编译器里叫循环对齐（alignment）或平移（shifting）。", "神经网络的逐元素算子几乎都是'读同一个位置'（距离 0），所以总能融合；卷积、注意力这类读'邻居'的算子融合起来就要小心，" "FlashAttention 本质上就是一次精心设计的融合（第 11 讲）。", ]) if __name__ == "__main__": demo1_iteration_space() deps_by_example = demo2_dependences() demo3_legality(deps_by_example) demo4_fusion()