【AI系统】AI 编译器基本架构

AI 编译器基本架构

在上一篇文章中将 AI 编译器的发展大致分为了 3 个阶段，分别为 1）朴素编译器、2）专用编译器以及 3）通用编译器。

本文作为上一篇文章 AI 编译器架构的一个延续，着重讨论 AI 编译器的通用架构。首先将回顾现有 AI 编译器架构（以 PyTorch 作为标杆），随后引出通用 AI 编译器的架构模型，并进一步介绍其 IR 中间表达层、前端优化以及后端优化的细节，最后以图的形式展示现有 AI 编译器全栈产品。

AI 编译器架构回顾

现有 AI 编译器架构即是专用 AI 编译器的架构：在表达上以 PyTorch 作为标杆，对静态图做转换，转换到计算图层 IR 进行优化；性能上希望打开计算图和算子的边界，进行重新组合优化以发挥芯片尽可能多的算力。

现有 AI 编译器架构图如下图所示。

此编译器接受的高级语言为 Python，编译器前端会对 Python 代码进行解析，解析器会将高层次的代码转换为一个中间表示（IR），以便进一步处理。这里编译器前端会生成 Graph IR 传递给 Graph Optimizer（图优化器）。

Graph Optimizer 接收到 Graph IR 后，会对解析后的计算图进行优化。优化的目的是减少计算图的冗余部分，提高执行效率。这可能包括算子融合、常量折叠等技术。Graph Optimizer 在优化完成后会向 Ops Optimizer（操作优化器）传递一个 Tensor IR。

Ops Optimizer 接收到 Tensor IR 后，其会针对每个算子进行具体的性能优化，例如重排计算顺序、内存优化等。所有的中间表达都传递至后端之后，后端会生成不同的硬件代码以及可执行程序。

AI 编译器通用架构

在回顾完现有 AI 编译器架构后，来看看一个理想化的 AI 编译器通用架构应该是什么样的。

推荐阅读 AI 编译器的综述，名称为 The Deep Learning Compiler: A Comprehensive Survey。其中有一副插图展示了一个通用 AI 编译器的完整架构，涵盖从模型输入到在不同硬件平台上执行的整个流程。它分为编译器前端（Compiler Frontend）和编译器后端（Compiler Backend）两个主要部分。下面将结合此图对通用 AI 编译器进行初步分析。

编译器前端

编译器前端（Compiler Frontend）主要负责接收和处理来自不同 AI 框架的模型，并将其转换为通用的中间表示（IR），进行初步优化。

编译器前端的组成集中展示在上图中间靠左部分。输入的神经网络模型格式可以来自多种框架（如 TensorFlow、PyTorch 等）；这些模型通过符号表示的转换（如 TVM、nGraph 等）生成计算图；高级 IR/图 IR（设备无关）包含了 DAG（有向无环图）和基于 let-binding 的表示形式以及张量计算；计算图经过多种优化，如代数简化、操作融合、操作下沉、公共子表达式消除（CSE）、死代码消除（DCE）和静态内存规划等，得到初步优化的计算图；随后通过模式匹配和图重写等方法进一步优化，最终生成优化后的计算图；同时，编译器前端也提供了调试工具（如 IR dumping）可以以文本或 DAG 形式呈现。

Input Format of DL Models（输入格式）：支持多种 AI 框架，如 TensorFlow、PyTorch、Caffe2、MXNet、飞桨（PaddlePaddle）和 ONNX 等。

Transformation（转换）：