TensorRT-Plugin编写

1 TensorRT Plugin 初识

实现原生不支持的算子 是 Plugin 最基础的能力，当然它还可以做更多事情，比如手动融合 TensorRT 没有自动融合的层 或 块。总结来说，TensorRT Plugin 的功能主要有以下几点：
(1) 实现 TensorRT 原生不支持的算子或 块；
(2) 手动融合 TensorRT 没有自动融合的算子 或块；
(3) 替换你认为性能不够的算子或块；

TensorRT Plugin 从实现形式来说，需要自己使用 CUDA C 来编写 CUDA Kernel，然后以 .so 的形式 嵌入到咱们网络的 TensorRT network 构建中，这可不是一件容易的事情。我认为其中的难度主要体现在：
(1) 写出高性能的 CUDA C Kernel 本身并不是一件简单的事情；
(2) 自己写的 Kernel，若没有经过系统测试，并没有原生算子那么可靠；

2. 插件编写

首先按照官方plugin的排布方式，下面以scatternd作为例子。为了验证编写的插件是有效的，将包含scatternd的onnx中的scatternd的名字修改为scatternd_test，修改后的onnx在转换tensorRT模型的时候，就会因为找不到scatternd_test从而报错。
修改onnx里面节点信息的代码如下：

import onnx

onnx_model = onnx.load("***.onnx")
graph = onnx_model.graph
node  = graph.node

for i in range(len(node)):
    if node[i].name == "ScatterND_148": 
        node[i].op_type = "ScatterND_test"

# onnx.checker.check_model(onnx_model)
onnx.save(onnx_model, 'out.onnx')

准备一个自己的插件:scatterPlugin.cpp、scatterPlugin.h和scatterLayer.cu，copy并paste官方代码，名字替换成自己的。
我们需要写两个类：
scatterPlugin，继承IPluginV2DynamicExt，是插件类，用于写插件具体的实现（插件类继承IPluginV2DynamicExt才可以支持动态尺寸）
scatterndPluginCreator，继承BaseCreator，是插件工厂类，用于根据需求创建该插件

2.1.插件类

Plugin类中构造函数

构造函数一般设置为三个：
第一个用于在parse阶段，PluginCreator用于创建该插件时调用的构造函数，需要传递权重信息以及参数。
第二个用于在clone阶段，复制这个plugin时会用到的构造函数。
第三个用于在deserialize阶段，用于将序列化好的权重和参数传入该plugin并创建onnx。

getNbOutputs

插件op返回多少个Tensor，比如MyCustomPlugin这个操作只输出一个Tensor(也就是一个output)，所以直接return 1：

int ScatterND_test::getNbOutputs() const noexcept
{
    // Plugin layer has 1 output
    return 1;
}

initialize初始化函数

该函数在该插件准备开始run之前执行。主要初始化一些提前开辟空间的参数，一般是一些cuda操作需要的参数(例如conv操作需要执行卷积操作，我们就需要提前开辟weight和bias的显存)，假如我们的算子需要这些参数，则在这里需要提前开辟显存。

int ScatterND_test::initialize() noexcept
{
    return 0;
}

getOutputDataType

返回结果的类型，一般来说我们插件op返回结果类型与输入类型一致：

// Return the DataType of the plugin output at the requested index
DataType ScatterND_test::getOutputDataType(int index, const nvinfer1::DataType* inputTypes, int nbInputs) const noexcept
{
    PLUGIN_ASSERT(index == 0);
    return inputTypes[dataTensorIdx];
}

getWorkspaceSize

这个函数需要返回这个插件op需要中间显存变量的实际数据大小(bytesize)，这个是通过TensorRT的接口去获取，是比较规范的方式。
我们需要在这里确定这个op需要多大的显存空间去运行，在实际运行的时候就可以直接使用TensorRT开辟好的空间而不是自己去申请显存空间。

size_t ScatterND_test::getWorkspaceSize(const PluginTensorDesc* inputs, int32_t nbInputs, const PluginTensorDesc* outputs,int32_t nbOutputs) const noexcept
{
    int32_t nSlices = calculateNumSlices(inputs[indexTensorIdx].dims);
    //transformCoeffs + transformed indices
    return outputs[0].dims.MAX_DIMS * sizeof(int32_t) + nSlices * sizeof(int32_t);
}

enqueue

实际插件op的执行函数，我们自己实现的cuda操作就放到这里(当然C++写的op也可以放进来，不过因为是CPU执行，速度就比较慢了)，与往常一样接受输入inputs产生输出outputs，传给相应的指针就可以。

