关于人工智能:冯思远Apache-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 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 的局限性：

1. Huge Gap between Relay and TIR。TVM 最大问题是 relay 到 TIR 的编译范式太平缓；
2. 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 上海站的演讲整顿内容。接下来，本次流动其余嘉宾分享的具体内容也将陆续在「HyperAI 超神经」公众号中公布推出，欢送继续关注！

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

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

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