Morgan Woods

真真夜夜的炼金工坊

AI编译原理4

感谢 up 主 ZOMI 酱:https://space.bilibili.com/517221395

为什么需要 AI 编译器#

面临的问题#

挑战类别描述案例
算子挑战越来越多新算子被提出,导致算子库的开发、维护、优化和测试工作量指数上升1. 硬件不仅需要实现新算子,还需要结合硬件进行特性优化和测试,以充分发挥硬件性能。 例如,对于 Convolution 运算,需要将其转换为 GEMM 矩阵乘法;对于新提出的 Swish 算子,硬件需要新增 Swish 对应实现。
2. 硬件供应商还会发布优化库,但需要提供开发类似的优化库,这会增加大量算子优化、封装的工作,并过度依赖库,无法有效利用专用硬件芯片能力
优化挑战专用加速芯片的爆发导致性能可移植性成为一种刚需。1. 大多数 NPU 使用 ASIC,在神经网络场景对计算、存储和数据搬运做了特殊的指令优化,以提升 AI 相关计算的性能
2. 不同厂商提供 XPU 的 ISA 千奇百怪,缺乏如 GCC、LLVM 等编译工具链,使得针对 CPU 和 GPU 已有的优化算子库和针对语言的优化 Pass 很难短期移植到 NPU 上

传统编译器与 AI 编译器#

img

传统编译器与 AI 编译器主要有两大差异:

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

    1. AI 编译器一般会有 high-level IR,用来抽象描述深度学习模型中的运算,如:Convolution、Matmul 等,甚至部分会有 Transformer 带有图的结构。
    2. 传统编译器相对而言 low-level IR,用于描述基本指令运算,如 load、store 等。有了 high-level IR,AI 编译器在描述深度学习模型类 DSL 更加方便。
  2. 优化策略:AI 编译器面向 AI 领域,优化时引入更多领域特定知识,从而进行更 high-level,更加 aggressive 优化手段。如:

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

AI 编译器的发展阶段#

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

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

什么是 AI 编译器#

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

AI 编译器的发展阶段#

img

Stage I: 朴素的 AI 编译器#

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

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

在表达上:

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

在性能上:

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

img

Stage II 专用的 AI 编译器#

img

在表达上:

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

在性能上:

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

img
img

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

Stage III 通用 AI 编译器#

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

img


AI 编译器的通用架构#

img

IR 中间表达#

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

  • 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 训练和推理。特定优化原语可以满足计算要求时,使用优化的内核库可显著提高性能,否则可能会受到进一步优化的约束。
      img

AI 编译器的挑战与思考#

img

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 运算的基石,并确保在推理计算与训练及推理计算之间取得理想的效率平衡点。

挑战#

  1. 动态 Shape 和动态计算图

    • 现状:主流 AI 编译器主要针对特定静态 shape 输入进行编译优化,对包含控制流语义的动态计算图支持有限。
    • 问题:AI 应用场景中存在大量需要动态计算图的任务需求。尽管前端可以尝试将计算图改写为静态计算图,或将适合编译的部分子图展开以供优化,但这些方法不能解决所有问题。
    • 实例:某些 AI 任务(如金字塔结构的检测模型)无法通过人工改写来静态化,编译器在这些情况下难以有效优化。
  2. Python 编译静态化

    方法描述
    Python JIT 虚拟机如 PyPy 或 CPython,试图在 Python 解释执行的基础上增加 JIT 编译加速能力,但兼容性存在问题。
    修饰符方式如 torch.jit.script、torch.jit.trace、torch.fx、LazyTensor、TorchDynamo。尽管这些手段已经广泛应用,但仍缺乏一个完备的方案。

    Python JIT 虚拟机

    img

    修饰符方式

    img

    • AI 编译器在 Python 静态化方面的挑战
      • 类型推导:从 Python 的动态类型到编译器 IR 的静态类型。
      • 控制流表达:如 if、else、while、for 等控制语句的表达。
      • 灵活的语法和数据类型转换:针对 Slice、Dict、Index 等操作的处理。
      • JIT 编译性能:无论是 Tracing Based 还是 AST Transform,都需要额外的编译开销。
  3. 发挥硬件性能,特别是 DSA 类芯片

  • 现状:DSA(专用硬件架构)在 AI 训练和推理中的重要性日益增加,如 CPU 的 SIMD 单元、GPU 的 SIMT 架构、华为昇腾 Ascend 的 Cube 核等。

  • 挑战

    • 性能优化依赖图算融合:需要图层和算子层独立优化,无法充分发挥芯片性能,需要图算融合优化,如子图切分、子图内垂直融合优化和水平并行优化。
    • 优化复杂度提升:涉及标量、向量、张量和加速指令的多级存储结构,导致 Kernel 实现的调度、分块、向量化、张量化复杂。
  • 解决方案

    • 打开图和算子的边界进行重新组合优化,以充分发挥芯片性能。
    • 融合多种优化手段:垂直融合优化(如 Buffer Fusion)和水平并行优化(如 Data Parallel)。
    • 自动生成重新组合优化后的子图的 Kernel 代码,如调度、分块、向量化。
  1. 处理神经网络特性:自动微分、自动并行等
