深度学习编译器中的算子调度与性能建模

前面看 MLC 的时候，我有个误区：看到 MatMul 就觉得它是矩阵乘法，看到 Conv2D 就觉得它是卷积，好像知道算子名字就差不多知道它怎么跑了。

后来我发现这只是知道“它要算什么”。离“它在某块硬件上怎么跑得快”，中间还隔着 shape、layout、schedule、内存层级和硬件指令。

这篇不是深度学习编译器教程，只是我给自己补的一层理解：算子定义语义，调度改执行方式，性能模型帮编译器判断哪个实现更值得用。

这个理解对我比较有用，是因为后面做端侧部署时，问题经常不是“模型能不能转格式”，而是更底下的东西：这个 shape 有没有合适 kernel，layout 转换会不会抵消收益，中间张量会不会把内存顶满，目标硬件到底擅长哪类计算。

算子像合同，不像施工图

深度学习模型最后会落到一堆算子上：

1
2
3
4
5
6
7
8
Conv2D
MatMul / GEMM
BatchNorm
ReLU
Softmax
LayerNorm
Add
Transpose

这些名字更像语义合同。它们说清楚输入输出的数学关系，但没有告诉你循环怎么排、数据怎么搬、线程怎么分。

拿矩阵乘法来说，最朴素的三层循环当然能算对：

1
2
3
4
5
6
7
for (int i = 0; i < M; i++) {
    for (int j = 0; j < N; j++) {
        for (int k = 0; k < K; k++) {
            C[i][j] += A[i][k] * B[k][j];
        }
    }
}

这段代码的问题不是“数学不对”，而是它没有表达任何硬件意识：cache 怎么复用，SIMD 怎么用，多线程怎么分块，GPU 上 block 和 shared memory 怎么配，端侧 NPU 又有哪些 kernel 和 buffer 限制。

所以我现在更愿意把算子看成一个上层接口。它告诉编译器要做哪类计算；至于怎么把这个计算铺到硬件上，是调度要解决的事情。

调度改的是执行形态

调度源语不会改变计算结果，但会改变循环结构、内存访问和并行方式。

源语	作用	我自己的理解
`split` / `tile`	拆循环、分块	不要一口吃完整个矩阵，先切成能复用的小块
`reorder`	调整循环顺序	让连续访问更连续，少一点 cache miss
`vectorize`	向量化	让一条 SIMD 指令处理多个元素
`unroll`	展开循环	少一点循环控制，多一点指令级并行
`parallel`	并行化	把不同块分给不同线程或计算单元
`cache_read` / `cache_write`	缓存读写	把热点数据放到更近的存储里
`compute_at`	控制中间计算位置	决定中间结果在哪层循环里生成
`inline`	内联简单计算	减少不必要的中间结果

这些词看起来很“编译器”，但背后的目标很朴素：让数据少搬一点，让计算单元少等一点。

比如对 MatMul 做 tile，代码形态可能变成这样：

 1
 2
 3
 4
 5
 6
 7
 8
 9
10
11
for (int io = 0; io < M; io += 32) {
    for (int jo = 0; jo < N; jo += 32) {
        for (int i = io; i < io + 32; i++) {
            for (int j = jo; j < jo + 32; j++) {
                for (int k = 0; k < K; k++) {
                    C[i][j] += A[i][k] * B[k][j];
                }
            }
        }
    }
}

这个例子还很粗糙。真实系统里还会继续调循环顺序、向量宽度、展开因子、线程绑定、shared memory 使用方式。重点不是 32 这个数字，而是它把一个大矩阵乘法改成了一组更适合缓存和并行的小问题。

调度空间很快会失控

一个矩阵乘法看起来简单，但只要开始调度，可选项马上变多：

tile size 取多少；
i/j/k 三层循环怎么排序；
SIMD 宽度怎么用；
哪一层循环展开；
哪一层并行；
A、B、C 哪些数据要进 cache；
中间结果什么时候写回；
数据布局要不要变。

这些选择不是互相独立的。tile size 变了，cache 复用会变；循环顺序变了，访问模式会变；layout 变了，前后算子的连接成本也会变。

这就是 AutoTVM、Ansor、Triton autotune 这类系统有意义的地方。它们本质上都在回答一个问题：

同一个计算，在这块硬件上，哪种执行方式更划算？

这里的“划算”也不只是快。编译时间、内存峰值、部署体积、动态 shape 支持，都会影响最后选择。

性能模型不能只看 FLOPs

以前我看算子性能，很容易先看计算量。比如一个大 MatMul 有多少 MACs，一个卷积有多少 FLOPs。

