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,312 @@
|
||||
/*
|
||||
* 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 <algorithm>
|
||||
#include <iostream>
|
||||
#include <iterator>
|
||||
#include <map>
|
||||
#include <memory>
|
||||
#include <string_view>
|
||||
|
||||
#include "common/serialize.hpp"
|
||||
#include "scatterElementsPluginKernel.h"
|
||||
#include "scatterElementsPluginLegacy.h"
|
||||
|
||||
namespace nvinfer1::plugin
|
||||
{
|
||||
|
||||
std::unordered_map<std::string, ReductionType> const kREDUCE_STR_TO_ENUM{
|
||||
{"add", ReductionType::kSUM},
|
||||
{"mean", ReductionType::kMEAN},
|
||||
{"mul", ReductionType::kMUL},
|
||||
{"min", ReductionType::kMIN},
|
||||
{"max", ReductionType::kMAX},
|
||||
};
|
||||
|
||||
namespace
|
||||
{
|
||||
constexpr char const* kSCATTER_ELEMENTS_NAME{"ScatterElements"};
|
||||
constexpr char const* kSCATTER_ELEMENTS_VERSION{"1"};
|
||||
} // namespace
|
||||
|
||||
ScatterElementsPluginV2::ScatterElementsPluginV2(ReductionType reduction, int32_t dim)
|
||||
: mReduction(reduction)
|
||||
, mAxis(dim)
|
||||
{
|
||||
}
|
||||
|
||||
ScatterElementsPluginV2::ScatterElementsPluginV2(std::string const& reduction, int32_t dim)
|
||||
: mReduction(kREDUCE_STR_TO_ENUM.at(reduction))
|
||||
, mAxis(dim)
|
||||
{
|
||||
}
|
||||
|
||||
ScatterElementsPluginV2::ScatterElementsPluginV2(void const* serialData, size_t serialLength)
|
||||
{
|
||||
deserialize_value(&serialData, &serialLength, &mReduction);
|
||||
deserialize_value(&serialData, &serialLength, &mAxis);
|
||||
}
|
||||
|
||||
int32_t ScatterElementsPluginV2::getNbOutputs() const noexcept
|
||||
{
|
||||
return 1;
|
||||
}
|
||||
|
||||
int32_t ScatterElementsPluginV2::initialize() noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
char const* ScatterElementsPluginV2::getPluginType() const noexcept
|
||||
{
|
||||
return kSCATTER_ELEMENTS_NAME;
|
||||
}
|
||||
|
||||
char const* ScatterElementsPluginV2::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSCATTER_ELEMENTS_VERSION;
|
||||
}
|
||||
|
||||
DimsExprs ScatterElementsPluginV2::getOutputDimensions(
|
||||
int32_t index, DimsExprs const* inputs, int32_t nbInputs, IExprBuilder& exprBuilder) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(nbInputs == 3);
|
||||
PLUGIN_VALIDATE(inputs);
|
||||
PLUGIN_VALIDATE(index <= kOUTPUT_TENSOR_IDX);
|
||||
// both outputs are of the same size
|
||||
DimsExprs out(inputs[kDATA_TENSOR_IDX]);
|
||||
return out;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return DimsExprs();
|
||||
}
|
||||
|
||||
int32_t ScatterElementsPluginV2::enqueue(PluginTensorDesc const* inputDesc, PluginTensorDesc const* outputDesc,
|
||||
void const* const* inputs, void* const* outputs, void* workspace, cudaStream_t stream) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputDesc[kINDICES_TENSOR_IDX].type == DataType::kINT64);
|
||||
|
||||
runScatterElementsKernel(outputs[kOUTPUT_TENSOR_IDX], inputs[kDATA_TENSOR_IDX], inputs[kUPDATES_TENSOR_IDX],
|
||||
inputs[kINDICES_TENSOR_IDX], outputDesc[kOUTPUT_TENSOR_IDX], inputDesc[kDATA_TENSOR_IDX],
|
||||
inputDesc[kUPDATES_TENSOR_IDX], inputDesc[kINDICES_TENSOR_IDX], mAxis, mReduction, stream);
|
||||
return 0;
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
return -1;
|
||||
}
|
||||
}
|
||||
|
||||
size_t ScatterElementsPluginV2::getSerializationSize() const noexcept
|
||||
{
|
||||
auto ret = serialized_size(mReduction) + serialized_size(mAxis);
|
||||
return ret;
|
||||
}
|
||||
|
||||
void ScatterElementsPluginV2::serialize(void* buffer) const noexcept
|
||||
{
|
||||
serialize_value(&buffer, mReduction);
|
||||
serialize_value(&buffer, mAxis);
|
||||
}
|
||||
|
||||
bool ScatterElementsPluginV2::supportsFormatCombination(
|
||||
int32_t pos, PluginTensorDesc const* inOut, int32_t nbInputs, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inOut && pos < (nbInputs + nbOutputs));
|
||||
|
||||
if (inOut[pos].format != PluginFormat::kLINEAR)
|
||||
{
|
||||
return false;
|
||||
}
|
||||
|
||||
auto mytype = inOut[pos].type;
|
||||
auto firsttype = inOut[kDATA_TENSOR_IDX].type;
|
||||
|
||||
// Only INT64 is supported for indices
|
||||
return pos == kINDICES_TENSOR_IDX ? (mytype == DataType::kINT64)
|
||||
: (mytype == firsttype)
|
||||
&& (mytype == DataType::kFLOAT || mytype == DataType::kHALF
|
||||
|| (hasBfloat16AtomicAdd() && mytype == DataType::kBF16) || mytype == DataType::kINT32
|
||||
|| mytype == DataType::kINT64);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
void ScatterElementsPluginV2::terminate() noexcept {}
|
||||
|
||||
void ScatterElementsPluginV2::destroy() noexcept
|
||||
{
|
||||
// This gets called when the network containing plugin is destroyed
|
||||
delete this;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* ScatterElementsPluginV2::clone() const noexcept
|
||||
{
|
||||
auto plugin = std::make_unique<ScatterElementsPluginV2>(mReduction, mAxis);
|
||||
plugin->setPluginNamespace(mNamespace.c_str());
|
||||
return plugin.release();
|
||||
}
|
||||
|
||||
void ScatterElementsPluginV2::configurePlugin(
|
||||
DynamicPluginTensorDesc const* in, int32_t nbInputs, DynamicPluginTensorDesc const* out, int32_t nbOutputs) noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(nbInputs == 3);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
}
|
||||
|
||||
DataType ScatterElementsPluginV2::getOutputDataType(
|
||||
int32_t index, DataType const* inputTypes, int32_t nbInputs) const noexcept
|
||||
{
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(inputTypes && nbInputs == 3 && index == kOUTPUT_TENSOR_IDX);
|
||||
}
|
||||
catch (std::exception const& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return inputTypes[kDATA_TENSOR_IDX];
|
||||
}
|
||||
|
||||
size_t ScatterElementsPluginV2::getWorkspaceSize(
|
||||
PluginTensorDesc const* inputs, int32_t nbInputs, PluginTensorDesc const* outputs, int32_t nbOutputs) const noexcept
|
||||
{
|
||||
return 0;
|
||||
}
|
||||
|
||||
void ScatterElementsPluginV2::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
|
||||
char const* ScatterElementsPluginV2::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
//
|
||||
// ScatterElementsPluginV2Creator
|
||||
//
|
||||
|
||||
ScatterElementsPluginV2Creator::ScatterElementsPluginV2Creator()
|
||||
{
|
||||
gPluginAttributes.clear();
|
||||
gPluginAttributes.emplace_back(PluginField("reduction"));
|
||||
gPluginAttributes.emplace_back(PluginField("axis"));
|
||||
gFC.nbFields = gPluginAttributes.size();
|
||||
gFC.fields = gPluginAttributes.data();
|
||||
}
|
||||
|
||||
char const* ScatterElementsPluginV2Creator::getPluginName() const noexcept
|
||||
{
|
||||
return kSCATTER_ELEMENTS_NAME;
|
||||
}
|
||||
|
||||
char const* ScatterElementsPluginV2Creator::getPluginVersion() const noexcept
|
||||
{
|
||||
return kSCATTER_ELEMENTS_VERSION;
|
||||
}
|
||||
|
||||
PluginFieldCollection const* ScatterElementsPluginV2Creator::getFieldNames() noexcept
|
||||
{
|
||||
return &gFC;
|
||||
}
|
||||
|
||||
char const* ScatterElementsPluginV2Creator::getPluginNamespace() const noexcept
|
||||
{
|
||||
return mNamespace.c_str();
|
||||
}
|
||||
|
||||
void ScatterElementsPluginV2Creator::setPluginNamespace(char const* libNamespace) noexcept
|
||||
{
|
||||
mNamespace = libNamespace;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* ScatterElementsPluginV2Creator::createPlugin(
|
||||
char const* name, PluginFieldCollection const* fc) noexcept
|
||||
{
|
||||
std::string reductionArg;
|
||||
int32_t axisArg = 0;
|
||||
|
||||
try
|
||||
{
|
||||
PLUGIN_VALIDATE(fc != nullptr);
|
||||
auto fields = fc->fields;
|
||||
|
||||
std::set<std::string> requiredFields{"reduction"};
|
||||
plugin::validateRequiredAttributesExist(requiredFields, fc);
|
||||
|
||||
using namespace std::string_view_literals;
|
||||
for (int32_t i = 0; i < fc->nbFields; ++i)
|
||||
{
|
||||
PLUGIN_VALIDATE(fields[i].name != nullptr);
|
||||
PLUGIN_VALIDATE(fields[i].data != nullptr);
|
||||
if (fields[i].name == "axis"sv)
|
||||
{
|
||||
auto data = static_cast<int32_t const*>(fields[i].data);
|
||||
axisArg = *data;
|
||||
}
|
||||
else if (fields[i].name == "reduction"sv)
|
||||
{
|
||||
auto data = static_cast<char const*>(fields[i].data);
|
||||
reductionArg = std::string(data);
|
||||
}
|
||||
}
|
||||
|
||||
PLUGIN_VALIDATE(kREDUCE_STR_TO_ENUM.find(reductionArg) != kREDUCE_STR_TO_ENUM.end(),
|
||||
(reductionArg + ": invalid value for 'reduction' plugin argument").c_str());
|
||||
|
||||
auto plugin = std::make_unique<ScatterElementsPluginV2>(reductionArg, axisArg);
|
||||
plugin->setPluginNamespace(mNamespace.c_str());
|
||||
return plugin.release();
|
||||
}
|
||||
catch (std::exception& e)
|
||||
{
|
||||
caughtError(e);
|
||||
}
|
||||
return nullptr;
|
||||
}
|
||||
|
||||
IPluginV2DynamicExt* ScatterElementsPluginV2Creator::deserializePlugin(
|
||||
char const* name, void const* serialData, size_t serialLength) noexcept
|
||||
{
|
||||
auto plugin = std::make_unique<ScatterElementsPluginV2>(serialData, serialLength);
|
||||
plugin->setPluginNamespace(mNamespace.c_str());
|
||||
return plugin.release();
|
||||
}
|
||||
|
||||
} // namespace nvinfer1::plugin
|
||||
Reference in New Issue
Block a user