AI编译器 - 小记


比较不错的教程


问题

  • AI 编译器跟传统编译器有什么区别吗?
  • AI 框架跟AI编译器什么关系?是AI框架包含AI编译器 or AI框架就是一个AI编译器?
  • AI 领域真的需要编译器吗?那为什么 PyTorch 动态图模式没有编译器的概念?
  • 技术投入比来看,神经网络编译器和人工算子实现,哪个性价比更高?


为什么需要AI 编译器

在这里插入图片描述


LLVM Enables Collaboration & Reuse
在这里插入图片描述


在这里插入图片描述

在这里插入图片描述


Challenge

越来越多新算子被提出,算子库的开发、维护、优化和测试工作量指数上升;
专用加速芯片爆发导致性能可移植性成为一种刚需;


Operator

  • 越来越多新算子被提出,算子库的开发、维护、优化和测试工作量指数上升:
    • 增加新算子,硬件不仅需要实现,还需要结合硬件进行特性优化和测试,尽量充分发挥硬件性能。以 Convolution 运算为例,需要将 Convolution 操作转换为 GEMM 矩阵乘法;算法新提出新的 Swish 算子,硬件需要新增 Swish 对应实现。
    • 硬件供应商还会有针对性发布优化库(如 MKL-DNN 和 CuDNN)。但是对于专用硬件,需要提供开发类似的优化库,不仅会增加大量算子优化、封装的工作,还会过于依赖库无法有效利用专用硬件芯片能力。

Optimization

  • 专用加速芯片爆发导致性能可移植性成为一种刚需:
    • 大多数 NPU 使用 ASIC,在神经网络场景对计算、存储和数据搬运 做了特殊的指令优化,使得对 AI 相关的计算会提升性能,如 NVIDIA 虽然作为 GPGPU,但是 DGX 系列提供专用的 Tensor Code。
    • 不同厂商提供 XPU 的 ISA%2C%20is%20called%20an%20implementation)(Instruction Set Architecture)千奇百怪,一般缺乏如 GCCLLVM 等编译工具链,使得针对 CPU 和 GPU已有的优化算子库和针对语言的优化 Pass 很难短期移植到 NPU 上。

传统编译器和 AI编译器

  • 目标相同: 通过自动化方式进行程序优化和代码生成,从而降低对不同硬件的手工优化;
  • 优化方式类似: 在编译优化层通过统一 IR 执行不同的Pass进行优化,从而提高执行性能;
  • 软件结构栈类似: 分成前端、优化、后端三段式,IR 解耦前端和后端使得模块化表示;
  • AI 编译器依赖传统编译器: AI编译器对 Graph IR 进行优化后,将优化后的 IR 转换成传统编译器 IR,最后依赖传统编译器针进行机器码生成。


差别:

在这里插入图片描述


在这里插入图片描述


传统编译器:

  • 输入高级语言输出低级语言
  • 主要问题是降低编程难度,其次是优化程序性能

AI 编译器

  • 输入计算图/算子,输出低级语言
  • 主要优化程序性能,其次是降低编程难度

在这里插入图片描述


IR 差异

IR 差异:AI 编译器的 IR 与传统编译器的IR所抽象出来的概念和意义并不相同。

  • AI编译器一般会有 high-level IR,用来抽象描述深度学习模型中的运算,如:Convolution、Matmul 等,
    甚至部分会有 Transformer 带有图的结构。
  • 传统编译器相对而言 low-level IR,用于描述基本指令运算,如 load、store 等。有了high-level IR,AI编
    译器在描述深度学习模型类 DSL 更加方便。

优化策略差异

AI 编译器面向AI领域,优化时引入更多领域特定知识,从而进行更 high-level,更加
aggressive 优化手段。如:

  • AI编译器在 high-level IR 执行算子融合,传统编译器执行类似 loop fusion 时候,往往更加保守。缺点是可能会导致调试执行信息跟踪难;
  • AI编译器可以降低计算精度,比如int8、fp16、bf16等,因为深度学习对计算精度不那么敏感。但传统编译器一般不执行改变变量类型和精度等优化。


什么是AI编译器


推理场景:输入 AI 框架训练出来的模型文件,输出能够在不同硬件高效执行的程序;
训练场景:输入高级语言表示的神经网络代码,输出能够在不同硬件高效执行的程序;


