TVM调研

相关术语

概念	描述
IR(Intermediate Representation)	神经网络的中间表示
Relay	一种函数式语言（functional language）和神经网络的中间表示（IR）
Relax	TVM的下一代图层级中间表示（IR）
TE(Tensor Expression)	张量表达式
TIR(Tensor Intermediate Representation)	张量级的中间表示，TVM 的低层次中间表示
TOPI(Tensor Operator Inventory)	张量算子清单
BYOC(Bring Your Own Codegen)	嵌入式和专用目标

TVM执行流程



1. 1. 导入模型：可以支持从tf，pytorch，或者onnx框架中导入模型。

2. 2. 模型转换到relay：relay是一个高级模型描述语言和IR，可以针对模型做些图级别的优化。

3. 3. 转换到TE：通过relay做过高级优化后，relay会用FuseOps pass去把模型分解成一个个子图，并且把子图转化到低阶表达TE中去。TE已经是一个描述张量计算的dsl。TE阶段也会提供一些调度原语去做更底层的循环优化，比如loop tile、vectorization、并行、loop unroll和loop fusion。tvm中还封装了一些常见的张量计算像conv2d、transpose这样。

4. 4. 自动优化调度：在tvm中，可以使用 AutoTVM 和 AutoScheduler 去搜索出一个张量计算的最佳优化调度方式，每个调度（schedule）会针对算子或者张量运算指定一种低级循环优化。自动优化共有两种方式：

• • AutoTVM：基于模板的 auto-tuning 模块。它运行搜索算法为用户定义的模板中的可调节旋钮找到最佳值。对于常见的运算符，其模板已经在TOPI中提供。

  • AutoScheduler：无模板的auto-tuning 模块。它不需要预先定义的 schedule 模板。相反，它通过分析计算的定义自动生成搜索空间。然后，它在生成的搜索空间中搜索最佳 schedule。

1. 5. 模型编译：选择最优的模型编译配置。通过上一步的优化调度搜索，各个子图最佳的优化方案会保存在一份JSON文件中。

2. 6. 转换到TIR：获取到最佳的优化调度空间后，TE的子图会被转化（lower）到TIR上面并且会执行更底层的循环优化。优化后的TIR会被转换到目标平台的编译器上。最后由编译器产生机器码，codegen包含以下编译器后端：

   1. LLVM：支持X86，arm，AMDPGU，nvptx

      NVCC：nv的编译器

                        BYOC：自定义编译器框架定义的嵌入式平台

1. 7. 编译到机器码：由编译器把code编译成机器码到目标平台运行。

TVM python API使用流程

使用TVMC Python流程

1. 环境导入：导入本地的TVM环境

import warnings
from tvm.driver import tvmc

warnings.filterwarnings('ignore')

2. 模型加载：将机器学习模型从受支持的框架转换为TVM的高级图表示语言Relay。支持的框架有Keras、ONNX、Tensorflow、TFLite 和 PyTorch

import onnx

model_path = 'resnet50-v2-7.onnx'
onnx_model = onnx.load(model_path)

model = tvmc.load(model_path, shape_dict={"data": [1, 3, 224, 224]}) # Step 1: Load
model.summary()    # 查看Relay

3. Tune(可选&推荐)：通过调优可以进一步提高运行速度，该过程是使用机器学习来查看模型（函数）中的每个运算，并试图找到更快的方法来运行它

# Step 3/1.5: Optional Tune
# 可以是 "cuda"
tvmc.tune(model, target="llvm")

4. 编译：将模型从Relay转换为目标机器可以理解的较低级语言

# Step 4: Compile
package = tvmc.compile(model, target="llvm", tuning_records="records.log")

5. 运行：编译后的包现在可以在硬件目标上运行，设备输入选项有：CPU、Cuda、CL、Metal 和 Vulkan

#Step 5: Run
results = tvmc.run(package, device="cpu")

参考链接

开始使用 TVMC Python：TVM 的高级 API — TVM 文档

https://www.cnblogs.com/whiteBear/p/16483688.html

NIO @ TVMCon 2023 | Chenfan Blog