Bring Your Own Codegen

TVM 的 Bring Your Own Codegen（BYOC）框架允许你将模型的部分内容卸载到自定义后端——硬件加速器、推理库，或你自己的 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_CODEGEN=ON和USE_EXAMPLE_NPU_RUNTIME=ON构建；TensorRT 小节需要USE_TENSORRT_CODEGEN=ON、USE_TENSORRT_RUNTIME=ON和USE_CUDA=ON，外加 CUDA GPU 以及匹配的 TensorRT 安装（来自 NVIDIA 的pip install tensorrt包或 TensorRT 压缩包）；最终部署小节还需要 PyTorch。当其后端不可用时，各小节会优雅降级。

BYOC 流程概览

BYOC 通过四个步骤把自定义后端接入 TVM 的编译流水线：

1. 注册 pattern—— 描述后端能处理哪些 Relax op 序列。
2. 划分图（partition）—— 将匹配到的算子归组为 composite function。
3. 运行 codegen—— 将每个 composite 降到后端表示（本教程中两个后端都是 JSON graph）。
4. 执行—— runtime 将每个 composite 派发到后端。

步骤 1 和 2 是纯 Python，可在任意环境运行；步骤 3 和 4 需要将后端的 codegen 与 runtime 编入 TVM，因此下文的 build-and-run 单元都做了可用性守卫。

Step 1：Import 后端以注册其 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_codegen=tvm.get_global_func("relax.ext.example_npu",True)has_example_npu_runtime=tvm.get_global_func("runtime.ExampleNPUJSONRuntimeCreate",True)has_example_npu=has_example_npu_codegenandhas_example_npu_runtime has_tensorrt_codegen=tvm.get_global_func("relax.ext.tensorrt",True)isnotNone_is_trt_runtime_enabled=tvm.get_global_func("relax.is_tensorrt_runtime_enabled",True)has_tensorrt=(has_tensorrt_codegenand_is_trt_runtime_enabledisnotNoneand_is_trt_runtime_enabled())has_cuda=tvm.cuda(0).exist

Step 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():conv=relax.op.nn.conv2d(data,weight)out=relax.op.nn.relu(conv)R.output(out)returnout

Step 3：为示例 NPU 做 Partition

FuseOpsByPattern将匹配已注册 pattern 的算子归组为 composite function；MergeCompositeFunctions再把绑定到同一后端的相邻 composite 合并为一次外部调用。有两个标志控制划分：

bind_constants=False让权重保持为函数参数，从而由 host 管理参数。（下文的 TensorRT 做相反选择：它把权重绑定为常量，因为会将其 bake 进 engine。）

annotate_codegen=True用带后端名称标签的函数包装每个匹配到的 composite——RunCodegen依据该标签路由。（后续的MergeCompositeFunctions在归组 composite 时也会附上该标签，因此下文的partition_for_tensorrt可以不设该标志。）

示例 NPU 注册了一个融合的conv2d + relupattern，优先级高于独立的conv2dpattern，因此这两个算子会折叠成单个example_npu.conv2d_relu_fusedcomposite——可在打印出的 module 中看到它。

npu_patterns=get_patterns_with_prefix("example_npu")npu_mod=FuseOpsByPattern(npu_patterns,bind_constants=False,annotate_codegen=True)(ConvReLU)npu_mod=MergeCompositeFunctions()(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), dtype="float32"), weight: R.Tensor((16, 3, 3, 3), dtype="float32")) -> R.Tensor((1, 16, 30, 30), dtype="float32"): 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), dtype="float32"), weight_1: R.Tensor((16, 3, 3, 3), dtype="float32")) -> R.Tensor((1, 16, 30, 30), dtype="float32"): R.func_attr({"Composite": "example_npu.conv2d_relu_fused"}) conv: R.Tensor((1, 16, 30, 30), dtype="float32") = R.nn.conv2d(data_1, weight_1, strides=[1, 1], padding=[0, 0, 0, 0], dilation=[1, 1], groups=1, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype=None) gv: R.Tensor((1, 16, 30, 30), dtype="float32") = R.nn.relu(conv) return gv output: R.Tensor((1, 16, 30, 30), dtype="float32") = local_func(data, weight) return output @R.function def main(data: R.Tensor((1, 3, 32, 32), dtype="float32"), weight: R.Tensor((16, 3, 3, 3), dtype="float32")) -> R.Tensor((1, 16, 30, 30), dtype="float32"): cls = Module with R.dataflow(): gv: R.Tensor((1, 16, 30, 30), dtype="float32") = cls.fused_relax_nn_conv2d_relax_nn_relu_example_npu_example_npu(data, weight) R.output(gv) return gv

