深度学习领域的“Linux”——TVM编译器
深度学习领域的“Linux”——TVM编译器
“IT有得聊”是机械工业出版社旗下IT专业资讯和服务平台，致力于帮助读者在广义的IT领域里，掌握更专业、更实用的知识与技能，快速提升职场竞争力。 点击蓝色微信名可快速关注我们!
人工智能（Artificial Intelligence，AI）已经在全世界信息产业中获得广泛应用。深度学习模型推动了AI技术革命，如 TensorFlow、PyTorch、MXNet、Caffe等。大多数现有的系统框架只针对小范围的服务器级 GPU进行过优化，因此需要做很多的优化努力，以便在汽车、手机端、物联网设备及专用加速器（FPGA、ASIC）等其他平台上部署。随着深度学习模型和硬件后端数量的增加，TVM构建了一种基于中间表示 (IR)的统一解决方案。TVM不仅能自动优化深度学习模型，还提供了跨平台的高效开源部署框架。有了TVM的帮助，只需要很少的定制工作，就可以轻松地在手机、嵌入式设备甚至浏览器上运行深度学习模型。TVM 还为多种硬件平台上的深度学习计算提供了统一的优化框架，包括一些有自主研发计算原语的专用加速器。TVM是一个深度学习编译器，所有人都能随时随地使用开源框架学习研发。围绕TVM形成了多元化社区，社区成员包括硬件供应商、编译器工程师和机器学习研究人员等，共同构建了一个统一的可编程软件堆栈，丰富了整个机器学习技术生态系统。TVM是一个新型的AI编译器，广泛应用于各种产品研发中，在企业与学术研究中有很大的影响。但是，目前市面上有关TVM的书还很少，长期从事人工智能芯片设计的吴建明博士，尤其擅长TVM/LLVM编译器，由他编写的《TVM编译器原理与实践》填补了TVM技术领域图书空缺。
《TVM编译器原理与实践》
作者：吴建明 吴一昊
在本文中，我们截取书中的部分内容，将大家比较关注的TVM基本原理进行简要阐述。
TVM是一个开源的、端到端的机器学习编译器框架，旨在将机器学习算法自动编译成可供下层硬件执行的机器语言， 从而利用多种类型的算力。也就是将深度学习模型进行高效推理、内存管理与线程调度优化后，再通过LLVM部署在CPU、GPU、FPGA、ARM等硬件设备上。
TVM的整体架构
TVM是将深度学习模型编译为针对目标硬件的可执行模块。TVM堆栈是一个端到端的编译堆栈，用于将深度学习工作负载部署到各种硬件设备上。TVM可以表述为一种把深度学习模型分发到各种硬件设备上的、端到端的解决方案。也可以使用 TVM工程中的LLVM后端代码生成器，直接将深度学习框架编译为在目标设备上运行的裸机代码。编译器支持直接接收AI深度学习框架的模型，如TensorFlow、Pytorch、Caffe、MxNet等，同时也支持一些模型的中间格式，如ONNX、CoreML( 全称Core Machine Learning，机器学习核心) 等。TVM Relay直接将这些模型编译成Graph IR( Intermediate Representation，中间表示)，同时优化这些Graph IR， 再输出优化后的Graph IR，然后编译成特定后端，最后生成可以识别的机器码， 完成模型推理部署。如在CPU上，TVM Relay会输出LLVM可以是别的IR，再通过LLVM编译为 CPU上可执行的机器码。现有的AI 深度学习框架主要有以下几种：1.Google的TensorFlow框架。2.Facebook的Pytorch框架。3.亚马逊公司的MxNet框架.4.百度的Paddle框架。5.旷视科技的MegEngine框架。6. 华为的Mindspore计算平台。7. 一流科技的OneFlow框架。图1所示为TVM整体框架。
图1.TVM整体框架下面对TVM整体框架进行简单解析:
1.输入 (load) 不同的框架模型 (Model) 文件，使用Relay将Model文件转换成IR 计算图，
随后进行优化计算，如算子融合、剪枝、转换、量化等。
2.TVM优化调度张量运算，将代码调度与计算进行分离。
3.生成不同后端设备的相应代码，用LLVM生成包括如ARM、JavaScript、x86等系统代码，或用OpenCL、Metal与CUDA等生成设备代码。通过这种IR堆栈表示，能对深度学习模型进行端到端优化，将算子转换成符合编译规则的文件。
4.在高级图IR( Intermediate
Representation，中间表示) 中优化常见的深度学习工作负载。
5.转换计算图，以便进行内存最小化与数据布局优化，并对不同硬件后端进行融合。
6.提供从前端深度学习框架到裸机硬件的端到端编译的pipeline(管道)。
TVM与人工编译器框架比较
在TensorFlow等现有框架使用的传统工作流程中，GPU （Graphics Processing Unit) 能对深度神经网络 (矩阵乘法与卷积) 中的原始运算符进行图形渲染，同时进行性能优化。
另外，可以采用基于TVM编译的方法。TVM自动从TensorFlow、Keras、Pytorch、MXNet及ONNX等高级框架中提取模型，先使用机器学习方法自动生成编译代码，再使用SPIR-V格式计算算子与图形渲染，最后将生成的代码打包成可部署的模块。SPIR-V是描述Vulkan使用的着色器语言，SPIR-V的格式是二进制表示， SPIR-V函数以控制流图 (Control Flow Graph，CFG) 的形式存在，SPIR-V数据结构保留了高级语言的层级关系。
基于编译方法的一个显著优点是基础架构的重用。通过重用基础架构，优化原生平台 ( 如CUDA、Metal和OpenCL) 的GPU内核，以便支持网络模型。如果GPU API到本机API的映射是有效的，只需很少的定制工作， 就可以得到类似的性能。AutoTVM基础架构允许为特定模型定制计算算子与图形渲染，并能生成最佳性能。图2为人工编译器框架 (SPIR-V) 与TVM对比。
图2.人工编译器框架 (SPIR-V) 与TVM对比
TVM已经提供了SPIR-V生成器，用于算子与图形渲染，并使用LLVM生成设备及主机程序。TVM已经有VulKan的SPIR-V目标，使用LLVM生成主机代码。可以仅将二者重新生成设备和主机程序，主要挑战是运行时，需要一个运行时来加载着色器代码，并使主机代码能够正确地与着色器通信。TVM具有最低的基于C++的运行时。构建了一个最小的Web运行时库，将生成的着色器和主机驱动代码链接，并生成一个WASM文件。
第一代TVM
第一代计算图称为NNVM (Neural Network Virtual Machine)。NNCM是亚马逊公司和华盛顿大学合作发布的开源的、端到端的深度学习编译器，支持将包括 TensorFlow、Pytorch、MCNet、Caffe2、CoreML等在内的深度学习模型编译部署到硬件上，并提供多级别的优化。图3所示为第一代TVM NNVM的演变。NNVM编译器应用图级和张量级优化，以获得最佳性能。采用与现有深度学习框架不同的方法， 后者将图形优化与部署运行时打包在一起。NNVM编译器采用了编译器的传统知识， 将优化与实际部署运行时分开。这种方法提供了实质性的优化，但仍使运行时保持轻量级。编译后的模块仅取决于最小的TVM运行时，部署在Raspberry Pi 或移动设备上时仅需300KB左右。
图3.第一代TVM NNVM的演变
但是，NNVM只支持静态图计算， 而且还存在以下几个缺陷：
１.对分支跳转、循环等控制流支持不力。
２.对计算图输入图支持不力，如word2vec (Word2vec是一种将单词转为向量的方法) 等。
第二代TVM
第二代TVM计算图称为Relay，它的图计算层变为Relay VM。Relay和NNVM的主要区别是，Relay IR除了支持dataflow (静态图)，还能够更好地处理controlflow(动态图)。Relay不仅是一种计算图的中间表示，也支持自动微分。静态图的风格是先定义后执行，动态图的风格是计算定义与执行不分开，进行实时计算。图4所示为第二代TVM Relay的演变。
图4.第二代TVM Relay的演变
第二代TVM Relay架构的主要特点为：
１.最高层级支持主流的深度学习前端框架， 如TensorFlow、MXNet、Pytorch等。
２)Relay IR支持可微分，该层级进行图融合、数据重排等图优化操作。
３) 基于tensor张量化计算图，并根据后端进行硬件原语级优化，AutoTVM根据优化目标探索搜索空间，找到最优解。
４) 后端支持ARM、GPU、CPU、NPU、SOC、ASIC、FPGA等不同设备平台进行编译运行。
《TVM编译器原理与实践》
作者：吴建明 吴一昊
撰 稿 人：计旭
责任编辑：张淑谦
审 核 人：曹新宇
