【AI编译器入门③】编译前端：ONNX、TorchScript与动态shape处理

【AI编译器入门③】编译前端：ONNX、TorchScript与动态shape处理

本文是「AI编译器」系列的第三篇。前端是AI编译器的入口------把各种格式的模型变成统一的IR。听起来像"格式转换"这么简单？才不是！前端还要处理算子映射、动态shape、自定义算子注册......稍有不慎，整个编译链路就会失败。本文把前端的三大核心问题逐一拆解。

---

一、写在前面：前端不只是格式转换

很多人觉得前端就是"把ONNX转成Relay"------就像把Word转成PDF一样，格式转换而已。

大错特错！

前端要处理的事情多着呢：

• 图结构提取：节点、边、输入输出关系
• 算子映射：ONNX的Conv对应TVM的什么算子？语义完全一致吗？
• 权重处理：initializer怎么转成TVM的常量？
• 动态shape：哪些维度是固定的？哪些是符号变量？
• 自定义算子：TVM不支持的算子怎么办？
• 类型推导：每个算子的输出shape和dtype是什么？

格式转换只是最表层的工作。前端的质量，直接决定了后面中端优化和后端代码生成能不能顺利进行。

先上一张前端Pipeline图：

---

二、前端Pipeline总览

ONNX / TorchScript / TensorFlow 模型
  → 前端解析
        ├── 图结构提取：遍历节点，建立输入输出拓扑
        ├── 算子映射：查映射表，把框架算子转成TVM算子
        ├── 权重处理：initializer/parameter → TVM Constant
        ├── 动态shape处理：确定哪些维度是符号变量
        └── 自定义算子注册：处理不支持的算子
  → 输出：Relay / Relax IRModule

简单说，前端就是个"翻译官"------把各种方言（ONNX、TorchScript、TF）翻译成统一的普通话（TVM IR）。

---

三、ONNX图解析

ONNX是模型交换的事实标准------几乎所有框架都能导出ONNX，几乎所有部署工具都支持ONNX。

3.1 解析流程

onnx.load("model.onnx")
  → 提取GraphProto（ONNX的图结构）
  → 遍历所有node，建立拓扑关系
  → 对每个node，查convert_map表，找到对应的TVM算子
  → 处理initializer（权重常量）
  → 处理输入shape/dtype，创建Relay/Relax Var
  → 构建Function和IRModule
  → InferType（推导所有张量的类型信息）

3.2 代码示例

import onnx
from tvm import relay

# 加载ONNX模型
onnx_model = onnx.load("resnet18.onnx")

# 解析为Relay
mod, params = relay.frontend.from_onnx(
    onnx_model,
    shape={"input": [1, 3, 224, 224]},
    dtype={"input": "float32"},
)

# 补全类型信息
mod = relay.transform.InferType()(mod)
print(mod)

3.3 opset版本适配

ONNX的算子语义会随opset版本变化------比如opset 11和opset 13的Resize语义就不一样。

TVM前端需要针对不同opset版本做适配。比如opset 13的bias_add和opset 11的实现可能不同。

💡 常见坑：ONNX模型的opset版本太低或太高，都可能导致前端解析失败。太低的话缺少某些算子，太高的话TVM还没适配。

---

四、TorchScript图解析

TorchScript是PyTorch的静态图表示。如果你用的是PyTorch模型，可以直接通过TorchScript解析，不用绕ONNX。

4.1 解析流程

torch.jit.trace(model, example_input)  # 用示例输入trace一次
  → 得到Scripted Model（TorchScript）
  → 遍历graph nodes
  → 对每个ATen算子，查映射表
  → 处理parameters为Constant
  → 构建Relay/Relax Function

4.2 代码示例

import torch
from tvm import relay

# 准备PyTorch模型和示例输入
model = torch.hub.load("pytorch/vision", "resnet18", pretrained=True)
model.eval()
example_input = torch.randn(1, 3, 224, 224)

# 导出TorchScript
scripted = torch.jit.trace(model, example_input)

# 解析为Relay
mod, params = relay.frontend.from_pytorch(
    scripted,
    shape_list=[("input", example_input.shape)],
)

4.3 ONNX vs TorchScript：怎么选？

特性	ONNX	TorchScript
通用性	跨框架，所有框架都支持导出	主要面向PyTorch
动态控制流	较弱	原生支持if/loop（用script而不是trace）
自定义算子	需要自定义op	可以直接导出ATen算子
生态工具	可视化工具多（Netron等）	需要看ATen图

一般建议：

• 如果模型是PyTorch的，而且结构比较常规 → 走ONNX（工具多、生态成熟）
• 如果模型有复杂控制流或自定义算子 → 走TorchScript（信息更完整）

---

五、动态shape处理：前端的硬骨头

动态shape是大模型推理的刚需------batch size、sequence length随时在变。

但动态shape也是前端的难点------shape不确定，很多优化就没法做。

5.1 Relay的动态shape策略（老办法）

Relay本质上是静态图IR，对动态shape支持较弱。常用的workaround：

方案一：固定shape编译

