近年来，AI应用程序已经无处不在。比如：智能家居设备由自然语言处理（NLP）和语音识别模型驱动，自动驾驶技术以计算机视觉模型为支柱。通常这些AI模型会部署在云平台、专用计算设备以及物联网传感器的内置微型芯片。
因此，我们在进行AI应用落地时，需要将AI模型从研发阶段转而部署到多种多样的生产环境，同时，也需要相当多的繁重工作。即使对于我们最熟悉的环境（例如：在 GPU 上），部署包含非标准算子的深度学习模型仍然需要大量的工程。而为了解决这些繁琐的问题，AI 编译器应运而生。
本系列将分享 AI 编译器的技术原理，本文为该系列第一篇文章，将简要介绍 AI 编译器诞生的背景、与传统编译器的区别、AI 编译器前后端等。
AI 编译器产生的背景
早期神经网络部署的侧重点在于框架和算子库。神经网络由数据流图来表示，图上的节点就是算子（比如：Conv2D、BatchNorm、Softmax），节点之间的连接代表 Tensor。由于数据流图很直观，很多框架的 Runtime 采用了类似 Caffe 的方式，运行时通过一定的顺序（例如：直接 Post-Order DFS）分配 Tensor、调用算子库就行了，优化重点在于优化算子库的性能。
随着时间的发展，这种直观的部署方式，也逐渐暴露出一些问题。
- 越来越多的新算子被提出，算子库的开发和维护工作量越来越大。比如：提出一个新的 Swish 算子，算子库就要新增 Swish 的实现，还要有优化、测试。Swish 由一些基础的一元、二元算子组成。
- NPU 的爆发导致模型性能可移植性成为一种刚需。大多数 NPU 作为一种 ASIC 在神经网络场景对计算、存储和数据移动（拷贝、读写等）做了特殊优化，相对 CPU、GPU来说，对能效比要好很多。在移动端和 edge 端，越来越多的 NPU 开始出现。同时，NPU 的指令集架构（ISA）千奇百怪，一般也缺乏 GCC、LLVM 等工具链，使得已有的针对 CPU 和 GPU 优化的算子库，很难短期移植到 NPU 上，充分利用硬件的能力达到较好的性能。
- 更多可优化的点得到关注。早期 CPU 和 GPU 上带宽问题不是很明显，大家更多关注单个算子的性能。但在移动端和 edge 端的应用中，逐渐遇到了带宽跟不上算力的问题，在这些 target 上增大带宽，意味着功耗和成本的上升，利用算子间的融合和调度，节省带宽开始被重视起来。
AI 编译器
AI编译器是指将机器学习算法从开发阶段，通过变换和优化算法，使其变成部署状态。
- 开发形式 是指我们在开发机器学习模型时使用的形式。典型的开发形式包括用 PyTorch、TensorFlow 或 JAX 等通用框架编写的模型描述，以及与之相关的权重。
- 部署形式 是指执行机器学习应用程序所需的形式。它通常涉及机器学习模型的每个步骤的支撑代码、管理资源（例如内存）的控制器，以及与应用程序开发环境的接口（例如用于 android 应用程序的 java API）。
image.png
AI 编译的目标
- 集成与最小化依赖：部署过程通常涉及集成 (Integration)，即将必要的元素组合在一起以用于部署应用程序。 例如，如果我们想启用一个安卓相机应用程序来检测猫，我们将需要图像分类模型的必要代码，但不需要模型无关的其他部分（例如，我们不需要包括用于 NLP 应用程序的embedding table）。代码集成、最小化依赖项的能力能够减小应用的大小，并且可以使应用程序部署到的更多的环境。
- 利用硬件加速：每个部署环境都有自己的一套原生加速技术，并且其中许多是专门为机器学习开发的。机器学习编译的一个目标就是是利用硬件本身的特性进行加速。 我们可以通过构建调用原生加速库的部署代码或生成利用原生指令（如 TensorCore）的代码来做到这一点。
- 通用优化：有许多等效的方法可以运行相同的模型执行。AI 编译的通用优化形式是不同形式的优化，以最小化内存使用或提高执行效率的方式转换模型执行。
AI 编译器与传统编译器的区别与联系
传统编译器：输入高级语言，输出低级语言；而 AI 编译器：输入计算图/算子，输出低级语言。
相同点在于两者都做了类似语言转换的工作。不同点在于传统编译器解决的主要问题是降低编程难度，其次是优化程序性能。而 AI 编译器解决的主要问题是优化程序性能，其次是降低编程难度。
两者的联系
因为AI编译器出现的比较晚，所以在设计的时候，往往会借鉴传统编译器的思路：
- 两者的理念比较类似。都力求通过一种更加通用，更加自动化的方式进行程序优化和代码生成，从而降低手工优化的工作。
- 两者的软件结构比较类似。一般都分成前端，IR，后端等模块。前端负责将不同的语言的描述转换成统一的IR表示，后端会对IR表示进行优化，最终生成可执行的代码。IR层用来解耦前端和后端，降低集成的工作量。
- 两者的优化方式比较类似。编译器都会对代码进行一系列的优化，从而提高性能，或者减少内存占用等。AI编译器和传统编译器都是通过在IR上面，运行各种各样的pass进行优化的。而且AI编译器往往还会借鉴传统编译器中的一些pass，比如：常量折叠（constant folding）， 死代码消除（dead code elimination）等。
- AI编译器通常会依赖于传统编译器。AI编译器在IR上面对模型进行优化之后，通常会有 lowering 的过程，将优化后的 high-level IR 转换成传统编译器的 low-level IR；然后，依赖传统编译器做最终的机器码生成。
两者的区别
两者最根本的区别是应用场景的区别：
- AI编译器是把一个深度学习模型转换成 executable。这里可以把一个深度学习模型，理解成一段用DSL（Domain Specific Language）描述的代码，executable 就是一段用硬件能理解的机器码描述的代码。
- 传统编译器是把一段用高级语言编写的代码转换成 executable。这里的高级语言可能是 C/C++ 等。
而应用场景的区别导致了两者在设计上的不同。
- 两者的 IR 表达层次有区别。 AI编译器一般会有一套 high-level 的IR，用来更抽象的描述深度学习模型，常用的high-level的运算，比如：convolution，matmul等。传统编译器的IR更偏low-level，用于描述一些更加基本的运算，比如：load，store，arithmetic等。有了 high-level 的IR，AI 编译器在描述深度学习模型的时候会更加方便。
- 两者的优化策略有区别。AI 编译器因为是面向 AI 领域的，在优化时，可以引入更多领域特定的先验知识，从而进行更加 high-level，更加 aggressive 的优化。比如说：
- AI编译器可以在high-level的IR上面做 operator fusion 等，而传统编译器在做类似的 loop fusion 的时候，往往更加保守。
- AI编译器可以降低计算的精度，比如：int8, bf16等，因为深度学习模型，对计算精度不那么敏感。但传统编译器一般不会做这种优化。
神经网络就是一组矩阵计算；神经网编译器，就是将这组计算针对平台尽可能加速。因此，对神经网络模型的优化，就是尽量减少逻辑判断，最好是一算到底。同时，对内存要尽可能优化，降低内存占用。
AI 编译器架构
AI 编译器主要分为编译器前端和编译器后端，分别针对于硬件无关和硬件相关的处理。每一个部分都有自己的 IR (Intermediate Representation，中间表达)，每个部分也会对其进行优化，具体如下图所示。
- 编译器前端对应High-level IR：用于表示计算图，其出现主要是为了解决传统编译器中难以表达深度学习模型中的复杂运算这一问题，为了实现更高效的优化，所以新设计了一套 IR。
- 编译器后端对应Low-level IR： 能够在更细粒度的层面上表示模型，从而能够针对于硬件进行优化。
AI 编译器前端
AI编译器前端主要将用户代码进行解析翻译得到计算图IR，并对其进行设备信息无关的优化，此时的优化并不考虑程序执行的底层硬件信息，其基础结构如下如图所示。
image.png
AI编译器前端的独特之处主要在于对自动微分功能的支持。为了满足自动微分功能带来的新需求，机器学习框架需要在传统中间表示的基础上设计新的中间表示结构。此外，AI编译器前端也会进行编译优化。编译优化意在解决编译生成的中间表示的低效性，使得代码的长度变短，编译与运行的时间减少，执行期间处理器的能耗变低。常见的前端优化有算子融合、内存分配、常量折叠、公共子表达式消除、死代码消除、代数化简
AI 编译器后端
AI编译器后端的主要职责是对前端下发的IR做进一步的计算图优化，让其更加贴合硬件，并为IR中的计算节点选择在硬件上执行的算子，然后为每个算子的输入输出分配硬件内存，最终生成一个可以在硬件上执行的任务序列，其总体架构如下图所示。
image.png
编译器后端处于前端和硬件驱动层中间，主要负责计算图优化、算子选择和内存分配的任务。
首先，需要根据硬件设备的特性将IR图进行等价图变换，以便在硬件上能够找到对应的执行算子，该过程是计算图优化的重要步骤之一。
在完成计算图优化之后，就要进行算子选择过程，为每个计算节点选择执行算子。算子选择是在得到优化的IR图后选取最合适的目标设备算子的过程。针对用户代码所产生的IR往往可以映射成多种不同的硬件算子，但是这些不同硬件算子的执行效率往往有很大差别，如何根据前端IR选择出最高效的算子，是算子选择的核心问题。算子选择本质上是一个模式匹配问题。其最简单的方法就是每一个IR节点对应一个目标硬件的算子，但是这种方法往往对目标硬件的资源利用比较差。现有的编译器一般都对每一个IR节点提供了多个候选的算子，算子选择目标就是从中选择最优的一个算子作为最终执行在设备上的算子。总的来说，在机器学习系统中，对前端生成的IR图上的各个节点进行拆分和融合，让前端所表示的高层次IR逐步转换为可以在硬件设备上执行的低层次IR。得到了这种更加贴合硬件的IR后，对于每个单节点的IR可能仍然有很多种不同的选择，例如：可以选择不同的输入输出格式和数据类型，需要对IR图上每个节点选择出最为合适的算子，算子选择过程可以认为是针对IR图的细粒度优化过程，最终生成完整的算子序列。
最后，遍历算子序列，为每个算子分配相应的输入输出内存；然后，将算子加载到设备上执行计算。
结语
本文主要介绍 AI 编译器诞生的背景、AI编译器的简介，AI编译器与传统编译器的区别、AI 编译器的架构以及AI编译器前后端等。
码字不易，如果觉得有帮助，欢迎点赞收藏加关注。
参考文档
- AI编译器和前端技术
- 【AI编译器原理】系列来啦！我们要从入门到放弃！
- 什么是机器学习编译
- AI与传统编译器
AI编译器技术剖析（一）-概述 - 知乎