1.为什么需要深度学习编译器
深度学习编译器主要为解决不同框架下训练的模型 部署到指定的某些设备上时所遇到的一系列复杂的问题,即将各种深度学习训练框架的模型 部署到各种硬件所面临的问题;
首先深度学习领域,从训练框架看,当前可选的框架有pytorch、TensorFlow、Mxnet、paddle,oneflow、caffe/caffe2、mindspore等,具体选择哪个,不尽相同,但如果项目要部署落地,则面临很多问题,即从推理框架角度来看,无论我们选择何种训练框架训练模型,我们最终都是要将训练好的模型部署到实际场景的,在模型部署的时候我们会发现我们要部署的设备可能是五花八门的,例如Intel CPU/Nvidia GPU/Intel GPU/Arm CPU/Arm GPU/FPGA/NPU(华为海思)/BPU(地平线)/MLU(寒武纪),如果我们要手写一个用于推理的框架在所有可能部署的设备上都达到良好的性能并且易于使用是一件非常困难的事。
为了解决上面的问题,科学家为编译器抽象出了编译器前端,编译器中端,编译器后端等概念,并引入IR (Intermediate Representation)的概念。解释如下:
编译器前端:接收C/C++/Java等不同语言,进行代码生成,吐出IR
编译器中端:接收IR,进行不同编译器后端可以共享的优化,如常量替换,死代码消除,循环优化等,吐出优化后的IR
编译器后端:接收优化后的IR,进行不同硬件的平台相关优化与硬件指令生成,吐出目标文件
因此我们可以将各个深度学习框架训练出来的模型看做各种编程语言,传入深度学习编译器,之后吐出IR,由于深度学习的IR其实就是计算图,所以可以叫做Graph IR,针对这些Graph IR可以做一些计算图优化在吐出IR分发给各种硬件使用,这样就解决了上述很多繁琐的问题,如下图所示:
深度学习编译器VS传统编译器
图中与传统编译系统进行了对比,发现存在相似之处:如都有中间表示、前端和后端分别硬件无关优化和硬件相关优化以及最后的代码生成模块;
同时也有一定的区别:深度学习编译系统是针对神经网络这一特定领域的编译系统,采用了以Python为主的动态解释器语言的前端、多层IR(如图层/算子层/codegen)设计、以及面向神经网络的特定优化(如自动微分、量化/混合精度、大规模并行、张量运算、循环优化等)。
深度学习编译器的架构
-
前端:将现有深度学习框架中的深度学习模型作为输入,然后将模型转换为计算图表示。前端需要实现各种框架不同格式的转换,并进行了结合了通用编译系统的优化技术和深度学习特定的优化技术的硬件无关的计算图优化,优化图逻辑,减少了冗余,提高了图IR的效率。计算图优化包括代数简化、算子融合、算子下沉、CSE公共子表达式消除、DCE死代码消除、静态内存规划和布局转换等。在前端之后,将生成优化的计算图并将其传递到后端。
-
中间表示IR:分布在前端和后端,是介于源语言和目标语言之间的程序表示,能够极大程度地提高编译流程的可拓展性,同时也能降低优化流程对前端和后端的破坏
-
后端:将高级IR转换为低级IR,并执行特定于硬件的计算图优化、算子选择以及内存分配这三个任务。计算图优化是在不影响模型的数值特性的基础上,通过图变换达到减少资源开销、适配硬件的执行能力、提升执行性能的目的。例如与硬件无关的算子内存 IO 优化和为了适配特定硬件指令限制而做的子图变换。数据存在多种存储格式和计算精度,不同的存储格式和计算精度在不同场景下对算子计算性能有较大的影响,算子选择是为 IR 图中的每个计算节点选择一个最适合在设备上执行的算子。经过计算图优化和算子选择以及常用的特定于硬件的优化包括硬件内部映射、内存分配和获取、内存延迟隐藏、并行化以及面向循环的优化后。使用JIT或AOT编译优化的低级IR,以生成不同硬件目标的代码,运行时要实现算子选择和内存分配等技术。
深度学习编译系统整体架构图
此外,对应深度学习工作流,完整的深度学习编译框架还包括:编程接口、硬件加速器、数据处理、模型部署、分布式训练等模块。深度学习编译系统以加速器为核心,硬件加速器模块提供丰富的编程接口;数据处理模块负责数据的读取、存储以及预处理;分布式训练模块为应对单个机器上内存及算力资源不足,并行加速模型训练的过程;模型部署模块负责进行模型压缩、针对推理硬件平台模型算子优化、进行模型混淆设计保证模型安全等。
中间表示
<