""" ============================================================================= CS349E · 第 13 讲共享模块：sched_lab.py —— 调度、多面体、自动调优、稀疏编译 ============================================================================= 【这个模块要解决什么问题】 第 12 讲的前端把模型变成了"一组内核"。这一讲为每个内核生成高效的代码。四个部分： A. 迷你 Halide（图像 / 模板计算的流水线）：算法写成 Func，调度决定每个 Func 算在哪里—— root（先整体算好存起来）、inline（用到时现算）、tile（在输出的每个块里算需要的那一小块）、sliding（滑动窗口）。 边界推断（bounds inference）算出每个 Func 要算多大的区域；生成 C 代码，用前置讲 P4 的 cbuild 编译、计时。 B. 多面体模型的二维版：迭代空间 + 依赖距离向量；仿射变换（交换、倾斜）的合法性、可分块性、并行性； 按变换后的顺序真的执行一遍，和原顺序的结果比较；数 cache 缺失。 C. 自动调优：一个 GPU 矩阵乘法内核的配置空间（像 Triton 的 BLOCK_M/N/K、num_warps、num_stages、GROUP_M）， 用一个**解析模型**当"硬件"（明确说明：不是实测）；比较随机搜索、进化搜索、学习代价模型引导的搜索。 D. 稀疏张量编译（TACO 的思想）：每一维的存储格式（稠密 / 压缩）决定生成什么样的循环； 加法要"并集"合并、乘法要"交集"合并；生成代码并验证；C 实测稀疏 vs 稠密的交叉点。 【运行】python3 sched_lab.py # 自检 【依赖】numpy；A、D 的计时需要 C 编译器（../prereq4 的 cbuild）；B 的 02 脚本用 ../prereq4 的 loop_ir.lru_misses 数缺失； C 用 ../lecture10 的 kernel_lab.tile_order（块的启动顺序）。都是只读导入。 """ import itertools import math import os import sys from dataclasses import dataclass import numpy as np HERE = os.path.dirname(os.path.abspath(__file__)) sys.path.insert(0, os.path.join(HERE, "..", "prereq4")) sys.path.insert(0, os.path.join(HERE, "..", "lecture10")) from cbuild import compile_c, find_cc, ptr # noqa: E402 # ============================================================================= # A. 迷你 Halide # ============================================================================= class Expr: """表达式树的基类：+ − × ÷ 用运算符重载拼成 BinOp。""" def _b(self, o, op, rev=False): o = o if isinstance(o, Expr) else Const(float(o)) return BinOp(op, o, self) if rev else BinOp(op, self, o) def __add__(self, o): return self._b(o, "+") def __radd__(self, o): return self._b(o, "+", True) def __sub__(self, o): return self._b(o, "-") def __rsub__(self, o): return self._b(o, "-", True) def __mul__(self, o): return self._b(o, "*") def __rmul__(self, o): return self._b(o, "*", True) def __truediv__(self, o): return self._b(o, "/") class Const(Expr): def __init__(self, v): self.v = v class BinOp(Expr): def __init__(self, op, a, b): self.op, self.a, self.b = op, a, b class Call(Expr): """f(x + dx, y + dy)：调用另一个 Func 在偏移 (dx, dy) 处的值（Halide 里写作 f(x + dx, y + dy)）。""" def __init__(self, f, dx, dy): self.f, self.dx, self.dy = f, dx, dy class Func: """ 一个 Halide 风格的函数：定义在整个二维整数平面上的"图像" f(x, y) = 表达式。 【算法】expr 只说"怎么算一个点"，不说按什么顺序、在哪里算——那是调度的事（算法与调度分离，前置讲 P4 §9.1）。 【输入】expr 为 None 的 Func 是输入图像。 【调用】f(dx, dy) 表示 f 在 (x + dx, y + dy) 处的值。 """ def __init__(self, name, expr=None): self.name, self.expr = name, expr def __call__(self, dx=0, dy=0): return Call(self, dx, dy) def __repr__(self): return self.name def calls(e): """表达式里所有的 Call（按出现顺序）。""" if isinstance(e, Call): return [e] if isinstance(e, BinOp): return calls(e.a) + calls(e.b) return [] def n_ops(e): """表达式里的算术运算个数（每次求值的运算量）。""" return 1 + n_ops(e.a) + n_ops(e.b) if isinstance(e, BinOp) else 0 def bounds(pipeline, region, inline=()): """ 边界推断：给定输出要算的区域 (x0, x1, y0, y1)（半开区间），每个 Func 至少要算哪个区域。 【思路】从输出往回推：消费者的区域按它对生产者的调用偏移扩张（最小偏移到最大偏移），多个消费者取并集。 被内联的 Func 不存储，但它的生产者的区域仍按"内联后的调用"推断——效果和把它当成一层照常推断一样。 【返回】{Func: (x0, x1, y0, y1)} """ req = {pipeline[-1]: region} for f in reversed(pipeline): if f.expr is None or f not in req: continue x0, x1, y0, y1 = req[f] for c in calls(f.expr): r = (x0 + c.dx, x1 + c.dx, y0 + c.dy, y1 + c.dy) if c.f in req: o = req[c.f] r = (min(o[0], r[0]), max(o[1], r[1]), min(o[2], r[2]), max(o[3], r[3])) req[c.f] = r return req def _area(r): return (r[1] - r[0]) * (r[3] - r[2]) def evaluations(pipeline, schedule, W, H, tile=(64, 64)): """ 每个 Func 的表达式一共被求值多少次、中间结果要存多少个数。 【调度】schedule：{Func 名: "root" | "inline" | "tile" | "sliding"}（输出和输入不需要写） root ：按整个输出所需的区域算一遍，存起来 inline ：不存，每个消费者每次用到时现算——被调用几次就算几次 tile ：输出按 tile 分块，每块里算这一块所需的区域（块之间的重叠部分被重复计算） sliding：（只用于"输入 → 生产者 → 输出"的两级流水线）按行滑动，只保留最近 2r + 1 行，不重复计算 【返回】(evals {Func: 次数}, storage {Func: 同时存储的元素数}, 总运算量) """ out = pipeline[-1] evals = {out: W * H} storage = {} full = bounds(pipeline, (0, W, 0, H)) tx, ty = tile ntiles = (W // tx) * (H // ty) tb = bounds(pipeline, (0, tx, 0, ty)) for f in reversed(pipeline[:-1]): if f.expr is None: continue s = schedule[f.name] if s == "root": evals[f] = _area(full[f]) storage[f] = _area(full[f]) elif s == "tile": evals[f] = ntiles * _area(tb[f]) storage[f] = _area(tb[f]) elif s == "sliding": evals[f] = _area(full[f]) ys = [c.dy for c in calls(out.expr) if c.f is f] storage[f] = (max(ys) - min(ys) + 1) * (full[f][1] - full[f][0]) elif s == "inline": n = 0 for g in pipeline: if g.expr is not None and g in evals: n += evals[g] * sum(1 for c in calls(g.expr) if c.f is f) evals[f] = n storage[f] = 0 else: raise ValueError(s) ops = sum(evals[f] * n_ops(f.expr) for f in evals) return evals, storage, ops def numpy_reference(pipeline, img, R, W, H): """用 numpy 逐个 Func 按"root"的方式算出输出（作参照）。img 是带 R 圈边的输入 [(H + 2R), (W + 2R)]。""" full = bounds(pipeline, (0, W, 0, H)) val = {} inp = pipeline[0] def get(f, x0, x1, y0, y1): if f is inp: return img[y0 + R:y1 + R, x0 + R:x1 + R] r = full[f] return val[f][y0 - r[2]:y1 - r[2], x0 - r[0]:x1 - r[0]] def ev(e, r): if isinstance(e, Const): return e.v if isinstance(e, Call): return get(e.f, r[0] + e.dx, r[1] + e.dx, r[2] + e.dy, r[3] + e.dy) a, b = ev(e.a, r), ev(e.b, r) return {"+": a + b, "-": a - b, "*": a * b, "/": a / b}[e.op] for f in pipeline[1:]: val[f] = np.asarray(ev(f.expr, full[f]), dtype=np.float32) return val[pipeline[-1]] def gen_c(pipeline, schedule, R, tile=(64, 64)): """ 生成整个流水线的 C 代码：void pipe(const float* in, float* out, float* scratch, int W, int H)。 【输入布局】in 是 (H + 2R) × (W + 2R) 的图像（四周各多 R 个像素的边），坐标 (x, y) 存在 in[(y + R)·INW + x + R]。 【存储】root 的 Func 存在 scratch 里按整个区域分配的缓冲；tile 的 Func 每块复用一小块缓冲； sliding 的 Func 用一个 2r + 1 行的环形缓冲；inline 的 Func 被代入表达式。 【返回】C 源码字符串；scratch 至少要多大（float 个数）由 scratch_size 给出。 【边界】W、H 要是 tile 的整数倍；sliding 只支持"输入 → 生产者 → 输出"。 """ out = pipeline[-1] inp = pipeline[0] full_rel = {} tb = bounds(pipeline, (0, tile[0], 0, tile[1])) lines = ["void pipe(const float* restrict in, float* restrict out, float* restrict scratch, int W, int H) {", f" const int INW = W + {2 * R};"] # root 缓冲的偏移 off = 0 buf = {} roots = [f for f in pipeline[1:-1] if schedule[f.name] == "root"] tiles = [f for f in pipeline[1:-1] if schedule[f.name] == "tile"] slides = [f for f in pipeline[1:-1] if schedule[f.name] == "sliding"] # root 的区域依赖 W、H：用 bounds 在 (0, W, 0, H) 上的"扩张量"表示 ext = bounds(pipeline, (0, 0, 0, 0)) # 输出为空区域时各 Func 区域的扩张：(dx0, dx1, dy0, dy1) for f in roots: x0, x1, y0, y1 = ext[f] full_rel[f] = (x0, x1, y0, y1) lines.append(f" float* {f.name} = scratch + {off} + 0; /* 区域 [{x0}, W{x1:+d}) × [{y0}, H{y1:+d}) */") buf[f] = ("root", x0, y0, f"(W{x1 - x0:+d})") off_expr = f"(W{x1 - x0:+d})*(H{y1 - y0:+d})" lines.append(f" scratch += {off_expr};") for f in tiles: r = tb[f] lines.append(f" float {f.name}_t[{_area(r)}]; /* 每块：[{r[0]}, {r[1]}) × [{r[2]}, {r[3]})（相对块的左上角） */") buf[f] = ("tile", r[0], r[2], str(r[1] - r[0])) for f in slides: ys = [c.dy for c in calls(out.expr) if c.f is f] rr = max(ys) - min(ys) + 1 x0, x1, _, _ = ext[f] lines.append(f" float* {f.name} = scratch; /* 环形缓冲：{rr} 行 × (W{x1 - x0:+d}) */") buf[f] = ("slide", x0, min(ys), f"(W{x1 - x0:+d})", rr) def ex(e, xs, ys): if isinstance(e, Const): return f"{e.v:.9g}f" if isinstance(e, BinOp): return f"({ex(e.a, xs, ys)} {e.op} {ex(e.b, xs, ys)})" f, xx, yy = e.f, f"{xs}{e.dx:+d}", f"{ys}{e.dy:+d}" if f is inp: return f"in[({yy} + {R})*INW + ({xx} + {R})]" s = schedule.get(f.name) if s == "inline": return "(" + ex(f.expr, f"({xx})", f"({yy})") + ")" kind = buf[f][0] if kind == "root": _, x0, y0, w = buf[f] return f"{f.name}[(({yy}) - ({y0}))*{w} + (({xx}) - ({x0}))]" if kind == "tile": _, x0, y0, w = buf[f] return f"{f.name}_t[(({yy}) - ty - ({y0}))*{w} + (({xx}) - tx - ({x0}))]" _, x0, y0, w, rr = buf[f] return f"{f.name}[((((({yy})) % {rr}) + {rr}) % {rr})*{w} + (({xx}) - ({x0}))]" # root：按流水线顺序整体算 for f in roots: x0, x1, y0, y1 = ext[f] lines += [f" for (int y = {y0}; y < H{y1:+d}; y++)", f" for (int x = {x0}; x < W{x1:+d}; x++)", f" {f.name}[(y - ({y0}))*(W{x1 - x0:+d}) + (x - ({x0}))] = {ex(f.expr, 'x', 'y')};"] if slides: f = slides[0] _, x0, ymin, w, rr = buf[f] x0e, x1e, _, _ = ext[f] ymax = ymin + rr - 1 lines += [f" for (int yy = {ymin}; yy < H{ymax:+d}; yy++) {{", f" for (int x = {x0e}; x < W{x1e:+d}; x++)", f" {f.name}[(((yy % {rr}) + {rr}) % {rr})*{w} + (x - ({x0}))] = {ex(f.expr, 'x', 'yy')};", f" int y = yy - ({ymax});", " if (y >= 0)", " for (int x = 0; x < W; x++)", f" out[y*W + x] = {ex(out.expr, 'x', 'y')};", " }"] elif tiles: lines += [f" for (int ty = 0; ty < H; ty += {tile[1]})", f" for (int tx = 0; tx < W; tx += {tile[0]}) {{"] for f in tiles: r = tb[f] lines += [f" for (int y = ty + ({r[2]}); y < ty + ({r[3]}); y++)", f" for (int x = tx + ({r[0]}); x < tx + ({r[1]}); x++)", f" {f.name}_t[(y - ty - ({r[2]}))*{r[1] - r[0]} + (x - tx - ({r[0]}))] = {ex(f.expr, 'x', 'y')};"] lines += [f" for (int y = ty; y < ty + {tile[1]}; y++)", f" for (int x = tx; x < tx + {tile[0]}; x++)", f" out[y*W + x] = {ex(out.expr, 'x', 'y')};", " }"] else: lines += [" for (int y = 0; y < H; y++)", " for (int x = 0; x < W; x++)", f" out[y*W + x] = {ex(out.expr, 'x', 'y')};"] lines.append("}") return "\n".join(lines) def scratch_size(pipeline, schedule, W, H): """scratch 缓冲要多少个 float（root 的区域之和 + sliding 的环形缓冲）。""" ext = bounds(pipeline, (0, 0, 0, 0)) out = pipeline[-1] n = 0 for f in pipeline[1:-1]: x0, x1, y0, y1 = ext[f] if schedule[f.name] == "root": n += (W + x1 - x0) * (H + y1 - y0) elif schedule[f.name] == "sliding": ys = [c.dy for c in calls(out.expr) if c.f is f] n += (max(ys) - min(ys) + 1) * (W + x1 - x0) return max(n, 1) def run_pipeline_c(pipeline, schedule, img, R, W, H, tile=(64, 64), flags=("-O3",)): """编译并运行一个调度，返回 (输出, 可重复调用的函数)。没有 C 编译器时抛出 RuntimeError。""" lib = compile_c(gen_c(pipeline, schedule, R, tile), flags=flags) out = np.zeros((H, W), np.float32) scratch = np.zeros(scratch_size(pipeline, schedule, W, H), np.float32) img = np.ascontiguousarray(img, dtype=np.float32) def call(): lib.pipe(ptr(img), ptr(out), ptr(scratch), W, H) call() return out, call def blur_pipeline(): """Halide 论文的经典例子：先水平 3 点平均，再竖直 3 点平均。返回 (流水线, 边宽 R)。""" inp = Func("in") bx = Func("blurx", (inp(-1, 0) + inp(0, 0) + inp(1, 0)) * (1.0 / 3)) by = Func("blury", (bx(0, -1) + bx(0, 0) + bx(0, 1)) * (1.0 / 3)) return [inp, bx, by], 1 def box_chain(L): """L 级 3×3 方框模糊串起来（每级 9 个点的平均）。返回 (流水线, 边宽 R = L)。""" fs = [Func("in")] for k in range(L): p = fs[-1] e = None for dy in (-1, 0, 1): for dx in (-1, 0, 1): e = p(dx, dy) if e is None else e + p(dx, dy) fs.append(Func(f"s{k + 1}", e * (1.0 / 9))) return fs, L # ============================================================================= # B. 多面体模型（二维） # ============================================================================= def lex_positive(v): """字典序为正：第一个非零分量 > 0。""" for c in v: if c != 0: return c > 0 return False def classify(T, deps): """ 一个 2×2 的整数变换 T（新的迭代向量 = T · 旧的）对一组依赖距离向量 deps 是否： legal ：每个依赖变换后仍字典序为正（先产生、后使用的顺序不变） tileable ：每个依赖变换后的每个分量都 ≥ 0（"完全可交换"的循环带，矩形分块合法） parallel_inner：每个依赖变换后的第一个分量 > 0（外层循环带走了所有依赖，内层循环可以并行） 【返回】字典，另含 dist（变换后的依赖向量列表） """ T = np.asarray(T) dist = [tuple(int(v) for v in T @ np.asarray(d)) for d in deps] return {"legal": all(lex_positive(d) for d in dist), "tileable": all(min(d) >= 0 for d in dist) and all(lex_positive(d) for d in dist), "parallel_inner": all(d[0] > 0 for d in dist), "dist": dist} def search_transforms(deps, max_coef=2): """ 枚举所有元素在 [−max_coef, max_coef]、行列式为 ±1（幺模，unimodular：变换后的整数点和原来一一对应）的 2×2 矩阵， 返回 (T, 分类) 的列表，可分块的按"变换后依赖分量之和"从小到大排序（Pluto 的目标：依赖距离越短，局部性越好）。 """ res = [] rng = range(-max_coef, max_coef + 1) for a, b, c, d in itertools.product(rng, rng, rng, rng): if abs(a * d - b * c) != 1: continue T = ((a, b), (c, d)) cl = classify(T, deps) if cl["legal"]: res.append((T, cl)) res.sort(key=lambda r: (not r[1]["tileable"], sum(sum(abs(x) for x in d) for d in r[1]["dist"]))) return res def order_points(points, T=((1, 0), (0, 1)), tile=None): """ 按变换 T 之后的字典序排列迭代点；tile = B 时再做 B×B 的矩形分块（先按块号、再按块内坐标）。 【返回】排好序的原迭代点列表 """ T = np.asarray(T) P = np.asarray(points) Q = P @ T.T if tile is None: keys = (Q[:, 1], Q[:, 0]) else: keys = (Q[:, 1], Q[:, 0], Q[:, 1] // tile, Q[:, 0] // tile) idx = np.lexsort(keys) return [tuple(p) for p in P[idx]] def jacobi_run(order, N, steps, init): """ 一维 Jacobi 迭代：A[t+1][i] = (A[t][i−1] + A[t][i] + A[t][i+1]) / 3，迭代点 (t, i)，t ∈ [0, steps)，i ∈ [1, N]。 A 存成完整的 [steps + 1, N + 2] 数组（每个时间步一行，不复用存储），所以只有"流依赖"：(1, −1)、(1, 0)、(1, 1)。 按给定的顺序执行；读到"还没算出来"的值（NaN）说明顺序违反了依赖。 【返回】(最后一行, 读到未计算值的次数, 按执行顺序的访存地址) """ A = np.full((steps + 1, N + 2), np.nan) A[0] = init A[:, 0] = init[0] A[:, -1] = init[-1] bad = 0 addrs = [] W = N + 2 for t, i in order: v = (A[t, i - 1] + A[t, i] + A[t, i + 1]) / 3 if np.isnan(v): bad += 1 v = 0.0 A[t + 1, i] = v addrs += [(t * W + i - 1) * 8, (t * W + i) * 8, (t * W + i + 1) * 8, ((t + 1) * W + i) * 8] return A[steps], bad, addrs def seidel_run(order, N, init): """ 二维 Gauss–Seidel（原地更新）：A[i][j] = (A[i−1][j] + A[i][j−1] + A[i+1][j] + A[i][j+1]) / 4，i, j ∈ [1, N]。 原地更新带来流依赖 (1, 0)、(0, 1) 和反依赖 (1, 0)、(0, 1)（读 A[i+1][j] 必须在它被覆盖之前）。 【返回】最终的数组 """ A = init.copy() for i, j in order: A[i, j] = (A[i - 1, j] + A[i, j - 1] + A[i + 1, j] + A[i, j + 1]) / 4 return A def wavefront_parallelism(points, T): """按 T 变换后，第一个分量（外层循环）每个值下有多少个迭代点——内层循环能并行的宽度。返回 (外层步数, 平均宽度, 最大宽度)。""" Q = np.asarray(points) @ np.asarray(T).T _, counts = np.unique(Q[:, 0], return_counts=True) return len(counts), float(counts.mean()), int(counts.max()) # ============================================================================= # C. 自动调优：GPU 矩阵乘法的配置空间与一个解析的"硬件"模型 # ============================================================================= @dataclass(frozen=True) class GemmConfig: """Triton 风格的矩阵乘法配置：块大小 BM × BN × BK、每块的 warp 数、流水线级数、L2 分组。""" bm: int bn: int bk: int warps: int stages: int group: int def as_tuple(self): return (self.bm, self.bn, self.bk, self.warps, self.stages, self.group) SPACE = {"bm": (16, 32, 64, 128, 256), "bn": (16, 32, 64, 128, 256), "bk": (16, 32, 64, 128), "warps": (1, 2, 4, 8), "stages": (1, 2, 3, 4, 5), "group": (1, 4, 8)} def all_configs(): """配置空间的全部点（5 × 5 × 4 × 4 × 5 × 3 = 6000 个）。""" keys = list(SPACE) return [GemmConfig(*v) for v in itertools.product(*(SPACE[k] for k in keys))] @dataclass(frozen=True) class TCHardware: """ "模拟的硬件"：H100 式的 SM 参数（矩阵乘法用 FP16 Tensor Core、FP32 累加）。 【来源】SM 数、Tensor Core 的乘加 / 周期、共享内存、寄存器、显存带宽：第 8、10 讲； l2_bytes_per_clk（每 SM 每周期从 L2 拿多少字节）、latency（一次 L2 / 显存读的延迟）是**假设**， launch 是每次内核的固定开销（第 11 讲）。整个模型只是一个"有合理结构的代价函数"，不是 H100 的真实性能。 """ n_sm: int = 132 clock: float = 1.83e9 fma_per_clk: int = 2048 smem_bytes_per_clk: int = 128 l2_bytes_per_clk: float = 48.0 hbm: float = 3.35e12 smem_limit: int = 227 * 1024 regs: int = 65536 max_warps: int = 64 latency: float = 800e-9 launch: float = 4e-6 H100_TC = TCHardware() def dram_bytes_rect(M, N, K, bm, bn, concurrent, group=None, elem=2): """ 显存流量的波次模型（第 10 讲 kernel_lab.dram_bytes_tiled 的长方形版本）：同一波里同时在跑的块按 K 齐头并进， 同一时刻用到的 A、B 块在 L2 里只从显存读一次；每波读它碰到的 A 块行和 B 块列各 K 个数，最后写 C 一遍。 """ from kernel_lab import tile_order tm, tn = math.ceil(M / bm), math.ceil(N / bn) pm, pn = tile_order(tm, tn, group) W = max(1, concurrent) byts = 0 for s0 in range(0, tm * tn, W): rows = len(np.unique(pm[s0:s0 + W])) cols = len(np.unique(pn[s0:s0 + W])) byts += (min(rows * bm, M) + min(cols * bn, N)) * K * elem return byts + M * N * elem def gemm_time(cfg, M, N, K, hw=H100_TC, noise=0.0, rng=None): """ "测量"一个配置在 M×N×K 的 FP16 矩阵乘法上的时间（秒）；不能启动的配置返回 inf。 【模型】 1. 资源：共享内存 = stages × (BM + BN) × BK × 2 字节；warp 排成 wy × wx（4 → 2×2，8 → 4×2）， 每个 warp 算 (BM/wy) × (BN/wx)，至少 16 × 8（一条 mma）；每线程寄存器 ≈ 累加器 + 40，超过 255 不能启动。 每个 SM 能同时放几个块 = 共享内存、寄存器、warp 数、32 个块的限制取最小。 2. 每个 K 步（BK）每个 SM 上：Tensor Core 时间 = 乘加 / 2048；共享内存时间 = 各 warp 读操作数的字节 / 128 （第 10 讲 §8.3 的复用账）；L2 时间 = 搬进来的块 / 每 SM 的 L2 带宽。三者取最大。 warp 太少时 Tensor Core 用不满（每个 warp 最多用一个子分区的 1/4）。 3. 延迟：每个块同时最多 stages − 1 个 K 步在路上；Little 定律——一个 K 步的数据要 latency 才到， 如果"在途的数据量 / latency"跟不上消耗，按跟得上的速度算。stages = 1 时完全不重叠。 4. 波次：块数 / (SM 数 × 每 SM 的块数) 向上取整；最后一波没占满的部分也要花一整波的时间。 5. 显存：第 10 讲的波次模型（分组顺序 GROUP_M 决定 L2 复用），作为下限。 6. 补齐：M、N 不是块的整数倍时，补出来的部分也要算。 noise > 0 时乘一个对数正态的测量噪声（模拟真实测量的抖动）。 """ bm, bn, bk, w, st, g = cfg.as_tuple() smem = st * (bm + bn) * bk * 2 wy, wx = {1: (1, 1), 2: (1, 2), 4: (2, 2), 8: (4, 2)}[w] wm, wn = bm // wy, bn // wx if smem > hw.smem_limit or wm < 16 or wn < 8 or bm % wy or bn % wx: return math.inf regs = wm * wn // 32 + 40 if regs > 255: return math.inf threads = 32 * w per_sm = min(32, hw.max_warps // w, hw.smem_limit // smem if smem else 32, hw.regs // (regs * threads)) if per_sm < 1: return math.inf tm, tn = math.ceil(M / bm), math.ceil(N / bn) tiles = tm * tn active = min(per_sm, max(1, math.ceil(tiles / hw.n_sm))) ksteps = math.ceil(K / bk) # 每个 K 步、每个 SM（active 个块一起） fma = active * bm * bn * bk tc_rate = min(hw.fma_per_clk, active * min(w, 4) * hw.fma_per_clk / 4) t_tc = fma / tc_rate smem_bytes = active * w * (wm + wn) * bk * 2 t_smem = smem_bytes / hw.smem_bytes_per_clk load_bytes = active * (bm + bn) * bk * 2 t_l2 = load_bytes / hw.l2_bytes_per_clk t_step = max(t_tc, t_smem, t_l2) lat = hw.latency * hw.clock if st == 1: t_step = t_step + lat else: t_step = max(t_step, lat / (st - 1)) per_wave = ksteps * t_step + lat + bm * bn * 2 * active / hw.l2_bytes_per_clk waves = math.ceil(tiles / (hw.n_sm * active)) t = waves * per_wave / hw.clock dram = dram_bytes_rect(M, N, K, bm, bn, hw.n_sm * active, None if g == 1 else g) t = max(t, dram / hw.hbm) + hw.launch if noise > 0: t *= float(np.exp((rng or np.random.default_rng()).normal(0, noise))) return t def static_valid(cfg, hw=H100_TC): """ 不用测量就能判断的"能不能启动"：共享内存、每线程寄存器、warp 的块至少一条 mma（编译器 / Triton 在调优前就会剪掉这些配置）。 """ bm, bn, bk, w, st, _ = cfg.as_tuple() wy, wx = {1: (1, 1), 2: (1, 2), 4: (2, 2), 8: (4, 2)}[w] wm, wn = bm // wy, bn // wx return (st * (bm + bn) * bk * 2 <= hw.smem_limit and wm >= 16 and wn >= 8 and bm % wy == 0 and bn % wx == 0 and wm * wn // 32 + 40 <= 255) def features(cfg): """学习代价模型用的特征：各参数的 log2 和它们两两的乘积（二次特征），外加常数项。""" v = np.log2(np.asarray(cfg.as_tuple(), dtype=np.float64)) quad = [v[i] * v[j] for i in range(len(v)) for j in range(i, len(v))] return np.concatenate([[1.0], v, quad]) class RidgeModel: """ 学习的代价模型：在 log(时间) 上做岭回归（二次特征）。训练数据应只含能启动的配置（先用 static_valid 剪掉其余的）； 万一混进了 inf，目标值记为"已见过的最慢 × 2"。 【真实系统】AutoTVM 用 XGBoost、Ansor 用梯度提升树 / 小神经网络，特征从生成的程序里抽（循环长度、访存量……）。 """ def __init__(self, lam=1e-2): self.lam = lam self.w = None def fit(self, cfgs, times): X = np.stack([features(c) for c in cfgs]) finite = [t for t in times if np.isfinite(t)] cap = (max(finite) if finite else 1.0) * 2 y = np.log(np.array([t if np.isfinite(t) else cap for t in times])) A = X.T @ X + self.lam * np.eye(X.shape[1]) self.w = np.linalg.solve(A, X.T @ y) return self def predict(self, cfgs): X = np.stack([features(c) for c in cfgs]) return X @ self.w def mutate(cfg, rng): """随机改一个参数到相邻的取值（进化搜索的"变异"）。""" keys = list(SPACE) k = keys[rng.integers(len(keys))] vals = SPACE[k] i = vals.index(getattr(cfg, k)) j = min(len(vals) - 1, max(0, i + (1 if rng.random() < 0.5 else -1))) d = dict(zip(keys, cfg.as_tuple())) d[k] = vals[j] return GemmConfig(*d.values()) def tune(strategy, measure, budget, rng, pool=None): """ 在预算 budget 次"测量"内找最快的配置。 【策略】 "random" ：均匀随机抽配置（不重复） "evolution" ：先随机测 16 个，之后每次从目前最好的 8 个里挑一个变异 "model" ：先随机测 16 个；之后每轮用岭回归模型给 1000 个候选（随机 + 最好几个的变异）打分，测预测最快的 8 个 【参数】pool：候选配置（默认：全部能静态通过 static_valid 的配置）；measure(cfg) → 测到的时间（可以带噪声） 【返回】历史：每次测量后"按测到的时间，目前最快的配置"的列表（调用方可以用无噪声的时间评价它） """ pool = pool or [c for c in all_configs() if static_valid(c)] pool_set = set(pool) seen, hist, best = {}, [], None def meas(c): nonlocal best if c in seen or c not in pool_set: return seen[c] = measure(c) if best is None or seen[c] < seen[best]: best = c hist.append(best) if strategy == "random": for i in rng.permutation(len(pool))[:budget]: meas(pool[i]) return hist for i in rng.permutation(len(pool))[:16]: meas(pool[i]) while len(hist) < budget: top = sorted(seen, key=seen.get)[:8] if strategy == "evolution": c = mutate(top[rng.integers(len(top))], rng) tries = 0 while (c in seen or c not in pool_set) and tries < 20: c = mutate(c, rng) tries += 1 if c in seen or c not in pool_set: c = pool[rng.integers(len(pool))] meas(c) else: model = RidgeModel().fit(list(seen), list(seen.values())) cands = [pool[i] for i in rng.integers(0, len(pool), 800)] cands += [mutate(top[rng.integers(len(top))], rng) for _ in range(200)] cands = [c for c in dict.fromkeys(cands) if c not in seen and c in pool_set] pred = model.predict(cands) for i in np.argsort(pred)[:8]: if len(hist) >= budget: break meas(cands[i]) return hist[:budget] # ============================================================================= # D. 稀疏张量编译（TACO 的思想） # ============================================================================= def to_csr(A): """稠密矩阵 → CSR（第 0 维稠密、第 1 维压缩）：pos（每行的起止）、crd（列号）、vals。""" pos, crd, vals = [0], [], [] for row in A: nz = np.flatnonzero(row) crd.extend(nz.tolist()) vals.extend(row[nz].tolist()) pos.append(len(crd)) return {"pos": np.array(pos, np.int32), "crd": np.array(crd, np.int32), "vals": np.array(vals, np.float32), "shape": A.shape} def to_dcsr(A): """稠密矩阵 → DCSR（两维都压缩）：只存非空的行（pos0 / crd0），每行里再存非零的列。""" rows = [i for i in range(A.shape[0]) if np.any(A[i])] pos1, crd1, vals = [0], [], [] for i in rows: nz = np.flatnonzero(A[i]) crd1.extend(nz.tolist()) vals.extend(A[i][nz].tolist()) pos1.append(len(crd1)) return {"pos0": np.array([0, len(rows)], np.int32), "crd0": np.array(rows, np.int32), "pos1": np.array(pos1, np.int32), "crd1": np.array(crd1, np.int32), "vals": np.array(vals, np.float32), "shape": A.shape} def gen_spmv(fmt): """ 为 y(i) = A(i, j) · x(j) 生成 Python 代码，A 的格式由 fmt 决定： "dense"：for i: for j: ... "csr" ：for i（稠密）: for p in pos[i]..pos[i+1]（压缩）: j = crd[p] ... "dcsr" ：for p0 in 非空的行: i = crd0[p0]; for p1 in pos1[p0]..pos1[p0+1]: ... 【返回】源代码字符串（函数名 spmv）。TACO 对任意的张量表达式和格式组合都能这样生成，这里只演示三种。 """ if fmt == "dense": body = ["def spmv(A, x, y):", " n, m = A.shape", " for i in range(n):", " acc = 0.0", " for j in range(m):", " acc += A[i, j] * x[j]", " y[i] = acc"] elif fmt == "csr": body = ["def spmv(A, x, y):", " pos, crd, vals = A['pos'], A['crd'], A['vals']", " for i in range(A['shape'][0]): # 第 0 维：稠密，遍历所有行", " acc = 0.0", " for p in range(pos[i], pos[i + 1]): # 第 1 维：压缩，只遍历这一行的非零", " j = crd[p]", " acc += vals[p] * x[j]", " y[i] = acc"] elif fmt == "dcsr": body = ["def spmv(A, x, y):", " crd0, pos1, crd1, vals = A['crd0'], A['pos1'], A['crd1'], A['vals']", " for p0 in range(A['pos0'][0], A['pos0'][1]): # 第 0 维也压缩：只遍历非空的行", " i = crd0[p0]", " acc = 0.0", " for p1 in range(pos1[p0], pos1[p0 + 1]):", " j = crd1[p1]", " acc += vals[p1] * x[j]", " y[i] = acc"] else: raise ValueError(fmt) return "\n".join(body) def gen_elementwise(op): """ 为两个 CSR 矩阵的逐元素运算生成 Python 代码（结果也是 CSR）： op = "+"：并集合并（union）——任何一个有非零，结果就可能非零； op = "*"：交集合并（intersection）——两个都非零，结果才非零。 【合并格（merge lattice）的直觉】同时走两个有序的列号序列： 两边都没走完时：取较小的列号；相等就两边一起前进（+ 和 × 都要处理），不等就只前进较小的那一边（只有 + 要处理）； 一边走完以后：+ 要把另一边剩下的抄过去，× 直接结束。 【返回】源代码字符串（函数名 ew），返回 (pos, crd, vals) """ lines = ["def ew(A, B):", " n = A['shape'][0]", " pos, crd, vals = [0], [], []", " for i in range(n):", " pa, ea = A['pos'][i], A['pos'][i + 1]", " pb, eb = B['pos'][i], B['pos'][i + 1]", " while pa < ea and pb < eb: # 两边都没走完", " ja, jb = A['crd'][pa], B['crd'][pb]", " j = min(ja, jb)", " if ja == j and jb == j:", f" crd.append(j); vals.append(A['vals'][pa] {op} B['vals'][pb])", " pa += 1; pb += 1"] if op == "+": lines += [" elif ja == j:", " crd.append(j); vals.append(A['vals'][pa]); pa += 1", " else:", " crd.append(j); vals.append(B['vals'][pb]); pb += 1", " while pa < ea: # 只剩 A", " crd.append(A['crd'][pa]); vals.append(A['vals'][pa]); pa += 1", " while pb < eb: # 只剩 B", " crd.append(B['crd'][pb]); vals.append(B['vals'][pb]); pb += 1"] else: lines += [" elif ja == j:", " pa += 1 # 只有 A 有：乘积为 0，跳过", " else:", " pb += 1", " # 一边走完，乘积全是 0：不用处理剩下的"] lines += [" pos.append(len(crd))", " return pos, crd, vals"] return "\n".join(lines) def load(src, name): """执行生成的 Python 源码，取出函数。""" ns = {} exec(src, ns) return ns[name] SPMV_C = r""" void spmv_csr(const int* restrict pos, const int* restrict crd, const float* restrict vals, const float* restrict x, float* restrict y, int n) { for (int i = 0; i < n; i++) { float acc = 0.0f; for (int p = pos[i]; p < pos[i + 1]; p++) acc += vals[p] * x[crd[p]]; y[i] = acc; } } void gemv_dense(const float* restrict A, const float* restrict x, float* restrict y, int n, int m) { for (int i = 0; i < n; i++) { float acc = 0.0f; for (int j = 0; j < m; j++) acc += A[(long)i * m + j] * x[j]; y[i] = acc; } } """ # ============================================================================= # 自检 # ============================================================================= def _selftest(): # A：边界推断与运算量 pipe, R = blur_pipeline() inp, bx, _ = pipe b = bounds(pipe, (0, 64, 0, 64)) assert b[bx] == (0, 64, -1, 65) and b[inp] == (-1, 65, -1, 65) ev, _, _ = evaluations(pipe, {"blurx": "inline"}, 64, 64) assert ev[bx] == 3 * 64 * 64 ev2, _, _ = evaluations(pipe, {"blurx": "tile"}, 64, 64, tile=(32, 32)) assert ev2[bx] == 4 * 32 * 34 if find_cc() is not None: rng = np.random.default_rng(0) W = H = 64 img = rng.random((H + 2 * R, W + 2 * R)).astype(np.float32) ref = numpy_reference(pipe, img, R, W, H) for s in ("root", "inline", "tile", "sliding"): out, _ = run_pipeline_c(pipe, {"blurx": s}, img, R, W, H, tile=(16, 16)) assert np.abs(out - ref).max() < 1e-5, s ch, R2 = box_chain(3) img2 = rng.random((H + 2 * R2, W + 2 * R2)).astype(np.float32) ref2 = numpy_reference(ch, img2, R2, W, H) for s in ("root", "inline", "tile"): sched = {f.name: s for f in ch[1:-1]} out, _ = run_pipeline_c(ch, sched, img2, R2, W, H, tile=(16, 16)) assert np.abs(out - ref2).max() < 1e-5, s # B：Jacobi 的倾斜 deps = [(1, -1), (1, 0), (1, 1)] assert classify(((1, 0), (0, 1)), deps)["legal"] and not classify(((1, 0), (0, 1)), deps)["tileable"] assert classify(((1, 0), (1, 1)), deps)["tileable"] assert not classify(((0, 1), (1, 0)), deps)["legal"] N, S = 30, 10 init = np.random.default_rng(1).random(N + 2) pts = [(t, i) for t in range(S) for i in range(1, N + 1)] ref, bad0, _ = jacobi_run(pts, N, S, init) good, bad1, _ = jacobi_run(order_points(pts, ((1, 0), (1, 1)), tile=4), N, S, init) _, bad2, _ = jacobi_run(order_points(pts, ((1, 0), (0, 1)), tile=4), N, S, init) assert bad0 == 0 and bad1 == 0 and np.allclose(good, ref) and bad2 > 0 # C：最优配置有限、模型可拟合 cfgs = all_configs() assert len(cfgs) == 6000 ts = [gemm_time(c, 4096, 4096, 4096) for c in cfgs[:200]] assert any(np.isfinite(ts)) # D：稀疏 rng = np.random.default_rng(2) A = (rng.random((20, 30)) < 0.2) * rng.standard_normal((20, 30)) B = (rng.random((20, 30)) < 0.2) * rng.standard_normal((20, 30)) x = rng.standard_normal(30) for fmt, conv in (("dense", lambda a: a), ("csr", to_csr), ("dcsr", to_dcsr)): y = np.zeros(20) load(gen_spmv(fmt), "spmv")(conv(A), x, y) assert np.allclose(y, A @ x, atol=1e-5) 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(20): for p in range(pos[i], pos[i + 1]): C[i, crd[p]] = vals[p] assert np.allclose(C, ref, atol=1e-5) print("sched_lab 自检通过。") if __name__ == "__main__": print("== sched_lab 自检 ==") _selftest()