Step 4：在示例 NPU 上 Codegen、build 与运行

RunCodegen调用每个被注解的 composite 的后端 codegen，将其替换为后端 runtime module（此处为 NPU 的 JSON graph）；随后relax.build编译剩余的 host 侧程序并链接所有内容。由于该 runtime 是一个不做任何计算的桩，我们只对输出 shape 做断言——数值是未初始化的。

np.random.seed(0)data_np=np.random.randn(1,3,32,32).astype("float32")weight_np=np.random.randn(16,3,3,3).astype("float32")ifhas_example_npu:npu_mod=RunCodegen()(npu_mod)withtvm.transform.PassContext(opt_level=3):npu_exec=relax.build(npu_mod,tvm.target.Target("llvm"))npu_vm=relax.VirtualMachine(npu_exec,tvm.cpu())npu_out=npu_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打包了你手写运行的FuseOpsByPattern+MergeCompositeFunctions，并使用 TensorRT 自己的 pattern 表。

权重变为常量（bind_constants=True）：TensorRT 会把它们 bake 进所构建的 engine，因此在 partition 之前先绑定参数。

真实数值。TensorRT 会真正计算，因此我们为 CUDA 构建、在 GPU 上运行，并与普通的 CPU 构建交叉验证——而不仅仅是 shape。

下面的 build-and-run 单元仅在 TensorRT 与 CUDA 可用时执行。在仅 CPU 的文档构建中，它们不会产生输出。

trt_mod=relax.transform.BindParams("main",{"weight":weight_np})(ConvReLU)trt_mod=partition_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), dtype="float32")) -> R.Tensor((1, 16, 30, 30), dtype="float32"): R.func_attr({"Codegen": "tensorrt"}) # from tvm.script import relax as R @R.function def gv(data_1: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor((1, 16, 30, 30), dtype="float32"): R.func_attr({"Composite": "tensorrt.nn.conv2d"}) with R.dataflow(): gv_1: R.Tensor((1, 16, 30, 30), dtype="float32") = R.nn.conv2d(data_1, metadata["relax.expr.Constant"][0], strides=[1, 1], padding=[0, 0, 0, 0], dilation=[1, 1], groups=1, data_layout="NCHW", kernel_layout="OIHW", out_layout="NCHW", out_dtype=None) R.output(gv_1) return gv_1 lv: R.Tensor((1, 16, 30, 30), dtype="float32") = gv(data) # from tvm.script import relax as R @R.function def gv1(lv_1: R.Tensor((1, 16, 30, 30), dtype="float32")) -> R.Tensor((1, 16, 30, 30), dtype="float32"): R.func_attr({"Composite": "tensorrt.nn.relu"}) with R.dataflow(): gv_1: R.Tensor((1, 16, 30, 30), dtype="float32") = R.nn.relu(lv_1) R.output(gv_1) return gv_1 gv_1: R.Tensor((1, 16, 30, 30), dtype="float32") = gv1(lv) return gv_1 @R.function def main(data: R.Tensor((1, 3, 32, 32), dtype="float32")) -> R.Tensor((1, 16, 30, 30), dtype="float32"): cls = Module with R.dataflow(): gv: R.Tensor((1, 16, 30, 30), dtype="float32") = cls.fused_relax_nn_conv2d_relax_nn_relu_tensorrt(data) R.output(gv) return gv # Metadata omitted. Use show_meta=True in script() method to show it.

为 CUDA 构建，在 GPU 上运行，并与 CPU 参考结果比较。

