探索图编译器：实现高效的图计算

anqi · July 10, 2024, 9:06am

前言

在当前的数字化时代，图编译器（Graph Compiler）在AI、高性能计算等领域大放异彩，成为支撑算力根基的关键技术。图编译器通过把计算任务表示为计算图（Computational Graph），同时利用MLIR编译器框架将其分析并优化，将高层次的代码转换为图结构并部署在相应硬件平台，从而显著提升计算性能和效率。结合往期对Tenstorrent RISC-V 数据流芯片及相关软件栈的介绍，本文将系统阐述针对数据流芯片加速的图编译器的设计与实现。

背景知识

为了便于读者理解本文内容，首先介绍一些背景知识。

计算图

计算图（Computational Graph）是一种数据结构，用于表示和组织计算过程中的变量及其相互依赖关系。在计算图中，节点（顶点）通常代表计算操作，如加法、乘法、激活函数等，而边（有向边）表示节点之间的数据依赖关系，即一个操作的输出可以作为另一个操作的输入。这种图形化表示方法使得复杂的数学表达式和算法步骤的执行顺序变得直观易懂，并且是实现优化计算性能的关键工具。

图1 计算图简单示例

图1为计算图的简单示例，左侧展示了一个代码片段，由两个基本操作"+“和”*"组成，右侧是针对左侧代码分析得到的计算图。我们可以得到以下信息：
计算节点：

c = a + b，该操作使用a和b作为输入数据进行加法运算，c作为结果输出。
e = c * d，该操作使用c和d作为输入数据进行乘法运算，e作为结果输出。
数据流依赖：
数据从a和b开始，通过加法操作流向c，因此c依赖于a和b。
c的值与d结合进行乘法操作，最终得到e，因此e依赖于c和d。
以此类推，针对更复杂的程序，我们可以通过类似的分析，从代码中提取其中的图结构，为接下来的图编译器设计提供理论基础。

MLIR编译框架

对于所有编程语言的代码，我们都可以从中分析出数据流和控制流（程序执行过程中指令和操作的顺序和路径），这些关系使程序具有一定的图特性，也因此为图编译器的实现提供了分析基础。由此，我们可以合理猜想，是否可以设计出高效的图编译器，来应用于硬件仿真、物理粒子仿真和AI等多种领域？答案是非常乐观的。
为寻求统一各个领域的解决方案，我们考虑使用MLIR框架。MLIR（Multi-Level Intermediate Representation，多层中间表示）是Chris Lattner在谷歌开发的一种编译器基础框架，旨在为编译器和工具链提供灵活且高效的基础设施。它支持从高层抽象到低层机器指令的多层次表示，具有模块化和可扩展性的特点。MLIR在传统编译器、机器学习、深度学习和图计算等领域中应用广泛，通过一系列通用的编译优化技术提高代码性能，并拥有活跃的开源社区和丰富的生态系统。

图2 MLIR logo

MLIR为广大编译器开发者提供了许多基础功能，包括：

多层中间表示

抽象层次：支持从高层次的抽象表示（如张量运算）到低层次的机器指令表示，覆盖了不同的优化阶段和硬件目标。
中间表示：允许在同一框架内表示不同层次的IR（Intermediate Representation，中间表示），提供跨层次的优化和转换。

模块化和可扩展性

方言机制：MLIR通过方言机制（一种低层但与平台无关的IR表示）支持自定义操作和类型，使开发者可以轻松扩展和适应新的编程模型和硬件架构。
插件化架构：允许用户定义新的转换和优化模块，增强框架的灵活性和可定制性。

编译优化

通用优化：提供一系列通用的编译优化技术，如常量折叠、循环优化、内联扩展和死代码消除等。
跨层级优化：支持在不同抽象层级上进行优化，确保代码在高效执行的同时保持高层级的抽象。

图编译器介绍

图编译器（Graph Compiler）是一种在MLIR框架下，将程序代码转化为图结构并进行优化的编译工具。其核心思想是把程序通过图编译器转化为计算图相关的IR，使用图结构表示程序中的计算任务和数据依赖关系，即用节点表示计算任务，用边表示数据流，然后在此基础上进行代码优化与处理，最后将IR中计算程序和数据流依赖关系映射到数据流芯片中，充分发挥数据流芯片的加速性能。

基本功能

