[Hackathon NO.71] 为 Paddle-TRT 添加 pad3d 算子

paddle/fluid/framework/op_desc.cc

@@ -582,6 +582,10 @@ bool OpDesc::HasOutput(const std::string &name) const {
   return outputs_.find(name) != outputs_.end();
 }
 
+bool OpDesc::HasInput(const std::string &name) const {
+  return inputs_.find(name) != inputs_.end();
+}
+
 std::vector<std::string> OpDesc::OutputArgumentNames() const {
   std::vector<std::string> retv;
   for (auto &ipt : this->outputs_) {

paddle/fluid/framework/op_desc.h

@@ -76,6 +76,8 @@ class OpDesc {
 
   bool HasOutput(const std::string &name) const;
 
+  bool HasInput(const std::string &name) const;
+
   std::vector<std::string> OutputArgumentNames() const;
 
   void SetOutput(const std::string &param_name,

paddle/fluid/inference/api/analysis_predictor.cc

@@ -2419,6 +2419,9 @@ USE_TRT_CONVERTER(batch_norm);
 USE_TRT_CONVERTER(concat);
 USE_TRT_CONVERTER(dropout);
 USE_TRT_CONVERTER(pad);
+#if IS_TRT_VERSION_GE(8200)
+USE_TRT_CONVERTER(pad3d);
+#endif
 USE_TRT_CONVERTER(hard_sigmoid);
 USE_TRT_CONVERTER(hard_swish);
 USE_TRT_CONVERTER(split);

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

@@ -16,6 +16,7 @@ list(
   concat_op.cc
   dropout_op.cc
   group_norm_op.cc
+  pad3d_op.cc
   pad_op.cc
   split_op.cc
   square_op.cc

paddle/fluid/inference/tensorrt/convert/pad3d_op.cc

@@ -0,0 +1,183 @@
+/* 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/convert/op_converter.h"
+
+namespace paddle {
+namespace inference {
+namespace tensorrt {
+
+/*
+ * Pad3dOp.
+ */
+class Pad3dOpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+#if IS_TRT_VERSION_GE(8200)
+    VLOG(3) << "convert a pad3d op to tensorrt pad3d layer";
+
+    framework::OpDesc op_desc(op, nullptr);
+
+    // Declare inputs
+    auto* input = engine_->GetITensor(op_desc.Input("X")[0]);
+
+    nvinfer1::ITensor* paddings;
+    if (op_desc.HasInput("Paddings") && op_desc.Input("Paddings").size() > 0) {
+      paddings = engine_->GetITensor(op_desc.Input("Paddings")[0]);
+    } else {
+      std::vector<int> paddings_v =
+          PADDLE_GET_CONST(std::vector<int>, op_desc.GetAttr("paddings"));
+      paddings = Add1DConstantLayer(paddings_v);
+    }
+
+    float value{0.F};
+    if (op_desc.HasAttr("value")) {
+      value = PADDLE_GET_CONST(float, op_desc.GetAttr("value"));
+    }
+
+    std::string padding_mode = "constant";
+    if (op_desc.HasAttr("mode")) {
+      padding_mode = PADDLE_GET_CONST(std::string, op_desc.GetAttr("mode"));
+    }
+
+    const int input_dim = input->getDimensions().nbDims;
+    const int pad_size = paddings->getDimensions().d[0];
+    PADDLE_ENFORCE_EQ(input_dim * 2 - 4,
+                      pad_size,
+                      phi::errors::InvalidArgument(
+                          "Expected paddings size is %d, but received %d.",
+                          input_dim * 2 - 4,
+                          pad_size));
+    // convert paddle pad to tensorrt pad
+    std::vector<int> shuffle_index{4, 2, 0, 5, 3, 1};
+    std::vector<nvinfer1::ITensor*> shuffle_inputs;
+    for (int i = 0; i < pad_size; i++) {
+      shuffle_inputs.push_back(GetEleTensorOfShape(paddings, shuffle_index[i]));
+    }
+    paddings = Concat(shuffle_inputs);
+    auto* pre_zeros = Add1DConstantLayer(std::vector<int>(2, 0));
+    auto start_slice1 = nvinfer1::Dims{1, { 0 }};
+    auto start_slice2 = nvinfer1::Dims{1, { 3 }};
+    auto size_slice = nvinfer1::Dims{1, { 3 }};
+    auto stride_slice = nvinfer1::Dims{1, { 1 }};
+
+    auto* pre_pad =
+        TRT_ENGINE_ADD_LAYER(
+            engine_, Slice, *paddings, start_slice1, size_slice, stride_slice)
+            ->getOutput(0);
+    pre_pad = Concat(std::vector<nvinfer1::ITensor*>{pre_zeros, pre_pad});
+    auto* post_pad =
+        TRT_ENGINE_ADD_LAYER(
+            engine_, Slice, *paddings, start_slice2, size_slice, stride_slice)
+            ->getOutput(0);
+    post_pad = Concat(std::vector<nvinfer1::ITensor*>{pre_zeros, post_pad});
+
+    std::vector<int> zeros_v(input_dim, 0);
+    auto const zeros = Add1DConstantLayer(zeros_v);
+
+    nvinfer1::ITensor* start{};
+    nvinfer1::ITensor* size{};
+    // elementwise add zeros and pre_pad
+    start = TRT_ENGINE_ADD_LAYER(engine_,
+                                 ElementWise,
+                                 *zeros,
+                                 *pre_pad,
+                                 nvinfer1::ElementWiseOperation::kSUB)
+                ->getOutput(0);
+
+    auto const total_padding =
+        TRT_ENGINE_ADD_LAYER(engine_,
+                             ElementWise,
+                             *pre_pad,
+                             *post_pad,
+                             nvinfer1::ElementWiseOperation::kSUM)
+            ->getOutput(0);
+
+    auto* input_shape = Shape(input);
+    size = TRT_ENGINE_ADD_LAYER(engine_,
+                                ElementWise,
+                                *input_shape,
+                                *total_padding,
+                                nvinfer1::ElementWiseOperation::kSUM)
+               ->getOutput(0);
+    // add slice layer
+    nvinfer1::Dims stride;
+    stride.nbDims = input_dim;
+    std::fill_n(stride.d, input_dim, 1);
+    auto const& dummy = stride;
+    auto* slice_layer =
+        TRT_ENGINE_ADD_LAYER(engine_,
+                             Slice,
+                             *const_cast<nvinfer1::ITensor*>(input),
+                             dummy,
+                             dummy,
+                             stride);
+    slice_layer->setInput(1, *start);
+    slice_layer->setInput(2, *size);
+    if (padding_mode == "constant") {
+#if IS_TRT_VERSION_GE(8500)
+      slice_layer->setMode(nvinfer1::SampleMode::kFILL);
+#else
+      slice_layer->setMode(nvinfer1::SliceMode::kFILL);
+#endif
+      if (value != 0.F) {
+        nvinfer1::ITensor* fill_value = nullptr;
+        switch (input->getType()) {
+          case nvinfer1::DataType::kFLOAT:
+          case nvinfer1::DataType::kHALF:
+          case nvinfer1::DataType::kINT8: {
+            fill_value = Add1DConstantLayer(value);
+            break;
+          }
+          default: {
+            int value_int = static_cast<int>(value);
+            fill_value = Add1DConstantLayer(value_int);
+            break;
+          }
+        }
+        slice_layer->setInput(4, *fill_value);
+      }
+    } else if (padding_mode == "reflect") {
+#if IS_TRT_VERSION_GE(8500)
+      slice_layer->setMode(nvinfer1::SampleMode::kREFLECT);
+#else
+      slice_layer->setMode(nvinfer1::SliceMode::kREFLECT);
+#endif
+    } else if (padding_mode == "replicate") {
+#if IS_TRT_VERSION_GE(8500)
+      slice_layer->setMode(nvinfer1::SampleMode::kCLAMP);
+#else
+      slice_layer->setMode(nvinfer1::SliceMode::kCLAMP);
+#endif
+    } else {
+      PADDLE_THROW(paddle::platform::errors::Fatal("Unsupported mode: %s",
+                                                   padding_mode));
+    }
+
+    auto output_name = op_desc.Output("Out")[0];
+    RreplenishLayerAndOutput(slice_layer, "pad3d", {output_name}, test_mode);
+
+#else
+    VLOG(3) << "pad3d is not supported when TensorRT < 8.2";
+#endif
+  }
+};
+
+}  // namespace tensorrt
+}  // namespace inference
+}  // namespace paddle
+
+REGISTER_TRT_OP_CONVERTER(pad3d, Pad3dOpConverter);

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -1748,6 +1748,35 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
     }
 
