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:
@@ -0,0 +1,352 @@
|
||||
# Quickly Deployable TensorRT Python Plugins [Experimental]
|
||||
|
||||
This is a sample to showcase quickly deployable Python-based plugin definitions (QDPs) in TensorRT (TRT). QDPs are able to support a large majority of use cases for adding custom operators to TRT, and will be the recommended option when it becomes a stable feature in 10.9.
|
||||
|
||||
This sample contains several mini-samples that demonstrate a few common use cases.
|
||||
|
||||
# Contents
|
||||
- [Introduction](#introduction)
|
||||
- [Setting up the environment](#setting-up-the-environment)
|
||||
- [Implementing a Quickly Deployable Python (QDP) Plugin](#implementing-a-quickly-deployable-python-qdp-plugin)
|
||||
- [A Simple Plugin: Elementwise-Add](#a-simple-plugin-elementwise-add)
|
||||
- [Implementing in-place custom ops with I/O aliasing](#implementing-in-place-custom-ops-with-io-aliasing)
|
||||
- [An Op with data-dependent output shapes: Non-zero](#an-op-with-data-dependent-output-shapes-non-zero)
|
||||
- [Using multiple tactics and ONNX: Cirular padding](#using-multiple-tactics-and-onnx-cirular-padding)
|
||||
- [Providing an Ahead-of-Time (AOT) implementation for Circular padding](#providing-an-ahead-of-time-aot-implementation-for-circular-padding)
|
||||
- [Additional Resources](#additional-resources)
|
||||
- [License](#license)
|
||||
- [Changelog](#changelog)
|
||||
- [Known issues](#known-issues)
|
||||
|
||||
|
||||
# Introduction
|
||||
|
||||
While the regular TRT plugin interfaces are powerful in the flexibility and tunability they provide, for the vast majority of use cases, users will benefit from the simplicity offered by the QDP workflow.
|
||||
- The `tensorrt.plugin` module provides many intuitive APIs that drastically reduces the amount of boilerplate required to implement a plugin
|
||||
- The concept of plugin registration, plugin creators and the plugin registry is abstracted away
|
||||
- The stateless nature of QDPs eliminates the complications of having to comply with a predefined plugin lifecycle
|
||||
|
||||
|
||||
# Setting Up The 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/quickly_deployable_plugins
|
||||
pip3 install -r requirements.txt
|
||||
```
|
||||
|
||||
# Implementing a Quickly Deployable Python (QDP) Plugin
|
||||
|
||||
QDP definitions consist of a set of decorated functions that define properties and behaviors of the plugin.
|
||||
### `@tensorrt.plugin.register`
|
||||
Returns shape and type characteristics of output tensors, and any attributes the plugin needs to function.
|
||||
|
||||
### `@tensorrt.plugin.impl`
|
||||
Performs the plugin computation. The decorated python function is executed 'just in time', as a python callback during runtime.
|
||||
|
||||
### (Optional) `@tensorrt.plugin.aot_impl`
|
||||
The decorated function directly returns an 'ahead of time' compiled kernel, along with information required to invoke it at runtime by TRT. This is in contrast with the above `@tensorrt.plugin.impl` in that, the returned kernel is baked into the built TRT engine. This is beneficial, when we need an engine that is fully independent of the python runtime - and hence, can be executed solely in a standard TensorRT C++ runtime (through `trtexec`, for example).
|
||||
|
||||
### (Optional) `@tensorrt.plugin.autotune`
|
||||
Defines the different data types and formats (tensor layouts) supported by the plugin's IO and any tactics supported by the plugin. Defining this function allows TensorRT to "tune" the plugin during the engine build to find the most performant type/format and tactic combination on the target system.
|
||||
|
||||
The specifics of these functions will become clear through the following mini-samples.
|
||||
|
||||
# A Simple Plugin: Elementwise-Add
|
||||
|
||||
This mini-sample contains an elementwise addition plugin, where the computation is being performed with an OpenAI Triton kernel. Let's first take a look at the `tensorrt.plugin.register` function.
|
||||
|
||||
```python
|
||||
import tensorrt.plugin as trtp
|
||||
|
||||
@trtp.register("sample::elemwise_add_plugin")
|
||||
def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> trtp.TensorDesc:
|
||||
return inp0.like()
|
||||
```
|
||||
|
||||
The argument "sample::elemwise_add_plugin" defines the namespace ("sample") and name ("elemwise_add_plugin") of the plugin. Input arguments to the decorated function (`plugin_desc`) annotated with `trt.plugin.TensorDesc` denote the input tensors; all others are interpreted as plugin attributes (see the [TRT API Reference](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/tensorrt.plugin/trt_plugin_register.html) for a full list of allowed attribute types). The output signature is a `trt.plugin.TensorDesc` describing the output. `inp0.like()` returns a tensor descriptor with identical shape and type characteristics to `inp0`.
|
||||
|
||||
The computation function, decorated with `trt.plugin.impl`, receives `trt.plugin.Tensor`s for each input and output. In contrast to `TensorDesc`s, a `Tensor` references an underlying data buffer, directly accessible through `Tensor.data_ptr`. When working with Torch and OpenAI Triton kernels, it is easier to use `torch.as_tensor()` to zero-copy construct a `torch.Tensor` corresponding to the `trt.plugin.Tensor`.
|
||||
|
||||
This sample also showcases the effect of omitting/defining a `trt.plugin.autotune` function, which must return a list of `trt.plugin.AutoTuneCombination`s. In this case, we define a single combination `AutoTuneCombination("FP32|FP16, FP32|FP16")`; this indicates that the input and output must be either both FP32 or both FP16. See the TRT API Reference for a detailed description of the grammar underlying `AutoTuneCombination`s.
|
||||
|
||||
## Running the sample
|
||||
|
||||
```bash
|
||||
python3 qdp_runner.py add [--autotune] [-v]
|
||||
```
|
||||
|
||||
`--autotune` simulates having defined a `trt.plugin.autotune` function. Enabling verbose logging (`-v`) is recommended to see the effect of autotuning. It can be observed that the `trt.plugin.impl` function is invoked several times during the engine build process when autotune is enabled. With autotuning turned off, `trt.plugin.impl` is invoked only once (when inference is run after building the engine).
|
||||
|
||||
```bash
|
||||
$ python3 qdp_runner.py add --autotune -v
|
||||
...
|
||||
Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
|
||||
Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
|
||||
Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
|
||||
Executing for inp0.dtype=DataType.FLOAT and output[0].dtype=DataType.FLOAT
|
||||
Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
|
||||
Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
|
||||
Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
|
||||
Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
|
||||
[I] Finished engine building in 1.073 seconds
|
||||
Executing for inp0.dtype=DataType.HALF and output[0].dtype=DataType.HALF
|
||||
```
|
||||
|
||||
# Implementing in-place custom ops with I/O aliasing
|
||||
|
||||
In-place computations can be accomplished with TRT plugins via aliased I/O. i.e. An input that needs to be modified in-place can be represented by an input-output pair, where the output is aliased to the input. For example, if in-place addition is needed (instead of the out-of-place addition of the above sample), that can be achieved as below:
|
||||
```python
|
||||
import tensorrt.plugin as trtp
|
||||
|
||||
@trtp.register("sample::elemwise_add_plugin_")
|
||||
def add_plugin_desc_(inp0: trtp.TensorDesc) -> trtp.TensorDesc:
|
||||
return inp0.aliased()
|
||||
```
|
||||
|
||||
Note the use of `trt.plugin.TensorDesc.aliased()` to produce an output `TensorDesc` that is aliased to `inp0`.
|
||||
|
||||
To appreciate the effect of aliasing better, this sample adds two in-place add plugins chained together.
|
||||
|
||||
## Running the sample
|
||||
|
||||
Enabling verbose logging (`-v`) is recommended to see the effect of autotuning, which is always enabled.
|
||||
|
||||
```bash
|
||||
python3 qdp_runner.py inplace_add [--autotune] [-v]
|
||||
```
|
||||
|
||||
# An Op with data-dependent output shapes: Non-zero
|
||||
|
||||
Non-zero is an operation where the indices of the non-zero elements of the input tensor is found -- it has data-dependent output shapes (DDS). As such, typical shape calculations cannot be done with input shapes.
|
||||
|
||||
To handle DDS, the extent of each data-dependent output dimension must be expressed in terms of a *_size tensor_*, which is a scalar that communicates to TRT an upper-bound and an autotune value for that dimension, in terms of the input shapes. The TRT engine build may be optimized for the autotune value, but the extent of that dimension may stretch up to the upper-bound at runtime.
|
||||
|
||||
In this sample, we consider a 2D input tensor `inp0`; the output will be an $N x 2$ tensor (a set of $N$ 2D indices), where $N$ is the number of non-zero indices. At maximum, all elements could be non-zero, and so the upper-bound could be expressed as `upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]`. Note that `trt.plugin.TensorDesc.shape_expr` returns symbolic shape expressions for that tensor. Arithmetic operations on shape expressions are supported through standard Python binary operators (see [TRT Python API reference](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/tensorrt.plugin/Shape/ShapeExpr.html) for full list of supported operations).
|
||||
|
||||
On average, we can expect half of the input to be filled with zero, so a size tensor can be constructed with that as the autotune value:
|
||||
```python
|
||||
st = trtp.size_tensor(opt = upper_bound // 2, upper_bound = upper_bound)
|
||||
```
|
||||
|
||||
Now we're ready to construct the output shape. `st.expr()` returns a shape expression for the size tensor, so a tensor descriptor for the output shape can be constructed as `trt.plugin.from_shape_expr((st.expr(), 2), dtype=trt.int32)`. TRT requires that any size tensors also be made outputs of the plugin. Putting things together, we arrive at the following:
|
||||
|
||||
```python
|
||||
import tensorrt.plugin as trtp
|
||||
|
||||
@trtp.register("sample::non_zero_plugin")
|
||||
def non_zero_plugin_reg(
|
||||
inp0: trtp.TensorDesc,
|
||||
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
|
||||
upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]
|
||||
st = trtp.size_tensor(upper_bound // 2, upper_bound)
|
||||
return trtp.from_shape_expr((st.expr(), 2), dtype=trt.int32), st
|
||||
```
|
||||
|
||||
## Running the sample
|
||||
|
||||
Enabling verbose logging (`-v`) is recommended to see the effect of autotuning, which is always enabled.
|
||||
|
||||
```bash
|
||||
python3 qdp_runner.py non_zero [-v]
|
||||
```
|
||||
|
||||
# Using multiple tactics and ONNX: Cirular padding
|
||||
|
||||
This sample contains a circular padding plugin, which is useful for ops like circular convolution. It is equivalent to PyTorch's [torch.nn.CircularPad2d](https://pytorch.org/docs/stable/generated/torch.nn.CircularPad2d.html#torch.nn.CircularPad2d).
|
||||
|
||||
Refer [this section about circular padding plugin](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/pluginGuide.html#example-circular-padding-plugin) in the python plugin guide for more info.
|
||||
|
||||
## ONNX model with a plugin
|
||||
|
||||
It is often useful to run an ONNX node with a custom op through a TRT plugin that you have written. To allow the TRT ONNX parser to correctly recognize your plugin as being mapped to an ONNX node, ensure that
|
||||
- The `op` property of the node is exactly the same as your plugin name.
|
||||
- The node contains a string attribute called "plugin_namespace" with the namespace of your plugin.
|
||||
|
||||
In this sample, we define a plugin with the ID "sample::circ_pad_plugin", so if using ONNX Graphsurgeon, the custom op node can be constructed as follows:
|
||||
|
||||
```python
|
||||
import onnx_graphsurgeon as gs
|
||||
|
||||
var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
|
||||
var_y = gs.Variable(name="y", dtype=np.float32)
|
||||
|
||||
circ_pad_node = gs.Node(
|
||||
name="circ_pad_plugin",
|
||||
op="circ_pad_plugin",
|
||||
inputs=[var_x],
|
||||
outputs=[var_y],
|
||||
attrs={"pads": pads, "plugin_namespace": "sample"},
|
||||
)
|
||||
```
|
||||
|
||||
## Multiple tactics
|
||||
|
||||
Sometimes, you may have multiple kernels (or backends) that can be used to perform the computation of the plugin -- these are typically called *_tactics_*. If it cannot be predetermined which of these tactics may perform the fastest, it is possible to let TRT time the plugin for each tactic and determine which one is fastest.
|
||||
|
||||
Communicating the availability of multiple tactics can simply be done through the `trt.plugin.autotune` function.
|
||||
```python
|
||||
import tensorrt.plugin as trtp
|
||||
from enum import IntEnum
|
||||
|
||||
class Tactic(IntEnum):
|
||||
TORCH = 1
|
||||
TRITON = 2
|
||||
|
||||
@trt.plugin.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32], outputs: Tuple[trtp.TensorDesc]) -> List[trtp.AutoTuneCombination]:
|
||||
c = trtp.AutoTuneCombination()
|
||||
c.pos([0, 1], "FP32|FP16")
|
||||
c.tactics([int(Tactic.TORCH), int(Tactic.TRITON)])
|
||||
return [c]
|
||||
```
|
||||
|
||||
Note that we're using another way of constructing a `trt.plugin.AutoTuneCombination` here -- namely, through `pos(...)` to populate the type/format information and `tactics(...)` to specify the tactics. In this sample, we use an OpenAI Triton kernel and `torch.nn.functional.pad` as two methods to compute the circular padding.
|
||||
|
||||
Refer [this section](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/pluginGuide.html#example-plugins-with-multiple-backends-using-custom-tactics) in the Python plugin guide for more info.
|
||||
|
||||
## Loading and running a TRT engine containing a plugin
|
||||
|
||||
If you have a TRT engine built with a plugin, executing that engine only requires the plugin definitions for `trt.plugin.register` and `trt.plugin.impl` to be available in the module where the engine is being deserialized (note: the `trt.plugin.autotune` definition is not required to be present).
|
||||
|
||||
To simulate the loading of an engine, first run this sample with the `--save_engine` flag, followed by `--artifacts_dir [dir]` with a directory in which you wish the engine to be saved. Then run the sample again with `--load engine` and `--artifacts_dir` set to the same directory.
|
||||
|
||||
## Running the sample
|
||||
|
||||
```bash
|
||||
python3 qdp_runner.py circ_pad [--multi_tactic] [--save_engine] [--load_engine] --mode {onnx,inetdef} [--artifacts_dir ARTIFACTS_DIR] [-v]
|
||||
|
||||
options:
|
||||
--multi_tactic Enable multiple tactics.
|
||||
--save_engine Save engine to the artifacts_dir.
|
||||
--load_engine Load engine from the artifacts_dir. Ignores all other options.
|
||||
--artifacts_dir ARTIFACTS_DIR
|
||||
Whether to store (or retrieve) artifacts.
|
||||
--mode {onnx,inetdef} Whether to use ONNX parser or INetworkDefinition APIs to construct the network.
|
||||
-v, --verbose Enable verbose log output.
|
||||
```
|
||||
|
||||
# Providing an Ahead-of-Time (AOT) implementation for Circular padding
|
||||
|
||||
Let's extend the [above sample](#using-multiple-tactics-and-onnx-cirular-padding) by providing an AOT implementation for the same circular padding operation.
|
||||
Instead of specifying the OpenAI Triton Kernel callback to TRT through `@trt.plugin.impl`, we can directly
|
||||
compile the kernel ahead of time, and provide that to TRT under `@trt.plugin.aot_impl`.
|
||||
|
||||
Refer [this section](https://docs.nvidia.com/deeplearning/tensorrt/latest/_static/python-api/pluginGuide.html#providing-an-ahead-of-time-aot-implementation) in the Python plugin guide for more info.
|
||||
|
||||
## ONNX model with an AOT plugin
|
||||
|
||||
The same rules apply as mentioned in the above [ONNX model with a plugin](#onnx-model-with-a-plugin) section.
|
||||
In addition to that, if the plugin has an AOT implementation that we'd like to use, we can modify the ONNX node to communicate that to the TRT ONNX parser.
|
||||
This should be done by adding a bool attribute called "aot" to the ONNX node, and setting it to True.
|
||||
Note that, this is on top of making sure that the ONNX node has the appropriate `op` property and "plugin_namespace" attribute as mentioned [previously](#onnx-model-with-a-plugin).
|
||||
|
||||
Therefore, using ONNX Graphsurgeon, the custom op node that uses the AOT implementation of "sample::circ_pad_plugin" can be constructed similarly:
|
||||
|
||||
```python
|
||||
import onnx_graphsurgeon as gs
|
||||
|
||||
var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
|
||||
var_y = gs.Variable(name="y", dtype=np.float32)
|
||||
|
||||
circ_pad_aot_node = gs.Node(
|
||||
name="circ_pad_plugin_aot",
|
||||
op="circ_pad_plugin",
|
||||
inputs=[var_x],
|
||||
outputs=[var_y],
|
||||
attrs={"pads": pads, "plugin_namespace": "sample", "aot": True},
|
||||
)
|
||||
```
|
||||
|
||||
## Loading and running a TRT engine containing an AOT plugin
|
||||
|
||||
If you have a TRT engine built with an AOT plugin, the plugin computation is already part of the engine. Therefore, it does not require any Python modules or definitions to be present at runtime. This means that the engine can be executed on the standard TRT runtime, as part of any tool that is capable of deserializing and running the engine (like [trtexec](../../trtexec/README.md).
|
||||
|
||||
To simulate the loading of an engine, first run this sample with the `--save_engine` flag, followed by `--artifacts_dir [dir]` with a directory in which you wish the engine to be saved. Then run the sample again with `--load engine` and `--artifacts_dir` set to the same directory.
|
||||
|
||||
## Running the sample
|
||||
|
||||
```bash
|
||||
python3 qdp_runner.py circ_pad [--multi_tactic] [--aot] [--save_engine] [--load_engine] --mode {onnx,inetdef} [--artifacts_dir ARTIFACTS_DIR] [-v]
|
||||
|
||||
options:
|
||||
--multi_tactic Enable multiple tactics. Combined with --aot, the advertised tactics are AOT-compiled.
|
||||
--save_engine Save engine to the artifacts_dir.
|
||||
--load_engine Load engine from the artifacts_dir. Ignores all other options.
|
||||
--artifacts_dir ARTIFACTS_DIR
|
||||
Whether to store (or retrieve) artifacts.
|
||||
--mode {onnx,inetdef} Whether to use ONNX parser or INetworkDefinition APIs to construct the network.
|
||||
--aot Use the AOT implementation of the plugin.
|
||||
-v, --verbose Enable verbose log output.
|
||||
```
|
||||
|
||||
## Combining AOT with multiple tactics
|
||||
|
||||
The AOT path can advertise more than one tactic. Each tactic must be a precompiled GPU kernel, so a Torch-dispatched tactic (like the `Tactic.TORCH` variant used in the JIT multi-tactic example above) cannot participate. In this sample we instead expose two variants of the same Triton kernel that differ only in `BLOCK_SIZE`. TRT times both PTX blobs at build time and bakes the winner into the engine.
|
||||
|
||||
The autotune declaration lists both tactic IDs:
|
||||
|
||||
```python
|
||||
class Tactic(IntEnum):
|
||||
BLOCK_256 = 1
|
||||
BLOCK_1024 = 2
|
||||
|
||||
@trt.plugin.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(inp0, outputs):
|
||||
c = trtp.AutoTuneCombination()
|
||||
c.pos([0, 1], "FP32|FP16")
|
||||
c.tactics([int(Tactic.BLOCK_256), int(Tactic.BLOCK_1024)])
|
||||
return [c]
|
||||
```
|
||||
|
||||
The AOT impl switches on the `tactic` argument, compiles the Triton kernel with the matching `BLOCK_SIZE` constexpr, and returns the corresponding PTX plus a `KernelLaunchParams` whose `grid_x` reflects that block size:
|
||||
|
||||
```python
|
||||
@trt.plugin.aot_impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_aot_impl(inp0, pads, outputs, tactic):
|
||||
block_size = {1: 256, 2: 1024}[tactic]
|
||||
src = triton.compiler.ASTSource(fn=circ_pad_kernel, signature=..., constexprs={"BLOCK_SIZE": block_size})
|
||||
compiled = triton.compile(src)
|
||||
launch_params = trtp.KernelLaunchParams()
|
||||
launch_params.grid_x = trtp.cdiv(outputs[0].shape_expr.numel(), block_size)
|
||||
launch_params.block_x = compiled.metadata.num_warps * 32
|
||||
launch_params.shared_mem = compiled.metadata.shared
|
||||
extra_args = trtp.SymIntExprs.from_tuple(...) # symbolic int32 kernel args. see source for more details.
|
||||
return compiled.metadata.name.encode(), compiled.asm["ptx"].encode(), launch_params, extra_args
|
||||
```
|
||||
|
||||
Run it with both flags:
|
||||
|
||||
```bash
|
||||
python3 qdp_runner.py circ_pad --mode inetdef --multi_tactic --aot -v
|
||||
```
|
||||
|
||||
Verbose logs show TRT timing both tactics during engine build, then a single winner is serialized into the engine.
|
||||
|
||||
# Additional Resources
|
||||
|
||||
**Python Plugin Guide**
|
||||
- [pluginGuide.md](../../../documentation/python/pluginGuide.md)
|
||||
|
||||
**`tensorrt.plugin` API reference**
|
||||
- [`tensorrt.plugin` module API reference](https://docs.nvidia.com/deeplearning/tensorrt/api/python_api/infer/tensorrt.plugin/index.html)
|
||||
|
||||
**Developer Guide**
|
||||
- [Extending TensorRT with Custom Layers](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#extending)
|
||||
|
||||
# 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
|
||||
- May 2026: Added multi-tactic AOT subsection for circular padding.
|
||||
- October 2025: Migrate to strongly typed APIs.
|
||||
- August 2025: Removed support for Python versions < 3.10.
|
||||
- December 2024: Added section on AOT Plugins, added contents section
|
||||
- October 2024: Initial release of this sample
|
||||
|
||||
# Known issues
|
||||
|
||||
There are no known issues in this sample
|
||||
@@ -0,0 +1,78 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
@triton.jit
|
||||
def add_kernel(x_ptr, y_ptr, n_elements, BLOCK_SIZE: tl.constexpr):
|
||||
pid = tl.program_id(0)
|
||||
offsets = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
mask = offsets < n_elements
|
||||
x = tl.load(x_ptr + offsets, mask=mask)
|
||||
tl.store(y_ptr + offsets, x + 1, mask=mask)
|
||||
|
||||
|
||||
@triton.jit
|
||||
def circ_pad_kernel(
|
||||
# input tensor
|
||||
X,
|
||||
# extra scalar args in between input and output tensors
|
||||
# for kernel signature to be compatible with AOT plugin impl
|
||||
all_pads_0,
|
||||
all_pads_2,
|
||||
all_pads_4,
|
||||
all_pads_6,
|
||||
orig_dims_0,
|
||||
orig_dims_1,
|
||||
orig_dims_2,
|
||||
orig_dims_3,
|
||||
Y_shape_1,
|
||||
Y_shape_2,
|
||||
Y_shape_3,
|
||||
X_len,
|
||||
Y_len,
|
||||
# output tensor
|
||||
Y,
|
||||
BLOCK_SIZE: tl.constexpr,
|
||||
):
|
||||
pid = tl.program_id(0)
|
||||
i = pid * BLOCK_SIZE + tl.arange(0, BLOCK_SIZE)
|
||||
|
||||
mask_y = i < Y_len
|
||||
|
||||
i3 = i % Y_shape_3
|
||||
i2 = (i // Y_shape_3) % Y_shape_2
|
||||
i1 = (i // Y_shape_3 // Y_shape_2) % Y_shape_1
|
||||
i0 = i // Y_shape_3 // Y_shape_2 // Y_shape_1
|
||||
|
||||
j0 = (i0 - all_pads_0 + orig_dims_0) % orig_dims_0
|
||||
j1 = (i1 - all_pads_2 + orig_dims_1) % orig_dims_1
|
||||
j2 = (i2 - all_pads_4 + orig_dims_2) % orig_dims_2
|
||||
j3 = (i3 - all_pads_6 + orig_dims_3) % orig_dims_3
|
||||
|
||||
load_idx = (
|
||||
orig_dims_3 * orig_dims_2 * orig_dims_1 * j0
|
||||
+ orig_dims_3 * orig_dims_2 * j1
|
||||
+ orig_dims_3 * j2
|
||||
+ j3
|
||||
)
|
||||
mask_x = load_idx < X_len
|
||||
|
||||
x = tl.load(X + load_idx, mask=mask_x)
|
||||
|
||||
tl.store(Y + i, x, mask=mask_y)
|
||||
@@ -0,0 +1,386 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import tensorrt as trt
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from typing import Tuple, List, Union
|
||||
|
||||
import tensorrt.plugin as trtp
|
||||
import numpy.typing as npt
|
||||
|
||||
import logging
|
||||
|
||||
logging.basicConfig(level=logging.INFO)
|
||||
logging.getLogger("QuicklyDeployablePlugins").setLevel(logging.INFO)
|
||||
|
||||
########## Elemwise-add plugin definition ##########
|
||||
|
||||
|
||||
@trtp.register("sample::elemwise_add_plugin")
|
||||
def add_plugin_desc(inp0: trtp.TensorDesc, block_size: int) -> trtp.TensorDesc:
|
||||
return inp0.like()
|
||||
|
||||
|
||||
# Helper to simulate defining/omitting an autotune definition for the plugin
|
||||
def register_autotune():
|
||||
# Type annotations can be omitted for autotune and impl definitions, but will be checked for consistency if added
|
||||
@trtp.autotune("sample::elemwise_add_plugin")
|
||||
def add_plugin_autotune(
|
||||
inp0: trtp.TensorDesc, outputs: Tuple[trtp.TensorDesc]
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16")]
|
||||
|
||||
|
||||
@trtp.impl("sample::elemwise_add_plugin")
|
||||
def add_plugin_impl(
|
||||
inp0: trtp.Tensor, block_size: int, outputs: Tuple[trtp.Tensor], stream: int
|
||||
) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
n = inp0.numel()
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp0_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_t = torch.as_tensor(outputs[0], device="cuda")
|
||||
|
||||
import triton
|
||||
from oait_kernels import add_kernel
|
||||
|
||||
add_kernel[(triton.cdiv(n, block_size),)](inp0_t, out_t, n, BLOCK_SIZE=block_size)
|
||||
|
||||
|
||||
########## In-place elemwise-add plugin definition ##########
|
||||
|
||||
|
||||
@trtp.register("sample::elemwise_add_plugin_")
|
||||
def add_plugin_desc_(inp0: trtp.TensorDesc, delta: int) -> trtp.TensorDesc:
|
||||
return inp0.aliased()
|
||||
|
||||
|
||||
@trtp.autotune("sample::elemwise_add_plugin_")
|
||||
def add_plugin_autotune_(inp0, outputs) -> List[trtp.AutoTuneCombination]:
|
||||
return [
|
||||
trtp.AutoTuneCombination("FP32, FP32", "LINEAR*HWC"),
|
||||
trtp.AutoTuneCombination("FP32|FP16, FP32|FP16", "LINEAR"),
|
||||
]
|
||||
|
||||
|
||||
@trtp.impl("sample::elemwise_add_plugin_")
|
||||
def add_plugin_impl_(inp0, delta: int, outputs, stream) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp0_t = torch.as_tensor(inp0, device="cuda")
|
||||
inp0_t.add_(delta)
|
||||
|
||||
|
||||
########## Non-zero plugin (DDS) ##########
|
||||
|
||||
|
||||
@trtp.register("sample::non_zero_plugin")
|
||||
def non_zero_plugin_reg(
|
||||
inp0: trtp.TensorDesc,
|
||||
) -> Tuple[trtp.TensorDesc, trtp.TensorDesc]:
|
||||
upper_bound = inp0.shape_expr[0] * inp0.shape_expr[1]
|
||||
st = trtp.size_tensor(upper_bound // 2, upper_bound)
|
||||
st.dtype = trt.int64
|
||||
return trtp.from_shape_expr((st.expr(), 2), dtype=trt.int32), st
|
||||
|
||||
|
||||
@trtp.autotune("sample::non_zero_plugin")
|
||||
def non_zero_plugin_autotune(inp0, outputs) -> List[trtp.AutoTuneCombination]:
|
||||
return [trtp.AutoTuneCombination("FP32|FP16, INT32, INT64")]
|
||||
|
||||
|
||||
@trtp.impl("sample::non_zero_plugin")
|
||||
def non_zero_plugin_impl(inp0, outputs, stream) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp0_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_1 = torch.as_tensor(outputs[1], device="cuda").reshape((-1,))
|
||||
|
||||
out = torch.nonzero(inp0_t)
|
||||
|
||||
out0 = torch.as_tensor(outputs[0].aliased(out.shape), device="cuda")
|
||||
out0.copy_(out)
|
||||
out_1.copy_(torch.Tensor([out.shape[0]]))
|
||||
|
||||
|
||||
########## Circular padding plugin ########
|
||||
|
||||
|
||||
@trtp.register("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_desc(
|
||||
inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32]
|
||||
) -> trtp.TensorDesc:
|
||||
ndim = inp0.ndim
|
||||
out_desc = inp0.like()
|
||||
|
||||
for i in range(np.size(pads) // 2):
|
||||
out_desc.shape_expr[ndim - i - 1] += int(pads[i * 2] + pads[i * 2 + 1])
|
||||
|
||||
return out_desc
|
||||
|
||||
|
||||
# Helper to define a multi-tactic implementation of the plugin
|
||||
def enable_multi_tactic_circ_pad():
|
||||
|
||||
from enum import IntEnum
|
||||
|
||||
class Tactic(IntEnum):
|
||||
TORCH = 1
|
||||
TRITON = 2
|
||||
|
||||
@trtp.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(
|
||||
inp0: trtp.TensorDesc,
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
c = trtp.AutoTuneCombination()
|
||||
c.pos([0, 1], "FP32|FP16")
|
||||
c.tactics([int(Tactic.TORCH), int(Tactic.TRITON)])
|
||||
return [c]
|
||||
|
||||
@trtp.impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_impl(
|
||||
inp0: trtp.Tensor,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.Tensor],
|
||||
stream: int,
|
||||
tactic: int,
|
||||
) -> None:
|
||||
|
||||
log = logging.getLogger("QuicklyDeployablePlugins")
|
||||
log.debug(
|
||||
f"Executing for inp0: dtype={inp0.dtype},format={inp0.format} and output[0]: dtype={outputs[0].dtype},format={outputs[0].format}"
|
||||
)
|
||||
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_t = torch.as_tensor(outputs[0], device="cuda")
|
||||
|
||||
if tactic == Tactic.TORCH:
|
||||
out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
|
||||
out_t.copy_(out)
|
||||
elif tactic == Tactic.TRITON:
|
||||
N = inp0.ndim
|
||||
all_pads = np.zeros((N * 2,), dtype=np.int32)
|
||||
out_dims = trtp.Shape(tuple(inp0.shape))
|
||||
|
||||
for i in range(np.size(pads) // 2):
|
||||
out_dims[N - i - 1] += pads[i * 2] + pads[i * 2 + 1]
|
||||
all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
|
||||
all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
|
||||
|
||||
all_pads = all_pads.tolist()
|
||||
|
||||
block_size = 256
|
||||
num_blocks = tuple(
|
||||
[int((np.prod(out_dims) + block_size - 1) // block_size)]
|
||||
)
|
||||
|
||||
from oait_kernels import circ_pad
|
||||
|
||||
circ_pad[num_blocks](
|
||||
inp_t,
|
||||
all_pads[0],
|
||||
all_pads[2],
|
||||
all_pads[4],
|
||||
all_pads[6],
|
||||
inp0.shape[0],
|
||||
inp0.shape[1],
|
||||
inp0.shape[2],
|
||||
inp0.shape[3],
|
||||
int(out_dims[1]),
|
||||
int(out_dims[2]),
|
||||
int(out_dims[3]),
|
||||
inp0.numel(),
|
||||
out_dims.numel(),
|
||||
out_t,
|
||||
BLOCK_SIZE=block_size,
|
||||
)
|
||||
|
||||
|
||||
# Shared AOT compilation body for the circ_pad plugin: build SymInt args,
|
||||
# compile the Triton kernel with the given BLOCK_SIZE, and pack launch params.
|
||||
def _compile_circ_pad_aot(
|
||||
inp0: trtp.TensorDesc,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
block_size: int,
|
||||
) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
|
||||
|
||||
N = inp0.ndim
|
||||
all_pads = np.zeros((N * 2,), dtype=np.int32)
|
||||
inp_dims = inp0.shape_expr
|
||||
out_dims = outputs[0].shape_expr
|
||||
|
||||
for i in range(np.size(pads) // 2):
|
||||
all_pads[N * 2 - 2 * i - 2] = pads[i * 2]
|
||||
all_pads[N * 2 - 2 * i - 1] = pads[i * 2 + 1]
|
||||
|
||||
all_pads = all_pads.tolist()
|
||||
|
||||
# Representing all int32-scalar-kernel-inputs as symbolic expressions.
|
||||
# These inputs are either constants or derivatives of input/output shapes (that may be dynamic).
|
||||
# The symbolic expressions are resolved after the full shape context becomes available at runtime.
|
||||
extra_args = trtp.SymIntExprs.from_tuple(
|
||||
[
|
||||
trtp.SymInt32(e)
|
||||
for e in [
|
||||
all_pads[0],
|
||||
all_pads[2],
|
||||
all_pads[4],
|
||||
all_pads[6],
|
||||
inp_dims[0],
|
||||
inp_dims[1],
|
||||
inp_dims[2],
|
||||
inp_dims[3],
|
||||
out_dims[1],
|
||||
out_dims[2],
|
||||
out_dims[3],
|
||||
inp_dims.numel(),
|
||||
out_dims.numel(),
|
||||
]
|
||||
]
|
||||
)
|
||||
|
||||
type_str = "fp32" if inp0.dtype == trt.float32 else "fp16"
|
||||
|
||||
from oait_kernels import circ_pad_kernel
|
||||
import triton
|
||||
|
||||
src = triton.compiler.ASTSource(
|
||||
fn=circ_pad_kernel,
|
||||
signature={
|
||||
"X": f"*{type_str}",
|
||||
"all_pads_0": "i32",
|
||||
"all_pads_2": "i32",
|
||||
"all_pads_4": "i32",
|
||||
"all_pads_6": "i32",
|
||||
"orig_dims_0": "i32",
|
||||
"orig_dims_1": "i32",
|
||||
"orig_dims_2": "i32",
|
||||
"orig_dims_3": "i32",
|
||||
"Y_shape_1": "i32",
|
||||
"Y_shape_2": "i32",
|
||||
"Y_shape_3": "i32",
|
||||
"X_len": "i32",
|
||||
"Y_len": "i32",
|
||||
"Y": f"*{type_str}",
|
||||
},
|
||||
constexprs={"BLOCK_SIZE": block_size},
|
||||
)
|
||||
|
||||
compiled_kernel = triton.compile(src)
|
||||
launch_params = trtp.KernelLaunchParams()
|
||||
launch_params.grid_x = trtp.cdiv(out_dims.numel(), block_size)
|
||||
launch_params.block_x = compiled_kernel.metadata.num_warps * 32
|
||||
launch_params.shared_mem = compiled_kernel.metadata.shared
|
||||
|
||||
return (
|
||||
compiled_kernel.metadata.name.encode(),
|
||||
compiled_kernel.asm["ptx"].encode(),
|
||||
launch_params,
|
||||
extra_args,
|
||||
)
|
||||
|
||||
|
||||
# Helper to define a single tactic implementation of the plugin
|
||||
def enable_single_tactic_circ_pad():
|
||||
@trtp.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(
|
||||
inp0: trtp.TensorDesc,
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
|
||||
return [trtp.AutoTuneCombination("FP32|FP16, FP32|FP16")]
|
||||
|
||||
@trtp.impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_impl(
|
||||
inp0: trtp.Tensor,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.Tensor],
|
||||
stream: int,
|
||||
) -> None:
|
||||
with torch.cuda.stream(torch.cuda.ExternalStream(stream)):
|
||||
inp_t = torch.as_tensor(inp0, device="cuda")
|
||||
out_t = torch.as_tensor(outputs[0], device="cuda")
|
||||
|
||||
out = torch.nn.functional.pad(inp_t, pads.tolist(), mode="circular")
|
||||
out_t.copy_(out)
|
||||
|
||||
@trtp.aot_impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_aot_impl(
|
||||
inp0: trtp.TensorDesc, pads: npt.NDArray[np.int32], outputs: Tuple[trtp.TensorDesc], tactic: int
|
||||
) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
|
||||
return _compile_circ_pad_aot(inp0, pads, outputs, block_size=256)
|
||||
|
||||
|
||||
# Helper to define a multi-tactic AOT implementation of the plugin.
|
||||
# Each tactic precompiles the same Triton kernel with a different BLOCK_SIZE,
|
||||
# so TRT times two PTX variants at build time and bakes the winner into the engine.
|
||||
def enable_multi_tactic_aot_circ_pad():
|
||||
|
||||
from enum import IntEnum
|
||||
|
||||
class Tactic(IntEnum):
|
||||
BLOCK_256 = 1
|
||||
BLOCK_1024 = 2
|
||||
|
||||
block_size_by_tactic = {
|
||||
int(Tactic.BLOCK_256): 256,
|
||||
int(Tactic.BLOCK_1024): 1024,
|
||||
}
|
||||
|
||||
@trtp.autotune("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_autotune(
|
||||
inp0: trtp.TensorDesc,
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
) -> List[trtp.AutoTuneCombination]:
|
||||
c = trtp.AutoTuneCombination()
|
||||
c.pos([0, 1], "FP32|FP16")
|
||||
c.tactics([int(Tactic.BLOCK_256), int(Tactic.BLOCK_1024)])
|
||||
return [c]
|
||||
|
||||
@trtp.aot_impl("sample::circ_pad_plugin")
|
||||
def circ_pad_plugin_aot_impl(
|
||||
inp0: trtp.TensorDesc,
|
||||
pads: npt.NDArray[np.int32],
|
||||
outputs: Tuple[trtp.TensorDesc],
|
||||
tactic: int,
|
||||
) -> Tuple[Union[str, bytes], Union[str, bytes], trtp.KernelLaunchParams, trtp.SymExprs]:
|
||||
block_size = block_size_by_tactic[tactic]
|
||||
logging.getLogger("QuicklyDeployablePlugins").debug(
|
||||
f"aot_impl invoked: tactic={tactic} -> BLOCK_SIZE={block_size}"
|
||||
)
|
||||
return _compile_circ_pad_aot(inp0, pads, outputs, block_size=block_size)
|
||||
@@ -0,0 +1,363 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2024-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
import tensorrt as trt
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
TrtRunner,
|
||||
create_network,
|
||||
engine_from_network,
|
||||
network_from_onnx_path,
|
||||
bytes_from_engine,
|
||||
engine_from_bytes,
|
||||
)
|
||||
|
||||
from polygraphy.backend.common import bytes_from_path
|
||||
from polygraphy import cuda
|
||||
|
||||
import onnx_graphsurgeon as gs
|
||||
import onnx
|
||||
import os
|
||||
import argparse
|
||||
|
||||
import tensorrt.plugin as trtp
|
||||
|
||||
import qdp_defs
|
||||
import logging
|
||||
|
||||
def run_add(enable_autotune=False):
|
||||
|
||||
if enable_autotune:
|
||||
qdp_defs.register_autotune()
|
||||
|
||||
BLOCK_SIZE = 256
|
||||
|
||||
builder, network = create_network(strongly_typed=True)
|
||||
x = torch.randint(10, (10, 3, 32, 32), dtype=torch.float32, device="cuda")
|
||||
|
||||
# Populate network
|
||||
i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
|
||||
|
||||
out = network.add_plugin(
|
||||
trtp.op.sample.elemwise_add_plugin(i_x, block_size=BLOCK_SIZE)
|
||||
)
|
||||
out.get_output(0).name = "y"
|
||||
network.mark_output(tensor=out.get_output(0))
|
||||
|
||||
builder.create_builder_config()
|
||||
|
||||
engine = engine_from_network(
|
||||
(builder, network),
|
||||
CreateConfig(),
|
||||
)
|
||||
|
||||
with TrtRunner(engine, "trt_runner") as runner:
|
||||
outputs = runner.infer(
|
||||
{
|
||||
"x": x,
|
||||
},
|
||||
copy_outputs_to_host=False,
|
||||
)
|
||||
|
||||
if torch.allclose(x + 1, outputs["y"]):
|
||||
print("Inference result is correct!")
|
||||
else:
|
||||
print("Inference result is incorrect!")
|
||||
|
||||
|
||||
def run_inplace_add():
|
||||
builder, network = create_network(strongly_typed=True)
|
||||
x = torch.ones((10, 3, 32, 32), dtype=torch.float32, device="cuda")
|
||||
|
||||
x_clone = x.clone()
|
||||
|
||||
i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
|
||||
|
||||
# Amounts to elementwise-add in the first and second plugins
|
||||
deltas = (2, 4)
|
||||
|
||||
out0 = network.add_plugin(trtp.op.sample.elemwise_add_plugin_(i_x, delta=deltas[0]))
|
||||
out1 = network.add_plugin(
|
||||
trtp.op.sample.elemwise_add_plugin_(out0.get_output(0), delta=deltas[1])
|
||||
)
|
||||
out1.get_output(0).name = "y"
|
||||
network.mark_output(tensor=out1.get_output(0))
|
||||
|
||||
builder.create_builder_config()
|
||||
|
||||
# Enable preview feature for aliasing plugin I/O
|
||||
config = CreateConfig(
|
||||
preview_features=[trt.PreviewFeature.ALIASED_PLUGIN_IO_10_03]
|
||||
)
|
||||
|
||||
engine = engine_from_network(
|
||||
(builder, network),
|
||||
config,
|
||||
)
|
||||
|
||||
context = engine.create_execution_context()
|
||||
|
||||
stream = cuda.Stream()
|
||||
|
||||
context.set_tensor_address("x", x.data_ptr())
|
||||
context.set_tensor_address("y", x.data_ptr())
|
||||
context.execute_async_v3(stream.ptr)
|
||||
stream.synchronize()
|
||||
|
||||
if torch.allclose(x, x_clone + sum(deltas), atol=1e-2):
|
||||
print("Inference result is correct!")
|
||||
else:
|
||||
print("Inference result is incorrect!")
|
||||
print(x[0][0][0][:10])
|
||||
print(x_clone[0][0][0][:10])
|
||||
|
||||
|
||||
def run_non_zero():
|
||||
builder, network = create_network(strongly_typed=True)
|
||||
inp_shape = (128, 128)
|
||||
|
||||
X = np.random.normal(size=inp_shape).astype(trt.nptype(trt.DataType.FLOAT))
|
||||
|
||||
# Zero out some random indices
|
||||
indices = np.random.choice(
|
||||
np.prod(inp_shape),
|
||||
replace=False,
|
||||
size=np.random.randint(0, np.prod(inp_shape) + 1),
|
||||
)
|
||||
X[np.unravel_index(indices, inp_shape)] = 0
|
||||
|
||||
# Populate network
|
||||
i_x = network.add_input(name="X", dtype=trt.DataType.FLOAT, shape=inp_shape)
|
||||
|
||||
out = network.add_plugin(trtp.op.sample.non_zero_plugin(i_x))
|
||||
out.get_output(0).name = "Y"
|
||||
network.mark_output(tensor=out.get_output(0))
|
||||
|
||||
builder.create_builder_config()
|
||||
|
||||
engine = engine_from_network(
|
||||
(builder, network),
|
||||
config=CreateConfig(),
|
||||
)
|
||||
|
||||
Y_ref = np.transpose(np.nonzero(X))
|
||||
|
||||
with TrtRunner(engine, "trt_runner") as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
Y = Y[np.lexsort(np.fliplr(Y).T)]
|
||||
|
||||
if np.allclose(Y, Y_ref, atol=1e-3):
|
||||
print("Inference result is correct!")
|
||||
else:
|
||||
print("Inference result is incorrect!")
|
||||
|
||||
|
||||
def check_artifacts_dir_exists(artifacts_dir):
|
||||
if not os.path.exists(artifacts_dir):
|
||||
raise ValueError(f"artifacts_dir '{artifacts_dir}' does not exist")
|
||||
|
||||
|
||||
def run_circ_pad(
|
||||
enable_multi_tactic=False, mode="onnx", artifacts_dir=None, save_or_load_engine=None, aot=False
|
||||
):
|
||||
|
||||
if enable_multi_tactic and aot:
|
||||
qdp_defs.enable_multi_tactic_aot_circ_pad()
|
||||
elif enable_multi_tactic:
|
||||
qdp_defs.enable_multi_tactic_circ_pad()
|
||||
else:
|
||||
qdp_defs.enable_single_tactic_circ_pad()
|
||||
|
||||
inp_shape = (10, 3, 32, 32)
|
||||
x = np.random.normal(size=inp_shape).astype(trt.nptype(trt.DataType.FLOAT))
|
||||
|
||||
pads = np.array((1, 1, 1, 1), dtype=np.int32)
|
||||
|
||||
if save_or_load_engine is not None and save_or_load_engine is False:
|
||||
check_artifacts_dir_exists(artifacts_dir)
|
||||
engine_path = os.path.join(artifacts_dir, "circ_pad.engine")
|
||||
engine = engine_from_bytes(bytes_from_path(engine_path))
|
||||
else:
|
||||
if mode == "inetdef":
|
||||
builder, network = create_network(strongly_typed=True)
|
||||
i_x = network.add_input(name="x", dtype=trt.DataType.FLOAT, shape=x.shape)
|
||||
out = network.add_plugin(trtp.op.sample.circ_pad_plugin(i_x, pads=pads), aot = aot)
|
||||
out.get_output(0).name = "y"
|
||||
network.mark_output(tensor=out.get_output(0))
|
||||
|
||||
engine = engine_from_network(
|
||||
(builder, network),
|
||||
CreateConfig(),
|
||||
)
|
||||
elif mode == "onnx":
|
||||
if artifacts_dir is None:
|
||||
raise ValueError("'artifacts_dir' must be specified in onnx mode")
|
||||
|
||||
check_artifacts_dir_exists(artifacts_dir)
|
||||
|
||||
onnx_path = os.path.join(artifacts_dir, "circ_pad.onnx")
|
||||
var_x = gs.Variable(name="x", shape=inp_shape, dtype=np.float32)
|
||||
var_y = gs.Variable(name="y", dtype=np.float32)
|
||||
circ_pad_node = gs.Node(
|
||||
name="circ_pad_plugin 0",
|
||||
op="circ_pad_plugin",
|
||||
inputs=[var_x],
|
||||
outputs=[var_y],
|
||||
attrs={"pads": pads, "plugin_namespace": "sample", "aot": aot},
|
||||
)
|
||||
graph = gs.Graph(
|
||||
nodes=[circ_pad_node], inputs=[var_x], outputs=[var_y], opset=16
|
||||
)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
engine = engine_from_network(
|
||||
network_from_onnx_path(onnx_path, strongly_typed=True), CreateConfig()
|
||||
)
|
||||
else:
|
||||
raise ValueError(f"Unknown mode {mode}")
|
||||
|
||||
if save_or_load_engine is not None and save_or_load_engine is True:
|
||||
check_artifacts_dir_exists(artifacts_dir)
|
||||
engine_path = os.path.join(artifacts_dir, "circ_pad.engine")
|
||||
with open(engine_path, "wb") as f:
|
||||
f.write(bytes_from_engine(engine))
|
||||
|
||||
Y_ref = np.pad(x, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
|
||||
with TrtRunner(engine, "trt_runner") as runner:
|
||||
outputs = runner.infer({"x": x})
|
||||
Y = outputs["y"]
|
||||
|
||||
if np.allclose(Y, Y_ref, atol=1e-2):
|
||||
print("Inference result is correct!")
|
||||
else:
|
||||
print("Inference result is incorrect!")
|
||||
|
||||
|
||||
def setup_add_sample(subparsers):
|
||||
subparser = subparsers.add_parser("add", help="'add' sample help")
|
||||
subparser.add_argument("--autotune", action="store_true", help="Enable autotuning")
|
||||
subparser.add_argument("--aot", action="store_true", help="Use the AOT implementation of the plugin")
|
||||
subparser.add_argument(
|
||||
"-v", "--verbose", action="store_true", help="Enable more verbose log output"
|
||||
)
|
||||
|
||||
|
||||
def setup_inplace_add_sample(subparsers):
|
||||
subparser = subparsers.add_parser("inplace_add", help="inplace_add sample help")
|
||||
subparser.add_argument(
|
||||
"-v", "--verbose", action="store_true", help="Enable more verbose log output"
|
||||
)
|
||||
|
||||
|
||||
def setup_non_zero_sample(subparsers):
|
||||
subparser = subparsers.add_parser("non_zero", help="non_zero sample help")
|
||||
subparser.add_argument(
|
||||
"-v", "--verbose", action="store_true", help="Enable more verbose log output"
|
||||
)
|
||||
|
||||
|
||||
def setup_circ_pad_sample(subparsers):
|
||||
subparser = subparsers.add_parser("circ_pad", help="circ_pad sample help.")
|
||||
subparser.add_argument(
|
||||
"--multi_tactic", action="store_true", help="Enable multiple tactics."
|
||||
)
|
||||
subparser.add_argument(
|
||||
"--save_engine", action="store_true", help="Save engine to the artifacts_dir."
|
||||
)
|
||||
subparser.add_argument(
|
||||
"--load_engine",
|
||||
action="store_true",
|
||||
help="Load engine from the artifacts_dir. Ignores all other options.",
|
||||
)
|
||||
subparser.add_argument(
|
||||
"--artifacts_dir",
|
||||
type=str,
|
||||
help="Whether to store (or retrieve) artifacts.",
|
||||
)
|
||||
subparser.add_argument(
|
||||
"--mode",
|
||||
type=str,
|
||||
choices=["onnx", "inetdef"],
|
||||
help="Whether to use ONNX parser or INetworkDefinition APIs to construct the network.",
|
||||
)
|
||||
subparser.add_argument("--aot", action="store_true", help="Use the AOT implementation of the plugin.")
|
||||
subparser.add_argument(
|
||||
"-v", "--verbose", action="store_true", help="Enable verbose log output."
|
||||
)
|
||||
|
||||
return subparser
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
parser = argparse.ArgumentParser()
|
||||
|
||||
parser = argparse.ArgumentParser(description="Main script help")
|
||||
subparsers = parser.add_subparsers(dest="sample", help="Mode help", required=True)
|
||||
|
||||
setup_add_sample(subparsers)
|
||||
setup_inplace_add_sample(subparsers)
|
||||
circ_pad_subparser = setup_circ_pad_sample(subparsers)
|
||||
setup_non_zero_sample(subparsers)
|
||||
|
||||
args = parser.parse_args()
|
||||
|
||||
if args.verbose:
|
||||
logging.getLogger("QuicklyDeployablePlugins").setLevel(logging.DEBUG)
|
||||
|
||||
if args.sample == "add":
|
||||
run_add(args.autotune)
|
||||
if args.sample == "inplace_add":
|
||||
run_inplace_add()
|
||||
if args.sample == "non_zero":
|
||||
run_non_zero()
|
||||
if args.sample == "circ_pad":
|
||||
if args.mode == "onnx":
|
||||
if args.artifacts_dir is None:
|
||||
parser.error(
|
||||
"circ_pad: argument --mode: When mode is 'onnx', artifacts_dir is required"
|
||||
)
|
||||
|
||||
save_or_load_engine = None
|
||||
|
||||
if args.load_engine is True:
|
||||
if args.save_engine is True:
|
||||
parser.error(
|
||||
"circ_pad: save_engine and load_engine cannot be specified at the same time. First save_engine and load_engine separately."
|
||||
)
|
||||
else:
|
||||
if args.multi_tactic is True or args.mode is not None:
|
||||
print(
|
||||
"warning circ_pad: when load_engine is specified, all other options except 'artifacts_dir' is ignored."
|
||||
)
|
||||
|
||||
save_or_load_engine = False
|
||||
else:
|
||||
if args.mode is None:
|
||||
circ_pad_subparser.print_help()
|
||||
parser.error(
|
||||
"circ_pad: '--mode' option is required."
|
||||
)
|
||||
|
||||
if args.save_engine is True:
|
||||
save_or_load_engine = True
|
||||
|
||||
run_circ_pad(args.multi_tactic, args.mode, args.artifacts_dir, save_or_load_engine, args.aot)
|
||||
@@ -0,0 +1,12 @@
|
||||
triton==3.5.0; (platform_system != "Windows")
|
||||
torch
|
||||
--extra-index-url https://pypi.ngc.nvidia.com
|
||||
polygraphy>=0.50.1
|
||||
colored
|
||||
numpy==1.26.4
|
||||
onnx==1.18.0; platform_system == "Windows"
|
||||
--extra-index-url https://pypi.ngc.nvidia.com
|
||||
onnx-graphsurgeon
|
||||
pyyaml==6.0.3
|
||||
requests==2.32.4
|
||||
tqdm==4.66.4
|
||||
Reference in New Issue
Block a user