int32_t ScatterND_test::enqueue(const PluginTensorDesc* inputDesc, const PluginTensorDesc* outputDesc,
    const void* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
{
    int32_t transformCoeff[nvinfer1::Dims::MAX_DIMS];
    std::memset(transformCoeff, 0, sizeof(int32_t) * outputDesc[0].dims.MAX_DIMS);
    Dims IndexDims = inputDesc[indexTensorIdx].dims;

    Dims dataDims = inputDesc[dataTensorIdx].dims;

    int32_t indexRank = IndexDims.d[IndexDims.nbDims - 1];
    PLUGIN_ASSERT(indexRank <= dataDims.nbDims);

    int32_t nSlices = calculateNumSlices(IndexDims);
    int32_t rowSize = 1;
    int32_t copySize = calculateCopySize(dataDims);
    int32_t elementSizeInBytes = 1;
    switch (inputDesc->type)
    {
    case DataType::kFLOAT:
    case DataType::kINT32:
        elementSizeInBytes = 4;
        break;
    case DataType::kHALF:
        elementSizeInBytes = 2;
        break;
    case DataType::kINT8:
    case DataType::kUINT8:   // *******************
    case DataType::kBOOL:
        elementSizeInBytes = 1;
        break;
    }
    
    for (int i = indexRank; i < dataDims.nbDims; i++)
    {
        rowSize *= dataDims.d[i];
    }
    
    calculateTransformCoeff(dataDims, indexRank, transformCoeff);

    scatterNDInference_test(stream, transformCoeff, 
    dataDims.nbDims, 
    indexRank, 
    nSlices, 
    rowSize,  
    copySize, 
    elementSizeInBytes,  
    inputs[indexTensorIdx],
    inputs[updateTensorIdx],
    inputs[dataTensorIdx],
    outputs[0],
    workspace );    

    return 0;
}

注意事项：我们默认写的.cu是fp32的，TensorRT在fp16运行模式下，运行到不支持fp16的插件op时，会自动切换到fp32模式，等插件op运行完再切换回来。

getOutputDimensions

TensorRT支持Dynamic-shape的时候，batch这一维度必须是explicit的，也就是说，TensorRT处理的维度从以往的三维[3,-1,-1]变成了[1,3,-1,-1]。最新的onnx-tensorrt也必须设置explicit的batchsize，而且这个batch维度在getOutputDimensions中是可以获取到的。

nvinfer1::DimsExprs ScatterND_test::getOutputDimensions(int32_t outputIndex, const DimsExprs* inputs, int32_t nbInputs, IExprBuilder& exprBuilder)  noexcept
{
    //output should have same dimensions as data tensor
    DimsExprs ret = inputs[dataTensorIdx];
    return ret;
}

set/getPluginNamespace

为这个插件设置namespace名字，如果不设置则默认是""，需要注意的是同一个namespace下的plugin如果名字相同会冲突

// Set plugin namespace
void ScatterND_test::setPluginNamespace(const char* pluginNamespace) noexcept
{
    mPluginNamespace = pluginNamespace;
}

const char* ScatterND_test::getPluginNamespace() const noexcept
{
    return mPluginNamespace.c_str();
}

configurePlugin

配置这个插件op，判断输入和输出类型数量是否正确。官方还提到通过这个配置信息可以告知TensorRT去选择合适的算法(algorithm)去调优这个模型。

但自动调优目前还没有尝试过，我们一般自己写的plugin执行代码都是定死的，所谓的调优步骤可能更多地针对官方的op。

void ScatterND_test::configurePlugin(const DynamicPluginTensorDesc* in, int32_t nbInputs, const DynamicPluginTensorDesc* out, int32_t nbOutputs) noexcept
{
}

clone

将这个plugin对象克隆一份给TensorRT的builder、network或者engine。这个成员函数会调用上述说到的第二个构造函数：

IPluginV2DynamicExt* ScatterND_test::clone() const noexcept
{
    try
    {
        // Create a new instance
        ScatterND_test* plugin = new ScatterND_test();
        plugin->setPluginNamespace(mPluginNamespace.c_str());
        return plugin;
    }
    catch (std::exception const& e)
    {
        nvinfer1::plugin::caughtError(e);
    }
    return nullptr;
}

getSerializationSize
返回序列化时需要写多少字节到buffer中。

size_t ScatterND_test::getSerializationSize() const noexcept
{
    
    return 0;
}

supportsFormatCombination

TensorRT调用此方法以判断pos索引的输入/输出是否支持inOut[pos].format和inOut[pos].type指定的格式/数据类型。

bool ScatterND_test::supportsFormatCombination(
    int32_t pos, const PluginTensorDesc* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
{
    PLUGIN_ASSERT(pos < 4);
    PLUGIN_ASSERT(nbInputs == 3);
    PLUGIN_ASSERT(nbOutputs == 1);
    const PluginTensorDesc& desc = inOut[pos];
    bool ret = false;
    switch (pos)
    {
    case dataTensorIdx:
    case updateTensorIdx:
        ret = ((desc.type == DataType::kFLOAT || desc.type == DataType::kINT32)
            && desc.format == TensorFormat::kLINEAR);
        break;
    case indexTensorIdx:
        ret = (desc.type == DataType::kINT32 && desc.format == TensorFormat::kLINEAR);
        break;
    case 3:
        ret = ((desc.type == DataType::kFLOAT || desc.type == DataType::kINT32) && desc.format == TensorFormat::kLINEAR);
        break;
    }
    return ret;
}

serialize

把需要用的数据按照顺序序列化到buffer里头。

void ScatterND_test::serialize(void* buffer) const noexcept
{
    return;
}

attachToContext

如果这个op使用到了一些其他东西，例如cublas handle，可以直接借助TensorRT内部提供的cublas handle:

void ScatterND_test::attachToContext(cudnnContext* cudnn, cublasContext* cublas, IGpuAllocator* gpuAllocator) noexcept
{
    return;    
}

2.2.插件工厂类

构造函数

创建一个空的mPluginAttributes初始化mFC

ScatterNDPluginCreator_test::ScatterNDPluginCreator_test()
{
    mFC.nbFields = 0;
}

createPlugin

这个成员函数作用是通过PluginFieldCollection去创建plugin，将op需要的权重和参数一个一个取出来，然后调用上文提到的第一个构造函数：

IPluginV2* ScatterNDPluginCreator_test::createPlugin(const char* name, const PluginFieldCollection* fc) noexcept
{
    try
    {
        ScatterND_test* obj = new ScatterND_test();
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }
    catch (std::exception const& e)
    {
        nvinfer1::plugin::caughtError(e);
    }
    return nullptr;
}

deserializePlugin

这个函数会被onnx-tensorrt的一个叫做TRT_PluginV2的转换op调用，这个op会读取onnx模型的data数据将其反序列化到network中。

IPluginV2* ScatterNDPluginCreator_test::deserializePlugin(
    const char* name, const void* serialData, size_t serialLength) noexcept
{
    try
    {
        // This object will be deleted when the network is destroyed, which will
        // call Normalize::destroy()
        ScatterND_test* obj = new ScatterND_test();
        obj->setPluginNamespace(mNamespace.c_str());
        return obj;
    }
    catch (std::exception const& e)
    {
        nvinfer1::plugin::caughtError(e);
    }
    return nullptr;
}

3. 关于plugin的注册

注册插件主要有两种方法，第一种基于tensorrt oss来更新libnvinfer_plugin.so，第二种基于REGISTER_TENSORRT_PLUGIN来注册插件

3.1 基于tensorrt oss注册插件

1. 首先是插件包的下载

git clone -b release/8.2 https://github.com/nvidia/TensorRT
cd TensorRT/
git submodule update --init --recursive
export TRT_SOURCE=`pwd`
cd $TRT_SOURCE
mkdir -p build && cd build

注释：
1）git clone中的-b 后面添加的是tensorrt版本，由于我的tensorrt版本是tensorrt8.2.1
2）git submodule update --init --recursive 会根据.gitmodules 文件里面的子仓库信息，来下载子仓库
3）PX4环境git submodule update --init --recursive失败的解决办法

2. cmake生成makefile

cmake .. -DGPU_ARCHS=75  -DTRT_LIB_DIR=/home/***/software/TensorRT-8.2.1.8/lib -DCMAKE_C_COMPILER=/usr/bin/gcc -DTRT_BIN_DIR=`pwd`/out

注释：
1）DGPU_ARCHS：是根据GPU算力来决定的，1660ti对应的是75的算力，具体的算力可以自己去查。
2）DTRT_LIB_DIR：这个参数填写的是正常下载Tensorrt解压包的对应的lib的位置，例如在我电脑上位置为home/***/software/TensorRT-8.2.1.8/lib
3）DTRT_BIN_DIR:这个参数对应的是编译生成的新的Tensorrt-lib的位置，在make过后会在out文件夹下新的Tensorrt的lib文件，在后续的使用中直接在环境变量链接该out文件夹就可以使用最新的Tensorrt-lib。
3. 编译插件

make nvinfer_plugin -j$(nproc)

4. 安装所有插件

sudo make install

这一步能跳过，因为第三步就已经生成了nvinfer_plugin的库，直接调到下一步。

5. 拷贝so文件
   在build文件中，找到libnvinfer_plugin.so、libnvinfer_plugin.so.8、libnvinfer_plugin.so.8.2.,将这三个文件直接拷贝到tensorrt的lib目录上home/**/software/TensorRT-8.2.1.8/lib

TensorRT OSS 在linux环境下编译安装

3.2 REGISTER_TENSORRT_PLUGIN来注册插件

在官方上看是否存在有已实现的插件。如果有的话，就参考它的实现，注意实现如下：
主要的修改
1）在h文件、cpp文件和cu文件，将里面的头文件对应的文件夹拷贝到项目中
2）在plugin的.cpp文件中,添加REGISTER_TENSORRT_PLUGIN(ScatterNDPluginCreator_test),同时
3）遇到部分代码可以注释掉，看一下有没有影响 。例如case DataType::kUINT8这行代码，会导致报 ‘nvinfer1::DataType’; did you mean ‘kINT8’的错误，所以我注释了，该插件也能正常使用。