从AI编译器的视角看硬件加速器设计

引言

最近在B站看ZOMI酱的《AI System》课程，里面讲到AI编译器的部分让我感觉有一些收获，所以记录一下目前的一些理解。

一、中间表达与计算图

ZOMI酱首先对传统编译器和AI编译器之间的区别做了介绍，这种对比是一种很好的理解方式，能帮助你后续不断回顾和梳理所有的知识，提升学习效率，但除此之外，我还想补充一种理解来更好地诠释AI编译器。我认为，要了解AI编译器就首先要理解一个核心的概念——计算图，它是AI算法在编译中的一种关键的中间表示（IR）。不知你是否思考过一个问题，开发人员能想到的IR形式肯定不止计算图一种，为什么选择它作为AI的IR呢？这可能要从AI算法和其他软件程序的区别说起。你在数据结构里见过的算法，基本都会有一个共同的特点，即算法的执行主要基于控制逻辑，代码中if-else语句、循环语句、switch等语句较多，影响算法走向的判断条件多，算法原理也常常涉及堆栈、队列、链表等复杂操作。然而AI算法则恰好相反，如果我们回想一下神经网络的形态，就会发现它的整体结构和计算细节非常明确，算法程序中控制逻辑很少，基本都是由各种具体的网络层按一定先后顺序连接成通路，虽然代码中会有用来实现卷积或矩阵乘法的for循环，但这些循环体本质上只是在表达一种大量计算的过程，而这些计算自身都是非常基本的乘累加操作（MAC），这与其他领域算法中那些控制过程依然有明显的区别，再加上AI算法中核心的计算过程本身只依赖于输入数据和模型参数，一般不会涉及键盘输入、文件读取等IO操作，这些都使得AI算法尤其适合用计算图来表示，且这种计算图能够由特定的算子来构成。因此，AI编译器的设计其实充分遵循了AI算法在程序形态上的特性，明确这一点后，面对AI编译器的前后端里围绕计算图和算子的各种优化时应该就会感到更好理解了。

注：AI编译器的计算图与ONNX模型
ONNX是一种用来标准地描述计算图的文件，PyTorch中可以导出AI算法对应的ONNX，并使用netron（https://netron.app/ ）来进行可视化。打开ONNX后，计算图中的节点会被呈现为彩色的方框，这些方框里会提供比如该层的类型、特征通道数、输入和输出的尺寸等详细信息，而方框之间的箭头就指示了数据的流向，从而能让我们直观地了解网络结构。在我参与过的研究中，ONNX模型经常被使用，个人认为把ONNX模型可视化出来从而方便理解算法计算过程是一个很好的方法，但是也有其局限性，即这个方法只适用于像DNN这样的计算过程比较规律性的算法，一堆卷积层/注意力层堆叠，如果算法的计算过程更加复杂，显示成ONNX的样子可能依然是眼花缭乱的，特别是一些不属于常规算子而是对数据做其他处理的过程。不管怎么说，阅读算法源码，包括用单步调试之类的方式一点一点看计算过程，应该是免不了的功夫。

二、AI编译器与AI芯片

AI编译器对算法所做的处理，说到底是要使算法的计算过程最终由硬件执行出来，因此AI编译器所做的一些事情实际上也就指出了AI芯片的设计者需要考虑的问题。比如AI编译器的前端部分会对数据排布方式和存储空间的分配进行调控，二者都是影响算法能否正确、高效地部署到硬件上的关键点。对于刚了解AI芯片领域的同学来说，往往容易陷入一个误区，即不经过系统级视角去建立一些整体性的认识，就直接沉浸在微观的思考里，比如你可能会在了解了卷积和矩阵乘法的运算细节以后，就开始想“如何用电路来实现卷积/矩阵乘法”，最后虽然能想出一些具体方案，但是你很可能是在不知不觉间模糊地假设整个过程的输入数据是“一个一个地”进入模块，或是其他来自于你的直觉的方式，如果有人问你诸如此类的问题——为何数据要以这种方式输入？你的模块完成一次计算所需的时间是多少、其大小跟你预设的数据排布与输入方式有何关系？当矩阵的各个维度的具体值很大或很小或甚至不可预测时，你的模块在功能设计、计算时间、所需资源等方面会受到什么影响？如果后续要做出一整个芯片，芯片总线的位宽该取多少、模块端口的位宽对总线位宽的有效利用率又有多少？当模块完成计算需要输出结果时，那些结果该放到哪里，你如何保证这些结果在写入存储时不会覆盖属于其他模块的数据呢？这时你大概就会变得不知道如何回答，因为你从一开始构想那个模块时，心里想的就仅仅是“用电路实现卷积/矩阵乘法”，而其实自己对芯片的数据排布、互连带宽、内存管理等问题没有什么概念。因此，我们通过对AI编译器的学习，就可以反过来思考为什么硬件需要编译器完成像特定的数据排布和合理的内存分配等前置的处理，从而学会在芯片设计当中做出一些必要的、合理的规划。

