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

发布时间：2026/10/29 1:56:16
【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 代码示例

importonnxfromtvmimportrelay# 加载ONNX模型onnx_model=onnx.load("resnet18.onnx")# 解析为Relaymod,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 代码示例

importtorchfromtvmimportrelay# 准备PyTorch模型和示例输入model=torch.hub.load("pytorch/vision","resnet18",pretrained=True)model.eval()example_input=torch.randn(1,3,224,224)# 导出TorchScriptscripted=torch.jit.trace(model,example_input)# 解析为Relaymod,params=relay.frontend.from_pytorch(scripted,shape_list=[("input",example_input.shape)],)

4.3 ONNX vs TorchScript：怎么选？

特性ONNXTorchScript
通用性跨框架，所有框架都支持导出主要面向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——用符号变量表示未知维度。

fromtvmimportrelaxfromtvm.relax.frontendimportnnclassDecodeStep(nn.Module):defforward(self,x:nn.Tensor):# x的shape是 [batch, seq_len, hidden]# seq_len是动态符号变量returnnn.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：

fromtvmimportrelax# 注册一个自定义算子@relax.op.register_op("relax.custom_op")defcustom_op(x):returnrelax.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
  • 更多内容持续更新中……