AI编译器的前端优化策略

背景

        工作领域是AI芯片工具链相关，很多相关知识的概念都是跟着项目成长建立起来，但是比较整个技术体系在脑海中都不太系统，比如项目参与中涉及到了很多AI编译器开发相关内容，东西比较零碎，工作中也没有太多时间去做复盘与查漏补缺。但是最近比较闲，发现了一个宝藏级的B站博主，系统的讲了很多AI芯片领域的知识，并把课程资源开源维护，极力推荐大家多多关注。在这里当个搬运工，传播一下。

     

总结

        AI系统框架图及AI编译器前端的位置示意如下：

        AI编译器的前端优化可以总结为神经网络相关的优化和代码/计算层面的优化两个方向。后面针对具体的点展开细说。 

分述

图层IR

本节包含：

• 计算图构成
• 静态计算图和动态计算图
• 对AI编译器的作用

1）计算图的构成

• 用来表示深度学习网络模型在训练与推理过程中计算逻辑与状态
• AI框架在后端会将Python构建网络模型前向与反向梯度计算以计算图的形式来表示
• 由基本数据结构张量 (Tensor) 和基本运算单元算子 (Operator) 构成
• 计算图中常用节点来表示算子，节点间的有向线段来表示张量状态，同时也描述了
• 计算间的依赖关系

2） 静态计算图和动态计算图

        静态图模式下：使用前端语言定义模型形成完整的程序表达后，不使用前端语言解释器进行执行，将前端描述的完整模型交给AI框架。AI框架先对API代码分析，获取网络层之间的连接拓扑
关系以及参数变量设置、损失函数等，接着用静态数据结构来描述拓扑结构及其他神经网络模
型组件。

        动态图采用前端语言自身的解释器对代码进行解析，利用AI框架本身的算子分发功能，算子会即刻执行并输出结果。动态图模式采用用户友好的命令式编程范式，使用前端语言构建神经网
络模型更加简洁。

特性	静态图	动态图
即时获取中间结果	否	是
代码调试难易	难	简单
控制流实现方式	特定的语法	前端语言语法
性能	优化策略多，性能更佳	图优化受限，性能较差
内存占用	内存占用少	内存占用相对较多
部署能力	可直接部署	不可直接部署

动态图转换静态图的方式

• 基于追踪转换 ：以动态图模式执行并记录调度的算子，构建和保存为静态图模型。
• 基于源码转换 ：分析前端代码将动态图代码自动转为静态图代码，底层使用静态图执行运行。

3）计算图对AI编译器的作用

图优化

• 计算图的出现允许 AI 框架在执行前看到深度学习模型定义全局信息
• 计算图作为 AI 框架中的高层中间表示，可以通过图优化 Pass 去化简计算图或提高执行效率

图层IR的好处

方便底层编译优化

• 统一表示来描述神经网络训练的全过程
• 可以序列化保存，不需要再次编译前端源代码
• 将神经网络模型中间表示转换为不同硬件代码
• 直接部署在硬件上，提供高效的推理服务
• 编译期可对计算过程的数据依赖进行分析：
  ◦ 简化数据流图
  ◦ 动态和静态内存优化
  ◦ 预计算算子间的调度策略
  ◦ 改善运行时Runtime性能

分层优化，便于扩展

• 切分出三个解耦的优化层
  ◦ 计算图优化
  ◦ 运行时调度优化
  ◦ 算子/内核执行优化
• 新网络模型结构/新训练算法，扩展步骤：                                                                                 ◦ 计算图层添加新算子
  ◦ 针对不同硬件内核，实现计算优化
  ◦ 注册算子和内核函数，运行时派发硬件执行

算子融合

本节包括：

• 算子融合方式
• 算子融合例子
• 融合的规则和算法

总结：

• 算子的融合方式有横向融合和纵向融合，但根据AI模型结构和算子的排列，可以衍生出更多不同的融