【AI编译器入门⑥】端到端编译流程实战：ONNX → TVM → TensorRT全链路

【AI编译器入门⑥】端到端编译流程实战：ONNX → TVM → TensorRT全链路

本文是「AI编译器」系列的第六篇。前面我们分别讲了前端、中端、后端，今天把它们串起来——从ONNX模型开始，经过TVM中端优化，再选择不同的后端（CUDA/CUTLASS/TensorRT），最终生成可部署产物。本文带你完整走一遍端到端编译链路。

一、写在前面：为什么要串起来看？

学了前端、中端、后端，你可能会问：它们是怎么配合的？实际项目中完整的编译流程是什么样的？

单独看每一块都很清楚，但串起来才知道：

• onnx-simplifier为什么要先跑一遍？
• 中端Pass的顺序为什么要这样排？
• TensorRT partition和CUTLASS partition有什么不同？
• 动态batch怎么支持？

今天我们就把整个链路打通，从一个ONNX模型开始，一直到生成可部署的.so文件。

先上一张端到端流程图：

二、完整链路总览

ONNX model
  ↓ Step 1：onnx-simplifier化简
ONNX简化模型
  ↓ Step 2：TVM前端解析
Relay/Relax IR
  ↓ Step 3：中端Pass优化
    DynamicToStatic → FoldConstant → SimplifyInference
    → FuseOps → DeadCodeElimination → InferType
优化后的IR
  ↓ Step 4：后端选择
    ├── TVM CUDA Codegen
    ├── TVM + CUTLASS partition
    └── TVM + TensorRT partition
  ↓ Step 5：编译 + 打包
runtime artifacts（.so / .trt）
  ↓ Step 6：部署运行
推理服务

我们一步步来看。

三、Step 1：ONNX化简（onnx-simplifier）

3.1 为什么要先化简？

ONNX模型从PyTorch/TensorFlow导出时，常常带了很多冗余节点：

• 无意义的Identity
• 常量shape的Reshape/Squeeze
• 成对的Cast
• 死代码分支

这些冗余节点不仅影响性能，还可能导致TVM前端解析失败或中端优化效果打折扣。

所以在喂给TVM之前，先在ONNX层做一轮"大扫除"——onnx-simplifier就是干这个的。

3.2 怎么用？

python -m onnxsim model.onnx model_sim.onnx

onnx-simplifier做了什么？

• 常量折叠（Constant Folding）
• 形状推导（Shape Inference）
• 死代码消除
• 冗余节点清理

💡 经验之谈：编译ONNX模型前，先跑一遍onnx-simplifier，几乎是标配操作。花不了多少时间，但能避免很多奇怪的问题。

四、Step 2：TVM前端解析

4.1 ONNX导入

import onnx
from tvm import relay

onnx_model = onnx.load("model_sim.onnx")
mod, params = relay.frontend.from_onnx(
    onnx_model,
    shape={"input": [1, 3, 224, 224]},
    dtype={"input": "float32"},
)
mod = relay.transform.InferType()(mod)

4.2 这一步可能遇到的坑

1. 算子不支持：ONNX模型里有TVM不认识的算子 → 需要注册自定义算子
2. opset版本不兼容：opset太高或太低 → 升级TVM或转换opset
3. 动态shape处理：batch/sequence length是动态的 → 用Relax或bucket编译
4. 数值精度问题：某些算子语义有细微差异 → 需要对齐验证

五、Step 3：中端Pass优化

5.1 标准Pass Pipeline

from tvm.relay import transform as T

mod = T.DynamicToStatic()(mod)       # 动态算子静态化（如果可能）
mod = T.FoldConstant()(mod)           # 常量折叠
mod = T.SimplifyInference()(mod)      # 推理化简（BN/Dropout等）
mod = T.FuseOps(fuse_opt_level=2)(mod) # 算子融合
mod = T.DeadCodeElimination()(mod)    # 死代码消除
mod = T.InferType()(mod)              # 重新推导类型

5.2 为什么是这个顺序？

Pass的顺序很重要——顺序不对，效果会大打折扣：

1. 先FoldConstant：常量折叠后，图变小了，后续Pass处理更快
2. 再SimplifyInference：把BN等训练算子化简，为融合创造条件
3. 再FuseOps：图干净了，融合效果更好
4. 最后DCE：前面的Pass可能产生死代码，最后清理一遍

🎯 Pass顺序不是随便排的——每一步都在为下一步创造条件。顺序错了，可能需要多跑好几轮才能收敛。

5.3 布局适配

如果后端是TensorRT或CUTLASS，还需要根据后端偏好调整布局：

mod = T.ConvertLayout({"nn.conv2d": ["NCHW", "OIHW"]})(mod)

注意：布局转换要和后端对齐，不然可能引入额外的transpose开销。

六、Step 4：后端选择

TVM支持多种后端，我们来看看怎么选、怎么用。

6.1 后端一：纯CUDA Codegen

最简单的方式——全部走TVM自己的CUDA Codegen：

target = tvm.target.cuda(model="a100")
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target=target, params=params)
lib.export_library("model_cuda.so")

适合场景：自定义算子多、需要深度定制、调试阶段。

6.2 后端二：TVM + CUTLASS

核心算子走CUTLASS，其余走TVM Codegen：

from tvm.contrib.cutlass import partition_for_cutlass

mod = partition_for_cutlass(mod, params)
with tvm.target.cuda(model="a100"):
    lib = relay.build(mod, target="cuda", params=params)
lib.export_library("model_cutlass.so")

原理：

1. 遍历图，识别MatMul/Conv/BMM等可offload的算子
2. 标记为CUTLASS子图
3. PartitionGraph切分
4. CUTLASS子图生成CUTLASS kernel，其余走TVM Codegen

