由于AI芯片的特殊性和高度定制化，为了兼容硬件的多样性，AI模型必须能被高效地映射到各种AI芯片上。AI编译器将深度学习框架描述的AI模型作为输入，将为各种AI芯片生成的优化代码作为输出。AI编译器的目标是通过编译优化的方法将深度学习框架产生的AI模型转化为与特定架构的AI芯片适配的可执行机器码。
主流的AI硬件厂商都提供了高度优化的算子库或者推理引擎来实现高效计算，但是依赖算子库的缺点是算子库覆盖的硬件范围有限，而且库函数的更新落后于模型的发展速度，无法充分发挥AI芯片算力。算子与硬件算子库高度耦合会导致算子不可分解，从而影响算子进一步优化，使其难以在硬件平台上高效执行。为了解决这些缺点，业界设计了AI编译器。AI编译器可以在代码生成期间生成对库函数的调用，这样工作量就可以从编译器转移到算子开发上了。相反，如果编译器有强大的代码生成能力，则对内核函数的优化的依赖也会减小。

AI编译器结构

常规的编译器包含三个部分：前端、优化器和后端。前端主要负责将输入的源代码解析为抽象语法树，做语法分析、语义分析等；优化器在前端的基础上进行优化；后端会尽可能利用目标机器上的特殊指令优化中间代码，并将其转化为指定硬件平台的机器码。该三部分的数据传输是通过IR(Intermediate Representation)传输的，IR可以保证编译器的跨平台。
AI编译器一般采用分层的设计，通过多级IR设计，实现针对AI模型的特定优化设计。高阶IR服务于前端，执行硬件无关的优化，低阶IR服务于后端，执行硬件相关的编译、优化和代码生成操作。高阶IR也叫图IR，旨在抽象计算和控制流，其目的是建立控制流及算子与数据之间的依赖关系，并为图优化提供接口。高阶IR还需要对数据张量和算子进行支持。低阶IR主要设计用于各种硬件相关的优化和代码生成。低阶IR会更加细粒度的反应硬件特性。低阶IR还需要兼容第三方工具链，利用已有编译工具完成通用优化和代码生成。
AI编译器的输入是深度学习框架的模型，然后AI编译器将输入转化为高阶IR。AI编译器前端会结合通用编译器优化和AI特定的图优化方法对计算图优化生成优化后的计算图，该部分的优化是硬件平台无关的优化。AI编译器的后端优化是在高阶IR转化为低阶IR后，利用硬件的先验信息，通过定制化优化pass将优化后的代码实现映射为AI芯片的执行指令。

TVM架构

TVM是一个端到端的全栈编译器，提供端到端的编译优化。TVM以AI模型作为输入，首先将其转化为计算图，然后执行高级数据流重写，为计算图生成优化图。算子级优化模块为优化图中每个融合算子生成高效代码，并以声明式张量表达式指定算子。TVM会给指定硬件目标算子建立可能的优化集合，然后使用基于机器学习的代价模型搜索优化算子。最后将生成的代码打包到可部署模块中。
代码生成是将TIR表示编译成目标硬件平台的机器码。TVM支持多种后端，包括但不限于：
LLVM：可以针对任意微处理器架构生成代码，包括标准x86和ARM处理器，AMDGPU和NVPTX代码生成，以及LLVM支持的任何其他平台；专门的编译器，如NVCC（NVIDIA的编译器）；嵌入式和专用目标，通过TVM的Bring Your Own Codegen（BYOC）框架实现。因此，最终的输出是针对特定硬件平台优化的机器码，它可以是一个可执行文件、动态链接库（DLL）、静态链接库（LIB）或其他硬件平台支持的格式。这种格式的代码可以直接在目标硬件上运行，以实现高效的模型推理。

TVM图级优化

TVM可以对计算图做各种图优化，按照优化范围可以分为局部优化和全局优化。局部优化包括算子融合、代数简化、常量折叠，是图级优化的重点。全局优化是在整个计算图中搜索特定特征，并对这些特征执行优化操作，例如死代码消除、公共子表达式消除。
高阶图优化的最好结果只能达到与算子库优化相同的性能。随着引入的算子越来越多，可融合的内核数量急剧膨胀，并且随着硬件后端种类不断增加，开发者将无法依赖算子库提供融合能力。考虑到每种硬件后端上算子手工实现优化内核会因工作量太大而不切实际。因此，TVM提供了代码生成方法为指定的AI模型算子生成可能的实现。在代码生成期间生成对库函数的调用，这样工作量就可以从编译器转移到算子开发上了。

