[Paddle-TRT ]add argsort op into TRT GenericPlugin (#59409)

lizexu123 · web-flow · commit 6e4163a50b60 · 2023-12-09T01:42:44.000+08:00
diff --git a/paddle/fluid/inference/tensorrt/dynamic_shape_infermeta.cc b/paddle/fluid/inference/tensorrt/dynamic_shape_infermeta.cc
@@ -845,6 +845,21 @@ nvinfer1::DimsExprs ScatterInferMeta(
   return ref_dims;
 }
 
+nvinfer1::DimsExprs ArgsortInferMeta(
+    int output_index,
+    const nvinfer1::DimsExprs* inputs,
+    int nb_inputs,
+    nvinfer1::IExprBuilder& expr_builder,  // NOLINT
+    const framework::OpDesc& op_desc) {
+  const nvinfer1::DimsExprs input_dims = inputs[0];
+  nvinfer1::DimsExprs output;
+  output.nbDims = input_dims.nbDims;
+  for (int i = 0; i < input_dims.nbDims; ++i) {
+    output.d[i] = input_dims.d[i];
+  }
+  return output;
+}
+
 nvinfer1::DimsExprs SolveInferMeta(
     int output_index,
     const nvinfer1::DimsExprs* inputs,
@@ -878,6 +893,7 @@ PD_REGISTER_DYNAMIC_INFER_META_FN(p_norm, PNormInferMeta);
 PD_REGISTER_DYNAMIC_INFER_META_FN(memory_efficient_attention,
                                   MemoryEfficientAttentionInferMeta);
 PD_REGISTER_DYNAMIC_INFER_META_FN(pad, PadInferMeta);
+PD_REGISTER_DYNAMIC_INFER_META_FN(argsort, ArgsortInferMeta);
 PD_REGISTER_DYNAMIC_INFER_META_FN(scatter, ScatterInferMeta);
 PD_REGISTER_DYNAMIC_INFER_META_FN(solve, SolveInferMeta);
 }  // namespace tensorrt
diff --git a/paddle/fluid/inference/tensorrt/dynamic_shape_infermeta_registry.h b/paddle/fluid/inference/tensorrt/dynamic_shape_infermeta_registry.h
@@ -34,6 +34,7 @@ USE_TRT_DYNAMIC_INFER_META_FN(conv2d_transpose);
 USE_TRT_DYNAMIC_INFER_META_FN(memory_efficient_attention);
 USE_TRT_DYNAMIC_INFER_META_FN(p_norm);
 USE_TRT_DYNAMIC_INFER_META_FN(pad);
+USE_TRT_DYNAMIC_INFER_META_FN(argsort);
 USE_TRT_DYNAMIC_INFER_META_FN(scatter);
 USE_TRT_DYNAMIC_INFER_META_FN(solve);
 }  // namespace tensorrt
