[MLC-01] Machine Learning Compilation Overview

这是陈天奇《Machine Learning Compilation》课程第一讲的学习笔记。课程从一个很实际的问题开始：模型在 PyTorch、TensorFlow 或 JAX 中能跑，并不代表它已经能以合适的延迟、内存和依赖规模部署到目标环境。

1. MLC 在编译什么#

机器学习编译可以概括为：把模型的开发形式变换成适合部署的形式。

开发形式
  模型定义 + 权重 + 框架 API
          |
          | IR 表示、程序变换、优化、代码生成或库调用
          v
部署形式
  可执行函数 + 运行时 + 资源管理 + 应用接口

它和传统编译器相似，都要保持程序语义并生成目标环境可执行的结果；但 MLC 的输入是张量计算图、权重、算子库和硬件约束，输出也未必是汇编。很多时候，编译结果是对已有高性能库的调用、一个部署包或一个能被运行时加载的模块。

2. 三个直接目标#

课程将 MLC 的目标归纳为三类：

这三类目标彼此并不独立。算子融合既减少了中间张量分配，也可能让数据更好地留在缓存或片上存储中；选择一个库实现既是集成问题，也会影响性能上限。

3. 张量与张量函数#

模型执行可以被看成张量函数的组合：

其中 $X$ 是输入张量，$W$ 是权重，$Y$ 是输出或中间结果。linear、relu、softmax 可以是元张量函数；整个模型也可以被看作一个更大的张量函数。

一个重点是：同一张量函数可以有多种等价实现。

linear -> relu
      |
      +-> 两次独立 kernel
      +-> 一个 fused kernel
      +-> 调用外部库
      +-> 生成目标设备专用代码

MLC 的工作正是在这些等价实现之间做选择和变换。

4. 抽象与实现#

课程里最值得反复记住的一句话是：抽象描述做什么，实现描述怎么做。

同一个 linear_relu 可以被表示为计算图节点，也可以被表示为循环嵌套、缓冲区读写和向量指令。前者保留高层组合关系，后者暴露循环、内存访问和并行方式。不同层次的表示决定了编译器能安全做哪些变换。

计算图        适合：算子融合、常量折叠、全图布局
张量程序      适合：tile、reorder、vectorize、memory scope
目标指令/库    适合：最终执行与硬件对接

因此，MLC 不是单一 pass，而是围绕不同抽象层组织的变换链。

5. 与推理 serving 的联系#

在大模型 serving 中，模型图、attention kernel、KV Cache 管理和调度器看起来是不同组件，但它们同样受“抽象与实现”的约束：高层请求调度不应直接绑定某个 kernel；kernel 优化又必须理解张量 shape、布局和内存层级。

学习 MLC 的价值在于建立一条从模型语义到执行细节的连续视角，而不是只把性能问题归结为“换一个更快的算子”。

6. 本讲总结#

MLC = 将张量函数从开发形式变换到部署形式。
 
核心对象：张量、张量函数、IR、运行时、目标硬件。
核心方法：在保持语义的前提下，选择或生成更合适的实现。

参考#

• Machine Learning Compilation, Summer 2022
• 课程原始学习材料