计算与调度

TVM提供了张量表达式语言，为自动代码生成提供支持。张量表达式语言的每个计算操作分为两个部分：第一个部分指定输出张量形状，第二部分描述张量中每个元素的计算规则。张量表达式并没有指定执行细节，这为不同硬件中的优化实现提供可能性，开发者可以通过lambda表达式快速定义计算，而无需实现新函数。不同实现方法性能存在巨大差异，因此，TVM要求开发者指定如何执行计算，如访问数据的顺序、多线程并行的方式。TVM中将这种计算的实现称为调度(严格意义上来说是执行计算的实现的规划)，由tvm.te.schedule.Schedule对象表示，其中包含若干阶段，每个阶段对应一个描述调度方式的操作。
TVM将计算和调度分离，同一种计算可以通过不同的调度实现。计算只定义了结果的计算方式，计算结果与运行平台无关。计算的实现由调度决定，调度取决于硬件，但不可影响结果的正确性。调度可以表示（为）张量表达式到底层代码的特定映射。TVM提供了许多调度原语为各种后端实现高性能代码。

自动调优框架

开发者需要对调度进行优化才能更好发挥硬件性能。这种优化可以是基于经验的手动优化，但是需要为每种硬件提供调度优化策略和调度相关参数，这样就会形成一个巨大的算子实现搜索空间。TVM提供了自动调度优化框架，可以为算子提供高性能底层代码实现。自动调优过程分两步：第一步定义搜索空间，第二部运行搜索算法。定义搜索空间可以看作是对调度的参数化，可将固定的调度改为定义调度策略的可调调度模板，通过该可调调度模板，从候选参数中选择在不同目标硬件上最优的参数组合。
TVM提供了四种调优器在搜索空间中找到最优调度，

• RandomTuner：随机顺序遍历配置空间；
• GridSearchTuner：以网络搜索顺序遍历配置空间；
• GATuner：用遗传算法搜索配置空间，该方法无代价模型，只能在真实机器上测试；
• XGBTuner：该优化器的代价模型基于模拟退火算法，使用XGBoost算法训练模型，然后在配置空间中测试最优配置。

在TVM中，搜索的主要是以下几个方面：

• 调度参数：这些参数定义了如何在目标硬件上执行计算。它们包括循环的重排、并行化、向量化、展开、融合等循环变换。
• 内存访问模式：不同的内存访问模式对性能有显著影响。搜索空间可能包括不同的内存访问模式，例如全局内存访问、共享内存访问、寄存器访问等。
• 算子实现：对于某些操作，可能有多种实现方式。搜索空间可能包含不同的算子实现，以找到最适合特定硬件的版本。
• 资源分配：这涉及到如何在硬件上分配计算资源，例如线程数、块数等。
• 精度配置：在某些情况下，可以通过降低操作的数值精度来提高性能，搜索空间可能包括不同的精度配置。
• 混合精度：在搜索空间中，还可以探索使用不同精度（如FP32和FP16）的组合来平衡性能和精度。
• 自动调度算法：TVM提供了AutoTVM和AutoScheduler两种自动调优模块。AutoTVM是基于模板的调优模块，而AutoScheduler（Ansor）是无模板的调优模块，它们都通过搜索算法来探索搜索空间。

TVM两级IR

TVM有两级IR：Relay IR和张量级IR。Relay IR作为一种高阶图级IR，旨在最大限度利用函数式编程语言、类型系统和编译器技术的研究成果，改善模型部署开销。Relay IR的表达同时采用基于DAG的IR和基于let-binding的IR。DAG即有向无环图，可以明确表示算子之间的依赖关系，但是缺乏计算范围的定义，可能会出现二义性问题。let-binding通过在有限范围内为某些函数提供let表达式来消除语义二义性。高级IR有助于执行通用优化如内存重用、布局转化、自动分区等。
高阶IR实现包括数据表示和算子实现。TVM的数据表示包括张量数据表示、形状表示、数据布局和边界推断等。TVM中数据通常以张量形式组织，通过占位符表示，持有明确的形状信息。通过占位符，开发者可以在不考虑具体元素的情况下操作计算图，实现了计算的定义与执行的分离。
数据布局是作为算子的参数写入的，以便算子的实现和减少编译开销。数据布局包括维度顺序、分片、填充和跨距等。不同的硬件最优数据不同，执行性能也不同。Relay算子受数据布局影响程度各不相同，布局转换本身会消耗资源引入开销，TVM将Relay算子分为对布局无感、对布局轻度敏感和对布局重度敏感三类。TVM提供了ConvertLayout Pass，希望能以最少的数据布局转化次数改变整个图的数据布局。
在高阶IR的算子实现中，开发者需要描述计算和调度，声明输入输出的形状。TVM添加新算子各层IR都需要改动，图级IR需要为新算子设置计算规则，张量级IR需要为每个硬件平台实现对应内核。
低阶IR以更细粒度的表示形式描述计算，通过提供计算调优和内存访问接口实现目标相关的优化。经过Relay优化和降级之后，优化的算子通过TIR降级为C++/CUDA，或者降级为LLVM IR，然后通过NVCC或者LLVM等后端优化器和代码生成器产生机器码。

