chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Waiting to run
Docker Image CI / build-ubuntu2004 (push) Waiting to run
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,171 @@
|
||||
# Python-based TRT Plugins
|
||||
|
||||
This is a sample to showcase Python-based plugin definitions in TRT. No changes to existing TRT APIs have been made
|
||||
to deliver this feature, so using the updated bindings should not break any existing code.
|
||||
|
||||
## Introduction
|
||||
|
||||
Until TRT 9.1, plugin implementations could only be done through the TRT C++ API. To use a plugin in a Python app, one had to
|
||||
- Implement plugin in C++ and build into a shared library
|
||||
- Load plugin lib and register plugin creator (statically or dynamically)
|
||||
- Retrieve plugin creator and create plugin instance through the respective Python API
|
||||
|
||||
The following design considerations were followed in creating bindings to allow Python-based plugin definitions:
|
||||
- Zero additional C++ code shall be required to implement, integrate and run a plugin within TensorRT
|
||||
- Offer the flexibility to implement the kernel(s) for the plugin through any method of choice
|
||||
- Many libraries have sprung up to provide CUDA kernel support with AOT/JIT compilation
|
||||
- Numba, OpenAI Triton, CuPy etc.
|
||||
- Could even do without explicit kernels (e.g. leverage PyTorch functional op)
|
||||
|
||||
- Will only support `IPluginV2DynamicExt` and `IPluginV3`-based plugins
|
||||
- Other plugin interfaces (except `IPluginV2IOExt`) are deprecated since TRT 8.5
|
||||
|
||||
With these bindings, plugins can be implemented and integrated to TRT purely with Python.
|
||||
|
||||
## Setting Up The Build Environment
|
||||
|
||||
To build and install the bindings, follow the instructions in `$TRT_OSSPATH/python/README.md`.
|
||||
|
||||
Then install the requisite packages
|
||||
```bash
|
||||
cd $TRT_OSSPATH/samples/python/trt_python_plugin
|
||||
pip3 install -r requirements.txt
|
||||
```
|
||||
Install `cupy-cuda11x` instead if testing on a CUDA 11.x environment.
|
||||
|
||||
# TensorRT Plugin API for Python
|
||||
|
||||
Implementing a TRT plugin in Python is similar to C++ in that implementation of `IPluginV2DynamicExt`+`IPluginCreator` or `IPluginV3`+`IPluginCreatorV3One` is necessary. Refer to the TensorRT Python API reference for a concise description.
|
||||
|
||||
## Differences in C++ and Python APIs for `IPluginV2DynamicExt`
|
||||
The interface methods in Python have mostly similar APIs to their C++ counterparts, except for `serialize()` and `enqueue()`.
|
||||
- While the C++ API for `serialize()` is `void serialize (void *buffer)` where the plugin writes to the passed-in `buffer`, the Python API is `serialize(self) -> bytes`, where the implementation of the method is expected to return a bytes object containing a serialized representation of the plugin object.
|
||||
- In `enqueue()`, the device pointers for input and output tensors are passed as their `intptr_t` casts. Since these buffers are created and owned by TRT, care must be taken when writing to them from the Python side.
|
||||
- No bindings yet for `attachToContext()` and `detachFromContext()` which are not pure virtual.
|
||||
|
||||
# Running the sample: Circular padding plugin
|
||||
|
||||
This sample contains a circular padding plugin, where the `enqueue` has been implemented with various frameworks for writing kernels or executing GPU ops (torch).
|
||||
|
||||
Each script accepts a command-line argument to choose precision from either FP32 or FP16. e.g.
|
||||
```bash
|
||||
python3 circ_pad_plugin_cuda_python.py --precision fp32 # fp32 or fp16
|
||||
```
|
||||
|
||||
## Circular padding
|
||||
|
||||
Circular padding is useful for ops like circular convolution in deep learning. The following image denotes how the original image (red) is circular padded once (green) and twice (blue):
|
||||
|
||||
|
||||
|
||||
The plugin shall have the following characteristics:
|
||||
- Input: 4-dimensional input (e.g. NxCxHxW)
|
||||
- Attribute(s): m-dimensional parameter `pads` where $m$ is even and $m/2 \le 4$. `pads` denotes the amount of padding to apply before and after each of the $m/2$ last dimensions of the input tensor.
|
||||
- Output: Padded tensor. Shape depends on `pads`.
|
||||
|
||||
## Baseline: Using a C++ plugin
|
||||
|
||||
To establish a baseline, we first demonstrate a C++ plugin implementing circular padding. The relevant files can be found in the `circ_plugin_cpp` folder: the included `CMakeLists.txt` can be used to build the shared library `libcirc_pad_plugin.so` / `circ_pad_plugin.dll`.
|
||||
|
||||
```bash
|
||||
cd $TRT_OSSPATH/samples/python/trt_python_plugin
|
||||
mkdir build && pushd build
|
||||
cmake .. && make -j
|
||||
popd
|
||||
python3 circ_pad_plugin_cpp.py --plugin-lib build/libcirc_pad_plugin.so
|
||||
```
|
||||
|
||||
## Python plugin: cuda-python
|
||||
|
||||
The cuda-python based implementation can be found in `circ_pad_plugin_cuda_python.py`. `cuda.nvrtc` is used to JIT compile a C/C++-based kernel, which is provided as a string. The compiled kernel is then launched through cuda-python's `cuda.cuLaunchKernel`.
|
||||
|
||||
`circ_pad_plugin_cuda_python.py` demonstrates an ONNX-based workflow: `circ_pad_plugin_inetdef_cuda_python.py` demonstrates a workflow where the model is constructed through `INetworkDefinition`.
|
||||
|
||||
## Python plugin: CuPy
|
||||
|
||||
The CuPy-based implementation can be found in `circ_pad_plugin_cupy.py`. CuPy's `RawKernel` class has been used to provide the C/C++-based kernel implementation as a string. CuPy will JIT compile the kernel.
|
||||
|
||||
## Python plugin: Triton (valid only on Linux)
|
||||
|
||||
The same plugin can be implemented with a Triton-based kernel as well. The only other change would be to `enqueue`. The entire implementation can be found in `circ_pad_plugin_triton.py`.
|
||||
|
||||
Some remarks:
|
||||
- Triton also allows for JIT-able kernels.
|
||||
- CuPy device arrays cannot be passed into Triton kernels directly -- only Torch arrays are accepted. However, we can use `torch.as_tensor()` to get around this constraint.
|
||||
- Triton does not seem to allow the specification of a CUDA stream.
|
||||
|
||||
## Python plugin: Numba
|
||||
|
||||
The Numba implementation can be found in `circ_pad_plugin_numba.py`. Some remarks:
|
||||
- Numba also allows for JIT-able kernels.
|
||||
- CuPy device arrays can be passed into Numba kernels without issue since CuPy arrays implement `__cuda_array_interface__`.
|
||||
|
||||
## Python plugin: Torch
|
||||
|
||||
The flexibility of the `enqueue()` interface means that it is not always necessary to implement a custom kernel. In this case, PyTorch's [torch.nn.functional.pad](https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html) offers the exact same capability we want, so we can use that inside `enqueue()`, as in `circ_pad_plugin_torch.py`.
|
||||
|
||||
## Python plugin: Multi-tactic, Multi-plugin (based on IPluginV3)
|
||||
|
||||
The entire implementation can be found in `circ_pad_plugin_multi_tactic.py`.
|
||||
|
||||
### Custom tactics
|
||||
|
||||
When multiple options are available to compute the same op, and it's not possible to reliably predict which one will be faster for the expected input shapes/types or the target platform,
|
||||
it is useful to ask TensorRT to time all available options during the build stage. In V2 plugins, TensorRT would only time different type/format combinations supported by the plugin, but
|
||||
V3 plugins allow users to specify any number of custom tactics to time also (in addition to type/format combinations).
|
||||
|
||||
In this example, we specify two custom tactics: PyTorch's [torch.nn.functional.pad](https://pytorch.org/docs/stable/generated/torch.nn.functional.pad.html) and a custom kernel written using OpenAI Triton.
|
||||
|
||||
It is possible to advertise tactics specific to a format combination. e.g. In this sample, we can support both tactics for FP32 I/O, and only support the OpenAI Triton tactic for FP16 I/O. To achieve this, return in `get_valid_tactics()` the set of tactics `T(f)` supported by the plugin for the format combination `f` indicated by the immediately preceding call to `configure_plugin()`. To enable this behavior in this sample, pass the flag `--per-format-tactics`.
|
||||
|
||||
### Multiple plugins instances
|
||||
|
||||
Imagine that you expect to have multiple instances of the same plugin in your network, which would operate on separate inputs, but where the input and output shapes/formats, as well
|
||||
as other determining plugin attributes would be the same. With V2 plugins, TensorRT would time all such plugin instances during the engine build -- however, this would be inefficient because the only salient difference between those instances are the values of the input tensors.
|
||||
|
||||
To communicate to TensorRT that you would like the timing for similar plugin instances to be cached, V3 plugins allow for the specification of a timing cache ID. The timing cache ID
|
||||
should only capture timing determinants extraneous to plugin I/O, like their shapes and formats. Typically, this would be the values of any plugin attributes that might be different
|
||||
between the plugin instances.
|
||||
|
||||
In this example,
|
||||
- The shape of the `pads` parameter affects timing, but only as far as it affects the output shape. Therefore, the timing cache ID could be an empty string.
|
||||
- We consider a scenario where there are two circular padding plugin instances with identical configurations. Therefore, only a single instance should be timed by TensorRT.
|
||||
This can be verified by inspecting the log.
|
||||
|
||||
# Limitations
|
||||
|
||||
- Plugins cannot be serialized into the engine (in contrast to `IBuilderConfig::setPluginsToSerialize()`)
|
||||
- Plugin class and Plugin Creator class must exist in the module where the engine is deserialized
|
||||
- The engine / ONNX model cannot be run from outside Python (e.g. with `trtexec`)
|
||||
- This functionality is possible to implement but comes at the cost of embedding the Python interpreter to the TRT runtime / the binary loading the engine
|
||||
- (For `IPluginV2DynamicExt` only) No bindings yet for `attachToContext()` and `detachFromContext()` which are not pure virtual.
|
||||
- `circ_pad_plugin_torch.py` may work on aarch64 platforms but is unsupported.
|
||||
|
||||
# FAQ
|
||||
|
||||
1. What are the performance impacts of a Python-based plugin versus a C++ one?
|
||||
|
||||
In preliminary testing, the Python overhead was found to be very minimal to negligible. In fact, if the kernels were compiled AOT (instead of JIT) the CuPY and Triton
|
||||
versions of the plugin were as performant as the C++ one. However, with Numba, there seems to be a significant kernel launch overhead.
|
||||
|
||||
2. Can I deploy a TRT engine including a Python plugin in a runtime environment without Python?
|
||||
|
||||
No. There is no way to fully embed a Python plugin into the engine that allows for it to be executed without the need for Python during inference time.
|
||||
|
||||
This design principle is what allows for the `enqueue()` to be implemented in any framework of choice.
|
||||
|
||||
# License
|
||||
|
||||
For terms and conditions for use, reproduction, and distribution, see the [TensorRT Software License Agreement](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sla/index.html) documentation.
|
||||
|
||||
# Changelog
|
||||
|
||||
October 2025: Migrate to strongly typed APIs.
|
||||
|
||||
August 2025: Removed support for Python versions < 3.10.
|
||||
|
||||
July 2023: Initial release of this sample
|
||||
|
||||
# Known issues
|
||||
|
||||
There are no known issues in this sample
|
||||
Reference in New Issue
Block a user