wehub-resource-sync
292 lines
8.8 KiB
C++
292 lines
8.8 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 "pillarScatter.h"
|
|
#include "common/templates.h"
|
|
#include <memory>
|
|
#include <string_view>
|
|
|
|
namespace nvinfer1::plugin
|
|
{
|
|
|
|
static char const* const kPLUGIN_VERSION{"1"};
|
|
static char const* const kPLUGIN_NAME{"PillarScatterPlugin"};
|
|
|
|
PillarScatterPlugin::PillarScatterPlugin(size_t h, size_t w)
|
|
: feature_y_size_(h)
|
|
, feature_x_size_(w)
|
|
{
|
|
}
|
|
|
|
PillarScatterPlugin::PillarScatterPlugin(void const* data, size_t length)
|
|
{
|
|
auto const* d = toPointer<char const>(data);
|
|
feature_y_size_ = readFromBuffer<size_t>(d);
|
|
feature_x_size_ = readFromBuffer<size_t>(d);
|
|
}
|
|
|
|
nvinfer1::IPluginV2DynamicExt* PillarScatterPlugin::clone() const noexcept
|
|
{
|
|
try
|
|
{
|
|
auto plugin = std::make_unique<PillarScatterPlugin>(feature_y_size_, feature_x_size_);
|
|
plugin->setPluginNamespace(mNamespace.c_str());
|
|
return plugin.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
nvinfer1::DimsExprs PillarScatterPlugin::getOutputDimensions(int32_t outputIndex, nvinfer1::DimsExprs const* inputs,
|
|
int32_t nbInputs, nvinfer1::IExprBuilder& exprBuilder) noexcept
|
|
{
|
|
PLUGIN_ASSERT(outputIndex == 0);
|
|
nvinfer1::DimsExprs output;
|
|
auto batch_size = inputs[0].d[0];
|
|
output.nbDims = 4;
|
|
output.d[0] = batch_size;
|
|
output.d[1] = inputs[0].d[2];
|
|
output.d[2] = exprBuilder.constant(feature_y_size_);
|
|
output.d[3] = exprBuilder.constant(feature_x_size_);
|
|
return output;
|
|
}
|
|
|
|
bool PillarScatterPlugin::supportsFormatCombination(
|
|
int32_t pos, nvinfer1::PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
|
{
|
|
PLUGIN_ASSERT(nbInputs == 3);
|
|
PLUGIN_ASSERT(nbOutputs == 1);
|
|
PluginTensorDesc const& in = inOut[pos];
|
|
if (pos == 0)
|
|
{
|
|
return (in.type == nvinfer1::DataType::kFLOAT || in.type == nvinfer1::DataType::kHALF)
|
|
&& (in.format == TensorFormat::kLINEAR);
|
|
}
|
|
if (pos == 1)
|
|
{
|
|
return (in.type == nvinfer1::DataType::kINT32) && (in.format == TensorFormat::kLINEAR);
|
|
}
|
|
if (pos == 2)
|
|
{
|
|
return (in.type == nvinfer1::DataType::kINT32) && (in.format == TensorFormat::kLINEAR);
|
|
}
|
|
if (pos == 3)
|
|
{
|
|
return (in.type == inOut[0].type) && (in.format == TensorFormat::kLINEAR);
|
|
}
|
|
return false;
|
|
}
|
|
|
|
void PillarScatterPlugin::configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in, int32_t nbInputs,
|
|
nvinfer1::DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
|
{
|
|
return;
|
|
}
|
|
|
|
size_t PillarScatterPlugin::getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs, int32_t nbInputs,
|
|
nvinfer1::PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
|
{
|
|
return 0;
|
|
}
|
|
int32_t PillarScatterPlugin::enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
|
|
nvinfer1::PluginTensorDesc const* /* outputDesc */, void const* const* inputs, void* const* outputs,
|
|
void* /* workspace */, cudaStream_t stream) noexcept
|
|
{
|
|
try
|
|
{
|
|
PLUGIN_VALIDATE(inputDesc != nullptr && inputs != nullptr && outputs != nullptr);
|
|
|
|
int32_t batchSize = inputDesc[0].dims.d[0];
|
|
int32_t maxPillarNum = inputDesc[0].dims.d[1];
|
|
int32_t numFeatures = inputDesc[0].dims.d[2];
|
|
|
|
nvinfer1::DataType inputType = inputDesc[0].type;
|
|
|
|
auto coords_data = static_cast<uint32_t const*>(inputs[1]);
|
|
auto params_data = static_cast<uint32_t const*>(inputs[2]);
|
|
|
|
uint32_t featureY = feature_y_size_;
|
|
uint32_t featureX = feature_x_size_;
|
|
|
|
int32_t status = -1;
|
|
|
|
if (inputType == nvinfer1::DataType::kHALF)
|
|
{
|
|
auto pillar_features_data = static_cast<half const*>(inputs[0]);
|
|
auto spatial_feature_data = static_cast<half*>(outputs[0]);
|
|
cudaMemsetAsync(
|
|
spatial_feature_data, 0, batchSize * numFeatures * featureY * featureX * sizeof(half), stream);
|
|
status = pillarScatterKernelLaunch<half>(batchSize, maxPillarNum, numFeatures, pillar_features_data,
|
|
coords_data, params_data, featureX, featureY, spatial_feature_data, stream);
|
|
}
|
|
else if (inputType == nvinfer1::DataType::kFLOAT)
|
|
{
|
|
auto const* pillar_features_data = static_cast<float const*>(inputs[0]);
|
|
auto* spatial_feature_data = static_cast<float*>(outputs[0]);
|
|
cudaMemsetAsync(
|
|
spatial_feature_data, 0, batchSize * numFeatures * featureY * featureX * sizeof(float), stream);
|
|
status = pillarScatterKernelLaunch<float>(batchSize, maxPillarNum, numFeatures, pillar_features_data,
|
|
coords_data, params_data, featureX, featureY, spatial_feature_data, stream);
|
|
}
|
|
PLUGIN_ASSERT(status == STATUS_SUCCESS);
|
|
return status;
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return -1;
|
|
}
|
|
|
|
nvinfer1::DataType PillarScatterPlugin::getOutputDataType(
|
|
int32_t index, nvinfer1::DataType const* inputTypes, int32_t nbInputs) const noexcept
|
|
{
|
|
return inputTypes[0];
|
|
}
|
|
|
|
char const* PillarScatterPlugin::getPluginType() const noexcept
|
|
{
|
|
return kPLUGIN_NAME;
|
|
}
|
|
|
|
char const* PillarScatterPlugin::getPluginVersion() const noexcept
|
|
{
|
|
return kPLUGIN_VERSION;
|
|
}
|
|
|
|
int32_t PillarScatterPlugin::getNbOutputs() const noexcept
|
|
{
|
|
return 1;
|
|
}
|
|
|
|
int32_t PillarScatterPlugin::initialize() noexcept
|
|
{
|
|
return 0;
|
|
}
|
|
|
|
void PillarScatterPlugin::terminate() noexcept {}
|
|
|
|
size_t PillarScatterPlugin::getSerializationSize() const noexcept
|
|
{
|
|
return 3 * sizeof(size_t);
|
|
}
|
|
|
|
void PillarScatterPlugin::serialize(void* buffer) const noexcept
|
|
{
|
|
char* d = reinterpret_cast<char*>(buffer);
|
|
writeToBuffer<size_t>(d, feature_y_size_);
|
|
writeToBuffer<size_t>(d, feature_x_size_);
|
|
}
|
|
|
|
void PillarScatterPlugin::destroy() noexcept
|
|
{
|
|
delete this;
|
|
}
|
|
|
|
void PillarScatterPlugin::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* PillarScatterPlugin::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
|
|
PillarScatterPluginCreator::PillarScatterPluginCreator()
|
|
{
|
|
mPluginAttributes.clear();
|
|
mPluginAttributes.emplace_back(PluginField("dense_shape", nullptr, PluginFieldType::kINT32, 1));
|
|
mFC.nbFields = mPluginAttributes.size();
|
|
mFC.fields = mPluginAttributes.data();
|
|
}
|
|
|
|
char const* PillarScatterPluginCreator::getPluginName() const noexcept
|
|
{
|
|
return kPLUGIN_NAME;
|
|
}
|
|
|
|
char const* PillarScatterPluginCreator::getPluginVersion() const noexcept
|
|
{
|
|
return kPLUGIN_VERSION;
|
|
}
|
|
|
|
PluginFieldCollection const* PillarScatterPluginCreator::getFieldNames() noexcept
|
|
{
|
|
return &mFC;
|
|
}
|
|
|
|
IPluginV2* PillarScatterPluginCreator::createPlugin(char const* name, PluginFieldCollection const* fc) noexcept
|
|
{
|
|
try
|
|
{
|
|
PluginField const* fields = fc->fields;
|
|
int32_t nbFields = fc->nbFields;
|
|
int32_t targetH = 0;
|
|
int32_t targetW = 0;
|
|
using namespace std::string_view_literals;
|
|
plugin::validateRequiredAttributesExist({"dense_shape"}, fc);
|
|
for (int32_t i = 0; i < nbFields; ++i)
|
|
{
|
|
if (fields[i].name == "dense_shape"sv)
|
|
{
|
|
int32_t const* ts = static_cast<int32_t const*>(fields[i].data);
|
|
targetH = ts[0];
|
|
targetW = ts[1];
|
|
PLUGIN_VALIDATE(targetH > 0 && targetW > 0);
|
|
}
|
|
}
|
|
auto plugin = std::make_unique<PillarScatterPlugin>(targetH, targetW);
|
|
return plugin.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
IPluginV2* PillarScatterPluginCreator::deserializePlugin(
|
|
char const* name, void const* serialData, size_t serialLength) noexcept
|
|
{
|
|
try
|
|
{
|
|
// This object will be deleted when the network is destroyed,
|
|
auto plugin = std::make_unique<PillarScatterPlugin>(serialData, serialLength);
|
|
return plugin.release();
|
|
}
|
|
catch (std::exception const& e)
|
|
{
|
|
caughtError(e);
|
|
}
|
|
return nullptr;
|
|
}
|
|
|
|
void PillarScatterPluginCreator::setPluginNamespace(char const* libNamespace) noexcept
|
|
{
|
|
mNamespace = libNamespace;
|
|
}
|
|
|
|
char const* PillarScatterPluginCreator::getPluginNamespace() const noexcept
|
|
{
|
|
return mNamespace.c_str();
|
|
}
|
|
} // namespace nvinfer1::plugin
|