冯思远：Apache TVM 与机器学习编译发展

By 超神经

内容一览：本文整理自 Apache TVM PMC、上海交通大学博士生冯思远在 2023 Meet TVM 上海站的演讲分享，主题为「TVM 与机器学习编译发展」，HyperAI超神经做了不改变原意的删减整理。

关键词：TVM   2023 Meet TVM   机器学习编译

下午好，欢迎大家今天来参加 2023 Meet TVM。作为 Apache TVM PMC，由我来给大家做关于 TVM 的发展以及 TVM 未来 Unity 框架的分享。

Apache TVM  Evolution

首先为什么会有 MLC (Machine Learning Compilation)？随着 AI 模型的不断扩展，实际生产应用中会有涌现更多的需求，很多应用的第一层 AI Application（如图所示）是共用的，包括 ResNet、BERT、Stable Diffusion 等模型。

第二层 Scenario 就不一样了，开发者需要在不同的场景去部署这些模型，最开始是云计算、高性能计算，需要用 GPU 进行加速。随着 AI 领域的加速，最重要的任务就是把它带进千家万户，也就是 Personal PC 个人电脑、Mobile 手机以及 Edge 设备。

然而不同的场景有不同的需求，包括降低成本、提升性能。像 Out of the Box 需要保证用户打开网页或者下载一个应用，可以立马使用，手机需要省电，Edge 需要在没有 OS 的硬件上运行起来，有些时候还需要在低功耗、低算力的芯片上把技术跑起来。这些就是大家在不同应用上遇到的困难，怎么解决？

MLC 领域对此是有共识的，即 Muli-Level IR Design。核心有三层，第一层 Graph-Level IR，中间层 Tensor-Level IR 以及下一层 Hardware-Level IR。这几层是必须的，因为模型是一个 Graph，中间层是 Tensor-Level IR，MLC 的核心就是优化 Tensor Computing。底下这两层 Hardware-Level IR 和 Hardware 是相互绑定的，就是说 TVM 不会去涉及直接生成汇编指令这一层，因为中间会有一些更细的优化技巧，这一层交由厂商或者编译器去解决。

ML Compiler 在设计之初有一下目标：

Dependency Minimization

第一点，最小化依赖部署。为什么现在 AI 应用没有真正落地，正是因为部署的门槛太高。运行过 ChatGPT 的人比 Stable Diffusion 多并不是因为 Stable Diffusion 不够厉害，而是因为 ChatGPT 提供了一个开箱可用的环境。在我看来，使用 Stable Diffusion，你需要先从 GitHub 上下载一个模型，然后开一个 GPU 服务器把它部署起来，但是 ChatGPT 则是开箱即用。开箱即用关键的一点就是最小化依赖，所有人、所有环境上都可以用。

Various Hardware Support

第二点，能够支持不同硬件。多样化硬件部署在发展早期不是一个最重要的题，但随着国内外 AI 芯片的发展，它会变得越来越重要，尤其是国内目前的环境以及国内芯片公司的现状，需要我们对各个硬件都有一个很好的支持。

Compilation Optimization

第三点，通用编译优化。通过上几层IR的编译可以优化性能，包括提升运行效率、减少内存的占用等。

现在大家大多将边编译边优化作为最重要的点，但对整个社区来说，前两点是关键的。因为这是从编译器角度来看的，并且这两点都是从零到一的突破，优化性能往往是锦上添花。

回到演讲主题，我将 TVM 的发展划分为四个阶段，仅代表个人观点。

TensorIR Abstraction

Stage 1：在这个阶段，TVM 在 CPU、 GPU 上做推理的优化及加速，GPU 特指 SIMT 的硬件部分。这个阶段很多云计算厂商因为发现在 CPU 和 GPU 上都能加速开始用 TVM。为什么？我前面讲到 CPU 和 GPU 没有 Tensorization Support。TVM 初代 TE Schedule 基于 Halide，没有良好的 Tensorization Support ，所以 TVM 后续发展以 Halide 发展的技术路线，包括 Auto TVM 和 Ansor，都对 Tensorization 的支持不友好。

首先来看硬件发展过程，从 CPU 到 GPU 是 2015、2016 年左右，从 GPU 到 TPU 是 2019 年左右。为了做 Tensorization 支持，TVM 先分析了 Tensorized Programs 特征。

