图优化：算子调度 Op Scheduling

可视化演示

展示HLO指令依赖图拓扑排序调度、内存约束调度峰值显存优化及多流水线并行执行重叠三种核心模式，直观呈现AI编译器图优化调度机制。

核心概念

总厨拿到一张「做菜流程网络图」时，第一反应不是立刻开火，而是先看全局：哪些工序必须先做，哪些工序可以穿插进行，哪些工具用完后要马上腾出来。这张「动线图」就是调度的雏形。AI 模型的计算也能画成同样的一张图：节点是算子，边是数据流向。算子调度 Op Scheduling，就是调度器给每个算子安排「开始时间」和「使用哪个资源」的过程。

在 PyTorch 或 XLA 这类编译器中，计算图先被捕获出来，然后调度器在图上做分析。它不改写算子的数学含义，只决定执行节奏。比如两个卷积之间如果夹着一个激活函数，调度器可以选择把它们放进同一个计算内核，也可以选择分开执行；选择不同，访存次数和 GPU 利用率就不同。这就像总厨决定让同一口锅「焯水后立刻爆香」，还是先洗锅再另起炉灶。

形式化地说，调度可以看作一个从算子集合到开始时间的映射函数：

总厨排动线图时，不需要改菜谱本身，只需要让每个厨师清楚「我什么时候干什么」。算子调度也是如此：不改算子，只改执行节奏。这个节奏，往往决定了模型是「飞快的」还是「卡顿的」。

数据依赖

一盘番茄炒蛋，必须先打蛋、切番茄，再下锅。总厨绝不会让炒锅先点火，然后干等蛋液。菜谱里的「先后箭头」是不能违背的顺序。计算图里的箭头，就是数据依赖。

如果一条边从算子 u 指向算子 v，意味着 u 的输出是 v 的输入。u 是生产者，v 是消费者。调度器必须让生产者先算完，消费者才能开工；一旦顺序颠倒，v 就会读到未初始化的数据。这种依赖关系让计算图天然成为有向无环图（DAG）——有环意味着算子永远在等一个永不出现的输入。

合法调度必须满足：对每一条依赖边，上游算子的完成时间不晚于下游算子的开始时间。这个条件写出来是：

总厨会把菜谱整理成一张「必须先做什么」的清单，让任何工序都不能插队到未完成的前一步之前。这种清单在编译器中就叫拓扑排序，它是调度方案的「草稿」，保证大方向不违反依赖。

硬件资源

动线图画得再漂亮，也要看厨房条件：几个灶台、几口锅、几个帮厨。总厨不能安排五道菜同时开火，因为灶台只有三个。AI 硬件同样有资源上限：GPU 的流式多处理器数量、内存带宽、显存容量，都是调度器要面对的「灶台和锅」。

调度器在排时间表时，不只是看依赖是否满足，还要保证任何时刻所有正在运行的算子占用的资源总量不超过硬件上限。如果两个算子没有依赖关系，但一共需要三块大显存，而设备只剩两块，它们就只能错开执行。也就是说，并行不是「想并行就能并行」，资源约束会直接压缩并行度。

这个约束可以写成：

资源约束像厨房的物理天花板，总厨只能在「天花板」下追求效率，而不是纸上谈兵。调度器设计得再好，也不能超出硬件能力去空想并行。

并行分析

蒸米饭和切配菜互不依赖，总厨可以让一个帮厨去蒸饭，另一个去切菜；而「焯水 → 爆香 → 下菜」必须一条线走完。并行分析就是找出计算图里那些没有依赖关系的算子，把它们放进同一个「并行组」。

两个算子只要能找到一条依赖路径互相连通，它们就不能真正并行；反之，若它们在两条不同的分支上，且没有汇合到一个必须等待的同步点，调度器就可以让它们同时执行。关键路径就是图中从起点到终点最长的那条依赖链。即使资源无限多，模型也不可能比关键路径更快跑完，因此它是整体耗时的下界。

计算关键路径长度有一个递推式：

总厨识别出「同时进行才最省时间」的步骤，也知道什么时候必须等所有帮厨完工才能上下一道工序。关键路径就是那串最不能拖延的菜，而其他工序都可以在它周围灵活穿插。

内存优化

一口锅焯完水，洗干净又可以拿来炒菜；如果每道步骤都买一口新锅，厨房再大也会被锅占满。计算图里的中间张量也存在显存 buffer 中：张量从被生产出来到被最后一个消费者读完，这一段存活时间称为生命周期（live interval）。

如果两个中间张量的生命周期完全不重叠，调度器就可以让它们复用同一块显存 buffer，而不是同时占两份空间。这也是为什么调度顺序会直接影响显存峰值：一个看似合理的执行顺序，可能让两个大张量「撞车」；而稍微调整先后，就能让它们错峰使用同一块内存。

内存复用条件可以写成：