高层代码映射到图结构：将C/C++、Verilog等语言的代码以及AI模型等，通过Clang、MLIR、CIRCT和IREE等编译器转化为我们所需的MLIR中间表示，之后利用图编译器将此IR转化为图分析结构中，从而实现程序到图的映射。
图优化：在编译过程中，图编译器会对生成的图结构表示IR进行多种优化。图编译器利用MLIR框架的工具，通过操作IR执行一系列图操作和优化，如节点合并和路径优化等，有助于最大限度地利用硬件资源，确保程序在特定平台上的高效运行。
代码生成：图编译器将图描述IR的可执行部分转换为低层次的执行代码，同时将节点依赖关系以特定形式存储。这种方法不仅分离了调度信息和可执行代码，还便于数据流芯片高效地将程序部署到执行单元中。

优势

高效性：通过图结构和优化技术，显著提升程序的执行效率。依托于现在成熟的图领域相关研究，我们能够得到许多高效的图分割算法以及优化算法，同时结合代码特性，分割出负载更加均衡，依赖关系更少的图，以提高程序执行的并行性，从而图编译器的高效性，提高数据流加速性能。
通用性：支持多种计算任务，适应不同领域的需求。由于MLIR强大的可扩展性和高级抽象能力，MLIR框架为图编译器赋予灵活性。例如Xilinx设计AIE方言，用于描述和编译专为Xilinx AI引擎架构设计的计算任务；IREE基于MLIR框架，将机器学习模型降低为统一的IR，用于优化、编译和执行机器学习模型和硬件中的计算图；CIRCT通过利用MLIR的多级中间表示和优化功能，帮助开发者高效地设计、优化和验证复杂的硬件电路。由此可见，通过MLIR，我们可以获得多个领域的中间表示（IR），并利用MLIR框架的多级优化和转换能力，将这些不同的IR转化为统一的中间表示，这种统一的表示形式使得我们能够接入图编译器，让图编译器更具灵活性。

MAGIC图编译器的设计与实现

在往期文章探索 Tenstorrent 的 AI 开发软件栈：TT-BUDA中，我们展示过 Tenstorrent 的图编译器，但适用范围有限，仅适合AI模型的推理部署。为了扩大 Tenstorrent 数据流芯片的应用范围并提高优化效果，兆松自研了一款图编译器 MAGIC，技术栈如图3所示，包括传统前端软件、AI和硬件领域接入，MLIR各方言的转换与优化，以及后端对各种硬件平台的适配，充分展现了MAGIC图编译器的通用性和发展潜力。

图3 MAGIC图编译器技术栈

本节我们将结合 Tenstorrent 数据流芯片，对 MAGIC 图编译器的设计与实现做部分介绍（本文着重介绍 AI 部分，其余部分会在其他主题文章中阐明）。

图结构与图切分

在图编译器中，程序经过多层转换后被表示为Graph等方言，然后整个CFG（控制流图）会被切分为多个子图，如图4所示。其中：

节点（Nodes）：一个基本块即为一个节点。
子图（Subgraphs）：包含若干节点，用于映射到单个或多个Tensix核。
边（Edges）：表示节点间或子图间的数据流关系，描述计算单元的输出如何作为下个计算单元的输入。

图4 图结构

Netlist设计

图5 transformer 映射到 tensix cores

为了把程序在数据流芯片上高效的运行起来, 首先需要通过图编译器把计算图转化为物理图，然后映射到 Tensix core 计算阵列上执行，如图5所示。

图6 图编译器与数据流芯片的联系

Tenstorrent 把物理图称之为 Netlist（Netlist是连接图编译器前端和后端的桥梁，描述了程序是如何在芯片上映射和执行的，包括数据依赖关系、硬件配置、数据类型等信息，在往期文章深入解析 TT-BUDA 架构及编译执行过程中有详细介绍），然后再传入 Buda Backend 进行处理，如图6所示。为了充分发挥数据流芯片的潜力，兆松对 Netlist 做了定制。

方言设计

为了优雅的解决计算图的映射问题，我们提出并设计了一种专门表示图结构的方言——Graph方言，和描述算子的方言——MK（Magic Kernel）方言。
MK方言适用于AI领域的图编译，为图编译器提供统一的算子表示，通过将来自不同方言的算子映射到MK方言，方便把粗粒度的算子逐步降低（partial lowering）到细粒度的算子。而Graph方言适用于通用领域，目的是方便各种方言lowering到Graph方言后做统一分析与优化。
此处我们仅介绍Graph方言设计思路，该方言的设计目标是提供一种标准化的、通用的表示形式，使得图编译器能够高效地处理各种图结构。方言的设计思想如下：

