[Paddle-TRT]Support argsort by TopK (#66811)

ckl117 · web-flow · commit 34e2654cca4c · 2024-08-07T10:28:03.000+08:00
* trt argsort
diff --git a/paddle/fluid/inference/api/analysis_predictor.cc b/paddle/fluid/inference/api/analysis_predictor.cc
@@ -3610,6 +3610,7 @@ USE_TRT_CONVERTER(lookup_table)
 USE_TRT_CONVERTER(lookup_table_v2)
 USE_TRT_CONVERTER(expand_v2)
 USE_TRT_CONVERTER(expand_as_v2)
+USE_TRT_CONVERTER(argsort)
 USE_TRT_CONVERTER(take_along_axis)
 USE_TRT_CONVERTER(skip_groupnorm_act)
 USE_TRT_CONVERTER(preln_groupnorm_act)
diff --git a/paddle/fluid/inference/tensorrt/convert/CMakeLists.txt b/paddle/fluid/inference/tensorrt/convert/CMakeLists.txt
@@ -107,6 +107,7 @@ list(
   skip_groupnorm_act_op.cc
   preln_groupnorm_act_op.cc
   expand_v2_op.cc
+  argsort_op.cc
   cumsum_op.cc
   temporal_shift_op.cc
   einsum_op.cc
diff --git a/paddle/fluid/inference/tensorrt/convert/argsort_op.cc b/paddle/fluid/inference/tensorrt/convert/argsort_op.cc
@@ -0,0 +1,92 @@
+/* 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
+https://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 {
+
+/*
+ * Argsort Op
+ */
+class ArgsortOpConverter : public OpConverter {
+ public:
+  void operator()(const framework::proto::OpDesc& op,
+                  const framework::Scope& scope,
+                  bool test_mode) override {
+    VLOG(3) << "convert a argsort op to tensorrt layer";
+
+    framework::OpDesc op_desc(op, nullptr);
+    std::string input_name = op_desc.Input("X").front();
+    std::string output_name = op_desc.Output("Out").front();
+    std::string indices_name = op_desc.Output("Indices").front();
+    int axis = PADDLE_GET_CONST(int, op_desc.GetAttr("axis"));
+    bool descending = PADDLE_GET_CONST(bool, op_desc.GetAttr("descending"));
+    auto* input_tensor = engine_->GetITensor(input_name);
+    nvinfer1::Dims input_tensor_dims = input_tensor->getDimensions();
+    nvinfer1::TopKOperation operation = nvinfer1::TopKOperation::kMIN;
+    if (descending) {
+      operation = nvinfer1::TopKOperation::kMAX;
+    }
+    if (axis < 0) {
+      axis += input_tensor_dims.nbDims;
+    }
+    nvinfer1::DataType in_type = input_tensor->getType();
+    bool need_cast = in_type != nvinfer1::DataType::kFLOAT ? true : false;
+    int x_rank = input_tensor->getDimensions().nbDims;
+    if (x_rank == 1) {
+      nvinfer1::Dims unsqueeze_shape;
+      unsqueeze_shape.nbDims = 2;
+      unsqueeze_shape.d[0] = 1;
+      unsqueeze_shape.d[1] = -1;
+      input_tensor = Reshape(input_tensor, unsqueeze_shape);
+      axis = 1;
+    }
+
+    if (need_cast) {
+      auto* cast_layer1 =
+          TRT_ENGINE_ADD_LAYER(engine_, Identity, *input_tensor);
+      cast_layer1->setOutputType(0, nvinfer1::DataType::kFLOAT);
+      cast_layer1->getOutput(0)->setType(nvinfer1::DataType::kFLOAT);
+      input_tensor = cast_layer1->getOutput(0);
+    }
+
+    auto* layer = TRT_ENGINE_ADD_LAYER(
+        engine_, TopK, *input_tensor, operation, 1, 1 << axis);
+    auto* shape = Shape(input_tensor);
+    auto* k_tensor = GetEleTensorOfShape(shape, axis, true);
+    layer->setInput(1, *k_tensor);
+    auto* Out = layer->getOutput(0);
+    auto* Indices = layer->getOutput(1);
+
+    if (x_rank == 1) {
+      nvinfer1::Dims squeeze_shape;
+      squeeze_shape.nbDims = 1;
+      squeeze_shape.d[0] = -1;
+      Out = Reshape(Out, squeeze_shape);
+      Indices = Reshape(Indices, squeeze_shape);
+    }
+
+    auto* out_layer = TRT_ENGINE_ADD_LAYER(engine_, Identity, *Out);
+    out_layer->setOutputType(0, in_type);
+    out_layer->getOutput(0)->setType(in_type);
+    ReplenishLayerAndOutput(out_layer, "argsort", {output_name}, test_mode);
+
+    auto* indices_layer = TRT_ENGINE_ADD_LAYER(engine_, Identity, *Indices);
+    indices_layer->setOutputType(0, Indices->getType());
+    indices_layer->getOutput(0)->setType(Indices->getType());
+    ReplenishLayerAndOutput(
+        indices_layer, "argsort", {indices_name}, test_mode);
+  }
+};
+
+}  // namespace paddle::inference::tensorrt
+
+REGISTER_TRT_OP_CONVERTER(argsort, ArgsortOpConverter);
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -2777,6 +2777,32 @@ struct SimpleOpTypeSetTeller : public Teller {
       }
     }
 
