chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,182 @@
|
||||
#
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
#
|
||||
|
||||
if(NOT TRT_SAFETY_INFERENCE_ONLY)
|
||||
add_executable(sample_plugin_safe_build
|
||||
sampleSafePluginBuild.cpp
|
||||
maxPoolPlugin.cpp
|
||||
maxPoolPluginCreator.cpp
|
||||
maxPoolPluginRuntime.cpp
|
||||
maxPoolPluginRuntimeCreator.cpp
|
||||
maxPoolKernel.cu
|
||||
)
|
||||
target_link_libraries(sample_plugin_safe_build PRIVATE
|
||||
trt_samples_common
|
||||
TRT_SAMPLES::tensorrt
|
||||
TRTSAFE::nvinfer_safe_shared
|
||||
)
|
||||
# Link ONNX parser if available
|
||||
if(TARGET nvonnxparser)
|
||||
target_link_libraries(sample_plugin_safe_build PRIVATE nvonnxparser)
|
||||
endif()
|
||||
if(TRT_OUT_DIR)
|
||||
set_target_properties(sample_plugin_safe_build
|
||||
PROPERTIES
|
||||
RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
|
||||
)
|
||||
endif()
|
||||
|
||||
add_dependencies(tensorrt_samples sample_plugin_safe_build)
|
||||
|
||||
installLibraries(
|
||||
TARGETS sample_plugin_safe_build
|
||||
OPTIONAL
|
||||
COMPONENT internal
|
||||
)
|
||||
endif()
|
||||
|
||||
add_executable(sample_plugin_safe_infer
|
||||
sampleSafePluginInfer.cpp
|
||||
maxPoolPluginRuntime.cpp
|
||||
maxPoolPluginRuntimeCreator.cpp
|
||||
maxPoolKernel.cu
|
||||
)
|
||||
|
||||
if(TRT_SAFETY_INFERENCE_ONLY)
|
||||
target_link_libraries(sample_plugin_safe_infer PRIVATE trt_global_definitions tensorrt_headers)
|
||||
target_include_directories(sample_plugin_safe_infer PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/../common
|
||||
)
|
||||
else()
|
||||
target_link_libraries(sample_plugin_safe_infer PRIVATE
|
||||
trt_samples_common
|
||||
TRTSAFE::nvinfer_safe_shared
|
||||
)
|
||||
endif()
|
||||
|
||||
if(TRT_OUT_DIR)
|
||||
set_target_properties(sample_plugin_safe_infer
|
||||
PROPERTIES
|
||||
RUNTIME_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
|
||||
)
|
||||
endif()
|
||||
|
||||
add_dependencies(tensorrt_samples sample_plugin_safe_infer)
|
||||
|
||||
installLibraries(
|
||||
TARGETS sample_plugin_safe_infer
|
||||
OPTIONAL
|
||||
COMPONENT internal
|
||||
)
|
||||
|
||||
# ==============================================================================
|
||||
# PLUGIN LIBRARIES FOR TRTEXEC INTEGRATION
|
||||
# ==============================================================================
|
||||
|
||||
if(NOT TRT_SAFETY_INFERENCE_ONLY)
|
||||
# Builder Plugin Library - for trtexec --safe (includes all plugin components)
|
||||
add_library(sample_safe_plugin_build_lib SHARED
|
||||
maxPoolPlugin.cpp
|
||||
maxPoolPluginRuntime.cpp
|
||||
maxPoolPluginCreator.cpp
|
||||
maxPoolPluginRuntimeCreator.cpp
|
||||
maxPoolPluginCreatorInterface.cpp
|
||||
maxPoolKernel.cu
|
||||
)
|
||||
|
||||
target_compile_definitions(sample_safe_plugin_build_lib PRIVATE
|
||||
GEN_PLUGIN_LIB=1 # Enable getSafetyPluginCreator export
|
||||
)
|
||||
|
||||
target_link_libraries(sample_safe_plugin_build_lib PRIVATE
|
||||
trt_samples_common
|
||||
TRT_SAMPLES::tensorrt
|
||||
TRTSAFE::nvinfer_safe_shared
|
||||
tensorrt_headers
|
||||
)
|
||||
|
||||
# Link ONNX parser if available
|
||||
if(TARGET nvonnxparser)
|
||||
target_link_libraries(sample_safe_plugin_build_lib PRIVATE nvonnxparser)
|
||||
endif()
|
||||
|
||||
target_include_directories(sample_safe_plugin_build_lib PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/../common
|
||||
)
|
||||
|
||||
if(TRT_OUT_DIR)
|
||||
set_target_properties(sample_safe_plugin_build_lib PROPERTIES
|
||||
OUTPUT_NAME "sample_safe_plugin_v3"
|
||||
LIBRARY_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
|
||||
)
|
||||
else()
|
||||
set_target_properties(sample_safe_plugin_build_lib PROPERTIES
|
||||
OUTPUT_NAME "sample_safe_plugin_v3"
|
||||
)
|
||||
endif()
|
||||
|
||||
add_dependencies(tensorrt_samples sample_safe_plugin_build_lib)
|
||||
|
||||
installLibraries(
|
||||
TARGETS sample_safe_plugin_build_lib
|
||||
OPTIONAL
|
||||
COMPONENT internal
|
||||
)
|
||||
endif()
|
||||
|
||||
# Safe Runtime Plugin Library - for trtexec_safe (runtime components only)
|
||||
add_library(sample_safe_plugin_runtime_lib SHARED
|
||||
maxPoolPluginRuntime.cpp
|
||||
maxPoolPluginRuntimeCreator.cpp
|
||||
maxPoolPluginRuntimeCreatorInterface.cpp
|
||||
maxPoolKernel.cu
|
||||
)
|
||||
|
||||
target_compile_definitions(sample_safe_plugin_runtime_lib PRIVATE
|
||||
GEN_PLUGIN_LIB=1 # Enable getSafetyPluginCreator export
|
||||
)
|
||||
|
||||
if(TRT_SAFETY_INFERENCE_ONLY)
|
||||
target_link_libraries(sample_safe_plugin_runtime_lib PRIVATE trt_global_definitions tensorrt_headers)
|
||||
target_include_directories(sample_safe_plugin_runtime_lib PRIVATE
|
||||
${CMAKE_CURRENT_SOURCE_DIR}/../common
|
||||
)
|
||||
else()
|
||||
target_link_libraries(sample_safe_plugin_runtime_lib PRIVATE
|
||||
trt_samples_common
|
||||
TRTSAFE::nvinfer_safe_shared
|
||||
)
|
||||
endif()
|
||||
|
||||
if(TRT_OUT_DIR)
|
||||
set_target_properties(sample_safe_plugin_runtime_lib PROPERTIES
|
||||
OUTPUT_NAME "sample_safe_plugin_v3_safe"
|
||||
LIBRARY_OUTPUT_DIRECTORY "${TRT_OUT_DIR}"
|
||||
)
|
||||
else()
|
||||
set_target_properties(sample_safe_plugin_runtime_lib PROPERTIES
|
||||
OUTPUT_NAME "sample_safe_plugin_v3_safe"
|
||||
)
|
||||
endif()
|
||||
|
||||
add_dependencies(tensorrt_samples sample_safe_plugin_runtime_lib)
|
||||
|
||||
installLibraries(
|
||||
TARGETS sample_safe_plugin_runtime_lib
|
||||
OPTIONAL
|
||||
COMPONENT internal
|
||||
)
|
||||
Executable
+270
@@ -0,0 +1,270 @@
|
||||
# Plugin Sample For TensorRT Safety
|
||||
|
||||
|
||||
**Table Of Contents**
|
||||
- [Description](#description)
|
||||
- [Workflow for safety plugin deployment](#workflow-for-safety-plugin-deployment)
|
||||
- [How does this sample work?](#how-does-this-sample-work)
|
||||
* [Register safe plugin creator](#register-safe-plugin-creator)
|
||||
- [Prerequisites](#prerequisites)
|
||||
- [Running the sample](#running-the-sample)
|
||||
* [Tool command line arguments](#tool-command-line-arguments)
|
||||
* [When to use remoteAutoTuningConfig](#when-to-use-remoteautotuningconfig)
|
||||
- [Shared plugin library and trtexec](#shared-plugin-library-and-trtexec)
|
||||
* [Creation of safety plugin library](#creation-of-safety-plugin-library)
|
||||
* [Using the plugin library with trtexec and trtexec\_safe](#using-the-plugin-library-with-trtexec-and-trtexec_safe)
|
||||
- [Additional resources](#additional-resources)
|
||||
- [License](#license)
|
||||
- [Changelog](#changelog)
|
||||
- [Known issues](#known-issues)
|
||||
|
||||
## Description
|
||||
|
||||
This sample, sampleSafePluginV3, consists of two parts: build and infer. The build part of this sample demonstrates how to use the builder in safety for networks that include plugins. The inference part of this sample demonstrates how to use the safe runtime and engine for networks that include plugins.
|
||||
|
||||
The build part builds a safe version of a TensorRT engine and saves it into a binary file, then the infer part loads the prebuilt safe engine and performs inference on an input image. The infer part uses the safety header proxy, with the following build steps demonstrating how to build it against the safe runtime for deployment and development.
|
||||
|
||||
## Workflow for safety plugin deployment
|
||||
|
||||
The deployment of the safety plugin on QNX safety involves a series of steps designed to ensure reliability. The recommended workflow is as follows:
|
||||
|
||||
- Build safety plugin for QNX standard platform and proxy plugin for X86 platform.
|
||||
- Build safety engine from X86 platform through remote auto-tuning on QNX safety platform.
|
||||
- Run safety engine with safety plugin on QNX standard platform with proxy runtime to perform necessary testing.
|
||||
- Build safety plugin for QNX safety platform.
|
||||
- Deploy safety engine on QNX safety platform with safety plugin.
|
||||
|
||||
## How does this sample work?
|
||||
|
||||
This sample uses an ONNX model that was trained on the [MNIST dataset](https://github.com/NVIDIA/DIGITS/blob/master/docs/GettingStarted.md).
|
||||
|
||||
Specifically, this sample:
|
||||
- Registers a Max Pooling plugin
|
||||
- Performs the basic setup and initialization of TensorRT using the ONNX parser
|
||||
- Imports a trained ONNX model using the ONNX parser
|
||||
- Preprocesses the input and stores the result in a managed buffer
|
||||
- Builds a safe engine
|
||||
- Serializes and deserializes the engine
|
||||
- Uses the engine to perform inference on an input image
|
||||
|
||||
To verify whether the engine is operating correctly, this sample picks a 28x28 image of a digit at random and runs inference on it using the engine it created. The output of the network is a probability distribution on the digit, showing which digit is likely that in the image.
|
||||
|
||||
### Register safe plugin creator
|
||||
|
||||
`ISafePluginCreatorV3One` will be used by TensorRT's builder for engine creation. You would need to register the safe plugin creator into the safePluginRegistry for engine generation and validation. The current sample registers MaxPoolCreator in namespace "" to parse and verify model `mnist_safe_plugin.onnx`.
|
||||
```
|
||||
auto safePluginRegistry = nvinfer2::safe::getSafePluginRegistry(g_recorder);
|
||||
safePluginRegistry->registerCreator(maxPoolPluginCreator, "", g_recorder);
|
||||
```
|
||||
|
||||
## Prerequisites
|
||||
1. Preparing sample data
|
||||
|
||||
See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README.
|
||||
|
||||
## Running the sample
|
||||
|
||||
1. Verify that the MNIST images are in `$TRT_DATADIR/mnist` and the ONNX network `mnist_safe_plugin.onnx` is in `$TRT_DATADIR/safe_plugin`.
|
||||
|
||||
2. Compile the sample by following the build instructions in the [TensorRT README](https://github.com/NVIDIA/TensorRT/). This will build the sample binaries, including `sample_plugin_safe_build` and `sample_plugin_safe_infer`.
|
||||
|
||||
3. The options that control compiling against safety runtime are summarized in the following table.
|
||||
| Compile Option | Default |Description|
|
||||
| ------------------------------- | ------- |---------- |
|
||||
|TRT_SAFETY_INFERENCE_ONLY | OFF |When enabled, build the infer part only, skip compiling the builder part.|
|
||||
|
||||
|
||||
4. Run the sample to build a TensorRT safe engine.
|
||||
```
|
||||
./sample_plugin_safe_build [--datadir=/path/to/data/dir/] [--remoteAutoTuningConfig=<config>] [--cpuOnly]
|
||||
```
|
||||
|
||||
This sample generates `safe_plugin.engine`, which is a binary file that contains the serialized engine data.
|
||||
|
||||
This sample reads ONNX model to build the network:
|
||||
|
||||
- `mnist_safe_plugin.onnx` - The ONNX model that contains the network design with maxPoolPlugin, version 1, namespace ""
|
||||
|
||||
**Note:** By default, this sample expects these files to be in either the `data/samples/safe_plugin/` or `data/safe_plugin/` directories. The list of default directories can be changed by adding one or more paths with `--datadir=/new/path/` as a command line argument.
|
||||
|
||||
5. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following:
|
||||
```
|
||||
&&&& RUNNING TensorRT.sample_safe_plugin_build [TensorRT v101200] [b5] # sample_plugin_safe_build --datadir=data/samples/safe_plugin/
|
||||
[04/23/2025-08:07:23] [I] Building a GPU inference engine for MNIST with plugins
|
||||
[04/23/2025-08:07:24] [I] [TRT] [MemUsageChange] Init CUDA: CPU +26, GPU +0, now: CPU 35, GPU 422 (MiB)
|
||||
[04/23/2025-08:07:36] [I] [TRT] [MemUsageChange] Init builder kernel library: CPU +1639, GPU +8, now: CPU 1874, GPU 430 (MiB)
|
||||
[04/23/2025-08:07:36] [I] [TRT] ----------------------------------------------------------------
|
||||
[04/23/2025-08:07:36] [I] [TRT] Input filename: data/samples/safe_plugin/mnist_safe_plugin.onnx
|
||||
[04/23/2025-08:07:36] [I] [TRT] ONNX IR version: 0.0.7
|
||||
[04/23/2025-08:07:36] [I] [TRT] Opset version: 8
|
||||
[04/23/2025-08:07:36] [I] [TRT] Producer name:
|
||||
[04/23/2025-08:07:36] [I] [TRT] Producer version:
|
||||
[04/23/2025-08:07:36] [I] [TRT] Domain:
|
||||
[04/23/2025-08:07:36] [I] [TRT] Model version: 0
|
||||
[04/23/2025-08:07:36] [I] [TRT] Doc string:
|
||||
[04/23/2025-08:07:36] [I] [TRT] ----------------------------------------------------------------
|
||||
[04/23/2025-08:07:36] [I] [TRT] Searching for plugin: MaxPoolPlugin, plugin_version: 1, plugin_namespace:
|
||||
[04/23/2025-08:07:36] [W] [TRT] onnxOpImporters.cpp:6641: Attribute pType not found in plugin node! Ensure that the plugin creator has a default value defined or the engine may fail to build.
|
||||
[04/23/2025-08:07:36] [I] [TRT] Successfully created plugin: MaxPoolPlugin
|
||||
[04/23/2025-08:07:40] [I] [TRT] Local timing cache in use. Profiling results in this builder pass will not be stored.
|
||||
[04/23/2025-08:07:40] [I] [TRT] Compiler backend is used during engine build.
|
||||
[04/23/2025-08:07:57] [I] [TRT] Detected 1 inputs and 1 output network tensors.
|
||||
[04/23/2025-08:07:57] [I] [TRT] Total Host Persistent Memory: 0 bytes
|
||||
[04/23/2025-08:07:57] [I] [TRT] Total Device Persistent Memory: 0 bytes
|
||||
[04/23/2025-08:07:57] [I] [TRT] Max Scratch Memory: 0 bytes
|
||||
[04/23/2025-08:07:57] [I] [TRT] Total Activation Memory: 0 bytes
|
||||
[04/23/2025-08:07:57] [I] [TRT] Total Weights Memory: 0 bytes
|
||||
[04/23/2025-08:07:57] [I] [TRT] Engine generation completed in 16.9161 seconds.
|
||||
[04/23/2025-08:07:57] [I] [TRT] [MemUsageStats] Peak memory usage of TRT CPU/GPU memory allocators: CPU 0 MiB, GPU 1 MiB
|
||||
&&&& PASSED TensorRT.sample_safe_plugin_build [TensorRT v101200] [b5] # sample_plugin_safe_build --datadir=data/samples/safe_plugin/
|
||||
```
|
||||
|
||||
This output shows that the sample ran successfully: `PASSED`.
|
||||
|
||||
6. Run the sample to perform inference on the digit:
|
||||
`./sample_plugin_safe_infer`
|
||||
|
||||
**Note:** This sample expects `./sample_plugin_safe_build` has been run to generate a safe engine file. It loads input image from `data/samples/mnist` directory, and walks back 10 directories to locate the image.
|
||||
|
||||
7. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following; ASCII rendering of the input image with digit 0:
|
||||
```
|
||||
Input:
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@+ :@@@@@@@@
|
||||
@@@@@@@@@@@@@@%= :. --%@@@@@
|
||||
@@@@@@@@@@@@@%. -@= - :@@@@@
|
||||
@@@@@@@@@@@@@: -@@#%@@ #@@@@
|
||||
@@@@@@@@@@@@: #@@@@@@@-#@@@@
|
||||
@@@@@@@@@@@= #@@@@@@@@=%@@@@
|
||||
@@@@@@@@@@= #@@@@@@@@@:@@@@@
|
||||
@@@@@@@@@+ -@@@@@@@@@%.@@@@@
|
||||
@@@@@@@@@::@@@@@@@@@@+-@@@@@
|
||||
@@@@@@@@-.%@@@@@@@@@@.*@@@@@
|
||||
@@@@@@@@ *@@@@@@@@@@@ *@@@@@
|
||||
@@@@@@@% %@@@@@@@@@%.-@@@@@@
|
||||
@@@@@@@:*@@@@@@@@@+. %@@@@@@
|
||||
@@@@@@# @@@@@@@@@# .*@@@@@@@
|
||||
@@@@@@# @@@@@@@@= +@@@@@@@@
|
||||
@@@@@@# @@@@@@%. .+@@@@@@@@@
|
||||
@@@@@@# @@@@@*. -%@@@@@@@@@@
|
||||
@@@@@@# --- =@@@@@@@@@@@@
|
||||
@@@@@@# *%@@@@@@@@@@@@@
|
||||
@@@@@@@%: -=%@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
@@@@@@@@@@@@@@@@@@@@@@@@@@@@
|
||||
|
||||
Set address of Input3 on device at 7fcb26612600
|
||||
Set address of Output on device at 7fcb26613400
|
||||
Output:
|
||||
Prob 0 0.9998 Class 0: **********
|
||||
Prob 1 0.0000 Class 1:
|
||||
Prob 2 0.0000 Class 2:
|
||||
Prob 3 0.0000 Class 3:
|
||||
Prob 4 0.0000 Class 4:
|
||||
Prob 5 0.0000 Class 5:
|
||||
Prob 6 0.0002 Class 6:
|
||||
Prob 7 0.0000 Class 7:
|
||||
Prob 8 0.0000 Class 8:
|
||||
Prob 9 0.0000 Class 9:
|
||||
&&&& PASSED TensorRT.sample_plugin_safe_infer # sample_plugin_safe_infer
|
||||
```
|
||||
This output shows that the sample ran successfully: `PASSED`.
|
||||
|
||||
### Tool command line arguments
|
||||
|
||||
To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
|
||||
|
||||
```bash
|
||||
sample_plugin_safe_build --help
|
||||
sample_plugin_safe_infer --help
|
||||
```
|
||||
|
||||
### When to use remoteAutoTuningConfig
|
||||
|
||||
The `--remoteAutoTuningConfig` parameter is designed for **cross-platform development scenarios** where you need to:
|
||||
|
||||
**Primary Use Case - Cross-Platform Building:**
|
||||
- **Build on Host Platform**: Compile and build TensorRT engines on a development machine (e.g., Linux x86_64)
|
||||
- **Auto-tune on Target Platform**: Perform kernel auto-tuning on the actual deployment target (e.g., QNX aarch64)
|
||||
|
||||
Use `--cpuOnly` with `--remoteAutoTuningConfig` to build the engine without a local GPU on the build host:
|
||||
```bash
|
||||
./sample_plugin_safe_build --remoteAutoTuningConfig=<config> --cpuOnly
|
||||
```
|
||||
|
||||
**Typical Scenarios:**
|
||||
- **QNX Development**: Building engines on Linux development machines but deploying on QNX automotive platforms
|
||||
|
||||
**Important Technical Limitation:**
|
||||
- **QNX Safety Devices**: QNX safety platforms do **NOT** support engine building operations. All engine construction must be performed on development platforms (Linux/QNX standard), making remote auto-tuning essential for safety deployments.
|
||||
|
||||
**Security Consideration:**
|
||||
- Use `PROMPT` as password on the command line in order to be interactively prompted for the password.
|
||||
This avoids leaking the password into the shell history. This functionality is only supported on
|
||||
Linux x86 -- the remote autotuning flow runs autotuning on an x86 Linux host which communicates
|
||||
over SSH to the aarch64 QNX machine.
|
||||
|
||||
## Shared plugin library and trtexec
|
||||
|
||||
Safety plugins must be built into shared libraries and dynamically loaded in applications. As shown in this sample, users may declare and register a plugin creator from their application directly, or define a plugin creator registering function within their plugin libraries themselves. In other words, we present a method that supports the manual registration of safety plugins from within a shared plugin library. Safety plugin libraries adhering to this protocol can be utilized with both trtexec and trtexec_safe.
|
||||
|
||||
### Creation of safety plugin library
|
||||
|
||||
After implementing the plugin creator, create an API entry function `getSafetyPluginCreator` to pull in the plugin creator. The entry function takes `pluginNamespace` and `pluginName` as arguments. If the input plugin namespace and name are valid, the function shall construct a `pluginCreator` instance, and return a pointer to the instance.
|
||||
```
|
||||
extern "C" __attribute__((visibility("default"))) nvinfer2::safe::IPluginCreatorInterface* getSafetyPluginCreator(
|
||||
char const* pluginNamespace, char const* pluginName)
|
||||
```
|
||||
|
||||
Please refer to `maxPoolPluginCreatorInterface.cpp` and `maxPoolPluginRuntimeCreatorInterface.cpp` for sample implementations.
|
||||
|
||||
### Using the plugin library with trtexec and trtexec_safe
|
||||
|
||||
Both `trtexec` and `trtexec_safe` support the `--safetyPlugins` argument, where you specify safety plugins to load from a safety plugin library that implements the manual-registration protocol. The sample plugin and its creator interface will be built into the shared libraries libsample_safe_plugin_v3.so(Linux x86_64 platform) for BUILD and RUNTIME capability and libsample_safe_plugin_v3_safe.so(QNX-safe platform) for RUNTIME only capability.
|
||||
|
||||
Sample command to build, validate, and save the safety engine on Linux x86_64 platform.
|
||||
```
|
||||
trtexec --onnx=$TRT_DATADIR/safe_plugin/mnist_safe_plugin.onnx --safe --skipInference --fp16 --consistency --safetyPlugins=libsample_safe_plugin_v3.so[::MaxPoolPlugin] --saveEngine=sample.engine
|
||||
```
|
||||
Sample command to run the pre-built engine on the QNX-safe platform.
|
||||
```
|
||||
trtexec_safe --loadEngine=sample.engine --safetyPlugins=libsample_safe_plugin_v3_safe.so[::MaxPoolPlugin]
|
||||
```
|
||||
|
||||
## Additional resources
|
||||
|
||||
The following resources provide a deeper understanding about sampleSafePluginV3.
|
||||
|
||||
**Dataset**
|
||||
- Sample data available in [TensorRT GitHub Releases](https://github.com/NVIDIA/TensorRT/releases).
|
||||
|
||||
**Documentation**
|
||||
- [NVIDIA’s TensorRT Documentation Library](https://docs.nvidia.com/deeplearning/sdk/tensorrt-archived/index.html)
|
||||
|
||||
## 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
|
||||
|
||||
- Apr. 2025
|
||||
- This is the first release of the `README.md` file and sample.
|
||||
|
||||
Jul. 2025
|
||||
This sample was updated for the TRT 10.13.1 safety release.
|
||||
|
||||
Dec. 2025
|
||||
This sample was updated to use the CMake-based build system.
|
||||
|
||||
Apr. 2026
|
||||
This sample was updated to add the `--cpuOnly` build option for remote auto-tuning workflows without requiring a local GPU on the build host.
|
||||
|
||||
## Known issues
|
||||
|
||||
There are no known issues in this sample.
|
||||
@@ -0,0 +1,153 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include "maxPoolKernel.h"
|
||||
|
||||
#include <cstdint>
|
||||
|
||||
template <typename T>
|
||||
__device__ __forceinline__ const T& max(const T& a, const T& b);
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ const half& max(const half& a, const half& b)
|
||||
{
|
||||
#if __CUDA_ARCH__ >= 700 || !defined(__CUDA_ARCH__)
|
||||
return __hgt(a, b) ? a : b;
|
||||
#else
|
||||
return (static_cast<float>(a) > static_cast<float>(b)) ? a : b;
|
||||
#endif
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ const float& max(const float& a, const float& b)
|
||||
{
|
||||
return (a > b) ? a : b;
|
||||
}
|
||||
|
||||
template <>
|
||||
__device__ __forceinline__ int8_t const& max(int8_t const& a, int8_t const& b)
|
||||
{
|
||||
return (a > b) ? a : b;
|
||||
}
|
||||
|
||||
// Cuda kernel to find maximum in the kernelsize matrix
|
||||
template <typename T>
|
||||
__global__ void maxKernel(
|
||||
int32_t B, int32_t C, int32_t H, int32_t W,
|
||||
const T* input,
|
||||
T* output, int32_t kernsize, int32_t stride, int32_t pad)
|
||||
{
|
||||
// Total input volume
|
||||
int32_t const N = B * C * H * W;
|
||||
int32_t out_id;
|
||||
int32_t b, c, h, w;
|
||||
|
||||
int32_t const H_out = (H + 2 * pad - kernsize) / stride + 1;
|
||||
int32_t const W_out = (W + 2 * pad - kernsize) / stride + 1;
|
||||
|
||||
// Index in the output tensor
|
||||
out_id = blockIdx.x * blockDim.x + threadIdx.x;
|
||||
if (out_id > B * C * H_out * W_out - 1)
|
||||
{
|
||||
return;
|
||||
}
|
||||
T maxim = static_cast<T>(0);
|
||||
|
||||
// Output index of batch
|
||||
b = out_id / (C * H_out * W_out);
|
||||
int32_t const temp = out_id % (C * H_out * W_out);
|
||||
|
||||
// Output index of channels
|
||||
c = temp / (H_out * W_out);
|
||||
int32_t const x = temp % (H_out * W_out);
|
||||
|
||||
// Output index of height
|
||||
h = x / W_out; // row major format
|
||||
|
||||
// Output index of width
|
||||
w = x % W_out;
|
||||
|
||||
// Index in input tensor considering stride
|
||||
int32_t k = (b * C * H * W) + (c * (H * W)) + (h * stride * W) + (w * stride);
|
||||
|
||||
maxim = input[k];
|
||||
// Find maximum value in the kernelsize matrix
|
||||
for (int32_t i = k; i < k + kernsize; i++)
|
||||
{
|
||||
for (int32_t j = 0; j < kernsize; j++)
|
||||
{
|
||||
if ((i + (j * W)) < N)
|
||||
maxim = max(maxim, input[i + (j * W)]);
|
||||
}
|
||||
}
|
||||
|
||||
output[out_id] = maxim;
|
||||
}
|
||||
|
||||
int32_t maxPoolFloat(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, const void* input, void* output,
|
||||
int32_t kernsize, int32_t stride, int32_t pad)
|
||||
{
|
||||
int32_t const blocksize = 512;
|
||||
// Compute number of entries in output
|
||||
int32_t const K = batch_size * C * ((H - kernsize + 2 * pad) / stride + 1) * ((W - kernsize + 2 * pad) / stride + 1);
|
||||
int32_t const g = ((K + blocksize - 1) / blocksize);
|
||||
|
||||
maxKernel<float><<<g, blocksize, 0, stream>>>(
|
||||
batch_size, C, H, W, static_cast<const float*>(input), static_cast<float*>(output), kernsize, stride, pad);
|
||||
auto retVal = cudaStreamSynchronize(stream);
|
||||
if (retVal != cudaSuccess)
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t maxPoolHalf(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, const void* input, void* output,
|
||||
int32_t kernsize, int32_t stride, int32_t pad)
|
||||
{
|
||||
int32_t const blocksize = 512;
|
||||
// Compute number of entries in output
|
||||
int32_t const K = batch_size * C * ((H - kernsize + 2 * pad) / stride + 1) * ((W - kernsize + 2 * pad) / stride + 1);
|
||||
int32_t const g = ((K + blocksize - 1) / blocksize);
|
||||
|
||||
maxKernel<half><<<g, blocksize, 0, stream>>>(
|
||||
batch_size, C, H, W, static_cast<const half*>(input), static_cast<half*>(output), kernsize, stride, pad);
|
||||
auto retVal = cudaStreamSynchronize(stream);
|
||||
if (retVal != cudaSuccess)
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t maxPoolInt8(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, const void* input, void* output,
|
||||
int32_t kernsize, int32_t stride, int32_t pad)
|
||||
{
|
||||
int32_t const blocksize = 512;
|
||||
// Compute number of entries in output
|
||||
int32_t const K = batch_size * C * ((H - kernsize + 2 * pad) / stride + 1) * ((W - kernsize + 2 * pad) / stride + 1);
|
||||
int32_t const g = ((K + blocksize - 1) / blocksize);
|
||||
|
||||
maxKernel<int8_t><<<g, blocksize, 0, stream>>>(
|
||||
batch_size, C, H, W, static_cast<int8_t const*>(input), static_cast<int8_t*>(output), kernsize, stride, pad);
|
||||
auto retVal = cudaStreamSynchronize(stream);
|
||||
if (retVal != cudaSuccess)
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
@@ -0,0 +1,30 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 1993-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#ifndef MAX_POOL_KERNEL_H
|
||||
#define MAX_POOL_KERNEL_H
|
||||
|
||||
#include <cstdint>
|
||||
#include <cuda_fp16.h>
|
||||
|
||||
int32_t maxPoolFloat(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, void const* input,
|
||||
void* output, int32_t kernsize, int32_t stride, int32_t pad);
|
||||
int32_t maxPoolHalf(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, void const* input,
|
||||
void* output, int32_t kernsize, int32_t stride, int32_t pad);
|
||||
int32_t maxPoolInt8(cudaStream_t stream, int32_t batch_size, int32_t C, int32_t H, int32_t W, void const* input,
|
||||
void* output, int32_t kernsize, int32_t stride, int32_t pad);
|
||||
|
||||
#endif // MAX_POOL_KERNEL_H
|
||||
@@ -0,0 +1,211 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cstdlib>
|
||||
#include <exception>
|
||||
#include <iostream>
|
||||
#include <string_view>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolKernel.h"
|
||||
#include "maxPoolPlugin.h"
|
||||
#include "safeCommon.h"
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
|
||||
IPluginV3* MaxPoolPlugin::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
return new MaxPoolPlugin{mParams};
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
std::cerr << "Error cloning MaxPoolPlugin: " << e.what() << std::endl;
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
IPluginCapability* MaxPoolPlugin::getCapabilityInterface(PluginCapabilityType type) noexcept
|
||||
{
|
||||
switch (type)
|
||||
{
|
||||
case PluginCapabilityType::kBUILD: return static_cast<IPluginV3OneSafeBuildMSS*>(this);
|
||||
case PluginCapabilityType::kRUNTIME: return static_cast<IPluginV3OneSafeRuntime*>(this);
|
||||
case PluginCapabilityType::kCORE: return static_cast<IPluginV3OneSafeCore*>(this);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
int32_t MaxPoolPlugin::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
int32_t MaxPoolPlugin::configurePlugin(
|
||||
TensorDescriptor const* in, int32_t nbInputs, TensorDescriptor const* out, int32_t nbOutputs) noexcept
|
||||
{
|
||||
SAFE_ASSERT(in && nbInputs == 1);
|
||||
SAFE_ASSERT(out && nbOutputs == 1);
|
||||
SAFE_ASSERT(in[0].dataType == out[0].dataType);
|
||||
|
||||
mParams.dtype = in[0].dataType;
|
||||
mParams.C = in[0].shape.d[1];
|
||||
mParams.H = in[0].shape.d[2];
|
||||
mParams.W = in[0].shape.d[3];
|
||||
mParams.H_out = out[0].shape.d[2];
|
||||
mParams.W_out = out[0].shape.d[3];
|
||||
|
||||
switch (mParams.dtype)
|
||||
{
|
||||
case nvinfer1::DataType::kINT8: mParams.dtypeBytes = 1; break;
|
||||
case nvinfer1::DataType::kHALF: mParams.dtypeBytes = 2; break;
|
||||
case nvinfer1::DataType::kFLOAT: mParams.dtypeBytes = 4; break;
|
||||
default:
|
||||
{
|
||||
mRecorder->reportError(
|
||||
nvinfer1::ErrorCode::kFAILED_EXECUTION, "Failed to execute due to unavailable precision.");
|
||||
return 1;
|
||||
}
|
||||
}
|
||||
return 0;
|
||||
}
|
||||
|
||||
bool MaxPoolPlugin::supportsFormatCombination(
|
||||
int32_t pos, TensorDescriptor const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
// For this method inputs are numbered 0..(nbInputs-1) and outputs are
|
||||
// numbered nbInputs..(nbInputs+nbOutputs-1). Using this numbering, pos is
|
||||
// an index into InOut, where 0 <= pos < nbInputs+nbOutputs.
|
||||
if (!(nbInputs == 1 && nbOutputs == 1 && pos < nbInputs + nbOutputs))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check if the data type is supported
|
||||
bool const supportedDataType = (inOut[pos].dataType == nvinfer1::DataType::kFLOAT)
|
||||
|| (inOut[pos].dataType == nvinfer1::DataType::kHALF) || (inOut[pos].dataType == nvinfer1::DataType::kINT8);
|
||||
|
||||
if (!supportedDataType)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
// Check if the format is supported (no vectorization)
|
||||
bool const supportedFormat = (inOut[pos].vectorizedDim == -1);
|
||||
|
||||
if (!supportedFormat)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
// For output tensors, ensure they match the input data type
|
||||
if (pos >= nbInputs) // This is an output tensor
|
||||
{
|
||||
// Output must match input data type
|
||||
return inOut[pos].dataType == inOut[0].dataType && inOut[pos].vectorizedDim == inOut[0].vectorizedDim;
|
||||
}
|
||||
|
||||
// For input tensors, just check that the type and format are supported
|
||||
return true;
|
||||
}
|
||||
|
||||
int32_t MaxPoolPlugin::getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
SAFE_ASSERT(inputTypes && nbInputs == 1);
|
||||
outputTypes[0] = inputTypes[0];
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t MaxPoolPlugin::getOutputShapes(
|
||||
Dims const* inputs, int32_t nbInputs, Dims* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
// Empty, will not be called
|
||||
return 0;
|
||||
}
|
||||
|
||||
int32_t MaxPoolPlugin::getSymbolicOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs* outputs,
|
||||
int32_t nbOutputs, IExprBuilder& exprBuilder) const noexcept
|
||||
{
|
||||
SAFE_ASSERT(inputs && nbInputs == 1 && inputs[0U].nbDims == 4);
|
||||
SAFE_ASSERT(outputs && nbOutputs == 1);
|
||||
outputs[0U].nbDims = inputs[0U].nbDims;
|
||||
outputs[0U].d[0U] = inputs[0U].d[0U];
|
||||
outputs[0U].d[1U] = inputs[0U].d[1U];
|
||||
|
||||
// Calculate height: (input_height + pad_y*2 - kernel_y) / stride_y + 1
|
||||
auto const* padY2
|
||||
= exprBuilder.operation(DimensionOperation::kPROD, *exprBuilder.constant(mParams.Py), *exprBuilder.constant(2));
|
||||
SAFE_ASSERT(padY2);
|
||||
SAFE_ASSERT(inputs[0U].d[2U]);
|
||||
auto const* inputPlusPadY = exprBuilder.operation(DimensionOperation::kSUM, *inputs[0U].d[2U], *padY2);
|
||||
SAFE_ASSERT(inputPlusPadY);
|
||||
auto const* heightMinusKernel
|
||||
= exprBuilder.operation(DimensionOperation::kSUB, *inputPlusPadY, *exprBuilder.constant(mParams.Ky));
|
||||
SAFE_ASSERT(heightMinusKernel);
|
||||
auto const* heightDivStride
|
||||
= exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *heightMinusKernel, *exprBuilder.constant(mParams.Sy));
|
||||
SAFE_ASSERT(heightDivStride);
|
||||
outputs[0U].d[2U] = exprBuilder.operation(DimensionOperation::kSUM, *heightDivStride, *exprBuilder.constant(1));
|
||||
SAFE_ASSERT(outputs[0U].d[2U]);
|
||||
|
||||
// Calculate width: (input_width + pad_x*2 - kernel_x) / stride_x + 1
|
||||
auto const* padX2
|
||||
= exprBuilder.operation(DimensionOperation::kPROD, *exprBuilder.constant(mParams.Px), *exprBuilder.constant(2));
|
||||
SAFE_ASSERT(padX2);
|
||||
SAFE_ASSERT(inputs[0U].d[3U]);
|
||||
auto const* inputPlusPadX = exprBuilder.operation(DimensionOperation::kSUM, *inputs[0U].d[3U], *padX2);
|
||||
SAFE_ASSERT(inputPlusPadX);
|
||||
auto const* widthMinusKernel
|
||||
= exprBuilder.operation(DimensionOperation::kSUB, *inputPlusPadX, *exprBuilder.constant(mParams.Kx));
|
||||
SAFE_ASSERT(widthMinusKernel);
|
||||
auto const* widthDivStride
|
||||
= exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *widthMinusKernel, *exprBuilder.constant(mParams.Sx));
|
||||
SAFE_ASSERT(widthDivStride);
|
||||
outputs[0U].d[3U] = exprBuilder.operation(DimensionOperation::kSUM, *widthDivStride, *exprBuilder.constant(1));
|
||||
SAFE_ASSERT(outputs);
|
||||
|
||||
return 0;
|
||||
}
|
||||
|
||||
size_t MaxPoolPlugin::getWorkspaceSize(
|
||||
TensorDescriptor const* inputs, int32_t nbInputs, TensorDescriptor const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
// MaxPool doesn't need any workspace memory
|
||||
return 0;
|
||||
}
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
extern "C" nvinfer2::safe::consistency::IPluginChecker* getPluginChecker(char const* name)
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
if (name != nullptr && name == "MaxPoolPlugin1"sv)
|
||||
{
|
||||
auto checker = std::make_unique<nvinfer2::safe::consistency::MaxPoolPluginChecker>();
|
||||
return checker.release();
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
@@ -0,0 +1,101 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef TRT_MAX_POOL_PLUGIN_H
|
||||
#define TRT_MAX_POOL_PLUGIN_H
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "NvInferConsistency.h"
|
||||
#include "NvInferPluginBase.h"
|
||||
#include "NvInferRuntime.h"
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
#include "maxPoolPluginRuntime.h"
|
||||
|
||||
using namespace nvinfer2::safe;
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
|
||||
class MaxPoolPlugin : public MaxPoolPluginRuntime, public IPluginV3OneSafeBuildMSS
|
||||
{
|
||||
public:
|
||||
MaxPoolPlugin(PoolParameters const& params)
|
||||
: MaxPoolPluginRuntime(params)
|
||||
{
|
||||
}
|
||||
|
||||
~MaxPoolPlugin() override = default;
|
||||
|
||||
// IPluginV3 Methods
|
||||
|
||||
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
|
||||
|
||||
IPluginV3* clone() noexcept override;
|
||||
|
||||
// IPluginV3OneSafeBuildMSS Methods
|
||||
|
||||
int32_t getNbOutputs() const noexcept override;
|
||||
|
||||
int32_t configurePlugin(
|
||||
TensorDescriptor const* in, int32_t nbInputs, TensorDescriptor const* out, int32_t nbOutputs) noexcept override;
|
||||
|
||||
bool supportsFormatCombination(
|
||||
int32_t pos, TensorDescriptor const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override;
|
||||
|
||||
int32_t getOutputDataTypes(
|
||||
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override;
|
||||
|
||||
int32_t getOutputShapes(
|
||||
Dims const* inputs, int32_t nbInputs, Dims* outputs, int32_t nbOutputs) const noexcept override;
|
||||
|
||||
int32_t getSymbolicOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs* outputs, int32_t nbOutputs,
|
||||
IExprBuilder& exprBuilder) const noexcept override;
|
||||
|
||||
size_t getWorkspaceSize(TensorDescriptor const* inputs, int32_t nbInputs, TensorDescriptor const* outputs,
|
||||
int32_t nbOutputs) const noexcept override;
|
||||
};
|
||||
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
namespace nvinfer2::safe::consistency
|
||||
{
|
||||
|
||||
class MaxPoolPluginChecker : public IPluginChecker
|
||||
{
|
||||
public:
|
||||
bool validate(std::vector<nvinfer2::safe::TensorDescriptor> const& /*Inputs*/,
|
||||
std::vector<nvinfer2::safe::TensorDescriptor> const& /*Outputs*/,
|
||||
nvinfer1::PluginFieldCollection* /*fc*/) noexcept override
|
||||
{
|
||||
// Always return true
|
||||
return true;
|
||||
}
|
||||
};
|
||||
|
||||
} // namespace nvinfer2::safe::consistency
|
||||
|
||||
extern "C" nvinfer2::safe::consistency::IPluginChecker* getPluginChecker(char const* name);
|
||||
|
||||
#endif // TRT_MAX_POOL_PLUGIN_H
|
||||
@@ -0,0 +1,100 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cstring>
|
||||
#include <exception>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <sstream>
|
||||
#include <string_view>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolPluginCreator.h"
|
||||
|
||||
namespace nvinfer1::plugin
|
||||
{
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
using namespace std::string_view_literals;
|
||||
|
||||
//! Like C++23's `container.append_range(std::span(data, count))`.
|
||||
template <typename Container, typename T>
|
||||
void appendRange(Container& container, T* data, int64_t count)
|
||||
{
|
||||
container.insert(container.end(), data, data + count);
|
||||
}
|
||||
|
||||
//! \return PoolParameters parsed from the given `PluginField`s.
|
||||
//! \param fields, nbFields: A span of fields.
|
||||
[[nodiscard]] PoolParameters parsePoolParameters(nvinfer1::PluginField const* const fields, int32_t const nbFields)
|
||||
{
|
||||
std::vector<int32_t> kernelShape;
|
||||
std::vector<int32_t> strides;
|
||||
std::vector<int32_t> pads;
|
||||
PoolingType pType{PoolingType::kMAX}; // Default to MAX pooling
|
||||
for (int32_t i{0}; i < nbFields; ++i)
|
||||
{
|
||||
auto const& field = fields[i];
|
||||
std::string_view const attrName = field.name;
|
||||
if (attrName == "kernel_shape"sv)
|
||||
{
|
||||
appendRange(kernelShape, static_cast<int32_t const*>(field.data), field.length);
|
||||
}
|
||||
else if (attrName == "strides"sv)
|
||||
{
|
||||
appendRange(strides, static_cast<int32_t const*>(field.data), field.length);
|
||||
}
|
||||
else if (attrName == "pads"sv)
|
||||
{
|
||||
appendRange(pads, static_cast<int32_t const*>(field.data), field.length);
|
||||
}
|
||||
else if (attrName == "pType"sv && field.data != nullptr)
|
||||
{
|
||||
pType = *static_cast<PoolingType const*>(field.data);
|
||||
}
|
||||
}
|
||||
|
||||
PoolParameters result{};
|
||||
result.pType = pType;
|
||||
result.Sx = strides.at(0);
|
||||
result.Sy = strides.at(1);
|
||||
result.Kx = kernelShape.at(0);
|
||||
result.Ky = kernelShape.at(1);
|
||||
result.Px = pads.at(0);
|
||||
result.Py = pads.at(1);
|
||||
return result;
|
||||
}
|
||||
} // namespace
|
||||
|
||||
IPluginV3* MaxPoolCreator::createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
// The build phase and the deserialization phase are handled differently.
|
||||
switch (phase)
|
||||
{
|
||||
case TensorRTPhase::kBUILD:
|
||||
return std::make_unique<MaxPoolPlugin>(parsePoolParameters(fc->fields, fc->nbFields)).release();
|
||||
case TensorRTPhase::kRUNTIME:
|
||||
return std::make_unique<MaxPoolPlugin>(*static_cast<PoolParameters const*>(fc->fields[0].data)).release();
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
} // namespace nvinfer1::plugin
|
||||
@@ -0,0 +1,49 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef TENSORRT_MAX_POOL_PLUGIN_CREATOR_H
|
||||
#define TENSORRT_MAX_POOL_PLUGIN_CREATOR_H
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "NvInferRuntime.h"
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolPlugin.h"
|
||||
#include "maxPoolPluginRuntimeCreator.h"
|
||||
|
||||
using namespace nvinfer2::safe;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
// Plugin factory class.
|
||||
class MaxPoolCreator : public MaxPoolRuntimeCreator
|
||||
{
|
||||
public:
|
||||
MaxPoolCreator() = default;
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
};
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
#ifdef GEN_PLUGIN_LIB
|
||||
extern "C" IPluginCreatorInterface* getSafetyPluginCreator(char const* pluginNamespace, char const* pluginName);
|
||||
#endif
|
||||
#endif // TENSORRT_MAX_POOL_PLUGIN_CREATOR_H
|
||||
@@ -0,0 +1,39 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <string_view>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolPluginCreator.h"
|
||||
|
||||
#ifdef GEN_PLUGIN_LIB
|
||||
extern "C" __attribute__((visibility("default"))) IPluginCreatorInterface* getSafetyPluginCreator(
|
||||
char const* pluginNamespace, char const* pluginName)
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
if (pluginName == "MaxPoolPlugin"sv && pluginNamespace == ""sv)
|
||||
{
|
||||
auto creator = std::make_unique<nvinfer1::plugin::MaxPoolCreator>();
|
||||
return creator.release();
|
||||
}
|
||||
|
||||
return nullptr;
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,154 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cstdlib>
|
||||
#include <cstring>
|
||||
#include <exception>
|
||||
#include <iostream>
|
||||
#include <vector>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolKernel.h"
|
||||
#include "maxPoolPluginRuntime.h"
|
||||
#include "safeCommon.h"
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
|
||||
MaxPoolPluginRuntime::MaxPoolPluginRuntime(PoolParameters const& params)
|
||||
: mParams{params}
|
||||
, mRecorder{nullptr}
|
||||
{
|
||||
initFieldsToSerialize();
|
||||
}
|
||||
|
||||
void MaxPoolPluginRuntime::initFieldsToSerialize()
|
||||
{
|
||||
// Serialize MaxPoolParameters.
|
||||
mDataToSerialize.clear();
|
||||
mDataToSerialize.emplace_back(
|
||||
nvinfer1::PluginField("parameters", &mParams, PluginFieldType::kUNKNOWN, sizeof(PoolParameters)));
|
||||
mFCToSerialize.nbFields = mDataToSerialize.size();
|
||||
mFCToSerialize.fields = mDataToSerialize.data();
|
||||
}
|
||||
|
||||
IPluginCapability* MaxPoolPluginRuntime::getCapabilityInterface(PluginCapabilityType type) noexcept
|
||||
{
|
||||
switch (type)
|
||||
{
|
||||
case PluginCapabilityType::kRUNTIME: return static_cast<IPluginV3OneSafeRuntime*>(this);
|
||||
case PluginCapabilityType::kCORE: return static_cast<IPluginV3OneSafeCore*>(this);
|
||||
case PluginCapabilityType::kBUILD: return nullptr; // MaxPoolPluginRuntime does not support build capability
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV3* MaxPoolPluginRuntime::clone() noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
return new MaxPoolPluginRuntime{mParams};
|
||||
}
|
||||
catch (...)
|
||||
{
|
||||
return nullptr;
|
||||
}
|
||||
}
|
||||
|
||||
AsciiChar const* MaxPoolPluginRuntime::getPluginName() const noexcept
|
||||
{
|
||||
return kMAX_POOL_PLUGIN_NAME;
|
||||
}
|
||||
|
||||
AsciiChar const* MaxPoolPluginRuntime::getPluginVersion() const noexcept
|
||||
{
|
||||
return kMAX_POOL_PLUGIN_VERSION;
|
||||
}
|
||||
|
||||
AsciiChar const* MaxPoolPluginRuntime::getPluginNamespace() const noexcept
|
||||
{
|
||||
return kMAX_POOL_PLUGIN_NAMESPACE;
|
||||
}
|
||||
|
||||
int32_t MaxPoolPluginRuntime::enqueue(TensorDescriptor const* inputDesc, TensorDescriptor const* outputDesc,
|
||||
void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
int32_t ret = 0;
|
||||
SAFE_ASSERT(inputs);
|
||||
SAFE_ASSERT(outputs);
|
||||
int32_t batch_size = inputDesc[0].shape.d[0];
|
||||
SAFE_ASSERT(batch_size > 0);
|
||||
switch (mParams.dtype)
|
||||
{
|
||||
case nvinfer1::DataType::kFLOAT:
|
||||
ret = maxPoolFloat(stream, batch_size, mParams.C, mParams.H, mParams.W, inputs[0], outputs[0], mParams.Kx,
|
||||
mParams.Sx, mParams.Px);
|
||||
break;
|
||||
case nvinfer1::DataType::kHALF:
|
||||
ret = maxPoolHalf(stream, batch_size, mParams.C, mParams.H, mParams.W, inputs[0], outputs[0], mParams.Kx,
|
||||
mParams.Sx, mParams.Px);
|
||||
break;
|
||||
case nvinfer1::DataType::kINT8:
|
||||
ret = maxPoolInt8(stream, batch_size, mParams.C, mParams.H, mParams.W, inputs[0], outputs[0], mParams.Kx,
|
||||
mParams.Sx, mParams.Px);
|
||||
break;
|
||||
default:
|
||||
{
|
||||
if (mRecorder)
|
||||
{
|
||||
mRecorder->reportError(
|
||||
nvinfer1::ErrorCode::kFAILED_EXECUTION, "Failed to execute due to unavailable precision.");
|
||||
}
|
||||
ret = 1;
|
||||
}
|
||||
}
|
||||
return ret;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
std::cerr << e.what() << '\n';
|
||||
}
|
||||
return -1;
|
||||
}
|
||||
|
||||
int32_t MaxPoolPluginRuntime::initResource(ISafePluginResourceContext const* context) noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* MaxPoolPluginRuntime::getFieldsToSerialize() noexcept
|
||||
{
|
||||
return &mFCToSerialize;
|
||||
}
|
||||
|
||||
ISafeRecorder* MaxPoolPluginRuntime::getSafeRecorder() const noexcept
|
||||
{
|
||||
return mRecorder;
|
||||
}
|
||||
|
||||
void MaxPoolPluginRuntime::setSafeRecorder(ISafeRecorder& recorder) noexcept
|
||||
{
|
||||
mRecorder = &recorder;
|
||||
}
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
@@ -0,0 +1,158 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#ifndef TRT_MAX_POOL_RUNTIME_PLUGIN_H
|
||||
#define TRT_MAX_POOL_RUNTIME_PLUGIN_H
|
||||
|
||||
#include <memory>
|
||||
#include <string>
|
||||
#include <vector>
|
||||
|
||||
#include <cuda_runtime.h>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
// In IPluginV3 interface, the plugin name, version, and name space must be
|
||||
// specified for the plugin and plugin creator exactly the same.
|
||||
constexpr char const* const kMAX_POOL_PLUGIN_NAME{"MaxPoolPlugin"};
|
||||
constexpr char const* const kMAX_POOL_PLUGIN_VERSION{"1"};
|
||||
constexpr char const* const kMAX_POOL_PLUGIN_NAMESPACE{""};
|
||||
|
||||
using namespace nvinfer2::safe;
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
|
||||
using AsciiChar = char;
|
||||
|
||||
enum class PoolingType : int32_t
|
||||
{
|
||||
kMAX = 0, // Maximum over elements
|
||||
kAVERAGE = 1, // Average over elements. If the tensor is padded, the count includes the padding
|
||||
kMAX_AVERAGE_BLEND = 2 // Blending between max and average pooling: (1-blendFactor)*maxPool + blendFactor*avgPool
|
||||
};
|
||||
|
||||
struct PoolParameters
|
||||
{
|
||||
// Input dimensions
|
||||
int32_t C;
|
||||
int32_t H;
|
||||
int32_t W;
|
||||
// Output dimensions
|
||||
int32_t H_out;
|
||||
int32_t W_out;
|
||||
// Kernel size,
|
||||
int32_t Kx;
|
||||
int32_t Ky;
|
||||
// Stride
|
||||
int32_t Sx;
|
||||
int32_t Sy;
|
||||
// Padding
|
||||
int32_t Px;
|
||||
int32_t Py;
|
||||
// Pooling Function
|
||||
PoolingType pType;
|
||||
|
||||
// Extra Info
|
||||
nvinfer1::DataType dtype;
|
||||
size_t dtypeBytes;
|
||||
int32_t nbInputs;
|
||||
int64_t nbIOProfile;
|
||||
std::shared_ptr<Dims> inputDims;
|
||||
|
||||
// Default Constructor
|
||||
PoolParameters()
|
||||
{
|
||||
// To do: Populate Parameters from fc object w/ hard code
|
||||
pType = PoolingType::kMAX;
|
||||
// stride
|
||||
Sx = 3;
|
||||
Sy = 3;
|
||||
// kernel size
|
||||
Kx = 3;
|
||||
Ky = 3;
|
||||
// padding
|
||||
Px = 0;
|
||||
Py = 0;
|
||||
}
|
||||
|
||||
// Parameterized Constructor
|
||||
PoolParameters(nvinfer1::plugin::PoolingType pooling_type, int32_t stride_x, int32_t stride_y, int32_t kernel_x,
|
||||
int32_t kernel_y, int32_t padding_x, int32_t padding_y)
|
||||
{
|
||||
|
||||
pType = pooling_type;
|
||||
// stride
|
||||
Sx = stride_x;
|
||||
Sy = stride_y;
|
||||
// kernel size
|
||||
Kx = kernel_x;
|
||||
Ky = kernel_y;
|
||||
// padding
|
||||
Px = padding_x;
|
||||
Py = padding_y;
|
||||
}
|
||||
};
|
||||
|
||||
class MaxPoolPluginRuntime : public IPluginV3, public IPluginV3OneSafeCore, public IPluginV3OneSafeRuntime
|
||||
{
|
||||
public:
|
||||
MaxPoolPluginRuntime(PoolParameters const& params);
|
||||
|
||||
~MaxPoolPluginRuntime() override = default;
|
||||
|
||||
// IPluginV3 Methods
|
||||
|
||||
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override;
|
||||
|
||||
// IPluginV3OneSafeCore Methods
|
||||
|
||||
AsciiChar const* getPluginName() const noexcept override;
|
||||
|
||||
AsciiChar const* getPluginVersion() const noexcept override;
|
||||
|
||||
AsciiChar const* getPluginNamespace() const noexcept override;
|
||||
|
||||
// IPluginV3OneSafeRuntime Methods
|
||||
|
||||
int32_t enqueue(TensorDescriptor const* inputDesc, TensorDescriptor const* outputDesc, void const* const* inputs,
|
||||
void* const* outputs, void* workspace, cudaStream_t stream) noexcept override;
|
||||
|
||||
int32_t initResource(ISafePluginResourceContext const* context) noexcept;
|
||||
|
||||
IPluginV3* clone() noexcept override;
|
||||
|
||||
PluginFieldCollection const* getFieldsToSerialize() noexcept override;
|
||||
|
||||
void setSafeRecorder(ISafeRecorder&) noexcept;
|
||||
|
||||
ISafeRecorder* getSafeRecorder() const noexcept;
|
||||
|
||||
protected:
|
||||
void initFieldsToSerialize();
|
||||
|
||||
// TensorRT plugin parameters.
|
||||
PoolParameters mParams;
|
||||
std::vector<nvinfer1::PluginField> mDataToSerialize;
|
||||
nvinfer1::PluginFieldCollection mFCToSerialize;
|
||||
ISafeRecorder* mRecorder;
|
||||
};
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
#endif // TRT_MAX_POOL_RUNTIME_PLUGIN_H
|
||||
@@ -0,0 +1,99 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <cstring>
|
||||
#include <exception>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <mutex>
|
||||
#include <sstream>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolPluginRuntimeCreator.h"
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
|
||||
// Plugin creator
|
||||
MaxPoolRuntimeCreator::MaxPoolRuntimeCreator()
|
||||
{
|
||||
// Declare the ONNX attributes that the ONNX parser will collect from the
|
||||
// ONNX model that contains the IdentityConv node.
|
||||
|
||||
mPluginAttributes.clear();
|
||||
mPluginAttributes.emplace_back(nvinfer1::PluginField("kernel_shape", nullptr, PluginFieldType::kINT32, 2));
|
||||
mPluginAttributes.emplace_back(nvinfer1::PluginField("strides", nullptr, PluginFieldType::kINT32, 2));
|
||||
mPluginAttributes.emplace_back(nvinfer1::PluginField("pads", nullptr, PluginFieldType::kINT32, 2));
|
||||
mPluginAttributes.emplace_back(nvinfer1::PluginField("pType", nullptr, PluginFieldType::kINT32, 1));
|
||||
|
||||
mFC.nbFields = mPluginAttributes.size();
|
||||
mFC.fields = mPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* MaxPoolRuntimeCreator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return kMAX_POOL_PLUGIN_NAMESPACE;
|
||||
}
|
||||
|
||||
char const* MaxPoolRuntimeCreator::getPluginName() const noexcept
|
||||
{
|
||||
return kMAX_POOL_PLUGIN_NAME;
|
||||
}
|
||||
|
||||
char const* MaxPoolRuntimeCreator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kMAX_POOL_PLUGIN_VERSION;
|
||||
}
|
||||
|
||||
nvinfer1::PluginFieldCollection const* MaxPoolRuntimeCreator::getFieldNames() noexcept
|
||||
{
|
||||
// This is only used in the build phase.
|
||||
return &mFC;
|
||||
}
|
||||
|
||||
ISafeRecorder* MaxPoolRuntimeCreator::getSafeRecorder() const noexcept
|
||||
{
|
||||
return mRecorder;
|
||||
}
|
||||
|
||||
void MaxPoolRuntimeCreator::setSafeRecorder(ISafeRecorder& recorder) noexcept
|
||||
{
|
||||
mRecorder = &recorder;
|
||||
}
|
||||
|
||||
IPluginV3* MaxPoolRuntimeCreator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept
|
||||
{
|
||||
// MaxPoolRuntimeCreator only handles RUNTIME phase.
|
||||
// BUILD phase should use MaxPoolCreator instead.
|
||||
if (phase == TensorRTPhase::kRUNTIME)
|
||||
{
|
||||
nvinfer1::PluginField const* fields{fc->fields};
|
||||
|
||||
PoolParameters params{*static_cast<PoolParameters const*>(fields[0].data)};
|
||||
|
||||
auto plugin = std::make_unique<MaxPoolPluginRuntime>(params);
|
||||
return plugin.release();
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
@@ -0,0 +1,67 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#ifndef TENSORRT_MAX_POOL_RUNTIME_PLUGIN_CREATOR_H
|
||||
#define TENSORRT_MAX_POOL_RUNTIME_PLUGIN_CREATOR_H
|
||||
|
||||
#include <vector>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolPluginRuntime.h"
|
||||
|
||||
using namespace nvinfer2::safe;
|
||||
|
||||
namespace nvinfer1
|
||||
{
|
||||
namespace plugin
|
||||
{
|
||||
// Plugin factory class.
|
||||
class MaxPoolRuntimeCreator : public ISafePluginCreatorV3One
|
||||
{
|
||||
public:
|
||||
MaxPoolRuntimeCreator();
|
||||
|
||||
~MaxPoolRuntimeCreator() override = default;
|
||||
|
||||
char const* getPluginNamespace() const noexcept override;
|
||||
|
||||
char const* getPluginName() const noexcept override;
|
||||
|
||||
char const* getPluginVersion() const noexcept override;
|
||||
|
||||
nvinfer1::PluginFieldCollection const* getFieldNames() noexcept override;
|
||||
|
||||
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override;
|
||||
|
||||
ISafeRecorder* getSafeRecorder() const noexcept override;
|
||||
|
||||
void setSafeRecorder(ISafeRecorder&) noexcept override;
|
||||
|
||||
protected:
|
||||
nvinfer1::PluginFieldCollection mFC;
|
||||
std::vector<nvinfer1::PluginField> mPluginAttributes;
|
||||
ISafeRecorder* mRecorder;
|
||||
};
|
||||
} // namespace plugin
|
||||
} // namespace nvinfer1
|
||||
|
||||
#ifdef GEN_PLUGIN_LIB
|
||||
extern "C" IPluginCreatorInterface* getSafetyPluginCreator(char const* pluginNamespace, char const* pluginName);
|
||||
#endif
|
||||
#endif // TENSORRT_MAX_POOL_RUNTIME_PLUGIN_CREATOR_H
|
||||
@@ -0,0 +1,39 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
|
||||
#include <string_view>
|
||||
|
||||
#include "NvInferSafePlugin.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
|
||||
#include "maxPoolPluginRuntimeCreator.h"
|
||||
|
||||
#ifdef GEN_PLUGIN_LIB
|
||||
extern "C" __attribute__((visibility("default"))) IPluginCreatorInterface* getSafetyPluginCreator(
|
||||
char const* pluginNamespace, char const* pluginName)
|
||||
{
|
||||
using namespace std::string_view_literals;
|
||||
if (pluginName == "MaxPoolPlugin"sv && pluginNamespace == ""sv)
|
||||
{
|
||||
auto creator = std::make_unique<nvinfer1::plugin::MaxPoolRuntimeCreator>();
|
||||
return creator.release();
|
||||
}
|
||||
|
||||
return nullptr;
|
||||
}
|
||||
#endif
|
||||
|
||||
@@ -0,0 +1,407 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#define DEFINE_TRT_ENTRYPOINTS 1
|
||||
#define DEFINE_TRT_ONNX_PARSER_ENTRYPOINT 0
|
||||
#define DEFINE_TRT_BUILDER_ENTRYPOINT 0
|
||||
#define DEFINE_TRT_REFITTER_ENTRYPOINT 0
|
||||
#define DEFINE_TRT_RUNTIME_ENTRYPOINT 0
|
||||
|
||||
#include "NvInfer.h"
|
||||
#include "NvInferSafeRuntime.h"
|
||||
#include "NvOnnxParser.h"
|
||||
|
||||
#include "argsParser.h"
|
||||
#include "buffers.h"
|
||||
#include "logger.h"
|
||||
#include "maxPoolPluginCreator.h"
|
||||
#include "parserOnnxConfig.h"
|
||||
#include "safeCommon.h"
|
||||
#include "safeErrorRecorder.h"
|
||||
#include "sampleUtils.h"
|
||||
#include <fstream>
|
||||
#include <ios>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
|
||||
std::string const gSampleName = "TensorRT.sample_safe_plugin_build";
|
||||
|
||||
using namespace nvinfer1;
|
||||
|
||||
static sample::SampleSafeRecorder g_recorder{nvinfer2::safe::Severity::kDEBUG};
|
||||
|
||||
namespace
|
||||
{
|
||||
|
||||
//!
|
||||
//! \brief The SampleSafePluginBuildArgs struct stores the additional arguments required by the sample
|
||||
//!
|
||||
struct SampleSafePluginBuildArgs : public samplesCommon::Args
|
||||
{
|
||||
std::string onnx{"mnist_safe_plugin_ds.onnx"};
|
||||
std::string engineFileName{"safe_plugin.engine"};
|
||||
std::string remoteAutoTuningConfig{};
|
||||
int32_t maxAuxStreams{0};
|
||||
bool cpuOnly{false};
|
||||
};
|
||||
|
||||
//!
|
||||
//! \brief This function parses arguments specific to the sample
|
||||
//!
|
||||
bool parseSampleSafePluginBuildArgs(SampleSafePluginBuildArgs& args, int32_t argc, char* argv[])
|
||||
{
|
||||
using namespace samplesSafeCommon;
|
||||
for (int32_t i = 1; i < argc; ++i)
|
||||
{
|
||||
std::string const arg = argv[i];
|
||||
if (auto value = parseString(arg, "saveEngine"))
|
||||
{
|
||||
if (!sample::validateNonEmpty(*value, "Engine filename"))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
args.engineFileName = std::move(*value);
|
||||
}
|
||||
else if (auto value = parseString(arg, "remoteAutoTuningConfig"))
|
||||
{
|
||||
if (!sample::validateNonEmpty(*value, "Remote auto tuning config")
|
||||
|| !sample::validateRemoteAutoTuningConfig(*value))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
args.remoteAutoTuningConfig = std::move(*value);
|
||||
}
|
||||
else if (auto const value = parseString(arg, "datadir", 'd'))
|
||||
{
|
||||
if (!sample::validateNonEmpty(*value, "Data directory path"))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
args.dataDirs.push_back(sample::normalizeDirectoryPath(*value));
|
||||
}
|
||||
else if (auto value = parseString(arg, "onnx"))
|
||||
{
|
||||
args.onnx = std::move(*value);
|
||||
}
|
||||
else if (auto const value = parseString(arg, "maxAuxStreams"))
|
||||
{
|
||||
args.maxAuxStreams = std::stoi(*value);
|
||||
if (args.maxAuxStreams < 0)
|
||||
{
|
||||
sample::gLogError << "Number of auxiliary streams must be >= 0, got: " << arg << "\n";
|
||||
return false;
|
||||
}
|
||||
}
|
||||
else if (parseBool(arg, "help", 'h'))
|
||||
{
|
||||
args.help = true;
|
||||
}
|
||||
else if (parseBool(arg, "cpuOnly"))
|
||||
{
|
||||
args.cpuOnly = true;
|
||||
}
|
||||
else
|
||||
{
|
||||
sample::gLogError << "Invalid Argument: " << arg << "\n";
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief The SampleSafePluginBuildParams struct stores the additional parameters required by the sample
|
||||
//!
|
||||
struct SampleSafePluginBuildParams : public samplesCommon::OnnxSampleParams
|
||||
{
|
||||
std::string engineFileName{};
|
||||
std::string remoteAutoTuningConfig{};
|
||||
bool std{false};
|
||||
int32_t maxAuxStreams{0};
|
||||
};
|
||||
|
||||
//!
|
||||
//! \brief Initialize members of the params struct using the command line args.
|
||||
//!
|
||||
SampleSafePluginBuildParams initializeSampleParams(SampleSafePluginBuildArgs const& args)
|
||||
{
|
||||
SampleSafePluginBuildParams params;
|
||||
if (args.dataDirs.empty()) // Use default directories if user hasn't provided directory paths.
|
||||
{
|
||||
params.dataDirs.push_back("data/");
|
||||
params.dataDirs.push_back("data/safe_plugin/");
|
||||
params.dataDirs.push_back("data/samples/safe_plugin/");
|
||||
}
|
||||
else // Use the data directory provided by the user.
|
||||
{
|
||||
params.dataDirs = args.dataDirs;
|
||||
}
|
||||
|
||||
params.onnxFileName = args.onnx;
|
||||
params.batchSize = 1;
|
||||
params.engineFileName = args.engineFileName;
|
||||
params.remoteAutoTuningConfig = args.remoteAutoTuningConfig;
|
||||
params.maxAuxStreams = args.maxAuxStreams;
|
||||
return params;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief The SampleSafePlugin class implements the sample.
|
||||
//!
|
||||
//! \details It creates the network using a trained ONNX MNIST classification model.
|
||||
//!
|
||||
class SampleSafePlugin
|
||||
{
|
||||
public:
|
||||
explicit SampleSafePlugin(SampleSafePluginBuildParams const& params)
|
||||
: mParams(params)
|
||||
{
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Builds the network engine.
|
||||
//!
|
||||
bool build();
|
||||
|
||||
private:
|
||||
//!
|
||||
//! \brief Uses an ONNX parser to create the MNIST Network and marks the
|
||||
//! output layers.
|
||||
//!
|
||||
bool constructNetwork(nvonnxparser::IParser* parser);
|
||||
|
||||
SampleSafePluginBuildParams mParams; //!< The parameters for the sample.
|
||||
|
||||
nvinfer1::Dims mInputDims; //!< The dimensions of the input to the network.
|
||||
|
||||
nvinfer1::plugin::MaxPoolCreator maxPoolPluginCreator{};
|
||||
};
|
||||
|
||||
//!
|
||||
//! \brief Creates the network, configures the builder and creates the network engine.
|
||||
//!
|
||||
//! \details This function creates the MNIST network by parsing the ONNX model and builds
|
||||
//! the engine that will be used to run MNIST.
|
||||
//!
|
||||
//! \return true if the engine was created successfully and false otherwise.
|
||||
//!
|
||||
bool SampleSafePlugin::build()
|
||||
{
|
||||
// Register custom plugin creator for Max pooling before building
|
||||
|
||||
auto safePluginRegistry = nvinfer2::safe::getSafePluginRegistry(g_recorder);
|
||||
if (!safePluginRegistry)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
safePluginRegistry->registerCreator(maxPoolPluginCreator, "", g_recorder);
|
||||
|
||||
auto builder = std::unique_ptr<nvinfer1::IBuilder>(nvinfer1::createInferBuilder(sample::gLogger.getTRTLogger()));
|
||||
if (!builder)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
NetworkDefinitionCreationFlags flags
|
||||
= (1U << static_cast<uint32_t>(NetworkDefinitionCreationFlag::kSTRONGLY_TYPED));
|
||||
auto network = std::unique_ptr<nvinfer1::INetworkDefinition>(builder->createNetworkV2(flags));
|
||||
if (!network)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
auto config = std::unique_ptr<nvinfer1::IBuilderConfig>(builder->createBuilderConfig());
|
||||
if (!config)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
auto parser
|
||||
= std::unique_ptr<nvonnxparser::IParser>(nvonnxparser::createParser(*network, sample::gLogger.getTRTLogger()));
|
||||
if (!parser)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
auto constructed = constructNetwork(parser.get());
|
||||
if (!constructed)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
// Set the input shape for the whole neural network by adding optimization profiles
|
||||
constexpr int64_t kBATCH_SIZE0 = 1;
|
||||
auto profile0 = builder->createOptimizationProfile();
|
||||
ASSERT(profile0);
|
||||
profile0->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kMIN, Dims4{kBATCH_SIZE0, 1, 28, 28});
|
||||
profile0->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kOPT, Dims4{kBATCH_SIZE0, 1, 28, 28});
|
||||
profile0->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kMAX, Dims4{kBATCH_SIZE0, 1, 28, 28});
|
||||
config->addOptimizationProfile(profile0);
|
||||
constexpr int64_t kBATCH_SIZE1 = 5;
|
||||
auto profile1 = builder->createOptimizationProfile();
|
||||
ASSERT(profile1);
|
||||
profile1->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kMIN, Dims4{kBATCH_SIZE1, 1, 28, 28});
|
||||
profile1->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kOPT, Dims4{kBATCH_SIZE1, 1, 28, 28});
|
||||
profile1->setDimensions(network->getInput(0)->getName(), OptProfileSelector::kMAX, Dims4{kBATCH_SIZE1, 1, 28, 28});
|
||||
config->addOptimizationProfile(profile1);
|
||||
config->setEngineCapability(nvinfer1::EngineCapability::kSAFETY);
|
||||
config->setMaxAuxStreams(mParams.maxAuxStreams);
|
||||
|
||||
// Set remote auto tuning config if provided
|
||||
if (!mParams.remoteAutoTuningConfig.empty())
|
||||
{
|
||||
config->setRemoteAutoTuningConfig(mParams.remoteAutoTuningConfig.c_str());
|
||||
}
|
||||
|
||||
auto buffer = std::unique_ptr<nvinfer1::IHostMemory>(builder->buildSerializedNetwork(*network, *config));
|
||||
if (!buffer)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
ASSERT(network->getNbInputs() == 1);
|
||||
mInputDims = network->getInput(0)->getDimensions();
|
||||
ASSERT(mInputDims.nbDims == 4);
|
||||
|
||||
// Save the engine
|
||||
std::string const engineFile = mParams.engineFileName;
|
||||
std::ofstream file(engineFile, std::ios::binary);
|
||||
if (!file)
|
||||
{
|
||||
sample::gLogError << "Failed to open file to save engine: " << engineFile << std::endl;
|
||||
return false;
|
||||
}
|
||||
file.write(reinterpret_cast<char const*>(buffer->data()), buffer->size());
|
||||
file.close();
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Uses an ONNX parser to create the MNIST Network and marks the
|
||||
//! output layers.
|
||||
//!
|
||||
//! \param parser ONNX parser used to parse the network
|
||||
//!
|
||||
bool SampleSafePlugin::constructNetwork(nvonnxparser::IParser* parser)
|
||||
{
|
||||
return parser->parseFromFile(locateFile(mParams.onnxFileName, mParams.dataDirs).c_str(),
|
||||
static_cast<int32_t>(sample::gLogger.getReportableSeverity()));
|
||||
}
|
||||
} // namespace
|
||||
|
||||
//!
|
||||
//! \brief Prints the help information for running this sample.
|
||||
//!
|
||||
void printHelpInfo()
|
||||
{
|
||||
SampleSafePluginBuildArgs const defArgs{};
|
||||
std::cout << R"(Usage: sample_plugin_safe_build [options]
|
||||
Options:
|
||||
--help, -h Print this message and exit.
|
||||
--datadir=DIR, -d=DIR Search for data in DIR. This option can be passed multiple times
|
||||
to add multiple search directories. If omitted, default data dirs are:
|
||||
data/samples/mnist/, data/mnist/
|
||||
--verbose Use verbose logging.
|
||||
--saveEngine=FILE Save the serialized engine into FILE (default = )"
|
||||
<< defArgs.engineFileName << R"().
|
||||
--onnx=FILE Load ONNX from FILE. (default = )"
|
||||
<< defArgs.onnx << R"().
|
||||
--remoteAutoTuningConfig=CONFIG
|
||||
Set remote auto tuning configuration in the following format:
|
||||
protocol://username[:password]@hostname[:port]?param1=value1¶m2=value2
|
||||
--maxAuxStreams=N Limit the number of auxiliary streams to N (default = )"
|
||||
<< defArgs.maxAuxStreams << R"().
|
||||
--cpuOnly Build the engine with CPU-only mode. Requires --remoteAutoTuningConfig.
|
||||
No local GPU is required on the build machine.
|
||||
|
||||
Examples:
|
||||
sample_plugin_safe_build \
|
||||
--remoteAutoTuningConfig=ssh://user:pass@192.0.2.100:22?remote_exec_path=/opt/tensorrt/bin&remote_lib_path=/opt/tensorrt/lib
|
||||
)";
|
||||
}
|
||||
|
||||
int main(int argc, char** argv)
|
||||
{
|
||||
SampleSafePluginBuildArgs args;
|
||||
bool const argsOK = parseSampleSafePluginBuildArgs(args, argc, argv);
|
||||
if (!argsOK)
|
||||
{
|
||||
printHelpInfo();
|
||||
return EXIT_FAILURE;
|
||||
}
|
||||
if (args.help)
|
||||
{
|
||||
printHelpInfo();
|
||||
return EXIT_SUCCESS;
|
||||
}
|
||||
|
||||
// Log remoteAutoTuningConfig usage
|
||||
if (!args.remoteAutoTuningConfig.empty())
|
||||
{
|
||||
sample::gLogInfo << "Remote auto tuning config specified: "
|
||||
<< sample::sanitizeRemoteAutoTuningConfig(args.remoteAutoTuningConfig) << std::endl;
|
||||
sample::gLogInfo << "This is a safety sample and will build in remote mode automatically." << std::endl;
|
||||
}
|
||||
|
||||
if (args.cpuOnly)
|
||||
{
|
||||
if (args.remoteAutoTuningConfig.empty())
|
||||
{
|
||||
sample::gLogError << "--cpuOnly requires --remoteAutoTuningConfig to be specified." << std::endl;
|
||||
printHelpInfo();
|
||||
return EXIT_FAILURE;
|
||||
}
|
||||
sample::gLogInfo << "Setting CPU-only mode" << std::endl;
|
||||
if (!samplesSafeCommon::applyCpuOnlyMode())
|
||||
{
|
||||
return EXIT_FAILURE;
|
||||
}
|
||||
}
|
||||
|
||||
if (!args.cpuOnly && !samplesCommon::isSmSafe())
|
||||
{
|
||||
sample::gLogInfo << "Skip safe mode test on unsupported platforms." << std::endl;
|
||||
return EXIT_SUCCESS;
|
||||
}
|
||||
|
||||
// Create sanitized argv for logging to avoid exposing credentials in test reports
|
||||
auto sanitizedArgs = sample::sanitizeArgv(argc, argv);
|
||||
std::vector<char const*> sanitizedArgv;
|
||||
sanitizedArgv.reserve(sanitizedArgs.size());
|
||||
for (auto const& s : sanitizedArgs)
|
||||
{
|
||||
sanitizedArgv.push_back(s.c_str());
|
||||
}
|
||||
|
||||
auto sampleTest
|
||||
= sample::gLogger.defineTest(gSampleName, static_cast<int32_t>(sanitizedArgv.size()), sanitizedArgv.data());
|
||||
|
||||
sample::gLogger.reportTestStart(sampleTest);
|
||||
|
||||
SampleSafePluginBuildParams params = initializeSampleParams(args);
|
||||
|
||||
SampleSafePlugin sample(params);
|
||||
sample::gLogInfo << "Building a GPU inference engine for MNIST with plugins" << std::endl;
|
||||
|
||||
if (!sample.build())
|
||||
{
|
||||
return sample::gLogger.reportFail(sampleTest);
|
||||
}
|
||||
|
||||
return sample::gLogger.reportPass(sampleTest);
|
||||
}
|
||||
@@ -0,0 +1,415 @@
|
||||
/*
|
||||
* SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
* SPDX-License-Identifier: Apache-2.0
|
||||
*
|
||||
* Licensed under the Apache License, Version 2.0 (the "License");
|
||||
* you may not use this file except in compliance with the License.
|
||||
* You may obtain a copy of the License at
|
||||
*
|
||||
* http://www.apache.org/licenses/LICENSE-2.0
|
||||
*
|
||||
* Unless required by applicable law or agreed to in writing, software
|
||||
* distributed under the License is distributed on an "AS IS" BASIS,
|
||||
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
* See the License for the specific language governing permissions and
|
||||
* limitations under the License.
|
||||
*/
|
||||
#include "NvInferSafeRuntime.h"
|
||||
#include "maxPoolPluginRuntimeCreator.h"
|
||||
#include "safeCommon.h"
|
||||
#include "safeErrorRecorder.h"
|
||||
#include <fstream>
|
||||
#include <ios>
|
||||
#include <iostream>
|
||||
#include <memory>
|
||||
#include <random>
|
||||
#include <string_view>
|
||||
#include <type_traits>
|
||||
|
||||
using namespace nvinfer1;
|
||||
using namespace samplesSafeCommon;
|
||||
std::string const gSampleName = "TensorRT.sample_safe_plugin_infer";
|
||||
|
||||
static sample::SampleSafeRecorder g_recorder{nvinfer2::safe::Severity::kINFO};
|
||||
|
||||
//!
|
||||
//! \brief The SampleSafeMNISTInferArgs struct stores the additional arguments required by the sample
|
||||
//!
|
||||
class SampleSafePluginInferArgs
|
||||
{
|
||||
public:
|
||||
std::string engineFileName{"safe_plugin.engine"};
|
||||
bool help{false};
|
||||
};
|
||||
|
||||
//!
|
||||
//! \brief This function parses arguments specific to the sample
|
||||
//!
|
||||
bool parseSampleSafePluginInferArgs(SampleSafePluginInferArgs& args, int32_t const argc, char const* const argv[])
|
||||
{
|
||||
for (int32_t i = 1; i < argc; ++i)
|
||||
{
|
||||
std::string const arg = argv[i];
|
||||
if (auto value = parseString(arg, "loadEngine"))
|
||||
{
|
||||
args.engineFileName = std::move(*value);
|
||||
}
|
||||
else if (parseBool(arg, "help", 'h'))
|
||||
{
|
||||
args.help = true;
|
||||
}
|
||||
else
|
||||
{
|
||||
SAFE_LOG << "Invalid Argument: " << arg << "\n";
|
||||
return false;
|
||||
}
|
||||
}
|
||||
return true;
|
||||
}
|
||||
|
||||
nvinfer2::safe::TypedArray createTypedArray(
|
||||
void* const ptr, DataType type, uint64_t bufferSize, nvinfer2::safe::ISafeRecorder& recorder)
|
||||
{
|
||||
switch (type)
|
||||
{
|
||||
case DataType::kFLOAT: return nvinfer2::safe::TypedArray(static_cast<float*>(ptr), bufferSize);
|
||||
case DataType::kHALF: return nvinfer2::safe::TypedArray(static_cast<nvinfer2::safe::half_t*>(ptr), bufferSize);
|
||||
case DataType::kINT32: return nvinfer2::safe::TypedArray(static_cast<int32_t*>(ptr), bufferSize);
|
||||
case DataType::kINT8: return nvinfer2::safe::TypedArray(static_cast<int8_t*>(ptr), bufferSize);
|
||||
default:
|
||||
{
|
||||
SAFE_LOG << "Invalid tensor DataType encountered." << std::endl;
|
||||
return nvinfer2::safe::TypedArray{};
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Allocate memory and memset it to zero using safe CUDA-compatible APIs.
|
||||
//!
|
||||
void* allocateAndMemset(uint64_t sizeInBytes, nvinfer2::safe::ISafeRecorder& recorder)
|
||||
{
|
||||
void* deviceBuf{nullptr};
|
||||
cudaStream_t stream;
|
||||
CUDA_CALL(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), recorder);
|
||||
CUDA_CALL(cudaMalloc(&deviceBuf, sizeInBytes), recorder);
|
||||
CUDA_CALL(cudaMemsetAsync(deviceBuf, 0, sizeInBytes, stream), recorder);
|
||||
CUDA_CALL(cudaStreamSynchronize(stream), recorder);
|
||||
CUDA_CALL(cudaStreamDestroy(stream), recorder);
|
||||
return deviceBuf;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Helper function to get the volume.
|
||||
//!
|
||||
inline int64_t volume(nvinfer1::Dims const& d)
|
||||
{
|
||||
return std::accumulate(d.d, d.d + d.nbDims, 1L, std::multiplies<int64_t>());
|
||||
}
|
||||
|
||||
///!
|
||||
//! \brief Loads the enginePlanFile from engineFile and returns it.
|
||||
//!
|
||||
std::vector<char> loadEnginePlanFile(std::string const& engineFile, int32_t& size)
|
||||
{
|
||||
std::string const& filename = engineFile;
|
||||
std::vector<char> engineStream;
|
||||
std::ifstream file(filename, std::ios::binary);
|
||||
if (!file.good())
|
||||
{
|
||||
SAFE_LOG << "Could not open input engine file or file is empty. File name: " << filename << std::endl;
|
||||
return {};
|
||||
}
|
||||
file.seekg(0, std::ifstream::end);
|
||||
size = file.tellg();
|
||||
file.seekg(0, std::ifstream::beg);
|
||||
engineStream.resize(size);
|
||||
file.read(engineStream.data(), size);
|
||||
file.close();
|
||||
|
||||
return engineStream;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Reads the input data, preprocesses, and stores the result in a managed buffer.
|
||||
//!
|
||||
bool processInput(void* const input, int32_t const inputFileIdx, int64_t const kBATCH_SIZE, int64_t const index)
|
||||
{
|
||||
constexpr int32_t kINPUT_H{28};
|
||||
constexpr int32_t kINPUT_W{28};
|
||||
|
||||
// Read the digit file according to the inputFileIdx.
|
||||
std::vector<uint8_t> fileData(static_cast<size_t>(kINPUT_H * kINPUT_W));
|
||||
std::vector<std::string> dataDirs;
|
||||
dataDirs.push_back("data/samples/mnist/");
|
||||
readPGMFile(locateFile(std::to_string(inputFileIdx) + ".pgm", dataDirs), fileData.data(), kINPUT_H, kINPUT_W);
|
||||
|
||||
// Print ASCII representation of digit.
|
||||
SAFE_LOG << "Input:\n";
|
||||
for (int32_t i = 0; i < kINPUT_H * kINPUT_W; i++)
|
||||
{
|
||||
SAFE_LOG << (" .:-=+*#%@"[fileData[i] / 26]) << (((i + 1) % kINPUT_W) ? "" : "\n");
|
||||
}
|
||||
SAFE_LOG << std::endl;
|
||||
|
||||
float* const hostInputBuffer = static_cast<float*>(input) + index * kINPUT_H * kINPUT_W;
|
||||
static_cast<void>(std::copy(fileData.begin(), fileData.end(), hostInputBuffer));
|
||||
// Normalize to 0-1 with background at 0
|
||||
static_cast<void>(std::transform(hostInputBuffer, hostInputBuffer + kINPUT_H * kINPUT_W, hostInputBuffer,
|
||||
[](float v) noexcept -> float { return 1.0F - v / 255.0F; }));
|
||||
|
||||
return true;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Verifies that the output is correct and prints it.
|
||||
//!
|
||||
bool verifyOutput(void* const output, std::vector<int32_t> const& groundTruthDigits, int64_t const batchSize)
|
||||
{
|
||||
bool result{true};
|
||||
constexpr int32_t kDIGITS{10};
|
||||
for (int64_t j = 0; j < batchSize; ++j)
|
||||
{
|
||||
float* const prob = static_cast<float*>(output) + j * kDIGITS;
|
||||
|
||||
// Print histogram of the output distribution.
|
||||
SAFE_LOG << "Output:" << std::endl;
|
||||
float val{0.0F};
|
||||
int32_t idx{0};
|
||||
|
||||
// Calculate Softmax
|
||||
float sum{0.0F};
|
||||
for (int32_t i = 0; i < kDIGITS; i++)
|
||||
{
|
||||
prob[i] = exp(prob[i]);
|
||||
sum += prob[i];
|
||||
}
|
||||
|
||||
for (int32_t i = 0; i < kDIGITS; i++)
|
||||
{
|
||||
prob[i] /= sum;
|
||||
if (val < prob[i])
|
||||
{
|
||||
val = prob[i];
|
||||
idx = i;
|
||||
}
|
||||
|
||||
SAFE_LOG << " Prob " << i << " " << std::fixed << std::setw(5) << std::setprecision(4) << prob[i]
|
||||
<< " Class " << i << ": "
|
||||
<< std::string(static_cast<int32_t>(std::floor(prob[i] * 10 + 0.5F)), '*') << std::endl;
|
||||
}
|
||||
result &= (idx == groundTruthDigits[j]) && (val > 0.9F);
|
||||
}
|
||||
return result;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Set I/O tensor buffer.
|
||||
//!
|
||||
void setTensorBuffer(nvinfer2::safe::ITRTGraph* graph, nvinfer2::safe::ISafeRecorder& recorder,
|
||||
std::string const& tensorName, void*& tensorAddress)
|
||||
{
|
||||
nvinfer2::safe::TensorDescriptor desc;
|
||||
SAFE_API_CALL(graph->getIOTensorDescriptor(desc, tensorName.c_str()), recorder);
|
||||
|
||||
void* deviceBuf = allocateAndMemset(desc.sizeInBytes, recorder);
|
||||
tensorAddress = deviceBuf;
|
||||
nvinfer2::safe::TypedArray tensor = createTypedArray(deviceBuf, desc.dataType, desc.sizeInBytes, recorder);
|
||||
SAFE_API_CALL(graph->setIOTensorAddress(tensorName.c_str(), tensor), recorder);
|
||||
SAFE_LOG << "Set address of " << tensorName << " on device at " << std::hex << (uint64_t) deviceBuf << std::dec
|
||||
<< std::endl;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Runs the TensorRT inference engine for this sample.
|
||||
//!
|
||||
//! \details This function is the main execution function of the sample. It allocates
|
||||
//! the buffer, sets inputs, executes the engine, and verifies the output.
|
||||
//!
|
||||
bool doInference(SampleSafePluginInferArgs const& args)
|
||||
{
|
||||
// Create the engine by loading from a local saved plan
|
||||
int32_t engineFileSize = 0;
|
||||
auto engineFile = loadEnginePlanFile(args.engineFileName, engineFileSize);
|
||||
SAFE_ASSERT(engineFileSize != 0);
|
||||
|
||||
// Inference
|
||||
nvinfer1::plugin::MaxPoolRuntimeCreator creator;
|
||||
ITRTGraph* graph = nullptr;
|
||||
|
||||
getSafePluginRegistry(g_recorder)->registerCreator(creator, "", g_recorder);
|
||||
createTRTGraph(graph, engineFile.data(), engineFileSize, g_recorder, true, nullptr);
|
||||
SAFE_ASSERT(graph != nullptr);
|
||||
|
||||
// Setup as many auxiliary streams as the graph requires - destroyed at scope end.
|
||||
auto auxStreamsDeleter = samplesSafeCommon::setUpAuxStreamsOn(*graph, g_recorder);
|
||||
|
||||
bool outputCorrect = true;
|
||||
int64_t nbIOProfile = 0;
|
||||
SAFE_API_CALL(graph->getNbIOProfiles(nbIOProfile), g_recorder);
|
||||
SAFE_ASSERT(nbIOProfile == 2);
|
||||
|
||||
auto descToString = [](nvinfer2::safe::TensorDescriptor const& desc) {
|
||||
std::stringstream ss;
|
||||
ss << desc.tensorName << " {";
|
||||
for (int64_t i = 0; i < desc.shape.nbDims; ++i)
|
||||
{
|
||||
ss << desc.shape.d[i];
|
||||
if (i < desc.shape.nbDims - 1)
|
||||
{
|
||||
ss << ", ";
|
||||
}
|
||||
}
|
||||
ss << "}";
|
||||
return ss.str();
|
||||
};
|
||||
|
||||
for (int64_t k = 0; k < nbIOProfile; ++k)
|
||||
{
|
||||
graph->setIOProfile(k);
|
||||
// Memory Config
|
||||
int64_t nbIOs{};
|
||||
SAFE_API_CALL(graph->getNbIOTensors(nbIOs), g_recorder);
|
||||
|
||||
// This sample only has one input and one output.
|
||||
SAFE_ASSERT(nbIOs == 2);
|
||||
constexpr int32_t inputIndex{0};
|
||||
constexpr int32_t outputIndex{1};
|
||||
|
||||
// Get the binding dimensions according to the input/output index.
|
||||
char const* inputBindingName = nullptr;
|
||||
char const* outputBindingName = nullptr;
|
||||
nvinfer2::safe::TensorDescriptor inputDesc;
|
||||
nvinfer2::safe::TensorDescriptor outputDesc;
|
||||
|
||||
graph->getIOTensorName(inputBindingName, inputIndex);
|
||||
graph->getIOTensorName(outputBindingName, outputIndex);
|
||||
graph->getIOTensorDescriptor(inputDesc, inputBindingName);
|
||||
graph->getIOTensorDescriptor(outputDesc, outputBindingName);
|
||||
SAFE_ASSERT(inputDesc.ioMode == nvinfer1::TensorIOMode::kINPUT);
|
||||
SAFE_ASSERT(outputDesc.ioMode == nvinfer1::TensorIOMode::kOUTPUT);
|
||||
SAFE_ASSERT(inputDesc.shape.nbDims > 0);
|
||||
SAFE_LOG << "Set IO profile to " << k << std::endl;
|
||||
SAFE_LOG << descToString(inputDesc) << std::endl;
|
||||
int64_t kBATCH_SIZE = inputDesc.shape.d[0];
|
||||
SAFE_ASSERT(0 < kBATCH_SIZE && kBATCH_SIZE <= 9);
|
||||
// Create host buffers
|
||||
std::vector<void*> hostBuffers(nbIOs, nullptr);
|
||||
|
||||
hostBuffers[inputIndex] = malloc(inputDesc.sizeInBytes);
|
||||
hostBuffers[outputIndex] = malloc(outputDesc.sizeInBytes);
|
||||
|
||||
std::vector<int32_t> groundTruthDigits(kBATCH_SIZE);
|
||||
for (int64_t j = 0; j < kBATCH_SIZE; ++j)
|
||||
{
|
||||
// Pick a random digit to try to infer.
|
||||
std::random_device rd;
|
||||
std::default_random_engine generator{rd()};
|
||||
std::uniform_int_distribution<int32_t> distribution(0, 9);
|
||||
int32_t const digit = distribution(generator);
|
||||
groundTruthDigits[j] = digit;
|
||||
// Read the input data into the managed buffers.
|
||||
if (!processInput(hostBuffers[inputIndex], digit, kBATCH_SIZE, j))
|
||||
{
|
||||
return false;
|
||||
}
|
||||
}
|
||||
std::vector<void*> buffers(nbIOs, nullptr);
|
||||
// Set input tensor values
|
||||
for (int64_t i = 0; i < nbIOs; ++i)
|
||||
{
|
||||
char const* tensor;
|
||||
SAFE_API_CALL(graph->getIOTensorName(tensor, i), g_recorder);
|
||||
setTensorBuffer(graph, g_recorder, tensor, buffers[i]);
|
||||
}
|
||||
|
||||
// Initialize main stream
|
||||
cudaStream_t stream;
|
||||
CUDA_CALL(cudaStreamCreateWithFlags(&stream, cudaStreamNonBlocking), g_recorder);
|
||||
|
||||
// Asynchronously copy data from host input buffers to device input buffers.
|
||||
CUDA_CALL(cudaMemcpyAsync(buffers[inputIndex], hostBuffers[inputIndex], inputDesc.sizeInBytes,
|
||||
cudaMemcpyHostToDevice, stream),
|
||||
g_recorder);
|
||||
|
||||
// Run the graph
|
||||
SAFE_API_CALL(graph->executeAsync(stream), g_recorder);
|
||||
|
||||
// Asynchronously copy data from device output buffers to host output buffers.
|
||||
CUDA_CALL(cudaMemcpyAsync(hostBuffers[outputIndex], buffers[outputIndex], outputDesc.sizeInBytes,
|
||||
cudaMemcpyDeviceToHost, stream),
|
||||
g_recorder);
|
||||
|
||||
graph->sync();
|
||||
|
||||
// Check and print the output of the inference.
|
||||
outputCorrect &= verifyOutput(hostBuffers[outputIndex], groundTruthDigits, kBATCH_SIZE);
|
||||
|
||||
// free host&device buffers
|
||||
free(hostBuffers[inputIndex]);
|
||||
free(hostBuffers[outputIndex]);
|
||||
CUDA_CALL(cudaFree(buffers[inputIndex]), g_recorder);
|
||||
CUDA_CALL(cudaFree(buffers[outputIndex]), g_recorder);
|
||||
}
|
||||
|
||||
destroyTRTGraph(graph);
|
||||
return outputCorrect;
|
||||
}
|
||||
|
||||
//!
|
||||
//! \brief Prints the help information for running this sample.
|
||||
//!
|
||||
void printHelpInfo()
|
||||
{
|
||||
SampleSafePluginInferArgs const defArgs{};
|
||||
std::cout << R"(Usage: sample_safe_plugin_infer [options]
|
||||
Options:
|
||||
--help, -h Print this message and exit.
|
||||
--loadEngine=FILE Load serialized engine from FILE (default = )"
|
||||
<< defArgs.engineFileName << R"().
|
||||
)";
|
||||
}
|
||||
|
||||
int main(int32_t argc, char** argv)
|
||||
{
|
||||
safetyCompliance::setPromgrAbility();
|
||||
SampleSafePluginInferArgs args;
|
||||
bool const argsOK = parseSampleSafePluginInferArgs(args, argc, argv);
|
||||
if (!argsOK)
|
||||
{
|
||||
printHelpInfo();
|
||||
return EXIT_FAILURE;
|
||||
}
|
||||
if (args.help)
|
||||
{
|
||||
printHelpInfo();
|
||||
return EXIT_SUCCESS;
|
||||
}
|
||||
|
||||
// Initialize SafeCuda before any other Cuda APIs are called. This may be skipped if createInferRuntime() is
|
||||
// called first as per DEEPLRN_RES_116
|
||||
safetyCompliance::initSafeCuda();
|
||||
|
||||
if (!samplesSafeCommon::isSmSafe())
|
||||
{
|
||||
SAFE_LOG << "Skip safe mode test on unsupported platforms." << std::endl;
|
||||
return EXIT_SUCCESS;
|
||||
}
|
||||
|
||||
TestResult result = TestResult::kPASSED;
|
||||
try
|
||||
{
|
||||
if (!doInference(args))
|
||||
{
|
||||
result = TestResult::kFAILED;
|
||||
}
|
||||
}
|
||||
catch (std::runtime_error& e)
|
||||
{
|
||||
SAFE_LOG << e.what() << std::endl;
|
||||
result = TestResult::kFAILED;
|
||||
}
|
||||
|
||||
reportTestResult("TensorRT.sample_plugin_safe_infer", result, argc, argv);
|
||||
return EXIT_SUCCESS;
|
||||
}
|
||||
Reference in New Issue
Block a user