AI编译器的后端优化策略

背景

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

​

     

总结

后端优化与前端优化的区别
• 前端优化：输入计算图，关注计算图整体拓扑结构，而不关心算子的具体实现。在 AI 编译器的
前端优化，对算子节点进行融合、消除、化简等操作，使计算图的计算和存储开销最小。
• 后端优化：关注算子节点的内部具体实现，针对具体实现使得性能达到最优。重点关心节点的
输入，输出，内存循环方式和计算的逻辑。 

AI编译器后端部分：

• 生成低级IR；
• 后端优化；
• 代码生成；

​

1）生成低级IR：不同 AI 编译器内部低级 IR 形式和定义不同，但是对于同一算子，算法的原理
实质相同。对于每个具体的算子，需要用AI编译器底层的接口来定义算法，再由编译器来生成
内部的低级IR。

2）后端优化：针对不同的硬件架构/微架构，不同的算法实现的方式有不同的性能，目的是找
到算子的最优实现方式，达到最优性能。同一算子不同形态如Conv1x1、 Conv3x3、 Conv7x7
都会有不同的循环优化方法。

3）代码生成：对优化后的低级 IR 转化为机器指令执行，现阶段最广泛的借助成熟的编译工具
来实现，非AI 编译器的核心内容。如把低级IR 转化成为 LLVM、NVCC 等成功编译工具的输入
形式，然后调用其生成机器指令。

优化策略

算子概念

算子：深度学习算法由一个个计算单元组成，我们称这些计算单元为算子（Operator，简称Op）。算子是一个函数空间到函数空间上的映射O：X→Y；从广义上讲，对任何函数进行某一项操作都可以认为是一个算子。于AI 框架而言，所开发的算子是网络模型中涉及到的计算函数。

算法：算法（Algorithm）是指解题方案的准确而完整的描述，是一系列解决问题的清晰指令，算法代表着用系统的方法描述解决问题的策略机制。

算子分类：

• 访存密集型：如 RNN 训练任务，由于 RNN 网络模型结构的计算密度更低，瓶颈转移到host端的 Op Launch 上，因此kernel之间甚至出现了大量空白。对于访存密集型算子部分的工作来说，新硬件带来了更大的性能优化空间。
• 计算密集型：计算密度是指一个程序在单位访存量下所需的计算量，单位是 FLOPs/Byte，计算密度较大，程序性能受硬件最大计算峰值（下文简称为算力）限制，称为计算密集型程序。

 算子优化的挑战：

• 优化手段多样：要在不同情况下权衡优化及其对应参数，对于优化专家来说也是相当耗费精力
• 通用性与移植性：不同类型的硬件架构差异，使得优化方法要考虑的因素也有很大不同
• 优化间相互影响：各种优化之间可能会相互制约，相互影响。这意味着找到最优的优化方法组合与序列就是一个困难的组合优化问题，甚至是 NP 问题。

算子库：
• 目的：针对访存密集型和计算密集型的算子进行优化。对于一个算子，要实现正确的逻辑计算
或许不难，但要结合硬件能力达到高性能就比较难了，要达到极致的性能更是难上加难。
• 业界一个最为常见的方式是将预置的算子实现封装成计算库。如 CuDNN、CuBLAS、OpenBLA
S、Eigen 这样优秀的计算库。

算子计算与调度

• 计算是算子的定义，回答算子是什么，如何通过具体的算法得到正确的定义结果。
• 调度是算子的执行策略和具体实现，回答系统具体如何执行算子的计算定义。
• 同一算子会有不同的实现方式（即不同的调度方式），但是只会有一种计算形态（具体定义）。
• 计算实现（Function / Expression）和计算在硬件单元上的调度（Schedule）是分离。

        算子调度具体执行的所有可能的调度方式称为调度空间，AI 编译器优化的目的就是为算子提供一种最优的调度，使得算子在硬件上运行时间最优。

​

循环优化

循环展开Loop Unrolling

• 对循环进行展开，以便每次迭代多次使用加载的值，使得一个时钟周期的流水线上尽可能满负荷计算。在流水线中，会因为指令顺序安排不合理而导致NPU等待空转，影响流水线效率。循环展开为编译器进行指令调度带来了更大的空间。

​

循环分块 Loop tiling
• 由于内存空间有限，代码访问的数据量过大时，无法一次性将所需要的数据加载到设备内存，
循环分块能有效提高NPU cache 上的访存效率，改善数据局部性。
• 如果分块应用于外部循环，会增加计算的空间和时间局部性；分块应与缓存块一起作用，可以
提高流水线的效率。

