// 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