简介(未完成)
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结构在未来是否会被替代,目前也没有定论。
编译器简介
- 前端优化,AI 编译器分为多层架构,最顶层由各种 AI 训练框架编写的神经网络模型架构,一般由 Python 编写,常见的 AI 训练框架有 PyTorch、MindSpore、PaddlePaddle 等。在导入 AI 编译器时需要用对应框架的 converter 功能转换为 AI 编译器统一的 Graph IR,并在计算图级别由 Graph Optimizer 进行计算图级优化,也叫前端优化。前端优化主要的计算图优化包括图算融合、数据排布转换、内存优化、死代码消除,这些优化是硬件无关的通用优化。前端优化针对计算图整体拓扑结构优化,不关心算子的具体实现。主要优化流程为对算子节点进行融合、消除、化简,使得计算图的计算和存储开销最小。
- 在得到优化后的计算图后,将其转换为TensorIR,送入OpsOptimizer进行算子级优化,也叫后端优化,这类优化是硬件相关的,主要包括循环优化、算子融合、tiling、张量化。针对单个算子的内部具体实现优化,使得算子的性能达到最优。主要优化流程为对算子节点的输入、输出、内存循环方式和计算逻辑进行编排与转换。在算子级优化结束后,即进入代码生成阶段。后端优化的流程一般分为三步:
- 生成低级 IR:将高级或计算图级别 IR(Graph IR)转换为低级 IR(Tensor IR)。
- 进行后端优化,并将 IR 转换为更低级的 IR。针对不同的硬件架构/微架构,不同的算法实现的方式有不同的性能,目的是找到算子的最优实现方式,达到最优性能。同一算子不同形态如 Conv1x1、 Conv3x3、 Conv7x7 都会有不同的循环优化方法。实现方式多种多样,可以人工凭借经验手写算子实现,也可以通过自动调优搜索一个高性能实现。
- 代码生成:根据硬件进行代码生成。对优化后的低级 IR 转化为机器指令执行,现阶段最广泛的做法为借助成熟的编译工具来实现,代码生成不是 AI 编译器的核心内容。如把低级 IR 转化成为 LLVM、NVCC 等编译工具的输入形式,然后调用其生成机器指令。
算子优化的挑战:算子根据其计算形式的特点可分为访存密集型与计算密集型。
- 访存密集(Memory-Bound)型。指的是在执行过程中主要涉及大量内存读取和写入操作的计算任务。这类算子通常需要频繁地从内存中读取数据,执行一些简单的计算操作,然后将结果写回内存。访存密集型算子的性能受限于内存带宽和访问延迟,而不太受计算能力的限制。如 RNN 训练任务,其网络结构的计算密度很低,因此瓶颈转移到 host 端的 Op Launch 上,算子的计算 kernel 之间出现大量空白。
- 计算密集(Compute-Bound)型。指的是在执行过程中主要涉及大量的计算操作,而对内存的访问相对较少的计算任务。这类算子主要依赖于 CPU 或 GPU 的计算能力,并且往往对内存带宽和访问延迟的需求不是特别高。一些数值计算密集型的算法,比如矩阵乘法、卷积运算、复杂的数学函数计算等,通常被认为是计算密集型的操作。 【AI系统】算子手工优化
示例
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值,并将它们相加到结果中