问题
什么是训练场景?什么是推理场景吗?
搞 AI 编译器为什么要了解算法呢?
搞 AI 算子为什么要了解编译器?



  • Python 为主的动态解释器语言前端
  • 多层 IR 设计,包括图编译、算子编译、代码生成
  • 面向神经网络、深度学习的特定优化
  • DSA 芯片架构的支持

在这里插入图片描述



AI编译器对比

在这里插入图片描述


XLA:优化机器学习编译器

XLA(加速线性代数)是一种针对特定领域的线性代数编译器,能够加快 TensorFlow 模型的运行速度,而且可能完全不需要更改源代码。

XLA:图层下发子图中的算子打开成小算子,基于小算子组成的子图进行编译优化,包括 buffer fusion、水平融合等,关键是大算子怎样打开、小算子如何重新融合、新的大算子如何生成,整体设计主要通过HLO/LLO/LLVM IR 实现,所有 Pass 规则都是手工提前指定。

在这里插入图片描述


TVM:端到端深度学习编译器

为了使得各种硬件后端的计算图层级和算子层级优化成为可能,TVM 从现有框架中取得 DL 程序的高层级表示,并产生多硬件平台后端上低层级的优化代码,其目标是展示与人工调优的竞争力。

TVM:分为Relay和TVM两层,Relay关注图层,TVM关注算子层,拿到前端子图进行优化,Relay关注算子间融合、TVM关注新算子和kernel生成,区别在于TVM 开放架构,Relay目标是可以接入各种前端,TVM也是一个可以独立使用的算子开发和编译的工具,算子实现方面采用 Compute(设计计算逻辑)和 Schedule(指定调度优化逻辑)分离方案。

在这里插入图片描述


TENSOR COMPREHENSIONS:神经网络语言编译器

TensorComprehensions 是一种可以构建 just in time(JIT)系统的语言,程序员可通过该语言用高级编程语言去高效的实现 GPU 等底层代码。
TC:算子计算逻辑的较容易实现,但 Schedule 开发难,既要了解算法逻辑又要熟悉硬件体系架构,此外图算边界打开后,小算子融合后,会生成新算子和 kernel,新算子 Schedule 很难生成,传统方法定义 Schedule 模板;TC 希望通过 Polyhedral model 实现 auto schedule。


在这里插入图片描述


nGraph:兼容所有框架的深度学习系统编译器

nGraph 的运作应作为深度学习框架当中更为复杂的 DNN 运算的基石,并确保在推理计算与训练及推理计算之间取得理想的效率平衡点。


在这里插入图片描述


AI 编译器的发展


Stage I 朴素的AI编译器

TensorFlow 早期版本,基于神经网络的编程模型,主要进行了graph 图和ops 算子两层抽象。

  • 图层:通过声明式的编程方式,以静态图方式执行,执行前进行硬件无关和硬件相关的编译优化。硬件无关的优化,如表达式化简、常量折叠、自动微分等;硬件相关的优化包括算子融合、内存分配等。
  • 算子层:通常采用手写 kernel 的方式,如在 NVIDIA GPU 上基于 CUDA kernel 实现大量的 .cu 算子或者依赖于 CuDNN 算子优化库。

表达上:

  • 静态图的表达式非 Python 原生,开发者主要通过框架提供Python API 显示构图,易用性上不好;

性能上:

  • DSA 专用加速芯片出现加剧了性能上的挑战;
  • 算子层提供的算子粒度和边界提前确定后,无法充分发挥硬件的性能;
  • 硬件厂商的提供的算子优化库也未必最优
    • 1)模型和 shape 确定情况下,可能还有更优算子实现;
    • 2)在 SIMT 和 SIMD 架构中,Scheduling、Tilling 都有有很大的空间。

在这里插入图片描述


Stage II 专用的AI编译器

在这里插入图片描述


表达上:以 PyTorch 为标杆的表达转换到计算图层 IR 进行优化。
性能上:打开计算图和算子的边界,进行重新组合优化,发挥芯片的算力。


表达上:

  • PyTorch 灵活表达 API 方式成为 AI 框架参考标杆,图层的神经网络编译器主要就是考虑如何把类 PyTorch 的表达转换到图层的IR进行优化。
  • 类PyTorch的Python原生表达,静态化转换;
  • AI专用编译器架构,打开图算边界进行融合优化;

