ZOMI的AI编译原理4

为什么需要AI编译器

面临的问题

传统编译器与AI编译器

1. IR 差异：AI 编译器的 IR 与传统编译器的IR所抽象出来的概念和意义并不相同。
   • AI编译器一般会有 high-level IR，用来抽象描述深度学习模型中的运算，如：Convolution、Matmul 等，甚至部分会有 Transformer 带有图的结构。
   • 传统编译器相对而言 low-level IR，用于描述基本指令运算，如 load、store 等。有了high-level IR，AI编译器在描述深度学习模型类 DSL 更加方便。
2. 优化策略：AI 编译器面向AI领域，优化时引入更多领域特定知识，从而进行更 high-level，更加aggressive 优化手段。如：
   • AI编译器在 high-level IR 执行算子融合，传统编译器执行类似 loop fusion 时候，往往更加保守。缺点是可能会导致调试执行信息跟踪难；
   • AI编译器可以降低计算精度，比如int8、fp16、bf16等，因为深度学习对计算精度不那么敏感。但传统编译器一般不执行改变变量类型和精度等优化。

AI编译器的发展阶段

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

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

什么是AI编译器

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

AI编译器的发展阶段

development history: Stage I 朴素的AI编译器

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

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

表达上：

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

性能上：

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

development history: Stage II 专用的AI编译器

表达上：

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

性能上：

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

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

development history: Stage III 通用AI编译器

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

AI编译器的通用架构

AI编译器的挑战与思考