TVMScript元调度原语组合编程

可视化演示

展示TVMScript五大调度原语分类、原语组合优化流水线(blockize→tensorize→compute_at)、以及元调度自动搜索空间探索，切换模式查看不同调度层面。

编程模型

回到我们的中央厨房。总调度师手上有一本厚厚的标准菜谱，也就是 TVM 的 TIR（Tensor Intermediate Representation）——一份用循环、读写、计算写成的“底层工序表”。传统做法是让厨师逐行背诵菜谱，而 TVMScript 给了你一支魔法笔：直接用 Python 语法，在这份工序表上做批注、划线、重组。你写的不是新菜谱，而是一张“怎么执行菜谱”的便签。

TVMScript 的调度脚本长什么样？你创建一个调度器对象 S = Schedule(mod)，它就像一个贴身助理，把你对循环结构的每一次修改都记录在案。通过 S.get_block("compute") 选中某个算子块，再通过 S.split(...)、S.vectorize(...) 等原语发出指令。整个过程完全在 Python 层面完成，不需要手写底层 IR，也不需要重写算法——你只改变“执行方式”，不改变“计算结果”。

用数学语言表达，调度器是在原计算图 $\mathcal{G}$ 上做循环结构映射：

这节告诉我们：调度原语组合的威力再大，也不能改变计算结果。正如总调度师可以随意安排灶台与厨师，但端上桌的必须还是那道红烧肉。

常用原语

厨房里的总调度师有哪些“常用指令”？第一个是循环分块（split）：把“一百份肉一锅炖”拆成“每十份一组，分十批处理”。第二个是向量化（vectorize）：让一个厨师同时拿四个锅铲，一次翻四份肉——对应 CPU/GPU 的 SIMD 指令。第三个是循环融合（fuse）：把相邻两道工序放在同一段流水线上，避免厨师来回跑动。第四个是重排序（reorder）：调整多个循环的先后顺序，让数据访问更顺路。

在 TVMScript 中，这些指令一行一个。例如对一个深度学习算子（如矩阵乘加偏置）：

这就像在厨房操作手册上写下：“第 4 组到第 7 组肉，用四锅连炒，先摆盘后开火。”每一步原语都对应一个数学变换——把一段循环轴拆成两个嵌套轴：

分块、向量化、融合这些原语，本质上都是“重新编排循环结构”。总调度师不必换厨具、不必换食材，只需在流水线上调整节奏，就能让效率天差地别。

sch = tvm.tir.Schedule(mod)block = sch.get_block("compute")i, j = sch.get_loops(block)i_outer, i_inner = sch.split(i, factors=[None, 4])sch.vectorize(i_inner)sch.reorder(j, i_outer, i_inner)

组合策略

单个原语只是“一招鲜”，真正厉害的是组合。中央厨房的效率巅峰，来自把 split、vectorize、fuse 和 reorder 嵌套使用。比如一个复杂算子，有行循环和列循环两层。你可以先把外层行循环 split 成大块，再把内层列循环 vectorize，最后把大块与某个相邻算子 fuse 起来，减少一次全量数据搬运。

组合策略的关键经验是：先切后融、多级嵌套、语义保持。先切是为了暴露并行度；融合是为了减少中间结果的落地；到每一步都要用 TVM 的 T.parallel、T.vectorized 等标注来表达你期望的硬件执行方式。看一个典型的组合示例：

这个组合在数学上对应一个多维循环空间的分块与重排。设原始迭代空间为 $I \times J$，经过两级 split 后变成四层嵌套：

组合策略的诀窍在心中：只有像搭积木一样理解每个原语的“变形效果”，才能面对新算子时信手拈来。总调度师不是记死菜谱，而是根据食材形状灵活调整切法。

i_outer, i_inner = sch.split(i, factors=[None, 8])j_outer, j_inner = sch.split(j, factors=[None, 8])sch.reorder(i_outer, j_outer, i_inner, j_inner)sch.vectorize(j_inner)sch.parallel(i_outer)

缓存优化

为什么最熟练的厨师也会把常用调料放在手边？因为数据搬运比计算更贵。在 AI 算子里，从主存（冰箱）取数往往比计算慢一个数量级。TVMScript 提供了 cache_read 和 cache_write 两组原语，相当于在灶台旁临时搭一个“调料台”（共享内存/寄存器），把反复使用的数据先搬近一点。

cache_read(block, 0, "shared") 的意思是：把计算块的第 0 个输入读到共享内存缓存中。cache_write(block, 0, "local") 则是先写到本地寄存器、攒够一批再统一写回主存，减少“每次一小勺、跑一趟冰箱”的尴尬。多级缓存原语的组合使用，让数据在寄存器→共享内存→全局内存之间形成高效流水线。

用公式表达，设一个计算块总共需要访存 $M$ 次，缓存命中后只有 $M_h$ 次走快速路径，那么访存代价比为：

总调度师最高明的操作，往往不是炒得更快，而是少跑几趟冰箱。缓存优化的本质是时间换空间（数据搬运的减少），这也是元调度原语中性价比最高的一类。

并行协作

中央厨房不可能只有一位厨师。一块 GPU 上有数百个“厨师”，分属不同的“班组”（线程块）。TVMScript 的 bind 原语，就是把这些循环层绑定到具体的线程或块上：sch.bind(loop, "blockIdx.x") 表示外层循环交给不同班组并行跑，sch.bind(inner_loop, "threadIdx.x") 表示内层循环由班组内多位厨师协作完成。

