在尝试将pytorch接入aie平台的过程中，发现自己对于AI编译器在机器学习系统中的定位所知甚少。由于先前经验大多只停留在将高级语言转换成硬件相关的ir代码，而忽略了后端以及运行时系统协作方面的知识，遇到瓶颈。因此本博客通过参考机器学习系统书籍，以及目前主流AI编译器的后端和运行时系统的协作，查漏补缺，补全自己的知识盲点。
AI编译器后端和运行时架构概览
通过参考mlsys open book和tensorflow xla解读，对机器学习编译有了新的理解。机器学习编译技术需要与运行时系统相互写作，才能将一个机器学习模型完整的跑在GPU/NPU上。下图是AI系统的一个软件栈示意图：
上述优化流程，是针对静态图的全图编译运行流程（tensorflow XLA这种动态圈图有一定区别）。按照软件栈层级关系，可以分为编译层和运行时系统层。
- 编译层：编译层的输入是整个静态计算图或是部分子图，因此首先需要编译器前端用户代码进行解析翻译得到计算图IR，并对其进行设备信息无关的优化，此时的优化并不考虑程序执行的底层硬件信息。编译器后端的主要职责是对前端下发的IR做进一步的计算图优化，让其更加贴合硬件，并为IR中的计算节点选择在硬件上执行的算子，然后为每个算子的输入输出分配硬件内存，最终生成一个可以在硬件上执行的任务序列。
- 运行时系统层：运行时系统设计一般均有一个虚拟机VM用以支持调度和快速部署（IREE架构的VM和Relay VM）。其输入是编译层生成的可在硬件上执行的任务序列，如果考虑跨硬件部署，则需要考虑使用函数类型擦除技术统一接口。然后，VM将任务序列转换成硬件设备api和host端代码，最终分别生成可执行。
后续会详细解读每一层的流程，以及可参考文献。
编译器后端
后端优化定位
重点参考zomi编译后端视频。声明：如下内容以及图片slides基本来源于该视频。
上图是编译后端在整个AI编译框架的位置。大致总结一下前端优化和后端优化的区别：前端优化输入是计算图，关注计算图整体拓扑结构，而不关心算子的具体实现。其优化手段主要有对于算子节点进行融合、消除、化简等操作，是计算图计算和存储开销小。后端优化则专注于算子节点的内部具体实现，针对具体实现使得性能达到最优，重点关心节点的输入，输出内存循环方式和计算的逻辑。
后端优化干了什么？
后端优化主要做了三件事：
低级代码生成
如上图所示，是TVM的架构流程。Tensor IR是将计算图用张量表示。低级代码生成阶段是将该张量表示进一步降级。
后端优化
底层硬件不同，对于conv等算子的排布优化也会不同。后端优化主要针对各个算子做循环优化排布。
代码生成
针对不同backend生成机器码。在AI编译器中，以GPU为例，一种思路是GPU端直接生成cuda代码，CPU端生成LLVM，并由llvm基础设施编译可执行。另一个思路是生成NVVM IR，交由nvvm进行编译。
所以，后端优化其实本质就是对于算子的编译优化。目前算子编译优化的挑战如下：
算子库因此应允而生：
算子库能够很好地压榨硬件性能，但是上层模型的变动和底层硬件的变动，算子库无法快速适配这些变化。因此AI编译器和算子库协作开发，由AI编译器提供灵活性，减少一部分算子库开发的人力投入。编译角度的后端优化可以分为两个路线：
- 使用机器学习等模型进行auto tuning
- 利用数学表达式，做多面体优化
这两个部分后续会专门开专题来解读。
运行时系统
这一章节关于运行时系统的描述，主要参考Lei chat blog，本章的描述均基于单节点来讨论，服务器等多节点暂不涉及。
首先，我们需要明确，什么是机器学习系统的runtime system。runtime system的主要工作有两点：
- 管理资源（内存资源等）
- 调度执行
根据实际语义，runtime亦有不同。多卡分布式训练场景下，runtime负责将训练任务分配调度到数据中心的不同节点。单节点上，runtime负责派发张量计算到特定的加速器。本章节关注单节点runtime系统，其一般运行在CPU上来管理资源并调度机器学习任务负载到一个或多个加速器上。
运行时抽象
运行时系统给用户提供了三个抽象：
主机设备分离
摘自博客的解释：
虚拟机抽象
这一部分是我们理解整个runtime的重点，具体的实际例子可以参考IREE 和tvm runtime设计细节。针对机器学习系统以及我们的cpu加速器协同架构，我们会有如下insights：
- 首先我们观察主机设备架构，会发现相比设备端，主机端的系统架构是比较统一的（ARM64， x86_64和RISC-V）。
- 观察机器学习的任务，发现机器学习任务一般是完美循环（定义见llvm官方文档）。因此访存方面用逻辑偏移即可，而非物理指针。
- 机器学习模型部署在端侧，因此轻量化也是一大考量。
基于上述insights，主流ai编译器均定义一套主机端虚拟指令集和虚拟机来实现调度。计算密集任务放在加速器上，存储资源追踪管理以及维度形状计算等逻辑（适合cpu端计算的逻辑）放在虚拟机上。如下图所示是IREE的虚拟机架构图：
上述虚拟机设计，还有一大好处，那就是可以很自然的完成从mlir系统到虚拟机指令ISA的编译转换。因此整个调度控制流程可以自然地通过编译器来控制，进一步提供可优化点。
HAL硬件抽象层
针对边缘/客户端设备的多样化硬件生态，最佳实践是基于现代图形/计算API（如Vulkan、DirectX 12、Metal）构建硬件抽象层（HAL）。这些API提供低开销、跨平台支持、显式控制等优势，能统一管理不同厂商的GPU/加速器（如NVIDIA/AMD/Apple）。虽然这些API直接使用复杂（如代码冗长），但通过编译器自动生成调用逻辑，可绕过开发难度，同时保留底层控制权。这种设计既能满足高性能需求（通过低层级优化），又能覆盖多平台/硬件的兼容性（如Windows/iOS/Android），为碎片化设备提供统一且高效的计算支持。如下是iree的HAL，VM架构图。
上述是runtime system的三个抽象，而实际生产中，runtime系统也要充分考虑性能问题：
在GPU资源管理与同步中，核心挑战在于平衡计算吞吐量与系统整体效率。由于GPU内存性能常落后于算力提升，单纯优化单内核代码（如矩阵乘法）难以发挥理论峰值，需避免运行时因频繁创建/销毁资源（如缓冲区、同步对象）或跨组件数据拷贝/强制等待导致的性能损耗。为此，可借鉴图形应用（如游戏引擎）的成熟实践：预分配资源池减少动态开销、预录异步命令链最大化GPU并行度、细粒度同步基于资源生命周期。这些策略可通过编译器自动化实现（如依赖分析与指令提升）。IREE运行时采用异步优先设计，支持跨组件无缝传递缓冲区与同步原语，利用操作系统底层机制（如Linux的dma_fence）避免硬性边界，既保障ML模型的高效执行，又能嵌入复杂应用（如图像处理管线）而不成为性能瓶颈，兼顾性能与系统协同。
上述关于运行时系统的描述，是比较普遍的，下面结合Open MLsys一书，详细解读不同运行时系统的分类。
经过算子选择与内存分配之后，计算任务可以通过运行时完成计算的调度与在硬件上的执行。根据是否将算子编译为计算图，计算的调度可以分为单算子调度与计算图调度两种方式，例如在MindSpore中分别提供了PyNative模式和Graph模式。而根据硬件提供的能力差异，计算图的执行方式又可以分为逐算子下发执行的交互式执行以及将整个计算图或者部分子图一次性下发到硬件的下沉式执行两种模式。 —- 《Open Mlysy》
单算子调度
一行行算子line by line的交由python运行时执行，好处是方便debug和看到即时效果，坏处是只能顺序执行，丧失很多并行性，同时也无法支持编译优化。
计算图调度
在一个典型的异构计算环境中，主要存在CPU、GPU以及NPU等多种计算设备，因此一张计算图可以由运行在不同设备上的算子组成为异构计算图。
所述计算图由如下几类异构硬件对应的算子组成：
- CPU算子：由C++语言编写实现并在主机上通过CPU执行的算子，CPU计算的性能取决于是否能够充分利用CPU多核心的计算能力。
- GPU算子：以英伟达GPU芯片为例，通过在主机侧将GPU Kernel逐个下发到GPU设备上，由GPU芯片执行算子的计算逻辑，由于芯片上具备大量的并行执行单元，可以为高度并行的算法提供强大的加速能力。
- NPU算子：以华为Ascend芯片为例， Ascend是一个高度集成的SoC芯片，NPU的优势是支持将部分或整个计算图下沉到芯片中完成计算，计算过程中不与Host发生交互，因此具备较高的计算性能。
- Python算子：在执行模式上与CPU算子类似，都是由主机上的CPU执行计算，区别在于计算逻辑是由Python语言的运行时通过Python解释器解释执行。
一般主流架构均提供了指定算子所在运行设备的能力，此处可以参考mindspore的具体实现。
完成计算图中算子对应设备的标记以后，计算图已经准备好被调度与执行，根据硬件能力的差异，可以将异构计算图的执行分为三种模式，分别是逐算子交互式执行，整图下沉执行与子图下沉执行。
- 交互式执行主要针对CPU和GPU的场景，计算图中的算子按照输入和输出的依赖关系被逐个调度与执行；
- 而整图下沉执行模式主要是针对NPU芯片而言，这类芯片主要的优势是能够将整个神经网络的计算图一次性下发到设备上，无需借助主机的CPU能力而独立完成计算图中所有算子的调度与执行，减少了主机和芯片的交互次数，借助NPU的张量加速能力，提高了计算效率和性能；
- 子图下沉执行模式是前面两种执行模式的结合，由于计算图自身表达的灵活性，对于复杂场景的计算图在NPU芯片上进行整图下沉执行的效率不一定能达到最优，因此可以将对于NPU芯片执行效率低下的部分分离出来，交给CPU或者GPU等执行效率更高的设备处理，而将部分更适合NPU计算的子图下沉到NPU进行计算，这样可以兼顾性能和灵活性两方面。
虽然在计算图上可以充分表达算子间的并发关系，在实际代码中会产生由于并发而引起的一些不预期的副作用场景
如上图所示，虚线是operator之间的隐含依赖关系，但是在计算图中难以表示。因此在实际执行计算图的过程中，会分为交互式下发和下沉式（执行方式区分）。
算子交互执行
框架的运行时根据计算图中算子的依赖关系，按照某种执行序（例如广度优先序）逐个将算子下发到硬件上执行。我们处理的计算图一般为异构计算图（计算图中的硬件有多种）。解读异构计算图交互执行之前，先解读同构计算图，异构计算图需要先通过子图切分为同构计算图。
同构计算图有两种执行方式，如下表所示：
| 执行方式 | 串行执行 | 并行执行 |
|---|---|---|
| 算子执行顺序 | 固定 | 不固定 |
| 算子执行线程 | 单线程 | 多线程 |
| 所需执行资源 | 较低 | 较高 |
而对于异构计算图，一般来说计算图的优化都是基于非异构计算图来实现的，要求计算图中的算子为同一设备上的，方便算子间的融合替换等优化操作，因此需要将一张异构计算图切分为多个非异构计算图，这里切分就比较灵活了，可以定义各种切分规则，一般按照产生尽量少的子图的切分规则来切分，尽量将多的同一设备上的算子放在一张子图中。如下图是一种可能的子图切分方式：
切分出子图后，下一步需要考虑的是子图执行。共可以分为两种策略：
子图拆分执行：将切分后的多个子图分开执行，即一个子图执行完再执行另一个子图，上一个子图的输出数据会传输给下一个子图的输入数据，并且下一个子图需要将输入数据拷贝为本图的device数据，如Graph_2_GPU需要将Graph_1_CPU的输出数据从CPU拷贝到GPU，反过来Graph_3_CPU需要将Graph2GPU的输出数据从GPU拷贝到CPU，子图之间互相切换执行有一定的开销。
子图合并执行：将切分后的多个子图进行合并，合并为一个整体的DAG执行，如 图7.5.12所示，通过算子的设备属性来插入拷贝算子以实现不同设备上的算子数据传输，并且拷贝算子也是进入整图中的，从而形成一个大的整图执行，减少子图之间的切换执行开销。
下表中对于这两种执行方式进行了对比：
| 执行方式 | 子图拆分 | 子图合并 |
|---|---|---|
| 异构数据传输 | 子图之间拷贝 | 算子之间拷贝 |
| 执行额外开销 | 子图切换执行开销 | 无 |
| 执行并发粒度 | 子图并发 | 算子原生并发 |
总结一下，主要是并发粒度的区别。子图切换执行开销是比较大的，因此子图合并是比较普适的做法。
mindspore的异构执行是子图合并，而非异构执行是并行执行，因此为子图合并并行执行（任意计算架构都有两个象限）。
下沉式执行
下沉式执行是通过专用芯片的SoC架构，将整个或部分计算图一次性调度到芯片上以完成全量数据的计算。
可重构芯片编译器
先前一章节阐述的是目前主流成熟的AI编译器后端和运行时架构，而面向FPGA或是AIE这样的可重构硬件，往往无法直接复用ASIC NPU的编译流程（缺乏成熟的runtime系统）。本章节结合Efficient processing for DNN的chapter6，讲解如何将一个dataflow graph映射到这些硬件上去。
上图摘自cornell的机器学习硬件课程slides。在众多参考资料中，将dataflow映射到FPGA或是特定NPU（如AIE），被称为Mapping，而不是Compiler。如下引用Efficient processing for DNN一书中的一段话，讲明两者的类比关系：
In conventional computer systems, the compiler translates the program into machine-readable binary codes for execution; in the processing of DNNs, the mapper translates the desired DNN layer computation (i.e., problem specification) along with its shape and size4 into a hardwarecompatible mapping for execution. While the compiler usually optimizes just for performance, the mapper will typically optimize for performance and/or energy efficiency.
总结一下，有如下异同点：
- 输入方面，传统编译是程序，而DNN加速器映射是问题特性以及输入的shape。
- Compiler的结果是可执行程序，而mapper是硬件的配置文件。
- 需要的硬件信息方面，编译器需要知道处理器的架构和微架构，而mapper需要知道DNN加速器的数据流和约束。
上图是计算图编译器的一个典型架构。
NPU/GPU典型编译器解读
TVM 后端架构&Runtime
运行时框架图
TVM编译器的一大优势是支持多硬件平台/设备部署，因此其运行时系统设计是重中之重。其运行时系统可以大体拆解为两个部分：
- 每个硬件继承自ModuleNode类，利用PackedFunc技术将特定device的设备API封装成统一格式的指令（类似bladeDISC等编译器的类型擦除设计）。
- 一个Executor调度每一个指令并运行。该Executor分为GraphExecutor（针对静态图）和Relay VM（针对动态图）。
如下图所示分别是TVM运行时框架和Relay VM示意图：
IREE runtime设计
这部分内容主要参考IREE runtime 设计slides。IREE项目整体的设计思路十分完整庞大，设计巧妙，需要单独开一章节讲解。
BladeDISC后端架构&Runtime
BladeDISC的key feature
BladeDISC项目在编译后端和runtime上，有许多巧思（动态shape引发的编译和运行时协作代码生成）。具体地参考BladeDISC papaer和BladeDISC github project，以及BladeDISC内部分享。
后端代码生成
BladeDISC是tensorflow的backend来使用，因此其代码生成是针对融合后的算子做代码生成即可。运行时可以依赖于tensorflow的runtime调度器，launch 代码生成的kernel。
在后端代码生成过程中，由于动态shape问题，主要有如下考虑：
针对存储密集子图，考虑指令重排和重叠来减少存储overhead。由于shape是动态的，传统的unroll操作可以抵消动态shape的部分影响。
针对存储密集子图，多代码生成和运行时预测机制，来确保针对特定shape选择到最合适的生成代码。
如上图所示，CPU端在做GPU kernel launch的时候，由于需要运行时speculation，所以会有一定延迟（microsecond），该时间GPU kernel可以overlap掉。多代码生成会考虑如下三个因素：
- 向量化版本 or 非向量化版本
- 隐藏broadcast删除 or 不删除
- 针对行归约操作，one block one row（行少列多，高并行，低tail latency） 或是 one warp one row（行多列少，低调度开销，高吞吐）的选择
针对计算密集子图，有如下insight：
计算密集操作后续一般紧跟elementwise操作，称elementwise为epilog。计算密集算子计算量相对更大，内存获取也更多，因此针对这种组合，主要针对计算密集做代码生成策略选择，epilog沿用即可。
计算密集算子的bottleneck是动态shape下，不同shape的算子编写性能区别很大，因此编译时做面向不同类型shape多代码生成，运行时预测选择。其流程如图所示：
可以看到，BladeDISC基于CUTLASS构建数据集，训练一个决策树模型，运行时决策出选用哪个生成代码。
Runtime设计：RAL
上图是RAL的完整设计，有如下特点：
- 有别于其他AI编译器（TVM和IREE），BladeDISC通过整体代码生成（张量计算生成device代码，运行调度生成主机host代码），减少虚拟机的解释执行开销，同时也增加了编译和运行时协同优化可能。
- 前端，tensorflow和pytorch都有对应的IO来接入。后端，每个设备都有抽象，能够支持内存管理，任务调度同步等api。
- 不论是主机端还是设备端，最后都依靠llvm来生成对应架构的二进制文件。
根据阿里PAI团队的slides做进一步理解：
上述slides阐述了如何对于编译生成的.exe做runtime，后续会结合代码进一步挖掘。
整体流程
可以清晰看到，分为独立模式和插件模式（深度集成Tensorflow和pytorch框架）。
AI芯片软硬件接口设计
前面几个章节介绍了AI编译器后端和运行时如何协作，同时还举例多种AI编译器的后端/运行时例子（IREE,TVM,可重构芯片编译器）。然而，运行时如何驱动底层driver完成整个计算流程尚不清楚。理解清楚这一部分内容，需要首先对于AI芯片软硬件接口有一定了解。