TPATTensorRT 插件自動生成工具

TPAT 是由騰訊聯(lián)合英偉達開發(fā)的 TensorRT 插件自動生成工具。

TensorRT 是當前應用最廣的 GPU 推理框架，但由于支持的算子數(shù)量有限，用戶面臨手寫插件以支持算子的痛點。而TPAT 能夠支持開放神經(jīng)網(wǎng)絡交換 (ONNX) 格式所有的算子，端到端生成 TensorRT 插件，在解放人力成本的同時，性能對比手寫毫不遜色。

TPAT 實現(xiàn)了 TensorRT 插件的全自動生成，TensorRT 的部署和上線能基本流程化不再需要人工參與。手寫插件的步驟將由 TPAT 代替，TPAT 全自動生成一個算子插件耗時僅需要 30-60 分鐘的時間（該時間用于搜索算子的高性能 CUDA Kernel）。

架構設計

TPAT 接受用戶輸入的 ONNX-Model，指定需要生成 TensorRT Plugin 的算子和 batch size，基于 TVM 深度學習編譯器，對固定形狀的算子進行 AutoTune，自動生成高性能的 CUDA Kernel. 將 CUDA Kernel 和 Runtime 必要的參數(shù)填充進 TensorRT Plugin 模板，生成動態(tài)鏈接庫，可以直接加載到 TensorRT 運行。

TPAT 部分算子性能數(shù)據(jù)

使用 TPAT 自動生成 TensorRT-7.2 不支持的算子，并且用 TPAT 優(yōu)化 TensorRT-7.2 原生實現(xiàn)性能較差的算子；

對比手寫 Plugin

優(yōu)化 TensorRT 原生算子

TPAT 使用案例

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

• 輸入包括了 onehot 算子的 ONNX_Model、Onehot 算子的名字、batch_size
• TPAT借助 TVM的Relay 和 AutoScheduler 組件，生成高性能的 CUDA Kernel；
• 經(jīng)過模板填充后直接生成可用的 onehot 算子 Plugin 的動態(tài)鏈接庫。