TPAT —— TensorRT 插件自动生成工具

TPAT 是由腾讯联合英伟达开发的 TensorRT 插件自动生成工具。

TensorRT 是当前应用最广的 GPU 推理框架，但由于支持的算子数量有限，用户面临手写插件以支持算子的痛点。而TPAT 能够支持开放神经网络交换 (ONNX) 格式所有的算子，端到端生成 TensorRT 插件，在解放人力成本的同时，性能对比手写毫不逊色。

TPAT 实现了 TensorRT 插件的全自动生成，TensorRT 的部署和上线能基本流程化不再需要人工参与。手写插件的步骤将由 TPAT 代替，TPAT 全自动生成一个算子插件耗时仅需要 30-60 分钟的时间（该时间用于搜索算子的高性能 CUDA Kernel）。

架构设计

TPAT 接受用户输入的 ONNX-Model，指定需要生成 TensorRT Plugin 的算子和 batch size，基于 TVM 深度学习编译器，对固定形状的算子进行 AutoTune，自动生成高性能的 CUDA Kernel. 将 CUDA Kernel 和 Runtime 必要的参数填充进 TensorRT Plugin 模板，生成动态链接库，可以直接加载到 TensorRT 运行。

TPAT 部分算子性能数据

使用 TPAT 自动生成 TensorRT-7.2 不支持的算子，并且用 TPAT 优化 TensorRT-7.2 原生实现性能较差的算子；

对比手写 Plugin

优化 TensorRT 原生算子

TPAT 使用案例

使用 TPAT 支持 Onehot 算子（TensorRT-7.2.2.3）

• 输入包括了 onehot 算子的 ONNX_Model、Onehot 算子的名字、batch_size
• TPAT借助 TVM的Relay 和 AutoScheduler 组件，生成高性能的 CUDA Kernel；
• 经过模板填充后直接生成可用的 onehot 算子 Plugin 的动态链接库。