TPAT : TensorRT Plugin Autogen Tool

Qian Qiu Tencent AI Lab
Meng Wang NVIDIA DevTech
作者信息未提供

目录

• 背景 (Background)

  • ONNX
  • TensorRT
  • 基于ONNX解析器的传统TensorRT工作流
  • 带有插件的TensorRT工作流
  • TVM

• TPAT：TensorRT 插件自动生成工具 (TPAT: TensorRT Plugin Autogen Tool)

  • TPAT亮点
  • TPAT架构
  • TPAT步骤

  • 步骤1：从ONNX到TVM Relay

    • 步骤1：从ONNX到TVM Relay（解决方案）

  • 步骤2：从TVM到CUDA源代码

  • 步骤3：从TVM运行时到CUDA核函数参数

    • 复杂案例：一个ONNX节点 -> 多个CUDA核函数

  • 步骤4：自动填充TensorRT插件模板

  • 步骤5：修改ONNX图

    • 复杂案例：多个ONNX节点需要插件
    • 子图自动生成
    • LayerNormalization子图自动生成

  • TPAT 示例

• 结果 (Result)

  • 与手写插件比较
  • 优化 TensorRT 算子
  • 支持的算子

• 总结 (Summary)

• 招聘信息

背景

ONNX

ONNX (Open Neural Network Exchange) 是一种开放的、用于机器学习模型的中间格式。
详情请见：https://github.com/onnx/onnx

Page 3 展示了ONNX的生态系统，包括与PyTorch、Chainer、Caffe2、MXNet、Cognitive Toolkit、XGBoost、PaddlePaddle等各种框架的连接，以及在CPU、GPU、FPGA、ASIC等不同硬件上的部署。
更多信息可查阅：https://azure.microsoft.com/es-es/blog/onnx-runtime-for-inferencing-machine-learning-models-now-in-preview/

TensorRT

TensorRT 是最流行的GPU推理引擎，在将模型部署到TensorRT后可以显著加速。

TensorRT的问题：
* TensorRT无法完全覆盖所有的运算符。
* 当运算符不受支持或其性能不足时，需要手动编写插件。
* 编写插件需要GPU和CUDA编程知识和经验。

Page 4 展示了TensorRT的内部优化流程，包括层与张量融合、精度校准、内核自动生成、训练神经网络、动态张量内存和多流执行。

基于ONNX解析器的传统TensorRT工作流

Page 5 描述了传统的TensorRT工作流：PyTorch/TensorFlow模型经过ONNX转换，然后由ONNX解析器处理。如果遇到不支持的运算符，需要手动编写插件（通常需要数天的工作量），之后通过trtexec工具部署到TensorRT。
* 编写插件是此工作流中最耗时的部分。
* 实现一个插件通常需要数天的工作量。

带有插件的TensorRT工作流

利用TensorRT ONNX解析器的回退（Fallback）机制：
* 实现TensorRT插件。
* 为插件构建一个独立的库。
* 预加载该库，ONNX解析器将自动尝试将不支持的层/操作作为插件导入（通过FallbackPluginImporter）。

优点与缺点：
* 优点：易于实现整个过程的自动化。
* 缺点：插件的定义必须与ONNX运算符保持一致，即名称、输入/输出。

TVM

TVM是一个编译器，它提供图级别和运算符级别的优化，以实现深度学习工作负载在不同硬件后端上的性能可移植性。
它提供了一个端到端的编译和优化堆栈，允许部署用高级框架指定的深度学习工作负载。

Page 7 展示了TVM的工作流：从TF/PyTorch/ONNX等框架，通过Relay（高级IR）到TE（计算定义），再通过AutoTVM/AutoScheduler（自动调优模块）生成TE+Schedule（优化规范），最终生成TIR（低级IR）和机器码。
关键思想： 结合TensorRT和TVM工作流。
更多信息可查阅：https://tvm.apache.org/docs/tutorial/introduction.html#sphx-glr-tutorial-introduction-py

TPAT：TensorRT 插件自动生成工具

TPAT利用TVM的强大功能自动生成CUDA内核，然后自动填充TensorRT插件模板。

Page 9 描述了TPAT的工作流：PyTorch/TensorFlow模型经过ONNX转换，然后由ONNX解析器处理。TVM在此阶段介入，自动生成插件（耗时30-60分钟），之后通过trtexec工具部署到TensorRT。
使用TPAT的TensorRT工作流：
* 不再需要手动编写插件。
* TensorRT插件完全自动化生成。

TPAT亮点

• 广泛覆盖： 支持所有TF/PyTorch/ONNX运算符。
• 全自动化： 为用户指定的运算符生成插件代码。
• 高性能： 性能优于手动编写的插件。

