Automatic plugin generation using custom kernels and host code

[bowang007]

TL;DR

We want to automatically generate the plugin and converters given the kernel and host code. Users can include a custom kernel in a TensorRT engine using Torch-TensorRT and users don't need to write the plugin themselves. Torch-TRT does everything for us.

Goal(s)

Allow users to use custom kernels Torch-TensorRT engines without the effort of writing tensorrt plugin. Increase the model performance with graph breaks in model.

Usecases

• Automatic TensorRT plugin generation
• Performance increases

Proposed APIs/UX

As what is demonstrated here, the workflow for a plugin is usually shown here

In this case, user would need to provide the kernel code and then write the plugin according to the needs. What we would do here is:

1. Introduce some code generation utilities in Torch-TensorRT. If we take a look at the tutorial shown above, we could find there is a Plugin example, which is also demonstrated in TensorRT repo here. This could be a template, once we provide the kernel code, Torch-TensorRT will analyze the input tensor shape, output tensor shape and etc. After getting all required information, the template could be used to generate the plugin according to each kernel. In this process, some techniques such as PyTorch fake tensor, inference in PyTorch could be applied here.

2. Introduce the tensor shape parsing system in Torch-TensorRT to analyze the required information to generate the plugin.