TVM 学习笔记（一） 初识 TVM

引言 - 什么是机器学习编译

• 机器学习编译 (machine learning compilation, MLC) 是指，将机器学习算法从开发阶段，通过变换和优化算法，使其变成部署状态。
  • 开发形式：是指我们在开发机器学习模型时使用的形式。典型的开发形式包括用 PyTorch、TensorFlow 或 JAX 等通用框架编写的模型描述，以及与之相关的权重。
  • 部署形式：是指执行机器学习应用程序所需的形式。它通常涉及机器学习模型的每个步骤的支撑代码、管理资源（例如内存）的控制器，以及与应用程序开发环境的接口。

为什么需要深度学习编译器

• 从训练框架角度来看，Google的TensorFlow和FaceBook的Pytorch是全球主流的深度学习框架，另外亚马逊的MxNet，百度的Paddle，旷视的MegEngine，华为的Mindspore以及一流科技的OneFlow也逐渐在被更多人接受和使用。这么多训练框架，我们究竟应该选择哪个？如果追求易用性，可能你会选择Pytorch，如果追求项目部署落地，可能你会选择TensorFlow，如果追求分布式训练最快可能你会体验OneFlow。

• 从推理框架角度来看，无论我们选择何种训练框架训练模型，我们最终都是要将训练好的模型部署到实际场景的，在模型部署的时候我们会发现我们要部署的设备可能是五花八门的，例如Intel CPU/Nvidia GPU/Intel GPU/Arm CPU/Arm GPU/FPGA/NPU(华为海思)/BPU(地平线)/MLU(寒武纪)，如果我们要手写一个用于推理的框架在所有可能部署的设备上都达到良好的性能并且易于使用是一件非常困难的事。

一般要部署模型到一个指定设备上，我们一般会使用硬件厂商自己推出的一些前向推理框架，例如在Intel的CPU/GPU上就使用OpenVINO，在Arm的CPU/GPU上使用NCNN/MNN等，在Nvidia GPU上使用TensorRT。虽然针对不同的硬件设备我们使用特定的推理框架进行部署是最优的，但这也同时存在问题，比如一个开发者训练了一个模型需要在多个不同类型的设备上进行部署，那么开发者需要将训练的模型分别转换到特定框架可以读取的格式，并且还要考虑各个推理框架算子实现是否完全对齐的问题，然后在不同平台部署时还容易出现的问题是开发者训练的模型在一个硬件上可以高效推理，部署到另外一个硬件上性能骤降。并且从探索ONNX模型转换来看，不同框架间模型转换工作也是阻碍各种训练框架模型快速落地的一大原因。

深度学习编译器与传统编译器

现代传统编译器（LLVM为例）包含编译器前端，编译器中端，编译器后端等概念，并引入IR (Intermediate Representation)的概念。解释如下： - 编译器前端：接收C/C++/Java等不同语言，进行代码生成，吐出IR - 编译器中端：接收IR，进行不同编译器后端可以共享的优化，如常量替换，死代码消除，循环优化等，吐出优化后的IR - 编译器后端：接收优化后的IR，进行不同硬件的平台相关优化与硬件指令生成，吐出目标文件

传统编译器编译流程

受到传统编译器的启发，深度学习编译器被提出，我们可以将各个训练框架训练出来的模型看作各种编程语言，然后将这些模型传入深度学习编译器之后吐出IR，由于深度学习的IR其实就是计算图，所以可以直接叫作Graph IR。针对这些Graph IR可以做一些计算图优化再吐出IR分发给各种硬件使用。这样，深度学习编译器的过程就和传统的编译器类似，可以解决上面提到的很多繁琐的问题。

TVM

TVM 是一个用于 CPU、GPU 和机器学习加速器的开源机器学习编译器框架，旨在让机器学习工程师能够在任何硬件后端上高效地优化和运行计算。

从这个图中我们可以看到，TVM架构的核心部分就是Relay编译器。Relay编译器支持直接接收深度学习框架的模型，如TensorFlow/Pytorch/Caffe/MxNet等，同时也支持一些模型的中间格式如ONNX、CoreML。这些模型被Relay直接编译成Graph IR，然后这些Graph IR被再次优化，吐出优化后的Graph IR，最后对于不同的后端这些Graph IR都会被编译为特定后端可以识别的机器码完成模型推理。比如对于CPU，NNVM就吐出LLVM可以识别的IR，再通过LLVM编译器编译为机器码到CPU上执行。

TVM 环境配置

环境要求：GCC 7.1，CMake 3.18 或更高版本，推荐使用 LLVM 构建 TVM 库以启用所有功能，如需使用 CUDA，请确保 CUDA 工具包的版本至少在 8.0 以上，Python推荐使用3.8.X+版本。

git clone --recursive https://github.com/apache/tvm tvm
cd tvm
mkdir build
cp cmake/config.cmake build
cd build
cmake ..
make -j4
export TVM_HOME=/path/to/tvm
export PYTHONPATH=$TVM_HOME/python:${PYTHONPATH}

在此之前需要安装llvm，直接在官网上下载已经编译好的文件。需要开启 CUDA 等功能见官方文档修改相应的config.cmake。