TRTorch plugin support

[peri044]

TRTorch plugin support

Currently we have the following plugins

1. Interpolate plugin (uses both pytorch aten and TRT layers)

2. Normalization plugin (in PR, uses pytorch aten calls)

3. adaptive max pooling (in PR, uses pytorch aten calls)

4. gelu plugin (in PR, uses TRT plugins)

5. Instance norm plugin (in PR, uses TRT plugins)

initLibNvInferPlugins should be used to initialize TensorRT plugins within TRTorch API.

Suggested approach

Build a new library called libtrtorch_plugins.so. It is the equivalent for libnvinfer_plugins.so in TRT.

Networks with TRT plugins or plugins with custom cuda kernels

Users should use libtrtorch_plugins.so if they are using existing TRT plugins in their network or if they implemented a plugin with their own custom cuda kernel (not using pytorch aten kernels as in plugins 1, 2, 3)

Plugins 4 and 5 should be basically compiled into libtrtorch_plugins.so. More custom kernels in the future will integrated into this library.

Benefit of this:

• Users can avoid the dependency on libtrtorch.so (is this true ?) and libtorch.so during runtime execution. For example, if a user serializes TRT engine (using TRTorch API), for standalone execution of the engine, we need TRT libraries and libtrtorch_plugins.so.
• Lightweight core library. Standard ops would be included in libtrtorch.so but network specific widely used ops (for eg: NMS, self-attention layers etc) would fall into libtrtorch_plugins.so
• Not a benefit but we can provide a standalone custom kernel example through TRTorch just like elu sample converter app. This can be demonstrated by integrating with libtrtorch_plugins.so for users to get started on custom ops.

Networks with plugins which use pytorch aten kernels

This usecase requires user to be dependent on torch, trtorch and tensorrt libraries. This would not require linking with libtrtorch_plugins.so.

Implementation

Currently the entry points for converting a Torchscript module into TRT engine are

• trtorch::ConvertGraphToTRTEngine
• trtorch::CompileGraph (which internally calls trtorch::ConvertGraphToTRTEngine)

trtorch::core::runtime::execute_engine consumes inputs and pre-built serialized_engine . So during runtime, we also need to initialize nvinfer plugins.

TensorRT plugins can be accessed via plugin registry. Sample registration code looks as follows.

std::unordered_map<std::string, nvinfer1::IPluginCreator*> mPluginRegistry;
initLibNvInferPlugins(&logger, "");
int numCreators = 0;
auto tmpList = getPluginRegistry()->getPluginCreatorList(&numCreators);
for (int k = 0; k < numCreators; ++k)
   {
     if (!tmpList[k])
        {
          std::cout << "Plugin Creator for plugin " << k << " is a nullptr." << std::endl;
          continue;
        }
     std::string pluginName = tmpList[k]->getPluginName();
     mPluginRegistry[pluginName] = tmpList[k];
     LOG_DEBUG("Register plugin: " << pluginName);
    }
LOG_DEBUG("Number of plugin: " << mPluginRegistry.size());

We can have a separate component under core/plugins which builds libtrtorch_plugins.so. By default, when we build TRTorch package we shall build libtrtorch.so and libtrtorch_plugins.so . However, we can have a flag that can disable build for plugins if not required (optional).

Code under core/plugins:

register_plugins.cpp (high level implementation. May differ in actual code.)

namespace trtorch{
    namespace plugins{
        void initializePlugins(){
            initLibNvInferPlugins(&logger, "");
            // Any external plugins (with cuda kernels) can be initialized here as well
            auto tmpList = getPluginRegistry()->getPluginCreatorList(&numCreators);
            for (int k = 0; k < numCreators; ++k)
               {
                 if (!tmpList[k])
                    {
                      std::cout << "Plugin Creator for plugin " << k << " is a nullptr." << std::endl;
                      continue;
                    }
                 std::string pluginName = tmpList[k]->getPluginName();
                 mPluginRegistry[pluginName] = tmpList[k];
                 LOG_DEBUG("Register plugin: " << pluginName);
                }
            // Add any external plugins to mPluginRegistry if required.
        }
        
    }
}

End-End workflow:

For a network with Gelu plugin (4) or an instance norm plugin (5), TRTorch compilation would be something like this