在AI System课程中，华为昇腾AI芯片的矩阵分块计算数据流设计与NC1HWC0格式的数据排布之间的对应是个很好的例子，我们在自己的设计中也需要像这样明确每个模块所需的数据排布方式，或是思考不同的数据排布方式对于同一个模块的设计是否在计算速度快慢、消耗资源多少等问题上有影响，确定更好的设计方案，并及时发现有没有必要引入一些数据重排模块来方便每个计算模块获得的数据能是自己想要的格式。同时，内存分配这件事也提醒我们要学会管理存储空间，以便计算过程能井然有序，事实上，在AI芯片中的每个模块一般都要有属于自己的指令，指令中包含的信息根据模块的不同而很多样，比如矩阵乘法模块的指令里可能会指明本次指令所发起的一轮计算里要处理的矩阵的尺寸，池化模块的指令里可能会有某个字段或比特要用来指示模块要按最大池化还是平均池化的模式来工作，但无论什么模块，指令中通常都会包含输入和输出数据在加速器的公用存储空间中的基地址，从而告诉模块在接下来的一轮工作中，数据是“从哪来、到哪去”的，而具体怎么给模块们设定这些基地址，就是一个内存分配的问题，你可以简单地把存储空间静态划分成不同部分给不同模块去使用，也可通过某些方式实现动态的管理（如从硬件上实现链表管理等，这需要一定的Verilog能力）。理解了这些问题，你也就会更明白如何给自己的Verilog设计跑demo——例如可以写一些python脚本来把从PyTorch中导出的网络真实数据（特征图、权重等）按你在硬件设计时设想好的格式去排布出来，把所得的txt文件（或者其他特定格式的文件）的内容导入到存储模块中，这种做法虽然好理解，甚至本来就是我们熟知的方式，但重新理解一下，这其实就相当于简单地模拟了编译器要干的事情。至于计算过程所对应的指令序列，最直接的方式就是按照计算过程手工编排出来，制作成txt等文件再导入硬件中。当然，如果你有能力直接做出一个自动化的程序来生成指令序列，那也是极好的，毕竟手工编实在是太费劲了。

注：AI加速器的基本架构与调度方式
这里说到的关于每个模块基于指令而工作的内容可能有点令人困惑，因为我们只是从计算模块的视角谈论它们需要怎样的指令，没有从整个加速器的视角解释模块们是如何接受调度的。其实这里的讨论主要是针对一种通用的加速器架构——“Overlay架构”而言的，这种架构代表着一种类似于通用处理器的工作方式：它一般的组成是中央控制模块+算子模块+存储模块+总线模块，其中算子模块（也即计算模块）是把算法里涉及的所有算子实现为一套模块，比如MAC模块专门用来做卷积和矩阵乘法，池化模块专门用来做池化，而对于算法里的每个层或者步骤，中央控制模块会给相应的算子模块派发指令从而调用算子模块去工作，指令会像上面提到的那样包含每次调用所对应的层或步骤的信息，比如矩阵尺寸、通道数等，一般还得有输入和输出数据的基地址以便算子模块能正确地读写存储。这样的工作方式让Overlay架构具有很好的通用性，因为只要全面地实现了算子，无论网络结构如何，最终都可以通过对算子模块的调度来完成计算。