TPAT架构

Page 11 展示了TPAT的架构：
用户提供来自框架（如Tensorflow、PyTorch）的模型、运算符名称、ONNX模型和批次大小作为输入。这些输入进入TVM流程：ONNX模型转换为Relay，经过TE/TOPI、AutoScheduler和Best Schedule，最终到达TVM Runtime。TVM Runtime生成CUDA源代码。CUDA源代码用于填充插件模板，生成TensorRT插件。此插件通过trtexec工具加载到TensorRT引擎，最终实现TensorRT推理。

TPAT步骤

给定一个包含需要TensorRT插件的节点的ONNX模型：
1. 从ONNX到TVM Relay。
2. 从TVM到CUDA源代码。
3. 从TVM运行时到CUDA内核参数。
4. 自动填充TensorRT插件模板。
5. 修改ONNX图。
使用构建的插件将修改后的ONNX模型转换为TensorRT。

步骤1：从ONNX到TVM Relay

• TVM支持ONNX到Relay的转换。
• TVM将为整个导入的ONNX图自动生成内核。
• 但只有指定的ONNX节点需要TVM自动调优，例如右图中OneHot节点。

Page 13 展示了一个示例图，其中input0经过MatMul、Cast和OneHot等操作。其中MatMul和Cast等节点可以转换为Relay IR并生成CUDA源代码。特别指出OneHot节点需要TVM的自动调优。

步骤1：从ONNX到TVM Relay（解决方案）

解决方案： 将指定的ONNX节点提取为一个小的ONNX模型，然后将其转换为TVM。
ONNX graphsurgeon是一个很好的工具，可以隔离子模型并使此步骤自动化。

Page 14 进一步阐述了解决方案，通过ONNX graphsurgeon工具将需要自动调优的OneHot节点从主图中分离出来，形成一个独立的子模型，然后将该子模型转换为Relay IR并生成CUDA源代码。

步骤2：从TVM到CUDA源代码

TVM自动生成CUDA内核源代码。

Page 15 再次展示了TVM的工作流，强调了从Relay（高级IR）经过一系列优化和调度后，最终生成机器码（其中包含了CUDA源代码）。

步骤3：从TVM运行时到CUDA核函数参数

• 仅有CUDA核函数源代码是不够的。

  • __global__ void tvmgen_default_fused_add_kernel0(float* T_add, float* placeholder) { ... }

• CUDA核函数启动参数（如网格大小和块大小）也是必需的。

