TensorRT Plugins in FX

[narendasan]

Plugins in FX

TL;DR

We ship a number of plugins which implement certain key operations in TensorRT for use in TorchScript frontend converters. We would like to be able to use these plugins and new ones in FX as well.

Goal(s)

Allow users to leverage the torchtrt_plugins.so library from FX and provide an area for FX frontend developers to put plugins.

Usecases

Proposed APIs / UX

We already ship libtorchtrt_plugins.so as part of the torch-tensorrt package. So the end user UX would not be any different. The developer UX would involve exposing these plugins in python (likely through standard TensorRT python APIs) so that they may be called and used.

• https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Plugin/IPluginRegistry.html
• https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html#ipluginv2layer
• https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#example1_add_custom_layer_python

Example Workflow

import tensorrt as trt
import numpy as np

TRT_LOGGER = trt.Logger()

trt.init_libnvinfer_plugins(TRT_LOGGER, '')
PLUGIN_CREATORS = trt.get_plugin_registry().plugin_creator_list

def get_trt_plugin(plugin_name):
    plugin = None
    for plugin_creator in PLUGIN_CREATORS:
        if plugin_creator.name == plugin_name:
            lrelu_slope_field = trt.PluginField("neg_slope", np.array([0.1], dtype=np.float32), trt.PluginFieldType.FLOAT32)
            field_collection = trt.PluginFieldCollection([lrelu_slope_field])
            plugin = plugin_creator.create_plugin(name=plugin_name, field_collection=field_collection)
    return plugin

@tensorrt_converter(torch.lrelu)
def mul(network, target, args, kwargs, layer_name):
    if len(kwargs) == 0:
        lhs_val, rhs_val = args
    else:
        lhs_val, rhs_val = kwargs["input"], kwargs["other"]

    if not all(isinstance(arg, trt.tensorrt.ITensor) for arg in \[lhs_val, rhs_val\]):
        raise RuntimeError("lrelu() received an input that is not part of the TensorRT region!")
        
    lrelu = network.add_plugin_v2(inputs=[lhs_val, rhs_val], plugin=get_trt_plugin("LReLU_TRT"))
    lrelu.name = layer_name
    return lrelu.get_output(0)

Limitations

Users who use converters based on python could not use the --fx-only build mode since there is now a dependency on C++ libraries.

Internal Implementation

Design

There shouldn't be too much invloved other than perhaps some convienence functions to load in the torch_tensorrt plugin registery in addition to the TensorRT one.

Extensions Required to Core API implementations

N/A

Data Structures

N/A

Details specific for TorchScript Support

N/A

Details specific for FX support

See above.

Implementation Phases

Prototype - S

Implement a converter using on of the torch_tensorrt_plugin plugins, such as one of the norms.

MVP (<TARGET RELEASE VERSION>) - S

Converters for the few plugins we ship and utilities to easily use the library from python.

[peri044]

1. So according to the prototype, Implement a converter using on of the torch_tensorrt_plugin plugins, such as one of the norms.
   If we have such a converter merged into fx path, --fx-only install option is eliminated correct ?

2. Another alternative is implementing plugins purely in python. I'm not sure if it's possible today. I don't see the methods of the plugin exposed in the documentation atleast to do that https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/Graph/Layers.html?highlight=ipluginv2#ipluginv2layer

[narendasan]

No you would disable the converted based on if the plugin was available similar to the torchvision converters

[narendasan]

I think implementing plugins in python would be hard