+    if (op_type == "argsort") {
+      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;
+      }
+      if (!desc.HasAttr("descending") || !desc.HasAttr("axis")) {
+        VLOG(3) << op_type << " needs attributes: descending and axis.";
+      }
+      auto x_var_name = desc.Input("X")[0];
+      auto* x_var_desc = block->FindVarRecursive(x_var_name);
+      std::vector<int64_t> shape = x_var_desc->GetShape();
+      int axis = PADDLE_GET_CONST(int, desc.GetAttr("axis"));
+      if (axis < 0) {
+        axis += shape.size();
+      }
+      if (shape[axis] > 3840 || shape[axis] < 0) {
+        VLOG(3) << op_type << " shape[" << axis << "] = " << shape[axis]
+                << " is invalid, it should less than 3840 and greater than "
+                   "zero in TensorRT.";
+        return false;
+      }
+    }
+
     if (op_type == "unbind") {
       if (!with_dynamic_shape) {
         VLOG(3) << "the unbind does not support "
@@ -3098,6 +3124,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "quantize_linear",
       "dequantize_linear",
       "share_data",
+      "argsort",
       "bitwise_and",
       "bitwise_or",
       "size"};
@@ -3272,6 +3299,7 @@ struct SimpleOpTypeSetTeller : public Teller {
       "quantize_linear",
       "dequantize_linear",
       "share_data",
+      "argsort",
       "bitwise_and",
       "bitwise_or",
       "size"};
diff --git a/test/ir/inference/test_trt_convert_argsort.py b/test/ir/inference/test_trt_convert_argsort.py
@@ -28,66 +28,67 @@ def is_program_valid(self, program_config: ProgramConfig) -> bool:
         return True
 
     def sample_program_configs(self):
-        def generate_input1(batch):
+        def generate_input1():
             if self.dims == 4:
-                return np.ones([batch, 3, 3, 24]).astype(np.float32)
+                return np.random.random([1, 3, 3, 4000]).astype(np.float32)
             elif self.dims == 3:
-                return np.ones([batch, 3, 24]).astype(np.float16)
+                return np.random.random([1, 3, 24]).astype(np.int32)
             elif self.dims == 2:
-                return np.ones([batch, 24]).astype(np.float16)
+                return np.random.random([1, 24]).astype(np.float32)
+            else:
+                return np.random.random([24]).astype(np.float32)
 
-        for dims in [2, 3, 4]:
-            for batch in [1, 6, 9]:
-                for axis in [-1, 0, 1]:
-                    for descending in [False, True]:
-                        self.dims = dims
-                        ops_config = [
-                            {
-                                "op_type": "argsort",
-                                "op_inputs": {"X": ["input_data"]},
-                                "op_outputs": {
-                                    "Out": ["output_data"],
-                                    "Indices": ["indices_data"],
-                                },
-                                "op_attrs": {
-                                    "axis": axis,
-                                    "descending": descending,
-                                },
-                            }
-                        ]
-                        ops = self.generate_op_config(ops_config)
-                        program_config = ProgramConfig(
-                            ops=ops,
-                            weights={},
-                            inputs={
-                                "input_data": TensorConfig(
-                                    data_gen=partial(generate_input1, batch)
-                                )
+        for dims in [1, 2, 3, 4]:
+            for axis in [-1, 0]:
+                for descending in [False, True]:
+                    self.dims = dims
+                    ops_config = [
+                        {
+                            "op_type": "argsort",
+                            "op_inputs": {"X": ["input_data"]},
+                            "op_outputs": {
+                                "Out": ["output_data"],
+                                "Indices": ["indices_data"],
                             },
-                            outputs=["output_data", "indices_data"],
-                        )
-                        yield program_config
+                            "op_attrs": {
+                                "axis": axis,
+                                "descending": descending,
+                            },
+                        }
+                    ]
+                    ops = self.generate_op_config(ops_config)
+                    program_config = ProgramConfig(
+                        ops=ops,
+                        weights={},
+                        inputs={
+                            "input_data": TensorConfig(
+                                data_gen=partial(generate_input1)
+                            )
+                        },
+                        outputs=["output_data", "indices_data"],
+                    )
+                    yield program_config
 
     def sample_predictor_configs(
         self, program_config
     ) -> (paddle_infer.Config, List[int], float):
         def generate_dynamic_shape(attrs):
             if self.dims == 4:
                 self.dynamic_shape.min_input_shape = {
-                    "input_data": [1, 3 - 1, 3 - 1, 24 - 1],
+                    "input_data": [1, 3, 3, 4000],
                 }
                 self.dynamic_shape.max_input_shape = {
-                    "input_data": [9, 3 + 1, 3 + 1, 24 + 1],
+                    "input_data": [9, 3, 3, 4000],
                 }
                 self.dynamic_shape.opt_input_shape = {
-                    "input_data": [6, 3, 3, 24],
+                    "input_data": [6, 3, 3, 4000],
                 }
             elif self.dims == 3:
                 self.dynamic_shape.min_input_shape = {
-                    "input_data": [1, 3 - 1, 24 - 1],
+                    "input_data": [1, 3, 24],
                 }
                 self.dynamic_shape.max_input_shape = {
-                    "input_data": [9, 3 + 1, 24 + 1],
+                    "input_data": [9, 3, 24],
                 }
                 self.dynamic_shape.opt_input_shape = {
                     "input_data": [6, 3, 24],
@@ -102,6 +103,16 @@ def generate_dynamic_shape(attrs):
                 self.dynamic_shape.opt_input_shape = {
                     "input_data": [6, 24],
                 }
+            else:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [24],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [25],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [24],
+                }
 
         def clear_dynamic_shape():
             self.dynamic_shape.max_input_shape = {}
@@ -111,18 +122,12 @@ def clear_dynamic_shape():
         attrs = [
             program_config.ops[i].attrs for i in range(len(program_config.ops))
         ]
-        self.trt_param.max_batch_size = 9
         self.trt_param.workspace_size = 1073741824
-        # for static_shape
-        clear_dynamic_shape()
-
         # for dynamic_shape
         generate_dynamic_shape(attrs)
         self.trt_param.precision = paddle_infer.PrecisionType.Float32
-        program_config.set_input_type(np.float32)
         yield self.create_inference_config(), (1, 3), 1e-5
         self.trt_param.precision = paddle_infer.PrecisionType.Half
-        program_config.set_input_type(np.float16)
         yield self.create_inference_config(), (1, 3), 1e-3
 
     def test(self):