Optimized loop nests with thread binding

第一，需要一个 loop testing，这是所有 Tensorized Program 必需的，底下有 Multi-dimensional data load，这和 CMT、CPU 不一样，它以张量而不是以标量为单位储存和计算的。

Multi-dimensional data load/store into specialized memory buffer

第二，它储存在一个特殊 memory buffer 里。

Opaque tensorized computation body 16x16 matrix multiplication

第三，会有一个硬件池允许计算。以下的 Tensor Primitive 为例子，计算 16* 16 的矩阵乘法，这个计算不再会表达成以标量组合的计算模式，而是会以一个指令去计算成账号单元。

依据以上三点 Tensorized Program 的定性分析，TVM 引入了 Computational Block。Block 是一个计算单元，其最外层有 nesting，中间有 iteration iterator、dependency relations，最底下是 body，其概念就是把内外层的计算分离，即 Isolate the internal computation tensorized computation。

Stage 2：在这个阶段，TVM 做 Auto-Tensorrization。具体如何实现的，这里一个例子详细展开。

Auto-Tensorization

输入 Input Program 和 Tensor Intrinsic，结果表明：TensorIR 和 TensorRT 在 GPU 上基本上打平，而在一些标准模型性能并不是很好。因为标准模型是 ML Perf 的标准指标，NVIDIA 工程师会花很多时间去做。在标准模型上表现优于 TensorRT 是比较少见的，这也是相当于打赢了业界最先进的技术。

这是 ARM 自研 CPU 上的性能比较，TensorIR 和 ArmComputeLib 在 End-to-End 上可以比 Ansor 和 PyTorch 快约 2 倍。性能不是最关键的，Auto-Tensorization 思想才最为核心。

Stage 3：An End-to-End ML Compiler for Tensorized Hardware。这个阶段可以将它扔到 GPU 或者已经支持的加速器芯片上，接下来有自动调优、模型导入，是一套自洽的系统。这个阶段 TVM 的核心是 End-to-End，可以开发一个直接使用的模型，但自定义难度很高。

接下来，关于 Relax、Unity 的发展和思考我会讲得慢一点，原因是：

个人认为 Relax、Unity 比较重要；

尚处于实验阶段，很多东西只有思想，缺少 End-to-End demo 及完整代码。

Apache TVM Stack 的局限性：

Huge Gap between Relay and TIR。TVM 最大问题是 relay 到 TIR 的编译范式太陡峭；

Fixed Pipeline for Most Hardware。TVM 标准的流程是 Relay 到 TIR 到？编译下去，实际上很多硬件要么只支持 BYOC，要么想通过 BYOC+TIR，Relay 支持的并不好，Either TIR or Library。以 GPU 加速为例，给 Relay 底层是固定的东西，要么写 CUDA 做 Auto Tuning，要么给 TensorRT 走 BYOC，要不走 cuBLAS 调第三方库。虽然有诸多选择，但都是一个二选一的问题。这个问题影响比较大，而且在 Relay 上不好解决。

解决方案：TVM Unity。

Apache TVM Unity

把 Relax 和 TIR 这两层 IR 当成一个整体，融合成了 Graph-TIR 编程范式。融合的形式：以最简单的 Linear 模型为例，这种情况下整个 lR 是可控可编程的，用 Graph-TIR 这一层的语言解决了 lower 过于陡峭的问题，高层算子可以逐步修改，甚至可以自行改成任何 BYOC 或 Function Calls。

1. 支持自定义解决 building Pipeline 的问题。原本 TVM Pipeline 是 Relay 额外的 TIR，然后 TIR 做 Tuning，Tuning 完成后传递给 LLVM 或 CUDA，这是一套 fix 的 building Pipeline，现在 Pipeline 发生了改变，

2. Choose the Best Perf among All Paths。开发者可以在 library 或 TIR 里做选择，调任意东西。这是 Unity 解决的最重要的问题，社区认为这是 unified ML compilation solution。

Misunderstanding

TVM 和 MLIR 是竞争关系

实际上，TVM 和 MLIR 没有明确的、同层次的竞争关系，TVM 专注于 MLC 机器学习编译，MLIR 强调 Multi-Level，也可以利用其特性来做 ML 编译。开发者使用 MLIR 做机器学习编译，一方面是因为 MLIR 与PyTorch等框架有原生对接，另一方面是因为在 Unity 之前， TVM 的定制化能力比较弱。

TVM = Inference Engine for CPU/GPU

TVM 从来不是一个 inference engine，它能够做编译， 开发者可以用其来做推理加速。TVM 是 compiler infrastructure，但不是 inference engine，「TVM 只能用来加速」这种想法是错误的。TVM 之所以能用来加速，本质原因是 Compiler，比 PyTorch 等 Eager Mode 的执行方式要快。

TVM = Auto Tuning

Relax 出现之前，大家对 TVM 的第一反应是，能够通过 Auto Tuning 获得更好的性能。TVM 接下来的发展方向是淡化这个概念。TVM 提供了各种途径以达到更好的性能、定制化整套编译流程。TVM Unity 要做的就是提供一套架构，把各种优势组合起来。

Next Step

下一步 TVM 要做 Cross-Layer Machine Learning Compiler Infrastructure 以及致力于成为 customizable Building Pipeline for Different Backends，在不同的硬件上支持可定制化，这一点需要结合各种方式、集各家所长。

Q & A

Q 1：TVM 有没有计划优化大模型？

A 1：在大模型方面，我们有一些初步的想法。目前 TVM 已经开始做分布式推理以及一些简单训练，但距离真正落地还需要一段时间。

Q 2：Relax 后续在 Dynamic Shape 方面会有哪些支持和演进？

A 2：Relay vm 支持 Dynamic shape，但没有生成 Symbolic Deduction，例如 nmk 的矩阵乘法其输出是 n 和 m，但在 Relay 的表达方式是 3 个 nmk 统称为 Any，即未知 dimension，其输出也是未知 dimension。Relay VM 能运行这些tasks，但是在编译阶段会丢失一些信息，所以 Relax 解决了这些问题，这就是 Relax 针对 Relay 在 Dynamic Shape 上的改进。

Q 3：TVM 的优化和 Device 的优化的结合问题。用 Graph 直接生成指令 的话，Device 的优化中 TE、TIR 好像不太用到。用 BYOC 的话，好像也是跳过 TE、TIR。分享中提到 Relax 可能会有一些定制化，好像可以解决这一问题。

A 3：实际上已经有很多硬件厂商走通了 TIR 这条路，还有一些厂商没有关注相关的技术，选择的还是 BYOC 方法。BYOC 不是严格意义上的 Compilation，在 Building Pattern 上是有限制的。综合来说，不是企业用不上社区技术，而是依据自身情况做出了不一样的选择。

Q 4：TVM Unity 的出现是否会涉及到较高的迁移成本？从 TVM PMC 的角度来看，如何帮助用户顺畅过渡到 TVM Unity？

A 4：TVM 社区并没有放弃 Relay，只是新增了 Relax 这个选项，因此老版本会持续演进，但是为了使用一些新功能，可能需要做一些代码及版本的迁移。

在Relax完整发布后，社区会提供的迁移教程和一定的工具支持，支持 Relay 模型直接导入到 Relax。但是，基于 Relay 开发的定制版本迁移到 Relax，还是需要一定的工作量，这种工作量对于公司十几人的团队，大概需要一个月。

Q 5：TensorIR 在 Tensor 这块相比之前有很大的进展，但我注意到 TensorIR 主要还是针对编程模型如 SIMT、SIMD 这些成熟的编程方式做的，TensorIR 在现在很多的新的 AI 芯片、编程模型这里有没有一些进展?

A 5：从社区角度来说，TensorIR 之所以做 SIMT 模型，是因为现在只能利用到 SIMT 硬件，很多厂商的硬件和指令集都不开源，能够接触到大厂的硬件基本只有 CPU、GPU 和部分手机 SoC，其他的厂商社区基本接触不到，所以不能根据它们的编程模型做。此外，即使社区和企业合作做出来了类似层级的 TIR ，也无法开源，这是出于商业运营的考量。

以上为冯思远在 2023 Meet TVM 上海站的演讲整理内容。接下来，本次活动其他嘉宾分享的详细内容也将陆续在本公众号中发布推出，欢迎持续关注！

获取 PPT：关注微信公众号「HyperAI超神经」，后台回复关键字TVM 上海，获取完整 PPT。

TVM 中文文档：https://tvm.hyper.ai/

GitHub 地址：https://github.com/apache/tvm