并行协作还需要同步原语——厨师们做完自己的菜不能直接上桌，要等所有同事都报数完成再统一出餐，否则顾客会拿到半熟的菜。在 CUDA 里对应 __syncthreads()，在 TVMScript 中通过 T.sync() 标注。计算每个线程负责的迭代份数，是并行绑定的核心数学：

并行编组做得好，一百份菜能同时出锅；做得不好，一半厨师在等，一半厨师在挤。TVMScript 的 bind 原语，就是帮你把“谁干什么”写到操作手册上。

性能调优

调度原语组合得再好，最终都要用数据说话。性能调优的第一件事是建立指标：吞吐率（每秒处理多少元素）、耗时（毫秒）、带宽利用率。TVMScript 配合 TVM 的 autotuner，可以在一个参数搜索空间内自动尝试不同的原语组合。

瓶颈分析的核心是判断一个算子到底“卡在计算上”还是“卡在访存上”。为此我们引入经典的 Roofline 模型：以算术强度（每字节访存对应多少次计算）为横轴，可达性能为纵轴，形成一条“屋顶曲线”。当算子的算术强度低时，性能等于访存带宽 $\beta$ 乘以强度 $I$；当强度超过拐点时，受限于峰值计算能力 $\pi$：

调优不是玄学，而是“找到短板、用原语补长”。总调度师盯着下单速度与上菜速度的差值，就知道该加灶台还是该改传菜路线——TVMScript 的调度原语就是那支可以动态调整的笔。

实战落地

有了前面的武器，我们看一个真实案例：优化一个矩阵乘法算子 C = A @ B。传统写法是三层朴素循环，性能平平。用 TVMScript 元调度原语组合改造，大致分四步：

这段脚本等于在厨房手册上写下：“把矩阵切片成 16×16 的小瓷砖，8 个 k 值一次算完，行与行并行翻炒。”当这些原语组合起来，矩阵乘的缓存命中率和并行度都会大幅提升，最终收益可对比：

实战落地时，TVMScript 的价值在于：把“调优经验”沉淀为可读的 Python 脚本，而不是散落的编译器黑魔法。总调度师把一次成功的大赛经验写进手册，以后每道菜都能复用这套打法——这就是元调度原语组合编程的魅力。

mod = tvm.script.from_source(matmul_tir)sch = tvm.tir.Schedule(mod)block = sch.get_block("C")i, j, k = sch.get_loops(block)i0, i1 = sch.split(i, factors=[None, 16])j0, j1 = sch.split(j, factors=[None, 16])k0, k1 = sch.split(k, factors=[None, 8])sch.reorder(i0, j0, i1, j1, k1)sch.vectorize(j1)sch.parallel(i0)

AI 应用场景

Python 演示

# -*- coding: utf-8 -*-# 本示例展示了 TVMScript 元调度原语组合的基本流程# 注意: 无真实 TVM 环境时也可运行, 核心 API 调用以中文伪注释说明 import math class LoopInfo:    def __init__(self, name, extent):        self.name = name        self.extent = extent        self.bound = None        self.child = None class ScheduleDemo:    def __init__(self, loop):        self.root = loop        self.calls = []     def get_loops(self):        loops = []        node = self.root        while node is not None:            loops.append(node)            node = node.child        return loops     def split(self, loop, factors):        outer = int(loop.extent // factors[1])        inner = int(factors[1])        self.calls.append("split({0}, {1}) -> outer {2} x inner {3}".format(loop.name, factors, outer, inner))        return LoopInfo(loop.name + "_o", outer), LoopInfo(loop.name + "_i", inner)     def vectorize(self, loop):        self.calls.append("vectorize({0}, width={1})".format(loop.name, loop.extent))     def parallel(self, loop):        self.calls.append("parallel({0})".format(loop.name))     def run(self):        print("=== 初始循环结构 ===")        i = self.get_loops()[0]        print("for {0} in range({1}): # 串行执行".format(i.name, i.extent))        io, ii = self.split(i, [None, 8])        self.vectorize(ii)        self.parallel(io)        print("=== 调度原语执行记录 ===")        for c in self.calls:            print("  " + c)        total = i.extent        per = ii.extent        groups = math.ceil(total / per)        print("=== 优化效果 ===")        print("串行分组数: {0} 组, 每组 {1} 份, 预计加速 {2}x".format(groups, per, groups)) if __name__ == "__main__":    inner_loop = LoopInfo("j", 16)    outer_loop = LoopInfo("i", 64)    outer_loop.child = inner_loop    demo = ScheduleDemo(outer_loop)    demo.run()

1. class ScheduleDemo:用 Python 类模拟 TVMScript 的 Schedule 调度器对象，无需安装 TVM 也能理解调度流程。
2. self.calls = []记录每次调度原语的调用痕迹，类似 TVM 中 sch.split(...) 的日志。
3. def split(self, loop, factors)实现 split 循环分块原语，将 64 份迭代拆成 8 组 × 8 份。
4. def vectorize(self, loop)模拟向量化原语，表示内层 8 份迭代用 SIMD 一次算完。
5. def parallel(self, loop)模拟并行原语，把外层 8 组分配给不同线程并行执行。
6. groups = math.ceil(total / per)计算分组数，这是加速比的理论上限依据。
7. demo.run()组合触发整套调度流程，输出原语记录与优化效果。

常见误区

面试题