TVM调度原语空间枚举

可视化演示

展示TVM三大调度原语：split/fuse循环拆分合并、cache_read多级内存层次数据搬运、以及rfactor归约分解并行化策略。

调度原语总览

想象后厨刚接到一张“全家福套餐”订单：蒸鱼、炒时蔬、炖汤，好几道菜要同时端上桌。主厨拿到菜单后第一件事不是点火，而是画一张工序表——谁先备菜、哪个灶台烧什么、炖汤和炒菜怎么错开用火。这张工序表，在 TVM 里就叫“调度”（schedule）。调度不改变菜谱本身，只改变做菜的方式；对应到编译器里，调度不改变计算结果，只改变算子的循环结构与数据流向。

深度学习模型在交给 GPU 之前，会被编译器拆成一张计算图（computational graph），图里的每个节点是一个算子，比如卷积、矩阵乘法。每个算子最终会写成一层层嵌套的循环去执行。调度要做的事，就是决定这些循环怎么组织、数据怎么流动。TVM 把调度的动作抽象成一批原语（primitive），按性质分成三大类：循环变换类（split、fuse、reorder）、数据缓存类（cache_read、cache_write）、归约分解类（rfactor）。任何一种调度方案，都是这些原语按一定顺序组合的结果。

把所有合法组合收集起来，就得到调度空间（schedule space）。数学上可以这样描述：给定一张计算图 \( G \)，再给一组合法的原语序列 \(\pi\)，每应用一次就得到一个调度方案；全部方案构成集合 \(\mathcal{S}\)。

就像主厨手里的刀、铲、锅，原语本身没有好坏，关键在于怎么组合。认识了工具清单，接下来就一把一把地试。

循环变换原语

炖一锅肉要 40 分钟，但灶台只空出 10 分钟。聪明的做法是把肉切成四小块，分给四个灶同时炖。“切块”这个动作，对应调度里的 split 原语。split 把一层循环 \([0, N)\) 按因子 \(f\) 拆成两层：外层负责分块数量，内层负责块内迭代。

与 split 相对的是 fuse，它把紧邻的两层循环合并成一层，减少循环启动和同步开销；reorder 则交换嵌套循环的先后顺序，让访存更连续。三者配合，循环结构几乎可以“揉圆捏扁”。有了切块、合并、调序这三板斧，主厨就能按硬件的脾气重塑工序——这为后面的数据缓存和并行归约打下了基础。

数据缓存原语

主厨发现每次炒菜都去仓库拿盐太慢了，于是在灶台边放一个调料架，开工前一次性把常用调料摆上。这个调料架就是 cache_read 原语。它在靠近计算单元的地方创建一块局部缓冲区，计算前把要用的数据从主存拷进来，之后计算全部命中这块缓冲区，避免反复访问慢速主存。

副作用也要考虑：缓存数据需要搬运时间，搬运本身不是免费的。cache_write 则负责另一个方向——计算先把结果攒在缓冲区里，攒够再一次写回主存，避免频繁的小写操作。TVM 中缓存分多级，如 global、shared、local，对应后厨的仓库、备菜台、手边调料碟；用哪一级要看数据复用距离和硬件结构。

从性能模型看，一次数据访问的平均时间由“命中缓存”和“未命中”两部分加权得到：

调料架摆好后，炒菜不再满后厨跑，但架子本身也占地方。缓存换来速度，也要付出空间和搬运成本——这正是后面调度搜索要权衡的地方。

归约分解原语

洗 100 筐菜，如果只有一个帮厨，他只能从头洗到尾；主厨叫来十个人，每人洗 10 筐，最后把洗好的菜归拢到一处。“先分工再汇总”，就是 rfactor 的思路。归约循环做的是 sum、max、min 这类“把一组值合成一个值”的运算，天然满足结合律：先算部分和，再算部分和的和，结果不变。

以双重求和为例，交换内外层求和的顺序不会改变最终答案：

帮厨们并行洗菜靠的是“菜怎么洗都还是那些菜”——结合律。rfactor 能安全分解的底气也正在于此。它把一个归约循环拆成两段：多线程先各自算局部归约，再做一次最终合并。

调度组合空间