diff --git a/paddle/fluid/inference/tensorrt/op_teller.cc b/paddle/fluid/inference/tensorrt/op_teller.cc
@@ -1974,8 +1974,8 @@ struct SimpleOpTypeSetTeller : public Teller {
 #elif !IS_TRT_VERSION_GE(8600)
         const auto x_shape = x_var_desc->GetShape();
         if (x_shape.empty()) {
-          VLOG(3)
-              << "BOOL type does not support 0 dim input when TensorRT < 8.6.";
+          VLOG(3) << "BOOL type does not support 0 dim input when TensorRT < "
+                     "8.6.";
           return false;
         }
 #endif
@@ -1993,8 +1993,8 @@ struct SimpleOpTypeSetTeller : public Teller {
         return false;
       }
       if (desc.HasAttr("allow_out_of_range")) {
-        VLOG(3)
-            << "allow_out_of_range one_hot/one_hot_v2 op is not supported now.";
+        VLOG(3) << "allow_out_of_range one_hot/one_hot_v2 op is not "
+                   "supported now.";
         if (PADDLE_GET_CONST(bool, desc.GetAttr("allow_out_of_range")))
           return false;
       }
@@ -2156,7 +2156,8 @@ struct SimpleOpTypeSetTeller : public Teller {
       if (with_dynamic_shape) {
         return true;
       }
-      // Static shape does not support the input tensors: Shape and ShapeTensor
+      // Static shape does not support the input tensors: Shape and
+      // ShapeTensor
       auto reshape_inputs = desc.Inputs();
       if (reshape_inputs.find("Shape") != reshape_inputs.end()) {
         if (!desc.Input("Shape").empty()) {
@@ -3192,6 +3193,14 @@ struct GenericPluginTeller : public Teller {
         return false;
       }
     }
+    // TODO(lizexu123): the tensorrt version on Windows supports low-level
+    // and inconsistent supportformation
+    if (op_type == "argsort") {
+      if (!with_dynamic_shape) {
+        VLOG(3) << "Ops(" << op_type << ") do not support static shape yet.";
+        return false;
+      }
+    }
     if (use_no_calib_int8) {
       return false;
     } else {
diff --git a/paddle/fluid/inference/tensorrt/plugin/generic_plugin.cu b/paddle/fluid/inference/tensorrt/plugin/generic_plugin.cu
@@ -24,6 +24,7 @@
 #include "paddle/phi/core/compat/op_utils.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 {
@@ -403,6 +404,24 @@ bool GenericPlugin::supportsFormatCombination(
     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)
@@ -447,6 +466,11 @@ nvinfer1::DataType GenericPlugin::getOutputDataType(
   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];
 }
 
@@ -552,6 +576,7 @@ int GenericPlugin::enqueue(const nvinfer1::PluginTensorDesc* input_desc,
                            void* workspace,
                            cudaStream_t stream) TRT_NOEXCEPT {
   platform::CUDAPlace place(platform::GetCurrentDeviceId());
+  platform::DeviceContextPool& pool = platform::DeviceContextPool::Instance();
   // TODO(inference): generic plugin do not support INT8 precision now.
   auto nvType2PhiType =
       [&](nvinfer1::DataType nv_dtype) -> std::pair<phi::DataType, int> {
@@ -633,6 +658,28 @@ int GenericPlugin::enqueue(const nvinfer1::PluginTensorDesc* input_desc,
   CHECK_EQ(phi_kernel_contexts_[data_type]->InputsSize(), getNbInputs());
   CHECK_EQ(phi_kernel_contexts_[data_type]->OutputsSize(), 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;
 }
 
diff --git a/test/ir/inference/CMakeLists.txt b/test/ir/inference/CMakeLists.txt
@@ -33,6 +33,7 @@ if(WIN32)
   list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES "test_groupnorm_act_pass_fuse_pass")
   list(REMOVE_ITEM TEST_TRT_IR_PASSES "test_trt_convert_trans_layernorm")
   list(REMOVE_ITEM TEST_TRT_CONVERTER "test_trt_convert_trans_layernorm")
+  list(REMOVE_ITEM TEST_TRT_CONVERTER "test_trt_convert_argsort")
   list(REMOVE_ITEM TEST_INFERENCE_IR_PASSES "test_trt_convert_trans_layernorm")
   list(REMOVE_ITEM TEST_TRT_IR_PASSES "test_trt_convert_fused_token_prune")
   list(REMOVE_ITEM TEST_TRT_CONVERTER "test_trt_convert_fused_token_prune")
diff --git a/test/ir/inference/test_trt_convert_argsort.py b/test/ir/inference/test_trt_convert_argsort.py
@@ -0,0 +1,133 @@
+# 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
+#
+#     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.
+
+import unittest
+from functools import partial
+from typing import List
+
+import numpy as np
+from program_config import ProgramConfig, TensorConfig
+from trt_layer_auto_scan_test import TrtLayerAutoScanTest
+
+import paddle.inference as paddle_infer
+
+
+class TrtConvertArgsort(TrtLayerAutoScanTest):
+    def is_program_valid(self, program_config: ProgramConfig) -> bool:
+        return True
+
+    def sample_program_configs(self):
+        def generate_input1(batch):
+            if self.dims == 4:
+                return np.ones([batch, 3, 3, 24]).astype(np.float32)
+            elif self.dims == 3:
+                return np.ones([batch, 3, 24]).astype(np.float16)
+            elif self.dims == 2:
+                return np.ones([batch, 24]).astype(np.float16)
+
+        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)
+                                )
+                            },
+                            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],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [9, 3 + 1, 3 + 1, 24 + 1],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [6, 3, 3, 24],
+                }
+            elif self.dims == 3:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [1, 3 - 1, 24 - 1],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [9, 3 + 1, 24 + 1],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [6, 3, 24],
+                }
+            elif self.dims == 2:
+                self.dynamic_shape.min_input_shape = {
+                    "input_data": [1, 24],
+                }
+                self.dynamic_shape.max_input_shape = {
+                    "input_data": [9, 24],
+                }
+                self.dynamic_shape.opt_input_shape = {
+                    "input_data": [6, 24],
+                }
+
+        def clear_dynamic_shape():
+            self.dynamic_shape.max_input_shape = {}
+            self.dynamic_shape.min_input_shape = {}
+            self.dynamic_shape.opt_input_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):
+        self.run_test()
+
+
+if __name__ == "__main__":
+    unittest.main()
