wehub-resource-sync
392 lines
11 KiB
C++
392 lines
11 KiB
C++
/*
|
|
* 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 "reorgPlugin.h"
|
|
|
|
#include <memory>
|
|
#include <string_view>
|
|
|
|
namespace nvinfer1::plugin
|
|
{
|
|
static char const* const kREORG_PLUGIN_STATIC_VERSION{"1"};
|
|
static char const* const kREORG_PLUGIN_DYNAMIC_VERSION{"2"};
|
|
static char const* const kREORG_PLUGIN_NAME{"Reorg_TRT"};
|
|
|
|
template <class TBaseClass>
|
|
Reorg<TBaseClass>::Reorg(int32_t strideValue)
|
|
: stride(strideValue)
|
|
{
|
|
}
|
|
|
|
template <class TBaseClass>
|
|
int32_t Reorg<TBaseClass>::getNbOutputs() const noexcept
|
|
{
|
|
return 1;
|
|
}
|
|
|
|
template <class TBaseClass>
|
|
int32_t Reorg<TBaseClass>::initialize() noexcept
|
|
{
|
|
return STATUS_SUCCESS;
|
|
}
|
|
|
|
template <class TBaseClass>
|
|
void Reorg<TBaseClass>::terminate() noexcept
|
|
{
|
|
}
|
|
|
|
template <class TBaseClass>
|
|
char const* Reorg<TBaseClass>::getPluginType() const noexcept
|
|
{
|
|
return kREORG_PLUGIN_NAME;
|
|
}
|
|
|
|
template <class TBaseClass>
|
|
void Reorg<TBaseClass>::destroy() noexcept
|
|
{
|
|
delete this;
|
|
}
|
|
|
|
// Set plugin namespace
|
|
template <class TBaseClass>
|
|
void Reorg<TBaseClass>::setPluginNamespace(char const* pluginNamespace) noexcept
|
|
{
|
|
mPluginNamespace = pluginNamespace;
|
|
}
|
|
|
|
template <class TBaseClass>
|
|
char const* Reorg<TBaseClass>::getPluginNamespace() const noexcept
|
|
{
|
|
return mPluginNamespace.c_str();
|
|
}
|
|
|
|
// Return the DataType of the plugin output at the requested index
|
|
template <class TBaseClass>
|
|
DataType Reorg<TBaseClass>::getOutputDataType(
|
|
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
|
|
{
|
|
// Only 1 input and 1 output from the plugin layer
|
|
PLUGIN_ASSERT(index == 0);
|
|
|
|
// Only DataType::kFLOAT is acceptable by the plugin layer
|
|
return DataType::kFLOAT;
|
|
}
|
|
|
|
// Attach the plugin object to an execution context and grant the plugin the access to some context resource.
|
|
template <class TBaseClass>
|
|
void Reorg<TBaseClass>::attachToContext(
|
|
cudnnContext* cudnnContext, cublasContext* cublasContext, IGpuAllocator* gpuAllocator) noexcept
|
|
{
|
|
}
|
|
|
|
// Detach the plugin object from its execution context.
|
|
template <class TBaseClass>
|
|
void Reorg<TBaseClass>::detachFromContext() noexcept
|
|
{
|
|
}
|
|
|
|
ReorgDynamic::ReorgDynamic(int32_t stride)
|
|
: Reorg<IPluginV2DynamicExt>(stride)
|
|
{
|
|
}
|
|
|
|
ReorgDynamic::ReorgDynamic(void const* buffer, size_t length)
|
|
{
|
|
char const* d = reinterpret_cast<char const*>(buffer);
|
|
char const* a = d;
|
|
stride = read<int32_t>(d);
|
|
PLUGIN_VALIDATE(d == a + length);
|
|
}
|
|
|
|
char const* ReorgDynamic::getPluginVersion() const noexcept
|
|
{
|
|
return kREORG_PLUGIN_DYNAMIC_VERSION;
|
|
}
|
|
|
|
size_t ReorgDynamic::getSerializationSize() const noexcept
|
|
{
|
|
// stride
|
|
return sizeof(int32_t);
|
|
}
|
|
|
|
size_t ReorgDynamic::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
|
|
PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
void ReorgDynamic::serialize(void* buffer) const noexcept
|
|
{
|
|
char *d = reinterpret_cast<char*>(buffer), *a = d;
|
|
write(d, stride);
|
|
PLUGIN_ASSERT(d == a + getSerializationSize());
|
|
}
|
|
|
|
DimsExprs ReorgDynamic::getOutputDimensions(
|
|
int32_t outputIndex, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
|
|
{
|
|
PLUGIN_ASSERT(nbInputs == 1);
|
|
PLUGIN_ASSERT(outputIndex == 0);
|
|
DimsExprs output{3, {}};
|
|
auto const* strideExpr = exprBuilder.constant(stride);
|
|
auto const* strideSquareExpr = exprBuilder.constant(stride * stride);
|
|
output.d[0] = exprBuilder.operation(DimensionOperation::kPROD, *inputs[0].d[0], *strideSquareExpr);
|
|
output.d[1] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[0].d[1], *strideExpr);
|
|
output.d[2] = exprBuilder.operation(DimensionOperation::kFLOOR_DIV, *inputs[0].d[2], *strideExpr);
|
|
return output;
|
|
}
|
|
|
|
bool ReorgDynamic::supportsFormatCombination(
|
|
int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
|
{
|
|
PLUGIN_ASSERT(pos >= 0 && pos <= 1);
|
|
PLUGIN_ASSERT(nbInputs == 1);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
return (inOut[pos].type == DataType::kFLOAT && inOut[pos].format == PluginFormat::kLINEAR);
|
|
}
|
|
|
|
void ReorgDynamic::configurePlugin(
|
|
DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
PLUGIN_ASSERT(nbInputs == 1);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
PLUGIN_ASSERT(in->desc.type == DataType::kFLOAT);
|
|
PLUGIN_ASSERT(out->desc.type == DataType::kFLOAT);
|
|
PLUGIN_ASSERT(in->desc.format == PluginFormat::kLINEAR);
|
|
PLUGIN_ASSERT(out->desc.format == PluginFormat::kLINEAR);
|
|
PLUGIN_ASSERT(stride > 0);
|
|
|
|
int32_t H = in->desc.dims.d[2];
|
|
int32_t W = in->desc.dims.d[3];
|
|
PLUGIN_ASSERT(H % stride == 0);
|
|
PLUGIN_ASSERT(W % stride == 0);
|
|
}
|
|
|
|
int32_t ReorgDynamic::enqueue(nvinfer1::PluginTensorDesc const* inputDesc, nvinfer1::PluginTensorDesc const* outputDesc,
|
|
void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
|
{
|
|
void const* inputData = inputs[0];
|
|
void* outputData = outputs[0];
|
|
int32_t const N = inputDesc[0].dims.d[0];
|
|
int32_t const C = inputDesc[0].dims.d[1];
|
|
int32_t const H = inputDesc[0].dims.d[2];
|
|
int32_t const W = inputDesc[0].dims.d[3];
|
|
pluginStatus_t status = reorgInference(stream, N, C, H, W, stride, inputData, outputData);
|
|
return status;
|
|
}
|
|
|
|
IPluginV2DynamicExt* ReorgDynamic::clone() const noexcept
|
|
{
|
|
try
|
|
{
|
|
auto plugin = std::make_unique<ReorgDynamic>(stride);
|
|
plugin->setPluginNamespace(mPluginNamespace.c_str());
|
|
return plugin.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
ReorgStatic::ReorgStatic(int32_t stride)
|
|
: Reorg<IPluginV2Ext>(stride)
|
|
{
|
|
}
|
|
|
|
ReorgStatic::ReorgStatic(int32_t C, int32_t H, int32_t W, int32_t stride)
|
|
: Reorg<IPluginV2Ext>(stride)
|
|
, C(C)
|
|
, H(H)
|
|
, W(W)
|
|
{
|
|
}
|
|
|
|
ReorgStatic::ReorgStatic(void const* buffer, size_t length)
|
|
{
|
|
char const* d = reinterpret_cast<char const*>(buffer);
|
|
char const* a = d;
|
|
C = read<int32_t>(d);
|
|
H = read<int32_t>(d);
|
|
W = read<int32_t>(d);
|
|
stride = read<int32_t>(d);
|
|
PLUGIN_VALIDATE(d == a + length);
|
|
}
|
|
|
|
char const* ReorgStatic::getPluginVersion() const noexcept
|
|
{
|
|
return kREORG_PLUGIN_STATIC_VERSION;
|
|
}
|
|
|
|
size_t ReorgStatic::getWorkspaceSize(int32_t maxBatchSize) const noexcept
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
size_t ReorgStatic::getSerializationSize() const noexcept
|
|
{
|
|
// C, H, W, stride
|
|
return sizeof(int32_t) * 4;
|
|
}
|
|
|
|
void ReorgStatic::serialize(void* buffer) const noexcept
|
|
{
|
|
char *d = reinterpret_cast<char*>(buffer), *a = d;
|
|
write(d, C);
|
|
write(d, H);
|
|
write(d, W);
|
|
write(d, stride);
|
|
PLUGIN_ASSERT(d == a + getSerializationSize());
|
|
}
|
|
|
|
Dims ReorgStatic::getOutputDimensions(int32_t index, Dims const* inputs, int32_t nbInputDims) noexcept
|
|
{
|
|
PLUGIN_ASSERT(nbInputDims == 1);
|
|
PLUGIN_ASSERT(index == 0);
|
|
return Dims3(inputs[0].d[0] * stride * stride, inputs[0].d[1] / stride, inputs[0].d[2] / stride);
|
|
}
|
|
|
|
int32_t ReorgStatic::enqueue(
|
|
int32_t batchSize, void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
|
{
|
|
void const* inputData = inputs[0];
|
|
void* outputData = outputs[0];
|
|
pluginStatus_t status = reorgInference(stream, batchSize, C, H, W, stride, inputData, outputData);
|
|
return status;
|
|
}
|
|
|
|
bool ReorgStatic::supportsFormat(DataType type, PluginFormat format) const noexcept
|
|
{
|
|
return (type == DataType::kFLOAT && format == PluginFormat::kLINEAR);
|
|
}
|
|
|
|
IPluginV2Ext* ReorgStatic::clone() const noexcept
|
|
{
|
|
try
|
|
{
|
|
auto plugin = std::make_unique<ReorgStatic>(C, H, W, stride);
|
|
plugin->setPluginNamespace(mPluginNamespace.c_str());
|
|
return plugin.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
// Configure the layer with input and output data types.
|
|
void ReorgStatic::configurePlugin(Dims const* inputDims, int32_t nbInputs, Dims const* outputDims, int32_t nbOutputs,
|
|
DataType const* inputTypes, DataType const* outputTypes, bool const* inputIsBroadcast,
|
|
bool const* outputIsBroadcast, PluginFormat floatFormat, int32_t maxBatchSize) noexcept
|
|
{
|
|
PLUGIN_ASSERT(*inputTypes == DataType::kFLOAT && floatFormat == PluginFormat::kLINEAR);
|
|
PLUGIN_ASSERT(nbInputs == 1);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
PLUGIN_ASSERT(stride > 0);
|
|
C = inputDims[0].d[0];
|
|
H = inputDims[0].d[1];
|
|
W = inputDims[0].d[2];
|
|
PLUGIN_ASSERT(H % stride == 0);
|
|
PLUGIN_ASSERT(W % stride == 0);
|
|
}
|
|
|
|
template <class TPluginClass>
|
|
ReorgPluginCreator<TPluginClass>::ReorgPluginCreator()
|
|
{
|
|
mPluginAttributes.clear();
|
|
mPluginAttributes.emplace_back(PluginField("stride", nullptr, PluginFieldType::kINT32, 1));
|
|
|
|
mFC.nbFields = mPluginAttributes.size();
|
|
mFC.fields = mPluginAttributes.data();
|
|
}
|
|
|
|
template <class TPluginClass>
|
|
char const* ReorgPluginCreator<TPluginClass>::getPluginName() const noexcept
|
|
{
|
|
return kREORG_PLUGIN_NAME;
|
|
}
|
|
|
|
template <class TPluginClass>
|
|
char const* ReorgPluginCreator<TPluginClass>::getPluginVersion() const noexcept
|
|
{
|
|
if (std::is_same_v<TPluginClass, ReorgStatic>)
|
|
{
|
|
return kREORG_PLUGIN_STATIC_VERSION;
|
|
}
|
|
else if (std::is_same_v<TPluginClass, ReorgDynamic>)
|
|
{
|
|
return kREORG_PLUGIN_DYNAMIC_VERSION;
|
|
}
|
|
return "";
|
|
}
|
|
|
|
template <class TPluginClass>
|
|
PluginFieldCollection const* ReorgPluginCreator<TPluginClass>::getFieldNames() noexcept
|
|
{
|
|
return &mFC;
|
|
}
|
|
|
|
template <class TPluginClass>
|
|
IPluginV2Ext* ReorgPluginCreator<TPluginClass>::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
|
|
{
|
|
try
|
|
{
|
|
using namespace std::string_view_literals;
|
|
PluginField const* fields = fc->fields;
|
|
PLUGIN_VALIDATE(fc->nbFields == 1);
|
|
PLUGIN_VALIDATE(fields[0].type == PluginFieldType::kINT32);
|
|
PLUGIN_VALIDATE(fields[0].name == "stride"sv);
|
|
stride = static_cast<int32_t>(*(static_cast<int32_t const*>(fields[0].data)));
|
|
|
|
PLUGIN_VALIDATE(stride > 0);
|
|
|
|
auto obj = std::make_unique<TPluginClass>(stride);
|
|
obj->setPluginNamespace(mNamespace.c_str());
|
|
return obj.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
template <class TPluginClass>
|
|
IPluginV2Ext* ReorgPluginCreator<TPluginClass>::deserializePlugin(
|
|
char const* name, void const* serialData, size_t serialLength) noexcept
|
|
{
|
|
try
|
|
{
|
|
// This object will be deleted when the network is destroyed, which will
|
|
// call ReorgPlugin::destroy()
|
|
auto obj = std::make_unique<TPluginClass>(serialData, serialLength);
|
|
obj->setPluginNamespace(mNamespace.c_str());
|
|
return obj.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
template class ReorgPluginCreator<ReorgStatic>;
|
|
template class ReorgPluginCreator<ReorgDynamic>;
|
|
|
|
} // namespace nvinfer1::plugin
|