Bring Your Own CodegenTVM 的 Bring Your Own CodegenBYOC框架允许你将模型的部分内容卸载到自定义后端——硬件加速器、推理库或你自己的 kernel——而其余部分仍由 TVM 编译。本教程分为两部分BYOC 如何工作—— 我们先用一个捆绑的、无需硬件的示例 NPU 后端讲解流程再用同一流程驱动真实的生产后端 NVIDIA TensorRT。两者都运行一个小型手写模型以便每一步都清晰可见两者之间唯一变化的是后端而这种对比正是本课要点。部署真实模型—— 随后将其投入实际使用把一个真实的 PyTorchnn.Module从 export 经 TensorRT 跑到 GPU 上。示例 NPU 是一个教学用桩stub其 runtime 会记录 NPU 会做出的调度决策内存层级、执行引擎、融合但不做真实计算因此输出 buffer 保持未初始化。所以在 NPU 相关小节我们只检查 shape不检查数值——它的职责是让每一步 BYOC 都可见没有任何隐藏。TensorRT 则用同一流程做真实计算因此我们会将其结果与参考结果交叉验证。前置条件示例 NPU 小节需要 TVM 以USE_EXAMPLE_NPU_CODEGENON和USE_EXAMPLE_NPU_RUNTIMEON构建TensorRT 小节需要USE_TENSORRT_CODEGENON、USE_TENSORRT_RUNTIMEON和USE_CUDAON外加 CUDA GPU 以及匹配的 TensorRT 安装来自 NVIDIA 的pip install tensorrt包或 TensorRT 压缩包最终部署小节还需要 PyTorch。当其后端不可用时各小节会优雅降级。BYOC 流程概览BYOC 通过四个步骤把自定义后端接入 TVM 的编译流水线注册 pattern—— 描述后端能处理哪些 Relax op 序列。划分图partition—— 将匹配到的算子归组为 composite function。运行 codegen—— 将每个 composite 降到后端表示本教程中两个后端都是 JSON graph。执行—— runtime 将每个 composite 派发到后端。步骤 1 和 2 是纯 Python可在任意环境运行步骤 3 和 4 需要将后端的 codegen 与 runtime 编入 TVM因此下文的 build-and-run 单元都做了可用性守卫。Step 1Import 后端以注册其 pattern导入后端模块会将其 pattern 注册到 TVM 的全局 registry。Pattern 注册与 C 构建无关——只有 codegen 和 runtime 需要把后端编进去——因此我们探测每个后端并据此守卫 build-and-run 单元。importosimporttempfileimportnumpyasnpimporttvmimporttvm.relax.backend.contrib.example_npufromtvmimportrelaxfromtvm.relax.backend.contrib.tensorrtimportpartition_for_tensorrtfromtvm.relax.backend.pattern_registryimportget_patterns_with_prefixfromtvm.relax.transformimportFuseOpsByPattern,MergeCompositeFunctions,RunCodegenfromtvm.scriptimportrelaxasR has_example_npu_codegentvm.get_global_func(relax.ext.example_npu,True)has_example_npu_runtimetvm.get_global_func(runtime.ExampleNPUJSONRuntimeCreate,True)has_example_npuhas_example_npu_codegenandhas_example_npu_runtime has_tensorrt_codegentvm.get_global_func(relax.ext.tensorrt,True)isnotNone_is_trt_runtime_enabledtvm.get_global_func(relax.is_tensorrt_runtime_enabled,True)has_tensorrt(has_tensorrt_codegenand_is_trt_runtime_enabledisnotNoneand_is_trt_runtime_enabled())has_cudatvm.cuda(0).existStep 2定义模型单个卷积后接 ReLU。该模型同时用于两个后端。tvm.script.ir_moduleclassConvReLU:R.functiondefmain(data:R.Tensor((1,3,32,32),float32),weight:R.Tensor((16,3,3,3),float32),)-R.Tensor((1,16,30,30),float32):withR.dataflow():convrelax.op.nn.conv2d(data,weight)outrelax.op.nn.relu(conv)R.output(out)returnoutStep 3为示例 NPU 做 PartitionFuseOpsByPattern将匹配已注册 pattern 的算子归组为 composite functionMergeCompositeFunctions再把绑定到同一后端的相邻 composite 合并为一次外部调用。有两个标志控制划分bind_constantsFalse让权重保持为函数参数从而由 host 管理参数。下文的 TensorRT 做相反选择它把权重绑定为常量因为会将其 bake 进 engine。annotate_codegenTrue用带后端名称标签的函数包装每个匹配到的 composite——RunCodegen依据该标签路由。后续的MergeCompositeFunctions在归组 composite 时也会附上该标签因此下文的partition_for_tensorrt可以不设该标志。示例 NPU 注册了一个融合的conv2d relupattern优先级高于独立的conv2dpattern因此这两个算子会折叠成单个example_npu.conv2d_relu_fusedcomposite——可在打印出的 module 中看到它。npu_patternsget_patterns_with_prefix(example_npu)npu_modFuseOpsByPattern(npu_patterns,bind_constantsFalse,annotate_codegenTrue)(ConvReLU)npu_modMergeCompositeFunctions()(npu_mod)print(After partitioning for the example NPU:)print(npu_mod)After partitioning for the example NPU: # from tvm.script import ir as I # from tvm.script import relax as R I.ir_module class Module: R.function def fused_relax_nn_conv2d_relax_nn_relu_example_npu_example_npu(data: R.Tensor((1, 3, 32, 32), dtypefloat32), weight: R.Tensor((16, 3, 3, 3), dtypefloat32)) - R.Tensor((1, 16, 30, 30), dtypefloat32): R.func_attr({Codegen: example_npu}) # from tvm.script import relax as R R.function def local_func(data_1: R.Tensor((1, 3, 32, 32), dtypefloat32), weight_1: R.Tensor((16, 3, 3, 3), dtypefloat32)) - R.Tensor((1, 16, 30, 30), dtypefloat32): R.func_attr({Composite: example_npu.conv2d_relu_fused}) conv: R.Tensor((1, 16, 30, 30), dtypefloat32) R.nn.conv2d(data_1, weight_1, strides[1, 1], padding[0, 0, 0, 0], dilation[1, 1], groups1, data_layoutNCHW, kernel_layoutOIHW, out_layoutNCHW, out_dtypeNone) gv: R.Tensor((1, 16, 30, 30), dtypefloat32) R.nn.relu(conv) return gv output: R.Tensor((1, 16, 30, 30), dtypefloat32) local_func(data, weight) return output R.function def main(data: R.Tensor((1, 3, 32, 32), dtypefloat32), weight: R.Tensor((16, 3, 3, 3), dtypefloat32)) - R.Tensor((1, 16, 30, 30), dtypefloat32): cls Module with R.dataflow(): gv: R.Tensor((1, 16, 30, 30), dtypefloat32) cls.fused_relax_nn_conv2d_relax_nn_relu_example_npu_example_npu(data, weight) R.output(gv) return gvStep 4在示例 NPU 上 Codegen、build 与运行RunCodegen调用每个被注解的 composite 的后端 codegen将其替换为后端 runtime module此处为 NPU 的 JSON graph随后relax.build编译剩余的 host 侧程序并链接所有内容。由于该 runtime 是一个不做任何计算的桩我们只对输出 shape 做断言——数值是未初始化的。np.random.seed(0)data_npnp.random.randn(1,3,32,32).astype(float32)weight_npnp.random.randn(16,3,3,3).astype(float32)ifhas_example_npu:npu_modRunCodegen()(npu_mod)withtvm.transform.PassContext(opt_level3):npu_execrelax.build(npu_mod,tvm.target.Target(llvm))npu_vmrelax.VirtualMachine(npu_exec,tvm.cpu())npu_outnpu_vm[main](tvm.runtime.tensor(data_np,tvm.cpu()),tvm.runtime.tensor(weight_np,tvm.cpu()))assertnpu_out.numpy().shape(1,16,30,30)print(Example NPU run completed. Output shape:,npu_out.numpy().shape)else:print(Example NPU backend unavailable; skipping its build and run.)Example NPU backend unavailable; skipping its build and run.同一流程打到真实后端TensorRT上面的 Step 1–4 就是全部机制。把它们对准真实后端时变化很少因此不再重复走一遍这里只说明对 NVIDIA TensorRT 有何不同一次调用完成 Partition。partition_for_tensorrt打包了你手写运行的FuseOpsByPatternMergeCompositeFunctions并使用 TensorRT 自己的 pattern 表。权重变为常量bind_constantsTrueTensorRT 会把它们 bake 进所构建的 engine因此在 partition 之前先绑定参数。真实数值。TensorRT 会真正计算因此我们为 CUDA 构建、在 GPU 上运行并与普通的 CPU 构建交叉验证——而不仅仅是 shape。下面的 build-and-run 单元仅在 TensorRT 与 CUDA 可用时执行。在仅 CPU 的文档构建中它们不会产生输出。trt_modrelax.transform.BindParams(main,{weight:weight_np})(ConvReLU)trt_modpartition_for_tensorrt(trt_mod)print(After partition_for_tensorrt:)print(trt_mod)After partition_for_tensorrt: # from tvm.script import ir as I # from tvm.script import relax as R I.ir_module class Module: R.function def fused_relax_nn_conv2d_relax_nn_relu_tensorrt(data: R.Tensor((1, 3, 32, 32), dtypefloat32)) - R.Tensor((1, 16, 30, 30), dtypefloat32): R.func_attr({Codegen: tensorrt}) # from tvm.script import relax as R R.function def gv(data_1: R.Tensor((1, 3, 32, 32), dtypefloat32)) - R.Tensor((1, 16, 30, 30), dtypefloat32): R.func_attr({Composite: tensorrt.nn.conv2d}) with R.dataflow(): gv_1: R.Tensor((1, 16, 30, 30), dtypefloat32) R.nn.conv2d(data_1, metadata[relax.expr.Constant][0], strides[1, 1], padding[0, 0, 0, 0], dilation[1, 1], groups1, data_layoutNCHW, kernel_layoutOIHW, out_layoutNCHW, out_dtypeNone) R.output(gv_1) return gv_1 lv: R.Tensor((1, 16, 30, 30), dtypefloat32) gv(data) # from tvm.script import relax as R R.function def gv1(lv_1: R.Tensor((1, 16, 30, 30), dtypefloat32)) - R.Tensor((1, 16, 30, 30), dtypefloat32): R.func_attr({Composite: tensorrt.nn.relu}) with R.dataflow(): gv_1: R.Tensor((1, 16, 30, 30), dtypefloat32) R.nn.relu(lv_1) R.output(gv_1) return gv_1 gv_1: R.Tensor((1, 16, 30, 30), dtypefloat32) gv1(lv) return gv_1 R.function def main(data: R.Tensor((1, 3, 32, 32), dtypefloat32)) - R.Tensor((1, 16, 30, 30), dtypefloat32): cls Module with R.dataflow(): gv: R.Tensor((1, 16, 30, 30), dtypefloat32) cls.fused_relax_nn_conv2d_relax_nn_relu_tensorrt(data) R.output(gv) return gv # Metadata omitted. Use show_metaTrue in script() method to show it.为 CUDA 构建在 GPU 上运行并与 CPU 参考结果比较。ifhas_tensorrtandhas_cuda:devtvm.cuda(0)withtvm.transform.PassContext(opt_level3):trt_execrelax.build(RunCodegen()(trt_mod),cuda)trt_outrelax.VirtualMachine(trt_exec,dev)[main](tvm.runtime.tensor(data_np,dev)).numpy()cpu_modrelax.transform.LegalizeOps()(relax.transform.BindParams(main,{weight:weight_np})(ConvReLU))cpu_execrelax.build(cpu_mod,llvm)cpu_outrelax.VirtualMachine(cpu_exec,tvm.cpu())[main](tvm.runtime.tensor(data_np,tvm.cpu())).numpy()np.testing.assert_allclose(trt_out,cpu_out,rtol1e-2,atol1e-2)print(TensorRT output shape:,trt_out.shape,- matches the CPU reference.)真实后端还会暴露教学桩没有的旋钮。通过relax.ext.tensorrt.options配置设置use_fp16可让 TensorRT 选用 FP16 kernel用少许精度换取速度流程的其余部分不变。其他选项由环境变量驱动TVM_TENSORRT_USE_INT8启用带 calibration 的 INT8TVM_TENSORRT_MAX_WORKSPACE_SIZE限制 build workspaceTVM_TENSORRT_CACHE_DIR将已构建的 engine 缓存到磁盘以便跨运行复用。ifhas_tensorrtandhas_cuda:fp16_modpartition_for_tensorrt(relax.transform.BindParams(main,{weight:weight_np})(ConvReLU))withtvm.transform.PassContext(opt_level3,config{relax.ext.tensorrt.options:{use_fp16:True}}):fp16_execrelax.build(RunCodegen()(fp16_mod),cuda)fp16_outrelax.VirtualMachine(fp16_exec,tvm.cuda(0))[main](tvm.runtime.tensor(data_np,tvm.cuda(0))).numpy()np.testing.assert_allclose(fp16_out,cpu_out,rtol5e-2,atol5e-2)print(TensorRT FP16 output shape:,fp16_out.shape,- matches within FP16 tolerance.)Example NPU vs TensorRT 一览同一套四步流程两个后端AspectExample NPU教学桩TensorRT真实后端Runtime记录决策不做计算构建并运行 nvinfer engineOutput未初始化检查 shape真实数值与 CPU 交叉验证Weightsbind_constantsFalsebind_constantsTruebake 进去Partition两步 Pass手写partition_for_tensorrt一次调用用 TensorRT 部署 PyTorch 模型上面全部使用手写IRModule以便每个 op 都可见。实践中你从训练好的模型出发。最后一节对真实的 PyTorchnn.Module端到端运行同一套partition_for_tensorrt流程export 它用 PyTorch frontend 导入 Relax权重以常量形式进入——正是 TensorRT bake 进 engine 所需的形式partition为 CUDA 构建并将 GPU 结果与 PyTorch 自身输出对照。除 frontend 导入外唯一差异是导入后的程序以 tuple 返回输出因此对单个结果张量取索引[0]partition-build-run 流程本身不变。本节额外需要 PyTorch。try:importtorchfromtorchimportnn has_torchTrueexceptImportError:has_torchFalseifhas_torchandhas_tensorrtandhas_cuda:fromtvm.relax.frontend.torchimportfrom_exported_programclassSmallConvNet(nn.Module):def__init__(self):super().__init__()self.conv1nn.Conv2d(3,8,3)self.conv2nn.Conv2d(8,16,3)self.poolnn.MaxPool2d(2)defforward(self,x):xtorch.relu(self.conv1(x))xself.pool(x)xtorch.relu(self.conv2(x))returnx torch_modelSmallConvNet().eval()example_inputtorch.randn(1,3,32,32)withtorch.no_grad():torch_reftorch_model(example_input).numpy()exportedtorch.export.export(torch_model,(example_input,))torch_modfrom_exported_program(exported)torch_modpartition_for_tensorrt(torch_mod)print(After importing and partitioning the PyTorch model:)print(torch_mod)torch_devtvm.cuda(0)withtvm.transform.PassContext(opt_level3):torch_execrelax.build(RunCodegen()(torch_mod),cuda)deployedrelax.VirtualMachine(torch_exec,torch_dev)[main](tvm.runtime.tensor(example_input.numpy(),torch_dev))[0].numpy()np.testing.assert_allclose(deployed,torch_ref,rtol1e-2,atol1e-2)print(Deployed PyTorch model on TensorRT; output,deployed.shape,matches PyTorch.)真实部署会构建一次并复用产物。将编译后的 module export 为共享库之后再加载运行——可在全新进程中进行无需 PyTorch也无需重新构建。ifhas_torchandhas_tensorrtandhas_cuda:withtempfile.TemporaryDirectory()astmpdir:lib_pathos.path.join(tmpdir,deployed_trt.so)torch_exec.export_library(lib_path)loadedtvm.runtime.load_module(lib_path)reranrelax.VirtualMachine(loaded,torch_dev)[main](tvm.runtime.tensor(example_input.numpy(),torch_dev))[0].numpy()np.testing.assert_allclose(reran,torch_ref,rtol1e-2,atol1e-2)print(Reloaded the exported library and reran; output,reran.shape,still matches.)真实部署注意事项算子覆盖与回退。TensorRT 只卸载其 pattern 表中的算子见python/tvm/relax/backend/contrib/tensorrt.py任何不支持的算子会留在 host 上。打印划分后的 module查找带Codegen: tensorrt的函数即可看到卸载了什么。动态 shape。Builder 会为动态的 leadingbatch维度设置 optimization profile因此该集成可以服务以 symbolic batch size export 的模型。Engine 构建成本。第一次构建 TensorRT engine 很慢不是卡死。设置TVM_TENSORRT_CACHE_DIR可将已构建的 engine 缓存到磁盘后续运行跳过重建。下一步要以示例 NPU 为起点构建你自己的后端用你的硬件 SDK 调用替换src/runtime/extra/contrib/example_npu/example_npu_runtime.cc中的桩 runtime。在patterns.py中扩展你的硬件支持的算子。若后端需要非 JSON 的序列化格式在src/relax/backend/contrib/下添加 C codegen。参照ExampleNPU.cmake在cmake/modules/contrib/下添加 CMake 模块。若要研究完整的真实后端实现参见 TensorRT 集成python/tvm/relax/backend/contrib/tensorrt.py中的 pattern 表与partition_for_tensorrtsrc/relax/backend/contrib/tensorrt/中的 codegen以及src/runtime/extra/contrib/tensorrt/中的 runtime。