wehub/paddlepaddle--paddle
1
0
Fork 0
Code Issues Pull Requests Actions 2 Deployments Models Agent Notes Packages Projects Releases Wiki Activity
Files
main
paddlepaddle--paddle/paddle/fluid/inference/tensorrt/plugin/generic_plugin.cu
T
wehub-resource-sync e25996e7db chore: import upstream snapshot with attribution
2026-07-13 12:40:42 +08:00

731 lines
29 KiB
Plaintext
Raw Permalink Blame History

// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
//
// 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 <map>
#include "paddle/fluid/framework/op_kernel_type.h"
#include "paddle/fluid/framework/phi_utils.h"
#include "paddle/fluid/inference/tensorrt/dynamic_shape_infermeta_registry.h"
#include "paddle/fluid/inference/tensorrt/plugin/generic_plugin.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/common/data_type.h"
#include "paddle/phi/core/compat/op_utils.h"
#include "paddle/phi/core/framework/framework.pb.h"
#include "paddle/phi/core/kernel_context.h"
#include "paddle/phi/core/kernel_factory.h"
#include "paddle/phi/kernels/funcs/data_type_transform.h"
namespace paddle {
namespace inference {
namespace tensorrt {
namespace plugin {
GeneratePluginDataType ProtoTypeToGeneratePluginDataType(
framework::proto::VarType_Type proto_type) {
using framework::proto::VarType_Type;
switch (proto_type) {
case VarType_Type::VarType_Type_BOOL:
return GeneratePluginDataType::PLUGIN_BOOL;
case VarType_Type::VarType_Type_UINT8:
return GeneratePluginDataType::PLUGIN_UINT8;
case VarType_Type::VarType_Type_INT8:
return GeneratePluginDataType::PLUGIN_INT8;
case VarType_Type::VarType_Type_INT16:
return GeneratePluginDataType::PLUGIN_INT16;
case VarType_Type::VarType_Type_INT32:
return GeneratePluginDataType::PLUGIN_INT32;
case VarType_Type::VarType_Type_INT64:
return GeneratePluginDataType::PLUGIN_INT64;
case VarType_Type::VarType_Type_FP16:
return GeneratePluginDataType::PLUGIN_FP16;
case VarType_Type::VarType_Type_FP32:
return GeneratePluginDataType::PLUGIN_FP32;
case VarType_Type::VarType_Type_FP64:
return GeneratePluginDataType::PLUGIN_FP64;
case VarType_Type::VarType_Type_SIZE_T:
return GeneratePluginDataType::PLUGIN_SIZE_T;
case VarType_Type::VarType_Type_BF16:
return GeneratePluginDataType::PLUGIN_BF16;
case VarType_Type::VarType_Type_COMPLEX64:
return GeneratePluginDataType::PLUGIN_COMPLEX64;
case VarType_Type::VarType_Type_COMPLEX128:
return GeneratePluginDataType::PLUGIN_COMPLEX128;
default:
PADDLE_THROW(common::errors::Unimplemented(
"This data type is currently not supported"));
}
}
void BuildPhiKernelContextAttr(const framework::OpDesc& op_desc,
phi::KernelContext* kernel_context,
const phi::KernelSignature& signature,
const phi::Kernel* phi_kernel) {
if (!phi_kernel->IsValid()) {
return;
}
const phi::KernelArgsDef& args_def = phi_kernel->args_def();
const auto& attr_names = signature.attr_names;
const auto& attr_defs = args_def.attribute_defs();
PADDLE_ENFORCE_EQ(
attr_names.size(),
attr_defs.size(),
common::errors::InvalidArgument(
"The attr_names.size() should be equal to attr_defs.size()."));
framework::AttrReader attr_reader(op_desc.GetAttrMap());
for (size_t k = 0; k < attr_names.size(); ++k) {
auto attr_name = attr_names[k];
auto* attr_ptr = attr_reader.GetAttr(attr_name);
if (attr_ptr) {
switch (attr_defs[k].type_index) {
case phi::AttributeType::SCALAR: {
auto& attr = *attr_ptr;
switch (AttrTypeID(attr)) {
case framework::proto::AttrType::FLOAT:
kernel_context->EmplaceBackAttr(
phi::Scalar(PADDLE_GET_CONST(float, attr)));
break;
case framework::proto::AttrType::FLOAT64:
kernel_context->EmplaceBackAttr(
phi::Scalar(PADDLE_GET_CONST(double, attr)));
break;
case framework::proto::AttrType::INT:
kernel_context->EmplaceBackAttr(
phi::Scalar(PADDLE_GET_CONST(int, attr)));
break;
case framework::proto::AttrType::LONG:
kernel_context->EmplaceBackAttr(
phi::Scalar(PADDLE_GET_CONST(int64_t, attr)));
break;
case framework::proto::AttrType::STRING:
kernel_context->EmplaceBackAttr(
phi::Scalar(PADDLE_GET_CONST(std::string, attr)));
break;
case framework::proto::AttrType::SCALAR:
kernel_context->EmplaceBackAttr(phi::Scalar(
PADDLE_GET_CONST(paddle::experimental::Scalar, attr)));
break;
default:
PADDLE_THROW(common::errors::Unimplemented(
"Unsupported cast op attribute `%s` to Scalar when "
"ProtoAttr2PhiAttr.",
attr_name));
}
} break;
case phi::AttributeType::INT_ARRAY: {
auto& attr = *attr_ptr;
switch (AttrTypeID(attr)) {
case framework::proto::AttrType::INTS:
kernel_context->EmplaceBackAttr(std::move(
phi::IntArray(PADDLE_GET_CONST(std::vector<int32_t>, attr))));
break;
case framework::proto::AttrType::LONGS:
kernel_context->EmplaceBackAttr(std::move(
phi::IntArray(PADDLE_GET_CONST(std::vector<int64_t>, attr))));
break;
case framework::proto::AttrType::INT:
kernel_context->EmplaceBackAttr(
phi::IntArray({PADDLE_GET_CONST(int, attr)}));
break;
default:
PADDLE_THROW(common::errors::Unimplemented(
"Unsupported cast op attribute `%s` to IntArray when "
"ProtoAttr2PhiAttr.",
attr_name));
}
} break;
case phi::AttributeType::SCALARS: {
auto& attr = *attr_ptr;
switch (AttrTypeID(attr)) {
case framework::proto::AttrType::INTS: {
const auto& vec = PADDLE_GET_CONST(std::vector<int32_t>, attr);
std::vector<phi::Scalar> scalar_list;
scalar_list.reserve(vec.size());
for (const auto& val : vec) {
scalar_list.emplace_back(val);
}
kernel_context->EmplaceBackAttr(std::move(scalar_list));
} break;
case framework::proto::AttrType::LONGS: {
const auto& vec = PADDLE_GET_CONST(std::vector<int64_t>, attr);
std::vector<phi::Scalar> scalar_list;
scalar_list.reserve(vec.size());
for (const auto& val : vec) {
scalar_list.emplace_back(val);
}
kernel_context->EmplaceBackAttr(std::move(scalar_list));
} break;
case framework::proto::AttrType::FLOATS: {
const auto& vec = PADDLE_GET_CONST(std::vector<float>, attr);
std::vector<phi::Scalar> scalar_list;
scalar_list.reserve(vec.size());
for (const auto& val : vec) {
scalar_list.emplace_back(val);
}
kernel_context->EmplaceBackAttr(std::move(scalar_list));
} break;
case framework::proto::AttrType::FLOAT64S: {
const auto& vec = PADDLE_GET_CONST(std::vector<double>, attr);
std::vector<phi::Scalar> scalar_list;
scalar_list.reserve(vec.size());
for (const auto& val : vec) {
scalar_list.emplace_back(val);
}
kernel_context->EmplaceBackAttr(std::move(scalar_list));
} break;
case framework::proto::AttrType::SCALARS: {
const auto& vec = PADDLE_GET_CONST(
std::vector<paddle::experimental::Scalar>, attr);
std::vector<phi::Scalar> scalar_list;
scalar_list.reserve(vec.size());
for (const auto& val : vec) {
scalar_list.emplace_back(val);
}
kernel_context->EmplaceBackAttr(std::move(scalar_list));
} break;
default:
PADDLE_THROW(common::errors::Unimplemented(
"Unsupported cast op attribute `%s` to vector<Scalar> when "
"ProtoAttr2PhiAttr.",
attr_name));
}
} break;
default: {
auto& attr = *attr_ptr;
switch (attr_defs[k].type_index) {
case phi::AttributeType::FLOAT32:
kernel_context->EmplaceBackAttr(PADDLE_GET_CONST(float, attr));
break;
case phi::AttributeType::FLOAT64:
kernel_context->EmplaceBackAttr(PADDLE_GET_CONST(double, attr));
break;
case phi::AttributeType::INT32:
kernel_context->EmplaceBackAttr(PADDLE_GET_CONST(int, attr));
break;
case phi::AttributeType::BOOL:
kernel_context->EmplaceBackAttr(PADDLE_GET_CONST(bool, attr));
break;
case phi::AttributeType::INT64:
kernel_context->EmplaceBackAttr(PADDLE_GET_CONST(int64_t, attr));
break;
case phi::AttributeType::INT32S:
kernel_context->EmplaceBackAttr(
PADDLE_GET_CONST(std::vector<int>, attr));
break;
case phi::AttributeType::DATA_TYPE: {
auto data_type = phi::TransToPhiDataType(
static_cast<framework::proto::VarType::Type>(
PADDLE_GET_CONST(int, attr)));
kernel_context->EmplaceBackAttr(data_type);
} break;
case phi::AttributeType::STRING:
kernel_context->EmplaceBackAttr(
PADDLE_GET_CONST(std::string, attr));
break;
case phi::AttributeType::INT64S:
switch (AttrTypeID(attr)) {
case framework::proto::AttrType::LONGS:
kernel_context->EmplaceBackAttr(
PADDLE_GET_CONST(std::vector<int64_t>, attr));
break;
case framework::proto::AttrType::INTS: {
const auto& vector_int_attr =
PADDLE_GET_CONST(std::vector<int>, attr);
const std::vector<int64_t> vector_int64_attr(
vector_int_attr.begin(), vector_int_attr.end());
kernel_context->EmplaceBackAttr(vector_int64_attr);
} break;
default:
PADDLE_THROW(common::errors::Unimplemented(
"Unsupported cast op attribute `%s` to vector<int64_t> "
"when ProtoAttr2PhiAttr.",
attr_name));
}
break;
case phi::AttributeType::FLOAT32S:
kernel_context->EmplaceBackAttr(
PADDLE_GET_CONST(std::vector<float>, attr));
break;
case phi::AttributeType::STRINGS:
kernel_context->EmplaceBackAttr(
PADDLE_GET_CONST(std::vector<std::string>, attr));
break;
case phi::AttributeType::BOOLS:
kernel_context->EmplaceBackAttr(
PADDLE_GET_CONST(std::vector<bool>, attr));
break;
case phi::AttributeType::FLOAT64S:
kernel_context->EmplaceBackAttr(
PADDLE_GET_CONST(std::vector<double>, attr));
break;
default:
PADDLE_THROW(common::errors::Unimplemented(
"Unsupported cast op attribute `%s` when construct "
"ProtoAttr2PhiAttr.",
attr_name));
}
}
}
}
}
PADDLE_ENFORCE_EQ(attr_names.size(),
kernel_context->AttrsSize(),
common::errors::InvalidArgument(
"The attr_names.size() should be equal to "
"kernel_context->AttrsSize()."
"Received attr_names.size() = % d,"
"kernel_context->AttrsSize() = %d.",
attr_names.size(),
kernel_context->AttrsSize()));
}
GenericPlugin::GenericPlugin(
const paddle::framework::proto::OpDesc& proto_op_desc,
const InputOutPutVarInfo& in_out_info,
bool with_fp16) {
proto_op_desc_ = proto_op_desc;
op_desc_ = std::move(framework::OpDesc(proto_op_desc_, nullptr));
proto_op_desc_.SerializeToString(&op_meta_data_);
inputs_data_type_ = in_out_info.inputs_data_type;
outputs_data_type_ = in_out_info.outputs_data_type;
with_fp16_ = with_fp16;
}
GenericPlugin::GenericPlugin(
const paddle::framework::proto::OpDesc& proto_op_desc,
const std::vector<GeneratePluginDataType>& inputs_data_type,
const std::vector<GeneratePluginDataType>& outputs_data_type,
bool with_fp16) {
proto_op_desc_ = proto_op_desc;
op_desc_ = std::move(framework::OpDesc(proto_op_desc_, nullptr));
proto_op_desc_.SerializeToString(&op_meta_data_);
inputs_data_type_ = inputs_data_type;
outputs_data_type_ = outputs_data_type;
with_fp16_ = with_fp16;
}
GenericPlugin::GenericPlugin(void const* serial_data, size_t serial_length) {
DeserializeValue(&serial_data, &serial_length, &inputs_data_type_);
DeserializeValue(&serial_data, &serial_length, &outputs_data_type_);
DeserializeValue(&serial_data, &serial_length, &with_fp16_);
std::string op_meta_data((char*)(serial_data), serial_length); // NOLINT
op_meta_data_ = std::move(op_meta_data);
proto_op_desc_.ParseFromString(op_meta_data_);
op_desc_ = std::move(framework::OpDesc(proto_op_desc_, nullptr));
}
int GenericPlugin::getNbOutputs() const TRT_NOEXCEPT {
int res = 0;
for (auto& i : op_desc_.Outputs()) {
if (!i.second.empty()) res += i.second.size();
}
return res;
}
int GenericPlugin::getNbInputs() const TRT_NOEXCEPT {
int res = 0;
for (auto& i : op_desc_.Inputs()) {
if (!i.second.empty()) res += i.second.size();
}
return res;
}
nvinfer1::IPluginV2DynamicExt* GenericPlugin::clone() const TRT_NOEXCEPT {
nvinfer1::IPluginV2DynamicExt* plugin = new GenericPlugin(
proto_op_desc_, inputs_data_type_, outputs_data_type_, with_fp16_);
plugin->initialize();
return plugin;
}
void GenericPlugin::serialize(void* buffer) const TRT_NOEXCEPT {
// inputs_data_type_
SerializeValue(&buffer, inputs_data_type_);
// outputs_data_type_
SerializeValue(&buffer, outputs_data_type_);
// use fp16
SerializeValue(&buffer, with_fp16_);
// serialize op_meta_data_
std::memcpy(buffer, op_meta_data_.c_str(), op_meta_data_.size());
reinterpret_cast<char*&>(buffer) += op_meta_data_.size();
}
bool GenericPlugin::supportsFormatCombination(
int pos,
const nvinfer1::PluginTensorDesc* in_out,
int nb_inputs,
int nb_outputs) TRT_NOEXCEPT {
if (op_desc_.Type() == "gather_nd" || op_desc_.Type() == "yolo_box") {
if (pos == 0)
return (in_out[pos].type == nvinfer1::DataType::kFLOAT ||
(isFp16Supported() &&
in_out[pos].type == nvinfer1::DataType::kHALF)) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
if (pos == 1)
return (in_out[pos].type == nvinfer1::DataType::kINT32) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
// output
if (pos == 2 || pos == 3)
return in_out[0].type == in_out[pos].type &&
in_out[0].format == in_out[pos].format;
} else if (op_desc_.Type() == "scatter_nd_add") {
// input X
if (pos == 0)
return (in_out[pos].type == nvinfer1::DataType::kFLOAT ||
(isFp16Supported() &&
in_out[pos].type == nvinfer1::DataType::kHALF)) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
// input Index
if (pos == 1)
return (in_out[pos].type == nvinfer1::DataType::kINT32) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
// input Updates and output
if (pos == 2 || pos == 3)
return in_out[0].type == in_out[pos].type &&
in_out[0].format == in_out[pos].format;
} else if (op_desc_.Type() == "lookup_table_v2") {
if (pos == 0)
return (in_out[pos].type == nvinfer1::DataType::kINT32 &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR));
if (pos == 1)
return (in_out[pos].type == nvinfer1::DataType::kFLOAT) ||
((isFp16Supported() &&
in_out[pos].type == nvinfer1::DataType::kHALF)) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
// output
if (pos == 2)
return in_out[1].type == in_out[pos].type &&
in_out[1].format == in_out[pos].format;
} else if (op_desc_.Type() == "argsort") {
// input x
if (pos == 0) {
return ((in_out[pos].type == nvinfer1::DataType::kFLOAT ||
(isFp16Supported() &&
in_out[pos].type == nvinfer1::DataType::kHALF)) &&
in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
}
// output out
if (pos == 1) {
return (in_out[pos].type == in_out[0].type &&
in_out[pos].format == in_out[0].format);
}
// output indices
if (pos == 2) {
return (in_out[pos].type == nvinfer1::DataType::kINT32 &&
in_out[pos].format == in_out[0].format);
}
} else if (op_desc_.Type() == "scatter") {
// input X
if (pos == 0)
return (in_out[pos].type == nvinfer1::DataType::kFLOAT ||
(isFp16Supported() &&
in_out[pos].type == nvinfer1::DataType::kHALF)) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
// Ids
if (pos == 1)
return (in_out[pos].type == nvinfer1::DataType::kINT32) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR);
// 3:output 2:input Updates
if (pos == 3 || pos == 2)
return in_out[0].type == in_out[pos].type &&
in_out[0].format == in_out[pos].format;
} else if (op_desc_.Type() == "solve") {
// input X
if (pos == 0)
return in_out[pos].type == nvinfer1::DataType::kFLOAT &&
in_out[pos].format == nvinfer1::TensorFormat::kLINEAR;
// input Y
if (pos == 1)
return in_out[pos].type == nvinfer1::DataType::kFLOAT &&
in_out[pos].format == nvinfer1::TensorFormat::kLINEAR;
// output
if (pos == 2)
return in_out[0].type == in_out[pos].type &&
in_out[0].format == in_out[pos].format;
} else {
return (in_out[pos].type == nvinfer1::DataType::kFLOAT ||
(isFp16Supported() &&
in_out[pos].type == nvinfer1::DataType::kHALF)) &&
(in_out[pos].format == nvinfer1::TensorFormat::kLINEAR) &&
(in_out[0].type == in_out[pos].type);
}
}
nvinfer1::DataType GenericPlugin::getOutputDataType(
int index,
const nvinfer1::DataType* input_types,
int nb_inputs) const TRT_NOEXCEPT {
if (op_desc_.Type() == "lookup_table_v2") {
return input_types[1];
}
if (op_desc_.Type() == "argsort") {
if (index == 1) {
return nvinfer1::DataType::kINT32;
}
}
return input_types[0];
}
int GenericPlugin::initialize() TRT_NOEXCEPT {
std::string op_type = op_desc_.Type();
phi::KernelSignature phi_kernel_signature;
if (phi::OpUtilsMap::Instance().HasArgumentMappingFn(op_type)) {
const phi::ArgumentMappingFn* argument_mapping_func =
phi::OpUtilsMap::Instance().GetArgumentMappingFn(op_type);
PluginArgumentMappingContext argument_mapping_context(&op_desc_);
phi_kernel_signature = (*argument_mapping_func)(argument_mapping_context);
} else {
phi_kernel_signature =
phi::DefaultKernelSignatureMap::Instance().Get(op_type);
}
PADDLE_ENFORCE_EQ(
phi::KernelFactory::Instance().HasCompatiblePhiKernel(op_type),
true,
common::errors::Fatal("%s has no compatible phi kernel!",
op_type.c_str()));
phi::DeviceContextPool& pool = phi::DeviceContextPool::Instance();
GPUPlace place(platform::GetCurrentDeviceId());
auto* dev_ctx = static_cast<phi::GPUContext*>(pool.Get(place));
std::vector<phi::DataType> precision_types{phi::DataType::FLOAT32,
phi::DataType::FLOAT16};
for (auto& precision_type : precision_types) {
phi::KernelKey phi_kernel_key(
phi::Backend::GPU, phi::DataLayout::ANY, precision_type);
auto nv_dtype = PhiType2NvType(precision_type);
phi_kernels_[nv_dtype] = std::make_unique<phi::Kernel>(
phi::KernelFactory::Instance().SelectKernel(phi_kernel_signature.name,
phi_kernel_key));
if (phi_kernel_contexts_.find(nv_dtype) == phi_kernel_contexts_.end() ||
!phi_kernel_contexts_[nv_dtype]) {
phi_kernel_contexts_[nv_dtype] =
std::make_unique<phi::KernelContext>(dev_ctx);
BuildPhiKernelContextAttr(op_desc_,
phi_kernel_contexts_[nv_dtype].get(),
phi_kernel_signature,
phi_kernels_[nv_dtype].get());
}
}
PADDLE_ENFORCE_EQ(phi_kernels_[nvinfer1::DataType::kFLOAT]->IsValid() ||
phi_kernels_[nvinfer1::DataType::kHALF]->IsValid(),
true,
common::errors::Fatal("%s phi kernel is invalid!.",
phi_kernel_signature.name));
if (!dense_tensor_inputs_)
dense_tensor_inputs_ = new std::vector<phi::DenseTensor>(getNbInputs());
if (!dense_tensor_outputs_)
dense_tensor_outputs_ = new std::vector<phi::DenseTensor>(getNbOutputs());
return 0;
}
nvinfer1::DimsExprs GenericPlugin::getOutputDimensions(
int output_index,
const nvinfer1::DimsExprs* inputs,
int nb_inputs,
nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
PADDLE_ENFORCE_EQ(
output_index < getNbOutputs(),
true,
common::errors::InvalidArgument(
"The output_index should be less than getNbOutputs()."));
auto& dynamic_infermeta_factory = tensorrt::DynamicMetaFnFactory::Instance();
PADDLE_ENFORCE_EQ(dynamic_infermeta_factory.Contains(op_desc_.Type()),
true,
common::errors::InvalidArgument(
"The %s op has no dynamic plugin infershape function!",
op_desc_.Type().c_str()));
auto* infershape_func = dynamic_infermeta_factory.Get(op_desc_.Type());
return infershape_func(
output_index, inputs, nb_inputs, expr_builder, op_desc_);
}
void GenericPlugin::configurePlugin(
const nvinfer1::DynamicPluginTensorDesc* in,
int nb_inputs,
const nvinfer1::DynamicPluginTensorDesc* out,
int nb_outputs) TRT_NOEXCEPT {
PADDLE_ENFORCE_EQ(phi_kernels_[nvinfer1::DataType::kFLOAT]->IsValid() ||
phi_kernels_[nvinfer1::DataType::kHALF]->IsValid(),
true,
common::errors::Fatal("Sorry, phi kernel is invalid!"));
PADDLE_ENFORCE_EQ(nb_inputs == getNbInputs(),
true,
common::errors::InvalidArgument(
"The nb_inputs should be equal to getNbInputs()."));
PADDLE_ENFORCE_EQ(nb_outputs == getNbOutputs(),
true,
common::errors::InvalidArgument(
"The nb_outputs should be equal to getNbOutputs()."));
}
// Shutdown the layer. This is called when the engine is destroyed
void GenericPlugin::terminate() TRT_NOEXCEPT {
delete dense_tensor_inputs_;
delete dense_tensor_outputs_;
}
int GenericPlugin::enqueue(const nvinfer1::PluginTensorDesc* input_desc,
const nvinfer1::PluginTensorDesc* output_desc,
const void* const* inputs,
void* const* outputs,
void* workspace,
cudaStream_t stream) TRT_NOEXCEPT {
GPUPlace place(platform::GetCurrentDeviceId());
phi::DeviceContextPool& pool = phi::DeviceContextPool::Instance();
// TODO(inference): generic plugin do not support INT8 precision now.
auto nvType2PhiType =
[&](nvinfer1::DataType nv_dtype) -> std::pair<phi::DataType, int> {
const std::map<nvinfer1::DataType, std::pair<phi::DataType, int>> _map{
{nvinfer1::DataType::kFLOAT, {phi::DataType::FLOAT32, sizeof(float)}},
{nvinfer1::DataType::kHALF, {phi::DataType::FLOAT16, sizeof(half)}},
{nvinfer1::DataType::kINT32, {phi::DataType::INT32, sizeof(int32_t)}},
{nvinfer1::DataType::kBOOL, {phi::DataType::BOOL, sizeof(bool)}},
};
PADDLE_ENFORCE_EQ(
_map.count(nv_dtype),
true,
common::errors::InvalidArgument("Sorry, dtype [ %d ] is not supported.",
static_cast<int>(nv_dtype)));
return _map.at(nv_dtype);
};
nvinfer1::DataType data_type;
// input
if (op_desc_.Type() == "lookup_table_v2") {
data_type = input_desc[1].type;
} else {
data_type = input_desc[0].type;
}
PADDLE_ENFORCE_EQ((data_type == nvinfer1::DataType::kFLOAT) ||
(data_type == nvinfer1::DataType::kHALF),
true,
common::errors::InvalidArgument(
"The data_type should be kFLOAT or kHALF."));
phi_kernel_contexts_[data_type]->ClearInputOutput();
auto* dev_ctx = static_cast<phi::GPUContext*>(pool.Get(place));
phi_kernel_contexts_[data_type]->SetDeviceContext(dev_ctx);
for (int i = 0; i < getNbInputs(); i++) {
if (inputs_data_type_[i] == GeneratePluginDataType::PLUGIN_OPTIONAL) {
phi_kernel_contexts_[data_type]->EmplaceBackInput(nullptr);
continue;
}
auto const& input_dims = input_desc[i].dims;
std::vector<int> input_shape;
for (int j = 0; j < input_dims.nbDims; j++)
input_shape.push_back(input_dims.d[j]);
int input_numel = 1;
for (int k = 0; k < input_shape.size(); k++) input_numel *= input_shape[k];
auto data_type_and_size = nvType2PhiType(input_desc[i].type);
phi::DenseTensorMeta input_meta(data_type_and_size.first,
common::make_ddim(input_shape));
std::shared_ptr<phi::Allocation> input_alloc(
new phi::Allocation((void*)(inputs[i]), // NOLINT
input_numel * data_type_and_size.second,
place));
(*dense_tensor_inputs_)[i] =
std::move(phi::DenseTensor(input_alloc, input_meta));
phi_kernel_contexts_[data_type]->EmplaceBackInput(
&((*dense_tensor_inputs_)[i]));
}
// output
for (int i = 0; i < getNbOutputs(); i++) {
auto const& output_dims = output_desc[i].dims;
std::vector<int> output_shape;
for (int j = 0; j < output_dims.nbDims; j++)
output_shape.push_back(output_dims.d[j]);
int output_numel = 1;
for (int k = 0; k < output_shape.size(); k++)
output_numel *= output_shape[k];
auto data_type_and_size = nvType2PhiType(output_desc[i].type);
phi::DenseTensorMeta output_meta(data_type_and_size.first,
common::make_ddim(output_shape));
std::shared_ptr<phi::Allocation> output_alloc(
new phi::Allocation(reinterpret_cast<void*>(outputs[i]),
output_numel * data_type_and_size.second,
place));
(*dense_tensor_outputs_)[i] =
std::move(phi::DenseTensor(output_alloc, output_meta));
phi_kernel_contexts_[data_type]->EmplaceBackOutput(
&((*dense_tensor_outputs_)[i]));
}
PADDLE_ENFORCE_EQ(
phi_kernel_contexts_[data_type]->InputsSize(),
getNbInputs(),
common::errors::InvalidArgument(
"The phi_kernel_contexts_[data_type]->InputsSize() "
"should be equal to getNbInputs()."
"Received phi_kernel_contexts_[data_type]->InputsSize() "
"= %d, getNbInputs() = %d.",
phi_kernel_contexts_[data_type]->InputsSize(),
getNbInputs()));
PADDLE_ENFORCE_EQ(phi_kernel_contexts_[data_type]->OutputsSize(),
getNbOutputs(),
common::errors::InvalidArgument(
"The phi_kernel_contexts_[data_type]->OutputsSize() "
"should be equal to getNbOutputs()."));
(*phi_kernels_[data_type])(phi_kernel_contexts_[data_type].get());
if (op_desc_.Type() == "argsort") {
for (int i = 0; i < getNbOutputs(); i++) {
phi::DenseTensor& output_tensor = (*dense_tensor_outputs_)[i];
phi::DataType dtype = output_tensor.dtype();
if (dtype == phi::DataType::INT64) {
auto& int32_tensor = output_tensor;
auto ctx = pool.Get(output_tensor.place());
int32_tensor = phi::funcs::TransDataType(
reinterpret_cast<const phi::GPUContext&>(*ctx),
output_tensor,
phi::DataType::INT32);
paddle::memory::Copy(output_tensor.place(),
outputs[i],
output_tensor.place(),
int32_tensor.data<int32_t>(),
int32_tensor.numel() * sizeof(int),
nullptr);
}
}
}
return cudaGetLastError() != cudaSuccess;
}
} // namespace plugin
} // namespace tensorrt
} // namespace inference
} // namespace paddle
Reference in New Issue View Git Blame Copy Permalink