任务现状挑战
自动并行当前大模型训练面临内存墙和性能墙,需要复杂的并行策略。
Scale out:多维混合并行能力,包括数据并行、张量并行、流水线并行等。
Scale up:重计算、混合精度、异构并行等。
依赖手工配置切分策略,门槛高,效率低。
半自动并行可以解决部分效率问题,但要实现自动找到最优并行策略依赖于编译和凸优化问题的解决。
自动微分控制流动态图通过 Python 侧执行控制流,一旦循环次数多,性能会劣化;静态图的自动微分通过逻辑拼接或计算图展开方式解决,能在一定程度上解决性能问题,但方案仍需完善。
高阶微分:通过Jacobian matrix有效展开高阶前向和后向微分;通过Hessian matrix模拟快速计算高阶微分。
控制流动态图在处理复杂控制流时性能下降,尤其是在循环次数多的情况下。静态图方案需要进一步优化以提高性能和稳定性。
高阶微分:如何有效展开和计算高阶前向和后向微分,确保计算的准确性和效率。开发者需要灵活控制高阶微分图。
  1. 易用性与性能兼顾
方面描述
与 AI 框架的边界和对接不同 AI 框架对深度学习任务的抽象描述和 API 接口不同,各有特点。需要考虑如何在不完整支持所有算子的情况下透明支持用户的计算图描述。
对用户透明性问题部分 AI 编译器不是完全自动的编译工具,性能表现依赖于用户提供的高层抽象实现模板(如 TVM)。这降低了人工调优算子实现的人力成本,但现有抽象无法足够描述创新硬件体系结构所需的算子实现,需要对编译器架构足够熟悉才能进行二次开发或架构重构,门槛高,开发负担重。
编译开销AI 编译器作为性能优化工具,只有在编译开销相对于性能收益足够有优势时才有实用价值。在某些应用场景下,对于编译开销的要求较高,使用 AI 编译器可能阻碍快速完成模型调试和验证,增加开发和部署的难度和负担。
性能问题编译器的优化本质上是将人工优化方法或难以探究的优化方法通过泛化和抽象,以有限的编译开销替代手工优化的人力成本。深度学习编译器只有在性能上真正能够替代或超过人工优化才能发挥价值。
鲁棒性大多数 AI 编译器处于研究阶段,产品成熟度与工业级应用有较大差距。需要考虑能否顺利对计算图编译,确保计算结果正确性,以及对错误进行跟踪和调试的能力。

其他问题#

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

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

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

未来#

  • 编译器形态:分开推理和训练,AOT 和 JIT 两种编译方式并存
  • IR 形态:需要有类似于 MLIR 对 AI 统一的 IR 表达
  • 自动并行:提供跨机器、跨节点自动并行的编译优化能力
  • 自动微分:提供高阶微分的计算方式,并方便对图进行操作
  • kernel 自动生成:降低开发门槛,快速实现高效和高泛化的算子
Loading...
Ownership of this post data is guaranteed by blockchain and smart contracts to the creator alone.