TVM是目前搞人工智能编译器很火的一个框架,他可以支持很好的后端,同时可以很好的做指令调度和优化。
对于没接触过tvm的同学来说,可能tvm是个黑盒,但其实,tvm的模块化设计很好,我们只需要了解tvm的几大模块化功能,就大致能对tvm有个初步的认识了。
本文主要介绍一下tvm框架的几大功能,希望有兴趣的同学可以对tvm有更深的认识。
TVM 优化编译框架
下面的图表展示了TVM作为一个优化编译框架,是如何一步步将机器模型转换为二进制的机器语言的。
上图共有7步,下面每个序号标题对应上图的每个步骤,会详细的说明每个步骤都干了什么。
TF/Pytorch/ONNX
从Tensorflow, pytorch, Onnx 这种框架中导入模型。可以将模型从上述框架中将模型导入到TVM。需要注意的是,TVM对于每一个前端(框架)的支持是会随着开源计划的发展而变化的。如果你在导入模型到TVM过程中发现有问题,那么你可能需要尝试将模型转成ONNX。
Relay(High-level IR)
这一步将上一步的模型转成Relay, Relay是TVM的high-level模型语言。当一个模型被导入到TVM之后,就会被用Relay来表示,对神经网络来说,Relay是一种函数式的语言和中间表示,它支持:
传统的数据流表示
- Functional-style scoping,支持let表达式,语言特性丰富,同时融合了不同的其他特性,开发者可以使用上述两种编程方式进行混合编程
- relay 提供了图级别的优化pass来优化模型
TE(computation defination)
这一步进一步将上层的表示lower(层次下降, 注1)到TE(Tensor Expression)表示。
lower指的是将高层次的表示转换为低层次的表示。经过高层的优化后,Relay会运行一个融合算子(FuseOps)的pass 来将模型进行拆分,拆分成很多的小的子图,同时,将小的子图进一步层次下降到 TE表示。TE用于描述tensor计算,它提供了多种调度原语(schedule primitives)来进行低层次(low-level)的优化,比如拆分(tiling),向量化(vectorization),并行化(parallelization), 循环展开(unrolling),融合(fusion)等。
为了更好的将Relay的表示转换为TE的表示,TVM还提供了一种叫做TOPI(Tensor Operator Inventory)的通用算子模板库,它包含了很多已经定义好的算子模板(比如conv2d, transpose等)。
auto scheduler
使用auto-tuning机制如AutoTVM或AutoScheduler来搜索最佳的调度。调度指的是对于在TE中定义的一个算子(Operator)或者子图(subgraph),完成一些low-level的循环优化。在TVM中有两种auto-tuning模型。
AutoTVM: 这是一种基于模板开发的auto-tuning模型。在使用时,用户需要定义一个模板,AutoTVM会运行搜索算法来寻找最佳的值用于匹配可以替换的(tunable)的节点。对于一般的通用算子而言,他们的模板在TOPI中已经提供了。
AutoScheduler (a.k.a. Ansor)
这是一种没有模板的auto-tuning模型。它不需要预先定义好调度模板,相反的是,他会通过分析用户需要的计算(比如一个resnet50网络),来自动生成一个搜索空间,然后,在这个搜索空间中自动搜索并且获得一个最佳的调度。
TE + scheduler
这一步选择最佳的配置来进行模型编译。经过上面提到的tuning过程之后,会生成一个tuning的结果,以JSON的格式存储信息,这个步骤会为每一个子图选择最佳的调度。
TIR(low-level IR)
这一步,lower到TIR(Tensor IntermediateRepresentation)层,这是TVM的low-level的中间表示。经过了上面提到的tuning之后,TVM已经选择了最佳的配置(调度,schedule),这时,每一个TE子图会使用TIR表示,同时会被low-level的优化pass所优化。紧接着,优化后的TIR会进一步层次下降(lowered)到硬件平台的目标编译器上,这是最终生成优化模型代码的过程,这个代码可以部署到相关的设备上。TVM提供了多种不同的编译器后端来进行代码的编译和优化:
- LLVM,它可以运行在很多微处理器架构上,比如PC端的标准x86 和ARM 处理器,AMDGPU 和 NVPTX 以及其他的LLVM支持的处理器平台。
- 特殊编译器,如NVCC, 它是英伟达的编译器.
- 嵌入式特殊架构,可以通过TVM的Bring Your Own Codegen (BYOC) 框架来适配以支持。
机器码
编译成机器码(二进制)。在最后这一步,编译器会根据指定好的硬件平台编译出对应的二进制机器码。
TVM可以把模型编译成可链接的模型,可以和TVM的runtime提供的C接口配合使用,比如loadmodule等等。TVM也可以把所有的东西包括runtime都打包好成一个包,然后部署到相关的产品上。
上面的几个步骤,大概就是TVM做的事情,他可以很有效的进行模型优化和部署。感兴趣的可以去官网下载代码试一试。