// global initialization of plugins. Only do it when necessary.
trtorch::plugins::initializePlugins(...) 
// Either this or if you have an engine, comment this and uncomment the next line
auto engine = trtorch::ConvertGraphToTRTEngine(g)
// trt_mod = trtorch::core::runtime::execute_engine(...)

Question: How do we handle converters for gelu and instance norm ops ? A typical converter op requires you to add ctx->addPluginV2 . If we have such ops, this may strictly require libtrtorch_plugins.so to be linked with libtrtorch.so . A workaround can be implementing converters for plugins also in core/plugins but this would break the independency of libtrtorch_plugins.so with libtrtorch.so.

[peri044]

@narendasan What are your thoughts on this ?

[narendasan]

Networks with TRT plugins or plugins with custom cuda kernels

Users should use libtrtorch_plugins.so if they are using existing TRT plugins in their network or if they implemented a plugin with their own custom cuda kernel (not using pytorch aten kernels as in plugins 1, 2, 3)

Plugins 4 and 5 should be basically compiled into separate library. More custom kernels in the future will integrated into this library

Can you explain the separate library here?

[peri044]

Oh by separate library I meant libtrtorch_plugins.so itself. Edited the post !!

[narendasan]

Users should use libtrtorch_plugins.so if they are using existing TRT plugins in their network or if they implemented a plugin with their own custom cuda kernel (not using pytorch aten kernels as in plugins 1, 2, 3)

Ideally the user never needs to think about if libtrtorch_plugins.so is linked unless they are running standalone TRT engines compiled with TRTorch in which case they would need to link it in their TRT app. It should otherwise be linked properly to libtrtorch.so and libtrtorchrt.so for the user.

IMO the idea should be that libtrtorch_plugins.so should encapsulate all plugins used by TRTorch during compilation which comes in one of two flavors, implemented in libnvinfer_plugin.so or implemented in source as part of the compiler. External plugin implementations made by users should exist separately and through the custom converter system.

The main reason to have a separate .so is to support people doing separate TRT apps otherwise we could just compile it as part of libtrtorch.so and libtrtorchrt.so

[narendasan]

We could maybe build in the plugin library into libtrtorch.so and libtrtorchrt.so actually and just redistribute the plugins libtrtorch_plugins.so only for the TRT app case. So libtrtorch.so and libtrtorchrt.so will have the plugin library fully baked in statically

[narendasan]

The question I think then is how you init the plugin library. Id like to avoid users making an init call in all but the TRT App case. For compile time is super easy because we can just call it in the compile function call. For runtime its a bit harder since we dont have a clear entry point. We could do some sort of fancy init function on load thing or we could maybe store a flag in the runtime that checks if plugins have been setup and on the first deseralization of a program call the init? We would need to make sure that the compile and runtime dont conflict or at least make sure that initing the plugin library twice doesnt cause problems

[narendasan]

Actually the simple way to have them in sync is to have the flag in the plugin library and then calls to compile or deserialize both check then init the plugin module

[peri044]

This is the approach we are targeting

All the plugins would be under libtrtorch_plugins.so. The full library libtrtorch.so would also have the plugins as well.
The cases we plan to support are

Case 1: Users only full library. They have a network with pytorch or TensorRT plugins. Since every plugin is built into libtrtorch.so, this network should run fine.

Case 2: - Users only using runtime library and have a TRT engine.
a) This is where libtrtorch_plugins.so comes into place. Link this as a dependency to libtrtorchrt.so

​ b) Only linklibtrtorch_plugins.so when it is necessary by the use and invoke it by user. This makes it complex and we don't plan to support this. If such a usecase arises, we will revisit.
​

Case 3: Standalone TRT network depends on plugins. In this case, users are not using trtorch runtime library. The plugin can be of two different flavours

​ a) nvinfer plugins : they can directly use trtexec or general TRT inference code and load libnvinfer_plugins.so
​ b) pytorch plugins : In this case, they need to use libtrtorch_plugins.so

How do we initialize plugin library?

option 1 : have a global flag which checks if plugin is initialized in compile time. If so, we don't need to initialize in runtime

option 2: Check auto registration of converters. More details will be added on this.