性能上:
打开计算图和算子的边界,进行重新组合优化,发挥芯片的算力。计算图层下发子图中的算子打开成小算子,基于小算子组成的子图,进行编译优化,包括buffer fusion、水平融合等,关键是大算子怎样打开、小算子如何重新融等。


  • 表达分离:计算图层和算子层仍然分开,算法工程师主要关注图层的表达,算子表达和实现主要是框架开发者和芯片厂商提供。
  • 功能泛化:对灵活表达上的动静态图转换、动态 Shape、稀疏计算、分布式并行优化等复杂的需求难以满足。 平衡效率和性能:算子实现上在
  • Schedule、Tiling、Codegen 上缺乏自动化手段,门槛高,开发者既要了解算子计算逻辑,又要熟悉硬件体系架构。

Stage III 通用AI编译器

  • 图算统一表达,实现融合优化
  • 算子实现上自动 Schedule、Tiling、Codegen,降低开发门槛
  • 更泛化优化能力,实现动静统一、动态 Shape、稀疏性、高阶微分、自动并行等
  • 包括编译器、运行时,异构计算、边缘到数据中心都模块化表示和组合,并专注于可用性

现有 AI 编译器架构

表达上:以 PyTorch 为标杆的表达转换到计算图层 IR 进行优化。
性能上:打开计算图和算子的边界,进行重新组合优化,发挥芯片的算力。


在这里插入图片描述


IR 中间表达

编译器主要分为前后端,分别针对于硬件无关和硬件相关的处理。每一个部分都有自己的 IR (Intermediate Representation,中间表达),每个部分也会对进行优化:

  • High-level IR:用于表示计算图,其出现主要是为了解决传统编译器中难以表达深度学习模型中的复杂运算这一问题,为了实现更高效的优化所以新设计了一套 IR。
  • Low-level IR:能够在更细粒度的层面上表示模型,从而能够针对于硬件进行优化,文中将其分为了三类。

Frontend 前端优化

构造计算图后,前端将应用图级优化。因为图提供了计算全局概述,所以更容易在图级发现和执行许多优化。前端优化与硬件无关,这意味着可以将计算图优化应用于各种后端目标。前端优化分为三类:
1、节点级优化,如 Zero-dim-tensor elimination、Nop Elimination
2、块级优化,如代数简化、常量折叠、算子融合
3、数据流级优化,如Common sub-expression elimination、DCE


Backend 后端优化

特定硬件的优化
目标针对特定硬件体系结构获取高性能代码。1)低级IR转换为LLVM IR,利用LLVM基础结构生成优化的CPU/GPU代码。2)使用领域知识定制优化,这可以更有效地利用目标硬件。


自动调整
由于在特定硬件优化中用于参数调整的搜索空间巨大,因此有必要利用自动调整来确定最佳参数设置。1)Halide/TVM允许调度和计算表达分开,使用自动调节来得出较佳配置。2)应用多面体模型 Polyhedral model 进行参数调整。


优化内核库
厂商特定优化内核库,广泛用于各种硬件上的加速DL训练和推理。特定优化原语可以满足计算要求时,使用优化的内核库可显著提高性能,否则可能会受到进一步优化的约束。




summary

naive

  • 以计算图和算子抽象为主
  • 计算图中采用部分编译器技术

specific

  • 类PyTorch的Python 原生表达,静态化转换
  • AI专用编译器架构,打开图算边界进行融合优化

universal

  • 图算统一表达,使用融合优化
  • 算子自动生成,降低开发门槛
  • 针对神经网络,泛化优化能力
  • 模块化表示组合,提升可用性

cons

  • API构图,易用性差
  • 大量DSA异构芯片对性能挑战
  • 算子边界确定无法充分发挥性能

  • 图算表达分开表达
  • 神经网络的功能泛化
  • 算子实现 schedule 等缺乏自动化

在这里插入图片描述



AI 编译器的输入:计算图

计算图的产生:AI框架

在这里插入图片描述


AI 编译器的挑战

  • 动态 Shape问题
  • Python 编译静态化
  • 发挥硬件性能,特别是DSA类芯片
  • 处理神经网络特性:自动微分、自动并行等
  • 易用性与性能兼顾


Challenge I: Dynamic shape