这个指标有用，但只看它会误判。

一个很粗的估计是：

1
执行时间 ≈ 计算量 / 硬件峰值算力

真实执行经常卡在别的地方：

内存带宽；
cache 命中率；
数据复用；
SIMD 利用率；
kernel launch overhead；
layout conversion；
中间结果读写；
算子是否能融合；
硬件有没有对应特殊指令。

所以我现在会把性能问题写成这一串：

1
operator + shape + dtype + layout + schedule + hardware

少一个维度，结论都可能变。

比如 Conv2D 本身不是一个性能答案。1x1 Conv、3x3 Conv、depthwise conv、不同 batch、不同 layout，优化方向都可能不一样。一个 kernel 在桌面 GPU 上很快，不代表放到端侧芯片上也合适。

Roofline 给了一个判断边界

Roofline Model 里有个概念叫算术强度：

1
Arithmetic Intensity = 计算量 / 数据搬运量

我觉得它最有用的地方，是帮我先判断瓶颈大概在哪。

大规模 MatMul 往往更接近 compute-bound。它的计算量约是：

1
M × N × K 次 MAC

如果按 FLOPs 统计，通常约是：

1
2 × M × N × K

A、B 的元素会被反复复用，所以只要调度做得好，它比较容易把计算单元喂饱。这也是 GPU、NPU、TensorCore 都很重视 GEMM 的原因。

反过来，ReLU、Add 这类算子计算很轻。以 ReLU 为例：

1
y = max(x, 0)

每个元素只做一次比较，但输入要读，输出要写。它的瓶颈很容易变成访存。

这时候单独优化 ReLU 的计算指令意义有限。更实际的做法是融合。

融合不是为了看起来高级

比如：

1
Conv -> BatchNorm -> ReLU

如果三个算子分开跑，中间结果会反复写回内存、再读出来。

推理阶段可以先做 BN folding：

1
BN(Conv(x, W, b)) -> Conv(x, W', b')

然后再把 Conv + ReLU 融成一个 kernel。这样 Conv 算完以后直接做 ReLU，最终结果写回一次。

这里的重点不是“融合”这个词，而是它减少了中间张量的读写。对于 memory-bound 的部分，这往往比抠几条计算指令更有用。

性能模型有几种粗细

工程里不可能所有东西都靠完整搜索。不同系统会混着用几类方法。

手工规则最容易落地。

比如大规模 MatMul 走某个加速库，小算子回退 CPU，ReLU/Add 尽量融合，Reshape 尽量消除。规则可控，但维护成本高，对新 shape 和新硬件不一定泛化。

解析代价模型会估计计算时间和访存时间：

1
2
3
Compute_Time = Ops / Effective_Compute_Throughput
Memory_Time  = Bytes / Effective_Memory_Bandwidth
Time ≈ max(Compute_Time, Memory_Time) + Overhead

它不一定准，但能告诉你方向。比如某个卷积算出来访存时间明显大于计算时间，那就应该优先想 layout、融合、复用，而不是只盯着峰值算力。

实测加搜索更贴近真实硬件。

系统生成多个候选 schedule，编译后在目标硬件上跑，记录时间，再继续搜索。它的好处是能适应实际环境，坏处也明显：慢、复杂、需要目标硬件反馈。

编译器最后是在做选择题

把这些东西放一起看，深度学习编译器一直在做选择：

1
2
3
4
5
6
这个算子放 CPU、GPU 还是 NPU？
这个 shape 用哪个 kernel？
这个 layout 要不要转换？
这个中间结果要不要写回？
这个算子能不能融合？
这个 schedule 值不值得编译？

所以它不是简单地把 ONNX、TensorFlow 或 PyTorch 模型换个格式。格式转换只是入口，后面还有图优化、算子选择、调度搜索、性能建模、内存规划和代码生成。

这也接上了 MLC 第一章里讲的“抽象与实现”。算子、计算图、循环嵌套、kernel、运行库，都是同一个计算在不同层级上的表示。编译器要做的，就是在这些表示之间转换，并尽量找到适合目标环境的实现。

我现在的理解

这篇写完以后，我对“机器学习编译”的理解比之前具体了一点。

算子是语义层，它说清楚要算什么。

调度是实现层，它决定循环、并行、缓存和数据移动怎么组织。

性能模型是决策层，它判断某个实现放在某块硬件上到底值不值得。

以前我会把“优化模型部署”理解成换框架、转格式、开量化。现在看，真正难的地方在更下面：同一个计算，换一种执行组织方式，性能可能完全不一样。

参考资料：《机器学习编译》中文课程导论