MLC on 理塘丁真

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

Fri, 22 May 2026 00:00:00 +0800

前面看 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、运行库，都是同一个计算在不同层级上的表示。编译器要做的，就是在这些表示之间转换，并尽量找到适合目标环境的实现。

我现在的理解

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

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

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

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

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

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

陈天奇 MLC 课程第一章概述：机器学习编译到底在做什么

Thu, 21 May 2026 00:00:00 +0800

最近开始看陈天奇老师团队的 MLC 课程。第一章主要回答一个问题：机器学习编译到底在做什么，为什么它会成为机器学习系统里绕不开的一层。

这篇先整理课程导论里的核心概念，后面再逐章补张量程序抽象、自动程序优化、GPU 加速和计算图优化这些内容。

从开发形式到部署形式

我们平时训练或调试模型时，通常面对的是“开发形式”：比如用 PyTorch、TensorFlow 或 JAX 写出来的模型代码，以及模型训练得到的权重。这个形式适合研究、调试和快速迭代，但不一定适合直接放到手机、浏览器、车载设备、云端服务或专用芯片上运行。

真正上线时，需要的是“部署形式”：它不仅包括模型计算本身，还包括运行所需的库、内存管理、硬件调用、应用接口，以及和具体系统环境配套的执行方式。

机器学习编译做的事情，就是把开发阶段的模型，通过一系列变换和优化，转换成更适合目标环境运行的部署形态。

为什么叫“编译”

传统编译器会把高级语言代码转换成机器可以执行的程序。机器学习编译也有类似的味道：输入是高层模型描述，输出是可以部署和执行的模型程序或运行库。

不过 MLC 和传统编译并不完全一样。它不一定每次都生成底层机器码，有时只是把模型转换成对已有高性能库的调用；有时会做算子融合、图优化、内存优化；有时还会针对 GPU、NPU、TensorCore 等硬件特性生成更合适的实现。

所以 MLC 更像是一套面向机器学习部署的系统方法：在模型表达、计算图、张量程序和硬件实现之间不断做抽象转换。

机器学习编译的三个目标

课程里把 MLC 的目标概括成三类。

第一是集成与最小化依赖。部署一个模型时，我们只想带上真正需要的计算和运行组件，而不是把整个训练框架都塞进应用里。这样可以减小体积，也能让模型更容易部署到资源受限的设备上。

第二是利用硬件加速。不同环境有不同的硬件能力：CPU、GPU、移动端 NPU、云端加速器，甚至专用 AI 芯片。MLC 需要把模型计算映射到这些硬件擅长的执行方式上，例如调用原生加速库，或者生成能利用特殊指令的底层程序。

第三是通用优化。同一个模型可以有多种等价执行方式，但不同写法的性能差距可能很大。MLC 会通过计算图优化、算子融合、内存复用、调度优化等方式，让模型跑得更快，占用更少资源。

这三个目标之间不是完全分开的。比如算子融合既可以减少依赖和运行开销，也可能带来更好的硬件利用率。

张量和张量函数

理解 MLC 的关键，是先把模型执行看成张量之间的计算。

张量是模型中的输入、输出和中间结果，本质上是多维数组。图片、文本特征、隐藏层状态、权重矩阵，最后都可以落到张量这个统一表示上。

张量函数则描述张量之间如何计算。一个线性层、一个 ReLU、一个 Softmax 都可以看成张量函数；多个算子组合起来的一段计算，甚至整个模型，也可以看成更大的张量函数。

MLC 做优化时，经常不是只盯着单个算子，而是看一组张量函数之间能不能被重写、融合或映射到更高效的实现。

上图展示了一个典型例子：开发阶段的模型里，linear 和 relu 是两个独立计算；部署时可以把它们融合成一个 linear_relu。这样可以减少中间张量的读写，也给后端实现留下更多优化空间。

这就是我理解里的“图级算子融合”：它不是改变模型数学意义，而是在保持结果等价的前提下，换一种更适合执行的计算组织方式。

抽象与实现

第一章里很重要的一句话是：抽象描述“做什么”，实现描述“怎么做”。

在机器学习系统里，同一个张量函数可以有很多层表示：它可以是计算图中的一个节点，也可以是循环嵌套，也可以是调用某个库函数，还可以是生成出来的一段 GPU kernel。

MLC 的核心工作，就是在这些抽象层之间做转换。高层抽象方便表达模型，低层实现负责贴近硬件。编译器要做的是在不改变语义的前提下，找到更高效、更适合部署环境的实现。

这也是为什么机器学习编译不只是“让模型快一点”。它实际上连接了模型、框架、运行时、硬件和应用部署。

为什么要学 MLC

对机器学习工程师来说，MLC 能帮助我们理解模型从训练代码走向生产环境的过程。当模型部署变慢、占内存过大，或者目标平台不被主流框架很好支持时，只会调模型结构是不够的，还需要理解底层执行链路。

对算法同学来说，MLC 可以解释为什么同一个模型在不同平台上性能差别很大，也能帮助我们判断一个模型结构是否适合真实部署。

对硬件和系统方向来说，MLC 提供了一套把新硬件能力接入机器学习生态的方法。模型越来越大，硬件越来越多样，靠手写适配很难长期维护，编译和自动优化会越来越重要。

小结

MLC 第一章可以概括成一句话：机器学习编译研究的是如何把开发阶段的机器学习模型，转换成适合具体环境部署和高效执行的形式。

它关心的不只是模型本身，还包括依赖裁剪、硬件加速、计算图优化、张量程序变换和运行时集成。后面的课程，其实就是沿着这些问题继续往下拆。

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