动态 shape 和动态计算图:

  • 现阶段主流的AI编译器主要针对特定静态 shape 输入完成编译优化,对包含 Control Flow 语义的动态计算图只能提供有限的支持或者完全不能够支持。
  • AI 的应用场景却恰恰存在大量这一类的任务需求。AI编译器前端将计算图改写为静态计算图,或者将适合编译器部分子图展开交给编译器进行优化。
  • 部分 AI 任务并不能通过人工改写来静态化(如金字塔结构的检测模型),导致这些情况下编译器完全难以有效优化。

Python Static 静态化

通过 JIT,Python 程序执行静态编译优化,从而提升性能、方便部署。业界 JIT 通用做法:
1、Python JIT 虚拟机:期望在 Python 解释执行的基础上增加 JIT 编译加速的能力,典型的如 PyPy 或者 CPython,但由于前期设计问题 Python JIT 虚拟机兼容难。
2、修饰符方式:Python JIT 通过修饰符,进行部分 Python 语句加速。但无论是 torch.jit.script、torch.jit.trace、torch.fx、LazyTensor、TorchDynamo,几乎把能用的手段都用上,还是缺乏一个完备的方案。


Python 执行时,首先会将.py文件中的源代码编译成Python的byte code(字节码),然后再由Python Virtual Machine 来执行这些编译好的 byte code。实际上,Python Virtual Machine 是一种抽象层次更高 Virtual Machine。基于C的 Python 编译出的字节码文件,通常是.pyc格式。


在这里插入图片描述


在这里插入图片描述


CPython: CPython 混合了编译和解释功能,将Python源代码首先被编译成一系列中间字节码,然后由CPython虚拟机内部 while 循环不断匹配字节码并执行对应字节码指令 case 分支内部的多条C函数。


在这里插入图片描述


PyPy:基于RPython语言构建, 利用 JIT 即时编译来执行 Python 代码。不同于解释器,并不会逐行运行代码,而是在执行程序前先将部分代码编译成机器码。


在这里插入图片描述


AI框架静态化的方案普遍采用修饰符这套方案,这套方案细分下来也有两种不同的方法:
Tracing Based,如 PyTorch.fx
AST Transform, 如 PyTorch.jit.trace


在这里插入图片描述


在这里插入图片描述


AI编译器在 Python 静态化方面的挑战:

  • 类型推导:从 Python 动态类型到编译器 IR 静态类型
  • 控制流的表达:if、else、while、for等控制表达
  • 灵活的语和数据类型转换:针对 Slice、Dict、Index等操作
  • JIT的编译性能:Tracing Based 还是 AST Transform 都需要增加额外的编译开销

Challenge II: Toward DSA

DSA 专用硬件架构:

  • AI 训练和推理都是对性能敏感,所以在AI的场景中大量用到加速器,包括CPU SIMD单元、GPU SIMT架构、华为昇腾Ascend Cube核等专用架构, NV GPU H100架 DSA特征也逐步明显。
  • 因此 AI 编译器逐步成为发挥多样性算力的关键,特别是近期 Dataflow + SIMD 类 DSA芯片占比逐步提升的情况下。

Dataflow:

  • Dataflow 的执行调度更加能发挥芯片的性能,即芯片进行整图或者子图的调度,而不是像GPU那样,主流是 kernel by kernel的调度。

编译器在性能优化的难度和复杂度挑战变大:

  • 性能优化依赖图算融合优化,图层和算子层独立优化无法充分发挥芯片性能,需要图算融合优化,例如子图切分、子图内垂直融合优化和水平并行优化;
  • 优化复杂度提升,标量+向量+张量+加速指令、多级的存储结构,导致 Kernel 实现 Schedule、Tilling、Vectorization、Tensorization复杂;

当前 AI 编译器没有形成统一完善的方案:

  • 打开图和算子的边界,进行重新组合优化,充分发挥芯片性能;
  • 多种优化手段融合:垂直融合优化(Buffer Fusion等)和水平并行(Data Parallel等)优化;
  • 重新组合优化后的子图的Kernel代码自动生成,如scheduling、tilling、vectorizing;

Challenge III: Specific Optimizations

自动并行,当前大模型训练遇到内存墙、性能墙等依赖复杂的并行策略来解决:

  • Scale out:多维混合并行能力,数据并行、张量并行、流水线并行等
  • Scale up:重计算、混合精度、异构并行等