总厨会紧盯「清台」时机：一道菜刚用完某口锅，下一道菜立刻接管。内存复用就是这种「抢锅」的智慧，让有限的显存容纳更大的模型、更长的序列。

调度策略

总厨排菜讲究策略：是先做耗时大菜让它慢慢炖，同时穿插小菜，还是先全力把所有简单菜上齐？策略不同，客人的等待体验完全不同。算子调度同样需要策略，而不是把所有合法拓扑序都试一遍——那会慢得不可接受。

最经典的方法是列表调度（List Scheduling）：维护一个「就绪队列」，所有前驱都完成的算子进入队列，调度器每次从队列里按优先级挑出要执行的算子。优先级怎么定？常用的做法是看关键路径：越靠近出口的算子越要优先执行，因为它在最长的依赖链上，稍有拖延就会拉长整体时间。

优先级递推式可以定义成：

总厨的「策略」就是调度器的启发式规则。策略越聪明，同样的灶台就能做出更多菜；策略越粗糙，资源再满也有可能让关键步骤排队等空位。

工程实践

现代餐厅的中央厨房不再只靠厨师个人经验，而是用智能排菜系统统一调度多家分店订单。AI 编译器就是模型后厨的「智能排菜系统」。PyTorch 2.0 会先捕获 FX 计算图，再做算子调度和代码生成；XLA 会把高层算子图融合成底层内核，再在 GPU 或 TPU 上安排执行顺序。

工程上关注的不再是单个算子快不快，而是端到端推理快不快。调度器要在依赖、资源、内存之间做权衡：太激进的并行会让显存峰值飙升，太保守的串行又会让 GPU 大量空转。PyTorch 2.0 的 TorchInductor 与 XLA 都在做同一件事：把「后厨动线」自动化。

衡量调度收益最直观的指标是加速比：

从手工排菜到中央厨房智能调度，算子调度让 AI 模型的「后厨」从人肉指挥走向自动优化。这也是推理加速和部署优化最核心的技术支撑之一。

AI 应用场景

Python 演示

# 简单列表调度器：模拟总厨在资源受限的厨房里排菜class OpNode:    def __init__(self, name):        self.name = name        self.duration = 1        self.resources = 1  class OpGraph:    def __init__(self, names, edges):        self.nodes = [OpNode(n) for n in names]        self.index = {n: i for i, n in enumerate(names)}        self.successors = [[] for _ in names]        self.indegree = [0] * len(names)        for u, v in edges:            self.successors[self.index[u]].append(self.index[v])            self.indegree[self.index[v]] += 1  class ListScheduler:    def __init__(self, graph, resource_limit):        self.graph = graph        self.resource_limit = resource_limit        self.remaining = graph.indegree[:]        self.finish_time = [0] * len(graph.nodes)        self.ready = [i for i in range(len(graph.nodes)) if self.remaining[i] == 0]     def schedule(self):        t = 0        running = []        finished = 0        result = []        total = len(self.graph.nodes)        while finished < total:            available = self.resource_limit - sum(self.graph.nodes[i].resources for i, _, _ in running)            for i in self.ready[:]:                if available >= self.graph.nodes[i].resources:                    available -= self.graph.nodes[i].resources                    self.ready.remove(i)                    running.append((i, t, self.graph.nodes[i].duration))            if running:                t += 1                next_running = []                for i, start, dur in running:                    if t - start >= dur:                        self.finish_time[i] = t                        result.append((self.graph.nodes[i].name, start, t))                        finished += 1                        for succ in self.graph.successors[i]:                            self.remaining[succ] -= 1                            if self.remaining[succ] == 0:                                self.ready.append(succ)                    else:                        next_running.append((i, start, dur))                running = next_running            else:                t += 1        return result, t  if __name__ == "__main__":    names = ["a", "b", "c", "d", "e", "f"]    edges = [("a", "b"), ("a", "c"), ("b", "d"), ("c", "d"), ("d", "e"), ("d", "f")]    g = OpGraph(names, edges)    scheduler = ListScheduler(g, resource_limit=2)    plan, makespan = scheduler.schedule()    for name, start, end in plan:        print(name, start, end)    print("makespan:", makespan)

1. self.remaining = graph.indegree[:]复制每个算子的前驱数量，作为「还没完成的依赖」计数。
2. self.ready = [i for i in range(len(graph.nodes)) if self.remaining[i] == 0]找出所有不需要等待的算子，放入就绪队列。
3. for i in self.ready[:]:遍历就绪队列快照，因为循环里会修改队列本身。
4. running.append((i, t, self.graph.nodes[i].duration))记录算子 i 在 t 时刻开始执行，并占用资源。
5. if t - start >= dur:判断当前算子是否已经执行完自己的耗时。
6. self.remaining[succ] -= 1下游算子的等待条件减少一个。
7. result.append((self.graph.nodes[i].name, start, t))保存最终的调度计划。

常见误区

面试题