简介(未完成)
The Deep Learning Compiler: A Comprehensive Survey
【从零开始学深度学习编译器】十一,初识MLIR 未读。 Dive into Deep Learning Compiler摘要(1)
TVM 自底向上(一):基本框架和概念 未读 TVM 自底向上(二):TIR 的概念和编译原理 未读
第一视角:深度学习框架这几年 Imtermediate Representation+Pass的模式主要是从LLVM的架构上借鉴来的。在编译器上主要是用来解决把M个编程语言中任意一个编译到N个硬件设备中任意一个执行的问题。简单的解决方案是为每个编程语言和硬件单独写一个编译器。这需要M*N个编译器。显然这对于复杂的编译器开发来说,是非常高成本的。
Intermediate Representation是架构设计中抽象能力的典型体现。不同编程语言的层次不一样,或者仅仅是单纯的支持的功能有些差异。但是,这些编程语言终归需要在某种硬件指令集上执行。所以在编译的过程中,他们会在某个抽象层次上形成共性的表达。而IR+Pass的方法很好的利用了这一点。其基本思想是通过多层Pass (编译改写过程),逐渐的把不同语言的表达方式在某个层次上改写成统一的IR的表达方式。在这个过程中,表达方式逐渐接近底层的硬件。而IR和Pass可以很好的被复用,极大的降低了研发的成本。深度学习框架也有着非常类似的需求。
- 用户希望通过高层语言描述模型的执行逻辑,甚至是仅仅声明模型的结构,而不去关心模型如何在硬件上完成训练或者推理。
- 深度学习框架需要解决模型在多种硬件上高效执行的问题,其中包括协同多个CPU、GPU、甚至大规模分布式集群进行工作的问题。也包括优化内存、显存开销、提高执行速度的问题。
更具体的。前文说到需要能够自动的将用户声明的模型Program自动的在多张显卡上并行计算、需要将Program拆分到多个机器上进行分布式计算、还需要修改执行图来进行算子融合和显存优化。Paddle在一开始零散的开展了上面描述的工作,在分布式、多卡并行、推理加速、甚至是模型的压缩量化上各自进行模型的改写。这个过程非常容易产生重复性的工作,也很难统一设计模式,让团队不同的研发快速理解这些代码。意识到这些问题后,我写了一个Single Static Assignment(SSA)的Graph,然后把Program通过第一个基础Pass改写成了SSA Graph。然后又写了第二个Pass把SSA Graph改写成了可以多卡并行的SSA Graph。后面的事情就应该可以以此类推了。比如推理加速可以在这个基础上实现OpFusionPass, InferenceMemoryOptimizationPass, PruningPass等等,进而达到执行时推理加速的目的。分布式训练时则可以有DistributedTransPass。量化压缩则可以有ConvertToInt8Pass等等。这一套东西基本解决了上层Program声明到底层执行器的Compiler问题。
随着项目的复杂化,很多棘手的问题逐渐从深度学习的领域技术问题转变成了软件工程开发和团队管理分工的问题。随着团队的不断变化,自己有时候是作为一个leader的角色在处理问题,有的时候又是以一个independent contributor的角色在参与讨论。很庆幸自己经历过这么一段,有些问题在亲身经历后才能想得明白,想得开。时代有时候会把你推向风口浪尖,让你带船队扬帆起航,在更多的时候是在不断的妥协与摸索中寻找前进的方向。
《AI编译器开发指南》
- Python静态化,Python几乎已经成了AI编程的首选语言,但是Python本身存在性能差、部署场景受限等问题。Python静态化是解决Python语言的限制的重要手段。为此,AI框架的核心任务是将Python构建的AI模型转换为计算图的形式,下发给硬件执行。按照构建计算图的方式,AI框架分为了AOT和JIT两种编译方式:
- AoT(Ahead of Time)编译的模式,在程序执行前先进行构图和编译优化,在整图的基础上生成反向计算图。由于能进行全局的编译优化和整图下沉执行,静态图适合大规模部署,适合挖掘硬件的性能,但是编程和调试体验较差。
- JIT则是边执行边构图,一般通过Tracing的方式实现自动微分。动态图更符合算法开发人员的编程调试习惯,更好的兼容Python生态,但是大多数场景下动态图性能比静态图差,并且部署环境限制较多。
- Transformer时代,编译优化还是大kernel?在Transformer架构出现之后,图层的结构相对而言变简单了,不再有之前的LSTM、RNN等复杂结构,其很大程度上把大量模型结构进行了统一。因此,很多人也在探讨,当模型固定了以后,是不是可以简化算子层,只用一个大kernel,而不需要再通过编译器,同时还能进一步提升性能?
- 虽然Transformer的基础结构是统一的,但是在不同的场景中会有不同的变种。
- 我们无法把整图都做成一个大kernel,所以非大kernel部分还是需要进行编译优化的。
- Transformer结构在未来是否会被替代,目前也没有定论。
示例
torch.compile、TinyGrad和ONNX这样的编译器,可以将简单的Python代码融合成为针对你的硬件进行优化的kernel。例如,我可以编写如下函数:
s = torch.sin(x)
c = torch.cos(x)
return s + c
简单来说,这个函数需要:
- 为 s 分配 x.shape()大小的内存
- 对 x 进行线性扫描,计算每个元素的 sin 值
- 为 c 分配x.shape()大小的内存
- 对 x 进行线性扫描,计算每个元素的 cos 值
- 为结果张量分配 x.shape() 大小的内存
- 对 s 和 c 进行线性扫描,将它们相加到结果中
这些step都比较慢,且其中一些step需要跳过Python和本地代码间的界线,不利于加速。那如果我使用torch.compile 来编译这个函数会怎样呢?torch 就为CPU生成了一个经过优化的C++ kernel,将foo融合成一个单一的kernel。(如果用GPU运行这个kernel,torch将为GPU生成一个CUDA kernel。)接下来的step是:
- 为结果张量分配 x.shape() 大小的内存
- 对 x (in_ptr0)进行线性扫描,计算sin和cos值,并将它们相加到结果中