+    if (op_type == "pad3d") {
+#if !IS_TRT_VERSION_GE(8200)
+      VLOG(3) << "pad3d is not supported when TensorRT < 8.2";
+      return false;
+#endif
+      if (!with_dynamic_shape) {
+        VLOG(3) << "pad3d is not supported static shape";
+        return false;
+      }
+      if (!desc.HasAttr("paddings") && !desc.HasInput("Paddings")) {
+        return false;
+      }
+      if (desc.HasAttr("mode")) {
+        std::string mode = PADDLE_GET_CONST(std::string, desc.GetAttr("mode"));
+        if (mode != "constant" && mode != "reflect" && mode != "replicate") {
+          VLOG(3) << "The pad3d layer of TRT only support "
+                     "constant/reflect/replicate mode.";
+          return false;
+        }
+      }
+      if (desc.HasAttr("data_format")) {
+        std::string data_format =
+            PADDLE_GET_CONST(std::string, desc.GetAttr("data_format"));
+        if (data_format != "NCDHW") {
+          VLOG(3) << "The pad3d layer of TRT only support NCDHW data format.";
+          return false;
+        }
+      }
+    }
     if (op_type == "swish") {
       auto* block = desc.Block();
       if (block == nullptr) {
@@ -1764,7 +1793,6 @@ struct SimpleOpTypeSetTeller : public Teller {
         return false;
       }
     }
-
     if (op_type == "prelu") {
       if (desc.Input("X").size() != 1) {
         VLOG(3) << "Invalid input X's size of prelu TRT converter. "
@@ -2694,6 +2722,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "batch_norm",
       "concat",
       "tanh",
+      "pad3d",
       "pad",
       "elementwise_add",
       "elementwise_sub",
@@ -2849,6 +2878,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "batch_norm",
       "concat",
       "tanh",
+      "pad3d",
       "pad",
       "elementwise_add",
       "elementwise_sub",
@@ -2974,14 +3004,6 @@ struct GenericPluginTeller : public Teller {
       if (!desc.HasAttr("iou_aware") && !desc.HasAttr("iou_aware_factor"))
         return false;
     }
-    if (op_type == "pad3d") {
-      auto pad3d_inputs = desc.Inputs();
-      if (pad3d_inputs.find("Paddings") != pad3d_inputs.end()) {
-        if (desc.Input("Paddings").size() >= 1) {
-          return false;
-        }
-      }
-    }
     if (use_no_calib_int8) {
       return false;
     } else {