Asfiya Baig
@@ -0,0 +1,66 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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.
|
||||
#
|
||||
|
||||
# We need cmake >= 3.8, since 3.8 introduced CUDA as a first class language
|
||||
cmake_minimum_required(VERSION 3.8 FATAL_ERROR)
|
||||
project(CircPadPlugin LANGUAGES CXX CUDA)
|
||||
|
||||
if(NOT MSVC)
|
||||
# Enable all compile warnings
|
||||
set(CMAKE_CXX_FLAGS "${CMAKE_CXX_FLAGS} -Wall -Wno-long-long -pedantic -Wno-deprecated-declarations")
|
||||
set(CMAKE_CUDA_FLAGS "${CMAKE_CUDA_FLAGS} -Xcompiler -Wno-deprecated-declarations")
|
||||
endif()
|
||||
|
||||
# Sets variable to a value if variable is unset.
|
||||
macro(set_ifndef var val)
|
||||
if (NOT ${var})
|
||||
set(${var} ${val})
|
||||
endif()
|
||||
message(STATUS "Configurable variable ${var} set to ${${var}}")
|
||||
endmacro()
|
||||
|
||||
# -------- CONFIGURATION --------
|
||||
if(NOT MSVC)
|
||||
set_ifndef(TRT_LIB /usr/lib/x86_64-linux-gnu)
|
||||
set_ifndef(TRT_INCLUDE /usr/include/x86_64-linux-gnu)
|
||||
set_ifndef(CUDA_INC_DIR /usr/local/cuda/include)
|
||||
set_ifndef(CUDA_LIB_DIR /usr/local/cuda)
|
||||
endif()
|
||||
|
||||
message("\nThe following variables are derived from the values of the previous variables unless provided explicitly:\n")
|
||||
|
||||
find_library(_NVINFER_LIB nvinfer HINTS ${TRT_LIB} PATH_SUFFIXES lib lib64)
|
||||
set_ifndef(NVINFER_LIB ${_NVINFER_LIB})
|
||||
|
||||
find_library(_CUDA_LIB cuda HINTS ${CUDA_LIB_DIR} PATH_SUFFIXES lib/stubs lib64/stubs)
|
||||
set_ifndef(CUDA_LIB ${_CUDA_LIB})
|
||||
|
||||
# -------- BUILDING --------
|
||||
|
||||
add_library(circ_pad_plugin SHARED
|
||||
${CMAKE_SOURCE_DIR}/circ_plugin_cpp/circ_pad_plugin.cu
|
||||
)
|
||||
|
||||
target_include_directories(circ_pad_plugin
|
||||
PUBLIC ${CUDA_INC_DIR}
|
||||
PUBLIC ${TRT_INCLUDE}
|
||||
)
|
||||
|
||||
set_property(TARGET circ_pad_plugin PROPERTY CUDA_STANDARD 14)
|
||||
|
||||
target_link_libraries(circ_pad_plugin PRIVATE ${NVINFER_LIB})
|
||||
target_link_libraries(circ_pad_plugin PRIVATE ${CUDA_LIB})
|
||||
@@ -0,0 +1,165 @@
|
||||
# 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.
|
||||
|
||||
### 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.
|
||||
|
||||
# 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
|
||||
|
||||
July 2023: Initial release of this sample
|
||||
|
||||
# Known issues
|
||||
|
||||
There are no known issues in this sample
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 17 KiB |
@@ -0,0 +1,82 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 argparse
|
||||
import onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
import onnx
|
||||
import ctypes
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
EngineFromNetwork,
|
||||
NetworkFromOnnxPath,
|
||||
TrtRunner,
|
||||
)
|
||||
|
||||
def parseArgs():
|
||||
parser = argparse.ArgumentParser(description="Options for Circular Padding plugin C++ example")
|
||||
|
||||
parser.add_argument('--precision', type=str, default="fp32", choices=["fp32", "fp16"], help="Precision to use for plugin")
|
||||
parser.add_argument('--plugin-lib', type=str, help="Path to the Circular Padding plugin lib", required=True)
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
args = parseArgs()
|
||||
|
||||
handle = ctypes.CDLL(args.plugin_lib)
|
||||
if not handle:
|
||||
raise RuntimeError("Could not load Circular Padding plugin library")
|
||||
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
inp_shape = (10, 3, 32, 32)
|
||||
X = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
# create ONNX model
|
||||
onnx_path = "test_CircPadPlugin.onnx"
|
||||
inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
|
||||
Y = gs.Variable(name="Y", dtype=precision)
|
||||
myPluginNode = gs.Node(
|
||||
name="CircPadPlugin",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputA],
|
||||
outputs=[Y],
|
||||
attrs={"pads": pads},
|
||||
)
|
||||
graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
# build engine
|
||||
build_engine = EngineFromNetwork(
|
||||
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
|
||||
)
|
||||
|
||||
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
# Run
|
||||
with TrtRunner(build_engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
|
||||
if np.allclose(Y, Y_ref):
|
||||
print("Inference result correct!")
|
||||
else:
|
||||
print("Inference result incorrect!")
|
||||
@@ -0,0 +1,336 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
import onnx
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
EngineFromNetwork,
|
||||
NetworkFromOnnxPath,
|
||||
TrtRunner
|
||||
)
|
||||
from polygraphy.json import to_json, from_json
|
||||
|
||||
from utils import checkCudaErrors, KernelHelper, parseArgs, CudaCtxManager
|
||||
from cuda import cuda
|
||||
|
||||
circ_pad_half_kernel = r'''
|
||||
#include <cuda_fp16.h>
|
||||
extern "C" __global__
|
||||
void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int Y_len) {
|
||||
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
for(int i = index; i < Y_len; i += stride)
|
||||
{
|
||||
int i3 = i % Y_shape[3];
|
||||
int i2 = (i / Y_shape[3]) % Y_shape[2];
|
||||
int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
|
||||
int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
|
||||
|
||||
int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
|
||||
int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
|
||||
int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
|
||||
int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
|
||||
|
||||
Y[i] = X[
|
||||
orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
|
||||
+ orig_dims[3] * orig_dims[2] * j1
|
||||
+ orig_dims[3] * j2
|
||||
+ j3
|
||||
];
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
circ_pad_float_kernel = r'''
|
||||
extern "C" __global__
|
||||
void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int Y_len) {
|
||||
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
for(int i = index; i < Y_len; i += stride)
|
||||
{
|
||||
int i3 = i % Y_shape[3];
|
||||
int i2 = (i / Y_shape[3]) % Y_shape[2];
|
||||
int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
|
||||
int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
|
||||
|
||||
int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
|
||||
int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
|
||||
int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
|
||||
int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
|
||||
|
||||
Y[i] = X[
|
||||
orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
|
||||
+ orig_dims[3] * orig_dims[2] * j1
|
||||
+ orig_dims[3] * j2
|
||||
+ j3
|
||||
];
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
class CircPadPlugin(trt.IPluginV2DynamicExt):
|
||||
def __init__(self, fc=None):
|
||||
trt.IPluginV2DynamicExt.__init__(self)
|
||||
self.pads = []
|
||||
self.X_shape = []
|
||||
self.N = 0
|
||||
|
||||
self.all_pads_d = None
|
||||
self.orig_dims_d = None
|
||||
self.Y_shape_d = None
|
||||
|
||||
self.num_outputs = 1
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_type = "CircPadPlugin"
|
||||
self.plugin_version = "1"
|
||||
|
||||
self.cuDevice = None
|
||||
|
||||
if fc is not None:
|
||||
assert set([f.name for f in fc]) == set(["pads", "N"]), "Field collection invalid"
|
||||
for f in fc:
|
||||
if f.name == "pads":
|
||||
self.pads = f.data
|
||||
elif f.name == "N":
|
||||
self.N = int(f.data)
|
||||
|
||||
def initialize(self):
|
||||
err, self.cuDevice = cuda.cuDeviceGet(0)
|
||||
trt.get_plugin_registry().acquire_plugin_resource("cuda_ctx", CudaCtxManager(self.cuDevice))
|
||||
self.all_pads_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N * 2))
|
||||
self.orig_dims_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
|
||||
self.Y_shape_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
|
||||
|
||||
def get_output_datatype(self, index, input_types):
|
||||
return input_types[0]
|
||||
|
||||
def get_output_dimensions(self, output_index, inputs, exprBuilder):
|
||||
|
||||
output_dims = trt.DimsExprs(inputs[0])
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
|
||||
trt.DimensionOperation.SUM,
|
||||
inputs[0][len(output_dims) - i - 1],
|
||||
exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
|
||||
)
|
||||
|
||||
return output_dims
|
||||
|
||||
def serialize(self):
|
||||
return to_json({"pads": self.pads, "N": self.N})
|
||||
|
||||
def configure_plugin(self, inp, out):
|
||||
X_dims = inp[0].desc.dims
|
||||
self.X_shape = np.zeros((len(X_dims),))
|
||||
for i in range(len(X_dims)):
|
||||
self.X_shape[i] = X_dims[i]
|
||||
|
||||
all_pads = np.zeros((self.N * 2,), dtype=np.int32)
|
||||
orig_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
out_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
out_dims[self.N - i - 1] += self.pads[i * 2] + self.pads[i * 2 + 1]
|
||||
all_pads[self.N * 2 - 2 * i - 2] = self.pads[i * 2]
|
||||
all_pads[self.N * 2 - 2 * i - 1] = self.pads[i * 2 + 1]
|
||||
|
||||
# Copy vectors from host memory to device memory
|
||||
if self.all_pads_d:
|
||||
checkCudaErrors(cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes))
|
||||
if self.orig_dims_d:
|
||||
checkCudaErrors(cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes))
|
||||
if self.Y_shape_d:
|
||||
checkCudaErrors(cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes))
|
||||
|
||||
self.Y_len_d = np.prod(out_dims)
|
||||
|
||||
def supports_format_combination(self, pos, in_out, num_inputs):
|
||||
assert num_inputs == 1
|
||||
assert pos < len(in_out)
|
||||
|
||||
desc = in_out[pos]
|
||||
if desc.format != trt.TensorFormat.LINEAR:
|
||||
return False
|
||||
|
||||
# first input should be float16 or float32
|
||||
if pos == 0:
|
||||
return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
|
||||
|
||||
# output should have the same type as the input
|
||||
if pos == 1:
|
||||
return in_out[0].type == desc.type
|
||||
|
||||
assert False
|
||||
|
||||
def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
|
||||
|
||||
inp_dtype = trt.nptype(input_desc[0].type)
|
||||
|
||||
blockSize = 256
|
||||
numBlocks = int((np.prod(np.array(self.X_shape)) + blockSize - 1) // blockSize)
|
||||
|
||||
da = np.array([inputs[0]], dtype=np.uint64)
|
||||
dc = np.array([outputs[0]], dtype=np.uint64)
|
||||
|
||||
d_all_pads = np.array([int(self.all_pads_d)], dtype=np.uint64)
|
||||
d_orig_dims = np.array([int(self.orig_dims_d)], dtype=np.uint64)
|
||||
d_Y_shape = np.array([int(self.Y_shape_d)], dtype=np.uint64)
|
||||
Y_len = np.array(self.Y_len_d, dtype=np.uint32)
|
||||
|
||||
args = [da, d_all_pads, d_orig_dims, dc, d_Y_shape, Y_len]
|
||||
kernelArgs = np.array([arg.ctypes.data for arg in args], dtype=np.uint64)
|
||||
|
||||
stream_ptr = np.array([stream], dtype=np.uint64)
|
||||
|
||||
if inp_dtype == np.float32:
|
||||
kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice))
|
||||
_circ_pad_float_kernel = kernelHelper.getFunction(b'circ_pad_float')
|
||||
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_float_kernel,
|
||||
numBlocks, 1, 1,
|
||||
blockSize, 1, 1,
|
||||
0,
|
||||
stream_ptr,
|
||||
kernelArgs, 0))
|
||||
elif inp_dtype == np.float16:
|
||||
kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice))
|
||||
_circ_pad_half_kernel = kernelHelper.getFunction(b'circ_pad_half')
|
||||
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_half_kernel,
|
||||
numBlocks, 1, 1,
|
||||
blockSize, 1, 1,
|
||||
0,
|
||||
stream_ptr,
|
||||
kernelArgs, 0))
|
||||
else:
|
||||
raise ValueError("inp_dtype not valid")
|
||||
|
||||
def clone(self):
|
||||
cloned_plugin = CircPadPlugin()
|
||||
cloned_plugin.__dict__.update(self.__dict__)
|
||||
return cloned_plugin
|
||||
|
||||
def terminate(self):
|
||||
if self.all_pads_d:
|
||||
checkCudaErrors(cuda.cuMemFree(self.all_pads_d))
|
||||
if self.orig_dims_d:
|
||||
checkCudaErrors(cuda.cuMemFree(self.orig_dims_d))
|
||||
if self.Y_shape_d:
|
||||
checkCudaErrors(cuda.cuMemFree(self.Y_shape_d))
|
||||
|
||||
trt.get_plugin_registry().release_plugin_resource("cuda_ctx")
|
||||
|
||||
#
|
||||
# The following defaults take effect since the respective methods are not overriden
|
||||
#
|
||||
|
||||
# def get_serialization_size(self):
|
||||
# return len(to_json({"pads": self.pads}))
|
||||
|
||||
# def get_workspace_size(self, input_desc, output_desc):
|
||||
# return 0
|
||||
|
||||
# def destroy(self):
|
||||
# pass
|
||||
|
||||
|
||||
class CircPadPluginCreator(trt.IPluginCreator):
|
||||
def __init__(self):
|
||||
trt.IPluginCreator.__init__(self)
|
||||
self.name = "CircPadPlugin"
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_version = "1"
|
||||
self.field_names = trt.PluginFieldCollection([
|
||||
trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
|
||||
trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32)
|
||||
])
|
||||
|
||||
def create_plugin(self, name, fc):
|
||||
return CircPadPlugin(fc)
|
||||
|
||||
def deserialize_plugin(self, name, data):
|
||||
deserialized = CircPadPlugin()
|
||||
j = dict(from_json(data))
|
||||
deserialized.__dict__.update(j)
|
||||
return deserialized
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
args = parseArgs()
|
||||
|
||||
# Initialize CUDA Driver API
|
||||
err, = cuda.cuInit(0)
|
||||
|
||||
# Retrieve handle for device 0
|
||||
err, cuDevice = cuda.cuDeviceGet(0)
|
||||
|
||||
plg_registry = trt.get_plugin_registry()
|
||||
|
||||
# Create context
|
||||
plg_registry.acquire_plugin_resource("cuda_ctx", CudaCtxManager(cuDevice))
|
||||
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
|
||||
inp_shape = (100, 2, 32, 32)
|
||||
X = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
# Load standard plugins
|
||||
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
|
||||
trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="")
|
||||
|
||||
# Register plugin creator
|
||||
my_plugin_creator = CircPadPluginCreator()
|
||||
plg_registry.register_creator(my_plugin_creator, "")
|
||||
|
||||
# create ONNX model
|
||||
onnx_path = "test_CircPadPlugin.onnx"
|
||||
inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
|
||||
Y = gs.Variable(name="Y", dtype=precision)
|
||||
myPluginNode = gs.Node(
|
||||
name="CircPadPlugin",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputA],
|
||||
outputs=[Y],
|
||||
attrs={"pads": pads, "N": 4},
|
||||
)
|
||||
graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
# build engine
|
||||
build_engine = EngineFromNetwork(
|
||||
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
|
||||
)
|
||||
|
||||
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
# Run
|
||||
with TrtRunner(build_engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
|
||||
if np.allclose(Y, Y_ref):
|
||||
print("Inference result correct!")
|
||||
else:
|
||||
print("Inference result incorrect!")
|
||||
|
||||
plg_registry.release_plugin_resource("cuda_ctx")
|
||||
@@ -0,0 +1,290 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
import onnx
|
||||
import cupy as cp
|
||||
import time
|
||||
import pickle
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
EngineFromNetwork,
|
||||
NetworkFromOnnxPath,
|
||||
TrtRunner
|
||||
)
|
||||
|
||||
from polygraphy.json import to_json, from_json
|
||||
|
||||
from utils import volume, parseArgs
|
||||
|
||||
circ_pad_half_kernel = cp.RawKernel(r'''
|
||||
#include <cuda_fp16.h>
|
||||
extern "C" __global__
|
||||
void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int const* Y_len) {
|
||||
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
for(int i = index; i < *Y_len; i += stride)
|
||||
{
|
||||
int i3 = i % Y_shape[3];
|
||||
int i2 = (i / Y_shape[3]) % Y_shape[2];
|
||||
int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
|
||||
int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
|
||||
|
||||
int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
|
||||
int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
|
||||
int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
|
||||
int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
|
||||
|
||||
Y[i] = X[
|
||||
orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
|
||||
+ orig_dims[3] * orig_dims[2] * j1
|
||||
+ orig_dims[3] * j2
|
||||
+ j3
|
||||
];
|
||||
}
|
||||
}
|
||||
''', 'circ_pad_half')
|
||||
|
||||
circ_pad_float_kernel = cp.RawKernel(r'''
|
||||
extern "C" __global__
|
||||
void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int const* Y_len) {
|
||||
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
for(int i = index; i < *Y_len; i += stride)
|
||||
{
|
||||
int i3 = i % Y_shape[3];
|
||||
int i2 = (i / Y_shape[3]) % Y_shape[2];
|
||||
int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
|
||||
int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
|
||||
|
||||
int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
|
||||
int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
|
||||
int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
|
||||
int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
|
||||
|
||||
Y[i] = X[
|
||||
orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
|
||||
+ orig_dims[3] * orig_dims[2] * j1
|
||||
+ orig_dims[3] * j2
|
||||
+ j3
|
||||
];
|
||||
}
|
||||
}
|
||||
''', 'circ_pad_float')
|
||||
|
||||
class CircPadPlugin(trt.IPluginV2DynamicExt):
|
||||
def __init__(self, fc=None):
|
||||
trt.IPluginV2DynamicExt.__init__(self)
|
||||
self.pads = []
|
||||
self.X_shape = []
|
||||
|
||||
self.num_outputs = 1
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_type = "CircPadPlugin"
|
||||
self.plugin_version = "1"
|
||||
|
||||
if fc is not None:
|
||||
assert fc[0].name == "pads"
|
||||
self.pads = fc[0].data
|
||||
|
||||
def get_output_datatype(self, index, input_types):
|
||||
return input_types[0]
|
||||
|
||||
def get_output_dimensions(self, output_index, inputs, exprBuilder):
|
||||
|
||||
output_dims = trt.DimsExprs(inputs[0])
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
|
||||
trt.DimensionOperation.SUM,
|
||||
inputs[0][len(output_dims) - i - 1],
|
||||
exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
|
||||
)
|
||||
|
||||
return output_dims
|
||||
|
||||
def serialize(self):
|
||||
return to_json({"pads": self.pads})
|
||||
|
||||
def configure_plugin(self, inp, out):
|
||||
X_dims = inp[0].desc.dims
|
||||
self.X_shape = np.zeros((len(X_dims),))
|
||||
for i in range(len(X_dims)):
|
||||
self.X_shape[i] = X_dims[i]
|
||||
|
||||
N = len(self.X_shape)
|
||||
all_pads = np.zeros((N * 2,))
|
||||
orig_dims = np.array(self.X_shape)
|
||||
out_dims = np.array(self.X_shape)
|
||||
|
||||
for i in range(np.size(pads) // 2):
|
||||
out_dims[N - i - 1] += self.pads[i * 2] + self.pads[i * 2 + 1]
|
||||
all_pads[N * 2 - 2 * i - 2] = self.pads[i * 2]
|
||||
all_pads[N * 2 - 2 * i - 1] = self.pads[i * 2 + 1]
|
||||
|
||||
self.all_pads_d = cp.asarray(all_pads, dtype=cp.int32)
|
||||
self.orig_dims_d = cp.asarray(orig_dims, dtype=cp.int32)
|
||||
self.Y_shape_d = cp.asarray(out_dims, dtype=cp.int32)
|
||||
self.Y_len_d = cp.array([np.prod(out_dims)], dtype=cp.int32)
|
||||
|
||||
def supports_format_combination(self, pos, in_out, num_inputs):
|
||||
assert num_inputs == 1
|
||||
assert pos < len(in_out)
|
||||
|
||||
desc = in_out[pos]
|
||||
if desc.format != trt.TensorFormat.LINEAR:
|
||||
return False
|
||||
|
||||
# first input should be float16 or float32
|
||||
if pos == 0:
|
||||
return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
|
||||
|
||||
# output should have the same type as the input
|
||||
if pos == 1:
|
||||
return in_out[0].type == desc.type
|
||||
|
||||
assert False
|
||||
|
||||
def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
|
||||
|
||||
inp_dtype = trt.nptype(input_desc[0].type)
|
||||
|
||||
a_mem = cp.cuda.UnownedMemory(
|
||||
inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
|
||||
)
|
||||
c_mem = cp.cuda.UnownedMemory(
|
||||
outputs[0],
|
||||
volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
|
||||
self,
|
||||
)
|
||||
|
||||
a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
|
||||
c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
|
||||
|
||||
a = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
|
||||
c = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
|
||||
|
||||
cuda_stream = cp.cuda.ExternalStream(stream)
|
||||
|
||||
blockSize = 256
|
||||
numBlocks = int((np.prod(np.array(self.X_shape)) + blockSize - 1) // blockSize)
|
||||
|
||||
with cuda_stream:
|
||||
if inp_dtype == np.float32:
|
||||
circ_pad_float_kernel((numBlocks,), (blockSize,), (a, self.all_pads_d, self.orig_dims_d, c, self.Y_shape_d, self.Y_len_d))
|
||||
elif inp_dtype == np.float16:
|
||||
circ_pad_half_kernel((numBlocks,), (blockSize,), (a, self.all_pads_d, self.orig_dims_d, c, self.Y_shape_d, self.Y_len_d))
|
||||
else:
|
||||
raise ValueError("inp_dtype not valid")
|
||||
|
||||
def clone(self):
|
||||
cloned_plugin = CircPadPlugin()
|
||||
cloned_plugin.__dict__.update(self.__dict__)
|
||||
return cloned_plugin
|
||||
|
||||
#
|
||||
# The following defaults take effect since the respective methods are not overriden
|
||||
#
|
||||
|
||||
# def initialize(self):
|
||||
# pass
|
||||
|
||||
# def get_serialization_size(self):
|
||||
# return len(to_json({"pads": self.pads}))
|
||||
|
||||
# def get_workspace_size(self, input_desc, output_desc):
|
||||
# return 0
|
||||
|
||||
# def destroy(self):
|
||||
# pass
|
||||
|
||||
# def terminate(self):
|
||||
# pass
|
||||
|
||||
class CircPadPluginCreator(trt.IPluginCreator):
|
||||
def __init__(self):
|
||||
trt.IPluginCreator.__init__(self)
|
||||
|
||||
self.name = "CircPadPlugin"
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_version = "1"
|
||||
self.field_names = trt.PluginFieldCollection(
|
||||
[trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
|
||||
)
|
||||
|
||||
def create_plugin(self, name, fc):
|
||||
return CircPadPlugin(fc)
|
||||
|
||||
def deserialize_plugin(self, name, data):
|
||||
j = dict(from_json(data.decode("utf-8")))
|
||||
deserialized = CircPadPlugin()
|
||||
deserialized.__dict__.update(j)
|
||||
return deserialized
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
args = parseArgs()
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
|
||||
inp_shape = (100, 2, 32, 32)
|
||||
X = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
# Load standard plugins
|
||||
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
|
||||
trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="")
|
||||
|
||||
# Register plugin creator
|
||||
plg_registry = trt.get_plugin_registry()
|
||||
my_plugin_creator = CircPadPluginCreator()
|
||||
plg_registry.register_creator(my_plugin_creator, "")
|
||||
|
||||
# create ONNX model
|
||||
onnx_path = "test_CircPadPlugin.onnx"
|
||||
inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
|
||||
Y = gs.Variable(name="Y", dtype=precision)
|
||||
myPluginNode = gs.Node(
|
||||
name="CircPadPlugin",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputA],
|
||||
outputs=[Y],
|
||||
attrs={"pads": pads},
|
||||
)
|
||||
graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
# build engine
|
||||
build_engine = EngineFromNetwork(
|
||||
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
|
||||
)
|
||||
|
||||
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
# Run
|
||||
with TrtRunner(build_engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
|
||||
if np.allclose(Y, Y_ref):
|
||||
print("Inference result correct!")
|
||||
else:
|
||||
print("Inference result incorrect!")
|
||||
@@ -0,0 +1,336 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
TrtRunner,
|
||||
create_network,
|
||||
engine_from_network
|
||||
)
|
||||
|
||||
from polygraphy.json import to_json, from_json
|
||||
|
||||
from utils import checkCudaErrors, KernelHelper, parseArgs, CudaCtxManager
|
||||
from cuda import cuda
|
||||
|
||||
circ_pad_half_kernel = r'''
|
||||
#include <cuda_fp16.h>
|
||||
extern "C" __global__
|
||||
void circ_pad_half(half const* X, int const* all_pads, int const* orig_dims, half* Y, int const* Y_shape, int Y_len) {
|
||||
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
for(int i = index; i < Y_len; i += stride)
|
||||
{
|
||||
int i3 = i % Y_shape[3];
|
||||
int i2 = (i / Y_shape[3]) % Y_shape[2];
|
||||
int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
|
||||
int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
|
||||
|
||||
int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
|
||||
int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
|
||||
int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
|
||||
int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
|
||||
|
||||
Y[i] = X[
|
||||
orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
|
||||
+ orig_dims[3] * orig_dims[2] * j1
|
||||
+ orig_dims[3] * j2
|
||||
+ j3
|
||||
];
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
circ_pad_float_kernel = r'''
|
||||
extern "C" __global__
|
||||
void circ_pad_float(float const* X, int const* all_pads, int const* orig_dims, float* Y, int const* Y_shape, int Y_len) {
|
||||
int index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int stride = blockDim.x * gridDim.x;
|
||||
|
||||
for(int i = index; i < Y_len; i += stride)
|
||||
{
|
||||
int i3 = i % Y_shape[3];
|
||||
int i2 = (i / Y_shape[3]) % Y_shape[2];
|
||||
int i1 = (i / Y_shape[3] / Y_shape[2]) % Y_shape[1];
|
||||
int i0 = i / Y_shape[3] / Y_shape[2] / Y_shape[1];
|
||||
|
||||
int j0 = (i0 - all_pads[0] + orig_dims[0]) % orig_dims[0];
|
||||
int j1 = (i1 - all_pads[2] + orig_dims[1]) % orig_dims[1];
|
||||
int j2 = (i2 - all_pads[4] + orig_dims[2]) % orig_dims[2];
|
||||
int j3 = (i3 - all_pads[6] + orig_dims[3]) % orig_dims[3];
|
||||
|
||||
Y[i] = X[
|
||||
orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
|
||||
+ orig_dims[3] * orig_dims[2] * j1
|
||||
+ orig_dims[3] * j2
|
||||
+ j3
|
||||
];
|
||||
}
|
||||
}
|
||||
'''
|
||||
|
||||
class CircPadPlugin(trt.IPluginV2DynamicExt):
|
||||
def __init__(self, fc=None):
|
||||
trt.IPluginV2DynamicExt.__init__(self)
|
||||
self.pads = []
|
||||
self.X_shape = []
|
||||
self.N = 0
|
||||
|
||||
self.all_pads_d = None
|
||||
self.orig_dims_d = None
|
||||
self.Y_shape_d = None
|
||||
|
||||
self.num_outputs = 1
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_type = "CircPadPlugin"
|
||||
self.plugin_version = "1"
|
||||
|
||||
self.cuDevice = None
|
||||
|
||||
if fc is not None:
|
||||
assert set([f.name for f in fc]) == set(["pads", "N"]), "Field collection invalid"
|
||||
for f in fc:
|
||||
if f.name == "pads":
|
||||
self.pads = f.data
|
||||
elif f.name == "N":
|
||||
self.N = int(f.data)
|
||||
|
||||
def initialize(self):
|
||||
err, self.cuDevice = cuda.cuDeviceGet(0)
|
||||
trt.get_plugin_registry().acquire_plugin_resource("cuda_ctx", CudaCtxManager(self.cuDevice))
|
||||
self.all_pads_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N * 2))
|
||||
self.orig_dims_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
|
||||
self.Y_shape_d = checkCudaErrors(cuda.cuMemAlloc(np.int32().itemsize * self.N))
|
||||
|
||||
def get_output_datatype(self, index, input_types):
|
||||
return input_types[0]
|
||||
|
||||
def get_output_dimensions(self, output_index, inputs, exprBuilder):
|
||||
|
||||
output_dims = trt.DimsExprs(inputs[0])
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
|
||||
trt.DimensionOperation.SUM,
|
||||
inputs[0][len(output_dims) - i - 1],
|
||||
exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
|
||||
)
|
||||
|
||||
return output_dims
|
||||
|
||||
def serialize(self):
|
||||
return to_json({"pads": self.pads, "N": self.N})
|
||||
|
||||
def configure_plugin(self, inp, out):
|
||||
X_dims = inp[0].desc.dims
|
||||
self.X_shape = np.zeros((len(X_dims),))
|
||||
for i in range(len(X_dims)):
|
||||
self.X_shape[i] = X_dims[i]
|
||||
|
||||
all_pads = np.zeros((self.N * 2,), dtype=np.int32)
|
||||
orig_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
out_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
out_dims[self.N - i - 1] += self.pads[i * 2] + self.pads[i * 2 + 1]
|
||||
all_pads[self.N * 2 - 2 * i - 2] = self.pads[i * 2]
|
||||
all_pads[self.N * 2 - 2 * i - 1] = self.pads[i * 2 + 1]
|
||||
|
||||
# Copy vectors from host memory to device memory
|
||||
if self.all_pads_d:
|
||||
checkCudaErrors(cuda.cuMemcpyHtoD(self.all_pads_d, all_pads, all_pads.nbytes))
|
||||
if self.orig_dims_d:
|
||||
checkCudaErrors(cuda.cuMemcpyHtoD(self.orig_dims_d, orig_dims, orig_dims.nbytes))
|
||||
if self.Y_shape_d:
|
||||
checkCudaErrors(cuda.cuMemcpyHtoD(self.Y_shape_d, out_dims, out_dims.nbytes))
|
||||
|
||||
self.Y_len_d = np.prod(out_dims)
|
||||
|
||||
def supports_format_combination(self, pos, in_out, num_inputs):
|
||||
assert num_inputs == 1
|
||||
assert pos < len(in_out)
|
||||
|
||||
desc = in_out[pos]
|
||||
if desc.format != trt.TensorFormat.LINEAR:
|
||||
return False
|
||||
|
||||
# first input should be float16 or float32
|
||||
if pos == 0:
|
||||
return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
|
||||
|
||||
# output should have the same type as the input
|
||||
if pos == 1:
|
||||
return in_out[0].type == desc.type
|
||||
|
||||
assert False
|
||||
|
||||
def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
|
||||
|
||||
inp_dtype = trt.nptype(input_desc[0].type)
|
||||
|
||||
blockSize = 256
|
||||
numBlocks = int((np.prod(np.array(self.X_shape)) + blockSize - 1) // blockSize)
|
||||
|
||||
da = np.array([inputs[0]], dtype=np.uint64)
|
||||
dc = np.array([outputs[0]], dtype=np.uint64)
|
||||
|
||||
d_all_pads = np.array([int(self.all_pads_d)], dtype=np.uint64)
|
||||
d_orig_dims = np.array([int(self.orig_dims_d)], dtype=np.uint64)
|
||||
d_Y_shape = np.array([int(self.Y_shape_d)], dtype=np.uint64)
|
||||
Y_len = np.array(self.Y_len_d, dtype=np.uint32)
|
||||
|
||||
args = [da, d_all_pads, d_orig_dims, dc, d_Y_shape, Y_len]
|
||||
kernelArgs = np.array([arg.ctypes.data for arg in args], dtype=np.uint64)
|
||||
|
||||
stream_ptr = np.array([stream], dtype=np.uint64)
|
||||
|
||||
if inp_dtype == np.float32:
|
||||
kernelHelper = KernelHelper(circ_pad_float_kernel, int(self.cuDevice))
|
||||
_circ_pad_float_kernel = kernelHelper.getFunction(b'circ_pad_float')
|
||||
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_float_kernel,
|
||||
numBlocks, 1, 1,
|
||||
blockSize, 1, 1,
|
||||
0,
|
||||
stream_ptr,
|
||||
kernelArgs, 0))
|
||||
elif inp_dtype == np.float16:
|
||||
kernelHelper = KernelHelper(circ_pad_half_kernel, int(self.cuDevice))
|
||||
_circ_pad_half_kernel = kernelHelper.getFunction(b'circ_pad_half')
|
||||
checkCudaErrors(cuda.cuLaunchKernel(_circ_pad_half_kernel,
|
||||
numBlocks, 1, 1,
|
||||
blockSize, 1, 1,
|
||||
0,
|
||||
stream_ptr,
|
||||
kernelArgs, 0))
|
||||
else:
|
||||
raise ValueError("inp_dtype not valid")
|
||||
|
||||
def clone(self):
|
||||
cloned_plugin = CircPadPlugin()
|
||||
cloned_plugin.__dict__.update(self.__dict__)
|
||||
return cloned_plugin
|
||||
|
||||
def terminate(self):
|
||||
if self.all_pads_d:
|
||||
checkCudaErrors(cuda.cuMemFree(self.all_pads_d))
|
||||
if self.orig_dims_d:
|
||||
checkCudaErrors(cuda.cuMemFree(self.orig_dims_d))
|
||||
if self.Y_shape_d:
|
||||
checkCudaErrors(cuda.cuMemFree(self.Y_shape_d))
|
||||
|
||||
plg_registry.release_plugin_resource("cuda_ctx")
|
||||
|
||||
#
|
||||
# The following defaults take effect since the respective methods are not overriden
|
||||
#
|
||||
|
||||
# def get_serialization_size(self):
|
||||
# return len(to_json({"pads": self.pads}))
|
||||
|
||||
# def get_workspace_size(self, input_desc, output_desc):
|
||||
# return 0
|
||||
|
||||
# def destroy(self):
|
||||
# pass
|
||||
|
||||
|
||||
class CircPadPluginCreator(trt.IPluginCreator):
|
||||
def __init__(self):
|
||||
trt.IPluginCreator.__init__(self)
|
||||
self.name = "CircPadPlugin"
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_version = "1"
|
||||
self.field_names = trt.PluginFieldCollection([
|
||||
trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32),
|
||||
trt.PluginField("N", np.array([]), trt.PluginFieldType.INT32)
|
||||
])
|
||||
|
||||
def create_plugin(self, name, fc):
|
||||
return CircPadPlugin(fc)
|
||||
|
||||
def deserialize_plugin(self, name, data):
|
||||
deserialized = CircPadPlugin()
|
||||
j = dict(from_json(data))
|
||||
deserialized.__dict__.update(j)
|
||||
return deserialized
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
args = parseArgs()
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
|
||||
# Initialize CUDA Driver API
|
||||
err, = cuda.cuInit(0)
|
||||
|
||||
# Retrieve handle for device 0
|
||||
err, cuDevice = cuda.cuDeviceGet(0)
|
||||
|
||||
plg_registry = trt.get_plugin_registry()
|
||||
|
||||
# Create context
|
||||
plg_registry.acquire_plugin_resource("cuda_ctx", CudaCtxManager(cuDevice))
|
||||
|
||||
inp_shape = (100, 2, 32, 32)
|
||||
X = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
|
||||
# Load standard plugins (if needed)
|
||||
trt.init_libnvinfer_plugins(TRT_LOGGER, namespace="")
|
||||
|
||||
# Register plugin creator
|
||||
my_plugin_creator = CircPadPluginCreator()
|
||||
plg_registry.register_creator(my_plugin_creator, "")
|
||||
|
||||
# Create plugin object
|
||||
builder, network = create_network()
|
||||
plg_creator = plg_registry.get_plugin_creator("CircPadPlugin", "1", "")
|
||||
plugin_fields_list = [
|
||||
trt.PluginField("pads", np.array(pads, dtype=np.int32), trt.PluginFieldType.INT32),
|
||||
trt.PluginField("N", np.array([4], dtype=np.int32), trt.PluginFieldType.INT32),
|
||||
]
|
||||
pfc = trt.PluginFieldCollection(plugin_fields_list)
|
||||
plugin = plg_creator.create_plugin("CircPadPlugin", pfc)
|
||||
|
||||
# Populate network
|
||||
input_X = network.add_input(name="X", dtype=trt.float32 if precision==np.float32 else trt.float16, shape=X.shape)
|
||||
out = network.add_plugin_v2([input_X], plugin)
|
||||
out.get_output(0).name = "Y"
|
||||
network.mark_output(tensor=out.get_output(0))
|
||||
|
||||
# Build engine
|
||||
config = builder.create_builder_config()
|
||||
engine = engine_from_network((builder, network), CreateConfig(fp16=precision==trt.float16))
|
||||
|
||||
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
# Run
|
||||
with TrtRunner(engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
|
||||
if np.allclose(Y, Y_ref):
|
||||
print("Inference result correct!")
|
||||
else:
|
||||
print("Inference result incorrect!")
|
||||
|
||||
plg_registry.release_plugin_resource("cuda_ctx")
|
||||
@@ -0,0 +1,318 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2024 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 onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
import onnx
|
||||
import cupy as cp
|
||||
import logging
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
EngineFromNetwork,
|
||||
NetworkFromOnnxPath,
|
||||
TrtRunner,
|
||||
)
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
from enum import IntEnum
|
||||
|
||||
from polygraphy.json import to_json, from_json
|
||||
import torch
|
||||
|
||||
from utils import volume, parseArgs
|
||||
|
||||
logger = logging.getLogger("CircPadMultiTactic")
|
||||
|
||||
class Tactic(IntEnum):
|
||||
TORCH = 1
|
||||
TRITON = 2
|
||||
|
||||
@triton.jit
|
||||
def circ_pad(X,
|
||||
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,
|
||||
Y_shape_1, Y_shape_2, Y_shape_3,
|
||||
X_len, Y_len, 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)
|
||||
|
||||
class CircPadPlugin(trt.IPluginV3, trt.IPluginV3OneCore, trt.IPluginV3OneBuild, trt.IPluginV3OneRuntime):
|
||||
def __init__(self, fc=None, phase=None):
|
||||
trt.IPluginV3.__init__(self)
|
||||
trt.IPluginV3OneCore.__init__(self)
|
||||
trt.IPluginV3OneBuild.__init__(self)
|
||||
trt.IPluginV3OneRuntime.__init__(self)
|
||||
self.pads = []
|
||||
self.X_shape = []
|
||||
|
||||
self.num_outputs = 1
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_name = "CircPadPlugin"
|
||||
self.plugin_version = "1"
|
||||
|
||||
# Set the timing cache ID to prevent unnecessary timing of second plugin instance
|
||||
self.timing_cache_id = ""
|
||||
|
||||
self.tactic = None
|
||||
|
||||
if fc is not None:
|
||||
assert fc[0].name == "pads"
|
||||
self.pads = fc[0].data
|
||||
|
||||
if phase is not None:
|
||||
self.phase = phase
|
||||
|
||||
def get_capability_interface(self, type):
|
||||
return self
|
||||
|
||||
def get_output_data_types(self, input_types):
|
||||
return [input_types[0]]
|
||||
|
||||
def get_output_shapes(self, inputs, shape_inputs, exprBuilder):
|
||||
output_dims = trt.DimsExprs(inputs[0])
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
|
||||
trt.DimensionOperation.SUM,
|
||||
inputs[0][len(output_dims) - i - 1],
|
||||
exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
|
||||
)
|
||||
|
||||
return [output_dims]
|
||||
|
||||
def get_fields_to_serialize(self):
|
||||
return trt.PluginFieldCollection([
|
||||
trt.PluginField("pads", self.pads, trt.PluginFieldType.INT32)
|
||||
])
|
||||
|
||||
def configure_plugin(self, inp, out):
|
||||
pass
|
||||
|
||||
def on_shape_change(self, inp, out):
|
||||
X_dims = inp[0].dims
|
||||
self.X_shape = np.zeros((len(X_dims),))
|
||||
for i in range(len(X_dims)):
|
||||
self.X_shape[i] = X_dims[i]
|
||||
|
||||
def supports_format_combination(self, pos, in_out, num_inputs):
|
||||
assert num_inputs == 1
|
||||
assert pos < len(in_out)
|
||||
|
||||
desc = in_out[pos].desc
|
||||
if desc.format != trt.TensorFormat.LINEAR:
|
||||
return False
|
||||
|
||||
# first input should be float16 or float32
|
||||
if pos == 0:
|
||||
return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
|
||||
|
||||
# output should have the same type as the input
|
||||
if pos == 1:
|
||||
return in_out[0].desc.type == desc.type
|
||||
|
||||
assert False
|
||||
|
||||
def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
|
||||
|
||||
inp_dtype = trt.nptype(input_desc[0].type)
|
||||
|
||||
a_mem = cp.cuda.UnownedMemory(
|
||||
inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
|
||||
)
|
||||
c_mem = cp.cuda.UnownedMemory(
|
||||
outputs[0],
|
||||
volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
|
||||
self,
|
||||
)
|
||||
|
||||
a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
|
||||
c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
|
||||
|
||||
c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
|
||||
|
||||
if self.phase == trt.TensorRTPhase.BUILD:
|
||||
logger.info(f"Timing tactic: {self.tactic}")
|
||||
|
||||
if self.tactic == Tactic.TORCH:
|
||||
# Use PyTorch functional op - no need to write kernel
|
||||
a_d = cp.ndarray(tuple(input_desc[0].dims), dtype=inp_dtype, memptr=a_ptr)
|
||||
a_t = torch.as_tensor(a_d, device='cuda')
|
||||
out = torch.nn.functional.pad(a_t, self.pads.tolist(), mode='circular')
|
||||
cp.copyto(c_d, cp.reshape(cp.asarray(out), (-1,)))
|
||||
elif self.tactic == Tactic.TRITON:
|
||||
a_d = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
|
||||
a_t = torch.as_tensor(a_d, device='cuda')
|
||||
c_t = torch.as_tensor(c_d, device='cuda')
|
||||
|
||||
N = len(self.X_shape)
|
||||
all_pads = np.zeros((N * 2,), dtype=np.int32)
|
||||
orig_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
out_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
|
||||
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()
|
||||
orig_dims = orig_dims.tolist()
|
||||
out_dims = out_dims.tolist()
|
||||
|
||||
blockSize = 256
|
||||
numBlocks = tuple(int((np.prod(out_dims) + blockSize - 1) // blockSize))
|
||||
|
||||
circ_pad[numBlocks](a_t,
|
||||
all_pads[0], all_pads[2], all_pads[4], all_pads[6],
|
||||
orig_dims[0], orig_dims[1], orig_dims[2], orig_dims[3],
|
||||
c_t,
|
||||
out_dims[1], out_dims[2], out_dims[3],
|
||||
int(np.prod(orig_dims)), int(np.prod(out_dims)), BLOCK_SIZE=256
|
||||
)
|
||||
else:
|
||||
raise RuntimeError("Invalid tactic")
|
||||
|
||||
def attach_to_context(self, context):
|
||||
return self.clone()
|
||||
|
||||
def get_valid_tactics(self):
|
||||
return [int(Tactic.TORCH), int(Tactic.TRITON)]
|
||||
|
||||
def set_tactic(self, tactic):
|
||||
self.tactic = Tactic(tactic)
|
||||
|
||||
if self.phase == trt.TensorRTPhase.RUNTIME:
|
||||
logger.info(f"Best tactic chosen: {self.tactic}")
|
||||
|
||||
def clone(self):
|
||||
cloned_plugin = CircPadPlugin()
|
||||
cloned_plugin.__dict__.update(self.__dict__)
|
||||
return cloned_plugin
|
||||
|
||||
#
|
||||
# The following defaults take effect since the respective methods are not overriden
|
||||
#
|
||||
|
||||
# def get_workspace_size(self, input_desc, output_desc):
|
||||
# return 0
|
||||
|
||||
# def destroy(self):
|
||||
# pass
|
||||
|
||||
|
||||
class CircPadPluginCreator(trt.IPluginCreatorV3One):
|
||||
def __init__(self):
|
||||
trt.IPluginCreatorV3One.__init__(self)
|
||||
self.name = "CircPadPlugin"
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_version = "1"
|
||||
self.field_names = trt.PluginFieldCollection([
|
||||
trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)
|
||||
])
|
||||
|
||||
def create_plugin(self, name, fc, phase):
|
||||
return CircPadPlugin(fc, phase)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
logging.basicConfig()
|
||||
logger.setLevel(logging.INFO)
|
||||
|
||||
args = parseArgs()
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
|
||||
inp_shape = (10, 3, 32, 32)
|
||||
X_A = np.random.normal(size=inp_shape).astype(precision)
|
||||
X_B = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
# Register plugin creator
|
||||
plg_registry = trt.get_plugin_registry()
|
||||
my_plugin_creator = CircPadPluginCreator()
|
||||
plg_registry.register_creator(my_plugin_creator, "")
|
||||
|
||||
# create ONNX model
|
||||
onnx_path = "test_CircPadPlugin.onnx"
|
||||
inputA = gs.Variable(name="X_A", shape=inp_shape, dtype=precision)
|
||||
inputB = gs.Variable(name="X_B", shape=inp_shape, dtype=precision)
|
||||
Y_A = gs.Variable(name="Y_A", dtype=precision)
|
||||
Y_B = gs.Variable(name="Y_B", dtype=precision)
|
||||
myPluginNode_A = gs.Node(
|
||||
name="CircPadPlugin_A",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputA],
|
||||
outputs=[Y_A],
|
||||
attrs={"pads": pads},
|
||||
)
|
||||
myPluginNode_B = gs.Node(
|
||||
name="CircPadPlugin_B",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputB],
|
||||
outputs=[Y_B],
|
||||
attrs={"pads": pads},
|
||||
)
|
||||
|
||||
graph = gs.Graph(nodes=[myPluginNode_A, myPluginNode_B], inputs=[inputA, inputB], outputs=[Y_A, Y_B], opset=16)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
# build engine
|
||||
build_engine = EngineFromNetwork(
|
||||
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
|
||||
)
|
||||
|
||||
Y_A_ref = np.pad(X_A, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
Y_B_ref = np.pad(X_B, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
|
||||
# Run
|
||||
with TrtRunner(build_engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X_A": X_A, "X_B": X_B})
|
||||
Y_A_out = outputs["Y_A"]
|
||||
Y_B_out = outputs["Y_B"]
|
||||
|
||||
if np.allclose(Y_A_out, Y_A_ref):
|
||||
print("Inference result A correct!")
|
||||
else:
|
||||
print("Inference result A incorrect!")
|
||||
|
||||
if np.allclose(Y_B_out, Y_B_ref):
|
||||
print("Inference result B correct!")
|
||||
else:
|
||||
print("Inference result B incorrect!")
|
||||
@@ -0,0 +1,249 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
import onnx
|
||||
import cupy as cp
|
||||
from numba import cuda
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
EngineFromNetwork,
|
||||
NetworkFromOnnxPath,
|
||||
TrtRunner,
|
||||
)
|
||||
|
||||
from polygraphy.json import to_json, from_json
|
||||
from utils import volume, parseArgs
|
||||
|
||||
@cuda.jit
|
||||
def circ_pad(X, all_pads, orig_dims, Y, Y_shape, Y_len):
|
||||
index = cuda.blockIdx.x * cuda.blockDim.x + cuda.threadIdx.x
|
||||
stride = cuda.blockDim.x * cuda.gridDim.x
|
||||
|
||||
for i in range(index, Y_len, stride):
|
||||
i3 = int(i % Y_shape[3])
|
||||
i2 = int((i // Y_shape[3]) % Y_shape[2])
|
||||
i1 = int((i // Y_shape[3] // Y_shape[2]) % Y_shape[1])
|
||||
i0 = int(i // Y_shape[3] // Y_shape[2] // Y_shape[1])
|
||||
|
||||
j0 = int((i0 - all_pads[0]) % orig_dims[0])
|
||||
j1 = int((i1 - all_pads[2]) % orig_dims[1])
|
||||
j2 = int((i2 - all_pads[4]) % orig_dims[2])
|
||||
j3 = int((i3 - all_pads[6]) % orig_dims[3])
|
||||
|
||||
Y[i] = X[
|
||||
int(
|
||||
orig_dims[3] * orig_dims[2] * orig_dims[1] * j0
|
||||
+ orig_dims[3] * orig_dims[2] * j1
|
||||
+ orig_dims[3] * j2
|
||||
+ j3
|
||||
)
|
||||
]
|
||||
|
||||
class CircPadPlugin(trt.IPluginV2DynamicExt):
|
||||
def __init__(self, fc=None):
|
||||
trt.IPluginV2DynamicExt.__init__(self)
|
||||
self.pads = []
|
||||
self.X_shape = []
|
||||
|
||||
self.num_outputs = 1
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_type = "CircPadPlugin"
|
||||
self.plugin_version = "1"
|
||||
|
||||
if fc is not None:
|
||||
assert fc[0].name == "pads"
|
||||
self.pads = fc[0].data
|
||||
|
||||
def get_output_datatype(self, index, input_types):
|
||||
return input_types[0]
|
||||
|
||||
def get_output_dimensions(self, output_index, inputs, exprBuilder):
|
||||
|
||||
output_dims = trt.DimsExprs(inputs[0])
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
|
||||
trt.DimensionOperation.SUM,
|
||||
inputs[0][len(output_dims) - i - 1],
|
||||
exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
|
||||
)
|
||||
|
||||
return output_dims
|
||||
|
||||
def serialize(self):
|
||||
return to_json({"pads": self.pads})
|
||||
|
||||
def configure_plugin(self, inp, out):
|
||||
X_dims = inp[0].desc.dims
|
||||
self.X_shape = np.zeros((len(X_dims),))
|
||||
for i in range(len(X_dims)):
|
||||
self.X_shape[i] = X_dims[i]
|
||||
|
||||
def supports_format_combination(self, pos, in_out, num_inputs):
|
||||
assert num_inputs == 1
|
||||
assert pos < len(in_out)
|
||||
|
||||
desc = in_out[pos]
|
||||
if desc.format != trt.TensorFormat.LINEAR:
|
||||
return False
|
||||
|
||||
# first input should be float16 or float32
|
||||
if pos == 0:
|
||||
return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
|
||||
|
||||
# output should have the same type as the input
|
||||
if pos == 1:
|
||||
return in_out[0].type == desc.type
|
||||
|
||||
assert False
|
||||
|
||||
def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
|
||||
|
||||
inp_dtype = trt.nptype(input_desc[0].type)
|
||||
|
||||
a_mem = cp.cuda.UnownedMemory(
|
||||
inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
|
||||
)
|
||||
c_mem = cp.cuda.UnownedMemory(
|
||||
outputs[0],
|
||||
volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
|
||||
self,
|
||||
)
|
||||
|
||||
a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
|
||||
c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
|
||||
|
||||
a = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
|
||||
c = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
|
||||
|
||||
numba_stream = cuda.external_stream(stream)
|
||||
|
||||
N = len(self.X_shape)
|
||||
all_pads = np.zeros((N * 2,))
|
||||
orig_dims = np.array(self.X_shape)
|
||||
out_dims = np.array(self.X_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_d = cp.asarray(all_pads)
|
||||
orig_dims_d = cp.asarray(orig_dims)
|
||||
Y_shape_d = cp.asarray(out_dims)
|
||||
|
||||
blockSize = 256
|
||||
numBlocks = int((np.prod(out_dims) + blockSize - 1) // blockSize)
|
||||
|
||||
circ_pad[numBlocks, blockSize, numba_stream](
|
||||
a, all_pads_d, orig_dims_d, c, Y_shape_d, np.prod(out_dims)
|
||||
)
|
||||
|
||||
return 0
|
||||
|
||||
def clone(self):
|
||||
cloned_plugin = CircPadPlugin()
|
||||
cloned_plugin.__dict__.update(self.__dict__)
|
||||
return cloned_plugin
|
||||
|
||||
#
|
||||
# The following defaults take effect since the respective methods are not overriden
|
||||
#
|
||||
|
||||
# def initialize(self):
|
||||
# pass
|
||||
|
||||
# def get_serialization_size(self):
|
||||
# return len(to_json({"pads": self.pads}))
|
||||
|
||||
# def get_workspace_size(self, input_desc, output_desc):
|
||||
# return 0
|
||||
|
||||
# def destroy(self):
|
||||
# pass
|
||||
|
||||
# def terminate(self):
|
||||
# pass
|
||||
|
||||
|
||||
class CircPadPluginCreator(trt.IPluginCreator):
|
||||
def __init__(self):
|
||||
trt.IPluginCreator.__init__(self)
|
||||
self.name = "CircPadPlugin"
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_version = "1"
|
||||
self.field_names = trt.PluginFieldCollection(
|
||||
[trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
|
||||
)
|
||||
|
||||
def create_plugin(self, name, fc):
|
||||
return CircPadPlugin(fc)
|
||||
|
||||
def deserialize_plugin(self, name, data):
|
||||
j = dict(from_json(data.decode("utf-8")))
|
||||
deserialized = CircPadPlugin()
|
||||
deserialized.__dict__.update(j)
|
||||
return deserialized
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
args = parseArgs()
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
|
||||
inp_shape = (10, 3, 32, 32)
|
||||
X = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
# Register plugin creator
|
||||
plg_registry = trt.get_plugin_registry()
|
||||
my_plugin_creator = CircPadPluginCreator()
|
||||
plg_registry.register_creator(my_plugin_creator, "")
|
||||
|
||||
# create ONNX model
|
||||
onnx_path = "test_CircPadPlugin.onnx"
|
||||
inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
|
||||
Y = gs.Variable(name="Y", dtype=precision)
|
||||
myPluginNode = gs.Node(
|
||||
name="CircPadPlugin",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputA],
|
||||
outputs=[Y],
|
||||
attrs={"pads": pads},
|
||||
)
|
||||
graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
# build engine
|
||||
build_engine = EngineFromNetwork(
|
||||
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
|
||||
)
|
||||
|
||||
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
# Run
|
||||
with TrtRunner(build_engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
|
||||
if np.allclose(Y, Y_ref):
|
||||
print("Inference result correct!")
|
||||
else:
|
||||
print("Inference result incorrect!")
|
||||
@@ -0,0 +1,208 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
import onnx
|
||||
import cupy as cp
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
EngineFromNetwork,
|
||||
NetworkFromOnnxPath,
|
||||
TrtRunner,
|
||||
)
|
||||
|
||||
from polygraphy.json import to_json, from_json
|
||||
import torch
|
||||
|
||||
from utils import volume, parseArgs
|
||||
|
||||
class CircPadPlugin(trt.IPluginV2DynamicExt):
|
||||
def __init__(self, fc=None):
|
||||
trt.IPluginV2DynamicExt.__init__(self)
|
||||
self.pads = []
|
||||
self.X_shape = []
|
||||
|
||||
self.num_outputs = 1
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_type = "CircPadPlugin"
|
||||
self.plugin_version = "1"
|
||||
|
||||
if fc is not None:
|
||||
assert fc[0].name == "pads"
|
||||
self.pads = fc[0].data
|
||||
|
||||
def get_output_datatype(self, index, input_types):
|
||||
return input_types[0]
|
||||
|
||||
def get_output_dimensions(self, output_index, inputs, exprBuilder):
|
||||
|
||||
output_dims = trt.DimsExprs(inputs[0])
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
|
||||
trt.DimensionOperation.SUM,
|
||||
inputs[0][len(output_dims) - i - 1],
|
||||
exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
|
||||
)
|
||||
|
||||
return output_dims
|
||||
|
||||
def serialize(self):
|
||||
return to_json({"pads": self.pads})
|
||||
|
||||
def configure_plugin(self, inp, out):
|
||||
X_dims = inp[0].desc.dims
|
||||
self.X_shape = np.zeros((len(X_dims),))
|
||||
for i in range(len(X_dims)):
|
||||
self.X_shape[i] = X_dims[i]
|
||||
|
||||
def supports_format_combination(self, pos, in_out, num_inputs):
|
||||
assert num_inputs == 1
|
||||
assert pos < len(in_out)
|
||||
|
||||
desc = in_out[pos]
|
||||
if desc.format != trt.TensorFormat.LINEAR:
|
||||
return False
|
||||
|
||||
# first input should be float16 or float32
|
||||
if pos == 0:
|
||||
return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
|
||||
|
||||
# output should have the same type as the input
|
||||
if pos == 1:
|
||||
return in_out[0].type == desc.type
|
||||
|
||||
assert False
|
||||
|
||||
def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
|
||||
|
||||
inp_dtype = trt.nptype(input_desc[0].type)
|
||||
|
||||
a_mem = cp.cuda.UnownedMemory(
|
||||
inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
|
||||
)
|
||||
c_mem = cp.cuda.UnownedMemory(
|
||||
outputs[0],
|
||||
volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
|
||||
self,
|
||||
)
|
||||
|
||||
a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
|
||||
c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
|
||||
|
||||
a_d = cp.ndarray(tuple(input_desc[0].dims), dtype=inp_dtype, memptr=a_ptr)
|
||||
c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
|
||||
|
||||
a_t = torch.as_tensor(a_d, device='cuda')
|
||||
|
||||
# Use PyTorch functional op - no need to write kernel
|
||||
out = torch.nn.functional.pad(a_t, self.pads.tolist(), mode='circular')
|
||||
cp.copyto(c_d, cp.reshape(cp.asarray(out), (-1,)))
|
||||
|
||||
return 0
|
||||
|
||||
def clone(self):
|
||||
cloned_plugin = CircPadPlugin()
|
||||
cloned_plugin.__dict__.update(self.__dict__)
|
||||
return cloned_plugin
|
||||
|
||||
#
|
||||
# The following defaults take effect since the respective methods are not overriden
|
||||
#
|
||||
|
||||
# def initialize(self):
|
||||
# pass
|
||||
|
||||
# def get_serialization_size(self):
|
||||
# return len(to_json({"pads": self.pads}))
|
||||
|
||||
# def get_workspace_size(self, input_desc, output_desc):
|
||||
# return 0
|
||||
|
||||
# def destroy(self):
|
||||
# pass
|
||||
|
||||
# def terminate(self):
|
||||
# pass
|
||||
|
||||
|
||||
class CircPadPluginCreator(trt.IPluginCreator):
|
||||
def __init__(self):
|
||||
trt.IPluginCreator.__init__(self)
|
||||
self.name = "CircPadPlugin"
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_version = "1"
|
||||
self.field_names = trt.PluginFieldCollection(
|
||||
[trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
|
||||
)
|
||||
|
||||
def create_plugin(self, name, fc):
|
||||
return CircPadPlugin(fc)
|
||||
|
||||
def deserialize_plugin(self, name, data):
|
||||
j = dict(from_json(data.decode("utf-8")))
|
||||
deserialized = CircPadPlugin()
|
||||
deserialized.__dict__.update(j)
|
||||
return deserialized
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
args = parseArgs()
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
|
||||
inp_shape = (10, 3, 32, 32)
|
||||
X = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
# Register plugin creator
|
||||
plg_registry = trt.get_plugin_registry()
|
||||
my_plugin_creator = CircPadPluginCreator()
|
||||
plg_registry.register_creator(my_plugin_creator, "")
|
||||
|
||||
# create ONNX model
|
||||
onnx_path = "test_CircPadPlugin.onnx"
|
||||
inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
|
||||
Y = gs.Variable(name="Y", dtype=precision)
|
||||
myPluginNode = gs.Node(
|
||||
name="CircPadPlugin",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputA],
|
||||
outputs=[Y],
|
||||
attrs={"pads": pads},
|
||||
)
|
||||
graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
# build engine
|
||||
build_engine = EngineFromNetwork(
|
||||
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
|
||||
)
|
||||
|
||||
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
# Run
|
||||
with TrtRunner(build_engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
|
||||
if np.allclose(Y, Y_ref):
|
||||
print("Inference result correct!")
|
||||
else:
|
||||
print("Inference result incorrect!")
|
||||
@@ -0,0 +1,263 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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 onnx_graphsurgeon as gs
|
||||
import numpy as np
|
||||
import onnx
|
||||
import cupy as cp
|
||||
|
||||
import triton
|
||||
import triton.language as tl
|
||||
|
||||
import tensorrt as trt
|
||||
from polygraphy.backend.trt import (
|
||||
CreateConfig,
|
||||
EngineFromNetwork,
|
||||
NetworkFromOnnxPath,
|
||||
TrtRunner,
|
||||
)
|
||||
|
||||
from polygraphy.json import to_json, from_json
|
||||
import torch
|
||||
|
||||
from utils import volume, parseArgs
|
||||
|
||||
@triton.jit
|
||||
def circ_pad(X,
|
||||
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,
|
||||
Y_shape_1, Y_shape_2, Y_shape_3,
|
||||
X_len, Y_len, 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)
|
||||
|
||||
|
||||
class CircPadPlugin(trt.IPluginV2DynamicExt):
|
||||
def __init__(self, fc=None):
|
||||
trt.IPluginV2DynamicExt.__init__(self)
|
||||
self.pads = []
|
||||
self.X_shape = []
|
||||
|
||||
self.num_outputs = 1
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_type = "CircPadPlugin"
|
||||
self.plugin_version = "1"
|
||||
|
||||
if fc is not None:
|
||||
assert fc[0].name == "pads"
|
||||
self.pads = fc[0].data
|
||||
|
||||
def get_output_datatype(self, index, input_types):
|
||||
return input_types[0]
|
||||
|
||||
def get_output_dimensions(self, output_index, inputs, exprBuilder):
|
||||
|
||||
output_dims = trt.DimsExprs(inputs[0])
|
||||
|
||||
for i in range(np.size(self.pads) // 2):
|
||||
output_dims[len(output_dims) - i - 1] = exprBuilder.operation(
|
||||
trt.DimensionOperation.SUM,
|
||||
inputs[0][len(output_dims) - i - 1],
|
||||
exprBuilder.constant(self.pads[i * 2] + self.pads[i * 2 + 1]),
|
||||
)
|
||||
|
||||
return output_dims
|
||||
|
||||
def serialize(self):
|
||||
return to_json({"pads": self.pads})
|
||||
|
||||
def configure_plugin(self, inp, out):
|
||||
X_dims = inp[0].desc.dims
|
||||
self.X_shape = np.zeros((len(X_dims),))
|
||||
for i in range(len(X_dims)):
|
||||
self.X_shape[i] = X_dims[i]
|
||||
|
||||
def supports_format_combination(self, pos, in_out, num_inputs):
|
||||
assert num_inputs == 1
|
||||
assert pos < len(in_out)
|
||||
|
||||
desc = in_out[pos]
|
||||
if desc.format != trt.TensorFormat.LINEAR:
|
||||
return False
|
||||
|
||||
# first input should be float16 or float32
|
||||
if pos == 0:
|
||||
return desc.type == trt.DataType.FLOAT or desc.type == trt.DataType.HALF
|
||||
|
||||
# output should have the same type as the input
|
||||
if pos == 1:
|
||||
return in_out[0].type == desc.type
|
||||
|
||||
assert False
|
||||
|
||||
def enqueue(self, input_desc, output_desc, inputs, outputs, workspace, stream):
|
||||
|
||||
inp_dtype = trt.nptype(input_desc[0].type)
|
||||
|
||||
a_mem = cp.cuda.UnownedMemory(
|
||||
inputs[0], volume(input_desc[0].dims) * cp.dtype(inp_dtype).itemsize, self
|
||||
)
|
||||
c_mem = cp.cuda.UnownedMemory(
|
||||
outputs[0],
|
||||
volume(output_desc[0].dims) * cp.dtype(inp_dtype).itemsize,
|
||||
self,
|
||||
)
|
||||
|
||||
a_ptr = cp.cuda.MemoryPointer(a_mem, 0)
|
||||
c_ptr = cp.cuda.MemoryPointer(c_mem, 0)
|
||||
|
||||
a_d = cp.ndarray((volume(input_desc[0].dims)), dtype=inp_dtype, memptr=a_ptr)
|
||||
c_d = cp.ndarray((volume(output_desc[0].dims)), dtype=inp_dtype, memptr=c_ptr)
|
||||
|
||||
a_t = torch.as_tensor(a_d, device='cuda')
|
||||
c_t = torch.as_tensor(c_d, device='cuda')
|
||||
|
||||
N = len(self.X_shape)
|
||||
all_pads = np.zeros((N * 2,), dtype=np.int32)
|
||||
orig_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
out_dims = np.array(self.X_shape, dtype=np.int32)
|
||||
|
||||
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()
|
||||
orig_dims = orig_dims.tolist()
|
||||
out_dims = out_dims.tolist()
|
||||
|
||||
blockSize = 256
|
||||
numBlocks = (int((np.prod(out_dims) + blockSize - 1) // blockSize),)
|
||||
|
||||
circ_pad[numBlocks](a_t,
|
||||
all_pads[0], all_pads[2], all_pads[4], all_pads[6],
|
||||
orig_dims[0], orig_dims[1], orig_dims[2], orig_dims[3],
|
||||
c_t,
|
||||
out_dims[1], out_dims[2], out_dims[3],
|
||||
int(np.prod(orig_dims)), int(np.prod(out_dims)), BLOCK_SIZE=256
|
||||
)
|
||||
|
||||
return 0
|
||||
|
||||
def clone(self):
|
||||
cloned_plugin = CircPadPlugin()
|
||||
cloned_plugin.__dict__.update(self.__dict__)
|
||||
return cloned_plugin
|
||||
|
||||
#
|
||||
# The following defaults take effect since the respective methods are not overriden
|
||||
#
|
||||
|
||||
# def initialize(self):
|
||||
# pass
|
||||
|
||||
# def get_serialization_size(self):
|
||||
# return len(to_json({"pads": self.pads}))
|
||||
|
||||
# def get_workspace_size(self, input_desc, output_desc):
|
||||
# return 0
|
||||
|
||||
# def destroy(self):
|
||||
# pass
|
||||
|
||||
# def terminate(self):
|
||||
# pass
|
||||
|
||||
|
||||
class CircPadPluginCreator(trt.IPluginCreator):
|
||||
def __init__(self):
|
||||
trt.IPluginCreator.__init__(self)
|
||||
self.name = "CircPadPlugin"
|
||||
self.plugin_namespace = ""
|
||||
self.plugin_version = "1"
|
||||
self.field_names = trt.PluginFieldCollection(
|
||||
[trt.PluginField("pads", np.array([]), trt.PluginFieldType.INT32)]
|
||||
)
|
||||
|
||||
def create_plugin(self, name, fc):
|
||||
return CircPadPlugin(fc)
|
||||
|
||||
def deserialize_plugin(self, name, data):
|
||||
j = dict(from_json(data.decode("utf-8")))
|
||||
deserialized = CircPadPlugin()
|
||||
deserialized.__dict__.update(j)
|
||||
return deserialized
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
args = parseArgs()
|
||||
precision = np.float32 if args.precision == "fp32" else np.float16
|
||||
|
||||
inp_shape = (10, 3, 32, 32)
|
||||
X = np.random.normal(size=inp_shape).astype(precision)
|
||||
|
||||
pads = (1, 1, 1, 1)
|
||||
|
||||
# Register plugin creator
|
||||
plg_registry = trt.get_plugin_registry()
|
||||
my_plugin_creator = CircPadPluginCreator()
|
||||
plg_registry.register_creator(my_plugin_creator, "")
|
||||
|
||||
# create ONNX model
|
||||
onnx_path = "test_CircPadPlugin.onnx"
|
||||
inputA = gs.Variable(name="X", shape=inp_shape, dtype=precision)
|
||||
Y = gs.Variable(name="Y", dtype=precision)
|
||||
myPluginNode = gs.Node(
|
||||
name="CircPadPlugin",
|
||||
op="CircPadPlugin",
|
||||
inputs=[inputA],
|
||||
outputs=[Y],
|
||||
attrs={"pads": pads},
|
||||
)
|
||||
graph = gs.Graph(nodes=[myPluginNode], inputs=[inputA], outputs=[Y], opset=16)
|
||||
onnx.save(gs.export_onnx(graph), onnx_path)
|
||||
|
||||
# build engine
|
||||
build_engine = EngineFromNetwork(
|
||||
NetworkFromOnnxPath(onnx_path), CreateConfig(fp16=precision==np.float16)
|
||||
)
|
||||
|
||||
Y_ref = np.pad(X, [[0, 0], [0, 0], [pads[0], pads[1]], [pads[2], pads[3]]], "wrap")
|
||||
# Run
|
||||
with TrtRunner(build_engine, "trt_runner")as runner:
|
||||
outputs = runner.infer({"X": X})
|
||||
Y = outputs["Y"]
|
||||
|
||||
if np.allclose(Y, Y_ref):
|
||||
print("Inference result correct!")
|
||||
else:
|
||||
print("Inference result incorrect!")
|
||||
@@ -0,0 +1,415 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2023 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.
|
||||
*/
|
||||
|
||||
#include "NvInfer.h"
|
||||
|
||||
#include <algorithm>
|
||||
#include <cstdint>
|
||||
#include <cstring>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <numeric>
|
||||
#include <vector>
|
||||
|
||||
#include <cuda.h>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
static void caughtError(std::exception const& e)
|
||||
{
|
||||
std::cout << e.what() << std::endl;
|
||||
}
|
||||
|
||||
// Write values into buffer
|
||||
template <typename T>
|
||||
void write(char*& buffer, T const& val)
|
||||
{
|
||||
std::memcpy(buffer, &val, sizeof(T));
|
||||
buffer += sizeof(T);
|
||||
}
|
||||
|
||||
// Read values from buffer
|
||||
template <typename T>
|
||||
T read(char const*& buffer)
|
||||
{
|
||||
T val{};
|
||||
std::memcpy(&val, buffer, sizeof(T));
|
||||
buffer += sizeof(T);
|
||||
return val;
|
||||
}
|
||||
|
||||
#define ASSERT(condition) \
|
||||
do \
|
||||
{ \
|
||||
if (!(condition)) \
|
||||
{ \
|
||||
std::cout << "Assertion failure: " << #condition << std::endl; \
|
||||
abort(); \
|
||||
} \
|
||||
} while (0)
|
||||
|
||||
template <typename Dtype>
|
||||
struct CudaBind
|
||||
{
|
||||
size_t mSize;
|
||||
Dtype* mPtr;
|
||||
|
||||
CudaBind(size_t size)
|
||||
{
|
||||
mSize = size;
|
||||
ASSERT(!cudaMalloc((void**) &mPtr, sizeof(Dtype) * mSize));
|
||||
}
|
||||
|
||||
~CudaBind()
|
||||
{
|
||||
if (mPtr != nullptr)
|
||||
{
|
||||
ASSERT(!cudaFree(mPtr));
|
||||
mPtr = nullptr;
|
||||
}
|
||||
}
|
||||
};
|
||||
|
||||
static int64_t volume(Dims const& dims)
|
||||
{
|
||||
return std::accumulate(dims.d, dims.d + dims.nbDims, int64_t{1}, std::multiplies<int64_t>{});
|
||||
}
|
||||
|
||||
template <typename T>
|
||||
__global__ void circPadKernel(
|
||||
T const* x, int32_t const* allPads, int32_t const* origDims, T* y, int32_t const* yShape, int32_t yLen)
|
||||
{
|
||||
int32_t index = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
int32_t stride = blockDim.x * gridDim.x;
|
||||
|
||||
for (int32_t i = index; i < yLen; i += stride)
|
||||
{
|
||||
int32_t i3 = i % yShape[3];
|
||||
int32_t i2 = (i / yShape[3]) % yShape[2];
|
||||
int32_t i1 = (i / yShape[3] / yShape[2]) % yShape[1];
|
||||
int32_t i0 = i / yShape[3] / yShape[2] / yShape[1];
|
||||
|
||||
int32_t j0 = (i0 - allPads[0] + origDims[0]) % origDims[0];
|
||||
int32_t j1 = (i1 - allPads[2] + origDims[1]) % origDims[1];
|
||||
int32_t j2 = (i2 - allPads[4] + origDims[2]) % origDims[2];
|
||||
int32_t j3 = (i3 - allPads[6] + origDims[3]) % origDims[3];
|
||||
|
||||
y[i] = x[origDims[3] * origDims[2] * origDims[1] * j0 + origDims[3] * origDims[2] * j1 + origDims[3] * j2
|
||||
+ j3];
|
||||
}
|
||||
}
|
||||
|
||||
class CircPadPlugin : public nvinfer1::IPluginV2DynamicExt
|
||||
{
|
||||
public:
|
||||
CircPadPlugin() = default;
|
||||
|
||||
CircPadPlugin(std::vector<int32_t> pads)
|
||||
: mPads(pads)
|
||||
{
|
||||
}
|
||||
|
||||
CircPadPlugin(CircPadPlugin const& p) = default;
|
||||
|
||||
CircPadPlugin(void const* serialData, size_t length)
|
||||
{
|
||||
ASSERT(serialData != nullptr);
|
||||
|
||||
char const* d = static_cast<char const*>(serialData);
|
||||
char const* a = d;
|
||||
|
||||
int32_t padsSize = read<int32_t>(d);
|
||||
mPads.resize(padsSize);
|
||||
for (int i = 0; i < padsSize; ++i)
|
||||
{
|
||||
mPads[i] = read<int32_t>(d);
|
||||
}
|
||||
|
||||
ASSERT(d == a + length);
|
||||
}
|
||||
|
||||
int32_t getNbOutputs() const noexcept override
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
PluginTensorDesc const& desc = inOut[pos];
|
||||
if (desc.format != TensorFormat::kLINEAR)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
// first input should be float16 or float32
|
||||
if (pos == 0)
|
||||
{
|
||||
return (inOut[pos].type == nvinfer1::DataType::kFLOAT || inOut[pos].type == nvinfer1::DataType::kHALF);
|
||||
}
|
||||
|
||||
// output should have the same type as the input
|
||||
if (pos == 1)
|
||||
{
|
||||
return (inOut[pos].type == inOut[0].type);
|
||||
}
|
||||
|
||||
return false;
|
||||
}
|
||||
|
||||
void configureWithFormat(nvinfer1::Dims const*, int32_t, nvinfer1::Dims const*, int32_t, nvinfer1::DataType type,
|
||||
nvinfer1::PluginFormat floatFormat, int32_t) noexcept override
|
||||
{
|
||||
}
|
||||
|
||||
int32_t initialize() noexcept override
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
void terminate() noexcept override
|
||||
{
|
||||
mAllPadsPtr.reset();
|
||||
mOrigDimsPtr.reset();
|
||||
mOutDimsPtr.reset();
|
||||
}
|
||||
|
||||
int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs,
|
||||
void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
||||
{
|
||||
auto inpDType = inputDesc[0].type;
|
||||
|
||||
int32_t const blockSize = 256;
|
||||
int32_t const numBlocks = (volume(outputDesc[0].dims) + blockSize - 1) / blockSize;
|
||||
|
||||
ASSERT(inpDType == DataType::kFLOAT || inpDType == DataType::kHALF);
|
||||
|
||||
if (inpDType == DataType::kFLOAT)
|
||||
{
|
||||
circPadKernel<float><<<numBlocks, blockSize, 0, stream>>>(static_cast<float const*>(inputs[0]),
|
||||
mAllPadsPtr->mPtr, mOrigDimsPtr->mPtr, static_cast<float*>(outputs[0]), mOutDimsPtr->mPtr,
|
||||
volume(outputDesc[0].dims));
|
||||
}
|
||||
else if (inpDType == DataType::kHALF)
|
||||
{
|
||||
circPadKernel<half><<<numBlocks, blockSize, 0, stream>>>(static_cast<half const*>(inputs[0]),
|
||||
mAllPadsPtr->mPtr, mOrigDimsPtr->mPtr, static_cast<half*>(outputs[0]), mOutDimsPtr->mPtr,
|
||||
volume(outputDesc[0].dims));
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
size_t getSerializationSize() const noexcept override
|
||||
{
|
||||
return (mPads.size() + 1) * sizeof(int32_t);
|
||||
}
|
||||
|
||||
void serialize(void* buffer) const noexcept override
|
||||
{
|
||||
ASSERT(buffer != nullptr);
|
||||
char* d = static_cast<char*>(buffer);
|
||||
char* a = d;
|
||||
write(d, static_cast<int32_t>(mPads.size()));
|
||||
for (int i = 0; i < mPads.size(); ++i)
|
||||
{
|
||||
write(d, mPads[i]);
|
||||
}
|
||||
ASSERT(d == a + getSerializationSize());
|
||||
}
|
||||
|
||||
char const* getPluginType() const noexcept override
|
||||
{
|
||||
return "CircPadPlugin";
|
||||
}
|
||||
|
||||
char const* getPluginVersion() const noexcept override
|
||||
{
|
||||
return "1";
|
||||
}
|
||||
|
||||
nvinfer1::IPluginV2DynamicExt* clone() const noexcept override
|
||||
{
|
||||
return new CircPadPlugin(*this);
|
||||
}
|
||||
|
||||
void destroy() noexcept override
|
||||
{
|
||||
delete this;
|
||||
}
|
||||
|
||||
void setPluginNamespace(char const* libNamespace) noexcept override
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
|
||||
char const* getPluginNamespace() const noexcept override
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
DataType getOutputDataType(int index, nvinfer1::DataType const* inputTypes, int nbInputs) const noexcept
|
||||
{
|
||||
return inputTypes[0];
|
||||
}
|
||||
|
||||
DimsExprs getOutputDimensions(
|
||||
int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
nvinfer1::DimsExprs outDims{inputs[0]};
|
||||
int32_t nbOutDims = inputs[0].nbDims;
|
||||
|
||||
for (int32_t i = 0; i < mPads.size() / 2; ++i)
|
||||
{
|
||||
outDims.d[nbOutDims - i - 1] = exprBuilder.operation(nvinfer1::DimensionOperation::kSUM,
|
||||
*inputs[0].d[nbOutDims - i - 1], *exprBuilder.constant(mPads[i * 2] + mPads[i * 2 + 1]));
|
||||
}
|
||||
|
||||
return outDims;
|
||||
}
|
||||
|
||||
void configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out,
|
||||
int32_t nbOutputs) noexcept
|
||||
{
|
||||
mN = in[0].desc.dims.nbDims;
|
||||
|
||||
std::vector<int32_t> allPads(mN * 2);
|
||||
std::vector<int32_t> origDims(mN);
|
||||
std::vector<int32_t> outDims(mN);
|
||||
|
||||
for (int32_t i = 0; i < mN; ++i)
|
||||
{
|
||||
origDims[i] = in[0].desc.dims.d[i];
|
||||
outDims[i] = in[0].desc.dims.d[i];
|
||||
}
|
||||
|
||||
for (int32_t i = 0; i < mPads.size() / 2; ++i)
|
||||
{
|
||||
outDims[mN - i - 1] += mPads[i * 2] + mPads[i * 2 + 1];
|
||||
allPads[mN * 2 - 2 * i - 2] = mPads[i * 2];
|
||||
allPads[mN * 2 - 2 * i - 1] = mPads[i * 2 + 1];
|
||||
}
|
||||
|
||||
mAllPadsPtr = std::make_shared<CudaBind<int32_t>>(mN * 2);
|
||||
mOrigDimsPtr = std::make_shared<CudaBind<int32_t>>(mN);
|
||||
mOutDimsPtr = std::make_shared<CudaBind<int32_t>>(mN);
|
||||
|
||||
ASSERT(
|
||||
!cudaMemcpy(mAllPadsPtr->mPtr, &allPads.front(), allPads.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
|
||||
ASSERT(!cudaMemcpy(
|
||||
mOrigDimsPtr->mPtr, &origDims.front(), origDims.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
|
||||
ASSERT(
|
||||
!cudaMemcpy(mOutDimsPtr->mPtr, &outDims.front(), outDims.size() * sizeof(int32_t), cudaMemcpyHostToDevice));
|
||||
}
|
||||
|
||||
size_t getWorkspaceSize(PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs,
|
||||
int32_t nbOutputs) const noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
private:
|
||||
std::vector<int32_t> mPads{};
|
||||
int32_t mN{};
|
||||
std::shared_ptr<CudaBind<int32_t>> mAllPadsPtr{};
|
||||
std::shared_ptr<CudaBind<int32_t>> mOrigDimsPtr{};
|
||||
std::shared_ptr<CudaBind<int32_t>> mOutDimsPtr{};
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
class CircPadPluginCreator : public nvinfer1::IPluginCreator
|
||||
{
|
||||
public:
|
||||
CircPadPluginCreator()
|
||||
{
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(PluginField("pads", nullptr, PluginFieldType::kINT32, 1));
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* getPluginName() const noexcept
|
||||
{
|
||||
return "CircPadPlugin";
|
||||
}
|
||||
|
||||
char const* getPluginVersion() const noexcept
|
||||
{
|
||||
return "1";
|
||||
}
|
||||
|
||||
PluginFieldCollection const* getFieldNames() noexcept
|
||||
{
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
IPluginV2* createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
std::vector<int32_t> pads;
|
||||
|
||||
for (int32_t i = 0; i < fc->nbFields; i++)
|
||||
{
|
||||
std::string field_name(fc->fields[i].name);
|
||||
if (field_name.compare("pads") == 0)
|
||||
{
|
||||
pads.resize(fc->fields[i].length);
|
||||
auto const* padsPtr = static_cast<int32_t const*>(fc->fields[i].data);
|
||||
std::copy_n(padsPtr, fc->fields[i].length, pads.data());
|
||||
}
|
||||
}
|
||||
|
||||
return new CircPadPlugin(pads);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2* deserializePlugin(char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
return new CircPadPlugin(serialData, serialLength);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
void setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
|
||||
char const* getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
private:
|
||||
nvinfer1::PluginFieldCollection mFC;
|
||||
std::vector<nvinfer1::PluginField> mPluginAttributes;
|
||||
std::string mNamespace;
|
||||
};
|
||||
|
||||
REGISTER_TENSORRT_PLUGIN(CircPadPluginCreator);
|
||||
@@ -0,0 +1,15 @@
|
||||
cuda-python
|
||||
cupy-cuda12x
|
||||
numba
|
||||
triton; platform_system != "Windows"
|
||||
torch
|
||||
--extra-index-url https://pypi.ngc.nvidia.com
|
||||
polygraphy
|
||||
colored
|
||||
numpy==1.23.5; platform_system != "Windows"
|
||||
--extra-index-url https://pypi.ngc.nvidia.com
|
||||
onnx-graphsurgeon
|
||||
pywin32; platform_system == "Windows"
|
||||
pyyaml==6.0.1
|
||||
requests==2.31.0
|
||||
tqdm==4.66.1
|
||||
@@ -0,0 +1,28 @@
|
||||
---
|
||||
args:
|
||||
polygraphy:
|
||||
- '--extra-index-url https://pypi.ngc.nvidia.com'
|
||||
torch: []
|
||||
conditions:
|
||||
cuda-python:
|
||||
- cuda-python
|
||||
onnx-graphsurgeon:
|
||||
- onnx-graphsurgeon
|
||||
triton:
|
||||
- triton; platform_system != "Windows"
|
||||
numpy:
|
||||
- numpy==1.23.5; platform_system != "Windows"
|
||||
torch:
|
||||
- torch
|
||||
packages:
|
||||
- cuda-python
|
||||
- cupy-cuda12x
|
||||
- numba
|
||||
- triton
|
||||
- torch
|
||||
- polygraphy
|
||||
- colored
|
||||
- numpy
|
||||
- onnx-graphsurgeon
|
||||
- pywin32
|
||||
...
|
||||
@@ -0,0 +1,118 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 1993-2023 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.
|
||||
#
|
||||
|
||||
from cuda import cuda, cudart, nvrtc
|
||||
import numpy as np
|
||||
import os
|
||||
import argparse
|
||||
import threading
|
||||
|
||||
import tensorrt as trt
|
||||
|
||||
def parseArgs():
|
||||
parser = argparse.ArgumentParser(description="Options for Circular Padding plugin C++ example")
|
||||
parser.add_argument('--precision', type=str, default="fp32", choices=["fp32", "fp16"], help="Precision to use for plugin")
|
||||
|
||||
return parser.parse_args()
|
||||
|
||||
def volume(d):
|
||||
return np.prod(d)
|
||||
|
||||
# Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/helper_cuda.py
|
||||
def checkCudaErrors(result):
|
||||
def _cudaGetErrorEnum(error):
|
||||
if isinstance(error, cuda.CUresult):
|
||||
err, name = cuda.cuGetErrorName(error)
|
||||
return name if err == cuda.CUresult.CUDA_SUCCESS else "<unknown>"
|
||||
elif isinstance(error, cudart.cudaError_t):
|
||||
return cudart.cudaGetErrorName(error)[1]
|
||||
elif isinstance(error, nvrtc.nvrtcResult):
|
||||
return nvrtc.nvrtcGetErrorString(error)[1]
|
||||
else:
|
||||
raise RuntimeError('Unknown error type: {}'.format(error))
|
||||
if result[0].value:
|
||||
raise RuntimeError("CUDA error code={}({})".format(result[0].value, _cudaGetErrorEnum(result[0])))
|
||||
if len(result) == 1:
|
||||
return None
|
||||
elif len(result) == 2:
|
||||
return result[1]
|
||||
else:
|
||||
return result[1:]
|
||||
|
||||
# Taken from https://github.com/NVIDIA/cuda-python/blob/main/examples/common/common.py
|
||||
class KernelHelper:
|
||||
def __init__(self, code, devID):
|
||||
prog = checkCudaErrors(nvrtc.nvrtcCreateProgram(str.encode(code), b'sourceCode.cu', 0, [], []))
|
||||
CUDA_HOME = os.getenv('CUDA_HOME')
|
||||
if CUDA_HOME == None:
|
||||
CUDA_HOME = os.getenv('CUDA_PATH')
|
||||
if CUDA_HOME == None:
|
||||
raise RuntimeError('Environment variable CUDA_HOME or CUDA_PATH is not set')
|
||||
include_dirs = os.path.join(CUDA_HOME, 'include')
|
||||
|
||||
# Initialize CUDA
|
||||
checkCudaErrors(cudart.cudaFree(0))
|
||||
|
||||
major = checkCudaErrors(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMajor, devID))
|
||||
minor = checkCudaErrors(cudart.cudaDeviceGetAttribute(cudart.cudaDeviceAttr.cudaDevAttrComputeCapabilityMinor, devID))
|
||||
_, nvrtc_minor = checkCudaErrors(nvrtc.nvrtcVersion())
|
||||
use_cubin = (nvrtc_minor >= 1)
|
||||
prefix = 'sm' if use_cubin else 'compute'
|
||||
arch_arg = bytes(f'--gpu-architecture={prefix}_{major}{minor}', 'ascii')
|
||||
|
||||
try:
|
||||
opts = [b'--fmad=true', arch_arg, '--include-path={}'.format(include_dirs).encode('UTF-8'),
|
||||
b'--std=c++11', b'-default-device']
|
||||
checkCudaErrors(nvrtc.nvrtcCompileProgram(prog, len(opts), opts))
|
||||
except RuntimeError as err:
|
||||
logSize = checkCudaErrors(nvrtc.nvrtcGetProgramLogSize(prog))
|
||||
log = b' ' * logSize
|
||||
checkCudaErrors(nvrtc.nvrtcGetProgramLog(prog, log))
|
||||
print(log.decode())
|
||||
print(err)
|
||||
exit(-1)
|
||||
|
||||
if use_cubin:
|
||||
dataSize = checkCudaErrors(nvrtc.nvrtcGetCUBINSize(prog))
|
||||
data = b' ' * dataSize
|
||||
checkCudaErrors(nvrtc.nvrtcGetCUBIN(prog, data))
|
||||
else:
|
||||
dataSize = checkCudaErrors(nvrtc.nvrtcGetPTXSize(prog))
|
||||
data = b' ' * dataSize
|
||||
checkCudaErrors(nvrtc.nvrtcGetPTX(prog, data))
|
||||
|
||||
self.module = checkCudaErrors(cuda.cuModuleLoadData(np.char.array(data)))
|
||||
|
||||
def getFunction(self, name):
|
||||
return checkCudaErrors(cuda.cuModuleGetFunction(self.module, name))
|
||||
|
||||
class CudaCtxManager(trt.IPluginResource):
|
||||
def __init__(self, device = None):
|
||||
trt.IPluginResource.__init__(self)
|
||||
self.device = device
|
||||
self.cuda_ctx = None
|
||||
|
||||
def clone(self):
|
||||
cloned = CudaCtxManager()
|
||||
cloned.__dict__.update(self.__dict__)
|
||||
# Delay the CUDA ctx creation until clone()
|
||||
# since only a cloned resource is registered by TRT
|
||||
_, cloned.cuda_ctx = cuda.cuCtxCreate(0, self.device)
|
||||
return cloned
|
||||
|
||||
def release(self):
|
||||
checkCudaErrors(cuda.cuCtxDestroy(self.cuda_ctx))
|
||||
Reference in New Issue
Block a user