适合场景：核心GEMM/Conv多，想用CUTLASS的高性能，又要保留TVM的灵活性。

6.3 后端三：TVM + TensorRT

大子图走TensorRT，不支持的fallback到TVM：

from tvm.contrib.tensorrt import partition_for_tensorrt

mod = partition_for_tensorrt(mod, params)
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target="cuda", params=params)
lib.export_library("model_tensorrt.so")

原理：

1. 遍历图，检查每个算子是否被TensorRT支持
2. 连续支持的子图标记为TensorRT子图
3. PartitionGraph切分
4. TensorRT子图调用TensorRT builder生成engine

适合场景：标准模型（CNN/Transformer），追求极致性能，能接受黑盒。

6.4 直接TensorRT vs TVM+TensorRT

怎么选？

• 标准模型、只用NVIDIA、追求性能 → 直接TensorRT
• 有自定义算子、要跨平台、需要可调试 → TVM+TensorRT

七、Step 5：动态Batch推理支持

大模型推理中，batch size总是动态变化的。怎么支持动态batch？

7.1 方案一：Relax原生动态shape

编译期生成动态调度代码，运行期根据输入shape直接执行：

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

class Model(nn.Module):
    def forward(self, x: nn.Tensor):
        return nn.matmul(x, self.w)

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

优点：灵活，一个产物支持所有batch
缺点：性能比全静态编译差一些（因为没法做完全的常量折叠和布局优化）

7.2 方案二：Relay bucket编译

预编译多个bucket尺寸，运行期根据实际batch选最近的：

buckets = [1, 4, 8, 16]
libs = {}
for b in buckets:
    mod_b = relay.frontend.from_onnx(
        onnx_model, shape={"input": [b, 3, 224, 224]}
    )
    libs[b] = relay.build(mod_b, target="cuda")

优点：每个bucket都是全静态优化，性能高
缺点：产物多，占空间；非bucket大小有padding浪费

7.3 方案三：TensorRT动态batch

TensorRT支持explicit batch + optimization profile：

trt_config = {
    "use_explicit_batch": True,
    "max_batch_size": 16,
}
mod = partition_for_tensorrt(mod, params, **trt_config)

设置min/opt/max三个profile点，TensorRT会在范围内动态调整。

7.4 实际部署中怎么选？

通常是组合方案：

• 主要形状用bucket编译（保证性能）
• 特殊形状走Relax动态shape（保证灵活性）
• TensorRT子图用profile支持动态范围

没有银弹，要根据业务场景权衡。

八、完整代码示例（伪代码）

最后，给大家一个完整的端到端编译流程的代码框架：

import onnx
from tvm import relay
from tvm.contrib import tensorrt

# Step 1: ONNX化简（命令行做）
# python -m onnxsim model.onnx model_sim.onnx

# Step 2: 前端解析
onnx_model = onnx.load("model_sim.onnx")
mod, params = relay.frontend.from_onnx(
    onnx_model, shape={"input": [1, 3, 224, 224]}
)

# Step 3: 中端Pass优化
from tvm.relay import transform as T
mod = T.FoldConstant()(mod)
mod = T.SimplifyInference()(mod)
mod = T.FuseOps(fuse_opt_level=2)(mod)
mod = T.DeadCodeElimination()(mod)
mod = T.InferType()(mod)

# Step 4: TensorRT partition（可选）
mod = tensorrt.partition_for_tensorrt(mod, params)

# Step 5: 编译
with tvm.transform.PassContext(opt_level=3):
    lib = relay.build(mod, target="cuda", params=params)

# Step 6: 导出
lib.export_library("model.so")

九、常见误区

误区一：“ONNX是万能的中间格式”

ONNX只能表达静态图的算子连接关系，动态控制流、自定义算子、某些PyTorch特性转换时会丢失信息。复杂模型需要额外处理。

误区二：“TensorRT子图越多越好”

TensorRT子图和TVM子图之间切换有数据搬运overhead。如果子图太小或切分太碎，整体性能可能反而下降。

误区三：“动态batch只需要Relax”

Relax原生动态shape编译灵活，但某些后端（如TensorRT）对动态shape支持有限。实际部署中可能需要Relax + bucket + profile组合方案。

十、总结一下

今天这篇文章，我们走了一遍端到端编译流程：

1. ONNX化简：onnx-simplifier先做一轮"大扫除"，减轻后续压力
2. 前端解析：ONNX → Relay/Relax，注意算子支持和动态shape
3. 中端Pass：常量折叠 → 推理化简 → 算子融合 → DCE，顺序很重要
4. 后端选择：纯CUDA灵活、CUTLASS核心算子快、TensorRT整体性能高——各取所长
5. 动态batch：Relax原生动态、Relay bucket、TensorRT profile——通常组合使用

下一篇预告：我们来认识一个面向多平台自动化部署的编译框架——Laser。它是怎么实现"一份模型，多平台同出"的？HPC Plugin机制又是什么？

📚 本系列导航：

• 第1篇：AI编译器基础：三段式Pipeline与全链路协同
• 第2篇：TVM架构拆解：Relay、Relax、TIR与Runtime
• 第3篇：编译前端：ONNX、TorchScript与动态shape处理
• 第4篇：编译中端：图优化Pass全解析
• 第5篇：编译后端：CUDA Codegen、CUTLASS与TensorRT
• 第6篇：端到端编译流程：ONNX → TVM → TensorRT全链路实战（本文）
• 第7篇：Laser编译框架：多平台自动化部署与HPC Plugin
• 更多内容持续更新中……