torchtrtc custom torch operators and custom torch converters

[andi4191]

Background

torchtrtc is a CLI frontend for Torch-TensorRT. It is meant to provide access to (primarily TorchScript) Torch-TensorRT APIs from the command line for simple compilation tasks.

Problem addressed

There are two classes of problems this RFC would address.

1. Custom Torch Operators

TorchScript allows for the definition of custom operators to extend to core PyTorch library. A familiar example of this would be operators like NMS defined by Torchvision, but many PyTorch ecosystem libraries ship custom operators (including Torch-TensorRT). When a model using custom operators is being deserialized, definitions of these operators must be registered with PyTorch beforehand. Currently torchtrtc has no mechanism to load these custom operators and therefore cannot deserialize TorchScript modules that use them.

2. Custom Torch-TensorRT Converters

Torch-TensorRT ships with a core set of converters which is not exhaustive of the entire PyTorch operator library. There is also currently no plan to support any 3rd party operators in the Torch-TensorRT converter library. As such, Torch-TensorRT exposes mechanisms to extend the converter library so that developers can add functionality as needed. Currently, these extensions cannot be used by torchtrtc.

Proposal

1. Custom Operators:

A new flag called --custom-torch-ops (repeatable) <path to shared object library/DLL containing custom torch ops> should be added to torchtrtc. The idea of this argument would be to take a path or paths to libraries defining custom ops and load them prior to deserializing the target TorchScript program.

2. Custom Torch-TensorRT Converters

A new flag called --custom-converters (repeatable) <path to shared object library/DLL containing custom converters> should be added to torchtrtc. The idea of this argument would be to take a path or paths to libraries defining custom converters and load them prior to compiling the target TorchScript module.

Implementation

Since both PyTorch operators and Torch-TensorRT converters are registered via a process called "self-registering classes", extensions are automatically loaded as part of the init step of loading a library. Therefore even though two new APIs are proposed, under the hood they use the same process of simply DL_OPEN'ing the provided shared object files prior to loading and compiling the provided program.