一眼就懂的TVM7大功能

简介: 简明扼要的介绍TVM的7大功能

TVM是目前搞人工智能编译器很火的一个框架,他可以支持很好的后端,同时可以很好的做指令调度和优化。

对于没接触过tvm的同学来说,可能tvm是个黑盒,但其实,tvm的模块化设计很好,我们只需要了解tvm的几大模块化功能,就大致能对tvm有个初步的认识了。

本文主要介绍一下tvm框架的几大功能,希望有兴趣的同学可以对tvm有更深的认识。

TVM 优化编译框架

下面的图表展示了TVM作为一个优化编译框架,是如何一步步将机器模型转换为二进制的机器语言的。

image.png

上图共有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做的事情,他可以很有效的进行模型优化和部署。感兴趣的可以去官网下载代码试一试。

TVM官网:https://github.com/apache/tvm

相关文章
|
2月前
|
网络协议 NoSQL Anolis
开局一张图,debug全靠瞪|内核问题定位与静态分析实战
本文描述了一个在AnolisOS 8.8操作系统上遇到的内核崩溃问题的调试过程。
|
缓存 前端开发 JavaScript
浅浅阅读umi中InitialState插件源码 - 杨磊
InitialState插件源码的简要介绍
1011 0
浅浅阅读umi中InitialState插件源码 - 杨磊
|
存储 索引 容器
灰太狼系列—打地鼠(内含源码) inscode中的直观运行
灰太狼系列—打地鼠(内含源码) inscode中的直观运行
|
存储 并行计算 程序员
【CUDA学习笔记】第二篇:CUDA C并行化编程【上半部分】(附案例代码下载方式)
【CUDA学习笔记】第二篇:CUDA C并行化编程【上半部分】(附案例代码下载方式)
414 0
|
缓存 前端开发 JavaScript
浅浅阅读umi中InitialState插件源码
InitialState插件源码的简要介绍
659 1
浅浅阅读umi中InitialState插件源码
|
程序员 C++
超实用技巧,代码跟踪器
超实用技巧,代码跟踪器
超实用技巧,代码跟踪器
Verilog小练习- 看时序写代码--学习自测(完结篇)
Verilog小练习- 看时序写代码--学习自测(完结篇)
134 0
Verilog小练习- 看时序写代码--学习自测(完结篇)
|
机器人 开发工具 C++
探究--gazebo里 关节是如何动起来的____撸源码探究ros gazebo里的 硬件抽象构建
探究--gazebo里 关节是如何动起来的____撸源码探究ros gazebo里的 硬件抽象构建
探究--gazebo里 关节是如何动起来的____撸源码探究ros  gazebo里的 硬件抽象构建
|
前端开发 JavaScript
精读《如何编译前端项目与组件》
# 1 引言 说到前端编译方案,也就是如何打包项目,如何编译组件,可选方案有很多,比如: - 通过 webpack / parcel / gulp 构建项目。 - 通过 parcel / gulp / babel 构建组件。 如果你喜欢零配置的 parcel,那么项目和组件都可以拿它来编译。 如果你业务比较复杂,需要使用 webpack 做深度定制,那么常见组合是:项目
1061 0