为每个常见的batch size单独编译一个产物。比如batch=1、batch=4、batch=8各编一个。

优点 ：性能最好（完全静态优化）

缺点：产物多，占空间；不常见的batch size要重新编译

方案二：bucket编译

预编译几个"桶"的尺寸，运行时选最近的那个桶。比如输入batch=5，就用bucket=8的产物（多余的padding掉）。

优点 ：产物数量可控

缺点：有padding浪费，性能略有损失

方案三：VM模式

用Relay VM（虚拟机）执行，支持动态shape。

优点 ：灵活，真正的动态

缺点：性能比全静态编译差不少

5.2 Relax的原生动态shape（新办法）

Relax从设计之初就支持动态shape------用符号变量表示未知维度。

from tvm import relax
from tvm.relax.frontend import nn

class DecodeStep(nn.Module):
    def forward(self, x: nn.Tensor):
        # x的shape是 [batch, seq_len, hidden]
        # seq_len是动态符号变量
        return nn.softmax(x, axis=-1)

model = DecodeStep()
mod, params = model.export_tvm(
    spec={"forward": {"x": nn.spec.Tensor([1, "seq_len", 768], "float32")}}
)

看到了吗？"seq_len"是个字符串，代表符号变量------编译期不知道具体是多少，运行时才知道。

5.3 动态shape的难点

动态shape听起来美好，但实现起来有很多挑战：

挑战	说明
内存规划难	编译期不知道具体尺寸，峰值显存没法精确估算
算子融合受限	动态维度可能破坏某些融合模式
后端支持不一	某些后端（如TensorRT）对动态shape支持有限，需要特殊处理
调度更复杂	TIR层的循环分块、线程绑定都要支持动态维度

这也是为什么Relax虽然好，但生态还在完善中------动态shape的全链路支持是个大工程。

---

六、自定义算子注册：遇到不支持的算子怎么办？

TVM不可能支持所有算子。遇到不支持的算子怎么办？自己注册一个。

6.1 算子映射表

前端维护了一个映射表，把框架算子名映射到TVM算子：

ONNX算子名          TVM算子
Conv           →  nn.conv2d
MatMul         →  nn.matmul / nn.dense
BatchNorm      →  nn.batch_norm
Softmax        →  nn.softmax
...

如果ONNX模型里有一个TVM不认识的算子，前端解析就会报错。

6.2 怎么注册自定义算子（Relax版）

Relax中注册自定义算子，通常通过call_pure_packed或call_tir：

from tvm import relax

# 注册一个自定义算子
@relax.op.register_op("relax.custom_op")
def custom_op(x):
    return relax.call_pure_packed(
        "tvm.contrib.custom_op",  # 底层的函数名
        x,
        sinfo_args=relax.TensorStructInfo(x.shape, x.dtype),  # 输出结构信息
    )

# 在前端中注册映射
relax.frontend.torch.op_mapping["aten::custom_op"] = custom_op

要点：

• 上层（Relax）定义算子的接口和语义
• 底层（PackedFunc或TIR）实现具体的计算
• 前端把框架算子映射到Relax算子

6.3 注册算子的四层完整链路

一个自定义算子，要在四层都注册才能跑通：

1. 前端层：把框架算子名映射到TVM算子
2. IR层：定义算子的语义、属性、类型推导
3. Compute/Schedule层：定义算子的计算和调度
4. 后端层：生成对应的硬件代码

只写kernel不注册映射？前端解析时就失败了，根本走不到后端。

---

七、常见误区

误区一："前端只是格式转换"

错！前端还要处理shape推断、dtype统一、权重常量转换、动态维度符号化、算子语义对齐。格式转换只是最表层的工作。

误区二："动态shape直接交给Runtime就行"

完全动态shape在编译期无法做内存规划和算子融合，会牺牲性能。Relax的做法是保留部分编译期知识（符号关系），而不是完全Runtime化。

误区三："自定义算子只需要写kernel"

错！自定义算子需要在前端、IR、compute/schedule、后端四层都注册。只写kernel，前端解析不过，根本用不了。

---

八、总结一下

今天这篇文章，我们讲了AI编译器的前端：

1. 前端不只是格式转换------还要处理算子映射、动态shape、自定义算子等
2. 两种主流输入：ONNX（通用、生态好）和TorchScript（PyTorch原生、控制流好）
3. 动态shape处理：Relay用bucket/VM workaround，Relax原生支持符号变量
4. 自定义算子注册：从前端到后端四层都要注册，缺一不可

下一篇预告：我们深入编译中端------图优化Pass。DCE、常量折叠、算子融合、布局转换......这些Pass都是干什么的？能带来多大收益？

---

📚 本系列导航：

• 第1篇：AI编译器基础：三段式Pipeline与全链路协同
• 第2篇：TVM架构拆解：Relay、Relax、TIR与Runtime
• 第3篇：编译前端：ONNX、TorchScript与动态shape处理（本文）
• 第4篇：编译中端：图优化Pass全解析
• 第5篇：编译后端：CUDA Codegen、CUTLASS与TensorRT
• 更多内容持续更新中......