三、编译优化与硬件优化的对比

除了会做那些为了调度硬件而必须做的准备之外，AI编译器还会从一些角度对算法的计算过程和硬件调度方式做出优化，从而让算法在硬件上执行时能有更高的性能或能效。例如在AI编译器的前端优化中，算子融合是一个有趣的步骤，它寻找算法中是否存在能够合为一体的相邻算子，通过把它们合并来减少中间结果的访存开销。除此之外，公共表达式消除、循环优化、访存延迟隐藏等也都是有效的方法。然而，把AI编译器的前后端优化都看过了之后，我们可能会遇到一个困惑——AI编译器似乎可以直接从软件调度层面完成像数据流的特殊设计等许多事情，那么AI芯片自身的优化是否又有必要呢？我觉得这个问题至少可以从以下2个层面来考虑。

一方面，软件只是负责调度硬件，算法的硬件部署最终能达到什么效果，还是要看硬件自身的条件。比如，AI编译器可以决定芯片里乘法器、加法器等功能电路如何被调用，但是并不决定这些电路本身的速度、能耗等属性，在不同芯片里，乘法器的流水线级数、工作的频率等都可能不同，所以即使是同样的调度过程，放在不同的硬件上实际执行出来，所得到的性能与能耗等也会有区别；再比如，AI编译器可以通过循环展开来增强算法执行时的性能，但具体能展开到什么程度，取决于目标硬件有多少的资源，就像同样基于CUDA生态，用不同型号的GPU来跑同样的算法，性能肯定是不同的。如今有一些先进的集成电路技术其实就是在从这个层面上解决一些问题，例如存算一体的出现改变了AI芯片仍然基于冯·诺依曼架构的局面，让芯片的工作方式本身能够天然地具有更高的能效；三维集成的工艺改变了芯片的集成架构，同样的访存过程，用三维的TSV互连来代替传统的二维接口，也就会直接实现更低的延迟、功耗和更高的带宽。这样的提升是软件调度本身做不到的，实际上是它们给软件调度提供了新的问题求解空间，软件调度优化真正的意义是尽可能地发挥目标硬件的能力，是逼近极限而非改变极限，硬件本身才代表着“现实”。

另一方面，从数据流的设计来说，在硬件上直接进行定制化的架构级创新，往往会得到比单纯优化软件调度更好的收益。例如，脉动阵列的计算过程非常好理解，我们是不是其实可以用软件调度的方式模拟出脉动阵列的数据流呢？把GPU里的线程之类的小结构视为脉动阵列里的PE，再从软件上按脉动阵列的原理去调度它们，是否也是可行的呢？我个人认为此类想法或多或少都是可行的，但得到的效果跟真正的脉动阵列电路会有本质区别，因为即使大体的数据流能够被软件调度所模仿出来，但数据的计算与流动的过程所涉及到的电路运转细节仍然跳不出GPU自身的硬件微结构所带来的限制，如“PE”间的数据传输，在GPU里可能实际上要经过共享内存访问、线程间信息同步等一系列过程，而真实的脉动阵列电路里的PE则具备直接的互连和简单的控制逻辑，所以相比专门设计出来的脉动阵列电路，GPU模拟的数据流在延迟、能耗等问题上仍然会有很大的差距。因此，虽然软件调度可以在一定程度上实现一些数据流的特殊设计，但是只有像定制化的AI芯片那样把概念上的数据流实现为特定的微结构，乃至特定的版图布局方式等，才能避免很多隐藏在底层的冗余开销，把计算的延迟、能效和面积优化到极致。定制化的硬件加速器还具备的一个好处是——由于详细的数据流等方面可能都已经能由加速器自身控制完成，所以AI编译器要负责的调度在一定程度上被简化，软硬件各司其职，高效工作。当然，硬件定制化程度的提高也就意味着市面上各种加速器会丰富多样，于是就像ZOMI酱所说，如何解决“软件碎片化”问题，让AI编译器能支持算法在多样化硬件平台上的部署，保证软硬件交互过程的高效性和可扩展性，是对于AI基础软件或AI基础设施（基础软件+底层硬件）的巨大挑战。