Scale Out/Up 最大挑战就是效率墙:

  • 依赖手工配置切分策略,对开发者来说,门槛高,效率低;
  • 半自动并行解决一部分效率,真正要解放依赖编译+凸优化问题,自动找到最优并行策略。

面向 HPC 场景自动微分的要求更高:

  • 控制流:动态图通过Python侧执行控制流,一旦循环次数多的话,性能劣化;静态图的自动微分通过逻辑拼接或者计算图展开方式来解决,能解决一定性能问题但是方案有待完善。
  • 高阶微分:高阶的前向微分和后向微分如何通过 Jacobian matrix 有效展开,高阶微分通过 Hessian matrix 如何模拟快速计算,开发者如何灵活对高阶微分的图进行控制?

Challenge IIII: 易用性与性能兼顾

与AI框架的边界和对接:

  • 不同AI框架对深度学习任务的抽象描述和 API 接口不同,语义和机制上有各自的特点。需要考虑如何在不保证所有算子被完整支持的情况下透明化的支持用户的计算图描述。

对用户透明性问题:

  • 部分 AI 编译器并非完全自动的编译工具,性能表现依赖于用户提供的高层抽象的实现模版,如TVM。主要是为算子开发工程师提供效率工具,降低用户人工调优各种算子实现的人力成本。
  • 现有抽象却常常无法足够描述创新的硬件体系结构上所需要的算子实现。需要对编译器架构足够熟悉的情况下对其进行二次开发甚至架构上的重构,门槛及开发负担仍然很高。

编译开销:

  • AI 编译器作为性能优化工具,只有在编译开销对比带来的性能收益有足够优势才有实用价值。部分应用场景下对于编译开销的要求较高。对于开发者而言,使用AI编译器阻碍其快速的完成模型的调试和验证工作,极大增加开发和部署的难度和负担。

性能问题:

  • 编译器的优化本质上是将人工的优化方法,或者人力不易探究到的优化方法通过泛化性的沉淀和抽象,以有限的编译开销来替代手工优化的人力成本。
    深度学习编译器只有在性能上真正能够代替或者超过人工优化才能真正发挥价值。

鲁棒性:

  • AI 编译器大多数处于研究型,产品成熟度距离工业级应用有较大差距。是否能够顺利对计算图编译,计算结果的正确性,对错误进行跟踪调试都需要考虑。

问题:

图算能否统一表达,统一编译优化,形成通用的 AI 编译器?

当前的AI框架下,图层和算子层是分开表达和优化,算法工程师主要是接触图层表达,AI框架、芯片、kernel开发工程师主要是负责算子的表达,未来能否会有IR打破图算之间的GAP,未来在AI+科学计算、AI+大数据等新场景驱动下,使得计算图层和算子层不再清晰,能否统一AI的编译优化?


完全的自动并行是否可行?

自动并行能根据用户输入的串行网络模型和提供的集群资源信息自动进行分布式训练,通过采用统一分布式计算图和统一资源图设计可支持任意并行策略和各类硬件集群资源上分布式训练,并且还能利用基于全局代价模型的规划器来自适应为训练任务选择硬件感知的并行策略。实际上自动并行是一个策略搜索问题,策略搜索能够在有限的搜索空间找到一个次有的答案,但是真正意义上的自动并行能否做到需要进一步思考和验证。


AI芯片需要编译器吗?AI芯片需要AI编译器吗?

AI芯片对于编译器的依赖取决于芯片本身的设计。越灵活的芯片对于编译器的依赖会越大。在AI芯片设计之初,有CISC风格把优化在芯片内部解决。但是随着专用领域的演化,为了支持更加灵活的需求,AI芯片本身会在保留张量指令集和特殊内存结构的前提下越来越灵活。未来的架构师需要芯片和系统协同设计,自动化也会越来越多地被应用到专用芯片中去。


AI 编译器的未来

1、编译器形态: 分开推理和训练, AOT 和 JIT 两种编译方式并存
2、IR 形态: 需要有类似于 MLIR 对 AI 统一的 IR 表达
3、自动并行: 提供跨机器、跨节点自动并行的编译优化能力
4、自动微分:提供高阶微分的计算方式,并方便对图进行操作
5、kernel 自动生成: 降低开发门槛,快速实现高效和高泛化的算子


2023-02-24

评论 2
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值