运行时定制

TVM运行时是部署和执行已编译模块的主要方式，其目标是提供可与前端语言交互的API集合。TVM运行时有两个基础模块：PackedFunc和ModuleNode。TVM代码生成将已编译对象封装为Module对象，并以PackedFunc对象形式返回其中的已编译函数。自定义硬件平台后端应派生ModuleNode子类实现各自的运行时模块，并在其中添加目标相关的运行时API调用。

• 使用TVM运行时：编译后的模型可以生成为一个可执行文件或者库文件，然后通过TVM运行时来加载和执行。TVM运行时是一个轻量级的模块，提供了C API以及Python和Rust等语言的绑定，用于动态加载和执行编译后的模型。
• 集成到应用程序：编译后的模型可以被集成到应用程序中，例如集成到移动应用或嵌入式设备中。这种方式下，应用程序可以直接调用编译后的模型执行推理。
• 使用Graph Executor：TVM提供了Graph Executor，它是一个执行引擎，用于执行编译后的模型。Graph Executor可以加载编译后的模型，并执行推理过程。
• 使用RPC服务：TVM支持远程过程调用（RPC）服务，允许模型在一台服务器上编译，然后通过网络在另一台设备上执行。这种方式适用于模型部署在不同的硬件平台或分布式系统中。
• 使用WebAssembly：TVM还支持将模型编译为WebAssembly格式，使得模型可以在Web浏览器中运行，无需后端服务器。
• 使用TVM微控制器运行时：对于微控制器和嵌入式系统，TVM提供了专门的运行时系统，用于在资源受限的设备上执行编译后的模型。
• 使用TensorRT：对于NVIDIA GPU，TVM可以与TensorRT集成，利用TensorRT的高性能推理优化来执行编译后的模型。
• 使用OpenCL：在支持OpenCL的设备上，TVM可以编译模型并利用OpenCL运行时来执行模型。
• 使用Metal：对于苹果设备，TVM可以编译模型并使用Metal API来执行模型。

TVM前后端优化

计算图到目标代码的过程中可以分解若干个独立的步骤，这些步骤成为pass。RelayIR和TIR都包含一系列pass，为了使这些pass正确顺序执行，需要一个管理器安排pass执行顺序，在TVM中称作TVM pass基础架构。通过pass基础架构可以管理不同层次IR的优化pass。

前端优化

前端优化时通过遍历Relay IR，并从中捕获特定计算特征，以此指导图重写和图转换。前端优化与硬件无关只适用于计算图。
优化pass可以分两种：函数级优化和模块级优化。函数级优化实现Relay或TIR模块中各函数内部优化，没有全局信息，不能添加或删除函数。如常量折叠、死代码消除。模块级优化用于实现过程间优化和分析，作用于IRModule对象上，可以对任何函数体做修改。如移除无用函数。

后端优化

TVM后端优化不仅针对低阶IR做优化，更针对不同硬件目标，结合硬件特性和后端优化技术，实现高效代码生成。后端优化可通过不同方式实现。一种是针对自定义加速器硬件，利用AI模型特性设计定制化的后端优化方法。第二种是将低阶IR转化为LLVM IR，利用LLVM基础结构生成优化的硬件可执行机器码。第三种是将低阶IR转化为内核源代码如CUDA，在利用已有的编译器如NVCC生成可执行代码。

AI推理引擎与AI编译器的区别

简单来讲，推理引擎就是对前端模型进行解释执行，一般是JIT的编译方式，提供了优化的kernel实现和设计的IR，并且IR一般是一种。AI编译器同样完成了对前端模型的执行操作但是一般包含多级IR，每级IR有不同的优化称作Pass，可以前端模型转化为C代码或者其他可执行代码，也可以转化为某种IR然后交由Runtime来解释执行。