AI编译器技术浅析

随着深度学习应用和专用加速器的兴起，为了快速部署DNN网络，产生了两条主要的技术路线：手写算子和AI编译器（深度学习编译器）生成代码。

手写算子需要具备领域特定专业知识的工程师打磨调试出高性能的代码实现，然后再经过传统编译器生成可执行代码，经验丰富的工程师可以写出性能极佳的算子。但是手写算子的缺点也非常明显：其一，需要堆人力，并且很难覆盖深度学习领域的所有算子。其二，手写算子一般都是单算子实现或者和前后element-wise算子做简单融合，很难做很深的算子融合，在追求极致的推理场景，还是无法完全发挥硬件性能极限。

AI编译器致力于自动生成高效代码。我理解的AI编译器技术栈是这样的：

首先，它的输入可能是深度学习框架的模型，AI编译器做前端的解析，生成AI的编译器High-Level的Gpaph IR（图IR）。在Graph IR上做图级别的优化，这里既能借鉴传统编译器的优化算法进行优化（比如常量折叠、死代码删除等等），也需要增加领域特定的优化（比如融合、layout、量化等等），领域特定的优化很大程度上依赖于领域特定架构（DSA）的信息。

然后，再转化为low-level的张量级别的IR（Tensor IR），在这一层，主要利用DSA特点进行优化，比如循环优化、片上存储的tiling、算子的fusion、并行、张量化等等。

TVM的autoTVM引入模板匹配搜索的方法，借助HalideIR计算和调度分离的思想，并提供一些primitives（比如split、inline、tile等）来构建一个搜索空间，再这个搜索空间中得到最优解。这种方法本质上是利用DSA的存储、并行等特点来得到schedule组合，由于是模板匹配的方式，针对特定的算子和加速器，需要专家来实现模板。为了解决autoTVM的问题，引入了auto-schedule（a.k.a Ansor）机制，可以参考ANSOR：一种Auto-Scheduler方法

华为AKG引入了多面体技术，将HalideIR转换为Polyhedral IR，对loop fusion和loop tiling进行建模，充分发挥并行性和局部性。可以参考AKG-NPU上算子自动生成技术探索

最后，借助传统编译器生成可执行代码。

不管是auto-schedule还是Polyhedral技术，当前还不够理想，无法充分发挥DSA特点来生成高效代码，并且还没有很好支持dynamic shape。AI编译器如何充分利用DSA这个领域特点做文章（算子自动生成技术），是当前AI编译器面临的技术难点。

当前，DNN的部署很大程度上还依赖手写算子库，AI编译器在如何自动化发挥DSA特点进行自定生成高效代码方面还有很长的路要走，当前一个折中的做法，是将手写算子和AI编译结合起来，发挥二者的优势，以求能够快速部署。