Merge pull request PaddlePaddle#9 from ckl117/ADFM-PNC

支持overwrite = True时的scatter算子，减少子图数量

paddle/fluid/inference/api/analysis_predictor.cc

@@ -3368,6 +3368,7 @@ USE_TRT_CONVERTER(argsort)
 USE_TRT_CONVERTER(atan2)
 USE_TRT_CONVERTER(index_put)
 USE_TRT_CONVERTER(scatter_nd_add)
+USE_TRT_CONVERTER(scatter)
 USE_TRT_CONVERTER(p_norm)
 USE_TRT_CONVERTER(flip)
 USE_TRT_CONVERTER(share_data)

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

@@ -119,6 +119,7 @@ list(
   dequantize_linear_op.cc
   index_put_op.cc
   scatter_nd_add_op.cc
+  scatter_op.cc
   p_norm_op.cc
   share_data_op.cc)
 

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

@@ -0,0 +1,56 @@
+/* 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. */
+
+#include "paddle/fluid/inference/tensorrt/convert/op_converter.h"
+
+namespace paddle::inference::tensorrt {
+
+/*
+ * scatter Op
+ */
+class ScatterOpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(3) << "convert a scatter op to tensorrt layer";
+
+    framework::OpDesc op_desc(op, nullptr);
+    std::string input_name = op_desc.Input("X").front();
+    std::string index_name = op_desc.Input("Ids").front();
+    std::string update_name = op_desc.Input("Updates").front();
+    std::string output_name = op_desc.Output("Out").front();
+    auto* input_tensor = engine_->GetITensor(input_name);
+    auto* index_tensor = engine_->GetITensor(index_name);
+    auto* update_tensor = engine_->GetITensor(update_name);
+    auto input_dims = input_tensor->getDimensions();
+    auto index_dims = index_tensor->getDimensions();
+    if (input_dims.nbDims == index_dims.nbDims) {
+      index_tensor = Reshape(
+          index_tensor, Concat({Shape(index_tensor), Add1DConstantLayer(1)}));
+    }
+    auto* layer = TRT_ENGINE_ADD_LAYER(engine_,
+                                       Scatter,
+                                       *input_tensor,
+                                       *index_tensor,
+                                       *update_tensor,
+                                       nvinfer1::ScatterMode::kND);
+    layer->setAxis(0);
+    ReplenishLayerAndOutput(layer, "scatter", {output_name}, test_mode);
+  }
+};
+
+}  // namespace paddle::inference::tensorrt
+
+REGISTER_TRT_OP_CONVERTER(scatter, ScatterOpConverter);

paddle/fluid/inference/tensorrt/op_teller.cc

@@ -2889,6 +2889,26 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
     }
 
+    if (op_type == "scatter") {
+      if (!with_dynamic_shape) {
+        VLOG(3) << "the scatter does not support "
+                   "static shape yet";
+        return false;
+      }
+      auto* block = desc.Block();
+      if (block == nullptr) {
+        VLOG(3) << "The block desc is nullptr, we can't continue to analyze. "
+                   "Developers need to check whether block_desc is passed in "
+                   "the pass.";
+        return false;
+      }
+      bool overwrite = PADDLE_GET_CONST(bool, desc.GetAttr("overwrite"));
+      if (!overwrite) {
+        VLOG(3) << op_type << " op only support overwrite = True.";
+        return false;
+      }
+    }
+
     if (op_type == "p_norm") {
       if (!with_dynamic_shape) {
         VLOG(3) << "the p_norm does not support "
@@ -3182,6 +3202,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "atan2",
       "index_put",
       "scatter_nd_add",
+      "scatter",
       "p_norm",
       "assign",
       "flip",
@@ -3359,6 +3380,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "atan2",
       "index_put",
       "scatter_nd_add",
+      "scatter",
       "p_norm",
       "assign",
       "flip",