什么是AI框架、AI编译器、推理引擎

今天概念性的理解AI框架、AI编译器、推理引擎的关系。

一个深度学习框架有多种类型编译器：

• • 图优化器
• • JIT（Just-In-Time）编译器，动态从内置的算子库中选择最优的执行方式
• • 代码生成路径，如果算子库不支持某些操作，就需要使用编译器生成代码
• • 传统编译器，比如 LLVM

什么是图呢？它将深度学习模型抽象为一个图结构，用于表示模型的计算过程、数据流、控制流，图中的节点表示一个张量操作，比如GEMM；一些特殊的节点表示控制流，比如条件分支和循环；图中的边表示数据流动关系和依赖。

在编译器的语境中，表达式（expression）、原语函数（primitive function）、内核（compute kernel）在不同场景表达的含义并不一样，经常交替使用。

一个表达式由操作符（operator）和操作数（operand）组成，比如 sigmoid 表达式 y = 1 / (1 + exp(-x))，操作数是 x 和 y，而操作符则有 exp、add、div 等。

原语函数是编译器 IR 层或AI框架定义的基础单位，也叫做算子，如 exp(x)、add(x, y)，它们用于构建表达式，它们是构成图的原子节点。

内核（kernel）是算子（或融合后的算子）在特定硬件架构上的底层实现代码，比如 CUDA kernel 实现的 sigmoid 算子。

举个例子，sigmoid 本质上是由多个基础原语函数（算子）构成的数学表达式，由于使用非常频繁，所以编译器会将基础算子融合为一个单独的内核 kernel，可以避免中间张量频繁读写内存，提高运行效率。

另一个关键概念是编译器的调度（schedule），它描述一个算子或表达式的变换策略和执行顺序，比如如何分块、如何进行内存布局重排、是否融合。举个例子 GEMM 表达式是固定的，但不同的调度可能由多个内核（kernel）实现，会带来巨大的性能差异。

以 PyTorch 1.0 版本来理解AI框架、AI编译器，1.0 版本是动态图模式（Eager ），动态执行计算图，AI编译器的作用并不明显。

PyTorch 实现了上千个算子（operator），比如 torch.matmul，内置 JIT 图优化器，最核心的调度器（也叫运行时）是 ATen，它将每个张量操作分发（Dispatcher）到对应硬件的算子库（比如 cuDNN）。如果对应硬件没有对应的内核，则 PyTorch 会使用 TVM、XLA、LLVM 编译器动态生成代码，或者回退到 CPU 上的实现。

不同硬件的低级别算子库，比如 cuDNN、ROCm MIOpen 等，它们专门针对特定指令集架构上的深度学习操作或基础数学函数进行优化，通过和AI框架合作，共同优化框架中的图优化器，也会将常用的算子整合到底层库 API 中。

除了 AI 框架、AI 编译器，还有推理引擎，一般仅支持静态图模式，PyTorch 也可以算是一个推理引擎，但显然不适合在生产环境下应用模型推理。

推理引擎是仅用于推理的框架，专门优化和部署已经训练好的模型，一般追求极致的推理性能，专注于图优化，它后面也会对接AI编译器。

比如我介绍了大量 ONNX Runtime 内容，它就是一个通用型推理引擎，价值在于它提供了一个统一的运行时接口来执行符合ONNX标准的模型，并能灵活调用后端的优化引擎，比如 TensorRT。

不管是AI框架还是推理引擎，面临的一个核心问题是算子的 fallback 回退问题。

当前训练和部署模型的方式是使用带图优化器的框架或推理引擎，并依赖硬件底层库的内核实现，问题非常复杂。

如果有 F 个框架、M 种硬件平台、P 个算子、每个算子有 S 个调度策略、每个算子有 D 种格式（FP32/INT8/BF16），那复杂度就是 𝑂(FMPSD)。对于 AI 框架或推理引擎来说，优化一个算子组合非常费劲，比如卷积操作，有不同的变体，调度策略取决于不同批次大小、kernel size、通道数量以及不同的算法实现。

另外一个问题，图优化的融合算子硬件底层不一定有具体的实现，AI 框架或推理引擎辛辛苦苦的算子融合硬件不支持，考虑到不同硬件有不同的微架构、内存模型，一些设备还是异构架构，比如手机同时包含CPU、GPU、NPU，你怎么搞？

对于一个开发AI框架和推理引擎的开发者来说，在 PyTorch 实现一个算子，然后要为不同的硬件平台写内核，不同硬件平台的开发语言还不一样，所以现在AI编译器也发展起来了，它的核心目标就是针对不同的计算场景和不同的硬件，自动生成高效的代码。

AI编译器将特定框架的模型转化为一个图无关的中间表示IR，然后执行一系列与目标硬件无关的优化（Pass），比如算子融合、常量折叠，优化后生成一个高层IR；接着编译器将高层IR转化为一个没有控制流的低层IR，进行硬件相关的优化，比如针对GPU的数据布局转换；最后可以选择将低层 IR 调度到不同硬件上的内核库、或自动生成不同硬件平台可执行代码、或交给不同硬件平台对应编译器编译生成代码。

对于编译器来说，编译算子的挑战在于调度策略的空间是巨大的，找到最优解可能有几十亿种的选择，所以AI编译器也会内置算法策略，使用强化学习来改进编译器。

AI编译器另外一个挑战是动态图优化，静态图比较适合固定层数，每一层维度都固定的情况，比如 ResNet-50。但基于 transformer 的模型，输入长度都是可变的，具有相当复杂的控制流。

编译器和AI框架应支持以下特性：

• • 能够适用不同的AI框架和硬件，强调优化能力是可复用的
• • 支持各类型的张量，比如 int8、float16、fp4、fp8 等
• • 支持常见的算子，比如矩阵乘法、卷积操作等
• • Fallback 回退机制，编译器不支持的子图或表达式可以安全回退
• • 支持AOT（Ahead-of-Time）编译和JIT（Just-In-Time）编译
• • 支持静态图编译
• • 动态形状的支持，需要处理可变输入长度以及控制流
• • 跨设备、跨线程并行执行
• • 自动微分
• • 支持分布式训练集合通信原语