每道菜有几十种切法、十几种火候，把这些工序排列组合，后厨的可能方案多到数不清。调度也一样：原语种类不多，但参数和次序让组合数量急剧膨胀。每种原语都带参数：split 的因子可以取 2、4、8……；cache_read 可以指定缓存层级；rfactor 可以选分解轴。每一步选择都是在为调度空间“添一个维度”。

如果第 \(i\) 步有 \(C_i\) 种可选方案，总共 \(k\) 步，那么调度空间的大小遵循乘法原理：

更麻烦的是，并非所有组合都合法：拆分必须整除或处理余数、fuse 只对相邻循环有效、rfactor 只对归约循环安全、数据依赖不能被破坏。空间不是简单笛卡尔积，而是一棵带约束的枚举树。组合多到爆炸，后厨不可能把每道工序都试一遍，所以下一步要想办法：怎样聪明地枚举。

枚举策略方法

真正的主厨不会一小时试一种做法。他会先排除明显不行的方案——比如整块肉直接下锅——再对剩下候选做快速估算，只对最有希望的精细尝试。枚举调度空间，就是在以计算图为根、以原语为边的树上做遍历。最朴素的办法是深度优先搜索：逐个生成完整调度并评估，但空间太大时根本跑不完。

于是枚举器引入三个关键技巧：第一是合法性剪枝，碰到破坏数据依赖的分支立刻退回；第二是下界估计，如果某个前缀调度已经不可能比当前最优解更好，就整棵子树放弃；第三是启发式评估，用代价模型（访存量、并行度、缓存命中率）给候选调度打预估分。TVM 的 Ansor 和 MetaSchedule 正是沿着这条路做的。

剪枝决策可以写成：

至此，从一张菜单到一桌菜，从原始算子的循环到组合爆炸的调度空间，再到聪明的枚举搜索——TVM 的后厨就这么开张了。

AI 应用场景

Python 演示

"""简化版 TVM 调度原语空间枚举演示。"""from itertools import product  class Schedule:    """一次调度：循环大小为 size，记录已应用的原语。"""     def __init__(self, size):        self.size = size        self.loops = 1        self.stages = []     def split(self, factor):        if self.size % factor != 0:            return False        self.size //= factor        self.loops += 1        self.stages.append(('split', factor))        return True     def fuse(self):        if self.loops < 2:            return False        self.loops -= 1        self.stages.append(('fuse',))        return True     def cache_read(self):        self.stages.append(('cache_read',))        return True     def rfactor(self):        if self.loops < 2:            return False        self.stages.append(('rfactor',))        return True     def clone(self):        s = Schedule(self.size)        s.loops = self.loops        s.stages = self.stages[:]        return s     def describe(self):        if not self.stages:            return 'raw-loop'        return ' -> '.join(            f'{name}({arg})' if arg is not None else f'{name}()'            for name, arg in self.stages        )  OPTIONS = [    ('split', [2, 4]),    ('fuse', [None]),    ('cache_read', [None]),    ('rfactor', [None]),]  def enumerate_schedules(size, max_steps=3):    frontier = [Schedule(size)]    results, visited = [], set()     for _ in range(max_steps):        nxt = []        for sched in frontier:            for name, args in OPTIONS:                for arg in args:                    cand = sched.clone()                    ok = cand.split(arg) if name == 'split' else \                         cand.fuse() if name == 'fuse' else \                         cand.cache_read() if name == 'cache_read' else \                         cand.rfactor()                    if ok:                        key = cand.describe()                        if key not in visited:                            visited.add(key)                            results.append(cand)                            nxt.append(cand)        frontier = nxt     return results  if __name__ == '__main__':    found = enumerate_schedules(16, max_steps=3)    for idx, sched in enumerate(found, 1):        print(f'{idx:3d}. {sched.describe()}')    print(f'共 {len(found)} 个合法调度序列')

1. self.stages.append(('split', factor))每成功应用一个原语就记录下来，形成调度路径。
2. if self.size % factor != 0: return False不能整除就拒绝，保证拆分后的循环边界合法。
3. def clone(self)深拷贝当前调度状态，枚举时各分支互不干扰。
4. cand.split(arg) if name == 'split' else ...按原语名分派到对应方法，模拟调度器的派发逻辑。
5. if __name__ == '__main__':给脚本一个可执行入口，直接运行即可打印全部枚举结果。

常见误区

面试题