chore: import upstream snapshot with attribution

paddle/fluid/inference/tensorrt/plugin/CMakeLists.txt

@@ -0,0 +1,74 @@
+list(
+  APPEND
+  TRT_FILES
+  trt_plugin.cc
+  split_op_plugin.cu
+  elementwise_op_plugin.cu
+  gelu_op_plugin.cu
+  pool_op_plugin.cu
+  swish_op_plugin.cu
+  group_norm_op_plugin.cu
+  layer_norm_op_plugin.cu
+  instance_norm_op_plugin.cu
+  qkv_to_context_plugin.cu
+  hard_swish_op_plugin.cu
+  stack_op_plugin.cu
+  anchor_generator_op_plugin.cu
+  yolo_box_op_plugin.cu
+  yolo_box_head_op_plugin.cu
+  roi_align_op_plugin.cu
+  gather_nd_op_plugin.cu
+  mish_op_plugin.cu
+  pool3d_op_plugin.cu
+  deformable_conv_op_plugin.cu
+  matmul_op_int8_plugin.cu
+  multihead_matmul_roformer_plugin.cu
+  transformer_input_output_convert_plugin.cu
+  remove_padding_plugin.cu
+  recover_padding_plugin.cu
+  c_allreduce_op_plugin.cu
+  preln_residual_bias_plugin.cu
+  fused_token_prune_op_plugin.cu
+  layernorm_shift_partition_op.cu
+  reverse_roll_op_plugin.cu
+  preln_layernorm_shift_partition_op.cu
+  trans_layernorm_op_plugin.cu
+  merge_layernorm_op_plugin.cu
+  skip_merge_layernorm_op_plugin.cu
+  skip_groupnorm_act_op_plugin.cu
+  preln_groupnorm_act_op_plugin.cu
+  elementwiseadd_transpose_op_plugin.cu
+  generic_plugin.cu
+  custom_generic_plugin.cu
+  many_emb_layernorm_plugin.cu
+  many_emb_layernorm_kernel.cu
+  prompt_tuning_emb_layernorm_varseqlen_kernel_hface.cu
+  prompt_tuning_emb_layernorm_varseqlen_plugin.cu)
+
+list(APPEND TRT_FILES many_emb_layernorm_varseqlen_plugin.cu
+     many_emb_layernorm_varseqlen_kernel_mtron.cu
+     many_emb_layernorm_varseqlen_kernel_hface.cu)
+
+if(CUSPARSELT_FOUND)
+  list(APPEND TRT_FILES spmm_plugin.cu)
+endif()
+
+if(WIN32)
+  nv_library(
+    tensorrt_plugin
+    SRCS ${TRT_FILES}
+    DEPS phi
+         tensorrt_engine
+         tensor
+         common
+         tensorrt_dynamic_shape_infermeta_factory
+         tensorrt_plugin_arg_mapping_context
+         dynload_tensorrt)
+else()
+  nv_library(
+    tensorrt_plugin
+    SRCS ${TRT_FILES}
+    DEPS phi tensorrt_engine tensor common
+         tensorrt_dynamic_shape_infermeta_factory
+         tensorrt_plugin_arg_mapping_context)
+endif()

paddle/fluid/inference/tensorrt/plugin/anchor_generator_op_plugin.h

@@ -0,0 +1,326 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class AnchorGeneratorPlugin : public nvinfer1::IPluginV2Ext {
+ public:
+  explicit AnchorGeneratorPlugin(const nvinfer1::DataType,
+                                 const std::vector<float>& anchor_sizes,
+                                 const std::vector<float>& aspect_ratios,
+                                 const std::vector<float>& stride,
+                                 const std::vector<float>& variances,
+                                 const float offset,
+                                 const int height,
+                                 const int width,
+                                 const int num_anchors,
+                                 const int box_num);
+  AnchorGeneratorPlugin(const void* data, size_t length);
+  ~AnchorGeneratorPlugin() override;
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nb_input_dims) TRT_NOEXCEPT override;
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::TensorFormat format)
+      const TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(int max_batch_size) const TRT_NOEXCEPT override;
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_type,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+  bool isOutputBroadcastAcrossBatch(int output_index,
+                                    const bool* input_is_broadcast,
+                                    int nb_inputs) const TRT_NOEXCEPT override;
+  bool canBroadcastInputAcrossBatch(int input_index) const
+      TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::Dims* input_dims,
+                       int nb_inputs,
+                       const nvinfer1::Dims* output_dims,
+                       int nb_outputs,
+                       const nvinfer1::DataType* input_types,
+                       const nvinfer1::DataType* output_types,
+                       const bool* input_is_broadcast,
+                       const bool* output_is_broadcast,
+                       nvinfer1::PluginFormat float_format,
+                       int max_batch_size) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* clone() const TRT_NOEXCEPT override;
+
+ private:
+  template <typename T>
+  int enqueue_impl(int batch_size,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  nvinfer1::DataType data_type_;
+  std::vector<float> anchor_sizes_;
+  std::vector<float> aspect_ratios_;
+  std::vector<float> stride_;
+  std::vector<float> variances_;
+  float offset_;
+  void* anchor_sizes_device_;
+  void* aspect_ratios_device_;
+  void* stride_device_;
+  void* variances_device_;
+  int height_;
+  int width_;
+  int num_anchors_;
+  int box_num_;
+  std::string namespace_;
+};
+
+class AnchorGeneratorPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  AnchorGeneratorPluginCreator() = default;
+  ~AnchorGeneratorPluginCreator() override = default;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+REGISTER_TRT_PLUGIN_V2(AnchorGeneratorPluginCreator);
+
+class AnchorGeneratorPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit AnchorGeneratorPluginDynamic(const nvinfer1::DataType data_type,
+                                        const std::vector<float>& anchor_sizes,
+                                        const std::vector<float>& aspect_ratios,
+                                        const std::vector<float>& stride,
+                                        const std::vector<float>& variances,
+                                        const float offset,
+                                        const int num_anchors);
+  AnchorGeneratorPluginDynamic(void const* data, size_t length);
+  ~AnchorGeneratorPluginDynamic();
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+
+ private:
+  template <typename T>
+  int enqueue_impl(const nvinfer1::PluginTensorDesc* inputDesc,
+                   const nvinfer1::PluginTensorDesc* outputDesc,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  nvinfer1::DataType data_type_;
+  std::vector<float> anchor_sizes_;
+  std::vector<float> aspect_ratios_;
+  std::vector<float> stride_;
+  std::vector<float> variances_;
+  float offset_;
+  void* anchor_sizes_device_;
+  void* aspect_ratios_device_;
+  void* stride_device_;
+  void* variances_device_;
+  int num_anchors_;
+  std::string namespace_;
+};
+
+class AnchorGeneratorPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  AnchorGeneratorPluginDynamicCreator() = default;
+  ~AnchorGeneratorPluginDynamicCreator() override = default;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+class PIRAnchorGeneratorPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit PIRAnchorGeneratorPluginDynamic(
+      const nvinfer1::DataType data_type,
+      const std::vector<float>& anchor_sizes,
+      const std::vector<float>& aspect_ratios,
+      const std::vector<float>& stride,
+      const std::vector<float>& variances,
+      const float offset,
+      const int num_anchors);
+  PIRAnchorGeneratorPluginDynamic(void const* data, size_t length);
+  ~PIRAnchorGeneratorPluginDynamic();
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+
+ private:
+  template <typename T>
+  int enqueue_impl(const nvinfer1::PluginTensorDesc* inputDesc,
+                   const nvinfer1::PluginTensorDesc* outputDesc,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  nvinfer1::DataType data_type_;
+  std::vector<float> anchor_sizes_;
+  std::vector<float> aspect_ratios_;
+  std::vector<float> stride_;
+  std::vector<float> variances_;
+  float offset_;
+  void* anchor_sizes_device_;
+  void* aspect_ratios_device_;
+  void* stride_device_;
+  void* variances_device_;
+  int num_anchors_;
+  std::string namespace_;
+};
+
+class PIRAnchorGeneratorPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  PIRAnchorGeneratorPluginDynamicCreator() = default;
+  ~PIRAnchorGeneratorPluginDynamicCreator() override = default;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+REGISTER_TRT_PLUGIN_V2(AnchorGeneratorPluginDynamicCreator);
+REGISTER_TRT_PLUGIN_V2(PIRAnchorGeneratorPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/c_allreduce_op_plugin.cu

@@ -0,0 +1,234 @@
+// Copyright (c) 2021 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 <cstring>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/c_allreduce_op_plugin.h"
+#include "paddle/phi/core/distributed/comm_context_manager.h"
+#include "paddle/phi/core/distributed/utils.h"
+#include "paddle/phi/core/platform/collective_helper.h"
+#if defined(PADDLE_WITH_NCCL) || defined(PADDLE_WITH_RCCL)
+#include "paddle/phi/core/distributed/nccl_comm_context.h"
+#endif
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+#if defined(PADDLE_WITH_NCCL)
+inline ncclDataType_t NvInferDtypeToNCCLDType(nvinfer1::DataType type) {
+  if (type == nvinfer1::DataType::kFLOAT) {
+    return ncclFloat;
+  } else if (type == nvinfer1::DataType::kHALF) {
+    return ncclFloat16;
+  } else if (type == nvinfer1::DataType::kINT8) {
+    return ncclInt8;
+  } else if (type == nvinfer1::DataType::kINT32) {
+    return ncclInt32;
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "This datatype in nccl is not supported."));
+  }
+}
+#endif
+
+CAllReducePluginDynamic::CAllReducePluginDynamic(void const* serialData,
+                                                 size_t serialLength) {
+  DeserializeValue(&serialData, &serialLength, &ring_id_);
+  DeserializeValue(&serialData, &serialLength, &use_calc_stream_);
+  DeserializeValue(&serialData, &serialLength, &red_type_);
+  DeserializeValue(&serialData, &serialLength, &with_fp16_);
+}
+nvinfer1::IPluginV2DynamicExt* CAllReducePluginDynamic::clone() const
+    TRT_NOEXCEPT {
+  return new CAllReducePluginDynamic(
+      ring_id_, use_calc_stream_, red_type_, with_fp16_);
+}
+
+const char* CAllReducePluginDynamic::getPluginType() const TRT_NOEXCEPT {
+  return "c_allreduce_plugin_dynamic";
+}
+int CAllReducePluginDynamic::getNbOutputs() const TRT_NOEXCEPT { return 1; }
+int CAllReducePluginDynamic::initialize() TRT_NOEXCEPT { return 0; };
+
+size_t CAllReducePluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(ring_id_) + SerializedSize(use_calc_stream_) +
+         SerializedSize(red_type_);
+  +SerializedSize(with_fp16_);
+}
+
+void CAllReducePluginDynamic::serialize(void* buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, ring_id_);
+  SerializeValue(&buffer, use_calc_stream_);
+  SerializeValue(&buffer, red_type_);
+  SerializeValue(&buffer, with_fp16_);
+}
+
+nvinfer1::DimsExprs CAllReducePluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  return inputs[0];
+}
+
+bool CAllReducePluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of CAllReduce plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc& in = in_out[pos];
+  if (pos == 0 || pos == 1) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+}
+
+void CAllReducePluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+size_t CAllReducePluginDynamic::getWorkspaceSize(
+    const nvinfer1::PluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::PluginTensorDesc* outputs,
+    int nbOutputs) const TRT_NOEXCEPT {
+  return 0;
+}
+
+void CAllReducePluginDynamic::destroy() TRT_NOEXCEPT { delete this; }
+
+nvinfer1::DataType CAllReducePluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The CAllReduce Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+int CAllReducePluginDynamic::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 {
+#if defined(PADDLE_WITH_NCCL)
+  auto input_dims = input_desc[0].dims;
+  size_t numel = ProductDim(input_dims);
+
+  auto input_type = input_desc[0].type;
+  void* sendbuff = const_cast<void*>(inputs[0]);
+  void* recvbuff = outputs[0];
+  ncclDataType_t dtype = NvInferDtypeToNCCLDType(input_type);
+  ncclRedOp_t nccl_red_type = ncclSum;
+  switch (red_type_) {
+    case kRedSum:
+      nccl_red_type = ncclSum;
+      break;
+
+    case kRedMax:
+      nccl_red_type = ncclMax;
+      break;
+
+    case kRedMin:
+      nccl_red_type = ncclMin;
+      break;
+
+    case kRedProd:
+      nccl_red_type = ncclProd;
+      break;
+
+    default:
+      PADDLE_THROW(common::errors::InvalidArgument("Invalid reduce type: %d",
+                                                   red_type_));
+  }
+  const auto& comm_context_manager =
+      phi::distributed::CommContextManager::GetInstance();
+  PADDLE_ENFORCE_EQ(comm_context_manager.Has(std::to_string(ring_id_)),
+                    true,
+                    common::errors::InvalidArgument(
+                        "You choose to use new communication library. "
+                        "But ring_id(%d) is "
+                        "not found in comm_context_manager.",
+                        std::to_string(ring_id_)));
+  auto comm_ctx = static_cast<phi::distributed::NCCLCommContext*>(
+      comm_context_manager.Get(std::to_string(ring_id_)));
+  PADDLE_ENFORCE_NE(comm_ctx,
+                    nullptr,
+                    common::errors::Unavailable(
+                        "NCCLCommContext is nullptr, collective op should "
+                        "has ring_id attr."));
+  auto stream2 = comm_ctx->GetStream();
+  // ncclRedOp_t nccl_red_type = ncclSum;
+  // comm_ctx->AllReduce(&inputs[0], inputs[0], nccl_red_type, stream);
+  phi::dynload::ncclAllReduce(sendbuff,
+                              recvbuff,
+                              numel,
+                              dtype,
+                              nccl_red_type,
+                              comm_ctx->GetNcclComm(),
+                              stream2);
+  VLOG(3) << "new NCCLCommContext has ring_id_ " << ring_id_;
+#endif
+  return (cudaGetLastError() != cudaSuccess);
+}
+
+const char* CAllReducePluginDynamicCreator::getPluginName() const TRT_NOEXCEPT {
+  return "c_allreduce_plugin_dynamic";
+}
+
+const char* CAllReducePluginDynamicCreator::getPluginVersion() const
+    TRT_NOEXCEPT {
+  return "1";
+}
+
+nvinfer1::IPluginV2* CAllReducePluginDynamicCreator::deserializePlugin(
+    const char* name,
+    const void* serial_data,
+    size_t serial_length) TRT_NOEXCEPT {
+  auto plugin = new CAllReducePluginDynamic(serial_data, serial_length);
+  return plugin;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/c_allreduce_op_plugin.h

@@ -0,0 +1,112 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+#include <stdio.h>
+
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+enum ReduceType { kRedSum, kRedMax, kRedMin, kRedProd };
+
+class CAllReducePluginDynamic : public DynamicPluginTensorRT {
+ private:
+  int ring_id_;
+  bool use_calc_stream_;
+  ReduceType red_type_;
+
+ public:
+  explicit CAllReducePluginDynamic(const int ring_id,
+                                   const bool use_calc_stream,
+                                   const ReduceType red_type,
+                                   const bool with_fp16) {
+    ring_id_ = ring_id;
+    use_calc_stream_ = use_calc_stream;
+    red_type_ = red_type;
+    with_fp16_ = with_fp16;
+  }
+  CAllReducePluginDynamic(void const* serialData, size_t serialLength);
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override;
+};
+
+class CAllReducePluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override;
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override;
+};
+
+REGISTER_TRT_PLUGIN_V2(CAllReducePluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/common/bertCommon.h

@@ -0,0 +1,222 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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.
+
+#pragma once
+
+#include <cublas_v2.h>
+#include <cuda_fp16.h>
+#include <cuda_runtime_api.h>
+#include <algorithm>
+#include <cassert>
+#include <memory>
+#include <numeric>
+#include <stdexcept>
+#include <vector>
+#include "NvInfer.h"
+#include "NvInferRuntimeCommon.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/plugin.h"
+
+#define TRT_UNUSED (void)
+#define BERT_PRINT_DEBUG_MSG 0
+#if BERT_PRINT_DEBUG_MSG
+#define TRANSFORMER_DEBUG_MSG(msg) (gLogVerbose << (msg) << std::endl)
+#define BERT_DEBUG_VALUE(key, value) (gLogVerbose << key << value << std::endl)
+#else
+#define TRANSFORMER_DEBUG_MSG(msg) TRT_UNUSED(msg)
+#define BERT_DEBUG_VALUE(key, value) \
+  TRT_UNUSED(key);                   \
+  TRT_UNUSED(value)
+#endif
+using half = __half;
+constexpr uint32_t BDIM = 1;  // batch dimension
+constexpr uint32_t SDIM = 0;  // seq len dimension
+constexpr uint32_t HDIM = 2;  // hidden dimension
+constexpr int32_t kSM_53 = 53;
+constexpr int32_t kSM_70 = 70;
+constexpr int32_t kSM_72 = 72;
+constexpr int32_t kSM_75 = 75;
+constexpr int32_t kSM_80 = 80;
+constexpr int32_t kSM_86 = 86;
+constexpr int32_t kSM_87 = 87;
+constexpr size_t threadsPerCta128 = 2 * 2 * 32;
+constexpr size_t threadsPerCta384 = 1 * 8 * 32;
+
+// The number of xmmas in the M dimension. We use one uint32_t per XMMA in the M
+// dimension: (s + 16*warps_m - 1) / (16*warps_m);
+constexpr size_t xmmasM128 = 4;
+constexpr size_t xmmasM384 = 24;
+
+// Packed mask size per batch. Layout is XMMAS_M * THREADS_PER_CTA.
+constexpr size_t unfusedMaskSize = 1;
+constexpr size_t packedMaskSize64 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize96 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize128 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize384 = xmmasM384 * threadsPerCta384;
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+inline uint32_t getElementSize(nvinfer1::DataType t) noexcept {
+  switch (t) {
+    case nvinfer1::DataType::kINT32:
+      return 4;
+    case nvinfer1::DataType::kFLOAT:
+      return 4;
+    case nvinfer1::DataType::kHALF:
+      return 2;
+    case nvinfer1::DataType::kBOOL:
+    case nvinfer1::DataType::kINT8:
+      return 1;
+    default:
+      return 0;
+  }
+}
+
+inline int64_t getWeightsSize(const nvinfer1::Weights& w,
+                              nvinfer1::DataType type) {
+  return w.count * getElementSize(type);
+}
+
+template <typename T>
+inline void serFromDev(char** buffer, const T* data, size_t nbElem) {
+  const size_t len = sizeof(T) * nbElem;
+  cudaMemcpy(
+      *buffer, static_cast<const void*>(data), len, cudaMemcpyDeviceToHost);
+  *buffer += len;
+}
+
+template <typename T>
+struct CudaDeleter {
+  void operator()(T* buf) { cudaFree(buf); }
+};
+
+template <typename T>
+using cuda_unique_ptr = std::unique_ptr<T, CudaDeleter<T>>;
+
+template <typename T>
+using cuda_shared_ptr = std::shared_ptr<T>;
+
+template <typename T>
+void make_cuda_shared(cuda_shared_ptr<T>* ptr, void* cudaMem) {
+  ptr->reset(static_cast<T*>(cudaMem), CudaDeleter<T>());
+}
+
+struct WeightsWithOwnership : public nvinfer1::Weights {
+  WeightsWithOwnership() {
+    values = nullptr;
+    count = 0;
+  }
+  ~WeightsWithOwnership() { operator delete[](const_cast<void*>(values)); }
+
+  WeightsWithOwnership(const WeightsWithOwnership&) = delete;
+  WeightsWithOwnership operator=(const WeightsWithOwnership&) = delete;
+  WeightsWithOwnership(const WeightsWithOwnership&&) = delete;
+  WeightsWithOwnership operator=(const WeightsWithOwnership&&) = delete;
+
+  void convertAndCopy(const nvinfer1::Weights& src, nvinfer1::DataType type) {
+    this->type = type;
+    this->count = src.count;
+    if (type == nvinfer1::DataType::kFLOAT) {
+      auto destBuf = new float[src.count];
+      this->values = destBuf;
+      if (src.type == nvinfer1::DataType::kFLOAT) {
+        TRANSFORMER_DEBUG_MSG("Float Weights(Host) => Float Array(Host)");
+        std::copy_n(static_cast<const float*>(src.values), src.count, destBuf);
+      } else {
+        assert(src.type == nvinfer1::DataType::kHALF);
+        TRANSFORMER_DEBUG_MSG("Half Weights(Host) => Float Array(Host)");
+        const auto s = static_cast<const half*>(src.values);
+        auto d = static_cast<float*>(const_cast<void*>(this->values));
+        for (auto it = 0; it < src.count; it++) {
+          d[it] = __half2float(s[it]);
+        }
+      }
+    } else if (type == nvinfer1::DataType::kHALF) {
+      auto destBuf = new half[src.count];
+      this->values = destBuf;
+      if (src.type == nvinfer1::DataType::kHALF) {
+        TRANSFORMER_DEBUG_MSG("Half Weights(Host) => Half Array(Host)");
+        std::copy_n(static_cast<const half*>(src.values), src.count, destBuf);
+      } else {
+        assert(src.type == nvinfer1::DataType::kFLOAT);
+        TRANSFORMER_DEBUG_MSG("Float Weights(Host) => Half Array(Host)");
+        const auto s = static_cast<const float*>(src.values);
+        auto d = static_cast<half*>(const_cast<void*>(this->values));
+        for (auto it = 0; it < src.count; it++) {
+          d[it] = __float2half(s[it]);
+        }
+      }
+    } else {
+      throw std::runtime_error("Unsupported DataType specified for plugin.");
+    }
+  }
+
+  void convertAndCopy(const char** srcBuf,
+                      size_t count,
+                      nvinfer1::DataType type) noexcept {
+    this->type = type;
+    this->count = count;
+    const auto nbBytes = getWeightsSize(*this, type);
+    auto destBuf = new char[nbBytes];
+    this->values = destBuf;
+    std::copy_n(*srcBuf, nbBytes, destBuf);
+    *srcBuf += nbBytes;
+  }
+};
+
+template <typename T>
+inline void copyToDevice(WeightsWithOwnership* hostWeights,
+                         size_t nbBytes,
+                         cuda_unique_ptr<T>* cudaWeights) {
+  if (hostWeights->values) {
+    void* cudaMem{nullptr};
+    cudaMalloc(&cudaMem, nbBytes);
+    cudaMemcpy(cudaMem, hostWeights->values, nbBytes, cudaMemcpyHostToDevice);
+    cudaWeights->reset(static_cast<T*>(cudaMem));
+  }
+}
+
+inline nvinfer1::DataType fieldTypeToDataType(
+    const nvinfer1::PluginFieldType ftype) {
+  switch (ftype) {
+    case nvinfer1::PluginFieldType::kFLOAT32: {
+      TRANSFORMER_DEBUG_MSG("PluginFieldType is Float32");
+      return nvinfer1::DataType::kFLOAT;
+    }
+    case nvinfer1::PluginFieldType::kFLOAT16: {
+      TRANSFORMER_DEBUG_MSG("PluginFieldType is Float16");
+      return nvinfer1::DataType::kHALF;
+    }
+    case nvinfer1::PluginFieldType::kINT32: {
+      TRANSFORMER_DEBUG_MSG("PluginFieldType is Int32");
+      return nvinfer1::DataType::kINT32;
+    }
+    case nvinfer1::PluginFieldType::kINT8: {
+      TRANSFORMER_DEBUG_MSG("PluginFieldType is Int8");
+      return nvinfer1::DataType::kINT8;
+    }
+    default:
+      TRANSFORMER_DEBUG_MSG("PluginFieldType is Float32");
+      return nvinfer1::DataType::kFLOAT;
+  }
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/common/common.cuh

@@ -0,0 +1,301 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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.
+
+#pragma once
+
+#include <cublas_v2.h>
+#include <cub/cub.cuh>
+#include "paddle/phi/core/platform/device_context.h"
+
+using kv_float = cub::KeyValuePair<float, float>;
+using kv_half = cub::KeyValuePair<half, half>;
+using kv_half2 = cub::KeyValuePair<half2, half2>;
+
+template <typename T>
+__device__ inline T rsqrt(const T& x);
+
+template <>
+__device__ inline float rsqrt(const float& x) {
+  return rsqrtf(x);
+}
+
+namespace cub {
+__host__ __device__ inline kv_float operator+(const kv_float& a,
+                                              const kv_float& b) {
+  return kv_float(a.key + b.key, a.value + b.value);
+}
+}  // namespace cub
+
+// Half Operations
+__device__ inline half2 __hadd2_with_fallback(const half2 a, const half2 b) {
+#if __CUDA_ARCH__ >= 530
+  return __hadd2(a, b);
+#else
+  float2 out{};
+  out.x = __half2float(a.x) + __half2float(b.x);
+  out.y = __half2float(a.y) + __half2float(b.y);
+  return __float22half2_rn(out);
+#endif
+}
+#if __CUDA_ARCH__ < 530
+template <typename T>
+__device__ inline T operator+(const T& a, const T& b);
+template <typename T>
+__device__ inline T operator*(const T& a, const T& b);
+template <>
+__device__ inline half2 operator+(const half2& a, const half2& b) {
+  return __hadd2_with_fallback(a, b);
+}
+template <>
+__device__ inline half2 operator*(const half2& a, const half2& b) {
+  float2 out{};
+  out.x = __half2float(a.x) * __half2float(b.x);
+  out.y = __half2float(a.y) * __half2float(b.y);
+  return __float22half2_rn(out);
+}
+template <typename T>
+__device__ inline T operator+(const T& a, const T& b);
+template <typename T>
+__device__ inline T operator/(const T& a, const T& b);
+template <typename T>
+__device__ inline T& operator+=(T& a, const T& b);  // NOLINT
+template <typename T>
+__device__ inline T operator-(const T& a, const T& b);
+template <typename T>
+__device__ inline T operator*(const T& a, const T& b);
+template <>
+__device__ inline half operator+(const half& a, const half& b) {
+  return __float2half(__half2float(a) + __half2float(b));
+}
+template <>
+__device__ inline half& operator+=(half& a, const half& b) {  // NOLINT
+  a = __float2half(__half2float(a) + __half2float(b));
+  return a;
+}
+template <>
+__device__ inline half operator-(const half& a, const half& b) {
+  return __float2half(__half2float(a) - __half2float(b));
+}
+template <>
+__device__ inline half operator*(const half& a, const half& b) {
+  return __float2half(__half2float(a) * __half2float(b));
+}
+template <>
+__device__ inline half operator/(const half& a, const half& b) {
+  return __float2half(__half2float(a) / __half2float(b));
+}
+#endif
+
+template <>
+__device__ inline half rsqrt(const half& x) {
+#if __CUDA_ARCH__ >= 530
+  return hrsqrt(x);
+#else
+  return __float2half(rsqrt(__half2float(x)));
+#endif
+}
+
+__device__ inline kv_half operator+(const kv_half& a, const kv_half& b) {
+  const half2 a2 = __halves2half2(a.key, a.value);
+  const half2 b2 = __halves2half2(b.key, b.value);
+  const half2 res = __hadd2_with_fallback(a2, b2);
+  return kv_half(res.x, res.y);
+}
+
+__device__ inline kv_half2 operator+(const kv_half2& a, const kv_half2& b) {
+  return kv_half2(__hadd2_with_fallback(a.key, b.key),
+                  __hadd2_with_fallback(a.value, b.value));
+}
+// Helper Functions
+template <typename T>
+using kvp = cub::KeyValuePair<T, T>;
+template <typename T, typename R, typename P, int TPB>
+__device__ inline void layerNorm(const kvp<R>& threadData,
+                                 const int ld,
+                                 const int offset,
+                                 const P* beta,
+                                 const P* gamma,
+                                 T* output) {
+  // Assuming threadData is already divided by ld
+  using BlockReduce = cub::BlockReduce<kvp<R>, TPB>;
+  __shared__ typename BlockReduce::TempStorage temp_storage;
+  __shared__ R mu;      // mean
+  __shared__ R rsigma;  // 1 / std.dev.
+  const auto sumKV = BlockReduce(temp_storage).Reduce(threadData, cub::Sum());
+  if (threadIdx.x == 0) {
+    mu = sumKV.key;
+    rsigma = rsqrt(sumKV.value - mu * mu);
+  }
+  __syncthreads();
+  for (int i = threadIdx.x; i < ld; i += TPB) {
+    const int idx = offset + i;
+    const R val = output[idx];
+    const R g(gamma[i]);
+    const R b(beta[i]);
+    output[idx] = g * (val - mu) * rsigma + b;
+  }
+}
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+// Helper Functions for multihead related plugins
+template <typename T>
+__global__ void transpose(T* src,
+                          T* dst,
+                          const int batch_size,
+                          const int seq_len,
+                          const int head_num,
+                          const int size_per_head) {
+  int batch_id = blockIdx.x / (head_num * seq_len);
+  int seq_id = blockIdx.x % seq_len;
+  int head_id = (blockIdx.x % (head_num * seq_len)) / seq_len;
+  dst[batch_id * (head_num * seq_len * size_per_head) +
+      seq_id * head_num * size_per_head + head_id * size_per_head +
+      threadIdx.x] = src[blockIdx.x * size_per_head + threadIdx.x];
+}
+
+template <typename T>
+__global__ void TransposeQkvKernel(const int H, const T* input, T* output) {
+  // Input: BxSx3xNxH
+  // Bias: 3xSxB
+  // Output: 3xBxNxSxH
+  int n = threadIdx.y;
+  int s = blockIdx.x;
+  int b = blockIdx.y;
+  int m = blockIdx.z;
+
+  const int N = blockDim.y;
+  const int S = gridDim.x;
+  const int B = gridDim.y;
+
+  const int NH = N * H;
+  const int NHS = NH * S;
+  const int in_offset = n * H + m * NH + s * 3 * NH + b * NHS * 3;
+  const int out_offset = s * H + n * S * H + b * NHS + m * NHS * B;
+
+  const int i = threadIdx.x;
+  output[out_offset + i] = input[in_offset + i];
+}
+
+inline void TransposeQKV(const int batch,
+                         const int seq_len,
+                         const int head_size,
+                         const int head_num,
+                         const float* input,
+                         float* output,
+                         cudaStream_t stream) {
+  int scratch_size = batch * head_num * seq_len * seq_len;
+  const dim3 grid(seq_len, batch, 3);
+  if (head_size % 4 == 0 && scratch_size % 4 == 0) {
+    const int h = head_size / 4;
+    const float4* input4 = reinterpret_cast<const float4*>(input);
+    float4* output4 = reinterpret_cast<float4*>(output);
+    const dim3 block(h, head_num, 1);
+    // limit h * head_num to max block size(1024).
+    PADDLE_ENFORCE_LE(h * head_num,
+                      1024,
+                      common::errors::InvalidArgument(
+                          "head_num (%d) * head_size (%d) should <= %d",
+                          head_num,
+                          head_size,
+                          1024 * 4));
+    TransposeQkvKernel<float4><<<grid, block, 0, stream>>>(h, input4, output4);
+  } else if (head_size % 2 == 0 && scratch_size % 2 == 0) {
+    const int h = head_size / 2;
+    const float2* input2 = reinterpret_cast<const float2*>(input);
+    float2* output2 = reinterpret_cast<float2*>(output);
+    const dim3 block(h, head_num, 1);
+    // limit h * head_num to max block size(1024).
+    PADDLE_ENFORCE_LE(h * head_num,
+                      1024,
+                      common::errors::InvalidArgument(
+                          "head_num (%d) * head_size (%d) should <= %d",
+                          head_num,
+                          head_size,
+                          1024 * 2));
+    TransposeQkvKernel<float2><<<grid, block, 0, stream>>>(h, input2, output2);
+  } else {
+    const dim3 block(head_size, head_num, 1);
+    // limit head_size * head_num to max block size(1024).
+    PADDLE_ENFORCE_LE(head_size * head_num,
+                      1024,
+                      common::errors::InvalidArgument(
+                          "head_num (%d) * head_size (%d) should <= %d",
+                          head_num,
+                          head_size,
+                          1024));
+    TransposeQkvKernel<float>
+        <<<grid, block, 0, stream>>>(head_size, input, output);
+  }
+}
+
+inline void TransposeQKV(const int batch,
+                         const int seq_len,
+                         const int head_size,
+                         const int head_num,
+                         const half* input,
+                         half* output,
+                         cudaStream_t stream) {
+  int scratch_size = batch * head_num * seq_len * seq_len;
+  const dim3 grid(seq_len, batch, 3);
+  if (head_size % 8 == 0 && scratch_size % 8 == 0) {
+    int h = head_size / 8;
+    const int4* input4 = reinterpret_cast<const int4*>(input);
+    int4* output4 = reinterpret_cast<int4*>(output);
+    dim3 block(h, head_num, 1);
+    // limit h * head_num to max block size(1024).
+    PADDLE_ENFORCE_LE(h * head_num,
+                      1024,
+                      common::errors::InvalidArgument(
+                          "head_num (%d) * head_size (%d) should <= %d",
+                          head_num,
+                          head_size,
+                          1024 * 8));
+    TransposeQkvKernel<int4><<<grid, block, 0, stream>>>(h, input4, output4);
+  } else if (head_size % 2 == 0 && scratch_size % 2 == 0) {
+    const int h = head_size / 2;
+    const half2* input2 = reinterpret_cast<const half2*>(input);
+    half2* output2 = reinterpret_cast<half2*>(output);
+    const dim3 block(h, head_num, 1);
+    // limit h * head_num to max block size(1024).
+    PADDLE_ENFORCE_LE(h * head_num,
+                      1024,
+                      common::errors::InvalidArgument(
+                          "head_num (%d) * head_size (%d) should <= %d",
+                          head_num,
+                          head_size,
+                          1024 * 2));
+    TransposeQkvKernel<half2><<<grid, block, 0, stream>>>(h, input2, output2);
+  } else {
+    const dim3 block(head_size, head_num, 1);
+    // limit head_size * head_num to max block size(1024).
+    PADDLE_ENFORCE_LE(head_size * head_num,
+                      1024,
+                      common::errors::InvalidArgument(
+                          "head_num (%d) * head_size (%d) should <= %d",
+                          head_num,
+                          head_size,
+                          1024));
+    TransposeQkvKernel<half>
+        <<<grid, block, 0, stream>>>(head_size, input, output);
+  }
+}
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/common/plugin.h

@@ -0,0 +1,62 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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.
+
+#pragma once
+
+#include <cuda_runtime.h>
+#include <cstring>
+#include <iostream>
+#include <memory>
+#include <sstream>
+#include <string>
+#include "NvInfer.h"
+#include "NvInferPlugin.h"
+
+typedef enum {
+  STATUS_SUCCESS = 0,
+  STATUS_FAILURE = 1,
+  STATUS_BAD_PARAM = 2,
+  STATUS_NOT_SUPPORTED = 3,
+  STATUS_NOT_INITIALIZED = 4
+} pluginStatus_t;
+
+namespace nvinfer1 {
+
+class BasePlugin : public IPluginV2 {
+ protected:
+  void setPluginNamespace(const char* libNamespace) noexcept override {
+    mNamespace = libNamespace;
+  }
+  const char* getPluginNamespace() const noexcept override {
+    return mNamespace.c_str();
+  }
+  std::string mNamespace;
+};
+
+class BaseCreator : public IPluginCreator {
+ public:
+  void setPluginNamespace(const char* libNamespace) noexcept override {
+    mNamespace = libNamespace;
+  }
+  const char* getPluginNamespace() const noexcept override {
+    return mNamespace.c_str();
+  }
+
+ protected:
+  std::string mNamespace;
+};
+
+}  // namespace nvinfer1

paddle/fluid/inference/tensorrt/plugin/common/serialize.h

@@ -0,0 +1,119 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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.
+
+#pragma once
+
+#include <cassert>
+#include <cstring>
+#include <type_traits>
+#include <vector>
+
+#include <iostream>
+using std::cerr;
+using std::cout;
+using std::endl;
+
+template <typename T>
+inline void serialize_value(void** buffer, T const& value);
+
+template <typename T>
+inline void deserialize_value(void const** buffer,
+                              size_t* buffer_size,
+                              T* value);
+
+template <typename T, class Enable = void>
+struct Serializer {};
+
+template <typename T>
+struct Serializer<T,
+                  typename std::enable_if<std::is_arithmetic<T>::value ||
+                                          std::is_enum<T>::value ||
+                                          std::is_pod<T>::value>::type> {
+  static size_t serialized_size(T const&) { return sizeof(T); }
+  static void serialize(void** buffer, T const& value) {
+    ::memcpy(*buffer, &value, sizeof(T));
+    reinterpret_cast<char*&>(*buffer) += sizeof(T);
+  }
+  static void deserialize(void const** buffer, size_t* buffer_size, T* value) {
+    assert(*buffer_size >= sizeof(T));
+    ::memcpy(value, *buffer, sizeof(T));
+    reinterpret_cast<char const*&>(*buffer) += sizeof(T);
+    *buffer_size -= sizeof(T);
+  }
+};
+
+template <>
+struct Serializer<const char*> {
+  static size_t serialized_size(const char* value) { return strlen(value) + 1; }
+  static void serialize(void** buffer, const char* value) {
+    ::strcpy(static_cast<char*>(*buffer), value);  // NOLINT
+    reinterpret_cast<char*&>(*buffer) += strlen(value) + 1;
+  }
+  static void deserialize(void const** buffer,
+                          size_t* buffer_size,
+                          const char** value) {
+    *value = static_cast<char const*>(*buffer);
+    size_t data_size = strnlen(*value, *buffer_size) + 1;
+    assert(*buffer_size >= data_size);
+    reinterpret_cast<char const*&>(*buffer) += data_size;
+    *buffer_size -= data_size;
+  }
+};
+
+template <typename T>
+struct Serializer<std::vector<T>,
+                  typename std::enable_if<std::is_arithmetic<T>::value ||
+                                          std::is_enum<T>::value ||
+                                          std::is_pod<T>::value>::type> {
+  static size_t serialized_size(std::vector<T> const& value) {
+    return sizeof(value.size()) + value.size() * sizeof(T);
+  }
+  static void serialize(void** buffer, std::vector<T> const& value) {
+    serialize_value(buffer, value.size());
+    size_t nbyte = value.size() * sizeof(T);
+    ::memcpy(*buffer, value.data(), nbyte);
+    reinterpret_cast<char*&>(*buffer) += nbyte;
+  }
+  static void deserialize(void const** buffer,
+                          size_t* buffer_size,
+                          std::vector<T>* value) {
+    size_t size;
+    deserialize_value(buffer, buffer_size, &size);
+    value->resize(size);
+    size_t nbyte = value->size() * sizeof(T);
+    assert(*buffer_size >= nbyte);
+    ::memcpy(value->data(), *buffer, nbyte);
+    reinterpret_cast<char const*&>(*buffer) += nbyte;
+    *buffer_size -= nbyte;
+  }
+};
+
+template <typename T>
+inline size_t serialized_size(T const& value) {
+  return Serializer<T>::serialized_size(value);
+}
+
+template <typename T>
+inline void serialize_value(void** buffer, T const& value) {
+  return Serializer<T>::serialize(buffer, value);
+}
+
+template <typename T>
+inline void deserialize_value(void const** buffer,
+                              size_t* buffer_size,
+                              T* value) {
+  return Serializer<T>::deserialize(buffer, buffer_size, value);
+}

paddle/fluid/inference/tensorrt/plugin/custom_generic_plugin.cu

@@ -0,0 +1,668 @@
+// Copyright (c) 2023 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 "paddle/fluid/inference/tensorrt/plugin/custom_generic_plugin.h"
+#include "paddle/common/enforce.h"
+#include "paddle/fluid/framework/op_kernel_type.h"
+#include "paddle/fluid/framework/phi_utils.h"
+#include "paddle/fluid/inference/tensorrt/op_teller.h"
+#include "paddle/phi/api/include/tensor.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"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+void validate(const std::string& op_type,
+              const std::string& datatype,
+              const std::string& tensor_format) {
+  std::unordered_set<std::string> supports_dtypes = {
+      "float32", "float16", "int8", "int32"};
+  std::unordered_set<std::string> supports_tensor_formats = {
+      "LINEAR", "CHW32", "CHW2", "HWC8", "CHW4"};
+  supports_tensor_formats.insert("DHWC8");
+  supports_tensor_formats.insert("HWC16");
+  // refer to
+  // https://docs.nvidia.com/deeplearning/tensorrt/developer-guide/index.html#ipluginv2
+  PADDLE_ENFORCE_GE(supports_dtypes.count(datatype),
+                    0,
+                    common::errors::InvalidArgument(
+                        "custom op [%s] has unsupported datatype: [%s], "
+                        "now only support: [float32, float16, int8, int32].",
+                        op_type,
+                        datatype));
+  PADDLE_ENFORCE_GE(
+      supports_tensor_formats.count(tensor_format),
+      0,
+      common::errors::InvalidArgument(
+          "custom op [%s] has unsupported tensor format: [%s], "
+          "now only support: [LINEAR, CHW32, CHW2, HWC8, CHW4, DHWC8(TensorRT "
+          "7.2 and after), HWC16(TensorRT 8.0 and after)].",
+          op_type,
+          tensor_format));
+
+  if (datatype == "float32") {
+    std::unordered_set<std::string> supports_formats_tmp = {"LINEAR", "CHW32"};
+    PADDLE_ENFORCE_GE(
+        supports_formats_tmp.count(tensor_format),
+        0,
+        common::errors::InvalidArgument(
+            "custom op [%s]: float32 only supports [LINEAR, CHW32], "
+            "but got tensor format: [%s], ",
+            op_type,
+            tensor_format));
+  }
+  if (datatype == "float16") {
+    std::unordered_set<std::string> supports_formats_tmp = {
+        "LINEAR", "CHW2", "HWC8", "CHW4"};
+    supports_formats_tmp.insert("DHWC8");
+    supports_formats_tmp.insert("HWC16");
+    PADDLE_ENFORCE_GE(supports_formats_tmp.count(tensor_format),
+                      0,
+                      common::errors::InvalidArgument(
+                          "custom op [%s]: float16 only supports [LINEAR, "
+                          "CHW2, HWC8, CHW4, DHWC8(TensorRT 7.2 and after), "
+                          "HWC16(TensorRT 8.0 and after)], "
+                          "but got tensor format: [%s], ",
+                          op_type,
+                          tensor_format));
+  }
+  if (datatype == "int8") {
+    std::unordered_set<std::string> supports_formats_tmp = {
+        "LINEAR", "CHW32", "CHW4"};
+    PADDLE_ENFORCE_GE(
+        supports_formats_tmp.count(tensor_format),
+        0,
+        common::errors::InvalidArgument(
+            "custom op [%s]: int8 only supports [LINEAR, CHW32, CHW4], "
+            "but got tensor format: [%s], ",
+            op_type,
+            tensor_format));
+  }
+  if (datatype == "int32") {
+    std::unordered_set<std::string> supports_formats_tmp = {"LINEAR"};
+    PADDLE_ENFORCE_GE(supports_formats_tmp.count(tensor_format),
+                      0,
+                      common::errors::InvalidArgument(
+                          "custom op [%s]: int32 only supports [LINEAR], "
+                          "but got tensor format: [%s], ",
+                          op_type,
+                          tensor_format));
+  }
+}
+
+std::vector<std::pair<std::string, std::string>> parseConfig(
+    const std::string& op_type, const std::string& config) {
+  std::vector<std::pair<std::string, std::string>> res;
+  size_t start = 0;
+  size_t seg = config.find("+", start);
+  while (seg != std::string::npos) {
+    std::string dtype_format = config.substr(start, seg - start);
+    size_t split_pos = dtype_format.find(":");
+    std::string dtype = dtype_format.substr(0, split_pos);
+    std::string format;
+    if (split_pos == std::string::npos) {
+      format = "LINEAR";
+    } else {
+      format = dtype_format.substr(split_pos + 1);
+    }
+    transform(dtype.begin(), dtype.end(), dtype.begin(), ::tolower);
+    transform(format.begin(), format.end(), format.begin(), ::toupper);
+    validate(op_type, dtype, format);
+    res.emplace_back(dtype, format);
+    start = seg + 1;
+    seg = config.find("+", start);
+  }
+  std::string dtype_format = config.substr(start);
+  size_t split_pos = dtype_format.find(":");
+  std::string dtype = dtype_format.substr(0, split_pos);
+  std::string format;
+  if (split_pos == std::string::npos) {
+    format = "LINEAR";
+  } else {
+    format = dtype_format.substr(split_pos + 1);
+  }
+  transform(dtype.begin(), dtype.end(), dtype.begin(), ::tolower);
+  transform(format.begin(), format.end(), format.begin(), ::toupper);
+  validate(op_type, dtype, format);
+  res.emplace_back(dtype, format);
+  return res;
+}
+
+nvinfer1::DataType getTrtDtype(std::string dtype) {
+  if (dtype == "float32") {
+    return nvinfer1::DataType::kFLOAT;
+  } else if (dtype == "float16") {
+    return nvinfer1::DataType::kHALF;
+  } else if (dtype == "int8") {
+    return nvinfer1::DataType::kINT8;
+  } else if (dtype == "int32") {
+    return nvinfer1::DataType::kINT32;
+  } else {
+    PADDLE_THROW(
+        common::errors::Unimplemented("Unsupported data type [%s]", dtype));
+  }
+}
+
+nvinfer1::TensorFormat getTrtTensorFormat(std::string tensor_format) {
+  if (tensor_format == "LINEAR") {
+    return nvinfer1::TensorFormat::kLINEAR;
+  } else if (tensor_format == "CHW32") {
+    return nvinfer1::TensorFormat::kCHW32;
+  } else if (tensor_format == "CHW2") {
+    return nvinfer1::TensorFormat::kCHW2;
+  } else if (tensor_format == "HWC8") {
+    return nvinfer1::TensorFormat::kHWC8;
+  } else if (tensor_format == "CHW4") {
+    return nvinfer1::TensorFormat::kCHW4;
+  } else if (tensor_format == "DHWC8") {
+    return nvinfer1::TensorFormat::kDHWC8;
+  } else if (tensor_format == "HWC16") {
+    return nvinfer1::TensorFormat::kHWC16;
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented("Unsupported tensor format [%s]",
+                                               tensor_format));
+  }
+}
+
+GenerateCustomGenericPluginDataType
+ProtoTypeToGenerateCustomGenericPluginDataType(
+    framework::proto::VarType_Type proto_type) {
+  using framework::proto::VarType_Type;
+  switch (proto_type) {
+    case VarType_Type::VarType_Type_BOOL:
+      return GenerateCustomGenericPluginDataType::PLUGIN_BOOL;
+    case VarType_Type::VarType_Type_UINT8:
+      return GenerateCustomGenericPluginDataType::PLUGIN_UINT8;
+    case VarType_Type::VarType_Type_INT8:
+      return GenerateCustomGenericPluginDataType::PLUGIN_INT8;
+    case VarType_Type::VarType_Type_INT16:
+      return GenerateCustomGenericPluginDataType::PLUGIN_INT16;
+    case VarType_Type::VarType_Type_INT32:
+      return GenerateCustomGenericPluginDataType::PLUGIN_INT32;
+    case VarType_Type::VarType_Type_INT64:
+      return GenerateCustomGenericPluginDataType::PLUGIN_INT64;
+    case VarType_Type::VarType_Type_FP16:
+      return GenerateCustomGenericPluginDataType::PLUGIN_FP16;
+    case VarType_Type::VarType_Type_FP32:
+      return GenerateCustomGenericPluginDataType::PLUGIN_FP32;
+    case VarType_Type::VarType_Type_FP64:
+      return GenerateCustomGenericPluginDataType::PLUGIN_FP64;
+    case VarType_Type::VarType_Type_SIZE_T:
+      return GenerateCustomGenericPluginDataType::PLUGIN_SIZE_T;
+    case VarType_Type::VarType_Type_BF16:
+      return GenerateCustomGenericPluginDataType::PLUGIN_BF16;
+    case VarType_Type::VarType_Type_COMPLEX64:
+      return GenerateCustomGenericPluginDataType::PLUGIN_COMPLEX64;
+    case VarType_Type::VarType_Type_COMPLEX128:
+      return GenerateCustomGenericPluginDataType::PLUGIN_COMPLEX128;
+    default:
+      PADDLE_THROW(common::errors::Unimplemented(
+          "This data type is currently not supported"));
+  }
+}
+
+CustomGenericPlugin::CustomGenericPlugin(
+    const paddle::framework::proto::OpDesc& proto_op_desc,
+    const InputOutPutVarInfo& in_out_info,
+    bool with_fp16) {
+  proto_op_desc_ = proto_op_desc;
+  op_desc_ = 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;
+}
+
+CustomGenericPlugin::CustomGenericPlugin(
+    const paddle::framework::proto::OpDesc& proto_op_desc,
+    const std::vector<GenerateCustomGenericPluginDataType>& inputs_data_type,
+    const std::vector<GenerateCustomGenericPluginDataType>& outputs_data_type,
+    bool with_fp16) {
+  proto_op_desc_ = proto_op_desc;
+  op_desc_ = 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;
+}
+
+CustomGenericPlugin::CustomGenericPlugin(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_ = framework::OpDesc(proto_op_desc_, nullptr);
+}
+
+int CustomGenericPlugin::getNbOutputs() const TRT_NOEXCEPT {
+  int res = 0;
+  for (auto& i : op_desc_.Outputs()) {
+    if (!i.second.empty()) res += i.second.size();
+  }
+  return res;
+}
+
+int CustomGenericPlugin::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* CustomGenericPlugin::clone() const TRT_NOEXCEPT {
+  nvinfer1::IPluginV2DynamicExt* plugin = new CustomGenericPlugin(
+      proto_op_desc_, inputs_data_type_, outputs_data_type_, with_fp16_);
+  plugin->initialize();
+  return plugin;
+}
+
+void CustomGenericPlugin::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 CustomGenericPlugin::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  auto& op_meta_info_map = OpMetaInfoMap::Instance();
+  const auto& meta_info_map = op_meta_info_map.GetMap();
+  auto& op_info = meta_info_map.at(op_desc_.Type()).front();
+  auto& supports_format_config =
+      OpMetaInfoHelper::GetTrtSupportsFormatConfig(op_info);
+  PADDLE_ENFORCE_NE(supports_format_config.empty(),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The %s op has no tensorrt plugin "
+                        "supportsFormatCombination config!"
+                        "Please use SetTrtSupportsFormatConfig to set.",
+                        op_desc_.Type().c_str()));
+  // generate support format combination function by config
+  size_t input_num = OpMetaInfoHelper::GetInputs(op_info).size();
+  size_t output_num = OpMetaInfoHelper::GetOutputs(op_info).size();
+  std::vector<std::vector<std::pair<std::string, std::string>>>
+      format_combinations;
+  for (auto& config : supports_format_config) {
+    auto format_combination = parseConfig(op_desc_.Type(), config);
+    PADDLE_ENFORCE_EQ(input_num + output_num,
+                      format_combination.size(),
+                      common::errors::InvalidArgument(
+                          "Expected %d format_combination, but got %d.",
+                          input_num + output_num,
+                          format_combination.size()));
+    format_combinations.emplace_back(format_combination);
+  }
+
+  bool is_supported = false;
+  for (size_t i = 0; i < input_num + output_num; ++i) {
+    if (i < input_num) {
+      if (pos == i) {
+        for (auto& format_combination : format_combinations) {
+          is_supported |=
+              (in_out[pos].type == getTrtDtype(format_combination[i].first) &&
+               in_out[pos].format ==
+                   getTrtTensorFormat(format_combination[i].second));
+        }
+      }
+    } else {
+      if (pos == i) {
+        for (auto& format_combination : format_combinations) {
+          bool is_supported_tmp = true;
+          for (size_t j = 0; j < input_num; ++j) {
+            is_supported_tmp &=
+                (in_out[j].type == getTrtDtype(format_combination[j].first) &&
+                 in_out[j].format ==
+                     getTrtTensorFormat(format_combination[j].second));
+          }
+          is_supported_tmp &=
+              (in_out[pos].type == getTrtDtype(format_combination[i].first) &&
+               in_out[pos].format ==
+                   getTrtTensorFormat(format_combination[i].second));
+          is_supported |= is_supported_tmp;
+        }
+      }
+    }
+  }
+  return is_supported;
+}
+
+nvinfer1::DataType CustomGenericPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NE(
+      input_types,
+      nullptr,
+      common::errors::Unavailable("Input type should not be nullptr."));
+  return input_types[0];
+}
+
+int CustomGenericPlugin::initialize() TRT_NOEXCEPT {
+  if (!tensor_inputs_)
+    tensor_inputs_ = new std::vector<paddle::Tensor>(getNbInputs());
+  if (!tensor_outputs_)
+    tensor_outputs_ = new std::vector<paddle::Tensor>(getNbOutputs());
+  return 0;
+}
+
+nvinfer1::DimsExprs CustomGenericPlugin::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_LT(
+      output_index,
+      getNbOutputs(),
+      common::errors::InvalidArgument(
+          "Output index (%d) must be less than the number of outputs (%d).",
+          output_index,
+          getNbOutputs()));
+  auto& op_meta_info_map = OpMetaInfoMap::Instance();
+  const auto& meta_info_map = op_meta_info_map.GetMap();
+  auto& op_info = meta_info_map.at(op_desc_.Type()).front();
+  auto& infer_shape_fn = OpMetaInfoHelper::GetTrtInferShapeFn(op_info);
+  PADDLE_ENFORCE_NE(infer_shape_fn,
+                    nullptr,
+                    common::errors::InvalidArgument(
+                        "The %s op has no getOutputDimensions function!"
+                        "Please use SetTrtInferShapeFn to set.",
+                        op_desc_.Type().c_str()));
+  std::vector<paddle::any> custom_attrs;
+  auto& attrs = op_desc_.GetAttrMap();
+  auto& op_attrs_names = OpMetaInfoHelper::GetAttrs(op_info);
+  for (auto& op_attrs_name : op_attrs_names) {
+    auto attr_name_and_type = paddle::ParseAttrStr(op_attrs_name);
+    auto attr_name = attr_name_and_type[0];
+    auto attr_type_str = attr_name_and_type[1];
+    if (attr_type_str == "bool") {
+      custom_attrs.emplace_back(PADDLE_GET_CONST(bool, attrs.at(attr_name)));
+    } else if (attr_type_str == "int") {
+      custom_attrs.emplace_back(PADDLE_GET_CONST(int, attrs.at(attr_name)));
+    } else if (attr_type_str == "float") {
+      custom_attrs.emplace_back(PADDLE_GET_CONST(float, attrs.at(attr_name)));
+    } else if (attr_type_str == "int64_t") {
+      custom_attrs.emplace_back(PADDLE_GET_CONST(int64_t, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::string") {
+      custom_attrs.emplace_back(
+          PADDLE_GET_CONST(std::string, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<int>") {
+      custom_attrs.emplace_back(
+          PADDLE_GET_CONST(std::vector<int>, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<float>") {
+      custom_attrs.emplace_back(
+          PADDLE_GET_CONST(std::vector<float>, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<int64_t>") {
+      custom_attrs.emplace_back(
+          PADDLE_GET_CONST(std::vector<int64_t>, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<std::string>") {
+      custom_attrs.emplace_back(
+          PADDLE_GET_CONST(std::vector<std::string>, attrs.at(attr_name)));
+    } else {
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Unsupported `%s` type value as custom attribute now. "
+          "Supported data types include `bool`, `int`, `float`, "
+          "`int64_t`, `std::string`, `std::vector<int>`, "
+          "`std::vector<float>`, `std::vector<int64_t>`, "
+          "`std::vector<std::string>`, Please check whether "
+          "the attribute data type and data type string are matched.",
+          attr_type_str));
+    }
+  }
+  return infer_shape_fn(
+      {output_index, nb_inputs}, inputs, expr_builder, custom_attrs);
+}
+
+void CustomGenericPlugin::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* in,
+    int nb_inputs,
+    const nvinfer1::DynamicPluginTensorDesc* out,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nb_inputs,
+                    getNbInputs(),
+                    common::errors::InvalidArgument(
+                        "Number of inputs (%d) does not match the "
+                        "expected number of inputs (%d).",
+                        nb_inputs,
+                        getNbInputs()));
+
+  PADDLE_ENFORCE_EQ(nb_outputs,
+                    getNbOutputs(),
+                    common::errors::InvalidArgument(
+                        "Number of outputs (%d) does not match the "
+                        "expected number of outputs (%d).",
+                        nb_outputs,
+                        getNbOutputs()));
+}
+
+// Shutdown the layer. This is called when the engine is destroyed
+void CustomGenericPlugin::terminate() TRT_NOEXCEPT {
+  delete tensor_inputs_;
+  delete tensor_outputs_;
+}
+
+int CustomGenericPlugin::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());
+  // TODO(inference): custom generic plugin do not support INT8 precision now.
+  auto protoType2PhiType =
+      [&](GenerateCustomGenericPluginDataType proto_type,
+          nvinfer1::DataType nv_dtype) -> std::pair<phi::DataType, int> {
+    if (proto_type == GenerateCustomGenericPluginDataType::PLUGIN_FP16) {
+      return {phi::DataType::FLOAT16, sizeof(half)};
+    } else if (proto_type == GenerateCustomGenericPluginDataType::PLUGIN_FP32) {
+      if (isFp16Supported() && nv_dtype == nvinfer1::DataType::kHALF) {
+        return {phi::DataType::FLOAT16, sizeof(half)};
+      } else {
+        return {phi::DataType::FLOAT32, sizeof(float)};
+      }
+    } else if (proto_type ==
+               GenerateCustomGenericPluginDataType::PLUGIN_INT64) {
+      return {phi::DataType::INT64, sizeof(int64_t)};
+    } else if (proto_type ==
+               GenerateCustomGenericPluginDataType::PLUGIN_INT32) {
+      return {phi::DataType::INT32, sizeof(int32_t)};
+    } else if (proto_type == GenerateCustomGenericPluginDataType::PLUGIN_BOOL) {
+      return {phi::DataType::BOOL, sizeof(bool)};
+    } else {
+      PADDLE_ENFORCE_EQ(
+          false,
+          true,
+          common::errors::InvalidArgument("Precision is not supported."));
+    }
+  };
+
+  nvinfer1::DataType 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 must be either kFLOAT or "
+                                      "kHALF, but received data type %d.",
+                                      static_cast<int>(data_type)));
+
+  paddle::CustomOpKernelContext kernel_ctx;
+  // input
+  for (int i = 0; i < getNbInputs(); i++) {
+    if (inputs_data_type_[i] ==
+        GenerateCustomGenericPluginDataType::PLUGIN_OPTIONAL) {
+      (*tensor_inputs_)[i] = paddle::Tensor();
+      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 : input_shape) input_numel *= k;
+
+    auto data_type_and_size =
+        protoType2PhiType(inputs_data_type_[i], data_type);
+
+    phi::DenseTensorMeta input_meta(data_type_and_size.first,
+                                    phi::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));
+    (*tensor_inputs_)[i] = paddle::Tensor(
+        std::make_shared<phi::DenseTensor>(input_alloc, input_meta));
+    kernel_ctx.EmplaceBackInput(std::move((*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 : output_shape) output_numel *= k;
+
+    auto data_type_and_size =
+        protoType2PhiType(outputs_data_type_[i], data_type);
+    phi::DenseTensorMeta output_meta(data_type_and_size.first,
+                                     phi::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));
+    (*tensor_outputs_)[i] = paddle::Tensor(
+        std::make_shared<phi::DenseTensor>(output_alloc, output_meta));
+    kernel_ctx.EmplaceBackOutput(std::move((*tensor_outputs_)[i]));
+  }
+
+  auto& op_meta_info_map = OpMetaInfoMap::Instance();
+  const auto& meta_info_map = op_meta_info_map.GetMap();
+  auto& op_info = meta_info_map.at(op_desc_.Type()).front();
+  auto& op_attrs_names = OpMetaInfoHelper::GetAttrs(op_info);
+  auto& attrs = op_desc_.GetAttrMap();
+  for (auto& op_attrs_name : op_attrs_names) {
+    auto attr_name_and_type = paddle::ParseAttrStr(op_attrs_name);
+    auto attr_name = attr_name_and_type[0];
+    auto attr_type_str = attr_name_and_type[1];
+    if (attr_type_str == "bool") {
+      kernel_ctx.EmplaceBackAttr(PADDLE_GET_CONST(bool, attrs.at(attr_name)));
+    } else if (attr_type_str == "int") {
+      kernel_ctx.EmplaceBackAttr(PADDLE_GET_CONST(int, attrs.at(attr_name)));
+    } else if (attr_type_str == "float") {
+      kernel_ctx.EmplaceBackAttr(PADDLE_GET_CONST(float, attrs.at(attr_name)));
+    } else if (attr_type_str == "int64_t") {
+      kernel_ctx.EmplaceBackAttr(
+          PADDLE_GET_CONST(int64_t, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::string") {
+      kernel_ctx.EmplaceBackAttr(
+          PADDLE_GET_CONST(std::string, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<int>") {
+      kernel_ctx.EmplaceBackAttr(
+          PADDLE_GET_CONST(std::vector<int>, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<float>") {
+      kernel_ctx.EmplaceBackAttr(
+          PADDLE_GET_CONST(std::vector<float>, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<int64_t>") {
+      kernel_ctx.EmplaceBackAttr(
+          PADDLE_GET_CONST(std::vector<int64_t>, attrs.at(attr_name)));
+    } else if (attr_type_str == "std::vector<std::string>") {
+      kernel_ctx.EmplaceBackAttr(
+          PADDLE_GET_CONST(std::vector<std::string>, attrs.at(attr_name)));
+    } else {
+      PADDLE_THROW(common::errors::Unimplemented(
+          "Unsupported `%s` type value as custom attribute now. "
+          "Supported data types include `bool`, `int`, `float`, "
+          "`int64_t`, `std::string`, `std::vector<int>`, "
+          "`std::vector<float>`, `std::vector<int64_t>`, "
+          "`std::vector<std::string>`, Please check whether "
+          "the attribute data type and data type string are matched.",
+          attr_type_str));
+    }
+  }
+  auto kernel_fn = OpMetaInfoHelper::GetKernelFn(op_info);
+  kernel_ctx.UpdatePlainOutputs(OpMetaInfoHelper::GetInputs(op_info),
+                                OpMetaInfoHelper::GetOutputs(op_info),
+                                OpMetaInfoHelper::GetInplaceMap(op_info));
+  kernel_fn(&kernel_ctx);
+  kernel_ctx.AssignInplaceOutputs();
+
+  // sync output tensor data into TensorRT output
+  auto* calc_outs = kernel_ctx.AllMutableOutput();
+  for (int i = 0; i < getNbOutputs(); i++) {
+    auto calc_out =
+        std::dynamic_pointer_cast<phi::DenseTensor>(calc_outs->at(i).impl());
+    if (reinterpret_cast<void*>(calc_out->data()) !=
+        reinterpret_cast<void*>(outputs[i])) {
+      LOG_FIRST_N(WARNING, 1)
+          << "You created new Tensor(s) in custom operator(s) used as "
+             "output(s), "
+             "we will do cudaMemcpy to synchronize the output(s) "
+             "address needed by TensorRT plugin. "
+             "Inplace operation is highly recommended for better performance.";
+      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 : output_shape) output_numel *= k;
+
+      auto data_type_and_size =
+          protoType2PhiType(outputs_data_type_[i], data_type);
+      phi::DenseTensorMeta output_meta(data_type_and_size.first,
+                                       phi::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));
+      phi::DenseTensor dense_output =
+          std::move(phi::DenseTensor(output_alloc, output_meta));
+      cudaMemcpy(static_cast<void*>(dense_output.data()),
+                 static_cast<void*>(calc_out->data()),
+                 output_numel * data_type_and_size.second,
+                 cudaMemcpyDeviceToDevice);
+    }
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/custom_generic_plugin.h

@@ -0,0 +1,186 @@
+// Copyright (c) 2023 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.
+
+#pragma once
+
+#include <NvInfer.h>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/op_desc.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/framework/type_defs.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/helper.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin_utils.h"
+#include "paddle/fluid/inference/tensorrt/plugin_arg_mapping_context.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/phi/api/ext/op_meta_info.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_context.h"
+#include "paddle/phi/core/memory/allocation/cuda_allocator.h"
+#include "paddle/phi/core/platform/device_context.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+enum class GenerateCustomGenericPluginDataType {
+  PLUGIN_BOOL,
+  PLUGIN_UINT8,
+  PLUGIN_INT8,
+  PLUGIN_INT16,
+  PLUGIN_INT32,
+  PLUGIN_INT64,
+  PLUGIN_FP16,
+  PLUGIN_FP32,
+  PLUGIN_FP64,
+  PLUGIN_BF16,
+  PLUGIN_SIZE_T,
+  PLUGIN_COMPLEX64,
+  PLUGIN_COMPLEX128,
+  PLUGIN_OPTIONAL,
+};
+
+GenerateCustomGenericPluginDataType
+ProtoTypeToGenerateCustomGenericPluginDataType(
+    framework::proto::VarType_Type proto_type);
+
+class CustomGenericPlugin : public DynamicPluginTensorRT {
+ public:
+  struct InputOutPutVarInfo {
+    std::vector<GenerateCustomGenericPluginDataType> inputs_data_type;
+    std::vector<GenerateCustomGenericPluginDataType> outputs_data_type;
+  };
+
+ public:
+  CustomGenericPlugin() = default;
+
+  CustomGenericPlugin(const paddle::framework::proto::OpDesc& proto_op_desc,
+                      const InputOutPutVarInfo& in_out_info,
+                      bool with_fp16_ = false);
+
+  CustomGenericPlugin(
+      const paddle::framework::proto::OpDesc& proto_op_desc,
+      const std::vector<GenerateCustomGenericPluginDataType>& inputs_data_type,
+      const std::vector<GenerateCustomGenericPluginDataType>& outputs_data_type,
+      bool with_fp16_ = false);
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  CustomGenericPlugin(void const* serialData, size_t serialLength);
+
+  // IPluginV2 method
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "custom_generic_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override;
+
+  int getNbInputs() const TRT_NOEXCEPT;
+
+  // Initialize the layer for execution.
+  int initialize() TRT_NOEXCEPT override;
+
+  // Shutdown the layer. This is called when the engine is destroyed
+  void terminate() TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override{};
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    size_t sum = 0;
+    sum += SerializedSize(inputs_data_type_);
+    sum += SerializedSize(outputs_data_type_);
+    sum += SerializedSize(with_fp16_);
+    sum += op_meta_data_.size();
+    return sum;
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  // The Func in IPluginV2
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override;
+
+  int 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 override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  bool isFp16Supported() { return with_fp16_; }
+
+ private:
+  bool with_fp16_{false};
+  std::string op_meta_data_;
+  framework::proto::OpDesc proto_op_desc_;
+  framework::OpDesc op_desc_;
+
+ private:
+  std::vector<paddle::Tensor>* tensor_inputs_{nullptr};
+  std::vector<paddle::Tensor>* tensor_outputs_{nullptr};
+
+ private:
+  std::vector<GenerateCustomGenericPluginDataType> inputs_data_type_;
+  std::vector<GenerateCustomGenericPluginDataType> outputs_data_type_;
+};
+
+class CustomGenericPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "custom_generic_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2DynamicExt* deserializePlugin(const char* name,
+                                                   const void* serial_data,
+                                                   size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new CustomGenericPlugin(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(CustomGenericPluginCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/deformable_conv_op_plugin.h

@@ -0,0 +1,418 @@
+/* Copyright (c) 2021 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. */
+
+#pragma once
+
+#include <cstdio>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/phi/backends/dynload/cublas.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class DeformableConvPlugin : public nvinfer1::IPluginV2Ext {
+ public:
+  explicit DeformableConvPlugin(const nvinfer1::DataType data_type,
+                                const nvinfer1::Weights& weights,
+                                const std::vector<int>& kernel_dims,
+                                const std::vector<int>& strides,
+                                const std::vector<int>& paddings,
+                                const std::vector<int>& dilations,
+                                const int groups,
+                                const int deformable_groups,
+                                const int im2col_step,
+                                const bool with_fp16);
+  explicit DeformableConvPlugin(const nvinfer1::DataType data_type,
+                                const nvinfer1::Weights& weights,
+                                const std::vector<int>& kernel_dims,
+                                const std::vector<int>& strides,
+                                const std::vector<int>& paddings,
+                                const std::vector<int>& dilations,
+                                const int groups,
+                                const int deformable_groups,
+                                const int im2col_step,
+                                const std::vector<int>& input_dim,
+                                const std::vector<int>& offset_dim,
+                                const std::vector<int>& mask_dim,
+                                const std::vector<int>& output_dim,
+                                const bool with_fp16);
+  DeformableConvPlugin(const void* data, size_t length);
+  ~DeformableConvPlugin() override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nb_input_dims) TRT_NOEXCEPT override;
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::TensorFormat format)
+      const TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(int max_batch_size) const TRT_NOEXCEPT override;
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_type,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+  bool isOutputBroadcastAcrossBatch(int output_index,
+                                    const bool* input_is_broadcast,
+                                    int nb_inputs) const TRT_NOEXCEPT override;
+  bool canBroadcastInputAcrossBatch(int input_index) const
+      TRT_NOEXCEPT override;
+
+  void attachToContext(cudnnContext* cudnnContext,
+                       cublasContext* cublasContext,
+                       nvinfer1::IGpuAllocator* gpuAllocator)
+      TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::Dims* input_dims,
+                       int nb_inputs,
+                       const nvinfer1::Dims* output_dims,
+                       int nb_outputs,
+                       const nvinfer1::DataType* input_types,
+                       const nvinfer1::DataType* output_types,
+                       const bool* input_is_broadcast,
+                       const bool* output_is_broadcast,
+                       nvinfer1::PluginFormat float_format,
+                       int max_batch_size) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* clone() const TRT_NOEXCEPT override;
+
+ private:
+  template <typename T>
+  int enqueue_impl(int batch_size,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  nvinfer1::Weights copyToDevice(const void* hostData, size_t count);
+  void serializeFromDevice(void** hostBuffer,
+                           const nvinfer1::Weights& deviceWeights) const;
+  nvinfer1::Weights deserializeToDevice(const void** hostBuffer, size_t count);
+
+  bool with_fp16_;
+  nvinfer1::DataType data_type_;
+  nvinfer1::Weights weights_;
+  std::vector<int> kernel_dims_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  std::vector<int> dilations_;
+  int groups_;
+  int deformable_groups_;
+  int im2col_step_;
+  std::string namespace_;
+
+  std::vector<int> input_dim_;
+  std::vector<int> offset_dim_;
+  std::vector<int> mask_dim_;
+  std::vector<int> output_dim_;
+
+  cublasHandle_t cublasHandle_;
+};
+
+class DeformableConvPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  DeformableConvPluginCreator() = default;
+  ~DeformableConvPluginCreator() override = default;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+REGISTER_TRT_PLUGIN_V2(DeformableConvPluginCreator);
+
+class DeformableConvPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit DeformableConvPluginDynamic(const nvinfer1::DataType data_type,
+                                       const nvinfer1::Weights& weights,
+                                       const std::vector<int>& kernel_dims,
+                                       const std::vector<int>& strides,
+                                       const std::vector<int>& paddings,
+                                       const std::vector<int>& dilations,
+                                       const int groups,
+                                       const int deformable_groups,
+                                       const int im2col_step,
+                                       const bool with_fp16);
+  DeformableConvPluginDynamic(const void* data, size_t length);
+  ~DeformableConvPluginDynamic() override;
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new DeformableConvPluginDynamic(data_type_,
+                                           weights_,
+                                           kernel_dims_,
+                                           strides_,
+                                           paddings_,
+                                           dilations_,
+                                           groups_,
+                                           deformable_groups_,
+                                           im2col_step_,
+                                           with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "deformable_conv_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+ private:
+  nvinfer1::Weights copyToDevice(const void* hostData, size_t count);
+  void serializeFromDevice(void** hostBuffer,
+                           const nvinfer1::Weights& deviceWeights) const;
+  nvinfer1::Weights deserializeToDevice(const void** hostBuffer, size_t count);
+  template <typename T>
+  int enqueue_impl(const nvinfer1::PluginTensorDesc* inputDesc,
+                   const nvinfer1::PluginTensorDesc* outputDesc,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  bool with_fp16_;
+  nvinfer1::DataType data_type_;
+  nvinfer1::Weights weights_;
+  std::vector<int> kernel_dims_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  std::vector<int> dilations_;
+  int groups_;
+  int deformable_groups_;
+  int im2col_step_;
+  std::string namespace_;
+
+  cublasHandle_t cublasHandle_;
+};
+
+class DeformableConvPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    namespace_ = lib_namespace;
+  }
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return namespace_.c_str();
+  }
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "deformable_conv_plugin_dynamic";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+class PIRDeformableConvPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit PIRDeformableConvPluginDynamic(const nvinfer1::DataType data_type,
+                                          const nvinfer1::Weights& weights,
+                                          const std::vector<int>& kernel_dims,
+                                          const std::vector<int>& strides,
+                                          const std::vector<int>& paddings,
+                                          const std::vector<int>& dilations,
+                                          const int groups,
+                                          const int deformable_groups,
+                                          const int im2col_step,
+                                          const bool with_fp16);
+  PIRDeformableConvPluginDynamic(const void* data, size_t length);
+  ~PIRDeformableConvPluginDynamic() override;
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new PIRDeformableConvPluginDynamic(data_type_,
+                                              weights_,
+                                              kernel_dims_,
+                                              strides_,
+                                              paddings_,
+                                              dilations_,
+                                              groups_,
+                                              deformable_groups_,
+                                              im2col_step_,
+                                              with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "deformable_conv_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+ private:
+  nvinfer1::Weights copyToDevice(const void* hostData, size_t count);
+  void serializeFromDevice(void** hostBuffer,
+                           const nvinfer1::Weights& deviceWeights) const;
+  nvinfer1::Weights deserializeToDevice(const void** hostBuffer, size_t count);
+  template <typename T>
+  int enqueue_impl(const nvinfer1::PluginTensorDesc* inputDesc,
+                   const nvinfer1::PluginTensorDesc* outputDesc,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  bool with_fp16_;
+  nvinfer1::DataType data_type_;
+  nvinfer1::Weights weights_;
+  std::vector<int> kernel_dims_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  std::vector<int> dilations_;
+  int groups_;
+  int deformable_groups_;
+  int im2col_step_;
+  std::string namespace_;
+
+  cublasHandle_t cublasHandle_;
+};
+
+class PIRDeformableConvPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  PIRDeformableConvPluginDynamicCreator() = default;
+  ~PIRDeformableConvPluginDynamicCreator() override = default;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+REGISTER_TRT_PLUGIN_V2(PIRDeformableConvPluginDynamicCreator);
+
+REGISTER_TRT_PLUGIN_V2(DeformableConvPluginDynamicCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.cu

@@ -0,0 +1,346 @@
+/* Copyright (c) 2018 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 <glog/logging.h>
+
+#include "paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+namespace details {
+template <typename T>
+struct Add {
+  __device__ T operator()(const T &a, const T &b) const { return a + b; }
+};
+
+template <typename T>
+struct Mul {
+  __device__ T operator()(const T &a, const T &b) const { return a * b; }
+};
+
+template <typename T>
+struct Div {
+  __device__ T operator()(const T &a, const T &b) const { return a / b; }
+};
+
+template <typename T>
+struct Sub {
+  __device__ T operator()(const T &a, const T &b) const { return a - b; }
+};
+
+template <typename T>
+struct Pow {
+  __device__ T operator()(const T &a, const T &b) const {
+    return static_cast<T>(::powf(static_cast<float>(a), static_cast<float>(b)));
+  }
+};
+
+}  // namespace details
+
+template <typename T, typename Operator>
+__global__ void elementwise_kernel(const size_t total,
+                                   const T *x_data,
+                                   const T *y_data,
+                                   T *out_data,
+                                   int pre,
+                                   int n,
+                                   int post,
+                                   Operator op) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (tid < total) {
+    int idx = tid / post % n;
+#if __CUDA_ARCH__ >= 350
+    out_data[tid] = op(__ldg(x_data + tid), __ldg(y_data + idx));
+#else
+    out_data[tid] = op(x_data[tid], y_data[idx]);
+#endif
+  }
+}
+
+nvinfer1::Dims ElementWisePlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims *input_dims, int num_inputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "There is only one output in TRT elementwise "
+                        "op plugin, but got output index: %d.",
+                        index));
+  PADDLE_ENFORCE_EQ(
+      num_inputs,
+      2,
+      common::errors::InvalidArgument("There are 2 inputs in TRT elementwise "
+                                      "op plugin, but got input number: %d.",
+                                      num_inputs));
+  PADDLE_ENFORCE_NOT_NULL(
+      input_dims,
+      common::errors::InvalidArgument(
+          "The input dims of TRT elementwise op plugin should not be null."));
+  return input_dims[0];
+}
+
+int ElementWisePlugin::initialize() TRT_NOEXCEPT {
+  axis_ = (axis_ == -1) ? dims_x_.nbDims - dims_y_.nbDims : axis_;
+  int trimmed_nb_dims = dims_y_.nbDims;
+  for (; trimmed_nb_dims > 0; --trimmed_nb_dims) {
+    if (dims_y_.d[trimmed_nb_dims - 1] != 1) {
+      break;
+    }
+  }
+  dims_y_.nbDims = trimmed_nb_dims;
+
+  PADDLE_ENFORCE_GE(dims_x_.nbDims,
+                    dims_y_.nbDims + axis_,
+                    common::errors::InvalidArgument(
+                        "We expect [number of x dims] >= [number of y dims + "
+                        "axis] in TRT elementwise op plugin, but got [number "
+                        "of x dims] = %d, [number of y dims + axis] = %d.",
+                        dims_x_.nbDims,
+                        dims_y_.nbDims + axis_));
+  PADDLE_ENFORCE_LT(
+      axis_,
+      dims_x_.nbDims,
+      common::errors::InvalidArgument("We expect [axis] < [number of x dims] "
+                                      "in TRT elementwise op plugin, but got "
+                                      "[axis] = %d, [number of x dims] = %d.",
+                                      axis_,
+                                      dims_x_.nbDims));
+
+  prev_size_ = 1;
+  midd_size_ = 1;
+  post_size_ = 1;
+  for (int i = 0; i < axis_; ++i) {
+    prev_size_ *= dims_x_.d[i];
+  }
+
+  for (int i = 0; i < dims_y_.nbDims; ++i) {
+    PADDLE_ENFORCE_EQ(dims_x_.d[i + axis_],
+                      dims_y_.d[i],
+                      common::errors::InvalidArgument(
+                          "Broadcast dimension mismatch. The dims of input Y "
+                          "should be a subsequence of X."));
+    midd_size_ *= dims_y_.d[i];
+  }
+
+  for (int i = axis_ + dims_y_.nbDims; i < dims_x_.nbDims; ++i) {
+    post_size_ *= dims_x_.d[i];
+  }
+  return 0;
+}
+
+int ElementWisePlugin::enqueue(int batch_size,
+                               const void *const *inputs,
+                               void *const *outputs,
+                               void *workspace,
+                               cudaStream_t stream) TRT_NOEXCEPT {
+  const float *x = reinterpret_cast<const float *>(inputs[0]);
+  const float *y = reinterpret_cast<const float *>(inputs[1]);
+  float *out = reinterpret_cast<float *>(outputs[0]);
+
+  int num = batch_size * prev_size_ * midd_size_ * post_size_;
+  int thread = 256;
+  int block = (num + thread - 1) / thread;
+  if (type_ == "add") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(num,
+                                                     x,
+                                                     y,
+                                                     out,
+                                                     prev_size_,
+                                                     batch_size * midd_size_,
+                                                     post_size_,
+                                                     details::Add<float>());
+  } else if (type_ == "mul") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(num,
+                                                     x,
+                                                     y,
+                                                     out,
+                                                     prev_size_,
+                                                     batch_size * midd_size_,
+                                                     post_size_,
+                                                     details::Mul<float>());
+  } else if (type_ == "div") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(num,
+                                                     x,
+                                                     y,
+                                                     out,
+                                                     prev_size_,
+                                                     batch_size * midd_size_,
+                                                     post_size_,
+                                                     details::Div<float>());
+  } else if (type_ == "sub") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(num,
+                                                     x,
+                                                     y,
+                                                     out,
+                                                     prev_size_,
+                                                     batch_size * midd_size_,
+                                                     post_size_,
+                                                     details::Sub<float>());
+  } else if (type_ == "pow") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(num,
+                                                     x,
+                                                     y,
+                                                     out,
+                                                     prev_size_,
+                                                     batch_size * midd_size_,
+                                                     post_size_,
+                                                     details::Pow<float>());
+  } else {
+    PADDLE_THROW(common::errors::Fatal(
+        "The %s type elementwise is not implemented in trt plugin.", type_));
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+int ElementwisePluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+size_t ElementwisePluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(type_.c_str()) + SerializedSize(axis_);
+}
+
+void ElementwisePluginDynamic::serialize(void *buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, type_.c_str());
+  SerializeValue(&buffer, axis_);
+}
+
+nvinfer1::DimsExprs ElementwisePluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  return inputs[0];
+}
+
+bool ElementwisePluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    return (in.type == nvinfer1::DataType::kFLOAT) &&
+           (in.format == nvinfer1::TensorFormat::kLINEAR);
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType ElementwisePluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Elementwise Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+int ElementwisePluginDynamic::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 {
+  auto x_dims = input_desc[0].dims;
+  auto y_dims = input_desc[1].dims;
+  int axis = (axis_ == -1) ? x_dims.nbDims - y_dims.nbDims : axis_;
+  int batch_size = x_dims.d[0];
+
+  int prev_size = 1;
+  int midd_size = 1;
+  int post_size = 1;
+  for (int i = 0; i < axis; ++i) {
+    prev_size *= x_dims.d[i];
+  }
+
+  int trimmed_nb_dims = y_dims.nbDims;
+  for (; trimmed_nb_dims > 0; --trimmed_nb_dims) {
+    if (y_dims.d[trimmed_nb_dims - 1] != 1) {
+      break;
+    }
+  }
+
+  for (int i = 0; i < trimmed_nb_dims; ++i) {
+    PADDLE_ENFORCE_EQ(x_dims.d[i + axis],
+                      y_dims.d[i],
+                      common::errors::InvalidArgument(
+                          "Broadcast dimension mismatch found in trt "
+                          "elementwise plugin's x and y input."));
+    midd_size *= y_dims.d[i];
+  }
+
+  for (int i = axis + trimmed_nb_dims; i < x_dims.nbDims; ++i) {
+    post_size *= x_dims.d[i];
+  }
+
+  const float *x = static_cast<const float *>(inputs[0]);
+  const float *y = static_cast<const float *>(inputs[1]);
+
+  float *out = static_cast<float *>(outputs[0]);
+
+  int num = prev_size * midd_size * post_size;
+  int thread = 256;
+  int block = (num + thread - 1) / thread;
+  if (type_ == "add") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(
+        num, x, y, out, prev_size, midd_size, post_size, details::Add<float>());
+  } else if (type_ == "mul") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(
+        num, x, y, out, prev_size, midd_size, post_size, details::Mul<float>());
+  } else if (type_ == "div") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(
+        num, x, y, out, prev_size, midd_size, post_size, details::Div<float>());
+  } else if (type_ == "sub") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(
+        num, x, y, out, prev_size, midd_size, post_size, details::Sub<float>());
+  } else if (type_ == "pow") {
+    elementwise_kernel<<<block, thread, 0, stream>>>(
+        num, x, y, out, prev_size, midd_size, post_size, details::Pow<float>());
+  } else {
+    PADDLE_THROW(common::errors::Unimplemented(
+        "Paddle-TRT only support elementwise "
+        "operation: {add, mul, div, sub, pow} currently, "
+        "but got %s.",
+        type_));
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/elementwise_op_plugin.h

@@ -0,0 +1,231 @@
+/* Copyright (c) 2018 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. */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class ElementWisePlugin : public PluginTensorRT {
+ public:
+  ElementWisePlugin(std::string type,
+                    nvinfer1::Dims const& dims_x,
+                    nvinfer1::Dims const& dims_y,
+                    int axis)
+      : type_(type),
+        dims_x_(dims_x),
+        dims_y_(dims_y),
+        axis_(axis),
+        prev_size_(1),
+        midd_size_(1),
+        post_size_(1) {}
+
+  ElementWisePlugin(void const* serial_data, size_t serial_length) {
+    deserializeBase(serial_data, serial_length);
+    const char* elementwise_type;
+    DeserializeValue(&serial_data, &serial_length, &elementwise_type);
+    type_ = std::string(elementwise_type);
+    DeserializeValue(&serial_data, &serial_length, &dims_x_);
+    DeserializeValue(&serial_data, &serial_length, &dims_y_);
+    DeserializeValue(&serial_data, &serial_length, &axis_);
+    DeserializeValue(&serial_data, &serial_length, &prev_size_);
+    DeserializeValue(&serial_data, &serial_length, &midd_size_);
+    DeserializeValue(&serial_data, &serial_length, &post_size_);
+  }
+
+  ElementWisePlugin* clone() const TRT_NOEXCEPT override {
+    return new ElementWisePlugin(type_, dims_x_, dims_y_, axis_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "elementwise_plugin";
+  }
+
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* input_dims,
+                                     int num_inputs) TRT_NOEXCEPT override;
+
+  int initialize() TRT_NOEXCEPT override;
+
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(type_.c_str()) +
+           SerializedSize(dims_x_) + SerializedSize(dims_y_) +
+           SerializedSize(axis_) + SerializedSize(prev_size_) +
+           SerializedSize(midd_size_) + SerializedSize(post_size_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, type_.c_str());
+    SerializeValue(&buffer, dims_x_);
+    SerializeValue(&buffer, dims_y_);
+    SerializeValue(&buffer, axis_);
+    SerializeValue(&buffer, prev_size_);
+    SerializeValue(&buffer, midd_size_);
+    SerializeValue(&buffer, post_size_);
+  }
+
+ protected:
+  std::string type_;
+  nvinfer1::Dims dims_x_;
+  nvinfer1::Dims dims_y_;
+  int axis_;
+  int prev_size_;
+  int midd_size_;
+  int post_size_;
+};
+
+class ElementWisePluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "elementwise_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new ElementWisePlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(ElementWisePluginCreator);
+
+class ElementwisePluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit ElementwisePluginDynamic(const std::string& type, int axis)
+      : type_(type), axis_(axis) {}
+  ElementwisePluginDynamic(void const* serialData, size_t serialLength) {
+    const char* elementwise_type;
+    DeserializeValue(&serialData, &serialLength, &elementwise_type);
+    type_ = std::string(elementwise_type);
+    DeserializeValue(&serialData, &serialLength, &axis_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new ElementwisePluginDynamic(type_, axis_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "elementwise_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  std::string type_;
+  int axis_;
+};
+
+class ElementwisePluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  ElementwisePluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "elementwise_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new ElementwisePluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+
+REGISTER_TRT_PLUGIN_V2(ElementwisePluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/elementwiseadd_transpose_op_plugin.cu

@@ -0,0 +1,260 @@
+/* Copyright (c) 2023 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 "paddle/fluid/inference/tensorrt/plugin/elementwiseadd_transpose_op_plugin.h"
+#include <glog/logging.h>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/kernels/elementwise_add_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+int ElementwiseAddTransposePluginDynamic::initialize() TRT_NOEXCEPT {
+  return 0;
+}
+
+size_t ElementwiseAddTransposePluginDynamic::getSerializationSize() const
+    TRT_NOEXCEPT {
+  return SerializedSize(axis_) + SerializedSize(output_shape_);
+}
+
+void ElementwiseAddTransposePluginDynamic::serialize(void *buffer) const
+    TRT_NOEXCEPT {
+  SerializeValue(&buffer, axis_);
+  SerializeValue(&buffer, output_shape_);
+}
+
+nvinfer1::DimsExprs ElementwiseAddTransposePluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 4;
+  ret.d[0] = inputs[0].d[0];
+  ret.d[1] = inputs[0].d[2];
+  ret.d[2] = expr_builder.constant(output_shape_[1]);
+  ret.d[3] = expr_builder.constant(output_shape_[2]);
+  return ret;
+}
+
+bool ElementwiseAddTransposePluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument("The input of elementwiseadd_transpose "
+                                      "plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  // (in_out && pos < (nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  // input 0
+  if (pos == 0) {
+    return (in.type == nvinfer1::DataType::kHALF ||
+            in.type == nvinfer1::DataType::kFLOAT) &&
+           (in.format == nvinfer1::TensorFormat::kLINEAR);
+  }
+  // input 1
+  if (pos == 1) {
+    return (in.type == in_out[0].type) &&
+           (in.format == nvinfer1::TensorFormat::kLINEAR);
+  }
+  // output 0
+  if (pos == 2) {
+    bool support_format = in.format == nvinfer1::TensorFormat::kLINEAR ||
+                          in.format == nvinfer1::TensorFormat::kHWC8;
+
+    return (in.type == in_out[0].type) && (support_format);
+  }
+}
+void ElementwiseAddTransposePluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *input_desc,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *output_desc,
+    int nbOutputs) TRT_NOEXCEPT {
+  const auto &x_dims = input_desc[0].desc.dims;
+  const auto &y_dims = input_desc[1].desc.dims;
+  const auto &out_dims = output_desc[0].desc.dims;
+  const auto &x_type = input_desc[0].desc.type;
+  std::vector<int> x_shape;
+  int x_numel = 1;
+  for (int i = 0; i < x_dims.nbDims; i++) {
+    x_shape.push_back(x_dims.d[i]);
+    x_numel *= x_dims.d[i];
+  }
+  std::vector<int> y_shape;
+  int y_numel = 1;
+  for (int i = 0; i < y_dims.nbDims; i++) {
+    y_shape.push_back(y_dims.d[i]);
+    y_numel *= y_dims.d[i];
+  }
+  std::vector<int> out_shape;
+  int out_numel = 1;
+  for (int i = 0; i < out_dims.nbDims; i++) {
+    out_shape.push_back(out_dims.d[i]);
+    out_numel *= out_dims.d[i];
+  }
+  x_numel_ = x_numel;
+  y_numel_ = y_numel;
+  out_numel_ = out_numel;
+  if (x_numel <= 0) {
+    return;
+  }
+  ele_out_tensor_.Resize(common::make_ddim(x_shape));
+  phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+  GPUPlace place(platform::GetCurrentDeviceId());
+  auto *device_context = static_cast<phi::GPUContext *>(pool.Get(place));
+  const phi::GPUContext &dev_ctx = *device_context;
+
+  if (x_type == nvinfer1::DataType::kFLOAT) {
+    x_meta_ = phi::DenseTensorMeta(phi::DataType::FLOAT32,
+                                   common::make_ddim(x_shape));
+    y_meta_ = phi::DenseTensorMeta(phi::DataType::FLOAT32,
+                                   common::make_ddim(y_shape));
+    out_meta_ = phi::DenseTensorMeta(phi::DataType::FLOAT32,
+                                     common::make_ddim(out_shape));
+    dev_ctx.template Alloc<float>(&ele_out_tensor_, x_numel * sizeof(float));
+  } else if (x_type == nvinfer1::DataType::kHALF) {
+    x_meta_ = phi::DenseTensorMeta(phi::DataType::FLOAT16,
+                                   common::make_ddim(x_shape));
+    y_meta_ = phi::DenseTensorMeta(phi::DataType::FLOAT16,
+                                   common::make_ddim(y_shape));
+    out_meta_ = phi::DenseTensorMeta(phi::DataType::FLOAT16,
+                                     common::make_ddim(out_shape));
+    dev_ctx.template Alloc<phi::dtype::float16>(
+        &ele_out_tensor_, x_numel * sizeof(phi::dtype::float16));
+  }
+}
+nvinfer1::DataType ElementwiseAddTransposePluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Elementwise Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+int ElementwiseAddTransposePluginDynamic::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 {
+  phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+  GPUPlace place(platform::GetCurrentDeviceId());
+  auto *device_context = static_cast<phi::GPUContext *>(pool.Get(place));
+  const phi::GPUContext &dev_ctx = *device_context;
+
+  auto input_type = input_desc[0].type;
+  auto output_format = output_desc[0].format;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. elementwiseadd_transpose-->fp32";
+    const float *x = static_cast<const float *>(inputs[0]);
+    const float *y = static_cast<const float *>(inputs[1]);
+    float *out = static_cast<float *>(outputs[0]);
+    VLOG(1) << "TRT Plugin format selected. elementwiseadd_transpose-->kLINEAR";
+    std::shared_ptr<phi::Allocation> x_alloc(new phi::Allocation(
+        static_cast<void *>(const_cast<float *>(x)),  // NOLINT
+        x_numel_ * sizeof(float),
+        place));
+    std::shared_ptr<phi::Allocation> y_alloc(new phi::Allocation(
+        static_cast<void *>(const_cast<float *>(y)),  // NOLINT
+        y_numel_ * sizeof(float),
+        place));
+    std::shared_ptr<phi::Allocation> out_alloc(
+        new phi::Allocation(static_cast<void *>(out),  // NOLINT
+                            out_numel_ * sizeof(float),
+                            place));
+    const phi::DenseTensor x_tensor = phi::DenseTensor(x_alloc, x_meta_);
+    const phi::DenseTensor y_tensor = phi::DenseTensor(y_alloc, y_meta_);
+    phi::DenseTensor out_tensor = phi::DenseTensor(out_alloc, out_meta_);
+    phi::AddKernel<float, phi::GPUContext>(
+        dev_ctx, x_tensor, y_tensor, &ele_out_tensor_);
+    phi::TransposeKernel<float, phi::GPUContext>(
+        dev_ctx, ele_out_tensor_, std::vector<int>{0, 2, 1}, &out_tensor);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. elementwiseadd_transpose-->fp16";
+    const half *x = static_cast<const half *>(inputs[0]);
+    const half *y = static_cast<const half *>(inputs[1]);
+    half *out = static_cast<half *>(outputs[0]);
+    if (output_format == nvinfer1::PluginFormat::kLINEAR) {
+      VLOG(1)
+          << "TRT Plugin format selected. elementwiseadd_transpose-->kLINEAR";
+      std::shared_ptr<phi::Allocation> x_alloc(new phi::Allocation(
+          static_cast<void *>(const_cast<half *>(x)),  // NOLINT
+          x_numel_ * sizeof(half),
+          place));
+      std::shared_ptr<phi::Allocation> y_alloc(new phi::Allocation(
+          static_cast<void *>(const_cast<half *>(y)),  // NOLINT
+          y_numel_ * sizeof(half),
+          place));
+
+      std::shared_ptr<phi::Allocation> out_alloc(
+          new phi::Allocation(static_cast<void *>(out),  // NOLINT
+                              out_numel_ * sizeof(half),
+                              place));
+      const phi::DenseTensor x_tensor = phi::DenseTensor(x_alloc, x_meta_);
+      const phi::DenseTensor y_tensor = phi::DenseTensor(y_alloc, y_meta_);
+      phi::DenseTensor out_tensor = phi::DenseTensor(out_alloc, out_meta_);
+      phi::AddKernel<phi::dtype::float16, phi::GPUContext>(
+          dev_ctx, x_tensor, y_tensor, &ele_out_tensor_);
+      phi::TransposeKernel<phi::dtype::float16, phi::GPUContext>(
+          dev_ctx, ele_out_tensor_, std::vector<int>{0, 2, 1}, &out_tensor);
+    } else if (output_format == nvinfer1::PluginFormat::kHWC8) {
+      VLOG(1) << "TRT Plugin format selected. elementwiseadd_transpose-->kHWC8";
+      std::shared_ptr<phi::Allocation> x_alloc(new phi::Allocation(
+          static_cast<void *>(const_cast<half *>(x)),  // NOLINT
+          x_numel_ * sizeof(half),
+          place));
+      std::shared_ptr<phi::Allocation> y_alloc(new phi::Allocation(
+          static_cast<void *>(const_cast<half *>(y)),  // NOLINT
+          y_numel_ * sizeof(half),
+          place));
+      std::shared_ptr<phi::Allocation> out_alloc(
+          new phi::Allocation(static_cast<void *>(out),  // NOLINT
+                              out_numel_ * sizeof(half),
+                              place));
+      const phi::DenseTensor x_tensor = phi::DenseTensor(x_alloc, x_meta_);
+      const phi::DenseTensor y_tensor = phi::DenseTensor(y_alloc, y_meta_);
+      phi::DenseTensor out_tensor = phi::DenseTensor(out_alloc, out_meta_);
+      phi::AddKernel<phi::dtype::float16, phi::GPUContext>(
+          dev_ctx, x_tensor, y_tensor, &out_tensor);
+    }
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/elementwiseadd_transpose_op_plugin.h

@@ -0,0 +1,150 @@
+/* Copyright (c) 2023 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. */
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class ElementwiseAddTransposePluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit ElementwiseAddTransposePluginDynamic(int axis,
+                                                std::vector<int> output_shape)
+      : axis_(axis), output_shape_(output_shape) {}
+  ElementwiseAddTransposePluginDynamic(void const* serialData,
+                                       size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &axis_);
+    DeserializeValue(&serialData, &serialLength, &output_shape_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new ElementwiseAddTransposePluginDynamic(axis_, output_shape_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "elementwise_add_transpose_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* input_desc,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* output_desc,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  int axis_;
+  std::vector<int> output_shape_;
+  phi::DenseTensorMeta x_meta_;
+  phi::DenseTensorMeta y_meta_;
+  phi::DenseTensorMeta out_meta_;
+  phi::DenseTensor ele_out_tensor_;
+  int x_numel_ = -1;
+  int y_numel_ = -1;
+  int out_numel_ = -1;
+};
+
+class ElementwiseAddTransposePluginDynamicCreator
+    : public nvinfer1::IPluginCreator {
+ public:
+  ElementwiseAddTransposePluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "elementwise_add_transpose_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin =
+        new ElementwiseAddTransposePluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+
+REGISTER_TRT_PLUGIN_V2(ElementwiseAddTransposePluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/fused_token_prune_op_plugin.cu

@@ -0,0 +1,657 @@
+// 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 <vector>
+
+#include "cub/cub.cuh"
+
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+#include "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/phi/backends/gpu/gpu_primitives.h"
+#include "paddle/phi/core/platform/device_context.h"
+
+#include "paddle/fluid/inference/tensorrt/plugin/fused_token_prune_op_plugin.h"
+#include "paddle/phi/kernels/funcs/fused_token_prune_utils.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+__global__ void compute_token_length(const int32_t* src,
+                                     int32_t* dst,
+                                     float scale) {
+  int32_t it = threadIdx.x;
+  dst[it] = max(static_cast<int>((src[it + 1] - src[it]) * scale), 1);
+}
+
+template <typename T>
+__global__ void fill_index_padding_score(int32_t* token_index,
+                                         const T* scores,
+                                         int32_t sequence_length,
+                                         T* padding_scores) {
+  int padding_scores_it = threadIdx.x + blockIdx.x * blockDim.x;
+  int scores_it = threadIdx.x + blockIdx.x * sequence_length;
+  token_index[padding_scores_it] = threadIdx.x;
+  if (threadIdx.x < sequence_length) {
+    padding_scores[padding_scores_it] = scores[scores_it];
+  } else {
+    padding_scores[padding_scores_it] = 0;
+  }
+}
+
+template <typename T, int BLOCK_THREADS, int ITEMS_PER_THREAD>
+__global__ void general_topk_pair_sort(T* in_keys, int32_t* in_out_values) {
+  typedef cub::BlockRadixSort<T, BLOCK_THREADS, ITEMS_PER_THREAD, int>
+      BlockRadixSort;
+  typedef cub::
+      BlockLoad<T, BLOCK_THREADS, ITEMS_PER_THREAD, cub::BLOCK_LOAD_TRANSPOSE>
+          BlockLoadKey;
+  typedef cub::
+      BlockLoad<int, BLOCK_THREADS, ITEMS_PER_THREAD, cub::BLOCK_LOAD_TRANSPOSE>
+          BlockLoadValue;
+  typedef cub::
+      BlockStore<T, BLOCK_THREADS, ITEMS_PER_THREAD, cub::BLOCK_STORE_TRANSPOSE>
+          BlockStoreKey;
+  typedef cub::BlockStore<int,
+                          BLOCK_THREADS,
+                          ITEMS_PER_THREAD,
+                          cub::BLOCK_STORE_TRANSPOSE>
+      BlockStoreValue;
+
+  __shared__ union {
+    typename BlockRadixSort::TempStorage sort;
+    typename BlockLoadKey::TempStorage loadkey;
+    typename BlockLoadValue::TempStorage loadvalue;
+    typename BlockStoreKey::TempStorage storekey;
+    typename BlockStoreValue::TempStorage storevalue;
+  } temp_storage;
+
+  int block_offset = blockIdx.x * BLOCK_THREADS * ITEMS_PER_THREAD;
+
+  T thread_keys[ITEMS_PER_THREAD];
+  int thread_values[ITEMS_PER_THREAD];
+  BlockLoadKey(temp_storage.loadkey).Load(in_keys + block_offset, thread_keys);
+  BlockLoadValue(temp_storage.loadvalue)
+      .Load(in_out_values + block_offset, thread_values);
+  __syncthreads();
+
+  BlockRadixSort(temp_storage.sort).SortDescending(thread_keys, thread_values);
+  __syncthreads();
+
+  BlockStoreValue(temp_storage.storevalue)
+      .Store(in_out_values + block_offset, thread_values);
+}
+
+__global__ void varlen_prune_token_change_order(
+    const half* tokens,
+    const int32_t* token_pos,
+    const int32_t padding_token_length,
+    const int32_t* token_index,
+    half* output) {
+  int batch = blockIdx.x;
+  int token_it = batch * gridDim.y + blockIdx.y;
+  int pre_value_it =
+      token_it * gridDim.z * blockDim.x + blockIdx.z * blockDim.x + threadIdx.x;
+  int token_index_it = batch * padding_token_length + blockIdx.y;
+
+  if (token_index[token_index_it] < token_pos[batch + 1] - token_pos[batch]) {
+    output[(token_index[token_index_it] + token_pos[batch]) * gridDim.z *
+               blockDim.x +
+           blockIdx.z * blockDim.x + threadIdx.x] = tokens[pre_value_it];
+  }
+}
+
+template <typename T>
+__global__ void prune_token_change_order(const T* tokens,
+                                         int32_t new_sequence_length,
+                                         const int32_t padding_token_length,
+                                         const int32_t* token_index,
+                                         T* output) {
+  int batch = blockIdx.x;
+  int token_it = batch * gridDim.y + blockIdx.y;
+  int pre_value_it =
+      token_it * gridDim.z * blockDim.x + blockIdx.z * blockDim.x + threadIdx.x;
+  int token_index_it = batch * padding_token_length + blockIdx.y;
+
+  if (token_index[token_index_it] < new_sequence_length) {
+    output[(batch * new_sequence_length + token_index[token_index_it]) *
+               gridDim.z * blockDim.x +
+           blockIdx.z * blockDim.x + threadIdx.x] = tokens[pre_value_it];
+  }
+}
+
+template <typename T>
+__global__ void prune_token_keep_order(const T* tokens,
+                                       int32_t pre_sequence_length,
+                                       int32_t new_sequence_length,
+                                       const int32_t padding_token_length,
+                                       const int32_t* token_index,
+                                       T* output0,
+                                       int32_t* output1) {
+  int batch = blockIdx.x;
+  int index = 0;
+  for (int i = 0; i < pre_sequence_length; ++i) {
+    if (token_index[batch * padding_token_length + i] < new_sequence_length) {
+      output0[(batch * new_sequence_length + index) * gridDim.y * blockDim.x +
+              blockIdx.y * blockDim.x + threadIdx.x] =
+          tokens[(batch * pre_sequence_length + i) * gridDim.y * blockDim.x +
+                 blockIdx.y * blockDim.x + threadIdx.x];
+      output1[batch * new_sequence_length + index] = i;
+      index++;
+    }
+  }
+}
+
+nvinfer1::DimsExprs FusedTokenPrunePluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  auto x_dims = inputs[1], new_mask_dims = inputs[3];
+  if (flag_varseqlen_) {
+    // max sum of seqlen: ceil(sum / scale) + n -1 >= for(i=0;i<n;i++) {sum +=
+    // floor(num(i) / scale)} auto
+    // pruned_sum_length=std::ceil(inputs[4].d[0]*new_mask_dims.d[2]/inputs[6].d[1])+
+    // inputs[1].d[0] - 1;
+    auto pruned_sum_length = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUB,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kSUM,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kCEIL_DIV,
+                *expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
+                                        *inputs[4].d[0],
+                                        *new_mask_dims.d[2]),
+                *inputs[6].d[1]),
+            *inputs[1].d[0]),
+        *expr_builder.constant(1));
+    if (output_index == 0) {
+      nvinfer1::DimsExprs ret;
+      ret.nbDims = 4;
+      ret.d[0] = pruned_sum_length;
+      ret.d[1] = x_dims.d[2];
+      ret.d[2] = expr_builder.constant(1);
+      ret.d[3] = expr_builder.constant(1);
+      return ret;
+    } else if (output_index == 1) {
+      nvinfer1::DimsExprs ret;
+      ret.nbDims = 2;
+      ret.d[0] = new_mask_dims.d[0];
+      ret.d[1] = new_mask_dims.d[2];
+      return ret;
+    } else if (output_index == 2) {
+      // word id
+      nvinfer1::DimsExprs ret;
+      ret.nbDims = 1;
+      ret.d[0] = pruned_sum_length;
+      return ret;
+    } else if (output_index == 3) {
+      // pos id
+      nvinfer1::DimsExprs ret = inputs[5];
+      return ret;
+    } else if (output_index == 4) {
+      // mask id
+      nvinfer1::DimsExprs ret;
+      ret.nbDims = 2;
+      ret.d[0] = inputs[6].d[0];
+      ret.d[1] = new_mask_dims.d[2];
+      return ret;
+    }
+  } else {
+    if (output_index == 0) {
+      nvinfer1::DimsExprs ret = x_dims;
+      ret.d[1] = new_mask_dims.d[2];
+      return ret;
+    } else {
+      nvinfer1::DimsExprs ret;
+      ret.nbDims = 2;
+      ret.d[0] = new_mask_dims.d[0];
+      ret.d[1] = new_mask_dims.d[2];
+      return ret;
+    }
+  }
+}
+
+bool FusedTokenPrunePluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc& in = in_out[pos];
+  if (flag_varseqlen_) {
+    if (pos <= 3 || pos == 7) {
+      if (with_fp16_) {
+        return (in.type == nvinfer1::DataType::kHALF) &&
+               (in.format == nvinfer1::TensorFormat::kLINEAR);
+      } else {
+        PADDLE_THROW(
+            common::errors::Fatal("The FusedTokenPrune TRT Plugin's input type "
+                                  "should be half for varseqlen."));
+      }
+    } else if (pos == 6 || pos == 11) {  // mask_id, mask_id_out
+      return (in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return in.type == nvinfer1::DataType::kINT32 &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    }
+  } else {
+    if (pos == 0) {
+      if (with_fp16_) {
+        return (in.type == nvinfer1::DataType::kHALF) &&
+               (in.format == nvinfer1::TensorFormat::kLINEAR);
+      } else {
+        return (in.type == nvinfer1::DataType::kFLOAT) &&
+               (in.format == nvinfer1::TensorFormat::kLINEAR);
+      }
+    } else if (pos <= 4) {
+      const nvinfer1::PluginTensorDesc& prev = in_out[0];
+      return in.type == prev.type && in.format == prev.format;
+    } else {
+      return in.type == nvinfer1::DataType::kINT32 &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    }
+  }
+}
+
+nvinfer1::DataType FusedTokenPrunePluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  if (flag_varseqlen_) {
+    if (index == 0) {
+      return nvinfer1::DataType::kHALF;
+    } else if (index == 4) {  // mask id
+      return input_types[6];
+    } else {
+      // index = 1,2,3
+      return nvinfer1::DataType::kINT32;
+    }
+  } else {
+    if (index == 0) {
+      return input_types[1];
+    } else {
+      // index = 1
+      return nvinfer1::DataType::kINT32;
+    }
+  }
+}
+
+size_t FusedTokenPrunePluginDynamic::getWorkspaceSize(
+    const nvinfer1::PluginTensorDesc* inputs,
+    int nb_inputs,
+    const nvinfer1::PluginTensorDesc* outputs,
+    int nb_outputs) const TRT_NOEXCEPT {
+  auto attn_dims = inputs[0].dims;
+  auto x_dims = inputs[1].dims;
+  auto new_mask_dims = inputs[3].dims;
+  auto bsz = attn_dims.d[0], nb_head = attn_dims.d[1],
+       max_seq_len = attn_dims.d[2];
+
+  int slimmed_x_len = new_mask_dims.d[2];
+  int total = bsz * nb_head * max_seq_len * max_seq_len;
+  size_t size = total * sizeof(float);
+  size += bsz * max_seq_len * sizeof(float);
+  size += bsz * max_seq_len * sizeof(int32_t);
+  size += bsz * max_seq_len * sizeof(float);
+  size += bsz * max_seq_len * sizeof(int32_t);
+  size += (bsz + 1) * sizeof(int);
+  size += bsz * slimmed_x_len * sizeof(int32_t);
+  return size;
+}
+
+int FusedTokenPrunePluginDynamic::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 {
+  if (flag_varseqlen_) {
+    if (!(input_desc[0].type == nvinfer1::DataType::kHALF &&
+          input_desc[1].type == nvinfer1::DataType::kHALF)) {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Token_prune'type must half for varseqlen"));
+    }
+    float scale =
+        static_cast<float>(input_desc[3].dims.d[2]) / input_desc[2].dims.d[2];
+    const int32_t* input5 =
+        static_cast<const int32_t*>(inputs[5]);            // pre pos id
+    int32_t* output3 = static_cast<int32_t*>(outputs[3]);  // new pos id
+    half* output0 = static_cast<half*>(outputs[0]);
+    const int32_t B = input_desc[1].dims.d[0];  // batches
+    const int32_t max_sequence_length =
+        input_desc[1].dims.d[1];                     // max sequence length
+    const int32_t length = input_desc[1].dims.d[2];  // hidden size
+    const half* scores = static_cast<const half*>(inputs[0]);  // reduce sum
+    const half* tokens = static_cast<const half*>(inputs[1]);
+    int32_t padding_token_length;
+    if (max_sequence_length <= 64) {
+      padding_token_length = 64;
+    } else if (max_sequence_length <= 128) {
+      padding_token_length = 128;
+    } else if (max_sequence_length <= 256) {
+      padding_token_length = 256;
+    } else if (max_sequence_length <= 384) {
+      padding_token_length = 384;
+    } else if (max_sequence_length <= 512) {
+      padding_token_length = 512;
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Token_prune'token_length must <= 512"));
+    }
+
+    // 1. Compute the token length after pruning.
+    compute_token_length<<<1, B, 0, stream>>>(
+        input5, pruned_token_lengths_, scale);
+
+    // 2. Padding scores
+    fill_index_padding_score<half><<<B, padding_token_length, 0, stream>>>(
+        token_index_,
+        scores,
+        max_sequence_length,
+        static_cast<half*>(padding_scores_));
+
+    // 3. compute new pos id
+    // Determine temporary device storage requirements
+    size_t temp_storage_bytes = 0;
+    cub::DeviceScan::ExclusiveSum(
+        NULL, temp_storage_bytes, pruned_token_lengths_, output3, B + 1);
+    // Allocate temporary storage
+
+    GPUPlace place(platform::GetCurrentDeviceId());
+    auto d_temp_storage = phi::memory_utils::Alloc(place, temp_storage_bytes);
+
+    // Run exclusive prefix sum
+    cub::DeviceScan::ExclusiveSum(d_temp_storage->ptr(),
+                                  temp_storage_bytes,
+                                  pruned_token_lengths_,
+                                  output3,
+                                  B + 1);
+
+    // 4. sort scores
+    if (padding_token_length == 64) {
+      general_topk_pair_sort<half, 32, 2><<<B, 32, 0, stream>>>(
+          static_cast<half*>(padding_scores_), token_index_);  // 64
+    } else if (padding_token_length == 128) {
+      general_topk_pair_sort<half, 32, 4><<<B, 32, 0, stream>>>(
+          static_cast<half*>(padding_scores_), token_index_);  // 128
+    } else if (padding_token_length == 256) {
+      general_topk_pair_sort<half, 64, 4><<<B, 64, 0, stream>>>(
+          static_cast<half*>(padding_scores_), token_index_);  // 256
+    } else if (padding_token_length == 384) {
+      general_topk_pair_sort<half, 96, 4><<<B, 96, 0, stream>>>(
+          static_cast<half*>(padding_scores_), token_index_);  // 384
+    } else {
+      general_topk_pair_sort<half, 128, 4><<<B, 128, 0, stream>>>(
+          static_cast<half*>(padding_scores_), token_index_);  // 512
+    }
+
+    // 5. compute output
+    int32_t num_threads;
+    if (length < 1024) {
+      num_threads = length;
+    } else {
+      if (length % 512 == 0) {
+        num_threads = 512;
+      } else if (length % 256 == 0) {
+        num_threads = 256;
+      } else if (length % 128 == 0) {
+        num_threads = 128;
+      } else if (length % 64 == 0) {
+        num_threads = 64;
+      } else if (length % 32 == 0) {
+        num_threads = 32;
+      } else if (length % 16 == 0) {
+        num_threads = 16;
+      } else if (length % 8 == 0) {
+        num_threads = 8;
+      } else if (length % 4 == 0) {
+        num_threads = 4;
+      } else if (length % 2 == 0) {
+        num_threads = 2;
+      } else {
+        num_threads = 1;
+      }
+    }
+    const dim3 num_blocks(
+        B,
+        max_sequence_length,
+        length / num_threads);  // batches, max_sequence_length, vector_length
+    varlen_prune_token_change_order<<<num_blocks, num_threads, 0, stream>>>(
+        tokens, output3, padding_token_length, token_index_, output0);
+  } else {
+    auto input_type = input_desc[0].type;
+    const int32_t B = input_desc[1].dims.d[0];  // batches
+    const int32_t pre_sequence_length = input_desc[1].dims.d[1];
+    const int32_t new_sequence_length = input_desc[3].dims.d[2];  // new mask
+    const int32_t length = input_desc[1].dims.d[2];               // hidden size
+    if (input_type == nvinfer1::DataType::kFLOAT) {
+      VLOG(1) << "TRT Plugin DataType selected. FusedTokenPrune-->fp32";
+      const float* scores = static_cast<const float*>(inputs[0]);  // reduce sum
+      const float* tokens = static_cast<const float*>(inputs[1]);  // X
+      float* output0 = static_cast<float*>(outputs[0]);
+      int32_t* output1 = static_cast<int32_t*>(outputs[1]);
+      int32_t padding_token_length;
+      if (pre_sequence_length <= 64) {
+        padding_token_length = 64;
+      } else if (pre_sequence_length <= 128) {
+        padding_token_length = 128;
+      } else if (pre_sequence_length <= 256) {
+        padding_token_length = 256;
+      } else if (pre_sequence_length <= 384) {
+        padding_token_length = 384;
+      } else if (pre_sequence_length <= 512) {
+        padding_token_length = 512;
+      } else {
+        PADDLE_THROW(common::errors::InvalidArgument(
+            "Token_prune'token_length must <= 512"));
+      }
+
+      // 1. Padding scores
+      fill_index_padding_score<float><<<B, padding_token_length, 0, stream>>>(
+          token_index_,
+          scores,
+          pre_sequence_length,
+          static_cast<float*>(padding_scores_));
+
+      // 2. sort scores
+      if (padding_token_length == 64) {
+        general_topk_pair_sort<float, 32, 2><<<B, 32, 0, stream>>>(
+            static_cast<float*>(padding_scores_), token_index_);  // 64
+      } else if (padding_token_length == 128) {
+        general_topk_pair_sort<float, 32, 4><<<B, 32, 0, stream>>>(
+            static_cast<float*>(padding_scores_), token_index_);  // 128
+      } else if (padding_token_length == 256) {
+        general_topk_pair_sort<float, 64, 4><<<B, 64, 0, stream>>>(
+            static_cast<float*>(padding_scores_), token_index_);  // 256
+      } else if (padding_token_length == 384) {
+        general_topk_pair_sort<float, 96, 4><<<B, 96, 0, stream>>>(
+            static_cast<float*>(padding_scores_), token_index_);  // 384
+      } else {
+        general_topk_pair_sort<float, 128, 4><<<B, 128, 0, stream>>>(
+            static_cast<float*>(padding_scores_), token_index_);  // 512
+      }
+
+      // 3. compute output
+      int32_t num_threads;
+      if (length < 1024) {
+        num_threads = length;
+      } else {
+        if (length % 512 == 0) {
+          num_threads = 512;
+        } else if (length % 256 == 0) {
+          num_threads = 256;
+        } else if (length % 128 == 0) {
+          num_threads = 128;
+        } else if (length % 64 == 0) {
+          num_threads = 64;
+        } else if (length % 32 == 0) {
+          num_threads = 32;
+        } else if (length % 16 == 0) {
+          num_threads = 16;
+        } else if (length % 8 == 0) {
+          num_threads = 8;
+        } else if (length % 4 == 0) {
+          num_threads = 4;
+        } else if (length % 2 == 0) {
+          num_threads = 2;
+        } else {
+          num_threads = 1;
+        }
+      }
+      if (keep_order_) {
+        const dim3 num_blocks(B, length / num_threads);
+        prune_token_keep_order<float>
+            <<<num_blocks, num_threads, 0, stream>>>(tokens,
+                                                     pre_sequence_length,
+                                                     new_sequence_length,
+                                                     padding_token_length,
+                                                     token_index_,
+                                                     output0,
+                                                     output1);
+      } else {
+        const dim3 num_blocks(B, pre_sequence_length, length / num_threads);
+        prune_token_change_order<float>
+            <<<num_blocks, num_threads, 0, stream>>>(tokens,
+                                                     new_sequence_length,
+                                                     padding_token_length,
+                                                     token_index_,
+                                                     output0);
+      }
+    } else if (input_type == nvinfer1::DataType::kHALF) {
+      VLOG(1) << "TRT Plugin DataType selected. FusedTokenPrune-->fp16";
+      const half* scores = static_cast<const half*>(inputs[0]);  // reduce sum
+      const half* tokens = static_cast<const half*>(inputs[1]);  // X
+      half* output0 = static_cast<half*>(outputs[0]);
+      int32_t* output1 = static_cast<int32_t*>(outputs[1]);
+      int32_t padding_token_length;
+      if (pre_sequence_length <= 64) {
+        padding_token_length = 64;
+      } else if (pre_sequence_length <= 128) {
+        padding_token_length = 128;
+      } else if (pre_sequence_length <= 256) {
+        padding_token_length = 256;
+      } else if (pre_sequence_length <= 384) {
+        padding_token_length = 384;
+      } else if (pre_sequence_length <= 512) {
+        padding_token_length = 512;
+      } else {
+        PADDLE_THROW(common::errors::InvalidArgument(
+            "Token_prune'token_length must <= 512"));
+      }
+
+      // 1. Padding scores
+      fill_index_padding_score<half><<<B, padding_token_length, 0, stream>>>(
+          token_index_,
+          scores,
+          pre_sequence_length,
+          static_cast<half*>(padding_scores_));
+
+      // 2. sort scores
+      if (padding_token_length == 64) {
+        general_topk_pair_sort<half, 32, 2><<<B, 32, 0, stream>>>(
+            static_cast<half*>(padding_scores_), token_index_);  // 64
+      } else if (padding_token_length == 128) {
+        general_topk_pair_sort<half, 32, 4><<<B, 32, 0, stream>>>(
+            static_cast<half*>(padding_scores_), token_index_);  // 128
+      } else if (padding_token_length == 256) {
+        general_topk_pair_sort<half, 64, 4><<<B, 64, 0, stream>>>(
+            static_cast<half*>(padding_scores_), token_index_);  // 256
+      } else if (padding_token_length == 384) {
+        general_topk_pair_sort<half, 96, 4><<<B, 96, 0, stream>>>(
+            static_cast<half*>(padding_scores_), token_index_);  // 384
+      } else {
+        general_topk_pair_sort<half, 128, 4><<<B, 128, 0, stream>>>(
+            static_cast<half*>(padding_scores_), token_index_);  // 512
+      }
+
+      // 3. compute output
+      int32_t num_threads;
+      if (length < 1024) {
+        num_threads = length;
+      } else {
+        if (length % 512 == 0) {
+          num_threads = 512;
+        } else if (length % 256 == 0) {
+          num_threads = 256;
+        } else if (length % 128 == 0) {
+          num_threads = 128;
+        } else if (length % 64 == 0) {
+          num_threads = 64;
+        } else if (length % 32 == 0) {
+          num_threads = 32;
+        } else if (length % 16 == 0) {
+          num_threads = 16;
+        } else if (length % 8 == 0) {
+          num_threads = 8;
+        } else if (length % 4 == 0) {
+          num_threads = 4;
+        } else if (length % 2 == 0) {
+          num_threads = 2;
+        } else {
+          num_threads = 1;
+        }
+      }
+      if (keep_order_) {
+        const dim3 num_blocks(B, length / num_threads);
+        prune_token_keep_order<half>
+            <<<num_blocks, num_threads, 0, stream>>>(tokens,
+                                                     pre_sequence_length,
+                                                     new_sequence_length,
+                                                     padding_token_length,
+                                                     token_index_,
+                                                     output0,
+                                                     output1);
+      } else {
+        const dim3 num_blocks(B, pre_sequence_length, length / num_threads);
+        prune_token_change_order<half>
+            <<<num_blocks, num_threads, 0, stream>>>(tokens,
+                                                     new_sequence_length,
+                                                     padding_token_length,
+                                                     token_index_,
+                                                     output0);
+      }
+    } else {
+      PADDLE_THROW(
+          common::errors::Fatal("The FusedTokenPrune TRT Plugin's input type "
+                                "should be float or half."));
+    }
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/fused_token_prune_op_plugin.h

@@ -0,0 +1,206 @@
+// 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.
+
+#pragma once
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class FusedTokenPrunePluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit FusedTokenPrunePluginDynamic(bool with_fp16,
+                                        bool keep_first_token,
+                                        bool keep_order,
+                                        bool flag_varseqlen)
+      : with_fp16_(with_fp16),
+        keep_first_token_(keep_first_token),
+        keep_order_(keep_order),
+        flag_varseqlen_(flag_varseqlen) {}
+  FusedTokenPrunePluginDynamic(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+    DeserializeValue(&serial_data, &serial_length, &keep_first_token_);
+    DeserializeValue(&serial_data, &serial_length, &keep_order_);
+    DeserializeValue(&serial_data, &serial_length, &flag_varseqlen_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    FusedTokenPrunePluginDynamic* ptr = new FusedTokenPrunePluginDynamic(
+        with_fp16_, keep_first_token_, keep_order_, flag_varseqlen_);
+    ptr->max_batches_ = max_batches_;
+    ptr->max_token_length_ = max_token_length_;
+    ptr->pruned_token_lengths_ = pruned_token_lengths_;
+    ptr->token_index_ = token_index_;
+    ptr->padding_scores_ = padding_scores_;
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "fused_token_prune_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override {
+    if (flag_varseqlen_) {
+      return 5;
+    } else {
+      return 2;
+    }
+  }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(with_fp16_) + SerializedSize(keep_first_token_) +
+           SerializedSize(keep_order_) + SerializedSize(flag_varseqlen_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, with_fp16_);
+    SerializeValue(&buffer, keep_first_token_);
+    SerializeValue(&buffer, keep_order_);
+    SerializeValue(&buffer, flag_varseqlen_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override {
+    max_batches_ = in[1].max.d[0];
+    max_token_length_ = in[1].max.d[1];
+    int32_t padding_token_length;
+    if (max_token_length_ <= 64) {
+      padding_token_length = 64;
+    } else if (max_token_length_ <= 128) {
+      padding_token_length = 128;
+    } else if (max_token_length_ <= 256) {
+      padding_token_length = 256;
+    } else if (max_token_length_ <= 384) {
+      padding_token_length = 384;
+    } else if (max_token_length_ <= 512) {
+      padding_token_length = 512;
+    } else {
+      try {
+        PADDLE_THROW(common::errors::InvalidArgument(
+            "Token_prune'token_length(max) must <= 512"));
+      } catch (std::exception& e) {
+      }
+    }
+    try {
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaMalloc(
+          &pruned_token_lengths_, (max_batches_ + 1) * sizeof(int32_t)));
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaMalloc(&token_index_,
+                     max_batches_ * padding_token_length * sizeof(int32_t)));
+      int32_t type_size = 4;
+      if (in[0].desc.type == nvinfer1::DataType::kHALF) {
+        type_size = 2;
+      } else {
+        type_size = 4;
+      }
+      PADDLE_ENFORCE_GPU_SUCCESS(cudaMalloc(
+          &padding_scores_, max_batches_ * padding_token_length * type_size));
+    } catch (std::exception& e) {
+    }
+  }
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nb_inputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nb_outputs) const TRT_NOEXCEPT override;
+
+  int 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 override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  bool with_fp16_;
+  bool keep_first_token_;
+  bool keep_order_;
+  bool flag_varseqlen_;
+  int32_t* pruned_token_lengths_;
+  int32_t* token_index_;
+  int32_t max_batches_;
+  int32_t max_token_length_;
+  void* padding_scores_;
+};
+
+class FusedTokenPrunePluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  FusedTokenPrunePluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "fused_token_prune_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new FusedTokenPrunePluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_;
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+REGISTER_TRT_PLUGIN_V2(FusedTokenPrunePluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/gather_nd_op_plugin.cu

@@ -0,0 +1,262 @@
+// Copyright (c) 2021 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 <cuda_fp16.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <functional>
+#include <numeric>
+#include <sstream>
+
+#include "NvInferRuntimeCommon.h"
+#include "paddle/fluid/inference/tensorrt/plugin/gather_nd_op_plugin.h"
+#include "paddle/phi/backends/gpu/gpu_helper.h"
+#include "paddle/phi/common/place.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T, typename IndexT = int>
+__global__ void GatherNdCUDAKernel(const T* input,
+                                   const int32_t* input_dims,
+                                   const IndexT* indices,
+                                   T* output,
+                                   int32_t remain_size,
+                                   int32_t slice_size,
+                                   int32_t end_size) {
+  CUDA_KERNEL_LOOP(i, remain_size * slice_size) {
+    int indices_i = i / slice_size;
+    int slice_i = i - indices_i * slice_size;  // offset inside the slice
+    IndexT gather_i = 0;
+    int32_t temp = slice_size;
+    for (int32_t j = end_size - 1; j >= 0; --j) {
+      auto index_value = indices[indices_i * end_size + j];
+      PADDLE_ENFORCE(
+          index_value >= -input_dims[j] && index_value < input_dims[j],
+          "The index is out of bounds, "
+          "please check whether the dimensions of index and "
+          "input meet the requirements. It should "
+          "be less than [%d] and greater or equal to [%d], but received [%d]",
+          input_dims[j],
+          -input_dims[j],
+          index_value);
+      if (index_value < 0) {
+        index_value += input_dims[j];
+      }
+      gather_i += (index_value * temp);
+      temp *= input_dims[j];
+    }
+    IndexT input_i = gather_i + slice_i;
+    *(output + i) = *(input + input_i);
+  }
+}
+
+int GatherNdPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+size_t GatherNdPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(with_fp16_);
+}
+
+void GatherNdPluginDynamic::serialize(void* buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, with_fp16_);
+}
+
+nvinfer1::DimsExprs GatherNdPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      nb_inputs,
+      2,
+      common::errors::InvalidArgument(
+          "The gather_nd plugin should have 2 input, but got %d.", nb_inputs));
+  PADDLE_ENFORCE_EQ(
+      output_index,
+      0,
+      common::errors::InvalidArgument("When GetOutputDimensions in gather_nd "
+                                      "plugin, the output_index should be 0."));
+
+  nvinfer1::DimsExprs x_dims = inputs[0];
+  nvinfer1::DimsExprs index_dims = inputs[1];
+
+  int32_t x_dims_size = x_dims.nbDims;
+  int32_t index_dims_size = index_dims.nbDims;
+
+  // TODO(wilber): The result dims should be Index.shape[:-1] +
+  // X.shape[Index.shape[-1]:], but the trt DimsExprs is an expression we can't
+  // get the actual value. So we only support one scenario: input_dims.size ==
+  // index_dims.size.
+  nvinfer1::DimsExprs ret(x_dims);
+  for (int i = 0; i < index_dims_size - 1; ++i) {
+    ret.d[i] = index_dims.d[i];
+  }
+
+  return ret;
+}
+
+bool GatherNdPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of gather_nd plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc& in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kFLOAT ||
+              in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  } else if (pos == 1) {
+    return in.type == nvinfer1::DataType::kINT32 &&
+           in.format == nvinfer1::TensorFormat::kLINEAR;
+  } else if (pos == 2) {
+    return in.type == in_out[0].type &&
+           in.format == nvinfer1::TensorFormat::kLINEAR;
+  }
+
+  return true;
+}
+
+nvinfer1::DataType GatherNdPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  return input_types[0];
+}
+
+int GatherNdPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  auto index_dims = input_desc[1].dims;
+  auto input_dims_size = input_dims.nbDims;
+  auto index_dims_size = index_dims.nbDims;
+
+  std::vector<int32_t> input_shape, index_shape, out_shape;
+  for (int i = 0; i < input_dims.nbDims; i++)
+    input_shape.push_back(input_dims.d[i]);
+  for (int i = 0; i < index_dims.nbDims; i++)
+    index_shape.push_back(index_dims.d[i]);
+  // The out_shape is
+  //   Index.shape[:-1] + X.shape[Index.shape[-1]:]
+  for (int i = 0; i < index_dims_size - 1; ++i) {
+    out_shape.emplace_back(index_shape[i]);
+  }
+  for (int i = index_shape[index_dims_size - 1]; i < input_dims_size; ++i) {
+    out_shape.emplace_back(input_shape[i]);
+  }
+
+  // final dim
+  int end_size = index_shape[index_dims_size - 1];
+  // remain dim
+  std::vector<int> remain_ddim(index_shape.begin(), index_shape.end() - 1);
+  int remain_numel = std::accumulate(
+      remain_ddim.begin(), remain_ddim.end(), 1, std::multiplies<int>());
+  // slice size
+  int slice_size = 1;
+  for (int i = end_size; i < input_dims_size; ++i) {
+    slice_size *= input_shape[i];
+  }
+
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. gather_nd-->fp32";
+
+    const float* p_input = static_cast<const float*>(inputs[0]);
+    const int32_t* p_index = static_cast<const int32_t*>(inputs[1]);
+    float* p_output = static_cast<float*>(outputs[0]);
+
+    if (input_dims_data_ == nullptr) {
+      cudaMalloc(&input_dims_data_, input_shape.size() * sizeof(int));
+    }
+    cudaMemcpyAsync(input_dims_data_,
+                    input_shape.data(),
+                    sizeof(int) * input_shape.size(),
+                    cudaMemcpyHostToDevice,
+                    stream);
+
+    int block = 512;
+    int n = slice_size * remain_numel;
+    int grid = (n + block - 1) / block;
+
+    GatherNdCUDAKernel<float, int32_t>
+        <<<grid, block, 0, stream>>>(p_input,
+                                     input_dims_data_,
+                                     p_index,
+                                     p_output,
+                                     remain_numel,
+                                     slice_size,
+                                     end_size);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. gather_nd-->fp16";
+
+    const half* p_input = static_cast<const half*>(inputs[0]);
+    const int32_t* p_index = static_cast<const int32_t*>(inputs[1]);
+    half* p_output = static_cast<half*>(outputs[0]);
+
+    if (input_dims_data_ == nullptr) {
+      cudaMalloc(&input_dims_data_, input_shape.size() * sizeof(int));
+    }
+    cudaMemcpyAsync(input_dims_data_,
+                    input_shape.data(),
+                    sizeof(int) * input_shape.size(),
+                    cudaMemcpyHostToDevice,
+                    stream);
+
+    int block = 512;
+    int n = slice_size * remain_numel;
+    int grid = (n + block - 1) / block;
+
+    GatherNdCUDAKernel<half, int32_t>
+        <<<grid, block, 0, stream>>>(p_input,
+                                     input_dims_data_,
+                                     p_index,
+                                     p_output,
+                                     remain_numel,
+                                     slice_size,
+                                     end_size);
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/gather_nd_op_plugin.h

@@ -0,0 +1,143 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class GatherNdPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit GatherNdPluginDynamic(bool with_fp16) { with_fp16_ = with_fp16; }
+
+  GatherNdPluginDynamic(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+  }
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new GatherNdPluginDynamic(with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "gather_nd_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override {
+    if (input_dims_data_) {
+      cudaFree(input_dims_data_);
+    }
+    delete this;
+  }
+
+ private:
+  int32_t* input_dims_data_{nullptr};
+};
+
+class GatherNdPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  GatherNdPluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "gather_nd_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new GatherNdPluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+
+REGISTER_TRT_PLUGIN_V2(GatherNdPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/gelu_op_plugin.cu

@@ -0,0 +1,221 @@
+// Copyright (c) 2018 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 <cassert>
+#include <cstring>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/gelu_op_plugin.h"
+#include "paddle/phi/common/float16.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+// constants for approximating the normal cdf
+static const float kA = 1.41421356237309504;  // sqrt(2)
+
+static const float kAT = 0.5;
+static const float kBT = 0.7978845608028654;    // sqrt(2.0/M_PI)
+static const float kCT = 0.035677408136300125;  // 0.044715 * sqrt(2.0/M_PI)
+
+bool GeluPlugin::supportsFormat(
+    nvinfer1::DataType type, nvinfer1::PluginFormat format) const TRT_NOEXCEPT {
+  if (with_fp16_) {
+    return ((type == nvinfer1::DataType::kFLOAT ||
+             type == nvinfer1::DataType::kHALF) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  } else {
+    return ((type == nvinfer1::DataType::kFLOAT) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  }
+}
+
+nvinfer1::Dims GeluPlugin::getOutputDimensions(int index,
+                                               const nvinfer1::Dims* in_dims,
+                                               int nb_inputs) TRT_NOEXCEPT {
+  assert(nb_inputs == 1);
+  assert(index < this->getNbOutputs());
+  nvinfer1::Dims const& input_dims = in_dims[0];
+  nvinfer1::Dims output_dims = input_dims;
+  return output_dims;
+}
+
+template <typename T, unsigned TPB>
+__global__ void gelu_kernel(const T a, int n, const T* input, T* output) {
+  const int idx = blockIdx.x * TPB + threadIdx.x;
+  if (idx < n) {
+    const T in = input[idx];
+    const T cdf = 0.5f * (1.0f + erff(in * 0.5f * a));
+    output[idx] = in * cdf;
+  }
+}
+
+template <typename T>
+__device__ T do_tanh(T a);
+
+template <>
+__device__ float do_tanh<float>(float a) {
+  return tanf(a);
+}
+
+template <>
+__device__ half do_tanh<half>(half a) {
+  const float tmp = tanhf(__half2float(a));
+  return __float2half(tmp);
+}
+
+// the kernel below is not aligned with fluid fp32 forward ones, use it for
+// fp16.
+template <typename T, unsigned TPB>
+__global__ void no_exact_gelu_kernel(
+    const T a, const T b, const T c, int n, const T* input, T* output) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const int idx = blockIdx.x * TPB + threadIdx.x;
+  if (idx < n) {
+    const T in = input[idx];
+    const T tmp = in * (c * in * in + b);
+    const T cdf = a + a * do_tanh<T>(tmp);
+    output[idx] = in * cdf;
+  }
+#endif
+}
+
+int GeluPlugin::enqueue(int batch_size,
+                        const void* const* inputs,
+                        void* const* outputs,
+                        void*,
+                        cudaStream_t stream) TRT_NOEXCEPT {
+  const auto& input_dims = this->getInputDims(0);
+  int num = batch_size;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    num *= input_dims.d[i];
+  }
+  const int block_size = 256;
+  const int grid_size = (num + block_size - 1) / block_size;
+
+  auto type = getDataType();
+  if (type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Gelu-->fp32";
+    const float* input = static_cast<const float*>(inputs[0]);
+    float* output = static_cast<float*>(outputs[0]);
+    gelu_kernel<float, block_size>
+        <<<grid_size, block_size, 0, stream>>>(kA, num, input, output);
+  } else if (type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Gelu-->fp16";
+    const half* input = static_cast<const half*>(inputs[0]);
+    half* output = static_cast<half*>(outputs[0]);
+    no_exact_gelu_kernel<half, block_size>
+        <<<grid_size, block_size, 0, stream>>>(
+            kAT, kBT, kCT, num, input, output);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The Gelu TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+nvinfer1::DimsExprs GeluPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  return inputs[0];
+}
+
+bool GeluPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc& in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kFLOAT ||
+              in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc& prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType GeluPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Gelu Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+int GeluPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  size_t num = ProductDim(input_dims);
+  const int block_size = 256;
+  const int grid_size = (num + block_size - 1) / block_size;
+
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Gelu-->fp32";
+    const float* input = static_cast<const float*>(inputs[0]);
+    float* output = static_cast<float*>(outputs[0]);
+    gelu_kernel<float, block_size>
+        <<<grid_size, block_size, 0, stream>>>(kA, num, input, output);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Gelu-->fp16";
+    const half* input = static_cast<const half*>(inputs[0]);
+    half* output = static_cast<half*>(outputs[0]);
+    no_exact_gelu_kernel<half, block_size>
+        <<<grid_size, block_size, 0, stream>>>(
+            kAT, kBT, kCT, num, input, output);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The Gelu TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/gelu_op_plugin.h

@@ -0,0 +1,172 @@
+// Copyright (c) 2019 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.
+
+#pragma once
+#include <stdio.h>
+
+#include <cassert>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class GeluPlugin : public PluginTensorRT {
+ public:
+  explicit GeluPlugin(const bool with_fp16) { with_fp16_ = with_fp16; }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  GeluPlugin(void const* serial_data, size_t serial_length) {
+    deserializeBase(serial_data, serial_length);
+  }
+
+  ~GeluPlugin() {}
+  GeluPlugin* clone() const TRT_NOEXCEPT override {
+    return new GeluPlugin(with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "gelu_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override;
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nb_input_dims) TRT_NOEXCEPT override;
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize();
+  }
+
+  // TRT will call this func  to serialize the configuration of TRT
+  // It should not be called by users.
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+  }
+};
+
+class GeluPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "gelu_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new GeluPlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(GeluPluginCreator);
+
+class GeluPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit GeluPluginDynamic(const bool with_fp16) { with_fp16_ = with_fp16; }
+  GeluPluginDynamic(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+  }
+
+  ~GeluPluginDynamic() {}
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new GeluPluginDynamic(with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "gelu_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nb_inputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nb_outputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int 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 override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+};
+
+class GeluPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "gelu_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new GeluPluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+};
+REGISTER_TRT_PLUGIN_V2(GeluPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/generic_plugin.cu

@@ -0,0 +1,730 @@
+// 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

paddle/fluid/inference/tensorrt/plugin/generic_plugin.h

@@ -0,0 +1,197 @@
+// Copyright (c) 2024 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.
+
+#pragma once
+
+#include <NvInfer.h>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/op_desc.h"
+#include "paddle/fluid/framework/op_registry.h"
+#include "paddle/fluid/framework/operator.h"
+#include "paddle/fluid/framework/scope.h"
+#include "paddle/fluid/framework/type_defs.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/helper.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin_utils.h"
+#include "paddle/fluid/inference/tensorrt/plugin_arg_mapping_context.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/core/kernel_context.h"
+#include "paddle/phi/core/memory/allocation/cuda_allocator.h"
+#include "paddle/phi/core/platform/device_context.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+enum class GeneratePluginDataType {
+  PLUGIN_BOOL,
+  PLUGIN_UINT8,
+  PLUGIN_INT8,
+  PLUGIN_INT16,
+  PLUGIN_INT32,
+  PLUGIN_INT64,
+  PLUGIN_FP16,
+  PLUGIN_FP32,
+  PLUGIN_FP64,
+  PLUGIN_BF16,
+  PLUGIN_SIZE_T,
+  PLUGIN_COMPLEX64,
+  PLUGIN_COMPLEX128,
+  PLUGIN_OPTIONAL,
+};
+
+GeneratePluginDataType ProtoTypeToGeneratePluginDataType(
+    framework::proto::VarType_Type proto_type);
+
+void BuildPhiKernelContextAttr(const framework::OpDesc& op_desc,
+                               phi::KernelContext* kernel_context,
+                               const phi::KernelSignature& signature,
+                               const phi::Kernel* phi_kernel);
+
+class GenericPlugin : public DynamicPluginTensorRT {
+ public:
+  struct InputOutPutVarInfo {
+    std::vector<GeneratePluginDataType> inputs_data_type;
+    std::vector<GeneratePluginDataType> outputs_data_type;
+  };
+
+ public:
+  GenericPlugin() {}
+
+  GenericPlugin(const paddle::framework::proto::OpDesc& proto_op_desc,
+                const InputOutPutVarInfo& in_out_info,
+                bool with_fp16_ = false);
+
+  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_ = false);
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  GenericPlugin(void const* serialData, size_t serialLength);
+
+  // IPluginV2 method
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "generic_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override;
+
+  int getNbInputs() const TRT_NOEXCEPT;
+
+  // Initialize the layer for execution.
+  int initialize() TRT_NOEXCEPT override;
+
+  // Shutdown the layer. This is called when the engine is destroyed
+  void terminate() TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override{};
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    size_t sum = 0;
+    sum += SerializedSize(inputs_data_type_);
+    sum += SerializedSize(outputs_data_type_);
+    sum += SerializedSize(with_fp16_);
+    sum += op_meta_data_.size();
+    return sum;
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  // The Func in IPluginV2
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override;
+
+  int 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 override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  bool isFp16Supported() {
+    auto half_dtype = nvinfer1::DataType::kHALF;
+    return with_fp16_ &&
+           !(phi_kernels_.find(half_dtype) == phi_kernels_.end()) &&
+           phi_kernels_[half_dtype]->IsValid();
+  }
+
+ private:
+  std::string op_meta_data_;
+  framework::proto::OpDesc proto_op_desc_;
+  framework::OpDesc op_desc_;
+
+ private:
+  std::unordered_map<nvinfer1::DataType, std::unique_ptr<phi::Kernel>>
+      phi_kernels_;
+  std::unordered_map<nvinfer1::DataType, std::unique_ptr<phi::KernelContext>>
+      phi_kernel_contexts_;
+
+  std::vector<phi::DenseTensor>* dense_tensor_inputs_{nullptr};
+  std::vector<phi::DenseTensor>* dense_tensor_outputs_{nullptr};
+
+ private:
+  std::vector<GeneratePluginDataType> inputs_data_type_;
+  std::vector<GeneratePluginDataType> outputs_data_type_;
+};
+
+class GenericPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "generic_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2DynamicExt* deserializePlugin(const char* name,
+                                                   const void* serial_data,
+                                                   size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new GenericPlugin(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(GenericPluginCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/group_norm_op_plugin.cu

@@ -0,0 +1,830 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+SPDX-FileCopyrightText: Copyright (c) 1993-2022 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 "paddle/fluid/inference/tensorrt/plugin/group_norm_op_plugin.h"
+#include "paddle/phi/kernels/group_norm_kernel.h"
+
+#include <cub/cub.cuh>
+#include "paddle/common/layout.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+using DataLayout = phi::DataLayout;
+
+template <typename T>
+static inline T divUp(T m, T n) {
+  return (m + n - 1) / n;
+}
+
+static inline __device__ __host__ float sigmoid(float x) {
+  return 1.F / (1.F + expf(-x));
+}
+
+struct GroupSums {
+  // Is it the 1st element of the group?
+  int32_t flag;
+  // The sum.
+  float sum;
+  // The sum of squares.
+  float sumSq;
+};
+
+struct GroupSumsOp {
+  inline __device__ GroupSums operator()(GroupSums const &a,
+                                         GroupSums const &b) {
+    GroupSums dst;
+    dst.sum = b.flag ? b.sum : (a.sum + b.sum);
+    dst.sumSq = b.flag ? b.sumSq : (a.sumSq + b.sumSq);
+    dst.flag = a.flag + b.flag;
+    return dst;
+  }
+};
+
+static int64_t findMaxDivisor(int64_t n, int64_t maxAllowedDivisor) {
+  int64_t maxDivisor = -1;
+  for (int64_t i = 1; i <= std::sqrt(n); i++) {
+    if (n % i == 0) {
+      int64_t divisor1 = n / i;
+      int64_t divisor2 = i;
+
+      if (divisor1 > maxDivisor && divisor1 < maxAllowedDivisor) {
+        maxDivisor = divisor1;
+      }
+      if (divisor2 > maxDivisor && divisor2 < maxAllowedDivisor) {
+        maxDivisor = divisor2;
+      }
+    }
+  }
+  return maxDivisor;
+}
+
+template <int tTHREADS_PER_BLOCK>
+__global__ void groupNormNCHW32SumKernelQDQ(
+    const GroupNormNDHWCParams<__half> params) {
+  // The object in charge of doing the sums for the different blocks.
+  typedef cub::BlockScan<GroupSums, tTHREADS_PER_BLOCK> BlockScan;
+
+  // Allocate shared memory for BlockScan.
+  __shared__ typename BlockScan::TempStorage tempStorage;
+  // Allocate shared memory for the groups. We could reduce the amount of shared
+  // memory reserved.
+  __shared__ float2 smem[tTHREADS_PER_BLOCK];
+
+  // The instance in the batch.
+  int32_t ni = blockIdx.z;
+  // The channel loaded by that thread (2 channels per thread for int8x2).
+  int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * 2;
+
+  // The first activation loaded by that block.
+  int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
+  // The last activation loaded by that block.
+  int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
+
+  // The sums.
+  float sum = 0.F;
+  float sumSq = 0.F;
+
+  const int8_t *src_ptr = reinterpret_cast<const int8_t *>(params.srcX);
+
+  // nchw32 layout
+  // batch offset + channel offset
+  int64_t nc_offset = ni * params.dhwc + ci / 32 * params.dhw * 32 + ci % 32;
+
+  // Iterate over the activations to compute the sums.
+  for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
+    // The offset.
+    int64_t offset = nc_offset + static_cast<int64_t>(dhwi) * 32;
+
+    // Fetch two channels per thread.
+    __half2 h2(0, 0);
+    if (ci < params.c) {
+      int8_t tmp_in[2];
+      *reinterpret_cast<int16_t *>(tmp_in) =
+          *reinterpret_cast<int16_t const *>(&src_ptr[offset]);
+      h2.x = params.dqScaleIn * tmp_in[0];
+      h2.y = params.dqScaleIn * tmp_in[1];
+    }
+
+    // Extract the two half values.
+    float2 f2 = __half22float2(h2);
+
+    // Update the sum.
+    sum += f2.x + f2.y;
+    // Update the sum of squares.
+    sumSq += f2.x * f2.x + f2.y * f2.y;
+  }
+
+  // The group that thread works on and the channel in the group (modulus).
+  int32_t gi = threadIdx.x * 2 / params.cPerGroup;
+  int32_t cj = threadIdx.x * 2 - params.cPerGroup * gi;
+
+  // The data for the summations.
+  GroupSums inp{cj == 0 ? 1 : 0, sum, sumSq};
+
+  // Do the segmented scan.
+  GroupSums out;
+  BlockScan(tempStorage).InclusiveScan(inp, out, GroupSumsOp());
+
+  // Store the results for the groups in shared memory (to produce coalesced
+  // stores later).
+  // 2 channels per thread
+  if (cj == params.cPerGroup - 2) {
+    smem[gi] = make_float2(out.sum, out.sumSq);
+  }
+
+  // Make sure the data is in shared memory.
+  __syncthreads();
+
+  // The global group index.
+  int32_t gj = blockIdx.x * params.groupsPerBlock + threadIdx.x;
+
+  // Threads that have nothing left to do, exit.
+  if (threadIdx.x >= params.groupsPerBlock || gj >= params.groups) {
+    return;
+  }
+
+  // The first threads (those storing to global memory, load the values).
+  float2 sums = smem[threadIdx.x];
+
+  // Store to global memory.
+  atomicAdd(&params.redBuffer[(2 * ni + 0) * params.groups + gj], sums.x);
+  atomicAdd(&params.redBuffer[(2 * ni + 1) * params.groups + gj], sums.y);
+}
+
+void groupNormNCHW32SumQDQ(const GroupNormNDHWCParams<__half> &params,
+                           cudaStream_t stream) {
+  dim3 grid;
+
+  // The number of blocks to compute all the channels.
+  grid.x = divUp(params.c, params.cPerBlock);
+  // The number of blocks to compute all the activations in a given instance.
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
+  // The number of instances.
+  grid.z = params.n;
+
+  switch (params.cPerBlock) {
+    case 320:
+      groupNormNCHW32SumKernelQDQ<160><<<grid, 160, 0, stream>>>(params);
+      break;
+    case 480:
+      groupNormNCHW32SumKernelQDQ<256><<<grid, 256, 0, stream>>>(params);
+      break;
+    case 256:
+      groupNormNCHW32SumKernelQDQ<128><<<grid, 128, 0, stream>>>(params);
+      break;
+    case 128:
+      groupNormNCHW32SumKernelQDQ<64><<<grid, 64, 0, stream>>>(params);
+      break;
+    case 8:
+      groupNormNCHW32SumKernelQDQ<4><<<grid, 4, 0, stream>>>(params);
+      break;
+  }
+}
+
+template <int tTHREADS_PER_BLOCK>
+__global__ void groupNormNCHW32ScaleKernelQDQ(
+    const GroupNormNDHWCParams<__half> params) {
+  // The instance in the batch.
+  int32_t ni = blockIdx.z;
+  // The channel loaded by that thread (2 channels per thread for F16x2).
+  int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * 2;
+  // The group that thread works on and the channel in the group (modulus).
+  int32_t gi = ci / params.cPerGroup;
+
+  const int8_t *src_ptr = reinterpret_cast<const int8_t *>(params.srcX);
+  int8_t *dst_ptr = reinterpret_cast<int8_t *>(params.dst);
+
+  // Load the sum and sum of squares for the group.
+  float sum = 0.F, sumSq = 0.F;
+  if (gi < params.groups) {
+    sum = params.redBuffer[(2 * ni + 0) * params.groups + gi];
+    sumSq = params.redBuffer[(2 * ni + 1) * params.groups + gi];
+  }
+
+  // Load gamma/beta.
+  float2 gammaF2, betaF2;
+  if (ci < params.c) {
+    gammaF2 = __half22float2(*reinterpret_cast<half2 const *>(
+        reinterpret_cast<half const *>(params.gamma) + ci));
+    betaF2 = __half22float2(*reinterpret_cast<half2 const *>(
+        reinterpret_cast<half const *>(params.beta) + ci));
+  }
+
+  // Compute the mean.
+  float mean = sum * params.invDHWC;
+  // Compute the variance.
+  float var = sumSq * params.invDHWC - (mean * mean);
+  // Compute the inverse of the stddev.
+  float invStdDev = rsqrtf(var + params.eps);
+
+  // The first activation loaded by that block.
+  int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
+  // The last activation loaded by that block.
+  int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
+
+  // nchw32 layout
+  int64_t c_offset = ci / 32 * params.dhw * 32 + ci % 32;
+
+  // Iterate over the activations to compute the sums.
+  for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
+    // The src/dst offset.
+    int64_t offset = static_cast<int64_t>(ni) * params.dhwc + c_offset +
+                     static_cast<int64_t>(dhwi) * 32;
+
+    // Fetch two channels per thread.
+    __half2 h2(0, 0);
+    if (ci < params.c) {
+      int8_t tmp_in[2];
+      *reinterpret_cast<int16_t *>(tmp_in) =
+          *reinterpret_cast<int16_t const *>(&src_ptr[offset]);
+      h2.x = params.dqScaleIn * tmp_in[0];
+      h2.y = params.dqScaleIn * tmp_in[1];
+    }
+
+    // Extract the two half values.
+    float2 f2 = __half22float2(h2);
+
+    // Normalize the channels.
+    f2.x = (f2.x - mean) * invStdDev;
+    f2.y = (f2.y - mean) * invStdDev;
+
+    // Scale by gamma and add beta.
+    f2.x = gammaF2.x * f2.x + betaF2.x;
+    f2.y = gammaF2.y * f2.y + betaF2.y;
+
+    // Apply Silu if needed.
+    if (params.withSilu) {
+      f2.x = f2.x * sigmoid(f2.x);
+      f2.y = f2.y * sigmoid(f2.y);
+    }
+
+    // Store the scaled values.
+    if (ci < params.c) {
+      int8_t tmp_in[2];
+      int32_t tmpq0 = __float2int_rn(params.inv_qScale * f2.x);
+      int32_t tmpq1 = __float2int_rn(params.inv_qScale * f2.y);
+      tmpq0 = max(-128, tmpq0);
+      tmpq0 = min(127, tmpq0);
+      tmpq1 = max(-128, tmpq1);
+      tmpq1 = min(127, tmpq1);
+      tmp_in[0] = tmpq0;
+      tmp_in[1] = tmpq1;
+      *reinterpret_cast<int16_t *>(&dst_ptr[offset]) =
+          *reinterpret_cast<int16_t *>(tmp_in);
+    }
+  }
+}
+
+void groupNormNCHW32ScaleQDQ(const GroupNormNDHWCParams<__half> &params,
+                             cudaStream_t stream) {
+  dim3 grid;
+
+  // The number of blocks to compute all the channels.
+  grid.x = divUp(params.c, params.cPerBlock);
+  // The number of blocks to compute all the activations in a given instance.
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
+  // The number of instances.
+  grid.z = params.n;
+
+  switch (params.cPerBlock) {
+    case 320:
+      groupNormNCHW32ScaleKernelQDQ<160><<<grid, 160, 0, stream>>>(params);
+      break;
+    case 480:
+      groupNormNCHW32ScaleKernelQDQ<256><<<grid, 256, 0, stream>>>(params);
+      break;
+    case 256:
+      groupNormNCHW32ScaleKernelQDQ<128><<<grid, 128, 0, stream>>>(params);
+      break;
+    case 128:
+      groupNormNCHW32ScaleKernelQDQ<64><<<grid, 64, 0, stream>>>(params);
+      break;
+    case 8:
+      groupNormNCHW32ScaleKernelQDQ<4><<<grid, 4, 0, stream>>>(params);
+      break;
+    default:
+      PADDLE_THROW(
+          common::errors::Fatal("The function groupNormNCHW32ScaleQDQ of "
+                                "GroupNorm TRT Plugin encounter error"));
+  }
+}
+
+int GroupNormPlugin::initialize() TRT_NOEXCEPT {
+  if (!with_fp16_) {
+    // if use fp32
+    cudaMalloc(&scale_gpu_, sizeof(float) * scale_.size());
+    cudaMalloc(&bias_gpu_, sizeof(float) * bias_.size());
+    cudaMemcpy(scale_gpu_,
+               scale_.data(),
+               scale_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+    cudaMemcpy(bias_gpu_,
+               bias_.data(),
+               bias_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+  } else {
+    // if use fp16
+    std::vector<half> scale_half(scale_.size());
+    std::vector<half> bias_half(bias_.size());
+    for (int i = 0; i < scale_.size(); ++i) {
+      scale_half[i] = static_cast<half>(scale_[i]);
+    }
+    for (int i = 0; i < bias_.size(); ++i) {
+      bias_half[i] = static_cast<half>(bias_[i]);
+    }
+    cudaMalloc(&scale_gpu_, sizeof(half) * scale_half.size());
+    cudaMalloc(&bias_gpu_, sizeof(half) * bias_half.size());
+    cudaMemcpy(scale_gpu_,
+               scale_half.data(),
+               scale_half.size() * sizeof(half),
+               cudaMemcpyHostToDevice);
+    cudaMemcpy(bias_gpu_,
+               bias_half.data(),
+               bias_half.size() * sizeof(half),
+               cudaMemcpyHostToDevice);
+  }
+  return 0;
+}
+
+bool GroupNormPlugin::supportsFormat(
+    nvinfer1::DataType type, nvinfer1::PluginFormat format) const TRT_NOEXCEPT {
+  if (with_fp16_) {
+    return ((type == nvinfer1::DataType::kHALF) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  } else {
+    return ((type == nvinfer1::DataType::kFLOAT) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  }
+}
+
+nvinfer1::Dims GroupNormPlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims *inputDims, int nbInputs) TRT_NOEXCEPT {
+  return inputDims[0];
+}
+
+int GroupNormPlugin::enqueue(int batch_size,
+                             const void *const *inputs,
+                             void *const *outputs,
+                             void *workspace,
+                             cudaStream_t stream) TRT_NOEXCEPT {
+  const auto &input_dims = this->getInputDims(0);
+  int groups = groups_;
+  float eps = eps_;
+  std::vector<int> input_shape;
+  input_shape.push_back(batch_size);
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    input_shape.push_back(input_dims.d[i]);
+  }
+  const auto input_ddim = common::make_ddim(input_shape);
+
+  int C = input_shape[1];
+
+  PADDLE_ENFORCE_EQ(
+      C,
+      scale_.size(),
+      common::errors::InvalidArgument(
+          "scale's size should be equal to the channel number in groupnorm,"
+          "but got channel number:%d, scale's size:%d.",
+          C,
+          scale_.size()));
+  PADDLE_ENFORCE_EQ(
+      C,
+      bias_.size(),
+      common::errors::InvalidArgument(
+          "bias's size should be equal to the channel number in groupnorm,"
+          "but got channel number:%d, bias's size:%d.",
+          C,
+          bias_.size()));
+  float *mean_d = static_cast<float *>(workspace);
+  float *variance_d = mean_d + input_shape[0] * groups_;
+  float *temp_variance_d = variance_d + input_shape[0] * groups_;
+  auto input_type = getDataType();
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. GroupNorm-->fp32";
+    const float *input = static_cast<const float *>(inputs[0]);
+    float *output = static_cast<float *>(outputs[0]);
+    phi::GroupNormDirectCUDAFunctor<float> group_norm;
+    group_norm(stream,
+               input,
+               input_shape,
+               reinterpret_cast<float *>(bias_gpu_),
+               reinterpret_cast<float *>(scale_gpu_),
+               temp_variance_d,
+               groups_,
+               eps_,
+               output,
+               mean_d,
+               variance_d,
+               DataLayout::kNCHW);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. GroupNorm-->fp16";
+    const half *input = static_cast<const half *>(inputs[0]);
+    half *output = static_cast<half *>(outputs[0]);
+    phi::GroupNormDirectCUDAFunctor<half, float> group_norm;
+    group_norm(stream,
+               input,
+               input_shape,
+               reinterpret_cast<const half *>(bias_gpu_),
+               reinterpret_cast<const half *>(scale_gpu_),
+               temp_variance_d,
+               groups_,
+               eps_,
+               output,
+               mean_d,
+               variance_d,
+               DataLayout::kNCHW);
+  } else {
+    PADDLE_THROW(common::errors::Fatal(
+        "The GroupNorm TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+nvinfer1::DimsExprs GroupNormPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputDims,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  return inputDims[0];
+}
+
+bool GroupNormPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of groupnorm plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+
+  bool int8_support = in.type == nvinfer1::DataType::kINT8 &&
+                      in.format == nvinfer1::PluginFormat::kCHW32;
+  bool fp16_support =
+      (in.type == nvinfer1::DataType::kHALF) &&
+      ((!with_silu_ && in.format == nvinfer1::PluginFormat::kLINEAR) ||
+       in.format == nvinfer1::PluginFormat::kHWC8);
+
+  if (pos == 0) {
+    if (with_int8_) {
+      return int8_support || fp16_support;
+    } else if (with_fp16_) {
+      return fp16_support;
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType GroupNormPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The groupnorm Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  PADDLE_ENFORCE_EQ((input_types[0] == nvinfer1::DataType::kFLOAT ||
+                     input_types[0] == nvinfer1::DataType::kHALF),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The input type should be half or float"));
+
+  return input_types[0];
+}
+int GroupNormPluginDynamic::initialize() TRT_NOEXCEPT {
+  if (with_fp16_ == false) {
+    // if use fp32
+    cudaMalloc(&scale_gpu_, sizeof(float) * scale_.size());
+    cudaMalloc(&bias_gpu_, sizeof(float) * bias_.size());
+    cudaMemcpy(scale_gpu_,
+               scale_.data(),
+               scale_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+    cudaMemcpy(bias_gpu_,
+               bias_.data(),
+               bias_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+  } else {
+    // if use fp16
+    std::vector<half> scale_half(scale_.size());
+    std::vector<half> bias_half(bias_.size());
+    for (int i = 0; i < scale_.size(); ++i) {
+      scale_half[i] = static_cast<half>(scale_[i]);
+    }
+    for (int i = 0; i < bias_.size(); ++i) {
+      bias_half[i] = static_cast<half>(bias_[i]);
+    }
+    cudaMalloc(&scale_gpu_, sizeof(half) * scale_.size());
+    cudaMalloc(&bias_gpu_, sizeof(half) * bias_.size());
+    cudaMemcpy(scale_gpu_,
+               scale_half.data(),
+               scale_half.size() * sizeof(half),
+               cudaMemcpyHostToDevice);
+    cudaMemcpy(bias_gpu_,
+               bias_half.data(),
+               bias_half.size() * sizeof(half),
+               cudaMemcpyHostToDevice);
+  }
+  return 0;
+}
+
+int GroupNormPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  int groups = groups_;
+  float eps = eps_;
+
+  std::vector<int> input_shape;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    input_shape.push_back(input_dims.d[i]);
+  }
+
+  const auto input_ddim = common::make_ddim(input_shape);
+
+  int C = input_shape[1];
+  int64_t image_size = input_shape[2] * input_shape[3];
+  int batchSize = input_shape[0];
+
+  PADDLE_ENFORCE_EQ(
+      C,
+      scale_.size(),
+      common::errors::InvalidArgument(
+          "scale's size should be equal to the channel number in groupnorm,"
+          "but got feature_size:%d, scale's size:%d.",
+          C,
+          scale_.size()));
+  PADDLE_ENFORCE_EQ(
+      C,
+      bias_.size(),
+      common::errors::InvalidArgument(
+          "bias's size should be equal to the channel number in groupnorm,"
+          "but got feature_size:%d, bias's size:%d.",
+          C,
+          bias_.size()));
+
+  float *mean_d = static_cast<float *>(workspace);
+  float *variance_d = mean_d + input_shape[0] * groups_;
+  float *temp_variance_d = variance_d + input_shape[0] * groups_;
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. GroupNorm-->fp32";
+    const float *input = reinterpret_cast<const float *>(inputs[0]);
+    float *output = static_cast<float *>(outputs[0]);
+    phi::GroupNormDirectCUDAFunctor<float, float> group_norm;
+    group_norm(stream,
+               input,
+               input_shape,
+               reinterpret_cast<float *>(bias_gpu_),
+               reinterpret_cast<float *>(scale_gpu_),
+               temp_variance_d,
+               groups,
+               eps,
+               output,
+               mean_d,
+               variance_d,
+               DataLayout::kNCHW);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. GroupNorm-->fp16";
+    const half *input = reinterpret_cast<const half *>(inputs[0]);
+    half *output = static_cast<half *>(outputs[0]);
+    if (input_desc[0].format == nvinfer1::PluginFormat::kLINEAR) {
+      phi::GroupNormDirectCUDAFunctor<half, float> group_norm;
+      group_norm(stream,
+                 input,
+                 input_shape,
+                 reinterpret_cast<half *>(bias_gpu_),
+                 reinterpret_cast<half *>(scale_gpu_),
+                 temp_variance_d,
+                 groups,
+                 eps,
+                 output,
+                 mean_d,
+                 variance_d,
+                 DataLayout::kNCHW);
+    } else if (input_desc[0].format == nvinfer1::PluginFormat::kHWC8) {
+      int32_t cPerBlock = 320;
+      int32_t maxBlocksPerDHW = 1024;
+      switch (input_desc[0].dims.d[1]) {
+        case 960:
+        case 1920:
+          cPerBlock = 480;
+          break;
+        case 512:
+        case 256:
+          cPerBlock = 256;
+          break;
+        case 128:
+          cPerBlock = 128;
+          break;
+        default:
+          cPerBlock = 320;
+      }
+      if (cPerBlock > input_desc[0].dims.d[1]) {
+        cPerBlock = 8;
+      }
+      auto d_dim = input_desc[0].dims.nbDims;
+      params_.n = input_desc[0].dims.d[0];
+      if (d_dim == 3) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = 1;
+        params_.w = input_desc[0].dims.d[2];
+      } else if (d_dim == 4) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = input_desc[0].dims.d[2];
+        params_.w = input_desc[0].dims.d[3];
+      } else {
+        // d_dim == 5
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = input_desc[0].dims.d[2];
+        params_.h = input_desc[0].dims.d[3];
+        params_.w = input_desc[0].dims.d[4];
+      }
+
+      params_.withSilu = with_silu_;
+      params_.dst = static_cast<half *>(outputs[0]);
+      params_.srcX = static_cast<half const *>(inputs[0]);
+      params_.gamma = reinterpret_cast<half *>(scale_gpu_);
+      params_.beta = reinterpret_cast<half *>(bias_gpu_);
+      params_.redBuffer = static_cast<float *>(workspace);
+      params_.var_data = nullptr;
+      // params_.n = input_desc[0].dims.d[0];
+      // params_.h = input_desc[0].dims.d[2];
+      // params_.w = input_desc[0].dims.d[3];
+      // params_.c = input_desc[0].dims.d[1];
+      params_.groups = groups_;
+      params_.dhw = static_cast<int64_t>(params_.d) * params_.h * params_.w;
+      const int64_t blocksPerDHW =
+          findMaxDivisor(params_.dhw, static_cast<int64_t>(maxBlocksPerDHW));
+      params_.dhwPerBlock = divUp(params_.dhw, blocksPerDHW);
+      params_.cPerBlock = cPerBlock;
+      params_.cPerGroup = params_.c / params_.groups;
+      params_.dhwc = params_.dhw * params_.c;
+      params_.invDHWC =
+          1.F / static_cast<float>(params_.dhw * params_.cPerGroup);
+      params_.groupsPerBlock = cPerBlock / params_.cPerGroup;
+      params_.eps = eps_;
+      params_.var_data = nullptr;
+
+      cudaMemsetAsync(params_.redBuffer,
+                      0,
+                      2 * sizeof(float) * params_.n * groups_,
+                      stream);
+
+      phi::groupNormNDHWCSum<half> ndhwc_sum;
+      ndhwc_sum(&params_, stream);
+      phi::groupNormNDHWCScale<half> ndhwc_scale;
+      ndhwc_scale(params_, stream);
+    } else {
+      PADDLE_THROW(common::errors::Fatal(
+          "The Groupnorm TRT Plugin's only support nchw or nhwc8 input"));
+    }
+  } else if (input_type == nvinfer1::DataType::kINT8) {
+    const int8_t *input = reinterpret_cast<const int8_t *>(inputs[0]);
+    int8_t *output = static_cast<int8_t *>(outputs[0]);
+
+    if (input_desc[0].format == nvinfer1::PluginFormat::kCHW32) {
+      int32_t cPerBlock = 320;
+      int32_t maxBlocksPerDHW = 1024;
+      switch (input_desc[0].dims.d[1]) {
+        case 960:
+        case 1920:
+          cPerBlock = 480;
+          break;
+        case 512:
+        case 256:
+          cPerBlock = 256;
+          break;
+        case 128:
+          cPerBlock = 128;
+          break;
+        default:
+          cPerBlock = 320;
+      }
+      if (cPerBlock > input_desc[0].dims.d[1]) {
+        cPerBlock = 8;
+      }
+      auto d_dim = input_desc[0].dims.nbDims;
+      params_.n = input_desc[0].dims.d[0];
+      if (d_dim == 3) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = 1;
+        params_.w = input_desc[0].dims.d[2];
+      } else if (d_dim == 4) {
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = 1;
+        params_.h = input_desc[0].dims.d[2];
+        params_.w = input_desc[0].dims.d[3];
+      } else {
+        // d_dim == 5
+        params_.c = input_desc[0].dims.d[1];
+        params_.d = input_desc[0].dims.d[2];
+        params_.h = input_desc[0].dims.d[3];
+        params_.w = input_desc[0].dims.d[4];
+      }
+      params_.withSilu = with_silu_;
+      params_.dst = static_cast<half *>(outputs[0]);
+      params_.srcX = static_cast<half const *>(inputs[0]);
+
+      params_.gamma = scale_gpu_;
+      params_.beta = bias_gpu_;
+      params_.redBuffer = static_cast<float *>(workspace);
+      // params_.n = input_desc[0].dims.d[0];
+      // params_.h = input_desc[0].dims.d[2];
+      // params_.w = input_desc[0].dims.d[3];
+      // params_.c = input_desc[0].dims.d[1];
+      params_.groups = groups_;
+      params_.dhw = static_cast<int64_t>(params_.d) * params_.h * params_.w;
+      const int64_t blocksPerDHW =
+          findMaxDivisor(params_.dhw, static_cast<int64_t>(maxBlocksPerDHW));
+      params_.dhwPerBlock = divUp(params_.dhw, blocksPerDHW);
+      params_.cPerBlock = cPerBlock;
+      params_.cPerGroup = params_.c / params_.groups;
+      params_.dhwc = params_.dhw * params_.c;
+      params_.invDHWC =
+          1.F / static_cast<float>(params_.dhw * params_.cPerGroup);
+      params_.groupsPerBlock = cPerBlock / params_.cPerGroup;
+      PADDLE_ENFORCE_EQ(
+          cPerBlock % params_.cPerGroup,
+          0,
+          common::errors::InvalidArgument(
+              "cPerBlock should be multiple of params_.cPerGroup, "
+              "now cPerBlock is %d, params_.cPerGroup is %d",
+              cPerBlock,
+              params_.cPerGroup));
+      PADDLE_ENFORCE_EQ(
+          params_.cPerGroup % 2,
+          0,
+          common::errors::InvalidArgument(
+              "params_.cPerGroup should be a even number, but received %d",
+              params_.cPerGroup));
+      params_.eps = eps_;
+      params_.dqScaleIn = input_desc[0].scale;
+      params_.inv_qScale = 1.f / output_desc[0].scale;
+
+      // Just used for TensorRTDynamicShapeGNTes in test_dynamic_engine.cc
+      // Do not Edit it
+      // params_.dqScaleIn = 1.f;
+      // params_.inv_qScale = 1 / 0.05f;
+
+      cudaMemsetAsync(params_.redBuffer,
+                      0,
+                      2 * sizeof(float) * params_.n * groups_,
+                      stream);
+      groupNormNCHW32SumQDQ(params_, stream);
+      groupNormNCHW32ScaleQDQ(params_, stream);
+    } else {
+      PADDLE_THROW(common::errors::Fatal(
+          "The Groupnorm TRT Plugin only support nchw32 input"));
+    }
+  } else {
+    // input not float
+    PADDLE_THROW(common::errors::Fatal(
+        "The Groupnorm TRT Plugin's only support fp32, fp16 or int8 input"));
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/group_norm_op_plugin.h

@@ -0,0 +1,308 @@
+/* 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. */
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/phi/kernels/group_norm_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+using phi::GroupNormNDHWCParams;
+class GroupNormPlugin : public PluginTensorRT {
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(scale_) +
+           SerializedSize(bias_) + SerializedSize(eps_) +
+           SerializedSize(groups_) + SerializedSize(mean_shape_) +
+           SerializedSize(variance_shape_) + SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, groups_);
+    SerializeValue(&buffer, mean_shape_);
+    SerializeValue(&buffer, variance_shape_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  GroupNormPlugin(const float* scale,
+                  const int scale_num,
+                  const float* bias,
+                  const int bias_num,
+                  float eps,
+                  int groups,
+                  std::vector<int64_t> mean_shape,
+                  std::vector<int64_t> variance_shape,
+                  bool with_fp16)
+      : groups_(groups),
+        eps_(eps),
+        mean_shape_(mean_shape),
+        variance_shape_(variance_shape),
+        with_fp16_(with_fp16) {
+    scale_.resize(scale_num);
+    bias_.resize(bias_num);
+    std::copy(scale, scale + scale_num, scale_.data());
+    std::copy(bias, bias + bias_num, bias_.data());
+  }
+  GroupNormPlugin(void const* serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &groups_);
+    DeserializeValue(&serialData, &serialLength, &mean_shape_);
+    DeserializeValue(&serialData, &serialLength, &variance_shape_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  ~GroupNormPlugin() {}
+  int initialize() TRT_NOEXCEPT override;
+  GroupNormPlugin* clone() const TRT_NOEXCEPT override {
+    auto* ptr = new GroupNormPlugin(scale_.data(),
+                                    scale_.size(),
+                                    bias_.data(),
+                                    bias_.size(),
+                                    eps_,
+                                    groups_,
+                                    mean_shape_,
+                                    variance_shape_,
+                                    with_fp16_);
+    ptr->scale_gpu_ = scale_gpu_;
+    ptr->bias_gpu_ = bias_gpu_;
+    return ptr;
+  }
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "groupnorm_plugin";
+  }
+  size_t getWorkspaceSize(int max_batch_size) const TRT_NOEXCEPT {
+    return 3 * max_batch_size * groups_;
+  }
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override;
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+
+  int enqueue(int batchSize,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override {
+    if (bias_gpu_) {
+      cudaFree(bias_gpu_);
+      bias_gpu_ = nullptr;
+    }
+    if (scale_gpu_) {
+      cudaFree(scale_gpu_);
+      scale_gpu_ = nullptr;
+    }
+  };
+
+ private:
+  std::vector<float> scale_;
+  std::vector<float> bias_;
+  void* scale_gpu_;
+  void* bias_gpu_;
+  int groups_;
+  float eps_;
+  std::vector<int64_t> mean_shape_;
+  std::vector<int64_t> variance_shape_;
+  bool with_fp16_;
+};
+class GroupNormPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "groupnorm_plugin";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new GroupNormPlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(GroupNormPluginCreator);
+
+class GroupNormPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  GroupNormPluginDynamic(const float* scale,
+                         const int scale_num,
+                         const float* bias,
+                         const int bias_num,
+                         float eps,
+                         int groups,
+                         std::vector<int64_t> mean_shape,
+                         std::vector<int64_t> variance_shape,
+                         bool with_silu,
+                         bool with_fp16,
+                         bool with_int8)
+      : groups_(groups),
+        eps_(eps),
+        mean_shape_(mean_shape),
+        variance_shape_(variance_shape),
+        with_silu_(with_silu),
+        with_fp16_(with_fp16),
+        with_int8_(with_int8) {
+    scale_.resize(scale_num);
+    bias_.resize(bias_num);
+    std::copy(scale, scale + scale_num, scale_.data());
+    std::copy(bias, bias + bias_num, bias_.data());
+  }
+
+  GroupNormPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &groups_);
+    DeserializeValue(&serialData, &serialLength, &mean_shape_);
+    DeserializeValue(&serialData, &serialLength, &variance_shape_);
+    DeserializeValue(&serialData, &serialLength, &with_silu_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+    DeserializeValue(&serialData, &serialLength, &with_int8_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    auto* ptr = new GroupNormPluginDynamic(scale_.data(),
+                                           scale_.size(),
+                                           bias_.data(),
+                                           bias_.size(),
+                                           eps_,
+                                           groups_,
+                                           mean_shape_,
+                                           variance_shape_,
+                                           with_silu_,
+                                           with_fp16_,
+                                           with_int8_);
+    ptr->scale_gpu_ = scale_gpu_;
+    ptr->bias_gpu_ = bias_gpu_;
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "groupnorm_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(scale_) + SerializedSize(bias_) +
+           SerializedSize(eps_) + SerializedSize(groups_) +
+           SerializedSize(mean_shape_) + SerializedSize(variance_shape_) +
+           SerializedSize(with_silu_) + SerializedSize(with_fp16_) +
+           +SerializedSize(with_int8_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, groups_);
+    SerializeValue(&buffer, mean_shape_);
+    SerializeValue(&buffer, variance_shape_);
+    SerializeValue(&buffer, with_silu_);
+    SerializeValue(&buffer, with_fp16_);
+    SerializeValue(&buffer, with_int8_);
+  }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 3 * inputs[0].dims.d[0] * groups_ * sizeof(float);
+  }
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+  void terminate() TRT_NOEXCEPT override {
+    if (bias_gpu_) {
+      cudaFree(bias_gpu_);
+      bias_gpu_ = nullptr;
+    }
+    if (scale_gpu_) {
+      cudaFree(scale_gpu_);
+      scale_gpu_ = nullptr;
+    }
+  };
+
+ private:
+  std::vector<float> scale_;
+  std::vector<float> bias_;
+  void* scale_gpu_ = nullptr;
+  void* bias_gpu_ = nullptr;
+  int groups_;
+  float eps_;
+  std::vector<int64_t> mean_shape_;
+  std::vector<int64_t> variance_shape_;
+  GroupNormNDHWCParams<half> params_;
+  bool with_silu_;
+  bool with_fp16_;
+  bool with_int8_;
+};
+class GroupNormPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "groupnorm_plugin_dynamic";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new GroupNormPluginDynamic(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(GroupNormPluginDynamicCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/hard_swish_op_plugin.cu

@@ -0,0 +1,161 @@
+// Copyright (c) 2020 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 <cassert>
+#include <cstring>
+
+#include "paddle/fluid/inference/tensorrt/plugin/hard_swish_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+nvinfer1::Dims HardSwishPlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims *in_dims, int nb_inputs) TRT_NOEXCEPT {
+  assert(nb_inputs == 1);
+  assert(index < this->getNbOutputs());
+  nvinfer1::Dims const &input_dims = in_dims[0];
+  nvinfer1::Dims output_dims = input_dims;
+  return output_dims;
+}
+
+template <typename T>
+__device__ T kMax(T a, T b) {
+  return a > b ? a : b;
+}
+
+template <typename T>
+__device__ T kMin(T a, T b) {
+  return a < b ? a : b;
+}
+
+template <typename T, unsigned TPB>
+__global__ void hard_swish_kernel(float threshold,
+                                  float scale,
+                                  float offset,
+                                  int n,
+                                  const T *input,
+                                  T *output) {
+  const int idx = blockIdx.x * TPB + threadIdx.x;
+  if (idx < n) {
+    const T in = input[idx];
+    output[idx] = in / scale * kMin<T>(kMax<T>(in + offset, 0), threshold);
+  }
+}
+
+int HardSwishPlugin::enqueue(int batch_size,
+                             const void *const *inputs,
+                             void *const *outputs,
+                             void *,
+                             cudaStream_t stream) TRT_NOEXCEPT {
+  const auto &input_dims = this->getInputDims(0);
+  int num = batch_size;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    num *= input_dims.d[i];
+  }
+  float threshold = threshold_;
+  float scale = scale_;
+  float offset = offset_;
+
+  const int block_size = 256;
+  const int grid_size = (num + block_size - 1) / block_size;
+
+  const float *input = static_cast<const float *>(inputs[0]);
+  float *output = static_cast<float *>(outputs[0]);
+  hard_swish_kernel<float, block_size><<<grid_size, block_size, 0, stream>>>(
+      threshold, scale, offset, num, input, output);
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+nvinfer1::DimsExprs HardSwishPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  return inputs[0];
+}
+
+int HardSwishPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  int num = 1;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    num *= input_dims.d[i];
+  }
+  float threshold = threshold_;
+  float scale = scale_;
+  float offset = offset_;
+  const int block_size = 256;
+  const int grid_size = (num + block_size - 1) / block_size;
+  const float *input = static_cast<const float *>(inputs[0]);
+  float *output = static_cast<float *>(outputs[0]);
+  hard_swish_kernel<float, block_size><<<grid_size, block_size, 0, stream>>>(
+      threshold, scale, offset, num, input, output);
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+nvinfer1::DataType HardSwishPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Elementwise Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+bool HardSwishPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    return (in.type == nvinfer1::DataType::kFLOAT) &&
+           (in.format == nvinfer1::TensorFormat::kLINEAR);
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/hard_swish_op_plugin.h

@@ -0,0 +1,215 @@
+// Copyright (c) 2019 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.
+
+#pragma once
+#include <stdio.h>
+
+#include <cassert>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class HardSwishPlugin : public PluginTensorRT {
+ public:
+  HardSwishPlugin(const float threshold, const float scale, const float offset)
+      : threshold_(threshold), scale_(scale), offset_(offset) {}
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  HardSwishPlugin(void const* serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &threshold_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &offset_);
+  }
+
+  ~HardSwishPlugin() {}
+  HardSwishPlugin* clone() const TRT_NOEXCEPT override {
+    return new HardSwishPlugin(threshold_, scale_, offset_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "hard_swish_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+  int enqueue(int batchSize,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(threshold_) +
+           SerializedSize(scale_) + SerializedSize(offset_);
+  }
+
+  // TRT will call this func  to serialize the configuration of TRT
+  // It should not be called by users.
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, threshold_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, offset_);
+  }
+
+ protected:
+  float threshold_;
+  float scale_;
+  float offset_;
+};
+
+class HardSwishPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "hard_swish_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new HardSwishPlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(HardSwishPluginCreator);
+
+class HardSwishPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  HardSwishPluginDynamic(const float threshold,
+                         const float scale,
+                         const float offset)
+      : threshold_(threshold), scale_(scale), offset_(offset) {}
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  HardSwishPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &threshold_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &offset_);
+  }
+  ~HardSwishPluginDynamic() {}
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new HardSwishPluginDynamic(threshold_, scale_, offset_);
+  }
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "hard_swish_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(threshold_) + SerializedSize(scale_) +
+           SerializedSize(offset_);
+  }
+
+  // TRT will call this func  to serialize the configuration of TRT
+  // It should not be called by users.
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, threshold_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, offset_);
+  }
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ protected:
+  float threshold_;
+  float scale_;
+  float offset_;
+};
+
+class HardSwishPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  HardSwishPluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "hard_swish_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new HardSwishPluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+REGISTER_TRT_PLUGIN_V2(HardSwishPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/instance_norm_op_plugin.cu

@@ -0,0 +1,347 @@
+// Copyright (c) 2018 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 <stdio.h>
+
+#include <cassert>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/instance_norm_op_plugin.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_dnn.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+cudnnStatus_t convert_trt2cudnn_dtype(nvinfer1::DataType trt_dtype,
+                                      cudnnDataType_t *cudnn_dtype) {
+  switch (trt_dtype) {
+    case nvinfer1::DataType::kFLOAT:
+      *cudnn_dtype = CUDNN_DATA_FLOAT;
+      break;
+    case nvinfer1::DataType::kHALF:
+      *cudnn_dtype = CUDNN_DATA_HALF;
+      break;
+    default:
+      return CUDNN_STATUS_BAD_PARAM;
+  }
+  return CUDNN_STATUS_SUCCESS;
+}
+
+int InstanceNormPlugin::initialize() TRT_NOEXCEPT { return 0; }
+
+nvinfer1::Dims InstanceNormPlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims *inputDims, int nbInputs) TRT_NOEXCEPT {
+  assert(nbInputs == 1);
+  assert(index < this->getNbOutputs());
+  nvinfer1::Dims const &input_dims = inputDims[0];
+  nvinfer1::Dims output_dims = input_dims;
+  return output_dims;
+}
+
+bool InstanceNormPlugin::supportsFormat(
+    nvinfer1::DataType type, nvinfer1::PluginFormat format) const TRT_NOEXCEPT {
+  return ((type == nvinfer1::DataType::kFLOAT ||
+           type == nvinfer1::DataType::kHALF) &&
+          (format == nvinfer1::PluginFormat::kLINEAR));
+}
+
+int InstanceNormPlugin::enqueue(int batch_size,
+                                const void *const *inputs,
+                                void *const *outputs,
+                                void *workspace,
+                                cudaStream_t stream) TRT_NOEXCEPT {
+  const auto &input_dims = this->getInputDims(0);
+  int n = batch_size;
+  int c = input_dims.d[0];
+  int h = input_dims.d[1];
+  int w = input_dims.d[2];
+
+  scale_t.Resize(common::make_ddim({batch_size, c}));
+  bias_t.Resize(common::make_ddim({batch_size, c}));
+  int device_id;
+  cudaGetDevice(&device_id);
+  float *scale_d = scale_t.mutable_data<float>(GPUPlace(device_id));
+  float *bias_d = bias_t.mutable_data<float>(GPUPlace(device_id));
+
+  for (int i = 0; i < batch_size; i++) {
+    cudaMemcpyAsync(scale_d + i * c,
+                    scale_.data(),
+                    sizeof(float) * c,
+                    cudaMemcpyHostToDevice,
+                    stream);
+    cudaMemcpyAsync(bias_d + i * c,
+                    bias_.data(),
+                    sizeof(float) * c,
+                    cudaMemcpyHostToDevice,
+                    stream);
+  }
+  phi::dynload::cudnnSetTensor4dDescriptor(
+      b_desc_, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, n * c, 1, 1);
+
+  cudnnDataType_t cudnn_dtype;
+  nvinfer1::DataType data_type = getDataType();
+  convert_trt2cudnn_dtype(data_type, &cudnn_dtype);
+  phi::dynload::cudnnSetTensor4dDescriptor(
+      x_desc_, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w);
+  phi::dynload::cudnnSetTensor4dDescriptor(
+      y_desc_, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w);
+  float alpha = 1;
+  float beta = 0;
+  phi::dynload::cudnnSetStream(handle_, stream);
+
+  void const *x_ptr = inputs[0];
+  void *y_ptr = outputs[0];
+  phi::dynload::cudnnBatchNormalizationForwardTraining(
+      handle_,
+      CUDNN_BATCHNORM_SPATIAL_PERSISTENT,
+      &alpha,
+      &beta,
+      x_desc_,
+      x_ptr,
+      y_desc_,
+      y_ptr,
+      b_desc_,
+      scale_d,
+      bias_d,
+      1.,
+      nullptr,
+      nullptr,
+      eps_,
+      nullptr,
+      nullptr);
+  return cudaGetLastError() != cudaSuccess;
+}
+
+int InstanceNormPluginEnqueue(const nvinfer1::PluginTensorDesc *inputDesc,
+                              const nvinfer1::PluginTensorDesc *outputDesc,
+                              const void *const *inputs,
+                              void *const *outputs,
+                              void *workspace,
+                              const float *scale,
+                              const float *bias,
+                              float eps,
+                              cudnnTensorDescriptor_t x_desc_,
+                              cudnnTensorDescriptor_t y_desc_,
+                              cudnnTensorDescriptor_t b_desc_,
+                              cudnnHandle_t handle_,
+                              cudaStream_t stream) {
+  nvinfer1::Dims input_dims = inputDesc[0].dims;
+  int n = input_dims.d[0];
+  int c = input_dims.d[1];
+  int h = input_dims.d[2];
+  int w = input_dims.d[3];
+  phi::DenseTensor scale_t;
+  phi::DenseTensor bias_t;
+  scale_t.Resize(common::make_ddim({n, c}));
+  bias_t.Resize(common::make_ddim({n, c}));
+  int device_id;
+  cudaGetDevice(&device_id);
+  float *scale_d = scale_t.mutable_data<float>(GPUPlace(device_id));
+  float *bias_d = bias_t.mutable_data<float>(GPUPlace(device_id));
+
+  for (int i = 0; i < n; i++) {
+    cudaMemcpyAsync(scale_d + i * c,
+                    scale,
+                    sizeof(float) * c,
+                    cudaMemcpyHostToDevice,
+                    stream);
+    cudaMemcpyAsync(bias_d + i * c,
+                    bias,
+                    sizeof(float) * c,
+                    cudaMemcpyHostToDevice,
+                    stream);
+  }
+  phi::dynload::cudnnSetTensor4dDescriptor(
+      b_desc_, CUDNN_TENSOR_NCHW, CUDNN_DATA_FLOAT, 1, n * c, 1, 1);
+
+  cudnnDataType_t cudnn_dtype;
+  auto data_type = inputDesc[0].type;
+  convert_trt2cudnn_dtype(data_type, &cudnn_dtype);
+  phi::dynload::cudnnSetTensor4dDescriptor(
+      x_desc_, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w);
+  phi::dynload::cudnnSetTensor4dDescriptor(
+      y_desc_, CUDNN_TENSOR_NCHW, cudnn_dtype, 1, n * c, h, w);
+  float alpha = 1;
+  float beta = 0;
+  phi::dynload::cudnnSetStream(handle_, stream);
+
+  void const *x_ptr = inputs[0];
+  void *y_ptr = outputs[0];
+  phi::dynload::cudnnBatchNormalizationForwardTraining(
+      handle_,
+      CUDNN_BATCHNORM_SPATIAL_PERSISTENT,
+      &alpha,
+      &beta,
+      x_desc_,
+      x_ptr,
+      y_desc_,
+      y_ptr,
+      b_desc_,
+      scale_d,
+      bias_d,
+      1.,
+      nullptr,
+      nullptr,
+      eps,
+      nullptr,
+      nullptr);
+  return cudaGetLastError() != cudaSuccess;
+}
+
+int InstanceNormPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+nvinfer1::DimsExprs InstanceNormPluginDynamic::getOutputDimensions(
+    int index,
+    const nvinfer1::DimsExprs *inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  assert(nbInputs == 1);
+  assert(index < this->getNbOutputs());
+  nvinfer1::DimsExprs output(inputs[0]);
+  return output;
+}
+
+bool InstanceNormPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  assert(inOut && pos < (nbInputs + nbOutputs));
+  assert(pos == 0 || pos == 1);
+  return ((inOut[pos].type == nvinfer1::DataType::kFLOAT ||
+           inOut[pos].type == nvinfer1::DataType::kHALF) &&
+          (inOut[pos].format == nvinfer1::PluginFormat::kLINEAR) &&
+          inOut[pos].type == inOut[0].type);
+}
+
+int InstanceNormPluginDynamic::enqueue(
+    const nvinfer1::PluginTensorDesc *inputDesc,
+    const nvinfer1::PluginTensorDesc *outputDesc,
+    const void *const *inputs,
+    void *const *outputs,
+    void *workspace,
+    cudaStream_t stream) TRT_NOEXCEPT {
+  return InstanceNormPluginEnqueue(inputDesc,
+                                   outputDesc,
+                                   inputs,
+                                   outputs,
+                                   workspace,
+                                   scale_.data(),
+                                   bias_.data(),
+                                   eps_,
+                                   x_desc_,
+                                   y_desc_,
+                                   b_desc_,
+                                   handle_,
+                                   stream);
+}
+
+nvinfer1::DataType InstanceNormPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *inputTypes,
+    int nbInputs) const TRT_NOEXCEPT {
+  assert(inputTypes && nbInputs > 0 && index == 0);
+  return inputTypes[0];
+}
+
+void InstanceNormPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+int PIRInstanceNormPlugin::initialize() TRT_NOEXCEPT { return 0; }
+
+nvinfer1::DimsExprs PIRInstanceNormPlugin::getOutputDimensions(
+    int index,
+    const nvinfer1::DimsExprs *inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  assert(nbInputs == 1);
+  assert(index < this->getNbOutputs());
+  nvinfer1::DimsExprs output(inputs[0]);
+  return output;
+}
+
+bool PIRInstanceNormPlugin::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  assert(inOut && pos < (nbInputs + nbOutputs));
+  assert(pos == 0 || pos == 1);
+  return ((inOut[pos].type == nvinfer1::DataType::kFLOAT ||
+           inOut[pos].type == nvinfer1::DataType::kHALF) &&
+          (inOut[pos].format == nvinfer1::PluginFormat::kLINEAR) &&
+          inOut[pos].type == inOut[0].type);
+}
+
+int PIRInstanceNormPlugin::enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
+                                   const nvinfer1::PluginTensorDesc *outputDesc,
+                                   const void *const *inputs,
+                                   void *const *outputs,
+                                   void *workspace,
+                                   cudaStream_t stream) TRT_NOEXCEPT {
+  const float *scale_ = reinterpret_cast<const float *>(inputs[1]);
+  const float *bias_ = reinterpret_cast<const float *>(inputs[2]);
+  return InstanceNormPluginEnqueue(inputDesc,
+                                   outputDesc,
+                                   inputs,
+                                   outputs,
+                                   workspace,
+                                   scale_,
+                                   bias_,
+                                   eps_,
+                                   x_desc_,
+                                   y_desc_,
+                                   b_desc_,
+                                   handle_,
+                                   stream);
+}
+
+nvinfer1::DataType PIRInstanceNormPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *inputTypes,
+    int nbInputs) const TRT_NOEXCEPT {
+  assert(inputTypes && nbInputs > 0 && index == 0);
+  return inputTypes[0];
+}
+
+void PIRInstanceNormPlugin::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+nvinfer1::IPluginV2 *PIRInstanceNormPluginCreator::createPlugin(
+    const char *name, const nvinfer1::PluginFieldCollection *fc) TRT_NOEXCEPT {
+  float epsilon = 1e-8;
+  for (int i = 0; i < fc->nbFields; ++i) {
+    const std::string field_name(fc->fields[i].name);
+    if (field_name.compare("epsilon") == 0) {
+      epsilon = *static_cast<const float *>(fc->fields[i].data);
+    } else {
+      assert(false && "unknown plugin field name.");
+    }
+  }
+  return new PIRInstanceNormPlugin(epsilon);
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/instance_norm_op_plugin.h

@@ -0,0 +1,380 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class InstanceNormPlugin : public PluginTensorRT {
+ private:
+  float eps_;
+  std::vector<float> scale_;
+  std::vector<float> bias_;
+
+  phi::DenseTensor scale_t;
+  phi::DenseTensor bias_t;
+  cudnnHandle_t handle_;
+  cudnnTensorDescriptor_t x_desc_, y_desc_, b_desc_;
+
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(eps_) +
+           SerializedSize(scale_) + SerializedSize(bias_);
+  }
+
+  // TRT will call this func when we need to serialize the configuration of
+  // tensorrt.
+  // It should not be called by users.
+  void serialize(void *buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, bias_);
+  }
+
+  explicit InstanceNormPlugin(const float eps,
+                              const std::vector<float> scale,
+                              const std::vector<float> bias)
+      : eps_(eps), scale_(scale), bias_(bias) {
+    PADDLE_ENFORCE_EQ(scale.size(),
+                      bias.size(),
+                      common::errors::InvalidArgument(
+                          "The instanceNorm's scale and bias should be the "
+                          "same size. Got scale size = %d, but bias size = %d",
+                          scale.size(),
+                          bias.size()));
+    phi::dynload::cudnnCreate(&handle_);
+    phi::dynload::cudnnCreateTensorDescriptor(&x_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&y_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&b_desc_);
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  InstanceNormPlugin(void const *serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &bias_);
+
+    phi::dynload::cudnnCreate(&handle_);
+    phi::dynload::cudnnCreateTensorDescriptor(&x_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&y_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&b_desc_);
+  }
+
+  ~InstanceNormPlugin() {
+    phi::dynload::cudnnDestroy(handle_);
+    phi::dynload::cudnnDestroyTensorDescriptor(x_desc_);
+    phi::dynload::cudnnDestroyTensorDescriptor(y_desc_);
+    phi::dynload::cudnnDestroyTensorDescriptor(b_desc_);
+  }
+
+  int initialize() TRT_NOEXCEPT override;
+
+  InstanceNormPlugin *clone() const TRT_NOEXCEPT override {
+    return new InstanceNormPlugin(eps_, scale_, bias_);
+  }
+
+  const char *getPluginType() const TRT_NOEXCEPT override {
+    return "instance_norm";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims *inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+
+  int enqueue(int batchSize,
+              const void *const *inputs,
+              void *const *outputs,
+              void *workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override;
+};
+
+class InstanceNormPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char *getPluginName() const TRT_NOEXCEPT override {
+    return "instance_norm";
+  }
+
+  const char *getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2 *deserializePlugin(const char *name,
+                                         const void *serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new InstanceNormPlugin(serial_data, serial_length);
+  }
+};
+
+class InstanceNormPluginDynamic : public DynamicPluginTensorRT {
+ private:
+  float eps_;
+  std::vector<float> scale_;
+  std::vector<float> bias_;
+
+  cudnnHandle_t handle_;
+  cudnnTensorDescriptor_t x_desc_, y_desc_, b_desc_;
+
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(eps_) + SerializedSize(scale_) +
+           SerializedSize(bias_);
+  }
+
+  // TRT will call this func when we need to serialize the configuration of
+  // tensorrt.
+  // It should not be called by users.
+  void serialize(void *buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, bias_);
+  }
+
+  explicit InstanceNormPluginDynamic(const float eps,
+                                     const std::vector<float> scale,
+                                     const std::vector<float> bias)
+      : eps_(eps), scale_(scale), bias_(bias) {
+    PADDLE_ENFORCE_EQ(scale.size(),
+                      bias.size(),
+                      common::errors::InvalidArgument(
+                          "The instanceNorm's scale and bias should be the "
+                          "same size. Got scale size = %d, but bias size = %d",
+                          scale.size(),
+                          bias.size()));
+    phi::dynload::cudnnCreate(&handle_);
+    phi::dynload::cudnnCreateTensorDescriptor(&x_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&y_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&b_desc_);
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  InstanceNormPluginDynamic(void const *serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &bias_);
+
+    phi::dynload::cudnnCreate(&handle_);
+    phi::dynload::cudnnCreateTensorDescriptor(&x_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&y_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&b_desc_);
+  }
+
+  ~InstanceNormPluginDynamic() {
+    phi::dynload::cudnnDestroy(handle_);
+    phi::dynload::cudnnDestroyTensorDescriptor(x_desc_);
+    phi::dynload::cudnnDestroyTensorDescriptor(y_desc_);
+    phi::dynload::cudnnDestroyTensorDescriptor(b_desc_);
+  }
+
+  int initialize() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2DynamicExt *clone() const TRT_NOEXCEPT override {
+    return new InstanceNormPluginDynamic(eps_, scale_, bias_);
+  }
+
+  const char *getPluginType() const TRT_NOEXCEPT override {
+    return "instance_norm_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs *inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder &expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc *inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc *out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc *outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
+              const nvinfer1::PluginTensorDesc *outputDesc,
+              const void *const *inputs,
+              void *const *outputs,
+              void *workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType *inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+};
+
+class InstanceNormPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char *getPluginName() const TRT_NOEXCEPT override {
+    return "instance_norm_dynamic";
+  }
+
+  const char *getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2 *deserializePlugin(const char *name,
+                                         const void *serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new InstanceNormPluginDynamic(serial_data, serial_length);
+  }
+};
+
+class PIRInstanceNormPlugin : public DynamicPluginTensorRT {
+ private:
+  float eps_;
+  cudnnHandle_t handle_;
+  cudnnTensorDescriptor_t x_desc_, y_desc_, b_desc_;
+
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(eps_);
+  }
+
+  // TRT will call this func when we need to serialize the configuration of
+  // tensorrt.
+  // It should not be called by users.
+  void serialize(void *buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, eps_);
+  }
+
+  explicit PIRInstanceNormPlugin(const float eps) : eps_(eps) {
+    phi::dynload::cudnnCreate(&handle_);
+    phi::dynload::cudnnCreateTensorDescriptor(&x_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&y_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&b_desc_);
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  PIRInstanceNormPlugin(void const *serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &eps_);
+
+    phi::dynload::cudnnCreate(&handle_);
+    phi::dynload::cudnnCreateTensorDescriptor(&x_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&y_desc_);
+    phi::dynload::cudnnCreateTensorDescriptor(&b_desc_);
+  }
+
+  ~PIRInstanceNormPlugin() {
+    phi::dynload::cudnnDestroy(handle_);
+    phi::dynload::cudnnDestroyTensorDescriptor(x_desc_);
+    phi::dynload::cudnnDestroyTensorDescriptor(y_desc_);
+    phi::dynload::cudnnDestroyTensorDescriptor(b_desc_);
+  }
+
+  int initialize() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2DynamicExt *clone() const TRT_NOEXCEPT override {
+    return new PIRInstanceNormPlugin(eps_);
+  }
+
+  const char *getPluginType() const TRT_NOEXCEPT override {
+    return "pir_instance_norm";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs *inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder &expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc *inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc *in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc *out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc *inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc *outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc *inputDesc,
+              const nvinfer1::PluginTensorDesc *outputDesc,
+              const void *const *inputs,
+              void *const *outputs,
+              void *workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType *inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+};
+
+class PIRInstanceNormPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char *getPluginName() const TRT_NOEXCEPT override {
+    return "pir_instance_norm";
+  }
+
+  const char *getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2 *createPlugin(const char *name,
+                                    const nvinfer1::PluginFieldCollection *fc)
+      TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2 *deserializePlugin(const char *name,
+                                         const void *serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PIRInstanceNormPlugin(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(InstanceNormPluginCreator);
+REGISTER_TRT_PLUGIN_V2(InstanceNormPluginDynamicCreator);
+REGISTER_TRT_PLUGIN_V2(PIRInstanceNormPluginCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/layer_norm_op_plugin.cu

@@ -0,0 +1,371 @@
+// Copyright (c) 2018 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 <stdio.h>
+
+#include <cassert>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/layer_norm_op_plugin.h"
+#include "paddle/phi/kernels/layer_norm_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+int LayerNormPlugin::initialize() TRT_NOEXCEPT {
+  cudaMalloc(&bias_gpu_, sizeof(float) * bias_.size());
+  cudaMemcpy(bias_gpu_,
+             bias_.data(),
+             bias_.size() * sizeof(float),
+             cudaMemcpyHostToDevice);
+  cudaMalloc(&scale_gpu_, sizeof(float) * scale_.size());
+  cudaMemcpy(scale_gpu_,
+             scale_.data(),
+             scale_.size() * sizeof(float),
+             cudaMemcpyHostToDevice);
+  return 0;
+}
+
+void LayerNormPlugin::terminate() TRT_NOEXCEPT {
+  if (bias_gpu_) {
+    cudaFree(bias_gpu_);
+    bias_gpu_ = nullptr;
+  }
+  if (scale_gpu_) {
+    cudaFree(scale_gpu_);
+    scale_gpu_ = nullptr;
+  }
+}
+
+nvinfer1::Dims LayerNormPlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims *inputDims, int nbInputs) TRT_NOEXCEPT {
+  assert(nbInputs == 1);
+  assert(index < this->getNbOutputs());
+  nvinfer1::Dims const &input_dims = inputDims[0];
+  nvinfer1::Dims output_dims = input_dims;
+  return output_dims;
+}
+
+bool LayerNormPlugin::supportsFormat(
+    nvinfer1::DataType type, nvinfer1::PluginFormat format) const TRT_NOEXCEPT {
+  if (with_fp16_) {
+    return ((type == nvinfer1::DataType::kFLOAT ||
+             type == nvinfer1::DataType::kHALF) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  } else {
+    return ((type == nvinfer1::DataType::kFLOAT) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  }
+}
+
+int LayerNormPlugin::enqueue(int batch_size,
+                             const void *const *inputs,
+                             void *const *outputs,
+                             void *workspace,
+                             cudaStream_t stream) TRT_NOEXCEPT {
+  const auto &input_dims = this->getInputDims(0);
+  int begin_norm_axis = begin_norm_axis_;
+  float eps = eps_;
+
+  PADDLE_ENFORCE_EQ(1,
+                    mean_shape_.size(),
+                    common::errors::InvalidArgument(
+                        "Size of mean_shape vector should be equal to 1,"
+                        "but got Size of mean_shape vector:%d",
+                        mean_shape_.size()));
+  PADDLE_ENFORCE_EQ(1,
+                    variance_shape_.size(),
+                    common::errors::InvalidArgument(
+                        "Size of variance_shape vector should be equal to 1,"
+                        "but got Size of mean_shape vector:%d",
+                        mean_shape_.size()));
+
+  int64_t batched_mean_shape = mean_shape_[0] * input_dims.d[0];
+  int64_t batched_variance_shape = variance_shape_[0] * input_dims.d[0];
+
+  std::vector<int64_t> input_shape;
+  input_shape.push_back(batch_size);
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    input_shape.push_back(input_dims.d[i]);
+  }
+  const auto input_ddim = common::make_ddim(input_shape);
+  auto matrix_dim = common::flatten_to_2d(input_ddim, begin_norm_axis);
+  int feature_size = static_cast<int>(matrix_dim[1]);
+  PADDLE_ENFORCE_EQ(feature_size,
+                    scale_.size(),
+                    common::errors::InvalidArgument(
+                        "scale's size should be equal to the feature_size,"
+                        "but got feature_size:%d, scale's size:%d.",
+                        feature_size,
+                        scale_.size()));
+  PADDLE_ENFORCE_EQ(feature_size,
+                    bias_.size(),
+                    common::errors::InvalidArgument(
+                        "bias's size should be equal to the feature_size,"
+                        "but got feature_size:%d, bias's size:%d.",
+                        feature_size,
+                        bias_.size()));
+
+  int device_id;
+  cudaGetDevice(&device_id);
+  mean_t.Resize(common::make_ddim({batched_mean_shape}));
+  variance_t.Resize(common::make_ddim({batched_variance_shape}));
+  float *mean_d = mean_t.mutable_data<float>(GPUPlace(device_id));
+  float *variance_d = variance_t.mutable_data<float>(GPUPlace(device_id));
+  auto input_type = getDataType();
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. LayerNorm-->fp32";
+    const float *input = reinterpret_cast<const float *>(inputs[0]);
+    float *output = static_cast<float *>(outputs[0]);
+    phi::LayerNormDirectCUDAFunctor<float, float> layer_norm;
+    layer_norm(stream,
+               input,
+               input_shape,
+               bias_gpu_,
+               scale_gpu_,
+               output,
+               mean_d,
+               variance_d,
+               begin_norm_axis,
+               eps);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. LayerNorm-->fp16";
+    const half *input = reinterpret_cast<const half *>(inputs[0]);
+    half *output = static_cast<half *>(outputs[0]);
+    phi::LayerNormDirectCUDAFunctor<half, float> layer_norm;
+    layer_norm(stream,
+               input,
+               input_shape,
+               bias_gpu_,
+               scale_gpu_,
+               output,
+               mean_d,
+               variance_d,
+               begin_norm_axis,
+               eps);
+  } else {
+    PADDLE_THROW(common::errors::Fatal(
+        "The LayerNorm TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+int LayerNormPluginDynamic::initialize() TRT_NOEXCEPT {
+  cudaMalloc(&bias_gpu_, sizeof(float) * bias_.size());
+  cudaMemcpy(bias_gpu_,
+             bias_.data(),
+             bias_.size() * sizeof(float),
+             cudaMemcpyHostToDevice);
+  cudaMalloc(&scale_gpu_, sizeof(float) * scale_.size());
+  cudaMemcpy(scale_gpu_,
+             scale_.data(),
+             scale_.size() * sizeof(float),
+             cudaMemcpyHostToDevice);
+  return 0;
+}
+
+void LayerNormPluginDynamic::terminate() TRT_NOEXCEPT {
+  if (bias_gpu_) {
+    cudaFree(bias_gpu_);
+    bias_gpu_ = nullptr;
+  }
+  if (scale_gpu_) {
+    cudaFree(scale_gpu_);
+    scale_gpu_ = nullptr;
+  }
+}
+
+nvinfer1::DimsExprs LayerNormPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputDims,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  return inputDims[0];
+}
+
+bool LayerNormPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of layernorm plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return ((in.type == nvinfer1::DataType::kFLOAT ||
+               in.type == nvinfer1::DataType::kHALF) &&
+              (in.format == nvinfer1::PluginFormat::kLINEAR));
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+void LayerNormPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {
+  const auto &input_dims = in[0].desc.dims;
+  int statis_num = 1;
+  for (int i = 0; i < begin_norm_axis_; i++) {
+    statis_num *= input_dims.d[i];
+  }
+  mean_shape_[0] = statis_num;
+  variance_shape_[0] = statis_num;
+}
+
+nvinfer1::DataType LayerNormPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The LayerNormPlugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  PADDLE_ENFORCE_EQ((input_types[0] == nvinfer1::DataType::kFLOAT ||
+                     input_types[0] == nvinfer1::DataType::kHALF),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The input type should be half or float"));
+  return input_types[0];
+}
+
+int LayerNormPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  int begin_norm_axis = begin_norm_axis_;
+  float eps = eps_;
+
+  std::vector<int64_t> input_shape;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    input_shape.push_back(input_dims.d[i]);
+  }
+  // in dynamic shape
+  // the batch num should be involved in mean/variance shape
+  PADDLE_ENFORCE_EQ(1,
+                    mean_shape_.size(),
+                    common::errors::InvalidArgument(
+                        "Size of mean_shape vector should be equal to 1,"
+                        "but got Size of mean_shape vector:%d",
+                        mean_shape_.size()));
+  PADDLE_ENFORCE_EQ(1,
+                    variance_shape_.size(),
+                    common::errors::InvalidArgument(
+                        "Size of variance_shape vector should be equal to 1,"
+                        "but got Size of mean_shape vector:%d",
+                        mean_shape_.size()));
+  PADDLE_ENFORCE_GE(mean_shape_[0],
+                    0,
+                    common::errors::InvalidArgument(
+                        "The size of mean vector should be positive,"
+                        "but got:%d",
+                        mean_shape_[0]));
+  PADDLE_ENFORCE_GE(variance_shape_[0],
+                    0,
+                    common::errors::InvalidArgument(
+                        "The size of mean vector should be positive,"
+                        "but got:%d",
+                        variance_shape_[0]));
+
+  const auto input_ddim = common::make_ddim(input_shape);
+  auto matrix_dim = common::flatten_to_2d(input_ddim, begin_norm_axis);
+  int feature_size = static_cast<int>(matrix_dim[1]);
+  PADDLE_ENFORCE_EQ(feature_size,
+                    scale_.size(),
+                    common::errors::InvalidArgument(
+                        "scale's size should be equal to the feature_size,"
+                        "but got feature_size:%d, scale's size:%d.",
+                        feature_size,
+                        scale_.size()));
+  PADDLE_ENFORCE_EQ(feature_size,
+                    bias_.size(),
+                    common::errors::InvalidArgument(
+                        "bias's size should be equal to the feature_size,"
+                        "but got feature_size:%d, bias's size:%d.",
+                        feature_size,
+                        bias_.size()));
+
+  int device_id;
+  cudaGetDevice(&device_id);
+  mean_t.Resize(common::make_ddim(mean_shape_));
+  variance_t.Resize(common::make_ddim(variance_shape_));
+  float *mean_d = mean_t.mutable_data<float>(GPUPlace(device_id));
+  float *variance_d = variance_t.mutable_data<float>(GPUPlace(device_id));
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. LayerNorm-->fp32";
+    const float *input = reinterpret_cast<const float *>(inputs[0]);
+    float *output = static_cast<float *>(outputs[0]);
+    phi::LayerNormDirectCUDAFunctor<float, float> layer_norm;
+    layer_norm(stream,
+               input,
+               input_shape,
+               bias_gpu_,
+               scale_gpu_,
+               output,
+               mean_d,
+               variance_d,
+               begin_norm_axis,
+               eps);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. LayerNorm-->fp16";
+    const half *input = reinterpret_cast<const half *>(inputs[0]);
+    half *output = static_cast<half *>(outputs[0]);
+    phi::LayerNormDirectCUDAFunctor<half, float> layer_norm;
+    layer_norm(stream,
+               input,
+               input_shape,
+               bias_gpu_,
+               scale_gpu_,
+               output,
+               mean_d,
+               variance_d,
+               begin_norm_axis,
+               eps);
+  } else {
+    PADDLE_THROW(common::errors::Fatal(
+        "The LayerNorm TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/layer_norm_op_plugin.h

@@ -0,0 +1,294 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class LayerNormPlugin : public PluginTensorRT {
+  std::vector<float> bias_;
+  std::vector<float> scale_;
+  phi::DenseTensor mean_t;
+  phi::DenseTensor variance_t;
+  int begin_norm_axis_;
+  float eps_;
+  std::vector<int64_t> mean_shape_;
+  std::vector<int64_t> variance_shape_;
+
+  // data on devices
+  float* bias_gpu_{nullptr};
+  float* scale_gpu_{nullptr};
+
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(bias_) +
+           SerializedSize(scale_) + SerializedSize(begin_norm_axis_) +
+           SerializedSize(eps_) + SerializedSize(mean_shape_) +
+           SerializedSize(variance_shape_) + SerializedSize(with_fp16_);
+  }
+
+  // TRT will call this func when we need to serialize the configuration of
+  // tensorrt.
+  // It should not be called by users.
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, begin_norm_axis_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, mean_shape_);
+    SerializeValue(&buffer, variance_shape_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  LayerNormPlugin(const float* bias,
+                  const int bias_num,
+                  const float* scale,
+                  const int scale_num,
+                  int begin_norm_axis,
+                  float eps,
+                  std::vector<int64_t> mean_shape,
+                  std::vector<int64_t> variance_shape,
+                  bool with_fp16)
+      : begin_norm_axis_(begin_norm_axis),
+        eps_(eps),
+        mean_shape_(mean_shape),
+        variance_shape_(variance_shape) {
+    with_fp16_ = with_fp16;
+    bias_.resize(bias_num);
+    scale_.resize(scale_num);
+    std::copy(bias, bias + bias_num, bias_.data());
+    std::copy(scale, scale + scale_num, scale_.data());
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  LayerNormPlugin(void const* serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &begin_norm_axis_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &mean_shape_);
+    DeserializeValue(&serialData, &serialLength, &variance_shape_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  ~LayerNormPlugin() {}
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+
+  LayerNormPlugin* clone() const TRT_NOEXCEPT override {
+    auto ptr = new LayerNormPlugin(bias_.data(),
+                                   bias_.size(),
+                                   scale_.data(),
+                                   scale_.size(),
+                                   begin_norm_axis_,
+                                   eps_,
+                                   mean_shape_,
+                                   variance_shape_,
+                                   with_fp16_);
+    ptr->bias_gpu_ = bias_gpu_;
+    ptr->scale_gpu_ = scale_gpu_;
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "layernorm_plugin";
+  }
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override;
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+  int enqueue(int batchSize,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+};
+
+class LayerNormPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "layernorm_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new LayerNormPlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(LayerNormPluginCreator);
+
+class LayerNormPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  LayerNormPluginDynamic(const float* bias,
+                         const int bias_num,
+                         const float* scale,
+                         const int scale_num,
+                         int begin_norm_axis,
+                         float eps,
+                         std::vector<int64_t> mean_shape,
+                         std::vector<int64_t> variance_shape,
+                         bool with_fp16)
+      : begin_norm_axis_(begin_norm_axis),
+        eps_(eps),
+        mean_shape_(mean_shape),
+        variance_shape_(variance_shape) {
+    with_fp16_ = with_fp16;
+    bias_.resize(bias_num);
+    scale_.resize(scale_num);
+    std::copy(bias, bias + bias_num, bias_.data());
+    std::copy(scale, scale + scale_num, scale_.data());
+  }
+
+  LayerNormPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &begin_norm_axis_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &mean_shape_);
+    DeserializeValue(&serialData, &serialLength, &variance_shape_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    auto ptr = new LayerNormPluginDynamic(bias_.data(),
+                                          bias_.size(),
+                                          scale_.data(),
+                                          scale_.size(),
+                                          begin_norm_axis_,
+                                          eps_,
+                                          mean_shape_,
+                                          variance_shape_,
+                                          with_fp16_);
+    ptr->bias_gpu_ = bias_gpu_;
+    ptr->scale_gpu_ = scale_gpu_;
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "layernorm_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(bias_) + SerializedSize(scale_) +
+           SerializedSize(begin_norm_axis_) + SerializedSize(eps_) +
+           SerializedSize(mean_shape_) + SerializedSize(variance_shape_) +
+           SerializedSize(with_fp16_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, begin_norm_axis_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, mean_shape_);
+    SerializeValue(&buffer, variance_shape_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  std::vector<float> bias_;
+  std::vector<float> scale_;
+  phi::DenseTensor mean_t;
+  phi::DenseTensor variance_t;
+  int begin_norm_axis_;
+  float eps_;
+  std::vector<int64_t> mean_shape_;
+  std::vector<int64_t> variance_shape_;
+  // data on devices
+  float* bias_gpu_{nullptr};
+  float* scale_gpu_{nullptr};
+};
+
+class LayerNormPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "layernorm_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new LayerNormPluginDynamic(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(LayerNormPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/layernorm_shift_partition_op.cu

@@ -0,0 +1,685 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2019-2022, NVIDIA CORPORATION.  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 <stdio.h>
+
+#include <cassert>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/layernorm_shift_partition_op.h"
+#include "paddle/phi/kernels/layer_norm_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+#define FINAL_MASK 0xffffffff
+
+template <typename T>
+__inline__ __device__ T warpReduceSum(T val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1)
+    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
+  return val;
+}
+
+/* Calculate the sum of all elements in a block */
+template <typename T>
+__inline__ __device__ T blockReduceSum(T val) {
+  static __shared__ T shared[32];
+  int lane = threadIdx.x & 0x1f;
+  int wid = threadIdx.x >> 5;
+
+  val = warpReduceSum<T>(val);
+
+  if (lane == 0) shared[wid] = val;
+
+  __syncthreads();
+
+  // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
+  // blockDim.x is not divided by 32
+  val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
+  val = warpReduceSum<T>(val);
+
+  return val;
+}
+
+template <typename T>
+__global__ void layernorm_shift_partition(T *out,
+                                          const T *input,
+                                          const T *gamma,
+                                          const T *beta,
+                                          int batch,
+                                          int H,
+                                          int W,
+                                          int n,
+                                          int shift_size,
+                                          int window_size,
+                                          const float eps) {
+  int tid = threadIdx.x;
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  float local_out =
+      (tid < n) ? static_cast<float>(__ldg(input + bid * n + tid)) : 0.0f;
+#else
+  float local_out = (tid < n) ? static_cast<float>(input[bid * n + tid]) : 0.0f;
+#endif
+
+  mean = blockReduceSum<float>(local_out);
+  if (threadIdx.x == 0) {
+    s_mean = mean / n;
+  }
+  __syncthreads();
+
+  float diff = (tid < n) ? (local_out - s_mean) : 0.0f;
+  variance = blockReduceSum<float>(diff * diff);
+  if (threadIdx.x == 0) {
+    s_variance = variance / n + eps;
+  }
+  __syncthreads();
+
+  if (tid < n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+    out[output_bid * n + tid] =
+        (T)(((local_out - s_mean) * rsqrtf(s_variance)) *
+                static_cast<float>(__ldg(&gamma[tid])) +
+            static_cast<float>(__ldg(&beta[tid])));
+#else
+    out[output_bid * n + tid] =
+        (T)(((local_out - s_mean) * rsqrtf(s_variance)) *
+                static_cast<float>(gamma[tid]) +
+            static_cast<float>(beta[tid]));
+#endif
+  }
+}
+
+template <>
+__global__ void layernorm_shift_partition(half2 *out_ptr,
+                                          const half2 *input_ptr,
+                                          const half2 *gamma_ptr,
+                                          const half2 *beta_ptr,
+                                          int batch,
+                                          int H,
+                                          int W,
+                                          int n,
+                                          int shift_size,
+                                          int window_size,
+                                          const float eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  int tid = threadIdx.x;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  float2 local_out_fp2;
+
+  float local_out = 0.0f;
+  int id = bid * n + tid;
+  if (tid < n) {
+    local_out_fp2 = __half22float2(__ldg(input_ptr + id));
+    local_out += local_out_fp2.x;
+    local_out += local_out_fp2.y;
+  }
+
+  mean = blockReduceSum<float>(local_out);
+  if (threadIdx.x == 0) {
+    s_mean = mean / (n * 2);
+  }
+  __syncthreads();
+
+  if (tid < n) {
+    variance = (local_out_fp2.x - s_mean) * (local_out_fp2.x - s_mean);
+    variance += (local_out_fp2.y - s_mean) * (local_out_fp2.y - s_mean);
+  }
+  variance = blockReduceSum<float>(variance);
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / (n * 2) + eps);
+  }
+  __syncthreads();
+
+  if (tid < n) {
+    float2 gamma_val = __half22float2(__ldg(&gamma_ptr[tid]));
+    float2 beta_val = __half22float2(__ldg(&beta_ptr[tid]));
+    local_out_fp2.x =
+        (local_out_fp2.x - s_mean) * s_variance * gamma_val.x + beta_val.x;
+    local_out_fp2.y =
+        (local_out_fp2.y - s_mean) * s_variance * gamma_val.y + beta_val.y;
+    out_ptr[output_bid * n + tid] = __float22half2_rn(local_out_fp2);
+  }
+#endif
+}
+
+#define kITE 4
+template <typename T>
+__global__ void layernorm_shift_partition_v2(T *out,
+                                             const T *__restrict input,
+                                             const T *__restrict gamma,
+                                             const T *__restrict beta,
+                                             int batch,
+                                             int H,
+                                             int W,
+                                             int n,
+                                             int shift_size,
+                                             int window_size,
+                                             const float eps) {
+  // constexpr int kITE = 4;
+  const int tid = threadIdx.x;
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  const int offset = bid * n;
+  const int output_offset = output_bid * n;
+
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  float local_out[kITE];
+
+  float sum = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+      local_out[i] = static_cast<float>(__ldg(input + offset + col_id));
+#else
+      local_out[i] = static_cast<float>(input[offset + col_id]);
+#endif
+      sum += local_out[i];
+    }
+  }
+
+  mean = blockReduceSum<float>(sum);
+  if (tid == 0) {
+    s_mean = mean / n;
+  }
+  __syncthreads();
+
+  float var = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      float diff = local_out[i] - s_mean;
+      local_out[i] = diff;
+      var += diff * diff;
+    }
+  }
+
+  variance = blockReduceSum<float>(var);
+  if (tid == 0) {
+    s_variance = rsqrtf(variance / n + eps);
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+      out[output_offset + col_id] =
+          (T)(local_out[i] * s_variance *
+                  static_cast<float>(__ldg(&gamma[col_id])) +
+              static_cast<float>(__ldg(&beta[col_id])));
+#else
+      out[output_offset + col_id] =
+          (T)(local_out[i] * s_variance * static_cast<float>(gamma[col_id]) +
+              static_cast<float>(beta[col_id]));
+#endif
+    }
+  }
+}
+
+template <>
+__global__ void layernorm_shift_partition_v2(half2 *out_ptr,
+                                             const half2 *__restrict input_ptr,
+                                             const half2 *__restrict gamma_ptr,
+                                             const half2 *__restrict beta_ptr,
+                                             int batch,
+                                             int H,
+                                             int W,
+                                             int n,
+                                             int shift_size,
+                                             int window_size,
+                                             const float eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  // constexpr int ite = 4;
+  const int tid = threadIdx.x;
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  const int offset = bid * n;
+  const int output_offset = output_bid * n;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  half2 local_out_half2[kITE];
+  const half2 zero = {static_cast<half>(0.0f), static_cast<half>(0.0f)};
+
+  // float sum = 0.0f;
+  half2 sum = __float2half2_rn(0.0f);
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      local_out_half2[i] = __ldg(input_ptr + offset + col_id);
+      sum += local_out_half2[i];
+    }
+  }
+
+  mean = blockReduceSum<float>(static_cast<float>(sum.x + sum.y));
+  if (threadIdx.x == 0) {
+    s_mean = mean / (n * 2);
+  }
+  __syncthreads();
+
+  float var = 0.0f;
+  half2 s_mean_2 = __float2half2_rn(s_mean);
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      local_out_half2[i] = local_out_half2[i] - s_mean_2;
+      float v1 = static_cast<float>(local_out_half2[i].x);
+      float v2 = static_cast<float>(local_out_half2[i].y);
+      var += v1 * v1 + v2 * v2;
+    }
+  }
+
+  variance = blockReduceSum<float>(var);
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / (n * 2) + eps);
+  }
+  __syncthreads();
+
+  half2 s_var_2 = __float2half2_rn(s_variance);
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      out_ptr[output_offset + col_id] =
+          local_out_half2[i] * s_var_2 * __ldg(&gamma_ptr[col_id]) +
+          __ldg(&beta_ptr[col_id]);
+    }
+  }
+#endif
+}
+
+template <typename T>
+void invokeLayernormShiftPartition(T *out,
+                                   const T *input,
+                                   const T *gamma,
+                                   const T *beta,
+                                   int batch,
+                                   int H,
+                                   int W,
+                                   int n,
+                                   int shift_size,
+                                   int window_size,
+                                   const float eps,
+                                   cudaStream_t stream) {
+  dim3 grid(W, H, batch);
+  int blockSize = (n + 31) / 32 * 32;
+  if (blockSize >= 768) {
+    blockSize = ((blockSize / 4) + 31) / 32 * 32;
+    layernorm_shift_partition_v2<T><<<grid, blockSize, 0, stream>>>(
+        out, input, gamma, beta, batch, H, W, n, shift_size, window_size, eps);
+  } else {
+    layernorm_shift_partition<T><<<grid, blockSize, 0, stream>>>(
+        out, input, gamma, beta, batch, H, W, n, shift_size, window_size, eps);
+  }
+}
+
+template <>
+void invokeLayernormShiftPartition(half *out,
+                                   const half *input,
+                                   const half *gamma,
+                                   const half *beta,
+                                   int batch,
+                                   int H,
+                                   int W,
+                                   int n,
+                                   int shift_size,
+                                   int window_size,
+                                   const float eps,
+                                   cudaStream_t stream) {
+  dim3 grid(W, H, batch);
+  int blockSize = n / 2;
+  blockSize = (blockSize + 31) / 32 * 32;
+
+  if ((batch * H * W >= 512 && blockSize >= 768) || blockSize > 1024) {
+    blockSize = ((blockSize / 4) + 31) / 32 * 32;
+    layernorm_shift_partition_v2<<<grid, blockSize, 0, stream>>>(
+        reinterpret_cast<half2 *>(out),
+        (const half2 *)input,
+        (const half2 *)gamma,
+        (const half2 *)beta,
+        batch,
+        H,
+        W,
+        n / 2,
+        shift_size,
+        window_size,
+        eps);
+  } else {
+    layernorm_shift_partition<<<grid, blockSize, 0, stream>>>(
+        reinterpret_cast<half2 *>(out),
+        (const half2 *)input,
+        (const half2 *)gamma,
+        (const half2 *)beta,
+        batch,
+        H,
+        W,
+        n / 2,
+        shift_size,
+        window_size,
+        eps);
+  }
+}
+
+template <typename T>
+static void convertAndCopy(const std::vector<float> &host, T *dev) {
+  T *host_ptr = new T[host.size()];
+  std::transform(host.begin(), host.end(), host_ptr, [](float x) {
+    return static_cast<T>(x);
+  });
+  cudaMemcpy(dev, host_ptr, sizeof(T) * host.size(), cudaMemcpyHostToDevice);
+  delete host_ptr;
+}
+
+void LayernormShiftPartitionPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+LayernormShiftPartitionPluginDynamic::LayernormShiftPartitionPluginDynamic(
+    const float *gamma,
+    const float *beta,
+    const int param_num,
+    int shift_size,
+    int window_size,
+    int input_resolution,
+    float eps,
+    bool with_fp16,
+    std::shared_ptr<void> gamma_dev,
+    std::shared_ptr<void> beta_dev)
+    : with_fp16_(with_fp16),
+      window_size_(window_size),
+      shift_size_(shift_size),
+      input_resolution_(input_resolution),
+      eps_(eps),
+      param_num_(param_num),
+      gamma_dev_(gamma_dev),
+      beta_dev_(beta_dev) {
+  beta_.resize(param_num);
+  gamma_.resize(param_num);
+  std::copy(gamma, gamma + param_num, gamma_.data());
+  std::copy(beta, beta + param_num, beta_.data());
+  int type_size = with_fp16 ? sizeof(half) : sizeof(float);
+  if (gamma_dev_ == nullptr) {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    gamma_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16)
+      convertAndCopy(gamma_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(gamma_, reinterpret_cast<float *>(p));
+  }
+  if (beta_dev_ == nullptr) {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    beta_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16)
+      convertAndCopy(beta_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(beta_, reinterpret_cast<float *>(p));
+  }
+}
+
+LayernormShiftPartitionPluginDynamic::LayernormShiftPartitionPluginDynamic(
+    void const *serialData, size_t serialLength) {
+  DeserializeValue(&serialData, &serialLength, &beta_);
+  DeserializeValue(&serialData, &serialLength, &gamma_);
+  DeserializeValue(&serialData, &serialLength, &param_num_);
+  DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  DeserializeValue(&serialData, &serialLength, &shift_size_);
+  DeserializeValue(&serialData, &serialLength, &window_size_);
+  DeserializeValue(&serialData, &serialLength, &input_resolution_);
+  DeserializeValue(&serialData, &serialLength, &eps_);
+  int type_size = with_fp16_ ? sizeof(half) : sizeof(float);
+  {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    gamma_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16_)
+      convertAndCopy(gamma_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(gamma_, reinterpret_cast<float *>(p));
+  }
+  {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    beta_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16_)
+      convertAndCopy(beta_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(beta_, reinterpret_cast<float *>(p));
+  }
+}
+
+bool LayernormShiftPartitionPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument("The input of LayernormShiftPartition "
+                                      "plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return in.type == nvinfer1::DataType::kHALF &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    } else {
+      return in.type == nvinfer1::DataType::kFLOAT &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType LayernormShiftPartitionPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      index,
+      0,
+      common::errors::InvalidArgument(
+          "The LayernormShiftPartition only has one input, so the "
+          "index value should be 0, but get %d.",
+          index));
+  return input_types[0];
+}
+
+nvinfer1::DimsExprs LayernormShiftPartitionPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      output_index,
+      0,
+      common::errors::InvalidArgument(
+          "There is only one output of the LayernormShiftPartition, "
+          "so the index should be zero,"
+          "but it's (%d)",
+          output_index));
+  PADDLE_ENFORCE_EQ(
+      nb_inputs,
+      1,
+      common::errors::InvalidArgument(
+          "The Input of the LayernormShiftPartition should be 1, but we found "
+          "it has (%d) inputs",
+          nb_inputs));
+
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = expr_builder.operation(
+      nvinfer1::DimensionOperation::kFLOOR_DIV,
+      *expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
+                              *inputs[0].d[0],
+                              *inputs[0].d[1]),
+      *expr_builder.constant(window_size_ * window_size_));
+  ret.d[1] = expr_builder.constant(window_size_ * window_size_);
+  ret.d[2] = inputs[0].d[2];
+  return ret;
+}
+
+int LayernormShiftPartitionPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  auto input_type = input_desc[0].type;
+  int batch = input_dims.d[0];
+  int emb_dim = input_dims.d[2];
+  PADDLE_ENFORCE_EQ(
+      input_resolution_ * input_resolution_,
+      input_dims.d[1],
+      common::errors::InvalidArgument(
+          "The LayernormShiftPartition's input_resolution is wrong (%d)",
+          input_dims.d[1]));
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(3) << "TRT Plugin DataType selected. LayernormShiftPartition-->fp32";
+    invokeLayernormShiftPartition(
+        reinterpret_cast<float *>(outputs[0]),
+        reinterpret_cast<const float *>(inputs[0]),
+        reinterpret_cast<const float *>(gamma_dev_.get()),
+        reinterpret_cast<const float *>(beta_dev_.get()),
+        batch,
+        input_resolution_,
+        input_resolution_,
+        emb_dim,
+        shift_size_,
+        window_size_,
+        eps_,
+        stream);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(3) << "TRT Plugin DataType selected. LayernormShiftPartition-->half";
+    invokeLayernormShiftPartition(
+        reinterpret_cast<half *>(outputs[0]),
+        reinterpret_cast<const half *>(inputs[0]),
+        reinterpret_cast<const half *>(gamma_dev_.get()),
+        reinterpret_cast<const half *>(beta_dev_.get()),
+        batch,
+        input_resolution_,
+        input_resolution_,
+        emb_dim,
+        shift_size_,
+        window_size_,
+        eps_,
+        stream);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The LayerNorm TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/layernorm_shift_partition_op.h

@@ -0,0 +1,157 @@
+// 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class LayernormShiftPartitionPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  LayernormShiftPartitionPluginDynamic(
+      const float* gamma,
+      const float* beta,
+      const int param_num,
+      int shift_size,
+      int window_size,
+      int input_resolution,
+      float eps,
+      bool with_fp16,
+      std::shared_ptr<void> gamma_dev = nullptr,
+      std::shared_ptr<void> beta_dev = nullptr);
+
+  LayernormShiftPartitionPluginDynamic(void const* serialData,
+                                       size_t serialLength);
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new LayernormShiftPartitionPluginDynamic(gamma_.data(),
+                                                    beta_.data(),
+                                                    beta_.size(),
+                                                    shift_size_,
+                                                    window_size_,
+                                                    input_resolution_,
+                                                    eps_,
+                                                    with_fp16_,
+                                                    gamma_dev_,
+                                                    beta_dev_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "layernorm_shift_partition_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(beta_) + SerializedSize(gamma_) +
+           SerializedSize(param_num_) + SerializedSize(with_fp16_) +
+           SerializedSize(shift_size_) + SerializedSize(window_size_) +
+           SerializedSize(input_resolution_) + SerializedSize(eps_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, beta_);
+    SerializeValue(&buffer, gamma_);
+    SerializeValue(&buffer, param_num_);
+    SerializeValue(&buffer, with_fp16_);
+    SerializeValue(&buffer, shift_size_);
+    SerializeValue(&buffer, window_size_);
+    SerializeValue(&buffer, input_resolution_);
+    SerializeValue(&buffer, eps_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  bool with_fp16_;
+  std::vector<float> gamma_;
+  std::vector<float> beta_;
+  int window_size_;
+  int shift_size_;
+  int input_resolution_;
+  int param_num_;
+  float eps_;
+  std::shared_ptr<void> gamma_dev_;
+  std::shared_ptr<void> beta_dev_;
+};
+
+class LayernormShiftPartitionPluginDynamicCreator
+    : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "layernorm_shift_partition_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new LayernormShiftPartitionPluginDynamic(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(LayernormShiftPartitionPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_kernel.cu

@@ -0,0 +1,475 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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 <cuda.h>
+#include <cassert>
+#include <cstring>
+#include <iostream>
+#include <vector>
+
+#include "NvInfer.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/common.cuh"
+#include "paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernel_2(int32_t ld,
+                                     int32_t const* inputIds0,
+                                     int32_t const* inputIds1,
+                                     float const* beta,
+                                     float const* gamma,
+                                     T const* mIdsEmbDev0,
+                                     T const* mIdsEmbDev1,
+                                     int32_t IdsSize0,
+                                     int32_t IdsSize1,
+                                     T* output) {
+  T const rld = T(1.f) / T(ld);
+  cub::Sum pairSum;
+  int32_t const seqPos = blockIdx.y * gridDim.x + blockIdx.x;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds0)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds0)[seqPos] >= IdsSize0) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds0)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[1] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // offset into embeddings is given by wordId * hidden_size
+  int32_t const outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    int32_t const offset0 = word_id[0] * ld;
+    T val = mIdsEmbDev0[offset0 + it];
+    int32_t const offset1 = word_id[1] * ld;
+    val += mIdsEmbDev1[offset1 + it];
+
+    output[outOffset + it] = val;
+    T const rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernel_3(int32_t ld,
+                                     int32_t const* inputIds0,
+                                     int32_t const* inputIds1,
+                                     int32_t const* inputIds2,
+                                     float const* beta,
+                                     float const* gamma,
+                                     T const* mIdsEmbDev0,
+                                     T const* mIdsEmbDev1,
+                                     T const* mIdsEmbDev2,
+                                     int32_t IdsSize0,
+                                     int32_t IdsSize1,
+                                     int32_t IdsSize2,
+                                     T* output) {
+  T const rld = T(1.f) / T(ld);
+  cub::Sum pairSum;
+  int32_t const seqPos = blockIdx.y * gridDim.x + blockIdx.x;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds0)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds0)[seqPos] >= IdsSize0) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds0)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[1] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds2)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds2)[seqPos] >= IdsSize2) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[2] = static_cast<int32_t const*>(inputIds2)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // offset into embeddings is given by wordId * hidden_size
+  int32_t const outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    int32_t const offset0 = word_id[0] * ld;
+    T val = mIdsEmbDev0[offset0 + it];
+    int32_t const offset1 = word_id[1] * ld;
+    val += mIdsEmbDev1[offset1 + it];
+    int32_t const offset2 = word_id[2] * ld;
+    val += mIdsEmbDev2[offset2 + it];
+
+    output[outOffset + it] = val;
+    T const rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernel_4(int32_t ld,
+                                     int32_t const* inputIds0,
+                                     int32_t const* inputIds1,
+                                     int32_t const* inputIds2,
+                                     int32_t const* inputIds3,
+                                     float const* beta,
+                                     float const* gamma,
+                                     T const* mIdsEmbDev0,
+                                     T const* mIdsEmbDev1,
+                                     T const* mIdsEmbDev2,
+                                     T const* mIdsEmbDev3,
+                                     int32_t IdsSize0,
+                                     int32_t IdsSize1,
+                                     int32_t IdsSize2,
+                                     int32_t IdsSize3,
+                                     T* output) {
+  T const rld = T(1.f) / T(ld);
+  cub::Sum pairSum;
+  int32_t const seqPos = blockIdx.y * gridDim.x + blockIdx.x;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds0)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds0)[seqPos] >= IdsSize0) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds0)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[1] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds2)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds2)[seqPos] >= IdsSize2) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[2] = static_cast<int32_t const*>(inputIds2)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds3)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds3)[seqPos] >= IdsSize3) {
+      printf(
+          "Error!!!!!!(embLayerNormVarPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[3] = static_cast<int32_t const*>(inputIds3)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // offset into embeddings is given by wordId * hidden_size
+  int32_t const outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    int32_t const offset0 = word_id[0] * ld;
+    T val = mIdsEmbDev0[offset0 + it];
+    int32_t const offset1 = word_id[1] * ld;
+    val += mIdsEmbDev1[offset1 + it];
+    int32_t const offset2 = word_id[2] * ld;
+    val += mIdsEmbDev2[offset2 + it];
+    int32_t const offset3 = word_id[3] * ld;
+    val += mIdsEmbDev3[offset3 + it];
+
+    output[outOffset + it] = val;
+    T const rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T>
+int32_t embSkipLayerNorm_2(cudaStream_t stream,
+                           int32_t ld,
+                           int32_t B,
+                           int32_t S,
+                           int const* inputIds0,
+                           int const* inputIds1,
+                           int32_t nbLookupTables,
+                           float const* beta,
+                           float const* gamma,
+                           T const* mIdsEmbDev0,
+                           T const* mIdsEmbDev1,
+                           int32_t IdsSize0,
+                           int32_t IdsSize1,
+                           T* output) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * nbLookupTables;
+  embLayerNormKernel_2<T, tpb><<<grid, block, cache_size, stream>>>(ld,
+                                                                    inputIds0,
+                                                                    inputIds1,
+                                                                    beta,
+                                                                    gamma,
+                                                                    mIdsEmbDev0,
+                                                                    mIdsEmbDev1,
+                                                                    IdsSize0,
+                                                                    IdsSize1,
+                                                                    output);
+  return cudaPeekAtLastError();
+}
+
+template <typename T>
+int32_t embSkipLayerNorm_3(cudaStream_t stream,
+                           int32_t ld,
+                           int32_t B,
+                           int32_t S,
+                           int const* inputIds0,
+                           int const* inputIds1,
+                           int const* inputIds2,
+                           int32_t nbLookupTables,
+                           float const* beta,
+                           float const* gamma,
+                           T const* mIdsEmbDev0,
+                           T const* mIdsEmbDev1,
+                           T const* mIdsEmbDev2,
+                           int32_t IdsSize0,
+                           int32_t IdsSize1,
+                           int32_t IdsSize2,
+                           T* output) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * nbLookupTables;
+  embLayerNormKernel_3<T, tpb><<<grid, block, cache_size, stream>>>(ld,
+                                                                    inputIds0,
+                                                                    inputIds1,
+                                                                    inputIds2,
+                                                                    beta,
+                                                                    gamma,
+                                                                    mIdsEmbDev0,
+                                                                    mIdsEmbDev1,
+                                                                    mIdsEmbDev2,
+                                                                    IdsSize0,
+                                                                    IdsSize1,
+                                                                    IdsSize2,
+                                                                    output);
+  return cudaPeekAtLastError();
+}
+
+template <typename T>
+int32_t embSkipLayerNorm_4(cudaStream_t stream,
+                           int32_t ld,
+                           int32_t B,
+                           int32_t S,
+                           int const* inputIds0,
+                           int const* inputIds1,
+                           int const* inputIds2,
+                           int const* inputIds3,
+                           int32_t nbLookupTables,
+                           float const* beta,
+                           float const* gamma,
+                           T const* mIdsEmbDev0,
+                           T const* mIdsEmbDev1,
+                           T const* mIdsEmbDev2,
+                           T const* mIdsEmbDev3,
+                           int32_t IdsSize0,
+                           int32_t IdsSize1,
+                           int32_t IdsSize2,
+                           int32_t IdsSize3,
+                           T* output) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * nbLookupTables;
+  embLayerNormKernel_4<T, tpb><<<grid, block, cache_size, stream>>>(ld,
+                                                                    inputIds0,
+                                                                    inputIds1,
+                                                                    inputIds2,
+                                                                    inputIds3,
+                                                                    beta,
+                                                                    gamma,
+                                                                    mIdsEmbDev0,
+                                                                    mIdsEmbDev1,
+                                                                    mIdsEmbDev2,
+                                                                    mIdsEmbDev3,
+                                                                    IdsSize0,
+                                                                    IdsSize1,
+                                                                    IdsSize2,
+                                                                    IdsSize3,
+                                                                    output);
+  return cudaPeekAtLastError();
+}
+
+template int32_t embSkipLayerNorm_2<float>(cudaStream_t,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t const*,
+                                           int32_t const*,
+                                           int32_t,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           int32_t,
+                                           int32_t,
+                                           float*);
+
+template int32_t embSkipLayerNorm_3<float>(cudaStream_t,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t const*,
+                                           int32_t const*,
+                                           int32_t const*,
+                                           int32_t,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t,
+                                           float*);
+
+template int32_t embSkipLayerNorm_4<float>(cudaStream_t,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t const*,
+                                           int32_t const*,
+                                           int32_t const*,
+                                           int32_t const*,
+                                           int32_t,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           float const*,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t,
+                                           int32_t,
+                                           float*);
+
+template int32_t embSkipLayerNorm_2<half>(cudaStream_t,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t const*,
+                                          int32_t const*,
+                                          int32_t,
+                                          float const*,
+                                          float const*,
+                                          half const*,
+                                          half const*,
+                                          int32_t,
+                                          int32_t,
+                                          half*);
+
+template int32_t embSkipLayerNorm_3<half>(cudaStream_t,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t const*,
+                                          int32_t const*,
+                                          int32_t const*,
+                                          int32_t,
+                                          float const*,
+                                          float const*,
+                                          half const*,
+                                          half const*,
+                                          half const*,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t,
+                                          half*);
+
+template int32_t embSkipLayerNorm_4<half>(cudaStream_t,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t const*,
+                                          int32_t const*,
+                                          int32_t const*,
+                                          int32_t const*,
+                                          int32_t,
+                                          float const*,
+                                          float const*,
+                                          half const*,
+                                          half const*,
+                                          half const*,
+                                          half const*,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t,
+                                          int32_t,
+                                          half*);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_plugin.cu

@@ -0,0 +1,497 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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 "paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_plugin.h"
+#include <cuda.h>
+#include <cstring>
+#include <vector>
+#include "NvInfer.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+constexpr size_t threadsPerCta128 = 2 * 2 * 32;
+constexpr size_t threadsPerCta256 = 1 * 4 * 32;
+constexpr size_t threadsPerCta384 = 1 * 8 * 32;
+// The number of xmmas in the M dimension. We use one uint32_t per XMMA in the M
+// dimension: (s + 16*warps_m - 1) / (16*warps_m);
+constexpr size_t xmmasM128 = 4;
+constexpr size_t xmmasM256 = 16;
+constexpr size_t xmmasM384 = 24;
+// Packed mask size per batch. Layout is XMMAS_M * THREADS_PER_CTA.
+constexpr size_t packedMaskSize128 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize256 = xmmasM256 * threadsPerCta256;
+constexpr size_t packedMaskSize384 = xmmasM384 * threadsPerCta384;
+char const* EMB_LAYER_NORM_VERSION{"1"};
+char const* EMB_LAYER_NORM_NAME{"ManyEmbLayerNormPluginDynamic"};
+// Static class fields initialization
+nvinfer1::PluginFieldCollection EmbLayerNormPluginCreator::mFC{};
+std::vector<nvinfer1::PluginField> EmbLayerNormPluginCreator::mPluginAttributes;
+
+EmbLayerNormPlugin::EmbLayerNormPlugin(
+    std::string const& name,
+    nvinfer1::DataType const type,
+    nvinfer1::Weights const& beta,
+    nvinfer1::Weights const& gamma,
+    const std::vector<nvinfer1::Weights>& IdsEmb)
+    : mLayerName(name),
+      mLd(beta.count),
+      mType(type),
+      mIdsEmb_(IdsEmb),
+      nbLookupTables_(static_cast<int>(IdsEmb.size())) {
+  // Assuming Weights.count is the number of elements and not bytes
+  assert(beta.count == gamma.count);
+  mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
+  mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
+  copyToDevice(&mGamma, sizeof(float) * mGamma.count, &mGammaDev);
+  copyToDevice(&mBeta, sizeof(float) * mBeta.count, &mBetaDev);
+  for (size_t i = 0; i < mIdsEmb_.size(); ++i) {
+    assert(mIdsEmb_[i].count % mLd == 0);
+    mIdsVocabSize.push_back(int32_t(mIdsEmb_[i].count / mLd));
+    WeightsWithOwnership tem_weight;
+    tem_weight.convertAndCopy(mIdsEmb_[i], mType);
+    void* cudaMem{nullptr};
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaMalloc(&cudaMem, getWeightsSize(tem_weight, mType)));
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpy(cudaMem,
+                                          tem_weight.values,
+                                          getWeightsSize(tem_weight, mType),
+                                          cudaMemcpyHostToDevice));
+    mIdsEmbPtrs.push_back(cudaMem);
+  }
+}
+
+EmbLayerNormPlugin::EmbLayerNormPlugin(std::string const& name,
+                                       void const* data,
+                                       size_t length)
+    : mLayerName(name),
+      mGammaDev(nullptr),
+      mBetaDev(nullptr),
+      mIdsEmbPtrs{},
+      mIdsEmb_{} {
+  // Deserialize in the same order as serialization
+  deserialize_value(&data, &length, &mType);
+  deserialize_value(&data, &length, &mLd);
+  deserialize_value(&data, &length, &nbLookupTables_);
+  for (int32_t i = 0; i < nbLookupTables_; ++i) {
+    int32_t tem;
+    deserialize_value(&data, &length, &tem);
+    mIdsVocabSize.push_back(tem);
+  }
+  char const* d = static_cast<char const*>(data);
+  mBeta.convertAndCopy(&d, mLd, nvinfer1::DataType::kFLOAT);
+  mGamma.convertAndCopy(&d, mLd, nvinfer1::DataType::kFLOAT);
+  for (int32_t i = 0; i < nbLookupTables_; ++i) {
+    nvinfer1::Weights pre_tem_weight;
+    pre_tem_weight.type = mType;
+    pre_tem_weight.count = mLd * size_t(mIdsVocabSize[i]);
+    const auto nbBytes = mLd * size_t(mIdsVocabSize[i]) * getElementSize(mType);
+    auto destBuf = new char[nbBytes];
+    pre_tem_weight.values = destBuf;
+    std::copy_n(d, nbBytes, destBuf);
+    d += nbBytes;
+    mIdsEmb_.push_back(pre_tem_weight);
+  }
+}
+
+// IPluginV2DynamicExt Methods
+nvinfer1::IPluginV2DynamicExt* EmbLayerNormPlugin::clone() const noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormPlugin clone");
+  auto p = new EmbLayerNormPlugin(mLayerName, mType, mBeta, mGamma, mIdsEmb_);
+  p->setPluginNamespace(mNamespace.c_str());
+  return p;
+}
+
+nvinfer1::DimsExprs EmbLayerNormPlugin::getOutputDimensions(
+    int32_t outputIndex,
+    nvinfer1::DimsExprs const* inputs,
+    int32_t nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) noexcept {
+  assert(outputIndex == 0);
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = inputs[0].d[0];
+  ret.d[1] = inputs[0].d[1];
+  ret.d[2] = exprBuilder.constant(mLd);
+  return ret;
+}
+
+bool EmbLayerNormPlugin::supportsFormatCombination(
+    int32_t pos,
+    nvinfer1::PluginTensorDesc const* inOut,
+    int32_t nbInputs,
+    int32_t nbOutputs) noexcept {
+  assert(nbOutputs == 1);
+  nvinfer1::PluginTensorDesc const& prev = inOut[0];
+  nvinfer1::PluginTensorDesc const& desc = inOut[pos];
+  if (desc.format != nvinfer1::TensorFormat::kLINEAR) {
+    return false;
+  }
+  if (pos == 0) {
+    return desc.type == nvinfer1::DataType::kINT32;
+  }
+  if (0 < pos && pos < nbInputs) {
+    assert(desc.dims.nbDims == prev.dims.nbDims);
+    for (int i = 0; i < prev.dims.nbDims; ++i) {
+      assert(desc.dims.d[i] == prev.dims.d[i]);
+    }
+    return desc.type == prev.type;
+  }
+  if (pos == nbInputs) {  // output
+    return desc.type == mType && desc.dims.nbDims == 3 &&
+           desc.dims.d[0] == prev.dims.d[0] && desc.dims.d[1] == prev.dims.d[1];
+  }
+}
+
+void EmbLayerNormPlugin::configurePlugin(
+    nvinfer1::DynamicPluginTensorDesc const* inputs,
+    int32_t nbInputs,
+    nvinfer1::DynamicPluginTensorDesc const* outputs,
+    int32_t nbOutputs) noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormPlugin configurePlugin");
+  assert(static_cast<size_t>(outputs[0].desc.dims.d[2]) ==
+         static_cast<size_t>(mLd));
+  int32_t const B = inputs[0].desc.dims.d[0];
+  if (B > 0) {
+    assert(outputs[0].desc.dims.d[0] == B);
+  }
+  assert(outputs[0].desc.type == mType);
+}
+
+size_t EmbLayerNormPlugin::getWorkspaceSize(
+    nvinfer1::PluginTensorDesc const* inputs,
+    int32_t nbInputs,
+    nvinfer1::PluginTensorDesc const* outputs,
+    int32_t nbOutputs) const noexcept {
+  return 0;
+}
+
+int32_t EmbLayerNormPlugin::enqueue(
+    nvinfer1::PluginTensorDesc const* inputDesc,
+    nvinfer1::PluginTensorDesc const* outputDesc,
+    void const* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) noexcept {
+  int32_t batchSize = inputDesc[0].dims.d[0];
+  int32_t const maxSeqlen = inputDesc[0].dims.d[1];
+  if (maxSeqlen > 512) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "EmbLayerNormPlugin support maxSeqlen is 512"));
+  }
+  const float* beta = mBetaDev.get();
+  const float* gamma = mGammaDev.get();
+  if (mType == nvinfer1::DataType::kFLOAT) {
+    auto output = static_cast<float*>(outputs[0]);
+    if (nbLookupTables_ == 2) {
+      return embSkipLayerNorm_2<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          maxSeqlen,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          output);
+    } else if (nbLookupTables_ == 3) {
+      return embSkipLayerNorm_3<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          maxSeqlen,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          static_cast<float const*>(mIdsEmbPtrs[2]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          output);
+    } else if (nbLookupTables_ == 4) {
+      return embSkipLayerNorm_4<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          maxSeqlen,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          static_cast<int32_t const*>(inputs[3]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          static_cast<float const*>(mIdsEmbPtrs[2]),
+          static_cast<float const*>(mIdsEmbPtrs[3]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          mIdsVocabSize[3],
+          output);
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only support 2,3,4 lookup_tables fused "));
+    }
+  } else if (mType == nvinfer1::DataType::kHALF) {
+    auto output = static_cast<half*>(outputs[0]);
+    if (nbLookupTables_ == 2) {
+      return embSkipLayerNorm_2<half>(stream,
+                                      static_cast<int32_t>(mLd),
+                                      batchSize,
+                                      maxSeqlen,
+                                      static_cast<int32_t const*>(inputs[0]),
+                                      static_cast<int32_t const*>(inputs[1]),
+                                      nbLookupTables_,
+                                      beta,
+                                      gamma,
+                                      static_cast<half const*>(mIdsEmbPtrs[0]),
+                                      static_cast<half const*>(mIdsEmbPtrs[1]),
+                                      mIdsVocabSize[0],
+                                      mIdsVocabSize[1],
+                                      output);
+    } else if (nbLookupTables_ == 3) {
+      return embSkipLayerNorm_3<half>(stream,
+                                      static_cast<int32_t>(mLd),
+                                      batchSize,
+                                      maxSeqlen,
+                                      static_cast<int32_t const*>(inputs[0]),
+                                      static_cast<int32_t const*>(inputs[1]),
+                                      static_cast<int32_t const*>(inputs[2]),
+                                      nbLookupTables_,
+                                      beta,
+                                      gamma,
+                                      static_cast<half const*>(mIdsEmbPtrs[0]),
+                                      static_cast<half const*>(mIdsEmbPtrs[1]),
+                                      static_cast<half const*>(mIdsEmbPtrs[2]),
+                                      mIdsVocabSize[0],
+                                      mIdsVocabSize[1],
+                                      mIdsVocabSize[2],
+                                      output);
+    } else if (nbLookupTables_ == 4) {
+      return embSkipLayerNorm_4<half>(stream,
+                                      static_cast<int32_t>(mLd),
+                                      batchSize,
+                                      maxSeqlen,
+                                      static_cast<int32_t const*>(inputs[0]),
+                                      static_cast<int32_t const*>(inputs[1]),
+                                      static_cast<int32_t const*>(inputs[2]),
+                                      static_cast<int32_t const*>(inputs[3]),
+                                      nbLookupTables_,
+                                      beta,
+                                      gamma,
+                                      static_cast<half const*>(mIdsEmbPtrs[0]),
+                                      static_cast<half const*>(mIdsEmbPtrs[1]),
+                                      static_cast<half const*>(mIdsEmbPtrs[2]),
+                                      static_cast<half const*>(mIdsEmbPtrs[3]),
+                                      mIdsVocabSize[0],
+                                      mIdsVocabSize[1],
+                                      mIdsVocabSize[2],
+                                      mIdsVocabSize[3],
+                                      output);
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only support 2,3,4 lookup_tables fused "));
+    }
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Unsupported type error, expected [kHALF,kFLOAT]"));
+  }
+  return STATUS_SUCCESS;
+}
+
+// IPluginV2Ext Methods
+nvinfer1::DataType EmbLayerNormPlugin::getOutputDataType(
+    int32_t index,
+    nvinfer1::DataType const* inputTypes,
+    int32_t nbInputs) const noexcept {
+  assert(index == 0);
+  assert(mType == nvinfer1::DataType::kHALF ||
+         mType == nvinfer1::DataType::kFLOAT);
+  return mType;
+}
+
+// IPluginV2 Methods
+char const* EmbLayerNormPlugin::getPluginType() const noexcept {
+  return EMB_LAYER_NORM_NAME;
+}
+
+char const* EmbLayerNormPlugin::getPluginVersion() const noexcept {
+  return EMB_LAYER_NORM_VERSION;
+}
+
+int32_t EmbLayerNormPlugin::getNbOutputs() const noexcept { return 1; }
+
+int32_t EmbLayerNormPlugin::initialize() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormPlugin initialize");
+  return 0;
+}
+
+void EmbLayerNormPlugin::terminate() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormPlugin terminate");
+}
+
+size_t EmbLayerNormPlugin::getSerializationSize() const noexcept {
+  size_t const wordSize = getElementSize(mType);
+  return 2 * sizeof(float) * mLd                            // beta + gamma
+         + sizeof(mType)                                    //
+         + sizeof(mLd)                                      //
+         + mIdsVocabSize.size() * sizeof(mIdsVocabSize[0])  //
+         + wordSize * mLd *
+               accumulate(
+                   mIdsVocabSize.begin(), mIdsVocabSize.end(), 0)  // ids emb
+         + sizeof(nbLookupTables_);  // numbers of lookup_table
+}
+
+void EmbLayerNormPlugin::serialize(void* buffer) const noexcept {
+  serialize_value(&buffer, mType);
+  serialize_value(&buffer, mLd);
+  serialize_value(&buffer, nbLookupTables_);
+  for (size_t i = 0; i < mIdsVocabSize.size(); ++i) {
+    serialize_value(&buffer, mIdsVocabSize[i]);
+  }
+  char* d = static_cast<char*>(buffer);
+  size_t const wordSize = getElementSize(mType);
+  serFromDev(&d, mBetaDev.get(), mLd);
+  serFromDev(&d, mGammaDev.get(), mLd);
+  for (size_t i = 0; i < mIdsEmbPtrs.size(); ++i) {
+    serFromDev(&d,
+               static_cast<char*>(mIdsEmbPtrs[i]),
+               mLd * mIdsVocabSize[i] * wordSize);
+  }
+}
+
+void EmbLayerNormPlugin::destroy() noexcept {
+  // This gets called when the network containing plugin is destroyed
+  mBetaDev.reset(nullptr);
+  mGammaDev.reset(nullptr);
+  for (size_t i = 0; i < mIdsEmbPtrs.size(); ++i) {
+    cudaFree(mIdsEmbPtrs[i]);
+  }
+  delete this;
+}
+
+void EmbLayerNormPlugin::setPluginNamespace(char const* libNamespace) noexcept {
+  mNamespace = libNamespace;
+}
+
+char const* EmbLayerNormPlugin::getPluginNamespace() const noexcept {
+  return mNamespace.c_str();
+}
+
+EmbLayerNormPluginCreator::EmbLayerNormPluginCreator() = default;
+
+char const* EmbLayerNormPluginCreator::getPluginName() const noexcept {
+  return EMB_LAYER_NORM_NAME;
+}
+
+char const* EmbLayerNormPluginCreator::getPluginVersion() const noexcept {
+  return EMB_LAYER_NORM_VERSION;
+}
+
+nvinfer1::PluginFieldCollection const*
+EmbLayerNormPluginCreator::getFieldNames() noexcept {
+  return &mFC;
+}
+
+bool initialize_fields(nvinfer1::PluginFieldCollection const* fc,
+                       nvinfer1::Weights* beta,
+                       nvinfer1::Weights* gamma,
+                       std::vector<nvinfer1::Weights>* IdsEmb) {
+  bool output_fp16 = false;
+  for (int32_t i = 0; i < fc->nbFields; i++) {
+    std::string field_name(fc->fields[i].name);
+    if (field_name.compare("bert_embeddings_layernorm_beta") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
+      beta->values = fc->fields[i].data;
+      beta->count = fc->fields[i].length;
+      beta->type = fieldTypeToDataType(fc->fields[i].type);
+    }
+
+    if (field_name.compare("bert_embeddings_layernorm_gamma") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
+      gamma->values = fc->fields[i].data;
+      gamma->count = fc->fields[i].length;
+      gamma->type = fieldTypeToDataType(fc->fields[i].type);
+    }
+
+    if (field_name.compare("output_fp16") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building output_fp16...");
+      assert(fc->fields[i].type == nvinfer1::PluginFieldType::kINT32);
+      output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
+    }
+    if (field_name.compare("bert_embeddings_word_embeddings_" +
+                           std::to_string(i - 3)) == 0) {
+      TRANSFORMER_DEBUG_MSG(
+          ("bert_embeddings_word_embeddings_" + std::to_string(i - 3)).c_str());
+      nvinfer1::Weights tem;
+      tem.values = fc->fields[i].data;
+      tem.count = fc->fields[i].length;
+      tem.type = fieldTypeToDataType(fc->fields[i].type);
+      IdsEmb->push_back(tem);
+    }
+  }
+  return output_fp16;
+}
+
+nvinfer1::IPluginV2* EmbLayerNormPluginCreator::createPlugin(
+    char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVar createPlugin");
+  nvinfer1::Weights beta;
+  nvinfer1::Weights gamma;
+  std::vector<nvinfer1::Weights> IdsEmb;
+  bool output_fp16 = initialize_fields(fc, &beta, &gamma, &IdsEmb);
+  TRANSFORMER_DEBUG_MSG("Building the Plugin...");
+  EmbLayerNormPlugin* p = new EmbLayerNormPlugin(
+      name,
+      output_fp16 ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
+      beta,
+      gamma,
+      IdsEmb);
+  return p;
+}
+
+nvinfer1::IPluginV2* EmbLayerNormPluginCreator::deserializePlugin(
+    char const* name, void const* serialData, size_t serialLength) noexcept {
+  return new EmbLayerNormPlugin(name, serialData, serialLength);
+}
+
+void EmbLayerNormPluginCreator::setPluginNamespace(
+    char const* libNamespace) noexcept {
+  mNamespace = libNamespace;
+}
+
+char const* EmbLayerNormPluginCreator::getPluginNamespace() const noexcept {
+  return mNamespace.c_str();
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_plugin.h

@@ -0,0 +1,203 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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.
+
+#pragma once
+#include <cuda.h>
+#include "NvInferPlugin.h"
+#include "NvInferRuntime.h"
+
+#include "paddle/fluid/inference/tensorrt/plugin/common/bertCommon.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/plugin.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/serialize.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T>
+int32_t embSkipLayerNorm_2(cudaStream_t,
+                           int32_t,
+                           int32_t,
+                           int32_t,
+                           int32_t const*,
+                           int32_t const*,
+                           int32_t,
+                           float const*,
+                           float const*,
+                           T const*,
+                           T const*,
+                           int32_t,
+                           int32_t,
+                           T*);
+
+template <typename T>
+int32_t embSkipLayerNorm_3(cudaStream_t,
+                           int32_t,
+                           int32_t,
+                           int32_t,
+                           int32_t const*,
+                           int32_t const*,
+                           int32_t const*,
+                           int32_t,
+                           float const*,
+                           float const*,
+                           T const*,
+                           T const*,
+                           T const*,
+                           int32_t,
+                           int32_t,
+                           int32_t,
+                           T*);
+
+template <typename T>
+int32_t embSkipLayerNorm_4(cudaStream_t,
+                           int32_t,
+                           int32_t,
+                           int32_t,
+                           int32_t const*,
+                           int32_t const*,
+                           int32_t const*,
+                           int32_t const*,
+                           int32_t,
+                           float const*,
+                           float const*,
+                           T const*,
+                           T const*,
+                           T const*,
+                           T const*,
+                           int32_t,
+                           int32_t,
+                           int32_t,
+                           int32_t,
+                           T*);
+class EmbLayerNormPlugin : public nvinfer1::IPluginV2DynamicExt {
+ public:
+  EmbLayerNormPlugin(std::string const& name,
+                     nvinfer1::DataType const type,
+                     nvinfer1::Weights const& beta,
+                     nvinfer1::Weights const& gamma,
+                     const std::vector<nvinfer1::Weights>& ids_emb);
+
+  EmbLayerNormPlugin(std::string const& name, void const* data, size_t length);
+
+  EmbLayerNormPlugin() = delete;
+
+  nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int32_t outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int32_t nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+
+  void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in,
+                       int32_t nbInputs,
+                       nvinfer1::DynamicPluginTensorDesc const* out,
+                       int32_t nbOutputs) noexcept override;
+
+  int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
+                  nvinfer1::PluginTensorDesc const* outputDesc,
+                  void const* const* inputs,
+                  void* const* outputs,
+                  void* workspace,
+                  cudaStream_t stream) noexcept override;
+
+  int32_t initialize() noexcept override;
+
+  void terminate() noexcept override;
+
+  char const* getPluginVersion() const noexcept override;
+
+  bool supportsFormatCombination(int32_t pos,
+                                 nvinfer1::PluginTensorDesc const* inOut,
+                                 int32_t nbInputs,
+                                 int32_t nbOutputs) noexcept override;
+
+  size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs,
+                          int32_t nbInputs,
+                          nvinfer1::PluginTensorDesc const* outputs,
+                          int32_t nbOutputs) const noexcept override;
+
+  nvinfer1::DataType getOutputDataType(
+      int32_t index,
+      nvinfer1::DataType const* inputTypes,
+      int32_t nbInputs) const noexcept override;
+
+  char const* getPluginType() const noexcept override;
+
+  int32_t getNbOutputs() const noexcept override;
+
+  size_t getSerializationSize() const noexcept override;
+
+  void serialize(void* buffer) const noexcept override;
+
+  void destroy() noexcept override;
+
+  char const* getPluginNamespace() const noexcept override;
+
+  void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+ protected:
+  std::string const mLayerName;
+  std::string mNamespace;
+  cuda_unique_ptr<float> mGammaDev;
+  cuda_unique_ptr<float> mBetaDev;
+  std::vector<void*> mIdsEmbPtrs;
+  size_t mLd;  // leading dim = hidden size
+  std::vector<int32_t> mIdsVocabSize;
+  WeightsWithOwnership mBeta;
+  WeightsWithOwnership mGamma;
+  nvinfer1::DataType mType;
+  std::vector<nvinfer1::Weights> mIdsEmb_;
+  int32_t nbLookupTables_ = 0;
+};
+
+class EmbLayerNormPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  EmbLayerNormPluginCreator();
+
+  char const* getPluginName() const noexcept override;
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;
+
+  void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+  char const* getPluginNamespace() const noexcept override;
+
+  nvinfer1::IPluginV2* createPlugin(
+      char const* name,
+      const nvinfer1::PluginFieldCollection* fc) noexcept override;
+
+  char const* getPluginVersion() const noexcept override;
+
+  nvinfer1::IPluginV2* deserializePlugin(char const* name,
+                                         void const* serialData,
+                                         size_t serialLength) noexcept override;
+
+ protected:
+  static nvinfer1::PluginFieldCollection mFC;
+  static std::vector<nvinfer1::PluginField> mPluginAttributes;
+  std::string mNamespace;
+};
+
+REGISTER_TRT_PLUGIN_V2(EmbLayerNormPluginCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_varseqlen_kernel_hface.cu

@@ -0,0 +1,477 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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 <cuda.h>
+#include <cassert>
+#include <cstring>
+#include <iostream>
+#include <vector>
+
+#include "NvInfer.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/common.cuh"
+#include "paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_varseqlen_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernelHFace_2(int32_t ld,
+                                          int32_t const* inputIds0,
+                                          int32_t const* inputIds1,
+                                          float const* beta,
+                                          float const* gamma,
+                                          T const* mIdsEmbDev0,
+                                          T const* mIdsEmbDev1,
+                                          int32_t IdsSize0,
+                                          int32_t IdsSize1,
+                                          T* output) {
+  cub::Sum pairSum;
+  int32_t const s = blockIdx.x;
+  int32_t const b = blockIdx.y;
+  int32_t const sumS = inputIds0[b];
+  int32_t const s_b = inputIds0[b + 1] - sumS;
+  if (s >= s_b) {
+    return;  // This CTA has nothing to do
+  }
+  T const rld = T(1.f) / T(ld);
+  int32_t const seqPos = sumS + s;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // 2. load pos/tok/word embeddings and add them together
+  // offset into embeddings is given by wordId * hidden_size
+  int32_t const poffset = blockIdx.x * ld;
+  int32_t const outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    T p(mIdsEmbDev0[poffset + it]);  // pos id
+    T val = p;
+    int32_t const offset = word_id[0] * ld;
+    val += mIdsEmbDev1[offset + it];
+    output[outOffset + it] = val;
+
+    T const rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // 3. layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernelHFace_3(int32_t ld,
+                                          int32_t const* inputIds0,
+                                          int32_t const* inputIds1,
+                                          int32_t const* inputIds2,
+                                          float const* beta,
+                                          float const* gamma,
+                                          T const* mIdsEmbDev0,
+                                          T const* mIdsEmbDev1,
+                                          T const* mIdsEmbDev2,
+                                          int32_t IdsSize0,
+                                          int32_t IdsSize1,
+                                          int32_t IdsSize2,
+                                          T* output) {
+  cub::Sum pairSum;
+  int32_t const s = blockIdx.x;
+  int32_t const b = blockIdx.y;
+  int32_t const sumS = inputIds0[b];
+  int32_t const s_b = inputIds0[b + 1] - sumS;
+  if (s >= s_b) {
+    return;  // This CTA has nothing to do
+  }
+  T const rld = T(1.f) / T(ld);
+  int32_t const seqPos = sumS + s;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds2)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds2)[seqPos] >= IdsSize2) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[1] = static_cast<int32_t const*>(inputIds2)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // 2. load pos/tok/word embeddings and add them together
+  // offset into embeddings is given by wordId * hidden_size
+  int32_t const poffset = blockIdx.x * ld;
+  int32_t const outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    T p(mIdsEmbDev0[poffset + it]);  // pos id
+    T val = p;
+    int32_t const offset0 = word_id[0] * ld;
+    val += mIdsEmbDev1[offset0 + it];
+    int32_t const offset1 = word_id[1] * ld;
+    val += mIdsEmbDev2[offset1 + it];
+    output[outOffset + it] = val;
+
+    T const rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // 3. layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernelHFace_4(int32_t ld,
+                                          int32_t const* inputIds0,
+                                          int32_t const* inputIds1,
+                                          int32_t const* inputIds2,
+                                          int32_t const* inputIds3,
+                                          float const* beta,
+                                          float const* gamma,
+                                          T const* mIdsEmbDev0,
+                                          T const* mIdsEmbDev1,
+                                          T const* mIdsEmbDev2,
+                                          T const* mIdsEmbDev3,
+                                          int32_t IdsSize0,
+                                          int32_t IdsSize1,
+                                          int32_t IdsSize2,
+                                          int32_t IdsSize3,
+                                          T* output) {
+  cub::Sum pairSum;
+  int32_t const s = blockIdx.x;
+  int32_t const b = blockIdx.y;
+  int32_t const sumS = inputIds0[b];
+  int32_t const s_b = inputIds0[b + 1] - sumS;
+  if (s >= s_b) {
+    return;  // This CTA has nothing to do
+  }
+  T const rld = T(1.f) / T(ld);
+  int32_t const seqPos = sumS + s;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds2)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds2)[seqPos] >= IdsSize2) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[1] = static_cast<int32_t const*>(inputIds2)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds3)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds3)[seqPos] >= IdsSize3) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[2] = static_cast<int32_t const*>(inputIds3)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // 2. load pos/tok/word embeddings and add them together
+  // offset into embeddings is given by wordId * hidden_size
+  int32_t const poffset = blockIdx.x * ld;
+  int32_t const outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    T p(mIdsEmbDev0[poffset + it]);  // pos id
+    T val = p;
+    int32_t const offset0 = word_id[0] * ld;
+    val += mIdsEmbDev1[offset0 + it];
+    int32_t const offset1 = word_id[1] * ld;
+    val += mIdsEmbDev2[offset1 + it];
+    int32_t const offset2 = word_id[2] * ld;
+    val += mIdsEmbDev3[offset2 + it];
+    output[outOffset + it] = val;
+
+    T const rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // 3. layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+template <typename T>
+int32_t embSkipLayerNormHFace_2(cudaStream_t stream,
+                                int32_t ld,
+                                int32_t B,
+                                int32_t S,
+                                int const* inputIds0,
+                                int const* inputIds1,
+                                int32_t nbLookupTables,
+                                float const* beta,
+                                float const* gamma,
+                                T const* mIdsEmbDev0,
+                                T const* mIdsEmbDev1,
+                                int32_t IdsSize0,
+                                int32_t IdsSize1,
+                                T* output) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * (nbLookupTables - 1);
+  embLayerNormKernelHFace_2<T, tpb>
+      <<<grid, block, cache_size, stream>>>(ld,
+                                            inputIds0,
+                                            inputIds1,
+                                            beta,
+                                            gamma,
+                                            mIdsEmbDev0,
+                                            mIdsEmbDev1,
+                                            IdsSize0,
+                                            IdsSize1,
+                                            output);
+  return cudaPeekAtLastError();
+}
+
+template <typename T>
+int32_t embSkipLayerNormHFace_3(cudaStream_t stream,
+                                int32_t ld,
+                                int32_t B,
+                                int32_t S,
+                                int const* inputIds0,
+                                int const* inputIds1,
+                                int const* inputIds2,
+                                int32_t nbLookupTables,
+                                float const* beta,
+                                float const* gamma,
+                                T const* mIdsEmbDev0,
+                                T const* mIdsEmbDev1,
+                                T const* mIdsEmbDev2,
+                                int32_t IdsSize0,
+                                int32_t IdsSize1,
+                                int32_t IdsSize2,
+                                T* output) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * (nbLookupTables - 1);
+  embLayerNormKernelHFace_3<T, tpb>
+      <<<grid, block, cache_size, stream>>>(ld,
+                                            inputIds0,
+                                            inputIds1,
+                                            inputIds2,
+                                            beta,
+                                            gamma,
+                                            mIdsEmbDev0,
+                                            mIdsEmbDev1,
+                                            mIdsEmbDev2,
+                                            IdsSize0,
+                                            IdsSize1,
+                                            IdsSize2,
+                                            output);
+  return cudaPeekAtLastError();
+}
+
+template <typename T>
+int32_t embSkipLayerNormHFace_4(cudaStream_t stream,
+                                int32_t ld,
+                                int32_t B,
+                                int32_t S,
+                                int const* inputIds0,
+                                int const* inputIds1,
+                                int const* inputIds2,
+                                int const* inputIds3,
+                                int32_t nbLookupTables,
+                                float const* beta,
+                                float const* gamma,
+                                T const* mIdsEmbDev0,
+                                T const* mIdsEmbDev1,
+                                T const* mIdsEmbDev2,
+                                T const* mIdsEmbDev3,
+                                int32_t IdsSize0,
+                                int32_t IdsSize1,
+                                int32_t IdsSize2,
+                                int32_t IdsSize3,
+                                T* output) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * (nbLookupTables - 1);
+  embLayerNormKernelHFace_4<T, tpb>
+      <<<grid, block, cache_size, stream>>>(ld,
+                                            inputIds0,
+                                            inputIds1,
+                                            inputIds2,
+                                            inputIds3,
+                                            beta,
+                                            gamma,
+                                            mIdsEmbDev0,
+                                            mIdsEmbDev1,
+                                            mIdsEmbDev2,
+                                            mIdsEmbDev3,
+                                            IdsSize0,
+                                            IdsSize1,
+                                            IdsSize2,
+                                            IdsSize3,
+                                            output);
+  return cudaPeekAtLastError();
+}
+
+template int32_t embSkipLayerNormHFace_2<float>(cudaStream_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                int32_t,
+                                                int32_t,
+                                                float*);
+
+template int32_t embSkipLayerNormHFace_3<float>(cudaStream_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                float*);
+
+template int32_t embSkipLayerNormHFace_4<float>(cudaStream_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                float*);
+
+template int32_t embSkipLayerNormHFace_2<half>(cudaStream_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t,
+                                               float const*,
+                                               float const*,
+                                               half const*,
+                                               half const*,
+                                               int32_t,
+                                               int32_t,
+                                               half*);
+
+template int32_t embSkipLayerNormHFace_3<half>(cudaStream_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t,
+                                               float const*,
+                                               float const*,
+                                               half const*,
+                                               half const*,
+                                               half const*,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               half*);
+
+template int32_t embSkipLayerNormHFace_4<half>(cudaStream_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t,
+                                               float const*,
+                                               float const*,
+                                               half const*,
+                                               half const*,
+                                               half const*,
+                                               half const*,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               half*);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_varseqlen_kernel_mtron.cu

@@ -0,0 +1,496 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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 <cuda.h>
+#include <cassert>
+#include <cstring>
+#include <iostream>
+#include <vector>
+
+#include "NvInfer.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/common.cuh"
+#include "paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_varseqlen_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernelMTron_2(int32_t ld,
+                                          int32_t const* inputIds0,
+                                          int32_t const* inputIds1,
+                                          float const* beta,
+                                          float const* gamma,
+                                          T const* mIdsEmbDev0,
+                                          T const* mIdsEmbDev1,
+                                          int32_t IdsSize0,
+                                          int32_t IdsSize1,
+                                          T* output,
+                                          T* skip) {
+  cub::Sum pairSum;
+  int32_t const s = blockIdx.x;
+  int32_t const b = blockIdx.y;
+  int32_t const sumS = inputIds0[b];
+  int32_t const s_b = inputIds0[b + 1] - sumS;
+  if (s >= s_b) {
+    return;  // This CTA has nothing to do
+  }
+  T const rld = T(1.f) / T(ld);
+  const int32_t seqPos = sumS + s;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // 2. load pos/tok/word embeddings and add them together
+  // offset into embeddings is given by wordId * hidden_size
+  const int32_t poffset = blockIdx.x * ld;
+  const int32_t outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    T p(mIdsEmbDev0[poffset + it]);  // pos id
+    T val = p;
+    const int32_t offset = word_id[0] * ld;
+    val += mIdsEmbDev1[offset + it];
+    output[outOffset + it] = val;
+    skip[outOffset + it] = val;
+
+    const T rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // 3. layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernelMTron_3(int32_t ld,
+                                          int32_t const* inputIds0,
+                                          int32_t const* inputIds1,
+                                          int32_t const* inputIds2,
+                                          float const* beta,
+                                          float const* gamma,
+                                          T const* mIdsEmbDev0,
+                                          T const* mIdsEmbDev1,
+                                          T const* mIdsEmbDev2,
+                                          int32_t IdsSize0,
+                                          int32_t IdsSize1,
+                                          int32_t IdsSize2,
+                                          T* output,
+                                          T* skip) {
+  cub::Sum pairSum;
+  const int32_t s = blockIdx.x;
+  const int32_t b = blockIdx.y;
+  const int32_t sumS = inputIds0[b];
+  const int32_t s_b = inputIds0[b + 1] - sumS;
+  if (s >= s_b) {
+    return;  // This CTA has nothing to do
+  }
+  const T rld = T(1.f) / T(ld);
+  const int32_t seqPos = sumS + s;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds2)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds2)[seqPos] >= IdsSize2) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[1] = static_cast<int32_t const*>(inputIds2)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // 2. load pos/tok/word embeddings and add them together
+  // offset into embeddings is given by wordId * hidden_size
+  const int32_t poffset = blockIdx.x * ld;
+  const int32_t outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    T p(mIdsEmbDev0[poffset + it]);  // pos id
+    T val = p;
+    const int32_t offset0 = word_id[0] * ld;
+    val += mIdsEmbDev1[offset0 + it];
+    const int32_t offset1 = word_id[1] * ld;
+    val += mIdsEmbDev2[offset1 + it];
+    output[outOffset + it] = val;
+    skip[outOffset + it] = val;
+
+    const T rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // 3. layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+
+template <typename T, unsigned TPB>
+__global__ void embLayerNormKernelMTron_4(int32_t ld,
+                                          int32_t const* inputIds0,
+                                          int32_t const* inputIds1,
+                                          int32_t const* inputIds2,
+                                          int32_t const* inputIds3,
+                                          float const* beta,
+                                          float const* gamma,
+                                          T const* mIdsEmbDev0,
+                                          T const* mIdsEmbDev1,
+                                          T const* mIdsEmbDev2,
+                                          T const* mIdsEmbDev3,
+                                          int32_t IdsSize0,
+                                          int32_t IdsSize1,
+                                          int32_t IdsSize2,
+                                          int32_t IdsSize3,
+                                          T* output,
+                                          T* skip) {
+  cub::Sum pairSum;
+  const int32_t s = blockIdx.x;
+  const int32_t b = blockIdx.y;
+  const int32_t sumS = inputIds0[b];
+  const int32_t s_b = inputIds0[b + 1] - sumS;
+  if (s >= s_b) {
+    return;  // This CTA has nothing to do
+  }
+  const T rld = T(1.f) / T(ld);
+  const int32_t seqPos = sumS + s;
+  extern __shared__ int32_t word_id[];
+
+  if (threadIdx.x == 0) {
+    if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[0] = static_cast<int32_t const*>(inputIds1)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds2)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds2)[seqPos] >= IdsSize2) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[1] = static_cast<int32_t const*>(inputIds2)[seqPos];
+    }
+
+    if (static_cast<int32_t const*>(inputIds3)[seqPos] < 0 ||
+        static_cast<int32_t const*>(inputIds3)[seqPos] >= IdsSize3) {
+      printf(
+          "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+          "table: ID < 0 or ID > max ");
+      return;
+    } else {
+      word_id[2] = static_cast<int32_t const*>(inputIds3)[seqPos];
+    }
+  }
+  __syncthreads();
+
+  // 2. load pos/tok/word embeddings and add them together
+  // offset into embeddings is given by wordId * hidden_size
+  const int32_t poffset = blockIdx.x * ld;
+  const int32_t outOffset = seqPos * ld;
+  // the output offset is given by b * (S*hidden_size) + s * hidden_size
+  kvp<T> threadData(0, 0);
+
+  for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+    T p(mIdsEmbDev0[poffset + it]);  // pos id
+    T val = p;
+    const int32_t offset0 = word_id[0] * ld;
+    val += mIdsEmbDev1[offset0 + it];
+    const int32_t offset1 = word_id[1] * ld;
+    val += mIdsEmbDev2[offset1 + it];
+    const int32_t offset2 = word_id[2] * ld;
+    val += mIdsEmbDev3[offset2 + it];
+    output[outOffset + it] = val;
+    skip[outOffset + it] = val;
+
+    const T rldval = rld * val;
+    threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+  }
+
+  // 3. layer norm on the sum
+  layerNorm<T, T, float, TPB>(threadData, ld, outOffset, beta, gamma, output);
+}
+template <typename T>
+int32_t embSkipLayerNormMTron_2(cudaStream_t stream,
+                                int32_t ld,
+                                int32_t B,
+                                int32_t S,
+                                int32_t const* inputIds0,
+                                int32_t const* inputIds1,
+                                int32_t nbLookupTables,
+                                float const* beta,
+                                float const* gamma,
+                                T const* mIdsEmbDev0,
+                                T const* mIdsEmbDev1,
+                                int32_t IdsSize0,
+                                int32_t IdsSize1,
+                                T* output,
+                                T* skip) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * (nbLookupTables - 1);
+  embLayerNormKernelMTron_2<T, tpb>
+      <<<grid, block, cache_size, stream>>>(ld,
+                                            inputIds0,
+                                            inputIds1,
+                                            beta,
+                                            gamma,
+                                            mIdsEmbDev0,
+                                            mIdsEmbDev1,
+                                            IdsSize0,
+                                            IdsSize1,
+                                            output,
+                                            skip);
+  return cudaPeekAtLastError();
+}
+
+template <typename T>
+int32_t embSkipLayerNormMTron_3(cudaStream_t stream,
+                                int32_t ld,
+                                int32_t B,
+                                int32_t S,
+                                int32_t const* inputIds0,
+                                int32_t const* inputIds1,
+                                int32_t const* inputIds2,
+                                int32_t nbLookupTables,
+                                float const* beta,
+                                float const* gamma,
+                                T const* mIdsEmbDev0,
+                                T const* mIdsEmbDev1,
+                                T const* mIdsEmbDev2,
+                                int32_t IdsSize0,
+                                int32_t IdsSize1,
+                                int32_t IdsSize2,
+                                T* output,
+                                T* skip) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * (nbLookupTables - 1);
+  embLayerNormKernelMTron_3<T, tpb>
+      <<<grid, block, cache_size, stream>>>(ld,
+                                            inputIds0,
+                                            inputIds1,
+                                            inputIds2,
+                                            beta,
+                                            gamma,
+                                            mIdsEmbDev0,
+                                            mIdsEmbDev1,
+                                            mIdsEmbDev2,
+                                            IdsSize0,
+                                            IdsSize1,
+                                            IdsSize2,
+                                            output,
+                                            skip);
+  return cudaPeekAtLastError();
+}
+
+template <typename T>
+int32_t embSkipLayerNormMTron_4(cudaStream_t stream,
+                                int32_t ld,
+                                int32_t B,
+                                int32_t S,
+                                int32_t const* inputIds0,
+                                int32_t const* inputIds1,
+                                int32_t const* inputIds2,
+                                int32_t const* inputIds3,
+                                int32_t nbLookupTables,
+                                float const* beta,
+                                float const* gamma,
+                                T const* mIdsEmbDev0,
+                                T const* mIdsEmbDev1,
+                                T const* mIdsEmbDev2,
+                                T const* mIdsEmbDev3,
+                                int32_t IdsSize0,
+                                int32_t IdsSize1,
+                                int32_t IdsSize2,
+                                int32_t IdsSize3,
+                                T* output,
+                                T* skip) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * (nbLookupTables - 1);
+  embLayerNormKernelMTron_4<T, tpb>
+      <<<grid, block, cache_size, stream>>>(ld,
+                                            inputIds0,
+                                            inputIds1,
+                                            inputIds2,
+                                            inputIds3,
+                                            beta,
+                                            gamma,
+                                            mIdsEmbDev0,
+                                            mIdsEmbDev1,
+                                            mIdsEmbDev2,
+                                            mIdsEmbDev3,
+                                            IdsSize0,
+                                            IdsSize1,
+                                            IdsSize2,
+                                            IdsSize3,
+                                            output,
+                                            skip);
+  return cudaPeekAtLastError();
+}
+
+template int32_t embSkipLayerNormMTron_2<float>(cudaStream_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                int32_t,
+                                                int32_t,
+                                                float*,
+                                                float*);
+
+template int32_t embSkipLayerNormMTron_3<float>(cudaStream_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                float*,
+                                                float*);
+
+template int32_t embSkipLayerNormMTron_4<float>(cudaStream_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t const*,
+                                                int32_t,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                float const*,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                int32_t,
+                                                float*,
+                                                float*);
+
+template int32_t embSkipLayerNormMTron_2<half>(cudaStream_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t,
+                                               float const*,
+                                               float const*,
+                                               half const*,
+                                               half const*,
+                                               int32_t,
+                                               int32_t,
+                                               half*,
+                                               half*);
+
+template int32_t embSkipLayerNormMTron_3<half>(cudaStream_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t,
+                                               float const*,
+                                               float const*,
+                                               half const*,
+                                               half const*,
+                                               half const*,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               half*,
+                                               half*);
+
+template int32_t embSkipLayerNormMTron_4<half>(cudaStream_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t const*,
+                                               int32_t,
+                                               float const*,
+                                               float const*,
+                                               half const*,
+                                               half const*,
+                                               half const*,
+                                               half const*,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               int32_t,
+                                               half*,
+                                               half*);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_varseqlen_plugin.cu

@@ -0,0 +1,909 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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 "paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_varseqlen_plugin.h"
+#include <cuda.h>
+#include <cstring>
+#include <vector>
+#include "NvInfer.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+constexpr size_t threadsPerCta128 = 2 * 2 * 32;
+constexpr size_t threadsPerCta256 = 1 * 4 * 32;
+constexpr size_t threadsPerCta384 = 1 * 8 * 32;
+// The number of xmmas in the M dimension. We use one uint32_t per XMMA in the M
+// dimension: (s + 16*warps_m - 1) / (16*warps_m);
+constexpr size_t xmmasM128 = 4;
+constexpr size_t xmmasM256 = 16;
+constexpr size_t xmmasM384 = 24;
+// Packed mask size per batch. Layout is XMMAS_M * THREADS_PER_CTA.
+constexpr size_t packedMaskSize128 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize256 = xmmasM256 * threadsPerCta256;
+constexpr size_t packedMaskSize384 = xmmasM384 * threadsPerCta384;
+char const* EMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE{"1"};
+char const* EMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON{"2"};
+char const* EMB_LAYER_NORM_VAR_SEQLEN_NAME{
+    "ManyEmbLayerNormVarlenPluginDynamic"};
+// Static class fields initialization
+nvinfer1::PluginFieldCollection EmbLayerNormVarSeqlenPluginBaseCreator::mFC{};
+std::vector<nvinfer1::PluginField>
+    EmbLayerNormVarSeqlenPluginBaseCreator::mPluginAttributes;
+
+EmbLayerNormVarSeqlenPluginBase::EmbLayerNormVarSeqlenPluginBase(
+    std::string const& name,
+    nvinfer1::DataType const type,
+    nvinfer1::Weights const& beta,
+    nvinfer1::Weights const& gamma,
+    const std::vector<nvinfer1::Weights>& IdsEmb)
+    : mLayerName(name),
+      mLd(beta.count),
+      mType(type),
+      mIdsEmb_(IdsEmb),
+      nbLookupTables_(static_cast<int>(IdsEmb.size())) {
+  // Assuming Weights.count is the number of elements and not bytes
+  assert(beta.count == gamma.count);
+  mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
+  mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
+  copyToDevice(&mGamma, sizeof(float) * mGamma.count, &mGammaDev);
+  copyToDevice(&mBeta, sizeof(float) * mBeta.count, &mBetaDev);
+  for (size_t i = 0; i < mIdsEmb_.size(); ++i) {
+    assert(mIdsEmb_[i].count % mLd == 0);
+    mIdsVocabSize.push_back(int32_t(mIdsEmb_[i].count / mLd));
+    WeightsWithOwnership tem_weight;
+    tem_weight.convertAndCopy(mIdsEmb_[i], mType);
+    void* cudaMem{nullptr};
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaMalloc(&cudaMem, getWeightsSize(tem_weight, mType)));
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpy(cudaMem,
+                                          tem_weight.values,
+                                          getWeightsSize(tem_weight, mType),
+                                          cudaMemcpyHostToDevice));
+    mIdsEmbPtrs.push_back(cudaMem);
+  }
+}
+
+EmbLayerNormVarSeqlenPluginBase::EmbLayerNormVarSeqlenPluginBase(
+    std::string const& name, void const* data, size_t length)
+    : mLayerName(name),
+      mGammaDev(nullptr),
+      mBetaDev(nullptr),
+      mIdsEmbPtrs{},
+      mIdsEmb_{} {
+  // Deserialize in the same order as serialization
+  deserialize_value(&data, &length, &mType);
+  deserialize_value(&data, &length, &mLd);
+  deserialize_value(&data, &length, &nbLookupTables_);
+  for (int32_t i = 0; i < nbLookupTables_; ++i) {
+    int32_t tem;
+    deserialize_value(&data, &length, &tem);
+    mIdsVocabSize.push_back(tem);
+  }
+  char const* d = static_cast<char const*>(data);
+  mBeta.convertAndCopy(&d, mLd, nvinfer1::DataType::kFLOAT);
+  mGamma.convertAndCopy(&d, mLd, nvinfer1::DataType::kFLOAT);
+  for (int32_t i = 0; i < nbLookupTables_; ++i) {
+    nvinfer1::Weights pre_tem_weight;
+    pre_tem_weight.type = mType;
+    pre_tem_weight.count = mLd * size_t(mIdsVocabSize[i]);
+    const auto nbBytes = mLd * size_t(mIdsVocabSize[i]) * getElementSize(mType);
+    auto destBuf = new char[nbBytes];
+    pre_tem_weight.values = destBuf;
+    std::copy_n(d, nbBytes, destBuf);
+    d += nbBytes;
+    mIdsEmb_.push_back(pre_tem_weight);
+  }
+}
+
+EmbLayerNormVarSeqlenPluginHFace::EmbLayerNormVarSeqlenPluginHFace(
+    std::string const& name,
+    nvinfer1::DataType const type,
+    nvinfer1::Weights const& beta,
+    nvinfer1::Weights const& gamma,
+    const std::vector<nvinfer1::Weights>& IdsEmb)
+    : EmbLayerNormVarSeqlenPluginBase(name, type, beta, gamma, IdsEmb) {}
+
+EmbLayerNormVarSeqlenPluginHFace::EmbLayerNormVarSeqlenPluginHFace(
+    std::string const& name, void const* data, size_t length)
+    : EmbLayerNormVarSeqlenPluginBase(name, data, length) {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace deserialize");
+}
+
+EmbLayerNormVarSeqlenPluginMTron::EmbLayerNormVarSeqlenPluginMTron(
+    std::string const& name,
+    nvinfer1::DataType const type,
+    nvinfer1::Weights const& beta,
+    nvinfer1::Weights const& gamma,
+    const std::vector<nvinfer1::Weights>& IdsEmb)
+    : EmbLayerNormVarSeqlenPluginBase(name, type, beta, gamma, IdsEmb) {}
+
+EmbLayerNormVarSeqlenPluginMTron::EmbLayerNormVarSeqlenPluginMTron(
+    std::string const& name, void const* data, size_t length)
+    : EmbLayerNormVarSeqlenPluginBase(name, data, length) {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron deserialize");
+}
+
+// IPluginV2DynamicExt Methods
+nvinfer1::IPluginV2DynamicExt* EmbLayerNormVarSeqlenPluginHFace::clone()
+    const noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace clone");
+  auto p = new EmbLayerNormVarSeqlenPluginHFace(
+      mLayerName, mType, mBeta, mGamma, mIdsEmb_);
+  p->setPluginNamespace(mNamespace.c_str());
+  return p;
+}
+
+nvinfer1::IPluginV2DynamicExt* EmbLayerNormVarSeqlenPluginMTron::clone()
+    const noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron clone");
+  auto p = new EmbLayerNormVarSeqlenPluginMTron(
+      mLayerName, mType, mBeta, mGamma, mIdsEmb_);
+  p->setPluginNamespace(mNamespace.c_str());
+  return p;
+}
+
+nvinfer1::DimsExprs EmbLayerNormVarSeqlenPluginHFace::getOutputDimensions(
+    int32_t outputIndex,
+    nvinfer1::DimsExprs const* inputs,
+    int32_t nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) noexcept {
+  for (int i = 1; i < nbInputs - 1; ++i) {
+    assert(inputs[i].nbDims == 1);                 // seq length
+    assert(inputs[i].nbDims == inputs[1].nbDims);  // same shape
+  }
+  assert(inputs[0].nbDims == 1);  // pos_id: B+1
+  if (outputIndex == 0) {
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 4;
+    ret.d[0] = inputs[1].d[0];  // sum of seq length
+    ret.d[1] = exprBuilder.constant(mLd);
+    ret.d[2] = exprBuilder.constant(1);
+    ret.d[3] = exprBuilder.constant(1);
+    return ret;
+  } else if (outputIndex == 1) {
+    // This is a hack: we just report some mask size and rely the plugins to
+    // play nicely together.
+    //      At runtime, depending on the actual maxSeqlen, the size might be
+    //      different.
+    int32_t maskSize_ = packedMaskSize384;
+    auto maskSize = exprBuilder.constant(maskSize_);
+    auto fp16maskSize =
+        exprBuilder.operation(nvinfer1::DimensionOperation::kPROD,
+                              *maskSize,
+                              *exprBuilder.constant(2));
+    auto Bplus1 = inputs[0].d[0];  // pos_id
+    auto one = exprBuilder.constant(1);
+    auto B = exprBuilder.operation(
+        nvinfer1::DimensionOperation::kSUB, *Bplus1, *one);
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 2;
+    ret.d[0] = B;
+    ret.d[1] = fp16maskSize;
+    return ret;
+  } else {
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 1;
+    ret.d[0] = inputs[nbInputs - 1].d[1];  // mask id: max seqlen
+    return ret;
+  }
+}
+
+nvinfer1::DimsExprs EmbLayerNormVarSeqlenPluginMTron::getOutputDimensions(
+    int32_t outputIndex,
+    nvinfer1::DimsExprs const* inputs,
+    int32_t nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) noexcept {
+  // Input should be input ids and token ids and cumulative seqlens
+  // Output should be the embeddings tensor and mask indices
+  for (int i = 1; i < nbInputs - 1; ++i) {
+    assert(inputs[i].nbDims == 1);                 // seq length
+    assert(inputs[i].nbDims == inputs[1].nbDims);  // same shape
+  }
+  assert(inputs[0].nbDims == 1);  // pos_id: B+1
+  if (outputIndex == 0 || outputIndex == 1) {
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 4;
+    ret.d[0] = inputs[1].d[0];  // sum of seq length
+    ret.d[1] = exprBuilder.constant(mLd);
+    ret.d[2] = exprBuilder.constant(1);
+    ret.d[3] = exprBuilder.constant(1);
+    return ret;
+  } else {
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 1;
+    ret.d[0] = inputs[nbInputs - 1].d[1];  // mask id: max seqlen
+    return ret;
+  }
+}
+
+bool EmbLayerNormVarSeqlenPluginBase::supportsFormatCombination(
+    int32_t pos,
+    nvinfer1::PluginTensorDesc const* inOut,
+    int32_t nbInputs,
+    int32_t nbOutputs) noexcept {
+  assert(nbOutputs == 3);
+  nvinfer1::PluginTensorDesc const& desc = inOut[pos];
+  if (desc.format != nvinfer1::TensorFormat::kLINEAR) {
+    return false;
+  }
+  if (pos == 0) {  // pos_id
+    return desc.dims.nbDims == 1 && desc.type == nvinfer1::DataType::kINT32;
+  }
+  if (pos == 1) {  //  input_id
+    return desc.dims.nbDims == 1 && desc.type == nvinfer1::DataType::kINT32;
+  }
+  nvinfer1::PluginTensorDesc const& prev = inOut[1];  // input_ids
+  if (1 < pos &&
+      pos < (nbInputs - 1)) {  // other ids: check it's the same as input_ids
+    return desc.type == prev.type && desc.dims.nbDims == 1 &&
+           desc.dims.d[0] == prev.dims.d[0];
+  }
+  if (pos == nbInputs - 1) {  // mask id
+    return desc.type == mType;
+  }
+  // embedded sequence
+  if (pos == nbInputs) {
+    return desc.type == mType && desc.dims.nbDims == 4 &&
+           desc.dims.d[0] == inOut[1].dims.d[0] && desc.dims.d[2] == 1 &&
+           desc.dims.d[3] == 1;
+  }
+  // mask(HFace) or pre_layernorm_bias(MTron)
+  if (pos == nbInputs + 1) {
+    return desc.type == mType;
+  }
+  // max seqlen
+  if (pos == nbInputs + 2) {
+    return desc.type == mType;
+  }
+}
+
+void checkConfigurationInputs(nvinfer1::DynamicPluginTensorDesc const* inputs,
+                              int32_t nbInputs,
+                              nvinfer1::DynamicPluginTensorDesc const* outputs,
+                              int32_t nbOutputs) noexcept {
+  // Validate input arguments
+  assert(nbOutputs == 3);
+  assert(inputs[0].desc.dims.nbDims == 1);
+  assert(inputs[0].desc.type == nvinfer1::DataType::kINT32);
+  for (int i = 1; i < nbInputs - 1; ++i) {
+    assert(inputs[i].desc.dims.nbDims == 1);
+    assert(inputs[i].desc.dims.d[0] == inputs[1].desc.dims.d[0]);
+    assert(inputs[i].desc.type == nvinfer1::DataType::kINT32);
+  }
+  assert(outputs[0].desc.dims.nbDims == 4);
+  assert(static_cast<size_t>(outputs[0].desc.dims.d[0]) ==
+         static_cast<size_t>(inputs[1].desc.dims.d[0]));
+  assert(outputs[0].desc.dims.d[2] == 1);
+  assert(outputs[0].desc.dims.d[3] == 1);
+}
+
+void EmbLayerNormVarSeqlenPluginHFace::configurePlugin(
+    nvinfer1::DynamicPluginTensorDesc const* inputs,
+    int32_t nbInputs,
+    nvinfer1::DynamicPluginTensorDesc const* outputs,
+    int32_t nbOutputs) noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace configurePlugin");
+  checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
+  assert(static_cast<size_t>(outputs[0].desc.dims.d[1]) ==
+         static_cast<size_t>(mLd));
+  int32_t const B = inputs[0].desc.dims.d[0] - 1;
+  // check mask
+  assert(outputs[1].desc.dims.nbDims == 2);
+  if (B > 0) {
+    assert(outputs[1].desc.dims.d[0] == B);
+  }
+  assert((outputs[1].desc.dims.d[1] == 2 * packedMaskSize384) ||
+         (outputs[1].desc.dims.d[1] == 2 * packedMaskSize128) ||
+         (outputs[1].desc.dims.d[1] == 2 * packedMaskSize256));
+  assert(outputs[0].desc.type == mType);
+  assert(outputs[1].desc.type == nvinfer1::DataType::kHALF);
+}
+
+void EmbLayerNormVarSeqlenPluginMTron::configurePlugin(
+    nvinfer1::DynamicPluginTensorDesc const* inputs,
+    int32_t nbInputs,
+    nvinfer1::DynamicPluginTensorDesc const* outputs,
+    int32_t nbOutputs) noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron configurePlugin");
+  checkConfigurationInputs(inputs, nbInputs, outputs, nbOutputs);
+  assert(static_cast<size_t>(outputs[0].desc.dims.d[1]) ==
+         static_cast<size_t>(mLd));
+  assert(outputs[1].desc.dims.nbDims == 4);
+  assert(static_cast<size_t>(outputs[1].desc.dims.d[0]) ==
+         static_cast<size_t>(inputs[1].desc.dims.d[0]));
+  assert(static_cast<size_t>(outputs[1].desc.dims.d[1]) ==
+         static_cast<size_t>(mLd));
+  assert(outputs[1].desc.dims.d[2] == 1);
+  assert(outputs[1].desc.dims.d[3] == 1);
+
+  assert(outputs[0].desc.type == mType);
+  assert(outputs[1].desc.type == mType);
+}
+
+size_t EmbLayerNormVarSeqlenPluginBase::getWorkspaceSize(
+    nvinfer1::PluginTensorDesc const* inputs,
+    int32_t nbInputs,
+    nvinfer1::PluginTensorDesc const* outputs,
+    int32_t nbOutputs) const noexcept {
+  return 0;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginHFace::enqueue(
+    nvinfer1::PluginTensorDesc const* inputDesc,
+    nvinfer1::PluginTensorDesc const* outputDesc,
+    void const* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) noexcept {
+  int32_t batchSize = inputDesc[0].dims.d[0] - 1;
+  // read out the maximum sequence length from the dummy input
+  int32_t const maxSeqlen = inputDesc[nbLookupTables_].dims.d[1];
+  int32_t S = 512;
+  if (maxSeqlen <= 128) {
+    S = 128;
+  } else if (maxSeqlen <= 192) {
+    S = 192;
+  } else if (maxSeqlen <= 256) {
+    S = 256;
+  } else if (maxSeqlen <= 384) {
+    S = 384;
+  } else if (maxSeqlen <= 512) {
+    S = 512;
+  } else {
+    std::cerr << "fused_embedding_eltwise_layernorm'max sequence lengths is "
+                 "512 for VarSeqlen"
+              << std::endl;
+  }
+  const float* beta = mBetaDev.get();
+  const float* gamma = mGammaDev.get();
+  if (mType == nvinfer1::DataType::kFLOAT) {
+    auto output = static_cast<float*>(outputs[0]);
+    if (nbLookupTables_ == 2) {
+      return embSkipLayerNormHFace_2<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          output);
+    } else if (nbLookupTables_ == 3) {
+      return embSkipLayerNormHFace_3<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          static_cast<float const*>(mIdsEmbPtrs[2]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          output);
+    } else if (nbLookupTables_ == 4) {
+      return embSkipLayerNormHFace_4<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          static_cast<int32_t const*>(inputs[3]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          static_cast<float const*>(mIdsEmbPtrs[2]),
+          static_cast<float const*>(mIdsEmbPtrs[3]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          mIdsVocabSize[3],
+          output);
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only support 2,3,4 lookup_tables fused "));
+    }
+  } else if (mType == nvinfer1::DataType::kHALF) {
+    auto output = static_cast<half*>(outputs[0]);
+    if (nbLookupTables_ == 2) {
+      return embSkipLayerNormHFace_2<half>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<half const*>(mIdsEmbPtrs[0]),
+          static_cast<half const*>(mIdsEmbPtrs[1]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          output);
+    } else if (nbLookupTables_ == 3) {
+      return embSkipLayerNormHFace_3<half>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<half const*>(mIdsEmbPtrs[0]),
+          static_cast<half const*>(mIdsEmbPtrs[1]),
+          static_cast<half const*>(mIdsEmbPtrs[2]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          output);
+    } else if (nbLookupTables_ == 4) {
+      return embSkipLayerNormHFace_4<half>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          static_cast<int32_t const*>(inputs[3]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<half const*>(mIdsEmbPtrs[0]),
+          static_cast<half const*>(mIdsEmbPtrs[1]),
+          static_cast<half const*>(mIdsEmbPtrs[2]),
+          static_cast<half const*>(mIdsEmbPtrs[3]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          mIdsVocabSize[3],
+          output);
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only support 2,3,4 lookup_tables fused "));
+    }
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Unsupported type error, expected [kHALF,kFLOAT]"));
+  }
+  return STATUS_SUCCESS;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginMTron::enqueue(
+    nvinfer1::PluginTensorDesc const* inputDesc,
+    nvinfer1::PluginTensorDesc const* outputDesc,
+    void const* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) noexcept {
+  int32_t batchSize = inputDesc[0].dims.d[0] - 1;
+  // read out the maximum sequence length from the dummy input
+  int32_t const maxSeqlen = inputDesc[nbLookupTables_].dims.d[1];
+  int32_t S = 512;
+  if (maxSeqlen <= 128) {
+    S = 128;
+  } else if (maxSeqlen <= 192) {
+    S = 192;
+  } else if (maxSeqlen <= 256) {
+    S = 256;
+  } else if (maxSeqlen <= 384) {
+    S = 384;
+  } else if (maxSeqlen <= 512) {
+    S = 512;
+  } else {
+    std::cerr << "fused_embedding_eltwise_layernorm'max sequence lengths is "
+                 "512 for VarSeqlen"
+              << std::endl;
+  }
+  const float* beta = mBetaDev.get();
+  const float* gamma = mGammaDev.get();
+
+  if (mType == nvinfer1::DataType::kFLOAT) {
+    auto output = static_cast<float*>(outputs[0]);
+    auto skip = static_cast<float*>(outputs[1]);
+    if (nbLookupTables_ == 2) {
+      return embSkipLayerNormMTron_2<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          output,
+          skip);
+    } else if (nbLookupTables_ == 3) {
+      return embSkipLayerNormMTron_3<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          static_cast<float const*>(mIdsEmbPtrs[2]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          output,
+          skip);
+    } else if (nbLookupTables_ == 4) {
+      return embSkipLayerNormMTron_4<float>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          static_cast<int32_t const*>(inputs[3]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<float const*>(mIdsEmbPtrs[0]),
+          static_cast<float const*>(mIdsEmbPtrs[1]),
+          static_cast<float const*>(mIdsEmbPtrs[2]),
+          static_cast<float const*>(mIdsEmbPtrs[3]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          mIdsVocabSize[3],
+          output,
+          skip);
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only support 2,3,4 lookup_tables fused "));
+    }
+  } else if (mType == nvinfer1::DataType::kHALF) {
+    auto output = static_cast<half*>(outputs[0]);
+    auto skip = static_cast<half*>(outputs[1]);
+    if (nbLookupTables_ == 2) {
+      return embSkipLayerNormMTron_2<half>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<half const*>(mIdsEmbPtrs[0]),
+          static_cast<half const*>(mIdsEmbPtrs[1]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          output,
+          skip);
+    } else if (nbLookupTables_ == 3) {
+      return embSkipLayerNormMTron_3<half>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<half const*>(mIdsEmbPtrs[0]),
+          static_cast<half const*>(mIdsEmbPtrs[1]),
+          static_cast<half const*>(mIdsEmbPtrs[2]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          output,
+          skip);
+    } else if (nbLookupTables_ == 4) {
+      return embSkipLayerNormMTron_4<half>(
+          stream,
+          static_cast<int32_t>(mLd),
+          batchSize,
+          S,
+          static_cast<int32_t const*>(inputs[0]),
+          static_cast<int32_t const*>(inputs[1]),
+          static_cast<int32_t const*>(inputs[2]),
+          static_cast<int32_t const*>(inputs[3]),
+          nbLookupTables_,
+          beta,
+          gamma,
+          static_cast<half const*>(mIdsEmbPtrs[0]),
+          static_cast<half const*>(mIdsEmbPtrs[1]),
+          static_cast<half const*>(mIdsEmbPtrs[2]),
+          static_cast<half const*>(mIdsEmbPtrs[3]),
+          mIdsVocabSize[0],
+          mIdsVocabSize[1],
+          mIdsVocabSize[2],
+          mIdsVocabSize[3],
+          output,
+          skip);
+    } else {
+      PADDLE_THROW(common::errors::InvalidArgument(
+          "Only support 2,3,4 lookup_tables fused "));
+    }
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "Unsupported type error, expected [kHALF,kFLOAT]"));
+  }
+  return STATUS_SUCCESS;
+}
+
+// IPluginV2Ext Methods
+nvinfer1::DataType EmbLayerNormVarSeqlenPluginBase::getOutputDataType(
+    int32_t index,
+    nvinfer1::DataType const* inputTypes,
+    int32_t nbInputs) const noexcept {
+  assert(index == 0 || index == 1);
+  if (index == 0) {
+    assert(mType == nvinfer1::DataType::kHALF ||
+           mType == nvinfer1::DataType::kFLOAT);
+    return mType;
+  }
+  return nvinfer1::DataType::kHALF;
+}
+
+// IPluginV2 Methods
+char const* EmbLayerNormVarSeqlenPluginBase::getPluginType() const noexcept {
+  return EMB_LAYER_NORM_VAR_SEQLEN_NAME;
+}
+
+char const* EmbLayerNormVarSeqlenPluginHFace::getPluginVersion()
+    const noexcept {
+  return EMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
+}
+
+char const* EmbLayerNormVarSeqlenPluginMTron::getPluginVersion()
+    const noexcept {
+  return EMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginBase::getNbOutputs() const noexcept {
+  return 3;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginHFace::initialize() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace initialize");
+  return 0;
+}
+
+int32_t EmbLayerNormVarSeqlenPluginMTron::initialize() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron initialize");
+  return 0;
+}
+
+void EmbLayerNormVarSeqlenPluginHFace::terminate() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace terminate");
+}
+
+void EmbLayerNormVarSeqlenPluginMTron::terminate() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron terminate");
+}
+
+size_t EmbLayerNormVarSeqlenPluginBase::getSerializationSize() const noexcept {
+  size_t const wordSize = getElementSize(mType);
+  return 2 * sizeof(float) * mLd                            // beta + gamma
+         + sizeof(mType)                                    //
+         + sizeof(mLd)                                      //
+         + mIdsVocabSize.size() * sizeof(mIdsVocabSize[0])  //
+         + wordSize * mLd *
+               accumulate(
+                   mIdsVocabSize.begin(), mIdsVocabSize.end(), 0)  // ids emb
+         + sizeof(nbLookupTables_);  // numbers of lookup_table
+}
+
+void EmbLayerNormVarSeqlenPluginBase::serialize(void* buffer) const noexcept {
+  serialize_value(&buffer, mType);
+  serialize_value(&buffer, mLd);
+  serialize_value(&buffer, nbLookupTables_);
+  for (size_t i = 0; i < mIdsVocabSize.size(); ++i) {
+    serialize_value(&buffer, mIdsVocabSize[i]);
+  }
+  char* d = static_cast<char*>(buffer);
+  size_t const wordSize = getElementSize(mType);
+  serFromDev(&d, mBetaDev.get(), mLd);
+  serFromDev(&d, mGammaDev.get(), mLd);
+  for (size_t i = 0; i < mIdsEmbPtrs.size(); ++i) {
+    serFromDev(&d,
+               static_cast<char*>(mIdsEmbPtrs[i]),
+               mLd * mIdsVocabSize[i] * wordSize);
+  }
+}
+
+void EmbLayerNormVarSeqlenPluginBase::destroy() noexcept {
+  // This gets called when the network containing plugin is destroyed
+  mBetaDev.reset(nullptr);
+  mGammaDev.reset(nullptr);
+  for (size_t i = 0; i < mIdsEmbPtrs.size(); ++i) {
+    cudaFree(mIdsEmbPtrs[i]);
+  }
+  delete this;
+}
+
+void EmbLayerNormVarSeqlenPluginHFace::destroy() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginHFace destroy");
+  EmbLayerNormVarSeqlenPluginBase::destroy();
+}
+
+void EmbLayerNormVarSeqlenPluginMTron::destroy() noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenPluginMTron destroy");
+  EmbLayerNormVarSeqlenPluginBase::destroy();
+}
+
+void EmbLayerNormVarSeqlenPluginBase::setPluginNamespace(
+    char const* libNamespace) noexcept {
+  mNamespace = libNamespace;
+}
+
+char const* EmbLayerNormVarSeqlenPluginBase::getPluginNamespace()
+    const noexcept {
+  return mNamespace.c_str();
+}
+
+EmbLayerNormVarSeqlenPluginBaseCreator::
+    EmbLayerNormVarSeqlenPluginBaseCreator() = default;
+
+char const* EmbLayerNormVarSeqlenPluginBaseCreator::getPluginName()
+    const noexcept {
+  return EMB_LAYER_NORM_VAR_SEQLEN_NAME;
+}
+
+char const* EmbLayerNormVarSeqlenPluginHFaceCreator::getPluginVersion()
+    const noexcept {
+  return EMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
+}
+
+char const* EmbLayerNormVarSeqlenPluginMTronCreator::getPluginVersion()
+    const noexcept {
+  return EMB_LAYER_NORM_VAR_SEQLEN_VERSION_MTRON;
+}
+
+nvinfer1::PluginFieldCollection const*
+EmbLayerNormVarSeqlenPluginBaseCreator::getFieldNames() noexcept {
+  return &mFC;
+}
+
+bool initializeFields(nvinfer1::PluginFieldCollection const* fc,
+                      nvinfer1::Weights* beta,
+                      nvinfer1::Weights* gamma,
+                      std::vector<nvinfer1::Weights>* IdsEmb) {
+  bool output_fp16 = false;
+  for (int32_t i = 0; i < fc->nbFields; i++) {
+    std::string field_name(fc->fields[i].name);
+    if (field_name.compare("bert_embeddings_layernorm_beta") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
+      beta->values = fc->fields[i].data;
+      beta->count = fc->fields[i].length;
+      beta->type = fieldTypeToDataType(fc->fields[i].type);
+    }
+
+    if (field_name.compare("bert_embeddings_layernorm_gamma") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
+      gamma->values = fc->fields[i].data;
+      gamma->count = fc->fields[i].length;
+      gamma->type = fieldTypeToDataType(fc->fields[i].type);
+    }
+
+    if (field_name.compare("output_fp16") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building output_fp16...");
+      assert(fc->fields[i].type == nvinfer1::PluginFieldType::kINT32);
+      output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
+    }
+    if (field_name.compare("bert_embeddings_word_embeddings_" +
+                           std::to_string(i - 3)) == 0) {
+      TRANSFORMER_DEBUG_MSG(
+          ("bert_embeddings_word_embeddings_" + std::to_string(i - 3)).c_str());
+      nvinfer1::Weights tem;
+      tem.values = fc->fields[i].data;
+      tem.count = fc->fields[i].length;
+      tem.type = fieldTypeToDataType(fc->fields[i].type);
+      IdsEmb->push_back(tem);
+    }
+  }
+  return output_fp16;
+}
+
+nvinfer1::IPluginV2* EmbLayerNormVarSeqlenPluginHFaceCreator::createPlugin(
+    char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenHFace createPlugin");
+  nvinfer1::Weights beta;
+  nvinfer1::Weights gamma;
+  std::vector<nvinfer1::Weights> IdsEmb;
+  bool output_fp16 = initializeFields(fc, &beta, &gamma, &IdsEmb);
+  TRANSFORMER_DEBUG_MSG("Building the Plugin...");
+  EmbLayerNormVarSeqlenPluginHFace* p = new EmbLayerNormVarSeqlenPluginHFace(
+      name,
+      output_fp16 ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
+      beta,
+      gamma,
+      IdsEmb);
+  return p;
+}
+
+nvinfer1::IPluginV2* EmbLayerNormVarSeqlenPluginMTronCreator::createPlugin(
+    char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept {
+  TRANSFORMER_DEBUG_MSG("EmbLayerNormVarSeqlenMTron createPlugin");
+  nvinfer1::Weights beta;
+  nvinfer1::Weights gamma;
+  std::vector<nvinfer1::Weights> IdsEmb;
+  bool output_fp16 = initializeFields(fc, &beta, &gamma, &IdsEmb);
+  TRANSFORMER_DEBUG_MSG("Building the Plugin...");
+  EmbLayerNormVarSeqlenPluginMTron* p = new EmbLayerNormVarSeqlenPluginMTron(
+      name,
+      output_fp16 ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
+      beta,
+      gamma,
+      IdsEmb);
+  return p;
+}
+
+nvinfer1::IPluginV2* EmbLayerNormVarSeqlenPluginHFaceCreator::deserializePlugin(
+    char const* name, void const* serialData, size_t serialLength) noexcept {
+  return new EmbLayerNormVarSeqlenPluginHFace(name, serialData, serialLength);
+}
+
+nvinfer1::IPluginV2* EmbLayerNormVarSeqlenPluginMTronCreator::deserializePlugin(
+    char const* name, void const* serialData, size_t serialLength) noexcept {
+  return new EmbLayerNormVarSeqlenPluginMTron(name, serialData, serialLength);
+}
+
+void EmbLayerNormVarSeqlenPluginBaseCreator::setPluginNamespace(
+    char const* libNamespace) noexcept {
+  mNamespace = libNamespace;
+}
+
+char const* EmbLayerNormVarSeqlenPluginBaseCreator::getPluginNamespace()
+    const noexcept {
+  return mNamespace.c_str();
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/many_emb_layernorm_varseqlen_plugin.h

@@ -0,0 +1,338 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION &
+// AFFILIATES. 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.
+
+#pragma once
+#include <cuda.h>
+#include "NvInferPlugin.h"
+#include "NvInferRuntime.h"
+
+#include "paddle/fluid/inference/tensorrt/plugin/common/bertCommon.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/plugin.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/serialize.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T>
+int32_t embSkipLayerNormHFace_2(cudaStream_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t,
+                                float const*,
+                                float const*,
+                                T const*,
+                                T const*,
+                                int32_t,
+                                int32_t,
+                                T*);
+
+template <typename T>
+int32_t embSkipLayerNormHFace_3(cudaStream_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t,
+                                float const*,
+                                float const*,
+                                T const*,
+                                T const*,
+                                T const*,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                T*);
+
+template <typename T>
+int32_t embSkipLayerNormHFace_4(cudaStream_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t,
+                                float const*,
+                                float const*,
+                                T const*,
+                                T const*,
+                                T const*,
+                                T const*,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                T*);
+
+template <typename T>
+int32_t embSkipLayerNormMTron_2(cudaStream_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t,
+                                float const*,
+                                float const*,
+                                T const*,
+                                T const*,
+                                int32_t,
+                                int32_t,
+                                T*,
+                                T*);
+
+template <typename T>
+int32_t embSkipLayerNormMTron_3(cudaStream_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t,
+                                float const*,
+                                float const*,
+                                T const*,
+                                T const*,
+                                T const*,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                T*,
+                                T*);
+
+template <typename T>
+int32_t embSkipLayerNormMTron_4(cudaStream_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t const*,
+                                int32_t,
+                                float const*,
+                                float const*,
+                                T const*,
+                                T const*,
+                                T const*,
+                                T const*,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                int32_t,
+                                T*,
+                                T*);
+class EmbLayerNormVarSeqlenPluginBase : public nvinfer1::IPluginV2DynamicExt {
+ public:
+  EmbLayerNormVarSeqlenPluginBase(
+      std::string const& name,
+      nvinfer1::DataType const type,
+      nvinfer1::Weights const& beta,
+      nvinfer1::Weights const& gamma,
+      const std::vector<nvinfer1::Weights>& ids_emb);
+
+  EmbLayerNormVarSeqlenPluginBase(std::string const& name,
+                                  void const* data,
+                                  size_t length);
+
+  // It doesn't make sense to make EmbLayerNormVarSeqlenPlugin without
+  // arguments, so we delete default constructor.
+  EmbLayerNormVarSeqlenPluginBase() = delete;
+
+  // IPluginV2DynamicExt Methods
+  bool supportsFormatCombination(int32_t pos,
+                                 nvinfer1::PluginTensorDesc const* inOut,
+                                 int32_t nbInputs,
+                                 int32_t nbOutputs) noexcept override;
+  size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs,
+                          int32_t nbInputs,
+                          nvinfer1::PluginTensorDesc const* outputs,
+                          int32_t nbOutputs) const noexcept override;
+
+  // IPluginV2Ext Methods
+  nvinfer1::DataType getOutputDataType(
+      int32_t index,
+      nvinfer1::DataType const* inputTypes,
+      int32_t nbInputs) const noexcept override;
+
+  // IPluginV2 Methods
+  char const* getPluginType() const noexcept override;
+  int32_t getNbOutputs() const noexcept override;
+  size_t getSerializationSize() const noexcept override;
+  void serialize(void* buffer) const noexcept override;
+  void destroy() noexcept override;
+  char const* getPluginNamespace() const noexcept override;
+  void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+ protected:
+  std::string const mLayerName;
+  std::string mNamespace;
+  cuda_unique_ptr<float> mGammaDev;
+  cuda_unique_ptr<float> mBetaDev;
+  std::vector<void*> mIdsEmbPtrs;
+  size_t mLd;  // leading dim = hidden size
+  std::vector<int32_t> mIdsVocabSize;
+  WeightsWithOwnership mBeta;
+  WeightsWithOwnership mGamma;
+  nvinfer1::DataType mType;
+  std::vector<nvinfer1::Weights> mIdsEmb_;
+  int32_t nbLookupTables_ = 0;
+};
+
+class EmbLayerNormVarSeqlenPluginHFace
+    : public EmbLayerNormVarSeqlenPluginBase {
+ public:
+  EmbLayerNormVarSeqlenPluginHFace(
+      std::string const& name,
+      nvinfer1::DataType const type,
+      nvinfer1::Weights const& beta,
+      nvinfer1::Weights const& gamma,
+      const std::vector<nvinfer1::Weights>& ids_emb);
+
+  EmbLayerNormVarSeqlenPluginHFace(std::string const& name,
+                                   void const* data,
+                                   size_t length);
+
+  // It doesn't make sense to make EmbLayerNormVarSeqlenPlugin without
+  // arguments, so we delete default constructor.
+  EmbLayerNormVarSeqlenPluginHFace() = delete;
+
+  // IPluginV2DynamicExt Methods
+  nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int32_t outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int32_t nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+  void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in,
+                       int32_t nbInputs,
+                       nvinfer1::DynamicPluginTensorDesc const* out,
+                       int32_t nbOutputs) noexcept override;
+  int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
+                  nvinfer1::PluginTensorDesc const* outputDesc,
+                  void const* const* inputs,
+                  void* const* outputs,
+                  void* workspace,
+                  cudaStream_t stream) noexcept override;
+  // IPluginV2 Methods
+  int32_t initialize() noexcept override;
+  void terminate() noexcept override;
+  void destroy() noexcept override;
+  char const* getPluginVersion() const noexcept override;
+};
+
+class EmbLayerNormVarSeqlenPluginMTron
+    : public EmbLayerNormVarSeqlenPluginBase {
+ public:
+  EmbLayerNormVarSeqlenPluginMTron(
+      std::string const& name,
+      nvinfer1::DataType const type,
+      nvinfer1::Weights const& beta,
+      nvinfer1::Weights const& gamma,
+      const std::vector<nvinfer1::Weights>& ids_emb);
+
+  EmbLayerNormVarSeqlenPluginMTron(std::string const& name,
+                                   void const* data,
+                                   size_t length);
+
+  // It doesn't make sense to make EmbLayerNormVarSeqlenPlugin without
+  // arguments, so we delete default constructor.
+  EmbLayerNormVarSeqlenPluginMTron() = delete;
+
+  // IPluginV2DynamicExt Methods
+  nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int32_t outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int32_t nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+  void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in,
+                       int32_t nbInputs,
+                       nvinfer1::DynamicPluginTensorDesc const* out,
+                       int32_t nbOutputs) noexcept override;
+  int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
+                  nvinfer1::PluginTensorDesc const* outputDesc,
+                  void const* const* inputs,
+                  void* const* outputs,
+                  void* workspace,
+                  cudaStream_t stream) noexcept override;
+
+  // IPluginV2 Methods
+  int32_t initialize() noexcept override;
+  void terminate() noexcept override;
+  void destroy() noexcept override;
+  char const* getPluginVersion() const noexcept override;
+};
+
+class EmbLayerNormVarSeqlenPluginBaseCreator : public nvinfer1::IPluginCreator {
+ public:
+  EmbLayerNormVarSeqlenPluginBaseCreator();
+
+  char const* getPluginName() const noexcept override;
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;
+
+  void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+  char const* getPluginNamespace() const noexcept override;
+
+ protected:
+  static nvinfer1::PluginFieldCollection mFC;
+  static std::vector<nvinfer1::PluginField> mPluginAttributes;
+  std::string mNamespace;
+};
+
+class EmbLayerNormVarSeqlenPluginHFaceCreator
+    : public EmbLayerNormVarSeqlenPluginBaseCreator {
+ public:
+  nvinfer1::IPluginV2* createPlugin(
+      char const* name,
+      const nvinfer1::PluginFieldCollection* fc) noexcept override;
+  char const* getPluginVersion() const noexcept override;
+  nvinfer1::IPluginV2* deserializePlugin(char const* name,
+                                         void const* serialData,
+                                         size_t serialLength) noexcept override;
+};
+
+class EmbLayerNormVarSeqlenPluginMTronCreator
+    : public EmbLayerNormVarSeqlenPluginBaseCreator {
+ public:
+  nvinfer1::IPluginV2* createPlugin(
+      char const* name,
+      const nvinfer1::PluginFieldCollection* fc) noexcept override;
+  char const* getPluginVersion() const noexcept override;
+  nvinfer1::IPluginV2* deserializePlugin(char const* name,
+                                         void const* serialData,
+                                         size_t serialLength) noexcept override;
+};
+REGISTER_TRT_PLUGIN_V2(EmbLayerNormVarSeqlenPluginHFaceCreator);
+REGISTER_TRT_PLUGIN_V2(EmbLayerNormVarSeqlenPluginMTronCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/matmul_op_int8_plugin.h

@@ -0,0 +1,445 @@
+/* Copyright (c) 2021 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. */
+
+#pragma once
+#include <cassert>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/phi/backends/dynload/cublasLt.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class MatmulPlugin : public nvinfer1::IPluginV2IOExt {
+ public:
+  MatmulPlugin(nvinfer1::Dims const& dims_x,
+               nvinfer1::Dims const& dims_y,
+               bool transA,
+               bool transB,
+               float alpha)
+      : dims_x_(dims_x),
+        dims_y_(dims_y),
+        transB_(transA),
+        transA_(transB),
+        alpha_(alpha) {}
+
+  MatmulPlugin(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &dims_x_);
+    DeserializeValue(&serial_data, &serial_length, &dims_y_);
+    DeserializeValue(&serial_data, &serial_length, &transB_);
+    DeserializeValue(&serial_data, &serial_length, &transA_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_scale_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_one_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_zero_);
+    DeserializeValue(&serial_data, &serial_length, &batch_);
+    DeserializeValue(&serial_data, &serial_length, &k_);
+    DeserializeValue(&serial_data, &serial_length, &m_);
+    DeserializeValue(&serial_data, &serial_length, &n_);
+    DeserializeValue(&serial_data, &serial_length, &cublas_);
+    DeserializeValue(&serial_data, &serial_length, &type_);
+    DeserializeValue(&serial_data, &serial_length, &Adesc_);
+    DeserializeValue(&serial_data, &serial_length, &Bdesc_);
+    DeserializeValue(&serial_data, &serial_length, &Cdesc_);
+    DeserializeValue(&serial_data, &serial_length, &AtransformDesc_);
+    DeserializeValue(&serial_data, &serial_length, &BtransformDesc_);
+    DeserializeValue(&serial_data, &serial_length, &CtransformDesc_);
+    DeserializeValue(&serial_data, &serial_length, &Atransform_);
+    DeserializeValue(&serial_data, &serial_length, &Btransform_);
+    DeserializeValue(&serial_data, &serial_length, &Ctransform_);
+    DeserializeValue(&serial_data, &serial_length, &transformDescT_);
+    DeserializeValue(&serial_data, &serial_length, &transformDescN_);
+    DeserializeValue(&serial_data, &serial_length, &matmulDesc_);
+  }
+
+  virtual bool isOutputBroadcastAcrossBatch(int32_t output_index,
+                                            const bool* input_is_broadcasted,
+                                            int32_t nb_inputs) const
+      TRT_NOEXCEPT {
+    return false;
+  }
+
+  virtual bool canBroadcastInputAcrossBatch(int32_t input_index) const
+      TRT_NOEXCEPT {
+    return false;
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  size_t getWorkspaceSize(int) const TRT_NOEXCEPT override { return 0; }
+
+  void setPluginNamespace(const char* plugin_namespace) TRT_NOEXCEPT override {
+    name_space_ = plugin_namespace;
+  }
+
+  nvinfer1::IPluginV2IOExt* clone() const TRT_NOEXCEPT override {
+    MatmulPlugin* ptr =
+        new MatmulPlugin(dims_x_, dims_y_, transB_, transA_, alpha_);
+    ptr->setPluginNamespace(this->getPluginNamespace());
+    ptr->batch_ = batch_;
+    ptr->k_ = k_;
+    ptr->m_ = m_;
+    ptr->n_ = n_;
+    ptr->alpha_scale_ = alpha_scale_;
+    ptr->alpha_one_ = alpha_one_;
+    ptr->alpha_zero_ = alpha_zero_;
+    ptr->cublas_ = cublas_;
+    ptr->type_ = type_;
+    ptr->Adesc_ = Adesc_;
+    ptr->Bdesc_ = Bdesc_;
+    ptr->Cdesc_ = Cdesc_;
+    ptr->AtransformDesc_ = AtransformDesc_;
+    ptr->BtransformDesc_ = BtransformDesc_;
+    ptr->CtransformDesc_ = CtransformDesc_;
+    ptr->Atransform_ = Atransform_;
+    ptr->Btransform_ = Btransform_;
+    ptr->Ctransform_ = Ctransform_;
+    ptr->transformDescT_ = transformDescT_;
+    ptr->transformDescN_ = transformDescN_;
+    ptr->matmulDesc_ = matmulDesc_;
+    return ptr;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return name_space_.c_str();
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "matmul_int8_plugin";
+  }
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* input_dims,
+                                     int num_inputs) TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int32_t pos,
+                                 nvinfer1::PluginTensorDesc const* inOut,
+                                 int32_t nbInputs,
+                                 int32_t nbOutputs) const TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::PluginTensorDesc* in,
+                       int32_t nbInputs,
+                       const nvinfer1::PluginTensorDesc* out,
+                       int32_t nbOutputs) TRT_NOEXCEPT override;
+
+  /*
+    bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+        const TRT_NOEXCEPT override;
+  */
+  int initialize() TRT_NOEXCEPT { return 0; }
+  void terminate() TRT_NOEXCEPT;
+
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+  void attachToContext(cudnnContext* cudnnContext,
+                       cublasContext* cublasContext,
+                       nvinfer1::IGpuAllocator* gpuAllocator)
+      TRT_NOEXCEPT override;
+  void detachFromContext() TRT_NOEXCEPT override;
+
+ protected:
+  nvinfer1::Dims dims_x_;
+  nvinfer1::Dims dims_y_;
+  bool transB_;
+  bool transA_;
+  float alpha_;
+  void *alpha_scale_{nullptr}, *alpha_one_{nullptr}, *alpha_zero_{nullptr};
+  int batch_;
+  uint64_t k_;
+  uint64_t m_;
+  uint64_t n_;
+  cublasLtHandle_t cublas_{nullptr};
+  nvinfer1::DataType type_;
+  cublasLtMatrixLayout_t Adesc_{nullptr}, Bdesc_{nullptr}, Cdesc_{nullptr};
+  cublasLtMatrixLayout_t AtransformDesc_{nullptr}, BtransformDesc_{nullptr},
+      CtransformDesc_{nullptr};
+  int8_t *Atransform_{nullptr}, *Btransform_{nullptr}, *Ctransform_{nullptr};
+  cublasLtMatrixTransformDesc_t transformDescT_{nullptr},
+      transformDescN_{nullptr};
+  cublasLtMatmulDesc_t matmulDesc_{nullptr};
+  std::string name_space_;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(dims_x_) + SerializedSize(dims_y_) +
+           SerializedSize(transB_) + SerializedSize(transA_) +
+           SerializedSize(alpha_) + SerializedSize(alpha_scale_) +
+           SerializedSize(alpha_one_) + SerializedSize(alpha_zero_) +
+           SerializedSize(batch_) + SerializedSize(k_) + SerializedSize(m_) +
+           SerializedSize(n_) + SerializedSize(cublas_) +
+           SerializedSize(type_) + SerializedSize(Adesc_) +
+           SerializedSize(Bdesc_) + SerializedSize(Cdesc_) +
+           SerializedSize(AtransformDesc_) + SerializedSize(BtransformDesc_) +
+           SerializedSize(CtransformDesc_) + SerializedSize(Atransform_) +
+           SerializedSize(Btransform_) + SerializedSize(Ctransform_) +
+           SerializedSize(transformDescT_) + SerializedSize(transformDescN_) +
+           SerializedSize(matmulDesc_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, dims_x_);
+    SerializeValue(&buffer, dims_y_);
+    SerializeValue(&buffer, transB_);
+    SerializeValue(&buffer, transA_);
+    SerializeValue(&buffer, alpha_);
+    SerializeValue(&buffer, alpha_scale_);
+    SerializeValue(&buffer, alpha_one_);
+    SerializeValue(&buffer, alpha_zero_);
+    SerializeValue(&buffer, batch_);
+    SerializeValue(&buffer, k_);
+    SerializeValue(&buffer, m_);
+    SerializeValue(&buffer, n_);
+    SerializeValue(&buffer, cublas_);
+    SerializeValue(&buffer, type_);
+    SerializeValue(&buffer, Adesc_);
+    SerializeValue(&buffer, Bdesc_);
+    SerializeValue(&buffer, Cdesc_);
+    SerializeValue(&buffer, AtransformDesc_);
+    SerializeValue(&buffer, BtransformDesc_);
+    SerializeValue(&buffer, CtransformDesc_);
+    SerializeValue(&buffer, Atransform_);
+    SerializeValue(&buffer, Btransform_);
+    SerializeValue(&buffer, Ctransform_);
+    SerializeValue(&buffer, transformDescT_);
+    SerializeValue(&buffer, transformDescN_);
+    SerializeValue(&buffer, matmulDesc_);
+  }
+};
+
+class MatmulPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  MatmulPluginCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "matmul_int8_plugin";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2IOExt* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2IOExt* deserializePlugin(const char* name,
+                                              void const* serial_data,
+                                              size_t serial_length)
+      TRT_NOEXCEPT override {
+    MatmulPlugin* obj = new MatmulPlugin(serial_data, serial_length);
+    obj->setPluginNamespace(name);
+    return obj;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+REGISTER_TRT_PLUGIN_V2(MatmulPluginCreator);
+
+class MatmulPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  MatmulPluginDynamic(bool transA, bool transB, float alpha)
+      : transB_(transA), transA_(transB), alpha_(alpha) {}
+
+  MatmulPluginDynamic(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &transB_);
+    DeserializeValue(&serial_data, &serial_length, &transA_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_scale_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_one_);
+    DeserializeValue(&serial_data, &serial_length, &alpha_zero_);
+    DeserializeValue(&serial_data, &serial_length, &cublas_);
+    DeserializeValue(&serial_data, &serial_length, &Atransform_);
+    DeserializeValue(&serial_data, &serial_length, &Btransform_);
+    DeserializeValue(&serial_data, &serial_length, &Ctransform_);
+    DeserializeValue(&serial_data, &serial_length, &type_);
+  }
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    MatmulPluginDynamic* ptr =
+        new MatmulPluginDynamic(transB_, transA_, alpha_);
+    ptr->setPluginNamespace(this->getPluginNamespace());
+    ptr->alpha_scale_ = alpha_scale_;
+    ptr->alpha_one_ = alpha_one_;
+    ptr->alpha_zero_ = alpha_zero_;
+    ptr->cublas_ = cublas_;
+    ptr->Atransform_ = Atransform_;
+    ptr->Btransform_ = Btransform_;
+    ptr->Ctransform_ = Ctransform_;
+    ptr->type_ = type_;
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "matmul_int8_dynamic_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+
+  int initialize() TRT_NOEXCEPT { return 0; }
+  void terminate() TRT_NOEXCEPT;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* inputs,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* outputs,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  void attachToContext(cudnnContext* cudnnContext,
+                       cublasContext* cublasContext,
+                       nvinfer1::IGpuAllocator* gpuAllocator)
+      TRT_NOEXCEPT override;
+
+  void detachFromContext() TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ protected:
+  bool transB_;
+  bool transA_;
+  float alpha_;
+  void *alpha_scale_{nullptr}, *alpha_one_{nullptr}, *alpha_zero_{nullptr};
+  cublasLtHandle_t cublas_{nullptr};
+  nvinfer1::DataType type_;
+  int8_t *Atransform_{nullptr}, *Btransform_{nullptr}, *Ctransform_{nullptr};
+  std::string name_space_;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(transB_) + SerializedSize(transA_) +
+           SerializedSize(alpha_) + SerializedSize(alpha_scale_) +
+           SerializedSize(alpha_one_) + SerializedSize(alpha_zero_) +
+           SerializedSize(Atransform_) + SerializedSize(Btransform_) +
+           SerializedSize(Ctransform_) + SerializedSize(cublas_) +
+           SerializedSize(type_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, transB_);
+    SerializeValue(&buffer, transA_);
+    SerializeValue(&buffer, alpha_);
+    SerializeValue(&buffer, alpha_scale_);
+    SerializeValue(&buffer, alpha_one_);
+    SerializeValue(&buffer, alpha_zero_);
+    SerializeValue(&buffer, Atransform_);
+    SerializeValue(&buffer, Btransform_);
+    SerializeValue(&buffer, Ctransform_);
+    SerializeValue(&buffer, cublas_);
+    SerializeValue(&buffer, type_);
+  }
+};
+
+class MatmulPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  MatmulPluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "matmul_int8_dynamic_plugin";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         void const* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    MatmulPluginDynamic* obj =
+        new MatmulPluginDynamic(serial_data, serial_length);
+    obj->setPluginNamespace(name);
+    return obj;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+REGISTER_TRT_PLUGIN_V2(MatmulPluginDynamicCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/merge_layernorm_op_plugin.cu

@@ -0,0 +1,325 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+   Copyright (c) 2019-2022, NVIDIA CORPORATION.  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 <algorithm>
+
+#include "paddle/fluid/inference/tensorrt/plugin/merge_layernorm_op_plugin.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+#define FINAL_MASK 0xffffffff
+
+template <typename T>
+__global__ void merge_layernorm_v2(T *out,
+                                   const T *__restrict input,
+                                   const T *__restrict gamma,
+                                   const T *__restrict beta,
+                                   const float layernorm_eps,
+                                   int batch,
+                                   int H,
+                                   int W,
+                                   int n) {
+  // input is [batch, 2*H, 2*W, n/4]
+  // output is [batch, H, W, n]
+  // grid (W, H, batch)
+  // block (n)
+  const int kIte = 4;
+  const int tid = threadIdx.x;
+  const int W_idx = blockIdx.x;
+  const int H_idx = blockIdx.y;
+  const size_t batch_offset = blockIdx.z * H * W * n;
+  const int input_H_stride = W * n / 2;
+  const int output_H_stride = W * n;
+  const int n_4 = n >> 2;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  float local_out[kIte];
+  float sum = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kIte; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      int part_id = col_id / n_4;
+      int offset_in_W = part_id / 2;
+      int offset_in_H = part_id % 2;
+      size_t input_id = batch_offset +
+                        (2 * H_idx + offset_in_H) * input_H_stride +
+                        (2 * W_idx + offset_in_W) * n_4 + (col_id % n_4);
+      local_out[i] = static_cast<float>(__ldg(input + input_id));
+      sum += local_out[i];
+    }
+  }
+
+  mean = phi::funcs::BlockReduceSum<float>(sum, FINAL_MASK);
+  if (tid == 0) {
+    s_mean = mean / n;
+  }
+  __syncthreads();
+
+  float var = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kIte; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      local_out[i] = local_out[i] - s_mean;
+      var += local_out[i] * local_out[i];
+    }
+  }
+
+  variance = phi::funcs::BlockReduceSum<float>(var, FINAL_MASK);
+  if (tid == 0) {
+    s_variance = rsqrtf(variance / n + layernorm_eps);
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (int i = 0; i < kIte; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      size_t output_idx =
+          batch_offset + H_idx * output_H_stride + W_idx * n + col_id;
+      out[output_idx] =
+          static_cast<T>(local_out[i] * s_variance *
+                             static_cast<float>(__ldg(&gamma[col_id])) +
+                         static_cast<float>(__ldg(&beta[col_id])));
+    }
+  }
+}
+
+template <typename T>
+void invokeMergeLayernorm(T *output,
+                          const T *input,
+                          const T *gamma,
+                          const T *beta,
+                          float layernorm_eps,
+                          int batch,
+                          int H,
+                          int W,
+                          int n,
+                          cudaStream_t stream) {
+  if ((W % 2 != 0) || (H % 2 != 0)) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "H(W) of merge layernorm should be a multiple of 2."));
+  }
+  dim3 grid(W / 2, H / 2, batch);
+  int blockSize = (n + 31) / 32 * 32;
+  merge_layernorm_v2<T><<<grid, blockSize, 0, stream>>>(
+      output, input, gamma, beta, layernorm_eps, batch, H / 2, W / 2, n * 4);
+}
+
+template void invokeMergeLayernorm<float>(float *output,
+                                          const float *input,
+                                          const float *gamma,
+                                          const float *beta,
+                                          float layernorm_eps,
+                                          int batch,
+                                          int H,
+                                          int W,
+                                          int n,
+                                          cudaStream_t stream);
+
+template void invokeMergeLayernorm<half>(half *output,
+                                         const half *input,
+                                         const half *gamma,
+                                         const half *beta,
+                                         float layernorm_eps,
+                                         int batch,
+                                         int H,
+                                         int W,
+                                         int n,
+                                         cudaStream_t stream);
+
+template <typename T>
+static void convertAndCopy(const std::vector<float> &host, T *dev) {
+  T *host_ptr = new T[host.size()];
+  std::transform(host.begin(), host.end(), host_ptr, [](float x) {
+    return static_cast<T>(x);
+  });
+  cudaMemcpy(dev, host_ptr, sizeof(T) * host.size(), cudaMemcpyHostToDevice);
+  delete host_ptr;
+}
+
+void MergeLayernormPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+MergeLayernormPluginDynamic::MergeLayernormPluginDynamic(
+    const float *bias_d,
+    const size_t bias_num,
+    const float *scale_d,
+    const size_t scale_num,
+    const float eps,
+    const int begin_norm_axis,
+    const bool with_fp16,
+    std::shared_ptr<void> bias_device,
+    std::shared_ptr<void> scale_device)
+    : eps_(eps),
+      begin_norm_axis_(begin_norm_axis),
+      with_fp16_(with_fp16),
+      bias_device_(bias_device),
+      scale_device_(scale_device) {
+  bias_.resize(bias_num);
+  scale_.resize(scale_num);
+  std::copy(bias_d, bias_d + bias_num, bias_.data());
+  std::copy(scale_d, scale_d + scale_num, scale_.data());
+  int type_size = with_fp16_ ? sizeof(half) : sizeof(float);
+  if (bias_device_ == nullptr) {
+    void *p;
+    cudaMalloc(&p, bias_num * type_size);
+    bias_device_.reset(p, [](void *ptr) { cudaFree(ptr); });
+
+    if (with_fp16) {
+      convertAndCopy<half>(bias_, reinterpret_cast<half *>(p));
+    } else {
+      convertAndCopy<float>(bias_, reinterpret_cast<float *>(p));
+    }
+  }
+  if (scale_device_ == nullptr) {
+    void *p;
+    cudaMalloc(&p, scale_num * type_size);
+    scale_device_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16) {
+      convertAndCopy<half>(scale_, reinterpret_cast<half *>(p));
+    } else {
+      convertAndCopy<float>(scale_, reinterpret_cast<float *>(p));
+    }
+  }
+}
+
+bool MergeLayernormPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument("The input of MergeLayernorm "
+                                      "plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return in.type == nvinfer1::DataType::kHALF &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    } else {
+      return in.type == nvinfer1::DataType::kFLOAT &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType MergeLayernormPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The MergeLayernorm only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+nvinfer1::DimsExprs MergeLayernormPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = inputs[0].d[0];
+  ret.d[1] = expr_builder.operation(nvinfer1::DimensionOperation::kFLOOR_DIV,
+                                    *inputs[0].d[1],
+                                    *expr_builder.constant(4));
+  ret.d[2] = expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
+                                    *inputs[0].d[2],
+                                    *expr_builder.constant(4));
+  return ret;
+}
+
+int MergeLayernormPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  auto input_type = input_desc[0].type;
+  int batch = input_dims.d[0];
+  int input_resolution = static_cast<int>(std::sqrt(input_dims.d[1]));
+  int dim = static_cast<int>(input_dims.d[2]);
+  PADDLE_ENFORCE_EQ(
+      input_resolution * input_resolution,
+      input_dims.d[1],
+      common::errors::InvalidArgument(
+          "The MergeLayernorm TRT Plugin get invalid input_resolution %d",
+          input_resolution));
+
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(3) << "TRT Plugin DataType selected. MergeLayernorm-->fp32";
+    invokeMergeLayernorm<float>(
+        reinterpret_cast<float *>(outputs[0]),
+        reinterpret_cast<const float *>(inputs[0]),
+        reinterpret_cast<const float *>(scale_device_.get()),
+        reinterpret_cast<const float *>(bias_device_.get()),
+        eps_,
+        batch,
+        input_resolution,
+        input_resolution,
+        dim,
+        stream);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(3) << "TRT Plugin DataType selected. MergeLayernorm-->fp16";
+    invokeMergeLayernorm<half>(
+        reinterpret_cast<half *>(outputs[0]),
+        reinterpret_cast<const half *>(inputs[0]),
+        reinterpret_cast<const half *>(scale_device_.get()),
+        reinterpret_cast<const half *>(bias_device_.get()),
+        eps_,
+        batch,
+        input_resolution,
+        input_resolution,
+        dim,
+        stream);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The MergeLayernorm TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/merge_layernorm_op_plugin.h

@@ -0,0 +1,141 @@
+/* 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. */
+
+#pragma once
+#include <memory>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+class MergeLayernormPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  MergeLayernormPluginDynamic(const float* bias_d,
+                              const size_t bias_num,
+                              const float* scale_d,
+                              const size_t scale_num,
+                              const float eps,
+                              const int begin_norm_axis,
+                              const bool with_fp16,
+                              std::shared_ptr<void> bias_device = nullptr,
+                              std::shared_ptr<void> scale_device = nullptr);
+
+  MergeLayernormPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &begin_norm_axis_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new MergeLayernormPluginDynamic(bias_.data(),
+                                           bias_.size(),
+                                           scale_.data(),
+                                           scale_.size(),
+                                           eps_,
+                                           begin_norm_axis_,
+                                           with_fp16_,
+                                           bias_device_,
+                                           scale_device_);
+  }
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "merge_layernorm_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(bias_) + SerializedSize(scale_) +
+           SerializedSize(eps_) + SerializedSize(begin_norm_axis_) +
+           SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, begin_norm_axis_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  std::vector<float> bias_;
+  std::vector<float> scale_;
+  float eps_;
+  int begin_norm_axis_;
+  bool with_fp16_;
+  std::shared_ptr<void> bias_device_ = nullptr;
+  std::shared_ptr<void> scale_device_ = nullptr;
+};
+class MergeLayernormPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "merge_layernorm_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new MergeLayernormPluginDynamic(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(MergeLayernormPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/mish_op_plugin.cu

@@ -0,0 +1,250 @@
+// Copyright (c) 2021 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 <cstring>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/mish_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+int MishPlugin::initialize() TRT_NOEXCEPT { return 0; }
+
+bool MishPlugin::supportsFormat(
+    nvinfer1::DataType type, nvinfer1::PluginFormat format) const TRT_NOEXCEPT {
+  if (with_fp16_) {
+    return ((type == nvinfer1::DataType::kFLOAT ||
+             type == nvinfer1::DataType::kHALF) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  } else {
+    return ((type == nvinfer1::DataType::kFLOAT) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  }
+}
+
+nvinfer1::Dims MishPlugin::getOutputDimensions(int index,
+                                               const nvinfer1::Dims* in_dims,
+                                               int nb_inputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      nb_inputs,
+      1,
+      common::errors::InvalidArgument("We expect [number of inputs] == 1"
+                                      "in TRT Mish op plugin, but got "
+                                      "[number of inputs] = %d.",
+                                      nb_inputs));
+  PADDLE_ENFORCE_LT(
+      index,
+      this->getNbOutputs(),
+      common::errors::InvalidArgument("We expect [index] < [number of outputs]"
+                                      "in TRT Mish op plugin, but got "
+                                      "[index] = %d, [number of outputs] = %d.",
+                                      index,
+                                      this->getNbOutputs()));
+  nvinfer1::Dims const& input_dims = in_dims[0];
+  nvinfer1::Dims output_dims = input_dims;
+  return output_dims;
+}
+
+template <typename T>
+__device__ T kTanh(T x) {
+  return tanh(x);
+}
+
+template <>
+__device__ half kTanh<half>(half x) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const float tmp = tanhf(__half2float(x));
+  return __float2half(tmp);
+#endif
+}
+
+template <typename T>
+__device__ T kSoftplus(T x, T threshold) {
+  return x > threshold ? x : log(exp(x) + static_cast<T>(1.0f));
+}
+
+template <>
+__device__ half kSoftplus<half>(half x, half threshold) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  return x > threshold ? x : hlog(hexp(x) + static_cast<half>(1.0f));
+#endif
+}
+
+template <typename T>
+__global__ void mish_kernel(float threshold, int n, const T* input, T* output) {
+  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    const T in = input[idx];
+    output[idx] = in * kTanh<T>(kSoftplus<T>(in, static_cast<T>(threshold)));
+  }
+}
+
+template <>
+__global__ void mish_kernel<half>(float threshold,
+                                  int n,
+                                  const half* input,
+                                  half* output) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const int idx = blockIdx.x * blockDim.x + threadIdx.x;
+  if (idx < n) {
+    const half in = input[idx];
+    output[idx] =
+        in * kTanh<half>(kSoftplus<half>(in, static_cast<half>(threshold)));
+  }
+#endif
+}
+
+int MishPlugin::enqueue(int batchSize,
+                        const void* const* inputs,
+                        void* const* outputs,
+                        void* workspace,
+                        cudaStream_t stream) TRT_NOEXCEPT {
+  const auto& input_dims = this->getInputDims(0);
+  int num = batchSize;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    num *= input_dims.d[i];
+  }
+
+  const int block_size = 256;
+  const int grid_size = (num + block_size - 1) / block_size;
+
+  auto type = getDataType();
+  if (type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Mish-->fp32";
+    const float* input = static_cast<const float*>(inputs[0]);
+    float* output = static_cast<float*>(outputs[0]);
+    mish_kernel<float>
+        <<<grid_size, block_size, 0, stream>>>(threshold_, num, input, output);
+  } else if (type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Mish-->fp16";
+    const half* input = static_cast<const half*>(inputs[0]);
+    half* output = static_cast<half*>(outputs[0]);
+    mish_kernel<half>
+        <<<grid_size, block_size, 0, stream>>>(threshold_, num, input, output);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The Mish TRT Plugin's input type should be float or half."));
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+// Dynamic Plugin below.
+int MishPluginDynamic::initialize() TRT_NOEXCEPT {
+  getPluginNamespace();
+  return 0;
+}
+
+size_t MishPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(threshold_) + SerializedSize(with_fp16_);
+}
+
+void MishPluginDynamic::serialize(void* buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, threshold_);
+  SerializeValue(&buffer, with_fp16_);
+}
+
+nvinfer1::DimsExprs MishPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  return inputs[0];
+}
+
+bool MishPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of mish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc& in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kFLOAT ||
+              in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc& prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType MishPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Mish Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+int MishPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  size_t num = ProductDim(input_dims);
+  const int block_size = 256;
+  const int grid_size = (num + block_size - 1) / block_size;
+
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Mish-->fp32";
+    const float* input = static_cast<const float*>(inputs[0]);
+    float* output = static_cast<float*>(outputs[0]);
+    mish_kernel<float>
+        <<<grid_size, block_size, 0, stream>>>(threshold_, num, input, output);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Mish-->fp16";
+    const half* input = static_cast<const half*>(inputs[0]);
+    half* output = static_cast<half*>(outputs[0]);
+    mish_kernel<half>
+        <<<grid_size, block_size, 0, stream>>>(threshold_, num, input, output);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The Mish TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/mish_op_plugin.h

@@ -0,0 +1,185 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+#include <stdio.h>
+
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class MishPlugin : public PluginTensorRT {
+ private:
+  float threshold_;
+
+ protected:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(threshold_);
+  }
+
+  // TRT will call this func  to serialize the configuration of TRT
+  // It should not be called by users.
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, threshold_);
+  }
+
+ public:
+  explicit MishPlugin(const float threshold, const bool with_fp16)
+      : threshold_(threshold) {
+    with_fp16_ = with_fp16;
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  MishPlugin(void const* serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &threshold_);
+  }
+
+  ~MishPlugin() {}
+  MishPlugin* clone() const TRT_NOEXCEPT override {
+    return new MishPlugin(threshold_, with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "mish_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override;
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+  int enqueue(int batchSize,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+};
+
+class MishPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "mish_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new MishPlugin(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(MishPluginCreator);
+
+class MishPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit MishPluginDynamic(const float threshold, const bool with_fp16)
+      : threshold_(threshold) {
+    with_fp16_ = with_fp16;
+  }
+  MishPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &threshold_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new MishPluginDynamic(threshold_, with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "mish_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  float threshold_;
+};
+
+class MishPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "mish_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new MishPluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(MishPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/multihead_matmul_roformer_plugin.cu

@@ -0,0 +1,360 @@
+// 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 "paddle/fluid/inference/tensorrt/plugin/multihead_matmul_roformer_plugin.h"
+#include <stdio.h>
+#include <cassert>
+#include <cub/cub.cuh>  // NOLINT
+#include <vector>
+#include "glog/logging.h"
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/common.cuh"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin_utils.h"
+#include "paddle/phi/core/platform/device_context.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/multihead_matmul_functor.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+int MultiheadMatmulRoformerPlugin::initialize() TRT_NOEXCEPT { return 0; }
+
+nvinfer1::DimsExprs MultiheadMatmulRoformerPlugin::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  // input[0], (B, S, 3 * N * H, 1, 1)
+  // input[1], (B, head_num, seq_len, seq_len)
+  // output, (B, seq_len, hidden)
+  PADDLE_ENFORCE_EQ(output_index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "There is only one output of the EmbEltwiseLayernorm, "
+                        "so the index should be zero,"
+                        "but it's (%d)",
+                        output_index));
+  PADDLE_ENFORCE_EQ(
+      nb_inputs,
+      4,
+      common::errors::InvalidArgument(
+          "The Input of the EmbEltwiseLayernorm should be 3, but we found "
+          "it has (%d) inputs",
+          nb_inputs));
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = inputs[0].d[0];
+  ret.d[1] = inputs[0].d[1];
+  ret.d[2] = expr_builder.constant(head_size_ * head_number_);
+  return ret;
+}
+
+bool MultiheadMatmulRoformerPlugin::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kFLOAT ||
+              in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+
+  if (pos == 1) {
+    return in.type == prev.type && in.format == prev.format;
+  }
+
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType MultiheadMatmulRoformerPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      index,
+      0,
+      common::errors::InvalidArgument(
+          "The EmbEltwiseLayernorm Plugin only has one input, so the "
+          "index value should be 0, but get %d.",
+          index));
+  return input_types[0];
+}
+
+template <typename T>
+__global__ void apply_scale(T *data, T scale, int n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (tid < n) {
+    data[tid] = data[tid] * scale;
+  }
+#endif
+}
+
+template <typename T>
+__global__ void RotrayKernel(const T *inputact,
+                             const T *input1,
+                             const T *input2,
+                             T *output,
+                             const int nElement,
+                             const int lastdim) {
+  const int index = blockIdx.x * blockDim.x + threadIdx.x;
+  if (index >= nElement) return;
+  T left_elemul_out = input1[index] * inputact[index];
+  int col = index % lastdim;
+  int half_lastdim = lastdim / 2;
+  const int right_index = index - col + (col + half_lastdim) % lastdim;
+  output[index] = left_elemul_out + input2[index] * inputact[right_index];
+}
+
+inline int round_up(int seq_len, int multiple = 32) {
+  PADDLE_ENFORCE_GT(
+      multiple,
+      0,
+      common::errors::InvalidArgument(
+          "multiple should be a positive number, but it's (%d)", multiple));
+  return ((seq_len + multiple - 1) / multiple) * multiple;
+}
+
+template <typename T>
+__global__ void broadcast(const T *src,
+                          T *dst,
+                          const int seq_len,
+                          const int head_num) {
+  int batch_id = blockIdx.x / (head_num * seq_len);
+  int dst_offset = blockIdx.x * seq_len;
+  if (threadIdx.x < seq_len) {
+    dst[threadIdx.x + dst_offset] = src[threadIdx.x + batch_id * seq_len];
+  }
+}
+
+int MultiheadMatmulRoformerPlugin::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 {
+  auto input_dims = input_desc[0].dims;
+  int input_num = ProductDim(input_dims);
+  // input[0], (B, S, 3 * N * H, 1, 1)
+  int batch = input_dims.d[0];
+  int seq_len = input_dims.d[1];
+  phi::DenseTensor multihead_temp_tensor;
+  // masks
+  int scratch_size = batch * head_number_ * seq_len * seq_len * 1;
+
+  int device_id;
+  cudaGetDevice(&device_id);
+  multihead_temp_tensor.Resize({scratch_size + input_num});
+  // for roformer
+  phi::DenseTensor temp_roformer_tensor;
+  temp_roformer_tensor.Resize({input_num});
+
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. RoformerQkvToContext-->fp32";
+    auto *multihead_temp_data =
+        multihead_temp_tensor.mutable_data<float>(GPUPlace(device_id));
+    auto *temp_roformer_data =
+        temp_roformer_tensor.mutable_data<float>(  // NOLINT
+            GPUPlace(device_id));
+    auto *tmp_roformer_ptr = reinterpret_cast<float *>(temp_roformer_data);
+    auto *qkptr = multihead_temp_data;
+    auto *tptr = multihead_temp_data + scratch_size;
+
+    const float *input0_data = static_cast<const float *>(inputs[0]);
+    // fit to [batch, head_num, length, length] + [batch, 1, 1, length]
+    phi::DenseTensor temp_qk_bias_tensor;
+    float *qk_bias = const_cast<float *>(static_cast<const float *>(inputs[3]));
+    if (ProductDim(input_desc[3].dims) == (batch * seq_len)) {
+      temp_qk_bias_tensor.Resize({batch, head_number_, seq_len, seq_len});
+      auto *temp_qk_bias =
+          temp_qk_bias_tensor.mutable_data<float>(GPUPlace(device_id));
+      int grid = batch * head_number_ * seq_len;
+      int block = round_up(seq_len);
+      broadcast<<<grid, block, 0, stream>>>(
+          static_cast<const float *>(inputs[3]),
+          temp_qk_bias,
+          seq_len,
+          head_number_);
+      qk_bias = temp_qk_bias;
+    }
+    const float *input3_data = static_cast<const float *>(qk_bias);
+    // BxSx3xNxH => tptr: 3xBxNxSxH.
+    TransposeQKV(
+        batch, seq_len, head_size_, head_number_, input0_data, tptr, stream);
+    cudaMemcpy(tmp_roformer_ptr,  // dst
+               tptr,              // src
+               input_num * sizeof(float),
+               cudaMemcpyDeviceToDevice);
+    int n_q = seq_len * head_number_ * head_size_ * batch;
+    constexpr int threads = 128;
+    int blocks = (n_q + threads - 1) / threads;
+    const float *input_cos_data = static_cast<const float *>(inputs[1]);
+    const float *input_sin_data = static_cast<const float *>(inputs[2]);
+    RotrayKernel<<<blocks, threads, 0, stream>>>(tmp_roformer_ptr,
+                                                 input_cos_data,
+                                                 input_sin_data,
+                                                 tptr,
+                                                 n_q,
+                                                 head_size_);  // q
+    RotrayKernel<<<blocks, threads, 0, stream>>>(tmp_roformer_ptr + n_q,
+                                                 input_cos_data,
+                                                 input_sin_data,
+                                                 tptr + n_q,
+                                                 n_q,
+                                                 head_size_);  // k
+
+    auto *device_ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(GPUPlace(device_id)));
+
+    const phi::GPUContext &dev_ctx = *device_ctx;
+    phi::funcs::MultiheadGPUComputeFunctor<float> multihead_compute_func;
+    multihead_compute_func(dev_ctx,
+                           batch,
+                           seq_len,
+                           head_number_,
+                           head_size_,
+                           qkptr,
+                           input3_data,
+                           false,
+                           tptr,
+                           scale_,
+                           static_cast<float>(0.0));
+
+    int grid = batch * head_number_ * seq_len;
+    int block = head_size_;
+    float *output = static_cast<float *>(outputs[0]);
+    transpose<float><<<grid, block, 0, stream>>>(
+        tptr, output, batch, seq_len, head_number_, head_size_);
+
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. QkvToContext-->fp16";
+    auto *multihead_temp_data =
+        multihead_temp_tensor.mutable_data<int16_t>(  // NOLINT
+            GPUPlace(device_id));
+
+    auto *temp_roformer_data =
+        temp_roformer_tensor.mutable_data<int16_t>(  // NOLINT
+            GPUPlace(device_id));
+    half *tmp_roformer_ptr = reinterpret_cast<half *>(temp_roformer_data);
+    half *qkptr = reinterpret_cast<half *>(multihead_temp_data);
+    half *tptr = qkptr + scratch_size;
+
+    const half *input0_data = static_cast<const half *>(inputs[0]);
+    // fit to [batch, head_num, length, length] + [batch, 1, 1, length]
+    phi::DenseTensor temp_qk_bias_tensor;
+    half *qk_bias = const_cast<half *>(static_cast<const half *>(inputs[3]));
+    if (ProductDim(input_desc[3].dims) == (batch * seq_len)) {
+      temp_qk_bias_tensor.Resize({batch, head_number_, seq_len, seq_len});
+      auto *temp_qk_bias = reinterpret_cast<half *>(
+          temp_qk_bias_tensor.mutable_data<int16_t>(GPUPlace(device_id)));
+      int grid = batch * head_number_ * seq_len;
+      int block = round_up(seq_len);
+      broadcast<<<grid, block, 0, stream>>>(
+          static_cast<const half *>(inputs[3]),
+          temp_qk_bias,
+          seq_len,
+          head_number_);
+      qk_bias = temp_qk_bias;
+    }
+    const half *input3_data = static_cast<const half *>(qk_bias);
+    // BxSx3xNxH => tptr: 3xBxNxSxH.
+    TransposeQKV(
+        batch, seq_len, head_size_, head_number_, input0_data, tptr, stream);
+    cudaMemcpy(tmp_roformer_ptr,
+               tptr,
+               input_num * sizeof(half),
+               cudaMemcpyDeviceToDevice);
+
+    auto *device_ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(GPUPlace(device_id)));
+
+    int n_q = seq_len * head_number_ * head_size_ * batch;
+    constexpr int threads = 128;
+    int blocks = (n_q + threads - 1) / threads;
+
+    const half *input_cos_data = static_cast<const half *>(inputs[1]);
+    const half *input_sin_data = static_cast<const half *>(inputs[2]);
+    RotrayKernel<<<blocks, threads, 0, stream>>>(tmp_roformer_ptr,
+                                                 input_cos_data,
+                                                 input_sin_data,
+                                                 tptr,
+                                                 n_q,
+                                                 head_size_);  // q
+    RotrayKernel<<<blocks, threads, 0, stream>>>(tmp_roformer_ptr + n_q,
+                                                 input_cos_data,
+                                                 input_sin_data,
+                                                 tptr + n_q,
+                                                 n_q,
+                                                 head_size_);  // k
+
+    apply_scale<<<blocks, threads, 0, stream>>>(
+        tptr, static_cast<half>(scale_), n_q);
+
+    const phi::GPUContext &dev_ctx = *device_ctx;
+    phi::funcs::MultiheadGPUComputeFunctor<half> multihead_compute_func;
+    multihead_compute_func(dev_ctx,
+                           batch,
+                           seq_len,
+                           head_number_,
+                           head_size_,
+                           qkptr,
+                           input3_data,
+                           false,
+                           tptr,
+                           half(1.),
+                           half(0.0));
+
+    int grid = batch * head_number_ * seq_len;
+    int block = head_size_;
+    half *output = static_cast<half *>(outputs[0]);
+    transpose<half><<<grid, block, 0, stream>>>(
+        tptr, output, batch, seq_len, head_number_, head_size_);
+  } else {
+    PADDLE_THROW(common::errors::Fatal(
+        "The QKV TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/multihead_matmul_roformer_plugin.h

@@ -0,0 +1,161 @@
+// 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class MultiheadMatmulRoformerPlugin : public DynamicPluginTensorRT {
+ public:
+  explicit MultiheadMatmulRoformerPlugin(
+      int hidden, int head_number, int head_size, float scale, bool with_fp16)
+      : hidden_(hidden),
+        head_number_(head_number),
+        head_size_(head_size),
+        scale_(scale) {
+    with_fp16_ = with_fp16;
+  }
+
+  MultiheadMatmulRoformerPlugin(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &hidden_);
+    DeserializeValue(&serial_data, &serial_length, &head_number_);
+    DeserializeValue(&serial_data, &serial_length, &head_size_);
+    DeserializeValue(&serial_data, &serial_length, &scale_);
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new MultiheadMatmulRoformerPlugin(
+        hidden_, head_number_, head_size_, scale_, with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "multihead_matmul_roformer_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(hidden_) + SerializedSize(head_number_) +
+           SerializedSize(head_size_) + SerializedSize(scale_) +
+           SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, hidden_);
+    SerializeValue(&buffer, head_number_);
+    SerializeValue(&buffer, head_size_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nb_inputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nb_outputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  int hidden_;
+  int head_number_;
+  int head_size_;
+  float scale_;
+};
+
+class MultiheadMatmulRoformerPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  MultiheadMatmulRoformerPluginCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "multihead_matmul_roformer_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new MultiheadMatmulRoformerPlugin(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_;
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+REGISTER_TRT_PLUGIN_V2(MultiheadMatmulRoformerPluginCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/pool3d_op_plugin.cu

@@ -0,0 +1,719 @@
+// Copyright (c) 2021 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, softwarepool
+// 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 "paddle/fluid/inference/tensorrt/plugin/pool3d_op_plugin.h"
+#include "paddle/phi/kernels/funcs/pooling.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+size_t Pool3DPlugin::getSerializationSize() const TRT_NOEXCEPT {
+  return getBaseSerializationSize() + SerializedSize(ceil_mode_) +
+         SerializedSize(pool3d_type_) + SerializedSize(adaptive_) +
+         SerializedSize(ksize_) + SerializedSize(strides_) +
+         SerializedSize(paddings_) + SerializedSize(input_shape_) +
+         SerializedSize(output_shape_);
+}
+
+// TRT will call this func when we need to serialize the configuration of
+// tensorrt.
+void Pool3DPlugin::serialize(void *buffer) const TRT_NOEXCEPT {
+  serializeBase(buffer);
+  SerializeValue(&buffer, ceil_mode_);
+  SerializeValue(&buffer, pool3d_type_);
+  SerializeValue(&buffer, adaptive_);
+  SerializeValue(&buffer, ksize_);
+  SerializeValue(&buffer, strides_);
+  SerializeValue(&buffer, paddings_);
+  SerializeValue(&buffer, input_shape_);
+  SerializeValue(&buffer, output_shape_);
+}
+
+Pool3DPlugin *Pool3DPlugin::clone() const TRT_NOEXCEPT {
+  return new Pool3DPlugin(ceil_mode_,
+                          pool3d_type_,
+                          adaptive_,
+                          ksize_,
+                          strides_,
+                          paddings_,
+                          input_shape_);
+}
+
+const char *Pool3DPlugin::getPluginType() const TRT_NOEXCEPT {
+  return "pool3d_plugin";
+}
+
+int Pool3DPlugin::getNbOutputs() const TRT_NOEXCEPT { return 1; }
+
+int Pool3DPlugin::initialize() TRT_NOEXCEPT { return 0; }
+
+nvinfer1::DataType Pool3DPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  return input_types[0];
+}
+
+void Pool3DPlugin::destroy() TRT_NOEXCEPT { delete this; }
+
+nvinfer1::Dims Pool3DPlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims *inputDims, int nbInputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nbInputs,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The Pool3D Plugin only has one input, so the nbInputs "
+                        "value should be 1, but get %d.",
+                        nbInputs));
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Pool3D Plugin only has one input, so "
+                        "the index value should be 0, but get %d.",
+                        index));
+  PADDLE_ENFORCE_EQ(inputDims[0].nbDims,
+                    4,
+                    common::errors::InvalidArgument(
+                        "The Pool3D Plugin only has four Dimensions, so the "
+                        "nbDims value should be 4, but get %d.",
+                        inputDims[0].nbDims));
+
+  nvinfer1::Dims const &input_dims = inputDims[0];
+
+  nvinfer1::Dims output_dims = input_dims;
+
+  output_dims.d[1] = output_shape_[1];
+  output_dims.d[2] = output_shape_[2];
+  output_dims.d[3] = output_shape_[3];
+  return output_dims;
+}
+
+int Pool3DPlugin::enqueue(int batchSize,
+                          const void *const *inputs,
+                          void *const *outputs,
+                          void *workspace,
+                          cudaStream_t stream) TRT_NOEXCEPT {
+  int input_size = 0;
+  float const *idata = reinterpret_cast<float const *>(inputs[0]);
+  float *const *odatas = reinterpret_cast<float *const *>(outputs);
+
+  std::vector<int> input_shape = input_shape_;
+  std::vector<int> output_shape = output_shape_;
+  input_shape.insert(input_shape.begin(), batchSize);
+  output_shape.insert(output_shape.begin(), batchSize);
+
+  if (pool3d_type_ == Pool3DType::max) {
+    phi::funcs::MaxPool<float> pool_process;
+    phi::funcs::Pool3dDirectCUDAFunctor<phi::funcs::MaxPool<float>, float>
+        pool3d_forward;
+    pool3d_forward(idata,
+                   input_shape,
+                   output_shape,
+                   ksize_,
+                   strides_,
+                   paddings_,
+                   true,
+                   adaptive_,
+                   odatas[0],
+                   stream,
+                   pool_process);
+  } else if (pool3d_type_ == Pool3DType::avg) {
+    phi::funcs::AvgPool<float> pool_process;
+    phi::funcs::Pool3dDirectCUDAFunctor<phi::funcs::AvgPool<float>, float>
+        pool3d_forward;
+    pool3d_forward(idata,
+                   input_shape,
+                   output_shape,
+                   ksize_,
+                   strides_,
+                   paddings_,
+                   true,
+                   adaptive_,
+                   odatas[0],
+                   stream,
+                   pool_process);
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+// Dynamic Plugin below.
+
+Pool3DPluginDynamic::Pool3DPluginDynamic(void const *serialData,
+                                         size_t serialLength) {
+  DeserializeValue(&serialData, &serialLength, &ceil_mode_);
+  const char *pool3d_type;
+  DeserializeValue(&serialData, &serialLength, &pool3d_type);
+  pool3d_type_ = std::string(pool3d_type);
+  DeserializeValue(&serialData, &serialLength, &adaptive_);
+  DeserializeValue(&serialData, &serialLength, &ksize_);
+  DeserializeValue(&serialData, &serialLength, &strides_);
+  DeserializeValue(&serialData, &serialLength, &paddings_);
+  DeserializeValue(&serialData, &serialLength, &is_global_);
+}
+
+nvinfer1::IPluginV2DynamicExt *Pool3DPluginDynamic::clone() const TRT_NOEXCEPT {
+  return new Pool3DPluginDynamic(ceil_mode_,
+                                 pool3d_type_,
+                                 adaptive_,
+                                 ksize_,
+                                 strides_,
+                                 paddings_,
+                                 is_global_);
+}
+
+const char *Pool3DPluginDynamic::getPluginType() const TRT_NOEXCEPT {
+  return "pool3d_plugin_dynamic";
+}
+int Pool3DPluginDynamic::getNbOutputs() const TRT_NOEXCEPT { return 1; }
+
+int Pool3DPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+void Pool3DPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+size_t Pool3DPluginDynamic::getWorkspaceSize(
+    const nvinfer1::PluginTensorDesc *inputs,
+    int nbInputs,
+    const nvinfer1::PluginTensorDesc *outputs,
+    int nbOutputs) const TRT_NOEXCEPT {
+  return 0;
+}
+
+size_t Pool3DPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(ceil_mode_) + SerializedSize(pool3d_type_.c_str()) +
+         SerializedSize(adaptive_) + SerializedSize(ksize_) +
+         SerializedSize(strides_) + SerializedSize(paddings_) +
+         SerializedSize(is_global_);
+}
+
+void Pool3DPluginDynamic::serialize(void *buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, ceil_mode_);
+  SerializeValue(&buffer, pool3d_type_.c_str());
+  SerializeValue(&buffer, adaptive_);
+  SerializeValue(&buffer, ksize_);
+  SerializeValue(&buffer, strides_);
+  SerializeValue(&buffer, paddings_);
+  SerializeValue(&buffer, is_global_);
+}
+
+nvinfer1::DimsExprs Pool3DPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nb_inputs,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The Split plugin should be only one input."));
+
+  PADDLE_ENFORCE_EQ(
+      inputs[0].d[1]->isConstant(),
+      true,
+      common::errors::InvalidArgument("The channel dimension should be "
+                                      "static, but we found it's dynamic."));
+  nvinfer1::DimsExprs output(inputs[0]);
+  if (is_global_) {
+    output.d[2] = expr_builder.constant(1);
+    output.d[3] = expr_builder.constant(1);
+    output.d[4] = expr_builder.constant(1);
+    return output;
+  }
+  if (adaptive_) {
+    output.d[2] = expr_builder.constant(ksize_[0]);
+    output.d[3] = expr_builder.constant(ksize_[1]);
+    output.d[4] = expr_builder.constant(ksize_[2]);
+    return output;
+  }
+
+  auto stri_0 = expr_builder.constant(strides_[0]);
+  auto stri_1 = expr_builder.constant(strides_[1]);
+  auto stri_2 = expr_builder.constant(strides_[2]);
+  auto one_value = expr_builder.constant(1);
+
+  auto v0_tmp = expr_builder.constant(-ksize_[0] + 2 * paddings_[0]);
+  auto v1_tmp = expr_builder.constant(-ksize_[1] + 2 * paddings_[1]);
+  auto v2_tmp = expr_builder.constant(-ksize_[2] + 2 * paddings_[2]);
+
+  auto ceil_tmp =
+      expr_builder.constant(-ksize_[0] + 2 * paddings_[0] + strides_[0] - 1);
+  auto ceil1_tmp =
+      expr_builder.constant(-ksize_[1] + 2 * paddings_[1] + strides_[1] - 1);
+  auto ceil2_tmp =
+      expr_builder.constant(-ksize_[2] + 2 * paddings_[2] + strides_[2] - 1);
+
+  if (!ceil_mode_) {
+    output.d[2] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[2], *v0_tmp),
+            *stri_0),
+        *one_value);
+    output.d[3] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[3], *v1_tmp),
+            *stri_1),
+        *one_value);
+    output.d[4] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[4], *v2_tmp),
+            *stri_2),
+        *one_value);
+
+  } else {
+    output.d[2] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[2], *ceil_tmp),
+            *stri_0),
+        *one_value);
+    output.d[3] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
+                                    *inputs[0].d[3],
+                                    *ceil1_tmp),
+            *stri_1),
+        *one_value);
+    output.d[4] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
+                                    *inputs[0].d[4],
+                                    *ceil2_tmp),
+            *stri_2),
+        *one_value);
+  }
+
+  return output;
+}
+
+bool Pool3DPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  return ((in_out[pos].type == nvinfer1::DataType::kFLOAT) &&
+          in_out[pos].format == nvinfer1::PluginFormat::kLINEAR);
+}
+
+nvinfer1::DataType Pool3DPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Pool3D Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  PADDLE_ENFORCE_EQ(
+      (input_types[0] == nvinfer1::DataType::kFLOAT),
+      true,
+      common::errors::InvalidArgument("The input type should be float"));
+  return input_types[0];
+}
+
+int Pool3DPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  int n = input_dims.d[0];
+  int c = input_dims.d[1];
+  int d = input_dims.d[2];
+  int h = input_dims.d[3];
+  int w = input_dims.d[4];
+
+  const float *input = static_cast<const float *>(inputs[0]);
+  float *output = static_cast<float *>(outputs[0]);
+
+  std::vector<int> input_shape, output_shape;
+  for (int i = 0; i < input_dims.nbDims; i++)
+    input_shape.push_back(input_dims.d[i]);
+  output_shape = input_shape;
+
+  std::vector<int> ksize = ksize_;
+  std::vector<int> paddings = paddings_;
+  if (is_global_) {
+    ksize[0] = d;
+    ksize[1] = h;
+    ksize[2] = w;
+    paddings[0] = 0;
+    paddings[1] = 0;
+    paddings[2] = 0;
+    output_shape[2] = 1;
+    output_shape[3] = 1;
+    output_shape[4] = 1;
+  } else {
+    auto data_dim = CalcOutputSize(
+        {d, h, w}, ceil_mode_, adaptive_, ksize_, strides_, paddings_);
+    output_shape[2] = data_dim[0];
+    output_shape[3] = data_dim[1];
+    output_shape[4] = data_dim[2];
+  }
+
+  if (pool3d_type_ == "max") {
+    phi::funcs::MaxPool<float> pool_process;
+    phi::funcs::Pool3dDirectCUDAFunctor<phi::funcs::MaxPool<float>, float>
+        pool3d_forward;
+    pool3d_forward(input,
+                   input_shape,
+                   output_shape,
+                   ksize,
+                   strides_,
+                   paddings,
+                   true,
+                   adaptive_,
+                   output,
+                   stream,
+                   pool_process);
+  } else if (pool3d_type_ == "avg") {
+    phi::funcs::AvgPool<float> pool_process;
+    phi::funcs::Pool3dDirectCUDAFunctor<phi::funcs::AvgPool<float>, float>
+        pool3d_forward;
+    pool3d_forward(input,
+                   input_shape,
+                   output_shape,
+                   ksize,
+                   strides_,
+                   paddings,
+                   true,
+                   adaptive_,
+                   output,
+                   stream,
+                   pool_process);
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+PIRPool3DPluginDynamic::PIRPool3DPluginDynamic(void const *serialData,
+                                               size_t serialLength) {
+  DeserializeValue(&serialData, &serialLength, &ceil_mode_);
+  const char *pool3d_type;
+  DeserializeValue(&serialData, &serialLength, &pool3d_type);
+  pool3d_type_ = std::string(pool3d_type);
+  DeserializeValue(&serialData, &serialLength, &adaptive_);
+  DeserializeValue(&serialData, &serialLength, &ksize_);
+  DeserializeValue(&serialData, &serialLength, &strides_);
+  DeserializeValue(&serialData, &serialLength, &paddings_);
+  DeserializeValue(&serialData, &serialLength, &is_global_);
+}
+
+nvinfer1::IPluginV2DynamicExt *PIRPool3DPluginDynamic::clone() const
+    TRT_NOEXCEPT {
+  return new PIRPool3DPluginDynamic(ceil_mode_,
+                                    pool3d_type_,
+                                    adaptive_,
+                                    ksize_,
+                                    strides_,
+                                    paddings_,
+                                    is_global_);
+}
+
+const char *PIRPool3DPluginDynamic::getPluginType() const TRT_NOEXCEPT {
+  return "pir_pool3d_plugin_dynamic";
+}
+int PIRPool3DPluginDynamic::getNbOutputs() const TRT_NOEXCEPT { return 1; }
+
+int PIRPool3DPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+void PIRPool3DPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+size_t PIRPool3DPluginDynamic::getWorkspaceSize(
+    const nvinfer1::PluginTensorDesc *inputs,
+    int nbInputs,
+    const nvinfer1::PluginTensorDesc *outputs,
+    int nbOutputs) const TRT_NOEXCEPT {
+  return 0;
+}
+
+size_t PIRPool3DPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(ceil_mode_) + SerializedSize(pool3d_type_.c_str()) +
+         SerializedSize(adaptive_) + SerializedSize(ksize_) +
+         SerializedSize(strides_) + SerializedSize(paddings_) +
+         SerializedSize(is_global_);
+}
+
+void PIRPool3DPluginDynamic::serialize(void *buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, ceil_mode_);
+  SerializeValue(&buffer, pool3d_type_.c_str());
+  SerializeValue(&buffer, adaptive_);
+  SerializeValue(&buffer, ksize_);
+  SerializeValue(&buffer, strides_);
+  SerializeValue(&buffer, paddings_);
+  SerializeValue(&buffer, is_global_);
+}
+
+nvinfer1::DimsExprs PIRPool3DPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nb_inputs,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The Split plugin should be only one input."));
+
+  PADDLE_ENFORCE_EQ(
+      inputs[0].d[1]->isConstant(),
+      true,
+      common::errors::InvalidArgument("The channel dimension should be "
+                                      "static, but we found it's dynamic."));
+  nvinfer1::DimsExprs output(inputs[0]);
+  if (is_global_) {
+    output.d[2] = expr_builder.constant(1);
+    output.d[3] = expr_builder.constant(1);
+    output.d[4] = expr_builder.constant(1);
+    return output;
+  }
+  if (adaptive_) {
+    output.d[2] = expr_builder.constant(ksize_[0]);
+    output.d[3] = expr_builder.constant(ksize_[1]);
+    output.d[4] = expr_builder.constant(ksize_[2]);
+    return output;
+  }
+
+  auto stri_0 = expr_builder.constant(strides_[0]);
+  auto stri_1 = expr_builder.constant(strides_[1]);
+  auto stri_2 = expr_builder.constant(strides_[2]);
+  auto one_value = expr_builder.constant(1);
+
+  auto v0_tmp = expr_builder.constant(-ksize_[0] + 2 * paddings_[0]);
+  auto v1_tmp = expr_builder.constant(-ksize_[1] + 2 * paddings_[1]);
+  auto v2_tmp = expr_builder.constant(-ksize_[2] + 2 * paddings_[2]);
+
+  auto ceil_tmp =
+      expr_builder.constant(-ksize_[0] + 2 * paddings_[0] + strides_[0] - 1);
+  auto ceil1_tmp =
+      expr_builder.constant(-ksize_[1] + 2 * paddings_[1] + strides_[1] - 1);
+  auto ceil2_tmp =
+      expr_builder.constant(-ksize_[2] + 2 * paddings_[2] + strides_[2] - 1);
+
+  if (!ceil_mode_) {
+    output.d[2] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[2], *v0_tmp),
+            *stri_0),
+        *one_value);
+    output.d[3] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[3], *v1_tmp),
+            *stri_1),
+        *one_value);
+    output.d[4] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[4], *v2_tmp),
+            *stri_2),
+        *one_value);
+
+  } else {
+    output.d[2] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[2], *ceil_tmp),
+            *stri_0),
+        *one_value);
+    output.d[3] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
+                                    *inputs[0].d[3],
+                                    *ceil1_tmp),
+            *stri_1),
+        *one_value);
+    output.d[4] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
+                                    *inputs[0].d[4],
+                                    *ceil2_tmp),
+            *stri_2),
+        *one_value);
+  }
+
+  return output;
+}
+
+bool PIRPool3DPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  return ((in_out[pos].type == nvinfer1::DataType::kFLOAT) &&
+          in_out[pos].format == nvinfer1::PluginFormat::kLINEAR);
+}
+
+nvinfer1::DataType PIRPool3DPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Pool3D Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  PADDLE_ENFORCE_EQ(
+      (input_types[0] == nvinfer1::DataType::kFLOAT),
+      true,
+      common::errors::InvalidArgument("The input type should be float"));
+  return input_types[0];
+}
+
+int PIRPool3DPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  int n = input_dims.d[0];
+  int c = input_dims.d[1];
+  int d = input_dims.d[2];
+  int h = input_dims.d[3];
+  int w = input_dims.d[4];
+
+  const float *input = static_cast<const float *>(inputs[0]);
+  float *output = static_cast<float *>(outputs[0]);
+
+  std::vector<int> input_shape, output_shape;
+  for (int i = 0; i < input_dims.nbDims; i++)
+    input_shape.push_back(input_dims.d[i]);
+  output_shape = input_shape;
+
+  std::vector<int> ksize = ksize_;
+  std::vector<int> paddings = paddings_;
+  if (is_global_) {
+    ksize[0] = d;
+    ksize[1] = h;
+    ksize[2] = w;
+    paddings[0] = 0;
+    paddings[1] = 0;
+    paddings[2] = 0;
+    output_shape[2] = 1;
+    output_shape[3] = 1;
+    output_shape[4] = 1;
+  } else {
+    auto data_dim = CalcOutputSize(
+        {d, h, w}, ceil_mode_, adaptive_, ksize_, strides_, paddings_);
+    output_shape[2] = data_dim[0];
+    output_shape[3] = data_dim[1];
+    output_shape[4] = data_dim[2];
+  }
+
+  if (pool3d_type_ == "max") {
+    phi::funcs::MaxPool<float> pool_process;
+    phi::funcs::Pool3dDirectCUDAFunctor<phi::funcs::MaxPool<float>, float>
+        pool3d_forward;
+    pool3d_forward(input,
+                   input_shape,
+                   output_shape,
+                   ksize,
+                   strides_,
+                   paddings,
+                   true,
+                   adaptive_,
+                   output,
+                   stream,
+                   pool_process);
+  } else if (pool3d_type_ == "avg") {
+    phi::funcs::AvgPool<float> pool_process;
+    phi::funcs::Pool3dDirectCUDAFunctor<phi::funcs::AvgPool<float>, float>
+        pool3d_forward;
+    pool3d_forward(input,
+                   input_shape,
+                   output_shape,
+                   ksize,
+                   strides_,
+                   paddings,
+                   true,
+                   adaptive_,
+                   output,
+                   stream,
+                   pool_process);
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/pool3d_op_plugin.h

@@ -0,0 +1,397 @@
+// Copyright (c) 2021 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.
+
+#pragma once
+#include <stdio.h>
+
+#include <cassert>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+static std::vector<int> CalcOutputSize(const std::vector<int>& input_shape,
+                                       const bool& ceil_mode,
+                                       const bool& adaptive,
+                                       const std::vector<int>& ksize,
+                                       const std::vector<int>& strides,
+                                       const std::vector<int>& paddings) {
+  std::vector<int> output_shape = input_shape;
+  if (adaptive) {
+    output_shape[0] = ksize[0];
+    output_shape[1] = ksize[1];
+    output_shape[2] = ksize[2];
+  } else {
+    int output_d =
+        (input_shape[0] - ksize[0] + 2 * paddings[0]) / strides[0] + 1;
+    int output_h =
+        (input_shape[1] - ksize[1] + 2 * paddings[1]) / strides[1] + 1;
+    int output_w =
+        (input_shape[2] - ksize[2] + 2 * paddings[2]) / strides[2] + 1;
+    if (ceil_mode) {
+      output_d =
+          (input_shape[0] - ksize[0] + 2 * paddings[0] + strides[0] - 1) /
+              strides[0] +
+          1;
+      output_h =
+          (input_shape[1] - ksize[1] + 2 * paddings[1] + strides[1] - 1) /
+              strides[1] +
+          1;
+      output_w =
+          (input_shape[2] - ksize[2] + 2 * paddings[2] + strides[2] - 1) /
+              strides[2] +
+          1;
+    }
+    output_shape[0] = output_d;
+    output_shape[1] = output_h;
+    output_shape[2] = output_w;
+  }
+  return output_shape;
+}
+
+class Pool3DPlugin : public PluginTensorRTV2Ext {
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  // TRT will call this func when we need to serialize the configuration of
+  // tensorrt.
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  enum class Pool3DType {
+    max = 0,
+    avg,
+  };
+  Pool3DPlugin() {}
+  Pool3DPlugin(bool ceil_mode,
+               Pool3DType pool3d_type,
+               bool adaptive,
+               std::vector<int> ksize,
+               std::vector<int> strides,
+               std::vector<int> paddings,
+               std::vector<int> input_shape)
+      : ceil_mode_(ceil_mode),
+        pool3d_type_(pool3d_type),
+        adaptive_(adaptive),
+        ksize_(ksize),
+        strides_(strides),
+        paddings_(paddings),
+        input_shape_(input_shape) {
+    output_shape_ = input_shape_;
+    std::vector<int> output_shape =
+        CalcOutputSize({input_shape_[1], input_shape_[2], input_shape_[3]},
+                       ceil_mode_,
+                       adaptive_,
+                       ksize_,
+                       strides_,
+                       paddings_);
+    output_shape_[1] = output_shape[0];
+    output_shape_[2] = output_shape[1];
+    output_shape_[3] = output_shape[2];
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  Pool3DPlugin(void const* serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &ceil_mode_);
+    DeserializeValue(&serialData, &serialLength, &pool3d_type_);
+    DeserializeValue(&serialData, &serialLength, &adaptive_);
+    DeserializeValue(&serialData, &serialLength, &ksize_);
+    DeserializeValue(&serialData, &serialLength, &strides_);
+    DeserializeValue(&serialData, &serialLength, &paddings_);
+    DeserializeValue(&serialData, &serialLength, &input_shape_);
+    DeserializeValue(&serialData, &serialLength, &output_shape_);
+  }
+
+  Pool3DPlugin* clone() const TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  int getNbOutputs() const TRT_NOEXCEPT override;
+
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+
+  int initialize() TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override;
+
+  int enqueue(int batchSize,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+ private:
+  bool ceil_mode_;
+  Pool3DType pool3d_type_;
+  bool adaptive_;
+  std::vector<int> ksize_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  std::vector<int> input_shape_;
+  std::vector<int> output_shape_;
+};
+
+class Pool3DPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "pool3d_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new Pool3DPlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(Pool3DPluginCreator);
+
+class Pool3DPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  Pool3DPluginDynamic() {}
+  Pool3DPluginDynamic(const bool& ceil_mode,
+                      const std::string& pool3d_type,
+                      const bool& adaptive,
+                      const std::vector<int>& ksize,
+                      const std::vector<int>& strides,
+                      const std::vector<int>& paddings,
+                      const bool& is_global)
+      : ceil_mode_(ceil_mode),
+        pool3d_type_(pool3d_type),
+        adaptive_(adaptive),
+        ksize_(ksize),
+        strides_(strides),
+        paddings_(paddings),
+        is_global_(is_global) {}
+
+  Pool3DPluginDynamic(void const* serialData, size_t serialLength);
+  ~Pool3DPluginDynamic() {}
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  bool ceil_mode_;
+  std::string pool3d_type_;
+  bool adaptive_;
+  std::vector<int> ksize_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  bool is_global_;
+};
+
+class Pool3DPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "pool3d_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new Pool3DPluginDynamic(serial_data, serial_length);
+  }
+};
+
+class PIRPool3DPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  PIRPool3DPluginDynamic() {}
+  PIRPool3DPluginDynamic(const bool& ceil_mode,
+                         const std::string& pool3d_type,
+                         const bool& adaptive,
+                         const std::vector<int>& ksize,
+                         const std::vector<int>& strides,
+                         const std::vector<int>& paddings,
+                         const bool& is_global)
+      : ceil_mode_(ceil_mode),
+        pool3d_type_(pool3d_type),
+        adaptive_(adaptive),
+        ksize_(ksize),
+        strides_(strides),
+        paddings_(paddings),
+        is_global_(is_global) {}
+
+  PIRPool3DPluginDynamic(void const* serialData, size_t serialLength);
+  ~PIRPool3DPluginDynamic() {}
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  bool ceil_mode_;
+  std::string pool3d_type_;
+  bool adaptive_;
+  std::vector<int> ksize_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  bool is_global_;
+};
+
+class PIRPool3DPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "pir_pool3d_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PIRPool3DPluginDynamic(serial_data, serial_length);
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    bool ceil_mode = false;
+    std::string pool3d_type;
+    bool adaptive = false;
+    std::vector<int> ksize;
+    std::vector<int> strides;
+    std::vector<int> paddings;
+    bool is_global = false;
+
+    for (int i = 0; i < fc->nbFields; ++i) {
+      const nvinfer1::PluginField& field = fc->fields[i];
+      const std::string field_name(field.name);
+
+      if (field_name.compare("ceil_mode") == 0) {
+        ceil_mode = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("pool3d_type") == 0) {
+        pool3d_type = std::string(static_cast<const char*>(fc->fields[i].data),
+                                  fc->fields[i].length);
+      } else if (field_name.compare("adaptive") == 0) {
+        adaptive = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("ksize") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        ksize.insert(ksize.end(), data, data + length);
+      } else if (field_name.compare("strides") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        strides.insert(strides.end(), data, data + length);
+      } else if (field_name.compare("paddings") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        paddings.insert(paddings.end(), data, data + length);
+      } else if (field_name.compare("is_global") == 0) {
+        is_global = *static_cast<const bool*>(field.data);
+      } else {
+        assert(false && "Unknown plugin field name.");
+      }
+    }
+    return new PIRPool3DPluginDynamic(
+        ceil_mode, pool3d_type, adaptive, ksize, strides, paddings, is_global);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(Pool3DPluginDynamicCreator);
+REGISTER_TRT_PLUGIN_V2(PIRPool3DPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/pool_op_plugin.cu

@@ -0,0 +1,364 @@
+// Copyright (c) 2018 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 "paddle/fluid/inference/tensorrt/plugin/pool_op_plugin.h"
+#include "paddle/phi/kernels/funcs/pooling.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+nvinfer1::Dims PoolPlugin::getOutputDimensions(int index,
+                                               const nvinfer1::Dims *inputDims,
+                                               int nbInputs) TRT_NOEXCEPT {
+  assert(nbInputs == 1);
+  assert(index == 0);
+  assert(inputDims[0].nbDims == 3);
+  nvinfer1::Dims const &input_dims = inputDims[0];
+
+  nvinfer1::Dims output_dims = input_dims;
+
+  output_dims.d[1] = output_shape_[1];
+  output_dims.d[2] = output_shape_[2];
+  return output_dims;
+}
+
+size_t PoolPlugin::getSerializationSize() const TRT_NOEXCEPT {
+  return getBaseSerializationSize() + SerializedSize(ceil_mode_) +
+         SerializedSize(pool_type_) + SerializedSize(adaptive_) +
+         SerializedSize(exclusive_) + SerializedSize(ksize_) +
+         SerializedSize(strides_) + SerializedSize(paddings_) +
+         SerializedSize(real_paddings_) + SerializedSize(input_shape_) +
+         SerializedSize(output_shape_);
+}
+
+// TRT will call this func when we need to serialize the configuration of
+// tensorrt.
+void PoolPlugin::serialize(void *buffer) const TRT_NOEXCEPT {
+  serializeBase(buffer);
+  SerializeValue(&buffer, ceil_mode_);
+  SerializeValue(&buffer, pool_type_);
+  SerializeValue(&buffer, adaptive_);
+  SerializeValue(&buffer, exclusive_);
+  SerializeValue(&buffer, ksize_);
+  SerializeValue(&buffer, strides_);
+  SerializeValue(&buffer, paddings_);
+  SerializeValue(&buffer, real_paddings_);
+  SerializeValue(&buffer, input_shape_);
+  SerializeValue(&buffer, output_shape_);
+}
+
+PoolPlugin *PoolPlugin::clone() const TRT_NOEXCEPT {
+  return new PoolPlugin(ceil_mode_,
+                        pool_type_,
+                        adaptive_,
+                        exclusive_,
+                        ksize_,
+                        strides_,
+                        paddings_,
+                        input_shape_,
+                        real_paddings_);
+}
+
+int PoolPlugin::enqueue(int batchSize,
+                        const void *const *inputs,
+                        void *const *outputs,
+                        void *workspace,
+                        cudaStream_t stream) TRT_NOEXCEPT {
+  auto const &input_dims = this->getInputDims(0);
+  int input_size = 0;
+  float const *idata = reinterpret_cast<float const *>(inputs[0]);
+  float *const *odatas = reinterpret_cast<float *const *>(outputs);
+
+  std::vector<int> input_shape = input_shape_;
+  std::vector<int> output_shape = output_shape_;
+  input_shape.insert(input_shape.begin(), batchSize);
+  output_shape.insert(output_shape.begin(), batchSize);
+
+  if (pool_type_ == PoolType::max) {
+    phi::funcs::MaxPool<float> pool_process;
+    phi::funcs::Pool2dDirectCUDAFunctor<phi::funcs::MaxPool<float>, float>
+        pool2d_forward;
+    pool2d_forward(idata,
+                   input_shape,
+                   output_shape,
+                   ksize_,
+                   strides_,
+                   paddings_,
+                   true,
+                   false,
+                   odatas[0],
+                   stream,
+                   pool_process);
+  } else if (pool_type_ == PoolType::avg) {
+    phi::funcs::AvgPool<float> pool_process;
+    phi::funcs::Pool2dDirectCUDAFunctor<phi::funcs::AvgPool<float>, float>
+        pool2d_forward;
+    pool2d_forward(idata,
+                   input_shape,
+                   output_shape,
+                   ksize_,
+                   strides_,
+                   paddings_,
+                   exclusive_,
+                   adaptive_,
+                   odatas[0],
+                   stream,
+                   pool_process);
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+PoolPluginDynamic::PoolPluginDynamic(void const *serialData,
+                                     size_t serialLength) {
+  DeserializeValue(&serialData, &serialLength, &ceil_mode_);
+  const char *pool_type;
+  DeserializeValue(&serialData, &serialLength, &pool_type);
+  pool_type_ = std::string(pool_type);
+  DeserializeValue(&serialData, &serialLength, &adaptive_);
+  DeserializeValue(&serialData, &serialLength, &exclusive_);
+  DeserializeValue(&serialData, &serialLength, &ksize_);
+  DeserializeValue(&serialData, &serialLength, &strides_);
+  DeserializeValue(&serialData, &serialLength, &paddings_);
+  DeserializeValue(&serialData, &serialLength, &is_global_);
+}
+
+size_t PoolPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(ceil_mode_) + SerializedSize(pool_type_.c_str()) +
+         SerializedSize(adaptive_) + SerializedSize(exclusive_) +
+         SerializedSize(ksize_) + SerializedSize(strides_) +
+         SerializedSize(paddings_) + SerializedSize(is_global_);
+}
+
+void PoolPluginDynamic::serialize(void *buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, ceil_mode_);
+  SerializeValue(&buffer, pool_type_.c_str());
+  SerializeValue(&buffer, adaptive_);
+  SerializeValue(&buffer, exclusive_);
+  SerializeValue(&buffer, ksize_);
+  SerializeValue(&buffer, strides_);
+  SerializeValue(&buffer, paddings_);
+  SerializeValue(&buffer, is_global_);
+}
+
+nvinfer1::IPluginV2DynamicExt *PoolPluginDynamic::clone() const TRT_NOEXCEPT {
+  return new PoolPluginDynamic(ceil_mode_,
+                               pool_type_,
+                               adaptive_,
+                               exclusive_,
+                               ksize_,
+                               strides_,
+                               paddings_,
+                               is_global_);
+}
+
+nvinfer1::DimsExprs PoolPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nb_inputs,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The Split plugin should be only one input."));
+
+  nvinfer1::DimsExprs output(inputs[0]);
+  if (is_global_ && !adaptive_) {
+    output.d[2] = expr_builder.constant(1);
+    output.d[3] = expr_builder.constant(1);
+    return output;
+  }
+  if (is_global_ && adaptive_) {
+    return inputs[0];
+  }
+  if (adaptive_) {
+    output.d[2] = expr_builder.constant(ksize_[0]);
+    output.d[3] = expr_builder.constant(ksize_[1]);
+    return output;
+  }
+
+  auto stri_0 = expr_builder.constant(strides_[0]);
+  auto stri_1 = expr_builder.constant(strides_[1]);
+  auto one_value = expr_builder.constant(1);
+
+  auto v0_tmp = expr_builder.constant(-ksize_[0] + 2 * paddings_[0]);
+  auto v1_tmp = expr_builder.constant(-ksize_[1] + 2 * paddings_[1]);
+
+  auto ceil_tmp =
+      expr_builder.constant(-ksize_[0] + 2 * paddings_[0] + strides_[0] - 1);
+  auto ceil1_tmp =
+      expr_builder.constant(-ksize_[1] + 2 * paddings_[1] + strides_[1] - 1);
+
+  if (!ceil_mode_) {
+    output.d[2] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[2], *v0_tmp),
+            *stri_0),
+        *one_value);
+    output.d[3] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[3], *v1_tmp),
+            *stri_1),
+        *one_value);
+
+  } else {
+    output.d[2] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(
+                nvinfer1::DimensionOperation::kSUM, *inputs[0].d[2], *ceil_tmp),
+            *stri_0),
+        *one_value);
+    output.d[3] = expr_builder.operation(
+        nvinfer1::DimensionOperation::kSUM,
+        *expr_builder.operation(
+            nvinfer1::DimensionOperation::kFLOOR_DIV,
+            *expr_builder.operation(nvinfer1::DimensionOperation::kSUM,
+                                    *inputs[0].d[3],
+                                    *ceil1_tmp),
+            *stri_1),
+        *one_value);
+  }
+
+  return output;
+}
+
+bool PoolPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  return ((in_out[pos].type == nvinfer1::DataType::kFLOAT) &&
+          in_out[pos].format == nvinfer1::PluginFormat::kLINEAR);
+}
+
+nvinfer1::DataType PoolPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Pool Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  PADDLE_ENFORCE_EQ(
+      (input_types[0] == nvinfer1::DataType::kFLOAT),
+      true,
+      common::errors::InvalidArgument("The input type should be float"));
+  return input_types[0];
+}
+
+int PoolPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  int n = input_dims.d[0];
+  int c = input_dims.d[1];
+  int h = input_dims.d[2];
+  int w = input_dims.d[3];
+
+  const float *input = static_cast<const float *>(inputs[0]);
+  float *output = static_cast<float *>(outputs[0]);
+
+  std::vector<int> input_shape, output_shape;
+  for (int i = 0; i < input_dims.nbDims; i++)
+    input_shape.push_back(input_dims.d[i]);
+  output_shape = input_shape;
+
+  std::vector<int> ksize = ksize_;
+  std::vector<int> paddings = paddings_;
+  if (is_global_) {
+    ksize[0] = h;
+    ksize[1] = w;
+    paddings[0] = 0;
+    paddings[1] = 0;
+    output_shape[2] = 1;
+    output_shape[3] = 1;
+    if (adaptive_) {
+      output_shape[2] = h;
+      output_shape[3] = w;
+    }
+  } else {
+    auto data_dim = CalcOutputSize(
+        {h, w}, ceil_mode_, adaptive_, ksize_, strides_, paddings_);
+    output_shape[2] = data_dim[0];
+    output_shape[3] = data_dim[1];
+  }
+
+  if (pool_type_ == "max") {
+    phi::funcs::MaxPool<float> pool_process;
+    phi::funcs::Pool2dDirectCUDAFunctor<phi::funcs::MaxPool<float>, float>
+        pool2d_forward;
+    pool2d_forward(input,
+                   input_shape,
+                   output_shape,
+                   ksize,
+                   strides_,
+                   paddings,
+                   true,
+                   false,
+                   output,
+                   stream,
+                   pool_process);
+  } else if (pool_type_ == "avg") {
+    phi::funcs::AvgPool<float> pool_process;
+    phi::funcs::Pool2dDirectCUDAFunctor<phi::funcs::AvgPool<float>, float>
+        pool2d_forward;
+    pool2d_forward(input,
+                   input_shape,
+                   output_shape,
+                   ksize,
+                   strides_,
+                   paddings,
+                   exclusive_,
+                   adaptive_,
+                   output,
+                   stream,
+                   pool_process);
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/pool_op_plugin.h

@@ -0,0 +1,339 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+#include <stdio.h>
+
+#include <cassert>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+static std::vector<int> CalcOutputSize(const std::vector<int>& input_shape,
+                                       const bool& ceil_mode,
+                                       const bool& adaptive,
+                                       const std::vector<int>& ksize,
+                                       const std::vector<int>& strides,
+                                       const std::vector<int>& real_paddings) {
+  std::vector<int> output_shape = input_shape;
+  if (adaptive) {
+    output_shape[0] = ksize[0];
+    output_shape[1] = ksize[1];
+  } else {
+    int output_h = 0, output_w = 0;
+    if (ceil_mode) {
+      output_h = (input_shape[0] - ksize[0] + real_paddings[0] +
+                  real_paddings[1] + strides[0] - 1) /
+                     strides[0] +
+                 1;
+      output_w = (input_shape[1] - ksize[1] + real_paddings[2] +
+                  real_paddings[3] + strides[1] - 1) /
+                     strides[1] +
+                 1;
+    }
+    // TRT will use native layer when ceil_model=false
+    /*
+    else{
+      output_h = (input_shape[0] - ksize[0] + real_paddings[0] +
+    real_paddings[1]) / strides[0] + 1;
+      output_w = (input_shape[1] - ksize[1] + real_paddings[2] +
+    real_paddings[3]) / strides[1] + 1;
+    }
+    */
+    output_shape[0] = output_h;
+    output_shape[1] = output_w;
+  }
+  return output_shape;
+}
+
+class PoolPlugin : public PluginTensorRT {
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  enum class PoolType {
+    max = 0,
+    avg,
+  };
+  PoolPlugin() {}
+  PoolPlugin(bool ceil_mode,
+             PoolType pool_type,
+             bool adaptive,
+             bool exclusive,
+             std::vector<int> ksize,
+             std::vector<int> strides,
+             std::vector<int> paddings,
+             std::vector<int> input_shape,
+             std::vector<int> real_paddings)
+      : ceil_mode_(ceil_mode),
+        pool_type_(pool_type),
+        adaptive_(adaptive),
+        exclusive_(exclusive),
+        ksize_(ksize),
+        strides_(strides),
+        paddings_(paddings),
+        real_paddings_(real_paddings),
+        input_shape_(input_shape) {
+    output_shape_ = input_shape_;
+    std::vector<int> output_shape =
+        CalcOutputSize({input_shape_[1], input_shape_[2]},
+                       ceil_mode_,
+                       adaptive_,
+                       ksize_,
+                       strides_,
+                       real_paddings_);
+    output_shape_[1] = output_shape[0];
+    output_shape_[2] = output_shape[1];
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  PoolPlugin(void const* serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &ceil_mode_);
+    DeserializeValue(&serialData, &serialLength, &pool_type_);
+    DeserializeValue(&serialData, &serialLength, &adaptive_);
+    DeserializeValue(&serialData, &serialLength, &exclusive_);
+    DeserializeValue(&serialData, &serialLength, &ksize_);
+    DeserializeValue(&serialData, &serialLength, &strides_);
+    DeserializeValue(&serialData, &serialLength, &paddings_);
+    DeserializeValue(&serialData, &serialLength, &real_paddings_);
+    DeserializeValue(&serialData, &serialLength, &input_shape_);
+    DeserializeValue(&serialData, &serialLength, &output_shape_);
+  }
+
+  PoolPlugin* clone() const TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "pool_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  int enqueue(int batchSize,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+ private:
+  bool ceil_mode_;
+  PoolType pool_type_;
+  bool adaptive_;
+  bool exclusive_;
+  std::vector<int> ksize_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  std::vector<int> real_paddings_;
+  std::vector<int> input_shape_;
+  std::vector<int> output_shape_;
+};
+
+class PoolPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "pool_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PoolPlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(PoolPluginCreator);
+
+class PoolPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  PoolPluginDynamic() {}
+  PoolPluginDynamic(const bool& ceil_mode,
+                    const std::string& pool_type,
+                    const bool& adaptive,
+                    bool exclusive,
+                    const std::vector<int>& ksize,
+                    const std::vector<int>& strides,
+                    const std::vector<int>& paddings,
+                    const bool& is_global)
+      : ceil_mode_(ceil_mode),
+        pool_type_(pool_type),
+        adaptive_(adaptive),
+        exclusive_(exclusive),
+        ksize_(ksize),
+        strides_(strides),
+        paddings_(paddings),
+        is_global_(is_global) {}
+
+  PoolPluginDynamic(void const* serialData, size_t serialLength);
+  ~PoolPluginDynamic() {}
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "pool_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  bool ceil_mode_;
+  std::string pool_type_;
+  bool adaptive_;
+  bool exclusive_;
+  std::vector<int> ksize_;
+  std::vector<int> strides_;
+  std::vector<int> paddings_;
+  bool is_global_;
+};
+
+class PoolPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "pool_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PoolPluginDynamic(serial_data, serial_length);
+  }
+};
+
+class PIRPoolPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "pir_pool_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PoolPluginDynamic(serial_data, serial_length);
+  }
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    bool ceil_mode = false;
+    std::string pool_type;
+    bool adaptive = false;
+    bool exclusive = false;
+    std::vector<int> ksize;
+    std::vector<int> strides;
+    std::vector<int> paddings;
+    bool global_pooling = false;
+
+    for (int i = 0; i < fc->nbFields; ++i) {
+      const nvinfer1::PluginField& field = fc->fields[i];
+      const std::string field_name(field.name);
+      if (field_name.compare("ceil_mode") == 0) {
+        ceil_mode = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("pool_type") == 0) {
+        pool_type = std::string(static_cast<const char*>(fc->fields[i].data),
+                                fc->fields[i].length);
+      } else if (field_name.compare("adaptive") == 0) {
+        adaptive = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("exclusive") == 0) {
+        exclusive = *static_cast<const bool*>(field.data);
+      } else if (field_name.compare("ksize") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        ksize.insert(ksize.end(), data, data + length);
+      } else if (field_name.compare("strides") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        strides.insert(strides.end(), data, data + length);
+      } else if (field_name.compare("paddings") == 0) {
+        const int length = fc->fields[i].length;
+        const int* data = static_cast<const int*>(fc->fields[i].data);
+        paddings.insert(paddings.end(), data, data + length);
+      } else if (field_name.compare("global_pooling") == 0) {
+        global_pooling = *static_cast<const bool*>(field.data);
+      } else {
+        assert(false && "unknown plugin field name.");
+      }
+    }
+    return new PoolPluginDynamic(ceil_mode,
+                                 pool_type,
+                                 adaptive,
+                                 exclusive,
+                                 ksize,
+                                 strides,
+                                 paddings,
+                                 global_pooling);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(PoolPluginDynamicCreator);
+REGISTER_TRT_PLUGIN_V2(PIRPoolPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/preln_groupnorm_act_op_plugin.cu

@@ -0,0 +1,497 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES.
+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 "paddle/fluid/inference/tensorrt/plugin/preln_groupnorm_act_op_plugin.h"
+#include <cub/cub.cuh>
+#include "paddle/common/layout.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+nvinfer1::DimsExprs PrelnGroupnormActPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputDims,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  return inputDims[0];
+}
+
+bool PrelnGroupnormActPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of prelnGroupnormAct plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return ((in.type == nvinfer1::DataType::kHALF) &&
+              (in.format == nvinfer1::PluginFormat::kHWC8));
+    } else {
+      PADDLE_THROW(common::errors::Fatal(
+          "PrelnGroupnormAct TRT Plugin is fp16 only so far"));
+
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType PrelnGroupnormActPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  return input_types[0];
+}
+
+int PrelnGroupnormActPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+template <typename T>
+static inline T divUp(T m, T n) {
+  return (m + n - 1) / n;
+}
+
+static int64_t findMaxDivisor(int64_t n, int64_t maxAllowedDivisor) {
+  int64_t maxDivisor = -1;
+  for (int64_t i = 1; i <= std::sqrt(n); i++) {
+    if (n % i == 0) {
+      int64_t divisor1 = n / i;
+      int64_t divisor2 = i;
+
+      if (divisor1 > maxDivisor && divisor1 < maxAllowedDivisor) {
+        maxDivisor = divisor1;
+      }
+      if (divisor2 > maxDivisor && divisor2 < maxAllowedDivisor) {
+        maxDivisor = divisor2;
+      }
+    }
+  }
+  return maxDivisor;
+}
+
+static inline __device__ __host__ float sigmoid(float x) {
+  return 1.F / (1.F + expf(-x));
+}
+
+struct GroupSums {
+  // Is it the 1st element of the group?
+  int32_t flag;
+  // The sum.
+  float sum;
+  // The sum of squares.
+  float sumSq;
+};
+
+struct GroupSumsOp {
+  inline __device__ GroupSums operator()(GroupSums const &a,
+                                         GroupSums const &b) {
+    GroupSums dst;
+    dst.sum = b.flag ? b.sum : (a.sum + b.sum);
+    dst.sumSq = b.flag ? b.sumSq : (a.sumSq + b.sumSq);
+    dst.flag = a.flag + b.flag;
+    return dst;
+  }
+};
+
+template <int32_t tTHREADS_PER_BLOCK>
+__global__ void prelnGroupNormNDHWCSumKernel(
+    GroupNormNDHWCParams<__half> params) {
+  // The object in charge of doing the sums for the different blocks.
+  typedef cub::BlockScan<GroupSums, tTHREADS_PER_BLOCK> BlockScan;
+
+  // Allocate shared memory for BlockScan.
+  __shared__ typename BlockScan::TempStorage tempStorage;
+  // Allocate shared memory for the groups. We could reduce the amount of shared
+  // memory reserved.
+  __shared__ float2 smem[tTHREADS_PER_BLOCK];
+
+  // The instance in the batch.
+  int32_t ni = blockIdx.z;
+  // The channel loaded by that thread (2 channels per thread for F16x2).
+  int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * 2;
+
+  // The first activation loaded by that block.
+  int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
+  // The last activation loaded by that block.
+  int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
+
+  // The sums.
+  float sum = 0.F;
+  float sumSq = 0.F;
+
+  // Iterate over the activations to compute the sums.
+  for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
+    // The offset.
+    int64_t offset = static_cast<int64_t>(ni) * params.dhwc +
+                     static_cast<int64_t>(dhwi) * params.c + ci;
+    // Fetch two channels per thread.
+    __half2 h2(0, 0);
+    if (ci < params.c) {
+      // int64_t offsetY = static_cast<int64_t>(ni) * params.c + ci;
+      __half2 y = *reinterpret_cast<__half2 const *>(&params.srcY[offset]);
+      h2 = *reinterpret_cast<__half2 const *>(&params.srcX[offset]);
+#if __CUDA_ARCH__ >= 530
+      h2 = __hadd2(h2, y);
+#else
+      float2 out{};
+      out.x = __half2float(h2.x) + __half2float(y.x);
+      out.y = __half2float(h2.y) + __half2float(y.y);
+      h2 = __float22half2_rn(out);
+#endif
+      // elementwise_add
+      *reinterpret_cast<__half2 *>(&params.eleOut[offset]) = h2;
+    }
+
+    // Extract the two half values.
+    float2 f2 = __half22float2(h2);
+
+    // Update the sum.
+    sum += f2.x + f2.y;
+    // Update the sum of squares.
+    sumSq += f2.x * f2.x + f2.y * f2.y;
+  }
+
+  // The group that thread works on and the channel in the group (modulus).
+  int32_t gi = threadIdx.x * 2 / params.cPerGroup;
+  int32_t cj = threadIdx.x * 2 - params.cPerGroup * gi;
+
+  // The data for the summations.
+  GroupSums inp{cj == 0 ? 1 : 0, sum, sumSq};
+
+  // Do the segmented scan.
+  GroupSums out;
+  BlockScan(tempStorage).InclusiveScan(inp, out, GroupSumsOp());
+
+  // Store the results for the groups in shared memory (to produce coalesced
+  // stores later).
+  if (cj == params.cPerGroup - 2 /* 2 channels per thread */) {
+    smem[gi] = make_float2(out.sum, out.sumSq);
+  }
+
+  // Make sure the data is in shared memory.
+  __syncthreads();
+
+  // The global group index.
+  int32_t gj = blockIdx.x * params.groupsPerBlock + threadIdx.x;
+
+  // Threads that have nothing left to do, exit.
+  if (threadIdx.x >= params.groupsPerBlock || gj >= params.groups) {
+    return;
+  }
+
+  // The first threads (those storing to global memory, load the values).
+  float2 sums = smem[threadIdx.x];
+
+  // Store to global memory.
+  atomicAdd(&params.redBuffer[(2 * ni + 0) * params.groups + gj], sums.x);
+  atomicAdd(&params.redBuffer[(2 * ni + 1) * params.groups + gj], sums.y);
+}
+
+void prelnGroupNormNDHWCSum(GroupNormNDHWCParams<__half> const &params,
+                            cudaStream_t stream) {
+  // Make sure the values are as we expect.
+  PADDLE_ENFORCE_EQ(params.c % params.cPerBlock,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The groupNormNDHWCSum of prelnGroupnormAct Plugin got "
+                        "wrong parameters: "
+                        "params.c %% params.cPerBlock should be 0, but get %d.",
+                        params.c % params.cPerBlock));
+  PADDLE_ENFORCE_EQ(
+      params.dhw % params.dhwPerBlock,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCSum of prelnGroupnormAct Plugin got wrong "
+          "parameters: "
+          "params.dhw  %% params.dhwPerBlock should be 0, but get %d.",
+          params.dhw % params.dhwPerBlock));
+  // Make sure a group does not span multiple blocks.
+  PADDLE_ENFORCE_EQ(
+      params.cPerBlock % params.cPerGroup,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCSum of prelnGroupnormAct Plugin got wrong "
+          "parameters: "
+          "params.cPerBlock %% params.cPerGroup should be 0, but get %d.",
+          params.cPerBlock % params.cPerGroup));
+  dim3 grid;
+
+  // The number of blocks to compute all the channels.
+  grid.x = params.c / params.cPerBlock;
+  // The number of blocks to compute all the activations in a given instance.
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
+  // The number of instances.
+  grid.z = params.n;
+
+  switch (params.cPerBlock) {
+    case 320:
+      prelnGroupNormNDHWCSumKernel<160><<<grid, 160, 0, stream>>>(params);
+      break;
+    case 480:
+      prelnGroupNormNDHWCSumKernel<256><<<grid, 256, 0, stream>>>(params);
+      break;
+    case 256:
+      prelnGroupNormNDHWCSumKernel<128><<<grid, 128, 0, stream>>>(params);
+      break;
+    case 128:
+      prelnGroupNormNDHWCSumKernel<64><<<grid, 64, 0, stream>>>(params);
+      break;
+    case 8:
+      prelnGroupNormNDHWCSumKernel<4><<<grid, 4, 0, stream>>>(params);
+      break;
+    default:
+      PADDLE_THROW(common::errors::Fatal(
+          "The function groupNormNDHWCSum of prelnGroupnormAct TRT Plugin "
+          "encounter error"));
+  }
+}
+
+template <int32_t tTHREADS_PER_BLOCK>
+__global__ void prelnGroupNormNDHWCScaleKernel(
+    GroupNormNDHWCParams<__half> params) {
+  // The instance in the batch.
+  int32_t ni = blockIdx.z;
+  // The channel loaded by that thread (2 channels per thread for F16x2).
+  int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * 2;
+  // The group that thread works on and the channel in the group (modulus).
+  int32_t gi = ci / params.cPerGroup;
+
+  // Load the sum and sum of squares for the group.
+  float sum = 0.F, sumSq = 0.F;
+  if (gi < params.groups) {
+    sum = params.redBuffer[(2 * ni + 0) * params.groups + gi];
+    sumSq = params.redBuffer[(2 * ni + 1) * params.groups + gi];
+  }
+
+  // Load gamma/beta.
+  float2 gammaF2, betaF2;
+  if (ci < params.c) {
+    gammaF2 = *reinterpret_cast<float2 const *>(
+        reinterpret_cast<float const *>(params.gamma) + ci);
+    betaF2 = *reinterpret_cast<float2 const *>(
+        reinterpret_cast<float const *>(params.beta) + ci);
+  }
+
+  // Compute the mean.
+  float mean = sum * params.invDHWC;
+  // Compute the variance.
+  float var = sumSq * params.invDHWC - (mean * mean);
+  // Compute the inverse of the stddev.
+  float invStdDev = rsqrtf(var + params.eps);
+
+  // The first activation loaded by that block.
+  int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
+  // The last activation loaded by that block.
+  int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
+
+  // Iterate over the activations to compute the sums.
+  for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
+    // The src/dst offset.
+    int64_t offset = (int64_t)ni * params.dhwc + dhwi * params.c + ci;
+
+    // Fetch two channels per thread.
+    __half2 h2(0, 0);
+    if (ci < params.c) {
+      h2 = *reinterpret_cast<__half2 const *>(&params.eleOut[offset]);
+    }
+
+    // Extract the two half values.
+    float2 f2 = __half22float2(h2);
+
+    // Normalize the channels.
+    f2.x = (f2.x - mean) * invStdDev;
+    f2.y = (f2.y - mean) * invStdDev;
+
+    // Scale by gamma and add beta.
+    f2.x = gammaF2.x * f2.x + betaF2.x;
+    f2.y = gammaF2.y * f2.y + betaF2.y;
+
+    // Apply Silu if needed.
+    if (params.withSilu) {
+      f2.x = f2.x * sigmoid(f2.x);
+      f2.y = f2.y * sigmoid(f2.y);
+    }
+
+    // Store the scaled values.
+    if (ci < params.c) {
+      *reinterpret_cast<__half2 *>(&params.dst[offset]) = __float22half2_rn(f2);
+    }
+  }
+}
+
+void prelnGroupNormNDHWCScale(GroupNormNDHWCParams<__half> const &params,
+                              cudaStream_t stream) {
+  // Make sure the dimensions are aligned with what we expect.
+  PADDLE_ENFORCE_EQ(
+      params.c % params.cPerBlock,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCScale of prelnGroupnormAct Plugin got "
+          "wrong parameters: "
+          "params.c %% params.cPerBlock should be 0, but get %d.",
+          params.c % params.cPerBlock));
+  // Make sure a group does not span multiple blocks.
+  PADDLE_ENFORCE_EQ(
+      params.cPerBlock % params.cPerGroup,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCScale of prelnGroupnormAct Plugin got wrong "
+          "parameters: "
+          "params.cPerBlock %% params.cPerGroup should be 0, but get %d.",
+          params.cPerBlock % params.cPerGroup));
+  dim3 grid;
+
+  // The number of blocks to compute all the channels.
+  grid.x = params.c / params.cPerBlock;
+  // The number of blocks to compute all the activations in a given instance.
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
+  // The number of instances.
+  grid.z = params.n;
+
+  switch (params.cPerBlock) {
+    case 320:
+      prelnGroupNormNDHWCScaleKernel<160><<<grid, 160, 0, stream>>>(params);
+      break;
+    case 480:
+      prelnGroupNormNDHWCScaleKernel<256><<<grid, 256, 0, stream>>>(params);
+      break;
+    case 256:
+      prelnGroupNormNDHWCScaleKernel<128><<<grid, 128, 0, stream>>>(params);
+      break;
+    case 128:
+      prelnGroupNormNDHWCScaleKernel<64><<<grid, 64, 0, stream>>>(params);
+      break;
+    case 8:
+      prelnGroupNormNDHWCScaleKernel<4><<<grid, 4, 0, stream>>>(params);
+      break;
+    default:
+      PADDLE_THROW(common::errors::Fatal(
+          "The function groupNormNDHWCSum of prelnGroupnormAct TRT Plugin "
+          "encounter error"));
+  }
+}
+
+int PrelnGroupnormActPluginDynamic::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 {
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. prelnGroupnormAct-->fp32";
+    PADDLE_THROW(common::errors::Fatal(
+        "The prelnGroupnormAct TRT Plugin's only support fp16 input"));
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. prelnGroupnormAct-->fp16";
+
+    int32_t cPerBlock = 320;
+    int32_t maxBlocksPerDHW = 1024;
+
+    switch (input_desc[0].dims.d[1]) {
+      case 960:
+      case 1920:
+        cPerBlock = 480;
+        break;
+      case 512:
+      case 256:
+        cPerBlock = 256;
+        break;
+      case 128:
+        cPerBlock = 128;
+        break;
+      default:
+        cPerBlock = 320;
+    }
+    if (cPerBlock > input_desc[0].dims.d[1]) {
+      cPerBlock = 8;
+    }
+    auto d_dim = input_desc[0].dims.nbDims;
+    params_.n = input_desc[0].dims.d[0];
+    if (d_dim == 3) {
+      params_.c = input_desc[0].dims.d[1];
+      params_.d = 1;
+      params_.h = 1;
+      params_.w = input_desc[0].dims.d[2];
+    } else if (d_dim == 4) {
+      params_.c = input_desc[0].dims.d[1];
+      params_.d = 1;
+      params_.h = input_desc[0].dims.d[2];
+      params_.w = input_desc[0].dims.d[3];
+    } else {
+      // d_dim == 5
+      params_.c = input_desc[0].dims.d[1];
+      params_.d = input_desc[0].dims.d[2];
+      params_.h = input_desc[0].dims.d[3];
+      params_.w = input_desc[0].dims.d[4];
+    }
+    params_.withSilu = with_silu_;
+    params_.dst = static_cast<half *>(outputs[1]);
+    params_.eleOut = static_cast<half *>(outputs[0]);
+    params_.srcX = static_cast<half const *>(inputs[0]);
+    params_.srcY = static_cast<half const *>(inputs[1]);
+    params_.gamma = scale_gpu_.get();
+    params_.beta = bias_gpu_.get();
+    params_.redBuffer = static_cast<float *>(workspace);
+    // params_.n = input_desc[0].dims.d[0];
+    // params_.h = input_desc[0].dims.d[2];
+    // params_.w = input_desc[0].dims.d[3];
+    // params_.c = input_desc[0].dims.d[1];
+    params_.groups = groups_;
+    params_.dhw = static_cast<int64_t>(params_.d) * params_.h * params_.w;
+    const int64_t blocksPerDHW =
+        findMaxDivisor(params_.dhw, static_cast<int64_t>(maxBlocksPerDHW));
+    params_.dhwPerBlock = divUp(params_.dhw, blocksPerDHW);
+    params_.cPerBlock = cPerBlock;
+    params_.cPerGroup = params_.c / params_.groups;
+    params_.dhwc = params_.dhw * params_.c;
+    params_.invDHWC = 1.F / static_cast<float>(params_.dhw * params_.cPerGroup);
+    params_.groupsPerBlock = cPerBlock / params_.cPerGroup;
+    params_.eps = eps_;
+
+    cudaMemsetAsync(params_.redBuffer, 0, ws_, stream);
+    prelnGroupNormNDHWCSum(params_, stream);
+    prelnGroupNormNDHWCScale(params_, stream);
+
+  } else {
+    // input not fp16
+    PADDLE_THROW(common::errors::Fatal(
+        "The PrelnGroupnormAct TRT Plugin's only support fp16 input"));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/preln_groupnorm_act_op_plugin.h

@@ -0,0 +1,201 @@
+/* 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. */
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/phi/kernels/group_norm_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+using phi::GroupNormNDHWCParams;
+class PrelnGroupnormActPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  PrelnGroupnormActPluginDynamic(const float* scale,
+                                 const int scale_num,
+                                 const float* bias,
+                                 const int bias_num,
+                                 float eps,
+                                 int groups,
+                                 bool with_silu,
+                                 bool with_fp16,
+                                 std::shared_ptr<void> scale_gpu = nullptr,
+                                 std::shared_ptr<void> bias_gpu = nullptr)
+      : scale_gpu_(scale_gpu),
+        bias_gpu_(bias_gpu),
+        groups_(groups),
+        eps_(eps),
+        with_silu_(with_silu),
+        with_fp16_(with_fp16) {
+    scale_.resize(scale_num);
+    bias_.resize(bias_num);
+    std::copy(scale, scale + scale_num, scale_.data());
+    std::copy(bias, bias + bias_num, bias_.data());
+    if (scale_gpu_ == nullptr) {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), scale_num * sizeof(float));
+      scale_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(
+          p, scale_.data(), scale_num * sizeof(float), cudaMemcpyHostToDevice);
+    }
+    if (bias_gpu_ == nullptr) {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), bias_num * sizeof(float));
+      bias_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(
+          p, bias_.data(), bias_num * sizeof(float), cudaMemcpyHostToDevice);
+    }
+  }
+
+  PrelnGroupnormActPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &groups_);
+    DeserializeValue(&serialData, &serialLength, &with_silu_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+
+    {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), scale_.size() * sizeof(float));
+      scale_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(p,
+                 scale_.data(),
+                 scale_.size() * sizeof(float),
+                 cudaMemcpyHostToDevice);
+    }
+    {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), bias_.size() * sizeof(float));
+      bias_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(p,
+                 bias_.data(),
+                 bias_.size() * sizeof(float),
+                 cudaMemcpyHostToDevice);
+    }
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    auto* ptr = new PrelnGroupnormActPluginDynamic(scale_.data(),
+                                                   scale_.size(),
+                                                   bias_.data(),
+                                                   bias_.size(),
+                                                   eps_,
+                                                   groups_,
+                                                   with_silu_,
+                                                   with_fp16_,
+                                                   scale_gpu_,
+                                                   bias_gpu_);
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "preln_groupnorm_act_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 2; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(scale_) + SerializedSize(bias_) +
+           SerializedSize(eps_) + SerializedSize(groups_) +
+           SerializedSize(with_silu_) + SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, groups_);
+    SerializeValue(&buffer, with_silu_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {
+    // sizeof(float2) * maxBatchSize * maxNumberOfGroup. float2
+    // contains two buffers for sum and squared sum;
+    ws_ = sizeof(float) * 2 * in[0].max.d[0] * groups_;
+  }
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return ws_;
+  }
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+  void terminate() TRT_NOEXCEPT override{};
+
+ private:
+  size_t ws_;
+  std::vector<float> scale_;
+  std::vector<float> bias_;
+  std::shared_ptr<void> scale_gpu_;
+  std::shared_ptr<void> bias_gpu_;
+  GroupNormNDHWCParams<__half> params_;
+  int groups_;
+  float eps_;
+  bool with_silu_;
+  bool with_fp16_;
+};
+
+class PrelnGroupnormActPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "preln_groupnorm_act_plugin_dynamic";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PrelnGroupnormActPluginDynamic(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(PrelnGroupnormActPluginDynamicCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/preln_layernorm_shift_partition_op.cu

@@ -0,0 +1,714 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2019-2022, NVIDIA CORPORATION.  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 <stdio.h>
+
+#include <cassert>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/preln_layernorm_shift_partition_op.h"
+#include "paddle/phi/kernels/layer_norm_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+#define FINAL_MASK 0xffffffff
+
+template <typename T>
+__inline__ __device__ T warpReduceSum(T val) {
+#pragma unroll
+  for (int mask = 16; mask > 0; mask >>= 1)
+    val += __shfl_xor_sync(FINAL_MASK, val, mask, 32);
+  return val;
+}
+
+/* Calculate the sum of all elements in a block */
+template <typename T>
+__inline__ __device__ T blockReduceSum(T val) {
+  static __shared__ T shared[32];
+  int lane = threadIdx.x & 0x1f;
+  int wid = threadIdx.x >> 5;
+
+  val = warpReduceSum<T>(val);
+
+  if (lane == 0) shared[wid] = val;
+
+  __syncthreads();
+
+  // Modify from blockDim.x << 5 to blockDim.x / 32. to prevent
+  // blockDim.x is not divided by 32
+  val = (threadIdx.x < (blockDim.x / 32.f)) ? shared[lane] : (T)(0.0f);
+  val = warpReduceSum<T>(val);
+
+  return val;
+}
+
+template <typename T>
+__global__ void preln_layernorm_shift_partition(T *out0,
+                                                T *out1,
+                                                const T *input0,
+                                                const T *input1,
+                                                const T *gamma,
+                                                const T *beta,
+                                                int batch,
+                                                int H,
+                                                int W,
+                                                int n,
+                                                int shift_size,
+                                                int window_size,
+                                                const float eps) {
+  int tid = threadIdx.x;
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  const int index = bid * n + tid;
+  float local_out = 0;
+  if (tid < n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+    local_out = static_cast<float>(__ldg(input0 + index));
+#else
+    local_out = static_cast<float>(input0[index]);
+#endif
+
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+    local_out += static_cast<float>(__ldg(input1 + index));
+#else
+    local_out += static_cast<float>(input1[index]);
+#endif
+    out0[index] = local_out;
+  }
+
+  mean = blockReduceSum<float>(local_out);
+  if (threadIdx.x == 0) {
+    s_mean = mean / n;
+  }
+  __syncthreads();
+
+  float diff = (tid < n) ? (local_out - s_mean) : 0.0f;
+  variance = blockReduceSum<float>(diff * diff);
+  if (threadIdx.x == 0) {
+    s_variance = variance / n + eps;
+  }
+  __syncthreads();
+
+  if (tid < n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+    out1[output_bid * n + tid] =
+        (T)(((local_out - s_mean) * rsqrtf(s_variance)) *
+                static_cast<float>(__ldg(&gamma[tid])) +
+            static_cast<float>(__ldg(&beta[tid])));
+#else
+    out1[output_bid * n + tid] =
+        (T)(((local_out - s_mean) * rsqrtf(s_variance)) *
+                static_cast<float>(gamma[tid]) +
+            static_cast<float>(beta[tid]));
+#endif
+  }
+}
+
+template <>
+__global__ void preln_layernorm_shift_partition(half2 *out0_ptr,
+                                                half2 *out1_ptr,
+                                                const half2 *input0_ptr,
+                                                const half2 *input1_ptr,
+                                                const half2 *gamma_ptr,
+                                                const half2 *beta_ptr,
+                                                int batch,
+                                                int H,
+                                                int W,
+                                                int n,
+                                                int shift_size,
+                                                int window_size,
+                                                const float eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  int tid = threadIdx.x;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  float2 local_out_fp2;
+
+  float local_out = 0.0f;
+  int id = bid * n + tid;
+  if (tid < n) {
+    half2 tmp = __hadd2(__ldg(input0_ptr + id), __ldg(input1_ptr + id));
+    local_out_fp2 = __half22float2(tmp);
+    local_out += local_out_fp2.x;
+    local_out += local_out_fp2.y;
+
+    out0_ptr[id] = tmp;
+  }
+
+  mean = blockReduceSum<float>(local_out);
+  if (threadIdx.x == 0) {
+    s_mean = mean / (n * 2);
+  }
+  __syncthreads();
+
+  if (tid < n) {
+    variance = (local_out_fp2.x - s_mean) * (local_out_fp2.x - s_mean);
+    variance += (local_out_fp2.y - s_mean) * (local_out_fp2.y - s_mean);
+  }
+  variance = blockReduceSum<float>(variance);
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / (n * 2) + eps);
+  }
+  __syncthreads();
+
+  if (tid < n) {
+    float2 gamma_val = __half22float2(__ldg(&gamma_ptr[tid]));
+    float2 beta_val = __half22float2(__ldg(&beta_ptr[tid]));
+    local_out_fp2.x =
+        (local_out_fp2.x - s_mean) * s_variance * gamma_val.x + beta_val.x;
+    local_out_fp2.y =
+        (local_out_fp2.y - s_mean) * s_variance * gamma_val.y + beta_val.y;
+    out1_ptr[output_bid * n + tid] = __float22half2_rn(local_out_fp2);
+  }
+#endif
+}
+
+#define kITE 4
+template <typename T>
+__global__ void preln_layernorm_shift_partition_v2(T *out0,
+                                                   T *out1,
+                                                   const T *__restrict input0,
+                                                   const T *__restrict input1,
+                                                   const T *__restrict gamma,
+                                                   const T *__restrict beta,
+                                                   int batch,
+                                                   int H,
+                                                   int W,
+                                                   int n,
+                                                   int shift_size,
+                                                   int window_size,
+                                                   const float eps) {
+  // constexpr int kITE = 4;
+  const int tid = threadIdx.x;
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  const int offset = bid * n;
+  const int output_offset = output_bid * n;
+
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  float local_out[kITE];
+
+  float sum = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    int index = offset + col_id;
+    if (col_id < n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+      local_out[i] = static_cast<float>(__ldg(input0 + index));
+#else
+      local_out[i] = static_cast<float>(input0[index]);
+#endif
+
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+      local_out[i] += static_cast<float>(__ldg(input1 + index));
+#else
+      local_out[i] += static_cast<float>(input1[index]);
+#endif
+      out0[index] = local_out[i];
+      sum += local_out[i];
+    }
+  }
+
+  mean = blockReduceSum<float>(sum);
+  if (tid == 0) {
+    s_mean = mean / n;
+  }
+  __syncthreads();
+
+  float var = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      float diff = local_out[i] - s_mean;
+      local_out[i] = diff;
+      var += diff * diff;
+    }
+  }
+
+  variance = blockReduceSum<float>(var);
+  if (tid == 0) {
+    s_variance = rsqrtf(variance / n + eps);
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+      out1[output_offset + col_id] =
+          (T)(local_out[i] * s_variance *
+                  static_cast<float>(__ldg(&gamma[col_id])) +
+              static_cast<float>(__ldg(&beta[col_id])));
+#else
+      out1[output_offset + col_id] =
+          (T)(local_out[i] * s_variance * static_cast<float>(gamma[col_id]) +
+              static_cast<float>(beta[col_id]));
+#endif
+    }
+  }
+}
+
+template <>
+__global__ void preln_layernorm_shift_partition_v2(
+    half2 *out0_ptr,
+    half2 *out1_ptr,
+    const half2 *__restrict input0_ptr,
+    const half2 *__restrict input1_ptr,
+    const half2 *__restrict gamma_ptr,
+    const half2 *__restrict beta_ptr,
+    int batch,
+    int H,
+    int W,
+    int n,
+    int shift_size,
+    int window_size,
+    const float eps) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  // constexpr int ite = 4;
+  const int tid = threadIdx.x;
+  const int batch_offset = blockIdx.z * gridDim.y * gridDim.x;
+  const int bid = batch_offset + blockIdx.y * gridDim.x + blockIdx.x;
+  const int shifted_H_idx =
+      (shift_size != 0) ? ((blockIdx.y - shift_size + gridDim.y) % gridDim.y)
+                        : blockIdx.y;
+  const int shifted_W_idx =
+      (shift_size != 0) ? ((blockIdx.x - shift_size + gridDim.x) % gridDim.x)
+                        : blockIdx.x;
+  const int window_H_idx = shifted_H_idx / window_size;
+  const int window_W_idx = shifted_W_idx / window_size;
+  const int stride_of_window_H = W / window_size;
+  const int window_idx = window_H_idx * stride_of_window_H + window_W_idx;
+  const int idx_in_window = (shifted_H_idx % window_size) * window_size +
+                            (shifted_W_idx % window_size);
+  const int output_bid =
+      batch_offset + window_idx * window_size * window_size + idx_in_window;
+  const int offset = bid * n;
+  const int output_offset = output_bid * n;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  half2 local_out_half2[kITE];
+
+  // float sum = 0.0f;
+  half2 sum = __float2half2_rn(0.0f);
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      int index = offset + col_id;
+      local_out_half2[i] = __ldg(input0_ptr + index);
+      local_out_half2[i] =
+          __hadd2(local_out_half2[i], __ldg(input1_ptr + index));
+      out0_ptr[i] = local_out_half2[i];
+      sum += local_out_half2[i];
+    }
+  }
+
+  mean = blockReduceSum<float>(static_cast<float>(sum.x + sum.y));
+  if (threadIdx.x == 0) {
+    s_mean = mean / (n * 2);
+  }
+  __syncthreads();
+
+  float var = 0.0f;
+  half2 s_mean_2 = __float2half2_rn(s_mean);
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      local_out_half2[i] = local_out_half2[i] - s_mean_2;
+      float v1 = static_cast<float>(local_out_half2[i].x);
+      float v2 = static_cast<float>(local_out_half2[i].y);
+      var += v1 * v1 + v2 * v2;
+    }
+  }
+
+  variance = blockReduceSum<float>(var);
+  if (threadIdx.x == 0) {
+    s_variance = rsqrtf(variance / (n * 2) + eps);
+  }
+  __syncthreads();
+
+  half2 s_var_2 = __float2half2_rn(s_variance);
+#pragma unroll
+  for (int i = 0; i < kITE; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      out1_ptr[output_offset + col_id] =
+          local_out_half2[i] * s_var_2 * __ldg(&gamma_ptr[col_id]) +
+          __ldg(&beta_ptr[col_id]);
+    }
+  }
+#endif
+}
+
+template <typename T>
+void invokePrelnLayernormShiftPartition(T *out0,
+                                        T *out1,
+                                        const T *input0,
+                                        const T *input1,
+                                        const T *gamma,
+                                        const T *beta,
+                                        int batch,
+                                        int H,
+                                        int W,
+                                        int n,
+                                        int shift_size,
+                                        int window_size,
+                                        const float eps,
+                                        cudaStream_t stream) {
+  dim3 grid(W, H, batch);
+  int blockSize = (n + 31) / 32 * 32;
+  if (blockSize >= 768) {
+    blockSize = ((blockSize / 4) + 31) / 32 * 32;
+    preln_layernorm_shift_partition_v2<T>
+        <<<grid, blockSize, 0, stream>>>(out0,
+                                         out1,
+                                         input0,
+                                         input1,
+                                         gamma,
+                                         beta,
+                                         batch,
+                                         H,
+                                         W,
+                                         n,
+                                         shift_size,
+                                         window_size,
+                                         eps);
+  } else {
+    preln_layernorm_shift_partition<T>
+        <<<grid, blockSize, 0, stream>>>(out0,
+                                         out1,
+                                         input0,
+                                         input1,
+                                         gamma,
+                                         beta,
+                                         batch,
+                                         H,
+                                         W,
+                                         n,
+                                         shift_size,
+                                         window_size,
+                                         eps);
+  }
+}
+
+template <>
+void invokePrelnLayernormShiftPartition(half *out0,
+                                        half *out1,
+                                        const half *input0,
+                                        const half *input1,
+                                        const half *gamma,
+                                        const half *beta,
+                                        int batch,
+                                        int H,
+                                        int W,
+                                        int n,
+                                        int shift_size,
+                                        int window_size,
+                                        const float eps,
+                                        cudaStream_t stream) {
+  dim3 grid(W, H, batch);
+  int blockSize = n / 2;
+  blockSize = (blockSize + 31) / 32 * 32;
+
+  if ((batch * H * W >= 512 && blockSize >= 768) || blockSize > 1024) {
+    blockSize = ((blockSize / 4) + 31) / 32 * 32;
+    preln_layernorm_shift_partition_v2<<<grid, blockSize, 0, stream>>>(
+        reinterpret_cast<half2 *>(out0),
+        reinterpret_cast<half2 *>(out1),
+        (const half2 *)input0,
+        (const half2 *)input1,
+        (const half2 *)gamma,
+        (const half2 *)beta,
+        batch,
+        H,
+        W,
+        n / 2,
+        shift_size,
+        window_size,
+        eps);
+  } else {
+    preln_layernorm_shift_partition<<<grid, blockSize, 0, stream>>>(
+        reinterpret_cast<half2 *>(out0),
+        reinterpret_cast<half2 *>(out1),
+        (const half2 *)input0,
+        (const half2 *)input1,
+        (const half2 *)gamma,
+        (const half2 *)beta,
+        batch,
+        H,
+        W,
+        n / 2,
+        shift_size,
+        window_size,
+        eps);
+  }
+}
+
+template <typename T>
+static void convertAndCopy(const std::vector<float> &host, T *dev) {
+  T *host_ptr = new T[host.size()];
+  std::transform(host.begin(), host.end(), host_ptr, [](float x) {
+    return static_cast<T>(x);
+  });
+  cudaMemcpy(dev, host_ptr, sizeof(T) * host.size(), cudaMemcpyHostToDevice);
+  delete host_ptr;
+}
+
+void PrelnLnormShiftPartitionPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+PrelnLnormShiftPartitionPluginDynamic::PrelnLnormShiftPartitionPluginDynamic(
+    const float *gamma,
+    const float *beta,
+    const int param_num,
+    int shift_size,
+    int window_size,
+    int input_resolution,
+    float eps,
+    bool with_fp16,
+    std::shared_ptr<void> gamma_dev,
+    std::shared_ptr<void> beta_dev)
+    : with_fp16_(with_fp16),
+      window_size_(window_size),
+      shift_size_(shift_size),
+      input_resolution_(input_resolution),
+      eps_(eps),
+      param_num_(param_num),
+      gamma_dev_(gamma_dev),
+      beta_dev_(beta_dev) {
+  beta_.resize(param_num);
+  gamma_.resize(param_num);
+  std::copy(gamma, gamma + param_num, gamma_.data());
+  std::copy(beta, beta + param_num, beta_.data());
+  int type_size = with_fp16 ? sizeof(half) : sizeof(float);
+  if (gamma_dev_ == nullptr) {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    gamma_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16)
+      convertAndCopy(gamma_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(gamma_, reinterpret_cast<float *>(p));
+  }
+  if (beta_dev_ == nullptr) {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    beta_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16)
+      convertAndCopy(beta_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(beta_, reinterpret_cast<float *>(p));
+  }
+}
+
+PrelnLnormShiftPartitionPluginDynamic::PrelnLnormShiftPartitionPluginDynamic(
+    void const *serialData, size_t serialLength) {
+  DeserializeValue(&serialData, &serialLength, &beta_);
+  DeserializeValue(&serialData, &serialLength, &gamma_);
+  DeserializeValue(&serialData, &serialLength, &param_num_);
+  DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  DeserializeValue(&serialData, &serialLength, &shift_size_);
+  DeserializeValue(&serialData, &serialLength, &window_size_);
+  DeserializeValue(&serialData, &serialLength, &input_resolution_);
+  DeserializeValue(&serialData, &serialLength, &eps_);
+  int type_size = with_fp16_ ? sizeof(half) : sizeof(float);
+  {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    gamma_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16_)
+      convertAndCopy(gamma_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(gamma_, reinterpret_cast<float *>(p));
+  }
+  {
+    void *p;
+    cudaMalloc(reinterpret_cast<void **>(&p), param_num_ * type_size);
+    beta_dev_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16_)
+      convertAndCopy(beta_, reinterpret_cast<half *>(p));
+    else
+      convertAndCopy(beta_, reinterpret_cast<float *>(p));
+  }
+}
+
+bool PrelnLnormShiftPartitionPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return in.type == nvinfer1::DataType::kHALF &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    } else {
+      return in.type == nvinfer1::DataType::kFLOAT &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType PrelnLnormShiftPartitionPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  return input_types[0];
+}
+
+nvinfer1::DimsExprs PrelnLnormShiftPartitionPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  if (output_index == 0) return inputs[0];
+
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = expr_builder.operation(
+      nvinfer1::DimensionOperation::kFLOOR_DIV,
+      *expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
+                              *inputs[0].d[0],
+                              *inputs[0].d[1]),
+      *expr_builder.constant(window_size_ * window_size_));
+  ret.d[1] = expr_builder.constant(window_size_ * window_size_);
+  ret.d[2] = inputs[0].d[2];
+  return ret;
+}
+
+int PrelnLnormShiftPartitionPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  auto input_type = input_desc[0].type;
+  int batch = input_dims.d[0];
+  int emb_dim = input_dims.d[2];
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(3)
+        << "TRT Plugin DataType selected. PreLayernormShiftPartition-->fp32";
+    invokePrelnLayernormShiftPartition(
+        reinterpret_cast<float *>(outputs[0]),
+        reinterpret_cast<float *>(outputs[1]),
+        reinterpret_cast<const float *>(inputs[0]),
+        reinterpret_cast<const float *>(inputs[1]),
+        reinterpret_cast<const float *>(gamma_dev_.get()),
+        reinterpret_cast<const float *>(beta_dev_.get()),
+        batch,
+        input_resolution_,
+        input_resolution_,
+        emb_dim,
+        shift_size_,
+        window_size_,
+        eps_,
+        stream);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(3)
+        << "TRT Plugin DataType selected. PreLayernormShiftPartition-->half";
+    invokePrelnLayernormShiftPartition(
+        reinterpret_cast<half *>(outputs[0]),
+        reinterpret_cast<half *>(outputs[1]),
+        reinterpret_cast<const half *>(inputs[0]),
+        reinterpret_cast<const half *>(inputs[1]),
+        reinterpret_cast<const half *>(gamma_dev_.get()),
+        reinterpret_cast<const half *>(beta_dev_.get()),
+        batch,
+        input_resolution_,
+        input_resolution_,
+        emb_dim,
+        shift_size_,
+        window_size_,
+        eps_,
+        stream);
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/preln_layernorm_shift_partition_op.h

@@ -0,0 +1,158 @@
+// 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class PrelnLnormShiftPartitionPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  PrelnLnormShiftPartitionPluginDynamic(
+      const float* gamma,
+      const float* beta,
+      const int param_num,
+      int shift_size,
+      int window_size,
+      int input_resolution,
+      float eps,
+      bool with_fp16,
+      std::shared_ptr<void> gamma_dev = nullptr,
+      std::shared_ptr<void> beta_dev = nullptr);
+
+  PrelnLnormShiftPartitionPluginDynamic(void const* serialData,
+                                        size_t serialLength);
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new PrelnLnormShiftPartitionPluginDynamic(gamma_.data(),
+                                                     beta_.data(),
+                                                     beta_.size(),
+                                                     shift_size_,
+                                                     window_size_,
+                                                     input_resolution_,
+                                                     eps_,
+                                                     with_fp16_,
+                                                     gamma_dev_,
+                                                     beta_dev_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "prelnlnorm_shift_partition_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 2; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(beta_) + SerializedSize(gamma_) +
+           SerializedSize(param_num_) + SerializedSize(with_fp16_) +
+           SerializedSize(shift_size_) + SerializedSize(window_size_) +
+           SerializedSize(input_resolution_) + SerializedSize(eps_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, beta_);
+    SerializeValue(&buffer, gamma_);
+    SerializeValue(&buffer, param_num_);
+    SerializeValue(&buffer, with_fp16_);
+    SerializeValue(&buffer, shift_size_);
+    SerializeValue(&buffer, window_size_);
+    SerializeValue(&buffer, input_resolution_);
+    SerializeValue(&buffer, eps_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  bool with_fp16_;
+  std::vector<float> gamma_;
+  std::vector<float> beta_;
+  int window_size_;
+  int shift_size_;
+  int input_resolution_;
+  int param_num_;
+  float eps_;
+  std::shared_ptr<void> gamma_dev_;
+  std::shared_ptr<void> beta_dev_;
+};
+
+class PrelnLnormShiftPartitionPluginDynamicCreator
+    : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "prelnlnorm_shift_partition_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new PrelnLnormShiftPartitionPluginDynamic(serial_data,
+                                                     serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(PrelnLnormShiftPartitionPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/preln_residual_bias_plugin.h

@@ -0,0 +1,343 @@
+// 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+using half = phi::dtype::float16;
+
+class PrelnResidualBiasPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit PrelnResidualBiasPluginDynamic(const float* bias,
+                                          const float* scale,
+                                          const half* ele_bias,
+                                          int bias_size,
+                                          int scale_size,
+                                          int ele_bias_size,
+                                          const float eps,
+                                          bool with_fp16)
+      : bias_size_(bias_size),
+        scale_size_(scale_size),
+        ele_bias_size_(ele_bias_size),
+        eps_(eps) {
+    with_fp16_ = with_fp16;
+    bias_.resize(bias_size);
+    scale_.resize(scale_size);
+
+    fp16_bias_.resize(bias_size);
+    fp16_scale_.resize(scale_size);
+
+    fp16_ele_bias_.resize(ele_bias_size);
+    std::copy(ele_bias, ele_bias + ele_bias_size, fp16_ele_bias_.data());
+    std::copy(bias, bias + bias_size, bias_.data());
+    std::copy(scale, scale + scale_size, scale_.data());
+    for (int i = 0; i < bias_size; i++) {
+      fp16_bias_[i] = static_cast<half>(bias[i]);
+    }
+    for (int i = 0; i < scale_size; i++) {
+      fp16_scale_[i] = static_cast<half>(scale[i]);
+    }
+  }
+
+  explicit PrelnResidualBiasPluginDynamic(const float* bias,
+                                          const float* scale,
+                                          const float* ele_bias,
+                                          int bias_size,
+                                          int scale_size,
+                                          int ele_bias_size,
+                                          const float eps,
+                                          bool with_fp16)
+      : bias_size_(bias_size),
+        scale_size_(scale_size),
+        ele_bias_size_(ele_bias_size),
+        eps_(eps) {
+    with_fp16_ = with_fp16;
+    bias_.resize(bias_size);
+    scale_.resize(scale_size);
+
+    fp32_ele_bias_.resize(ele_bias_size);
+    std::copy(ele_bias, ele_bias + ele_bias_size, fp32_ele_bias_.data());
+    std::copy(bias, bias + bias_size, bias_.data());
+    std::copy(scale, scale + scale_size, scale_.data());
+  }
+
+  PrelnResidualBiasPluginDynamic(void const* serial_data,
+                                 size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &bias_);
+    DeserializeValue(&serial_data, &serial_length, &fp16_bias_);
+    DeserializeValue(&serial_data, &serial_length, &scale_);
+    DeserializeValue(&serial_data, &serial_length, &fp16_scale_);
+    DeserializeValue(&serial_data, &serial_length, &fp32_ele_bias_);
+    DeserializeValue(&serial_data, &serial_length, &fp16_ele_bias_);
+    DeserializeValue(&serial_data, &serial_length, &bias_size_);
+    DeserializeValue(&serial_data, &serial_length, &scale_size_);
+    DeserializeValue(&serial_data, &serial_length, &ele_bias_size_);
+    DeserializeValue(&serial_data, &serial_length, &eps_);
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+  }
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nb_inputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nb_outputs) const TRT_NOEXCEPT override;
+
+  int 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 override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+
+ private:
+  // bias for fp32 mode
+  std::vector<float> bias_;
+  // bias for fp16 mode
+  std::vector<half> fp16_bias_;
+  // scale for fp32 mode
+  std::vector<float> scale_;
+  // scale for fp16 mode
+  std::vector<half> fp16_scale_;
+  std::vector<float> fp32_ele_bias_;
+  std::vector<half> fp16_ele_bias_;
+
+  float* bias_gpu_{nullptr};
+  half* fp16_bias_gpu_{nullptr};
+  float* scale_gpu_{nullptr};
+  half* fp16_scale_gpu_{nullptr};
+  void* ele_bias_gpu_{nullptr};
+
+  int bias_size_;
+  int scale_size_;
+  int ele_bias_size_;
+
+  float eps_;
+  bool with_fp16_;
+};
+
+class PrelnResidualBiasPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override;
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override;
+};
+
+class PIRPrelnResidualBiasPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit PIRPrelnResidualBiasPluginDynamic(const float* bias,
+                                             const float* scale,
+                                             const half* ele_bias,
+                                             int bias_size,
+                                             int scale_size,
+                                             int ele_bias_size,
+                                             const float eps,
+                                             bool with_fp16)
+      : bias_size_(bias_size),
+        scale_size_(scale_size),
+        ele_bias_size_(ele_bias_size),
+        eps_(eps) {
+    with_fp16_ = with_fp16;
+    bias_.resize(bias_size);
+    scale_.resize(scale_size);
+
+    fp16_bias_.resize(bias_size);
+    fp16_scale_.resize(scale_size);
+
+    fp16_ele_bias_.resize(ele_bias_size);
+    std::copy(ele_bias, ele_bias + ele_bias_size, fp16_ele_bias_.data());
+    std::copy(bias, bias + bias_size, bias_.data());
+    std::copy(scale, scale + scale_size, scale_.data());
+    for (int i = 0; i < bias_size; i++) {
+      fp16_bias_[i] = static_cast<half>(bias[i]);
+    }
+    for (int i = 0; i < scale_size; i++) {
+      fp16_scale_[i] = static_cast<half>(scale[i]);
+    }
+  }
+
+  explicit PIRPrelnResidualBiasPluginDynamic(const float* bias,
+                                             const float* scale,
+                                             const float* ele_bias,
+                                             int bias_size,
+                                             int scale_size,
+                                             int ele_bias_size,
+                                             const float eps,
+                                             bool with_fp16)
+      : bias_size_(bias_size),
+        scale_size_(scale_size),
+        ele_bias_size_(ele_bias_size),
+        eps_(eps) {
+    with_fp16_ = with_fp16;
+    bias_.resize(bias_size);
+    scale_.resize(scale_size);
+
+    fp32_ele_bias_.resize(ele_bias_size);
+    std::copy(ele_bias, ele_bias + ele_bias_size, fp32_ele_bias_.data());
+    std::copy(bias, bias + bias_size, bias_.data());
+    std::copy(scale, scale + scale_size, scale_.data());
+  }
+
+  PIRPrelnResidualBiasPluginDynamic(void const* serial_data,
+                                    size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &bias_);
+    DeserializeValue(&serial_data, &serial_length, &fp16_bias_);
+    DeserializeValue(&serial_data, &serial_length, &scale_);
+    DeserializeValue(&serial_data, &serial_length, &fp16_scale_);
+    DeserializeValue(&serial_data, &serial_length, &fp32_ele_bias_);
+    DeserializeValue(&serial_data, &serial_length, &fp16_ele_bias_);
+    DeserializeValue(&serial_data, &serial_length, &bias_size_);
+    DeserializeValue(&serial_data, &serial_length, &scale_size_);
+    DeserializeValue(&serial_data, &serial_length, &ele_bias_size_);
+    DeserializeValue(&serial_data, &serial_length, &eps_);
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+  }
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nb_inputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nb_outputs) const TRT_NOEXCEPT override;
+
+  int 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 override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+
+ private:
+  // bias for fp32 mode
+  std::vector<float> bias_;
+  // bias for fp16 mode
+  std::vector<half> fp16_bias_;
+  // scale for fp32 mode
+  std::vector<float> scale_;
+  // scale for fp16 mode
+  std::vector<half> fp16_scale_;
+  std::vector<float> fp32_ele_bias_;
+  std::vector<half> fp16_ele_bias_;
+
+  float* bias_gpu_{nullptr};
+  half* fp16_bias_gpu_{nullptr};
+  float* scale_gpu_{nullptr};
+  half* fp16_scale_gpu_{nullptr};
+  void* ele_bias_gpu_{nullptr};
+
+  int bias_size_;
+  int scale_size_;
+  int ele_bias_size_;
+
+  float eps_;
+  bool with_fp16_;
+};
+
+class PIRPrelnResidualBiasPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override;
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override;
+};
+REGISTER_TRT_PLUGIN_V2(PrelnResidualBiasPluginDynamicCreator);
+REGISTER_TRT_PLUGIN_V2(PIRPrelnResidualBiasPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/prompt_tuning_emb_layernorm_varseqlen_kernel_hface.cu

@@ -0,0 +1,204 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+// AFFILIATES. 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 <cuda.h>
+#include <cassert>
+#include <cstring>
+#include <iostream>
+#include <vector>
+
+#include "NvInfer.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/common.cuh"
+#include "paddle/fluid/inference/tensorrt/plugin/prompt_tuning_emb_layernorm_varseqlen_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T, unsigned TPB>
+__global__ void prompt_tuning_embKernel(int32_t B,
+                                        int32_t ld,
+                                        int32_t const* inputIds0,
+                                        int32_t const* inputIds1,
+                                        int32_t const* inputIds2,
+                                        T const* dense_vector,
+                                        float const* beta,
+                                        float const* gamma,
+                                        T const* mIdsEmbDev0,
+                                        T const* mIdsEmbDev1,
+                                        T const* mIdsEmbDev2,
+                                        int32_t IdsSize0,
+                                        int32_t IdsSize1,
+                                        int32_t IdsSize2,
+                                        T* output,
+                                        int32_t* new_pos_id) {
+  cub::Sum pairSum;
+  int32_t const s = blockIdx.x;
+  int32_t const b = blockIdx.y;
+  int32_t const sumS = inputIds0[b];
+  int32_t const s_b = inputIds0[b + 1] - inputIds0[b];
+
+  int32_t const new_sumS = sumS + b;
+
+  // new pos_id: Add an id to each sentence
+  new_pos_id[b] = new_sumS;
+
+  // last id
+  if (b == B - 1) {
+    new_pos_id[B] = inputIds0[B] + B;
+  }
+
+  T const rld = T(1.f) / T(ld);
+  int32_t const seqPos = sumS + s;
+  int32_t const out_seqPos = new_sumS + s + 1;
+  int32_t const new_out_seqPos = new_sumS + s;
+
+  kvp<T> threadData(0, 0);
+
+  int32_t const new_outoffset = new_out_seqPos * ld;
+  int32_t const prompt_tuning_offset = new_sumS * ld;
+  int32_t const dense_vector_offset = b * ld;
+
+  if (s < s_b) {
+    extern __shared__ int32_t word_id[];
+    if (threadIdx.x == 0) {
+      if (static_cast<int32_t const*>(inputIds1)[seqPos] < 0 ||
+          static_cast<int32_t const*>(inputIds1)[seqPos] >= IdsSize1) {
+        printf(
+            "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+            "table: ID < 0 or ID > max ");
+        return;
+      } else {
+        word_id[0] = static_cast<int32_t const*>(inputIds1)[seqPos];
+      }
+
+      if (static_cast<int32_t const*>(inputIds2)[seqPos] < 0 ||
+          static_cast<int32_t const*>(inputIds2)[seqPos] >= IdsSize2) {
+        printf(
+            "Error!!!!!!(embLayerNormVarSeqlenPlugin): ID cannot be lookup "
+            "table: ID < 0 or ID > max ");
+        return;
+      } else {
+        word_id[1] = static_cast<int32_t const*>(inputIds2)[seqPos];
+      }
+    }
+    __syncthreads();
+
+    // 2. load pos/tok/word embeddings and add them together
+    // offset into embeddings is given by wordId * hidden_size
+    int32_t const poffset = blockIdx.x * ld;
+    int32_t const outoffset = out_seqPos * ld;
+
+    // the output offset is given by b * (S*hidden_size) + s * hidden_size
+
+    for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+      T p(mIdsEmbDev0[poffset + it]);  // pos id
+      T val = p;
+      int32_t const offset0 = word_id[0] * ld;
+      val += mIdsEmbDev1[offset0 + it];
+      int32_t const offset1 = word_id[1] * ld;
+      val += mIdsEmbDev2[offset1 + it];
+      output[outoffset + it] = val;
+      T const rldval = rld * val;
+      threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+    }
+    // 3. layer norm on the sum
+    layerNorm<T, T, float, TPB>(threadData, ld, outoffset, beta, gamma, output);
+  } else if (s == s_b) {
+    for (int32_t it = threadIdx.x; it < ld; it += TPB) {
+      T val = dense_vector[dense_vector_offset + it];
+      output[prompt_tuning_offset + it] = val;
+      T const rldval = rld * val;
+      threadData = pairSum(threadData, kvp<T>(rldval, rldval * val));
+      // 3. layer norm on the sum
+    }
+    layerNorm<T, T, float, TPB>(
+        threadData, ld, prompt_tuning_offset, beta, gamma, output);
+
+  } else {
+    return;  // This CTA has nothing to do
+  }
+}
+
+template <typename T>
+int32_t prompt_tuning_emb(cudaStream_t stream,
+                          int32_t ld,
+                          int32_t B,
+                          int32_t S,
+                          int const* inputIds0,
+                          int const* inputIds1,
+                          int const* inputIds2,
+                          T const* dense_vector,
+                          int32_t nbLookupTables,
+                          float const* beta,
+                          float const* gamma,
+                          T const* mIdsEmbDev0,
+                          T const* mIdsEmbDev1,
+                          T const* mIdsEmbDev2,
+                          int32_t IdsSize0,
+                          int32_t IdsSize1,
+                          int32_t IdsSize2,
+                          T* output,
+                          int32_t* new_pos_id) {
+  constexpr int32_t tpb = 256;
+  dim3 const grid(S, B, 1);
+  dim3 const block(tpb, 1, 1);
+  size_t cache_size = sizeof(int32_t) * (nbLookupTables - 1);
+  prompt_tuning_embKernel<T, tpb>
+      <<<grid, block, cache_size, stream>>>(B,
+                                            ld,
+                                            inputIds0,
+                                            inputIds1,
+                                            inputIds2,
+                                            dense_vector,
+                                            beta,
+                                            gamma,
+                                            mIdsEmbDev0,
+                                            mIdsEmbDev1,
+                                            mIdsEmbDev2,
+                                            IdsSize0,
+                                            IdsSize1,
+                                            IdsSize2,
+                                            output,
+                                            new_pos_id);
+  return cudaPeekAtLastError();
+}
+
+template int32_t prompt_tuning_emb<half>(cudaStream_t,
+                                         int32_t,
+                                         int32_t,
+                                         int32_t,
+                                         int32_t const*,
+                                         int32_t const*,
+                                         int32_t const*,
+                                         half const*,
+                                         int32_t,
+                                         float const*,
+                                         float const*,
+                                         half const*,
+                                         half const*,
+                                         half const*,
+                                         int32_t,
+                                         int32_t,
+                                         int32_t,
+                                         half*,
+                                         int32_t*);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/prompt_tuning_emb_layernorm_varseqlen_plugin.cu

@@ -0,0 +1,562 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+// AFFILIATES. 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 "paddle/fluid/inference/tensorrt/plugin/prompt_tuning_emb_layernorm_varseqlen_plugin.h"
+#include <cuda.h>
+#include <cstring>
+#include <vector>
+#include "NvInfer.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+constexpr size_t threadsPerCta128 = 2 * 2 * 32;
+constexpr size_t threadsPerCta256 = 1 * 4 * 32;
+constexpr size_t threadsPerCta384 = 1 * 8 * 32;
+// The number of xmmas in the M dimension. We use one uint32_t per XMMA in the M
+// dimension: (s + 16*warps_m - 1) / (16*warps_m);
+constexpr size_t xmmasM128 = 4;
+constexpr size_t xmmasM256 = 16;
+constexpr size_t xmmasM384 = 24;
+// Packed mask size per batch. Layout is XMMAS_M * THREADS_PER_CTA.
+constexpr size_t packedMaskSize128 = xmmasM128 * threadsPerCta128;
+constexpr size_t packedMaskSize256 = xmmasM256 * threadsPerCta256;
+constexpr size_t packedMaskSize384 = xmmasM384 * threadsPerCta384;
+char const* PTUNING_EMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE{"1"};
+char const* PTUNING_EMB_LAYER_NORM_VAR_SEQLEN_NAME{
+    "PromptTuningEmbLayerNormVarlenPluginDynamic"};
+// Static class fields initialization
+nvinfer1::PluginFieldCollection
+    TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator::mFC{};
+std::vector<nvinfer1::PluginField>
+    TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator::mPluginAttributes;
+
+TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::
+    TrtPromptTuningEmbLayerNormVarSeqlenPluginBase(
+        std::string const& name,
+        nvinfer1::DataType const type,
+        nvinfer1::Weights const& beta,
+        nvinfer1::Weights const& gamma,
+        const std::vector<nvinfer1::Weights>& IdsEmb)
+    : mLayerName(name),
+      mLd(beta.count),
+      mType(type),
+      mIdsEmb_(IdsEmb),
+      nbLookupTables_(static_cast<int>(IdsEmb.size())) {
+  // Assuming Weights.count is the number of elements and not bytes
+  assert(beta.count == gamma.count);
+  mBeta.convertAndCopy(beta, nvinfer1::DataType::kFLOAT);
+  mGamma.convertAndCopy(gamma, nvinfer1::DataType::kFLOAT);
+  copyToDevice(&mGamma, sizeof(float) * mGamma.count, &mGammaDev);
+  copyToDevice(&mBeta, sizeof(float) * mBeta.count, &mBetaDev);
+  for (size_t i = 0; i < mIdsEmb_.size(); ++i) {
+    assert(mIdsEmb_[i].count % mLd == 0);
+    mIdsVocabSize.push_back(int32_t(mIdsEmb_[i].count / mLd));
+    WeightsWithOwnership tem_weight;
+    tem_weight.convertAndCopy(mIdsEmb_[i], mType);
+    void* cudaMem{nullptr};
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaMalloc(&cudaMem, getWeightsSize(tem_weight, mType)));
+    PADDLE_ENFORCE_GPU_SUCCESS(cudaMemcpy(cudaMem,
+                                          tem_weight.values,
+                                          getWeightsSize(tem_weight, mType),
+                                          cudaMemcpyHostToDevice));
+    mIdsEmbPtrs.push_back(cudaMem);
+  }
+}
+
+TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::
+    TrtPromptTuningEmbLayerNormVarSeqlenPluginBase(std::string const& name,
+                                                   void const* data,
+                                                   size_t length)
+    : mLayerName(name),
+      mGammaDev(nullptr),
+      mBetaDev(nullptr),
+      mIdsEmbPtrs{},
+      mIdsEmb_{} {
+  // Deserialize in the same order as serialization
+  deserialize_value(&data, &length, &mType);
+  deserialize_value(&data, &length, &mLd);
+  deserialize_value(&data, &length, &nbLookupTables_);
+  for (int32_t i = 0; i < nbLookupTables_; ++i) {
+    int32_t tem;
+    deserialize_value(&data, &length, &tem);
+    mIdsVocabSize.push_back(tem);
+  }
+  char const* d = static_cast<char const*>(data);
+  mBeta.convertAndCopy(&d, mLd, nvinfer1::DataType::kFLOAT);
+  mGamma.convertAndCopy(&d, mLd, nvinfer1::DataType::kFLOAT);
+  for (int32_t i = 0; i < nbLookupTables_; ++i) {
+    nvinfer1::Weights pre_tem_weight;
+    pre_tem_weight.type = mType;
+    pre_tem_weight.count = mLd * size_t(mIdsVocabSize[i]);
+    const auto nbBytes = mLd * size_t(mIdsVocabSize[i]) * getElementSize(mType);
+    auto destBuf = new char[nbBytes];
+    pre_tem_weight.values = destBuf;
+    std::copy_n(d, nbBytes, destBuf);
+    d += nbBytes;
+    mIdsEmb_.push_back(pre_tem_weight);
+  }
+}
+
+TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::
+    TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace(
+        std::string const& name,
+        nvinfer1::DataType const type,
+        nvinfer1::Weights const& beta,
+        nvinfer1::Weights const& gamma,
+        const std::vector<nvinfer1::Weights>& IdsEmb)
+    : TrtPromptTuningEmbLayerNormVarSeqlenPluginBase(
+          name, type, beta, gamma, IdsEmb) {}
+
+TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::
+    TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace(std::string const& name,
+                                                    void const* data,
+                                                    size_t length)
+    : TrtPromptTuningEmbLayerNormVarSeqlenPluginBase(name, data, length) {
+  TRANSFORMER_DEBUG_MSG(
+      "TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace deserialize");
+}
+
+// IPluginV2DynamicExt Methods
+nvinfer1::IPluginV2DynamicExt*
+TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::clone() const noexcept {
+  TRANSFORMER_DEBUG_MSG(
+      "TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace clone");
+  auto p = new TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace(
+      mLayerName, mType, mBeta, mGamma, mIdsEmb_);
+  p->setPluginNamespace(mNamespace.c_str());
+  return p;
+}
+
+nvinfer1::DimsExprs
+TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::getOutputDimensions(
+    int32_t outputIndex,
+    nvinfer1::DimsExprs const* inputs,
+    int32_t nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) noexcept {
+  for (int i = 1; i < nbInputs - 2; ++i) {
+    assert(inputs[i].nbDims == 1);                 // seq length
+    assert(inputs[i].nbDims == inputs[1].nbDims);  // same shape
+  }
+  assert(inputs[0].nbDims == 1);  // pos_id: B+1
+  auto one = exprBuilder.constant(1);
+  auto Bplus1 = inputs[0].d[0];  // pos_id
+  auto B =
+      exprBuilder.operation(nvinfer1::DimensionOperation::kSUB, *Bplus1, *one);
+  if (outputIndex == 0) {
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 4;
+    ret.d[0] = exprBuilder.operation(nvinfer1::DimensionOperation::kSUM,
+                                     *inputs[1].d[0],
+                                     *B);  // sum of seq length
+    ret.d[1] = exprBuilder.constant(mLd);
+    ret.d[2] = exprBuilder.constant(1);
+    ret.d[3] = exprBuilder.constant(1);
+    return ret;
+  } else if (outputIndex == 1) {
+    // This is a hack: we just report some mask size and rely the plugins to
+    // play nicely together.
+    //      At runtime, depending on the actual maxSeqlen, the size might be
+    //      different.
+    int32_t maskSize_ = packedMaskSize384;
+    auto maskSize = exprBuilder.constant(maskSize_);
+    auto fp16maskSize =
+        exprBuilder.operation(nvinfer1::DimensionOperation::kPROD,
+                              *maskSize,
+                              *exprBuilder.constant(2));
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 2;
+    ret.d[0] = B;
+    ret.d[1] = fp16maskSize;
+    return ret;
+  } else if (outputIndex == 2) {
+    nvinfer1::DimsExprs ret;
+    ret.nbDims = 1;
+    ret.d[0] = exprBuilder.operation(nvinfer1::DimensionOperation::kSUM,
+                                     *inputs[nbInputs - 2].d[1],
+                                     *one);  // max seqlen
+    return ret;
+  } else if (outputIndex == 3) {
+    nvinfer1::DimsExprs ret = inputs[nbInputs - 2];  // new mask_id
+    ret.d[1] = exprBuilder.operation(
+        nvinfer1::DimensionOperation::kSUM, *inputs[nbInputs - 2].d[1], *one);
+    return ret;
+  } else if (outputIndex == 4) {
+    nvinfer1::DimsExprs ret = inputs[0];  // new pos_id
+    return ret;
+  }
+}
+
+bool TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::supportsFormatCombination(
+    int32_t pos,
+    nvinfer1::PluginTensorDesc const* inOut,
+    int32_t nbInputs,
+    int32_t nbOutputs) noexcept {
+  assert(nbOutputs == 5);
+  nvinfer1::PluginTensorDesc const& desc = inOut[pos];
+  if (desc.format != nvinfer1::TensorFormat::kLINEAR) {
+    return false;
+  }
+  if (pos == 0) {  // pos_id
+    return desc.dims.nbDims == 1 && desc.type == nvinfer1::DataType::kINT32;
+  }
+  if (pos == 1) {  //  input_id
+    return desc.dims.nbDims == 1 && desc.type == nvinfer1::DataType::kINT32;
+  }
+  nvinfer1::PluginTensorDesc const& prev = inOut[1];  // input_ids
+  if (1 < pos &&
+      pos < (nbInputs - 2)) {  // other ids: check it's the same as input_ids
+    return desc.type == prev.type && desc.dims.nbDims == 1 &&
+           desc.dims.d[0] == prev.dims.d[0];
+  }
+  if (pos == nbInputs - 2) {  // mask id
+    return desc.type == mType;
+  }
+  if (pos == nbInputs - 1) {  // dense vector
+    return desc.type == mType;
+  }
+  // embedded sequence
+  if (pos == nbInputs) {
+    return desc.type == mType && desc.dims.nbDims == 4 && desc.dims.d[2] == 1 &&
+           desc.dims.d[3] == 1;
+  }
+  // mask(HFace)
+  if (pos == nbInputs + 1) {
+    return desc.type == mType;
+  }
+  // max seqlen
+  if (pos == nbInputs + 2) {
+    return desc.type == mType;
+  }
+  // new mask_id
+  if (pos == nbInputs + 3) {
+    return desc.type == mType;
+  }
+  // new pos_id
+  if (pos == nbInputs + 4) {
+    return desc.dims.nbDims == 1 && desc.type == nvinfer1::DataType::kINT32;
+  }
+}
+
+void check_tensors(nvinfer1::DynamicPluginTensorDesc const* inputs,
+                   int32_t nbInputs,
+                   nvinfer1::DynamicPluginTensorDesc const* outputs,
+                   int32_t nbOutputs) noexcept {
+  // Validate input arguments
+  assert(nbOutputs == 5);
+  assert(inputs[0].desc.dims.nbDims == 1);
+  assert(inputs[0].desc.type == nvinfer1::DataType::kINT32);
+  for (int i = 1; i < nbInputs - 2; ++i) {
+    assert(inputs[i].desc.dims.nbDims == 1);
+    assert(inputs[i].desc.dims.d[0] == inputs[1].desc.dims.d[0]);
+    assert(inputs[i].desc.type == nvinfer1::DataType::kINT32);
+  }
+  assert(outputs[0].desc.dims.nbDims == 4);
+  assert(outputs[0].desc.dims.d[2] == 1);
+  assert(outputs[0].desc.dims.d[3] == 1);
+}
+
+void TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::configurePlugin(
+    nvinfer1::DynamicPluginTensorDesc const* inputs,
+    int32_t nbInputs,
+    nvinfer1::DynamicPluginTensorDesc const* outputs,
+    int32_t nbOutputs) noexcept {
+  TRANSFORMER_DEBUG_MSG(
+      "TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace configurePlugin");
+  check_tensors(inputs, nbInputs, outputs, nbOutputs);
+  assert(static_cast<size_t>(outputs[0].desc.dims.d[1]) ==
+         static_cast<size_t>(mLd));
+  int32_t const B = inputs[0].desc.dims.d[0] - 1;
+  // check mask
+  assert(outputs[1].desc.dims.nbDims == 2);
+  if (B > 0) {
+    assert(outputs[1].desc.dims.d[0] == B);
+  }
+  assert((outputs[1].desc.dims.d[1] == 2 * packedMaskSize384) ||
+         (outputs[1].desc.dims.d[1] == 2 * packedMaskSize128) ||
+         (outputs[1].desc.dims.d[1] == 2 * packedMaskSize256));
+  assert(outputs[0].desc.type == mType);
+  assert(outputs[1].desc.type == nvinfer1::DataType::kHALF);
+}
+
+size_t TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::getWorkspaceSize(
+    nvinfer1::PluginTensorDesc const* inputs,
+    int32_t nbInputs,
+    nvinfer1::PluginTensorDesc const* outputs,
+    int32_t nbOutputs) const noexcept {
+  return 0;
+}
+
+int32_t TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::enqueue(
+    nvinfer1::PluginTensorDesc const* inputDesc,
+    nvinfer1::PluginTensorDesc const* outputDesc,
+    void const* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) noexcept {
+  int32_t batchSize = inputDesc[0].dims.d[0] - 1;
+  // read out the maximum sequence length from the dummy input
+  int32_t const maxSeqlen = inputDesc[nbLookupTables_].dims.d[1] + 1;
+  int32_t S;
+  if (maxSeqlen <= 128) {
+    S = 128;
+  } else if (maxSeqlen <= 192) {
+    S = 192;
+  } else if (maxSeqlen <= 256) {
+    S = 256;
+  } else if (maxSeqlen <= 384) {
+    S = 384;
+  } else if (maxSeqlen <= 512) {
+    S = 512;
+  } else {
+    PADDLE_THROW(common::errors::Fatal("The max seqlenth is 512."));
+  }
+  const float* beta = mBetaDev.get();
+  const float* gamma = mGammaDev.get();
+
+  auto output = static_cast<half*>(outputs[0]);
+  auto new_pos_id = static_cast<int32_t*>(outputs[4]);
+  return prompt_tuning_emb<half>(stream,
+                                 static_cast<int32_t>(mLd),
+                                 batchSize,
+                                 S,
+                                 static_cast<int32_t const*>(inputs[0]),
+                                 static_cast<int32_t const*>(inputs[1]),
+                                 static_cast<int32_t const*>(inputs[2]),
+                                 static_cast<half const*>(inputs[4]),
+                                 nbLookupTables_,
+                                 beta,
+                                 gamma,
+                                 static_cast<half const*>(mIdsEmbPtrs[0]),
+                                 static_cast<half const*>(mIdsEmbPtrs[1]),
+                                 static_cast<half const*>(mIdsEmbPtrs[2]),
+                                 mIdsVocabSize[0],
+                                 mIdsVocabSize[1],
+                                 mIdsVocabSize[2],
+                                 output,
+                                 new_pos_id);
+}
+
+// IPluginV2Ext Methods
+nvinfer1::DataType
+TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::getOutputDataType(
+    int32_t index,
+    nvinfer1::DataType const* inputTypes,
+    int32_t nbInputs) const noexcept {
+  assert(mType == nvinfer1::DataType::kHALF);
+  if (index == 0) {
+    return mType;
+  } else if (index == 1) {
+    return mType;
+  } else if (index == 2) {
+    return mType;
+  } else if (index == 3) {
+    return mType;
+  } else if (index == 4) {
+    return nvinfer1::DataType::kINT32;
+  }
+}
+
+// IPluginV2 Methods
+char const* TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::getPluginType()
+    const noexcept {
+  return PTUNING_EMB_LAYER_NORM_VAR_SEQLEN_NAME;
+}
+
+char const* TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::getPluginVersion()
+    const noexcept {
+  return PTUNING_EMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
+}
+
+int32_t TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::getNbOutputs()
+    const noexcept {
+  return 5;
+}
+
+int32_t TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::initialize() noexcept {
+  TRANSFORMER_DEBUG_MSG(
+      "TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace initialize");
+  return 0;
+}
+
+void TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::terminate() noexcept {
+  TRANSFORMER_DEBUG_MSG(
+      "TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace terminate");
+}
+
+size_t TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::getSerializationSize()
+    const noexcept {
+  size_t const wordSize = getElementSize(mType);
+  return 2 * sizeof(float) * mLd                            // beta + gamma
+         + sizeof(mType)                                    //
+         + sizeof(mLd)                                      //
+         + mIdsVocabSize.size() * sizeof(mIdsVocabSize[0])  //
+         + wordSize * mLd *
+               accumulate(
+                   mIdsVocabSize.begin(), mIdsVocabSize.end(), 0)  // ids emb
+         + sizeof(nbLookupTables_);  // numbers of lookup_table
+}
+
+void TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::serialize(
+    void* buffer) const noexcept {
+  serialize_value(&buffer, mType);
+  serialize_value(&buffer, mLd);
+  serialize_value(&buffer, nbLookupTables_);
+  for (size_t i = 0; i < mIdsVocabSize.size(); ++i) {
+    serialize_value(&buffer, mIdsVocabSize[i]);
+  }
+  char* d = static_cast<char*>(buffer);
+  size_t const wordSize = getElementSize(mType);
+  serFromDev(&d, mBetaDev.get(), mLd);
+  serFromDev(&d, mGammaDev.get(), mLd);
+  for (size_t i = 0; i < mIdsEmbPtrs.size(); ++i) {
+    serFromDev(&d,
+               static_cast<char*>(mIdsEmbPtrs[i]),
+               mLd * mIdsVocabSize[i] * wordSize);
+  }
+}
+
+void TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::destroy() noexcept {
+  // This gets called when the network containing plugin is destroyed
+  mBetaDev.reset(nullptr);
+  mGammaDev.reset(nullptr);
+  for (size_t i = 0; i < mIdsEmbPtrs.size(); ++i) {
+    cudaFree(mIdsEmbPtrs[i]);
+  }
+  delete this;
+}
+
+void TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace::destroy() noexcept {
+  TRANSFORMER_DEBUG_MSG(
+      "TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace destroy");
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::destroy();
+}
+
+void TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::setPluginNamespace(
+    char const* libNamespace) noexcept {
+  mNamespace = libNamespace;
+}
+
+char const* TrtPromptTuningEmbLayerNormVarSeqlenPluginBase::getPluginNamespace()
+    const noexcept {
+  return mNamespace.c_str();
+}
+
+TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator::
+    TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator() = default;
+
+char const*
+TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator::getPluginName()
+    const noexcept {
+  return PTUNING_EMB_LAYER_NORM_VAR_SEQLEN_NAME;
+}
+
+char const*
+TrtPromptTuningEmbLayerNormVarSeqlenPluginHFaceCreator::getPluginVersion()
+    const noexcept {
+  return PTUNING_EMB_LAYER_NORM_VAR_SEQLEN_VERSION_HFACE;
+}
+
+nvinfer1::PluginFieldCollection const*
+TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator::
+    getFieldNames() noexcept {
+  return &mFC;
+}
+
+bool InitializeFields(nvinfer1::PluginFieldCollection const* fc,
+                      nvinfer1::Weights* beta,
+                      nvinfer1::Weights* gamma,
+                      std::vector<nvinfer1::Weights>* IdsEmb) {
+  bool output_fp16 = false;
+  for (int32_t i = 0; i < fc->nbFields; i++) {
+    std::string field_name(fc->fields[i].name);
+    if (field_name.compare("bert_embeddings_layernorm_beta") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building bert_embeddings_layernorm_beta...");
+      beta->values = fc->fields[i].data;
+      beta->count = fc->fields[i].length;
+      beta->type = fieldTypeToDataType(fc->fields[i].type);
+    }
+
+    if (field_name.compare("bert_embeddings_layernorm_gamma") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building bert_embeddings_layernorm_gamma...");
+      gamma->values = fc->fields[i].data;
+      gamma->count = fc->fields[i].length;
+      gamma->type = fieldTypeToDataType(fc->fields[i].type);
+    }
+
+    if (field_name.compare("output_fp16") == 0) {
+      TRANSFORMER_DEBUG_MSG("Building output_fp16...");
+      assert(fc->fields[i].type == nvinfer1::PluginFieldType::kINT32);
+      output_fp16 = static_cast<int32_t const*>(fc->fields[i].data)[0] != 0;
+    }
+    if (field_name.compare("bert_embeddings_word_embeddings_" +
+                           std::to_string(i - 3)) == 0) {
+      TRANSFORMER_DEBUG_MSG(
+          ("bert_embeddings_word_embeddings_" + std::to_string(i - 3)).c_str());
+      nvinfer1::Weights tem;
+      tem.values = fc->fields[i].data;
+      tem.count = fc->fields[i].length;
+      tem.type = fieldTypeToDataType(fc->fields[i].type);
+      IdsEmb->push_back(tem);
+    }
+  }
+  return output_fp16;
+}
+
+nvinfer1::IPluginV2*
+TrtPromptTuningEmbLayerNormVarSeqlenPluginHFaceCreator::createPlugin(
+    char const* name, nvinfer1::PluginFieldCollection const* fc) noexcept {
+  TRANSFORMER_DEBUG_MSG("PromptTuningEmbLayerNormVarSeqlenHFace createPlugin");
+  nvinfer1::Weights beta;
+  nvinfer1::Weights gamma;
+  std::vector<nvinfer1::Weights> IdsEmb;
+  bool output_fp16 = InitializeFields(fc, &beta, &gamma, &IdsEmb);
+  TRANSFORMER_DEBUG_MSG("Building the Plugin...");
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace* p =
+      new TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace(
+          name,
+          output_fp16 ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
+          beta,
+          gamma,
+          IdsEmb);
+  return p;
+}
+
+nvinfer1::IPluginV2*
+TrtPromptTuningEmbLayerNormVarSeqlenPluginHFaceCreator::deserializePlugin(
+    char const* name, void const* serialData, size_t serialLength) noexcept {
+  return new TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace(
+      name, serialData, serialLength);
+}
+
+void TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator::setPluginNamespace(
+    char const* libNamespace) noexcept {
+  mNamespace = libNamespace;
+}
+
+char const*
+TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator::getPluginNamespace()
+    const noexcept {
+  return mNamespace.c_str();
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/prompt_tuning_emb_layernorm_varseqlen_plugin.h

@@ -0,0 +1,189 @@
+// Copyright (c) 2023 PaddlePaddle Authors. All Rights Reserved.
+// SPDX-FileCopyrightText: Copyright (c) 1993-2023 NVIDIA CORPORATION &
+// AFFILIATES. 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.
+
+#pragma once
+#include <cuda.h>
+#include "NvInferPlugin.h"
+#include "NvInferRuntime.h"
+
+#include "paddle/fluid/inference/tensorrt/plugin/common/bertCommon.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/plugin.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/serialize.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T>
+int32_t prompt_tuning_emb(cudaStream_t,
+                          int32_t,
+                          int32_t,
+                          int32_t,
+                          int32_t const*,
+                          int32_t const*,
+                          int32_t const*,
+                          T const*,
+                          int32_t,
+                          float const*,
+                          float const*,
+                          T const*,
+                          T const*,
+                          T const*,
+                          int32_t,
+                          int32_t,
+                          int32_t,
+                          T*,
+                          int32_t*);
+class TrtPromptTuningEmbLayerNormVarSeqlenPluginBase
+    : public nvinfer1::IPluginV2DynamicExt {
+ public:
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginBase(
+      std::string const& name,
+      nvinfer1::DataType const type,
+      nvinfer1::Weights const& beta,
+      nvinfer1::Weights const& gamma,
+      const std::vector<nvinfer1::Weights>& ids_emb);
+
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginBase(std::string const& name,
+                                                 void const* data,
+                                                 size_t length);
+
+  // It doesn't make sense to make TrtPromptTuningEmbLayerNormVarSeqlenPlugin
+  // without arguments, so we delete default constructor.
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginBase() = delete;
+
+  // IPluginV2DynamicExt Methods
+  bool supportsFormatCombination(int32_t pos,
+                                 nvinfer1::PluginTensorDesc const* inOut,
+                                 int32_t nbInputs,
+                                 int32_t nbOutputs) noexcept override;
+  size_t getWorkspaceSize(nvinfer1::PluginTensorDesc const* inputs,
+                          int32_t nbInputs,
+                          nvinfer1::PluginTensorDesc const* outputs,
+                          int32_t nbOutputs) const noexcept override;
+
+  // IPluginV2Ext Methods
+  nvinfer1::DataType getOutputDataType(
+      int32_t index,
+      nvinfer1::DataType const* inputTypes,
+      int32_t nbInputs) const noexcept override;
+
+  // IPluginV2 Methods
+  char const* getPluginType() const noexcept override;
+  int32_t getNbOutputs() const noexcept override;
+  size_t getSerializationSize() const noexcept override;
+  void serialize(void* buffer) const noexcept override;
+  void destroy() noexcept override;
+  char const* getPluginNamespace() const noexcept override;
+  void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+ protected:
+  std::string const mLayerName;
+  std::string mNamespace;
+  cuda_unique_ptr<float> mGammaDev;
+  cuda_unique_ptr<float> mBetaDev;
+  std::vector<void*> mIdsEmbPtrs;
+  size_t mLd;  // leading dim = hidden size
+  std::vector<int32_t> mIdsVocabSize;
+  WeightsWithOwnership mBeta;
+  WeightsWithOwnership mGamma;
+  nvinfer1::DataType mType;
+  std::vector<nvinfer1::Weights> mIdsEmb_;
+  int32_t nbLookupTables_ = 0;
+};
+
+class TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace
+    : public TrtPromptTuningEmbLayerNormVarSeqlenPluginBase {
+ public:
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace(
+      std::string const& name,
+      nvinfer1::DataType const type,
+      nvinfer1::Weights const& beta,
+      nvinfer1::Weights const& gamma,
+      const std::vector<nvinfer1::Weights>& ids_emb);
+
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace(std::string const& name,
+                                                  void const* data,
+                                                  size_t length);
+
+  // It doesn't make sense to make TrtPromptTuningEmbLayerNormVarSeqlenPlugin
+  // without arguments, so we delete default constructor.
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginHFace() = delete;
+
+  // IPluginV2DynamicExt Methods
+  nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int32_t outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int32_t nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+  void configurePlugin(nvinfer1::DynamicPluginTensorDesc const* in,
+                       int32_t nbInputs,
+                       nvinfer1::DynamicPluginTensorDesc const* out,
+                       int32_t nbOutputs) noexcept override;
+  int32_t enqueue(nvinfer1::PluginTensorDesc const* inputDesc,
+                  nvinfer1::PluginTensorDesc const* outputDesc,
+                  void const* const* inputs,
+                  void* const* outputs,
+                  void* workspace,
+                  cudaStream_t stream) noexcept override;
+  // IPluginV2 Methods
+  int32_t initialize() noexcept override;
+  void terminate() noexcept override;
+  void destroy() noexcept override;
+  char const* getPluginVersion() const noexcept override;
+};
+
+class TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator
+    : public nvinfer1::IPluginCreator {
+ public:
+  TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator();
+
+  char const* getPluginName() const noexcept override;
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;
+
+  void setPluginNamespace(char const* pluginNamespace) noexcept override;
+
+  char const* getPluginNamespace() const noexcept override;
+
+ protected:
+  static nvinfer1::PluginFieldCollection mFC;
+  static std::vector<nvinfer1::PluginField> mPluginAttributes;
+  std::string mNamespace;
+};
+
+class TrtPromptTuningEmbLayerNormVarSeqlenPluginHFaceCreator
+    : public TrtPromptTuningEmbLayerNormVarSeqlenPluginBaseCreator {
+ public:
+  nvinfer1::IPluginV2* createPlugin(
+      char const* name,
+      const nvinfer1::PluginFieldCollection* fc) noexcept override;
+  char const* getPluginVersion() const noexcept override;
+  nvinfer1::IPluginV2* deserializePlugin(char const* name,
+                                         void const* serialData,
+                                         size_t serialLength) noexcept override;
+};
+
+REGISTER_TRT_PLUGIN_V2(TrtPromptTuningEmbLayerNormVarSeqlenPluginHFaceCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/qkv_to_context_plugin.cu

@@ -0,0 +1,533 @@
+// Copyright (c) 2018 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 <stdio.h>
+
+#include <cassert>
+#include <cub/cub.cuh>  // NOLINT
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/plugin/common/common.cuh"
+#include "paddle/fluid/inference/tensorrt/plugin/qkv_to_context_plugin.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin_utils.h"
+#include "paddle/phi/core/platform/device_context.h"
+#include "paddle/phi/kernels/funcs/blas/blas.h"
+#include "paddle/phi/kernels/funcs/multihead_matmul_functor.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+inline int round_up(int seq_len, int multiple = 32) {
+  PADDLE_ENFORCE_GT(
+      multiple,
+      0,
+      common::errors::InvalidArgument(
+          "multiple should be a positive number, but it's (%d)", multiple));
+  return ((seq_len + multiple - 1) / multiple) * multiple;
+}
+
+template <typename T>
+__global__ void reset_qk_bias(T *input, int real_seq_len, int seq_len) {
+  if (threadIdx.x < seq_len) {
+    int id = threadIdx.x + blockIdx.x * seq_len;
+    input[id] = threadIdx.x >= real_seq_len ? (T)0.0f : (T)1.0f;
+  }
+}
+
+template <typename T>
+__global__ void transpose_qkv_padding(
+    const T *src,  // (Batch, real_seq_len, 3 , head_num * size_per_head)
+    T *dst,        // (3 * batch * head_num * seq_len * size_per_head)
+    const int batch_size,
+    const int seq_len,
+    const int head_num,
+    const int size_per_head,
+    const int real_seq_len) {
+  // const dim3 grid(seq_len, batch, 3);
+  // const dim3 block(head_size, head_num, 1);
+  int qkv_id = blockIdx.z;
+  int batch_id = blockIdx.y;
+  int seq_id = blockIdx.x;
+  int head_id = threadIdx.y;
+  const int dst_offset =
+      qkv_id * batch_size * head_num * seq_len * size_per_head +
+      batch_id * head_num * seq_len * size_per_head +
+      head_id * seq_len * size_per_head + seq_id * size_per_head;
+  const int src_offset =
+      batch_id * real_seq_len * 3 * head_num * size_per_head +
+      seq_id * 3 * head_num * size_per_head +
+      qkv_id * head_num * size_per_head + head_id * size_per_head;
+  if (seq_id < real_seq_len) {
+    dst[threadIdx.x + dst_offset] = src[threadIdx.x + src_offset];
+  }
+}
+
+template <typename T>
+__global__ void transpose_qkv_unpadding(const T *src,
+                                        T *dst,
+                                        const int batch_size,
+                                        const int seq_len,
+                                        const int head_num,
+                                        const int size_per_head,
+                                        const int real_seq_len) {
+  int batch_id = blockIdx.y;
+  int seq_id = blockIdx.x;
+  int head_id = threadIdx.y;
+  const int src_offset = batch_id * head_num * seq_len * size_per_head +
+                         head_id * seq_len * size_per_head +
+                         seq_id * size_per_head;
+  const int dst_offset = batch_id * real_seq_len * head_num * size_per_head +
+                         seq_id * head_num * size_per_head +
+                         head_id * size_per_head;
+
+  dst[threadIdx.x + dst_offset] = src[threadIdx.x + src_offset];
+}
+
+#define LAUNCH_TRANSPOSE_KERNEL(TYPE, VECTOR_SIZE, PAD_TYPE)                \
+  do {                                                                      \
+    int h = head_size / VECTOR_SIZE;                                        \
+    const TYPE *input##VECTOR_SIZE = reinterpret_cast<const TYPE *>(input); \
+    TYPE *output##VECTOR_SIZE = reinterpret_cast<TYPE *>(output);           \
+    dim3 block(h, head_num, 1);                                             \
+    transpose_qkv_##PAD_TYPE<TYPE>                                          \
+        <<<grid, block, 0, stream>>>(input##VECTOR_SIZE,                    \
+                                     output##VECTOR_SIZE,                   \
+                                     batch,                                 \
+                                     seq_len,                               \
+                                     head_num,                              \
+                                     h,                                     \
+                                     real_seq_len);                         \
+  } while (0)
+
+inline void TransposePadding(const half *input,
+                             half *output,
+                             const int batch,
+                             const int seq_len,
+                             const int head_num,
+                             const int head_size,
+                             const int real_seq_len,
+                             cudaStream_t stream) {
+  const dim3 grid(seq_len, batch, 3);
+  if (head_size % 8 == 0) {
+    LAUNCH_TRANSPOSE_KERNEL(int4, 8, padding);
+  } else if (head_size % 2 == 0) {
+    LAUNCH_TRANSPOSE_KERNEL(half2, 2, padding);
+  } else {
+    LAUNCH_TRANSPOSE_KERNEL(half, 1, padding);
+  }
+}
+
+inline void TransposeUnPadding(const half *input,
+                               half *output,
+                               const int batch,
+                               const int seq_len,
+                               const int head_num,
+                               const int head_size,
+                               const int real_seq_len,
+                               cudaStream_t stream) {
+  const dim3 grid(real_seq_len, batch);
+  if (head_size % 8 == 0) {
+    LAUNCH_TRANSPOSE_KERNEL(int4, 8, unpadding);
+  } else if (head_size % 2 == 0) {
+    LAUNCH_TRANSPOSE_KERNEL(half2, 2, unpadding);
+  } else {
+    LAUNCH_TRANSPOSE_KERNEL(half, 1, unpadding);
+  }
+}
+
+int QkvToContextPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+nvinfer1::DimsExprs QkvToContextPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  // input[0], (B, S, 3 * N * H, 1, 1)
+  // input[1], (B, head_num, seq_len, seq_len)
+  // output, (B, seq_len, hidden)
+  PADDLE_ENFORCE_EQ(output_index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "There is only one output of the EmbEltwiseLayernorm, "
+                        "so the index should be zero,"
+                        "but it's (%d)",
+                        output_index));
+  PADDLE_ENFORCE_EQ(
+      nb_inputs,
+      2,
+      common::errors::InvalidArgument(
+          "The Input of the EmbEltwiseLayernorm should be 3, but we found "
+          "it has (%d) inputs",
+          nb_inputs));
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = inputs[0].d[0];
+  ret.d[1] = inputs[0].d[1];
+  ret.d[2] = expr_builder.constant(head_size_ * head_number_);
+  return ret;
+}
+
+void QkvToContextPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nb_inputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nb_outputs) TRT_NOEXCEPT {
+  auto input_dims = in[0].desc.dims;
+  int batch = input_dims.d[0];
+  int real_seq_len = input_dims.d[1];
+  int seq_len = round_up(real_seq_len, 8);
+  if (batch != -1 && real_seq_len != -1) {
+    int device_id = 0;
+    cudaGetDevice(&device_id);
+    auto *device_ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(GPUPlace(device_id)));
+    const phi::GPUContext &dev_ctx = *device_ctx;
+    auto stream = dev_ctx.stream();
+    tensor_.Resize({batch, seq_len, seq_len, head_number_});
+    if (in[0].desc.type == nvinfer1::DataType::kHALF) {
+      tensor_.Resize({batch, seq_len, seq_len, 1});
+      int blocks = batch * 1 * seq_len;
+      mask_half_ = reinterpret_cast<half *>(
+          tensor_.mutable_data<int16_t>(GPUPlace(device_id)));
+      reset_qk_bias<<<blocks, 1024, 0, stream>>>(
+          mask_half_, real_seq_len, seq_len);
+    } else if (in[0].desc.type == nvinfer1::DataType::kFLOAT) {
+      fake_qk_bias_ = reinterpret_cast<float *>(
+          tensor_.mutable_data<int32_t>(GPUPlace(device_id)));
+      int64_t size = sizeof(int32_t) * batch * seq_len * seq_len * head_number_;
+#ifdef PADDLE_WITH_HIP
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          hipMemsetAsync(fake_qk_bias_, 0, size, dev_ctx.stream()));
+#else
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaMemsetAsync(fake_qk_bias_, 0, size, dev_ctx.stream()));
+#endif
+    } else {
+      PADDLE_THROW(common::errors::Fatal(
+          "The QKV TRT Plugin's input type should be float or half."));
+    }
+  }
+}
+
+bool QkvToContextPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kFLOAT ||
+              in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+
+  if (pos == 1) {
+    return in.type == prev.type && in.format == prev.format;
+  }
+
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType QkvToContextPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      index,
+      0,
+      common::errors::InvalidArgument(
+          "The EmbEltwiseLayernorm Plugin only has one input, so the "
+          "index value should be 0, but get %d.",
+          index));
+  return input_types[0];
+}
+
+template <typename T>
+__global__ void apply_scale(T *data, T scale, int n) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  if (tid < n) {
+    data[tid] = data[tid] * scale;
+  }
+#endif
+}
+
+template <typename T>
+__global__ void broadcast(const T *src,
+                          T *dst,
+                          const int seq_len,
+                          const int head_num) {
+  int batch_id = blockIdx.x / (head_num * seq_len);
+  int dst_offset = blockIdx.x * seq_len;
+  if (threadIdx.x < seq_len) {
+    dst[threadIdx.x + dst_offset] = src[threadIdx.x + batch_id * seq_len];
+  }
+}
+
+template <typename T>
+__global__ void broadcast_batch_head_number(const T *src,
+                                            T *dst,
+                                            const int batch_size,
+                                            const int seq_len,
+                                            const int head_num) {
+  int batch_id = blockIdx.x % seq_len;
+  int dst_offset = blockIdx.x * seq_len;
+  if (threadIdx.x < seq_len) {
+    dst[threadIdx.x + dst_offset] = src[threadIdx.x + batch_id * seq_len];
+  }
+}
+
+int QkvToContextPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  int input_num = ProductDim(input_dims);
+  // input[0], (B, S, 3 * N * H, 1, 1)
+  int batch = input_dims.d[0];
+  int seq_len = input_dims.d[1];
+  phi::DenseTensor multihead_temp_tensor;
+  int scratch_size = batch * head_number_ * seq_len * seq_len * 1;
+
+  int device_id;
+  cudaGetDevice(&device_id);
+  multihead_temp_tensor.Resize({scratch_size + input_num});
+
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. QkvToContext-->fp32";
+    auto *multihead_temp_data =
+        multihead_temp_tensor.mutable_data<float>(GPUPlace(device_id));
+    auto *qkptr = multihead_temp_data;
+    auto *tptr = multihead_temp_data + scratch_size;
+
+    const float *input0_data = static_cast<const float *>(inputs[0]);
+    // fit to [batch, head_num, length, length] + [batch, 1, 1, length]
+    phi::DenseTensor temp_qk_bias_tensor;
+    float *qk_bias = const_cast<float *>(static_cast<const float *>(inputs[1]));
+    if (ProductDim(input_desc[1].dims) == (batch * seq_len)) {
+      temp_qk_bias_tensor.Resize({batch, head_number_, seq_len, seq_len});
+      auto *temp_qk_bias =
+          temp_qk_bias_tensor.mutable_data<float>(GPUPlace(device_id));
+      int grid = batch * head_number_ * seq_len;
+      int block = round_up(seq_len);
+      broadcast<<<grid, block, 0, stream>>>(
+          static_cast<const float *>(inputs[1]),
+          temp_qk_bias,
+          seq_len,
+          head_number_);
+      qk_bias = temp_qk_bias;
+    }
+    // fit to [batch, head_num, length, length] + [1, 1, length, length]
+    if (ProductDim(input_desc[1].dims) == (seq_len * seq_len)) {
+      temp_qk_bias_tensor.Resize({batch, head_number_, seq_len, seq_len});
+      auto *temp_qk_bias = reinterpret_cast<float *>(
+          temp_qk_bias_tensor.mutable_data<float>(GPUPlace(device_id)));
+      int grid = batch * head_number_ * seq_len;
+      int block = round_up(seq_len);
+      broadcast_batch_head_number<<<grid, block, 0, stream>>>(
+          static_cast<const float *>(inputs[1]),
+          temp_qk_bias,
+          batch,
+          seq_len,
+          head_number_);
+      qk_bias = temp_qk_bias;
+    }
+    // fake qk_bias
+    if (ProductDim(input_desc[1].dims) == ProductDim(input_desc[0].dims)) {
+      qk_bias = fake_qk_bias_;
+    }
+    const float *input1_data = static_cast<const float *>(qk_bias);
+    // BxSx3xNxH => tptr: 3xBxNxSxH.
+    TransposeQKV(
+        batch, seq_len, head_size_, head_number_, input0_data, tptr, stream);
+
+    auto *device_ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(GPUPlace(device_id)));
+
+    const phi::GPUContext &dev_ctx = *device_ctx;
+    phi::funcs::MultiheadGPUComputeFunctor<float> multihead_compute_func;
+    multihead_compute_func(dev_ctx,
+                           batch,
+                           seq_len,
+                           head_number_,
+                           head_size_,
+                           qkptr,
+                           input1_data,
+                           false,
+                           tptr,
+                           scale_,
+                           static_cast<float>(0.0));
+
+    int grid = batch * head_number_ * seq_len;
+    int block = head_size_;
+    float *output = static_cast<float *>(outputs[0]);
+    transpose<float><<<grid, block, 0, stream>>>(
+        tptr, output, batch, seq_len, head_number_, head_size_);
+
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. QkvToContext-->fp16";
+    int real_seq_len = seq_len;
+    int need_padding = false;
+    // fake qk_bias
+    if (ProductDim(input_desc[1].dims) == ProductDim(input_desc[0].dims)) {
+      seq_len = round_up(real_seq_len, 8);
+      scratch_size = batch * head_number_ * seq_len * seq_len * 1;
+      input_num = batch * seq_len * 3 * head_number_ * head_size_;
+      multihead_temp_tensor.Resize({scratch_size + input_num});
+      need_padding = (real_seq_len != seq_len) ? true : false;
+    }
+    auto *multihead_temp_data =
+        multihead_temp_tensor.mutable_data<int16_t>(  // NOLINT
+            GPUPlace(device_id));
+
+    half *qkptr = reinterpret_cast<half *>(multihead_temp_data);
+    half *tptr = qkptr + scratch_size;
+
+    const half *input0_data = static_cast<const half *>(inputs[0]);
+    // fit to [batch, head_num, length, length] + [batch, 1, 1, length]
+    phi::DenseTensor temp_qk_bias_tensor;
+    half *qk_bias = const_cast<half *>(static_cast<const half *>(inputs[1]));
+    if (ProductDim(input_desc[1].dims) == (batch * seq_len)) {
+      temp_qk_bias_tensor.Resize({batch, head_number_, seq_len, seq_len});
+      auto *temp_qk_bias = reinterpret_cast<half *>(
+          temp_qk_bias_tensor.mutable_data<int16_t>(GPUPlace(device_id)));
+      int grid = batch * head_number_ * seq_len;
+      int block = round_up(seq_len);
+      broadcast<<<grid, block, 0, stream>>>(
+          static_cast<const half *>(inputs[1]),
+          temp_qk_bias,
+          seq_len,
+          head_number_);
+      qk_bias = temp_qk_bias;
+    }
+    // fit to [batch, head_num, length, length] + [1, 1, length, length]
+    if (ProductDim(input_desc[1].dims) == (seq_len * seq_len)) {
+      temp_qk_bias_tensor.Resize({batch, head_number_, seq_len, seq_len});
+      auto *temp_qk_bias = reinterpret_cast<half *>(
+          temp_qk_bias_tensor.mutable_data<int16_t>(GPUPlace(device_id)));
+      int grid = batch * head_number_ * seq_len;
+      int block = round_up(seq_len);
+      broadcast_batch_head_number<<<grid, block, 0, stream>>>(
+          static_cast<const half *>(inputs[1]),
+          temp_qk_bias,
+          batch,
+          seq_len,
+          head_number_);
+      qk_bias = temp_qk_bias;
+    }
+    // padding:    mask_half_ = [1.0,....1.0...1.0....,0.0f]
+    // no_padding: mask_half_ = [1.0,....1.0,.........,1.0f]
+    bool bias_is_mask = false;
+    if (ProductDim(input_desc[1].dims) == ProductDim(input_desc[0].dims)) {
+      qk_bias = mask_half_;
+      bias_is_mask = true;
+    }
+    const half *input1_data = static_cast<const half *>(qk_bias);
+    // BxSx3xNxH => tptr: 3xBxNxSxH.
+    if (need_padding) {
+      PADDLE_ENFORCE_GPU_SUCCESS(
+          cudaMemsetAsync(tptr, 0, sizeof(half) * input_num, stream));
+      TransposePadding(input0_data,
+                       tptr,
+                       batch,
+                       seq_len,
+                       head_number_,
+                       head_size_,
+                       real_seq_len,
+                       stream);
+    } else {
+      TransposeQKV(
+          batch, seq_len, head_size_, head_number_, input0_data, tptr, stream);
+    }
+
+    auto *device_ctx = static_cast<phi::GPUContext *>(
+        phi::DeviceContextPool::Instance().Get(GPUPlace(device_id)));
+
+    int n_q = seq_len * head_number_ * head_size_ * batch;
+    constexpr int threads = 128;
+    int blocks = (n_q + threads - 1) / threads;
+
+    apply_scale<<<blocks, threads, 0, stream>>>(
+        tptr, static_cast<half>(scale_), n_q);
+
+    const phi::GPUContext &dev_ctx = *device_ctx;
+    phi::funcs::MultiheadGPUComputeFunctor<half> multihead_compute_func;
+    multihead_compute_func(dev_ctx,
+                           batch,
+                           seq_len,
+                           head_number_,
+                           head_size_,
+                           qkptr,
+                           input1_data,
+                           bias_is_mask,
+                           tptr,
+                           half(1.),
+                           half(0.0));
+
+    int grid = batch * head_number_ * seq_len;
+    int block = head_size_;
+    half *output = static_cast<half *>(outputs[0]);
+    if (need_padding) {
+      TransposeUnPadding(tptr,
+                         output,
+                         batch,
+                         seq_len,
+                         head_number_,
+                         head_size_,
+                         real_seq_len,
+                         stream);
+    } else {
+      transpose<half><<<grid, block, 0, stream>>>(
+          tptr, output, batch, seq_len, head_number_, head_size_);
+    }
+  } else {
+    PADDLE_THROW(common::errors::Fatal(
+        "The QKV TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/qkv_to_context_plugin.h

@@ -0,0 +1,178 @@
+// Copyright (c) 2020 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.
+
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class QkvToContextPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit QkvToContextPluginDynamic(
+      int hidden, int head_number, int head_size, float scale, bool with_fp16)
+      : hidden_(hidden),
+        head_number_(head_number),
+        head_size_(head_size),
+        scale_(scale) {
+    with_fp16_ = with_fp16;
+  }
+
+  QkvToContextPluginDynamic(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &hidden_);
+    DeserializeValue(&serial_data, &serial_length, &head_number_);
+    DeserializeValue(&serial_data, &serial_length, &head_size_);
+    DeserializeValue(&serial_data, &serial_length, &scale_);
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new QkvToContextPluginDynamic(
+        hidden_, head_number_, head_size_, scale_, with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "qkv_to_context_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(hidden_) + SerializedSize(head_number_) +
+           SerializedSize(head_size_) + SerializedSize(scale_) +
+           SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, hidden_);
+    SerializeValue(&buffer, head_number_);
+    SerializeValue(&buffer, head_size_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nb_inputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nb_outputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  int hidden_;
+  int head_number_;
+  int head_size_;
+  float scale_;
+  phi::DenseTensor tensor_;
+  half* mask_half_;
+  float* fake_qk_bias_;
+};
+
+class QkvToContextPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  QkvToContextPluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "qkv_to_context_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new QkvToContextPluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_;
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+REGISTER_TRT_PLUGIN_V2(QkvToContextPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/recover_padding_plugin.cu

@@ -0,0 +1,162 @@
+/* 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 "paddle/fluid/inference/tensorrt/plugin/recover_padding_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+__global__ void RecoverPaddingKernel(const half* input0,
+                                     const int32_t* input1,
+                                     half* output) {
+  int word_id = blockIdx.x * gridDim.y + blockIdx.y;
+  int32_t sequence_length = input1[blockIdx.x + 1] - input1[blockIdx.x];
+  if (blockIdx.y < sequence_length) {
+    output[word_id * gridDim.z * blockDim.x + blockIdx.z * blockDim.x +
+           threadIdx.x] =
+        input0[(input1[blockIdx.x] + blockIdx.y) * gridDim.z * blockDim.x +
+               blockIdx.z * blockDim.x + threadIdx.x];
+  } else {
+    output[word_id * gridDim.z * blockDim.x + blockIdx.z * blockDim.x +
+           threadIdx.x] = 0;
+  }
+}
+
+nvinfer1::DataType RecoverPaddingPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  return input_types[0];
+}
+
+nvinfer1::DimsExprs RecoverPaddingPlugin::getOutputDimensions(
+    int outputIndex,
+    const nvinfer1::DimsExprs* inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs output_dims{};
+  output_dims.nbDims = 3;
+  const auto* one = exprBuilder.constant(1);
+  output_dims.d[0] = exprBuilder.operation(
+      nvinfer1::DimensionOperation::kSUB, *inputs[1].d[0], *one);
+  output_dims.d[1] = inputs[2].d[1];
+  output_dims.d[2] = inputs[0].d[1];
+  return output_dims;
+}
+
+bool RecoverPaddingPlugin::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nbInputs,
+                    3,
+                    common::errors::InvalidArgument("Must have 3 inputs, "
+                                                    "but got %d input(s). ",
+                                                    nbInputs));
+  PADDLE_ENFORCE_EQ(nbOutputs,
+                    getNbOutputs(),
+                    common::errors::InvalidArgument("Must have 1 output, "
+                                                    "but got %d output(s). ",
+                                                    nbOutputs));
+  if (pos == 1) {  // PosId
+    return inOut[pos].type == nvinfer1::DataType::kINT32 &&
+           inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
+  } else if (pos == 2) {  // mask_id
+    return inOut[pos].type == nvinfer1::DataType::kFLOAT &&
+           inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
+  } else {
+    return inOut[pos].type == nvinfer1::DataType::kHALF &&
+           inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
+  }
+  // return (inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format
+  // == nvinfer1::TensorFormat::kLINEAR)||
+  // (inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format ==
+  // nvinfer1::TensorFormat::kLINEAR)||
+  // (inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format ==
+  // nvinfer1::TensorFormat::kCHW32);
+}
+
+void RecoverPaddingPlugin::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* outputs,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+void RecoverPaddingPlugin::attachToContext(
+    cudnnContext* cudnnContext,
+    cublasContext* cublasContext,
+    nvinfer1::IGpuAllocator* gpuAllocator) TRT_NOEXCEPT {}
+
+void RecoverPaddingPlugin::detachFromContext() TRT_NOEXCEPT {}
+
+void RecoverPaddingPlugin::terminate() TRT_NOEXCEPT {}
+
+int RecoverPaddingPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+                                  const nvinfer1::PluginTensorDesc* outputDesc,
+                                  const void* const* inputs,
+                                  void* const* outputs,
+                                  void* workspace,
+                                  cudaStream_t stream) TRT_NOEXCEPT {
+  const auto input0_desc = inputDesc[0];
+  const auto input1_desc = inputDesc[1];
+  const auto input2_desc = inputDesc[2];
+  const half* input0 = static_cast<const half*>(inputs[0]);
+  const int32_t* input1 =
+      static_cast<const int32_t*>(inputs[1]);  // pos_id_tensor
+  half* output = static_cast<half*>(outputs[0]);
+  const int32_t vector_length = input0_desc.dims.d[1];
+  int32_t num_threads;
+  if (vector_length < 1024) {
+    num_threads = vector_length;
+  } else {
+    if (vector_length % 512 == 0) {
+      num_threads = 512;
+    } else if (vector_length % 256 == 0) {
+      num_threads = 256;
+    } else if (vector_length % 128 == 0) {
+      num_threads = 128;
+    } else if (vector_length % 64 == 0) {
+      num_threads = 64;
+    } else if (vector_length % 32 == 0) {
+      num_threads = 32;
+    } else if (vector_length % 16 == 0) {
+      num_threads = 16;
+    } else if (vector_length % 8 == 0) {
+      num_threads = 8;
+    } else if (vector_length % 4 == 0) {
+      num_threads = 4;
+    } else if (vector_length % 2 == 0) {
+      num_threads = 2;
+    } else {
+      num_threads = 1;
+    }
+  }
+  const dim3 num_blocks(
+      input1_desc.dims.d[0] - 1,
+      input2_desc.dims.d[1],
+      vector_length / num_threads);  //  batches, max sequence length
+                                     //  (mask_id.dims.d[1]),
+                                     //  input.dims.d[1]/***
+  RecoverPaddingKernel<<<num_blocks, num_threads, 0, stream>>>(
+      input0, input1, output);
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/recover_padding_plugin.h

@@ -0,0 +1,141 @@
+/* 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. */
+
+#pragma once
+#include <cassert>
+#include <string>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class RecoverPaddingPlugin : public DynamicPluginTensorRT {
+ public:
+  RecoverPaddingPlugin() {}
+
+  RecoverPaddingPlugin(void const* serial_data, size_t serial_length) {}
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    RecoverPaddingPlugin* ptr = new RecoverPaddingPlugin();
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "recover_padding_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+
+  int initialize() TRT_NOEXCEPT { return 0; }
+  void terminate() TRT_NOEXCEPT;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* inputs,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* outputs,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  void attachToContext(cudnnContext* cudnnContext,
+                       cublasContext* cublasContext,
+                       nvinfer1::IGpuAllocator* gpuAllocator)
+      TRT_NOEXCEPT override;
+
+  void detachFromContext() TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ protected:
+  size_t getSerializationSize() const TRT_NOEXCEPT override { return 0; }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {}
+};
+
+class RecoverPaddingPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  RecoverPaddingPluginCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "recover_padding_plugin";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(
+      const char* name, const nvinfer1::PluginFieldCollection* plugin_field)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         void const* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    RecoverPaddingPlugin* obj =
+        new RecoverPaddingPlugin(serial_data, serial_length);
+    obj->setPluginNamespace(name);
+    return obj;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+};
+REGISTER_TRT_PLUGIN_V2(RecoverPaddingPluginCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/remove_padding_plugin.cu

@@ -0,0 +1,154 @@
+/* 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 "paddle/fluid/inference/tensorrt/plugin/remove_padding_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+__global__ void RemovePaddingKernel(const half* input0,
+                                    const int32_t* input1,
+                                    half* output) {
+  int word_id = blockIdx.x * gridDim.y + blockIdx.y;
+  int32_t sequence_length = input1[blockIdx.x + 1] - input1[blockIdx.x];
+  if (blockIdx.y < sequence_length) {
+    output[(input1[blockIdx.x] + blockIdx.y) * gridDim.z * blockDim.x +
+           blockIdx.z * blockDim.x + threadIdx.x] =
+        input0[word_id * gridDim.z * blockDim.x + blockIdx.z * blockDim.x +
+               threadIdx.x];
+  }
+}
+
+nvinfer1::DataType RemovePaddingPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  return input_types[0];
+}
+
+nvinfer1::DimsExprs RemovePaddingPlugin::getOutputDimensions(
+    int outputIndex,
+    const nvinfer1::DimsExprs* inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs output_dims{};
+  output_dims.nbDims = 4;
+  output_dims.d[0] = inputs[2].d[0];
+  output_dims.d[1] = inputs[0].d[2];
+  output_dims.d[2] = exprBuilder.constant(1);
+  output_dims.d[3] = exprBuilder.constant(1);
+
+  return output_dims;
+}
+
+bool RemovePaddingPlugin::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nbInputs,
+                    3,
+                    common::errors::InvalidArgument("Must have 3 inputs, "
+                                                    "but got %d input(s). ",
+                                                    nbInputs));
+  PADDLE_ENFORCE_EQ(nbOutputs,
+                    getNbOutputs(),
+                    common::errors::InvalidArgument("Must have 1 output, "
+                                                    "but got %d output(s). ",
+                                                    nbOutputs));
+  if (pos == 1 || pos == 2) {  // pos_id, work_id
+    return inOut[pos].type == nvinfer1::DataType::kINT32 &&
+           inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
+  } else {
+    return inOut[pos].type == nvinfer1::DataType::kHALF &&
+           inOut[pos].format == nvinfer1::TensorFormat::kLINEAR;
+  }
+  // return (inOut[pos].type == nvinfer1::DataType::kFLOAT && inOut[pos].format
+  // == nvinfer1::TensorFormat::kLINEAR)||
+  // (inOut[pos].type == nvinfer1::DataType::kHALF && inOut[pos].format ==
+  // nvinfer1::TensorFormat::kLINEAR)||
+  // (inOut[pos].type == nvinfer1::DataType::kINT8 && inOut[pos].format ==
+  // nvinfer1::TensorFormat::kCHW32);
+}
+
+void RemovePaddingPlugin::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* outputs,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+void RemovePaddingPlugin::attachToContext(cudnnContext* cudnnContext,
+                                          cublasContext* cublasContext,
+                                          nvinfer1::IGpuAllocator* gpuAllocator)
+    TRT_NOEXCEPT {}
+
+void RemovePaddingPlugin::detachFromContext() TRT_NOEXCEPT {}
+
+void RemovePaddingPlugin::terminate() TRT_NOEXCEPT {}
+
+int RemovePaddingPlugin::enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+                                 const nvinfer1::PluginTensorDesc* outputDesc,
+                                 const void* const* inputs,
+                                 void* const* outputs,
+                                 void* workspace,
+                                 cudaStream_t stream) TRT_NOEXCEPT {
+  const half* input0 = static_cast<const half*>(inputs[0]);
+  const int32_t* input1 =
+      static_cast<const int32_t*>(inputs[1]);  // pos_id_tensor
+  half* output = static_cast<half*>(outputs[0]);
+  const auto input0_desc = inputDesc[0];
+  const int32_t vector_length = input0_desc.dims.d[2];
+  int32_t num_threads;
+  if (vector_length < 1024) {
+    num_threads = vector_length;
+  } else {
+    if (vector_length % 512 == 0) {
+      num_threads = 512;
+    } else if (vector_length % 256 == 0) {
+      num_threads = 256;
+    } else if (vector_length % 128 == 0) {
+      num_threads = 128;
+    } else if (vector_length % 64 == 0) {
+      num_threads = 64;
+    } else if (vector_length % 32 == 0) {
+      num_threads = 32;
+    } else if (vector_length % 16 == 0) {
+      num_threads = 16;
+    } else if (vector_length % 8 == 0) {
+      num_threads = 8;
+    } else if (vector_length % 4 == 0) {
+      num_threads = 4;
+    } else if (vector_length % 2 == 0) {
+      num_threads = 2;
+    } else {
+      num_threads = 1;
+    }
+  }
+  const dim3 num_blocks(
+      input0_desc.dims.d[0],
+      input0_desc.dims.d[1],
+      vector_length /
+          num_threads);  //  batches, max sequence length, input0.dims.d[2]/***
+
+  RemovePaddingKernel<<<num_blocks, num_threads, 0, stream>>>(
+      input0, input1, output);
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/remove_padding_plugin.h

@@ -0,0 +1,141 @@
+/* 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. */
+
+#pragma once
+#include <cassert>
+#include <string>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class RemovePaddingPlugin : public DynamicPluginTensorRT {
+ public:
+  RemovePaddingPlugin() {}
+
+  RemovePaddingPlugin(void const* serial_data, size_t serial_length) {}
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    RemovePaddingPlugin* ptr = new RemovePaddingPlugin();
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "remove_padding_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+
+  int initialize() TRT_NOEXCEPT { return 0; }
+  void terminate() TRT_NOEXCEPT;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* inputs,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* outputs,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  void attachToContext(cudnnContext* cudnnContext,
+                       cublasContext* cublasContext,
+                       nvinfer1::IGpuAllocator* gpuAllocator)
+      TRT_NOEXCEPT override;
+
+  void detachFromContext() TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ protected:
+  size_t getSerializationSize() const TRT_NOEXCEPT override { return 0; }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {}
+};
+
+class RemovePaddingPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  RemovePaddingPluginCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "remove_padding_plugin";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(
+      const char* name, const nvinfer1::PluginFieldCollection* plugin_field)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         void const* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    RemovePaddingPlugin* obj =
+        new RemovePaddingPlugin(serial_data, serial_length);
+    obj->setPluginNamespace(name);
+    return obj;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+};
+REGISTER_TRT_PLUGIN_V2(RemovePaddingPluginCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/reverse_roll_op_plugin.cu

@@ -0,0 +1,260 @@
+// Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+// Copyright (c) 2019-2022, NVIDIA CORPORATION.  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 "paddle/fluid/inference/tensorrt/plugin/reverse_roll_op_plugin.h"
+#include <vector>
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+/*******************  invokeReverseRoll  ***********************/
+// src is [batch*window_num, window_len, dim]
+// dst is [batch, H, W, dim] + rolled
+// grid(W, H, batch)
+// block(min(1024, dim))
+
+template <typename T>
+__global__ void reverse_roll(T *dst,
+                             const T *src,
+                             const int batch,
+                             const int window_num,
+                             const int window_len,
+                             const int window_size,
+                             const int H,
+                             const int W,
+                             const int shift_size,
+                             const int dim) {
+  const int batch_idx = blockIdx.z;
+  const int H_idx_shifted = (blockIdx.y + shift_size) % H;
+  const int W_idx_shifted = (blockIdx.x + shift_size) % W;
+  const int H_idx = blockIdx.y;
+  const int W_idx = blockIdx.x;
+  const int window_idx =
+      H_idx / window_size * (W / window_size) + W_idx / window_size;
+  const int idx_in_window =
+      (H_idx % window_size) * window_size + (W_idx % window_size);
+  const int input_offset =
+      (batch_idx * window_num + window_idx) * window_len + idx_in_window;
+  const int output_offset = (batch_idx * H + H_idx_shifted) * W + W_idx_shifted;
+  for (int tid = threadIdx.x; tid < dim; tid += blockDim.x) {
+    dst[output_offset * dim + tid] = src[input_offset * dim + tid];
+  }
+}
+
+// src is [batch*window_num, window_len, dim]
+// dst is [batch, H, W, dim] + rolled
+// grid(W, H, batch)
+// block(min(1024, dim))
+template <typename T>
+void invokeReverseRoll(T *dst,
+                       const T *src,
+                       int batch,
+                       int window_num,
+                       int window_len,
+                       int window_size,
+                       int H,
+                       int W,
+                       int dim,
+                       int shift_size,
+                       cudaStream_t stream) {
+  dim3 grid(W, H, batch);
+  int blockSize = dim;
+  if (std::is_same<T, half>::value && (dim % 2 == 0)) {
+    blockSize = dim / 2;
+    if (blockSize > 1024) {
+      blockSize = 1024;
+    }
+    using T2 = half2;
+    reverse_roll<<<grid, blockSize, 0, stream>>>(
+        reinterpret_cast<T2 *>(dst),
+        reinterpret_cast<const T2 *>(src),
+        batch,
+        window_num,
+        window_len,
+        window_size,
+        H,
+        W,
+        shift_size,
+        dim / 2);
+  } else {
+    if (blockSize > 1024) {
+      blockSize = 1024;
+    }
+    reverse_roll<<<grid, blockSize, 0, stream>>>(dst,
+                                                 src,
+                                                 batch,
+                                                 window_num,
+                                                 window_len,
+                                                 window_size,
+                                                 H,
+                                                 W,
+                                                 shift_size,
+                                                 dim);
+  }
+}
+
+template void invokeReverseRoll(float *dst,
+                                const float *src,
+                                int batch,
+                                int window_num,
+                                int window_len,
+                                int window_size,
+                                int H,
+                                int W,
+                                int dim,
+                                int shift_size,
+                                cudaStream_t stream);
+
+template void invokeReverseRoll(half *dst,
+                                const half *src,
+                                int batch,
+                                int window_num,
+                                int window_len,
+                                int window_size,
+                                int H,
+                                int W,
+                                int dim,
+                                int shift_size,
+                                cudaStream_t stream);
+
+void ReverseRollPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+bool ReverseRollPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument("The input of ReverseRoll "
+                                      "plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return in.type == nvinfer1::DataType::kHALF &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    } else {
+      return in.type == nvinfer1::DataType::kFLOAT &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType ReverseRollPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The ReverseRoll only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+nvinfer1::DimsExprs ReverseRollPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(output_index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "There is only one output of the ReverseRoll, "
+                        "so the index should be zero,"
+                        "but it's (%d)",
+                        output_index));
+  PADDLE_ENFORCE_EQ(
+      nb_inputs,
+      1,
+      common::errors::InvalidArgument(
+          "The Input of the ReverseRoll should be 1, but we found "
+          "it has (%d) inputs",
+          nb_inputs));
+
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = expr_builder.operation(nvinfer1::DimensionOperation::kFLOOR_DIV,
+                                    *inputs[0].d[0],
+                                    *expr_builder.constant(window_num_));
+  ret.d[1] = expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
+                                    *inputs[0].d[1],
+                                    *expr_builder.constant(window_num_));
+  ret.d[2] = inputs[0].d[2];
+  return ret;
+}
+int ReverseRollPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  auto input_type = input_desc[0].type;
+  int batch = input_dims.d[0] / window_num_;
+  int dim = input_dims.d[2];
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(3) << "TRT Plugin DataType selected. ReverseRoll-->fp32";
+    invokeReverseRoll(reinterpret_cast<float *>(outputs[0]),
+                      reinterpret_cast<const float *>(inputs[0]),
+                      batch,
+                      window_num_,
+                      window_len_,
+                      window_size_,
+                      input_resolution_,
+                      input_resolution_,
+                      dim,
+                      shift_size_,
+                      stream);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(3) << "TRT Plugin DataType selected. ReverseRoll-->fp16";
+    invokeReverseRoll(reinterpret_cast<half *>(outputs[0]),
+                      reinterpret_cast<const half *>(inputs[0]),
+                      batch,
+                      window_num_,
+                      window_len_,
+                      window_size_,
+                      input_resolution_,
+                      input_resolution_,
+                      dim,
+                      shift_size_,
+                      stream);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The ReverseRoll TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/reverse_roll_op_plugin.h

@@ -0,0 +1,138 @@
+// 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.
+
+#pragma once
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class ReverseRollPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  ReverseRollPluginDynamic(int window_num,
+                           int window_len,
+                           int window_size,
+                           int input_resolution,
+                           int shift_size,
+                           bool with_fp16)
+      : window_num_(window_num),
+        window_len_(window_len),
+        window_size_(window_size),
+        input_resolution_(input_resolution),
+        shift_size_(shift_size),
+        with_fp16_(with_fp16) {}
+  ReverseRollPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &window_num_);
+    DeserializeValue(&serialData, &serialLength, &window_len_);
+    DeserializeValue(&serialData, &serialLength, &window_size_);
+    DeserializeValue(&serialData, &serialLength, &input_resolution_);
+    DeserializeValue(&serialData, &serialLength, &shift_size_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new ReverseRollPluginDynamic(window_num_,
+                                        window_len_,
+                                        window_size_,
+                                        input_resolution_,
+                                        shift_size_,
+                                        with_fp16_);
+  }
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "reverse_roll_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(window_num_) + SerializedSize(window_len_) +
+           SerializedSize(window_size_) + SerializedSize(input_resolution_) +
+           SerializedSize(shift_size_) + SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, window_num_);
+    SerializeValue(&buffer, window_len_);
+    SerializeValue(&buffer, window_size_);
+    SerializeValue(&buffer, input_resolution_);
+    SerializeValue(&buffer, shift_size_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  int window_num_;
+  int window_len_;
+  int window_size_;
+  int input_resolution_;
+  int shift_size_;
+  bool with_fp16_;
+};
+
+class ReverseRollPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "reverse_roll_dynamic";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new ReverseRollPluginDynamic(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(ReverseRollPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/roi_align_op_plugin.cu

@@ -0,0 +1,791 @@
+// Copyright (c) 2018 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 <cuda_fp16.h>
+#include <cuda_runtime.h>
+
+#include <algorithm>
+
+#include "paddle/fluid/inference/tensorrt/plugin/roi_align_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <class T>
+__inline__ __device__ T BilinearInterpolate(
+    const T* input_data, const int height, const int width, T y, T x) {
+  if (y < -1.f || y > height || x < -1.f || x > width) return 0;
+  y = y <= 0.f ? 0.f : y;
+  x = x <= 0.f ? 0.f : x;
+  int y_low = static_cast<int>(y);
+  int x_low = static_cast<int>(x);
+  int y_high;
+  int x_high;
+  if (y_low >= height - 1) {
+    y_high = y_low = height - 1;
+    y = static_cast<T>(y_low);
+  } else {
+    y_high = y_low + 1;
+  }
+  if (x_low >= width - 1) {
+    x_high = x_low = width - 1;
+    x = static_cast<T>(x_low);
+  } else {
+    x_high = x_low + 1;
+  }
+  T ly = y - y_low, lx = x - x_low;
+  T hy = 1.f - ly, hx = 1.f - lx;
+  T v1 = input_data[y_low * width + x_low];
+  T v2 = input_data[y_low * width + x_high];
+  T v3 = input_data[y_high * width + x_low];
+  T v4 = input_data[y_high * width + x_high];
+  T w1 = hy * hx, w2 = hy * lx, w3 = ly * hx, w4 = ly * lx;
+  T val = (w1 * v1 + w2 * v2 + w3 * v3 + w4 * v4);
+  return val;
+}
+
+template <typename T, typename OutT, bool USE_SMEM>
+__global__ void GPUROIAlignOpt(const int nthreads,
+                               const T* __restrict__ input_data,
+                               const T* __restrict__ input_rois,
+                               const float spatial_scale,
+                               const int channels,
+                               const int height,
+                               const int width,
+                               const int pooled_height,
+                               const int pooled_width,
+                               const int sampling_ratio,
+                               const int num_rois,
+                               const bool aligned,
+                               OutT* __restrict__ output_data) {
+  const int batch = blockIdx.x;
+  const int channel = blockIdx.y;
+  const T* offset_input_data =
+      input_data + (batch * channels + channel) * height * width;
+  extern __shared__ T s_input_data[];
+  if (USE_SMEM) {
+    for (int idx = threadIdx.x; idx < height * width; idx += blockDim.x) {
+      s_input_data[idx] = offset_input_data[idx];
+    }
+    __syncthreads();
+  }
+  for (int idx = threadIdx.x; idx < num_rois * pooled_height * pooled_width;
+       idx += blockDim.x) {
+    const int pw = idx % pooled_width;
+    const int ph = (idx / pooled_width) % pooled_height;
+    const int roi_idx = (idx / pooled_width / pooled_height) % num_rois;
+    const int n = batch * num_rois + roi_idx;
+    const float4 rois_offset = reinterpret_cast<const float4*>(input_rois)[n];
+    const T roi_offset = aligned ? static_cast<T>(0.5) : 0;
+    const T roi_xmin = rois_offset.x * spatial_scale - roi_offset;
+    const T roi_ymin = rois_offset.y * spatial_scale - roi_offset;
+    const T roi_xmax = rois_offset.z * spatial_scale - roi_offset;
+    const T roi_ymax = rois_offset.w * spatial_scale - roi_offset;
+
+    T roi_width = roi_xmax - roi_xmin;
+    T roi_height = roi_ymax - roi_ymin;
+    if (!aligned) {
+      roi_width = max(roi_width, static_cast<T>(1.));
+      roi_height = max(roi_height, static_cast<T>(1.));
+    }
+    const T bin_size_h =
+        static_cast<T>(roi_height) / static_cast<T>(pooled_height);
+    const T bin_size_w =
+        static_cast<T>(roi_width) / static_cast<T>(pooled_width);
+    const int roi_bin_grid_h = (sampling_ratio > 0)
+                                   ? sampling_ratio
+                                   : ceil(roi_height / pooled_height);
+    const int roi_bin_grid_w =
+        (sampling_ratio > 0) ? sampling_ratio : ceil(roi_width / pooled_width);
+    const T count = max(roi_bin_grid_h * roi_bin_grid_w, 1);
+    T output_val = 0.f;
+    for (int iy = 0; iy < roi_bin_grid_h; ++iy) {
+      const T y = roi_ymin + ph * bin_size_h +
+                  static_cast<T>(iy + .5f) * bin_size_h /
+                      static_cast<T>(roi_bin_grid_h);
+      for (int ix = 0; ix < roi_bin_grid_w; ++ix) {
+        const T x = roi_xmin + pw * bin_size_w +
+                    static_cast<T>(ix + .5f) * bin_size_w /
+                        static_cast<T>(roi_bin_grid_w);
+        if (USE_SMEM) {
+          T val = BilinearInterpolate<T>(s_input_data, height, width, y, x);
+          output_val += val;
+        } else {
+          T val =
+              BilinearInterpolate<T>(offset_input_data, height, width, y, x);
+          output_val += val;
+        }
+      }
+    }
+    output_val /= count;
+    const int out_offset =
+        batch * num_rois * channels * pooled_height * pooled_width +
+        roi_idx * channels * pooled_height * pooled_width +
+        channel * pooled_height * pooled_width + ph * pooled_width + pw;
+    output_data[out_offset] = static_cast<OutT>(output_val);
+  }
+}
+
+RoiAlignPluginDynamic::RoiAlignPluginDynamic(const nvinfer1::DataType data_type,
+                                             const int pooled_height,
+                                             const int pooled_width,
+                                             float spatial_scale,
+                                             int sampling_ratio,
+                                             bool aligned)
+    : data_type_(data_type),
+      pooled_height_(pooled_height),
+      pooled_width_(pooled_width),
+      spatial_scale_(spatial_scale),
+      sampling_ratio_(sampling_ratio),
+      aligned_(aligned) {
+  bool data_type_is_valid = data_type_ == nvinfer1::DataType::kFLOAT ||
+                            data_type_ == nvinfer1::DataType::kHALF;
+  PADDLE_ENFORCE_EQ(data_type_is_valid,
+                    true,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts kFLOAT(%d) or "
+                        "kHALF(%d) data type, but the received data type = %d",
+                        static_cast<int>(nvinfer1::DataType::kFLOAT),
+                        static_cast<int>(nvinfer1::DataType::kHALF),
+                        static_cast<int>(data_type_)));
+
+  PADDLE_ENFORCE_GT(pooled_height_,
+                    0,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts pooled_height "
+                        "greater than %d, but the received pooled_height = %d",
+                        0,
+                        pooled_height_));
+
+  PADDLE_ENFORCE_GT(pooled_width_,
+                    0,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts pooled_width greater "
+                        "than %d, but the received pooled_width = %d",
+                        0,
+                        pooled_height_));
+
+  PADDLE_ENFORCE_GT(spatial_scale_,
+                    0.f,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts spatial_scale "
+                        "greater than %f, but the received spatial_scale = %f",
+                        0,
+                        spatial_scale_));
+
+  int smem_per_block = -1;
+  int device = -1;
+  cudaGetDevice(&device);
+
+  PADDLE_ENFORCE_GE(
+      device,
+      0,
+      common::errors::InvalidArgument(
+          "The cuda device ID should be greater than %d, but device ID is %d",
+          0,
+          device));
+
+  cudaDeviceGetAttribute(
+      &smem_per_block, cudaDevAttrMaxSharedMemoryPerBlock, device);
+  smem_per_block_ = smem_per_block;
+}
+
+RoiAlignPluginDynamic::RoiAlignPluginDynamic(void const* data, size_t length) {
+  DeserializeValue(&data, &length, &data_type_);
+  DeserializeValue(&data, &length, &pooled_height_);
+  DeserializeValue(&data, &length, &pooled_width_);
+  DeserializeValue(&data, &length, &spatial_scale_);
+  DeserializeValue(&data, &length, &sampling_ratio_);
+  DeserializeValue(&data, &length, &aligned_);
+  int smem_per_block = -1;
+  int device = -1;
+  cudaGetDevice(&device);
+  PADDLE_ENFORCE_GE(
+      device,
+      0,
+      common::errors::InvalidArgument(
+          "The cuda device ID should be greater than %d, but device ID is %d",
+          0,
+          device));
+  cudaDeviceGetAttribute(
+      &smem_per_block, cudaDevAttrMaxSharedMemoryPerBlock, device);
+  smem_per_block_ = smem_per_block;
+}
+
+nvinfer1::IPluginV2DynamicExt* RoiAlignPluginDynamic::clone() const
+    TRT_NOEXCEPT {
+  auto* plugin = new RoiAlignPluginDynamic(data_type_,
+                                           pooled_height_,
+                                           pooled_width_,
+                                           spatial_scale_,
+                                           sampling_ratio_,
+                                           aligned_);
+  plugin->setPluginNamespace(namespace_.c_str());
+  return plugin;
+}
+
+nvinfer1::DimsExprs RoiAlignPluginDynamic::getOutputDimensions(
+    int outputIndex,
+    const nvinfer1::DimsExprs* inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs ret{};
+  ret.nbDims = 4;
+  ret.d[0] = inputs[1].d[0];  // roi
+  ret.d[1] = inputs[0].d[1];  // X
+  ret.d[2] = exprBuilder.constant(pooled_height_);
+  ret.d[3] = exprBuilder.constant(pooled_width_);
+  return ret;
+}
+
+bool RoiAlignPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  if (inOut[pos].format != nvinfer1::TensorFormat::kLINEAR) {
+    return false;
+  }
+  if (pos < 2) {  // input
+    return inOut[pos].type == nvinfer1::DataType::kFLOAT;
+  }
+  return inOut[pos].type == data_type_;
+}
+
+void RoiAlignPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+size_t RoiAlignPluginDynamic::getWorkspaceSize(
+    const nvinfer1::PluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::PluginTensorDesc* outputs,
+    int nbOutputs) const TRT_NOEXCEPT {
+  return 0;
+}
+
+template <typename T, typename OutT>
+int RoiAlignPluginDynamic::enqueue_impl(
+    const nvinfer1::PluginTensorDesc* inputDesc,
+    const nvinfer1::PluginTensorDesc* outputDesc,
+    const void* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) {
+  auto in_dims = inputDesc[0].dims;
+  auto rois_dims = inputDesc[1].dims;
+  auto out_dims = outputDesc[0].dims;
+
+  int rois_num = rois_dims.d[0];
+  if (rois_num == 0) return cudaGetLastError() != cudaSuccess;
+
+  int batch = in_dims.d[0];
+  int channels = in_dims.d[1];
+  int height = in_dims.d[2];
+  int width = in_dims.d[3];
+
+  int output_size =
+      out_dims.d[0] * out_dims.d[1] * out_dims.d[2] * out_dims.d[3];
+
+  const dim3 blocks(batch, channels);
+  const int threads = 512;
+
+  if (smem_per_block_ < width * height * sizeof(T)) {
+    GPUROIAlignOpt<T, OutT, false>
+        <<<blocks, threads, 0, stream>>>(output_size,
+                                         static_cast<const T*>(inputs[0]),
+                                         static_cast<const T*>(inputs[1]),
+                                         spatial_scale_,
+                                         channels,
+                                         height,
+                                         width,
+                                         pooled_height_,
+                                         pooled_width_,
+                                         sampling_ratio_,
+                                         rois_num / batch,
+                                         aligned_,
+                                         static_cast<OutT*>(outputs[0]));
+  } else {
+    GPUROIAlignOpt<T, OutT, false>
+        <<<blocks, threads, width * height * sizeof(T), stream>>>(
+            output_size,
+            static_cast<const T*>(inputs[0]),
+            static_cast<const T*>(inputs[1]),
+            spatial_scale_,
+            channels,
+            height,
+            width,
+            pooled_height_,
+            pooled_width_,
+            sampling_ratio_,
+            rois_num / batch,
+            aligned_,
+            static_cast<OutT*>(outputs[0]));
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+int RoiAlignPluginDynamic::enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+                                   const nvinfer1::PluginTensorDesc* outputDesc,
+                                   const void* const* inputs,
+                                   void* const* outputs,
+                                   void* workspace,
+                                   cudaStream_t stream) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(outputDesc[0].type,
+                    data_type_,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlignPluginDynamic expects outputDesc[0].type "
+                        "equal to data_type_"));
+
+  if (data_type_ == nvinfer1::DataType::kHALF) {
+    return enqueue_impl<float, half>(
+        inputDesc, outputDesc, inputs, outputs, workspace, stream);
+  }
+  return enqueue_impl<float, float>(
+      inputDesc, outputDesc, inputs, outputs, workspace, stream);
+}
+
+nvinfer1::DataType RoiAlignPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* inputTypes,
+    int nbInputs) const TRT_NOEXCEPT {
+  return inputTypes[0];
+}
+
+const char* RoiAlignPluginDynamic::getPluginType() const TRT_NOEXCEPT {
+  return "roi_align_plugin_dynamic";
+}
+
+const char* RoiAlignPluginDynamic::getPluginVersion() const TRT_NOEXCEPT {
+  return "2";
+}
+
+int RoiAlignPluginDynamic::getNbOutputs() const TRT_NOEXCEPT { return 1; }
+
+int RoiAlignPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+void RoiAlignPluginDynamic::terminate() TRT_NOEXCEPT {}
+
+size_t RoiAlignPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  size_t serialize_size = 0;
+  serialize_size += SerializedSize(data_type_);
+  serialize_size += SerializedSize(pooled_height_);
+  serialize_size += SerializedSize(pooled_width_);
+  serialize_size += SerializedSize(spatial_scale_);
+  serialize_size += SerializedSize(sampling_ratio_);
+  serialize_size += SerializedSize(aligned_);
+  return serialize_size;
+}
+
+void RoiAlignPluginDynamic::serialize(void* buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, data_type_);
+  SerializeValue(&buffer, pooled_height_);
+  SerializeValue(&buffer, pooled_width_);
+  SerializeValue(&buffer, spatial_scale_);
+  SerializeValue(&buffer, sampling_ratio_);
+  SerializeValue(&buffer, aligned_);
+}
+
+void RoiAlignPluginDynamic::destroy() TRT_NOEXCEPT {}
+
+RoiAlignPluginDynamicCreator::RoiAlignPluginDynamicCreator() = default;
+
+void RoiAlignPluginDynamicCreator::setPluginNamespace(const char* lib_namespace)
+    TRT_NOEXCEPT {
+  namespace_ = std::string(lib_namespace);
+}
+
+const char* RoiAlignPluginDynamicCreator::getPluginNamespace() const
+    TRT_NOEXCEPT {
+  return namespace_.c_str();
+}
+
+const char* RoiAlignPluginDynamicCreator::getPluginName() const TRT_NOEXCEPT {
+  return "roi_align_plugin_dynamic";
+}
+
+const char* RoiAlignPluginDynamicCreator::getPluginVersion() const
+    TRT_NOEXCEPT {
+  return "2";
+}
+
+const nvinfer1::PluginFieldCollection*
+RoiAlignPluginDynamicCreator::getFieldNames() TRT_NOEXCEPT {
+  return &field_collection_;
+}
+
+nvinfer1::IPluginV2Ext* RoiAlignPluginDynamicCreator::createPlugin(
+    const char* name, const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT {
+  const nvinfer1::PluginField* fields = fc->fields;
+  return nullptr;
+}
+
+nvinfer1::IPluginV2Ext* RoiAlignPluginDynamicCreator::deserializePlugin(
+    const char* name,
+    const void* serial_data,
+    size_t serial_length) TRT_NOEXCEPT {
+  auto plugin = new RoiAlignPluginDynamic(serial_data, serial_length);
+  plugin->setPluginNamespace(namespace_.c_str());
+  return plugin;
+}
+
+PIRRoiAlignPluginDynamic::PIRRoiAlignPluginDynamic(
+    const nvinfer1::DataType data_type,
+    const int pooled_height,
+    const int pooled_width,
+    float spatial_scale,
+    int sampling_ratio,
+    bool aligned)
+    : data_type_(data_type),
+      pooled_height_(pooled_height),
+      pooled_width_(pooled_width),
+      spatial_scale_(spatial_scale),
+      sampling_ratio_(sampling_ratio),
+      aligned_(aligned) {
+  bool data_type_is_valid = data_type_ == nvinfer1::DataType::kFLOAT ||
+                            data_type_ == nvinfer1::DataType::kHALF;
+  PADDLE_ENFORCE_EQ(data_type_is_valid,
+                    true,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts kFLOAT(%d) or "
+                        "kHALF(%d) data type, but the received data type = %d",
+                        static_cast<int>(nvinfer1::DataType::kFLOAT),
+                        static_cast<int>(nvinfer1::DataType::kHALF),
+                        static_cast<int>(data_type_)));
+
+  PADDLE_ENFORCE_GT(pooled_height_,
+                    0,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts pooled_height "
+                        "greater than %d, but the received pooled_height = %d",
+                        0,
+                        pooled_height_));
+
+  PADDLE_ENFORCE_GT(pooled_width_,
+                    0,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts pooled_width greater "
+                        "than %d, but the received pooled_width = %d",
+                        0,
+                        pooled_height_));
+
+  PADDLE_ENFORCE_GT(spatial_scale_,
+                    0.f,
+                    common::errors::InvalidArgument(
+                        "TRT RoiAlign plugin only accepts spatial_scale "
+                        "greater than %f, but the received spatial_scale = %f",
+                        0,
+                        spatial_scale_));
+
+  int smem_per_block = -1;
+  int device = -1;
+  cudaGetDevice(&device);
+
+  PADDLE_ENFORCE_GE(
+      device,
+      0,
+      common::errors::InvalidArgument(
+          "The cuda device ID should be greater than %d, but device ID is %d",
+          0,
+          device));
+
+  cudaDeviceGetAttribute(
+      &smem_per_block, cudaDevAttrMaxSharedMemoryPerBlock, device);
+  smem_per_block_ = smem_per_block;
+}
+
+PIRRoiAlignPluginDynamic::PIRRoiAlignPluginDynamic(void const* data,
+                                                   size_t length) {
+  DeserializeValue(&data, &length, &data_type_);
+  DeserializeValue(&data, &length, &pooled_height_);
+  DeserializeValue(&data, &length, &pooled_width_);
+  DeserializeValue(&data, &length, &spatial_scale_);
+  DeserializeValue(&data, &length, &sampling_ratio_);
+  DeserializeValue(&data, &length, &aligned_);
+  int smem_per_block = -1;
+  int device = -1;
+  cudaGetDevice(&device);
+  PADDLE_ENFORCE_GE(
+      device,
+      0,
+      common::errors::InvalidArgument(
+          "The cuda device ID should be greater than %d, but device ID is %d",
+          0,
+          device));
+  cudaDeviceGetAttribute(
+      &smem_per_block, cudaDevAttrMaxSharedMemoryPerBlock, device);
+  smem_per_block_ = smem_per_block;
+}
+
+nvinfer1::IPluginV2DynamicExt* PIRRoiAlignPluginDynamic::clone() const
+    TRT_NOEXCEPT {
+  auto* plugin = new PIRRoiAlignPluginDynamic(data_type_,
+                                              pooled_height_,
+                                              pooled_width_,
+                                              spatial_scale_,
+                                              sampling_ratio_,
+                                              aligned_);
+  plugin->setPluginNamespace(namespace_.c_str());
+  return plugin;
+}
+
+nvinfer1::DimsExprs PIRRoiAlignPluginDynamic::getOutputDimensions(
+    int outputIndex,
+    const nvinfer1::DimsExprs* inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs ret{};
+  ret.nbDims = 4;
+  ret.d[0] = inputs[1].d[0];  // roi
+  ret.d[1] = inputs[0].d[1];  // X
+  ret.d[2] = exprBuilder.constant(pooled_height_);
+  ret.d[3] = exprBuilder.constant(pooled_width_);
+  return ret;
+}
+
+bool PIRRoiAlignPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  if (inOut[pos].format != nvinfer1::TensorFormat::kLINEAR) {
+    return false;
+  }
+  if (pos < 2) {  // input
+    return inOut[pos].type == nvinfer1::DataType::kFLOAT;
+  }
+  return inOut[pos].type == data_type_;
+}
+
+void PIRRoiAlignPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+size_t PIRRoiAlignPluginDynamic::getWorkspaceSize(
+    const nvinfer1::PluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::PluginTensorDesc* outputs,
+    int nbOutputs) const TRT_NOEXCEPT {
+  return 0;
+}
+
+template <typename T, typename OutT>
+int PIRRoiAlignPluginDynamic::enqueue_impl(
+    const nvinfer1::PluginTensorDesc* inputDesc,
+    const nvinfer1::PluginTensorDesc* outputDesc,
+    const void* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) {
+  auto in_dims = inputDesc[0].dims;
+  auto rois_dims = inputDesc[1].dims;
+  auto out_dims = outputDesc[0].dims;
+  int rois_num = rois_dims.d[0];
+  if (rois_num == 0) return cudaGetLastError() != cudaSuccess;
+
+  int batch = in_dims.d[0];
+  int channels = in_dims.d[1];
+  int height = in_dims.d[2];
+  int width = in_dims.d[3];
+  int output_size =
+      out_dims.d[0] * out_dims.d[1] * out_dims.d[2] * out_dims.d[3];
+
+  const dim3 blocks(batch, channels);
+  const int threads = 512;
+
+  if (smem_per_block_ < width * height * sizeof(T)) {
+    GPUROIAlignOpt<T, OutT, false>
+        <<<blocks, threads, 0, stream>>>(output_size,
+                                         static_cast<const T*>(inputs[0]),
+                                         static_cast<const T*>(inputs[1]),
+                                         spatial_scale_,
+                                         channels,
+                                         height,
+                                         width,
+                                         pooled_height_,
+                                         pooled_width_,
+                                         sampling_ratio_,
+                                         rois_num / batch,
+                                         aligned_,
+                                         static_cast<OutT*>(outputs[0]));
+  } else {
+    GPUROIAlignOpt<T, OutT, false>
+        <<<blocks, threads, width * height * sizeof(T), stream>>>(
+            output_size,
+            static_cast<const T*>(inputs[0]),
+            static_cast<const T*>(inputs[1]),
+            spatial_scale_,
+            channels,
+            height,
+            width,
+            pooled_height_,
+            pooled_width_,
+            sampling_ratio_,
+            rois_num / batch,
+            aligned_,
+            static_cast<OutT*>(outputs[0]));
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+int PIRRoiAlignPluginDynamic::enqueue(
+    const nvinfer1::PluginTensorDesc* inputDesc,
+    const nvinfer1::PluginTensorDesc* outputDesc,
+    const void* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      outputDesc[0].type,
+      data_type_,
+      common::errors::InvalidArgument(
+          "TRT PIRRoiAlignPluginDynamic expects outputDesc[0].type "
+          "equal to data_type_"));
+
+  if (data_type_ == nvinfer1::DataType::kHALF) {
+    return enqueue_impl<float, half>(
+        inputDesc, outputDesc, inputs, outputs, workspace, stream);
+  }
+  return enqueue_impl<float, float>(
+      inputDesc, outputDesc, inputs, outputs, workspace, stream);
+}
+
+nvinfer1::DataType PIRRoiAlignPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* inputTypes,
+    int nbInputs) const TRT_NOEXCEPT {
+  return inputTypes[0];
+}
+
+const char* PIRRoiAlignPluginDynamic::getPluginType() const TRT_NOEXCEPT {
+  return "pir_roi_align_plugin_dynamic";
+}
+
+const char* PIRRoiAlignPluginDynamic::getPluginVersion() const TRT_NOEXCEPT {
+  return "2";
+}
+
+int PIRRoiAlignPluginDynamic::getNbOutputs() const TRT_NOEXCEPT { return 1; }
+
+int PIRRoiAlignPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+void PIRRoiAlignPluginDynamic::terminate() TRT_NOEXCEPT {}
+
+size_t PIRRoiAlignPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  size_t serialize_size = 0;
+  serialize_size += SerializedSize(data_type_);
+  serialize_size += SerializedSize(pooled_height_);
+  serialize_size += SerializedSize(pooled_width_);
+  serialize_size += SerializedSize(spatial_scale_);
+  serialize_size += SerializedSize(sampling_ratio_);
+  serialize_size += SerializedSize(aligned_);
+  return serialize_size;
+}
+
+void PIRRoiAlignPluginDynamic::serialize(void* buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, data_type_);
+  SerializeValue(&buffer, pooled_height_);
+  SerializeValue(&buffer, pooled_width_);
+  SerializeValue(&buffer, spatial_scale_);
+  SerializeValue(&buffer, sampling_ratio_);
+  SerializeValue(&buffer, aligned_);
+}
+
+void PIRRoiAlignPluginDynamic::destroy() TRT_NOEXCEPT {}
+
+PIRRoiAlignPluginDynamicCreator::PIRRoiAlignPluginDynamicCreator() = default;
+
+void PIRRoiAlignPluginDynamicCreator::setPluginNamespace(
+    const char* lib_namespace) TRT_NOEXCEPT {
+  namespace_ = std::string(lib_namespace);
+}
+
+const char* PIRRoiAlignPluginDynamicCreator::getPluginNamespace() const
+    TRT_NOEXCEPT {
+  return namespace_.c_str();
+}
+
+const char* PIRRoiAlignPluginDynamicCreator::getPluginName() const
+    TRT_NOEXCEPT {
+  return "pir_roi_align_plugin_dynamic";
+}
+
+const char* PIRRoiAlignPluginDynamicCreator::getPluginVersion() const
+    TRT_NOEXCEPT {
+  return "2";
+}
+
+const nvinfer1::PluginFieldCollection*
+PIRRoiAlignPluginDynamicCreator::getFieldNames() TRT_NOEXCEPT {
+  return &field_collection_;
+}
+
+nvinfer1::IPluginV2Ext* PIRRoiAlignPluginDynamicCreator::createPlugin(
+    const char* name, const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT {
+  const nvinfer1::PluginField* fields = fc->fields;
+  int type_id = -1;
+  int pooled_height = 1;
+  int pooled_width = 1;
+  float spatial_scale = 1.0;
+  int sampling_ratio = -1;
+  bool aligned = false;
+
+  for (int i = 0; i < fc->nbFields; ++i) {
+    const std::string field_name(fc->fields[i].name);
+    if (field_name.compare("type_id") == 0) {
+      type_id = *static_cast<const int*>(fc->fields[i].data);
+    } else if (field_name.compare("pooled_height") == 0) {
+      pooled_height = *static_cast<const int*>(fc->fields[i].data);
+    } else if (field_name.compare("pooled_width") == 0) {
+      pooled_width = *static_cast<const int*>(fc->fields[i].data);
+    } else if (field_name.compare("spatial_scale") == 0) {
+      spatial_scale = *static_cast<const float*>(fc->fields[i].data);
+    } else if (field_name.compare("sampling_ratio") == 0) {
+      sampling_ratio = *static_cast<const int*>(fc->fields[i].data);
+    } else if (field_name.compare("aligned") == 0) {
+      aligned = *static_cast<const bool*>(fc->fields[i].data);
+    } else {
+      assert(false && "unknown plugin field name.");
+    }
+  }
+  return new PIRRoiAlignPluginDynamic(
+      type_id ? nvinfer1::DataType::kHALF : nvinfer1::DataType::kFLOAT,
+      pooled_height,
+      pooled_width,
+      spatial_scale,
+      sampling_ratio,
+      aligned);
+}
+
+nvinfer1::IPluginV2Ext* PIRRoiAlignPluginDynamicCreator::deserializePlugin(
+    const char* name,
+    const void* serial_data,
+    size_t serial_length) TRT_NOEXCEPT {
+  auto plugin = new PIRRoiAlignPluginDynamic(serial_data, serial_length);
+  plugin->setPluginNamespace(namespace_.c_str());
+  return plugin;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/roi_align_op_plugin.h

@@ -0,0 +1,220 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class RoiAlignPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit RoiAlignPluginDynamic(const nvinfer1::DataType data_type,
+                                 const int pooled_height,
+                                 const int pooled_width,
+                                 float spatial_scale,
+                                 int sampling_ratio,
+                                 bool aligned);
+  RoiAlignPluginDynamic(void const* data, size_t length);
+  ~RoiAlignPluginDynamic() = default;
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+
+ private:
+  template <typename T, typename OutT>
+  int enqueue_impl(const nvinfer1::PluginTensorDesc* inputDesc,
+                   const nvinfer1::PluginTensorDesc* outputDesc,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+
+  nvinfer1::DataType data_type_;
+  int pooled_height_;
+  int pooled_width_;
+  float spatial_scale_;
+  int sampling_ratio_;
+  int smem_per_block_;
+  bool aligned_;
+  std::string namespace_;
+};
+
+class RoiAlignPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  RoiAlignPluginDynamicCreator();
+  ~RoiAlignPluginDynamicCreator() override = default;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+class PIRRoiAlignPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit PIRRoiAlignPluginDynamic(const nvinfer1::DataType data_type,
+                                    const int pooled_height,
+                                    const int pooled_width,
+                                    float spatial_scale,
+                                    int sampling_ratio,
+                                    bool aligned);
+  PIRRoiAlignPluginDynamic(void const* data, size_t length);
+  ~PIRRoiAlignPluginDynamic() = default;
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+
+ private:
+  template <typename T, typename OutT>
+  int enqueue_impl(const nvinfer1::PluginTensorDesc* inputDesc,
+                   const nvinfer1::PluginTensorDesc* outputDesc,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+
+  nvinfer1::DataType data_type_;
+  int pooled_height_;
+  int pooled_width_;
+  float spatial_scale_;
+  int sampling_ratio_;
+  int smem_per_block_;
+  bool aligned_;
+  std::string namespace_;
+};
+
+class PIRRoiAlignPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  PIRRoiAlignPluginDynamicCreator();
+  ~PIRRoiAlignPluginDynamicCreator() override = default;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+REGISTER_TRT_PLUGIN_V2(RoiAlignPluginDynamicCreator);
+REGISTER_TRT_PLUGIN_V2(PIRRoiAlignPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/skip_groupnorm_act_op_plugin.cu

@@ -0,0 +1,507 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+SPDX-FileCopyrightText: Copyright (c) 1993-2022 NVIDIA CORPORATION & AFFILIATES.
+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 "paddle/fluid/inference/tensorrt/plugin/skip_groupnorm_act_op_plugin.h"
+#include <cub/cub.cuh>
+#include "paddle/common/layout.h"
+#include "paddle/phi/backends/gpu/gpu_context.h"
+#include "paddle/phi/kernels/funcs/math_function.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+nvinfer1::DimsExprs SkipGroupnormActPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputDims,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  return inputDims[0];
+}
+
+bool SkipGroupnormActPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of SkipGroupnormAct plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return ((in.type == nvinfer1::DataType::kHALF) &&
+              (in.format == nvinfer1::PluginFormat::kHWC8));
+    } else {
+      PADDLE_THROW(common::errors::Fatal(
+          "SkipGroupnormAct TRT Plugin is fp16 only so far"));
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType SkipGroupnormActPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      index,
+      0,
+      common::errors::InvalidArgument(
+          "The SkipGroupnormAct Plugin only has one input, so the "
+          "index value should be 0, but get %d.",
+          index));
+  PADDLE_ENFORCE_EQ((input_types[0] == nvinfer1::DataType::kFLOAT ||
+                     input_types[0] == nvinfer1::DataType::kHALF),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The input type should be half or float"));
+
+  return input_types[0];
+}
+int SkipGroupnormActPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+template <typename T>
+static inline T divUp(T m, T n) {
+  return (m + n - 1) / n;
+}
+
+static int64_t findMaxDivisor(int64_t n, int64_t maxAllowedDivisor) {
+  int64_t maxDivisor = -1;
+  for (int64_t i = 1; i <= std::sqrt(n); i++) {
+    if (n % i == 0) {
+      int64_t divisor1 = n / i;
+      int64_t divisor2 = i;
+
+      if (divisor1 > maxDivisor && divisor1 < maxAllowedDivisor) {
+        maxDivisor = divisor1;
+      }
+      if (divisor2 > maxDivisor && divisor2 < maxAllowedDivisor) {
+        maxDivisor = divisor2;
+      }
+    }
+  }
+  return maxDivisor;
+}
+
+static inline __device__ __host__ float sigmoid(float x) {
+  return 1.F / (1.F + expf(-x));
+}
+
+struct GroupSums {
+  // Is it the 1st element of the group?
+  int32_t flag;
+  // The sum.
+  float sum;
+  // The sum of squares.
+  float sumSq;
+};
+
+struct GroupSumsOp {
+  inline __device__ GroupSums operator()(GroupSums const &a,
+                                         GroupSums const &b) {
+    GroupSums dst;
+    dst.sum = b.flag ? b.sum : (a.sum + b.sum);
+    dst.sumSq = b.flag ? b.sumSq : (a.sumSq + b.sumSq);
+    dst.flag = a.flag + b.flag;
+    return dst;
+  }
+};
+
+template <int32_t tTHREADS_PER_BLOCK>
+__global__ void skipGroupNormNDHWCSumKernel(
+    GroupNormNDHWCParams<__half> params) {
+  // The object in charge of doing the sums for the different blocks.
+  typedef cub::BlockScan<GroupSums, tTHREADS_PER_BLOCK> BlockScan;
+
+  // Allocate shared memory for BlockScan.
+  __shared__ typename BlockScan::TempStorage tempStorage;
+  // Allocate shared memory for the groups. We could reduce the amount of shared
+  // memory reserved.
+  __shared__ float2 smem[tTHREADS_PER_BLOCK];
+
+  // The instance in the batch.
+  int32_t ni = blockIdx.z;
+  // The channel loaded by that thread (2 channels per thread for F16x2).
+  int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * 2;
+
+  // The first activation loaded by that block.
+  int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
+  // The last activation loaded by that block.
+  int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
+
+  // The sums.
+  float sum = 0.F;
+  float sumSq = 0.F;
+
+  // Iterate over the activations to compute the sums.
+  for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
+    // The offset.
+    int64_t offset = static_cast<int64_t>(ni) * params.dhwc +
+                     static_cast<int64_t>(dhwi) * params.c + ci;
+    // Fetch two channels per thread.
+    __half2 h2(0, 0);
+    if (ci < params.c) {
+      // W = 1, H = 1
+      int64_t offsetY = static_cast<int64_t>(ni) * params.c + ci;
+      __half2 y = *reinterpret_cast<__half2 const *>(&params.srcY[offsetY]);
+      h2 = *reinterpret_cast<__half2 const *>(&params.srcX[offset]);
+#if __CUDA_ARCH__ >= 530
+      h2 = __hadd2(h2, y);
+#else
+      float2 out{};
+      out.x = __half2float(h2.x) + __half2float(y.x);
+      out.y = __half2float(h2.y) + __half2float(y.y);
+      h2 = __float22half2_rn(out);
+#endif
+      // elementwise_add
+      *reinterpret_cast<__half2 *>(&params.dst[offset]) = h2;
+    }
+
+    // Extract the two half values.
+    float2 f2 = __half22float2(h2);
+
+    // Update the sum.
+    sum += f2.x + f2.y;
+    // Update the sum of squares.
+    sumSq += f2.x * f2.x + f2.y * f2.y;
+  }
+
+  // The group that thread works on and the channel in the group (modulus).
+  int32_t gi = threadIdx.x * 2 / params.cPerGroup;
+  int32_t cj = threadIdx.x * 2 - params.cPerGroup * gi;
+
+  // The data for the summations.
+  GroupSums inp{cj == 0 ? 1 : 0, sum, sumSq};
+
+  // Do the segmented scan.
+  GroupSums out;
+  BlockScan(tempStorage).InclusiveScan(inp, out, GroupSumsOp());
+
+  // Store the results for the groups in shared memory (to produce coalesced
+  // stores later).
+  if (cj == params.cPerGroup - 2 /* 2 channels per thread */) {
+    smem[gi] = make_float2(out.sum, out.sumSq);
+  }
+
+  // Make sure the data is in shared memory.
+  __syncthreads();
+
+  // The global group index.
+  int32_t gj = blockIdx.x * params.groupsPerBlock + threadIdx.x;
+
+  // Threads that have nothing left to do, exit.
+  if (threadIdx.x >= params.groupsPerBlock || gj >= params.groups) {
+    return;
+  }
+
+  // The first threads (those storing to global memory, load the values).
+  float2 sums = smem[threadIdx.x];
+
+  // Store to global memory.
+  atomicAdd(&params.redBuffer[(2 * ni + 0) * params.groups + gj], sums.x);
+  atomicAdd(&params.redBuffer[(2 * ni + 1) * params.groups + gj], sums.y);
+}
+
+void skipGroupNormNDHWCSum(GroupNormNDHWCParams<__half> const &params,
+                           cudaStream_t stream) {
+  // Make sure the values are as we expect.
+  PADDLE_ENFORCE_EQ(params.c % params.cPerBlock,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The groupNormNDHWCSum of SkipGroupnormAct Plugin got "
+                        "wrong parameters: "
+                        "params.c %% params.cPerBlock should be 0, but get %d.",
+                        params.c % params.cPerBlock));
+  PADDLE_ENFORCE_EQ(
+      params.dhw % params.dhwPerBlock,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCSum of SkipGroupnormAct Plugin got wrong "
+          "parameters: "
+          "params.dhw  %% params.dhwPerBlock should be 0, but get %d.",
+          params.dhw % params.dhwPerBlock));
+  // Make sure a group does not span multiple blocks.
+  PADDLE_ENFORCE_EQ(
+      params.cPerBlock % params.cPerGroup,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCSum of SkipGroupnormAct Plugin got wrong "
+          "parameters: "
+          "params.cPerBlock %% params.cPerGroup should be 0, but get %d.",
+          params.cPerBlock % params.cPerGroup));
+  dim3 grid;
+
+  // The number of blocks to compute all the channels.
+  grid.x = params.c / params.cPerBlock;
+  // The number of blocks to compute all the activations in a given instance.
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
+  // The number of instances.
+  grid.z = params.n;
+
+  switch (params.cPerBlock) {
+    case 320:
+      skipGroupNormNDHWCSumKernel<160><<<grid, 160, 0, stream>>>(params);
+      break;
+    case 480:
+      skipGroupNormNDHWCSumKernel<256><<<grid, 256, 0, stream>>>(params);
+      break;
+    case 256:
+      skipGroupNormNDHWCSumKernel<128><<<grid, 128, 0, stream>>>(params);
+      break;
+    case 128:
+      skipGroupNormNDHWCSumKernel<64><<<grid, 64, 0, stream>>>(params);
+      break;
+    case 8:
+      skipGroupNormNDHWCSumKernel<4><<<grid, 4, 0, stream>>>(params);
+      break;
+    default:
+      PADDLE_THROW(common::errors::Fatal(
+          "The function groupNormNDHWCSum of SkipGroupnormAct TRT Plugin "
+          "encounter error"));
+  }
+}
+
+template <int32_t tTHREADS_PER_BLOCK>
+__global__ void skipGroupNormNDHWCScaleKernel(
+    GroupNormNDHWCParams<__half> params) {
+  // The instance in the batch.
+  int32_t ni = blockIdx.z;
+  // The channel loaded by that thread (2 channels per thread for F16x2).
+  int32_t ci = blockIdx.x * params.cPerBlock + threadIdx.x * 2;
+  // The group that thread works on and the channel in the group (modulus).
+  int32_t gi = ci / params.cPerGroup;
+
+  // Load the sum and sum of squares for the group.
+  float sum = 0.F, sumSq = 0.F;
+  if (gi < params.groups) {
+    sum = params.redBuffer[(2 * ni + 0) * params.groups + gi];
+    sumSq = params.redBuffer[(2 * ni + 1) * params.groups + gi];
+  }
+
+  // Load gamma/beta.
+  float2 gammaF2, betaF2;
+  if (ci < params.c) {
+    gammaF2 = *reinterpret_cast<float2 const *>(
+        reinterpret_cast<float const *>(params.gamma) + ci);
+    betaF2 = *reinterpret_cast<float2 const *>(
+        reinterpret_cast<float const *>(params.beta) + ci);
+  }
+
+  // Compute the mean.
+  float mean = sum * params.invDHWC;
+  // Compute the variance.
+  float var = sumSq * params.invDHWC - (mean * mean);
+  // Compute the inverse of the stddev.
+  float invStdDev = rsqrtf(var + params.eps);
+
+  // The first activation loaded by that block.
+  int64_t dhwBegin = static_cast<int64_t>(blockIdx.y) * params.dhwPerBlock;
+  // The last activation loaded by that block.
+  int64_t dhwEnd = min(dhwBegin + params.dhwPerBlock, params.dhw);
+
+  // Iterate over the activations to compute the sums.
+  for (int64_t dhwi = dhwBegin; dhwi < dhwEnd; ++dhwi) {
+    // The src/dst offset.
+    int64_t offset = (int64_t)ni * params.dhwc + dhwi * params.c + ci;
+
+    // Fetch two channels per thread.
+    __half2 h2(0, 0);
+    if (ci < params.c) {
+      h2 = *reinterpret_cast<__half2 const *>(&params.dst[offset]);
+    }
+
+    // Extract the two half values.
+    float2 f2 = __half22float2(h2);
+
+    // Normalize the channels.
+    f2.x = (f2.x - mean) * invStdDev;
+    f2.y = (f2.y - mean) * invStdDev;
+
+    // Scale by gamma and add beta.
+    f2.x = gammaF2.x * f2.x + betaF2.x;
+    f2.y = gammaF2.y * f2.y + betaF2.y;
+
+    // Apply Silu if needed.
+    if (params.withSilu) {
+      f2.x = f2.x * sigmoid(f2.x);
+      f2.y = f2.y * sigmoid(f2.y);
+    }
+
+    // Store the scaled values.
+    if (ci < params.c) {
+      *reinterpret_cast<__half2 *>(&params.dst[offset]) = __float22half2_rn(f2);
+    }
+  }
+}
+
+void skipGroupNormNDHWCScale(GroupNormNDHWCParams<__half> const &params,
+                             cudaStream_t stream) {
+  // Make sure the dimensions are aligned with what we expect.
+  PADDLE_ENFORCE_EQ(
+      params.c % params.cPerBlock,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCScale of SkipGroupnormAct Plugin got "
+          "wrong parameters: "
+          "params.c %% params.cPerBlock should be 0, but get %d.",
+          params.c % params.cPerBlock));
+  // Make sure a group does not span multiple blocks.
+  PADDLE_ENFORCE_EQ(
+      params.cPerBlock % params.cPerGroup,
+      0,
+      common::errors::InvalidArgument(
+          "The groupNormNDHWCScale of SkipGroupnormAct Plugin got wrong "
+          "parameters: "
+          "params.cPerBlock %% params.cPerGroup should be 0, but get %d.",
+          params.cPerBlock % params.cPerGroup));
+  dim3 grid;
+
+  // The number of blocks to compute all the channels.
+  grid.x = params.c / params.cPerBlock;
+  // The number of blocks to compute all the activations in a given instance.
+  grid.y = divUp(params.dhw, params.dhwPerBlock);
+  // The number of instances.
+  grid.z = params.n;
+
+  switch (params.cPerBlock) {
+    case 320:
+      skipGroupNormNDHWCScaleKernel<160><<<grid, 160, 0, stream>>>(params);
+      break;
+    case 480:
+      skipGroupNormNDHWCScaleKernel<256><<<grid, 256, 0, stream>>>(params);
+      break;
+    case 256:
+      skipGroupNormNDHWCScaleKernel<128><<<grid, 128, 0, stream>>>(params);
+      break;
+    case 128:
+      skipGroupNormNDHWCScaleKernel<64><<<grid, 64, 0, stream>>>(params);
+      break;
+    case 8:
+      skipGroupNormNDHWCScaleKernel<4><<<grid, 4, 0, stream>>>(params);
+      break;
+    default:
+      PADDLE_THROW(common::errors::Fatal(
+          "The function groupNormNDHWCSum of SkipGroupnormAct TRT Plugin "
+          "encounter error"));
+  }
+}
+
+int SkipGroupnormActPluginDynamic::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 {
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. SkipGroupnormAct-->fp32";
+    PADDLE_THROW(common::errors::Fatal(
+        "The SkipGroupnormAct TRT Plugin's only support fp16 input"));
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. SkipGroupnormAct-->fp16";
+    int32_t cPerBlock = 320;
+    int32_t maxBlocksPerDHW = 1024;
+
+    switch (input_desc[0].dims.d[1]) {
+      case 960:
+      case 1920:
+        cPerBlock = 480;
+        break;
+      case 512:
+      case 256:
+        cPerBlock = 256;
+        break;
+      case 128:
+        cPerBlock = 128;
+        break;
+      default:
+        cPerBlock = 320;
+    }
+    if (cPerBlock > input_desc[0].dims.d[1]) {
+      cPerBlock = 8;
+    }
+    auto d_dim = input_desc[0].dims.nbDims;
+    params_.n = input_desc[0].dims.d[0];
+    if (d_dim == 3) {
+      params_.c = input_desc[0].dims.d[1];
+      params_.d = 1;
+      params_.h = 1;
+      params_.w = input_desc[0].dims.d[2];
+    } else if (d_dim == 4) {
+      params_.c = input_desc[0].dims.d[1];
+      params_.d = 1;
+      params_.h = input_desc[0].dims.d[2];
+      params_.w = input_desc[0].dims.d[3];
+    } else {
+      // d_dim == 5
+      params_.c = input_desc[0].dims.d[1];
+      params_.d = input_desc[0].dims.d[2];
+      params_.h = input_desc[0].dims.d[3];
+      params_.w = input_desc[0].dims.d[4];
+    }
+    params_.withSilu = true;
+    params_.dst = static_cast<half *>(outputs[0]);
+    params_.srcX = static_cast<half const *>(inputs[0]);
+    params_.srcY = static_cast<half const *>(inputs[1]);
+    params_.gamma = scale_gpu_.get();
+    params_.beta = bias_gpu_.get();
+    params_.redBuffer = static_cast<float *>(workspace);
+    // params_.n = input_desc[0].dims.d[0];
+    // params_.h = input_desc[0].dims.d[2];
+    // params_.w = input_desc[0].dims.d[3];
+    // params_.c = input_desc[0].dims.d[1];
+    params_.groups = groups_;
+    params_.dhw = static_cast<int64_t>(params_.d) * params_.h * params_.w;
+    const int64_t blocksPerDHW =
+        findMaxDivisor(params_.dhw, static_cast<int64_t>(maxBlocksPerDHW));
+    params_.dhwPerBlock = divUp(params_.dhw, blocksPerDHW);
+    params_.cPerBlock = cPerBlock;
+    params_.cPerGroup = params_.c / params_.groups;
+    params_.dhwc = params_.dhw * params_.c;
+    params_.invDHWC = 1.F / static_cast<float>(params_.dhw * params_.cPerGroup);
+    params_.groupsPerBlock = cPerBlock / params_.cPerGroup;
+    params_.eps = eps_;
+
+    cudaMemsetAsync(params_.redBuffer, 0, ws_, stream);
+    skipGroupNormNDHWCSum(params_, stream);
+    skipGroupNormNDHWCScale(params_, stream);
+
+  } else {
+    // input not fp16
+    PADDLE_THROW(common::errors::Fatal(
+        "The SkipGroupnormAct TRT Plugin's only support fp16 input"));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/skip_groupnorm_act_op_plugin.h

@@ -0,0 +1,195 @@
+/* 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. */
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/phi/kernels/group_norm_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+using phi::GroupNormNDHWCParams;
+class SkipGroupnormActPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  SkipGroupnormActPluginDynamic(const float* scale,
+                                const int scale_num,
+                                const float* bias,
+                                const int bias_num,
+                                float eps,
+                                int groups,
+                                bool with_fp16,
+                                std::shared_ptr<void> scale_gpu = nullptr,
+                                std::shared_ptr<void> bias_gpu = nullptr)
+      : scale_gpu_(scale_gpu),
+        bias_gpu_(bias_gpu),
+        groups_(groups),
+        eps_(eps),
+        with_fp16_(with_fp16) {
+    scale_.resize(scale_num);
+    bias_.resize(bias_num);
+    std::copy(scale, scale + scale_num, scale_.data());
+    std::copy(bias, bias + bias_num, bias_.data());
+    if (scale_gpu_ == nullptr) {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), scale_num * sizeof(float));
+      scale_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(
+          p, scale_.data(), scale_num * sizeof(float), cudaMemcpyHostToDevice);
+    }
+    if (bias_gpu_ == nullptr) {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), bias_num * sizeof(float));
+      bias_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(
+          p, bias_.data(), bias_num * sizeof(float), cudaMemcpyHostToDevice);
+    }
+  }
+
+  SkipGroupnormActPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &groups_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+
+    {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), scale_.size() * sizeof(float));
+      scale_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(p,
+                 scale_.data(),
+                 scale_.size() * sizeof(float),
+                 cudaMemcpyHostToDevice);
+    }
+    {
+      void* p;
+      cudaMalloc(reinterpret_cast<void**>(&p), bias_.size() * sizeof(float));
+      bias_gpu_.reset(p, [](void* ptr) { cudaFree(ptr); });
+      cudaMemcpy(p,
+                 bias_.data(),
+                 bias_.size() * sizeof(float),
+                 cudaMemcpyHostToDevice);
+    }
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    auto* ptr = new SkipGroupnormActPluginDynamic(scale_.data(),
+                                                  scale_.size(),
+                                                  bias_.data(),
+                                                  bias_.size(),
+                                                  eps_,
+                                                  groups_,
+                                                  with_fp16_,
+                                                  scale_gpu_,
+                                                  bias_gpu_);
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "skip_groupnorm_act_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(scale_) + SerializedSize(bias_) +
+           SerializedSize(eps_) + SerializedSize(groups_) +
+           SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, groups_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {
+    // sizeof(float2) * maxBatchSize * maxNumberOfGroup. float2
+    // contains two buffers for sum and squared sum;
+    ws_ = sizeof(float) * 2 * in[0].max.d[0] * groups_;
+  }
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return ws_;
+  }
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+  void terminate() TRT_NOEXCEPT override{};
+
+ private:
+  size_t ws_;
+  std::vector<float> scale_;
+  std::vector<float> bias_;
+  std::shared_ptr<void> scale_gpu_;
+  std::shared_ptr<void> bias_gpu_;
+  GroupNormNDHWCParams<__half> params_;
+  int groups_;
+  float eps_;
+  bool with_fp16_;
+};
+
+class SkipGroupnormActPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "skip_groupnorm_act_plugin_dynamic";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new SkipGroupnormActPluginDynamic(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(SkipGroupnormActPluginDynamicCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/skip_merge_layernorm_op_plugin.cu

@@ -0,0 +1,340 @@
+/* Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
+   Copyright (c) 2019-2022, NVIDIA CORPORATION.  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 <algorithm>
+
+#include "paddle/fluid/inference/tensorrt/plugin/skip_merge_layernorm_op_plugin.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+#define FINAL_MASK 0xffffffff
+
+template <typename T>
+__global__ void merge_layernorm_v2(T *out,
+                                   const T *__restrict input0,
+                                   const T *__restrict input1,
+                                   const T *__restrict gamma,
+                                   const T *__restrict beta,
+                                   const float layernorm_eps,
+                                   int batch,
+                                   int H,
+                                   int W,
+                                   int n) {
+  // input is [batch, 2*H, 2*W, n/4]
+  // output is [batch, H, W, n]
+  // grid (W, H, batch)
+  // block (n)
+  const int kIte = 4;
+  const int tid = threadIdx.x;
+  const int W_idx = blockIdx.x;
+  const int H_idx = blockIdx.y;
+  const size_t batch_offset = blockIdx.z * H * W * n;
+  const int input_H_stride = W * n / 2;
+  const int output_H_stride = W * n;
+  const int n_4 = n >> 2;
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float mean = 0.0f;
+  float variance = 0.0f;
+  float local_out[kIte];
+  float sum = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kIte; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      int part_id = col_id / n_4;
+      int offset_in_W = part_id / 2;
+      int offset_in_H = part_id % 2;
+      size_t input_id = batch_offset +
+                        (2 * H_idx + offset_in_H) * input_H_stride +
+                        (2 * W_idx + offset_in_W) * n_4 + (col_id % n_4);
+      local_out[i] = static_cast<float>(__ldg(input0 + input_id));
+      local_out[i] += static_cast<float>(__ldg(input1 + input_id));
+      sum += local_out[i];
+    }
+  }
+
+  mean = phi::funcs::BlockReduceSum<float>(sum, FINAL_MASK);
+  if (tid == 0) {
+    s_mean = mean / n;
+  }
+  __syncthreads();
+
+  float var = 0.0f;
+#pragma unroll
+  for (int i = 0; i < kIte; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      local_out[i] = local_out[i] - s_mean;
+      var += local_out[i] * local_out[i];
+    }
+  }
+
+  variance = phi::funcs::BlockReduceSum<float>(var, FINAL_MASK);
+  if (tid == 0) {
+    s_variance = rsqrtf(variance / n + layernorm_eps);
+  }
+  __syncthreads();
+
+#pragma unroll
+  for (int i = 0; i < kIte; i++) {
+    int col_id = i * blockDim.x + tid;
+    if (col_id < n) {
+      size_t output_idx =
+          batch_offset + H_idx * output_H_stride + W_idx * n + col_id;
+      out[output_idx] =
+          static_cast<T>(local_out[i] * s_variance *
+                             static_cast<float>(__ldg(&gamma[col_id])) +
+                         static_cast<float>(__ldg(&beta[col_id])));
+    }
+  }
+}
+
+template <typename T>
+void invokeMergeLayernorm(T *output,
+                          const T *input0,
+                          const T *input1,
+                          const T *gamma,
+                          const T *beta,
+                          float layernorm_eps,
+                          int batch,
+                          int H,
+                          int W,
+                          int n,
+                          cudaStream_t stream) {
+  if ((W % 2 != 0) || (H % 2 != 0)) {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "H(W) of merge layernorm should be a multiple of 2."));
+  }
+  dim3 grid(W / 2, H / 2, batch);
+  int blockSize = (n + 31) / 32 * 32;
+  merge_layernorm_v2<T><<<grid, blockSize, 0, stream>>>(output,
+                                                        input0,
+                                                        input1,
+                                                        gamma,
+                                                        beta,
+                                                        layernorm_eps,
+                                                        batch,
+                                                        H / 2,
+                                                        W / 2,
+                                                        n * 4);
+}
+
+template void invokeMergeLayernorm<float>(float *output,
+                                          const float *input0,
+                                          const float *input1,
+                                          const float *gamma,
+                                          const float *beta,
+                                          float layernorm_eps,
+                                          int batch,
+                                          int H,
+                                          int W,
+                                          int n,
+                                          cudaStream_t stream);
+
+template void invokeMergeLayernorm<half>(half *output,
+                                         const half *input0,
+                                         const half *input1,
+                                         const half *gamma,
+                                         const half *beta,
+                                         float layernorm_eps,
+                                         int batch,
+                                         int H,
+                                         int W,
+                                         int n,
+                                         cudaStream_t stream);
+
+template <typename T>
+static void convertAndCopy(const std::vector<float> &host, T *dev) {
+  T *host_ptr = new T[host.size()];
+  std::transform(host.begin(), host.end(), host_ptr, [](float x) {
+    return static_cast<T>(x);
+  });
+  cudaMemcpy(dev, host_ptr, sizeof(T) * host.size(), cudaMemcpyHostToDevice);
+  delete host_ptr;
+}
+
+void SkipMergeLayernormPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+SkipMergeLayernormPluginDynamic::SkipMergeLayernormPluginDynamic(
+    const float *bias_d,
+    const size_t bias_num,
+    const float *scale_d,
+    const size_t scale_num,
+    const float eps,
+    const int begin_norm_axis,
+    const bool with_fp16,
+    std::shared_ptr<void> bias_device,
+    std::shared_ptr<void> scale_device)
+    : eps_(eps),
+      begin_norm_axis_(begin_norm_axis),
+      with_fp16_(with_fp16),
+      bias_device_(bias_device),
+      scale_device_(scale_device) {
+  bias_.resize(bias_num);
+  scale_.resize(scale_num);
+  std::copy(bias_d, bias_d + bias_num, bias_.data());
+  std::copy(scale_d, scale_d + scale_num, scale_.data());
+  int type_size = with_fp16_ ? sizeof(half) : sizeof(float);
+  if (bias_device_ == nullptr) {
+    void *p;
+    cudaMalloc(&p, bias_num * type_size);
+    bias_device_.reset(p, [](void *ptr) { cudaFree(ptr); });
+
+    if (with_fp16) {
+      convertAndCopy<half>(bias_, reinterpret_cast<half *>(p));
+    } else {
+      convertAndCopy<float>(bias_, reinterpret_cast<float *>(p));
+    }
+  }
+  if (scale_device_ == nullptr) {
+    void *p;
+    cudaMalloc(&p, scale_num * type_size);
+    scale_device_.reset(p, [](void *ptr) { cudaFree(ptr); });
+    if (with_fp16) {
+      convertAndCopy<half>(scale_, reinterpret_cast<half *>(p));
+    } else {
+      convertAndCopy<float>(scale_, reinterpret_cast<float *>(p));
+    }
+  }
+}
+
+bool SkipMergeLayernormPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument("The input of MergeLayernorm "
+                                      "plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return in.type == nvinfer1::DataType::kHALF &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    } else {
+      return in.type == nvinfer1::DataType::kFLOAT &&
+             in.format == nvinfer1::TensorFormat::kLINEAR;
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType SkipMergeLayernormPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The MergeLayernorm only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+nvinfer1::DimsExprs SkipMergeLayernormPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = inputs[0].d[0];
+  ret.d[1] = expr_builder.operation(nvinfer1::DimensionOperation::kFLOOR_DIV,
+                                    *inputs[0].d[1],
+                                    *expr_builder.constant(4));
+  ret.d[2] = expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
+                                    *inputs[0].d[2],
+                                    *expr_builder.constant(4));
+  return ret;
+}
+
+int SkipMergeLayernormPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  auto input_type = input_desc[0].type;
+  int batch = input_dims.d[0];
+  int input_resolution = static_cast<int>(std::sqrt(input_dims.d[1]));
+  int dim = static_cast<int>(input_dims.d[2]);
+  PADDLE_ENFORCE_EQ(
+      input_resolution * input_resolution,
+      input_dims.d[1],
+      common::errors::InvalidArgument(
+          "The MergeLayernorm TRT Plugin get invalid input_resolution %d",
+          input_resolution));
+
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(3) << "TRT Plugin DataType selected. MergeLayernorm-->fp32";
+    invokeMergeLayernorm<float>(
+        reinterpret_cast<float *>(outputs[0]),
+        reinterpret_cast<const float *>(inputs[0]),
+        reinterpret_cast<const float *>(inputs[1]),
+        reinterpret_cast<const float *>(scale_device_.get()),
+        reinterpret_cast<const float *>(bias_device_.get()),
+        eps_,
+        batch,
+        input_resolution,
+        input_resolution,
+        dim,
+        stream);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(3) << "TRT Plugin DataType selected. MergeLayernorm-->fp16";
+    invokeMergeLayernorm<half>(
+        reinterpret_cast<half *>(outputs[0]),
+        reinterpret_cast<const half *>(inputs[0]),
+        reinterpret_cast<const half *>(inputs[1]),
+        reinterpret_cast<const half *>(scale_device_.get()),
+        reinterpret_cast<const half *>(bias_device_.get()),
+        eps_,
+        batch,
+        input_resolution,
+        input_resolution,
+        dim,
+        stream);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The MergeLayernorm TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/skip_merge_layernorm_op_plugin.h

@@ -0,0 +1,141 @@
+/* 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. */
+
+#pragma once
+#include <memory>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+class SkipMergeLayernormPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  SkipMergeLayernormPluginDynamic(const float* bias_d,
+                                  const size_t bias_num,
+                                  const float* scale_d,
+                                  const size_t scale_num,
+                                  const float eps,
+                                  const int begin_norm_axis,
+                                  const bool with_fp16,
+                                  std::shared_ptr<void> bias_device = nullptr,
+                                  std::shared_ptr<void> scale_device = nullptr);
+
+  SkipMergeLayernormPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &begin_norm_axis_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new SkipMergeLayernormPluginDynamic(bias_.data(),
+                                               bias_.size(),
+                                               scale_.data(),
+                                               scale_.size(),
+                                               eps_,
+                                               begin_norm_axis_,
+                                               with_fp16_,
+                                               bias_device_,
+                                               scale_device_);
+  }
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "skip_merge_layernorm_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(bias_) + SerializedSize(scale_) +
+           SerializedSize(eps_) + SerializedSize(begin_norm_axis_) +
+           SerializedSize(with_fp16_);
+  }
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, begin_norm_axis_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  std::vector<float> bias_;
+  std::vector<float> scale_;
+  float eps_;
+  int begin_norm_axis_;
+  bool with_fp16_;
+  std::shared_ptr<void> bias_device_ = nullptr;
+  std::shared_ptr<void> scale_device_ = nullptr;
+};
+class SkipMergeLayernormPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "skip_merge_layernorm_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new SkipMergeLayernormPluginDynamic(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(SkipMergeLayernormPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/split_op_plugin.cu

@@ -0,0 +1,338 @@
+// Copyright (c) 2018 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 <cuda_fp16.h>
+#include <thrust/device_vector.h>
+
+#include <algorithm>
+
+#include "paddle/fluid/inference/tensorrt/plugin/split_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+template <typename T>
+__device__ int upper_bound(T const* vals, int n, T const& key) {
+  int i = 0;
+  while (n > 0) {
+    int m = n / 2;
+    int j = i + m;
+    if (!(key < vals[j])) {
+      i = j + 1;
+      n -= m + 1;
+    } else {
+      n = m;
+    }
+  }
+  return i;
+}
+
+nvinfer1::Dims SplitPlugin::getOutputDimensions(
+    int index, const nvinfer1::Dims* input_dims, int num_inputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(num_inputs,
+                    1,
+                    common::errors::InvalidArgument(
+                        "Invalid number of inputs of split TRT plugin. "
+                        "Expected 1, received %d.",
+                        num_inputs));
+  PADDLE_ENFORCE_LT(
+      index,
+      this->getNbOutputs(),
+      common::errors::InvalidArgument(
+          "Index of output should be less than the total number of outputs in "
+          "split TensorRT plugin. Received index = %d >= total outputs = %d",
+          index,
+          this->getNbOutputs()));
+
+  nvinfer1::Dims output_dims = input_dims[0];
+  output_dims.d[axis_] = output_length_.at(index);
+  return output_dims;
+}
+
+void SplitPlugin::shareData(const SplitPlugin* another) {
+  outer_rows_ = another->outer_rows_;
+  inner_cols_ = another->inner_cols_;
+  same_shape_ = another->same_shape_;
+  axis_shape_ = another->axis_shape_;
+  segment_offsets_ = another->segment_offsets_;
+}
+
+int SplitPlugin::initialize() TRT_NOEXCEPT {
+  PADDLE_ENFORCE_LE(axis_,
+                    nvinfer1::Dims::MAX_DIMS,
+                    common::errors::InvalidArgument(
+                        "Axis dimension exceeds max dimension in TensorRT. "
+                        "Received axis = %d > MAX_DIMS = %d",
+                        axis_,
+                        nvinfer1::Dims::MAX_DIMS));
+  // notice input dims is [C, H, W]
+  nvinfer1::Dims dims = this->getInputDims(0);
+  outer_rows_ = 1;
+  inner_cols_ = 1;
+  for (int i = 0; i < axis_; ++i) {
+    outer_rows_ *= dims.d[i];
+  }
+  for (int i = axis_ + 1; i < dims.nbDims; ++i) {
+    inner_cols_ *= dims.d[i];
+  }
+  same_shape_ = true;
+  std::vector<int> segment_offsets(1, 0);
+  for (int i = 0; i < this->getNbOutputs(); ++i) {
+    if (output_length_[i] != output_length_[0]) {
+      same_shape_ = false;
+    }
+    segment_offsets.push_back(segment_offsets.back() + output_length_[i]);
+  }
+  axis_shape_ = dims.d[axis_];
+  segment_offsets_ = std::move(segment_offsets);
+  return 0;
+}
+
+// nothing to release according to initialize
+void SplitPlugin::terminate() TRT_NOEXCEPT {}
+
+// The following part of the code refers to onnx-tensorrt
+// https://github.com/onnx/onnx-tensorrt/blob/master/Split.cu
+template <typename T>
+__global__ void split_kernel(int nsegment,
+                             int const* __restrict__ segment_offsets,
+                             T const* __restrict__ idata,
+                             T* const* odatas,
+                             int inner_cols,
+                             int axis_shape,
+                             int outer_rows) {
+  int x0 = threadIdx.x + blockIdx.x * blockDim.x;
+  int src_y0 = threadIdx.y + blockIdx.y * blockDim.y;
+  int z0 = threadIdx.z + blockIdx.z * blockDim.z;
+  for (int z = z0; z < outer_rows; z += blockDim.z * gridDim.z) {
+    for (int src_y = src_y0; src_y < axis_shape;
+         src_y += blockDim.y * gridDim.y) {
+      for (int x = x0; x < inner_cols; x += blockDim.x * gridDim.x) {
+        int segment = upper_bound(segment_offsets, nsegment, src_y) - 1;
+        int dst_y = src_y - segment_offsets[segment];
+        int dst_ny = segment_offsets[segment + 1] - segment_offsets[segment];
+        odatas[segment][x + inner_cols * (dst_y + dst_ny * z)] =
+            idata[x + inner_cols * (src_y + axis_shape * z)];
+      }
+    }
+  }
+}
+
+int SplitPlugin::enqueue(int batchSize,
+                         const void* const* inputs,
+                         void* const* outputs,
+                         void* workspace,
+                         cudaStream_t stream) TRT_NOEXCEPT {
+  // this two thrust variables declared here , not with in .h
+  // to avoid compiling error in cuda 11.6
+  thrust::device_vector<int> d_segment_offsets = segment_offsets_;
+  thrust::device_vector<float*> d_output_ptrs;
+  d_output_ptrs.resize(segment_offsets_.size(), nullptr);
+  const int* d_segment_offsets_ptr =
+      thrust::raw_pointer_cast(&d_segment_offsets[0]);
+  float const* input_ptr = reinterpret_cast<float const*>(inputs[0]);
+  float* const* h_odatas = reinterpret_cast<float* const*>(outputs);
+  float** output_ptrs = thrust::raw_pointer_cast(&d_output_ptrs[0]);
+  PADDLE_ENFORCE_GPU_SUCCESS(
+      cudaMemcpyAsync(output_ptrs,
+                      h_odatas,
+                      d_output_ptrs.size() * sizeof(float*),
+                      cudaMemcpyHostToDevice,
+                      stream));
+
+  int outer_rows = outer_rows_ * batchSize;
+
+  dim3 block(32, 16);
+  dim3 grid(std::min((inner_cols_ - 1) / block.x + 1, 65535u),
+            std::min((axis_shape_ - 1) / block.y + 1, 65535u),
+            std::min((outer_rows_ - 1) / block.z + 1, 65535u));
+
+  split_kernel<<<grid, block, 0, stream>>>(segment_offsets_.size(),
+                                           d_segment_offsets_ptr,
+                                           input_ptr,
+                                           output_ptrs,
+                                           inner_cols_,
+                                           axis_shape_,
+                                           outer_rows);
+  return cudaGetLastError() != cudaSuccess;
+}
+
+// Dynamic Plugin below.
+int SplitPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+size_t SplitPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(axis_) + SerializedSize(output_length_) +
+         SerializedSize(with_fp16_);
+}
+
+void SplitPluginDynamic::serialize(void* buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, axis_);
+  SerializeValue(&buffer, output_length_);
+  SerializeValue(&buffer, with_fp16_);
+}
+
+nvinfer1::DimsExprs SplitPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nb_inputs,
+                    1,
+                    common::errors::InvalidArgument(
+                        "The Split plugin should be only one input."));
+  PADDLE_ENFORCE_LT(output_index,
+                    output_length_.size(),
+                    common::errors::InvalidArgument(
+                        "When GetOutputDimensions, the index(%d) should not "
+                        "greater the num(%d) of the outputs.",
+                        output_index,
+                        output_length_.size()));
+
+  nvinfer1::DimsExprs output_dims = inputs[0];
+  output_dims.d[axis_] = expr_builder.constant(output_length_.at(output_index));
+
+  return output_dims;
+}
+
+bool SplitPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of split plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc& in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kFLOAT ||
+              in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc& prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType SplitPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  return input_types[0];
+}
+
+int SplitPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  int outer_rows = 1;
+  int inner_cols = 1;
+  // with batch
+  for (int i = 0; i < axis_; i++) {
+    outer_rows *= input_dims.d[i];
+  }
+
+  for (int i = axis_ + 1; i < input_dims.nbDims; i++) {
+    inner_cols *= input_dims.d[i];
+  }
+
+  std::vector<int> segment_offsets(1, 0);
+  for (int i = 0; i < this->getNbOutputs(); i++) {
+    segment_offsets.push_back(segment_offsets.back() + output_length_[i]);
+  }
+  int axis_shape = input_dims.d[axis_];
+  thrust::device_vector<int> d_segment_offsets = segment_offsets;
+  const int* d_segment_offsets_ptr =
+      thrust::raw_pointer_cast(&d_segment_offsets[0]);
+
+  dim3 block(32, 16);
+  dim3 grid(std::min((inner_cols - 1) / block.x + 1, 65535u),
+            std::min((axis_shape - 1) / block.y + 1, 65535u),
+            std::min((outer_rows - 1) / block.z + 1, 65535u));
+
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Split-->fp32";
+    thrust::device_vector<float*> d_output_ptrs;
+    d_output_ptrs.resize(this->getNbOutputs(), nullptr);
+
+    const float* input_ptr = static_cast<const float*>(inputs[0]);
+    float* const* h_odatas = reinterpret_cast<float* const*>(outputs);
+    float** output_ptrs = thrust::raw_pointer_cast(&d_output_ptrs[0]);
+
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaMemcpyAsync(output_ptrs,
+                        h_odatas,
+                        d_output_ptrs.size() * sizeof(float*),
+                        cudaMemcpyHostToDevice,
+                        stream));
+
+    split_kernel<<<grid, block, 0, stream>>>(d_segment_offsets.size(),
+                                             d_segment_offsets_ptr,
+                                             input_ptr,
+                                             output_ptrs,
+                                             inner_cols,
+                                             axis_shape,
+                                             outer_rows);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Split-->fp16";
+    thrust::device_vector<half*> d_output_ptrs;
+    d_output_ptrs.resize(this->getNbOutputs(), nullptr);
+
+    const half* input_ptr = static_cast<const half*>(inputs[0]);
+    half* const* h_odatas = reinterpret_cast<half* const*>(outputs);
+    half** output_ptrs = thrust::raw_pointer_cast(&d_output_ptrs[0]);
+
+    PADDLE_ENFORCE_GPU_SUCCESS(
+        cudaMemcpyAsync(output_ptrs,
+                        h_odatas,
+                        d_output_ptrs.size() * sizeof(half*),
+                        cudaMemcpyHostToDevice,
+                        stream));
+
+    split_kernel<<<grid, block, 0, stream>>>(d_segment_offsets.size(),
+                                             d_segment_offsets_ptr,
+                                             input_ptr,
+                                             output_ptrs,
+                                             inner_cols,
+                                             axis_shape,
+                                             outer_rows);
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/split_op_plugin.h

@@ -0,0 +1,264 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class SplitPlugin : public PluginTensorRTV2Ext {
+ public:
+  SplitPlugin() {}
+  SplitPlugin(int axis, std::vector<int> const& output_lengths, bool with_fp16)
+      : axis_(axis), same_shape_(true), output_length_(output_lengths) {
+    with_fp16_ = with_fp16;
+  }
+
+  SplitPlugin(void const* serial_data, size_t serial_length) {
+    deserializeBase(serial_data, serial_length);
+    DeserializeValue(&serial_data, &serial_length, &axis_);
+    DeserializeValue(&serial_data, &serial_length, &output_length_);
+  }
+
+  nvinfer1::IPluginV2Ext* clone() const TRT_NOEXCEPT override {
+    SplitPlugin* ptr = new SplitPlugin(axis_, output_length_, with_fp16_);
+    ptr->setPluginNamespace(this->getPluginNamespace());
+    ptr->shareData(this);
+    return ptr;
+  }
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override {
+    return input_types[0];
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "split_plugin_v2ext";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override {
+    return output_length_.size();
+  }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* input_dims,
+                                     int num_inputs) TRT_NOEXCEPT override;
+
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ protected:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(axis_) + SerializedSize(output_length_) +
+           getBaseSerializationSize();
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, axis_);
+    SerializeValue(&buffer, output_length_);
+  }
+
+  int axis_;
+  int outer_rows_;
+  int inner_cols_;
+  int axis_shape_;
+  bool same_shape_;
+  std::vector<int> output_length_;
+  std::vector<int> segment_offsets_;
+
+ private:
+  void shareData(const SplitPlugin* another);
+};
+
+class SplitPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  SplitPluginCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "split_plugin_v2ext";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    // not implemented
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new SplitPlugin(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+
+REGISTER_TRT_PLUGIN_V2(SplitPluginCreator);
+
+class SplitPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  SplitPluginDynamic(int axis,
+                     std::vector<int> const& output_lengths,
+                     bool with_fp16)
+      : axis_(axis), output_length_(output_lengths) {
+    with_fp16_ = with_fp16;
+  }
+
+  SplitPluginDynamic(void const* serial_data, size_t serial_length) {
+    DeserializeValue(&serial_data, &serial_length, &axis_);
+    DeserializeValue(&serial_data, &serial_length, &output_length_);
+    DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+  }
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new SplitPluginDynamic(axis_, output_length_, with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "split_plugin";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override {
+    return output_length_.size();
+  }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  int axis_;
+  std::vector<int> output_length_;
+};
+
+class SplitPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  SplitPluginDynamicCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "split_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    auto plugin = new SplitPluginDynamic(serial_data, serial_length);
+    return plugin;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+
+REGISTER_TRT_PLUGIN_V2(SplitPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/spmm_plugin.h

@@ -0,0 +1,180 @@
+/* Copyright (c) 2022, PaddlePaddle Authors, NVIDIA CORPORATION. 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.
+*/
+
+#pragma once
+
+#include <algorithm>
+#include <cassert>
+#include <iostream>
+#include <stdexcept>
+#include <vector>
+
+#include "NvInfer.h"
+#include "NvInferPlugin.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/phi/backends/dynload/cusparseLt.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class SpmmPluginDynamic : public nvinfer1::IPluginV2DynamicExt {
+ public:
+  enum class Activation { kNone, kRelu, kGelu };
+  SpmmPluginDynamic(const std::string& name,
+                    const nvinfer1::DataType precision,
+                    const int out_dim,
+                    const nvinfer1::Weights& weight,
+                    const nvinfer1::Weights& bias,
+                    Activation activation);
+  // The second constructor is for clone member function
+  SpmmPluginDynamic(const std::string& name,
+                    const nvinfer1::DataType precision,
+                    const int out_dim,
+                    const int k,
+                    const void* weight,
+                    size_t compressed_size,
+                    const void* bias,
+                    bool is_configured,
+                    const int m_max,
+                    const int optim_alg,
+                    Activation activation);
+  SpmmPluginDynamic(const std::string name, const void* data, size_t length);
+  SpmmPluginDynamic() = delete;
+  nvinfer1::IPluginV2DynamicExt* clone() const noexcept override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder) noexcept override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) noexcept override;
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) noexcept override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const noexcept override;
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) noexcept override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const noexcept override;
+  const char* getPluginType() const noexcept override;
+  const char* getPluginVersion() const noexcept override;
+  int getNbOutputs() const noexcept override;
+  int initialize() noexcept override;
+  void terminate() noexcept override;
+  size_t getSerializationSize() const noexcept override;
+  void serialize(void* buffer) const noexcept override;
+  void destroy() noexcept override;
+  void setPluginNamespace(const char* pluginNamespace) noexcept override;
+  const char* getPluginNamespace() const noexcept override;
+
+ private:
+  struct cusparseLtContext {
+    cusparseLtHandle_t handle;
+    cusparseLtMatDescriptor_t matA;
+    cusparseLtMatDescriptor_t matB;
+    cusparseLtMatDescriptor_t matC;
+    cusparseLtMatmulDescriptor_t matmul;
+    cusparseLtMatmulAlgSelection_t alg_sel;
+    cusparseLtMatmulPlan_t plan;
+    cusparseLtContext();
+    ~cusparseLtContext();
+    size_t workspace_size{0};
+    bool is_initialized{false};
+    int activation{0};
+    float relu_upper_bound{0};
+    float relu_threshold{0};
+    void init(int m,
+              int n,
+              int k,
+              cudaDataType_t type,
+              void* bias_ptr,
+              SpmmPluginDynamic::Activation activation);
+    void setAlgo(int id);
+    void destroy();
+    void compressMatB(int n,
+                      int k,
+                      cudaDataType_t type,
+                      void* src,
+                      void** dest,
+                      size_t* compressed_size);
+  };  // struct SpmmPluginDynamic::cusparseLtContext
+  const std::string layer_name_;
+  std::string namespace_;
+  nvinfer1::DataType precision_;
+  size_t precision_size_;
+  size_t
+      element_size_;  // size of weight (float if INT8 or FLOAT; half if HALF)
+  int out_dim_;
+  int k_;
+  int m_max_;
+  bool is_configured_;  // already get m, scale bias, and search the optim alg
+                        // or not
+  int optim_alg_;       // the index of optimal algorithm
+  float weight_scale_;  // record the weight scale from constructor
+  void* weight_compressed_;      // host compressed weight
+  void* weight_compressed_dev_;  //  device compressed weight
+  std::shared_ptr<void>
+      weight_compressed_dev_global_;  // shared pointer to the
+                                      // device compressed weight
+  size_t compressed_size_;            // size of compressed weight
+  bool has_bias_;                     // there is bias or not
+  void* bias_;                        // host bias
+  void* bias_dev_;                    // device bias
+  Activation activation_;             // record the activation type
+  cusparseLtContext spmm_context_;
+};  // class SpmmPluginDynamic
+
+class SpmmPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  SpmmPluginDynamicCreator();
+  const char* getPluginName() const noexcept override;
+  const char* getPluginVersion() const noexcept override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() noexcept override;
+  nvinfer1::IPluginV2* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) noexcept override;
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serialData,
+                                         size_t serialLength) noexcept override;
+  void setPluginNamespace(const char* pluginNamespace) noexcept override;
+  const char* getPluginNamespace() const noexcept override;
+
+ private:
+  static nvinfer1::PluginFieldCollection field_collection_;
+  static std::vector<nvinfer1::PluginField> plugin_attr_;
+  std::string namespace_;
+};  // class SpmmPluginDynamicCreator
+
+REGISTER_TRT_PLUGIN_V2(SpmmPluginDynamicCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/stack_op_plugin.cu

@@ -0,0 +1,284 @@
+// Copyright (c) 2018 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 <cassert>
+#include <cstring>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/plugin/stack_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+StackPluginDynamic::StackPluginDynamic(int axis, int num_stack, bool with_fp16)
+    : axis_(axis), num_stack_(num_stack) {
+  with_fp16_ = with_fp16;
+}
+
+StackPluginDynamic::StackPluginDynamic(void const* serial_data,
+                                       size_t serial_length) {
+  DeserializeValue(&serial_data, &serial_length, &axis_);
+  DeserializeValue(&serial_data, &serial_length, &num_stack_);
+  DeserializeValue(&serial_data, &serial_length, &with_fp16_);
+}
+
+StackPluginDynamic::~StackPluginDynamic() = default;
+
+nvinfer1::IPluginV2DynamicExt* StackPluginDynamic::clone() const TRT_NOEXCEPT {
+  return new StackPluginDynamic(axis_, num_stack_, with_fp16_);
+}
+
+const char* StackPluginDynamic::getPluginType() const TRT_NOEXCEPT {
+  return "stack_plugin";
+}
+
+int StackPluginDynamic::getNbOutputs() const TRT_NOEXCEPT { return 1; }
+
+int StackPluginDynamic::initialize() TRT_NOEXCEPT { return 0; }
+
+size_t StackPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  size_t serialize_size = 0;
+  serialize_size += SerializedSize(axis_);
+  serialize_size += SerializedSize(num_stack_);
+  serialize_size += SerializedSize(with_fp16_);
+  return serialize_size;
+}
+
+void StackPluginDynamic::serialize(void* buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, axis_);
+  SerializeValue(&buffer, num_stack_);
+  SerializeValue(&buffer, with_fp16_);
+}
+
+nvinfer1::DimsExprs StackPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs output(inputs[0]);
+  output.nbDims = inputs[0].nbDims + 1;
+
+  for (int i = inputs[0].nbDims; i > axis_; --i) {
+    output.d[i] = inputs[0].d[i - 1];
+  }
+  output.d[axis_] = expr_builder.constant(nb_inputs);
+  return output;
+}
+
+void StackPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* out,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+size_t StackPluginDynamic::getWorkspaceSize(
+    const nvinfer1::PluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::PluginTensorDesc* outputs,
+    int nbOutputs) const TRT_NOEXCEPT {
+  return num_stack_ * sizeof(uintptr_t);
+}
+
+void StackPluginDynamic::destroy() TRT_NOEXCEPT { delete this; }
+
+void StackPluginDynamic::terminate() TRT_NOEXCEPT {}
+
+bool StackPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of stack plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc& in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      return (in.type == nvinfer1::DataType::kHALF) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc& prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType StackPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      index,
+      0,
+      common::errors::InvalidArgument("The index should be equal to 0"));
+  return input_types[0];
+}
+
+template <typename T>
+__global__ void StackKernel(const T* const* input,
+                            T* output,
+                            int num_stack,
+                            int base_unit) {
+  int stack_id = blockIdx.y;
+  int lead_id = blockIdx.x;
+
+  for (int i = threadIdx.x; i < base_unit; i += blockDim.x) {
+    output[lead_id * num_stack * base_unit + stack_id * base_unit + i] =
+        input[stack_id][lead_id * base_unit + i];
+  }
+}
+
+int StackPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;  // (batch, seq, seq)
+  auto out_dims = output_desc[0].dims;   // (batch, num_head, seq, seq)
+  auto out_num_dims = out_dims.nbDims;
+
+  int base_unit = 1;
+  for (int i = axis_ + 1; i < out_num_dims; ++i) {
+    PADDLE_ENFORCE_GT(out_dims.d[i],
+                      0,
+                      common::errors::InvalidArgument(
+                          "Input dimensions should be greater than 0"));
+    base_unit *= out_dims.d[i];
+  }
+
+  int lead_unit = 1;
+  for (int i = 0; i < axis_; ++i) {
+    PADDLE_ENFORCE_GT(out_dims.d[i],
+                      0,
+                      common::errors::InvalidArgument(
+                          "Input dimensions should be greater than 0"));
+    lead_unit *= out_dims.d[i];
+  }
+
+  PADDLE_ENFORCE_EQ(
+      out_dims.d[axis_],
+      num_stack_,
+      common::errors::InvalidArgument("number of stack axis should be same"));
+
+  cudaMemcpyAsync(workspace,
+                  reinterpret_cast<const void* const>(inputs),
+                  sizeof(void*) * out_dims.d[axis_],
+                  cudaMemcpyHostToDevice,
+                  stream);
+
+  const int num_stacks = out_dims.d[axis_];
+  // lead_unit may be very large, so make it be blockIdx.x
+  dim3 num_blocks(lead_unit, num_stacks);
+  const int num_threads = 256;
+  auto infer_type = input_desc[0].type;
+
+  if (infer_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Stack-->fp32";
+    float* output = static_cast<float*>(outputs[0]);
+    StackKernel<float><<<num_blocks, num_threads, 0, stream>>>(
+        reinterpret_cast<const float* const*>(workspace),
+        output,
+        num_stacks,
+        base_unit);
+  } else if (infer_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Stack-->fp16";
+    __half* output = static_cast<__half*>(outputs[0]);
+    StackKernel<__half><<<num_blocks, num_threads, 0, stream>>>(
+        reinterpret_cast<const __half* const*>(workspace),
+        output,
+        num_stacks,
+        base_unit);
+  } else {
+    PADDLE_THROW(
+        common::errors::Fatal("The Stack TRT Plugin's input type only "
+                              "support float or half currently."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+StackPluginDynamicCreator::StackPluginDynamicCreator() = default;
+
+const char* StackPluginDynamicCreator::getPluginName() const TRT_NOEXCEPT {
+  return "stack_plugin";
+}
+
+const char* StackPluginDynamicCreator::getPluginVersion() const TRT_NOEXCEPT {
+  return "1";
+}
+
+const nvinfer1::PluginFieldCollection*
+StackPluginDynamicCreator::getFieldNames() TRT_NOEXCEPT {
+  return &field_collection_;
+}
+
+nvinfer1::IPluginV2* StackPluginDynamicCreator::createPlugin(
+    const char* name, const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT {
+  int axis = -1;
+  int num_stack = -1;
+  bool with_fp16 = false;
+
+  for (int i = 0; i < fc->nbFields; ++i) {
+    const std::string name(fc->fields[i].name);
+    if (name == "axis") {
+      axis = static_cast<const int*>(fc->fields[i].data)[0];
+    } else if (name == "num_stack") {
+      num_stack = static_cast<const int*>(fc->fields[i].data)[0];
+    } else if (name == "with_fp16") {
+      with_fp16 = static_cast<const bool*>(fc->fields[i].data)[0];
+    } else {
+      PADDLE_THROW(common::errors::Fatal("Meet an unknown plugin field '" +
+                                         name +
+                                         "' when creating stack op plugin."));
+    }
+  }
+  return new StackPluginDynamic(axis, num_stack, with_fp16);
+}
+
+nvinfer1::IPluginV2* StackPluginDynamicCreator::deserializePlugin(
+    const char* name,
+    const void* serial_data,
+    size_t serial_length) TRT_NOEXCEPT {
+  auto plugin = new StackPluginDynamic(serial_data, serial_length);
+  return plugin;
+}
+
+void StackPluginDynamicCreator::setPluginNamespace(const char* lib_namespace)
+    TRT_NOEXCEPT {
+  plugin_namespace_ = lib_namespace;
+}
+
+const char* StackPluginDynamicCreator::getPluginNamespace() const TRT_NOEXCEPT {
+  return plugin_namespace_.c_str();
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/stack_op_plugin.h

@@ -0,0 +1,103 @@
+// Copyright (c) 2019 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.
+
+#pragma once
+#include <stdio.h>
+
+#include <cassert>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class StackPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit StackPluginDynamic(int axis, int num_stack, bool with_fp16);
+  StackPluginDynamic(void const* serial_data, size_t serial_length);
+  ~StackPluginDynamic();
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override;
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+
+ private:
+  int axis_;
+  int num_stack_;
+};
+
+class StackPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  StackPluginDynamicCreator();
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+
+ private:
+  std::string plugin_namespace_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+REGISTER_TRT_PLUGIN_V2(StackPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/swish_op_plugin.cu

@@ -0,0 +1,234 @@
+// Copyright (c) 2018 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 <stdio.h>
+
+#include <cassert>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/swish_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+int SwishPlugin::initialize() TRT_NOEXCEPT { return 0; }
+void SwishPlugin::terminate() TRT_NOEXCEPT {}
+
+bool SwishPlugin::supportsFormat(
+    nvinfer1::DataType type, nvinfer1::PluginFormat format) const TRT_NOEXCEPT {
+  if (with_fp16_) {
+    return (type == nvinfer1::DataType::kFLOAT ||
+            type == nvinfer1::DataType::kHALF) &&
+           (format == nvinfer1::TensorFormat::kLINEAR);
+  }
+  return (type == nvinfer1::DataType::kFLOAT) &&
+         (format == nvinfer1::TensorFormat::kLINEAR);
+}
+
+nvinfer1::Dims SwishPlugin::getOutputDimensions(int index,
+                                                const nvinfer1::Dims *inputDims,
+                                                int nbInputs) TRT_NOEXCEPT {
+  assert(nbInputs == 1);
+  assert(index < this->getNbOutputs());
+  nvinfer1::Dims const &input_dims = inputDims[0];
+  nvinfer1::Dims output_dims = input_dims;
+  return output_dims;
+}
+
+template <typename T>
+__device__ T math_exp(T a);
+
+template <>
+__device__ half math_exp<half>(half a) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  return hexp(a);
+#endif
+}
+
+template <>
+__device__ float math_exp<float>(float a) {
+  return expf(a);
+}
+
+template <typename T>
+__global__ void swish_kernel(int num, const T *input, T *output, T beta) {
+  int index = blockIdx.x * blockDim.x + threadIdx.x;
+  if (index < num) {
+#if __CUDA_ARCH__ >= 350
+    output[index] =
+        __ldg(input + index) /
+        (static_cast<T>(1.0) + math_exp<T>(-beta * __ldg(input + index)));
+#else
+    output[index] = input[index] /
+                    (static_cast<T>(1.0) + math_exp<T>(-beta * input[index]));
+#endif
+  }
+}
+
+template <>
+__global__ void swish_kernel<half>(int num,
+                                   const half *input,
+                                   half *output,
+                                   half beta) {
+  int index = blockIdx.x * blockDim.x + threadIdx.x;
+  if (index < num) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+    output[index] =
+        __ldg(input + index) /
+        (static_cast<half>(1.0) + math_exp<half>(-beta * __ldg(input + index)));
+#endif
+  }
+}
+
+int SwishPlugin::enqueue(int batch_size,
+                         const void *const *inputs,
+                         void *const *outputs,
+                         void *workspace,
+                         cudaStream_t stream) TRT_NOEXCEPT {
+  const auto &input_dims = this->getInputDims(0);
+  int num = batch_size;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    num *= input_dims.d[i];
+  }
+  int threads = 1024;
+  int blocks = (num + threads - 1) / threads;
+  auto type = getDataType();
+  if (type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Swish-->fp32";
+    const float *input = reinterpret_cast<const float *>(inputs[0]);
+    float *output = reinterpret_cast<float *const *>(outputs)[0];
+    swish_kernel<<<blocks, threads, 0, stream>>>(num, input, output, beta_);
+  } else if (type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Swish-->fp16";
+    const half *input = reinterpret_cast<const half *>(inputs[0]);
+    half *output = reinterpret_cast<half *const *>(outputs)[0];
+    swish_kernel<<<blocks, threads, 0, stream>>>(
+        num, input, output, (half)beta_);
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The Swish TRT Plugin's input type should be float or half."));
+  }
+
+  return cudaGetLastError() != cudaSuccess;
+}
+
+int SwishPluginDynamic::initialize() TRT_NOEXCEPT {
+  getPluginNamespace();
+  return 0;
+}
+
+size_t SwishPluginDynamic::getSerializationSize() const TRT_NOEXCEPT {
+  return SerializedSize(beta_) + SerializedSize(with_fp16_);
+}
+
+void SwishPluginDynamic::serialize(void *buffer) const TRT_NOEXCEPT {
+  SerializeValue(&buffer, beta_);
+  SerializeValue(&buffer, with_fp16_);
+}
+
+nvinfer1::DimsExprs SwishPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  return inputs[0];
+}
+
+bool SwishPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of swish plugin should not be nullptr."));
+
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  (in_out && pos < (nb_inputs + nb_outputs));
+
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      bool res = (in.type == nvinfer1::DataType::kFLOAT ||
+                  in.type == nvinfer1::DataType::kHALF);
+      res = res && (in.format == nvinfer1::TensorFormat::kLINEAR);
+      return res;
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  const nvinfer1::PluginTensorDesc &prev = in_out[pos - 1];
+  // output
+  return in.type == prev.type && in.format == prev.format;
+}
+
+nvinfer1::DataType SwishPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(index,
+                    0,
+                    common::errors::InvalidArgument(
+                        "The Swish Plugin only has one input, so the "
+                        "index value should be 0, but get %d.",
+                        index));
+  return input_types[0];
+}
+
+int SwishPluginDynamic::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 {
+  auto input_dims = input_desc[0].dims;
+  size_t num = ProductDim(input_dims);
+  int threads = 1024;
+  int blocks = (num + threads - 1) / threads;
+
+  auto input_type = input_desc[0].type;
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. Swish-->fp32";
+    const float *input = static_cast<const float *>(inputs[0]);
+    float *output = static_cast<float *>(outputs[0]);
+    swish_kernel<float>
+        <<<blocks, threads, 0, stream>>>(num, input, output, beta_);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. Swish-->fp16";
+    const half *input = static_cast<const half *>(inputs[0]);
+    half *output = static_cast<half *>(outputs[0]);
+    swish_kernel<half><<<blocks, threads, 0, stream>>>(
+        num, input, output, static_cast<half>(beta_));
+  } else {
+    PADDLE_THROW(common::errors::InvalidArgument(
+        "The Swish TRT Plugin's input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/swish_op_plugin.h

@@ -0,0 +1,196 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class SwishPlugin : public PluginTensorRTV2Ext {
+ private:
+  float beta_;
+
+ public:
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(beta_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, beta_);
+  }
+
+  explicit SwishPlugin(const float beta, const bool with_fp16) : beta_(beta) {
+    with_fp16_ = with_fp16;
+  }
+
+  // It was used for tensorrt deserialization.
+  // It should not be called by users.
+  SwishPlugin(void const* serialData, size_t serialLength) {
+    deserializeBase(serialData, serialLength);
+    DeserializeValue(&serialData, &serialLength, &beta_);
+  }
+
+  ~SwishPlugin() {}
+
+  int initialize() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* clone() const TRT_NOEXCEPT override {
+    auto* plugin = new SwishPlugin(beta_, with_fp16_);
+    plugin->data_format_ = data_format_;
+    plugin->data_type_ = data_type_;
+    plugin->input_dims_ = input_dims_;
+    return plugin;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "swish_plugin";
+  }
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override {
+    return input_types[0];
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nbInputDims) TRT_NOEXCEPT override;
+
+  int enqueue(int batchSize,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  void terminate() TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override { delete this; }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "2"; }
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override;
+};
+
+class SwishPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "swish_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "2"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new SwishPlugin(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(SwishPluginCreator);
+
+class SwishPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit SwishPluginDynamic(const float beta, const bool with_fp16)
+      : beta_(beta) {
+    with_fp16_ = with_fp16;
+  }
+  SwishPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &beta_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    return new SwishPluginDynamic(beta_, with_fp16_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "swish_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+  int initialize() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override {}
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  float beta_;
+};
+
+class SwishPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "swish_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new SwishPluginDynamic(serial_data, serial_length);
+  }
+};
+REGISTER_TRT_PLUGIN_V2(SwishPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/trans_layernorm_op_plugin.cu

@@ -0,0 +1,553 @@
+// Copyright (c) 2023 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 <stdio.h>
+
+#include <cassert>
+#include <vector>
+
+#include "glog/logging.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trans_layernorm_op_plugin.h"
+#include "paddle/phi/kernels/funcs/math_cuda_utils.h"
+#include "paddle/phi/kernels/layer_norm_kernel.h"
+#include "paddle/phi/kernels/transpose_kernel.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+inline int getSMVersion() {
+  const int device = phi::backends::gpu::GetCurrentDeviceId();
+  const phi::gpuDeviceProp prop =
+      phi::backends::gpu::GetDeviceProperties(device);
+  return prop.major * 10 + prop.minor;
+}
+
+#define FINAL_MASK 0xffffffff
+
+template <int UNROLL_FACTOR>
+__global__ void GeneralResidualLayerNormOpt2(half2 *normed_output,
+                                             half2 *output,
+                                             const half2 *__restrict src,
+                                             const half2 *__restrict gamma,
+                                             const half2 *__restrict beta,
+                                             int m,
+                                             int n,
+                                             float epsilon) {
+#if CUDA_ARCH_FP16_SUPPORTED(__CUDA_ARCH__)
+  __shared__ float s_mean;
+  __shared__ float s_variance;
+  float x_sum = 0.0f;
+  float x2_sum = 0.0f;
+  const int b_offset = blockIdx.x * n;
+
+#pragma unroll UNROLL_FACTOR
+  for (int i = threadIdx.x; i < n; i += blockDim.x) {
+    const int index = b_offset + i;
+    float val_1 = 0.0f;
+    float val_2 = 0.0f;
+    half2 tmp;
+    tmp = __ldg(&src[index]);
+    val_1 += static_cast<float>(tmp.x);
+    val_2 += static_cast<float>(tmp.y);
+    output[index] = tmp;
+    x_sum += val_1 + val_2;
+    x2_sum += val_1 * val_1 + val_2 * val_2;
+  }
+  float sums[2];
+  sums[0] = x_sum;
+  sums[1] = x2_sum;
+  phi::funcs::BlockReduceSumV2<float, 2>(sums);
+  constexpr int Half2VecSize = 2;
+  if (threadIdx.x == 0) {
+    s_mean = sums[0] / n / Half2VecSize;
+    s_variance = rsqrtf(sums[1] / n / Half2VecSize - s_mean * s_mean + epsilon);
+  }
+  __syncthreads();
+
+  half2 mean_2 = __float2half2_rn(s_mean);
+  half2 var_2 = __float2half2_rn(s_variance);
+
+#pragma unroll UNROLL_FACTOR
+  for (int i = threadIdx.x; i < n; i += blockDim.x) {
+    const int index = b_offset + i;
+    half2 val = __hmul2(__hmul2(__hsub2(output[index], mean_2), var_2),
+                        __ldg(&gamma[i]));
+    if (beta) {
+      val = __hadd2(val, __ldg(&beta[i]));
+    }
+    normed_output[index] = val;
+  }
+#endif
+}
+
+#define HALF2_RESIDUAL_LAYERNORM_OPT2(UNROLL_FACTOR)                         \
+  GeneralResidualLayerNormOpt2<UNROLL_FACTOR>                                \
+      <<<rows, block, 0, stream>>>(reinterpret_cast<half2 *>(layernorm_dst), \
+                                   reinterpret_cast<half2 *>(dst),           \
+                                   (const half2 *)input,                     \
+                                   (const half2 *)fp16_scale_gpu_,           \
+                                   (const half2 *)fp16_bias_gpu_,            \
+                                   rows,                                     \
+                                   half_n,                                   \
+                                   eps);
+
+int TransLayerNormPluginDynamic::initialize() TRT_NOEXCEPT {
+  if (!with_fp16_) {
+    cudaMalloc(&bias_gpu_, sizeof(float) * bias_.size());
+    cudaMemcpy(bias_gpu_,
+               bias_.data(),
+               bias_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+    cudaMalloc(&scale_gpu_, sizeof(float) * scale_.size());
+    cudaMemcpy(scale_gpu_,
+               scale_.data(),
+               scale_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+  } else {
+    cudaMalloc(&bias_gpu_, sizeof(float) * bias_.size());
+    cudaMemcpy(bias_gpu_,
+               bias_.data(),
+               bias_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+    cudaMalloc(&scale_gpu_, sizeof(float) * scale_.size());
+    cudaMemcpy(scale_gpu_,
+               scale_.data(),
+               scale_.size() * sizeof(float),
+               cudaMemcpyHostToDevice);
+    std::vector<half> fp16_bias_;
+    std::vector<half> fp16_scale_;
+    fp16_bias_.resize(bias_.size());
+    fp16_scale_.resize(scale_.size());
+    for (int i = 0; i < bias_.size(); i++) {
+      fp16_bias_[i] = static_cast<half>(bias_[i]);
+    }
+    for (int i = 0; i < scale_.size(); i++) {
+      fp16_scale_[i] = static_cast<half>(scale_[i]);
+    }
+
+    cudaMalloc(&fp16_bias_gpu_, sizeof(half) * fp16_bias_.size());
+    cudaMemcpy(fp16_bias_gpu_,
+               fp16_bias_.data(),
+               fp16_bias_.size() * sizeof(half),
+               cudaMemcpyHostToDevice);
+    cudaMalloc(&fp16_scale_gpu_, sizeof(half) * fp16_scale_.size());
+    cudaMemcpy(fp16_scale_gpu_,
+               fp16_scale_.data(),
+               fp16_scale_.size() * sizeof(half),
+               cudaMemcpyHostToDevice);
+  }
+  return 0;
+}
+
+void TransLayerNormPluginDynamic::terminate() TRT_NOEXCEPT {
+  if (bias_gpu_) {
+    cudaFree(bias_gpu_);
+    bias_gpu_ = nullptr;
+  }
+  if (scale_gpu_) {
+    cudaFree(scale_gpu_);
+    scale_gpu_ = nullptr;
+  }
+  if (fp16_bias_gpu_) {
+    cudaFree(fp16_bias_gpu_);
+    fp16_bias_gpu_ = nullptr;
+  }
+  if (fp16_scale_gpu_) {
+    cudaFree(fp16_scale_gpu_);
+    fp16_scale_gpu_ = nullptr;
+  }
+}
+
+nvinfer1::DimsExprs TransLayerNormPluginDynamic::getOutputDimensions(
+    int output_index,
+    const nvinfer1::DimsExprs *inputDims,
+    int nb_inputs,
+    nvinfer1::IExprBuilder &expr_builder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs ret;
+  ret.nbDims = 3;
+  ret.d[0] = inputDims[0].d[0];
+  ret.d[1] = expr_builder.operation(nvinfer1::DimensionOperation::kPROD,
+                                    *inputDims[0].d[2],
+                                    *inputDims[0].d[3]);
+  ret.d[2] = inputDims[0].d[1];
+  return ret;
+}
+
+bool TransLayerNormPluginDynamic::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc *in_out,
+    int nb_inputs,
+    int nb_outputs) TRT_NOEXCEPT {
+  int feature_size = bias_.size();
+  PADDLE_ENFORCE_GE(
+      feature_size,
+      0,
+      common::errors::InvalidArgument(
+          "The feature size of layernorm feature_size must be positive,"
+          "but got:%d",
+          feature_size));
+
+  PADDLE_ENFORCE_NOT_NULL(
+      in_out,
+      common::errors::InvalidArgument(
+          "The input of layernorm plugin should not be nullptr."));
+  PADDLE_ENFORCE_LT(
+      pos,
+      nb_inputs + nb_outputs,
+      common::errors::InvalidArgument("The pos(%d) should be less than the "
+                                      "num(%d) of the input and the output.",
+                                      pos,
+                                      nb_inputs + nb_outputs));
+  const nvinfer1::PluginTensorDesc &in = in_out[pos];
+  if (pos == 0) {
+    if (with_fp16_) {
+      if (feature_size % 8 == 0) {
+        // now, we only support khwc8 for feature_size % 8 == 0
+        return ((in.type == nvinfer1::DataType::kHALF) &&
+                (in.format == nvinfer1::PluginFormat::kLINEAR ||
+                 in.format == nvinfer1::PluginFormat::kHWC8));
+      } else {
+        return ((in.type == nvinfer1::DataType::kHALF) &&
+                (in.format == nvinfer1::PluginFormat::kLINEAR));
+      }
+    } else {
+      return (in.type == nvinfer1::DataType::kFLOAT) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  if (pos == 1) {
+    if (with_fp16_) {
+      return (in.type == in_out[0].type &&
+              (in.format == nvinfer1::PluginFormat::kLINEAR));
+    } else {
+      return (in.type == in_out[0].type) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+  if (pos == 2) {
+    if (with_fp16_) {
+      return (in.type == in_out[0].type &&
+              (in.format == nvinfer1::PluginFormat::kLINEAR));
+    } else {
+      return (in.type == in_out[0].type) &&
+             (in.format == nvinfer1::TensorFormat::kLINEAR);
+    }
+  }
+}
+
+void TransLayerNormPluginDynamic::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc *in,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc *out,
+    int nbOutputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      begin_norm_axis_,
+      3,
+      common::errors::InvalidArgument(
+          "The transpose_LayerNorm Plugin only has begin_norm_axis_ = 3"
+          "but get %d.",
+          begin_norm_axis_));
+  const auto &input_dims = in[0].desc.dims;
+  int statis_num = input_dims.d[0] * input_dims.d[2] * input_dims.d[3];
+  mean_shape_[0] = statis_num;
+  variance_shape_[0] = statis_num;
+}
+
+nvinfer1::DataType TransLayerNormPluginDynamic::getOutputDataType(
+    int index,
+    const nvinfer1::DataType *input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(
+      nb_inputs,
+      1,
+      common::errors::InvalidArgument(
+          "The transpose_LayerNorm Plugin only has one input, so the "
+          "nb_inputs value should be 1, but get %d.",
+          nb_inputs));
+  PADDLE_ENFORCE_EQ((input_types[0] == nvinfer1::DataType::kFLOAT ||
+                     input_types[0] == nvinfer1::DataType::kHALF),
+                    true,
+                    common::errors::InvalidArgument(
+                        "The input type should be half or float"));
+  return input_types[0];
+}
+
+int TransLayerNormPluginDynamic::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 {
+  const auto &input_dims = input_desc[0].dims;
+  int begin_norm_axis = begin_norm_axis_;
+  float eps = eps_;
+
+  std::vector<int64_t> input_shape;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    input_shape.push_back(input_dims.d[i]);
+  }
+  int input_numel = 1;
+  for (int i = 0; i < input_dims.nbDims; i++) {
+    input_numel *= input_dims.d[i];
+  }
+  PADDLE_ENFORCE_EQ(1,
+                    mean_shape_.size(),
+                    common::errors::InvalidArgument(
+                        "Size of mean_shape vector should be equal to 1,"
+                        "but got Size of mean_shape vector:%d",
+                        mean_shape_.size()));
+  PADDLE_ENFORCE_EQ(1,
+                    variance_shape_.size(),
+                    common::errors::InvalidArgument(
+                        "Size of variance_shape vector should be equal to 1,"
+                        "but got Size of mean_shape vector:%d",
+                        mean_shape_.size()));
+  PADDLE_ENFORCE_GE(mean_shape_[0],
+                    0,
+                    common::errors::InvalidArgument(
+                        "The size of mean vector should be positive,"
+                        "but got:%d",
+                        mean_shape_[0]));
+  PADDLE_ENFORCE_GE(variance_shape_[0],
+                    0,
+                    common::errors::InvalidArgument(
+                        "The size of mean vector should be positive,"
+                        "but got:%d",
+                        variance_shape_[0]));
+
+  // transpose do not change numel
+  int trans_result_numel = input_numel;
+  std::vector<int64_t> trans_result_shape{
+      input_shape[0], input_shape[2], input_shape[3], input_shape[1]};
+
+  const auto input_ddim = common::make_ddim(input_shape);
+  int feature_size = static_cast<int>(input_ddim[1]);
+  PADDLE_ENFORCE_EQ(feature_size,
+                    scale_.size(),
+                    common::errors::InvalidArgument(
+                        "scale's size should be equal to the feature_size,"
+                        "but got feature_size:%d, scale's size:%d.",
+                        feature_size,
+                        scale_.size()));
+  PADDLE_ENFORCE_EQ(feature_size,
+                    bias_.size(),
+                    common::errors::InvalidArgument(
+                        "bias's size should be equal to the feature_size,"
+                        "but got feature_size:%d, bias's size:%d.",
+                        feature_size,
+                        bias_.size()));
+
+  int device_id = -1;
+  cudaGetDevice(&device_id);
+  PADDLE_ENFORCE_GE(
+      device_id,
+      0,
+      common::errors::InvalidArgument("device_id should be positive,"
+                                      "but got:%d",
+                                      device_id));
+
+  auto input_type = input_desc[0].type;
+
+  phi::DeviceContextPool &pool = phi::DeviceContextPool::Instance();
+  GPUPlace place(platform::GetCurrentDeviceId());
+  auto *device_context = static_cast<phi::GPUContext *>(pool.Get(place));
+  const phi::GPUContext &dev_ctx = *device_context;
+
+  mean_t.Resize(common::make_ddim(mean_shape_));
+  variance_t.Resize(common::make_ddim(variance_shape_));
+  float *mean_d =
+      dev_ctx.template Alloc<float>(&mean_t, mean_shape_[0] * sizeof(float));
+  float *variance_d = dev_ctx.template Alloc<float>(
+      &variance_t, variance_shape_[0] * sizeof(float));
+
+  if (input_type == nvinfer1::DataType::kFLOAT) {
+    VLOG(1) << "TRT Plugin DataType selected. trans_layernorm-->fp32";
+    VLOG(1) << "TRT Plugin format selected. trans_layernorm-->kLINEAR";
+    const float *input = reinterpret_cast<const float *>(inputs[0]);
+    float *layernorm_dst = static_cast<float *>(outputs[0]);
+    float *dst = static_cast<float *>(outputs[1]);
+    // transpose and norm do not change numel
+    int trans_result_numel = input_numel;
+    int norm_result_numel = input_numel;
+    phi::DenseTensorMeta input_meta(phi::DataType::FLOAT32,
+                                    common::make_ddim(input_shape));
+    phi::DenseTensorMeta bias_meta(phi::DataType::FLOAT32,
+                                   common::make_ddim({feature_size}));
+    phi::DenseTensorMeta scale_meta(phi::DataType::FLOAT32,
+                                    common::make_ddim({feature_size}));
+    phi::DenseTensorMeta trans_result_meta(
+        phi::DataType::FLOAT32, common::make_ddim(trans_result_shape));
+    phi::DenseTensorMeta norm_result_meta(
+        phi::DataType::FLOAT32, common::make_ddim(trans_result_shape));
+    std::shared_ptr<phi::Allocation> input_alloc(new phi::Allocation(
+        static_cast<void *>(const_cast<float *>(input)),  // NOLINT
+        input_numel * sizeof(float),
+        place));
+    std::shared_ptr<phi::Allocation> bias_alloc(
+        new phi::Allocation(static_cast<float *>(bias_gpu_),  // NOLINT
+                            feature_size * sizeof(float),
+                            place));
+    std::shared_ptr<phi::Allocation> scale_alloc(new phi::Allocation(
+        static_cast<float *>(scale_gpu_), feature_size * sizeof(float), place));
+    std::shared_ptr<phi::Allocation> trans_result_alloc(
+        new phi::Allocation(static_cast<float *>(dst),  // NOLINT
+                            trans_result_numel * sizeof(float),
+                            place));
+    std::shared_ptr<phi::Allocation> norm_result_alloc(
+        new phi::Allocation(static_cast<float *>(layernorm_dst),  // NOLINT
+                            norm_result_numel * sizeof(float),
+                            place));
+
+    const phi::DenseTensor input_tensor =
+        phi::DenseTensor(input_alloc, input_meta);
+    phi::DenseTensor bias_tensor = phi::DenseTensor(bias_alloc, bias_meta);
+    phi::DenseTensor scale_tensor = phi::DenseTensor(scale_alloc, scale_meta);
+    phi::DenseTensor trans_result_tensor =
+        phi::DenseTensor(trans_result_alloc, trans_result_meta);
+    phi::DenseTensor norm_result_tensor =
+        phi::DenseTensor(norm_result_alloc, norm_result_meta);
+
+    phi::TransposeKernel<float, phi::GPUContext>(dev_ctx,
+                                                 input_tensor,
+                                                 std::vector<int>{0, 2, 3, 1},
+                                                 &trans_result_tensor);
+    phi::LayerNormKernel<float, phi::GPUContext>(dev_ctx,
+                                                 trans_result_tensor,
+                                                 scale_tensor,
+                                                 bias_tensor,
+                                                 eps,
+                                                 begin_norm_axis,
+                                                 &norm_result_tensor,
+                                                 &mean_t,
+                                                 &variance_t);
+  } else if (input_type == nvinfer1::DataType::kHALF) {
+    VLOG(1) << "TRT Plugin DataType selected. trans_layernorm-->fp16";
+    const half *input = reinterpret_cast<const half *>(inputs[0]);
+    half *layernorm_dst = static_cast<half *>(outputs[0]);
+    half *dst = static_cast<half *>(outputs[1]);
+    if (input_desc[0].format == nvinfer1::PluginFormat::kLINEAR) {
+      VLOG(1) << "TRT Plugin format selected. trans_layernorm-->kLINEAR";
+      phi::DenseTensorMeta input_meta(phi::DataType::FLOAT16,
+                                      common::make_ddim(input_shape));
+      std::shared_ptr<phi::Allocation> input_alloc(new phi::Allocation(
+          static_cast<void *>(const_cast<half *>(input)),  // NOLINT
+          input_numel * sizeof(half),
+          place));
+      phi::DenseTensorMeta trans_result_meta(
+          phi::DataType::FLOAT16, common::make_ddim(trans_result_shape));
+      std::shared_ptr<phi::Allocation> trans_result_alloc(
+          new phi::Allocation(static_cast<void *>(dst),  // NOLINT
+                              trans_result_numel * sizeof(half),
+                              place));
+      const phi::DenseTensor input_tensor =
+          phi::DenseTensor(input_alloc, input_meta);
+      phi::DenseTensor trans_result_tensor =
+          phi::DenseTensor(trans_result_alloc, trans_result_meta);
+      phi::TransposeKernel<phi::dtype::float16, phi::GPUContext>(
+          dev_ctx,
+          input_tensor,
+          std::vector<int>{0, 2, 3, 1},
+          &trans_result_tensor);
+      phi::LayerNormDirectCUDAFunctor<half, float> layer_norm;
+      layer_norm(stream,
+                 dst,
+                 trans_result_shape,
+                 bias_gpu_,
+                 scale_gpu_,
+                 layernorm_dst,
+                 mean_d,
+                 variance_d,
+                 begin_norm_axis,
+                 eps);
+    } else if (input_desc[0].format == nvinfer1::PluginFormat::kHWC8) {
+      VLOG(1) << "TRT Plugin format selected. trans_layernorm-->kHWC8";
+      int sm = getSMVersion();
+      // sm >= 60 to support __ldg
+      if (sm >= 60) {
+        int64_t hidden = input_shape[1];
+        if (hidden % 2 == 0) {
+          const size_t rows =
+              static_cast<size_t>(input_shape[0] * input_shape[2] *
+                                  input_shape[3]);  // batch * seq_length
+          int half_n = hidden / 2;
+          int half_n_32 = (half_n + 31) / 32 * 32;
+          dim3 block(std::min(half_n_32, 512));
+          int rolls_per_thread = half_n / block.x;
+          int unroll_factor = 8;
+          while (unroll_factor > rolls_per_thread && unroll_factor > 1) {
+            unroll_factor /= 2;
+          }
+          switch (unroll_factor) {
+            case 1:
+              HALF2_RESIDUAL_LAYERNORM_OPT2(1);
+              break;
+            case 2:
+              HALF2_RESIDUAL_LAYERNORM_OPT2(2);
+              break;
+            case 4:
+              HALF2_RESIDUAL_LAYERNORM_OPT2(4);
+              break;
+            case 8:
+              HALF2_RESIDUAL_LAYERNORM_OPT2(8);
+              break;
+            default:
+              PADDLE_THROW(common::errors::Fatal(
+                  "Invalid UNROLL_FACTOR in transpose_layernorm trt plugin."));
+          }
+        } else {
+          cudaMemcpyAsync(
+              dst, input, input_numel * sizeof(half), cudaMemcpyDeviceToDevice);
+          phi::LayerNormDirectCUDAFunctor<half, float> layer_norm;
+          layer_norm(stream,
+                     input,
+                     trans_result_shape,
+                     bias_gpu_,
+                     scale_gpu_,
+                     layernorm_dst,
+                     mean_d,
+                     variance_d,
+                     begin_norm_axis,
+                     eps);
+        }
+      } else {
+        cudaMemcpyAsync(
+            dst, input, input_numel * sizeof(half), cudaMemcpyDeviceToDevice);
+        phi::LayerNormDirectCUDAFunctor<half, float> layer_norm;
+        layer_norm(stream,
+                   input,
+                   trans_result_shape,
+                   bias_gpu_,
+                   scale_gpu_,
+                   layernorm_dst,
+                   mean_d,
+                   variance_d,
+                   begin_norm_axis,
+                   eps);
+      }
+    }
+  } else {
+    PADDLE_THROW(
+        common::errors::Fatal("The TransLayerNormPluginDynamic TRT Plugin's "
+                              "input type should be float or half."));
+  }
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/trans_layernorm_op_plugin.h

@@ -0,0 +1,176 @@
+// 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.
+
+#pragma once
+
+#include <algorithm>
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/framework/tensor.h"
+#include "paddle/fluid/framework/tensor_util.h"
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+class TransLayerNormPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  TransLayerNormPluginDynamic(const float* bias,
+                              const int bias_num,
+                              const float* scale,
+                              const int scale_num,
+                              int begin_norm_axis,
+                              float eps,
+                              std::vector<int64_t> mean_shape,
+                              std::vector<int64_t> variance_shape,
+                              bool with_fp16)
+      : begin_norm_axis_(begin_norm_axis),
+        eps_(eps),
+        mean_shape_(mean_shape),
+        variance_shape_(variance_shape) {
+    with_fp16_ = with_fp16;
+    bias_.resize(bias_num);
+    scale_.resize(scale_num);
+    std::copy(bias, bias + bias_num, bias_.data());
+    std::copy(scale, scale + scale_num, scale_.data());
+  }
+
+  TransLayerNormPluginDynamic(void const* serialData, size_t serialLength) {
+    DeserializeValue(&serialData, &serialLength, &bias_);
+    DeserializeValue(&serialData, &serialLength, &scale_);
+    DeserializeValue(&serialData, &serialLength, &begin_norm_axis_);
+    DeserializeValue(&serialData, &serialLength, &eps_);
+    DeserializeValue(&serialData, &serialLength, &mean_shape_);
+    DeserializeValue(&serialData, &serialLength, &variance_shape_);
+    DeserializeValue(&serialData, &serialLength, &with_fp16_);
+  }
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    auto ptr = new TransLayerNormPluginDynamic(bias_.data(),
+                                               bias_.size(),
+                                               scale_.data(),
+                                               scale_.size(),
+                                               begin_norm_axis_,
+                                               eps_,
+                                               mean_shape_,
+                                               variance_shape_,
+                                               with_fp16_);
+    ptr->bias_gpu_ = bias_gpu_;
+    ptr->scale_gpu_ = scale_gpu_;
+    ptr->fp16_bias_gpu_ = fp16_bias_gpu_;
+    ptr->fp16_scale_gpu_ = fp16_scale_gpu_;
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "trans_layernorm_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 2; }
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return SerializedSize(bias_) + SerializedSize(scale_) +
+           SerializedSize(begin_norm_axis_) + SerializedSize(eps_) +
+           SerializedSize(mean_shape_) + SerializedSize(variance_shape_) +
+           SerializedSize(with_fp16_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    SerializeValue(&buffer, bias_);
+    SerializeValue(&buffer, scale_);
+    SerializeValue(&buffer, begin_norm_axis_);
+    SerializeValue(&buffer, eps_);
+    SerializeValue(&buffer, mean_shape_);
+    SerializeValue(&buffer, variance_shape_);
+    SerializeValue(&buffer, with_fp16_);
+  }
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ private:
+  std::vector<float> bias_;
+  std::vector<float> scale_;
+  phi::DenseTensor mean_t;
+  phi::DenseTensor variance_t;
+  int begin_norm_axis_;
+  float eps_;
+  std::vector<int64_t> mean_shape_;
+  std::vector<int64_t> variance_shape_;
+  // data on devices
+  float* bias_gpu_{nullptr};
+  float* scale_gpu_{nullptr};
+  half* fp16_bias_gpu_{nullptr};
+  half* fp16_scale_gpu_{nullptr};
+};
+
+class TransLayerNormPluginDynamicCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "trans_layernorm_plugin_dynamic";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new TransLayerNormPluginDynamic(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(TransLayerNormPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/transformer_input_output_convert_plugin.cu

@@ -0,0 +1,359 @@
+/* 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 "paddle/phi/common/memory_utils.h"
+#include "paddle/phi/core/platform/device/gpu/gpu_info.h"
+#include "paddle/phi/core/platform/device_context.h"
+
+#include "cub/cub.cuh"
+#include "paddle/fluid/inference/tensorrt/plugin/transformer_input_output_convert_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+__global__ void remove_padding_kernel(const half* input0,
+                                      const int32_t* input1,
+                                      half* output) {
+  int word_id = blockIdx.x * gridDim.y + blockIdx.y;
+  int32_t sequence_length = input1[blockIdx.x + 1] - input1[blockIdx.x];
+  if (blockIdx.y < sequence_length) {
+    output[(input1[blockIdx.x] + blockIdx.y) * gridDim.z * blockDim.x +
+           blockIdx.z * blockDim.x + threadIdx.x] =
+        input0[word_id * gridDim.z * blockDim.x + blockIdx.z * blockDim.x +
+               threadIdx.x];
+  }
+}
+
+__global__ void recover_padding_kernel(const half* input0,
+                                       const int32_t* input1,
+                                       half* output) {
+  int word_id = blockIdx.x * gridDim.y + blockIdx.y;
+  int32_t sequence_length = input1[blockIdx.x + 1] - input1[blockIdx.x];
+  if (blockIdx.y < sequence_length) {
+    output[word_id * gridDim.z * blockDim.x + blockIdx.z * blockDim.x +
+           threadIdx.x] =
+        input0[(input1[blockIdx.x] + blockIdx.y) * gridDim.z * blockDim.x +
+               blockIdx.z * blockDim.x + threadIdx.x];
+  } else {
+    output[word_id * gridDim.z * blockDim.x + blockIdx.z * blockDim.x +
+           threadIdx.x] = 0;
+  }
+}
+
+nvinfer1::DataType TransformerInputConvertPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  if (index == 0) {  // new input
+    return nvinfer1::DataType::kHALF;
+  } else if (index == 1) {  // mask
+    return nvinfer1::DataType::kHALF;
+  } else if (index == 2) {  // pos id
+    return nvinfer1::DataType::kINT32;
+  } else if (index == 3) {  // max_seqlen_tensor
+    return nvinfer1::DataType::kHALF;
+  }
+}
+
+nvinfer1::DimsExprs TransformerInputConvertPlugin::getOutputDimensions(
+    int outputIndex,
+    const nvinfer1::DimsExprs* inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) TRT_NOEXCEPT {
+  constexpr size_t threadsPerCta384 = 1 * 8 * 32;
+  constexpr size_t xmmasM384 = 24;
+  constexpr size_t packedMaskSize384 = xmmasM384 * threadsPerCta384;
+  int32_t maskSize_ = packedMaskSize384;
+  auto maskSize = exprBuilder.constant(maskSize_);
+  auto fp16maskSize = exprBuilder.operation(
+      nvinfer1::DimensionOperation::kPROD, *maskSize, *exprBuilder.constant(2));
+
+  auto one = exprBuilder.constant(1);
+  auto B = inputs[0].d[0];
+  auto MaxLength = inputs[0].d[1];
+  auto Hidden = inputs[0].d[2];
+
+  nvinfer1::DimsExprs output_dims;
+  if (outputIndex == 0) {  // new input
+    output_dims.nbDims = 4;
+    output_dims.d[0] = exprBuilder.operation(
+        nvinfer1::DimensionOperation::kPROD, *B, *MaxLength);
+    output_dims.d[1] = Hidden;
+    output_dims.d[2] = exprBuilder.constant(1);
+    output_dims.d[3] = exprBuilder.constant(1);
+  } else if (outputIndex == 1) {  // mask
+    output_dims.nbDims = 2;
+    output_dims.d[0] = B;
+    output_dims.d[1] = fp16maskSize;
+  } else if (outputIndex == 2) {  // pos id
+    output_dims.nbDims = 1;
+    output_dims.d[0] =
+        exprBuilder.operation(nvinfer1::DimensionOperation::kSUM, *B, *one);
+  } else if (outputIndex == 3) {  // max_seqlen_tensor
+    output_dims.nbDims = 1;
+    output_dims.d[0] = MaxLength;
+  }
+  return output_dims;
+}
+
+bool TransformerInputConvertPlugin::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nbInputs,
+                    2,
+                    common::errors::InvalidArgument(
+                        "TransformerInputConvertPlugin must have 2 inputs, "
+                        "but got %d input(s). ",
+                        nbInputs));
+  PADDLE_ENFORCE_EQ(nbOutputs,
+                    4,
+                    common::errors::InvalidArgument(
+                        "TransformerInputConvertPlugin must have 4 outputs, "
+                        "but got %d output(s). ",
+                        nbOutputs));
+  if (pos == 0) {  //  input
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kHALF;
+  } else if (pos == 1) {  //  reducesum_qk_bias
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kINT32;
+  } else if (pos == 2) {  // new input
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kHALF;
+  } else if (pos == 3) {  // mask
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kHALF;
+  } else if (pos == 4) {  // pos id
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kINT32;
+  } else if (pos == 5) {  // max_seqlen_tensor
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kHALF;
+  }
+}
+
+void TransformerInputConvertPlugin::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* outputs,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+void TransformerInputConvertPlugin::attachToContext(
+    cudnnContext* cudnnContext,
+    cublasContext* cublasContext,
+    nvinfer1::IGpuAllocator* gpuAllocator) TRT_NOEXCEPT {}
+
+void TransformerInputConvertPlugin::detachFromContext() TRT_NOEXCEPT {}
+
+void TransformerInputConvertPlugin::terminate() TRT_NOEXCEPT {}
+
+int TransformerInputConvertPlugin::enqueue(
+    const nvinfer1::PluginTensorDesc* inputDesc,
+    const nvinfer1::PluginTensorDesc* outputDesc,
+    const void* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) TRT_NOEXCEPT {
+  // input(no_varlen), reducesum_qk_bias, input(varlen), mask, pos_id,
+  // max_seqlen_tensor
+  const half* input0 = static_cast<const half*>(inputs[0]);  // input(no_varlen)
+  const int32_t* input1 =
+      static_cast<const int32_t*>(inputs[1]);            // reducesum_qk_bias
+  half* output0 = static_cast<half*>(outputs[0]);        // input(varlen)
+  int32_t* output2 = static_cast<int32_t*>(outputs[2]);  // pos_id
+  const auto input0_desc = inputDesc[0];
+  const int32_t B = input0_desc.dims.d[0];           // batches
+  const int32_t MaxLength = input0_desc.dims.d[1];   // max token length
+  const int32_t HiddenSize = input0_desc.dims.d[2];  // hidden size
+
+  // Determine temporary device storage requirements
+  size_t temp_storage_bytes = 0;
+  cub::DeviceScan::ExclusiveSum(
+      NULL, temp_storage_bytes, input1, output2, B + 1);
+  // Allocate temporary storage
+  GPUPlace place(platform::GetCurrentDeviceId());
+  auto d_temp_storage = phi::memory_utils::Alloc(place, temp_storage_bytes);
+  // Run exclusive prefix sum
+  cub::DeviceScan::ExclusiveSum(
+      d_temp_storage->ptr(), temp_storage_bytes, input1, output2, B + 1);
+  const int32_t vector_length = HiddenSize;
+  int32_t num_threads;
+  if (vector_length < 1024) {
+    num_threads = vector_length;
+  } else {
+    if (vector_length % 512 == 0) {
+      num_threads = 512;
+    } else if (vector_length % 256 == 0) {
+      num_threads = 256;
+    } else if (vector_length % 128 == 0) {
+      num_threads = 128;
+    } else if (vector_length % 64 == 0) {
+      num_threads = 64;
+    } else if (vector_length % 32 == 0) {
+      num_threads = 32;
+    } else if (vector_length % 16 == 0) {
+      num_threads = 16;
+    } else if (vector_length % 8 == 0) {
+      num_threads = 8;
+    } else if (vector_length % 4 == 0) {
+      num_threads = 4;
+    } else if (vector_length % 2 == 0) {
+      num_threads = 2;
+    } else {
+      num_threads = 1;
+    }
+  }
+  const dim3 num_blocks(
+      B,
+      MaxLength,
+      vector_length /
+          num_threads);  //  batches, max sequence length, input0.dims.d[2]/*
+  remove_padding_kernel<<<num_blocks, num_threads, 0, stream>>>(
+      input0, output2, output0);  // input(no_varlen), pos_id, input(varlen)
+  return cudaGetLastError() != cudaSuccess;
+}
+
+nvinfer1::DataType TransformerOutputConvertPlugin::getOutputDataType(
+    int index,
+    const nvinfer1::DataType* input_types,
+    int nb_inputs) const TRT_NOEXCEPT {
+  if (index == 0) {
+    return nvinfer1::DataType::kHALF;
+  }
+}
+
+nvinfer1::DimsExprs TransformerOutputConvertPlugin::getOutputDimensions(
+    int outputIndex,
+    const nvinfer1::DimsExprs* inputs,
+    int nbInputs,
+    nvinfer1::IExprBuilder& exprBuilder) TRT_NOEXCEPT {
+  nvinfer1::DimsExprs output_dims;
+  if (outputIndex == 0) {
+    output_dims = inputs[1];
+  }
+  return output_dims;
+}
+
+bool TransformerOutputConvertPlugin::supportsFormatCombination(
+    int pos,
+    const nvinfer1::PluginTensorDesc* inOut,
+    int nbInputs,
+    int nbOutputs) TRT_NOEXCEPT {
+  PADDLE_ENFORCE_EQ(nbInputs,
+                    3,
+                    common::errors::InvalidArgument(
+                        "TransformerOutputConvertPlugin must have 3 inputs, "
+                        "but got %d input(s). ",
+                        nbInputs));
+  PADDLE_ENFORCE_EQ(nbOutputs,
+                    1,
+                    common::errors::InvalidArgument(
+                        "TransformerOutputConvertPlugin must have 1 output, "
+                        "but got %d output(s). ",
+                        nbOutputs));
+  if (pos == 0) {  // qkv plugin output(varlen)
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kHALF;
+  } else if (pos == 1) {  // qkv plugin input(no_varlen)
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kHALF;
+  } else if (pos == 2) {  // pos id
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kINT32;
+  } else if (pos == 3) {  // qkv plugin output(no_varlen)
+    return inOut[pos].format == nvinfer1::TensorFormat::kLINEAR &&
+           inOut[pos].type == nvinfer1::DataType::kHALF;
+  }
+}
+
+void TransformerOutputConvertPlugin::configurePlugin(
+    const nvinfer1::DynamicPluginTensorDesc* inputs,
+    int nbInputs,
+    const nvinfer1::DynamicPluginTensorDesc* outputs,
+    int nbOutputs) TRT_NOEXCEPT {}
+
+void TransformerOutputConvertPlugin::attachToContext(
+    cudnnContext* cudnnContext,
+    cublasContext* cublasContext,
+    nvinfer1::IGpuAllocator* gpuAllocator) TRT_NOEXCEPT {}
+
+void TransformerOutputConvertPlugin::detachFromContext() TRT_NOEXCEPT {}
+
+void TransformerOutputConvertPlugin::terminate() TRT_NOEXCEPT {}
+
+int TransformerOutputConvertPlugin::enqueue(
+    const nvinfer1::PluginTensorDesc* inputDesc,
+    const nvinfer1::PluginTensorDesc* outputDesc,
+    const void* const* inputs,
+    void* const* outputs,
+    void* workspace,
+    cudaStream_t stream) TRT_NOEXCEPT {
+  const half* input0 =
+      static_cast<const half*>(inputs[0]);  // qkv plugin output(varlen)
+  const half* input1 =
+      static_cast<const half*>(inputs[1]);  // qkv plugin input(no_varlen)
+  const int32_t* input2 = static_cast<const int32_t*>(inputs[2]);  // pos id
+  half* output =
+      static_cast<half*>(outputs[0]);  // qkv plugin output(no_varlen)
+  const auto input1_desc = inputDesc[1];
+  const int32_t B = input1_desc.dims.d[0];           // batches
+  const int32_t MaxLength = input1_desc.dims.d[1];   // max token length
+  const int32_t HiddenSize = input1_desc.dims.d[2];  // hidden size
+
+  const int32_t vector_length = HiddenSize;
+  int32_t num_threads;
+  if (vector_length < 1024) {
+    num_threads = vector_length;
+  } else {
+    if (vector_length % 512 == 0) {
+      num_threads = 512;
+    } else if (vector_length % 256 == 0) {
+      num_threads = 256;
+    } else if (vector_length % 128 == 0) {
+      num_threads = 128;
+    } else if (vector_length % 64 == 0) {
+      num_threads = 64;
+    } else if (vector_length % 32 == 0) {
+      num_threads = 32;
+    } else if (vector_length % 16 == 0) {
+      num_threads = 16;
+    } else if (vector_length % 8 == 0) {
+      num_threads = 8;
+    } else if (vector_length % 4 == 0) {
+      num_threads = 4;
+    } else if (vector_length % 2 == 0) {
+      num_threads = 2;
+    } else {
+      num_threads = 1;
+    }
+  }
+  const dim3 num_blocks(
+      B,
+      MaxLength,
+      vector_length / num_threads);  //  batches, max sequence length
+                                     //  (mask_id.dims.d[1]),
+                                     //  input.dims.d[1]/*
+  recover_padding_kernel<<<num_blocks, num_threads, 0, stream>>>(
+      input0, input2, output);
+  return cudaGetLastError() != cudaSuccess;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/transformer_input_output_convert_plugin.h

@@ -0,0 +1,256 @@
+/* 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. */
+
+#pragma once
+#include <cassert>
+#include <string>
+
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class TransformerInputConvertPlugin : public DynamicPluginTensorRT {
+ public:
+  TransformerInputConvertPlugin() {}
+
+  TransformerInputConvertPlugin(void const* serial_data, size_t serial_length) {
+  }
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    TransformerInputConvertPlugin* ptr = new TransformerInputConvertPlugin();
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "transformer_input_convert_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 4; }
+
+  int initialize() TRT_NOEXCEPT { return 0; }
+  void terminate() TRT_NOEXCEPT;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* inputs,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* outputs,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  void attachToContext(cudnnContext* cudnnContext,
+                       cublasContext* cublasContext,
+                       nvinfer1::IGpuAllocator* gpuAllocator)
+      TRT_NOEXCEPT override;
+
+  void detachFromContext() TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ protected:
+  size_t getSerializationSize() const TRT_NOEXCEPT override { return 0; }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {}
+};
+
+class TransformerInputConvertPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  TransformerInputConvertPluginCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "transformer_input_convert_plugin";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(
+      const char* name, const nvinfer1::PluginFieldCollection* plugin_field)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         void const* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    TransformerInputConvertPlugin* obj =
+        new TransformerInputConvertPlugin(serial_data, serial_length);
+    obj->setPluginNamespace(name);
+    return obj;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+};
+
+class TransformerOutputConvertPlugin : public DynamicPluginTensorRT {
+ public:
+  TransformerOutputConvertPlugin() {}
+
+  TransformerOutputConvertPlugin(void const* serial_data,
+                                 size_t serial_length) {}
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override {
+    TransformerOutputConvertPlugin* ptr = new TransformerOutputConvertPlugin();
+    return ptr;
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "transformer_output_convert_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+
+  int initialize() TRT_NOEXCEPT { return 0; }
+  void terminate() TRT_NOEXCEPT;
+  nvinfer1::DimsExprs getOutputDimensions(
+      int outputIndex,
+      const nvinfer1::DimsExprs* inputs,
+      int nbInputs,
+      nvinfer1::IExprBuilder& exprBuilder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* inOut,
+                                 int nbInputs,
+                                 int nbOutputs) TRT_NOEXCEPT override;
+
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* inputs,
+                       int nbInputs,
+                       const nvinfer1::DynamicPluginTensorDesc* outputs,
+                       int nbOutputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  void attachToContext(cudnnContext* cudnnContext,
+                       cublasContext* cublasContext,
+                       nvinfer1::IGpuAllocator* gpuAllocator)
+      TRT_NOEXCEPT override;
+
+  void detachFromContext() TRT_NOEXCEPT override;
+
+  int enqueue(const nvinfer1::PluginTensorDesc* inputDesc,
+              const nvinfer1::PluginTensorDesc* outputDesc,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* inputTypes,
+                                       int nbInputs) const
+      TRT_NOEXCEPT override;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+ protected:
+  size_t getSerializationSize() const TRT_NOEXCEPT override { return 0; }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {}
+};
+
+class TransformerOutputConvertPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  TransformerOutputConvertPluginCreator() {}
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "transformer_output_convert_plugin";
+  }
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override {
+    return &field_collection_;
+  }
+
+  nvinfer1::IPluginV2* createPlugin(
+      const char* name, const nvinfer1::PluginFieldCollection* plugin_field)
+      TRT_NOEXCEPT override {
+    return nullptr;
+  }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         void const* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    TransformerOutputConvertPlugin* obj =
+        new TransformerOutputConvertPlugin(serial_data, serial_length);
+    obj->setPluginNamespace(name);
+    return obj;
+  }
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override {
+    plugin_namespace_ = lib_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return plugin_namespace_.c_str();
+  }
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+};
+
+REGISTER_TRT_PLUGIN_V2(TransformerInputConvertPluginCreator);
+REGISTER_TRT_PLUGIN_V2(TransformerOutputConvertPluginCreator);
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/trt_plugin.cc

@@ -0,0 +1,139 @@
+// Copyright (c) 2018 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 "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle::inference::tensorrt::plugin {
+
+inline void Serialize(void*& buffer,  // NOLINT
+                      const std::vector<nvinfer1::Dims>& input_dims,
+                      nvinfer1::DataType data_type,
+                      nvinfer1::PluginFormat data_format,
+                      bool with_fp16) {
+  SerializeValue(&buffer, input_dims);
+  SerializeValue(&buffer, data_type);
+  SerializeValue(&buffer, data_format);
+  SerializeValue(&buffer, with_fp16);
+}
+
+inline void Deserialize(void const*& serial_data,  // NOLINT
+                        size_t& serial_length,     // NOLINT
+                        std::vector<nvinfer1::Dims>* input_dims,
+                        nvinfer1::DataType* data_type,
+                        nvinfer1::PluginFormat* data_format,
+                        bool* with_fp16) {
+  DeserializeValue(&serial_data, &serial_length, input_dims);
+  DeserializeValue(&serial_data, &serial_length, data_type);
+  DeserializeValue(&serial_data, &serial_length, data_format);
+  DeserializeValue(&serial_data, &serial_length, with_fp16);
+}
+
+inline size_t SerializeSize(const std::vector<nvinfer1::Dims>& input_dims,
+                            nvinfer1::DataType data_type,
+                            nvinfer1::PluginFormat data_format,
+                            bool with_fp16) {
+  return (SerializedSize(input_dims) + SerializedSize(data_type) +
+          SerializedSize(data_format) + SerializedSize(with_fp16));
+}
+
+void PluginTensorRT::serializeBase(void*& buffer) const {
+  Serialize(buffer, input_dims_, data_type_, data_format_, with_fp16_);
+}
+
+void PluginTensorRT::deserializeBase(void const*& serial_data,
+                                     size_t& serial_length) {
+  Deserialize(serial_data,
+              serial_length,
+              &input_dims_,
+              &data_type_,
+              &data_format_,
+              &with_fp16_);
+}
+
+size_t PluginTensorRT::getBaseSerializationSize() const {
+  return SerializeSize(input_dims_, data_type_, data_format_, with_fp16_);
+}
+
+bool PluginTensorRT::supportsFormat(
+    nvinfer1::DataType type, nvinfer1::PluginFormat format) const TRT_NOEXCEPT {
+  return ((type == nvinfer1::DataType::kFLOAT) &&
+          (format == nvinfer1::PluginFormat::kLINEAR));
+}
+
+void PluginTensorRT::configureWithFormat(const nvinfer1::Dims* input_dims,
+                                         int num_inputs,
+                                         const nvinfer1::Dims* output_dims,
+                                         int num_outputs,
+                                         nvinfer1::DataType type,
+                                         nvinfer1::PluginFormat format,
+                                         int max_batch_size) TRT_NOEXCEPT {
+  data_type_ = type;
+  data_format_ = format;
+  input_dims_.assign(input_dims, input_dims + num_inputs);
+}
+
+void PluginTensorRTV2Ext::serializeBase(void*& buffer) const {
+  Serialize(buffer, input_dims_, data_type_, data_format_, with_fp16_);
+}
+
+void PluginTensorRTV2Ext::deserializeBase(void const*& serial_data,
+                                          size_t& serial_length) {
+  Deserialize(serial_data,
+              serial_length,
+              &input_dims_,
+              &data_type_,
+              &data_format_,
+              &with_fp16_);
+}
+
+size_t PluginTensorRTV2Ext::getBaseSerializationSize() const {
+  return SerializeSize(input_dims_, data_type_, data_format_, with_fp16_);
+}
+
+void PluginTensorRTV2Ext::configurePlugin(
+    const nvinfer1::Dims* input_dims,
+    int32_t nb_inputs,
+    const nvinfer1::Dims* output_dims,
+    int32_t nb_outputs,
+    const nvinfer1::DataType* input_types,
+    const nvinfer1::DataType* output_types,
+    const bool* input_is_broadcast,
+    const bool* output_is_broadcast,
+    nvinfer1::PluginFormat float_format,
+    int32_t max_batch_size) TRT_NOEXCEPT {
+  input_dims_.assign(input_dims, input_dims + nb_inputs);
+  data_format_ = float_format;
+  data_type_ = input_types[0];
+}
+
+const nvinfer1::PluginFieldCollection* TensorRTPluginCreator::getFieldNames()
+    TRT_NOEXCEPT {
+  return &field_collection_;
+}
+
+nvinfer1::IPluginV2* TensorRTPluginCreator::createPlugin(
+    const char* name, const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT {
+  return nullptr;
+}
+
+void TensorRTPluginCreator::setPluginNamespace(const char* lib_namespace)
+    TRT_NOEXCEPT {
+  plugin_namespace_ = lib_namespace;
+}
+
+const char* TensorRTPluginCreator::getPluginNamespace() const TRT_NOEXCEPT {
+  return plugin_namespace_.c_str();
+}
+
+}  // namespace paddle::inference::tensorrt::plugin

paddle/fluid/inference/tensorrt/plugin/trt_plugin.h

@@ -0,0 +1,377 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <NvInfer.h>
+
+#include <cstring>
+#include <string>
+#include <unordered_map>
+#include <utility>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/helper.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin_utils.h"
+#include "paddle/fluid/platform/enforce.h"
+#include "paddle/fluid/platform/tensorrt/trt_plugin.h"
+#include "paddle/phi/core/platform/profiler/event_tracing.h"
+
+namespace nvinfer1 {
+class ITensor;
+}  // namespace nvinfer1
+
+PD_DECLARE_bool(profile);
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+#if defined(_WIN32)
+#define UNUSED
+#define __builtin_expect(EXP, C) (EXP)
+#else
+#define UNUSED __attribute__((unused))
+#endif
+
+class PluginTensorRT;
+
+typedef std::function<PluginTensorRT*(const void*, size_t)>
+    PluginDeserializeFunc;
+
+typedef std::function<PluginTensorRT*(void)> PluginConstructFunc;
+
+// Deprecated. Do not inherit this class, please refer to PluginTensorRTV2Ext
+class PluginTensorRT : public nvinfer1::IPluginV2 {
+ public:
+  PluginTensorRT() : with_fp16_(false) {}
+
+  // It was used for TensorRT deserialization.
+  // It should not be called by users.
+  PluginTensorRT(const void* serialized_data, size_t length) {}
+
+  virtual ~PluginTensorRT() {}
+
+  nvinfer1::Dims const& getInputDims(int index) const {
+    return input_dims_.at(index);
+  }
+
+  nvinfer1::DataType getDataType() const { return data_type_; }
+
+  nvinfer1::PluginFormat getDataFormat() const { return data_format_; }
+
+  // IPluginV2
+  virtual const char* getPluginType() const TRT_NOEXCEPT = 0;
+
+  virtual const char* getPluginVersion() const TRT_NOEXCEPT { return "1"; }
+
+  int getNbOutputs() const TRT_NOEXCEPT { return 1; }
+
+  virtual nvinfer1::Dims getOutputDimensions(int index,
+                                             const nvinfer1::Dims* input_dims,
+                                             int num_inputs) TRT_NOEXCEPT = 0;
+
+  // Check format support. The default is FLOAT32 and kLINEAR.
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override;
+
+  // Configure the layer
+  void configureWithFormat(const nvinfer1::Dims* input_dims,
+                           int num_inputs,
+                           const nvinfer1::Dims* output_dims,
+                           int num_outputs,
+                           nvinfer1::DataType type,
+                           nvinfer1::PluginFormat format,
+                           int max_batch_size) TRT_NOEXCEPT override;
+
+  // Initialize the layer for execution.
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  // Shutdown the layer. This is called when the engine is destroyed
+  void terminate() TRT_NOEXCEPT override {}
+
+  // Find the workspace size required by the layer
+  size_t getWorkspaceSize(int) const TRT_NOEXCEPT override { return 0; }
+
+  // Execute the layer
+  virtual int enqueue(int batch_size,
+                      const void* const* inputs,
+                      void* const* outputs,
+                      void* workspace,
+                      cudaStream_t stream) TRT_NOEXCEPT = 0;
+
+  // Find the size of the serialization buffer required
+  virtual size_t getSerializationSize() const TRT_NOEXCEPT = 0;
+
+  // Serialize the layer config to buffer.
+  // TensorRT will call this func to serialize the configuration of TensorRT
+  // engine. It should not be called by users.
+  virtual void serialize(void* buffer) const TRT_NOEXCEPT = 0;
+
+  void destroy() TRT_NOEXCEPT override { delete this; }
+
+  virtual nvinfer1::IPluginV2* clone() const TRT_NOEXCEPT = 0;
+
+  void setPluginNamespace(const char* plugin_namespace) TRT_NOEXCEPT override {
+    namespace_ = plugin_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return namespace_.c_str();
+  }
+
+ protected:
+  // Deserialize input_dims, max_batch_size, data_type, data_format
+  void deserializeBase(void const*& serial_data,  // NOLINT
+                       size_t& serial_length);    // NOLINT
+  size_t getBaseSerializationSize() const;
+  // Serialize input_dims, max_batch_size, data_type, data_format
+  void serializeBase(void*& buffer) const;  // NOLINT
+
+  std::vector<nvinfer1::Dims> input_dims_;
+  nvinfer1::DataType data_type_;
+  nvinfer1::PluginFormat data_format_;
+
+  bool with_fp16_;
+
+ private:
+  std::string namespace_;
+};
+
+// TensorRT introduced IPluginV2Ext after 5.1, Paddle no longer supports
+// versions before 5.1
+class PluginTensorRTV2Ext : public nvinfer1::IPluginV2Ext {
+ public:
+  PluginTensorRTV2Ext() : with_fp16_(false) {}
+  PluginTensorRTV2Ext(const void* serialized_data, size_t length) {}
+
+  nvinfer1::Dims const& getInputDims(int index) const {
+    return input_dims_.at(index);
+  }
+  nvinfer1::DataType getDataType() const { return data_type_; }
+  nvinfer1::PluginFormat getDataFormat() const { return data_format_; }
+
+  // The Func in IPluginV2Ext
+  virtual nvinfer1::DataType getOutputDataType(
+      int index,
+      const nvinfer1::DataType* input_types,
+      int nb_inputs) const TRT_NOEXCEPT = 0;
+
+  virtual bool isOutputBroadcastAcrossBatch(int32_t output_index,
+                                            const bool* input_is_broadcasted,
+                                            int32_t nb_inputs) const
+      TRT_NOEXCEPT {
+    return false;
+  }
+
+  virtual bool canBroadcastInputAcrossBatch(int32_t input_index) const
+      TRT_NOEXCEPT {
+    return false;
+  }
+
+  void configurePlugin(const nvinfer1::Dims* input_dims,
+                       int32_t nb_inputs,
+                       const nvinfer1::Dims* output_dims,
+                       int32_t nb_outputs,
+                       const nvinfer1::DataType* input_types,
+                       const nvinfer1::DataType* output_types,
+                       const bool* input_is_broadcast,
+                       const bool* output_is_broadcast,
+                       nvinfer1::PluginFormat float_format,
+                       int32_t max_batch_size) TRT_NOEXCEPT override;
+
+  virtual IPluginV2Ext* clone() const TRT_NOEXCEPT = 0;
+
+  void attachToContext(cudnnContext*,
+                       cublasContext*,
+                       nvinfer1::IGpuAllocator*) TRT_NOEXCEPT override {}
+
+  void detachFromContext() TRT_NOEXCEPT override {}
+
+  // The Func in IPluginV2
+  virtual const char* getPluginType() const TRT_NOEXCEPT = 0;
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+  virtual int32_t getNbOutputs() const TRT_NOEXCEPT { return 1; }
+  virtual nvinfer1::Dims getOutputDimensions(int32_t index,
+                                             const nvinfer1::Dims* inputs,
+                                             int32_t nb_input) TRT_NOEXCEPT = 0;
+  // Check format support. The default is FLOAT32 and NCHW.
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::PluginFormat format)
+      const TRT_NOEXCEPT override {
+    return ((type == nvinfer1::DataType::kFLOAT) &&
+            (format == nvinfer1::PluginFormat::kLINEAR));
+  }
+  // Initialize the layer for execution.
+  // This is called when the engine is created.
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  // Shutdown the layer. This is called when the engine is destroyed
+  void terminate() TRT_NOEXCEPT override {}
+
+  // Find the workspace size required by the layer
+  size_t getWorkspaceSize(int) const TRT_NOEXCEPT override { return 0; }
+
+  // Execute the layer
+  virtual int enqueue(int batch_size,
+                      const void* const* inputs,
+                      void* const* outputs,
+                      void* workspace,
+                      cudaStream_t stream) TRT_NOEXCEPT = 0;
+
+  // Find the size of the serialization buffer required
+  virtual size_t getSerializationSize() const TRT_NOEXCEPT = 0;
+
+  // Serialize the layer config to buffer.
+  // TensorRT will call this func to serialize the configuration of TensorRT
+  // engine. It should not be called by users.
+  virtual void serialize(void* buffer) const TRT_NOEXCEPT = 0;
+
+  virtual void destroy() TRT_NOEXCEPT = 0;
+
+  void setPluginNamespace(const char* plugin_namespace) TRT_NOEXCEPT override {
+    name_space_ = plugin_namespace;
+  }
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return name_space_.c_str();
+  }
+
+ protected:
+  void deserializeBase(void const*& serial_data,  // NOLINT
+                       size_t& serial_length);    // NOLINT
+  size_t getBaseSerializationSize() const;
+  void serializeBase(void*& buffer) const;  // NOLINT
+
+ protected:
+  std::vector<nvinfer1::Dims> input_dims_;
+  nvinfer1::DataType data_type_;
+  nvinfer1::PluginFormat data_format_;
+  bool with_fp16_;
+
+ private:
+  std::string name_space_;
+};
+
+class DynamicPluginTensorRT : public nvinfer1::IPluginV2DynamicExt {
+ public:
+  DynamicPluginTensorRT() : with_fp16_(false) {}
+  DynamicPluginTensorRT(const void* serialized_data, size_t length) {}
+
+  // The Func in IPluginExt or IpluginExtV2
+  virtual const char* getPluginVersion() const TRT_NOEXCEPT { return "1"; }
+  virtual const char* getPluginType() const TRT_NOEXCEPT = 0;
+  int getNbOutputs() const TRT_NOEXCEPT { return 1; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  void terminate() TRT_NOEXCEPT override{};
+
+  virtual size_t getSerializationSize() const TRT_NOEXCEPT = 0;
+  virtual void serialize(void* buffer) const TRT_NOEXCEPT = 0;
+
+  // The Func in IPluginV2
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT = 0;
+  virtual nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder) TRT_NOEXCEPT = 0;  // NOLINT
+
+  virtual bool supportsFormatCombination(
+      int pos,
+      const nvinfer1::PluginTensorDesc* in_out,
+      int nb_inputs,
+      int nb_outputs) TRT_NOEXCEPT = 0;
+
+  virtual void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                               int nb_inputs,
+                               const nvinfer1::DynamicPluginTensorDesc* out,
+                               int nb_outputs) TRT_NOEXCEPT = 0;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nb_inputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nb_outputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  virtual int 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 = 0;
+
+  virtual nvinfer1::DataType getOutputDataType(
+      int index,
+      const nvinfer1::DataType* input_types,
+      int nb_inputs) const TRT_NOEXCEPT = 0;
+  void setPluginNamespace(const char* plugin_namespace) TRT_NOEXCEPT override {
+    name_space_ = plugin_namespace;
+  }
+  const char* getPluginNamespace() const TRT_NOEXCEPT override {
+    return name_space_.c_str();
+  }
+  virtual void destroy() TRT_NOEXCEPT = 0;
+
+ protected:
+  void deserializeBase(void const*& serial_data,  // NOLINT
+                       size_t& serial_length);    // NOLINT
+  size_t getBaseSerializationSize() const;
+  void serializeBase(void*& buffer) const;  // NOLINT
+  bool with_fp16_;
+
+ private:
+  std::string name_space_;
+  std::string plugin_base_;
+};
+
+class TensorRTPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  TensorRTPluginCreator() = default;
+
+  virtual const char* getPluginName() const TRT_NOEXCEPT = 0;
+
+  virtual const char* getPluginVersion() const TRT_NOEXCEPT = 0;
+
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2* createPlugin(const char* name,
+                                    const nvinfer1::PluginFieldCollection* fc)
+      TRT_NOEXCEPT override;
+
+  virtual nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                                 const void* serial_data,
+                                                 size_t serial_length)
+      TRT_NOEXCEPT = 0;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+
+ private:
+  std::string plugin_namespace_;
+  std::string plugin_name_;
+  nvinfer1::PluginFieldCollection field_collection_{0, nullptr};
+  std::vector<nvinfer1::PluginField> plugin_attributes_;
+};
+
+using TrtPluginRegistry = paddle::platform::TrtPluginRegistry;
+
+template <typename T>
+using TrtPluginRegistrarV2 = paddle::platform::TrtPluginRegistrarV2<T>;
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/trt_plugin_utils.h

@@ -0,0 +1,155 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+#include <cuda_fp16.h>
+#include <cstring>
+#include <string>
+#include <type_traits>
+#include <vector>
+
+#include "paddle/fluid/platform/enforce.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+// Some trt base classes lack of the destructor.
+// We use a assisted class to fix this.
+struct DeleteHelper {
+ protected:
+  virtual ~DeleteHelper() {}
+};
+
+template <typename T>
+inline void SerializeValue(void** buffer, T const& value);
+
+template <typename T>
+inline void DeserializeValue(void const** buffer,
+                             size_t* buffer_size,
+                             T* value);
+
+namespace details {
+
+template <typename T, class Enable = void>
+struct Serializer {};
+
+template <typename T>
+struct Serializer<
+    T,
+    typename std::enable_if<std::is_arithmetic<T>::value ||
+                            std::is_enum<T>::value || std::is_pod<T>::value ||
+                            std::is_same<T, half>::value>::type> {
+  static size_t SerializedSize(T const& value) { return sizeof(T); }
+
+  static void Serialize(void** buffer, T const& value) {
+    std::memcpy(*buffer, &value, sizeof(T));
+    reinterpret_cast<char*&>(*buffer) += sizeof(T);
+  }
+
+  static void Deserialize(void const** buffer, size_t* buffer_size, T* value) {
+    assert(*buffer_size >= sizeof(T));
+    std::memcpy(value, *buffer, sizeof(T));
+    reinterpret_cast<char const*&>(*buffer) += sizeof(T);
+    *buffer_size -= sizeof(T);
+  }
+};
+
+template <>
+struct Serializer<const char*> {
+  static size_t SerializedSize(const char* value) { return strlen(value) + 1; }
+
+  static void Serialize(void** buffer, const char* value) {
+    std::strcpy(static_cast<char*>(*buffer), value);  // NOLINT
+    reinterpret_cast<char*&>(*buffer) += strlen(value) + 1;
+  }
+
+  static void Deserialize(void const** buffer,
+                          size_t* buffer_size,
+                          const char** value) {
+    *value = static_cast<char const*>(*buffer);
+    size_t data_size = strnlen(*value, *buffer_size) + 1;
+    assert(*buffer_size >= data_size);
+    reinterpret_cast<char const*&>(*buffer) += data_size;
+    *buffer_size -= data_size;
+  }
+};
+
+template <typename T>
+struct Serializer<
+    std::vector<T>,
+    typename std::enable_if<std::is_arithmetic<T>::value ||
+                            std::is_enum<T>::value || std::is_pod<T>::value ||
+                            std::is_same<T, half>::value>::type> {
+  static size_t SerializedSize(std::vector<T> const& value) {
+    return sizeof(value.size()) + value.size() * sizeof(T);
+  }
+
+  static void Serialize(void** buffer, std::vector<T> const& value) {
+    SerializeValue(buffer, value.size());
+    size_t nbyte = value.size() * sizeof(T);
+    std::memcpy(*buffer, value.data(), nbyte);
+    reinterpret_cast<char*&>(*buffer) += nbyte;
+  }
+
+  static void Deserialize(void const** buffer,
+                          size_t* buffer_size,
+                          std::vector<T>* value) {
+    size_t size;
+    DeserializeValue(buffer, buffer_size, &size);
+    value->resize(size);
+    size_t nbyte = value->size() * sizeof(T);
+    PADDLE_ENFORCE_GE(*buffer_size,
+                      nbyte,
+                      common::errors::InvalidArgument(
+                          "Insufficient data in buffer, expect contains %d "
+                          "byte, but actually only contains %d byte.",
+                          *buffer_size,
+                          nbyte));
+    std::memcpy(value->data(), *buffer, nbyte);
+    reinterpret_cast<char const*&>(*buffer) += nbyte;
+    *buffer_size -= nbyte;
+  }
+};
+
+}  // namespace details
+
+template <typename T>
+inline size_t SerializedSize(T const& value) {
+  return details::Serializer<T>::SerializedSize(value);
+}
+
+template <typename T>
+inline void SerializeValue(void** buffer, T const& value) {
+  return details::Serializer<T>::Serialize(buffer, value);
+}
+
+template <typename T>
+inline void DeserializeValue(void const** buffer,
+                             size_t* buffer_size,
+                             T* value) {
+  return details::Serializer<T>::Deserialize(buffer, buffer_size, value);
+}
+
+template <typename T>
+inline void SerializeCudaPointer(void** buffer, T* value, int size) {
+  cudaMemcpy((*buffer), value, size * sizeof(T), cudaMemcpyDeviceToHost);
+  reinterpret_cast<char*&>(*buffer) += size * sizeof(T);
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/yolo_box_head_op_plugin.cu

@@ -0,0 +1,95 @@
+// 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 "paddle/fluid/inference/tensorrt/plugin/yolo_box_head_op_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+inline __device__ float SigmoidGPU(const float& x) {
+  return 1.0f / (1.0f + __expf(-x));
+}
+
+__global__ void YoloBoxHeadKernel(const float* input,
+                                  float* output,
+                                  const int grid_size_x,
+                                  const int grid_size_y,
+                                  const int class_num,
+                                  const int anchors_num) {
+  int x_id = blockIdx.x * blockDim.x + threadIdx.x;
+  int y_id = blockIdx.y * blockDim.y + threadIdx.y;
+  int z_id = blockIdx.z * blockDim.z + threadIdx.z;
+  if ((x_id >= grid_size_x) || (y_id >= grid_size_y) || (z_id >= anchors_num)) {
+    return;
+  }
+  const int grids_num = grid_size_x * grid_size_y;
+  const int bbindex = y_id * grid_size_x + x_id;
+
+  // objectness
+  output[bbindex + grids_num * (z_id * (5 + class_num) + 4)] =
+      SigmoidGPU(input[bbindex + grids_num * (z_id * (5 + class_num) + 4)]);
+  // x
+  output[bbindex + grids_num * (z_id * (5 + class_num) + 0)] =
+      SigmoidGPU(input[bbindex + grids_num * (z_id * (5 + class_num) + 0)]);
+  // y
+  output[bbindex + grids_num * (z_id * (5 + class_num) + 1)] =
+      SigmoidGPU(input[bbindex + grids_num * (z_id * (5 + class_num) + 1)]);
+  // w
+  output[bbindex + grids_num * (z_id * (5 + class_num) + 2)] =
+      __expf(input[bbindex + grids_num * (z_id * (5 + class_num) + 2)]);
+  // h
+  output[bbindex + grids_num * (z_id * (5 + class_num) + 3)] =
+      __expf(input[bbindex + grids_num * (z_id * (5 + class_num) + 3)]);
+  // Probabilities of classes
+  for (int i = 0; i < class_num; ++i) {
+    output[bbindex + grids_num * (z_id * (5 + class_num) + (5 + i))] =
+        SigmoidGPU(
+            input[bbindex + grids_num * (z_id * (5 + class_num) + (5 + i))]);
+  }
+}
+
+int YoloBoxHeadPlugin::enqueue(int batch_size,
+                               const void* const* inputs,
+                               void* const* outputs,
+                               void* workspace,
+                               cudaStream_t stream) TRT_NOEXCEPT {
+  const int h = input_dims_[0].d[1];
+  const int w = input_dims_[0].d[2];
+  const int grid_size_x = w;
+  const int grid_size_y = h;
+  const int anchors_num = anchors_.size() / 2;
+  const float* input_data = static_cast<const float*>(inputs[0]);
+  float* output_data = static_cast<float*>(outputs[0]);
+  const int volume = input_dims_[0].d[0] * h * w;
+  dim3 block(16, 16, 4);
+  dim3 grid((grid_size_x / block.x) + 1,
+            (grid_size_y / block.y) + 1,
+            (anchors_num / block.z) + 1);
+  for (int n = 0; n < batch_size; n++) {
+    YoloBoxHeadKernel<<<grid, block, 0, stream>>>(input_data + n * volume,
+                                                  output_data + n * volume,
+                                                  grid_size_x,
+                                                  grid_size_y,
+                                                  class_num_,
+                                                  anchors_num);
+  }
+  return 0;
+}
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/yolo_box_head_op_plugin.h

@@ -0,0 +1,105 @@
+// 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.
+
+#pragma once
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class YoloBoxHeadPlugin : public PluginTensorRT {
+ public:
+  explicit YoloBoxHeadPlugin(const std::vector<int>& anchors,
+                             const int class_num)
+      : anchors_(anchors), class_num_(class_num) {}
+
+  YoloBoxHeadPlugin(const void* data, size_t length) {
+    deserializeBase(data, length);
+    DeserializeValue(&data, &length, &anchors_);
+    DeserializeValue(&data, &length, &class_num_);
+  }
+
+  ~YoloBoxHeadPlugin() override{};
+
+  nvinfer1::IPluginV2* clone() const TRT_NOEXCEPT override {
+    return new YoloBoxHeadPlugin(anchors_, class_num_);
+  }
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "yolo_box_head_plugin";
+  }
+
+  int getNbOutputs() const TRT_NOEXCEPT override { return 1; }
+
+  int initialize() TRT_NOEXCEPT override { return 0; }
+
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nb_input_dims) TRT_NOEXCEPT override {
+    assert(index == 0);
+    assert(nb_input_dims == 1);
+    return inputs[0];
+  }
+
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+
+  size_t getSerializationSize() const TRT_NOEXCEPT override {
+    return getBaseSerializationSize() + SerializedSize(anchors_) +
+           SerializedSize(class_num_);
+  }
+
+  void serialize(void* buffer) const TRT_NOEXCEPT override {
+    serializeBase(buffer);
+    SerializeValue(&buffer, anchors_);
+    SerializeValue(&buffer, class_num_);
+  }
+
+ private:
+  std::vector<int> anchors_;
+  int class_num_;
+  std::string namespace_;
+};
+
+class YoloBoxHeadPluginCreator : public TensorRTPluginCreator {
+ public:
+  const char* getPluginName() const TRT_NOEXCEPT override {
+    return "yolo_box_head_plugin";
+  }
+
+  const char* getPluginVersion() const TRT_NOEXCEPT override { return "1"; }
+
+  nvinfer1::IPluginV2* deserializePlugin(const char* name,
+                                         const void* serial_data,
+                                         size_t serial_length)
+      TRT_NOEXCEPT override {
+    return new YoloBoxHeadPlugin(serial_data, serial_length);
+  }
+};
+
+REGISTER_TRT_PLUGIN_V2(YoloBoxHeadPluginCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle

paddle/fluid/inference/tensorrt/plugin/yolo_box_op_plugin.h

@@ -0,0 +1,352 @@
+// Copyright (c) 2018 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.
+
+#pragma once
+
+#include <string>
+#include <vector>
+
+#include "paddle/fluid/inference/tensorrt/engine.h"
+#include "paddle/fluid/inference/tensorrt/plugin/trt_plugin.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+namespace plugin {
+
+class YoloBoxPlugin : public nvinfer1::IPluginV2Ext {
+ public:
+  explicit YoloBoxPlugin(const nvinfer1::DataType data_type,
+                         const std::vector<int>& anchors,
+                         const int class_num,
+                         const float conf_thresh,
+                         const int downsample_ratio,
+                         const bool clip_bbox,
+                         const float scale_x_y,
+                         const bool iou_aware,
+                         const float iou_aware_factor,
+                         const int input_h,
+                         const int input_w);
+  YoloBoxPlugin(const void* data, size_t length);
+  ~YoloBoxPlugin() override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nb_input_dims) TRT_NOEXCEPT override;
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::TensorFormat format)
+      const TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(int max_batch_size) const TRT_NOEXCEPT override;
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  template <typename T>
+  int enqueue_impl(int batch_size,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_type,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+  bool isOutputBroadcastAcrossBatch(int output_index,
+                                    const bool* input_is_broadcast,
+                                    int nb_inputs) const TRT_NOEXCEPT override;
+  bool canBroadcastInputAcrossBatch(int input_index) const
+      TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::Dims* input_dims,
+                       int nb_inputs,
+                       const nvinfer1::Dims* output_dims,
+                       int nb_outputs,
+                       const nvinfer1::DataType* input_types,
+                       const nvinfer1::DataType* output_types,
+                       const bool* input_is_broadcast,
+                       const bool* output_is_broadcast,
+                       nvinfer1::PluginFormat float_format,
+                       int max_batch_size) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* clone() const TRT_NOEXCEPT override;
+
+ private:
+  nvinfer1::DataType data_type_;
+  std::vector<int> anchors_;
+  int* anchors_device_;
+  int class_num_;
+  float conf_thresh_;
+  int downsample_ratio_;
+  bool clip_bbox_;
+  float scale_x_y_;
+  int input_h_;
+  int input_w_;
+  bool iou_aware_;
+  float iou_aware_factor_;
+  std::string namespace_;
+};
+
+class YoloBoxPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  YoloBoxPluginCreator();
+  ~YoloBoxPluginCreator() override = default;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+class YoloBoxPluginDynamic : public DynamicPluginTensorRT {
+ public:
+  explicit YoloBoxPluginDynamic(const nvinfer1::DataType data_type,
+                                const std::vector<int>& anchors,
+                                const int class_num,
+                                const float conf_thresh,
+                                const int downsample_ratio,
+                                const bool clip_bbox,
+                                const float scale_x_y,
+                                const bool iou_aware,
+                                const float iou_aware_factor);
+  YoloBoxPluginDynamic(void const* serialData, size_t serialLength);
+  ~YoloBoxPluginDynamic() {}
+
+  nvinfer1::IPluginV2DynamicExt* clone() const TRT_NOEXCEPT override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override {
+    return "yolo_box_plugin_dynamic";
+  }
+  int getNbOutputs() const TRT_NOEXCEPT override { return 2; }
+  int initialize() TRT_NOEXCEPT override { return 0; }
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+
+  nvinfer1::DimsExprs getOutputDimensions(
+      int output_index,
+      const nvinfer1::DimsExprs* inputs,
+      int nb_inputs,
+      nvinfer1::IExprBuilder& expr_builder)  // NOLINT
+      TRT_NOEXCEPT override;
+
+  bool supportsFormatCombination(int pos,
+                                 const nvinfer1::PluginTensorDesc* in_out,
+                                 int nb_inputs,
+                                 int nb_outputs) TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::DynamicPluginTensorDesc* in,
+                       int nb_inputs,
+                       const nvinfer1::DynamicPluginTensorDesc* out,
+                       int nb_outputs) TRT_NOEXCEPT override;
+
+  size_t getWorkspaceSize(const nvinfer1::PluginTensorDesc* inputs,
+                          int nbInputs,
+                          const nvinfer1::PluginTensorDesc* outputs,
+                          int nbOutputs) const TRT_NOEXCEPT override {
+    return 0;
+  }
+
+  int 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 override;
+
+  template <typename T>
+  int enqueue_impl(const nvinfer1::PluginTensorDesc* input_desc,
+                   const nvinfer1::PluginTensorDesc* output_desc,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_types,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+
+ private:
+  nvinfer1::DataType data_type_;
+  std::vector<int> anchors_;
+  int* anchors_device_;
+  int class_num_;
+  float conf_thresh_;
+  int downsample_ratio_;
+  bool clip_bbox_;
+  float scale_x_y_;
+  bool iou_aware_;
+  float iou_aware_factor_;
+};
+
+class YoloBoxPluginDynamicCreator : public nvinfer1::IPluginCreator {
+ public:
+  YoloBoxPluginDynamicCreator();
+  ~YoloBoxPluginDynamicCreator() override = default;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+class PIRYoloBoxPlugin : public nvinfer1::IPluginV2Ext {
+ public:
+  explicit PIRYoloBoxPlugin(const nvinfer1::DataType data_type,
+                            const std::vector<int>& anchors,
+                            const int class_num,
+                            const float conf_thresh,
+                            const int downsample_ratio,
+                            const bool clip_bbox,
+                            const float scale_x_y,
+                            const bool iou_aware,
+                            const float iou_aware_factor,
+                            const int input_h,
+                            const int input_w);
+  PIRYoloBoxPlugin(const void* data, size_t length);
+  ~PIRYoloBoxPlugin() override;
+
+  const char* getPluginType() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  int getNbOutputs() const TRT_NOEXCEPT override;
+  nvinfer1::Dims getOutputDimensions(int index,
+                                     const nvinfer1::Dims* inputs,
+                                     int nb_input_dims) TRT_NOEXCEPT override;
+  bool supportsFormat(nvinfer1::DataType type, nvinfer1::TensorFormat format)
+      const TRT_NOEXCEPT override;
+  size_t getWorkspaceSize(int max_batch_size) const TRT_NOEXCEPT override;
+  int enqueue(int batch_size,
+              const void* const* inputs,
+              void* const* outputs,
+              void* workspace,
+              cudaStream_t stream) TRT_NOEXCEPT override;
+  template <typename T>
+  int enqueue_impl(int batch_size,
+                   const void* const* inputs,
+                   void* const* outputs,
+                   void* workspace,
+                   cudaStream_t stream);
+  int initialize() TRT_NOEXCEPT override;
+  void terminate() TRT_NOEXCEPT override;
+  size_t getSerializationSize() const TRT_NOEXCEPT override;
+  void serialize(void* buffer) const TRT_NOEXCEPT override;
+  void destroy() TRT_NOEXCEPT override;
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+
+  nvinfer1::DataType getOutputDataType(int index,
+                                       const nvinfer1::DataType* input_type,
+                                       int nb_inputs) const
+      TRT_NOEXCEPT override;
+  bool isOutputBroadcastAcrossBatch(int output_index,
+                                    const bool* input_is_broadcast,
+                                    int nb_inputs) const TRT_NOEXCEPT override;
+  bool canBroadcastInputAcrossBatch(int input_index) const
+      TRT_NOEXCEPT override;
+  void configurePlugin(const nvinfer1::Dims* input_dims,
+                       int nb_inputs,
+                       const nvinfer1::Dims* output_dims,
+                       int nb_outputs,
+                       const nvinfer1::DataType* input_types,
+                       const nvinfer1::DataType* output_types,
+                       const bool* input_is_broadcast,
+                       const bool* output_is_broadcast,
+                       nvinfer1::PluginFormat float_format,
+                       int max_batch_size) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* clone() const TRT_NOEXCEPT override;
+
+ private:
+  nvinfer1::DataType data_type_;
+  std::vector<int> anchors_;
+  int* anchors_device_;
+  int class_num_;
+  float conf_thresh_;
+  int downsample_ratio_;
+  bool clip_bbox_;
+  float scale_x_y_;
+  int input_h_;
+  int input_w_;
+  bool iou_aware_;
+  float iou_aware_factor_;
+  std::string namespace_;
+};
+
+class PIRYoloBoxPluginCreator : public nvinfer1::IPluginCreator {
+ public:
+  PIRYoloBoxPluginCreator();
+  ~PIRYoloBoxPluginCreator() override = default;
+
+  void setPluginNamespace(const char* lib_namespace) TRT_NOEXCEPT override;
+  const char* getPluginNamespace() const TRT_NOEXCEPT override;
+  const char* getPluginName() const TRT_NOEXCEPT override;
+  const char* getPluginVersion() const TRT_NOEXCEPT override;
+  const nvinfer1::PluginFieldCollection* getFieldNames() TRT_NOEXCEPT override;
+
+  nvinfer1::IPluginV2Ext* createPlugin(
+      const char* name,
+      const nvinfer1::PluginFieldCollection* fc) TRT_NOEXCEPT override;
+  nvinfer1::IPluginV2Ext* deserializePlugin(const char* name,
+                                            const void* serial_data,
+                                            size_t serial_length)
+      TRT_NOEXCEPT override;
+
+ private:
+  std::string namespace_;
+  nvinfer1::PluginFieldCollection field_collection_;
+};
+
+REGISTER_TRT_PLUGIN_V2(YoloBoxPluginCreator);
+REGISTER_TRT_PLUGIN_V2(PIRYoloBoxPluginCreator);
+REGISTER_TRT_PLUGIN_V2(YoloBoxPluginDynamicCreator);
+
+}  // namespace plugin
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle