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

本文首发自 HyperAI 超神经微信公众号~

下午好，欢迎大家今天来参加 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 16×16 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