ifhas_tensorrtandhas_cuda:dev=tvm.cuda(0)withtvm.transform.PassContext(opt_level=3):trt_exec=relax.build(RunCodegen()(trt_mod),"cuda")trt_out=relax.VirtualMachine(trt_exec,dev)["main"](tvm.runtime.tensor(data_np,dev)).numpy()cpu_mod=relax.transform.LegalizeOps()(relax.transform.BindParams("main",{"weight":weight_np})(ConvReLU))cpu_exec=relax.build(cpu_mod,"llvm")cpu_out=relax.VirtualMachine(cpu_exec,tvm.cpu())["main"](tvm.runtime.tensor(data_np,tvm.cpu())).numpy()np.testing.assert_allclose(trt_out,cpu_out,rtol=1e-2,atol=1e-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 的 INT8，TVM_TENSORRT_MAX_WORKSPACE_SIZE限制 build workspace，TVM_TENSORRT_CACHE_DIR将已构建的 engine 缓存到磁盘以便跨运行复用。）

ifhas_tensorrtandhas_cuda:fp16_mod=partition_for_tensorrt(relax.transform.BindParams("main",{"weight":weight_np})(ConvReLU))withtvm.transform.PassContext(opt_level=3,config={"relax.ext.tensorrt.options":{"use_fp16":True}}):fp16_exec=relax.build(RunCodegen()(fp16_mod),"cuda")fp16_out=relax.VirtualMachine(fp16_exec,tvm.cuda(0))["main"](tvm.runtime.tensor(data_np,tvm.cuda(0))).numpy()np.testing.assert_allclose(fp16_out,cpu_out,rtol=5e-2,atol=5e-2)print("TensorRT FP16 output shape:",fp16_out.shape,"- matches within FP16 tolerance.")

Example NPU vs 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_torch=TrueexceptImportError:has_torch=Falseifhas_torchandhas_tensorrtandhas_cuda:fromtvm.relax.frontend.torchimportfrom_exported_programclassSmallConvNet(nn.Module):def__init__(self):super().__init__()self.conv1=nn.Conv2d(3,8,3)self.conv2=nn.Conv2d(8,16,3)self.pool=nn.MaxPool2d(2)defforward(self,x):x=torch.relu(self.conv1(x))x=self.pool(x)x=torch.relu(self.conv2(x))returnx torch_model=SmallConvNet().eval()example_input=torch.randn(1,3,32,32)withtorch.no_grad():torch_ref=torch_model(example_input).numpy()exported=torch.export.export(torch_model,(example_input,))torch_mod=from_exported_program(exported)torch_mod=partition_for_tensorrt(torch_mod)print("After importing and partitioning the PyTorch model:")print(torch_mod)torch_dev=tvm.cuda(0)withtvm.transform.PassContext(opt_level=3):torch_exec=relax.build(RunCodegen()(torch_mod),"cuda")deployed=relax.VirtualMachine(torch_exec,torch_dev)["main"](tvm.runtime.tensor(example_input.numpy(),torch_dev))[0].numpy()np.testing.assert_allclose(deployed,torch_ref,rtol=1e-2,atol=1e-2)print("Deployed PyTorch model on TensorRT; output",deployed.shape,"matches PyTorch.")

真实部署会构建一次并复用产物。将编译后的 module export 为共享库，之后再加载运行——可在全新进程中进行，无需 PyTorch，也无需重新构建。

ifhas_torchandhas_tensorrtandhas_cuda:withtempfile.TemporaryDirectory()astmpdir:lib_path=os.path.join(tmpdir,"deployed_trt.so")torch_exec.export_library(lib_path)loaded=tvm.runtime.load_module(lib_path)reran=relax.VirtualMachine(loaded,torch_dev)["main"](tvm.runtime.tensor(example_input.numpy(),torch_dev))[0].numpy()np.testing.assert_allclose(reran,torch_ref,rtol=1e-2,atol=1e-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 会为动态的 leading（batch）维度设置 optimization profile，因此该集成可以服务以 symbolic batch size export 的模型。

Engine 构建成本。第一次构建 TensorRT engine 很慢（不是卡死）。设置TVM_TENSORRT_CACHE_DIR可将已构建的 engine 缓存到磁盘，后续运行跳过重建。

下一步

要以示例 NPU 为起点构建你自己的后端：

1. 用你的硬件 SDK 调用替换src/runtime/extra/contrib/example_npu/example_npu_runtime.cc中的桩 runtime。
2. 在patterns.py中扩展你的硬件支持的算子。
3. 若后端需要非 JSON 的序列化格式，在src/relax/backend/contrib/下添加 C++ codegen。
4. 参照ExampleNPU.cmake，在cmake/modules/contrib/下添加 CMake 模块。

若要研究完整的真实后端实现，参见 TensorRT 集成：python/tvm/relax/backend/contrib/tensorrt.py中的 pattern 表与partition_for_tensorrt，src/relax/backend/contrib/tensorrt/中的 codegen，以及src/runtime/extra/contrib/tensorrt/中的 runtime。