抽象图结构：提供一种标准化的表示方式，为图结构设计相应的op，包括对节点、子图、边、数据传输等元素的定义，以IR的形式表示计算任务以及任务间的数据依赖关系。
针对不同领域：由于不同领域的代码结构不同，设计不同的图抽象方法以及分割算法，来适应各个领域的加速任务。
抽象硬件架构：为了适应不同硬件平台并确保图结构的高效映射和优化，我们需要将硬件架构进行抽象。通过对硬件结构的底层实现进行分析，我们可以抽象出对应的操作，以此提高图编译器的通用性和灵活性，使其能够针对各种硬件平台进行优化。
图相关操作：利用MLIR框架的特性，通过设计pass来操作IR，以达到图划分及优化的目的。

编译流程

图7 编译流程

程序经过编译器前端的处理，代码会被转换为 MLIR 方言，其中包含SCF、MK、LLVM方言以及用于生成Netlist的Graph等相关方言。编译流程如图7所示，下面分为三个阶段进行阐述：

图切分：

程序以SCF、MK、LLVM和Graph等方言作为输入，用Graph IR作为图结构的载体，存储关于图的所有信息。
在MLIR框架下应用图切分算法，通过操作Graph IR实现对图的切分和优化，将大图切分为多个子图，每个子图代表一段可独立执行的程序。
划分后的子图由Graph方言表示为独立的程序单元，从Graph IR中获取子图及子图间边的信息，并生成Netlist（表征子图之间的依赖和硬件映射关系，详细介绍请见深入解析 TT-BUDA 架构及编译执行过程）。

MLIR + LLVM C Backend：

使用MLIR提供的工具，将图中需要放入计算核中的IR进一步转换为标准的LLVM IR。
利用LLVM C Backend工具，将MLIR IR中的程序转换为C程序。之所以要转换为C程序，是因为目前 Tenstorrent 的软件工具链的后端仅接受C语言作为输入。

BUDA Backend：

BUDA Backend 是TT-BUDA的编译器后端，它接收由前两个步骤生成的所有C程序和Netlist。
根据 Netlist 中的映射关系，Buda Backend把C程序部署到Tensix流处理器阵列上。

技术挑战

传统的多核CPU和GPU使用虚拟内存，而且核间是遵循缓存一致性的，SRAM 是由缓存控制器这个硬件来管理的。但 Tenstorrent 数据流芯片有很大不同：Tensix 核间是通过路由转发数据的，并没有缓存一致性；SRAM 与 NoC 对程序员或编译器是可见的，SRAM 直接使用物理地址，并没有使用全局地址空间做内存管理。
如果从软件定义网络的角度来理解，数据平面是数据流芯片, 控制平面就是图编译器。源于数据流芯片的上述特性，图编译器会面临一些技术挑战，下面基于兆松MAGIC图编译的研发经验对此做些介绍：

控制流的处理

控制流的概念：控制流描述了程序执行的顺序与路径。在应用程序代码中，控制流由条件语句（例如 if 语句）、循环语句（例如 for 循环、while 循环）和跳转语句（例如 goto 语句或函数调用）等控制结构来控制。
为什么要做控制流的处理呢？有两个直接原因。一个是在编译时无法确定分支跳转方向，在运行时才能确定；另外一个原因是，Tenstorrent 数据流芯片没有使用全局地址空间做内存管理，在硬件层面不支持直接分支跳转。基于以上原因，我们需要对控制流做处理。处理的方式主要有两种，第一种是把条件跳转及其分支放入同一Tensix核中计算，使核间不存在条件跳转关系；第二种是使用间接方式做分支跳转。

指针的处理
程序中常存在指针操作，在执行时指针会随着数据流在核间传输。但从往期文章Tenstorrent数据流芯片Grayskull 和 Wormhole解析可知，在 Tenstorrent 数据流芯片中每个Tensix核拥有独立的SRAM，不能通过全局内存地址访问，所以我们需要消除指针或地址在 Tensix 核间的传输。
大子图的处理
某些子图可能会超出计算核SRAM的容量，需要图编译器把过大的子图拆解为多个小子图。
除了以上技术难点，还要满足一些约束条件：

负载均衡：将计算任务均衡分布到多个计算单元上，避免单个计算单元过载，从而提高整体计算性能。
最小化数据传输：减少不同计算单元之间的数据传输量，降低通信开销，提升系统整体效率。

总结

在AI和高性能计算领域，RISC-V与数据流芯片是未来大趋势，图编译器在其中扮演着关键角色。通过图编译器把程序转换为计算图，然后进行一系列优化，最后映射部署到数据流芯片上高效执行。在本期文章中，我们系统阐述了图编译器的设计与实现，包括各种技术挑战与解决思路。我们希望通过图编译器和数据流芯片的密切配合，共同筑牢算力发展根基，推动算力普惠化。