""" ============================================================================= CS349E · 第 13 讲 · 05_sparse_compilation.py —— 稀疏张量编译：格式决定循环（TACO 的思想） ============================================================================= 【这个脚本要回答的问题】 稠密的循环嵌套可以随便分块、交换；稀疏张量只存非零元，循环必须"跟着数据结构走"。 TACO（Kjolstad et al. 2017）的想法：每个张量的每一维各选一种存储格式（稠密 / 压缩）， 编译器根据表达式和各维的格式**生成**遍历代码——加法要并集合并，乘法要交集合并。 演示 1：同一个矩阵的三种存储：稠密、CSR（第 0 维稠密、第 1 维压缩）、DCSR（两维都压缩）。 演示 2：为 y = A·x（SpMV）生成三种格式的代码；工作量各是多少；非常稀疏（很多空行）时 DCSR 的好处。 演示 3：逐元素的加法和乘法：并集合并 vs 交集合并（合并格的直觉），生成的代码和验证。 演示 4：在你的 CPU 上实测：CSR 的 SpMV vs 稠密的矩阵乘向量，密度多低时稀疏才划算；2:4 结构化稀疏为什么在普通代码里不快。 【关于数字】演示 4 是墙钟时间（单线程 C，-O3），随机器变化；观察按当次运行的数字生成。其余每次相同。 【运行】python3 05_sparse_compilation.py （约 5 秒） 【依赖】numpy；本目录的 sched_lab；../prereq4 的 cbuild（演示 4 需要 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 SPMV_C, gen_elementwise, gen_spmv, load, to_csr, to_dcsr # noqa: E402 from cbuild import best_time, compile_c, find_cc, ptr # noqa: E402 from console import fmt_bytes, fmt_time, print_header, print_observations, print_table # noqa: E402 import ctypes # noqa: E402 INT_P = ctypes.POINTER(ctypes.c_int) def iptr(a): return a.ctypes.data_as(INT_P) def demo1_formats(): print_header("演示 1：同一个矩阵的三种存储") A = np.array([[5, 0, 0, 1, 0, 0], [0, 0, 0, 0, 0, 0], [0, 2, 0, 0, 0, 3], [0, 0, 0, 0, 4, 0]], dtype=np.float32) print("A =\n" + "\n".join(" " + " ".join(f"{v:g}" for v in r) for r in A)) c = to_csr(A) d = to_dcsr(A) print(f"\nCSR（第 0 维稠密、第 1 维压缩）：pos = {c['pos'].tolist()}，crd = {c['crd'].tolist()}，vals = {c['vals'].tolist()}") print(f"DCSR（两维都压缩）：pos0 = {d['pos0'].tolist()}，crd0 = {d['crd0'].tolist()}，pos1 = {d['pos1'].tolist()}，" f"crd1 = {d['crd1'].tolist()}，vals = {d['vals'].tolist()}") n = 4096 rows = [] for dens in (0.5, 0.1, 0.01, 0.001): nnz = int(dens * n * n) rows.append([f"{dens:.1%}", fmt_bytes(4 * n * n), fmt_bytes(4 * (n + 1) + 8 * nnz), f"{(4 * n * n) / (4 * (n + 1) + 8 * nnz):.1f}"]) print_table([f"密度（{n}×{n}，FP32）", "稠密", "CSR（4 字节列号 + 4 字节值）", "稠密 / CSR"], rows) print_observations([ "压缩的一维用两个数组描述：pos[i] 到 pos[i+1] 是第 i 行的非零元在 crd / vals 里的范围，crd 是它们的列号。第 1 行没有非零元，pos[1] = pos[2]。" "DCSR 连行也压缩：只记非空的行（crd0 = [0, 2, 3]），适合绝大多数行都是空的\"超稀疏\"矩阵。", "TACO 把这些格式统一成\"每一维选一种格式\"：dense、compressed（还有 singleton、hashed 等）。CSR = (dense, compressed)，" "DCSR = (compressed, compressed)，CSC 是先列后行的 (dense, compressed)，COO 是 (compressed-nonunique, singleton)。", "存储：每个非零元要 8 字节（值 + 列号），比稠密的 4 字节多一倍——密度 50% 时 CSR 和稠密一样大，密度要远低于一半，压缩才划算。", ]) def demo2_spmv(): print_header("演示 2：y = A·x 的代码由 A 的格式决定") for fmt in ("csr", "dcsr"): print(f"--- 为 {fmt.upper()} 生成的代码 ---") print(gen_spmv(fmt)) rng = np.random.default_rng(0) rows = [] for name, n, m, p_row, dens in (("均匀稀疏 200×300，密度 5%", 200, 300, 1.0, 0.05), ("超稀疏：只有 2% 的行非空，非空行里密度 5%", 2000, 300, 0.02, 0.05)): A = ((rng.random((n, m)) < dens) * (rng.random((n, 1)) < p_row)) * rng.standard_normal((n, m)) x = rng.standard_normal(m) nnz = int(np.count_nonzero(A)) nonempty = int(np.count_nonzero(A.any(axis=1))) errs = [] for fmt, conv in (("dense", lambda a: a), ("csr", to_csr), ("dcsr", to_dcsr)): y = np.zeros(n) load(gen_spmv(fmt), "spmv")(conv(A), x, y) errs.append(np.abs(y - A @ x).max()) rows.append([name, nnz, f"{n * m:,}", f"{nnz + n:,}", f"{nnz + nonempty:,}", f"{max(errs):.1e}"]) print_table(["矩阵", "非零元", "稠密：内层迭代", "CSR：非零 + 行数", "DCSR：非零 + 非空行数", "三种代码的最大误差"], rows) print_observations([ "稠密的代码遍历所有 n × m 个位置；CSR 的代码外层遍历所有行（稠密的那一维），内层只遍历这一行的非零元——工作量 = 非零元个数 + 行数。", "超稀疏时，CSR 光是\"遍历所有行、发现是空的\"就占了大半；DCSR 外层也只遍历非空的行。同一个表达式，换一种格式，编译器生成另一种循环——" "这就是 TACO 的核心：代码不是为每种格式手写的，而是由\"表达式 + 每一维的格式\"推导出来的。", "x 是稠密向量，所以内层可以按 crd[p] 直接取 x[j]（随机访问）。如果 x 也是稀疏的，内层就要同时走 A 的一行和 x 的非零元——那就是演示 3 的合并。", ]) def demo3_merge(): print_header("演示 3：逐元素运算的合并——加法要并集，乘法要交集") print("--- C = A + B（两个 CSR）生成的代码 ---") print(gen_elementwise("+")) rng = np.random.default_rng(1) n, m = 300, 400 rows = [] for da, db in ((0.05, 0.05), (0.05, 0.5), (0.3, 0.3)): A = (rng.random((n, m)) < da) * rng.standard_normal((n, m)) B = (rng.random((n, m)) < db) * rng.standard_normal((n, m)) res = [] for op, ref in (("+", A + B), ("*", A * B)): pos, crd, vals = load(gen_elementwise(op), "ew")(to_csr(A), to_csr(B)) C = np.zeros_like(A) for i in range(n): for p in range(pos[i], pos[i + 1]): C[i, crd[p]] = vals[p] res.append((len(crd), np.abs(C - ref).max())) rows.append([f"{da:.0%} / {db:.0%}", int(np.count_nonzero(A)), int(np.count_nonzero(B)), res[0][0], res[1][0], f"{max(res[0][1], res[1][1]):.1e}"]) print_table(["A / B 的密度", "A 的非零", "B 的非零", "A + B 的非零", "A ∘ B 的非零", "最大误差"], rows) print_observations([ "两个有序的列号序列同时往前走：两边的列号相等时两边一起前进；不等时只前进较小的那一边。加法时\"只有一边有\"的元素也要输出（并集），" "一边走完以后还要把另一边剩下的抄过去；乘法时只有\"两边都有\"才输出（交集），一边走完就可以结束——乘法的代码更短、更快。", "TACO 用\"合并格\"（merge lattice）系统地推导这些分支：表达式里每个 + 是并集、每个 × 是交集，嵌套的表达式组合成一张格，" "每个格点对应生成代码里的一个 while 循环。它还会利用格式做优化：A ∘ B 里如果 B 是稠密的，就只遍历 A 的非零元、直接按下标取 B。", "结果的非零个数：加法约为两者之和减去重叠，乘法只有重叠部分（密度 5% × 5% 时只剩约 0.25%）。结果也要\"装配\"成压缩格式——" "TACO 把计算和装配（算 pos、crd）一起生成。", ]) def demo4_crossover(): print_header("演示 4：本机实测：CSR 的 SpMV vs 稠密的矩阵乘向量（n = 4096，FP32，单线程 C）") if find_cc() is None: print("没有 C 编译器，跳过。") return libs = {} for key, flags in (("plain", ("-O3",)), ("fast", ("-O3", "-ffast-math"))): lib = compile_c(SPMV_C, flags=flags) lib.spmv_csr.argtypes = [INT_P, INT_P, ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float), ctypes.POINTER(ctypes.c_float), ctypes.c_int] libs[key] = lib rng = np.random.default_rng(2) n = 4096 x = rng.standard_normal(n).astype(np.float32) y = np.zeros(n, np.float32) A0 = rng.standard_normal((n, n)).astype(np.float32) t_plain = best_time(lambda: libs["plain"].gemv_dense(ptr(A0), ptr(x), ptr(y), n, n), 5) t_dense = best_time(lambda: libs["fast"].gemv_dense(ptr(A0), ptr(x), ptr(y), n, n), 5) print(f"稠密 GEMV：只开 -O3 时 {fmt_time(t_plain)}；加 -ffast-math 后 {fmt_time(t_dense)}（下表用后者作基准，两种代码都用 -O3 -ffast-math 编译）") rows, cross, last_fast = [], None, None for name, mask in [(f"随机 {d:.1%}", rng.random((n, n)) < d) for d in (0.001, 0.01, 0.05, 0.1, 0.15, 0.2, 0.3, 0.5)] + \ [("2:4 结构化（每 4 个里留 2 个）", np.tile(np.array([1, 1, 0, 0], bool), (n, n // 4)))]: A = A0 * mask c = to_csr(A) pos, crd, vals = c["pos"], c["crd"], c["vals"] t = best_time(lambda: libs["fast"].spmv_csr(iptr(pos), iptr(crd), ptr(vals), ptr(x), ptr(y), n), 5) err = np.abs(y - A @ x).max() / max(1e-30, np.abs(A @ x).max()) rows.append([name, f"{np.count_nonzero(A) / A.size:.1%}", fmt_time(t), f"{t_dense / t:.2f}", f"{err:.1e}"]) if name.startswith("随机"): if t <= t_dense: last_fast = name elif cross is None: cross = name print_table(["稀疏模式", "密度", "CSR SpMV", "稠密 / CSR（> 1 表示稀疏更快）", "相对误差"], rows) print_observations([ "密度很低时 CSR 快得多：它只碰非零元。密度升高，CSR 每个非零元要多读一个列号、按列号随机访问 x，内层循环长度不规则、取 x 是\"gather\"，" f"编译器向量化不了；稠密的代码连续读、能向量化。在这台机器上，稀疏到\"{last_fast}\"还更快，到\"{cross}\"就不划算了。", f"一个陷阱：只开 -O3 时，稠密代码里 acc += A[i][j]·x[j] 的浮点归约不允许重新结合，编译器不向量化，慢了 {t_plain / t_dense:.0f} 倍——" "拿它当基准，会得出\"密度 50% 的稀疏也更快\"的错误结论（这个脚本的第一版就是这样，见 README）。比较两种实现时，基准本身必须是好的（前置讲 P4 §8.3；第 14 讲评测 LLM 写的内核时同样重要）。", "2:4 结构化稀疏（第 3 讲）密度正好 50%：用 CSR 这种通用格式存，比稠密还慢。它的价值在于**硬件**：Ampere 起的稀疏 Tensor Core 用 2 位的元数据" "记录每 4 个里哪 2 个非零，在矩阵乘法单元里直接跳过，理论上快 2 倍——格式、硬件和编译器要一起设计。", "LLM 里的\"稀疏\"大多是结构化的：MoE 的专家路由（按专家分组的稠密矩阵乘法）、分块稀疏注意力（FlashAttention 跳过整块，第 11 讲 §5.3）、" "分页 KV 的块表（第 7 讲）。它们都把不规则性挪到\"块\"的粒度上，块内仍然是稠密、规则的计算。", ]) if __name__ == "__main__": demo1_formats() demo2_spmv() demo3_merge() demo4_crossover()