• Loop Tiling 的目的是确保一个 Cache 在被用过以后，后面再用的时候其仍然在 cache 中。
• 实现思路：当一个数组总的数据量无法放在 cache 时，把总数据分成一个个 tile 去访问，令每
个 tile 都可以满足 Cache
• 具体做法：把一层内层循环分成 outer loop * inner loop。然后把 outer loop 移到更外层去，
从而确保 inner loop 一定能满足 Cache

在我之前写的博文CUDA编程的矩阵乘优化中，有一个优化思想就是循环分块。

循环重排 Loop Reorder
• 内外层循环重排，改善空间局部性，并最大限度地利用引入缓存的数据。对循环进行重新排序，
以最大程度地减少跨步并将访问模式与内存中的数据存储模式对齐

循环融合 Loop Fusion
• 循环融合将相邻或紧密间隔的循环融合在一起，减少的循环开销和增加的计算密度可改善软件
流水线，数据结构的缓存局部性增加。

循环拆分 Loop Split
• 拆分主要是将循环分成多个循环，可以在有条件的循环中使用，分为无条件循环和含条件循环。

指令优化

向量化 Vectorization

SIMD 寄存器先保持vec_sum的值，在reduction 步骤,，vec_sum 再逐个求总和。

张量化 Tensorization

• Volta 架构中，一个SM由8个FP64 Cuda Cores，16个INT32 Cuda Core，16个FP32 Cuda Core，和128个Tensor Core组成，一共有4个SM。

        主流CPU/GPU硬件厂商都提供了专门用于张量化计算的张量指令，如英伟达的张量核指令、
英特尔的VN。利用张量指令的一种方法是调用硬件厂商提供的算子库，如英伟达的cuBLAS和
cuDNN，以及英特尔的oneDNN等。然而，当模型中出现新的算子或需要进一步提高性能时，这种方法的局限性便显露无遗。

1. 新的硬件体系带来了超越向量运算的新指令集，调度必须使用这些指令才能从加速中获益
2. 张量计算基元的输入是多维的，具有固定的或可变的长度，并指示不同的数据布局
3. 新的 AI 加速器正以它们自己的张量指令出现

存储优化

访存延迟 Latency Hiding
• 延迟隐藏（Latency Hiding）是指将内存操作与计算重叠，最大限度地提高内存和计算资源利用
率的过程。

CPU：
• 延迟隐藏可以通过多线程，或者硬件隐式数据预取实现
GPU：
• 依赖于 Wrap Schedule 对多线程的管理调度和上下文切换实现
NPU/TPU：
• 采用解耦访问/执行（Decoupled Access/Execute，DAE）架构

内存分配

        传统编译器通过将内存逻辑划分为不同区段来提供程序员访问内存的权限（栈、堆、静态存储区），而每段空间都有各自的大小，一旦开发者在使用过程中将该大小耗尽，就会出现内存不足错误。
        AI芯片上内存分为：Shared Memory， Local Momory。如下是GPU chip的内存层级示意图：

自动调优

        AI编译器的Ato tuning是指对于给定的程序和目标架构，自动调优算法的方式找到最优的编译优化方法。需要考虑的问题是：使用哪些优化方法？选择什么参数集？用什么顺序应用优化方法以达到最佳性能？

步骤：

1） 参数化 Parameterization
2）成本模型 Cost Model
3）搜索算法 Search Algorithm

参数化 Parameterization：对调度优化问题进行建模，参数化优化空间一般由可参数化变换（Lo
op）的可能参数取值组合构成，因此需要调度原语进行参数化表示。Halide 将算法和调度解耦，
TVM 提供调度模板。

成本模型 Cost Model：用来评价某一参数化下的调度性能，根据对调度额评价来指导最搜索到
最优的调度策略。可以从运行时间、内存占用、编译后指令数来评价。实现方式主要有1）基
于NPU硬件的黑盒模型；2）基于模拟的预定义模型；3）ML-Base 模型，通过机器学习模型来
对调度性能进行预测

搜索算法 Search Algorithm：确定初始化和搜索空间后，在搜索空间找找到达到性能最优的参数
配置。常用的搜索算法有1）遗传算法、2）模拟退火算法、3）强化学习等。

系列：AI编译器的前端优化策略-CSDN博客