• TVM源代码经过修改，以提供必要的信息。
  • int tvm_test_add::enqueue(const nvinfer1::PluginTensorDesc* inputDesc, const nvinfer1::PluginTensorDesc* outputDesc, const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept { ... tvmgen_default_fused_add_kernel0<<<dimGrid, dimBlock, 0, stream>>>((float*)outputs[0], (float*)inputs[0]); }

复杂案例：一个ONNX节点 -> 多个CUDA核函数

• 转换流程：

  • ONNX Op (OneHot) 经过转换器到TVM-Relay (op.where, op.onehot)。
  • 通过融合通道，产生融合任务 (fused_less_add_where, fused_one_hot_kernel0)。
  • 最终生成两个CUDA核函数。

• 示例CUDA核函数代码：

  • __global__ void fused_less_add_where0(int* T_where, int* placeholder, int* placeholder1) { ... }
  • __global__ void fused_one_hot_kernel0(float* T_one_hot, int* placeholder, float* placeholder1, float* placeholder2) { ... }

• 问题：我们如何解释这些核函数参数（如placeholder）的含义，并将它们映射到TensorRT插件的输入/输出？

• 如果存在多个CUDA核函数，可能需要额外的设备内存来存储中间结果。
• 对于TensorRT插件，此额外的设备内存应通过workspace进行管理以获得更好的性能。
• workspace的总大小及相关偏移量可以从TVM运行时计算得出。

• 此外，这些CUDA核函数的输入/输出顺序应仔细处理。

• OneHot操作详情:

  • 生成基于输入的one-hot张量。
  • Indices：包含索引的张量。
  • Depth：指定one-hot张量中类数量的标量。
  • Values：rank 1张量 [off_value, on_value]，例如 [0, 1]。

• Host Function 到 Device Function 的映射:

  • Fused_less_add_where 对应 Kernel_0
  • Fused_one_hot 对应 Kernel_1

• TVM内存管理 -> TensorRT插件入队参数顺序:

EID	0	1	2	3	4	5
Size	131072	4	131072	4	4	3355432
Pointer	inputs[0]	workspace	workspace +4	workspace +131076	workspace +131080	outputs[0]
Meaning	input_0	depth	intermediate variable	on_value	off_value	output_0

• fused_less_add_where 的CUDA核函数调用及参数映射：
  • 调用：fused_less_add_where<<<dimGrid, dimBlock, 0, stream>>>((int*) (workspace + 4), (int*)inputs[0], (int*)workspace);
  • 核函数定义：__global__ void fused_less_add_where(int* T_where, int* placeholder, int* placeholder1)

	EID_2	EID_0	EID_1
Host function	workspace + 4	inputs[0]	workspace
Device function	T_where	placeholder	placeholder1
Meaning	intermediate variable	input_0	depth

• fused_one_hot 的CUDA核函数调用及参数映射：
  • 调用：fused_one_hot<<<dimGrid, dimBlock, 0, stream>>>((float*)outputs[0], (int*)(workspace + 4), (float*)(workspace + 131076), (float*)(workspace + 131080));
  • 核函数定义：__global__ void fused_one_hot(float* T_one_hot, int* placeholder, float* placeholder1, float* placeholder2)

	EID_5	EID_2	EID_3	EID_4
Host function	outputs[0]	workspace+4	workspace+131076	workspace+131080
Device function	T_one_hot	placeholder	placeholder1	placeholder2
Meaning	output_0	intermediate variable	on_value	off_value

步骤4：自动填充TensorRT插件模板

• Jinja2:

  • 一个快速且广泛使用的Python模板引擎。
  • 提供丰富的控制结构（循环和条件语句）。

• 将TensorRT插件接口视为字符串模板，将动态内容留给Jinja2引擎进行渲染。

• Jinja2工作流程:

  • Template 和 Data 输入到 Jinja2 engine。
  • Jinja2 engine 生成 Final document。

• 示例代码:

from jinja2 import Template
template = Template('Hello {{ name }}!')
outputText = template.render(name='John')
print(outputText)
# Hello John!

• 通过Jinja2自动生成TensorRT插件的代码示例：
  • template_plugin.h (模板文件，包含占位符，例如 {{plugin_output_number}})。
  • run_template.py (Python脚本，使用Jinja2加载模板并渲染，传入参数如 plugin_name, plugin_output_number, plugin_output_shape)。
  • generated_plugin.h (生成的插件头文件，占位符已被实际值填充，例如 return 1; 和 return nvinfer1::Dims3(1,16,16);)。

步骤5：修改ONNX图

• TensorRT ONNX解析器的回退机制要求插件定义必须与ONNX操作符保持一致，即名称、输入/输出。

• 我们需要使ONNX节点的op类型与插件名称对齐，并在以下情况下特别区分：

  • 相同的op类型但不同的TensorRT插件。
  • TensorRT ONNX解析器支持的op类型。
  • 其他特殊情况。

• TPAT允许用户为给定的ONNX节点指定TensorRT插件名称。否则，默认插件命名约定为 "tvm" + op_name。

• 对于原始ONNX图（左侧），调优节点是 one_hot (op name) 和 OneHot (op type)。

• op类型必须修改为与插件名称相同，即 tvm_one_hot。
• 再次利用 ONNX graphsurgeon 使此步骤自动化。
• 图示:
  • 左侧原始ONNX图：MatMul -> Cast -> OneHot -> Reshape
  • 右侧修改后的ONNX图：MatMul -> Cast -> tvm_one_hot -> Reshape
  • 修改后的图经过 TensorRT plugin -> ONNX parser -> TensorRT engine。

复杂案例：多个ONNX节点需要插件

• 维护一个ONNX节点及其对应的TensorRT插件的映射。
• 启用TensorRT插件缓存以避免重复的自动调优。
• 例如，中间的 OneHot 节点与右侧的相同，而左侧的则不同。
  • 中间和右侧：values [0, 1]
  • 左侧：values [-1, 1]

• 正如修改后的ONNX图所示，中间和右侧的 OneHot 节点共享相同的TensorRT插件，即 tvm_one_hot_1，而左侧节点对应 tvm_one_hot_2 插件。
• 在为所有ONNX节点自动生成TensorRT插件后，操作符类型会被修改。
• 图示:
  • 原始ONNX图（类似26页），但 OneHot 节点被替换为 tvm_one_hot_1 或 tvm_one_hot_2。

子图自动生成

• TensorRT子图融合:

  • 层融合对TensorRT的卓越性能贡献巨大。
  • 然而，TensorRT的层融合是静态的，并在内部维护。
  • TensorRT用户必须手动融合子图并编写TensorRT插件。

• TPAT能够为ONNX子图自动生成TensorRT插件。

• 性能高度依赖于TVM图级别和操作符级别优化的能力。

LayerNormalization子图自动生成

• 展示了一个复杂ONNX子图（包含Reshape, ReduceMean, Sub, Mul, Add, Pow, Div等操作）如何通过TPAT转换为TensorRT插件的过程。
• 流程图:
  • 左侧复杂ONNX子图 -> Relay IR -> CUDA source code -> TensorRT plugin -> RedSubMulRedAddPowDivMulAdd (融合后的操作) -> Add。

TPAT 示例

一个完整的 one_hot 操作到TensorRT引擎的TPAT流程示例，如下图所示（包含op name: one_hot 和 batch size: 256 的ONNX模型，经过Cast, OneHot, Reshape等操作）。

TPAT模块的工作流程如下：

1. ONNX模型输入。
2. 通过Relay and Autoscheduler进行优化。
3. 生成CUDA Kernel。
4. 填充Plugin Template。
5. 生成TensorRT Plugin。
6. 最终由Trtexec tool生成TensorRT engine。

下图展示了TPAT的工作流程示例，特别是如何将ONNX模型中的操作转换为优化的TensorRT插件，并包含了生成的CUDA内核代码片段，以及一个tvm_one_hot_enqueue函数的C++代码片段，说明了TPAT如何自动生成用于TensorRT插件的底层CUDA实现。

结果

与手写插件比较

该幻灯片展示了TPAT与手写插件在不同操作上的性能比较结果。实验设备为Nvidia Tesla T4。

主要结果如下：
* GatherNd 操作：对于 shape=(1, 128, 128, 224)，手写插件耗时 180.0 ms，TPAT 耗时 127.0 ms，提升 1.4x。
* GatherNd 操作：对于 shape=(1, 158, 128, 128)，手写插件耗时 21.1 ms，TPAT 耗时 17.5 ms，提升 1.2x。
* ScatterND-update 操作：对于 data=(32, 128, 128, 256)，手写插件耗时 0.36 ms (ScatterND-kernel)，TPAT 耗时 0.18 ms (ScatterND-kernel)，提升 2.0x。
* ScatterND-Add 操作：对于 data=(32, 128, 128, 256)，手写插件耗时 2.37 ms，TPAT 耗时 2.13 ms，提升 1.1x。
* OneHot 操作：对于 depth=16384，手写插件耗时 9.3 us，TPAT 耗时 11.6 us，性能比为 0.8x（即TPAT略慢）。

优化 TensorRT 算子

该幻灯片展示了TPAT对现有TensorRT算子的优化结果，对比了TensorRT 7.2.2.3与TPAT的性能。实验设备为Nvidia Tesla T4。

主要结果如下：
* Resize 操作：对于 NCHW(32, 32, 32, 44, 50) -> (32, 32, 176, 200)，TensorRT 7.2.2.3 耗时 1.450 ms，TPAT 耗时 1.079 ms，提升 1.38x。
* Resize 操作：对于 NCHW(32, 64, 11, 12) -> (32, 64, 176, 200)，TensorRT 7.2.2.3 耗时 2.759 ms，TPAT 耗时 2.171 ms，提升 1.27x。
* Reduce_max 操作：对于 input_shape=(256, 82, 128)，output_shape=(256, 128)，TensorRT 7.2.2.3 耗时 0.167 ms，TPAT 耗时 0.036 ms，提升 4.63x。

支持的算子

该幻灯片总结了TPAT支持的算子数量：
* TPAT 共支持 137 个算子。
* 其中有 17 个算子由 TPAT 支持，但 TensorRT 不支持。

下表列出了这些由 TPAT 支持但 TensorRT 不支持的算子：

表格中列出的算子包括：
* BitShift
* Compress
* ConcatFromSequence
* Einsum
* Hardmax
* IsInf
* MaxRoIPool
* MaxUnpool
* Mod
* NegativeLogLikelihoodLoss
* OneHot
* RandomNormal
* Reciprocal
* RoiAlign
* Shrink
* SoftmaxCrossEntropyLoss
* Xor

总结

本幻灯片对TPAT进行了总结，并展望了未来计划：
* TPAT（TensorRT Plugin Autogen Tool）消除了为不支持或效率低下的算子手动编写TensorRT插件的困扰。
* 它利用TVM的力量自动生成CUDA内核，并且是完全自动化的。
* 它基于TensorRT和TVM的ONNX工作流。
* TPAT即将开源。欢迎使用并提供反馈！

未来计划：
* 针对TensorRT的激进子图优化 (Radical Subgraph optimization for TensorRT)
* 动态形状支持 (Dynamic shape support)
* 半精度和INT8支持 (Half precision and INT8 support)

招聘信息

我们正在招聘！腾讯AI Lab深度学习编译器团队
联系方式：[email protected]