【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 1onnx-simplifier化简 ONNX简化模型 ↓ Step 2TVM前端解析 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 1ONNX化简onnx-simplifier3.1 为什么要先化简ONNX模型从PyTorch/TensorFlow导出时常常带了很多冗余节点无意义的Identity常量shape的Reshape/Squeeze成对的Cast死代码分支这些冗余节点不仅影响性能还可能导致TVM前端解析失败或中端优化效果打折扣。所以在喂给TVM之前先在ONNX层做一轮大扫除——onnx-simplifier就是干这个的。3.2 怎么用python-monnxsim model.onnx model_sim.onnxonnx-simplifier做了什么常量折叠Constant Folding形状推导Shape Inference死代码消除冗余节点清理经验之谈编译ONNX模型前先跑一遍onnx-simplifier几乎是标配操作。花不了多少时间但能避免很多奇怪的问题。四、Step 2TVM前端解析4.1 ONNX导入importonnxfromtvmimportrelay onnx_modelonnx.load(model_sim.onnx)mod,paramsrelay.frontend.from_onnx(onnx_model,shape{input:[1,3,224,224]},dtype{input:float32},)modrelay.transform.InferType()(mod)4.2 这一步可能遇到的坑算子不支持ONNX模型里有TVM不认识的算子 → 需要注册自定义算子opset版本不兼容opset太高或太低 → 升级TVM或转换opset动态shape处理batch/sequence length是动态的 → 用Relax或bucket编译数值精度问题某些算子语义有细微差异 → 需要对齐验证五、Step 3中端Pass优化5.1 标准Pass Pipelinefromtvm.relayimporttransformasT modT.DynamicToStatic()(mod)# 动态算子静态化如果可能modT.FoldConstant()(mod)# 常量折叠modT.SimplifyInference()(mod)# 推理化简BN/Dropout等modT.FuseOps(fuse_opt_level2)(mod)# 算子融合modT.DeadCodeElimination()(mod)# 死代码消除modT.InferType()(mod)# 重新推导类型5.2 为什么是这个顺序Pass的顺序很重要——顺序不对效果会大打折扣先FoldConstant常量折叠后图变小了后续Pass处理更快再SimplifyInference把BN等训练算子化简为融合创造条件再FuseOps图干净了融合效果更好最后DCE前面的Pass可能产生死代码最后清理一遍Pass顺序不是随便排的——每一步都在为下一步创造条件。顺序错了可能需要多跑好几轮才能收敛。5.3 布局适配如果后端是TensorRT或CUTLASS还需要根据后端偏好调整布局modT.ConvertLayout({nn.conv2d:[NCHW,OIHW]})(mod)注意布局转换要和后端对齐不然可能引入额外的transpose开销。六、Step 4后端选择TVM支持多种后端我们来看看怎么选、怎么用。6.1 后端一纯CUDA Codegen最简单的方式——全部走TVM自己的CUDA Codegentargettvm.target.cuda(modela100)withtvm.transform.PassContext(opt_level3):librelay.build(mod,targettarget,paramsparams)lib.export_library(model_cuda.so)适合场景自定义算子多、需要深度定制、调试阶段。6.2 后端二TVM CUTLASS核心算子走CUTLASS其余走TVM Codegenfromtvm.contrib.cutlassimportpartition_for_cutlass modpartition_for_cutlass(mod,params)withtvm.target.cuda(modela100):librelay.build(mod,targetcuda,paramsparams)lib.export_library(model_cutlass.so)原理遍历图识别MatMul/Conv/BMM等可offload的算子标记为CUTLASS子图PartitionGraph切分CUTLASS子图生成CUTLASS kernel其余走TVM Codegen适合场景核心GEMM/Conv多想用CUTLASS的高性能又要保留TVM的灵活性。6.3 后端三TVM TensorRT大子图走TensorRT不支持的fallback到TVMfromtvm.contrib.tensorrtimportpartition_for_tensorrt modpartition_for_tensorrt(mod,params)withtvm.transform.PassContext(opt_level3):librelay.build(mod,targetcuda,paramsparams)lib.export_library(model_tensorrt.so)原理遍历图检查每个算子是否被TensorRT支持连续支持的子图标记为TensorRT子图PartitionGraph切分TensorRT子图调用TensorRT builder生成engine适合场景标准模型CNN/Transformer追求极致性能能接受黑盒。6.4 直接TensorRT vs TVMTensorRT维度直接TensorRTTVM TensorRT性能最高全图TensorRT略低子图切换有overhead灵活性低黑盒高不支持的算子fallback到TVM跨平台只支持NVIDIATVM部分可跨平台可调试性差好TVM部分可以看IR、看生成代码怎么选标准模型、只用NVIDIA、追求性能 → 直接TensorRT有自定义算子、要跨平台、需要可调试 → TVMTensorRT七、Step 5动态Batch推理支持大模型推理中batch size总是动态变化的。怎么支持动态batch7.1 方案一Relax原生动态shape编译期生成动态调度代码运行期根据输入shape直接执行fromtvmimportrelaxfromtvm.relax.frontendimportnnclassModel(nn.Module):defforward(self,x:nn.Tensor):returnnn.matmul(x,self.w)modelModel()mod,paramsmodel.export_tvm(spec{forward:{x:nn.spec.Tensor([batch,768],float32)}})优点灵活一个产物支持所有batch缺点性能比全静态编译差一些因为没法做完全的常量折叠和布局优化7.2 方案二Relay bucket编译预编译多个bucket尺寸运行期根据实际batch选最近的buckets[1,4,8,16]libs{}forbinbuckets:mod_brelay.frontend.from_onnx(onnx_model,shape{input:[b,3,224,224]})libs[b]relay.build(mod_b,targetcuda)优点每个bucket都是全静态优化性能高缺点产物多占空间非bucket大小有padding浪费7.3 方案三TensorRT动态batchTensorRT支持explicit batch optimization profiletrt_config{use_explicit_batch:True,max_batch_size:16,}modpartition_for_tensorrt(mod,params,**trt_config)设置min/opt/max三个profile点TensorRT会在范围内动态调整。7.4 实际部署中怎么选通常是组合方案主要形状用bucket编译保证性能特殊形状走Relax动态shape保证灵活性TensorRT子图用profile支持动态范围没有银弹要根据业务场景权衡。八、完整代码示例伪代码最后给大家一个完整的端到端编译流程的代码框架importonnxfromtvmimportrelayfromtvm.contribimporttensorrt# Step 1: ONNX化简命令行做# python -m onnxsim model.onnx model_sim.onnx# Step 2: 前端解析onnx_modelonnx.load(model_sim.onnx)mod,paramsrelay.frontend.from_onnx(onnx_model,shape{input:[1,3,224,224]})# Step 3: 中端Pass优化fromtvm.relayimporttransformasT modT.FoldConstant()(mod)modT.SimplifyInference()(mod)modT.FuseOps(fuse_opt_level2)(mod)modT.DeadCodeElimination()(mod)modT.InferType()(mod)# Step 4: TensorRT partition可选modtensorrt.partition_for_tensorrt(mod,params)# Step 5: 编译withtvm.transform.PassContext(opt_level3):librelay.build(mod,targetcuda,paramsparams)# Step 6: 导出lib.export_library(model.so)九、常见误区误区一“ONNX是万能的中间格式”ONNX只能表达静态图的算子连接关系动态控制流、自定义算子、某些PyTorch特性转换时会丢失信息。复杂模型需要额外处理。误区二“TensorRT子图越多越好”TensorRT子图和TVM子图之间切换有数据搬运overhead。如果子图太小或切分太碎整体性能可能反而下降。误区三“动态batch只需要Relax”Relax原生动态shape编译灵活但某些后端如TensorRT对动态shape支持有限。实际部署中可能需要Relax bucket profile组合方案。十、总结一下今天这篇文章我们走了一遍端到端编译流程ONNX化简onnx-simplifier先做一轮大扫除减轻后续压力前端解析ONNX → Relay/Relax注意算子支持和动态shape中端Pass常量折叠 → 推理化简 → 算子融合 → DCE顺序很重要后端选择纯CUDA灵活、CUTLASS核心算子快、TensorRT整体性能高——各取所长动态batchRelax原生动态、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更多内容持续更新中……