wehub/nvidia--tensorrt
1
0
Fork 0
Code Issues Pull Requests Actions 29 Deployments Models Agent Notes Packages Projects Releases Wiki Activity

chore: import upstream snapshot with attribution
Docker Image CI / build-ubuntu2004 (push) Has been cancelled
Details

Browse Source
This commit is contained in:
wehub-resource-sync
2026-07-13 13:36:55 +08:00
commit c8a779b1bb
1887 changed files with 3245738 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
samples/sampleNonZeroPlugin/CMakeLists.txt
+28
View File
@@ -0,0 +1,28 @@
#
# 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.
#
add_executable(sample_non_zero_plugin
sampleNonZeroPlugin.cpp
nonZeroKernel.cu
)
target_link_libraries(sample_non_zero_plugin PRIVATE trt_samples_common TRT_SAMPLES::tensorrt)
add_dependencies(tensorrt_samples sample_non_zero_plugin)
installLibraries(
TARGETS sample_non_zero_plugin
OPTIONAL
COMPONENT internal
)
samples/sampleNonZeroPlugin/README.md
+184
View File
@@ -0,0 +1,184 @@
# NonZero Plugin for TensorRT using IPluginV3
**Table Of Contents**
- [Description](#description)
- [How does this sample work?](#how-does-this-sample-work)
* [Implementing a NonZero plugin using IPluginV3 interface](#implementing-a-nonzero-plugin-using-ipluginv3-interface)
* [Creating network and building the engine](#creating-network-and-building-the-engine)
* [Running inference](#running-inference)
- [Running the sample](#running-the-sample)
* [Sample `--help` options](#sample---help-options)
- [Additional resources](#additional-resources)
- [License](#license)
- [Changelog](#changelog)
- [Known issues](#known-issues)
## Description
This sample, sampleNonZeroPlugin, implements a plugin for the NonZero operation, customizable to output the non-zero indices in
either a row order (each set of indices in the same row) or column order format (each set of indices in the same column).
NonZero is an operation where the non-zero indices of the input tensor is found.
## How does this sample work?
This sample creates and runs a TensorRT engine built from a network containing a single NonZeroPlugin node. It demonstrates how
custom layers with data-dependent output shapes can be implemented and added to a TensorRT network.
Specifically, this sample:
- [Implements a TensorRT plugin for the NonZero operation](#implementing-a-nonzero-plugin-using-ipluginv3-interface)
- [Creates a network and builds an engine](#creating-network-and-building-the-engine)
- [Runs inference using the generated TensorRT network](#running-inference)
### Implementing a NonZero plugin using IPluginV3 interface
Until `IPluginV3` (and associated interfaces), TensorRT plugins could not have outputs whose shapes depended on the input values (they could only depend
on input shapes). `IPluginV3OneBuild` which exposes a build capability for `IPluginV3`, provides support for such data-dependent output shapes.
`NonZeroPlugin` in this sample is written to handle 2-D input tensors of shape $R \times C$. Assume that the tensor contains $K$ non-zero elements and that the
non-zero indices are required in a row ordering (each set of indices in its own row). Then the output shape would be $K \times 2$.
The output shapes are expressed to the TensorRT builder through the `IPluginV3OneBuild::getOutputShapes()` API. Expressing the second dimension of the output is
straightforward:
```
outputs[0].d[1] = exprBuilder.constant(2);
```
The extent of each data-dependent dimension in the plugin must be expressed in terms of a *_size tensor_*. A size tensor is a scalar output of
`DataType::kINT32` or `DataType::kINT64` that must be added as one of the plugin outputs. In this case, it is sufficient to declare one size tensor to denote the extent of the
first dimension of the non-zero indices output. To declare a size tensor, one must provide an upper-bound and optimum value for its extent as `IDimensionExpr`s. These can be formed through the `IExprBuilder` argument passed to the `IPluginV3OneBuild::getOutputShapes()` method.
- For unknown inputs, the upper-bound is the total number of elements in the input
```
auto upperBound = exprBuilder.operation(DimensionOperation::kPROD, *inputs[0].d[0], *inputs[0].d[1]);
```
- A good estimate for the optimum is that half of the elements are non-zero
```
auto optValue = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *upperBound, *exprBuilder.constant(2));
```
Now we can declare the size tensor using the `IExprBuilder::declareSizeTensor()` method, which also requires the specification of the output index at which the size tensor would reside. Let us place it after the non-zero indices output:
```
auto numNonZeroSizeTensor = exprBuilder.declareSizeTensor(1, *optValue, *upperBound);
```
Now we are ready to specify the extent of the first dimension of the non-zero indices output:
```
outputs[0].d[0] = numNonZeroSizeTensor;
```
and let's not forget to declare that the size tensor is a scalar (0-D):
```
outputs[1].nbDims = 0;
```
The `NonZeroPlugin` can also be configured to emit the non-zero indices in a column-order fashion through the `rowOrder` plugin attribute, by setting it to `0`.
In this case, the first output of the plugin will have shape $2 \times K$, and the output shape specification must be adjusted accordingly.
### Creating network and building the engine
To add the plugin to the network, the `INetworkDefinition::addPluginV3()` method must be used.
Similar to `IPluginCreator` used for V2 plugins, V3 plugins must be accompanied by the registration of a plugin creator implementing the `IPluginCreatorV3One`
interface.
### Running inference
As sample inputs, random images from MNIST dataset are selected and scaled to between `[0,1]`. The network will output both the non-zero indices,
as well as the non-zero count.
## Prerequisites
1. Preparing sample data
See [Preparing sample data](../README.md#preparing-sample-data) in the main samples README.
## Running the sample
1. Compile the sample by following build instructions in [TensorRT README](https://github.com/NVIDIA/TensorRT/).
2. Run the sample to build and run the MNIST engine from the ONNX model.
```
./sample_non_zero_plugin [-h or --help] [-d or --datadir=<path to data directory>] [--columnOrder] [--fp16]
```
3. Verify that the sample ran successfully. If the sample runs successfully you should see output similar to the following:
```
&&&& RUNNING TensorRT.sample_non_zero_plugin # ./sample_non_zero_plugin
...
[I] Input:
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.854902, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.858824, 0, 0, 0.0745098, 0, 0.564706, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.317647, 0, 0, 0.47451, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.0431373, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.854902, 0, 0, 0.145098
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.564706, 0, 0, 0.996078
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.282353
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.854902
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.854902, 0, 0, 0.145098, 0, 0.564706
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.564706, 0, 0, 0.996078, 0, 0
[I] 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0.282353, 0, 0
[I]
[I] Output:
[I] 2 14
[I] 3 9
[I] 3 12
[I] 3 14
[I] 4 9
[I] 4 12
[I] 5 12
[I] 8 12
[I] 8 15
[I] 9 12
[I] 9 15
[I] 10 15
[I] 13 15
[I] 14 10
[I] 14 13
[I] 14 15
[I] 15 10
[I] 15 13
[I] 16 13
&&&& PASSED TensorRT.sample_non_zero_plugin # ./sample_non_zero_plugin
```
### Sample `--help` options
To see the full list of available options and their descriptions, use the `-h` or `--help` command line option.
# Additional resources
The following resources provide a deeper understanding about the V3 TensorRT plugins and the NonZero operation:
**NonZero**
- [ONNX: NonZero](https://onnx.ai/onnx/operators/onnx__NonZero.html)
**TensorRT plugins**
- [Extending TensorRT with Custom Layers](https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#extending)
**Other documentation**
- [Introduction To NVIDIA’s TensorRT Samples](https://docs.nvidia.com/deeplearning/sdk/tensorrt-sample-support-guide/index.html#samples)
- [Working With TensorRT Using The C++ API](https://docs.nvidia.com/deeplearning/sdk/tensorrt-developer-guide/index.html#c_topics)
- [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
October 2025
Migrate to strongly typed APIs.
March 2024
This is the first version of this `README.md` file.
# Known issues
Windows users building this sample with Visual Studio with a CUDA version different from the TensorRT package will need to retarget the project to build against the installed CUDA version through the `Build Dependencies -> Build Customization` menu.
samples/sampleNonZeroPlugin/nonZeroKernel.cu
+83
View File
@@ -0,0 +1,83 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#include "nonZeroKernel.h"
inline __device__ int32_t isZero(float const& a)
{
return a == 0.F;
}
inline __device__ int32_t isZero(half const& a)
{
#if __CUDA_ARCH__ >= 530
return a == __float2half(0.F);
#else
return __half2float(a) == 0.F;
#endif
}
template <typename T>
__global__ void findNonZeroIndicesKernel(
T const* X, int32_t* indices, unsigned long long* count, unsigned long long const* K, int32_t R, int32_t C, int32_t rowOrder)
{
int32_t col = blockIdx.x * blockDim.x + threadIdx.x;
// Check if the column index is within bounds
if (col < C)
{
for (int32_t row = 0; row < R; ++row)
{
if (!isZero(X[row * C + col]))
{
unsigned long long index = atomicAdd(count, 1ULL); // Increment count atomically and get the previous value
if (indices)
{
if(rowOrder == 0)
{
indices[index] = row;
indices[index + *K] = col;
}
else
{
indices[2 * index] = row;
indices[2 * index + 1] = col;
}
}
}
}
}
}
template <typename T>
void nonZeroIndicesImpl(T const* X, int32_t* indices, int64_t* count, int64_t const* K, int32_t R, int32_t C,
bool rowOrder, cudaStream_t stream)
{
constexpr int32_t kBLOCK_SIZE = 256;
int32_t const blocksPerGrid = (C + kBLOCK_SIZE - 1) / kBLOCK_SIZE;
static_assert(sizeof(unsigned long long) == 8U, "unsigned long long must be 8 bytes in NVCC");
findNonZeroIndicesKernel<<<blocksPerGrid, kBLOCK_SIZE, 0, stream>>>(
X, indices, reinterpret_cast<unsigned long long*>(count), reinterpret_cast<unsigned long long const*>(K), R, C, static_cast<int32_t>(rowOrder));
}
#define NONZERO_SPECIALIZED_IMPL(T) \
template void nonZeroIndicesImpl<T>(T const* X, int32_t* indices, int64_t* count, int64_t const* K, int32_t R, \
int32_t C, bool rowOrder, cudaStream_t stream);
NONZERO_SPECIALIZED_IMPL(float)
NONZERO_SPECIALIZED_IMPL(half)
samples/sampleNonZeroPlugin/nonZeroKernel.h
+28
View File
@@ -0,0 +1,28 @@
/*
* SPDX-FileCopyrightText: Copyright (c) 1993-2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
* SPDX-License-Identifier: Apache-2.0
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
#ifndef SAMPLE_NONZERO_KERNEL_H
#define SAMPLE_NONZERO_KERNEL_H
#include <cuda_fp16.h>
#include <cstdint>
template <typename T>
void nonZeroIndicesImpl(T const* X, int32_t* indices, int64_t* count, int64_t const* K, int32_t R, int32_t C,
bool rowOrder, cudaStream_t stream);
#endif // SAMPLE_NONZERO_KERNEL_H
samples/sampleNonZeroPlugin/sampleNonZeroPlugin.cpp
+788
View File
@@ -0,0 +1,788 @@
/*
* 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.
*/
//!
//! sampleNonZeroPlugin.cpp
//! This file contains a sample demonstrating a plugin for NonZero.
//! It can be run with the following command line:
//! Command: ./sample_non_zero_plugin [-h or --help] [-d=/path/to/data/dir or --datadir=/path/to/data/dir]
//!
// Define TRT entrypoints used in common code
#define DEFINE_TRT_ENTRYPOINTS 1
#include "argsParser.h"
#include "buffers.h"
#include "common.h"
#include "logger.h"
#include "nonZeroKernel.h"
#include "parserOnnxConfig.h"
#include "NvInfer.h"
#include <cuda_runtime_api.h>
#include <cstdlib>
#include <fstream>
#include <iostream>
#include <random>
#include <sstream>
#include <string_view>
using namespace nvinfer1;
std::string const kSAMPLE_NAME = "TensorRT.sample_non_zero_plugin";
using half = __half;
void nonZeroIndicesHelper(nvinfer1::DataType type, void const* X, void* indices, void* count, void const* K, int32_t R,
int32_t C, bool rowOrder, cudaStream_t stream)
{
if (type == nvinfer1::DataType::kFLOAT)
{
nonZeroIndicesImpl<float>(static_cast<float const*>(X), static_cast<int32_t*>(indices),
static_cast<int64_t*>(count), static_cast<int64_t const*>(K), R, C, rowOrder, stream);
}
else if (type == nvinfer1::DataType::kHALF)
{
nonZeroIndicesImpl<half>(static_cast<half const*>(X), static_cast<int32_t*>(indices),
static_cast<int64_t*>(count), static_cast<int64_t const*>(K), R, C, rowOrder, stream);
}
else
{
ASSERT(false && "Unsupported data type");
}
}
class NonZeroPlugin : public IPluginV3, public IPluginV3OneCore, public IPluginV3OneBuild, public IPluginV3OneRuntime
{
public:
NonZeroPlugin(NonZeroPlugin const& p) = default;
NonZeroPlugin(bool rowOrder)
: mRowOrder(rowOrder)
{
initFieldsToSerialize();
}
void initFieldsToSerialize()
{
mDataToSerialize.clear();
mDataToSerialize.emplace_back(PluginField("rowOrder", &mRowOrder, PluginFieldType::kINT32, 1));
mFCToSerialize.nbFields = mDataToSerialize.size();
mFCToSerialize.fields = mDataToSerialize.data();
}
// IPluginV3 methods
IPluginCapability* getCapabilityInterface(PluginCapabilityType type) noexcept override
{
try
{
if (type == PluginCapabilityType::kBUILD)
{
return static_cast<IPluginV3OneBuild*>(this);
}
if (type == PluginCapabilityType::kRUNTIME)
{
return static_cast<IPluginV3OneRuntime*>(this);
}
ASSERT(type == PluginCapabilityType::kCORE);
return static_cast<IPluginV3OneCore*>(this);
}
catch (std::exception const& e)
{
sample::gLogError << e.what() << std::endl;
}
return nullptr;
}
IPluginV3* clone() noexcept override
{
auto clone = std::make_unique<NonZeroPlugin>(*this);
clone->initFieldsToSerialize();
return clone.release();
}
// IPluginV3OneCore methods
char const* getPluginName() const noexcept override
{
return "NonZeroPlugin";
}
char const* getPluginVersion() const noexcept override
{
return "0";
}
char const* getPluginNamespace() const noexcept override
{
return "";
}
// IPluginV3OneBuild methods
int32_t getNbOutputs() const noexcept override
{
return 2;
}
int32_t configurePlugin(DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out,
int32_t nbOutputs) noexcept override
{
return 0;
}
bool supportsFormatCombination(
int32_t pos, DynamicPluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept override
{
bool typeOk{false};
if (pos == 0)
{
typeOk = inOut[0].desc.type == DataType::kFLOAT || inOut[0].desc.type == DataType::kHALF;
}
else if (pos == 1)
{
typeOk = inOut[1].desc.type == DataType::kINT32;
}
else // pos == 2
{
// size tensor outputs must be NCHW INT64
typeOk = inOut[2].desc.type == DataType::kINT64;
}
return inOut[pos].desc.format == PluginFormat::kLINEAR && typeOk;
}
int32_t getOutputDataTypes(
DataType* outputTypes, int32_t nbOutputs, DataType const* inputTypes, int32_t nbInputs) const noexcept override
{
outputTypes[0] = DataType::kINT32;
outputTypes[1] = DataType::kINT64;
return 0;
}
int32_t getOutputShapes(DimsExprs const* inputs, int32_t nbInputs, DimsExprs const* shapeInputs,
int32_t nbShapeInputs, DimsExprs* outputs, int32_t nbOutputs, IExprBuilder& exprBuilder) noexcept override
{
// The input tensor must be 2-D
if (inputs[0].nbDims != 2)
{
return -1;
}
outputs[0].nbDims = 2;
auto upperBound = exprBuilder.operation(DimensionOperation::kPROD, *inputs[0].d[0], *inputs[0].d[1]);
// On average, we can assume that half of all elements will be non-zero
auto optValue = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *upperBound, *exprBuilder.constant(2));
auto numNonZeroSizeTensor = exprBuilder.declareSizeTensor(1, *optValue, *upperBound);
if (!mRowOrder)
{
outputs[0].d[0] = exprBuilder.constant(2);
outputs[0].d[1] = numNonZeroSizeTensor;
}
else
{
outputs[0].d[0] = numNonZeroSizeTensor;
outputs[0].d[1] = exprBuilder.constant(2);
}
// output at index 1 is a size tensor
outputs[1].nbDims = 0; // size tensors must be declared as 0-D
return 0;
}
// IPluginV3OneRuntime methods
int32_t enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc, void const* const* inputs,
void* const* outputs, void* workspace, cudaStream_t stream) noexcept override
{
int32_t const R = inputDesc[0].dims.d[0];
int32_t const C = inputDesc[0].dims.d[1];
auto type = inputDesc[0].type;
if (!(type == nvinfer1::DataType::kHALF || type == nvinfer1::DataType::kFLOAT))
{
sample::gLogError << "Unsupported: Sample only supports DataType::kHALF and DataType::FLOAT" << std::endl;
return -1;
}
cudaMemsetAsync(outputs[1], 0, sizeof(int64_t), stream);
if (!mRowOrder && workspace == nullptr)
{
sample::gLogError << "Unsupported: workspace is needed but is null" << std::endl;
return -1;
}
if (!mRowOrder)
{
// When constructing a column major output, the kernel needs to be aware of the total number of non-zero
// elements so as to write the non-zero indices at the correct places. Therefore, we will launch the kernel
// twice: first, only to calculate the total non-zero count, which will be stored in workspace; and
// then to actually write the non-zero indices to the outputs[0] buffer.
cudaMemsetAsync(workspace, 0, sizeof(int64_t), stream);
nonZeroIndicesHelper(type, inputs[0], nullptr, workspace, nullptr, R, C, mRowOrder, stream);
nonZeroIndicesHelper(type, inputs[0], outputs[0], outputs[1], workspace, R, C, mRowOrder, stream);
}
else
{
nonZeroIndicesHelper(type, inputs[0], outputs[0], outputs[1], 0, R, C, mRowOrder, stream);
}
return 0;
}
int32_t onShapeChange(
PluginTensorDesc const* in, int32_t nbInputs, PluginTensorDesc const* out, int32_t nbOutputs) noexcept override
{
return 0;
}
IPluginV3* attachToContext(IPluginResourceContext* context) noexcept override
{
return clone();
}
PluginFieldCollection const* getFieldsToSerialize() noexcept override
{
return &mFCToSerialize;
}
size_t getWorkspaceSize(DynamicPluginTensorDesc const* inputs, int32_t nbInputs,
DynamicPluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept override
{
return sizeof(int64_t);
}
private:
bool mRowOrder{true};
std::vector<nvinfer1::PluginField> mDataToSerialize;
nvinfer1::PluginFieldCollection mFCToSerialize;
};
class NonZeroPluginCreator : public nvinfer1::IPluginCreatorV3One
{
public:
NonZeroPluginCreator()
{
mPluginAttributes.clear();
mPluginAttributes.emplace_back(PluginField("rowOrder", nullptr, PluginFieldType::kINT32, 1));
mFC.nbFields = mPluginAttributes.size();
mFC.fields = mPluginAttributes.data();
}
char const* getPluginName() const noexcept override
{
return "NonZeroPlugin";
}
char const* getPluginVersion() const noexcept override
{
return "0";
}
PluginFieldCollection const* getFieldNames() noexcept override
{
return &mFC;
}
IPluginV3* createPlugin(char const* name, PluginFieldCollection const* fc, TensorRTPhase phase) noexcept override
{
try
{
using namespace std::string_view_literals;
bool rowOrder{true};
for (int32_t i = 0; i < fc->nbFields; ++i)
{
std::string_view const fieldName(fc->fields[i].name);
if (fieldName == "rowOrder"sv)
{
rowOrder = *static_cast<bool const*>(fc->fields[i].data);
}
}
return new NonZeroPlugin(rowOrder);
}
catch (std::exception const& e)
{
sample::gLogError << e.what() << std::endl;
}
return nullptr;
}
char const* getPluginNamespace() const noexcept override
{
return "";
}
private:
nvinfer1::PluginFieldCollection mFC;
std::vector<nvinfer1::PluginField> mPluginAttributes;
};
namespace
{
struct NonZeroParams : public samplesCommon::SampleParams
{
bool rowOrder{true};
};
} // namespace
//! \brief The SampleNonZeroPlugin class implements a NonZero plugin
//!
//! \details The plugin is able to output the non-zero indices in row major or column major order
//!
class SampleNonZeroPlugin
{
public:
SampleNonZeroPlugin(NonZeroParams const& params)
: mParams(params)
, mRuntime(nullptr)
, mEngine(nullptr)
{
mSeed = static_cast<uint32_t>(time(nullptr));
}
//!
//! \brief Function builds the network engine
//!
bool build();
//!
//! \brief Runs the TensorRT inference engine for this sample
//!
bool infer();
private:
NonZeroParams mParams; //!< The parameters for the sample.
//! The PluginCreator instance used to create NonZeroPlugin.
//! The instance needs to stay alive across build and infer stages so that the entry in PluginRegistry remains valid
//! throughout.
NonZeroPluginCreator mPluginCreator;
nvinfer1::Dims mInputDims; //!< The dimensions of the input to the network.
nvinfer1::Dims mOutputDims; //!< The dimensions of the output to the network.
std::shared_ptr<nvinfer1::IRuntime> mRuntime; //!< The TensorRT runtime used to deserialize the engine
std::shared_ptr<nvinfer1::ICudaEngine> mEngine; //!< The TensorRT engine used to run the network
uint32_t mSeed{};
//!
//! \brief Creates a TensorRT network and inserts a NonZero plugin
//!
bool constructNetwork(std::unique_ptr<nvinfer1::IBuilder>& builder,
std::unique_ptr<nvinfer1::INetworkDefinition>& network, std::unique_ptr<nvinfer1::IBuilderConfig>& config);
//!
//! \brief Reads the input and stores the result in a managed buffer
//!
bool processInput(samplesCommon::BufferManager const& buffers);
//!
//! \brief Verifies the result
//!
bool verifyOutput(samplesCommon::BufferManager const& buffers);
};
//!
//! \brief Creates the network, configures the builder and creates the network engine
//!
//! \details This function creates a network containing a NonZeroPlugin and builds
//! the engine that will be used to run the plugin (mEngine)
//!
//! \return true if the engine was created successfully and false otherwise
//!
bool SampleNonZeroPlugin::build()
{
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;
}
bool const registered = getPluginRegistry()->registerCreator(mPluginCreator, "");
ASSERT(registered && "Registration of NonZeroPluginCreator failed");
auto constructed = constructNetwork(builder, network, config);
if (!constructed)
{
return false;
}
// CUDA stream used for profiling by the builder.
auto profileStream = samplesCommon::makeCudaStream();
if (!profileStream)
{
return false;
}
config->setProfileStream(*profileStream);
std::unique_ptr<IHostMemory> plan{builder->buildSerializedNetwork(*network, *config)};
if (!plan)
{
return false;
}
mRuntime = std::shared_ptr<nvinfer1::IRuntime>(createInferRuntime(sample::gLogger.getTRTLogger()));
if (!mRuntime)
{
return false;
}
mEngine = std::shared_ptr<nvinfer1::ICudaEngine>(mRuntime->deserializeCudaEngine(plan->data(), plan->size()));
if (!mEngine)
{
return false;
}
ASSERT(network->getNbInputs() == 1);
mInputDims = network->getInput(0)->getDimensions();
ASSERT(mInputDims.nbDims == 2);
ASSERT(network->getNbOutputs() == 2);
mOutputDims = network->getOutput(0)->getDimensions();
ASSERT(mOutputDims.nbDims == 2);
return true;
}
//!
//! \brief Creates a network with a single custom layer containing the NonZero plugin and marks the
//! output layers
//!
//! \param network Pointer to the network that will be populated with the NonZero plugin
//!
//! \param builder Pointer to the engine builder
//!
bool SampleNonZeroPlugin::constructNetwork(std::unique_ptr<nvinfer1::IBuilder>& builder,
std::unique_ptr<nvinfer1::INetworkDefinition>& network, std::unique_ptr<nvinfer1::IBuilderConfig>& config)
{
std::default_random_engine generator(mSeed);
std::uniform_int_distribution<int32_t> distr(10, 25);
int32_t const R = distr(generator);
int32_t const C = distr(generator);
auto* in = network->addInput("Input", DataType::kFLOAT, {2, {R, C}});
ASSERT(in != nullptr);
if (mParams.fp16)
{
auto castLayer = network->addCast(*in, DataType::kHALF);
ASSERT(castLayer != nullptr);
in = castLayer->getOutput(0);
}
std::vector<PluginField> const vecPF{{"rowOrder", &mParams.rowOrder, PluginFieldType::kINT32, 1}};
PluginFieldCollection pfc{static_cast<int32_t>(vecPF.size()), vecPF.data()};
auto pluginCreator = static_cast<IPluginCreatorV3One*>(getPluginRegistry()->getCreator("NonZeroPlugin", "0", ""));
ASSERT(pluginCreator != nullptr && "NonZeroPluginCreator not properly registered to registry");
auto plugin = std::unique_ptr<IPluginV3>(pluginCreator->createPlugin("NonZeroPlugin", &pfc, TensorRTPhase::kBUILD));
ASSERT(plugin != nullptr && "NonZeroPlugin construction failed");
std::vector<ITensor*> inputsVec{in};
auto pluginNonZeroLayer = network->addPluginV3(inputsVec.data(), inputsVec.size(), nullptr, 0, *plugin);
ASSERT(pluginNonZeroLayer != nullptr);
ASSERT(pluginNonZeroLayer->getOutput(0) != nullptr);
ASSERT(pluginNonZeroLayer->getOutput(1) != nullptr);
pluginNonZeroLayer->getOutput(0)->setName("Output0");
pluginNonZeroLayer->getOutput(1)->setName("Output1");
network->markOutput(*(pluginNonZeroLayer->getOutput(0)));
network->markOutput(*(pluginNonZeroLayer->getOutput(1)));
return true;
}
//!
//! \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 and executes the engine.
//!
bool SampleNonZeroPlugin::infer()
{
// Since the data dependent output size cannot be inferred from the engine denote a sufficient size for the
// corresponding output buffer (along with the rest of the I/O tensors)
std::vector<int64_t> ioVolumes = {mInputDims.d[0] * mInputDims.d[1], mInputDims.d[0] * mInputDims.d[1] * 2, 1};
// Create RAII buffer manager object
samplesCommon::BufferManager buffers(mEngine, ioVolumes);
auto context = std::unique_ptr<nvinfer1::IExecutionContext>(mEngine->createExecutionContext());
if (!context)
{
return false;
}
for (int32_t i = 0, e = mEngine->getNbIOTensors(); i < e; ++i)
{
auto const name = mEngine->getIOTensorName(i);
context->setTensorAddress(name, buffers.getDeviceBuffer(name));
}
// Read the input data into the managed buffers
ASSERT(mParams.inputTensorNames.size() == 1);
if (!processInput(buffers))
{
return false;
}
// Create CUDA stream for the execution of this inference.
cudaStream_t stream;
CHECK(cudaStreamCreate(&stream));
// Memcpy from host input buffers to device input buffers
buffers.copyInputToDeviceAsync(stream);
bool status = context->enqueueV3(stream);
if (!status)
{
return false;
}
// Asynchronously copy data from device output buffers to host output buffers.
buffers.copyOutputToHostAsync(stream);
// Wait for the work in the stream to complete.
CHECK(cudaStreamSynchronize(stream));
// Release stream.
CHECK(cudaStreamDestroy(stream));
// Verify results
if (!verifyOutput(buffers))
{
return false;
}
return true;
}
//!
//! \brief Reads the input and stores the result in a managed buffer
//!
bool SampleNonZeroPlugin::processInput(samplesCommon::BufferManager const& buffers)
{
int32_t const inputH = mInputDims.d[0];
int32_t const inputW = mInputDims.d[1];
std::vector<uint8_t> fileData(inputH * inputW);
std::default_random_engine generator(mSeed);
std::uniform_int_distribution<int32_t> distr(0, 9);
auto const number = distr(generator);
samplesCommon::readPGMFile(
samplesCommon::locateFile(std::to_string(number) + ".pgm", mParams.dataDirs), fileData.data(), inputH, inputW);
float* hostDataBuffer = static_cast<float*>(buffers.getHostBuffer(mParams.inputTensorNames[0]));
for (int32_t i = 0; i < inputH * inputW; ++i)
{
auto const raw = 1.0 - float(fileData[i] / 255.0);
hostDataBuffer[i] = raw;
}
sample::gLogInfo << "Input:" << std::endl;
for (int32_t i = 0; i < inputH; ++i)
{
for (int32_t j = 0; j < inputW; ++j)
{
sample::gLogInfo << hostDataBuffer[i * inputW + j];
if (j < inputW - 1)
{
sample::gLogInfo << ", ";
}
}
sample::gLogInfo << std::endl;
}
sample::gLogInfo << std::endl;
return true;
}
//!
//! \brief Verify result
//!
//! \return whether the output correctly identifies all (and only) non-zero elements
//!
// NOLINTNEXTLINE(readability-function-cognitive-complexity)
bool SampleNonZeroPlugin::verifyOutput(samplesCommon::BufferManager const& buffers)
{
float* input = static_cast<float*>(buffers.getHostBuffer(mParams.inputTensorNames[0]));
int32_t* output = static_cast<int32_t*>(buffers.getHostBuffer(mParams.outputTensorNames[0]));
int64_t count = *static_cast<int64_t*>(buffers.getHostBuffer(mParams.outputTensorNames[1]));
std::vector<bool> covered(mInputDims.d[0] * mInputDims.d[1], false);
sample::gLogInfo << "Output:" << std::endl;
if (mParams.rowOrder)
{
for (int32_t i = 0; i < count; ++i)
{
for (int32_t j = 0; j < 2; ++j)
{
sample::gLogInfo << output[j + 2 * i] << " ";
}
sample::gLogInfo << std::endl;
}
}
else
{
for (int32_t i = 0; i < 2; ++i)
{
for (int32_t j = 0; j < count; ++j)
{
sample::gLogInfo << output[j + count * i] << " ";
}
sample::gLogInfo << std::endl;
}
}
if (!mParams.rowOrder)
{
for (int32_t i = 0; i < count; ++i)
{
auto const idx = output[i] * mInputDims.d[1] + output[i + count];
covered[idx] = true;
if (input[idx] == 0.F)
{
return false;
}
}
}
else
{
for (int32_t i = 0; i < count; ++i)
{
auto const idx = output[2 * i] * mInputDims.d[1] + output[2 * i + 1];
covered[idx] = true;
if (input[idx] == 0.F)
{
return false;
}
}
}
for (int32_t i = 0; i < static_cast<int32_t>(covered.size()); ++i)
{
if (!covered[i])
{
if (input[i] != 0.F)
{
return false;
}
}
}
return true;
}
//!
//! \brief Initializes members of the params struct using the command line args
//!
NonZeroParams initializeSampleParams(samplesCommon::Args const& args)
{
NonZeroParams params;
if (args.dataDirs.empty()) // Use default directories if user hasn't provided directory paths
{
params.dataDirs.push_back("data/mnist/");
params.dataDirs.push_back("data/samples/mnist/");
}
else // Use the data directory provided by the user
{
params.dataDirs = args.dataDirs;
}
params.inputTensorNames.push_back("Input");
params.outputTensorNames.push_back("Output0");
params.outputTensorNames.push_back("Output1");
params.fp16 = args.runInFp16;
params.rowOrder = args.rowOrder;
return params;
}
//!
//! \brief Prints the help information for running this sample
//!
void printHelpInfo()
{
std::cout << "Usage: ./sample_non_zero_plugin [-h or --help] [-d or --datadir=<path to data directory>]"
<< std::endl;
std::cout << "--help Display help information" << std::endl;
std::cout << "--datadir Specify path to a data directory, overriding the default. This option can be used "
"multiple times to add multiple directories. If no data directories are given, the default is to use "
"(data/samples/mnist/, data/mnist/)"
<< std::endl;
std::cout << "--fp16 Run in FP16 mode." << std::endl;
std::cout << "--columnOrder Run plugin in column major output mode." << std::endl;
}
int main(int argc, char** argv)
{
samplesCommon::Args args;
bool argsOK = samplesCommon::parseArgs(args, argc, argv);
if (!argsOK)
{
sample::gLogError << "Invalid arguments" << std::endl;
printHelpInfo();
return EXIT_FAILURE;
}
if (args.help)
{
printHelpInfo();
return EXIT_SUCCESS;
}
auto sampleTest = sample::gLogger.defineTest(kSAMPLE_NAME, argc, argv);
sample::gLogger.reportTestStart(sampleTest);
SampleNonZeroPlugin sample(initializeSampleParams(args));
sample::gLogInfo << "Building and running a GPU inference engine for NonZero plugin" << std::endl;
if (!sample.build())
{
return sample::gLogger.reportFail(sampleTest);
}
if (!sample.infer())
{
return sample::gLogger.reportFail(sampleTest);
}
return sample::gLogger.reportPass(sampleTest);
}